From 72f0f12b7caff3b8b1e09074d17f994807079752 Mon Sep 17 00:00:00 2001 From: rcfa <1831995+rcfa@users.noreply.github.com> Date: Wed, 30 Sep 2026 05:58:55 +0000 Subject: [PATCH 1/2] qwen4exp : use maskless block selection only where the selected-key kernels run Complete-block selection builds a maskless flash-attention op whose visibility lives only in the selected-cell list (src[5]). Only the selected-key kernels read that list, and they exist only in the HIP backend (RDNA3.5); the gate did not check the device. ggml-cuda aborts on the op when no QSA kernel takes it. The Vulkan backend accepted it and ran its dense kernel, which ignores src[5], so the QSA layers attended to the whole K/V view, unmasked, as soon as the context exceeded indexer_top_k + ratio - 1: - other sequences' cells: with several conversations in one unified cache, a conversation's output depended on another conversation's content (a 2048-token prompt, greedy, changed at the 2nd token next to related text, where it had been 6.6 nats ahead; unchanged next to unrelated text) - no sparsity: every cell instead of the top-k blocks test-backend-ops shows it directly: QSA_PREFILL_MASKLESS fails on Vulkan (0/6, ERR 0.48 to 0.98 against the FP64 reference). - The memory decides once per model whether its QSA layers all run on ROCm (LLAMA_QSA_SELECTED_KEY=0/1 overrides that), and qwen4exp_use_block_selection() requires it. Everywhere else the masked top-k path is used, which the dense kernels compute correctly (QSA_PREFILL and QSA_DECODE pass on Vulkan except their poison=1 cases, which fail on master as well). - Vulkan no longer claims flash attention with a selection list and no mask, so the op can never be computed wrongly there. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 4 ++++ src/llama-memory-hybrid-idx.cpp | 15 +++++++++++++++ src/llama-memory-hybrid-idx.h | 4 ++++ src/models/qwen4exp.cpp | 12 ++++++------ 4 files changed, 29 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 6d002063f15e..ec3aba2c4ca6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -16777,6 +16777,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_FLASH_ATTN_EXT: { bool coopmat2 = device->coopmat2; + // maskless selected-key attention (src[5] without a mask): no Vulkan kernel reads the list, the dense kernel would attend to all cells + if (op->src[5] != nullptr && op->src[3] == nullptr) { + return false; + } uint32_t HSK = op->src[1]->ne[0]; uint32_t HSV = op->src[2]->ne[0]; if ((HSK % 8) != 0 || (HSV % 8) != 0) { diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index c32477868eef..5fd0e1f56482 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -20,6 +20,7 @@ static bool hybrid_idx_no_recr(const llama_memory_recurrent * r) { #include #include #include +#include #include #include #include @@ -78,6 +79,20 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); }()) { + if (mem_idx) { + // the selected-key attention kernels exist only in the HIP backend (RDNA3.5); the QSA layers must all run there + bool all = true, any = false; + for (int il = 0; il < (int) model.hparams.n_layer_all; ++il) { + if (!model.hparams.has_kv(il) || !filter_idx(il) || model.hparams.dsv4_compress_ratios[il] <= 0) { continue; } + const char * reg = ggml_backend_reg_name(ggml_backend_dev_backend_reg(model.dev_layer(il))); + any = true; + all = all && std::strcmp(reg, "ROCm") == 0; + } + selected_key_attn = any && all && offload; + if (const char * e = getenv("LLAMA_QSA_SELECTED_KEY")) { selected_key_attn = atoi(e) != 0; } + LLAMA_LOG_INFO("%s: QSA attention: %s\n", __func__, selected_key_attn ? + "selected-key kernels (maskless block selection)" : "masked top-k (no selected-key kernels on this backend)"); + } if (!mem_idx || !offload || n_swa != 0) { return; } qsa_prefix = qsa_prefix_state(kv_size); const int layers = model.hparams.n_layer_all; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index a28c00cdc061..0d27ae3fb6d6 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -99,6 +99,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { void qsa_invalidate(); bool qsa_prefix_matches(const llama_ubatch & ubatch) const; bool qsa_fast(int il, const llama_ubatch & ubatch) const; + bool qsa_selected_key_attn() const { return selected_key_attn; } ggml_tensor * qsa_cache(ggml_context * ctx, int il, int64_t blocks) const; void qsa_fill_updates(ggml_tensor * members, ggml_tensor * positions, ggml_tensor * rows) const; void qsa_commit(int il) const; @@ -118,6 +119,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { ggml_tensor * tails, const llama_ubatch & ubatch, uint32_t ratio) const; bool incremental_qsa = false; + // the QSA layers run on ROCm, which has the selected-key kernels for the maskless block selection; other backends use the masked top-k + bool selected_key_attn = false; bool qsa_recover_pending = false; bool qsa_recover(llama_seq_id seq); qsa_prefix_state qsa_prefix; @@ -176,6 +179,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { bool qsa_prefix_matches(const llama_ubatch & u) const { return mem && mem->qsa_prefix_matches(u); } bool qsa_fast(int il, const llama_ubatch & u) const { return mem && mem->qsa_fast(il, u); } + bool qsa_selected_key_attn() const { return mem && mem->qsa_selected_key_attn(); } // cells the block metadata must cover: the KV view is sized by occupied cells, but a cache whose positions // run ahead of its cells (an MTP draft never receives the image cells an M-RoPE image pins to one position) // has blocks past that view. Never below get_n_kv(), padded like it so graph reuse keeps its cadence. diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index d071c5281acb..50a591bef940 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -827,12 +827,12 @@ static int64_t qwen4exp_query_strip(int64_t n_tokens, int64_t n_stream) { // append the query's own partial block as the tail (up to ratio-1 cells), -1 where nothing is visible. The // selected rows then carry the visibility themselves, so attention can run without a mask. Only the selected-key // attention kernels understand that layout (F16 K/V, head 256, single stream; qsa_decode takes 1..512 queries, -// qsa_prefill larger ubatches). They exist only in the HIP backend for RDNA3.5; the gate does not check the device, so -// on any other backend the resulting maskless op is not taken by a QSA kernel. +// qsa_prefill larger ubatches). They exist only in the HIP backend for RDNA3.5, so the selection needs selected_key_attn +// (the QSA layers run on ROCm). Other backends do not read the list: ggml-cuda aborts, Vulkan would attend to all cells. static bool qwen4exp_use_block_selection(bool blk_bias, int64_t n_stream, int64_t ratio, int64_t n_kv, const llama_ubatch & ubatch, const llama_cparams & cparams, const llama_hparams & hparams, - ggml_type type_k, ggml_type type_v) { - return blk_bias && cparams.causal_attn && n_stream==1 && ratio>1 && hparams.indexer_top_k%ratio==0 && + ggml_type type_k, ggml_type type_v, bool selected_key_attn) { + return selected_key_attn && blk_bias && cparams.causal_attn && n_stream==1 && ratio>1 && hparams.indexer_top_k%ratio==0 && n_kv>hparams.indexer_top_k+ratio-1 && n_kv<=16777216 && ubatch.token && cparams.flash_attn && cparams.offload_kqv && hparams.f_max_alibi_bias==0.0f && !hparams.attn_soft_cap && hparams.n_embd_head_k()==256 && hparams.n_embd_head_v()==256 && @@ -982,7 +982,7 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { // the selection mode and the strip bounds are baked into the graph, so a reused graph must agree on them const auto * attn = mctx->get_attn(); const bool blocks = attn && qwen4exp_use_block_selection(blk_bias, n_stream, ratio, n_kv, - params.ubatch, params.cparams, params.hparams, attn->type_k(), attn->type_v()); + params.ubatch, params.cparams, params.hparams, attn->type_k(), attn->type_v(), mctx->qsa_selected_key_attn()); const bool scalar = blocks && params.hparams.n_swa == 0 && mctx->qsa_scalar_visibility(params.ubatch, ratio); res &= (tail_idxs != nullptr) == scalar; res &= compact == scalar; @@ -1145,7 +1145,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); const auto * attn_ctx = mctx_hyb->get_attn(); const bool blocks = attn_ctx && qwen4exp_use_block_selection(blk_bias, n_stream, r, n_kv, ubatch, cparams, hparams, - attn_ctx->type_k(), attn_ctx->type_v()); + attn_ctx->type_k(), attn_ctx->type_v(), mctx_hyb->qsa_selected_key_attn()); const bool scalar = blocks && hparams.n_swa == 0 && mctx_hyb->qsa_scalar_visibility(ubatch, (uint32_t) r); qsa->compact = scalar; qsa->maskless = scalar; From 3bf53565b7d5d96f855eb33e9351e528d6f255ee Mon Sep 17 00:00:00 2001 From: rcfa <1831995+rcfa@users.noreply.github.com> Date: Wed, 30 Sep 2026 06:18:00 +0000 Subject: [PATCH 2/2] vulkan : selected-key flash attention for QSA block selection Complete-block selection builds a maskless FLASH_ATTN_EXT whose visibility is only in the selection list (src[5]). Vulkan had no kernel that reads the list: #115 makes it decline the op, and Vulkan then takes the masked per-cell top-k path, which is correct but slow at depth. flash_attn_sel.comp computes the op: one workgroup per query x K/V head, one 64-lane subgroup per query head of that K/V head (GQA 4/8/12, head 256). It sorts the query's list in shared memory and skips invalid and duplicate entries, as the CPU path does, then gathers K and V in blocks of 32 cells. Two lanes score one key, so the online softmax needs two subgroup reductions per block. With fewer than 64 workgroups the rows are split, and flash_attn_split_k_reduce combines the parts. supports_op takes a maskless selected-key op exactly when this kernel does. llama_memory_hybrid_idx asks each QSA layer's device with a probe op instead of checking for ROCm by name. HIP stays accepted; CUDA is not asked, because its supports_op ignores the list. Radeon 890M (gfx1150), Qwen3.8-Flash-Next UD-IQ4_XS: - test-backend-ops: QSA_PREFILL_MASKLESS 6/6 and QSA_DECODE_MASKLESS 25/25 (master: 0/6 and 0/25), NMSE <= 6.9e-13 against the FP64 reference - pp2048 at depth 12000 / 32000: 140.9 / 134.1 t/s (masked path, #115: 139.0 / 88.7) Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 --- .../ggml-vulkan/ggml-vulkan-push-constants.h | 12 + ggml/src/ggml-vulkan/ggml-vulkan-types.h | 1 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 128 ++++++++++- .../vulkan-shaders/flash_attn_sel.comp | 209 ++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + src/llama-memory-hybrid-idx.cpp | 21 +- src/llama-memory-hybrid-idx.h | 2 +- src/models/qwen4exp.cpp | 4 +- 8 files changed, 371 insertions(+), 7 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 13b8e73db151..7d8fa430d55c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h @@ -729,6 +729,18 @@ struct vk_op_flash_attn_top_k_push_constants { }; static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128); +struct vk_op_flash_attn_sel_push_constants { + uint32_t n_kv, n_sel, n_sort; + uint32_t nbq1, nbq2, nbq3; + uint32_t nbk1, nbk2, nbk3; + uint32_t nbv1, nbv2, nbv3; + uint32_t nbi1, nbi3; + uint32_t nb1, nb2, nb3; + float scale; + uint32_t n_split, chunk, n_head, n_tok, n_stream; +}; +static_assert(sizeof(vk_op_flash_attn_sel_push_constants) <= 128); + struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index ab50260ae6ac..b03930ab8232 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -1009,6 +1009,7 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_cm_small_f16[LI_NH_COUNT]; vk_pipeline pipeline_lightning_indexer_decode_cm_f16[LI_NH_COUNT]; vk_pipeline pipeline_flash_attn_top_k_f16; + vk_pipeline pipeline_flash_attn_sel_f16[3]; // maskless selected-key FA, GQA 4, 8, 12 vk_pipeline pipeline_flash_attn_top_k_cm_f16; vk_pipeline pipeline_flash_attn_gather_f16; vk_pipeline pipeline_flash_attn_gather_dq[GGML_TYPE_COUNT]; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ec3aba2c4ca6..b8b9c5c4c273 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3925,6 +3925,17 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "flash_attn_top_k_f16", flash_attn_top_k_f16_len, flash_attn_top_k_f16_data, "main", 6, sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, device->subgroup_size); + // shared memory: the sorted row (16 KiB), Q (1 KiB per head), a K and a V block (16 KiB each): at most GQA 12 + static const char * sel_names[3] = { "flash_attn_sel_f16_gqa4", "flash_attn_sel_f16_gqa8", "flash_attn_sel_f16_gqa12" }; + for (uint32_t i = 0; i < 3; ++i) { + const uint32_t gqa = 4 * (i + 1); + if (device->subgroup_shuffle && 64 * gqa <= device->properties.limits.maxComputeWorkGroupInvocations && + device->properties.limits.maxComputeSharedMemorySize >= 65536) { + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_sel_f16[i], + sel_names[i], flash_attn_sel_f16_len, flash_attn_sel_f16_data, "main", 5, + sizeof(vk_op_flash_attn_sel_push_constants), {1, 1, 1}, {64 * gqa, 64, gqa}, 1, true, true, 64); + } + } ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_f16, "flash_attn_gather_f16", flash_attn_gather_f16_len, flash_attn_gather_f16_data, "main", 5, sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, @@ -8478,6 +8489,113 @@ static void ggml_vk_perf_mark_subop(ggml_backend_vk_context * ctx, vk_context& s subctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++); } +// maskless selected-key attention (QSA block selection): the pipeline that takes this op, or nullptr +static vk_pipeline ggml_vk_flash_attn_sel_pipeline(const vk_device & device, const ggml_tensor * op) { + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; + const ggml_tensor * ids = op->src[5]; + if (ids == nullptr || op->src[3] != nullptr || op->src[4] != nullptr || ggml_get_op_params_i32(op, 4) != 0) { + return nullptr; + } + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f || op->type != GGML_TYPE_F32 || !ggml_is_contiguous(op) || + q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || ids->type != GGML_TYPE_I32 || + q->ne[0] != 256 || k->ne[0] != 256 || v->ne[0] != 256 || + k->ne[1] != v->ne[1] || k->ne[2] != v->ne[2] || k->ne[3] != q->ne[3] || v->ne[3] != q->ne[3] || + k->ne[2] < 1 || q->ne[2] % k->ne[2] != 0 || + ids->ne[0] < 1 || ids->ne[0] > 4096 || ids->ne[1] < q->ne[1] || ids->ne[2] != 1 || + (ids->ne[3] != 1 && ids->ne[3] != q->ne[3]) || + q->nb[0] != sizeof(float) || k->nb[0] != sizeof(ggml_fp16_t) || v->nb[0] != sizeof(ggml_fp16_t) || + ids->nb[0] != sizeof(int32_t)) { + return nullptr; + } + for (int i = 1; i < 4; ++i) { + if (q->nb[i] % sizeof(float) || k->nb[i] % sizeof(ggml_fp16_t) || v->nb[i] % sizeof(ggml_fp16_t) || + ids->nb[i] % sizeof(int32_t)) { + return nullptr; + } + } + // the shader indexes with 32-bit element offsets + if (ggml_nbytes(k) / sizeof(ggml_fp16_t) > UINT32_MAX || ggml_nbytes(v) / sizeof(ggml_fp16_t) > UINT32_MAX || + ggml_nbytes(q) / sizeof(float) > UINT32_MAX || ggml_nbytes(op) / sizeof(float) > UINT32_MAX) { + return nullptr; + } + switch (q->ne[2] / k->ne[2]) { + case 4: return device->pipeline_flash_attn_sel_f16[0]; + case 8: return device->pipeline_flash_attn_sel_f16[1]; + case 12: return device->pipeline_flash_attn_sel_f16[2]; + default: return nullptr; + } +} + +static bool ggml_vk_flash_attn_sel(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, ggml_tensor * dst) { + vk_pipeline pipeline = ggml_vk_flash_attn_sel_pipeline(ctx->device, dst); + if (!pipeline) { + return false; + } + const ggml_tensor * ids = dst->src[5]; + float scale = 1.0f; + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + uint32_t n_sort = 1; + while (n_sort < (uint32_t) ids->ne[0]) { + n_sort <<= 1; + } + // few queries (decode) give too few workgroups to fill the GPU: split each row, at least 64 keys per split + const uint32_t n_sel = (uint32_t) ids->ne[0]; + const uint32_t n_wg = (uint32_t) (q->ne[1] * k->ne[2] * q->ne[3]); + uint32_t n_split = n_wg < 64 ? std::min(CEIL_DIV(64u, n_wg), CEIL_DIV(n_sel, 64u)) : 1u; + const uint32_t chunk = ROUNDUP_POW2(CEIL_DIV(n_sel, std::max(n_split, 1u)), 32u); // whole key blocks + n_split = CEIL_DIV(n_sel, chunk); + const vk_op_flash_attn_sel_push_constants pc = { + (uint32_t) k->ne[1], (uint32_t) ids->ne[0], n_sort, + (uint32_t) (q->nb[1] / sizeof(float)), (uint32_t) (q->nb[2] / sizeof(float)), (uint32_t) (q->nb[3] / sizeof(float)), + (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[2] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (v->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (v->nb[2] / sizeof(ggml_fp16_t)), (uint32_t) (v->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (ids->nb[1] / sizeof(int32_t)), ids->ne[3] == 1 ? 0u : (uint32_t) (ids->nb[3] / sizeof(int32_t)), + (uint32_t) (dst->nb[1] / sizeof(float)), (uint32_t) (dst->nb[2] / sizeof(float)), (uint32_t) (dst->nb[3] / sizeof(float)), + scale, n_split, chunk, (uint32_t) q->ne[2], (uint32_t) q->ne[1], (uint32_t) q->ne[3], + }; + const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + const std::array grid = { (uint32_t) q->ne[1] * n_split, (uint32_t) k->ne[2], (uint32_t) q->ne[3] }; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + if (n_split == 1) { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, v), ggml_vk_tensor_subbuffer(ctx, ids), dst_buf }, + pc, grid); + return true; + } + const uint32_t D = 256; + const uint32_t NH = (uint32_t) q->ne[2]; + const uint32_t N = (uint32_t) q->ne[1]; + const uint32_t NS = (uint32_t) q->ne[3]; + const size_t split_size = ((size_t) D * NH * N * NS * n_split + (size_t) NH * 2 * N * NS * n_split) * sizeof(float); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); + if (ctx->prealloc_size_split_k < split_size) { + ctx->prealloc_size_split_k = split_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, v), ggml_vk_tensor_subbuffer(ctx, ids), split_buf }, + pc, grid); + ctx->prealloc_split_k_need_sync = true; + ggml_vk_sync_buffers(ctx, subctx); + const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = { D, NH, N, N, NS, n_split, 0 }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, + { split_buf, q_buf, dst_buf }, reduce_pc, { NH, D, N * NS }); + ctx->prealloc_split_k_need_sync = true; + return true; +} + static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { @@ -9223,6 +9341,12 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const assert(dst->type == GGML_TYPE_F32); assert(q->type == GGML_TYPE_F32); + if (dst->src[5] != nullptr && mask == nullptr) { + // supports_op only takes a maskless selected-key op when flash_attn_sel does + const bool done = ggml_vk_flash_attn_sel(ctx, subctx, q, k, v, dst); + GGML_ASSERT(done); + return; + } if (ggml_vk_flash_attn_top_k(ctx, subctx, q, k, v, mask, sinks, dst)) { return; } @@ -16777,9 +16901,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_FLASH_ATTN_EXT: { bool coopmat2 = device->coopmat2; - // maskless selected-key attention (src[5] without a mask): no Vulkan kernel reads the list, the dense kernel would attend to all cells + // maskless selected-key attention (src[5] without a mask): only flash_attn_sel reads the list, the dense kernel would attend to all cells if (op->src[5] != nullptr && op->src[3] == nullptr) { - return false; + return ggml_vk_flash_attn_sel_pipeline(device, op) != nullptr; } uint32_t HSK = op->src[1]->ne[0]; uint32_t HSV = op->src[2]->ne[0]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp new file mode 100644 index 000000000000..5a9ecd215ff4 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp @@ -0,0 +1,209 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require + +// Maskless selected-key flash attention (QSA complete-block selection): each query's ids row names the cells it +// may attend to, -1 for none. One workgroup per (query, K/V head, stream), one subgroup per query head of that K/V +// head, so each K/V row is read once per query. The row is sorted and duplicates are skipped, like the CPU path. +// Keys go in blocks of BK: two lanes compute one key's dot product (half the dims each), so the softmax needs two +// subgroup reductions per block instead of one per key. With n_split > 1 (few queries) each workgroup takes one +// chunk of the sorted row and writes a partial result in the layout of flash_attn_split_k_reduce.comp. + +layout(constant_id = 0) const uint WORKGROUP_SIZE = 768; +layout(constant_id = 1) const uint SUBGROUP_SIZE = 64; +layout(constant_id = 2) const uint GQA = 12; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer VBuf { float16_t data_v[]; }; +layout(binding = 3) readonly buffer IdsBuf { int data_ids[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_sel; + uint n_sort; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk2; + uint nbk3; + uint nbv1; + uint nbv2; + uint nbv3; + uint nbi1; + uint nbi3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint n_split; + uint chunk; + uint n_head; + uint n_tok; + uint n_stream; +} p; + +// The pipeline is created with subgroup size 64 and head size 256 only (see ggml_vk_flash_attn_sel_pipeline). +const uint HEAD_SIZE = 256; +const uint LANES = 64; +const uint PER_LANE = HEAD_SIZE / LANES; +const uint BK = 32; +const uint HALF = HEAD_SIZE / 2; +const uint MAX_SORT = 4096; +const uint NONE = 0xffffffffu; +const float FLT_MAX = 3.402823466e38; + +shared uint sorted_key[MAX_SORT]; +shared uint block_key[BK]; +shared float q_sh[GQA * HEAD_SIZE]; +shared float16_t kt_sh[HEAD_SIZE * BK]; // K block, transposed: [dim][key], so the per-lane key reads hit different banks +shared float16_t v_sh[BK * HEAD_SIZE]; // V block: [key][dim] +shared float p_sh[GQA * BK]; + +void main() { + const float NEG_INF = uintBitsToFloat(0xff800000); + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint sg = gl_SubgroupID; + const uint split = gl_WorkGroupID.x % p.n_split; + const uint query = gl_WorkGroupID.x / p.n_split; + const uint kv_head = gl_WorkGroupID.y; + const uint stream = gl_WorkGroupID.z; + const uint head = kv_head * GQA + sg; + + // No early return: every invocation must reach every barrier below. + + // the query's row, invalid entries as NONE, sorted ascending (NONE last) + const uint ids_base = stream * p.nbi3 + query * p.nbi1; + for (uint i = tid; i < p.n_sort; i += WORKGROUP_SIZE) { + uint key = NONE; + if (i < p.n_sel) { + const int id = data_ids[ids_base + i]; + if (id >= 0 && uint(id) < p.n_kv) { + key = uint(id); + } + } + sorted_key[i] = key; + } + const uint q_base = stream * p.nbq3 + head * p.nbq2 + query * p.nbq1; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + q_sh[sg * HEAD_SIZE + lane + i * LANES] = data_q[q_base + lane + i * LANES]; + } + barrier(); + for (uint size = 2; size <= p.n_sort; size <<= 1) { + for (uint stride = size >> 1; stride > 0; stride >>= 1) { + for (uint i = tid; i < p.n_sort; i += WORKGROUP_SIZE) { + const uint j = i ^ stride; + if (j > i) { + const uint a = sorted_key[i]; + const uint b = sorted_key[j]; + if ((a > b) == ((i & size) == 0)) { + sorted_key[i] = b; + sorted_key[j] = a; + } + } + } + barrier(); + } + } + + float acc[PER_LANE]; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + acc[i] = 0.0; + } + float row_max = NEG_INF; + float row_sum = 0.0; + + const uint k_base = stream * p.nbk3 + kv_head * p.nbk2; + const uint v_base = stream * p.nbv3 + kv_head * p.nbv2; + const uint kcol = lane % BK; // the key this lane scores + const uint khalf = lane / BK; // and which half of its dims + + const uint kb_begin = split * p.chunk; + const uint kb_end = min(p.n_sel, kb_begin + p.chunk); + for (uint kb = kb_begin; kb < kb_end; kb += BK) { + // sorted: once the block starts at NONE, so does everything after it (uniform, read from shared memory) + if (sorted_key[kb] == NONE) { + break; + } + if (tid < BK) { + const uint s = kb + tid; + uint key = s < kb_end ? sorted_key[s] : NONE; + if (key != NONE && s > 0 && sorted_key[s - 1] == key) { + key = NONE; // duplicate + } + block_key[tid] = key; + } + barrier(); + + for (uint idx = tid; idx < BK * HEAD_SIZE; idx += WORKGROUP_SIZE) { + const uint col = idx / HEAD_SIZE; + const uint dim = idx % HEAD_SIZE; + const uint key = block_key[col]; + kt_sh[dim * BK + col] = key != NONE ? data_k[k_base + key * p.nbk1 + dim] : float16_t(0.0); + v_sh[idx] = key != NONE ? data_v[v_base + key * p.nbv1 + dim] : float16_t(0.0); + } + barrier(); + + float s = 0.0; + const uint qoff = sg * HEAD_SIZE + khalf * HALF; + const uint koff = khalf * HALF * BK + kcol; + for (uint d = 0; d < HALF; ++d) { + s += q_sh[qoff + d] * float(kt_sh[koff + d * BK]); + } + s += subgroupShuffleXor(s, BK); + const bool valid = block_key[kcol] != NONE; + s = valid ? s * p.scale : NEG_INF; + + const float new_max = max(row_max, subgroupMax(s)); + const float old_scale = new_max == row_max ? 1.0 : exp(row_max - new_max); + const float weight = valid ? exp(s - new_max) : 0.0; + row_sum = row_sum * old_scale + subgroupAdd(khalf == 0 ? weight : 0.0); + row_max = new_max; + if (khalf == 0) { + p_sh[sg * BK + kcol] = weight; + } + subgroupMemoryBarrierShared(); + subgroupBarrier(); + + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + acc[i] *= old_scale; + } + for (uint j = 0; j < BK; ++j) { + const float w = p_sh[sg * BK + j]; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + acc[i] += w * float(v_sh[j * HEAD_SIZE + lane + i * LANES]); + } + } + barrier(); + } + + if (p.n_split == 1) { + const uint dst_base = stream * p.nb3 + query * p.nb2 + head * p.nb1; + const float inv_sum = row_sum == 0.0 ? 0.0 : 1.0 / row_sum; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + data_dst[dst_base + lane + i * LANES] = acc[i] * inv_sum; + } + } else { + // partial: unnormalized O, then L and M per head (an empty split writes M = -FLT_MAX, not -inf, so the + // reduce never computes exp(-inf - -inf)) + const uint slot = split + p.n_split * (query + p.n_tok * stream); + const uint o_base = HEAD_SIZE * p.n_head * slot + HEAD_SIZE * head; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + data_dst[o_base + lane + i * LANES] = acc[i]; + } + if (lane == 0) { + const uint lm = HEAD_SIZE * p.n_head * p.n_tok * p.n_stream * p.n_split + p.n_head * 2 * slot + head; + data_dst[lm] = row_sum; + data_dst[lm + p.n_head] = max(row_max, -FLT_MAX); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index f3e45215a7ef..23f43fab49c5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -869,6 +869,7 @@ void process_shaders() { string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); #endif string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); + string_to_spv("flash_attn_sel_f16", "flash_attn_sel.comp", {}); #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) string_to_spv("flash_attn_top_k_cm_f16", "flash_attn_top_k_cm.comp", {}); #endif diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 5fd0e1f56482..a976cfa1584f 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -29,6 +29,21 @@ static bool hybrid_idx_no_recr(const llama_memory_recurrent * r) { // llama_memory_hybrid_idx // +// can dev run maskless selected-key attention with the shapes of QSA layer il? (a probe op, as for weight buffer types) +static bool qsa_selected_key_supported(ggml_backend_dev_t dev, const llama_hparams & hp, int il) { + ggml_init_params params = { 8*ggml_tensor_overhead(), nullptr, true }; + ggml_context_ptr ctx { ggml_init(params) }; + const int64_t d = hp.n_embd_head_k(); + const int64_t n_sel = hp.indexer_top_k + hp.dsv4_compress_ratios[il] - 1; + ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, d, 1, hp.n_head(il), 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, d, 4*n_sel, hp.n_head_kv(il), 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, d, 4*n_sel, hp.n_head_kv(il), 1); + ggml_tensor * ids = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_I32, n_sel, 1, 1, 1); + ggml_tensor * fa = ggml_flash_attn_ext(ctx.get(), q, k, v, nullptr, 1.0f, 0.0f, 0.0f); + ggml_flash_attn_ext_add_top_k(fa, ids, 0); + return ggml_backend_dev_supports_op(dev, fa); +} + llama_memory_hybrid_idx::llama_memory_hybrid_idx( const llama_model & model, /* attn */ @@ -80,13 +95,15 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( nullptr, filter_idx, nullptr, nullptr, "idx_"); }()) { if (mem_idx) { - // the selected-key attention kernels exist only in the HIP backend (RDNA3.5); the QSA layers must all run there + // the QSA layers must all run where a kernel reads the selection list: HIP (RDNA3.5), or a backend that takes the probe op + // CUDA is not asked: its supports_op does not look at the list bool all = true, any = false; for (int il = 0; il < (int) model.hparams.n_layer_all; ++il) { if (!model.hparams.has_kv(il) || !filter_idx(il) || model.hparams.dsv4_compress_ratios[il] <= 0) { continue; } const char * reg = ggml_backend_reg_name(ggml_backend_dev_backend_reg(model.dev_layer(il))); any = true; - all = all && std::strcmp(reg, "ROCm") == 0; + all = all && (std::strcmp(reg, "ROCm") == 0 || + (std::strcmp(reg, "CUDA") != 0 && qsa_selected_key_supported(model.dev_layer(il), model.hparams, il))); } selected_key_attn = any && all && offload; if (const char * e = getenv("LLAMA_QSA_SELECTED_KEY")) { selected_key_attn = atoi(e) != 0; } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 0d27ae3fb6d6..f26209435dca 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -119,7 +119,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { ggml_tensor * tails, const llama_ubatch & ubatch, uint32_t ratio) const; bool incremental_qsa = false; - // the QSA layers run on ROCm, which has the selected-key kernels for the maskless block selection; other backends use the masked top-k + // the QSA layers run where a kernel reads the maskless block selection (HIP on RDNA3.5, Vulkan); otherwise the masked top-k bool selected_key_attn = false; bool qsa_recover_pending = false; bool qsa_recover(llama_seq_id seq); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 50a591bef940..6112baca5909 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -827,8 +827,8 @@ static int64_t qwen4exp_query_strip(int64_t n_tokens, int64_t n_stream) { // append the query's own partial block as the tail (up to ratio-1 cells), -1 where nothing is visible. The // selected rows then carry the visibility themselves, so attention can run without a mask. Only the selected-key // attention kernels understand that layout (F16 K/V, head 256, single stream; qsa_decode takes 1..512 queries, -// qsa_prefill larger ubatches). They exist only in the HIP backend for RDNA3.5, so the selection needs selected_key_attn -// (the QSA layers run on ROCm). Other backends do not read the list: ggml-cuda aborts, Vulkan would attend to all cells. +// qsa_prefill larger ubatches). The HIP backend (RDNA3.5) and Vulkan (flash_attn_sel) have them, so the selection needs +// selected_key_attn (all QSA layers run there). Dense kernels do not read the list: ggml-cuda aborts, Vulkan declines. static bool qwen4exp_use_block_selection(bool blk_bias, int64_t n_stream, int64_t ratio, int64_t n_kv, const llama_ubatch & ubatch, const llama_cparams & cparams, const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, bool selected_key_attn) {