From 541f63a63edc5c02c160c835c4ea8aa3bcfa5a79 Mon Sep 17 00:00:00 2001 From: Abdel Darwish <168705979+abdel-darwish-27@users.noreply.github.com> Date: Tue, 1 Sep 2026 07:23:16 +1000 Subject: [PATCH 1/4] qwen4exp: gather-based sparse attention for QSA decode At decode time build_attn_qsa built a full-n_kv mask (-INF everywhere with the top-k positions unmasked) and ran attention over the entire KV cache, so the indexer's top-k selection saved no attention compute. A kernel profile at 141K context shows flash_attn_ext_f16 at ~15 ms per QSA layer per token (~180 ms per token over 12 QSA layers); decode collapses from 53 tok/s shallow to ~7 tok/s at 141K on 2x RTX A6000. Add a decode-only gather path, taken for single-token-per-stream ubatches once the cache is at least twice the top-k width: - select winning *blocks* directly on the per-block indexer scores (the block bias already carries visibility), avoiding the O(n_kv) expansion of block scores to token scores and sorting n_blocks entries instead of n_kv - map winning blocks to cell indices via blk_cells and gather the selected cells' K/V (whole-cell rows; a cell's heads are contiguous in the cache) plus their kq_mask values, then attend densely over r*K_blk cells (2048 for Qwen3.8-Flash-Next), a multiple of 256 so flash attention padding holds - skip the O(n_kv) host-side cell_blk fill when the graph never references it QWEN4EXP_QSA_GATHER=0 restores the masked path (same binary A/B lever). Correctness: greedy outputs byte-identical to the masked path at 75K and 141K depth; mid-context needle retrieval passes in both modes at all tested depths. Decode throughput, UD-IQ4_XS, q8_0 KV, single stream (repeats within 0.1 t/s): depth masked gather 34K 17.4 19.4 (+11%) 68K 12.0 13.9 (+16%) 141K 7.1 8.8 (+23%) The remaining depth scaling in both modes is the indexer recomputing pooled block keys from the raw cache every layer per token; caching those incrementally is a follow-up. --- src/llama-memory-hybrid-idx.cpp | 11 +-- src/models/models.h | 6 +- src/models/qwen4exp.cpp | 123 ++++++++++++++++++++++++++++++-- 3 files changed, 129 insertions(+), 11 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d4e59d77e57..45c788e781a 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -350,7 +350,10 @@ void llama_memory_hybrid_idx_context::set_input_qsa( GGML_ASSERT(ratio > 0); GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); - GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); + // the block-gather decode path never references cell_blk in the graph, so the + // scheduler leaves it unallocated; skip its (O(n_kv) per ubatch) fill in that case + const bool has_cell_blk = cell_blk->buffer != nullptr; + GGML_ASSERT(!has_cell_blk || ggml_backend_buffer_is_host(cell_blk->buffer)); const int64_t n_kv = cell_blk->ne[0]; const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch @@ -361,7 +364,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( GGML_ASSERT(n_tokens % n_ns == 0); const int64_t n_tps = n_tokens/n_ns; // tokens per stream - int32_t * dst_cell_blk = (int32_t *) cell_blk->data; + int32_t * dst_cell_blk = has_cell_blk ? (int32_t *) cell_blk->data : nullptr; int32_t * dst_blk_cells = (int32_t *) blk_cells->data; int32_t * dst_blk_pos = (int32_t *) blk_pos->data; float * dst_bias = (float *) bias->data; @@ -385,7 +388,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream); - int32_t * cur_cell_blk = dst_cell_blk + s*n_kv; + int32_t * cur_cell_blk = has_cell_blk ? dst_cell_blk + s*n_kv : nullptr; int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks); // an incomplete block cannot be pooled; the bias below forces those tail cells in @@ -420,7 +423,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it // per-cell mode carries that -inf itself and only needs the gather in range - for (int64_t j = 0; j < n_kv; ++j) { + for (int64_t j = 0; has_cell_blk && j < n_kv; ++j) { if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) { blk_of[j] = -1; } diff --git a/src/models/models.h b/src/models/models.h index af60764c2f7..285e5af34bf 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2318,7 +2318,8 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * v_cur, ggml_tensor * top_k, float kq_scale, - int il); + int il, + bool gather = false); // the QSA cache layout inputs do not depend on the layer, only on its compress ratio, // so the layers sharing a ratio share one input set @@ -2331,7 +2332,8 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il); + int il, + bool gather = false); ggml_tensor * build_layer_attn_linear( llm_graph_input_rs * inp, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b5003..e2ddeea7090 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -472,7 +472,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il) { + int il, + bool gather) { const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); const int64_t idx_dim = hparams.indexer_head_size; @@ -581,6 +582,31 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( score = ggml_add(ctx0, score, inp->bias); } + // block-level top-k for the gather path: the scores are one-per-block anyway, so pick + // the winning blocks first and expand only those to cell indices via blk_cells. this + // sorts n_blocks instead of n_kv entries and skips every O(n_kv) op below. the cell + // count r*K_blk stays a multiple of 256, which keeps flash attention's padding happy. + // needs blk_bias: the per-block bias is what carries visibility into the selection. + if (gather && blk_bias && 256 % r == 0) { + const int64_t K_blk = std::min(n_blocks, + GGML_PAD(((int64_t) hparams.indexer_top_k + r - 1)/r, 256/r)); + + ggml_tensor * blk_idx = ggml_cont(ctx0, ggml_top_k(ctx0, score, K_blk)); // I32 [K_blk, n_tps, n_stream] + blk_idx = ggml_reshape_4d(ctx0, blk_idx, K_blk, 1, n_stream, 1); + + // blk_cells maps (block, slot-in-block) -> cell: view one row of r cells per block + ggml_tensor * bc = ggml_view_4d(ctx0, inp->blk_cells, r, n_blocks, 1, n_stream, + r*ggml_element_size(inp->blk_cells), + r*n_blocks*ggml_element_size(inp->blk_cells), + inp->blk_cells->nb[1], 0); + + ggml_tensor * cells = ggml_get_rows(ctx0, bc, blk_idx); // I32 [r, K_blk, 1, n_stream] + cells = ggml_reshape_4d(ctx0, cells, r*K_blk, n_tps, 1, n_stream); + cb(cells, "indexer_top_k_cells", il); + + return cells; + } + // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary ggml_tensor * expanded = ggml_get_rows(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); @@ -596,7 +622,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( cb(expanded, "indexer_score_tokens", il); // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail - const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); + // for the gather path the width is padded to a multiple of GGML_KQ_MASK_PAD so the + // gathered K/V/mask satisfy flash attention's padding requirements without extra ops; + // the gathered mask carries -INF for any surplus cells the padded top-k pulls in + const int64_t width = gather + ? std::min(n_kv, GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256)) + : std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); @@ -616,7 +647,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * v_cur, ggml_tensor * top_k, float kq_scale, - int il) { + int il, + bool gather) { // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer // has already scored with its own query in build_qsa_top_k, so top_k is unaffected. if (inp->self_k_rot) { @@ -648,6 +680,67 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * kq_mask = inp->get_kq_mask(); + // decode-time sparse gather: instead of masking the full cache (which makes flash + // attention scan all n_kv cells to use only indexer_top_k of them), copy the selected + // cells out of the cache and attend densely over just those. this bounds the per-token + // attention cost by the top-k width instead of the context depth. + // restricted to single-token-per-stream ubatches (decode), where every stream carries + // exactly one top-k list, so the gathered K/V stay one tensor per stream. + if (gather) { + ggml_tensor * kf = mctx_cur->get_k(ctx0, il); // [hd_k, n_head_kv, n_kv, ns] + ggml_tensor * vf = mctx_cur->get_v(ctx0, il); // [hd_v, n_head_kv, n_kv, ns] + + const int64_t hd_k = kf->ne[0]; + const int64_t hd_v = vf->ne[0]; + const int64_t n_h_kv = kf->ne[1]; + const int64_t n_kv = kf->ne[2]; + const int64_t ns = kf->ne[3]; + const int64_t n_topk = top_k->ne[0]; + + GGML_ASSERT(top_k->ne[1] == 1 && "QSA gather requires single-token-per-stream ubatches"); + GGML_ASSERT(kq_mask->ne[0] == n_kv); + + // the heads of a cell are contiguous in the cache, so a cell can be gathered as one row + GGML_ASSERT(kf->nb[2] == ggml_row_size(kf->type, hd_k*n_h_kv)); + GGML_ASSERT(vf->nb[2] == ggml_row_size(vf->type, hd_v*n_h_kv)); + + ggml_tensor * k_cells = ggml_view_4d(ctx0, kf, hd_k*n_h_kv, n_kv, 1, ns, + kf->nb[2], kf->nb[2]*n_kv, kf->nb[3], 0); + ggml_tensor * v_cells = ggml_view_4d(ctx0, vf, hd_v*n_h_kv, n_kv, 1, ns, + vf->nb[2], vf->nb[2]*n_kv, vf->nb[3], 0); + + // top_k [n_topk, 1, 1, ns] -> the index layout ggml_get_rows expects: [n_topk, 1, ns, 1] + ggml_tensor * idx = ggml_reshape_4d(ctx0, top_k, n_topk, 1, ns, 1); + + // get_rows dequantizes the cells to F32; build_attn_mha casts to F16 for flash attention + ggml_tensor * k_g = ggml_get_rows(ctx0, k_cells, idx); // F32 [hd_k*n_h_kv, n_topk, 1, ns] + ggml_tensor * v_g = ggml_get_rows(ctx0, v_cells, idx); // F32 [hd_v*n_h_kv, n_topk, 1, ns] + + k_g = ggml_reshape_4d(ctx0, k_g, hd_k, n_h_kv, n_topk, ns); + v_g = ggml_reshape_4d(ctx0, v_g, hd_v, n_h_kv, n_topk, ns); + cb(k_g, "qsa_k_gathered", il); + cb(v_g, "qsa_v_gathered", il); + + // gather the same cells' mask values: keeps -INF for any invalid cell the padded + // top-k width pulled in (e.g. when fewer than n_topk cells are visible) + ggml_tensor * m1 = ggml_view_4d(ctx0, kq_mask, 1, n_kv, 1, ns, + kq_mask->nb[0], kq_mask->nb[1], kq_mask->nb[3], 0); + ggml_tensor * m_g = ggml_get_rows(ctx0, m1, idx); // F32 [1, n_topk, 1, ns] + m_g = ggml_reshape_4d(ctx0, m_g, n_topk, 1, 1, ns); + m_g = ggml_cast(ctx0, m_g, GGML_TYPE_F16); // FA wants contiguous F16 + cb(m_g, "qsa_mask_gathered", il); + + ggml_tensor * cur = build_attn_mha(q_cur, k_g, v_g, nullptr, m_g, nullptr, nullptr, kq_scale, il); + cb(cur, "kqv_out", il); + + // the rotation is its own inverse, so undo it on the value side of the output + if (inp->self_v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); + } + + return cur; + } + // prepare new kq mask - starts filled with -INFINITY ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); @@ -702,7 +795,27 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( // indexer reads the same block input as q/k/v; no cache or no ratio means dense const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; - ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + // gather-based QSA decode: worth it once the cache is meaningfully deeper than the + // top-k width; below that the masked path costs about the same. QWEN4EXP_QSA_GATHER=0 + // disables it (A/B lever, and an escape hatch). + static const bool gather_enabled = [] { + const char * e = getenv("QWEN4EXP_QSA_GATHER"); + return e == nullptr || atoi(e) != 0; + }(); + + bool gather = false; + if (qsa && gather_enabled) { + const int64_t r = hparams.dsv4_compress_ratios[il]; + const int64_t n_kv = mctx_hyb->get_idx()->get_n_kv(); + const int64_t width = GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256); + + const int64_t n_stream = mctx_hyb->get_n_stream(); + + gather = n_tokens == n_stream && n_kv >= 2*width; + + } + + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il, gather) : nullptr; // Qwen3Next uses a single Q projection that outputs query + gate ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] @@ -755,7 +868,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; if (top_k) { - cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il); + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il, gather); } else { cur = build_attn(inp, nullptr, nullptr, nullptr, From aaa6297b14611f19642cb20fe9209874aeacf0d9 Mon Sep 17 00:00:00 2001 From: Abdel Darwish <168705979+abdel-darwish-27@users.noreply.github.com> Date: Tue, 1 Sep 2026 07:23:17 +1000 Subject: [PATCH 2/4] cuda: dispatch batched top-k to segmented sort instead of per-row DeviceTopK With CCCL >= 3.2 available, ggml_cuda_op_top_k ran cub::DeviceTopK::MaxPairs in a serial per-row loop. Decode-shaped calls (nrows == 1) are fine, but prompt processing hands this op hundreds of rows: profiling one 141K-token prefill of a qwen4exp model showed 903,702 DeviceTopK invocations (512-row batches x 12 QSA layers x ubatches), each launching 3-4 kernels. Keep DeviceTopK for nrows <= 4 and fall back to the segmented-sort path for larger batches until DeviceSegmentedTopK exists (NVIDIA/cccl#6391). --- ggml/src/ggml-cuda/top-k.cu | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd29333..1877eec68a0 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -36,7 +36,9 @@ static void top_k_cub(ggml_cuda_pool & pool, ncols, k, env)); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#endif // CUB_TOP_K_AVAILABLE + +#if defined(GGML_CUDA_USE_CUB) static int next_power_of_2(int x) { int n = 1; @@ -64,13 +66,19 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t k = dst->ne[0]; ggml_cuda_pool & pool = ctx.pool(); #ifdef CUB_TOP_K_AVAILABLE + // DeviceTopK selects rather than sorts, but only one row per call, so it wins for + // decode-shaped calls (a few rows) and loses badly to one segmented sort when a + // prompt batch brings hundreds of rows. // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented // https://github.com/NVIDIA/cccl/issues/6391 - // TODO: investigate if there exists a point where parallelized argsort is faster than sequential top-k - for (int i = 0; i < nrows; i++) { - top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); + if (nrows <= 4) { + for (int i = 0; i < nrows; i++) { + top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); + } + return; } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#endif // CUB_TOP_K_AVAILABLE +#if defined(GGML_CUDA_USE_CUB) // Fall back to argsort + copy const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); From ede17be6e3978fb5a41202f8bf4f022d7436a8e1 Mon Sep 17 00:00:00 2001 From: Abdel Darwish <168705979+abdel-darwish-27@users.noreply.github.com> Date: Tue, 1 Sep 2026 07:23:17 +1000 Subject: [PATCH 3/4] kv-cells: make for_each_token_in a bounded array scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit llama_kv_cache::get_prev_tokens (the n-gram/PLE predecessor lookup, called once per ubatch) iterated the `used` std::set — an RB-tree walk over every used cell — and then tested all LLAMA_MAX_SEQ (256) bits of each matching cell's sequence bitset. At 141K context this made a single call cost ~40 ms and consume 58% of total decode CPU time (perf: 50.9% llama_kv_cache::get_prev_tokens + 7.1% std::_Rb_tree_increment), leaving both GPUs ~15% utilized. Replace the tree walk with a contiguous scan of the pos/seq/ext arrays (same cells, same index order; empty cells have pos == -1) and hoist the queried seq ids out of the per-cell loop (queries carry a handful of ids, usually one). Decode at 141K context (Qwen3.8-Flash-Next UD-IQ4_XS, 2x RTX A6000): masked attention path: 9.0 -> 13.1 tok/s gather path: 11.6 -> 21.7 tok/s Outputs byte-identical before/after (greedy, fixed seed). This helps every model that uses get_prev_tokens (all PLE/n-gram architectures) at long context. --- src/llama-kv-cells.h | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 5167c037db3..8c072e94e68 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -314,18 +314,31 @@ class llama_kv_cells { // note: used by n-gram input embeddings to recover the tokens preceding a ubatch template void for_each_token_in(const std::bitset & seqs, llama_pos p0, llama_pos p1, F && f) const { - for (const auto & i : used) { - if (pos[i] < p0 || pos[i] >= p1) { - continue; + // scan the raw arrays instead of iterating `used`: walking that RB-tree touches + // every used cell through pointer chasing and dominates decode CPU time at long + // context (this runs once per ubatch). the array scan visits the same cells in + // the same (index) order with contiguous reads. empty cells have pos == -1. + const uint32_t n_cells = (uint32_t) pos.size(); + + // hoist the queried seq ids out of the loop: testing all LLAMA_MAX_SEQ bits for + // every cell costs n_cells*LLAMA_MAX_SEQ tests per call, which dominated decode + // CPU time at long context. queries carry a handful of seq ids (usually one). + llama_seq_id seq_ids[LLAMA_MAX_SEQ]; + uint32_t n_seq_ids = 0; + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (seqs.test(s)) { + seq_ids[n_seq_ids++] = s; } + } - const auto m = seq[i] & seqs; - if (m.none()) { + for (uint32_t i = 0; i < n_cells; ++i) { + if (pos[i] < 0 || pos[i] < p0 || pos[i] >= p1) { continue; } - for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { - if (m.test(s)) { + for (uint32_t k = 0; k < n_seq_ids; ++k) { + const llama_seq_id s = seq_ids[k]; + if (seq[i].test(s)) { f(s, pos[i], ext[i].tok); } } From e89368c241a7f6d6c0a38682048dbedc77d2bfea Mon Sep 17 00:00:00 2001 From: Abdel Darwish <168705979+abdel-darwish-27@users.noreply.github.com> Date: Tue, 1 Sep 2026 07:23:17 +1000 Subject: [PATCH 4/4] qwen4exp: compact per-cell mask for the gather path The gather path only reads one row of the attention kq_mask (to carry each selected cell's visibility into the gathered attention), but referencing it kept the whole FA-padded tensor alive: an O(n_kv x GGML_KQ_MASK_PAD) host fill plus an n_kv x 64 x 2-byte upload every decode step (~18 MB/token at 141K ctx, measured as the largest H2D stream during decode, with the staging copy attributed to the driver at ~14% of decode CPU). Add a compact F32 [n_kv, n_tps, n_stream] visibility row to the QSA input set, filled in set_input_qsa alongside the existing per-token pass and gathered in place of the kq_mask row. The attention kq_mask then goes unreferenced in gather graphs and is neither filled nor uploaded; llm_graph_input_mem_hybrid now skips it when unallocated, matching llm_graph_input_attn_kv. Upload drops 18 MB -> 1 MB per token. Outputs remain byte-identical to the masked path; mid-context needle retrieval passes at 68K and 141K. --- src/llama-graph.cpp | 6 +++++- src/llama-memory-hybrid-idx.cpp | 28 +++++++++++++++++++++++++++ src/llama-memory-hybrid-idx.h | 2 +- src/models/models.h | 1 + src/models/qwen4exp.cpp | 34 +++++++++++++++++++++++++++------ 5 files changed, 63 insertions(+), 8 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..bef74c1979d 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1087,7 +1087,11 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); - mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + // the mask can be left unallocated when the graph never attends over the full + // cache (e.g. the qwen4exp gather path reads its own compact per-cell mask) + if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) { + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + } if (inp_attn->self_k_rot) { mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot); diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 45c788e781a..98753e9ad9b 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -343,6 +343,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + ggml_tensor * mask_row, ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, @@ -435,6 +436,33 @@ void llama_memory_hybrid_idx_context::set_input_qsa( const llama_seq_id seq_id = ubatch->seq_id[i][0]; const llama_pos q = ubatch->pos[i]; + // per-cell visibility for the gather path: what row i of the attention + // kq_mask would say, minus SWA/alibi (the gather gate excludes those). + // unallocated (and skipped) when the graph takes the masked path. + if (mask_row != nullptr && mask_row->buffer != nullptr) { + float * dst_row = (float *) mask_row->data + i*n_kv; + + const bool is_2d = ubatch->is_pos_2d(); + + const llama_pos qx = is_2d ? ubatch->pos[i + ubatch->n_tokens*2] : 0; + const llama_pos qy = is_2d ? ubatch->pos[i + ubatch->n_tokens] : 0; + + for (int64_t j = 0; j < n_kv; ++j) { + float v = -INFINITY; + + if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { + v = 0.0f; + + // M-RoPE: image tokens can share the temporal position of the query + if (is_2d && cells.pos_get(j) == q && cells.ext_get(j).is_2d_gt(qx, qy)) { + v = -INFINITY; + } + } + + dst_row[j] = v; + } + } + // the tail is an incomplete block and is always visible, as in the reference const llama_pos tail_start = (q + 1)/r*r; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e3472646d0f..d5c88be823d 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -137,7 +137,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns] // the caller then adds the attention mask, the only part of the bias that varies within a block - void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, ggml_tensor * mask_row, ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; diff --git a/src/models/models.h b/src/models/models.h index 285e5af34bf..736288ffd86 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2317,6 +2317,7 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * mask_row, float kq_scale, int il, bool gather = false); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index e2ddeea7090..bbf0abff010 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -422,7 +422,7 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override { mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); + mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, mask_row, bias, ubatch, ratio, blk_bias); } bool can_reuse(const llm_graph_params & params) override { @@ -449,6 +449,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv); res &= bias->ne[1] == params.ubatch.n_tokens/n_stream; + res &= mask_row->ne[0] == n_kv; + res &= mask_row->ne[1] == params.ubatch.n_tokens/n_stream; + res &= mask_row->ne[2] == n_stream; + return res; } @@ -459,6 +463,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream] + // per-cell visibility for the gather path (0 or -INF); left unallocated by the + // masked path, which reads the attention kq_mask instead. F32 [n_kv, n_tps, n_stream] + ggml_tensor * mask_row = nullptr; + const llama_memory_hybrid_idx_context * mctx; const uint32_t ratio; @@ -512,7 +520,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream); + qsa->mask_row = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream); + ggml_set_input(qsa->mask_row); ggml_set_input(qsa->cell_blk); ggml_set_input(qsa->blk_cells); ggml_set_input(qsa->blk_pos); @@ -646,6 +656,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * mask_row, float kq_scale, int il, bool gather) { @@ -721,10 +732,15 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( cb(k_g, "qsa_k_gathered", il); cb(v_g, "qsa_v_gathered", il); - // gather the same cells' mask values: keeps -INF for any invalid cell the padded - // top-k width pulled in (e.g. when fewer than n_topk cells are visible) - ggml_tensor * m1 = ggml_view_4d(ctx0, kq_mask, 1, n_kv, 1, ns, - kq_mask->nb[0], kq_mask->nb[1], kq_mask->nb[3], 0); + // gather the same cells' visibility: keeps -INF for any invalid cell the padded + // top-k width pulled in (e.g. when fewer than n_topk cells are visible). + // reading the compact per-cell mask_row instead of the attention kq_mask leaves + // the latter unreferenced, so it is neither filled (O(n_kv * pad) host work) nor + // uploaded (n_kv * GGML_KQ_MASK_PAD * 2 bytes, ~18 MB/token at 141K ctx) + GGML_ASSERT(mask_row != nullptr); + GGML_ASSERT(mask_row->ne[0] == n_kv); + ggml_tensor * m1 = ggml_view_4d(ctx0, mask_row, 1, n_kv, 1, ns, + mask_row->nb[0], mask_row->nb[1], mask_row->nb[3], 0); ggml_tensor * m_g = ggml_get_rows(ctx0, m1, idx); // F32 [1, n_topk, 1, ns] m_g = ggml_reshape_4d(ctx0, m_g, n_topk, 1, 1, ns); m_g = ggml_cast(ctx0, m_g, GGML_TYPE_F16); // FA wants contiguous F16 @@ -868,7 +884,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; if (top_k) { - cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il, gather); + ggml_tensor * qsa_mask_row = nullptr; + if (gather) { + const auto it = qsa_inps.find((uint32_t) hparams.dsv4_compress_ratios[il]); + GGML_ASSERT(it != qsa_inps.end()); + qsa_mask_row = it->second->mask_row; + } + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, qsa_mask_row, kq_scale, il, gather); } else { cur = build_attn(inp, nullptr, nullptr, nullptr,