qwen4exp: reduce the generation slowdown as context grows

src/llama-kv-cache.cpp

@@ -1868,16 +1868,33 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
 
     for (uint32_t s = 0; s < n_stream; ++s) {
         // p_max inclusive: an embd token looks up cells at its own (shared) position
-        v_cells[s].for_each_token_in(seqs, 0, p_max + 1,
+        v_cells[s].for_each_token_in(seqs, w0, p_max + 1,
             [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
-                if (pos >= w0) {
-                    hist[key(seq_id, pos)] = tok;
-                } else if (pos > below[seq_id].first) {
-                    below[seq_id] = { pos, tok };
-                }
+                hist[key(seq_id, pos)] = tok;
             });
     }
 
+    // below[] only answers when an M-RoPE gap leaves the window empty, so it costs a pass of
+    // its own that contiguous positions never pay
+    bool below_ready = false;
+
+    const auto ensure_below = [&]() {
+        if (below_ready) {
+            return;
+        }
+
+        below_ready = true;
+
+        for (uint32_t s = 0; s < n_stream; ++s) {
+            v_cells[s].for_each_token_in(seqs, 0, w0,
+                [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
+                    if (pos > below[seq_id].first) {
+                        below[seq_id] = { pos, tok };
+                    }
+                });
+        }
+    };
+
     // the token at pos p, or the nearest earlier one when p falls in an M-RoPE gap
     const auto lookup = [&](llama_seq_id seq_id, llama_pos p) -> llama_token {
         for (llama_pos q = p; q >= w0; --q) {
@@ -1886,6 +1903,9 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
                 return it->second;
             }
         }
+
+        ensure_below();
+
         return below[seq_id].second;
     };
 

src/llama-kv-cells.h

@@ -10,6 +10,10 @@
 #include <set>
 #include <vector>
 
+#if defined(_MSC_VER)
+#include <intrin.h>
+#endif
+
 struct llama_kv_cell_ext {
     // 2D spatial positions, typically used for M-RoPE
     llama_pos x = 0;
@@ -31,6 +35,105 @@ struct llama_kv_cell_ext {
     }
 };
 
+// index of the lowest set bit, and of the highest; the word is never zero here
+#if defined(_MSC_VER)
+static inline uint32_t llama_kv_ctz64(uint64_t x) {
+    unsigned long r;
+    _BitScanForward64(&r, x);
+    return (uint32_t) r;
+}
+
+static inline uint32_t llama_kv_clz64(uint64_t x) {
+    unsigned long r;
+    _BitScanReverse64(&r, x);
+    return (uint32_t) r;
+}
+#else
+static inline uint32_t llama_kv_ctz64(uint64_t x) {
+    return (uint32_t) __builtin_ctzll(x);
+}
+
+static inline uint32_t llama_kv_clz64(uint64_t x) {
+    return (uint32_t) (63 - __builtin_clzll(x));
+}
+#endif
+
+// the cells tracked are a large fraction of the cache, so a bitmap is both smaller and faster
+// to walk than a tree node per index
+class llama_kv_idx_set {
+public:
+    void resize(uint32_t n) {
+        bits.assign((n + 63)/64, 0);
+        n_set = 0;
+    }
+
+    void clear() {
+        std::fill(bits.begin(), bits.end(), 0);
+        n_set = 0;
+    }
+
+    void insert(uint32_t i) {
+        uint64_t & w = bits[i/64];
+        const uint64_t b = 1ull << (i%64);
+
+        n_set += (w & b) == 0;
+        w |= b;
+    }
+
+    void erase(uint32_t i) {
+        uint64_t & w = bits[i/64];
+        const uint64_t b = 1ull << (i%64);
+
+        n_set -= (w & b) != 0;
+        w &= ~b;
+    }
+
+    bool contains(uint32_t i) const {
+        return (bits[i/64] >> (i%64)) & 1;
+    }
+
+    uint32_t size()  const { return n_set; }
+    bool     empty() const { return n_set == 0; }
+
+    uint32_t first() const {
+        for (size_t w = 0; w < bits.size(); ++w) {
+            if (bits[w]) {
+                return 64*w + llama_kv_ctz64(bits[w]);
+            }
+        }
+
+        return 0;
+    }
+
+    uint32_t last() const {
+        for (size_t w = bits.size(); w-- > 0; ) {
+            if (bits[w]) {
+                return 64*w + llama_kv_clz64(bits[w]);
+            }
+        }
+
+        return 0;
+    }
+
+    // ascending order, as the set it replaces guaranteed
+    template <class F>
+    void for_each(F && f) const {
+        for (size_t w = 0; w < bits.size(); ++w) {
+            uint64_t m = bits[w];
+
+            while (m) {
+                f((uint32_t) (64*w + llama_kv_ctz64(m)));
+                m &= m - 1;
+            }
+        }
+    }
+
+private:
+    std::vector<uint64_t> bits;
+
+    uint32_t n_set = 0;
+};
+
 // meta information about KV cells that can be part of multiple sequences at the same time
 // TODO: add unit tests
 class llama_kv_cells {
@@ -69,6 +172,7 @@ class llama_kv_cells {
         ext.resize(n);
         shift.resize(n);
         seq.resize(n);
+        used.resize(n);
 
         reset();
     }
@@ -87,13 +191,13 @@ class llama_kv_cells {
     // the index of the first cell that is used
     // return 0 if no cells are used
     uint32_t used_min() const {
-        return used.empty() ? 0 : *used.begin();
+        return used.empty() ? 0 : used.first();
     }
 
     // the index of the last cell that is used + 1
     // return 0 if no cells are used
     uint32_t used_max_p1() const {
-        return used.empty() ? 0 : *used.rbegin() + 1;
+        return used.empty() ? 0 : used.last() + 1;
     }
 
     bool get_has_shift() const {
@@ -314,22 +418,24 @@ class llama_kv_cells {
     // note: used by n-gram input embeddings to recover the tokens preceding a ubatch
     template<typename F>
     void for_each_token_in(const std::bitset<LLAMA_MAX_SEQ> & seqs, llama_pos p0, llama_pos p1, F && f) const {
-        for (const auto & i : used) {
+        used.for_each([&](uint32_t i) {
             if (pos[i] < p0 || pos[i] >= p1) {
-                continue;
+                return;
             }
 
             const auto m = seq[i] & seqs;
-            if (m.none()) {
-                continue;
-            }
 
-            for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+            // a cell carries a handful of sequences at most, so stop once they are all seen
+            // instead of walking the whole LLAMA_MAX_SEQ width
+            size_t left = m.count();
+
+            for (llama_seq_id s = 0; left > 0 && s < (llama_seq_id) LLAMA_MAX_SEQ; ++s) {
                 if (m.test(s)) {
                     f(s, pos[i], ext[i].tok);
+                    --left;
                 }
             }
-        }
+        });
     }
 
     // note: call only if the cell is not empty and the seq_id is not in the cell
@@ -486,7 +592,7 @@ class llama_kv_cells {
     bool has_shift = false;
 
     // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id)
-    std::set<uint32_t> used;
+    llama_kv_idx_set used;
 
     std::vector<llama_pos> pos;
 

src/models/models.h

@@ -2320,6 +2320,22 @@ struct llama_model_qwen4exp : public llama_model_base {
                           float   kq_scale,
                             int   il);
 
+        // attention over the whole window, masked down to the selection
+        ggml_tensor * build_qsa_scan(
+        llm_graph_input_attn_kv * inp,
+                    ggml_tensor * q_cur,
+                    ggml_tensor * top_k,
+                          float   kq_scale,
+                            int   il);
+
+        // attention over the selected cells alone, gathered into a window per query
+        ggml_tensor * build_qsa_gather(
+        llm_graph_input_attn_kv * inp,
+                    ggml_tensor * q_cur,
+                    ggml_tensor * top_k,
+                          float   kq_scale,
+                            int   il);
+
         // the QSA cache layout inputs do not depend on the layer, only on its compress ratio,
         // so the layers sharing a ratio share one input set
         std::map<uint32_t, llm_graph_input_qsa *> qsa_inps;

src/models/qwen4exp.cpp

@@ -579,9 +579,17 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
             ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream));
     score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);
     score = ggml_relu(ctx0, score);
-    score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
-    score = ggml_sum_rows(ctx0, score);
-    score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream);
+
+    // the heads sit side by side on ne[1] and there are few of them, so summing slices beats a
+    // transpose that would carry the whole block by token surface twice over
+    ggml_tensor * summed = nullptr;
+    for (int64_t h = 0; h < n_idx_h; ++h) {
+        ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream,
+                score->nb[2], score->nb[3], h*score->nb[1]);
+        summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice);
+    }
+
+    score = summed;
     cb(score, "indexer_score", il);
 
     // one value per block, so it is cheaper to bias here than after the cells are expanded
@@ -654,6 +662,35 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
         ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
     }
 
+    const int64_t width = top_k->ne[0];
+    const int64_t n_tps = top_k->ne[1];
+    const int64_t n_kv  = mctx_cur->get_n_kv();
+
+    // the two costs meet at 2*n_tps*width == n_kv; the margin keeps the windows small enough
+    // for the compute buffer of a decode graph. flash attention is what keeps the values as rows
+    ggml_tensor * cur = cparams.flash_attn && 4*n_tps*width < n_kv
+        ? build_qsa_gather(inp, q_cur, top_k, kq_scale, il)
+        : build_qsa_scan  (inp, q_cur, top_k, kq_scale, il);
+    cb(cur, "kqv_out", il);
+
+    // the rotation is its own inverse, so undo it on the value side of the output
+    if (inp->self_v_rot) {
+        cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
+    }
+
+    return cur;
+}
+
+// The window stays whole and the mask hides every cell the selection leaves out.
+// The mask build below copies the MLA sparse path in llm_graph_context::build_attn.
+ggml_tensor * llama_model_qwen4exp::graph::build_qsa_scan(
+        llm_graph_input_attn_kv * inp,
+        ggml_tensor *             q_cur,
+        ggml_tensor *             top_k,
+        float                     kq_scale,
+        int                       il) {
+    const auto * mctx_cur = inp->mctx;
+
     ggml_tensor * kq_mask = inp->get_kq_mask();
 
     // prepare new kq mask - starts filled with -INFINITY
@@ -686,15 +723,64 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
     ggml_tensor * k = mctx_cur->get_k(ctx0, il);
     ggml_tensor * v = mctx_cur->get_v(ctx0, il);
 
-    ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il);
-    cb(cur, "kqv_out", il);
+    return build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il);
+}
 
-    // the rotation is its own inverse, so undo it on the value side of the output
-    if (inp->self_v_rot) {
-        cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
-    }
+// Key and value traffic follows the budget instead of the whole cache. The queries ride the
+// stream axis of the attention, each carrying the window its own selection named.
+ggml_tensor * llama_model_qwen4exp::graph::build_qsa_gather(
+        llm_graph_input_attn_kv * inp,
+        ggml_tensor *             q_cur,
+        ggml_tensor *             top_k,
+        float                     kq_scale,
+        int                       il) {
+    const auto * mctx_cur = inp->mctx;
 
-    return cur;
+    const int64_t width    = top_k->ne[0];
+    const int64_t n_tps    = top_k->ne[1];
+    const int64_t n_stream = top_k->ne[3];
+    const int64_t n_q      = n_tps*n_stream;
+
+    ggml_tensor * k_all = mctx_cur->get_k(ctx0, il);
+    ggml_tensor * v_all = mctx_cur->get_v(ctx0, il);
+
+    const int64_t n_kv = k_all->ne[2];
+
+    // a cell holds its heads back to back, so one row of the gather is one whole cell
+    ggml_tensor * k_cells = ggml_view_3d(ctx0, k_all, k_all->ne[0]*k_all->ne[1], n_kv, n_stream,
+            k_all->nb[2], k_all->nb[3], 0);
+    ggml_tensor * v_cells = ggml_view_3d(ctx0, v_all, v_all->ne[0]*v_all->ne[1], n_kv, n_stream,
+            v_all->nb[2], v_all->nb[3], 0);
+
+    // a cell index names a cell of its own stream, and the ubatch lays the queries of a stream
+    // out contiguously, so the flat index of a window is the index of its query
+    ggml_tensor * idx_stream = ggml_reshape_2d(ctx0, top_k, width*n_tps, n_stream);
+
+    ggml_tensor * k_sel = ggml_get_rows(ctx0, k_cells, idx_stream);
+    ggml_tensor * v_sel = ggml_get_rows(ctx0, v_cells, idx_stream);
+
+    k_sel = ggml_reshape_4d(ctx0, k_sel, k_all->ne[0], k_all->ne[1], width, n_q);
+    v_sel = ggml_reshape_4d(ctx0, v_sel, v_all->ne[0], v_all->ne[1], width, n_q);
+    cb(k_sel, "qsa_k_sel", il);
+    cb(v_sel, "qsa_v_sel", il);
+
+    // gathering the attention mask at the selected cells leaves the same values the scan path
+    // would put there, so the window carries the reach of its query
+    ggml_tensor * kq_mask = inp->get_kq_mask();
+
+    GGML_ASSERT(kq_mask->nb[3] == kq_mask->nb[1]*n_tps);
+
+    ggml_tensor * mask_cells = ggml_view_3d(ctx0, kq_mask, 1, n_kv, n_q,
+            kq_mask->nb[0], kq_mask->nb[1], 0);
+
+    ggml_tensor * idx_query = ggml_reshape_3d(ctx0, top_k, width, n_q, 1);
+
+    ggml_tensor * mask = ggml_get_rows(ctx0, mask_cells, idx_query);
+
+    mask = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mask, width, 1, 1, n_q), GGML_TYPE_F16);
+    cb(mask, "qsa_mask_sel", il);
+
+    return build_attn_mha(q_cur, k_sel, v_sel, nullptr, mask, nullptr, nullptr, kq_scale, il);
 }
 
 ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(