From 593a196f7811028b170e7b7d4118d3b05914a882 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 13 Jul 2026 13:28:05 +0300 Subject: [PATCH] DSA: Better way to build the attention mask (#2119) * WIP: indexer_topk on CUDA * Forgot these * WIP * WIP * This seems to work * Minor * Fix bug. Fix suggested by @sayap using GLM-5.2 * GLM-DSA: much better PP long context performance (CUDA) * DSA: Better way to build the attention mask --- ggml/include/ggml.h | 6 +++ ggml/src/ggml-cuda.cu | 4 ++ ggml/src/ggml-cuda/indexer_topk.cu | 61 ++++++++++++++++++++++++++ ggml/src/ggml-cuda/indexer_topk.cuh | 2 + ggml/src/ggml.c | 35 ++++++++++++++- ggml/src/iqk/iqk_cpu_ops.cpp | 68 +++++++++++++++++++++++++++++ ggml/src/iqk/iqk_cpu_ops.h | 2 + src/graphs/build_deepseek2.cpp | 35 +-------------- src/llama-context.h | 1 - 9 files changed, 177 insertions(+), 37 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 40b2c7ca..d3c540da 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -705,6 +705,7 @@ extern "C" { GGML_OP_FUSED_RMS_RMS_ADD, GGML_OP_BLEND, GGML_OP_INDEXER_TOPK, + GGML_OP_MASK_TOPK, GGML_OP_SINKHORN, GGML_OP_COUNT, @@ -2406,6 +2407,11 @@ extern "C" { struct ggml_tensor * b, float c); + GGML_API struct ggml_tensor * ggml_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * mask, + struct ggml_tensor * topk); + // sort rows enum ggml_sort_order { diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index cd88309e..f11229b4 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -4137,6 +4137,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_INDEXER_TOPK: ggml_cuda_op_indexer_topk(ctx, dst); break; + case GGML_OP_MASK_TOPK: + ggml_cuda_op_indexer_mask(ctx, dst); + break; default: return false; } @@ -5041,6 +5044,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons op->src[3]->ne[0] == op->src[0]->ne[2]; case GGML_OP_DELTA_NET: case GGML_OP_INDEXER_TOPK: + case GGML_OP_MASK_TOPK: return true; case GGML_OP_SINKHORN: { const int sink_s = op->op_params[0]; diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu index f03ec23a..14b12386 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cu +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -214,3 +214,64 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds } } + +template +static __global__ void k_indexer_mask(int ne0, int ne1, int ne2, int ntopk, int ne11, + 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 mask_t * mask, const int * idx, mask_t * dst) { + int i1 = blockIdx.x; + int i3 = i1 / (ne1*ne2); i1 -= i3*ne1*ne2; + int i2 = i1 / (ne1); i1 -= i2*ne1; + + auto m = (const mask_t *)((const char *)mask + i1*nb01 + i2*nb02 + i3*nb03); + auto i = (const int *)((const char *)idx + i1*nb11 + i2*nb12 + i3*nb13); + auto d = (mask_t *)((char *)dst + i1*nb1 + i2*nb2 + i3*nb3); + + mask_t inf, zero; + if constexpr (std::is_same_v) { + inf = __float2half(-INFINITY); + zero = __float2half(0.0f); + } else { + inf = -INFINITY; + zero = 0.0f; + } + + if (i1 < ne11) { + for (int j = threadIdx.x; j < ne0; j += blockDim.x) d[j] = inf; + for (int j = threadIdx.x; j < ntopk; j += blockDim.x) d[i[j]] = zero; + for (int j = threadIdx.x; j < ne0; j += blockDim.x) d[j] += m[j]; + } else { + for (int j = threadIdx.x; j < ne0; j += blockDim.x) d[j] = m[j]; + } + +} + +void ggml_cuda_op_indexer_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + auto mask = dst->src[0]; + auto topk = dst->src[1]; + GGML_ASSERT(mask->ne[0] >= topk->ne[1]); + GGML_ASSERT(mask->ne[1] >= topk->ne[1] && mask->ne[2] == topk->ne[2] && mask->ne[3] == topk->ne[3]); + GGML_ASSERT(ggml_are_same_shape(mask, dst)); + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(topk->type == GGML_TYPE_I32); + GGML_ASSERT(mask->type == dst->type); + + int nrows = ggml_nrows(dst); + + if (dst->type == GGML_TYPE_F16) { + k_indexer_mask<<>>(dst->ne[0], dst->ne[1], dst->ne[2], topk->ne[0], topk->ne[1], + mask->nb[1], mask->nb[2], mask->nb[3], + topk->nb[1], topk->nb[2], topk->nb[3], + dst->nb[1], dst->nb[2], dst->nb[3], + (const half *)mask->data, (const int *)topk->data, (half *)dst->data); + } else { + k_indexer_mask<<>>(dst->ne[0], dst->ne[1], dst->ne[2], topk->ne[0], topk->ne[1], + mask->nb[1], mask->nb[2], mask->nb[3], + topk->nb[1], topk->nb[2], topk->nb[3], + dst->nb[1], dst->nb[2], dst->nb[3], + (const float *)mask->data, (const int *)topk->data, (float *)dst->data); + } + +} diff --git a/ggml/src/ggml-cuda/indexer_topk.cuh b/ggml/src/ggml-cuda/indexer_topk.cuh index 4f919e9a..e131d26d 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cuh +++ b/ggml/src/ggml-cuda/indexer_topk.cuh @@ -6,3 +6,5 @@ #include "common.cuh" void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_indexer_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 3fe4e6d4..9ec3886c 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4333,10 +4333,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "FUSED_RMS_RMS_ADD", "BLEND", "INDEXER_TOPK", + "MASK_TOPK", "SINKHORN", }; -static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); +static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -4456,11 +4457,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "rms(x1)+rms(x2)", "blend(a,b,c)", "indexer_topk(k, q, w, mask)", + "mask_topk(mask, topk)", "sinkhorn(x)", }; -static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); +static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -10230,6 +10232,29 @@ struct ggml_tensor * ggml_blend( return result; } +struct ggml_tensor * ggml_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * mask, + struct ggml_tensor * topk) { + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(topk->type == GGML_TYPE_I32); + // The mask may be padded along dim 1, so topk->ne[1] must be <= mask->ne[1] + if (topk->ne[1] > mask->ne[1] || topk->ne[2] != mask->ne[2] || topk->ne[3] != mask->ne[3]) { + printf("%s: Oops. topk is %ld x %ld x %ld x %ld, mask is %ld x %ld x %ld x %ld\n", __func__, + topk->ne[0], topk->ne[1], topk->ne[2], topk->ne[3], + mask->ne[0], mask->ne[1], mask->ne[2], mask->ne[3]); + } + GGML_ASSERT(topk->ne[1] <= mask->ne[1] && topk->ne[2] == mask->ne[2] && topk->ne[3] == mask->ne[3]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, mask); + result->src[0] = mask; + result->src[1] = topk; + result->op = GGML_OP_MASK_TOPK; + + return result; +} + + // ggml_argsort struct ggml_tensor * ggml_argsort( @@ -24943,6 +24968,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml GGML_ABORT("Fatal error"); } } break; + case GGML_OP_MASK_TOPK: + { + iqk_mask_topk(tensor, params->ith, params->nth); + } break; case GGML_OP_WIN_PART: { ggml_compute_forward_win_part(params, tensor); @@ -26005,6 +26034,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor case GGML_OP_SOLVE_TRI: case GGML_OP_DELTA_NET: case GGML_OP_INDEXER_TOPK: + case GGML_OP_MASK_TOPK: case GGML_OP_SINKHORN: { GGML_ABORT("fatal error"); // TODO: not implemented @@ -26743,6 +26773,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_SOLVE_TRI: case GGML_OP_DELTA_NET: case GGML_OP_INDEXER_TOPK: + case GGML_OP_MASK_TOPK: case GGML_OP_SINKHORN: { n_tasks = n_threads; diff --git a/ggml/src/iqk/iqk_cpu_ops.cpp b/ggml/src/iqk/iqk_cpu_ops.cpp index 0102f8cb..e50f51ea 100644 --- a/ggml/src/iqk/iqk_cpu_ops.cpp +++ b/ggml/src/iqk/iqk_cpu_ops.cpp @@ -1044,3 +1044,71 @@ void iqk_blend(struct ggml_tensor * dst, int ith, int nth) { } } } + +namespace { +inline void iqk_add_f16(int n, const ggml_half * x, ggml_half * y) { + int i = 0; +#if defined __AVX2__ && defined __F16C__ + for ( ; i + 7 < n; i += 8) { + auto vx32 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i))); + auto vy32 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(y + i))); + vy32 = _mm256_add_ps(vy32, vx32); + _mm_storeu_si128((__m128i *)(y + i), _mm256_cvtps_ph(vy32, _MM_FROUND_TO_NEAREST_INT)); + } +#endif + for (; i < n; ++i) { + float y32 = GGML_FP16_TO_FP32(y[i]) + GGML_FP16_TO_FP32(x[i]); + y[i] = GGML_FP32_TO_FP16(y32); + } +} +} + +void iqk_mask_topk(struct ggml_tensor * dst, int ith, int nth) { + auto mask = dst->src[0]; + auto topk = dst->src[1]; + GGML_ASSERT(mask->ne[0] >= topk->ne[1]); + GGML_ASSERT(mask->ne[1] >= topk->ne[1] && mask->ne[2] == topk->ne[2] && mask->ne[3] == topk->ne[3]); + GGML_ASSERT(ggml_are_same_shape(mask, dst)); + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(topk->type == GGML_TYPE_I32); + GGML_ASSERT(mask->type == dst->type); + + int nrows = ggml_nrows(dst); + int npt = (nrows + nth - 1)/nth; + int first = ith*npt; + int last = std::min(first + npt, nrows); + + auto hinf = GGML_FP32_TO_FP16(-INFINITY); + auto hzero = GGML_FP32_TO_FP16(0.0f); + + int n = dst->ne[0]; + int nidx = topk->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]; + auto idx = (const int *)((const char *)topk->data + topk->nb[1]*i1 + topk->nb[2]*i2 + topk->nb[3]*i3); + if (dst->type == GGML_TYPE_F16) { + auto x = (const ggml_half *)((const char *)mask->data + mask->nb[1]*i1 + mask->nb[2]*i2 + mask->nb[3]*i3); + auto y = (ggml_half *)((char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]); + if (i1 < topk->ne[1]) { + for (int j = 0; j < n; ++j) y[j] = hinf; + for (int j = 0; j < nidx; ++j) y[idx[j]] = hzero; + iqk_add_f16(n, x, y); + } else { + for (int j = 0; j < n; ++j) y[j] = x[j]; + } + } else { + auto x = (const float *)((const char *)mask->data + mask->nb[1]*i1 + mask->nb[2]*i2 + mask->nb[3]*i3); + auto y = (float *)((char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]); + if (i1 < topk->ne[1]) { + for (int j = 0; j < n; ++j) y[j] = -INFINITY; + for (int j = 0; j < nidx; ++j) y[idx[j]] = 0.0f; + for (int j = 0; j < n; ++j) y[j] += x[j]; + } else { + for (int j = 0; j < n; ++j) y[j] = x[j]; + } + } + } +} diff --git a/ggml/src/iqk/iqk_cpu_ops.h b/ggml/src/iqk/iqk_cpu_ops.h index 7d6eed74..5847dc22 100644 --- a/ggml/src/iqk/iqk_cpu_ops.h +++ b/ggml/src/iqk/iqk_cpu_ops.h @@ -43,6 +43,8 @@ void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth); void iqk_blend(struct ggml_tensor * dst, int ith, int nth); +void iqk_mask_topk(struct ggml_tensor * dst, int ith, int nth); + #ifdef __cplusplus } #endif diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index 77c8a64f..a8a300d2 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -670,34 +670,6 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_sparse_mask( return sparse; } -static ggml_tensor * build_deepseek2_dsa_fa_mask(const llama_context & lctx, ggml_context * ctx0, ggml_tensor * KQ_mask, ggml_tensor * sorted) { - GGML_ASSERT(KQ_mask && KQ_mask->type == GGML_TYPE_F16); - GGML_ASSERT(sorted && sorted->type == GGML_TYPE_I32); - GGML_ASSERT(KQ_mask->ne[1] >= sorted->ne[1]); - - int n_top_k = (int64_t) lctx.model.hparams.indexer_top_k; - if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k; - - int n_kv_local = KQ_mask->ne[0]; - if (n_top_k >= n_kv_local) { - return KQ_mask; - } - - GGML_ASSERT(sorted->ne[1] == lctx.inp_mask_inf->ne[1]); - auto top_k = ggml_view_2d(ctx0, sorted, n_top_k, sorted->ne[1], sorted->nb[1], 0); - auto mask32 = ggml_blend(ctx0, lctx.inp_mask_inf, top_k, 0.0f); - if (KQ_mask->ne[1] == mask32->ne[1]) { - auto mask16 = ggml_add(ctx0, KQ_mask, mask32); - return mask16; - } - auto kq1 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], mask32->ne[1], KQ_mask->nb[1], 0); - auto kq2 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], KQ_mask->ne[1] - mask32->ne[1], KQ_mask->nb[1], mask32->ne[1]*KQ_mask->nb[1]); - kq1 = ggml_add(ctx0, kq1, mask32); - auto mask16 = ggml_concat(ctx0, kq1, kq2, 1); - return mask16; -} - - // Adapt the (F32, unpadded {n_kv, n_tokens}) sparse mask for ggml_flash_attn_ext, which on this fork // requires the mask to be F16, contiguous, and padded in ne[1] to GGML_PAD(n_queries, GGML_KQ_MASK_PAD) // (build_inp_KQ_mask creates the dense -fa 1 mask exactly that way). We: @@ -829,7 +801,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( dsa_last_full_sorted = build_deepseek2_dsa_indexer(gf, il, q, cur, KQ_mask, inp_pos); if (dsa_last_full_sorted && !dsa_fast_path) { if (lctx.cparams.flash_attn) { - last_sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, dsa_last_full_sorted); + last_sparse_mask_fa = ggml_indexer_mask(ctx0, KQ_mask, dsa_last_full_sorted); } else { last_sparse_mask = build_deepseek2_dsa_sparse_mask(dsa_last_full_sorted, KQ_mask); } @@ -1293,11 +1265,6 @@ ggml_cgraph * llm_build_context::build_deepseek2() { mask = ggml_reshape_2d(ctx0, mask, n_padded, KQ_mask->ne[1]); ggml_build_forward_expand(gf, mask); dsa_tg_fast_mask = mask; - } else { - auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens); - minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY); - ggml_build_forward_expand(gf, minus_inf); - lctx.inp_mask_inf = minus_inf; } } diff --git a/src/llama-context.h b/src/llama-context.h index 6c8c0185..c5cf1a85 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -414,7 +414,6 @@ struct llama_context { struct ggml_tensor * inp_mtp_states = nullptr; struct ggml_tensor * inp_mtp_carry = nullptr; // F32 [n_embd, nextn-1] per-head hidden at the last committed position struct ggml_tensor * inp_dsa_sink = nullptr; // F32 [n_kv, n_tokens] per-sequence attention-sink boost for DSA indexer top-k - struct ggml_tensor * inp_mask_inf = nullptr; struct openpangu_swa_window_view_state { bool active = false;