From bccc54b136b4ea4d240a2c9d2cea71e32e56dbfc Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 18 Jul 2026 10:46:01 -0300 Subject: [PATCH 1/4] common: optimize DeepSeek V4 sparse attention Replace dense compressed-attention scans during DeepSeek V4 prefill with a backend-neutral sparse path that packs the sliding window and Lightning Indexer selections into a compact working set. - Add DSV4_SPARSE_PACK with a CPU reference and backend capability probes. - Extend Flash Attention with broadcast masks and row-indexed sinks. - Select sparse prefill graphs when supported while preserving dense fallback. - Test sparse packing and tiled and padded Flash Attention paths. Assisted-by: Codex --- ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 22 +++++++++- ggml/src/ggml-backend-meta.cpp | 3 +- ggml/src/ggml-cpu/ggml-cpu.c | 5 +++ ggml/src/ggml-cpu/ops.cpp | 70 +++++++++++++++++++++++++++--- ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml.c | 68 +++++++++++++++++++++++++++-- src/llama-context.cpp | 15 +++++++ src/llama-cparams.h | 2 + src/llama-graph.h | 1 + src/models/deepseek4.cpp | 53 +++++++++++++++++++++++ tests/test-backend-ops.cpp | 78 +++++++++++++++++++++++++++++++--- 12 files changed, 303 insertions(+), 19 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 276aea00ea1..a88eab44011 100644 --- a/ggml/include/ggml-rpc.h +++ b/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 1 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 35f0c44ec42..83de623960d 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -574,6 +574,7 @@ extern "C" { 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 +2435,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 +2608,17 @@ extern "C" { struct ggml_tensor * weights, struct ggml_tensor * mask); + // 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. diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index a5a3a58ad05..c8a345cf846 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -986,7 +986,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } break; 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: { diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 491316f7491..b471a24fbdd 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2076,6 +2076,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); @@ -2259,6 +2263,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { 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; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 42ec809ce52..df042d23ac1 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/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; @@ -11229,6 +11230,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( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 4c1642a6760..cb1629e0bfc 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -109,6 +109,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p 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); +void ggml_compute_forward_dsv4_sparse_pack(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 59191c663eb..ae904c6a1ad 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1083,6 +1083,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "DSV4_HC_COMB", "DSV4_HC_PRE", "DSV4_HC_POST", + "DSV4_SPARSE_PACK", "UNARY", @@ -1100,7 +1101,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1198,6 +1199,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "dsv4_hc_comb(mixes, scale, base)", "dsv4_hc_pre(x, weights)", "dsv4_hc_post(x, residual, post, comb)", + "dsv4_sparse_pack(raw_k, comp_k, raw_mask, comp_mask, comp_idx)", "unary(x)", @@ -1215,7 +1217,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -5416,7 +5418,7 @@ struct ggml_tensor * ggml_flash_attn_ext( if (mask) { GGML_ASSERT(mask->type == GGML_TYPE_F16); - GGML_ASSERT(ggml_is_contiguous(mask)); + GGML_ASSERT(ggml_is_contiguous_rows(mask)); //GGML_ASSERT(ggml_can_repeat_rows(mask, qk)); GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); @@ -5479,6 +5481,18 @@ void ggml_flash_attn_ext_add_sinks( a->src[4] = sinks; } +void ggml_flash_attn_ext_add_sinks_rows( + struct ggml_tensor * a, + struct ggml_tensor * sinks) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[4] == NULL); + GGML_ASSERT(a->src[0]->ne[1] == sinks->ne[0]); + GGML_ASSERT(sinks->type == GGML_TYPE_F32); + + a->src[4] = sinks; + ggml_set_op_params_i32(a, 4, 1); +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( @@ -6344,6 +6358,54 @@ struct ggml_tensor * ggml_lightning_indexer( return result; } +// ggml_dsv4_sparse_pack + +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) { + GGML_ASSERT(raw_k->type == GGML_TYPE_F16); + GGML_ASSERT(comp_k->type == GGML_TYPE_F16); + GGML_ASSERT(raw_mask->type == GGML_TYPE_F16); + GGML_ASSERT(comp_mask->type == GGML_TYPE_F16); + GGML_ASSERT(comp_idx->type == GGML_TYPE_I32); + + const int64_t d = raw_k->ne[0]; + const int64_t n_stream = raw_k->ne[3]; + const int64_t nq = raw_mask->ne[1]; + const int64_t nk = n_raw + comp_idx->ne[0]; + + GGML_ASSERT(n_raw > 0 && n_raw <= raw_k->ne[2]); + GGML_ASSERT(comp_k->ne[0] == d); + GGML_ASSERT(raw_k->ne[1] == 1 && comp_k->ne[1] == 1); + GGML_ASSERT(comp_k->ne[3] == n_stream); + GGML_ASSERT(raw_mask->ne[0] == raw_k->ne[2]); + GGML_ASSERT(comp_mask->ne[0] == comp_k->ne[2]); + GGML_ASSERT(comp_mask->ne[1] == nq); + GGML_ASSERT(raw_mask->ne[2] == 1 && comp_mask->ne[2] == 1); + GGML_ASSERT(raw_mask->ne[3] == n_stream && comp_mask->ne[3] == n_stream); + GGML_ASSERT(comp_idx->ne[1] == nq && comp_idx->ne[2] == 1); + GGML_ASSERT(comp_idx->ne[3] == n_stream); + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, + nk*(d + 1), nq*n_stream); + + result->op = GGML_OP_DSV4_SPARSE_PACK; + result->src[0] = raw_k; + result->src[1] = comp_k; + result->src[2] = raw_mask; + result->src[3] = comp_mask; + result->src[4] = comp_idx; + + ggml_set_op_params_i32(result, 0, n_raw); + + return result; +} + // ggml_dsv4_hc_comb struct ggml_tensor * ggml_dsv4_hc_comb( diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 19cca7df1e9..c53e66c35c6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -79,6 +79,12 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = { /*.n_tokens_per_seq =*/ 1, }; +static const llm_fused_op_probe llm_fused_op_dsv4_sparse_probe = { + /*.op =*/ LLM_FUSED_OP_DSV4_SPARSE_PACK, + /*.name =*/ "fused DeepSeek V4 sparse attention packing", + /*.n_tokens_per_seq =*/ 512, +}; + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -259,6 +265,9 @@ llama_context::llama_context( cparams.fused_dsv4_hc_post = true; cparams.auto_fhc = true; + cparams.fused_dsv4_sparse = true; + cparams.auto_fdsv4_sparse = true; + // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -572,6 +581,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post); cparams.auto_fhc = false; } + + if (cparams.auto_fdsv4_sparse) { + LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 sparse attention support:\n", func); + resolve(llm_fused_op_dsv4_sparse_probe, cparams.fused_dsv4_sparse); + cparams.auto_fdsv4_sparse = false; + } } void llama_context::sched_reserve() { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed85..933ba2c8612 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -47,6 +47,8 @@ struct llama_cparams { bool fused_dsv4_hc_comb; bool fused_dsv4_hc_post; bool auto_fhc; + bool fused_dsv4_sparse; // use indexed sparse attention packing for DSV4 prefill + bool auto_fdsv4_sparse; bool no_perf; bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP] bool op_offload; diff --git a/src/llama-graph.h b/src/llama-graph.h index 160e2941355..e15d822a256 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -46,6 +46,7 @@ enum llm_fused_op { LLM_FUSED_OP_DSV4_HC_PRE, LLM_FUSED_OP_DSV4_HC_COMB, LLM_FUSED_OP_DSV4_HC_POST, + LLM_FUSED_OP_DSV4_SPARSE_PACK, }; enum llm_ffn_op_type : int { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index e68dc49b6df..2eae0105aa8 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -776,6 +776,59 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( cb(k_all, "csa_k_all", il); ggml_tensor * raw_mask = inp_attn->get_kq_mask(); + + // The Lightning Indexer selects a different compressed working set for every + // token. Reinterpret the 64 attention heads as query rows of one MQA problem + // per token so Metal's tiled flash-attention kernel consumes only those keys. + // Dense FA remains faster for decode/small caches; auto probing forces this + // branch once during setup so unsupported layer devices can disable it. + const bool sparse_probe = cparams.auto_fdsv4_sparse; + const bool sparse_prefill = q->ne[2] >= 8 && n_csa > (int64_t) hparams.indexer_top_k; + if (cparams.fused_dsv4_sparse && cparams.flash_attn && + raw_k->type == GGML_TYPE_F16 && csa_k->type == GGML_TYPE_F16 && + (sparse_probe || sparse_prefill)) { + const int64_t n_stream = csa_k->ne[3]; + const int64_t nq = q->ne[2]/n_stream; + const int64_t nt = q->ne[2]; + const int64_t n_head = q->ne[1]; + const int64_t n_raw = std::min(hparams.n_swa, raw_k->ne[2]); + + GGML_ASSERT(q->ne[0] == raw_k->ne[0]); + GGML_ASSERT(raw_k->ne[1] == 1 && csa_k->ne[1] == 1); + GGML_ASSERT(raw_k->ne[3] == n_stream); + GGML_ASSERT(raw_mask->ne[1] == nq && raw_mask->ne[3] == n_stream); + GGML_ASSERT(inp_csa.kq_mask->ne[1] == nq && inp_csa.kq_mask->ne[3] == n_stream); + + ggml_tensor * packed = ggml_dsv4_sparse_pack(ctx0, raw_k, csa_k, raw_mask, + inp_csa.kq_mask, top_k, n_raw); + cb(packed, "csa_sparse_pack", il); + res->add_fused_node({LLM_FUSED_OP_DSV4_SPARSE_PACK, packed, il}); + + const int64_t nk = n_raw + top_k->ne[0]; + ggml_tensor * k_sel = ggml_view_4d(ctx0, packed, q->ne[0], nk, 1, nt, + q->ne[0]*sizeof(ggml_fp16_t), q->ne[0]*nk*sizeof(ggml_fp16_t), packed->nb[1], 0); + ggml_tensor * mask_sel = ggml_view_4d(ctx0, packed, nk, 1, 1, nt, + nk*sizeof(ggml_fp16_t), nk*sizeof(ggml_fp16_t), packed->nb[1], + q->ne[0]*nk*sizeof(ggml_fp16_t)); + cb(k_sel, "csa_sparse_k", il); + cb(mask_sel, "csa_sparse_mask", il); + + ggml_tensor * q_fa = ggml_reshape_4d(ctx0, q, q->ne[0], n_head, 1, nt); + ggml_tensor * out = ggml_flash_attn_ext(ctx0, q_fa, k_sel, k_sel, mask_sel, kq_scale, + hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + ggml_flash_attn_ext_add_sinks_rows(out, sinks); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, out, il}); + out = ggml_reshape_2d(ctx0, out, q->ne[0]*n_head, nt); + ggml_build_forward_expand(gf, out); + + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_csa_lid_sparse", il); + return out; + } + ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il); ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8cb59893586..a5831a9f45f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6824,9 +6824,10 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_K; const ggml_type type_V; std::array permute; + const bool sinks_rows; std::string vars() override { - return VARS_TO_STR14(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute); + return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, sinks_rows); } double max_nmse_err() override { @@ -6842,9 +6843,9 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, - ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}) + ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, bool sinks_rows = false) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute) {} + type_K(type_K), type_V(type_V), permute(permute), sinks_rows(sinks_rows) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -6892,18 +6893,22 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * m = nullptr; if (mask) { - m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nr23[1]); + m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, sinks_rows ? 1 : nb, 1, nr23[1]); ggml_set_name(m, "m"); } ggml_tensor * s = nullptr; if (sinks) { - s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, q->ne[2]); + s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, sinks_rows ? q->ne[1] : q->ne[2]); ggml_set_name(s, "s"); } ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap); - ggml_flash_attn_ext_add_sinks(out, s); + if (sinks_rows) { + ggml_flash_attn_ext_add_sinks_rows(out, s); + } else { + ggml_flash_attn_ext_add_sinks(out, s); + } ggml_flash_attn_ext_set_prec (out, prec); ggml_set_name(out, "out"); @@ -7355,6 +7360,62 @@ struct test_lightning_indexer : public test_case { } }; +// GGML_OP_DSV4_SPARSE_PACK +struct test_dsv4_sparse_pack : public test_case { + const int64_t nb; + const int64_t ns; + + std::string vars() override { return VARS_TO_STR2(nb, ns); } + + test_dsv4_sparse_pack(int64_t nb = 3, int64_t ns = 2) : nb(nb), ns(ns) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + constexpr int64_t d = 512; + constexpr int64_t nr = 11; + constexpr int64_t nc = 17; + constexpr int64_t kr = 7; + constexpr int64_t kc = 13; + + ggml_tensor * raw_k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, 1, nr, ns); + ggml_tensor * cmp_k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, 1, nc, ns); + ggml_tensor * raw_m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, nr, nb, 1, ns); + ggml_tensor * cmp_m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, nc, nb, 1, ns); + ggml_tensor * cmp_i = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, kc, nb, 1, ns); + + ggml_set_name(raw_k, "raw_k"); + ggml_set_name(cmp_k, "comp_k"); + ggml_set_name(raw_m, "raw_mask"); + ggml_set_name(cmp_m, "comp_mask"); + ggml_set_name(cmp_i, "comp_idx"); + + ggml_tensor * out = ggml_dsv4_sparse_pack(ctx, raw_k, cmp_k, raw_m, cmp_m, cmp_i, kr); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + const int32_t limit = 17; + std::vector data(ggml_nelements(t)); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = (int32_t) ((i*7 + 3) % limit); + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + } else { + init_tensor_uniform(t, -2.0f, 2.0f); + } + } + + ggml_tensor * raw_m = ggml_get_tensor(ctx, "raw_mask"); + std::vector mask(ggml_nelements(raw_m)); + for (size_t i = 0; i < mask.size(); ++i) { + mask[i] = ggml_fp32_to_fp16((i % 5) == 0 ? -INFINITY : (float) (i % 7)); + } + ggml_backend_tensor_set(raw_m, mask.data(), 0, mask.size()*sizeof(ggml_fp16_t)); + } +}; + // Deserializable generic test case struct input_tensor { ggml_type type; @@ -9370,6 +9431,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2048, 2, 1, 3}, k)); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2049, 2, 1, 3}, k)); } + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2049, 2, 1, 1}, 512)); // exhaustive top_k tests //for (int i = 1; i < 9999; ++i) { @@ -9594,6 +9656,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_Q4_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q2_0)); test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {1, 1}, 640, 64, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {1, 1}, 641, 8, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true)); // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). @@ -9737,6 +9801,8 @@ static std::vector> make_test_cases_eval() { } } + test_cases.emplace_back(new test_dsv4_sparse_pack()); + return test_cases; } #ifdef _MSC_VER From b5892a94a891478aa323aaed16063fd4557082c0 Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 18 Jul 2026 10:46:17 -0300 Subject: [PATCH 2/4] metal: optimize DeepSeek V4 sparse attention Implement the sparse DeepSeek V4 prefill path on Metal so Flash Attention consumes compact per-token key sets instead of scanning the full compressed cache. - Add a DSV4_SPARSE_PACK kernel for raw-window and selected keys. - Add an exact radix TOP_K specialization for 512 Lightning Indexer results. - Extend tiled Flash Attention with broadcast masks and row-indexed sinks. - Register backend support for the fused packing operation. Assisted-by: Codex --- ggml/src/ggml-metal/ggml-metal-device.cpp | 15 +- ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-device.m | 15 ++ ggml/src/ggml-metal/ggml-metal-impl.h | 24 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 89 ++++++++- ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/ggml-metal.metal | 231 ++++++++++++++++++++-- 7 files changed, 359 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index c153bd82177..8fe3a391f10 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -66,6 +66,7 @@ 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_SPARSE_PACK: op_str = "dsv4_sparse_pack"; break; default: GGML_ABORT("fatal error"); }; @@ -1330,6 +1331,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib) { + const char * name = "kernel_top_k_radix_f32_i32"; + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_TOP_K); @@ -1464,6 +1474,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( // do bounds checks for the mask? const bool bc_mask = op->src[3] && (op->src[3]->ne[1] % 8 != 0); + const bool scan_mask = has_mask && op->src[3]->ne[1] != 1; snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d", "flash_attn_ext", @@ -1471,7 +1482,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( dk, dv); - snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d", + snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_scanm=%d_ns10=%d_ns20=%d_nsg=%d", base, has_mask, has_sinks, @@ -1479,6 +1490,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( has_scap, has_kvpad, bc_mask, + scan_mask, ns10, ns20, nsg); @@ -1494,6 +1506,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT + 4); ggml_metal_cv_set_bool(cv, bc_mask, FC_FLASH_ATTN_EXT + 10); + ggml_metal_cv_set_bool(cv, scan_mask, FC_FLASH_ATTN_EXT + 11); ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT + 20); ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT + 21); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 7e1deeaa210..1a2daf47a91 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -144,6 +144,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix (ggml_metal_library_t lib); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse ); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 2dc6eb8fdbc..c8933dc408a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1375,6 +1375,21 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op->src[2]) && ggml_is_contiguous_rows(op->src[3]); + case GGML_OP_DSV4_SPARSE_PACK: + return op->src[0]->type == GGML_TYPE_F16 && + op->src[1]->type == GGML_TYPE_F16 && + op->src[2]->type == GGML_TYPE_F16 && + op->src[3]->type == GGML_TYPE_F16 && + op->src[4]->type == GGML_TYPE_I32 && + op->type == GGML_TYPE_F16 && + op->src[0]->ne[0] == 512 && + ggml_get_op_params_i32(op, 0) <= 128 && + ggml_get_op_params_i32(op, 0) + op->src[4]->ne[0] <= 128 + 512 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]) && + ggml_is_contiguous_rows(op->src[2]) && + ggml_is_contiguous_rows(op->src[3]) && + ggml_is_contiguous_rows(op->src[4]); case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index e173b91c0c5..561826387e3 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -401,6 +401,7 @@ typedef struct { float m0; float m1; int32_t n_head_log2; + int32_t sinks_rows; float logit_softcap; } ggml_metal_kargs_flash_attn_ext; @@ -436,6 +437,7 @@ typedef struct { float m0; float m1; int32_t n_head_log2; + int32_t sinks_rows; float logit_softcap; } ggml_metal_kargs_flash_attn_ext_vec; @@ -1238,6 +1240,28 @@ typedef struct { uint64_t nb_d2; } ggml_metal_kargs_dsv4_hc_post; +typedef struct { + int32_t n_embd; + int32_t n_batch; + int32_t n_raw; + int32_t n_raw_k; + int32_t n_comp; + uint64_t nb_rk2; + uint64_t nb_rk3; + uint64_t nb_ck2; + uint64_t nb_ck3; + uint64_t nb_rm0; + uint64_t nb_rm1; + uint64_t nb_rm3; + uint64_t nb_cm0; + uint64_t nb_cm1; + uint64_t nb_cm3; + uint64_t nb_ci0; + uint64_t nb_ci1; + uint64_t nb_ci3; + uint64_t nb_d1; +} ggml_metal_kargs_dsv4_sparse_pack; + typedef struct { int32_t ne00; int32_t ne01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index c5d7619c12f..007e4282ca2 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -330,6 +330,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_dsv4_hc(ctx, idx); } break; + case GGML_OP_DSV4_SPARSE_PACK: + { + n_fuse = ggml_metal_op_dsv4_sparse_pack(ctx, idx); + } break; case GGML_OP_SOFT_MAX: { n_fuse = ggml_metal_op_soft_max(ctx, idx); @@ -1508,6 +1512,56 @@ int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_dsv4_sparse_pack(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + GGML_ASSERT(op->op == GGML_OP_DSV4_SPARSE_PACK); + + const ggml_tensor * raw_k = op->src[0]; + const ggml_tensor * comp_k = op->src[1]; + const ggml_tensor * raw_mask = op->src[2]; + const ggml_tensor * comp_mask = op->src[3]; + const ggml_tensor * comp_idx = op->src[4]; + + ggml_metal_kargs_dsv4_sparse_pack args = { + /*.n_embd =*/ (int32_t) raw_k->ne[0], + /*.n_batch =*/ (int32_t) raw_mask->ne[1], + /*.n_raw =*/ ggml_get_op_params_i32(op, 0), + /*.n_raw_k =*/ (int32_t) raw_k->ne[2], + /*.n_comp =*/ (int32_t) comp_idx->ne[0], + /*.nb_rk2 =*/ raw_k->nb[2], + /*.nb_rk3 =*/ raw_k->nb[3], + /*.nb_ck2 =*/ comp_k->nb[2], + /*.nb_ck3 =*/ comp_k->nb[3], + /*.nb_rm0 =*/ raw_mask->nb[0], + /*.nb_rm1 =*/ raw_mask->nb[1], + /*.nb_rm3 =*/ raw_mask->nb[3], + /*.nb_cm0 =*/ comp_mask->nb[0], + /*.nb_cm1 =*/ comp_mask->nb[1], + /*.nb_cm3 =*/ comp_mask->nb[3], + /*.nb_ci0 =*/ comp_idx->nb[0], + /*.nb_ci1 =*/ comp_idx->nb[1], + /*.nb_ci3 =*/ comp_idx->nb[3], + /*.nb_d1 =*/ op->nb[1], + }; + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_base(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + for (int i = 0; i < 5; ++i) { + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[i]), i + 1); + } + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 6); + + // Four 256-thread packers can reside per core; a 512-thread group only + // allows two and exposes the random selected-row reads to more latency. + const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2939,8 +2993,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { GGML_ASSERT(ne12 == ne22); GGML_ASSERT(!op->src[3] || op->src[3]->type == GGML_TYPE_F16); - GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] && - "the Flash-Attention Metal kernel requires the mask to be at least n_queries big"); + GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] == 1 || op->src[3]->ne[1] >= op->src[0]->ne[1]); float scale; float max_bias; @@ -2958,6 +3011,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { const bool has_sinks = op->src[4] != NULL; const bool has_bias = max_bias != 0.0f; const bool has_scap = logit_softcap != 0.0f; + const bool sinks_rows = ggml_get_op_params_i32(op, 4); const uint32_t n_head = op->src[0]->ne[2]; const int32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head)); @@ -3035,7 +3089,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { need_sync = true; } - if (has_mask) { + if (has_mask && op->src[3]->ne[1] != 1) { assert(ggml_metal_op_flash_attn_ext_extra_blk(op) != 0); ggml_metal_kargs_flash_attn_ext_blk args0 = { @@ -3131,6 +3185,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.m0 =*/ m0, /*.m1 =*/ m1, /*.n_head_log2 =*/ n_head_log2, + /*.sinks_rows =*/ sinks_rows, /*.logit_softcap =*/ logit_softcap, }; @@ -3269,6 +3324,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.m0 =*/ m0, /*.m1 =*/ m1, /*.n_head_log2 =*/ n_head_log2, + /*.sinks_rows =*/ sinks_rows, /*.logit_softcap =*/ logit_softcap, }; @@ -4871,6 +4927,33 @@ int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + if (ne0 == 512 && ne00 > 1024) { + ggml_metal_kargs_argsort args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, + /*.ne03 =*/ ne03, + /*.nb00 =*/ nb00, + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb02, + /*.nb03 =*/ nb03, + /*.ne0 =*/ ne0, + /*.ne1 =*/ ne1, + /*.ne2 =*/ ne2, + /*.ne3 =*/ ne3, + /*.top_k =*/ ne0, + }; + + auto pipeline = ggml_metal_library_get_pipeline_top_k_radix(lib); + GGML_ASSERT(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline) >= 512); + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2); + ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, 512, 1, 1); + return 1; + } + auto pipeline = ggml_metal_library_get_pipeline_top_k(lib, op); // bitonic sort requires the number of elements to be power of 2 diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index b03b59e0bd9..0f1c5d893ba 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -54,6 +54,7 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx); int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_dsv4_sparse_pack (ggml_metal_op_t ctx, int idx); int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx); int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx); int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 7d12cb0fe39..c904429be52 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6032,6 +6032,119 @@ kernel void kernel_argsort_f32_i32( template [[host_name("kernel_argsort_f32_i32_asc")]] kernel argsort_t kernel_argsort_f32_i32; template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32; +// DSV4's Lightning Indexer selects 512 entries from a much longer score row. +// A full bitonic sort spends most of its time ordering entries that are thrown +// away. Find the exact 512th score with an MSD radix selection, collect that +// partition, then sort only the retained 512 indices. +kernel void kernel_top_k_radix_f32_i32( + constant ggml_metal_kargs_argsort & args, + device const char * src0, + device int32_t * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort radix = 16; + constexpr ushort n_top = 512; + + threadgroup atomic_uint histogram[radix]; + threadgroup atomic_uint n_selected; + threadgroup uint prefix; + threadgroup uint rank; + threadgroup int32_t selected[n_top]; + + const int i01 = tgpig.x; + const int i02 = tgpig.y; + const int i03 = tgpig.z; + device const float * row = (device const float *) (src0 + + args.nb01*i01 + args.nb02*i02 + args.nb03*i03); + + if (tiitg == 0) { + prefix = 0; + rank = n_top - 1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Flip IEEE-754 keys into monotonically increasing unsigned integers. + for (int shift = 28; shift >= 0; shift -= 4) { + if (tiitg < radix) { + atomic_store_explicit(&histogram[tiitg], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint mask = shift == 28 ? 0u : 0xffffffffu << (shift + 4); + for (int i = tiitg; i < args.ne00; i += ntg.x) { + const uint bits = as_type(row[i]); + const uint key = bits ^ (uint(int(bits) >> 31) | 0x80000000u); + if ((key & mask) == prefix) { + atomic_fetch_add_explicit(&histogram[(key >> shift) & 0xfu], 1u, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiitg == 0) { + uint r = rank; + for (int digit = radix - 1; digit >= 0; --digit) { + const uint count = atomic_load_explicit(&histogram[digit], memory_order_relaxed); + if (r < count) { + prefix |= uint(digit) << shift; + rank = r; + break; + } + r -= count; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tiitg == 0) { + atomic_store_explicit(&n_selected, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int i = tiitg; i < args.ne00; i += ntg.x) { + const uint bits = as_type(row[i]); + const uint key = bits ^ (uint(int(bits) >> 31) | 0x80000000u); + if (key > prefix) { + const uint pos = atomic_fetch_add_explicit(&n_selected, 1u, memory_order_relaxed); + if (pos < n_top) { + selected[pos] = i; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int i = tiitg; i < args.ne00; i += ntg.x) { + const uint bits = as_type(row[i]); + const uint key = bits ^ (uint(int(bits) >> 31) | 0x80000000u); + if (key == prefix) { + const uint pos = atomic_fetch_add_explicit(&n_selected, 1u, memory_order_relaxed); + if (pos < n_top) { + selected[pos] = i; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Restore TOP_K's descending-order contract for the retained partition. + const int col = tiitg; + for (int k = 2; k <= n_top; k *= 2) { + for (int j = k/2; j > 0; j /= 2) { + const int ixj = col ^ j; + if (ixj > col) { + const float lhs = row[selected[col]]; + const float rhs = row[selected[ixj]]; + if (((col & k) == 0 && lhs < rhs) || ((col & k) != 0 && lhs > rhs)) { + SWAP(selected[col], selected[ixj]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + dst += args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03; + dst[col] = selected[col]; +} + typedef void (argsort_merge_t)( constant ggml_metal_kargs_argsort_merge & args, device const char * src0, @@ -6346,20 +6459,19 @@ kernel void kernel_flash_attn_ext_blk( char res = i0*C + C > args.ne30 ? 1 : 0; - device const half * mask_src = (device const half *) (mask + (i1*Q)*args.nb31 + i2*args.nb32 + i3*args.nb33) + i0*C + tiisg; - // detailed check of the elements of the block if ((C > NW || Q > 1) && res == 0) { half mmin = MAXHALF; half mmax = -MAXHALF; - FOR_UNROLL (short j = 0; j < Q; ++j) { + const short nq_mask = args.ne31 == 1 ? 1 : Q; + FOR_UNROLL (short j = 0; j < nq_mask; ++j) { + device const half * mask_src = (device const half *) (mask + + ((i1*Q + j)%args.ne31)*args.nb31 + i2*args.nb32 + i3*args.nb33) + i0*C + tiisg; FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) { mmin = min(mmin, mask_src[ii*NW]); mmax = max(mmax, mask_src[ii*NW]); } - - mask_src += args.nb31/2; } mmin = simd_min(mmin); @@ -6389,6 +6501,7 @@ constant bool FC_flash_attn_ext_has_scap [[function_constant(FC_FLASH_ATTN_EXT constant bool FC_flash_attn_ext_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT + 4)]]; constant bool FC_flash_attn_ext_bc_mask [[function_constant(FC_FLASH_ATTN_EXT + 10)]]; +constant bool FC_flash_attn_ext_scan_mask [[function_constant(FC_FLASH_ATTN_EXT + 11)]]; //constant float FC_flash_attn_ext_scale [[function_constant(FC_FLASH_ATTN_EXT + 10)]]; //constant float FC_flash_attn_ext_max_bias [[function_constant(FC_FLASH_ATTN_EXT + 11)]]; @@ -6499,7 +6612,7 @@ void kernel_flash_attn_ext_impl( FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { const short j = jj*NSG + sgitg; - pm2[jj] = (device const half2 *) ((device const char *) mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + pm2[jj] = (device const half2 *) ((device const char *) mask + ((iq1 + j)%args.ne31)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); } { @@ -6603,7 +6716,7 @@ void kernel_flash_attn_ext_impl( const short j = jj*NSG + sgitg; pm2[jj] = (device const half2 *) ((device const half *) mask + - (iq1 + j)*C + + ((iq1 + j)%args.ne31)*C + (iq2%args.ne32)*(C*args.ne31) + (iq3%args.ne33)*(C*args.ne31*args.ne32)); } @@ -6616,7 +6729,7 @@ void kernel_flash_attn_ext_impl( // read the mask into shared mem if (FC_flash_attn_ext_has_mask) { - blk_cur = blk[ic0]; + blk_cur = FC_flash_attn_ext_scan_mask ? blk[ic0] : 1; if (blk_cur == 0) { FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { @@ -6631,7 +6744,7 @@ void kernel_flash_attn_ext_impl( const short j = jj*NSG + sgitg; if (FC_flash_attn_ext_bc_mask) { - sm2[j*SH + tiisg] = (iq1 + j) < args.ne31 ? pm2[jj][tiisg] : half2(-MAXHALF, -MAXHALF); + sm2[j*SH + tiisg] = args.ne31 == 1 || (iq1 + j) < args.ne31 ? pm2[jj][tiisg] : half2(-MAXHALF, -MAXHALF); } else { sm2[j*SH + tiisg] = pm2[jj][tiisg]; } @@ -6994,7 +7107,8 @@ void kernel_flash_attn_ext_impl( const short j = jj*NSG + sgitg; const float m = M[jj]; - const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2; + const int sink_idx = args.sinks_rows ? iq1 + j : iq2; + const float s = tiisg == 0 ? ((device const float *) sinks)[sink_idx] : -FLT_MAX/2; M[jj] = simd_max(max(M[jj], s)); @@ -7393,7 +7507,7 @@ kernel void kernel_flash_attn_ext_vec( const short ty = tiisg/NL; // pointer to the mask - device const half * pm = (device const half *) (mask + iq1*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + device const half * pm = (device const half *) (mask + (iq1%args.ne31)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); float slope = 1.0f; @@ -7433,7 +7547,7 @@ kernel void kernel_flash_attn_ext_vec( } } else { pm = (device const half *) (mask) + - iq1*C + + (iq1%args.ne31)*C + (iq2%args.ne32)*(C*args.ne31) + (iq3%args.ne33)*(C*args.ne31*args.ne32); } @@ -7638,7 +7752,8 @@ kernel void kernel_flash_attn_ext_vec( if (FC_flash_attn_ext_vec_has_sinks && sgitg == 0 && iwg == 0) { const float m = M; - const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2; + const int sink_idx = args.sinks_rows ? iq1 : iq2; + const float s = tiisg == 0 ? ((device const float *) sinks)[sink_idx] : -FLT_MAX/2; M = simd_max(max(M, s)); @@ -11601,3 +11716,93 @@ kernel void kernel_dsv4_hc_post_f32( *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst]; } } + +// One threadgroup owns one token. The Lightning-Indexer selection is shared by +// all 64 attention heads, so every cache row is fetched once and packed next to +// a head-broadcast mask. This layout lets the regular tiled FA kernel treat the +// heads as query rows without materializing any dense top-k mask. +kernel void kernel_dsv4_sparse_pack( + constant ggml_metal_kargs_dsv4_sparse_pack & args, + device const char * raw_k, + device const char * comp_k, + device const char * raw_mask, + device const char * comp_mask, + device const char * comp_idx, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort ntg[[threads_per_threadgroup]]) { + constexpr int max_selected = 128 + 512; + threadgroup int selected_idx[max_selected]; + threadgroup half selected_mask[max_selected]; + + const int iq = tgpig % args.n_batch; + const int is = tgpig / args.n_batch; + const int nk = args.n_raw + args.n_comp; + + device half * out = (device half *) (dst + (uint64_t) tgpig*args.nb_d1); + device half * out_k = out; + device half * out_m = out + args.n_embd*nk; + + // Selection indices and masks are shared by every embedding lane/head. Load + // them once per token instead of issuing up to 512 identical device reads. + // The raw SWA mask has at most n_raw finite entries. Avoid a separate F32 + // cast + top-k dispatch by compacting those entries directly here. + if (tiitg == 0) { + int n = 0; + for (int idx = 0; idx < args.n_raw_k && n < args.n_raw; ++idx) { + const half m = *(device const half *) (raw_mask + + (uint64_t) idx*args.nb_rm0 + (uint64_t) iq*args.nb_rm1 + (uint64_t) is*args.nb_rm3); + if (isfinite(m)) { + selected_idx[n] = idx; + selected_mask[n] = m; + ++n; + } + } + for (; n < args.n_raw; ++n) { + selected_idx[n] = -1; + selected_mask[n] = -INFINITY; + } + } + for (int i = tiitg; i < args.n_comp; i += ntg) { + const int oi = args.n_raw + i; + const int idx = *(device const int *) (comp_idx + + (uint64_t) i*args.nb_ci0 + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); + selected_idx[oi] = idx; + selected_mask[oi] = *(device const half *) (comp_mask + + (uint64_t) idx*args.nb_cm0 + (uint64_t) iq*args.nb_cm1 + (uint64_t) is*args.nb_cm3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int i = 0; i < args.n_raw; ++i) { + const int idx = selected_idx[i]; + if (idx >= 0) { + device const half * src = (device const half *) (raw_k + + (uint64_t) idx*args.nb_rk2 + (uint64_t) is*args.nb_rk3); + for (int e = tiitg; e < args.n_embd; e += ntg) { + out_k[i*args.n_embd + e] = src[e]; + } + } else { + for (int e = tiitg; e < args.n_embd; e += ntg) { + out_k[i*args.n_embd + e] = 0.0h; + } + } + if (tiitg == 0) { + out_m[i] = selected_mask[i]; + } + } + + for (int i = 0; i < args.n_comp; ++i) { + const int oi = args.n_raw + i; + const int idx = selected_idx[oi]; + device const half * src = (device const half *) (comp_k + + (uint64_t) idx*args.nb_ck2 + (uint64_t) is*args.nb_ck3); + for (int e = tiitg; e < args.n_embd; e += ntg) { + out_k[oi*args.n_embd + e] = src[e]; + } + if (tiitg == 0) { + out_m[oi] = selected_mask[oi]; + } + } + +} From 2d0857f3e45a92c34fddab3082e8a197388ef2dc Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 18 Jul 2026 10:46:39 -0300 Subject: [PATCH 3/4] common: optimize DeepSeek V4 token generation Reduce DeepSeek V4 decode overhead with fused compressor and mask operations, avoid redundant Lightning Indexer work, and bound long-context attention by gathering only selected compressed keys. - Add DSV4_COMPRESS and DSV4_TOP_K_MASK with CPU references and capability probes. - Skip Lightning Indexer and TOP_K when every compressed row is selected. - Reuse sparse packing for deep single-token decode with dense fallbacks. - Test fused compression, mask construction, and compact Flash Attention. Assisted-by: Codex --- ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 25 +++++ ggml/src/ggml-backend-meta.cpp | 2 + ggml/src/ggml-cpu/ggml-cpu.c | 10 ++ ggml/src/ggml-cpu/ops.cpp | 126 ++++++++++++++++++++++ ggml/src/ggml-cpu/ops.h | 2 + ggml/src/ggml.c | 80 +++++++++++++- src/llama-context.cpp | 23 ++++ src/llama-cparams.h | 3 + src/llama-graph.h | 2 + src/models/deepseek4.cpp | 188 +++++++++++++++++++++++---------- tests/test-backend-ops.cpp | 131 +++++++++++++++++++++-- 12 files changed, 527 insertions(+), 69 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index a88eab44011..57aab90d592 100644 --- a/ggml/include/ggml-rpc.h +++ b/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 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 83de623960d..f371002d090 100644 --- a/ggml/include/ggml.h +++ b/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( diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index c8a345cf846..ce0521b6b7b 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/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: diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index b471a24fbdd..3e284329a88 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/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: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index df042d23ac1..800dd1487c8 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/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) { diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index cb1629e0bfc..bcc8815c9e5 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/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); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index ae904c6a1ad..be7bad8b0a5 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1080,6 +1080,8 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "SOLVE_TRI", "GATED_DELTA_NET", "LIGHTNING_INDEXER", + "DSV4_COMPRESS", + "DSV4_TOP_K_MASK", "DSV4_HC_COMB", "DSV4_HC_PRE", "DSV4_HC_POST", @@ -1101,7 +1103,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1196,6 +1198,8 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "A X = B, A triangular, solve X", "gated_delta_net(q, k, v, g, beta, s)", "lightning_indexer(q, k, weights, mask)", + "dsv4_compress(kv_state, score_state, read_idxs)", + "dsv4_top_k_mask(raw_mask, comp_mask, comp_idx)", "dsv4_hc_comb(mixes, scale, base)", "dsv4_hc_pre(x, weights)", "dsv4_hc_post(x, residual, post, comb)", @@ -1217,7 +1221,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6358,6 +6362,75 @@ struct ggml_tensor * ggml_lightning_indexer( return result; } +// ggml_dsv4_compress + +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) { + 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(ratio > 0); + GGML_ASSERT(kv_state->ne[0] == score_state->ne[0]); + GGML_ASSERT(kv_state->ne[1] == score_state->ne[1]); + GGML_ASSERT(kv_state->ne[2] == 1 && kv_state->ne[3] == 1); + GGML_ASSERT(score_state->ne[2] == 1 && score_state->ne[3] == 1); + GGML_ASSERT(read_idxs->ne[1] == 1 && read_idxs->ne[2] == 1 && read_idxs->ne[3] == 1); + + const int64_t n_read_per_block = (overlap ? 2 : 1)*ratio; + GGML_ASSERT(read_idxs->ne[0] % n_read_per_block == 0); + + const int64_t n_blocks = read_idxs->ne[0]/n_read_per_block; + const int64_t n_embd = overlap ? kv_state->ne[0]/2 : kv_state->ne[0]; + + GGML_ASSERT(n_blocks > 0 && n_embd > 0); + GGML_ASSERT(kv_state->ne[0] == (overlap ? 2 : 1)*n_embd); + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_blocks); + + ggml_set_op_params_i32(result, 0, ratio); + ggml_set_op_params_i32(result, 1, overlap ? 1 : 0); + + result->op = GGML_OP_DSV4_COMPRESS; + result->src[0] = kv_state; + result->src[1] = score_state; + result->src[2] = read_idxs; + + return result; +} + +// ggml_dsv4_top_k_mask + +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) { + 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(raw_mask->ne[2] == 1 && comp_mask->ne[2] == 1 && comp_idx->ne[2] == 1); + GGML_ASSERT(raw_mask->ne[1] == comp_mask->ne[1]); + GGML_ASSERT(raw_mask->ne[1] == comp_idx->ne[1]); + GGML_ASSERT(raw_mask->ne[3] == comp_mask->ne[3]); + GGML_ASSERT(raw_mask->ne[3] == comp_idx->ne[3]); + GGML_ASSERT(comp_idx->ne[0] <= comp_mask->ne[0]); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, + raw_mask->ne[0] + comp_mask->ne[0], raw_mask->ne[1], 1, raw_mask->ne[3]); + + result->op = GGML_OP_DSV4_TOP_K_MASK; + result->src[0] = raw_mask; + result->src[1] = comp_mask; + result->src[2] = comp_idx; + + return result; +} + // ggml_dsv4_sparse_pack struct ggml_tensor * ggml_dsv4_sparse_pack( @@ -6379,7 +6452,8 @@ struct ggml_tensor * ggml_dsv4_sparse_pack( const int64_t nq = raw_mask->ne[1]; const int64_t nk = n_raw + comp_idx->ne[0]; - GGML_ASSERT(n_raw > 0 && n_raw <= raw_k->ne[2]); + GGML_ASSERT(n_raw >= 0 && n_raw <= raw_k->ne[2]); + GGML_ASSERT(nk > 0); GGML_ASSERT(comp_k->ne[0] == d); GGML_ASSERT(raw_k->ne[1] == 1 && comp_k->ne[1] == 1); GGML_ASSERT(comp_k->ne[3] == n_stream); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c53e66c35c6..6e87f7d45ac 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -61,6 +61,18 @@ static const llm_fused_op_probe llm_fused_op_lid_probe = { /*.n_tokens_per_seq =*/ 1, }; +static const llm_fused_op_probe llm_fused_op_dsv4_compress_probe = { + /*.op =*/ LLM_FUSED_OP_DSV4_COMPRESS, + /*.name =*/ "fused DeepSeek V4 compressor", + /*.n_tokens_per_seq =*/ 1, +}; + +static const llm_fused_op_probe llm_fused_op_dsv4_top_k_mask_probe = { + /*.op =*/ LLM_FUSED_OP_DSV4_TOP_K_MASK, + /*.name =*/ "fused DeepSeek V4 top-k mask", + /*.n_tokens_per_seq =*/ 1, +}; + static const llm_fused_op_probe llm_fused_op_dsv4_hc_pre_probe = { /*.op =*/ LLM_FUSED_OP_DSV4_HC_PRE, /*.name =*/ "fused DeepSeek V4 HC pre", @@ -260,6 +272,10 @@ llama_context::llama_context( cparams.fused_lid = true; cparams.auto_flid = true; + cparams.fused_dsv4_compress = true; + cparams.fused_dsv4_top_k_mask = true; + cparams.auto_fdsv4_aux = true; + cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; cparams.fused_dsv4_hc_post = true; @@ -587,6 +603,13 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 resolve(llm_fused_op_dsv4_sparse_probe, cparams.fused_dsv4_sparse); cparams.auto_fdsv4_sparse = false; } + + if (cparams.auto_fdsv4_aux) { + LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 auxiliary ops support:\n", func); + resolve(llm_fused_op_dsv4_compress_probe, cparams.fused_dsv4_compress); + resolve(llm_fused_op_dsv4_top_k_mask_probe, cparams.fused_dsv4_top_k_mask); + cparams.auto_fdsv4_aux = false; + } } void llama_context::sched_reserve() { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 933ba2c8612..122d303d207 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -43,6 +43,9 @@ struct llama_cparams { bool auto_fgdn; bool fused_lid; // use fused lightning indexer bool auto_flid; + bool fused_dsv4_compress; + bool fused_dsv4_top_k_mask; + bool auto_fdsv4_aux; bool fused_dsv4_hc_pre; bool fused_dsv4_hc_comb; bool fused_dsv4_hc_post; diff --git a/src/llama-graph.h b/src/llama-graph.h index e15d822a256..b68087db3f3 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -43,6 +43,8 @@ enum llm_fused_op { LLM_FUSED_OP_GDN_AR, LLM_FUSED_OP_GDN_CH, LLM_FUSED_OP_LIGHTNING_INDEXER, + LLM_FUSED_OP_DSV4_COMPRESS, + LLM_FUSED_OP_DSV4_TOP_K_MASK, LLM_FUSED_OP_DSV4_HC_PRE, LLM_FUSED_OP_DSV4_HC_COMB, LLM_FUSED_OP_DSV4_HC_POST, diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 2eae0105aa8..86c1b393080 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -479,21 +479,28 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state( GGML_ASSERT(state_read_idxs->ne[0] == DSV4_HCA_RATIO*n_blocks); GGML_ASSERT(n_embd_head >= n_embd_head_rope); - ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs); - kv = ggml_reshape_3d(ctx0, kv, n_embd_head, DSV4_HCA_RATIO, n_blocks); - cb(kv, name, il); + ggml_tensor * comp = nullptr; + if (cparams.fused_dsv4_compress) { + comp = ggml_dsv4_compress( + ctx0, kv_state, score_state, state_read_idxs, DSV4_HCA_RATIO, false); + res->add_fused_node({LLM_FUSED_OP_DSV4_COMPRESS, comp, il}); + } else { + ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, DSV4_HCA_RATIO, n_blocks); + cb(kv, name, il); - ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs); - score = ggml_reshape_3d(ctx0, score, n_embd_head, DSV4_HCA_RATIO, n_blocks); - cb(score, name, il); + ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs); + score = ggml_reshape_3d(ctx0, score, n_embd_head, DSV4_HCA_RATIO, n_blocks); + cb(score, name, il); - ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3)); - ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3)); + ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); - ggml_tensor * weights = ggml_soft_max(ctx0, scores); - ggml_tensor * comp = ggml_mul(ctx0, values, weights); - comp = ggml_sum_rows(ctx0, comp); - comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + ggml_tensor * weights = ggml_soft_max(ctx0, scores); + comp = ggml_mul(ctx0, values, weights); + comp = ggml_sum_rows(ctx0, comp); + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + } cb(comp, name, il); comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); @@ -540,44 +547,49 @@ ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_sta GGML_ASSERT(score_state->ne[0] == 2*n_embd_head); GGML_ASSERT(n_embd_head >= n_embd_head_rope); - kv_state = dsv4_append_zero_row(ctx0, kv_state, false); - score_state = dsv4_append_zero_row(ctx0, score_state, true); - - const int64_t n_read = ratio*n_blocks; - - ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs); - ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs); - - ggml_tensor * kv_prev = ggml_cont(ctx0, - ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0)); - kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks); - cb(kv_prev, name, il); - - ggml_tensor * score_prev = ggml_cont(ctx0, - ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0)); - score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks); - cb(score_prev, name, il); - - ggml_tensor * kv_cur = ggml_cont(ctx0, - ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], - n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head))); - kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks); - - ggml_tensor * score_cur = ggml_cont(ctx0, - ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], - n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head))); - score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks); - - ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1); - ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1); - - values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3)); - scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3)); - - ggml_tensor * weights = ggml_soft_max(ctx0, scores); - ggml_tensor * comp = ggml_mul(ctx0, values, weights); - comp = ggml_sum_rows(ctx0, comp); - comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + ggml_tensor * comp = nullptr; + if (cparams.fused_dsv4_compress) { + comp = ggml_dsv4_compress( + ctx0, kv_state, score_state, state_read_idxs, (int32_t) ratio, true); + res->add_fused_node({LLM_FUSED_OP_DSV4_COMPRESS, comp, il}); + } else { + kv_state = dsv4_append_zero_row(ctx0, kv_state, false); + score_state = dsv4_append_zero_row(ctx0, score_state, true); + + const int64_t n_read = ratio*n_blocks; + ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs); + ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs); + + ggml_tensor * kv_prev = ggml_cont(ctx0, + ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0)); + kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks); + cb(kv_prev, name, il); + + ggml_tensor * score_prev = ggml_cont(ctx0, + ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0)); + score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks); + cb(score_prev, name, il); + + ggml_tensor * kv_cur = ggml_cont(ctx0, + ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], + n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head))); + kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks); + + ggml_tensor * score_cur = ggml_cont(ctx0, + ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], + n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head))); + score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks); + + ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1); + ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1); + values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3)); + scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3)); + + ggml_tensor * weights = ggml_soft_max(ctx0, scores); + comp = ggml_mul(ctx0, values, weights); + comp = ggml_sum_rows(ctx0, comp); + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + } cb(comp, name, il); comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); @@ -744,8 +756,6 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( const auto & inp_csa = inp_dsv4->get_csa(); GGML_ASSERT(inp_csa.kq_mask); - ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); - ggml_tensor * k_rot = inp_attn->self_k_rot; if (k_rot) { q = llama_mul_mat_hadamard(ctx0, q, k_rot); @@ -777,16 +787,66 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * raw_mask = inp_attn->get_kq_mask(); + // Selecting every compressed row is equivalent to using the visibility + // mask directly. Avoid building the Lightning Indexer and TOP_K graph until + // the compressed cache grows beyond the selection size. Auto probes are + // the only exception because they must materialize their fused ops. + ggml_tensor * top_k = nullptr; + if (cparams.auto_fdsv4_aux || cparams.auto_fdsv4_sparse || + n_csa > (int64_t) hparams.indexer_top_k) { + top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); + } + + // At sufficiently deep single-token decode, gather only the Lightning + // Indexer's selected compressed keys. This preserves the regular head layout + // used by Metal's vector Flash Attention while making its KV work fixed-size. + const bool gather_decode = cparams.fused_dsv4_sparse && cparams.flash_attn && + raw_k->type == GGML_TYPE_F16 && csa_k->type == GGML_TYPE_F16 && + q->ne[2] == 1 && csa_k->ne[3] == 1 && top_k && + top_k->ne[1] == 1 && top_k->ne[3] == 1 && + n_csa >= 2*(int64_t) hparams.indexer_top_k; + if (gather_decode) { + ggml_tensor * packed = ggml_dsv4_sparse_pack( + ctx0, raw_k, csa_k, raw_mask, inp_csa.kq_mask, top_k, 0); + cb(packed, "csa_gathered_pack", il); + res->add_fused_node({LLM_FUSED_OP_DSV4_SPARSE_PACK, packed, il}); + + const int64_t nk = top_k->ne[0]; + ggml_tensor * gathered = ggml_view_4d(ctx0, packed, csa_k->ne[0], 1, nk, 1, + csa_k->ne[0]*sizeof(ggml_fp16_t), csa_k->ne[0]*sizeof(ggml_fp16_t), packed->nb[1], 0); + cb(gathered, "csa_gathered_k", il); + + ggml_tensor * k_sel = ggml_concat(ctx0, raw_k, gathered, 2); + cb(k_sel, "csa_k_selected", il); + + ggml_tensor * comp_mask = ggml_view_4d(ctx0, packed, nk, 1, 1, 1, + nk*sizeof(ggml_fp16_t), nk*sizeof(ggml_fp16_t), packed->nb[1], + csa_k->ne[0]*nk*sizeof(ggml_fp16_t)); + cb(comp_mask, "csa_gathered_mask", il); + + ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, comp_mask, 0); + cb(kq_mask, "csa_lid_kq_mask", il); + + ggml_tensor * out = build_attn_mha( + q, k_sel, k_sel, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_csa_lid_gathered", il); + return out; + } + // The Lightning Indexer selects a different compressed working set for every - // token. Reinterpret the 64 attention heads as query rows of one MQA problem - // per token so Metal's tiled flash-attention kernel consumes only those keys. - // Dense FA remains faster for decode/small caches; auto probing forces this - // branch once during setup so unsupported layer devices can disable it. + // prefill token. Reinterpret the 64 attention heads as query rows of one MQA + // problem per token so Metal's tiled flash-attention kernel consumes only + // those keys. Auto probing forces this branch once during setup so unsupported + // layer devices can disable it. const bool sparse_probe = cparams.auto_fdsv4_sparse; const bool sparse_prefill = q->ne[2] >= 8 && n_csa > (int64_t) hparams.indexer_top_k; if (cparams.fused_dsv4_sparse && cparams.flash_attn && raw_k->type == GGML_TYPE_F16 && csa_k->type == GGML_TYPE_F16 && (sparse_probe || sparse_prefill)) { + GGML_ASSERT(top_k); const int64_t n_stream = csa_k->ne[3]; const int64_t nq = q->ne[2]/n_stream; const int64_t nt = q->ne[2]; @@ -829,9 +889,21 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( return out; } - ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il); + ggml_tensor * csa_mask = inp_csa.kq_mask; + ggml_tensor * kq_mask = nullptr; + if (top_k && cparams.fused_dsv4_top_k_mask) { + kq_mask = ggml_dsv4_top_k_mask(ctx0, raw_mask, csa_mask, top_k); + cb(kq_mask, "csa_top_k_mask", il); + res->add_fused_node({LLM_FUSED_OP_DSV4_TOP_K_MASK, kq_mask, il}); + } else { + if (top_k) { + csa_mask = build_top_k_mask(csa_mask, top_k, "csa_top_k_mask", il); + } else { + cb(csa_mask, "csa_top_k_mask", il); + } + kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); + } - ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a5831a9f45f..61d33fdd6b1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3760,6 +3760,99 @@ struct test_snake_fuse : public test_case { }; +struct test_dsv4_compress : public test_case { + const int64_t n_embd; + const int64_t n_rows; + const int64_t n_blocks; + const int32_t ratio; + const bool overlap; + + std::string vars() override { + return VARS_TO_STR5(n_embd, n_rows, n_blocks, ratio, overlap); + } + + double max_nmse_err() override { return 1e-6; } + + test_dsv4_compress(int64_t n_embd, int64_t n_rows, int64_t n_blocks, int32_t ratio, bool overlap) + : n_embd(n_embd), n_rows(n_rows), n_blocks(n_blocks), ratio(ratio), overlap(overlap) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t n_state = (overlap ? 2 : 1)*n_embd; + const int64_t n_read = (overlap ? 2 : 1)*ratio*n_blocks; + + ggml_tensor * kv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_state, n_rows); + ggml_tensor * score = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_state, n_rows); + ggml_tensor * idx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_read); + ggml_set_name(kv, "kv"); + ggml_set_name(score, "score"); + ggml_set_name(idx, "idx"); + + ggml_tensor * out = ggml_dsv4_compress(ctx, kv, score, idx, ratio, overlap); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + std::vector data(ggml_nelements(t)); + for (size_t i = 0; i < data.size(); ++i) { + // Include the synthetic zero/-inf row sentinel used by the + // production read plans without materializing that row. + data[i] = (i % 11 == 0) ? (int32_t) n_rows : (int32_t) ((7*i + 3) % n_rows); + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + } else { + init_tensor_uniform(t, -2.0f, 2.0f); + } + } + } +}; + +struct test_dsv4_top_k_mask : public test_case { + const int64_t n_raw; + const int64_t n_comp; + const int64_t n_select; + const int64_t n_query; + const int64_t n_stream; + + std::string vars() override { + return VARS_TO_STR5(n_raw, n_comp, n_select, n_query, n_stream); + } + + test_dsv4_top_k_mask(int64_t n_raw, int64_t n_comp, int64_t n_select, int64_t n_query, int64_t n_stream) + : n_raw(n_raw), n_comp(n_comp), n_select(n_select), n_query(n_query), n_stream(n_stream) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * raw = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_raw, n_query, 1, n_stream); + ggml_tensor * comp = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_comp, n_query, 1, n_stream); + ggml_tensor * idx = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_select, n_query, 1, n_stream); + ggml_set_name(raw, "raw"); + ggml_set_name(comp, "comp"); + ggml_set_name(idx, "idx"); + + ggml_tensor * out = ggml_dsv4_top_k_mask(ctx, raw, comp, idx); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + std::vector data(ggml_nelements(t)); + for (int64_t row = 0; row < n_query*n_stream; ++row) { + for (int64_t i = 0; i < n_select; ++i) { + data[row*n_select + i] = (int32_t) ((7*i + 11*row + 3) % n_comp); + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); + } else { + init_tensor_uniform(t, -2.0f, 2.0f); + } + } + } +}; + struct test_dsv4_hc : public test_case { static constexpr int64_t hc = 4; @@ -7364,17 +7457,17 @@ struct test_lightning_indexer : public test_case { struct test_dsv4_sparse_pack : public test_case { const int64_t nb; const int64_t ns; + const int64_t kr; - std::string vars() override { return VARS_TO_STR2(nb, ns); } + std::string vars() override { return VARS_TO_STR3(nb, ns, kr); } - test_dsv4_sparse_pack(int64_t nb = 3, int64_t ns = 2) : nb(nb), ns(ns) {} + test_dsv4_sparse_pack(int64_t nb = 3, int64_t ns = 2, int64_t kr = 7) : nb(nb), ns(ns), kr(kr) {} ggml_tensor * build_graph(ggml_context * ctx) override { constexpr int64_t d = 512; constexpr int64_t nr = 11; - constexpr int64_t nc = 17; - constexpr int64_t kr = 7; - constexpr int64_t kc = 13; + const int64_t nc = kr == 0 ? 517 : 17; + const int64_t kc = kr == 0 ? 512 : 13; ggml_tensor * raw_k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, 1, nr, ns); ggml_tensor * cmp_k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, 1, nc, ns); @@ -7396,7 +7489,7 @@ struct test_dsv4_sparse_pack : public test_case { void initialize_tensors(ggml_context * ctx) override { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { if (t->type == GGML_TYPE_I32) { - const int32_t limit = 17; + const int32_t limit = kr == 0 ? 517 : 17; std::vector data(ggml_nelements(t)); for (size_t i = 0; i < data.size(); ++i) { data[i] = (int32_t) ((i*7 + 3) % limit); @@ -8127,6 +8220,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_snake_fuse(type, { 64, 32, 2, 3})); // ne[2] > 1 and ne[3] > 1 } + test_cases.emplace_back(new test_dsv4_compress(512, 17, 3, 4, true)); + test_cases.emplace_back(new test_dsv4_compress(512, 137, 2, 128, false)); + test_cases.emplace_back(new test_dsv4_top_k_mask(17, 137, 64, 3, 2)); + test_cases.emplace_back(new test_dsv4_top_k_mask(128, 2500, 512, 1, 1)); + test_cases.emplace_back(new test_dsv4_hc_comb(1, 1)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); @@ -9658,6 +9756,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {1, 1}, 640, 64, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {1, 1}, 641, 8, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true)); + for (int kv : { 640, 2628, 5128, }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {64, 1}, kv, 1, true, true, + 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). @@ -9802,6 +9904,7 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_dsv4_sparse_pack()); + test_cases.emplace_back(new test_dsv4_sparse_pack(1, 1, 0)); return test_cases; } @@ -9813,6 +9916,15 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; + // DeepSeek V4 token-generation shapes. + test_cases.emplace_back(new test_dsv4_hc_comb(1, 4)); + test_cases.emplace_back(new test_dsv4_hc_pre(4096, 1)); + test_cases.emplace_back(new test_dsv4_hc_post(4096, 1)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2628, 1, 1, 1}, 512)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {5128, 1, 1, 1}, 512)); + test_cases.emplace_back(new test_dsv4_top_k_mask(128, 2500, 512, 1, 1)); + test_cases.emplace_back(new test_dsv4_top_k_mask(128, 5000, 512, 1, 1)); + // Conv2d: K=CRS=NPQ=4096 matmul performance uint32_t iwh_idx = 0; uint32_t kwh_idx = 1; @@ -10019,6 +10131,13 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // DeepSeek V4 compressed sparse attention decode: one 512-wide KV head is + // shared by 64 query heads and the visibility mask is broadcast per row. + for (int kv : { 2628, 5128, }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {64, 1}, kv, 1, true, true, + 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } + for (int kv : { 4096, 8192, 16384, }) { for (int hs : { 64, 128, }) { for (int nr : { 1, 4, }) { From f65a7006bf2fc26c88b13280de12a7a2825fa144 Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 18 Jul 2026 10:46:57 -0300 Subject: [PATCH 4/4] metal: optimize DeepSeek V4 token generation Fuse compressor and attention-mask construction on Metal and keep long-context decode bounded by gathering only selected compressed keys. - Add kernels for DSV4_COMPRESS and DSV4_TOP_K_MASK. - Skip K and V loads for fully masked vector Flash Attention groups. - Extend DSV4_SPARSE_PACK with a parallel compressed-only decode mode. - Preserve compact F16 masks without intermediate conversions. Assisted-by: Codex --- ggml/src/ggml-metal/ggml-metal-device.cpp | 13 +- ggml/src/ggml-metal/ggml-metal-device.m | 19 +++ ggml/src/ggml-metal/ggml-metal-impl.h | 30 ++++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 107 ++++++++++++- ggml/src/ggml-metal/ggml-metal-ops.h | 2 + ggml/src/ggml-metal/ggml-metal.metal | 186 ++++++++++++++++++++-- 6 files changed, 342 insertions(+), 15 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 8fe3a391f10..d813064d9c6 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/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); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c8933dc408a..d27b655f2ac 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/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 && diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 561826387e3..aac5418bbe9 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/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; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 007e4282ca2..dcdd99e493d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -324,6 +324,14 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_lightning_indexer(ctx, idx); } break; + case GGML_OP_DSV4_COMPRESS: + { + n_fuse = ggml_metal_op_dsv4_compress(ctx, idx); + } break; + case GGML_OP_DSV4_TOP_K_MASK: + { + n_fuse = ggml_metal_op_dsv4_top_k_mask(ctx, idx); + } break; case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: @@ -1381,6 +1389,97 @@ int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_dsv4_compress(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + GGML_ASSERT(op->op == GGML_OP_DSV4_COMPRESS); + + const ggml_tensor * kv_state = op->src[0]; + const ggml_tensor * score_state = op->src[1]; + const ggml_tensor * read_idxs = op->src[2]; + + const int32_t ratio = ggml_get_op_params_i32(op, 0); + const int32_t overlap = ggml_get_op_params_i32(op, 1); + const int32_t n_read = (overlap ? 2 : 1)*ratio; + + 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(op->type == GGML_TYPE_F32); + GGML_ASSERT(n_read <= 128); + + ggml_metal_kargs_dsv4_compress args = { + /*.n_embd =*/ (int32_t) op->ne[0], + /*.n_blocks =*/ (int32_t) op->ne[1], + /*.n_rows =*/ (int32_t) kv_state->ne[1], + /*.ratio =*/ ratio, + /*.overlap =*/ overlap, + /*.nb_k0 =*/ kv_state->nb[0], + /*.nb_k1 =*/ kv_state->nb[1], + /*.nb_s0 =*/ score_state->nb[0], + /*.nb_s1 =*/ score_state->nb[1], + /*.nb_i0 =*/ read_idxs->nb[0], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + }; + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_base(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(kv_state), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(score_state), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(read_idxs), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + ggml_metal_encoder_set_threadgroup_memory_size(enc, n_read*sizeof(int32_t), 0); + + const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + ggml_metal_encoder_dispatch_threadgroups( + enc, (args.n_embd + nth - 1)/nth, args.n_blocks, 1, nth, 1, 1); + + return 1; +} + +int ggml_metal_op_dsv4_top_k_mask(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + GGML_ASSERT(op->op == GGML_OP_DSV4_TOP_K_MASK); + + const ggml_tensor * raw_mask = op->src[0]; + const ggml_tensor * comp_mask = op->src[1]; + const ggml_tensor * comp_idx = op->src[2]; + + ggml_metal_kargs_dsv4_top_k_mask args = { + /*.n_raw =*/ (int32_t) raw_mask->ne[0], + /*.n_comp =*/ (int32_t) comp_mask->ne[0], + /*.n_select =*/ (int32_t) comp_idx->ne[0], + /*.n_query =*/ (int32_t) raw_mask->ne[1], + /*.nb_rm1 =*/ raw_mask->nb[1], + /*.nb_rm3 =*/ raw_mask->nb[3], + /*.nb_cm1 =*/ comp_mask->nb[1], + /*.nb_cm3 =*/ comp_mask->nb[3], + /*.nb_ci1 =*/ comp_idx->nb[1], + /*.nb_ci3 =*/ comp_idx->nb[3], + /*.nb_d1 =*/ op->nb[1], + /*.nb_d3 =*/ op->nb[3], + }; + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_base(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(raw_mask), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comp_mask), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comp_idx), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + + const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + ggml_metal_encoder_dispatch_threadgroups( + enc, args.n_query*raw_mask->ne[3], 1, 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -1557,7 +1656,13 @@ int ggml_metal_op_dsv4_sparse_pack(ggml_metal_op_t ctx, int idx) { // Four 256-thread packers can reside per core; a 512-thread group only // allows two and exposes the random selected-row reads to more latency. const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); - ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + if (args.n_raw == 0) { + // Decode gathers independent compressed rows. Giving each row its own + // threadgroup exposes enough parallelism to hide the random-read latency. + ggml_metal_encoder_dispatch_threadgroups(enc, args.n_comp, op->ne[1], 1, nth, 1, 1); + } else { + ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + } return 1; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 0f1c5d893ba..b5aeb63b3c2 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -54,6 +54,8 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx); int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_dsv4_compress (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_dsv4_top_k_mask (ggml_metal_op_t ctx, int idx); int ggml_metal_op_dsv4_sparse_pack (ggml_metal_op_t ctx, int idx); int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx); int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index c904429be52..e62ad0d8305 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -7378,6 +7378,7 @@ constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_ constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]]; constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]]; constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]]; +constant bool FC_flash_attn_ext_vec_sparse_mask [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]]; //constant float FC_flash_attn_ext_vec_scale [[function_constant(FC_FLASH_ATTN_EXT_VEC + 10)]]; //constant float FC_flash_attn_ext_vec_max_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 11)]]; @@ -7576,6 +7577,21 @@ kernel void kernel_flash_attn_ext_vec( // each simdgroup processes 1 query and NE (NW/NL) cache elements FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + bool active = true; + if (FC_flash_attn_ext_vec_sparse_mask) { + active = false; + FOR_UNROLL (short ii = 0; ii < NE; ++ii) { + active |= sm[cc*NE + ii] > -MAXHALF; + } + } + + // The mask decision is uniform across the simdgroup. Avoid + // loading a full K row when all NE logits in this group are + // masked (notably DSV4's 512-of-N compressed selection). + if (!active) { + continue; + } + if (is_same::value) { FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { mqk[cc] += dot((float4) pk4[cc*NE*NS10/4 + ii*NL], (float4) pq4[ii*NL]); @@ -7686,12 +7702,34 @@ kernel void kernel_flash_attn_ext_vec( const auto sst = ss + ty; FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + bool active = true; + if (FC_flash_attn_ext_vec_sparse_mask) { + active = false; + FOR_UNROLL (short jj = 0; jj < NE; ++jj) { + active |= ss[cc*NE + jj] != 0.0f; + } + } + if (!active) { + continue; + } + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { lo[ii] += o4_t(float4(pv4[cc*NE*NS20/4 + ii*NL])*float4(sst[cc*NE])); } } } else { FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + bool active = true; + if (FC_flash_attn_ext_vec_sparse_mask) { + active = false; + FOR_UNROLL (short jj = 0; jj < NE; ++jj) { + active |= ss[cc*NE + jj] != 0.0f; + } + } + if (!active) { + continue; + } + device const vd4_t * pv4 = (device const vd4_t *) (v + ((ic + NE*cc + ty)*args.nb21)); FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { @@ -11544,6 +11582,107 @@ kernel void kernel_lightning_indexer( } } +// Fuse the get_rows, softmax, multiply and reduction sequence used by the DSV4 +// compressors. One threadgroup owns a contiguous embedding tile for one output +// block, so its row plan is loaded once into threadgroup memory and all state +// reads remain coalesced across embedding lanes. +kernel void kernel_dsv4_compress( + constant ggml_metal_kargs_dsv4_compress & args, + device const char * kv_state, + device const char * score_state, + device const char * read_idxs, + device char * dst, + threadgroup int32_t * idxs [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int n_read = (args.overlap ? 2 : 1)*args.ratio; + const int ib = tgpig.y; + + for (int j = tiitg; j < n_read; j += ntg.x) { + const bool cur_half = args.overlap && j >= args.ratio; + const int jr = cur_half ? j - args.ratio : j; + const int idx_pos = (cur_half ? args.ratio*args.n_blocks : 0) + ib*args.ratio + jr; + idxs[j] = *(device const int32_t *) (read_idxs + idx_pos*args.nb_i0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int i0 = tgpig.x*ntg.x + tiitg; + if (i0 >= args.n_embd) { + return; + } + + float score_max = -INFINITY; + for (int j = 0; j < n_read; ++j) { + const int idx = idxs[j]; + if (idx < 0 || idx >= args.n_rows) { + continue; + } + + const bool cur_half = args.overlap && j >= args.ratio; + const int i_src = (cur_half ? args.n_embd : 0) + i0; + const float score = *(device const float *) (score_state + i_src*args.nb_s0 + idx*args.nb_s1); + score_max = max(score_max, score); + } + + float sum_v = 0.0f; + float sum_w = 0.0f; + if (score_max != -INFINITY) { + for (int j = 0; j < n_read; ++j) { + const int idx = idxs[j]; + if (idx < 0 || idx >= args.n_rows) { + continue; + } + + const bool cur_half = args.overlap && j >= args.ratio; + const int i_src = (cur_half ? args.n_embd : 0) + i0; + const float score = *(device const float *) (score_state + i_src*args.nb_s0 + idx*args.nb_s1); + const float weight = exp(score - score_max); + const float value = *(device const float *) (kv_state + i_src*args.nb_k0 + idx*args.nb_k1); + sum_v += value*weight; + sum_w += weight; + } + } + + *(device float *) (dst + i0*args.nb_d0 + ib*args.nb_d1) = sum_w > 0.0f ? sum_v/sum_w : 0.0f; +} + +// Materialize raw visibility and the Lightning-Indexer selection in one pass. +// TOP_K indices are unique, so selected rows can overwrite the initial -INF +// fill without atomics after a single threadgroup barrier. +kernel void kernel_dsv4_top_k_mask( + constant ggml_metal_kargs_dsv4_top_k_mask & args, + device const char * raw_mask, + device const char * comp_mask, + device const char * comp_idx, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort ntg[[threads_per_threadgroup]]) { + const int iq = tgpig % args.n_query; + const int is = tgpig / args.n_query; + + device const half * rm = (device const half *) (raw_mask + (uint64_t) iq*args.nb_rm1 + (uint64_t) is*args.nb_rm3); + device const half * cm = (device const half *) (comp_mask + (uint64_t) iq*args.nb_cm1 + (uint64_t) is*args.nb_cm3); + device const int * ci = (device const int *) (comp_idx + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); + device half * out = (device half *) (dst + (uint64_t) iq*args.nb_d1 + (uint64_t) is*args.nb_d3); + + for (int i = tiitg; i < args.n_raw; i += ntg) { + out[i] = rm[i]; + } + for (int i = tiitg; i < args.n_comp; i += ntg) { + out[args.n_raw + i] = -INFINITY; + } + threadgroup_barrier(mem_flags::mem_device); + + for (int i = tiitg; i < args.n_select; i += ntg) { + const int idx = ci[i]; + if (idx >= 0 && idx < args.n_comp) { + out[args.n_raw + idx] = cm[idx]; + } + } +} + typedef decltype(kernel_lightning_indexer) kernel_lightning_indexer_t; template [[host_name("kernel_lightning_indexer_f32")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; @@ -11717,10 +11856,9 @@ kernel void kernel_dsv4_hc_post_f32( } } -// One threadgroup owns one token. The Lightning-Indexer selection is shared by -// all 64 attention heads, so every cache row is fetched once and packed next to -// a head-broadcast mask. This layout lets the regular tiled FA kernel treat the -// heads as query rows without materializing any dense top-k mask. +// Prefill assigns one threadgroup per token. Decode's compressed-only mode +// assigns one threadgroup per selected row to expose enough random-read +// parallelism. The packed masks remain adjacent to the F16 key storage. kernel void kernel_dsv4_sparse_pack( constant ggml_metal_kargs_dsv4_sparse_pack & args, device const char * raw_k, @@ -11729,18 +11867,40 @@ kernel void kernel_dsv4_sparse_pack( device const char * comp_mask, device const char * comp_idx, device char * dst, - uint tgpig[[threadgroup_position_in_grid]], + uint3 tgpig[[threadgroup_position_in_grid]], ushort tiitg[[thread_index_in_threadgroup]], - ushort ntg[[threads_per_threadgroup]]) { + ushort3 ntg[[threads_per_threadgroup]]) { constexpr int max_selected = 128 + 512; threadgroup int selected_idx[max_selected]; threadgroup half selected_mask[max_selected]; - const int iq = tgpig % args.n_batch; - const int is = tgpig / args.n_batch; + if (args.n_raw == 0) { + const int i = tgpig.x; + const int it = tgpig.y; + const int iq = it % args.n_batch; + const int is = it / args.n_batch; + const int idx = *(device const int *) (comp_idx + + (uint64_t) i*args.nb_ci0 + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); + + device const half * src = (device const half *) (comp_k + + (uint64_t) idx*args.nb_ck2 + (uint64_t) is*args.nb_ck3); + device half * out = (device half *) (dst + (uint64_t) it*args.nb_d1); + for (int e = tiitg; e < args.n_embd; e += ntg.x) { + out[i*args.n_embd + e] = src[e]; + } + if (tiitg == 0) { + out[args.n_embd*args.n_comp + i] = *(device const half *) (comp_mask + + (uint64_t) idx*args.nb_cm0 + (uint64_t) iq*args.nb_cm1 + (uint64_t) is*args.nb_cm3); + } + return; + } + + const int it = tgpig.x; + const int iq = it % args.n_batch; + const int is = it / args.n_batch; const int nk = args.n_raw + args.n_comp; - device half * out = (device half *) (dst + (uint64_t) tgpig*args.nb_d1); + device half * out = (device half *) (dst + (uint64_t) it*args.nb_d1); device half * out_k = out; device half * out_m = out + args.n_embd*nk; @@ -11764,7 +11924,7 @@ kernel void kernel_dsv4_sparse_pack( selected_mask[n] = -INFINITY; } } - for (int i = tiitg; i < args.n_comp; i += ntg) { + for (int i = tiitg; i < args.n_comp; i += ntg.x) { const int oi = args.n_raw + i; const int idx = *(device const int *) (comp_idx + (uint64_t) i*args.nb_ci0 + (uint64_t) iq*args.nb_ci1 + (uint64_t) is*args.nb_ci3); @@ -11779,11 +11939,11 @@ kernel void kernel_dsv4_sparse_pack( if (idx >= 0) { device const half * src = (device const half *) (raw_k + (uint64_t) idx*args.nb_rk2 + (uint64_t) is*args.nb_rk3); - for (int e = tiitg; e < args.n_embd; e += ntg) { + for (int e = tiitg; e < args.n_embd; e += ntg.x) { out_k[i*args.n_embd + e] = src[e]; } } else { - for (int e = tiitg; e < args.n_embd; e += ntg) { + for (int e = tiitg; e < args.n_embd; e += ntg.x) { out_k[i*args.n_embd + e] = 0.0h; } } @@ -11797,7 +11957,7 @@ kernel void kernel_dsv4_sparse_pack( const int idx = selected_idx[oi]; device const half * src = (device const half *) (comp_k + (uint64_t) idx*args.nb_ck2 + (uint64_t) is*args.nb_ck3); - for (int e = tiitg; e < args.n_embd; e += ntg) { + for (int e = tiitg; e < args.n_embd; e += ntg.x) { out_k[oi*args.n_embd + e] = src[e]; } if (tiitg == 0) {