qwen4exp : use maskless block selection only where the selected-key kernels run

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) {

src/llama-memory-hybrid-idx.cpp

@@ -20,6 +20,7 @@ static bool hybrid_idx_no_recr(const llama_memory_recurrent * r) {
 #include <algorithm>
 #include <cassert>
 #include <cmath>
+#include <cstdlib>
 #include <cstring>
 #include <iterator>
 #include <stdexcept>
@@ -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;

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.

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;