DSA: Better way to build the attention mask

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 {

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];

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 <typename mask_t>
+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<mask_t, half>) {
+        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<<<nrows, 256, 0, ctx.stream()>>>(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<<<nrows, 256, 0, ctx.stream()>>>(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);
+    }
+
+}

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);

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;

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];
+            }
+        }
+    }
+}

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

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;
         }
     }
 

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;