From e74cec36ede841d35df95dda62e8071dcf5378d4 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/2] 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 d4e59d77e570..45c788e781ab 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 9b87a40d5af9..2a493583b0b0 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 abf6a0502fbf..81fd71812ea1 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -480,7 +480,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; @@ -589,6 +590,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); @@ -604,7 +630,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)); @@ -624,7 +655,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) { @@ -656,6 +688,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); @@ -710,7 +803,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 ] @@ -763,7 +876,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 3bda5d023ac375a7047ccd547b098d9f58cc3d74 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/2] 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 72db486cacea..1f8adc04c580 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 45c788e781ab..98753e9ad9b6 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 e3472646d0f6..d5c88be823d7 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 2a493583b0b0..f69681d5d1dd 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 81fd71812ea1..e6a702c8a4c5 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -430,7 +430,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 { @@ -457,6 +457,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; } @@ -467,6 +471,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; @@ -520,7 +528,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); @@ -654,6 +664,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) { @@ -729,10 +740,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 @@ -876,7 +892,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,