metal: Improve deepseek v4 token generation

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    5
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    3
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -571,9 +571,12 @@ extern "C" {
         GGML_OP_SOLVE_TRI,
         GGML_OP_GATED_DELTA_NET,
         GGML_OP_LIGHTNING_INDEXER,
+        GGML_OP_DSV4_COMPRESS,
+        GGML_OP_DSV4_TOP_K_MASK,
         GGML_OP_DSV4_HC_COMB,
         GGML_OP_DSV4_HC_PRE,
         GGML_OP_DSV4_HC_POST,
+        GGML_OP_DSV4_SPARSE_PACK,
 
         GGML_OP_UNARY,
 
@@ -2434,8 +2437,14 @@ extern "C" {
             const struct ggml_tensor * a);
 
     GGML_API void ggml_flash_attn_ext_add_sinks(
-            struct ggml_tensor * a,
-            struct ggml_tensor * sinks);
+        struct ggml_tensor * a,
+        struct ggml_tensor * sinks);
+
+    // Like ggml_flash_attn_ext_add_sinks(), but indexes sinks by the query-row
+    // dimension instead of the attention-head dimension.
+    GGML_API void ggml_flash_attn_ext_add_sinks_rows(
+        struct ggml_tensor * a,
+        struct ggml_tensor * sinks);
 
     // TODO: needs to be adapted to ggml_flash_attn_ext
     GGML_API struct ggml_tensor * ggml_flash_attn_back(
@@ -2601,6 +2610,40 @@ extern "C" {
         struct ggml_tensor  * weights,
         struct ggml_tensor  * mask);
 
+    // DeepSeek V4 compressor weighted reduction.
+    //
+    // kv_state, score_state: [overlap ? 2*n_embd : n_embd, n_rows]
+    // read_idxs:             [(overlap ? 2 : 1)*ratio*n_blocks]
+    // res:                   [n_embd, n_blocks]
+    GGML_API struct ggml_tensor * ggml_dsv4_compress(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * kv_state,
+            struct ggml_tensor  * score_state,
+            struct ggml_tensor  * read_idxs,
+            int32_t               ratio,
+            bool                  overlap);
+
+    // Builds the raw + selected-compressed F16 attention mask in one pass.
+    // raw_mask:  [n_raw, n_query, 1, n_stream]
+    // comp_mask: [n_comp, n_query, 1, n_stream]
+    // comp_idx:  [n_select, n_query, 1, n_stream]
+    GGML_API struct ggml_tensor * ggml_dsv4_top_k_mask(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * raw_mask,
+            struct ggml_tensor  * comp_mask,
+            struct ggml_tensor  * comp_idx);
+
+    // Packs per-token raw-window and Lightning-Indexer selections into the
+    // strided K + mask storage consumed by DeepSeek V4 sparse flash attention.
+    GGML_API struct ggml_tensor * ggml_dsv4_sparse_pack(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * raw_k,
+            struct ggml_tensor  * comp_k,
+            struct ggml_tensor  * raw_mask,
+            struct ggml_tensor  * comp_mask,
+            struct ggml_tensor  * comp_idx,
+            int64_t               n_raw);
+
     // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
     // In short these operations are replacements for the original residual connection (x = transformer(x) + x)
     // using a richer representation through streams.

ggml/src/ggml-backend-meta.cpp

@@ -984,9 +984,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_GATED_DELTA_NET: {
                 split_state = handle_gated_delta_net(src_ss);
             } break;
+            case GGML_OP_DSV4_COMPRESS:
+            case GGML_OP_DSV4_TOP_K_MASK:
             case GGML_OP_DSV4_HC_COMB:
             case GGML_OP_DSV4_HC_PRE:
-            case GGML_OP_DSV4_HC_POST: {
+            case GGML_OP_DSV4_HC_POST:
+            case GGML_OP_DSV4_SPARSE_PACK: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
             case GGML_OP_UNARY: {

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2064,6 +2064,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_lightning_indexer(params, tensor);
             } break;
+        case GGML_OP_DSV4_COMPRESS:
+            {
+                ggml_compute_forward_dsv4_compress(params, tensor);
+            } break;
+        case GGML_OP_DSV4_TOP_K_MASK:
+            {
+                ggml_compute_forward_dsv4_top_k_mask(params, tensor);
+            } break;
         case GGML_OP_DSV4_HC_COMB:
             {
                 ggml_compute_forward_dsv4_hc_comb(params, tensor);
@@ -2076,6 +2084,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_dsv4_hc_post(params, tensor);
             } break;
+        case GGML_OP_DSV4_SPARSE_PACK:
+            {
+                ggml_compute_forward_dsv4_sparse_pack(params, tensor);
+            } break;
         case GGML_OP_MAP_CUSTOM1:
             {
                 ggml_compute_forward_map_custom1(params, tensor);
@@ -2256,9 +2268,12 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_COUNT_EQUAL:
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_GATED_DELTA_NET:
+        case GGML_OP_DSV4_COMPRESS:
+        case GGML_OP_DSV4_TOP_K_MASK:
         case GGML_OP_DSV4_HC_COMB:
         case GGML_OP_DSV4_HC_PRE:
         case GGML_OP_DSV4_HC_POST:
+        case GGML_OP_DSV4_SPARSE_PACK:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/ops.cpp

@@ -8528,6 +8528,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
     memcpy(&scale,         (float *) dst->op_params + 0, sizeof(float));
     memcpy(&max_bias,      (float *) dst->op_params + 1, sizeof(float));
     memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float));
+    const bool sinks_rows = ggml_get_op_params_i32(dst, 4);
 
     if (logit_softcap != 0) {
         scale /= logit_softcap;
@@ -8572,7 +8573,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
             memset(VKQ32, 0, DV*sizeof(float));
         }
 
-        const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + iq1*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL;
+        const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + (iq1%mask->ne[1])*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL;
 
         // k indices
         const int ik3 = iq3 / rk3;
@@ -8664,7 +8665,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
 
         // sinks - apply only on the first kv-chunk
         if (sinks && ic_start == 0) {
-            const float s = ((float *)((char *) sinks->data))[h];
+            const float s = ((float *)((char *) sinks->data))[sinks_rows ? iq1 : h];
 
             float ms = 1.0f;
             float vs = 1.0f;
@@ -8764,6 +8765,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
     memcpy(&scale,         (float *) dst->op_params + 0, sizeof(float));
     memcpy(&max_bias,      (float *) dst->op_params + 1, sizeof(float));
     memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float));
+    const bool sinks_rows = ggml_get_op_params_i32(dst, 4);
 
     if (logit_softcap != 0) {
         scale /= logit_softcap;
@@ -8853,7 +8855,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
             if (mask) {
                 bool can_skip = true;
                 for (int tq = 0; tq < tile_rows; tq++) {
-                    const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + (iq1 + tq)*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]);
+                    const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + ((iq1 + tq)%mask->ne[1])*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]);
                     for (int tk = 0; tk < kv_tile; tk++) {
                         mask32[tq * KV_TILE_SZ + tk] = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]);
                         if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) {
@@ -8956,9 +8958,8 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
 
         // sinks (apply only to valid rows in the tile)
         if (sinks) {
-            const float s = ((float *)((char *) sinks->data))[h];
-
             for (int tq = 0; tq < tile_rows; tq++) {
+                const float s = ((float *)((char *) sinks->data))[sinks_rows ? iq1 + tq : h];
                 float ms = 1.0f;
                 float vs = 1.0f;
 
@@ -10945,6 +10946,132 @@ void ggml_compute_forward_gated_delta_net(
 }
 
 
+// ggml_compute_forward_dsv4_compress
+
+void ggml_compute_forward_dsv4_compress(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * kv_state    = dst->src[0];
+    const ggml_tensor * score_state = dst->src[1];
+    const ggml_tensor * read_idxs   = dst->src[2];
+
+    GGML_ASSERT(kv_state->type    == GGML_TYPE_F32);
+    GGML_ASSERT(score_state->type == GGML_TYPE_F32);
+    GGML_ASSERT(read_idxs->type   == GGML_TYPE_I32);
+    GGML_ASSERT(dst->type         == GGML_TYPE_F32);
+
+    const int32_t ratio   = ggml_get_op_params_i32(dst, 0);
+    const bool    overlap = ggml_get_op_params_i32(dst, 1) != 0;
+    const int64_t n_embd   = dst->ne[0];
+    const int64_t n_blocks = dst->ne[1];
+    const int64_t n_rows   = kv_state->ne[1];
+    const int64_t n_read   = (overlap ? 2 : 1)*ratio;
+
+    GGML_ASSERT(ratio > 0 && n_blocks > 0);
+    GGML_ASSERT(kv_state->ne[0] == (overlap ? 2 : 1)*n_embd);
+    GGML_ASSERT(kv_state->ne[0] == score_state->ne[0]);
+    GGML_ASSERT(kv_state->ne[1] == score_state->ne[1]);
+    GGML_ASSERT(read_idxs->ne[0] == n_read*n_blocks);
+
+    GGML_TENSOR_LOCALS(size_t, nbk, kv_state,    nb);
+    GGML_TENSOR_LOCALS(size_t, nbs, score_state, nb);
+    GGML_TENSOR_LOCALS(size_t, nbi, read_idxs,   nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,         nb);
+
+    const int64_t nr  = n_embd*n_blocks;
+    const int64_t dr  = (nr + params->nth - 1)/params->nth;
+    const int64_t ir0 = dr*params->ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t i0 = ir % n_embd;
+        const int64_t ib = ir / n_embd;
+
+        float score_max = -INFINITY;
+        for (int64_t j = 0; j < n_read; ++j) {
+            const bool cur_half = overlap && j >= ratio;
+            const int64_t jr = cur_half ? j - ratio : j;
+            const int64_t idx_pos = (cur_half ? ratio*n_blocks : 0) + ib*ratio + jr;
+            const int32_t idx = *(const int32_t *) ((const char *) read_idxs->data + idx_pos*nbi0);
+
+            GGML_ASSERT(idx >= 0 && idx <= n_rows);
+            if (idx == n_rows) {
+                continue;
+            }
+
+            const int64_t i_src = (cur_half ? n_embd : 0) + i0;
+            const float score = *(const float *) ((const char *) score_state->data + i_src*nbs0 + idx*nbs1);
+            score_max = MAX(score_max, score);
+        }
+
+        float sum_v = 0.0f;
+        float sum_w = 0.0f;
+        if (score_max != -INFINITY) {
+            for (int64_t j = 0; j < n_read; ++j) {
+                const bool cur_half = overlap && j >= ratio;
+                const int64_t jr = cur_half ? j - ratio : j;
+                const int64_t idx_pos = (cur_half ? ratio*n_blocks : 0) + ib*ratio + jr;
+                const int32_t idx = *(const int32_t *) ((const char *) read_idxs->data + idx_pos*nbi0);
+
+                if (idx == n_rows) {
+                    continue;
+                }
+
+                const int64_t i_src = (cur_half ? n_embd : 0) + i0;
+                const float score = *(const float *) ((const char *) score_state->data + i_src*nbs0 + idx*nbs1);
+                const float weight = expf(score - score_max);
+                const float value = *(const float *) ((const char *) kv_state->data + i_src*nbk0 + idx*nbk1);
+                sum_v += value*weight;
+                sum_w += weight;
+            }
+        }
+
+        *(float *) ((char *) dst->data + i0*nbd0 + ib*nbd1) = sum_w > 0.0f ? sum_v/sum_w : 0.0f;
+    }
+}
+
+// ggml_compute_forward_dsv4_top_k_mask
+
+void ggml_compute_forward_dsv4_top_k_mask(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * raw_mask  = dst->src[0];
+    const ggml_tensor * comp_mask = dst->src[1];
+    const ggml_tensor * comp_idx  = dst->src[2];
+
+    GGML_ASSERT(raw_mask->type == GGML_TYPE_F16);
+    GGML_ASSERT(comp_mask->type == GGML_TYPE_F16);
+    GGML_ASSERT(comp_idx->type == GGML_TYPE_I32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F16);
+
+    const int64_t n_raw  = raw_mask->ne[0];
+    const int64_t n_comp = comp_mask->ne[0];
+    const int64_t n_sel  = comp_idx->ne[0];
+    const int64_t nq     = raw_mask->ne[1];
+    const int64_t nrows  = nq*raw_mask->ne[3];
+    const ggml_fp16_t neg_inf = GGML_CPU_FP32_TO_FP16(-INFINITY);
+
+    for (int64_t row = params->ith; row < nrows; row += params->nth) {
+        const int64_t iq = row % nq;
+        const int64_t is = row / nq;
+        ggml_fp16_t * out = (ggml_fp16_t *) ((char *) dst->data + iq*dst->nb[1] + is*dst->nb[3]);
+
+        for (int64_t i = 0; i < n_raw; ++i) {
+            out[i] = *(const ggml_fp16_t *) ((const char *) raw_mask->data +
+                    i*raw_mask->nb[0] + iq*raw_mask->nb[1] + is*raw_mask->nb[3]);
+        }
+        std::fill(out + n_raw, out + n_raw + n_comp, neg_inf);
+
+        for (int64_t i = 0; i < n_sel; ++i) {
+            const int32_t idx = *(const int32_t *) ((const char *) comp_idx->data +
+                    i*comp_idx->nb[0] + iq*comp_idx->nb[1] + is*comp_idx->nb[3]);
+            GGML_ASSERT(idx >= 0 && idx < n_comp);
+            out[n_raw + idx] = *(const ggml_fp16_t *) ((const char *) comp_mask->data +
+                    idx*comp_mask->nb[0] + iq*comp_mask->nb[1] + is*comp_mask->nb[3]);
+        }
+    }
+}
+
 // ggml_compute_forward_dsv4_hc_comb
 
 static void ggml_dsv4_hc_comb_norm_cols(float * comb, float eps) {
@@ -11229,6 +11356,65 @@ void ggml_compute_forward_dsv4_hc_post(
     }
 }
 
+// ggml_compute_forward_dsv4_sparse_pack
+
+void ggml_compute_forward_dsv4_sparse_pack(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * raw_k     = dst->src[0];
+    const ggml_tensor * comp_k    = dst->src[1];
+    const ggml_tensor * raw_mask  = dst->src[2];
+    const ggml_tensor * comp_mask = dst->src[3];
+    const ggml_tensor * comp_idx  = dst->src[4];
+
+    const int64_t d      = raw_k->ne[0];
+    const int64_t nq     = raw_mask->ne[1];
+    const int64_t nt     = dst->ne[1];
+    const int64_t nr     = ggml_get_op_params_i32(dst, 0);
+    const int64_t nc     = comp_idx->ne[0];
+    const int64_t nk     = nr + nc;
+
+    GGML_ASSERT(dst->type == GGML_TYPE_F16);
+
+    for (int64_t it = params->ith; it < nt; it += params->nth) {
+        const int64_t iq = it % nq;
+        const int64_t is = it / nq;
+        ggml_fp16_t * out = (ggml_fp16_t *) ((char *) dst->data + it*dst->nb[1]);
+        ggml_fp16_t * out_k = out;
+        ggml_fp16_t * out_m = out + d*nk;
+
+        int64_t ir = 0;
+        for (int64_t idx = 0; idx < raw_k->ne[2] && ir < nr; ++idx) {
+            const ggml_fp16_t m = *(const ggml_fp16_t *) ((const char *) raw_mask->data +
+                    idx*raw_mask->nb[0] + iq*raw_mask->nb[1] + is*raw_mask->nb[3]);
+            if (!std::isfinite(GGML_CPU_FP16_TO_FP32(m))) {
+                continue;
+            }
+            memcpy(out_k + ir*d, (const char *) raw_k->data + idx*raw_k->nb[2] + is*raw_k->nb[3],
+                    d*sizeof(ggml_fp16_t));
+            out_m[ir] = m;
+            ++ir;
+        }
+        for (; ir < nr; ++ir) {
+            memset(out_k + ir*d, 0, d*sizeof(ggml_fp16_t));
+            out_m[ir] = GGML_CPU_FP32_TO_FP16(-INFINITY);
+        }
+
+        for (int64_t i = 0; i < nc; ++i) {
+            const int64_t oi = nr + i;
+            const int32_t idx = *(const int32_t *) ((const char *) comp_idx->data +
+                    i*comp_idx->nb[0] + iq*comp_idx->nb[1] + is*comp_idx->nb[3]);
+            GGML_ASSERT(idx >= 0 && idx < comp_k->ne[2]);
+            memcpy(out_k + oi*d, (const char *) comp_k->data + idx*comp_k->nb[2] + is*comp_k->nb[3],
+                    d*sizeof(ggml_fp16_t));
+            const ggml_fp16_t m = *(const ggml_fp16_t *) ((const char *) comp_mask->data +
+                    idx*comp_mask->nb[0] + iq*comp_mask->nb[1] + is*comp_mask->nb[3]);
+            out_m[oi] = m;
+        }
+
+    }
+}
+
 // ggml_compute_forward_rwkv_wkv7
 
 static void ggml_compute_forward_rwkv_wkv7_f32(