From 618450d07ca8c5dd2966988a43f887b058663f4c Mon Sep 17 00:00:00 2001 From: Pascal Date: Fri, 28 Aug 2026 05:56:32 +0200 Subject: [PATCH 1/6] qwen4exp: sum indexer heads by slices The indexer score reduction went through a permute, a cont and a sum_rows over the head axis. That left sum_rows with ne0 = 4, one block per row for a four element reduction, and the transpose copied the whole block by token surface twice on the way in. The heads are adjacent on ne[1], so each one is a strided view and the sum is a short chain of adds, the same shape as the pooling loop above it. Measured on RTX PRO 6000, UD-Q4_K_XL, fa on, per_layer_token_embd on CPU, llama-bench r=2: pp2048 @ d32768 2163 -> 2356 t/s pp2048 @ d65536 1497 -> 1666 t/s tg32 @ d32768 67 -> 69 t/s tg32 @ d65536 45 -> 48 t/s Greedy output is unchanged token for token on an 80k token retrieval probe. --- src/models/qwen4exp.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index abf6a0502fbf..154c82d40819 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -579,9 +579,17 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream)); score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream); score = ggml_relu(ctx0, score); - score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); - score = ggml_sum_rows(ctx0, score); - score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); + + // the heads sit side by side on ne[1] and there are few of them, so summing slices beats a + // transpose that would carry the whole block by token surface twice over + ggml_tensor * summed = nullptr; + for (int64_t h = 0; h < n_idx_h; ++h) { + ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream, + score->nb[2], score->nb[3], h*score->nb[1]); + summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice); + } + + score = summed; cb(score, "indexer_score", il); // one value per block, so it is cheaper to bias here than after the cells are expanded From 39efb6390e13fe52796314664f78a620bb3f952f Mon Sep 17 00:00:00 2001 From: Pascal Date: Fri, 28 Aug 2026 21:38:35 +0200 Subject: [PATCH 2/6] kv-cells: stop the sequence scan once all sequences are seen for_each_token_in() walked the full LLAMA_MAX_SEQ width for every used cell. get_prev_tokens(), its only caller, passes p0 = 0, so no cell is filtered out by position: the scan visits every cell of the cache to keep the few positions of the n-gram window. A cell carries a handful of sequences at most, so the scan stops once their count is reached. Behaviour is unchanged. The cost this removes is proportional to the number of used cells, so it shows up at long context and is invisible on short prompts. It applies to any n-gram model that resolves predecessors from the KV cells, not just qwen4exp. --- src/llama-kv-cells.h | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 5167c037db3e..9469b9a28a71 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -320,13 +320,15 @@ class llama_kv_cells { } const auto m = seq[i] & seqs; - if (m.none()) { - continue; - } - for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + // a cell carries a handful of sequences at most, so stop once they are all seen + // instead of walking the whole LLAMA_MAX_SEQ width + size_t left = m.count(); + + for (llama_seq_id s = 0; left > 0 && s < (llama_seq_id) LLAMA_MAX_SEQ; ++s) { if (m.test(s)) { f(s, pos[i], ext[i].tok); + --left; } } } From d7ca8ea59435faa675537ed3eb353c660a238b05 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sat, 29 Aug 2026 05:27:47 +0200 Subject: [PATCH 3/6] qwen4exp: attend the selected cells instead of masking the window QSA names about n_top_k cells per query, but build_attn_qsa turned that selection into a mask over the whole window and handed the full K and V to flash attention. On CUDA the only skip is flash_attn_mask_to_KV_max, a trailing cut, and the selection always keeps the current tail, so the cut never fires: the selection saved no work at all. The selected cells are now gathered into a window of their own, one per query, and the queries ride the stream axis of the attention so each carries the window its own selection named. The mask comes from gathering the attention mask at those cells, which leaves the same values the mask path would put there. The scan reads the window once for all the queries of a stream while the gather moves the selected cells twice, so the two meet at 2*n_tps*width == n_kv. Below that margin the scan still wins, which keeps prefill on the existing path; only decode and short chunks take the gather. Flash attention is required, since the gather reads the value side as rows. The win grows with depth, as the scan cost follows n_kv while the gather cost does not: nothing at 55k context, and generation goes from 52.05 to 55.55 t/s at 132k. Prefill is unchanged. Retrieval from a 132k token context is unaffected. --- src/models/models.h | 16 +++++++ src/models/qwen4exp.cpp | 95 ++++++++++++++++++++++++++++++++++++++--- 2 files changed, 104 insertions(+), 7 deletions(-) diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af9..7180051bf600 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2320,6 +2320,22 @@ struct llama_model_qwen4exp : public llama_model_base { float kq_scale, int il); + // attention over the whole window, masked down to the selection + ggml_tensor * build_qsa_scan( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + + // attention over the selected cells alone, gathered into a window per query + ggml_tensor * build_qsa_gather( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + // 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 std::map qsa_inps; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 154c82d40819..079a175cdc2e 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -662,6 +662,37 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } + const int64_t width = top_k->ne[0]; + const int64_t n_tps = top_k->ne[1]; + const int64_t n_kv = mctx_cur->get_n_kv(); + + // the scan reads the window once for all the queries of a stream, the gather moves the + // selected cells twice, so the two meet at 2*n_tps*width == n_kv. the margin below keeps + // the win clear and the windows small enough to stay in the compute buffer of a decode + // graph. flash attention keeps the value side as rows, which is what the gather reads + ggml_tensor * cur = cparams.flash_attn && 4*n_tps*width < n_kv + ? build_qsa_gather(inp, q_cur, top_k, kq_scale, il) + : build_qsa_scan (inp, q_cur, top_k, 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; +} + +// The window stays whole and the mask hides every cell the selection leaves out. +// The mask build below copies the MLA sparse path in llm_graph_context::build_attn. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_scan( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + const auto * mctx_cur = inp->mctx; + ggml_tensor * kq_mask = inp->get_kq_mask(); // prepare new kq mask - starts filled with -INFINITY @@ -694,15 +725,65 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il); - cb(cur, "kqv_out", il); + return build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, 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); - } +// The selected cells are gathered into a window of their own, one per query, so the key and +// value traffic follows the budget instead of the whole cache. The queries then ride the stream +// axis of the attention: each one carries the window its own selection named. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_gather( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + const auto * mctx_cur = inp->mctx; - return cur; + const int64_t width = top_k->ne[0]; + const int64_t n_tps = top_k->ne[1]; + const int64_t n_stream = top_k->ne[3]; + const int64_t n_q = n_tps*n_stream; + + ggml_tensor * k_all = mctx_cur->get_k(ctx0, il); + ggml_tensor * v_all = mctx_cur->get_v(ctx0, il); + + const int64_t n_kv = k_all->ne[2]; + + // a cell holds its heads back to back, so one row of the gather is one whole cell + ggml_tensor * k_cells = ggml_view_3d(ctx0, k_all, k_all->ne[0]*k_all->ne[1], n_kv, n_stream, + k_all->nb[2], k_all->nb[3], 0); + ggml_tensor * v_cells = ggml_view_3d(ctx0, v_all, v_all->ne[0]*v_all->ne[1], n_kv, n_stream, + v_all->nb[2], v_all->nb[3], 0); + + // a cell index names a cell of its own stream, and the ubatch lays the queries of a stream + // out contiguously, so the flat index of a window is the index of its query + ggml_tensor * idx_stream = ggml_reshape_2d(ctx0, top_k, width*n_tps, n_stream); + + ggml_tensor * k_sel = ggml_get_rows(ctx0, k_cells, idx_stream); + ggml_tensor * v_sel = ggml_get_rows(ctx0, v_cells, idx_stream); + + k_sel = ggml_reshape_4d(ctx0, k_sel, k_all->ne[0], k_all->ne[1], width, n_q); + v_sel = ggml_reshape_4d(ctx0, v_sel, v_all->ne[0], v_all->ne[1], width, n_q); + cb(k_sel, "qsa_k_sel", il); + cb(v_sel, "qsa_v_sel", il); + + // gathering the attention mask at the selected cells leaves the same values the scan path + // would put there, so the window carries the reach of its query + ggml_tensor * kq_mask = inp->get_kq_mask(); + + GGML_ASSERT(kq_mask->nb[3] == kq_mask->nb[1]*n_tps); + + ggml_tensor * mask_cells = ggml_view_3d(ctx0, kq_mask, 1, n_kv, n_q, + kq_mask->nb[0], kq_mask->nb[1], 0); + + ggml_tensor * idx_query = ggml_reshape_3d(ctx0, top_k, width, n_q, 1); + + ggml_tensor * mask = ggml_get_rows(ctx0, mask_cells, idx_query); + + mask = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mask, width, 1, 1, n_q), GGML_TYPE_F16); + cb(mask, "qsa_mask_sel", il); + + return build_attn_mha(q_cur, k_sel, v_sel, nullptr, mask, nullptr, nullptr, kq_scale, il); } ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( From 28e37926cecff7e9c8a7c2011fb98311458b3daa Mon Sep 17 00:00:00 2001 From: Pascal Date: Sat, 29 Aug 2026 08:16:06 +0200 Subject: [PATCH 4/6] kv-cache: scan only the n-gram window in get_prev_tokens The scan ran over the whole position range, from 0 up to p_max, so every cell of the cache passed the position filter and paid the sequence mask work. The full range was only needed for below[], which holds the nearest token before the window and is read solely when a lookup finds nothing in [w0, p], that is on an M-RoPE gap. The main scan now covers the window alone, and the below[] pass is deferred to the first lookup that falls through. Contiguous positions never trigger it, so text decoding pays one windowed pass instead of a full one. Measured on Qwen3.8-Flash-Next at 55k context on an RTX PRO 6000, warm runs with the first one discarded: generation 74.4 -> 76.3 t/s Prompt processing is unaffected, since the scan is amortised over the whole ubatch there. What is left of this call is the traversal of the used-cell set itself, which the range restriction cannot avoid. --- src/llama-kv-cache.cpp | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 8fafcd15304e..33a4fc24c6db 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1868,16 +1868,33 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st for (uint32_t s = 0; s < n_stream; ++s) { // p_max inclusive: an embd token looks up cells at its own (shared) position - v_cells[s].for_each_token_in(seqs, 0, p_max + 1, + v_cells[s].for_each_token_in(seqs, w0, p_max + 1, [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) { - if (pos >= w0) { - hist[key(seq_id, pos)] = tok; - } else if (pos > below[seq_id].first) { - below[seq_id] = { pos, tok }; - } + hist[key(seq_id, pos)] = tok; }); } + // below[] only answers when an M-RoPE gap leaves the window empty, so it costs a pass of + // its own that contiguous positions never pay + bool below_ready = false; + + const auto ensure_below = [&]() { + if (below_ready) { + return; + } + + below_ready = true; + + for (uint32_t s = 0; s < n_stream; ++s) { + v_cells[s].for_each_token_in(seqs, 0, w0, + [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) { + if (pos > below[seq_id].first) { + below[seq_id] = { pos, tok }; + } + }); + } + }; + // the token at pos p, or the nearest earlier one when p falls in an M-RoPE gap const auto lookup = [&](llama_seq_id seq_id, llama_pos p) -> llama_token { for (llama_pos q = p; q >= w0; --q) { @@ -1886,6 +1903,9 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st return it->second; } } + + ensure_below(); + return below[seq_id].second; }; From b88668b47d0631768f6c4125369d9d832cb8d559 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sat, 29 Aug 2026 09:10:03 +0200 Subject: [PATCH 5/6] kv-cells: keep the used-cell set as a bitmap used was a std::set, so a full pass over the used cells chased one tree node per index with no locality. The set is dense, a large fraction of the cache, which makes a bitmap both smaller and far cheaper to walk: one word per 64 cells and a trailing-zero count per set bit. llama_kv_idx_set exposes the same operations the class relied on, and visits indices in ascending order like the set did. insert, erase and size stay O(1); first and last scan the words and are called once per ubatch. Measured on Qwen3.8-Flash-Next at 55k context on an RTX PRO 6000, warm runs with the first one discarded: generation 76.2 -> 77.4 t/s The gain scales with the number of used cells, so it grows with context and is invisible on short ones. This touches the KV cache for every model, though only the n-gram path walks the set often enough to notice. --- src/llama-kv-cells.h | 117 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 111 insertions(+), 6 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 9469b9a28a71..ad1c91b72766 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -10,6 +10,10 @@ #include #include +#if defined(_MSC_VER) +#include +#endif + struct llama_kv_cell_ext { // 2D spatial positions, typically used for M-RoPE llama_pos x = 0; @@ -31,6 +35,106 @@ struct llama_kv_cell_ext { } }; +// index of the lowest set bit, and of the highest; the word is never zero here +#if defined(_MSC_VER) +static inline uint32_t llama_kv_ctz64(uint64_t x) { + unsigned long r; + _BitScanForward64(&r, x); + return (uint32_t) r; +} + +static inline uint32_t llama_kv_clz64(uint64_t x) { + unsigned long r; + _BitScanReverse64(&r, x); + return (uint32_t) r; +} +#else +static inline uint32_t llama_kv_ctz64(uint64_t x) { + return (uint32_t) __builtin_ctzll(x); +} + +static inline uint32_t llama_kv_clz64(uint64_t x) { + return (uint32_t) (63 - __builtin_clzll(x)); +} +#endif + +// a dense set of cell indices, kept as a bitmap so that a full pass reads words instead of +// chasing tree nodes. the cells it tracks are a large fraction of the cache, so the bitmap is +// both smaller and faster to walk than a node per index +class llama_kv_idx_set { +public: + void resize(uint32_t n) { + bits.assign((n + 63)/64, 0); + n_set = 0; + } + + void clear() { + std::fill(bits.begin(), bits.end(), 0); + n_set = 0; + } + + void insert(uint32_t i) { + uint64_t & w = bits[i/64]; + const uint64_t b = 1ull << (i%64); + + n_set += (w & b) == 0; + w |= b; + } + + void erase(uint32_t i) { + uint64_t & w = bits[i/64]; + const uint64_t b = 1ull << (i%64); + + n_set -= (w & b) != 0; + w &= ~b; + } + + bool contains(uint32_t i) const { + return (bits[i/64] >> (i%64)) & 1; + } + + uint32_t size() const { return n_set; } + bool empty() const { return n_set == 0; } + + uint32_t first() const { + for (size_t w = 0; w < bits.size(); ++w) { + if (bits[w]) { + return 64*w + llama_kv_ctz64(bits[w]); + } + } + + return 0; + } + + uint32_t last() const { + for (size_t w = bits.size(); w-- > 0; ) { + if (bits[w]) { + return 64*w + llama_kv_clz64(bits[w]); + } + } + + return 0; + } + + // visits the indices in ascending order + template + void for_each(F && f) const { + for (size_t w = 0; w < bits.size(); ++w) { + uint64_t m = bits[w]; + + while (m) { + f((uint32_t) (64*w + llama_kv_ctz64(m))); + m &= m - 1; + } + } + } + +private: + std::vector bits; + + uint32_t n_set = 0; +}; + // meta information about KV cells that can be part of multiple sequences at the same time // TODO: add unit tests class llama_kv_cells { @@ -69,6 +173,7 @@ class llama_kv_cells { ext.resize(n); shift.resize(n); seq.resize(n); + used.resize(n); reset(); } @@ -87,13 +192,13 @@ class llama_kv_cells { // the index of the first cell that is used // return 0 if no cells are used uint32_t used_min() const { - return used.empty() ? 0 : *used.begin(); + return used.empty() ? 0 : used.first(); } // the index of the last cell that is used + 1 // return 0 if no cells are used uint32_t used_max_p1() const { - return used.empty() ? 0 : *used.rbegin() + 1; + return used.empty() ? 0 : used.last() + 1; } bool get_has_shift() const { @@ -314,9 +419,9 @@ 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) { + used.for_each([&](uint32_t i) { if (pos[i] < p0 || pos[i] >= p1) { - continue; + return; } const auto m = seq[i] & seqs; @@ -331,7 +436,7 @@ class llama_kv_cells { --left; } } - } + }); } // note: call only if the cell is not empty and the seq_id is not in the cell @@ -488,7 +593,7 @@ class llama_kv_cells { bool has_shift = false; // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id) - std::set used; + llama_kv_idx_set used; std::vector pos; From db40b22d69f098feb3eb583e1386dd33e63703e6 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sat, 29 Aug 2026 21:59:01 +0200 Subject: [PATCH 6/6] nits --- src/llama-kv-cells.h | 7 +++---- src/models/qwen4exp.cpp | 11 ++++------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index ad1c91b72766..2914d35f5100 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -58,9 +58,8 @@ static inline uint32_t llama_kv_clz64(uint64_t x) { } #endif -// a dense set of cell indices, kept as a bitmap so that a full pass reads words instead of -// chasing tree nodes. the cells it tracks are a large fraction of the cache, so the bitmap is -// both smaller and faster to walk than a node per index +// the cells tracked are a large fraction of the cache, so a bitmap is both smaller and faster +// to walk than a tree node per index class llama_kv_idx_set { public: void resize(uint32_t n) { @@ -116,7 +115,7 @@ class llama_kv_idx_set { return 0; } - // visits the indices in ascending order + // ascending order, as the set it replaces guaranteed template void for_each(F && f) const { for (size_t w = 0; w < bits.size(); ++w) { diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 079a175cdc2e..f6dcdcce0703 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -666,10 +666,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( const int64_t n_tps = top_k->ne[1]; const int64_t n_kv = mctx_cur->get_n_kv(); - // the scan reads the window once for all the queries of a stream, the gather moves the - // selected cells twice, so the two meet at 2*n_tps*width == n_kv. the margin below keeps - // the win clear and the windows small enough to stay in the compute buffer of a decode - // graph. flash attention keeps the value side as rows, which is what the gather reads + // the two costs meet at 2*n_tps*width == n_kv; the margin keeps the windows small enough + // for the compute buffer of a decode graph. flash attention is what keeps the values as rows ggml_tensor * cur = cparams.flash_attn && 4*n_tps*width < n_kv ? build_qsa_gather(inp, q_cur, top_k, kq_scale, il) : build_qsa_scan (inp, q_cur, top_k, kq_scale, il); @@ -728,9 +726,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_scan( return build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il); } -// The selected cells are gathered into a window of their own, one per query, so the key and -// value traffic follows the budget instead of the whole cache. The queries then ride the stream -// axis of the attention: each one carries the window its own selection named. +// Key and value traffic follows the budget instead of the whole cache. The queries ride the +// stream axis of the attention, each carrying the window its own selection named. ggml_tensor * llama_model_qwen4exp::graph::build_qsa_gather( llm_graph_input_attn_kv * inp, ggml_tensor * q_cur,