qwen4exp: gather-based sparse attention for QSA decode (+ cuda top-k batching fix)

ggml/src/ggml-cuda/top-k.cu

@@ -36,7 +36,9 @@ static void top_k_cub(ggml_cuda_pool & pool,
                          ncols, k, env));
 }
 
-#elif defined(GGML_CUDA_USE_CUB)  // CUB_TOP_K_AVAILABLE
+#endif                            // CUB_TOP_K_AVAILABLE
+
+#if defined(GGML_CUDA_USE_CUB)
 
 static int next_power_of_2(int x) {
     int n = 1;
@@ -64,13 +66,19 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const int64_t    k     = dst->ne[0];
     ggml_cuda_pool & pool  = ctx.pool();
 #ifdef CUB_TOP_K_AVAILABLE
+    // DeviceTopK selects rather than sorts, but only one row per call, so it wins for
+    // decode-shaped calls (a few rows) and loses badly to one segmented sort when a
+    // prompt batch brings hundreds of rows.
     // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented
     // https://github.com/NVIDIA/cccl/issues/6391
-    // TODO: investigate if there exists a point where parallelized argsort is faster than sequential top-k
-    for (int i = 0; i < nrows; i++) {
-        top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream);
+    if (nrows <= 4) {
+        for (int i = 0; i < nrows; i++) {
+            top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream);
+        }
+        return;
     }
-#elif defined(GGML_CUDA_USE_CUB)  // CUB_TOP_K_AVAILABLE
+#endif // CUB_TOP_K_AVAILABLE
+#if defined(GGML_CUDA_USE_CUB)
     // Fall back to argsort + copy
     const int    ncols_pad      = next_power_of_2(ncols);
     const size_t shared_mem     = ncols_pad * sizeof(int);

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-kv-cells.h

@@ -314,18 +314,31 @@ 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) {
-            if (pos[i] < p0 || pos[i] >= p1) {
-                continue;
+        // scan the raw arrays instead of iterating `used`: walking that RB-tree touches
+        // every used cell through pointer chasing and dominates decode CPU time at long
+        // context (this runs once per ubatch). the array scan visits the same cells in
+        // the same (index) order with contiguous reads. empty cells have pos == -1.
+        const uint32_t n_cells = (uint32_t) pos.size();
+
+        // hoist the queried seq ids out of the loop: testing all LLAMA_MAX_SEQ bits for
+        // every cell costs n_cells*LLAMA_MAX_SEQ tests per call, which dominated decode
+        // CPU time at long context. queries carry a handful of seq ids (usually one).
+        llama_seq_id seq_ids[LLAMA_MAX_SEQ];
+        uint32_t n_seq_ids = 0;
+        for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+            if (seqs.test(s)) {
+                seq_ids[n_seq_ids++] = s;
             }
+        }
 
-            const auto m = seq[i] & seqs;
-            if (m.none()) {
+        for (uint32_t i = 0; i < n_cells; ++i) {
+            if (pos[i] < 0 || pos[i] < p0 || pos[i] >= p1) {
                 continue;
             }
 
-            for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
-                if (m.test(s)) {
+            for (uint32_t k = 0; k < n_seq_ids; ++k) {
+                const llama_seq_id s = seq_ids[k];
+                if (seq[i].test(s)) {
                     f(s, pos[i], ext[i].tok);
                 }
             }

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,