diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index f7e99c761dc0..a4903e0c07df 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 @@ -51,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; @@ -58,9 +72,56 @@ 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_"); + }()), + 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 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; + + 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 +200,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 +216,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 +226,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 +236,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 +246,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 +256,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 +292,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 +310,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 +359,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 +372,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 +382,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 +695,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 +818,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/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..03d73506ea3b 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -508,9 +508,13 @@ 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 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) {} @@ -518,7 +522,8 @@ class 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 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; @@ -543,7 +552,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; @@ -571,6 +580,14 @@ bool llm_graph_input_qsa::can_reuse(const llm_graph_params & 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; } @@ -607,21 +624,59 @@ 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; + + 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); + + // 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); + 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->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_set_input(qsa->cell_blk); + ggml_set_input(qsa->bias); - llm_graph_input_qsa * inp = qsa.get(); - res->add_input(std::move(qsa)); + // 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); + } // 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); @@ -634,31 +689,63 @@ 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); - // 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); + // 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; - // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows - ggml_tensor * pooled = nullptr; + 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, 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. + // 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_cont(ctx0, - 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; + ggml_tensor * slice = ggml_view_3d(ctx0, members, idx_dim, n_recomp, n_stream, + members->nb[2], members->nb[3], i*members->nb[1]); + fresh = fresh ? ggml_add(ctx0, fresh, slice) : slice; } - pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); - cb(pooled, "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 - 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);