qwen4exp: compute QSA cache-layout inputs once per ubatch, not per layer

src/llama-kv-cells.h

@@ -50,6 +50,8 @@ class llama_kv_cells {
         for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {
             seq_pos[s].clear();
         }
+
+        generation++;
     }
 
     void reset_shift() {
@@ -100,6 +102,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());
@@ -164,6 +172,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;
 
@@ -195,6 +205,8 @@ class llama_kv_cells {
     void set(const std::vector<uint32_t> & 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];
 
@@ -227,6 +239,8 @@ class llama_kv_cells {
         assert(i < pos.size());
         assert(pos[i] != -1);
 
+        generation++;
+
         seq_pos_rm(i);
         seq[i].reset();
 
@@ -245,6 +259,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]);
 
@@ -265,6 +281,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();
@@ -338,6 +356,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]);
     }
@@ -424,13 +444,18 @@ class llama_kv_cells {
         assert(pos[i] == -1);
         assert(seq[i].none());
 
+        generation++;
+
         pos[i] = p;
 
         used.insert(i);
     }
 
     void ext_set(uint32_t i, llama_kv_cell_ext p) {
         assert(i < ext.size());
+
+        generation++;
+
         ext[i] = p;
     }
 
@@ -441,6 +466,8 @@ class llama_kv_cells {
         assert(i < pos.size());
         assert(pos[i] != -1);
 
+        generation++;
+
         seq_pos_rm(i);
 
         pos[i]   += d;
@@ -470,6 +497,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);
@@ -485,6 +514,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<uint32_t> used;
 

src/llama-memory-hybrid-idx.cpp

@@ -8,6 +8,7 @@
 #include <algorithm>
 #include <cassert>
 #include <cmath>
+#include <cstring>
 #include <iterator>
 #include <stdexcept>
 
@@ -366,8 +367,57 @@ 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;
 
-    // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio
-    // all mrope sections carry it: exact for text, approximate for images
+    // 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) {
         for (int64_t s = 0; s < n_ns; ++s) {
             for (int64_t b = 0; b < n_blocks; ++b) {
@@ -462,4 +512,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;
+    }
 }

src/llama-memory-hybrid-idx.h

@@ -85,6 +85,23 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid {
     llama_hparams hparams_idx;
 
     const std::unique_ptr<llama_kv_cache> 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<int32_t> cell_blk;
+        std::vector<int32_t> blk_cells;
+        std::vector<int32_t> blk_pos;
+        std::vector<float>   bias;
+    };
+
+    mutable qsa_memo memo_qsa;
 };
 
 class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {

src/models/models.h

@@ -2351,7 +2351,19 @@ struct llama_model_qwen4exp : public llama_model_base {
         // build_rs writes the state tensor in place, so one gather per cache tensor is reused
         std::map<ggml_tensor *, ggml_tensor *> rs_rows;
 
-        // one conv history per cache tensor: delta-net and PLE each have their own
+        // 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<uint32_t, qsa_inps> 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,
                     ggml_tensor * conv_states_all,

src/models/qwen4exp.cpp

@@ -489,52 +489,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;
 
-    // only the "which block is visible" half of the bias varies per block
-    // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells
-    // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out
-    // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there
-    const bool blk_bias = kq_mask != nullptr &&
-        kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream &&
-        cparams.causal_attn && !hparams.use_alibi;
-
-    // nothing above depends on the layer, so the layers sharing a ratio share one input set
-    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<llm_graph_input_qsa>(mctx_hyb, (uint32_t) r, blk_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<llm_graph_input_qsa>(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, blk_bias ? n_blocks : n_kv, n_tps, 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);
 
-        inp = qsa.get();
+        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));
-        qsa_inps.emplace((uint32_t) r, inp);
     }
 
+    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 members; r is small, so summing slices beats a transpose plus sum_rows
@@ -551,7 +548,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);
@@ -583,16 +580,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
 
     // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary
     ggml_tensor * expanded = ggml_get_rows(ctx0,
-            ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk);
+            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));
-
-    if (blk_bias) {
-        // flash attention keeps the mask in f16; the scores are f32
-        ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32);
-        expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream));
-    } else {
-        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: whole blocks plus the tail