From 558360ee368cae6e293e5377f4c0048b0ed34da1 Mon Sep 17 00:00:00 2001 From: alex Date: Sun, 30 Aug 2026 18:23:56 +0300 Subject: [PATCH] qwen4exp: compute QSA cache-layout inputs once per ubatch, not per layer set_input_qsa results depend only on the indexer cells and the ubatch, yet the graph asked for them once per attention layer (48x), each time running O(n_kv) serial host loops and filling a separate set of input tensors. At long context this dominated decode: GPU util sat at ~31% while host code burned tens of ms per token. Two changes: 1. Share one set of QSA input tensors (k_idxs, cell_blk, blk_cells, blk_pos, bias) across all layers with the same compress ratio (graph::qsa_shared). One fill + one H2D copy per ubatch instead of 48. 2. Memoize set_input_qsa on llama_memory_hybrid_idx, keyed by a fingerprint of per-stream cell generations (new llama_kv_cells::get_generation(), bumped by every mutator) and the ubatch tokens. Decode-step repeats hit the memo; prefill batches above a 256MB bias cap bypass it. Measured on Qwen3.8-Flash-Next UD-IQ4_XS, 2x RTX 3090 (NVLink), f16 KV: decode tps 32k 131k 176k 229k before 29.1 16.2 OOM - after 32.1 21.1 18.2 15.5 prefill tps 32k 131k 176k 229k before 336 120 - - after 492 315 260 220 --- src/llama-kv-cells.h | 32 ++++++++++++++++ src/llama-memory-hybrid-idx.cpp | 65 +++++++++++++++++++++++++++++++++ src/llama-memory-hybrid-idx.h | 17 +++++++++ src/models/models.h | 12 ++++++ src/models/qwen4exp.cpp | 52 ++++++++++++++++---------- 5 files changed, 159 insertions(+), 19 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index fddd31a0b21..59bc141303e 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -46,6 +46,8 @@ class llama_kv_cells { for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) { seq_pos[s].clear(); } + + generation++; } void reset_shift() { @@ -96,6 +98,12 @@ class llama_kv_cells { return has_shift; } + // incremented by every mutator above; used by llama_memory_hybrid_idx to detect + // that the cell contents changed and cached host-side data must be recomputed + uint32_t get_generation() const { + return generation; + } + // move cell isrc to idst (used during defrag) //void mv(uint32_t isrc, uint32_t idst) { // assert(isrc < pos.size()); @@ -160,6 +168,8 @@ class llama_kv_cells { void set(uint32_t i, const llama_kv_cells & other) { assert(i + other.pos.size() <= pos.size()); + generation++; + for (uint32_t j = 0; j < other.pos.size(); ++j) { const auto idx = i + j; @@ -191,6 +201,8 @@ class llama_kv_cells { void set(const std::vector & idxs, const llama_kv_cells & other) { assert(idxs.size() == other.pos.size()); + generation++; + for (uint32_t j = 0; j < other.pos.size(); ++j) { const auto idx = idxs[j]; @@ -223,6 +235,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + generation++; + seq_pos_rm(i); seq[i].reset(); @@ -241,6 +255,8 @@ class llama_kv_cells { assert(pos[i] != -1); assert(seq_id >= 0); + generation++; + seq[i].reset(seq_id); seq_pos_dec(seq_id, pos[i]); @@ -261,6 +277,8 @@ class llama_kv_cells { bool seq_keep(uint32_t i, llama_seq_id seq_id) { assert(i < pos.size()); + generation++; + if (seq[i].test(seq_id)) { seq_pos_rm(i); seq[i].reset(); @@ -311,6 +329,8 @@ class llama_kv_cells { assert(pos[i] != -1); assert(!seq[i].test(seq_id)); + generation++; + seq[i].set(seq_id); seq_pos_inc(seq_id, pos[i]); } @@ -397,6 +417,8 @@ class llama_kv_cells { assert(pos[i] == -1); assert(seq[i].none()); + generation++; + pos[i] = p; used.insert(i); @@ -404,6 +426,9 @@ class llama_kv_cells { void ext_set(uint32_t i, llama_kv_cell_ext p) { assert(i < ext.size()); + + generation++; + ext[i] = p; } @@ -414,6 +439,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + generation++; + seq_pos_rm(i); pos[i] += d; @@ -443,6 +470,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + generation++; + const llama_pos p_old = pos[i]; seq_pos_rm(i); @@ -458,6 +487,9 @@ class llama_kv_cells { private: bool has_shift = false; + // bumped on every mutation of the cell contents (see get_generation()) + uint32_t generation = 0; + // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id) std::set used; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 0974a67e4d4..49885e357c4 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -7,6 +7,7 @@ #include #include #include +#include // // llama_memory_hybrid_idx @@ -301,6 +302,55 @@ void llama_memory_hybrid_idx_context::set_input_qsa( int32_t * dst_blk_pos = (int32_t *) blk_pos->data; float * dst_bias = (float *) bias->data; + // the graph calls this once per attention layer with identical inputs and the + // result depends only on the indexer cells and the ubatch, so memoize it on the + // memory object (see qsa_memo). only memoize small batches (decode): during + // prefill the bias matrix is n_kv*n_tokens floats and the cost is amortized. + const size_t bytes_cell_blk = sizeof(int32_t)*n_kv*n_ns; + const size_t bytes_blk_cells = sizeof(int32_t)*r*n_blocks*n_ns; + const size_t bytes_blk_pos = sizeof(int32_t)*4*n_blocks*n_ns; + const size_t bytes_bias = sizeof(float)*n_kv*n_tokens; + + auto & memo = mem->memo_qsa; + + const bool use_memo = bytes_bias < (size_t) 256*1024*1024; + + uint64_t key = 0; + if (use_memo) { + key = 1469598103934665603ull; + const auto fnv = [&key](uint64_t v) { + key ^= v; + key *= 1099511628211ull; + }; + + fnv((uint64_t) n_kv); + fnv((uint64_t) n_ns); + fnv((uint64_t) n_blocks); + fnv((uint64_t) n_tokens); + fnv((uint64_t) r); + + for (int64_t s = 0; s < n_ns; ++s) { + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + + fnv((uint64_t) (int64_t) seq_of_stream); + fnv(mem->get_mem_idx()->get_cells(seq_of_stream).get_generation()); + } + + for (int64_t i = 0; i < n_tokens; ++i) { + fnv((uint64_t) (int64_t) ubatch->seq_id[i][0]); + fnv((uint64_t) (int64_t) ubatch->pos[i]); + } + + if (memo.valid && memo.key == key) { + memcpy(dst_cell_blk, memo.cell_blk.data(), bytes_cell_blk); + memcpy(dst_blk_cells, memo.blk_cells.data(), bytes_blk_cells); + memcpy(dst_blk_pos, memo.blk_pos.data(), bytes_blk_pos); + memcpy(dst_bias, memo.bias.data(), bytes_bias); + + return; + } + } + // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio. All three // mrope sections carry it: exact for text, approximate for images. Positions repeat per stream. for (int64_t sec = 0; sec < 4; ++sec) { @@ -377,4 +427,19 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } } + + if (use_memo) { + memo.cell_blk.resize(n_kv*n_ns); + memo.blk_cells.resize(r*n_blocks*n_ns); + memo.blk_pos.resize(4*n_blocks*n_ns); + memo.bias.resize(n_kv*n_tokens); + + memcpy(memo.cell_blk.data(), dst_cell_blk, bytes_cell_blk); + memcpy(memo.blk_cells.data(), dst_blk_cells, bytes_blk_cells); + memcpy(memo.blk_pos.data(), dst_blk_pos, bytes_blk_pos); + memcpy(memo.bias.data(), dst_bias, bytes_bias); + + memo.key = key; + memo.valid = true; + } } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index d5e75ef705c..cefb27a13ad 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -87,6 +87,23 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + +public: + // memo for set_input_qsa: its result depends only on the indexer cells content and + // the ubatch, but the graph asks for it once per attention layer with identical + // arguments. lives here (not in the short-lived batch context) so the memo survives + // across decode steps. keyed by a fingerprint of cells generations + ubatch tokens. + struct qsa_memo { + bool valid = false; + uint64_t key = 0; + + std::vector cell_blk; + std::vector blk_cells; + std::vector blk_pos; + std::vector bias; + }; + + mutable qsa_memo memo_qsa; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { diff --git a/src/models/models.h b/src/models/models.h index 6fe49d072e1..f804551d476 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2355,6 +2355,18 @@ struct llama_model_qwen4exp : public llama_model_base { // layer; both convolutions share this gather. std::map rs_rows; + // QSA cache-layout inputs depend only on the cells and the ubatch, so all + // layers sharing a compress ratio reuse one tensor set instead of building + // (and filling, and copying to the device) 48 identical copies + struct qsa_inps { + ggml_tensor * k_idxs; + ggml_tensor * cell_blk; + ggml_tensor * blk_cells; + ggml_tensor * blk_pos; + ggml_tensor * bias; + }; + std::map qsa_shared; + // conv history at an explicit offset: delta-net and PLE share the row ggml_tensor * build_conv_state_at( llm_graph_input_rs * inp, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 65e9b199107..638da2b9408 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -432,35 +432,49 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( GGML_ASSERT(n_tokens % n_stream == 0); const int64_t n_tps = n_tokens/n_stream; - auto qsa = std::make_unique(mctx_hyb, (uint32_t) r); - - 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, 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); + // the cache-layout inputs are identical for every layer with this ratio: + // build and register them once, reuse the tensors afterwards + auto sh = qsa_shared.find((uint32_t) r); + if (sh == qsa_shared.end()) { + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r); + + 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, 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); + + sh = qsa_shared.emplace((uint32_t) r, qsa_inps{ + qsa->k_idxs, qsa->cell_blk, qsa->blk_cells, qsa->blk_pos, qsa->bias, + }).first; + + res->add_input(std::move(qsa)); + } - llm_graph_input_qsa * inp = qsa.get(); - res->add_input(std::move(qsa)); + ggml_tensor * in_k_idxs = sh->second.k_idxs; + ggml_tensor * in_cell_blk = sh->second.cell_blk; + ggml_tensor * in_blk_cells = sh->second.blk_cells; + ggml_tensor * in_blk_pos = sh->second.blk_pos; + ggml_tensor * in_bias = sh->second.bias; // 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); k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens); cb(k_raw, "indexer_k_raw", il); - ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il)); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, in_k_idxs, il)); // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream]. 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); + ggml_tensor * members = ggml_get_rows(ctx0, k_all, in_blk_cells); members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); // mean over the block's members; compress_ratio is small, so summing slices beats @@ -478,7 +492,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // 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 = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); - pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, + pooled = ggml_rope_multi(ctx0, pooled, in_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); @@ -509,9 +523,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // budget is a whole number of blocks and members tie, so the cut still lands on // a block boundary. get_rows gathers rows, so scores are transposed first. ggml_tensor * expanded = ggml_get_rows(ctx0, - ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), in_cell_blk); expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); - expanded = ggml_add(ctx0, expanded, inp->bias); + expanded = ggml_add(ctx0, expanded, in_bias); cb(expanded, "indexer_score_tokens", il); // the reference returns indexer_top_k + compress_ratio - 1: a whole budget of