qwen4exp : gather-based sparse attention for QSA decode

src/llama-graph.cpp

@@ -1087,7 +1087,11 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
     mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
     mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
 
-    mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+    // the mask can be left unallocated when the graph never attends over the full
+    // cache (e.g. the qwen4exp gather path reads its own compact per-cell mask)
+    if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) {
+        mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+    }
 
     if (inp_attn->self_k_rot) {
         mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot);

src/llama-memory-hybrid-idx.cpp

@@ -343,14 +343,18 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
         ggml_tensor * cell_blk,
         ggml_tensor * blk_cells,
         ggml_tensor * blk_pos,
+        ggml_tensor * mask_row,
         ggml_tensor * bias,
         const llama_ubatch * ubatch,
         uint32_t ratio,
         bool blk_bias) const {
     GGML_ASSERT(ratio > 0);
     GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
 
-    GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
+    // the block-gather decode path never references cell_blk in the graph, so the
+    // scheduler leaves it unallocated; skip its (O(n_kv) per ubatch) fill in that case
+    const bool has_cell_blk = cell_blk->buffer != nullptr;
+    GGML_ASSERT(!has_cell_blk || ggml_backend_buffer_is_host(cell_blk->buffer));
 
     const int64_t n_kv     = cell_blk->ne[0];
     const int64_t n_ns     = cell_blk->ne[1];        // streams in this ubatch
@@ -361,7 +365,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
     GGML_ASSERT(n_tokens % n_ns == 0);
     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_cell_blk  = has_cell_blk ? (int32_t *) cell_blk->data : nullptr;
     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;
@@ -385,7 +389,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
         const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];
         const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream);
 
-        int32_t * cur_cell_blk  = dst_cell_blk  + s*n_kv;
+        int32_t * cur_cell_blk  = has_cell_blk ? dst_cell_blk + s*n_kv : nullptr;
         int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
 
         // an incomplete block cannot be pooled; the bias below forces those tail cells in
@@ -420,7 +424,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
 
         // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it
         // per-cell mode carries that -inf itself and only needs the gather in range
-        for (int64_t j = 0; j < n_kv; ++j) {
+        for (int64_t j = 0; has_cell_blk && j < n_kv; ++j) {
             if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) {
                 blk_of[j] = -1;
             }
@@ -432,6 +436,33 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
             const llama_seq_id seq_id = ubatch->seq_id[i][0];
             const llama_pos    q      = ubatch->pos[i];
 
+            // per-cell visibility for the gather path: what row i of the attention
+            // kq_mask would say, minus SWA/alibi (the gather gate excludes those).
+            // unallocated (and skipped) when the graph takes the masked path.
+            if (mask_row != nullptr && mask_row->buffer != nullptr) {
+                float * dst_row = (float *) mask_row->data + i*n_kv;
+
+                const bool is_2d = ubatch->is_pos_2d();
+
+                const llama_pos qx = is_2d ? ubatch->pos[i + ubatch->n_tokens*2] : 0;
+                const llama_pos qy = is_2d ? ubatch->pos[i + ubatch->n_tokens]   : 0;
+
+                for (int64_t j = 0; j < n_kv; ++j) {
+                    float v = -INFINITY;
+
+                    if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) {
+                        v = 0.0f;
+
+                        // M-RoPE: image tokens can share the temporal position of the query
+                        if (is_2d && cells.pos_get(j) == q && cells.ext_get(j).is_2d_gt(qx, qy)) {
+                            v = -INFINITY;
+                        }
+                    }
+
+                    dst_row[j] = v;
+                }
+            }
+
             // the tail is an incomplete block and is always visible, as in the reference
             const llama_pos tail_start = (q + 1)/r*r;
 

src/llama-memory-hybrid-idx.h

@@ -137,7 +137,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
     //   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
-    void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
+    void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, ggml_tensor * mask_row,
                        ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
                        bool blk_bias) const;
 

src/models/models.h

@@ -2317,8 +2317,10 @@ struct llama_model_qwen4exp : public llama_model_base {
                     ggml_tensor * k_cur,
                     ggml_tensor * v_cur,
                     ggml_tensor * top_k,
+                    ggml_tensor * mask_row,
                           float   kq_scale,
-                            int   il);
+                            int   il,
+                           bool   gather = false);
 
         // 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
@@ -2331,7 +2333,8 @@ struct llama_model_qwen4exp : public llama_model_base {
                     ggml_tensor * inp_pos,
                     ggml_tensor * kq_mask,
                             int * sections,
-                            int   il);
+                            int   il,
+                           bool   gather = false);
 
         ggml_tensor * build_layer_attn_linear(
              llm_graph_input_rs * inp,

src/models/qwen4exp.cpp

@@ -430,7 +430,7 @@ 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, mask_row, bias, ubatch, ratio, blk_bias);
     }
 
     bool can_reuse(const llm_graph_params & params) override {
@@ -457,6 +457,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i {
         res &= bias->ne[0]      == (blk_bias ? n_blocks : n_kv);
         res &= bias->ne[1]      == params.ubatch.n_tokens/n_stream;
 
+        res &= mask_row->ne[0]  == n_kv;
+        res &= mask_row->ne[1]  == params.ubatch.n_tokens/n_stream;
+        res &= mask_row->ne[2]  == n_stream;
+
         return res;
     }
 
@@ -467,6 +471,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i {
     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]
 
+    // per-cell visibility for the gather path (0 or -INF); left unallocated by the
+    // masked path, which reads the attention kq_mask instead. F32 [n_kv, n_tps, n_stream]
+    ggml_tensor * mask_row  = nullptr;
+
     const llama_memory_hybrid_idx_context * mctx;
     const uint32_t ratio;
 
@@ -480,7 +488,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
         ggml_tensor *                           inp_pos,
         ggml_tensor *                           kq_mask,
         int *                                   sections,
-        int                                     il) {
+        int                                     il,
+        bool                                    gather) {
     const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
 
     const int64_t idx_dim  = hparams.indexer_head_size;
@@ -519,7 +528,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
         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->mask_row  = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream);
 
+        ggml_set_input(qsa->mask_row);
         ggml_set_input(qsa->cell_blk);
         ggml_set_input(qsa->blk_cells);
         ggml_set_input(qsa->blk_pos);
@@ -589,6 +600,31 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
         score = ggml_add(ctx0, score, inp->bias);
     }
 
+    // block-level top-k for the gather path: the scores are one-per-block anyway, so pick
+    // the winning blocks first and expand only those to cell indices via blk_cells. this
+    // sorts n_blocks instead of n_kv entries and skips every O(n_kv) op below. the cell
+    // count r*K_blk stays a multiple of 256, which keeps flash attention's padding happy.
+    // needs blk_bias: the per-block bias is what carries visibility into the selection.
+    if (gather && blk_bias && 256 % r == 0) {
+        const int64_t K_blk = std::min<int64_t>(n_blocks,
+                GGML_PAD(((int64_t) hparams.indexer_top_k + r - 1)/r, 256/r));
+
+        ggml_tensor * blk_idx = ggml_cont(ctx0, ggml_top_k(ctx0, score, K_blk)); // I32 [K_blk, n_tps, n_stream]
+        blk_idx = ggml_reshape_4d(ctx0, blk_idx, K_blk, 1, n_stream, 1);
+
+        // blk_cells maps (block, slot-in-block) -> cell: view one row of r cells per block
+        ggml_tensor * bc = ggml_view_4d(ctx0, inp->blk_cells, r, n_blocks, 1, n_stream,
+                r*ggml_element_size(inp->blk_cells),
+                r*n_blocks*ggml_element_size(inp->blk_cells),
+                inp->blk_cells->nb[1], 0);
+
+        ggml_tensor * cells = ggml_get_rows(ctx0, bc, blk_idx);                  // I32 [r, K_blk, 1, n_stream]
+        cells = ggml_reshape_4d(ctx0, cells, r*K_blk, n_tps, 1, n_stream);
+        cb(cells, "indexer_top_k_cells", il);
+
+        return cells;
+    }
+
     // 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);
@@ -604,7 +640,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     cb(expanded, "indexer_score_tokens", il);
 
     // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail
-    const int64_t width = std::min<int64_t>(n_kv, (int64_t) hparams.indexer_top_k + r - 1);
+    // for the gather path the width is padded to a multiple of GGML_KQ_MASK_PAD so the
+    // gathered K/V/mask satisfy flash attention's padding requirements without extra ops;
+    // the gathered mask carries -INF for any surplus cells the padded top-k pulls in
+    const int64_t width = gather
+        ? std::min<int64_t>(n_kv, GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256))
+        : std::min<int64_t>(n_kv, (int64_t) hparams.indexer_top_k + r - 1);
 
     ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width));
 
@@ -623,8 +664,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
         ggml_tensor *             k_cur,
         ggml_tensor *             v_cur,
         ggml_tensor *             top_k,
+        ggml_tensor *             mask_row,
         float                     kq_scale,
-        int                       il) {
+        int                       il,
+        bool                      gather) {
     // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer
     // has already scored with its own query in build_qsa_top_k, so top_k is unaffected.
     if (inp->self_k_rot) {
@@ -656,6 +699,72 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
 
     ggml_tensor * kq_mask = inp->get_kq_mask();
 
+    // decode-time sparse gather: instead of masking the full cache (which makes flash
+    // attention scan all n_kv cells to use only indexer_top_k of them), copy the selected
+    // cells out of the cache and attend densely over just those. this bounds the per-token
+    // attention cost by the top-k width instead of the context depth.
+    // restricted to single-token-per-stream ubatches (decode), where every stream carries
+    // exactly one top-k list, so the gathered K/V stay one tensor per stream.
+    if (gather) {
+        ggml_tensor * kf = mctx_cur->get_k(ctx0, il); // [hd_k, n_head_kv, n_kv, ns]
+        ggml_tensor * vf = mctx_cur->get_v(ctx0, il); // [hd_v, n_head_kv, n_kv, ns]
+
+        const int64_t hd_k   = kf->ne[0];
+        const int64_t hd_v   = vf->ne[0];
+        const int64_t n_h_kv = kf->ne[1];
+        const int64_t n_kv   = kf->ne[2];
+        const int64_t ns     = kf->ne[3];
+        const int64_t n_topk = top_k->ne[0];
+
+        GGML_ASSERT(top_k->ne[1] == 1 && "QSA gather requires single-token-per-stream ubatches");
+        GGML_ASSERT(kq_mask->ne[0] == n_kv);
+
+        // the heads of a cell are contiguous in the cache, so a cell can be gathered as one row
+        GGML_ASSERT(kf->nb[2] == ggml_row_size(kf->type, hd_k*n_h_kv));
+        GGML_ASSERT(vf->nb[2] == ggml_row_size(vf->type, hd_v*n_h_kv));
+
+        ggml_tensor * k_cells = ggml_view_4d(ctx0, kf, hd_k*n_h_kv, n_kv, 1, ns,
+                kf->nb[2], kf->nb[2]*n_kv, kf->nb[3], 0);
+        ggml_tensor * v_cells = ggml_view_4d(ctx0, vf, hd_v*n_h_kv, n_kv, 1, ns,
+                vf->nb[2], vf->nb[2]*n_kv, vf->nb[3], 0);
+
+        // top_k [n_topk, 1, 1, ns] -> the index layout ggml_get_rows expects: [n_topk, 1, ns, 1]
+        ggml_tensor * idx = ggml_reshape_4d(ctx0, top_k, n_topk, 1, ns, 1);
+
+        // get_rows dequantizes the cells to F32; build_attn_mha casts to F16 for flash attention
+        ggml_tensor * k_g = ggml_get_rows(ctx0, k_cells, idx); // F32 [hd_k*n_h_kv, n_topk, 1, ns]
+        ggml_tensor * v_g = ggml_get_rows(ctx0, v_cells, idx); // F32 [hd_v*n_h_kv, n_topk, 1, ns]
+
+        k_g = ggml_reshape_4d(ctx0, k_g, hd_k, n_h_kv, n_topk, ns);
+        v_g = ggml_reshape_4d(ctx0, v_g, hd_v, n_h_kv, n_topk, ns);
+        cb(k_g, "qsa_k_gathered", il);
+        cb(v_g, "qsa_v_gathered", il);
+
+        // gather the same cells' visibility: keeps -INF for any invalid cell the padded
+        // top-k width pulled in (e.g. when fewer than n_topk cells are visible).
+        // reading the compact per-cell mask_row instead of the attention kq_mask leaves
+        // the latter unreferenced, so it is neither filled (O(n_kv * pad) host work) nor
+        // uploaded (n_kv * GGML_KQ_MASK_PAD * 2 bytes, ~18 MB/token at 141K ctx)
+        GGML_ASSERT(mask_row != nullptr);
+        GGML_ASSERT(mask_row->ne[0] == n_kv);
+        ggml_tensor * m1 = ggml_view_4d(ctx0, mask_row, 1, n_kv, 1, ns,
+                mask_row->nb[0], mask_row->nb[1], mask_row->nb[3], 0);
+        ggml_tensor * m_g = ggml_get_rows(ctx0, m1, idx);            // F32 [1, n_topk, 1, ns]
+        m_g = ggml_reshape_4d(ctx0, m_g, n_topk, 1, 1, ns);
+        m_g = ggml_cast(ctx0, m_g, GGML_TYPE_F16);                   // FA wants contiguous F16
+        cb(m_g, "qsa_mask_gathered", il);
+
+        ggml_tensor * cur = build_attn_mha(q_cur, k_g, v_g, nullptr, m_g, nullptr, nullptr, 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;
+    }
+
     // prepare new kq mask - starts filled with -INFINITY
     ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
 
@@ -710,7 +819,27 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
     // indexer reads the same block input as q/k/v; no cache or no ratio means dense
     const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0;
 
-    ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr;
+    // gather-based QSA decode: worth it once the cache is meaningfully deeper than the
+    // top-k width; below that the masked path costs about the same. QWEN4EXP_QSA_GATHER=0
+    // disables it (A/B lever, and an escape hatch).
+    static const bool gather_enabled = [] {
+        const char * e = getenv("QWEN4EXP_QSA_GATHER");
+        return e == nullptr || atoi(e) != 0;
+    }();
+
+    bool gather = false;
+    if (qsa && gather_enabled) {
+        const int64_t r     = hparams.dsv4_compress_ratios[il];
+        const int64_t n_kv  = mctx_hyb->get_idx()->get_n_kv();
+        const int64_t width = GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256);
+
+        const int64_t n_stream = mctx_hyb->get_n_stream();
+
+        gather = n_tokens == n_stream && n_kv >= 2*width;
+
+    }
+
+    ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il, gather) : nullptr;
 
     // Qwen3Next uses a single Q projection that outputs query + gate
     ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ]
@@ -763,7 +892,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
     const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
 
     if (top_k) {
-        cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il);
+        ggml_tensor * qsa_mask_row = nullptr;
+        if (gather) {
+            const auto it = qsa_inps.find((uint32_t) hparams.dsv4_compress_ratios[il]);
+            GGML_ASSERT(it != qsa_inps.end());
+            qsa_mask_row = it->second->mask_row;
+        }
+        cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, qsa_mask_row, kq_scale, il, gather);
     } else {
         cur = build_attn(inp,
                     nullptr, nullptr, nullptr,