From 4d4016f333b08924198182e9c1e35212d9b77c96 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH 1/6] qwen4exp: share the QSA input set per compress ratio build_qsa_top_k built a fresh llm_graph_input_qsa on every call, and it is called once per QSA layer -- twelve input sets per ubatch holding byte-identical data. set_input_qsa is an O(n_kv) per-cell scan (the file's own TODO measures it at 865 us at 33k context), so that was twelve scans where one would do: about 10 ms per step at 33k and 41 ms at 131k, all on the CPU and invisible to a GPU profiler. Key the inputs by compress ratio and reuse them across layers. The ratio is fixed per layer and the resolved layout depends on nothing else, so layers that share a ratio can share the whole set. This is not new work: upstream ggml-org has had it since the commit that added qwen4exp, and the code here is that code. This fork's independent port did not carry it, so it is picked back out and applied on its own. (cherry picked from commit 6c84c7d5d8833c6e0df69628f75a0f599797934e, "model: add Qwen3.8-Flash-Next (qwen4exp)", ggml-org/llama.cpp#27742) Measured on Strix Halo (gfx1151, Vulkan), Qwen3.8-Flash-Next UD-IQ4_XS, -ctk f16 -ctv f16, MTP n_max=3, ctx 262144, llama-server with prompt cache reuse, temp=0, n_predict=160: depth before after 2048 33.69 27.55 8192 47.29 48.10 32768 35.22 39.67 131072 25.05 28.19 Output is bit-identical at every depth. The swing at 2048 is speculative decoding landing on a different accept/reject path, not a regression. --- src/models/models.h | 6 ++++++ src/models/qwen4exp.cpp | 43 ++++++++++++++++++++++++++--------------- 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/src/models/models.h b/src/models/models.h index 3833065941e4..380860b6e22a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2316,6 +2316,8 @@ struct llama_model_qwen35 : public llama_model_base { struct llama_model_qwen4exp : public llama_model_base { llama_model_qwen4exp(const struct llama_model_params & params) : llama_model_base(params) {} + class llm_graph_input_qsa; + void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; @@ -2370,6 +2372,10 @@ struct llama_model_qwen4exp : public llama_model_base { 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; + // QSA: token indices this layer's queries may attend to, or nullptr for dense ggml_tensor * build_qsa_top_k( const llama_memory_hybrid_idx_context * mctx_hyb, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index eca27d604d3a..daa201732801 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -510,7 +510,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( // QSA attends to a budget of whole blocks of compress_ratio tokens, each scored by one // mean-pooled indexer key, plus the incomplete tail. set_input resolves the cache layout. -class llm_graph_input_qsa : public llm_graph_input_i { +class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { public: llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) : mctx(mctx), ratio(ratio), blk_bias(blk_bias) {} @@ -543,7 +543,7 @@ class llm_graph_input_qsa : public llm_graph_input_i { // it turns on the kq_mask matching those same three, and causal_attn / use_alibi are fixed for the // context. n_kv is padded, so this holds between padding steps. set_input still runs on the reuse // path, so only topology is certified here. -bool llm_graph_input_qsa::can_reuse(const llm_graph_params & params) { +bool llama_model_qwen4exp::llm_graph_input_qsa::can_reuse(const llm_graph_params & params) { const auto * m = static_cast(params.mctx); this->mctx = m; @@ -607,21 +607,32 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && cparams.causal_attn && !hparams.use_alibi; - auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + // nothing above depends on the layer, so the layers sharing a ratio share one input set. + // set_input_qsa is an O(n_kv) per-cell scan (about 865 us at 33k context) and every QSA + // layer was paying it for byte-identical data: 12 scans per ubatch instead of one. + llm_graph_input_qsa * inp = nullptr; - qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); - qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); - 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); - - ggml_set_input(qsa->cell_blk); - ggml_set_input(qsa->blk_cells); - ggml_set_input(qsa->blk_pos); - ggml_set_input(qsa->bias); - - llm_graph_input_qsa * inp = qsa.get(); - res->add_input(std::move(qsa)); + const auto it = qsa_inps.find((uint32_t) r); + if (it != qsa_inps.end()) { + inp = it->second; + } else { + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + + qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); + 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); + + ggml_set_input(qsa->cell_blk); + ggml_set_input(qsa->blk_cells); + ggml_set_input(qsa->blk_pos); + ggml_set_input(qsa->bias); + + inp = qsa.get(); + res->add_input(std::move(qsa)); + qsa_inps.emplace((uint32_t) r, inp); + } // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); From f347860db68eb0b493f29ec6568e8b147d165256 Mon Sep 17 00:00:00 2001 From: Masahito Suzuki Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH 2/6] qwen4exp: drop the ggml_cont on the QSA block-mean slices The block mean cut `members` into r strided views and materialised each one before adding. ggml_add has no contiguity requirement -- both the Vulkan and CUDA backends gate it on type alone and their kernels index through the nb strides -- and ggml_dup_tensor already gives the sum a contiguous home, so no cont is needed for the accumulator either. Removing it drops r reads plus r writes of [idx_dim, n_blocks] f32 per layer per ubatch: about 34 MB at 33k context over 12 layers. In a profile the CONT ops were +2.07 ms (15%) of the graph-time increment from depth 2048 to 32768. The addition order is unchanged, so the arithmetic is identical. Measured on Strix Halo (gfx1151, Vulkan), ctx 262144, f16 KV, MTP n_max=3: 131072 goes 25.05 -> 28.26 t/s (+12.8%) against the base, output bit-identical. On a CUDA host with far more bandwidth (2x RTX 3090, IQ1_S, ctx 154624) the same patch is worth +2.7% (11.97 -> 12.29 at 131072). What it removes is memory traffic, so the gain tracks how scarce bandwidth is. --- src/models/qwen4exp.cpp | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index daa201732801..25794aeba1f8 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -649,12 +649,16 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); - // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows + // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows. + // the slices are strided views of members. ggml_add has no contiguity requirement (the Vulkan + // backend gates it on type alone) and ggml_dup_tensor gives the sum a contiguous home, so the + // per-slice ggml_cont was materialising the whole of members a second time for nothing: + // r reads plus r writes of [idx_dim, n_blocks] f32 per layer per ubatch, about 34 MB at 33k + // context and 12 layers. the addition order is unchanged, so the arithmetic is identical. ggml_tensor * pooled = nullptr; for (int64_t i = 0; i < r; ++i) { - ggml_tensor * slice = ggml_cont(ctx0, - ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, - members->nb[2], members->nb[3], i*members->nb[1])); + ggml_tensor * slice = ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, + members->nb[2], members->nb[3], i*members->nb[1]); pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; } pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); From 9081e592df90e9eee9f51e589e4ffd551e96f8a9 Mon Sep 17 00:00:00 2001 From: Masahito Suzuki Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH 3/6] qwen4exp: cache the pooled QSA indexer keys, recompute only the tail A full block never changes again: its cells are written once, and the pooling, normalisation and rotation that turn them into an indexer key are all determined by position. build_qsa_top_k rebuilt every block on every ubatch anyway. At 131k over 12 layers that is GET_ROWS 7.8 ms + adds 13.8 ms + RMS_NORM_MUL 4.1 ms + ROPE 3.0 ms, about a quarter of a 97 ms decode step. Add a cache holding one row per block. The graph recomputes the last n_recomp blocks, writes them back with ggml_set_rows, and the score matmul reads the whole cache. n_recomp is n_blocks while the cache is invalid, so the incremental and full paths are the same code. Four things this has to get right, each of which cost a real bug: - Stored at F32. Pooling is per block with no reduction across rows, so a block computed in a window of 66 gives bit for bit what it gave when all 32768 were computed inline. Storing at full width keeps that true end to end, which turns the correctness check from "the output is self consistent" into "the output matches the build without the cache". f16 would halve the score matmul's read -- 0.6 ms of 97 -- and is not worth giving up an exact reference for. - Single-stream caches only. Rows are addressed by block index with no stream offset. The test is on the cache, not the ubatch: a slot of a multi-stream cache also sees n_stream == 1, but its rows start at sinfo.s0. - The block table is host scratch, not a graph tensor. ggml-alloc gives data only to tensors some node reads. Once the graph reads the recompute window instead of the whole table, blk_cells and blk_pos are orphans with data == nullptr, and set_input writes through a null pointer. - Invalidation is a position watermark, not a flag. Blocks are cut on the position line, so an operation that only touches positions above some point leaves everything below it correct. Speculative decoding drops its rejected tail with seq_rm on every single step; treating that as "forget everything" makes the cache recompute the whole table each time -- correct output, no speedup, and nothing in an end-to-end timing says why. state_read invalidates too. The server's prompt cache restores a slot through that path and never goes through seq_rm; without the hook the cache keeps scoring the blocks of whatever prompt ran before. Measured on Strix Halo (gfx1151, Vulkan), Qwen3.8-Flash-Next UD-IQ4_XS, -ctk f16 -ctv f16, MTP n_max=3, ctx 262144: depth before after 2048 27.55 27.26 8192 49.41 49.70 32768 41.00 42.82 131072 32.46 41.59 Cumulative with the previous patches: 25.05 -> 41.59 t/s at 131072, +66% against the base. Output is bit-identical to the build without the cache at every depth, and via every path: ascending, descending (which shrinks the cache through seq_rm), a changed prefix (which goes through the server's prompt cache and its restore path), and back. --- src/llama-memory-hybrid-idx.cpp | 183 +++++++++++++++++++++++++++++++- src/llama-memory-hybrid-idx.h | 46 +++++++- src/models/qwen4exp.cpp | 108 +++++++++++++++---- 3 files changed, 311 insertions(+), 26 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index f7e99c761dc0..f94af9f11b2b 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -11,6 +11,8 @@ #include #include #include +#include + // // llama_memory_hybrid_idx @@ -61,6 +63,49 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( model, hparams_idx, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); + }()), + hparams_pool(model.hparams), + mem_pool(filter_idx == nullptr ? nullptr : [&] () -> llama_kv_cache * { + // the smallest compress ratio gives the most blocks, so size the cache for that + uint32_t r_min = 0; + for (uint32_t il = 0; il < model.hparams.n_layer(); ++il) { + const uint32_t r = model.hparams.dsv4_compress_ratios[il]; + if (r > 0 && (r_min == 0 || r < r_min)) { + r_min = r; + } + } + + if (r_min == 0) { + return nullptr; + } + + // one row per block, plus a spare that writes for unused recompute slots are sent to + const uint32_t n_blocks_max = (kv_size + r_min - 1)/r_min + 1; + + std::fill(hparams_pool.n_head_kv_arr.begin(), hparams_pool.n_head_kv_arr.end(), 1); + hparams_pool.n_embd_head_k_full = model.hparams.indexer_head_size; + + // nothing reads V here; make it the same width as K rather than the model's, so the + // allocation the cache makes for it is not several hundred megabytes of dead weight + hparams_pool.n_embd_head_v_full = model.hparams.indexer_head_size; + + // the rows are already rotated when they are written, so a K-shift must not touch them + hparams_pool.rope_type = LLAMA_ROPE_TYPE_NONE; + + LLAMA_LOG_INFO("%s: creating QSA pooled-key cache, size = %u blocks (ratio %u)\n", + __func__, n_blocks_max, r_min); + + // F32, not type_k. Pooling, normalisation and rotation are all per block with no + // reduction across rows, so a block computed in a window of 66 gives bit for bit what + // the same block gave when every block was computed inline. Storing the result at full + // width keeps that true end to end, which turns the correctness check for this cache + // from "the output is self consistent" into "the output matches the build without it". + // f16 would halve the score matmul's read of the pooled keys, worth about 0.6 ms of a + // 97 ms decode step at 131k context -- not worth giving up an exact reference for. + return new llama_kv_cache( + model, hparams_pool, GGML_TYPE_F32, GGML_TYPE_F32, v_trans, offload, unified, + n_blocks_max, n_seq_max, 1, 0, LLAMA_SWA_TYPE_NONE, + nullptr, filter_idx, nullptr, nullptr, "pool_"); }()) {} llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { @@ -139,6 +184,8 @@ llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lc } void llama_memory_hybrid_idx::clear(bool data) { + qsa_pool_invalidate(); + llama_memory_hybrid::clear(data); if (mem_idx) { @@ -153,6 +200,8 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po return false; } + qsa_pool_invalidate_from(p0); + if (mem_idx) { mem_idx->seq_rm(seq_id, p0, p1); } @@ -161,6 +210,8 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po } void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + qsa_pool_invalidate(); + llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1); if (mem_idx) { @@ -169,6 +220,8 @@ void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_i } void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { + qsa_pool_invalidate(); + llama_memory_hybrid::seq_keep(seq_id); if (mem_idx) { @@ -177,6 +230,8 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + qsa_pool_invalidate_from(p0); + llama_memory_hybrid::seq_add(seq_id, p0, p1, shift); if (mem_idx) { @@ -185,6 +240,8 @@ void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_p } void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + qsa_pool_invalidate(); + llama_memory_hybrid::seq_div(seq_id, p0, p1, d); if (mem_idx) { @@ -219,6 +276,11 @@ void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id se } void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // A restore replaces the cells wholesale and the pooled keys are derived from them. This is + // how the server's prompt cache brings a slot back, and it never goes through seq_rm: without + // this the pooled cache would keep scoring the blocks of whatever prompt ran before. + qsa_pool_invalidate(); + llama_memory_hybrid::state_read(io, seq_id, flags); // [TAG_HYBRID_IDX_STATE] must mirror the write order above. @@ -232,6 +294,46 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_ } +llama_kv_cache * llama_memory_hybrid_idx::get_mem_pool() const { + return mem_pool.get(); +} + +void llama_memory_hybrid_idx::qsa_pool_invalidate() const { + pool_valid_pos = 0; +} + +void llama_memory_hybrid_idx::qsa_pool_invalidate_from(llama_pos p0) const { + const uint32_t p = p0 < 0 ? 0 : (uint32_t) p0; + + if (p < pool_valid_pos) { + pool_valid_pos = p; + } +} + +void llama_memory_hybrid_idx::qsa_pool_validate(uint32_t n_pos) const { + pool_valid_pos = n_pos; +} + +uint32_t llama_memory_hybrid_idx::qsa_pool_n_recomp( + uint32_t ratio, uint32_t n_tokens, uint32_t n_kv, uint32_t n_pad_kv) const { + GGML_ASSERT(ratio > 0); + + const uint32_t n_blocks = (n_kv + ratio - 1)/ratio; + + if (mem_pool == nullptr) { + return n_blocks; + } + + // whatever sits above the watermark has to be rebuilt, and so does the tail this ubatch + // can reach: the blocks its own tokens land in, plus the blocks n_kv gains when it next + // grows by a padding step. Below both, nothing has moved. + const uint32_t stale = n_blocks - std::min(pool_valid_pos/ratio, n_blocks); + const uint32_t own = (n_tokens + ratio - 1)/ratio + 1; + const uint32_t grow = (n_pad_kv + ratio - 1)/ratio; + + return std::min(n_blocks, std::max(stale, own + grow)); +} + llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const { return mem_idx.get(); } @@ -241,6 +343,9 @@ void llama_memory_hybrid_idx::set_input_qsa( ggml_tensor * blk_cells, ggml_tensor * blk_pos, ggml_tensor * bias, + ggml_tensor * pool_idxs, + ggml_tensor * pool_cells, + ggml_tensor * pool_pos, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const { @@ -251,7 +356,9 @@ void llama_memory_hybrid_idx::set_input_qsa( const int64_t n_kv = cell_blk->ne[0]; const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch - const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns); + // not from blk_pos: with the pooled-key cache the graph reads the window, not the whole + // block table, so blk_cells and blk_pos are not graph tensors at all + const int64_t n_blocks = (n_kv + (int64_t) ratio - 1)/(int64_t) ratio; const int64_t n_tokens = ubatch->n_tokens; const int64_t r = ratio; @@ -259,10 +366,30 @@ void llama_memory_hybrid_idx::set_input_qsa( 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_blk_cells = (int32_t *) blk_cells->data; - int32_t * dst_blk_pos = (int32_t *) blk_pos->data; float * dst_bias = (float *) bias->data; + // The block table is an intermediate: cell_blk and bias are what the dense path feeds to + // the graph, and with the pooled-key cache only the window taken from the table is. An + // input tensor no node reads is never allocated by ggml-alloc, so when the graph has no + // use for these they are plain host scratch instead. + std::vector blk_cells_host; + std::vector blk_pos_host; + + int32_t * dst_blk_cells; + int32_t * dst_blk_pos; + + if (blk_cells != nullptr) { + GGML_ASSERT(blk_pos != nullptr); + dst_blk_cells = (int32_t *) blk_cells->data; + dst_blk_pos = (int32_t *) blk_pos->data; + } else { + GGML_ASSERT(blk_pos == nullptr); + blk_cells_host.resize((size_t) r*n_blocks*n_ns); + blk_pos_host .resize((size_t) 4*n_blocks*n_ns); + dst_blk_cells = blk_cells_host.data(); + dst_blk_pos = blk_pos_host .data(); + } + // a block is keyed on (sequence set, index bucket): a unified cache counts every sequence // from zero, so the bucket alone would pool two sequences into one block GGML_ASSERT(r <= 64); @@ -552,6 +679,40 @@ void llama_memory_hybrid_idx::set_input_qsa( } } } + + // the pooled-key window: the last n_recomp blocks. Everything below the window belongs to a + // full block whose cells were written once and never moved, so the cache still holds the + // right key for it -- pooling, normalisation and rotation are all position-determined. + if (pool_cells != nullptr) { + const int64_t n_recomp = pool_cells->ne[0]/r; + + GGML_ASSERT(n_ns == 1 && "the pooled-key path is single stream; the graph falls back otherwise"); + GGML_ASSERT(n_recomp > 0 && n_recomp <= n_blocks); + GGML_ASSERT(pool_idxs->ne[0] == n_recomp); + GGML_ASSERT(pool_pos->ne[0] == 4*n_recomp); + + int64_t * dst_pool_idxs = (int64_t *) pool_idxs->data; + int32_t * dst_pool_cells = (int32_t *) pool_cells->data; + int32_t * dst_pool_pos = (int32_t *) pool_pos->data; + + const int64_t b0 = n_blocks - n_recomp; + + for (int64_t i = 0; i < n_recomp; ++i) { + const int64_t b = b0 + i; + + dst_pool_idxs[i] = b; + + for (int64_t k = 0; k < r; ++k) { + dst_pool_cells[i*r + k] = dst_blk_cells[b*r + k]; + } + + for (int64_t sec = 0; sec < 4; ++sec) { + dst_pool_pos[sec*n_recomp + i] = dst_blk_pos[sec*n_blocks + b]; + } + } + + qsa_pool_validate((uint32_t) n_kv); + } } // @@ -641,10 +802,24 @@ void llama_memory_hybrid_idx_context::set_input_qsa( ggml_tensor * blk_cells, ggml_tensor * blk_pos, ggml_tensor * bias, + ggml_tensor * pool_idxs, + ggml_tensor * pool_cells, + ggml_tensor * pool_pos, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const { GGML_ASSERT(mem != nullptr); - mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); + mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, + pool_idxs, pool_cells, pool_pos, ubatch, ratio, blk_bias); +} + +llama_kv_cache * llama_memory_hybrid_idx_context::get_mem_pool() const { + return mem == nullptr ? nullptr : mem->get_mem_pool(); +} + +uint32_t llama_memory_hybrid_idx_context::qsa_pool_n_recomp( + uint32_t ratio, uint32_t n_tokens, uint32_t n_kv, uint32_t n_pad_kv) const { + GGML_ASSERT(mem != nullptr); + return mem->qsa_pool_n_recomp(ratio, n_tokens, n_kv, n_pad_kv); } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 7bf3a320059b..8f6520c1e9a4 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -77,6 +77,26 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer + // QSA pooled-key cache: one row per block, holding the mean-pooled, RMS-normalised and + // rotated indexer key that build_qsa_top_k scores against. A block that is full never + // changes again -- its cells are written once and pooling, normalisation and rotation are + // all position-determined -- so only the tail of the cache is recomputed per ubatch. + // Anything that moves cells (a shift, a removal, a copy, a state load) drops the lot. + llama_kv_cache * get_mem_pool() const; // nullptr when the model carries no indexer + + // how many trailing blocks the next graph must recompute. n_blocks while the cache is + // invalid, otherwise the tail this ubatch can touch: the blocks its own tokens land in, + // plus the blocks n_kv gains when it next grows by a padding step. + uint32_t qsa_pool_n_recomp(uint32_t ratio, uint32_t n_tokens, uint32_t n_kv, uint32_t n_pad_kv) const; + + // Blocks are cut on the position line, so what a cache-disturbing operation costs is a + // watermark, not a flag: everything below the first position it touches is still right. + // Speculative decoding drops its rejected tail with seq_rm on every single step, and + // treating that as "forget everything" made the cache recompute the whole table each time. + void qsa_pool_invalidate() const; // forget the lot + void qsa_pool_invalidate_from(llama_pos p0) const; // forget positions >= p0 + void qsa_pool_validate(uint32_t n_pos) const; // positions < n_pos are now pooled + // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache. // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout: // cell_blk I32 [n_kv, ns] block each cell belongs to @@ -85,8 +105,13 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { // 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 + // pool_* are null when the caller wants every block recomputed inline (the pre-cache path): + // pool_idxs I64 [n_recomp] rows of the pooled cache this ubatch rewrites + // pool_cells I32 [ratio*n_recomp] cells making up each rewritten block + // pool_pos I32 [4*n_recomp] mrope position rows of each rewritten block void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, - ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, + ggml_tensor * bias, ggml_tensor * pool_idxs, ggml_tensor * pool_cells, + ggml_tensor * pool_pos, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; private: @@ -100,6 +125,14 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + + // the pooled-key cache is addressed by block, so it needs its own hparams too + llama_hparams hparams_pool; + + const std::unique_ptr mem_pool; + + // pooled keys are correct for the blocks that cover positions below this + mutable uint32_t pool_valid_pos = 0; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -141,11 +174,20 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // nullptr with no indexer const llama_kv_cache_context * get_idx() const; + // the QSA pooled-key cache, and how many trailing blocks the graph must rewrite into it + llama_kv_cache * get_mem_pool() const; + uint32_t qsa_pool_n_recomp(uint32_t ratio, uint32_t n_tokens, uint32_t n_kv, uint32_t n_pad_kv) const; + // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified uint32_t get_n_stream() const; + // pool_* are null when the caller wants every block recomputed inline (the pre-cache path): + // pool_idxs I64 [n_recomp] rows of the pooled cache this ubatch rewrites + // pool_cells I32 [ratio*n_recomp] cells making up each rewritten block + // pool_pos I32 [4*n_recomp] mrope position rows of each rewritten block void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, - ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, + ggml_tensor * bias, ggml_tensor * pool_idxs, ggml_tensor * pool_cells, + ggml_tensor * pool_pos, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; private: diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 25794aeba1f8..e9f23af6ac9b 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -508,6 +508,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( return ggml_mul(ctx0, normalized, gated); } +// llama_kv_cache::get_n_kv pads n_kv to at least 256 cells, so n_kv grows in steps of 256 and +// the block count in steps of 256/ratio. The pooled-key window has to cover that jump. +static constexpr uint32_t QSA_N_PAD_KV = 256; + // QSA attends to a budget of whole blocks of compress_ratio tokens, each scored by one // mean-pooled indexer key, plus the incomplete tail. set_input resolves the cache layout. class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { @@ -518,7 +522,8 @@ 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, bias, + pool_idxs, pool_cells, pool_pos, ubatch, ratio, blk_bias); } bool can_reuse(const llm_graph_params & params) override; @@ -529,6 +534,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] 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] + // the pooled-key window; null when every block is recomputed inline + ggml_tensor * pool_idxs = nullptr; // I64 [n_recomp] + ggml_tensor * pool_cells = nullptr; // I32 [ratio*n_recomp] + ggml_tensor * pool_pos = nullptr; // I32 [4*n_recomp] const llama_memory_hybrid_idx_context * mctx; const uint32_t ratio; @@ -571,6 +580,14 @@ bool llama_model_qwen4exp::llm_graph_input_qsa::can_reuse(const llm_graph_params res &= bias != nullptr && bias->ne[0] == (blk_bias ? n_blocks : n_kv); res &= bias != nullptr && bias->ne[1] == n_tokens/n_stream; + // the window is sized from whether the pooled cache was valid when the graph was built, + // so a graph built over a valid cache must not be reused after something dropped it + if (pool_cells != nullptr) { + const int64_t want = (int64_t) ratio * + m->qsa_pool_n_recomp(ratio, (uint32_t) n_tokens, (uint32_t) n_kv, QSA_N_PAD_KV); + res &= pool_cells->buffer != nullptr && pool_cells->ne[0] == want; + } + return res; } @@ -618,17 +635,44 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( } else { auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + // the pooled-key cache is addressed by block with no stream offset, so it serves a + // single-stream cache only; everything else keeps recomputing every block inline. + // the test is on the cache, not on the ubatch: a slot of a multi-stream cache gets + // n_stream == 1 too, but its rows start at sinfo.s0 rather than at zero. + llama_kv_cache * mem_pool = mctx_hyb->get_mem_pool(); + + const bool use_pool = mem_pool != nullptr && mem_pool->get_n_stream() == 1; + qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); - 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); ggml_set_input(qsa->cell_blk); - ggml_set_input(qsa->blk_cells); - ggml_set_input(qsa->blk_pos); ggml_set_input(qsa->bias); + // ggml-alloc gives data only to tensors some node reads, so an input the graph has no + // use for keeps data == nullptr and set_input then writes through a null pointer. With + // the pooled-key cache the graph reads the recompute window, never the whole block + // table, so these two are not created and set_input_qsa keeps the table in host scratch. + if (!use_pool) { + 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); + + ggml_set_input(qsa->blk_cells); + ggml_set_input(qsa->blk_pos); + } else { + const int64_t n_recomp = mctx_hyb->qsa_pool_n_recomp( + (uint32_t) r, (uint32_t) n_tokens, (uint32_t) n_kv, QSA_N_PAD_KV); + + qsa->pool_idxs = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_recomp); + qsa->pool_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r*n_recomp); + qsa->pool_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_recomp); + + ggml_set_input(qsa->pool_idxs); + ggml_set_input(qsa->pool_cells); + ggml_set_input(qsa->pool_pos); + } + inp = qsa.get(); res->add_input(std::move(qsa)); qsa_inps.emplace((uint32_t) r, inp); @@ -645,35 +689,59 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * k_all = mctx_idx->get_k(ctx0, il); k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0); + // A full block never changes again: its cells are written once, and pooling, normalisation + // and rotation are all position-determined. So with the pooled-key cache only the tail of + // the block table is recomputed and written back; the rest is read straight out of the + // cache. n_recomp is n_blocks while the cache is invalid, which makes this the same code + // as the inline path, just writing its result out as well. + const int64_t n_recomp = inp->pool_cells ? inp->pool_cells->ne[0]/r : n_blocks; + + ggml_tensor * blk_src = inp->pool_cells ? inp->pool_cells : inp->blk_cells; + // gathers per stream: blk_cells row s indexes stream s's own cells - ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); - members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); + ggml_tensor * members = ggml_get_rows(ctx0, k_all, blk_src); + members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_recomp, n_stream); // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows. // the slices are strided views of members. ggml_add has no contiguity requirement (the Vulkan // backend gates it on type alone) and ggml_dup_tensor gives the sum a contiguous home, so the - // per-slice ggml_cont was materialising the whole of members a second time for nothing: - // r reads plus r writes of [idx_dim, n_blocks] f32 per layer per ubatch, about 34 MB at 33k - // context and 12 layers. the addition order is unchanged, so the arithmetic is identical. - ggml_tensor * pooled = nullptr; + // per-slice ggml_cont was materialising the whole of members a second time for nothing. + // the addition order is unchanged, so the arithmetic is identical. + ggml_tensor * fresh = nullptr; for (int64_t i = 0; i < r; ++i) { - ggml_tensor * slice = ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, + ggml_tensor * slice = ggml_view_3d(ctx0, members, idx_dim, n_recomp, n_stream, members->nb[2], members->nb[3], i*members->nb[1]); - pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; + fresh = fresh ? ggml_add(ctx0, fresh, slice) : slice; } - pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); - cb(pooled, "indexer_k_pooled", il); + fresh = ggml_scale(ctx0, fresh, 1.0f/(float) r); + cb(fresh, "indexer_k_pooled", il); // count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536 - pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks*n_stream, 1); - pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + fresh = ggml_reshape_3d(ctx0, fresh, idx_dim, n_recomp*n_stream, 1); + fresh = build_norm(fresh, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after. - pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); - pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, + fresh = ggml_reshape_3d(ctx0, fresh, idx_dim, 1, n_recomp*n_stream); + fresh = ggml_rope_multi(ctx0, fresh, inp->pool_pos ? inp->pool_pos : inp->blk_pos, nullptr, n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream); + + ggml_tensor * pooled = nullptr; + + if (inp->pool_cells) { + // write the window back, then read the whole block range out of the cache. the write is + // expanded first, the same ordering the KV cache relies on for its own write-then-read. + ggml_tensor * pk = mctx_hyb->get_mem_pool()->get_k_storage(il); + + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, + ggml_reshape_2d(ctx0, pk, pk->ne[0], pk->ne[1]*pk->ne[2]), + ggml_reshape_2d(ctx0, fresh, idx_dim, n_recomp*n_stream), + inp->pool_idxs)); + + pooled = ggml_view_3d(ctx0, pk, idx_dim, n_blocks, n_stream, pk->nb[1], pk->nb[2], 0); + } else { + pooled = ggml_reshape_3d(ctx0, fresh, idx_dim, n_blocks, n_stream); + } cb(pooled, "indexer_k", il); ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur); From fbe9e678e23b8c66522cd449fe64442d9cd0f6fd Mon Sep 17 00:00:00 2001 From: Masahito Suzuki Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH 4/6] qwen4exp: narrow the indexer cache's V Nothing reads it. build_qsa_top_k only ever calls cpy_k and get_k on that cache: the indexer scores blocks against a key and has no value side at all. llama_kv_cache allocates a V anyway (it skips one only for MLA), and at the model's own n_embd_head_v of 256 that is dead weight of n_head_kv(1) * 256 elements per cell per layer. Narrow it to a single element. The type has to go with it: a quantised row must be a whole number of blocks and one element of q8_0 is not, so ask for F32 and the whole V costs 4 bytes per cell per layer. Measured on Strix Halo at ctx 262144 with an f16 cache: 90444 MB -> 88810 MB after load, 92714 MB -> 91251 MB after generation, against a computed 1610 MB. Output unchanged at 2048 and 8192. On a 2x RTX 3090 host at ctx 154624 with a q8_0 cache it is worth 481 MiB, which is what let that host's context ceiling move from 157696 to 180224. --- src/llama-memory-hybrid-idx.cpp | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index f94af9f11b2b..289b40bcaf6c 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -53,6 +53,18 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + // Nothing reads this cache's V. build_qsa_top_k only ever calls cpy_k and get_k on it: + // the indexer scores blocks against a key, and has no value side at all. llama_kv_cache + // allocates a V anyway (it only skips one for MLA), and at the model's own value width + // that is n_head_kv(1) * n_embd_head_v(256) per cell per layer of dead weight -- 482 MiB + // at ctx=154624 over 12 QSA layers with a q8_0 cache, 1.6 GiB at ctx=262144 with f16. + // + // Narrow it to a single element. The type has to go with it: a quantised row must be a + // whole number of blocks, and one element of q8_0 is not, so ask for F32 below and the + // whole V costs 4 bytes per cell per layer. + hparams_idx.n_embd_head_v_full = 1; + hparams_idx.n_embd_head_v_swa = 1; + // the cached indexer keys are raw, rotation happens after pooling at read time, so a // K-shift must not rotate them while the stream copies in the same update still apply hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE; @@ -60,7 +72,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); return new llama_kv_cache( - model, hparams_idx, type_k, type_v, v_trans, offload, unified, + model, hparams_idx, type_k, GGML_TYPE_F32, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); }()), From b96f6a0fef0c7350372f0e1f723b7eaa0cc9462d Mon Sep 17 00:00:00 2001 From: Masahito Suzuki Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH 5/6] qwen4exp: narrow the pooled-key cache's V as well Same dead allocation, in the cache the pooled-key patch adds. The graph writes pooled keys with ggml_set_rows and reads them back as a view; it never asks for a value side. That cache was created with V at the key width, which still left indexer_head_size F32 elements per block per layer. One element instead. The K side is already F32, so the type stays. Measured on Strix Halo at ctx 262144, f16 KV, MTP n_max=3: llama_kv_cache: Vulkan0 KV buffer size 768.01 MiB -> 387.01 MiB GTT after load 83754 MiB -> 83373 MiB Both move by the same 381 MiB, and every other buffer is unchanged to the byte (attention KV 6144.00, indexer 780.00, RS 112.57, compute 8535.86 / 1599.59). Output sha is identical at 2048 / 8192 / 32768 and so are the speculative decoder's accept counts (147/81, 119/119, 121/115), so the draft path is unaffected too. --- src/llama-memory-hybrid-idx.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 289b40bcaf6c..a4903e0c07df 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -97,9 +97,13 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( std::fill(hparams_pool.n_head_kv_arr.begin(), hparams_pool.n_head_kv_arr.end(), 1); hparams_pool.n_embd_head_k_full = model.hparams.indexer_head_size; - // nothing reads V here; make it the same width as K rather than the model's, so the - // allocation the cache makes for it is not several hundred megabytes of dead weight - hparams_pool.n_embd_head_v_full = model.hparams.indexer_head_size; + // nothing reads this cache's V, for the same reason the indexer cache's V is dead: the + // graph writes pooled keys with set_rows and reads them back as a view, and never asks + // for a value side at all. Narrowing it to the key width still left indexer_head_size + // F32 elements per block per layer -- 384 MiB at ctx=262144 over 12 QSA layers. One + // element costs 4 bytes per block per layer instead. + hparams_pool.n_embd_head_v_full = 1; + hparams_pool.n_embd_head_v_swa = 1; // the rows are already rotated when they are written, so a K-shift must not touch them hparams_pool.rope_type = LLAMA_ROPE_TYPE_NONE; From a4e0ee7dbea56fd58769ba7604a5055f01fa766c Mon Sep 17 00:00:00 2001 From: Masahito Suzuki Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH 6/6] qwen4exp: drop the 1/r scale on the QSA block mean Its only consumer is the RMS norm below it, and RMS norm is scale invariant: rms(x*s) = x*s / sqrt(mean(x^2)*s^2 + eps) = x / sqrt(mean(x^2) + eps/s^2) so dropping the divide only moves the effective epsilon from eps to r^2*eps -- 1e-6 to 1.6e-5 against a mean square of order 1. The scale was a full read and write of [idx_dim, n_blocks] f32 per layer per ubatch: 2.25 ms of a 97 ms decode step at 131k context. This is the one patch in the series that is not bit-exact by construction. Measured on Strix Halo (gfx1151, Vulkan), ctx 262144, f16 KV, MTP n_max=3: 131072 goes 28.26 -> 32.46 t/s cumulative with the previous patch. --- src/models/qwen4exp.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index e9f23af6ac9b..03d73506ea3b 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -713,8 +713,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( members->nb[2], members->nb[3], i*members->nb[1]); fresh = fresh ? ggml_add(ctx0, fresh, slice) : slice; } - fresh = ggml_scale(ctx0, fresh, 1.0f/(float) r); - cb(fresh, "indexer_k_pooled", il); + // no ggml_scale by 1/r here: the only consumer is the RMS norm below, and RMS norm is + // scale invariant. rms(x*s) = x*s / sqrt(mean(x^2)*s^2 + eps) = x / sqrt(mean(x^2) + eps/s^2), + // so dropping the divide only moves the epsilon from eps to r^2*eps -- 1e-6 to 1.6e-5 against + // a mean square of order 1. the scale was a full read and write of [idx_dim, n_blocks] f32 per + // layer per ubatch: 2.25 ms of a 97 ms decode step at 131k context. + cb(fresh, "indexer_k_sum", il); // count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536 fresh = ggml_reshape_3d(ctx0, fresh, idx_dim, n_recomp*n_stream, 1);