From bccc54b136b4ea4d240a2c9d2cea71e32e56dbfc Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 18 Jul 2026 10:46:01 -0300 Subject: [PATCH 1/2] 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 276aea00ea1b..a88eab440113 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 35f0c44ec421..83de623960d0 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 a5a3a58ad054..c8a345cf846f 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 491316f74912..b471a24fbdd9 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 42ec809ce521..df042d23ac12 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 4c1642a67603..cb1629e0bfc0 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 59191c663eb0..ae904c6a1ad8 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 19cca7df1e9d..c53e66c35c6a 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 5018170ed85e..933ba2c86123 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 160e29413552..e15d822a2567 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 e68dc49b6dff..2eae0105aa82 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 8cb598935861..a5831a9f45f6 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/2] 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 c153bd82177b..8fe3a391f103 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 7e1deeaa2102..1a2daf47a914 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 2dc6eb8fdbc0..c8933dc408ac 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 e173b91c0c5c..561826387e3f 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 c5d7619c12fa..007e4282ca21 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 b03b59e0bd92..0f1c5d893baa 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 7d12cb0fe39b..c904429be527 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]; + } + } + +}