metal: deepseek v4 optimize 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    1
+#define RPC_PROTO_PATCH_VERSION    3
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 102, "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,6 +571,8 @@ 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,
@@ -2608,6 +2610,29 @@ 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(

ggml/src/ggml-backend-meta.cpp

@@ -984,6 +984,8 @@ 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:

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);
@@ -2260,6 +2268,8 @@ 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:

ggml/src/ggml-cpu/ops.cpp

@@ -10946,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) {

ggml/src/ggml-cpu/ops.h

@@ -106,6 +106,8 @@ void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, s
 void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_compress(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -66,6 +66,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_base(ggml
     const char * op_str = "undefined";
     switch (op) {
         case GGML_OP_ADD_ID: op_str = "add_id"; break;
+        case GGML_OP_DSV4_COMPRESS: op_str = "dsv4_compress"; break;
+        case GGML_OP_DSV4_TOP_K_MASK: op_str = "dsv4_top_k_mask"; break;
         case GGML_OP_DSV4_SPARSE_PACK: op_str = "dsv4_sparse_pack"; break;
         default: GGML_ABORT("fatal error");
     };
@@ -1541,19 +1543,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
     const int32_t ns10 = op->src[1]->nb[1]/op->src[1]->nb[0];
     const int32_t ns20 = op->src[2]->nb[1]/op->src[2]->nb[0];
 
+    // DSV4 decode has a single 512-wide KV head, 64 query heads and a
+    // one-row top-k mask. Specialize per-row mask skipping to this signature
+    // so the extra checks cannot affect ordinary causal-attention kernels.
+    const bool sparse_mask = has_mask && has_sinks && dk == 512 && dv == 512 &&
+        op->src[0]->ne[1] == 1 && op->src[0]->ne[2] == 64 &&
+        op->src[1]->ne[2] == 1 && op->src[3]->ne[1] == 1 && op->src[4]->ne[0] == 64;
+
     snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d",
             "flash_attn_ext_vec",
             ggml_type_name(op->src[1]->type),
             dk,
             dv);
 
-    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
+    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_mskip=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
             base,
             has_mask,
             has_sinks,
             has_bias,
             has_scap,
             has_kvpad,
+            sparse_mask,
             ns10,
             ns20,
             nsg, nwg);
@@ -1567,6 +1577,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
         ggml_metal_cv_set_bool(cv, has_bias,  FC_FLASH_ATTN_EXT_VEC + 2);
         ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_VEC + 3);
         ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4);
+        ggml_metal_cv_set_bool(cv, sparse_mask, FC_FLASH_ATTN_EXT_VEC + 5);
 
         ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
         ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1351,6 +1351,25 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                 ggml_is_contiguous_rows(op->src[0]) &&
                 ggml_is_contiguous_rows(op->src[1]) &&
                 ggml_is_contiguous_rows(op->src[2]);
+        case GGML_OP_DSV4_COMPRESS:
+            return op->src[0]->type == GGML_TYPE_F32 &&
+                op->src[1]->type == GGML_TYPE_F32 &&
+                op->src[2]->type == GGML_TYPE_I32 &&
+                op->type         == GGML_TYPE_F32 &&
+                (ggml_get_op_params_i32(op, 1) ? 2 : 1)*ggml_get_op_params_i32(op, 0) <= 128 &&
+                ggml_is_contiguous_rows(op->src[0]) &&
+                ggml_is_contiguous_rows(op->src[1]) &&
+                ggml_is_contiguous(op->src[2]);
+        case GGML_OP_DSV4_TOP_K_MASK:
+            return op->src[0]->type == GGML_TYPE_F16 &&
+                op->src[1]->type == GGML_TYPE_F16 &&
+                op->src[2]->type == GGML_TYPE_I32 &&
+                op->type         == GGML_TYPE_F16 &&
+                op->src[2]->ne[0] <= op->src[1]->ne[0] &&
+                ggml_is_contiguous(op->src[0]) &&
+                ggml_is_contiguous(op->src[1]) &&
+                ggml_is_contiguous(op->src[2]) &&
+                ggml_is_contiguous(op);
         case GGML_OP_DSV4_HC_PRE:
             return has_simdgroup_reduction &&
                 op->src[0]->type == GGML_TYPE_F32 &&

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -1210,6 +1210,36 @@ typedef struct {
     float    eps;
 } ggml_metal_kargs_dsv4_hc_comb;
 
+typedef struct {
+    int32_t  n_embd;
+    int32_t  n_blocks;
+    int32_t  n_rows;
+    int32_t  ratio;
+    int32_t  overlap;
+    uint64_t nb_k0;
+    uint64_t nb_k1;
+    uint64_t nb_s0;
+    uint64_t nb_s1;
+    uint64_t nb_i0;
+    uint64_t nb_d0;
+    uint64_t nb_d1;
+} ggml_metal_kargs_dsv4_compress;
+
+typedef struct {
+    int32_t  n_raw;
+    int32_t  n_comp;
+    int32_t  n_select;
+    int32_t  n_query;
+    uint64_t nb_rm1;
+    uint64_t nb_rm3;
+    uint64_t nb_cm1;
+    uint64_t nb_cm3;
+    uint64_t nb_ci1;
+    uint64_t nb_ci3;
+    uint64_t nb_d1;
+    uint64_t nb_d3;
+} ggml_metal_kargs_dsv4_top_k_mask;
+
 typedef struct {
     int32_t  n_embd;
     int32_t  n_tokens;