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;