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] 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;