From 862828d02a77c57871e13aaab93f14ad2078c69d Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Tue, 4 Aug 2026 07:02:51 +0000 Subject: [PATCH] Revert "DS4: faster long-context TG (#2201)" This reverts commit 707374b3c3c2630f35f5179cbe7a687b7499472b. --- ggml/include/ggml.h | 7 --- ggml/src/ggml-cuda.cu | 3 - ggml/src/ggml-cuda/getrows.cu | 82 +----------------------- ggml/src/ggml.c | 112 --------------------------------- src/graphs/build_deepseek4.cpp | 26 +++----- 5 files changed, 9 insertions(+), 221 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index b52ecd2e7..d16a2d02f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -1946,13 +1946,6 @@ extern "C" { struct ggml_tensor * b, struct ggml_tensor * c); - GGML_API struct ggml_tensor * ggml_get_rows_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool same_type, - bool dim0); - // a TD [n_embd, ne1, ne2, ne3] // b TS [n_embd, n_rows, ne02, ne03] | ne02 == ne2, ne03 == ne3 // c I64 [n_rows, ne11, ne12, 1] | c[i] in [0, ne1) diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index ad0db58e2..11187faf8 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -4841,9 +4841,6 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons } break; case GGML_OP_GET_ROWS: { - if (op->type == op->src[0]->type) { - return true; - } switch (op->src[0]->type) { case GGML_TYPE_F16: case GGML_TYPE_F32: diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu index 40c677c48..fd67dce1d 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu @@ -139,34 +139,6 @@ static void get_rows_cuda_float(const ggml_tensor * src0, const ggml_tensor * sr GGML_UNUSED(dst); } -template -static __global__ void k_get_rows_dim0(int n, - size_t nb01, size_t nb02, size_t nb03, size_t nb11, size_t nb12, size_t nb13, size_t nb1, size_t nb2, size_t nb3, - const data_t * __restrict__ src, const int * __restrict__ idx, data_t * __restrict__ dst) { - int i1 = blockIdx.x; - int i2 = blockIdx.y; - int i3 = blockIdx.z; - src += nb01*i1 + nb02*i2 + nb03*i3; - idx += nb11*i1 + nb12*i2 + nb13*i3; - dst += nb1*i1 + nb2*i2 + nb3*i3; - - for (int j = threadIdx.x; j < n; j += blockDim.x) dst[j] = src[idx[j]]; -} - -template -static __global__ void k_get_rows_dim1(int n, - size_t nb01, size_t nb02, size_t nb03, size_t nb11, size_t nb12, size_t nb1, size_t nb2, size_t nb3, - const data_t * __restrict__ src, const int * __restrict__ idx, data_t * __restrict__ dst) { - int i1 = blockIdx.x; - int i2 = blockIdx.y; - int i3 = blockIdx.z; - idx += nb11*i2 + nb12*i3; - src += nb01*idx[i1] + nb02*i2 + nb03*i3; - dst += nb1*i1 + nb2*i2 + nb3*i3; - - for (int j = threadIdx.x; j < n; j += blockDim.x) dst[j] = src[j]; -} - void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -175,60 +147,8 @@ void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { float * dst_d = (float *)dst->data; cudaStream_t stream = ctx.stream(); + GGML_ASSERT(src1->type == GGML_TYPE_I32); - - if (dst->type != GGML_TYPE_F32 || dst->op_params[0] == 1) { - constexpr int k_block_size = 512; - GGML_ASSERT(src0->type == dst->type); - if (dst->op_params[0] == 1) { - //printf("%s(%s, %s) - dim0\n", __func__, src0->name, ggml_type_name(src0->type)); - GGML_ASSERT(src0->ne[1] == src1->ne[1] && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3]); - GGML_ASSERT(src0->ne[0] >= src1->ne[0]); - GGML_ASSERT(dst->ne[0] == src1->ne[0]); - auto type_size = ggml_type_size(dst->type); - GGML_ASSERT(type_size == 2 || type_size == 4); - dim3 grid(dst->ne[1], dst->ne[2], dst->ne[3]); - if (type_size == 2) { - k_get_rows_dim0<<>>(dst->ne[0], - src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, - src1->nb[1]/type_size, src1->nb[2]/type_size, src1->nb[3]/type_size, - dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, - (const uint16_t *)src0->data, (const int *)src1->data, (uint16_t *)dst->data); - } else { - k_get_rows_dim0<<>>(dst->ne[0], - src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, - src1->nb[1]/type_size, src1->nb[2]/type_size, src1->nb[3]/type_size, - dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, - (const uint32_t *)src0->data, (const int *)src1->data, (uint32_t *)dst->data); - } - } else { - //printf("%s(%s, %s) - dim1\n", __func__, src0->name, ggml_type_name(src0->type)); - GGML_ASSERT(src0->ne[2] == src1->ne[1]); - GGML_ASSERT(src0->ne[3] == src1->ne[2]); - GGML_ASSERT(src0->ne[3] == 1); - GGML_ASSERT(dst->ne[0] == src0->ne[0] && dst->ne[1] == src1->ne[0] && dst->ne[2] == src1->ne[1] && dst->ne[2]); - - auto row_size = ggml_row_size(dst->type, dst->ne[0]); - GGML_ASSERT(row_size % 2 == 0); - auto type_size = row_size % 4 == 0 ? 4 : 2; - dim3 grid(dst->ne[1], dst->ne[2], dst->ne[3]); - if (type_size == 2) { - k_get_rows_dim1<<>>(dst->ne[0], - src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, - src1->nb[1]/type_size, src1->nb[2]/type_size, - dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, - (const uint16_t *)src0->data, (const int *)src1->data, (uint16_t *)dst->data); - } else { - k_get_rows_dim1<<>>(dst->ne[0], - src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, - src1->nb[1]/type_size, src1->nb[2]/type_size, - dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, - (const uint32_t *)src0->data, (const int *)src1->data, (uint32_t *)dst->data); - } - } - return; - } - GGML_ASSERT(dst->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32)); GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 0cba53103..46d53d043 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -8952,45 +8952,6 @@ struct ggml_tensor * ggml_get_rows( return result; } -struct ggml_tensor * ggml_get_rows_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool same_type, - bool dim0) { - if (dim0) { - GGML_ASSERT(same_type); - GGML_ASSERT(!ggml_is_quantized(a->type)); - GGML_ASSERT(ggml_type_size(a->type) % 2 == 0); // we do not support dim0 get_rows on int8, fp8 data - } - if (!dim0 && !same_type) { - return ggml_get_rows(ctx, a, b); - } - GGML_ASSERT(b->type == GGML_TYPE_I32); - if (dim0) { - GGML_ASSERT(a->ne[1] == b->ne[1] && a->ne[2] == b->ne[2] && a->ne[3] == b->ne[3]); - GGML_ASSERT(a->ne[0] >= b->ne[0]); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, b->ne[0], a->ne[1], a->ne[2], a->ne[3]); - result->op = GGML_OP_GET_ROWS; - result->op_params[0] = 1; - result->src[0] = a; - result->src[1] = b; - return result; - } - GGML_ASSERT(a->ne[2] == b->ne[1]); - GGML_ASSERT(a->ne[3] == b->ne[2]); - GGML_ASSERT(b->ne[3] == 1); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); - result->op = GGML_OP_GET_ROWS; - result->src[0] = a; - result->src[1] = b; - return result; - -} - - // ggml_get_rows_back struct ggml_tensor * ggml_get_rows_back( @@ -19736,85 +19697,12 @@ static void ggml_compute_forward_get_rows_f32( } } -static void ggml_compute_forward_get_rows_any( - const struct ggml_compute_params * params, - struct ggml_tensor * dst) { - - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(src1->type == GGML_TYPE_I32); - GGML_ASSERT(src0->type == dst->type); - - const int ith = params->ith; - const int nth = params->nth; - - int nrows = ggml_nrows(dst); - int npt = (nrows + nth - 1)/nth; - int first = npt*ith; - int last = MIN(first + npt, nrows); - - if (dst->op_params[0] == 1) { - size_t type_size = ggml_type_size(dst->type); - GGML_ASSERT(type_size == 2 || type_size == 4); - GGML_ASSERT(src0->ne[1] == src1->ne[1] && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3]); - GGML_ASSERT(src0->ne[0] >= src1->ne[0]); - GGML_ASSERT(dst->ne[0] == src1->ne[0]); - for (int ir = first; ir < last; ++ir) { - int i3 = ir/(dst->ne[1]*dst->ne[2]); - int i2 = (ir - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; - int i1 = ir - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; - const int * idx = (const int *)((const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3]); - const char * cx = (const char *)((const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]); - char * cy = ( char *)(( char *) dst->data + i1* dst->nb[1] + i2* dst->nb[2] + i3* dst->nb[3]); - if (type_size == 2) { - const uint16_t * x = (const uint16_t *)cx; - uint16_t * y = ( uint16_t *)cy; - for (int j = 0; j < (int)dst->ne[0]; ++j) { - y[j] = x[idx[j]]; - } - } else { - const uint32_t * x = (const uint32_t *)cx; - uint32_t * y = ( uint32_t *)cy; - for (int j = 0; j < (int)dst->ne[0]; ++j) { - y[j] = x[idx[j]]; - } - } - } - return; - } - - GGML_ASSERT(src0->ne[2] == src1->ne[1]); - GGML_ASSERT(src0->ne[3] == src1->ne[2]); - GGML_ASSERT(src0->ne[3] == 1); - GGML_ASSERT(dst->ne[0] == src0->ne[0] && dst->ne[1] == src1->ne[0] && dst->ne[2] == src1->ne[1] && dst->ne[2]); - - size_t row_size = ggml_row_size(dst->type, dst->ne[0]); - //GGML_ASSERT(row_size % 2 == 0); - //size_t type_size = row_size % 4 == 0 ? 4 : 2; - - for (int ir = first; ir < last; ++ir) { - int i3 = ir/(dst->ne[1]*dst->ne[2]); - int i2 = (ir - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; - int i1 = ir - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; - const int * idx = (const int *)((const char *)src1->data + i2*src1->nb[1] + i3*src1->nb[2]); - const char * cx = (const char *)src0->data + idx[i1]*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]; - char * cy = ( char *) dst->data + i1 * dst->nb[1] + i2* dst->nb[2] + i3* dst->nb[3]; - memcpy(cy, cx, row_size); - } -} - static void ggml_compute_forward_get_rows( const struct ggml_compute_params * params, struct ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; - if (dst->type != GGML_TYPE_F32) { - ggml_compute_forward_get_rows_any(params, dst); - return; - } - switch (src0->type) { case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index 37d1e68a5..b9909e5d6 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -1118,10 +1118,7 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu ggml_tensor * cache, const auto & extra_ctx, const std::string & tag, int n_swa_eff) { auto n_stream = std::max(1, lctx.dsv4.cache.n_stream); - auto extra_k = cache; - if (extra_k->ne[1] > 1) { - extra_k = dsv4_comp_get_k(ctx0, cache, extra_ctx, n_embd_head, cache->ne[1]/n_stream); - } + auto extra_k = dsv4_comp_get_k(ctx0, cache, extra_ctx, n_embd_head, cache->ne[1]/n_stream); if (cparams.flash_attn) { extra_mask = dsv4_pad_mask_tokens(ctx0, extra_mask, n_tokens); } @@ -1149,6 +1146,7 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu auto attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_mask, model.layers[il].attn_sinks, kq_scale, cb, il, n_swa_eff, gf); return attn; + //return std::make_pair(k_all, kq_mask); }; auto num_streams = [] (const auto & comp) { @@ -1162,24 +1160,16 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu lctx.dsv4.lid_plan.n_kv > 0 && !cparams.k_cache_hadamard) { auto csa_mask = lctx.dsv4.inputs.csa.kq_mask; - auto csa_kv = lctx.dsv4.cache.csa_k[il]; if (hparams.indexer_top_k < lctx.dsv4.inputs.csa.kq_mask->ne[0]) { auto top_k = dsv4_build_lid_top_k(ctx0, llm, qr, cur, inp_pos, il, gf, cb); - if (n_tokens == 1) { - // When we are dealing with a single token, we can just use ggml_get_rows_ext to get the - // selected rows from the CSA cache and setup the corresponding mask. This makes the - // raw_kv and csa_kv concetenation much less expensive for long context. - csa_kv = ggml_get_rows_ext(ctx0, csa_kv, top_k, true, false); - csa_kv = ggml_reshape_3d(ctx0, csa_kv, csa_kv->ne[0], 1, csa_kv->ne[1]); - csa_mask = ggml_get_rows_ext(ctx0, csa_mask, top_k, true, true); - } else { - csa_mask = build_top_k_mask(ctx0, dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.csa.kq_mask, nullptr, - lctx.dsv4.csa_plan.n_kv, n_tokens, num_streams(lctx.dsv4.csa_ctx), cb, il), top_k); - cb(csa_mask, "csa_mask", il); - } + csa_mask = build_top_k_mask(ctx0, + dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.csa.kq_mask, nullptr, + lctx.dsv4.csa_plan.n_kv, n_tokens, num_streams(lctx.dsv4.csa_ctx), cb, il), + top_k); + cb(csa_mask, "csa_mask", il); } int n_csa = hparams.n_swa + hparams.indexer_top_k; - attn = build_the_attn(raw_k, raw_mask, csa_mask, csa_kv, lctx.dsv4.csa_ctx, "csa", n_csa); + attn = build_the_attn(raw_k, raw_mask, csa_mask, lctx.dsv4.cache.csa_k[il], lctx.dsv4.csa_ctx, "csa", n_csa); cb(attn, "attn_csa", il); } else if (ratio == llama_context::dsv4_runtime::HCA_RATIO && lctx.dsv4.inputs.hca.kq_mask != nullptr &&