vulkan: sparse prefill flash attention for qwen4exp top-k masks

docs/fix-fa-push-constants-group-vs-dst-shape.md

@@ -0,0 +1,71 @@
+
+# Recipe: push constants carry the destination tensor's shape, never the dispatch tile's
+
+## Bug
+
+The Vulkan grouped union prefill path (`ggml_vk_flash_attn_union_groups`) processes the
+batch in groups of 64 query rows. It sized three things by the *group* while the shaders
+expect the *destination tensor's* shape:
+
+1. `vk_flash_attn_push_constants.ne1` was set to the group's row count `N`. `ne1` is the
+   head-to-head stride of the output, and `dst` is `[HSV, n_head_q, n_batch, ns]`:
+
+   ```glsl
+   // flash_attn.comp / flash_attn_cm1.comp / flash_attn_cm2.comp:
+   uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4;
+   data_ov4[o_offset + (iq2*HSV + (i*Br + row)*p.ne1*HSV)/4 + ...] = ...;
+   ```
+
+   It must be `q->ne[2]` (the head count). Passing the row count transposes the output
+   for every head but the first.
+
+2. The group's slice of `dst` advanced by `dst->nb[1]`. One batch row is `n_head_q*HSV`
+   wide, so the batch-row stride is `nb[2]`.
+
+3. The split-K reduce was dispatched as `{N, HSV, neq2}` with
+   `ne1 = N, ne2 = neq2` — the opposite of its own convention (`x` enumerates heads,
+   `z` the rows of the split buffer, `ne1` is the head stride). Compare the dense call
+   site in `ggml_vk_flash_attn`, which is the reference.
+
+## Symptom
+
+`ERR ≈ 1.0` (not a small numeric drift) for `nh != nb`, passing for `nh == nb`; at
+`nb=128` a `DEVICE_LOST` from the out-of-range writes. Split-K shapes (`nb=8`, `nb=16`
+at depth, where `split_k > 1`) failed *independently* of the `ne1` error, so fixing only
+the first one leaves a shape family still broken.
+
+## Fix
+
+- `ne1 = q->ne[2]`, `ne2 = rows` (the group's own row count in the split buffer, matching
+  the allocation and the reduce's `ne2`).
+- `dst_buf.offset += batch_off * dst->nb[2]`.
+- reduce: `{HSV, neq2, N, N, 1, split_k, false}` dispatched over `{neq2, HSV, N}`.
+
+## Debugging technique: bisect the *shape*, and read the rule off the result
+
+The error looked like a shader bug (garbage output), but the decisive observation was
+the *pattern* of pass/fail across shapes, not any single failure: `nh=64` passed while
+`nh=16/24/32/48` failed at fixed `nb=64`. That is a symmetry, and a symmetry says the
+host is passing a quantity that coincides with the right one only in the symmetric case.
+Sweeping `nb` separately from `nh` then inverts it (`nb=24, nh=32` fails while
+`nb=24, nh=24` passes), which pins the bug to the *pair* rather than either dimension.
+
+Two traps cost most of the time here:
+
+- **A verdict of "PASS" is only meaningful if the path was actually taken.** A dense
+  fallback reports PASS. The union gate is a *measurement* (the previous call's overlap
+  count), so on the first call it always declines, and `test-backend-ops` computes each
+  case once: an entire sweep "passed" while measuring dense. Confirm engagement from a
+  counter (`GGML_VK_FA_UNION_STATS=1`), not from a verdict or a timing.
+- **A regex `-p` filter over the case's `vars()` string silently matches several cases**
+  (or none), so "N/N tests passed" may belong to a case you did not mean to run. Assert
+  on the exact expected count (`1/1`), never on the presence of a `tests passed` line
+  (which also matches `0/1 tests passed`).
+
+## General rule
+
+A compute shader's push constants describe the **tensors and their logical indexing**;
+the dispatch geometry (tile coordinates, workgroup counts) is separate. Anything used
+inside an offset expression must be a tensor invariant. When a path tiles a tensor and
+reuses the same push-constant struct, the tiled quantity belongs only in the *dispatch*
+fields, never in the tensor-shape fields.

docs/fix-sparse-compact-subtot-overflow.md

@@ -0,0 +1,37 @@
+
+# Recipe: shared-memory per-subgroup tallies must be sized for the whole workgroup
+
+## Bug
+
+`flash_attn_sparse_compact.comp` (sparse FA mask compaction) kept a
+per-subgroup tally in `shared uint sub_tot[8]`, but the pipeline is
+created with `compact_wg = min(1024, ...)` threads = **16 subgroups**
+on wave-64 GPUs (RADV). Waves 8..15 wrote past the array.
+
+## Symptom
+
+Not a crash: each mask row's compacted index list silently shrank from
+~2051 entries to ~930, keeping ascending order (the lost waves' counts
+zero out, so slots shift and later writes collide/overwrite). The FA
+then attends over a pseudo-random half of the selected cells. The model
+stays locally coherent but loses global structure — long-context
+analysis hallucinates breaks that do not exist. Bit-diff tests at small
+context pass because the sparse path never fires below KV >= 4096 with
+the ratio gate.
+
+## Fix
+
+`shared uint sub_tot[32];` (max subgroups = 1024 threads / 32 lanes).
+General rule: any shared array indexed by `gl_SubgroupID` must be sized
+`max workgroup threads / min subgroup size` (32 on most GPUs, 16 on
+wave-64 AMD if the workgroup is <= 512), not by the value you see on
+your dev machine.
+
+## Debugging technique
+
+Synchronous or even queued-async host readbacks inside the dispatch
+path race the recording (commands are not submitted yet). Snapshot the
+debug state at dispatch (`g_sparse_dbg`), then read after the graph's
+fence in `ggml_backend_vk_graph_compute` (`GGML_VK_SPARSE_VALIDATE`).
+Validation criteria: list length == mask finite count, strictly
+ascending, all in-bounds.

docs/fix-sparse-fa-shared-tile-list.md

@@ -0,0 +1,85 @@
+# One index list per dispatch tile: the sparse-FA tiling contract
+
+## Symptom
+
+The sparse mask-compaction Flash Attention path returned plausible-looking but
+wrong answers on prefill shapes: a full-model recall probe over a 36k-token
+prompt answered "NO" to a question whose answer ("YES") was in the file, while
+the same model with `GGML_VK_FA_SPARSE_DISABLE=1` answered "YES". Decode and
+small-batch shapes were unaffected.
+
+## Root cause
+
+`flash_attn_base.glsl` resolves **one** sparse index list and **one** mask row
+for the whole dispatch tile, not per row:
+
+```glsl
+uint32_t qrow = (p.gqa_ratio > 1) ? gqa_iq1 : (i * Br);
+sparse_base = (((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 + qrow) * p.split_kv;
+```
+
+with `m_stride == 0` under GQA (every row of the tile shares one mask row, and
+`m_offset = gqa_iq1 * m_row_len`). The consumer shaders assume the same thing,
+and `flash_attn_cm1.comp` states it outright:
+
+```glsl
+// sparse is gqa-gated (m_stride == 0): all four rows share the value
+FLOAT_TYPE mv = FLOAT_TYPE(data_m[m_offset + kcol]);
+```
+
+That is correct only when the tile's rows are the GQA heads of a **single**
+query. The host folds GQA only when `N <= 8`, so any shape with `N > 8` runs
+with `gqa_ratio == 1`, where a `Br`-row tile (16, or 64 for cm2) spans 16
+*different* tokens: every row then attended the tile's first token's selection
+and read its mask row. Silent wrong answers, no crash.
+
+## How it was localized
+
+A full-model A/B is far too slow for this. It reproduces in seconds in
+`test-backend-ops` once the test actually sets the sparse hint (the pre-existing
+top-k test left `op_params[5]` at 0, so every "sparse" case was silently
+measuring dense):
+
+| picks per query row | ERR vs CPU | reading |
+| --- | --- | --- |
+| different per token (the real case) | 1.85 | broken |
+| identical within a 16-row tile | 0.035 | share-the-tile hypothesis |
+| identical across the whole batch | 0.00055 (pass) | only row 0 is right |
+
+The monotone fall of the error with the amount of sharing is the signature of a
+shared-tile read. Confirming it directly: forcing a single-row tile
+(`block_rows = 1`) makes the per-token case pass at nb = 64/128/512 — it is
+correct, but ~4x slower than dense (147 vs 607 GFLOPS at nb=512, depth 32768),
+so it is not a fix, it is a measurement that pins the cause.
+
+## Fix
+
+Gate the sparse path on the property the tiling actually needs, encoded in the
+host as `gqa_ratio > 1` rather than `N >= 64`, and write the contract down next
+to the gate. Prefill then declines to dense (correct, slower); decode keeps the
+sparse path.
+
+Note the direction of the tradeoff: with the shared tile the prefill numbers
+look *good* (it is the same work with 1/Br of the gathers) and the answers are
+wrong. Prefer the gate that is provably correct; a per-tile **union** consumer
+(one list per tile = the union of its rows' selections, since every row of the
+tile then legitimately shares it) is how prefill gets the speed back correctly.
+
+## Reusable checklist
+
+- A shared assumption in a tiled kernel is a correctness condition, not an
+  optimization detail. When a tile resolves one index/row/list for all its rows,
+  find the shape invariant that makes those rows interchangeable, and gate on
+  *that*, not on a proxy (`N`, batch size, depth).
+- Before trusting any sparse/perf measurement, verify the path was taken at all:
+  env-gated one-line `fprintf` is enough, and it must be removed before commit.
+  In this episode the test never set the sparse hint and an upstream comparison
+  was measuring its *dense fallback* — a wrong conclusion that survived several
+  rounds of reasoning.
+- Sweep the amount of sharing (per-token -> per-tile -> per-batch). If the error
+  collapses monotonically to zero as sharing grows, the bug is a shared read, not
+  a bad index or a stride error.
+- Check the generated artifacts get rebuilt: `ggml-vulkan-shaders.hpp` did not
+  list the `.comp` files in its `DEPENDS`, so shader edits were not compiled and
+  measurements silently described old code. `touch`/remove the generated header
+  after editing any shader.

docs/qwen4exp-mtp-sparse-draft.md

@@ -0,0 +1,87 @@
+# qwen4exp: sparse QSA for the MTP draft block
+
+## Symptom
+
+`llama-server --spec-type draft-mtp` on qwen4exp loses pp512 with context depth far below
+the no-draft server: 6936 / 4494 / 2235 / 790 t/s at d0 / 8192 / 32768 / 131072 on the
+micro model against 10631 / 6445 / 3943 / 1983 without the draft. `GGML_VK_PERF_LOGGER`
+shows why: every prefill ubatch runs a SECOND graph (the draft, `graph_mtp`), and its one
+attention layer does a DENSE full-context FA over the whole ubatch (508 rows at ub 512).
+That FA grows linearly with depth: ~37.6 ms of a ~49 ms draft window at d32768, ~390 ms
+per prefill at d131072.
+
+## Why the draft ran dense, and why that was wrong
+
+`graph_mtp` hardcoded `mctx_hyb = nullptr`, so `build_layer_attn`'s QSA gate
+(`mctx_hyb != nullptr && get_idx() != nullptr && dsv4_compress_ratios[il] > 0`) was false,
+and the MTP context carried a plain `llama_kv_cache` (the deepseek32 MTP pattern).
+
+The MTP block of this architecture is a full-attention QSA layer:
+
+1. The real sidecar (`Qwen3.8-Flash-Next-MTP-Q4_K_M.gguf`, `blk.48`) ships its own
+   indexer tensors (`indexer.{q_proj,k_proj,q_norm,k_norm}`) with values distinct from
+   every trunk indexer layer.
+2. The indexer norms sit far from their zero initialization, so they were trained: the
+   block's attention used the indexer during training.
+3. The reference config says `mtp.layer_types = ["full_attention"]`; in this architecture
+   every full-attention layer is a QSA layer.
+
+The sidecar's `compress_ratios[48] == 0` is a padding artifact: `llama-model-saver.cpp`
+writes the array out to `n_layer_all` from the trunk's array, whose tail is zero. It does
+NOT mean the draft runs dense.
+
+## What the fix does
+
+1. **Loader** (`load_arch_hparams`): when a nextn layer has a zero ratio but carries
+   `blk.<il>.indexer.q_proj.weight`, restore its ratio from the trunk's QSA layers. A
+   genuinely dense MTP block (no indexer tensors) stays dense.
+2. **Memory** (`create_memory`): the MTP context gets `llama_memory_hybrid_idx` with the
+   trunk's filters inverted - attention + indexer over `il >= n_layer()`, and a recurrent
+   filter that is always false. The old comment ("a hybrid memory with an empty recurrent
+   layer set fails its buffer allocation") is wrong: an all-false filter leaves the
+   recurrent cache without tensors and without buffers, and `find_slot` still succeeds on
+   its single spare cell.
+3. **Graph** (`graph_mtp`): build `build_inp_mem_hybrid()` and pass the hybrid-idx
+   context into `build_layer_attn`, like the mainline graph does.
+4. **Prefill-only gate**: `graph_mtp` enables QSA only for `n_tokens >= 16`. A decode or
+   verify batch of 1-4 rows pays more for the indexer pipeline (k_proj, pooling, top-k,
+   plus the O(n_kv) host scan in `set_input_qsa`) than sparse FA saves: measured on the
+   micro model, ~+2-3 ms per draft decode window against ~1.6-2 ms of dense FA, which is
+   a ~25% tg loss. Prefill batches amortize the pipeline and win big at depth.
+5. **Indexer keys stay written**: when the gate declines, `build_layer_attn` still runs
+   one `index_k_proj` matmul + `cpy_k` into the index cache. The pooled-key cache
+   recomputes every block above its watermark from the stored keys, so a hole in the key
+   cache would poison a later sparse read. The watermark only advances inside
+   `set_input_qsa`, so not touching the pool on dense ubatches is safe by construction.
+
+Also in this change: `llm_graph_input_mem_hybrid{,_k,_iswa}::set_input` skip an `s_copy`
+input that no node reads (ggml-alloc assigns no buffer to it). The MTP graph reads no
+recurrent state, and the old unconditional dereference aborted on it.
+
+## Micro-model measurements (llama-server, pp512 / tg128, same session)
+
+| depth | dense draft pp | QSA draft pp | dense draft tg | QSA draft tg |
+| --- | --- | --- | --- | --- |
+| 0 | 6258 | 6157 | 124 | 124 |
+| 8192 | 4118 | 4579 | 114 | 112 |
+| 32768 | 2120 | 2350 | 94 | 87 |
+| 131072 | 734 | 1019 | 53 | 49 |
+
+Prefill is at or above the dense draft at every depth and +39% at d131072; the draft's
+per-ubatch FA is depth-independent once the union takes (the sparse draft attends the
+compact union, not the whole cache). Generation time stays within a few percent of the
+dense draft.
+
+## Open question for the real model
+
+The micro model's weights are random, so the draft's indexer produces top-k selections
+with only 19-39% overlap between adjacent tokens; the union gate then declines at
+shallow/mid depth and the sparse path pays the indexer without the FA win there. The
+trunk's random indexer reaches 93% overlap and the real trunk measures ov=86, so trained
+weights should compact much tighter. Before calling this done on the real model:
+
+1. A/B acceptance rate and tg of the sparse draft against `GGML_VK_FA_SPARSE_DISABLE=1`
+   on the same server (the draft only proposes; the target verifies, so the text is
+   identical either way - only the acceptance rate moves).
+2. Check the union gate engages for the draft from shallow depth
+   (`GGML_VK_FA_UNION_STATS`), which is where the remaining prefill win lives.

ggml/include/ggml.h

@@ -2433,6 +2433,12 @@ extern "C" {
     GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec(
             const struct ggml_tensor * a);
 
+    // Use finite mask entries as a sparse K/V set. Set 0 to disable.
+    // n_kv_max must bound the number of finite entries in every mask row.
+    GGML_API void ggml_flash_attn_ext_set_sparse(
+            struct ggml_tensor * a,
+            int32_t              n_kv_max);
+
     GGML_API void ggml_flash_attn_ext_add_sinks(
             struct ggml_tensor * a,
             struct ggml_tensor * sinks);

ggml/src/ggml-cuda/fattn-common.cuh

@@ -718,6 +718,9 @@ static __global__ void flash_attn_mask_to_KV_max(
     KV_max[sequence*ne31 + jt] = KV_max_sj;
 }
 
+void ggml_cuda_flash_attn_ext_compact_mask(
+        const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream);
+
 template<int D, int ncols1, int ncols2> // D == head size
 __launch_bounds__(D, 1)
 static __global__ void flash_attn_stream_k_fixup_uniform(
@@ -972,8 +975,8 @@ static __global__ void flash_attn_combine_results(
 template <int DV, int ncols1, int ncols2>
 void launch_fattn(
     ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared,
-    const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE,
-    const bool stream_k_strided_eligible = false
+    const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse,
+    const int warp_size = WARP_SIZE, const bool stream_k_strided_eligible = false
 ) {
     constexpr int ncols = ncols1 * ncols2;
 
@@ -1089,10 +1092,20 @@ void launch_fattn(
     const int ntiles_z_gqa = ((gqa_ratio + ncols2 - 1) / ncols2);
     const int ntiles_dst   = ntiles_x * ntiles_z_gqa * K->ne[2] * Q->ne[3];
 
+    const int32_t n_kv_max = use_sparse ? ggml_get_op_params_i32(KQV, 5) : 0;
+    if (use_sparse) {
+        GGML_ASSERT(mask != nullptr);
+        GGML_ASSERT(n_kv_max > 0);
+        const size_t mask_rows = size_t(mask->ne[1]) * mask->ne[3];
+
+        KV_max.alloc(size_t(n_kv_max) * mask_rows);
+        ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, n_kv_max, main_stream);
+    }
+
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
     //     multiple sequences of possibly different lengths.
-    if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
+    if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
 
@@ -1115,7 +1128,8 @@ void launch_fattn(
     GGML_ASSERT(max_blocks_per_sm > 0);
     int parallel_blocks = max_blocks_per_sm;
 
-    const int ntiles_KV = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.
+    const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1];
+    const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.
 
     dim3 blocks_num;
     bool stream_k_strided = false;
@@ -1221,7 +1235,7 @@ void launch_fattn(
         !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr,
         scale, max_bias, m0, m1, n_head_log2, logit_softcap,
         Q->ne[0], ne01,     Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3],
-        K->ne[0], K->ne[1], K->ne[2], K->ne[3], nb11, nb12, nb13,
+        K->ne[0], n_kv, K->ne[2], K->ne[3], nb11, nb12, nb13,
         nb21, nb22, nb23,
         mask ? mask->ne[1] : 0, mask ? mask->ne[2] : 0, mask ? mask->ne[3] : 0,
         mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0