diff --git a/docs/fix-fa-push-constants-group-vs-dst-shape.md b/docs/fix-fa-push-constants-group-vs-dst-shape.md new file mode 100644 index 000000000000..a6ceeb6b8e7e --- /dev/null +++ b/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. diff --git a/docs/fix-sparse-compact-subtot-overflow.md b/docs/fix-sparse-compact-subtot-overflow.md new file mode 100644 index 000000000000..1048bf0bd37c --- /dev/null +++ b/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. diff --git a/docs/fix-sparse-fa-shared-tile-list.md b/docs/fix-sparse-fa-shared-tile-list.md new file mode 100644 index 000000000000..762454608db9 --- /dev/null +++ b/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. diff --git a/docs/qwen4exp-mtp-sparse-draft.md b/docs/qwen4exp-mtp-sparse-draft.md new file mode 100644 index 000000000000..90379a71624d --- /dev/null +++ b/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..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. diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 655b6085add0..83bda9bfef99 100644 --- a/ggml/include/ggml.h +++ b/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); diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index d50284a8efc0..e548723619df 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/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 // 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 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 diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 5b2199f74a84..2254aa4306b8 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -349,20 +349,22 @@ static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, return cp_async_available(cc) && ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2, cc) : 0; } -static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2) { +static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages( + const int DKQ, const int DV, const int ncols1, const int ncols2, const bool use_sparse) { #ifdef CP_ASYNC_AVAILABLE - return ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; + return ncols2 >= 2 && !use_sparse ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; #else - GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2); + GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2, use_sparse); return 0; #endif // CP_ASYNC_AVAILABLE } // ------------------------------------------------------------------------------------------------------------------ -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( - const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) { + const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, + const int i_sup, const int32_t * const __restrict__ indices) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); // K/V data is loaded with decreasing granularity for D for better memory bandwidth. // The minimum granularity is 16 bytes. @@ -371,6 +373,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( if constexpr (use_cp_async) { static_assert(warp_size == 32, "bad warp_size"); static_assert(!oob_check, "OOB check not compatible with cp_async"); + static_assert(!use_sparse, "sparse gather not compatible with cp_async"); constexpr int preload = 64; const unsigned int tile_KV_32 = ggml_cuda_cvta_generic_to_shared(tile_KV); @@ -432,8 +435,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, - !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + if constexpr (use_sparse) { + const int32_t index = i < i_sup ? indices[i] : -1; + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, + index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero); + } else { + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, + !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + } } } }; @@ -447,14 +456,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( } } -template +template static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const half * const __restrict__ mask_h, half * const __restrict__ tile_mask, - const int stride_mask, const int i_sup, const int j0, const uint3 ne01) { + const int stride_mask, const int i_sup, const int j0, const uint3 ne01, + const int32_t * const __restrict__ indices) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); if constexpr (use_cp_async) { static_assert(nbatch_fa <= 8*warp_size && nbatch_fa % 8 == 0, "bad nbatch_fa"); static_assert(!oob_check, "OOB check incompatible with cp_async"); + static_assert(!use_sparse, "sparse gather incompatible with cp_async"); constexpr int preload = nbatch_fa >= 32 ? nbatch_fa * sizeof(half) : 64; constexpr int cols_per_warp = 8*warp_size/nbatch_fa; constexpr int stride_j = nwarps * cols_per_warp; @@ -474,7 +485,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( cp_async_cg_16(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + i); } - } else if constexpr (oob_check) { + } else if constexpr (oob_check || use_sparse) { #pragma unroll for (int j1 = 0; j1 < ncols1; j1 += nwarps) { const int j_sram = j1 + threadIdx.y; @@ -488,7 +499,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( for (int i0 = 0; i0 < nbatch_fa; i0 += warp_size) { const int i = i0 + threadIdx.x; - tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + i] : half(0.0f); + if constexpr (use_sparse) { + const int32_t index = i < i_sup ? indices[i] : -1; + tile_mask[j_sram*(nbatch_fa + 8) + i] = index >= 0 ? mask_h[int64_t(j_vram)*stride_mask + index] : half(-INFINITY); + } else { + tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + i] : half(0.0f); + } } } } else if constexpr (nbatch_fa < 2*warp_size) { @@ -528,13 +544,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( } template static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ indices, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, @@ -566,13 +583,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); constexpr int stride_tile_K = nbatch_K2 + 4; constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4; const int k_VKQ_0 = kb0 * nbatch_fa; + const int32_t * const tile_indices = use_sparse ? indices + k_VKQ_0 : nullptr; #if defined(TURING_MMA_AVAILABLE) T_C_KQ KQ_C[nbatch_fa/(np*(cols_per_warp == 8 ? T_C_KQ::I : T_C_KQ::J))]; #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) @@ -588,13 +606,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool use_cp_async = true; cp_async_wait_all(); __syncthreads(); - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup, nullptr); } else { constexpr bool use_cp_async = nstages == 1; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + if constexpr (use_sparse) { + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, tile_indices); + } else { + flash_attn_ext_f16_load_mask + (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, nullptr); + } } } @@ -607,8 +630,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); + if constexpr (use_sparse) { + flash_attn_ext_f16_load_tile + (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup, tile_indices); + } else { + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup, nullptr); + } if (use_cp_async) { cp_async_wait_all(); } @@ -933,6 +961,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } if constexpr (nstages > 1) { + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading"); // Preload K tile for next iteration: constexpr bool use_cp_async = true; @@ -940,11 +969,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( __syncthreads(); if (!last_iter) { if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup, nullptr); } } @@ -959,8 +988,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); + if constexpr (use_sparse) { + flash_attn_ext_f16_load_tile + (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup, tile_indices); + } else { + flash_attn_ext_f16_load_tile + (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup, nullptr); + } if (use_cp_async) { cp_async_wait_all(); } @@ -1015,7 +1049,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, @@ -1113,12 +1147,13 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) -template +template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ indices, const float * const __restrict__ sinks_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, @@ -1158,7 +1193,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1257,37 +1292,38 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); constexpr bool use_cp_async = true; constexpr bool oob_check = false; constexpr int k_VKQ_sup = nbatch_fa; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); + flash_attn_ext_f16_load_mask + (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup, nullptr); } // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. - if constexpr (ncols2 == 1) { + if constexpr (ncols2 == 1 || use_sparse) { constexpr bool oob_check = true; for (; kb0 < kb0_stop-1; ++kb0) { constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } constexpr bool last_iter = true; const int k_VKQ_sup = ne11 - kb0*nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } else { @@ -1296,18 +1332,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } constexpr bool last_iter = true; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } @@ -1692,7 +1728,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, kb0_start, kb0_stop); @@ -1700,7 +1736,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } -template +template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, @@ -1731,7 +1767,8 @@ static __global__ void flash_attn_ext_f16( const char * GGML_CUDA_RESTRICT V = V_ptr; const char * GGML_CUDA_RESTRICT mask = mask_ptr; const char * GGML_CUDA_RESTRICT sinks = sinks_ptr; - const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; + const int * GGML_CUDA_RESTRICT KV_max = use_sparse ? nullptr : KV_max_ptr; + const int32_t * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr; float * GGML_CUDA_RESTRICT dst = dst_ptr; float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; @@ -1855,6 +1892,7 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt)*ne11 : nullptr; const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; @@ -1864,13 +1902,13 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } @@ -1901,6 +1939,7 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt)*ne11 : nullptr; const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; @@ -1910,8 +1949,8 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, @@ -1927,6 +1966,14 @@ static __global__ void flash_attn_ext_f16( #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) } +static constexpr bool ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse( + const int DKQ, const int DV, const int ncols1, const int ncols2) { + return (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8) || + (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16); +} + +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; @@ -1975,29 +2022,57 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml using fattn_kernel_ptr_t = fattn_kernel_t; #endif // defined(GGML_USE_HIP) fattn_kernel_t fattn_kernel; + bool use_sparse = false; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; #if defined(GGML_USE_HIP) if constexpr (DKQ == 64 && DV == 64 && ncols1 == 8 && ncols2 == 8) { fattn_kernel = stream_k_strided ? - flash_attn_ext_f16 : - flash_attn_ext_f16; + flash_attn_ext_f16 : + flash_attn_ext_f16; } else #endif +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) { + constexpr bool use_sparse_kernel = true; + fattn_kernel = flash_attn_ext_f16; + use_sparse = true; + + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } + } else + { + constexpr bool use_sparse_kernel = false; + fattn_kernel = flash_attn_ext_f16; + + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } + } + } else +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) { - fattn_kernel = flash_attn_ext_f16; - } + constexpr bool use_sparse_kernel = false; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) - static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; - if (!shared_memory_limit_raised[id]) { - CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); - shared_memory_limit_raised[id] = true; - } + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } #endif // !defined(GGML_USE_MUSA) + } } else { constexpr bool use_logit_softcap = true; - fattn_kernel = flash_attn_ext_f16; + constexpr bool use_sparse_kernel = false; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -2009,7 +2084,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml } launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host, stream_k_strided); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host, stream_k_strided); } diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index 2959c569b1c2..0e9d2f5739fa 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -1418,7 +1418,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, false, warp_size); return; } } @@ -1434,7 +1434,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, false, warp_size); return; } } @@ -1446,7 +1446,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, false, warp_size); return; } } @@ -1458,7 +1458,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, false, warp_size); return; } } @@ -1470,7 +1470,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, false, warp_size); return; } } @@ -1481,7 +1481,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); fattn_kernel_t fattn_kernel = flash_attn_tile; launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, warp_size); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, type_K == GGML_TYPE_F16, type_V == GGML_TYPE_F16, false, false, warp_size); return; } diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 69dd93686243..519b36b9ff49 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -540,7 +540,7 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false, false); } template diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index faa45e968c4d..d3edf525c54f 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -5,11 +5,147 @@ #include "fattn-vec.cuh" #include "fattn.cuh" +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +__launch_bounds__(256, 1) +static __global__ void flash_attn_mask_to_sparse_indices( + const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max, + const int64_t s31, const int64_t s33) { + ggml_cuda_pdl_sync(); + + constexpr int values_per_lane = 8; + const int tid = threadIdx.x; + const int warp = tid / WARP_SIZE; + const int lane = tid % WARP_SIZE; + const int sequence = blockIdx.y; + const int query = blockIdx.x; + + const half * mask = mask_ptr + sequence*s33 + query*s31; + int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + query)*n_kv_max; + + __shared__ int warp_offsets[256/WARP_SIZE]; + __shared__ int row_count; + __shared__ int chunk_count; + + if (tid == 0) { + row_count = 0; + } + __syncthreads(); + + for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) { + uint32_t selected_warp[values_per_lane]; + int warp_count = 0; +#pragma unroll + for (int item = 0; item < values_per_lane; ++item) { + const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; + const bool selected = i < ne30 && isfinite(__half2float(mask[i])); + selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected); + warp_count += __popc(selected_warp[item]); + } + + if (lane == 0) { + warp_offsets[warp] = warp_count; + } + __syncthreads(); + + if (tid == 0) { + int offset = 0; + for (int iw = 0; iw < 256/WARP_SIZE; ++iw) { + const int count = warp_offsets[iw]; + warp_offsets[iw] = offset; + offset += count; + } + chunk_count = offset; + } + __syncthreads(); + + const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1; + int warp_item_offset = 0; +#pragma unroll + for (int item = 0; item < values_per_lane; ++item) { + const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; + const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask); + if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) { + indices[dst] = i; + } + warp_item_offset += __popc(selected_warp[item]); + } + __syncthreads(); + + if (tid == 0) { + row_count += chunk_count; + } + __syncthreads(); + } + + ggml_cuda_pdl_lc(); + + const int count = row_count; + for (int i = count + tid; i < n_kv_max; i += blockDim.x) { + indices[i] = -1; + } + if (count > n_kv_max) { + if (tid == 0) { + printf("flash attention sparse mask row exceeds n_kv_max (%d > %d)\n", count, n_kv_max); + __trap(); + } + } +} +#endif + +void ggml_cuda_flash_attn_ext_compact_mask( + const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(mask, indices, n_kv_max, stream); + GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA"); +#else + const int64_t s31 = mask->nb[1] / sizeof(half); + const int64_t s33 = mask->nb[3] / sizeof(half); + const dim3 blocks_num(mask->ne[1], mask->ne[3], 1); + const dim3 block_dim(256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params(blocks_num, block_dim, 0, stream); + ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33); + CUDA_CHECK(cudaGetLastError()); +#endif +} + +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(ctx, dst); + return false; +#else + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * mask = dst->src[3]; + const int cc = ggml_cuda_info().devices[ctx.device].cc; + + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4); + return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && + mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f && + mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 && + K->ne[1] >= std::max(4096, 2LL*n_kv_max); +#endif +} + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * Q = dst->src[0]; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + return; + } + } +#endif + if constexpr (ncols2 <= 8) { if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 738da961a0f1..5d2daa1f827f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -658,19 +658,25 @@ static constexpr std::initializer_list snake_pattern { GGM GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; -// qwen4 QSA indexer: gather per-block scores to cells + add f16 mask (cast+reshape) + top-k, -// fused into one radix-select. The cast/reshape are elided; the raw f16 mask is read in-shader. -static constexpr std::initializer_list topk_qsa_pattern { GGML_OP_GET_ROWS, GGML_OP_PERMUTE, - GGML_OP_CONT, GGML_OP_CPY, - GGML_OP_RESHAPE, GGML_OP_ADD, - GGML_OP_TOP_K }; +// qwen4 QSA indexer: transpose the block score into cell rows, gather the cells, add the f16 +// mask (cast+reshape) and radix-select, all fused into one kernel. +// The pattern starts at the transpose in front of the gather: the kernel wants the values, not +// the reordered copy, so the copy is folded in and the score is read in its native block-major +// layout, where a row of consecutive blocks is contiguous. The mask cast/reshape are folded as +// well and the raw f16 mask is read in-shader. Every node still exists for the unfused +// fallback (small k) and for the other backends. +static constexpr std::initializer_list topk_qsa_pattern { GGML_OP_CONT, GGML_OP_GET_ROWS, + GGML_OP_PERMUTE, GGML_OP_CONT, + GGML_OP_CPY, GGML_OP_RESHAPE, + GGML_OP_ADD, GGML_OP_TOP_K }; static constexpr std::initializer_list> topk_qsa_edges { - { 1, 0, 0 }, // permute->src[0] == get_rows - { 2, 0, 1 }, // cont->src[0] == permute - { 4, 0, 3 }, // reshape->src[0] == cpy (mask cast) - { 5, 0, 2 }, // add->src[0] == cont - { 5, 1, 4 }, // add->src[1] == reshape - { 6, 0, 5 }, // top_k->src[0] == add + { 1, 0, 0 }, // get_rows->src[0] == the transpose of the block score + { 2, 0, 1 }, // permute->src[0] == get_rows + { 3, 0, 2 }, // cont->src[0] == permute + { 5, 0, 4 }, // reshape->src[0] == cpy (mask cast) + { 6, 0, 3 }, // add->src[0] == post-gather cont + { 6, 1, 5 }, // add->src[1] == reshape + { 7, 0, 6 }, // top_k->src[0] == add }; //node #978 ( SOFT_MAX): ffn_moe_probs-15 ( 0K) [Vulka ] use=2: ffn_moe_logits-15 ( 0K) [Vulka ] @@ -1141,6 +1147,8 @@ struct vk_device_struct { std::map, vk_pipeline> pipeline_fa_mask_opt; + vk_pipeline pipeline_fa_sparse_compact; + vk_pipeline pipeline_flash_attn_split_k_reduce; vk_pipeline pipeline_count_experts; vk_pipeline pipeline_mmid_row_lists; @@ -1971,13 +1979,30 @@ struct vk_op_dsv4_hc_post_push_constants { }; static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); +// Shared bitmap capacity in flash_attn_union.comp: 12288 words = 393216 compressed rows. +static constexpr uint32_t VK_FA_UNION_MAX_WORDS = 12288; + +// Query rows per union group. A GQA batch is priced per group, so this is also the batch the +// estimator is keyed by (see the gate in ggml_vk_flash_attn_gather_compact). +static constexpr uint32_t VK_FA_UNION_GROUP_ROWS = 64; +// Per-group slot in the union stat buffer: the scan writes four uints per group (padded compact +// rows, raw union size, candidate count, batch). One slot per group of the largest batch, because +// a prefill dispatches every group's scan and each writes its own slot - the slot is what makes +// scan(g+1) data-independent of FA(g). A buffer sized for one slot leaves all but the first group +// writing past its end. +static constexpr uint32_t VK_FA_UNION_STAT_SLOT = 16; +static constexpr uint32_t VK_FA_UNION_STAT_GROUPS = 256; // covers ub 16384 + struct vk_op_flash_attn_union_push_constants { uint32_t n_kv, n_kv_raw, n_batch, n_top_k, max_union, nbt1, max_words, pad_to, count_only; + uint32_t batch_off; // first top-k row of this group (grouped prefill); 0 otherwise }; // nbk1/nbk3 are in 4-byte WORDS, not elements: the gather relocates K rows verbatim and never // interprets what is in them, so it works for any type whose row is a whole number of words. struct vk_op_flash_attn_gather_union_push_constants { uint32_t n_kv, n_kv_raw, kv_c_max, nbk1, nbm1, n_batch, row_words; + uint32_t src_head_stride; // source KV-head stride in words; 0 for the single-head MLA row + uint32_t batch_off; // first mask row of this group (grouped prefill); 0 otherwise }; struct vk_op_flash_attn_gather_push_constants { uint32_t n_kv, n_kv_raw, n_top_k, kv_c; @@ -2174,6 +2199,16 @@ struct vk_op_flash_attn_mask_opt_push_constants { uint32_t nbd3; }; +struct vk_op_flash_attn_sparse_compact_push_constants { + uint32_t KV; + uint32_t nem1; + uint32_t nem2; + uint32_t nbm1; + uint32_t nbm2; + uint32_t nbm3; + uint32_t n_kv_max; +}; + // Allow pre-recording command buffers struct vk_staging_memcpy { vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {} @@ -2533,7 +2568,9 @@ struct ggml_backend_vk_context { // 0.24 at 8) and because a speculative decode varies the batch with the accept count, so a // single slot would be invalidated on nearly every step. This path caps the batch at 64. vk_buffer fa_union_stat; - float fa_union_est_ratio[64]; // union / candidates, decaying peak; 0 = unseeded + // Indexed by batch size, so a ub larger than this array would read past its end: the read + // site clamps, and 1024 covers the largest ub the server can be started with. + float fa_union_est_ratio[1024]; // union / candidates, decaying peak; 0 = unseeded uint64_t fa_union_declines; vk::Fence fence, almost_ready_fence; bool submit_pending {}; @@ -4119,7 +4156,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ } static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool aligned, bool f32acc, - bool use_mask, bool use_mask_opt, bool use_logit_softcap, ggml_type k_type, ggml_type v_type, + bool use_mask, bool use_mask_opt, bool use_logit_softcap, bool use_sparse, ggml_type k_type, ggml_type v_type, bool use_dynamic_kv = false) { const bool old_amd_windows = device->vendor_id == VK_VENDOR_ID_AMD && device->driver_id == vk::DriverId::eAmdProprietary && (device->architecture == AMD_GCN || device->architecture == AMD_RDNA1 || device->architecture == AMD_RDNA2); @@ -4128,7 +4165,8 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const (use_mask ? 2 : 0) | (use_logit_softcap ? 4 : 0) | (old_amd_windows ? 8 : 0) | - (use_dynamic_kv ? 16 : 0); + (use_sparse ? 16 : 0) | + (use_dynamic_kv ? 32 : 0); const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size; @@ -6126,6 +6164,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, it.second, "fa_mask_opt", fa_mask_opt_len, fa_mask_opt_data, "main", 2, sizeof(vk_op_flash_attn_mask_opt_push_constants), {1, 1, 1}, {128, 128 / device->subgroup_size, BrBc.first, BrBc.second}, 1, true, true, device->subgroup_size); } + { + // Large workgroup so the per-row KV scan parallelizes; capped to device limits. + const uint32_t compact_wg = std::min({1024u, device->properties.limits.maxComputeWorkGroupInvocations, device->properties.limits.maxComputeWorkGroupSize[0]}); + ggml_vk_create_pipeline(device, device->pipeline_fa_sparse_compact, "fa_sparse_compact", fa_sparse_compact_len, fa_sparse_compact_data, "main", 2, sizeof(vk_op_flash_attn_sparse_compact_push_constants), {1, 1, 1}, {compact_wg}, 1, true); + } + if (device->subgroup_clustered && device->subgroup_require_full_support) { ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4, "quantize_q8_1_x4", quantize_q8_1_x4_subgroup_len, quantize_q8_1_x4_subgroup_data, "main", 2, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1, true, true); } else { @@ -6403,10 +6447,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // large-k fallback: one workgroup per row, radix-select instead of a full sort. The QSA // variant (spec constant 1) additionally gathers the qwen4 indexer input on the fly. + // Spec constant 2 batches the emit scan; without subgroup shuffle the shader keeps the + // per-chunk path. { const uint32_t BLOCK_SIZE = 1u << std::min(10u, device->max_workgroup_size_log2); - ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_f32, "topk_radix_f32", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 0}, 1, true); - ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_qsa, "topk_radix_qsa", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 1}, 1, true); + const int32_t emit_w = device->subgroup_shuffle ? 8 : 1; + ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_f32, "topk_radix_f32", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 0, emit_w}, 1, true); + ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_qsa, "topk_radix_qsa", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 1, emit_w}, 1, true); } ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); @@ -11959,7 +12006,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const uint32_t q_stride = (uint32_t) (q->nb[1] / sizeof(float)); const bool aligned = raw_kv % tuning.block_cols == 0 && (q_stride & 7) == 0 && (k_stride & 7) == 0; const vk_fa_pipeline_state raw_state = get_fa_pipeline_state(ctx->device, tuning, D, D, aligned, f32acc, - true, false, false, GGML_TYPE_F16, GGML_TYPE_F16); + true, false, false, false, GGML_TYPE_F16, GGML_TYPE_F16); if (raw_state.path == FA_COOPMAT1 && ctx->device->pipeline_flash_attn_split_k_reduce) { vk_pipeline raw_pipeline; { @@ -12078,8 +12125,15 @@ struct vk_fa_compact_state { bool separate_v = false; // vc_buf holds V; otherwise V is read from kc_buf uint32_t v_row_bytes = 0; uint32_t v_row_elems = 0; + // Grouped prefill: the batch is processed in groups of group_rows query rows, each + // against the union of its own selections; the compact scratch is reused per group, so + // the allocation is one group's worst case. group_rows == 0 means the whole batch is + // compacted at once (the decode path). + uint32_t group_rows = 0; + uint32_t n_groups = 0; + uint32_t ul_words = 0; // per-group union index list capacity, words vk_subbuffer kv_buf; - vk_subbuffer kc_buf, mc_buf, vc_buf; + vk_subbuffer kc_buf, mc_buf, vc_buf, ul_buf; }; // Small host-visible buffer holding the last union count the device produced: @@ -12092,7 +12146,7 @@ static bool ggml_vk_fa_union_stat_init(ggml_backend_vk_context * ctx) { return ctx->fa_union_stat->ptr != nullptr; } try { - ctx->fa_union_stat = ggml_vk_create_buffer(ctx->device, 64, + ctx->fa_union_stat = ggml_vk_create_buffer(ctx->device, VK_FA_UNION_STAT_SLOT * VK_FA_UNION_STAT_GROUPS, {vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent}); } catch (const vk::SystemError &) { @@ -12101,7 +12155,7 @@ static bool ggml_vk_fa_union_stat_init(ggml_backend_vk_context * ctx) { if (ctx->fa_union_stat->ptr == nullptr) { return false; } - memset(ctx->fa_union_stat->ptr, 0, 64); + memset(ctx->fa_union_stat->ptr, 0, VK_FA_UNION_STAT_SLOT * VK_FA_UNION_STAT_GROUPS); return true; } @@ -12143,13 +12197,16 @@ static uint32_t ggml_vk_fa_union_estimate(ggml_backend_vk_context * ctx, uint32_ const uint32_t cand = stat[2]; const uint32_t nb_obs = stat[3]; - if (u > 0 && cand > 0 && nb_obs > 0 && nb_obs < 64) { + if (u > 0 && cand > 0 && nb_obs > 0 && nb_obs < 1024) { const float r = std::min(1.0f, (float) u / (float) cand); float & e = ctx->fa_union_est_ratio[nb_obs]; e = e > 0.0f ? 0.5f * r + 0.5f * e : r; } - const float ratio = ctx->fa_union_est_ratio[n_batch]; + // The read index is clamped: the array is indexed by batch size and a ub can reach its + // length, which would read the word after it. + const uint32_t nb = std::min(n_batch, (uint32_t) (sizeof(ctx->fa_union_est_ratio) / sizeof(ctx->fa_union_est_ratio[0])) - 1u); + const float ratio = ctx->fa_union_est_ratio[nb]; if (ratio <= 0.0f) { return 0; } @@ -12163,6 +12220,215 @@ static uint32_t ggml_vk_fa_union_estimate(ggml_backend_vk_context * ctx, uint32_ // compact contiguous scratch in prealloc_y, and let the ordinary dense FA below run // over the compacted K/V/mask. Correct by the same contract as the sparse shader: // the source mask carries the selection, and the gathered mask preserves it. +// Grouped union prefill dispatch. For each group of query rows: compact the union of that +// group's selections into the shared scratch, then run flash attention against it. Dispatches +// are ordered, so each group's union/gather may overwrite the previous group's compact set and +// the scratch only needs one group's worst case. Every group sees the same kv bound +// (st.kv_c, 256-padded); rows past a group's own union are zeroed K with an -inf mask, which +// is softmax-neutral, so a uniform bound is correct even though the unions differ. +static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, + ggml_tensor * dst, const vk_fa_compact_state & st) { + const ggml_tensor * top_k = dst->src[5]; + const int32_t n_kv_raw = ggml_get_op_params_i32(dst, 4); + const uint32_t n_batch_total = (uint32_t) q->ne[1]; + const uint32_t HSK = (uint32_t) k->ne[0]; + const uint32_t HSV = (uint32_t) v->ne[0]; + const uint32_t neq2 = (uint32_t) q->ne[2]; + const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16; + const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) neq2)); + const uint32_t mask_n_head_log2 = n_head_log2; // no sinks on this path + const float m0 = 1.0f, m1 = 1.0f; // max_bias == 0 (gated in gather_compact) + float scale = 1.0f; + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + + const uint32_t k_row_words = st.row_bytes / 4; + const uint32_t v_row_words = st.v_row_bytes / 4; + + // Per-group scratch slots. Every group gets its own union-list region and its own + // kv-count slot, so a group's scan has no data dependency on any other group's work: + // scan(g+1) is issued right after FA(g) with NO barrier between them, so it overlaps + // the FA on the GPU instead of paying a serialized 4 ms after it. The gathers still + // reuse the single compact K/V/mask region, so a barrier separates FA(g) from + // gather(g+1) (and scan(g+1) from gather(g+1), the sync below). + // NOTE: the kv-count slot MUST be 4 uints apart because the scan writes 4 words and the + // estimator reads slot 0 (group 0) as the latest measurement. + const size_t ul_slot_sz = (size_t) st.ul_words * sizeof(uint32_t); + const size_t kv_slot_sz = VK_FA_UNION_STAT_SLOT; + const size_t ul_base = st.ul_buf.offset; + + // Group 0's scan first: the first gather needs its list and count. + { + const uint32_t rows = std::min(st.group_rows, n_batch_total); + const vk_op_flash_attn_union_push_constants upc0 = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, rows, (uint32_t) top_k->ne[0], st.ul_words, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), VK_FA_UNION_MAX_WORDS, 256u, 0u, 0u, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, + { ggml_vk_tensor_subbuffer(ctx, top_k), + ggml_vk_subbuffer(ctx, st.ul_buf.buffer, ul_base), + ggml_vk_subbuffer(ctx, st.kv_buf.buffer, 0) }, upc0, { 1, 1, 1 }); + } + + for (uint32_t g = 0; g < st.n_groups; ++g) { + const uint32_t rows = std::min(st.group_rows, n_batch_total - g * st.group_rows); + const uint32_t batch_off = g * st.group_rows; + const vk_subbuffer ul_g = ggml_vk_subbuffer(ctx, st.ul_buf.buffer, ul_base + g * ul_slot_sz); + const vk_subbuffer kv_g = ggml_vk_subbuffer(ctx, st.kv_buf.buffer, g * kv_slot_sz); + + if (g > 0) { + // Independent of everything issued so far: own ul slot, own count slot, + // reads only the top-k tensor. No barrier: may overlap the previous FA. + const vk_op_flash_attn_union_push_constants upc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, rows, (uint32_t) top_k->ne[0], st.ul_words, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), VK_FA_UNION_MAX_WORDS, 256u, 0u, batch_off, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, + { ggml_vk_tensor_subbuffer(ctx, top_k), ul_g, kv_g }, upc, { 1, 1, 1 }); + } + ggml_vk_sync_buffers(ctx, subctx); + + const vk_op_flash_attn_gather_union_push_constants gkpc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, st.kv_c, + (uint32_t) (k->nb[1] / 4), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + rows, k_row_words, + (uint32_t) (k->nb[2] / 4), + batch_off, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, + { ggml_vk_tensor_subbuffer(ctx, k), ul_g, ggml_vk_tensor_subbuffer(ctx, mask), + st.kc_buf, st.mc_buf, kv_g }, gkpc, { st.kv_c, st.n_head_kv, 1 }); + const vk_op_flash_attn_gather_union_push_constants gvpc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, st.kv_c, + (uint32_t) (v->nb[1] / 4), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + rows, v_row_words, + (uint32_t) (v->nb[2] / 4), + batch_off, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, + { ggml_vk_tensor_subbuffer(ctx, v), ul_g, ggml_vk_tensor_subbuffer(ctx, mask), + st.vc_buf, st.mc_buf, kv_g }, gvpc, { st.kv_c, st.n_head_kv, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + + // flash attention against the group's compact set + const uint32_t N = rows; + vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, HSK, HSV, N, st.kv_c, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); + const uint32_t q_stride = (uint32_t) (q->nb[1] / ggml_type_size(q->type)); + const uint32_t alignment = tuning.block_cols; + bool aligned = (st.kv_c % alignment) == 0 && + (q_stride & 7) == 0 && (st.row_elems & 7) == 0 && (st.v_row_elems & 7) == 0; + if (((HSK | HSV) % 16) != 0 && tuning.path == FA_COOPMAT2) { + aligned = false; + } + vk_fa_pipeline_state fa_state = get_fa_pipeline_state(ctx->device, tuning, HSK, HSV, aligned, f32acc, + true, false, false, false, GGML_TYPE_F16, GGML_TYPE_F16, true); + vk_pipeline pipeline = nullptr; + { + std::lock_guard guard(ctx->device->compile_mutex); + auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16; + auto it = pipelines.find(fa_state); + if (it != pipelines.end()) { + pipeline = it->second; + } else { + pipelines[fa_state] = pipeline = std::make_shared(); + } + } + assert(pipeline); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const uint32_t Br = fa_state.Br; + const uint32_t Bc = fa_state.Bc; + GGML_ASSERT(Br == pipeline->wg_denoms[0]); + const uint32_t Tr = CEIL_DIV(N, Br); + const uint32_t workgroups_x = (uint32_t) N; + const uint32_t workgroups_y = neq2; + + uint32_t split_kv = st.kv_c; + uint32_t split_k = 1; + const uint32_t shader_core_count = ctx->device->shader_core_count ? ctx->device->shader_core_count : 16; + const uint32_t total_wgs = Tr * workgroups_y; + if (total_wgs < shader_core_count * 2) { + split_k = shader_core_count * 2 / total_wgs; + } + if (split_k > 1) { + split_kv = ROUNDUP_POW2(std::max(1u, st.kv_c / split_k), alignment); + split_k = CEIL_DIV(st.kv_c, split_kv); + } + + const uint64_t split_k_size = split_k > 1 + ? (HSV * (uint64_t) N * sizeof(float) + (uint64_t) N * sizeof(float) * 2) * split_k * neq2 : 0; + if (split_k_size > ctx->device->properties.limits.maxStorageBufferRange) { + GGML_ABORT("Requested preallocation size is too large"); + } + if (ctx->prealloc_size_split_k < split_k_size) { + ctx->prealloc_size_split_k = split_k_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (split_k > 1 && ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + + const uint32_t eff_nbk2 = st.kv_c * st.row_bytes; + const uint32_t eff_nbk3 = st.n_head_kv * st.kv_c * st.row_bytes; + const uint32_t eff_nbv2 = st.kv_c * st.v_row_bytes; + const uint32_t eff_nbv3 = st.n_head_kv * st.kv_c * st.v_row_bytes; + + vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); + q_buf.offset += batch_off * (uint64_t) q->nb[1]; + // dst is [HSV, n_head_q, n_batch, ns]: one batch row is n_head_q*HSV wide, so the + // group's slice starts nb[2] (not nb[1], which is the head stride) into the tensor. + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + dst_buf.offset += batch_off * (uint64_t) dst->nb[2]; + vk_subbuffer sinks_buf = q_buf; + const vk_subbuffer k_buf = st.kc_buf, v_buf = st.vc_buf, mask_buf = st.mc_buf; + + // ne1/ne2 are the destination's shape, not this group's: ne1 is the head COUNT and sets + // the head-to-head stride of the output (o_offset + iq2*HSV + row*ne1*HSV, dst being + // [HSV, n_head_q, n_batch, ns]), so it must be q->ne[2]; passing the group's row count + // there misplaces every head but the first and only looked right on the nh == nb shapes + // the tests happened to use. ne2 is the row count of the split buffer, which holds this + // group alone (ne3 == ns), and it must agree with the allocation below and with the + // reduce's own ne2. The group's row count is N, used by the tile math and the mask. + const vk_flash_attn_push_constants pc = { N, st.kv_c, + neq2, rows, 1, + neq2, 1, + st.n_head_kv, 1, + st.n_head_kv, 1, + N, 1, 1, + q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3], + st.row_elems, eff_nbk2, eff_nbk3, + st.v_row_elems, eff_nbv2, eff_nbv3, + scale, 0.0f, 0.0f, + mask_n_head_log2, m0, m1, + 1, split_kv, split_k }; + + if (split_k > 1) { + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); + const uint32_t dispatch_x = Tr * split_k * pipeline->wg_denoms[0]; + vk_subbuffer split_k_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, q_buf, kv_g}, + pc, { dispatch_x, workgroups_y, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + // Same convention as the dense call (see ggml_vk_flash_attn): x enumerates HEADS, + // z the split buffer's rows, ne1 is the head stride of both the split buffer and the + // destination. The group's split buffer and dst subbuffer both start at this group's + // first row, so both row counts are N here. + const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, neq2, N, N, 1, split_k, false }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, + {split_k_buf, sinks_buf, dst_buf}, pc2, { neq2, HSV, N }); + } else { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, q_buf, kv_g}, + pc, { workgroups_x, workgroups_y, 1 }); + } + } + ggml_vk_fa_union_stat_host_barrier(subctx); + ctx->prealloc_y_need_sync = true; +} + static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_context & subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, ggml_tensor * dst, vk_fa_compact_state & st) { @@ -12193,7 +12459,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ if ((gather_env && gather_env[0] == '0') || !top_k || !ctx->device->pipeline_flash_attn_gather_f16 || - q->ne[1] < 1 || q->ne[1] >= 64 || // 1..63: >=64 goes to the sparse prefill path + q->ne[1] < 1 || q->type != GGML_TYPE_F32 || !kv_word_addressable || !v_word_addressable || !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || q->ne[0] != k->ne[0] || k->ne[2] != v->ne[2] || n_head_kv == 0 || @@ -12247,26 +12513,64 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // source rows and would decline, while the union measures around 3300 and is well worth // compacting. kv_c stays the worst case for every allocation and dispatch bound, so a // wrong estimate is a slow step, never a wrong answer. - const uint32_t max_words = 12288; // shared bitmap capacity in flash_attn_union.comp + const uint32_t max_words = VK_FA_UNION_MAX_WORDS; const bool bitmap_fits = (uint64_t) ((k->ne[1] - n_kv_raw) + 31) / 32 <= max_words; static const char * union_env = getenv("GGML_VK_FA_TOPK_UNION"); - // The deduplicated union still assumes the MLA row (one KV head, V == K), so a GQA cache - // takes the per-token form below. That only costs it the small-batch/draft case. - if ((!union_env || union_env[0] != '0') && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && - n_head_kv == 1 && !separate_v && + // The union no longer assumes the MLA row: a GQA cache (separate V, several KV heads) is + // served by the same bitmap and index list, because the selection is per token, not per + // head; only the gather gains a head dimension. The dequant-on-gather variant stays + // MLA-only, so a GQA union runs verbatim f16 rows. + const bool gqa_form = n_head_kv != 1 || separate_v; + // The grouped union for a GQA cache is on by default (GGML_VK_FA_TOPK_UNION_GQA=0 opts + // out): the per-row sparse compaction it replaced declined on every prefill batch + // (quadratic in batch), so the grouped union is the only sparse prefill a GQA cache + // gets, and its gate already declines wherever it would not pay. + static const char * union_gqa_env = getenv("GGML_VK_FA_TOPK_UNION_GQA"); + const bool union_gqa_enabled = !(union_gqa_env && union_gqa_env[0] == '0'); + if (union_gqa_enabled && (!union_env || union_env[0] != '0') && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && + (!gqa_form || (v_word_addressable && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16)) && + (uint32_t) CEIL_DIV(n_batch, VK_FA_UNION_GROUP_ROWS) <= VK_FA_UNION_STAT_GROUPS && ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16 && ggml_vk_fa_union_stat_init(ctx)) { - const uint32_t max_union = n_cand; - const uint32_t kv_c_est = ggml_vk_fa_union_estimate(ctx, (uint32_t) n_kv_raw, n_batch, n_cand); + const uint32_t R = (uint32_t) (k->ne[1] - n_kv_raw); + // The economics are per group for a GQA cache: the batch is processed in groups of 64 + // query rows, each against the union of its own selections (see the branch below), so + // the estimate and the probe run with the group's shape. + // + // The group is min(64, n_batch), NOT 64. A speculative decode batch is 2-4 rows, so + // pricing it as a 64-row group reads past the end of the top-k tensor - the shader + // indexes rows [0, 64) of a tensor that has n_batch of them - and files the resulting + // garbage under the same estimate slot a prefill group reads. A poisoned slot drives + // the estimate up to the source size, the prefill gate then fails its k->ne[1] >= + // 2*kv_c_est test at every depth and prefill runs dense: the whole union path silently + // disappears for as long as decoding continues, which is a server with speculative + // decoding but not a benchmark. Keying by the real group size also keeps the two + // shapes' measurements in separate slots instead of overwriting each other. + const uint32_t union_batch = gqa_form ? std::min(VK_FA_UNION_GROUP_ROWS, n_batch) : n_batch; + const uint32_t union_cand = union_batch * (uint32_t) top_k->ne[0]; + const uint32_t max_union = std::min(union_cand, R); + const uint32_t kv_c_bound = (uint32_t) GGML_PAD(n_kv_raw + max_union, 256u); + const uint32_t kv_c_est = ggml_vk_fa_union_estimate(ctx, (uint32_t) n_kv_raw, union_batch, union_cand); // Two separate questions. Does the compact set fit under the gate at all, and does // deduplicating actually shrink it: with no overlap to exploit the union is the same // size as the per-token blocks and the scan is pure cost, measured at 1.2% of the op // at 512k depth. The worst-case bound on the source keeps a collapse in overlap to // roughly dense cost for the one step it takes the estimate to catch up. - const bool worth_it = kv_c_est != 0 && kv_c_est < kv_c && - (uint64_t) k->ne[1] >= 2ull * kv_c_est && - (uint64_t) k->ne[1] >= (uint64_t) kv_c; + // For the union the compact bound is the SOURCE size (one row per distinct cell, so + // never more than n_kv_raw + R), not the per-token kv_c above. + bool worth_it = kv_c_est != 0 && kv_c_est < kv_c_bound && + (uint64_t) k->ne[1] >= 2ull * kv_c_est; + // GGML_VK_FA_UNION_FORCE=1 admits the union without the estimate. The gate above is a + // measurement, so the call that produces it is also the call that must decline on it - + // which leaves the path with no deterministic coverage and no A/B arm, since + // test-backend-ops computes a case once and a timing comparison has to have the path + // taken. These are the two things this switch exists for; it costs only a slow step. + static const char * force_env = getenv("GGML_VK_FA_UNION_FORCE"); + static const bool union_force = force_env && force_env[0] != '\0' && force_env[0] != '0'; + if (union_force) { + worth_it = true; + } // GGML_VK_FA_UNION_STATS=N: report every Nth call (N=1 means every call) what the gate // decided and on what measurement. The alternative is inferring engagement from a @@ -12284,7 +12588,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ if ((calls++ % period) == 0) { fprintf(stderr, "[fa-union] n_kv=%lld n_kv_raw=%d n_batch=%u cand=%u " "union/cand=%.3f kv_c %u -> est %u %s (%llu/%llu taken)\n", - (long long) k->ne[1], n_kv_raw, n_batch, n_cand, (double) ctx->fa_union_est_ratio[n_batch], + (long long) k->ne[1], n_kv_raw, union_batch, union_cand, (double) ctx->fa_union_est_ratio[union_batch], kv_c, kv_c_est, worth_it ? "UNION" : "declined", (unsigned long long) taken, (unsigned long long) calls); } @@ -12303,8 +12607,8 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ ctx->fa_union_declines++; { const vk_op_flash_attn_union_push_constants ppc = { - (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, - (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 1u, + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, union_batch, (uint32_t) top_k->ne[0], max_union, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 1u, 0u, }; const vk_subbuffer stat_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); @@ -12321,6 +12625,55 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ goto union_unavailable; } + if (gqa_form) { + // Grouped union prefill. The batch is processed in groups of UNION_GROUP_ROWS + // query rows, each against the union of its own selections: consecutive tokens + // share most of their selection, so a group's union is a fraction of the full + // cache while a whole 512-row chunk's union approaches it. The compact scratch is + // reused per group (dispatches are ordered), so the allocation is one group's + // worst case, not the batch's. ggml_vk_flash_attn_union_groups runs the + // per-group union/gather/FA sequence; this branch only sizes and files the state. + const uint32_t group_rows = VK_FA_UNION_GROUP_ROWS; + const uint32_t g_max_union = max_union; + const uint32_t g_kv_c = kv_c_bound; + const uint32_t v_row_by = (uint32_t) ggml_row_size(v->type, v->ne[0]); + const size_t gkc_sz = (size_t) n_head_kv * g_kv_c * k_row_bytes; + const size_t gvc_sz = (size_t) n_head_kv * g_kv_c * v_row_by; + const size_t gmc_sz = (size_t) group_rows * g_kv_c * sizeof(ggml_fp16_t); + const size_t gul_sz = (size_t) g_max_union * sizeof(uint32_t); + // One union-list slot per group: the groups' scans are independent and run + // concurrently (see ggml_vk_flash_attn_union_groups), so they cannot share a + // region. The compact K/V/mask region stays single-slot (serialized gathers). + const size_t gul_all = gul_sz * CEIL_DIV(n_batch, group_rows); + const size_t gneed = gkc_sz + gvc_sz + gmc_sz + gul_all; + if (ctx->prealloc_size_y < gneed) { + ctx->prealloc_size_y = gneed; + ggml_vk_preallocate_buffers(ctx, subctx); + } + st.active = true; + st.dynamic_kv = true; + st.kv_c = g_kv_c; + st.n_batch = group_rows; + st.row_bytes = k_row_bytes; + st.row_elems = (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.dequantized = false; + st.n_head_kv = n_head_kv; + st.separate_v = separate_v; + st.v_row_bytes = v_row_by; + st.v_row_elems = (uint32_t) (v->ne[0] / ggml_blck_size(v->type)); + st.group_rows = group_rows; + st.n_groups = CEIL_DIV(n_batch, group_rows); + st.ul_words = g_max_union; + st.kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); + st.vc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, gkc_sz); + st.mc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, gkc_sz + gvc_sz); + st.ul_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, gkc_sz + gvc_sz + gmc_sz); + st.kv_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, st.n_groups); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_gather_union_f16, 2 * st.n_groups); + return true; + } + // Decoding on the way in makes the scratch f16 and hands flash attention its f16 path, // which is worth far more than the extra scratch bytes: the inline decode it replaces // costs a measured 0.15 us per KV row attended, every step. @@ -12353,7 +12706,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ const vk_op_flash_attn_union_push_constants upc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, - (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 0u, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 0u, 0u, }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, { ggml_vk_tensor_subbuffer(ctx, top_k), ul_buf, uc_buf }, upc, { 1, 1, 1 }); @@ -12365,6 +12718,8 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ dq ? (uint32_t) (k->nb[1] / ggml_type_size(k->type)) : (uint32_t) (k->nb[1] / 4), (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), n_batch, dq ? (uint32_t) k->ne[0] : k_row_words, + 0u, // src_head_stride: the MLA row has one head + 0u, // batch_off: the whole batch is one group here }; ggml_vk_dispatch_pipeline(ctx, subctx, gather_pipe, { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), @@ -12386,10 +12741,17 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ st.vc_buf = kc_buf; // V is the K latent here, so it reads the same scratch st.mc_buf = mc_buf; st.kv_buf = uc_buf; + st.group_rows = 0; return true; } union_unavailable:; + // The per-token form is quadratic in batch, so prefill batches never take it: with no + // union available, dense serves them (the cost model above guarantees it is never slower). + if (q->ne[1] >= 64) { + return false; + } + // Per-token blocks have no dedup, so this form really does cost kv_c: the gather writes // then re-reads ~the active bytes while dense reads the source KV once, so compaction only // pays when the source is comfortably larger than the active set. @@ -12576,6 +12938,12 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx // bindings and strides; every other decision then sizes itself to the compact KV. vk_fa_compact_state fa_compact; if (ggml_vk_flash_attn_gather_compact(ctx, subctx, q, k, v, mask, dst, fa_compact)) { + if (fa_compact.group_rows > 0) { + // Grouped union prefill: the per-group union/gather/FA sequence is its own + // dispatch plan, nothing of the single-dispatch tail below applies. + ggml_vk_flash_attn_union_groups(ctx, subctx, q, k, v, mask, dst, fa_compact); + return; + } KV = fa_compact.kv_c; nem0 = fa_compact.kv_c; nem1 = fa_compact.n_batch; @@ -12657,6 +13025,40 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc); + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + if (logit_softcap != 0) { + scale /= logit_softcap; + } + + // Sparse mask hint (op_params[5]): compact the <= n_kv_max finite positions and gather only those. + // A sparse dispatch resolves ONE index list and ONE mask row per TILE, not per row + // (flash_attn_base.glsl: `qrow = (gqa_ratio > 1) ? gqa_iq1 : i * Br`), so every row of a + // tile attends the selection and mask row of the tile's first row. That is only correct + // when those rows are the gqa heads of a single query, which is the gqa_ratio > 1 case; + // the host only folds GQA when N <= 8. Large-N (prefill) shapes therefore have + // gqa_ratio == 1 and must NOT take this path: measured with a single-row tile it is + // correct but ~4x slower than dense (147 vs 607 GFLOPS at nb=512, depth 32768), and with + // a multi-row tile it is fast and silently wrong. Give prefill a per-tile union consumer + // before re-enabling it here. + const int32_t n_kv_max = mask ? ggml_get_op_params_i32(dst, 5) : 0; + static const bool disable_sparse = getenv("GGML_VK_FA_SPARSE_DISABLE") != nullptr; + // cm2 dense is fast, so it needs a larger reduction to win. + const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : 2; + const bool use_sparse = !disable_sparse && n_kv_max > 0 && mask && + max_bias == 0.0f && logit_softcap == 0.0f && + k_type_eff == GGML_TYPE_F16 && v_type_eff == GGML_TYPE_F16 && + nem0 == KV && + (int64_t)KV >= std::max(4096, min_ratio * (int64_t)n_kv_max) && + gqa_ratio > 1 && + !fa_compact.active; + const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type)); uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type)); uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type)); @@ -12691,7 +13093,6 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t)); nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t)); } - const uint32_t alignment = tuning_params.block_cols; bool aligned = (KV % alignment) == 0 && // the "aligned" shader variant will forcibly align strides, for performance @@ -12702,23 +13103,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx aligned = false; } - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; - - memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); - memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); - - if (logit_softcap != 0) { - scale /= logit_softcap; - } - // Only use mask opt when the mask is fairly large. This hasn't been tuned extensively. - bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 + bool use_mask_opt = mask && !use_sparse && !fa_compact.active && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256); vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc, - mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff, + mask != nullptr, use_mask_opt, logit_softcap != 0, use_sparse, k_type_eff, v_type_eff, fa_compact.dynamic_kv); vk_pipeline pipeline = nullptr; @@ -12754,7 +13143,19 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const uint32_t Tr = CEIL_DIV(N, Br); // Try to use split_k when KV is large enough to be worth the overhead. - if (gqa_ratio > 1 && workgroups_x <= Br) { + // Sparse: split_kv carries n_kv_max, split_k partitions its blocks for occupancy. + if (use_sparse) { + split_kv = (uint32_t)n_kv_max; + const uint32_t total_blocks = CEIL_DIV((uint32_t)n_kv_max, Bc); + const uint32_t base_wgs = (gqa_ratio > 1 ? workgroups_x : Tr) * workgroups_y * workgroups_z; + if (base_wgs < shader_core_count * 2) { + split_k = shader_core_count * 2 / base_wgs; + } + split_k = std::max(1u, std::min(split_k, total_blocks)); + // Match the shader's per-split block count so no split is empty. + const uint32_t per_blocks = CEIL_DIV(total_blocks, split_k); + split_k = CEIL_DIV(total_blocks, per_blocks); + } else if (gqa_ratio > 1 && workgroups_x <= Br) { split_k = shader_core_count * 2 / (workgroups_x * workgroups_y * workgroups_z); } else if (gqa_ratio <= 1) { uint32_t total_wgs_no_split = Tr * workgroups_y * workgroups_z; @@ -12763,7 +13164,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } } - if (split_k > 1) { + if (!use_sparse && split_k > 1) { // Try to evenly split KV into split_k chunks, but it needs to be a multiple // of "align", so recompute split_k based on that. split_kv = ROUNDUP_POW2(std::max(1u, KV / split_k), alignment); @@ -12810,6 +13211,21 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } } + // Sparse index scratch reuses prealloc_y (mutually exclusive with mask opt). + const uint64_t sparse_idx_size = use_sparse + ? sizeof(int32_t) * (uint64_t)n_kv_max * nem1 * nem2 * nem3 + : 0; + if (use_sparse) { + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_fa_sparse_compact, 1); + if (ctx->prealloc_size_y < sparse_idx_size) { + ctx->prealloc_size_y = sparse_idx_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_y_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + } + const uint32_t n_head_kv = neq2; const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); @@ -12827,6 +13243,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf; + vk_subbuffer sparse_buf = use_sparse ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf; // Dequant+transpose quant K/V directly into a per-head-contiguous [HS, KV, n_head_kv, ns] f16 // scratch (dequant_*_transpose shader) so the f16 FA reads KV coalesced. One pass, no temp. @@ -12894,6 +13311,24 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const uint32_t eff_nbv2 = fa_compact.active ? fa_compact.kv_c * fa_compact.v_row_bytes : nbv2_eff; const uint32_t eff_nbv3 = fa_compact.active ? fa_compact.n_head_kv * fa_compact.kv_c * fa_compact.v_row_bytes : nbv3_eff; + if (use_sparse) + { + const vk_op_flash_attn_sparse_compact_push_constants sc_pc = { + KV, + nem1, + nem2, + (uint32_t)(mask->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t)(mask->nb[2] / sizeof(ggml_fp16_t)), + (uint32_t)(mask->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t)n_kv_max, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_fa_sparse_compact, + { mask_buf, sparse_buf }, sc_pc, + { nem1, nem2, nem3 }); + ggml_vk_sync_buffers(ctx, subctx); + } + const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, (uint32_t)neq2, (uint32_t)neq3, @@ -12926,7 +13361,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx vk_subbuffer split_k_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf, fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf, use_sparse ? sparse_buf : (fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf)}, pc, { dispatch_x, workgroups_y, workgroups_z }); ggml_vk_sync_buffers(ctx, subctx); @@ -12941,9 +13376,16 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx workgroups_x *= pipeline->wg_denoms[0]; } ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf, fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf, use_sparse ? sparse_buf : (fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf)}, pc, { workgroups_x, workgroups_y, workgroups_z }); } + + if (use_dequant_kv) { + ctx->prealloc_x_need_sync = true; + } + if (use_mask_opt || use_sparse) { + ctx->prealloc_y_need_sync = true; + } } static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) { @@ -15900,13 +16342,19 @@ static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, cons ctx->prealloc_x_need_sync = true; } +static bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b, bool elementwise); +static bool ggml_backend_buffer_is_vk(ggml_backend_buffer_t buffer); + static void ggml_vk_topk_qsa(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_cgraph * cgraph, int node_idx) { - const ggml_tensor * get_rows = cgraph->nodes[node_idx + 0]; + const ggml_tensor * pre_cont = cgraph->nodes[node_idx + 0]; + const ggml_tensor * get_rows = cgraph->nodes[node_idx + 1]; const ggml_tensor * add = cgraph->nodes[node_idx + ctx->num_additional_fused_ops - 1]; ggml_tensor * top_k = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; - const ggml_tensor * scores = get_rows->src[0]; // [n_tps, n_blocks, n_stream] - const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream] + // the transpose in front of the gather is folded away, so the kernel reads the block score + // in its native layout: the shape comes from the transpose dst, the storage from its src + const ggml_tensor * scores = pre_cont->src[0]->src[0]; // [n_blocks, n_tps, n_stream] + const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream] // raw f16 mask: follow the reshape/cpy chain back to the materialized input const ggml_tensor * mask = add->src[1]; @@ -15914,9 +16362,9 @@ static void ggml_vk_topk_qsa(ggml_backend_vk_context * ctx, vk_context& subctx, mask = mask->src[0]; } - const uint32_t n_tps = scores->ne[0]; - const uint32_t n_blocks = scores->ne[1]; - const uint32_t n_stream = scores->ne[2]; + const uint32_t n_tps = pre_cont->ne[0]; + const uint32_t n_blocks = pre_cont->ne[1]; + const uint32_t n_stream = pre_cont->ne[2]; const uint32_t n_kv = cell_blk->ne[0]; const uint32_t width = top_k->ne[0]; const uint32_t nrows = n_tps * n_stream; @@ -15924,10 +16372,42 @@ static void ggml_vk_topk_qsa(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline pipeline = ctx->device->pipeline_topk_radix_qsa; GGML_ASSERT(pipeline != nullptr); - // scratch holds the gathered+masked input, materialized once and reused across passes - const size_t scratch_size = size_t{ n_kv } * nrows * sizeof(float); - if (ctx->prealloc_size_x < scratch_size) { - ctx->prealloc_size_x = scratch_size; + // The kernel reads the block score and writes cell indices, and ggml-alloc does place the + // output in the memory of the block score, so with independently scheduled workgroups the + // kernel would overwrite cells it has not read yet. These three tensors are exactly what the + // kernel reads, so testing them against the output is the whole hazard: on a hit the indices + // go to private storage and a copy fills the real output after a barrier, on a miss the + // kernel writes the output directly and pays nothing for a case that cannot happen. + // The fusion guard is skipped for this fusion because this test asks the right question - + // what the kernel reads - while the guard asks about the pattern's elided intermediates, + // which this kernel never touches. + // Only a Vulkan buffer can be compared: ggml_vk_tensors_overlap reads the buffer context as + // its own, and the context of another buffer type is just its data pointer. Two distinct + // buffers never share storage, so a tensor outside a Vulkan buffer can only alias the output + // by sitting in the very same buffer - which is still an overlap, and goes private. + const auto overlaps = [](const ggml_tensor * a, const ggml_tensor * b) { + if (a->buffer == nullptr || b->buffer == nullptr) { + return true; + } + if (!ggml_backend_buffer_is_vk(a->buffer) || !ggml_backend_buffer_is_vk(b->buffer)) { + return a->buffer == b->buffer; + } + return ggml_vk_tensors_overlap(a, b, false); + }; + bool private_out = overlaps(scores, top_k) || overlaps(cell_blk, top_k) || overlaps(mask, top_k); + // Whether the allocator overlaps the two is a property of the graph, so no test case reaches + // the private route on its own and the guard exemption below would be untested. This admits + // it without the overlap, which can only cost a copy. + static const char * priv_env = getenv("GGML_VK_QSA_PRIV_FORCE"); + static const bool priv_force = priv_env && priv_env[0] != '\0' && priv_env[0] != '0'; + private_out = private_out || priv_force; + + // a descriptor offset must be a multiple of minStorageBufferOffsetAlignment + const size_t scratch_size = GGML_PAD(size_t{ n_kv } * nrows * sizeof(float), + ctx->device->properties.limits.minStorageBufferOffsetAlignment); + const size_t out_size = private_out ? size_t{ width } * nrows * sizeof(int32_t) : 0; + if (ctx->prealloc_size_x < scratch_size + out_size) { + ctx->prealloc_size_x = scratch_size + out_size; ggml_vk_preallocate_buffers(ctx, subctx); } if (ctx->prealloc_x_need_sync) { @@ -15940,12 +16420,35 @@ static void ggml_vk_topk_qsa(ggml_backend_vk_context * ctx, vk_context& subctx, std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]), 1, }; - vk_subbuffer scratch_buf { ctx->prealloc_x, 0, ctx->prealloc_x->size }; + vk_subbuffer scratch_buf { ctx->prealloc_x, 0, scratch_size }; + vk_subbuffer out_buf = private_out ? vk_subbuffer{ ctx->prealloc_x, scratch_size, out_size } + : ggml_vk_tensor_subbuffer(ctx, top_k); ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - { ggml_vk_tensor_subbuffer(ctx, scores), ggml_vk_tensor_subbuffer(ctx, top_k), + { ggml_vk_tensor_subbuffer(ctx, scores), out_buf, ggml_vk_tensor_subbuffer(ctx, cell_blk), ggml_vk_tensor_subbuffer(ctx, mask), scratch_buf }, pc, elements); + + if (private_out) { + // shader write -> transfer read + ggml_vk_sync_buffers(ctx, subctx); + + // ggml_vk_tensor_subbuffer rounds the offset down and grows the range for shader + // addressing, which a straight buffer copy must not do: use the exact offset + vk_buffer dst_buf = nullptr; + size_t dst_off = 0; + if (ctx->device->uma) { + ggml_vk_host_get(ctx->device, top_k->data, dst_buf, dst_off); + } + if (!dst_buf) { + auto dst_buf_ctx = (ggml_backend_vk_buffer_context *) top_k->buffer->context; + dst_buf = dst_buf_ctx->dev_buffer; + dst_off = vk_tensor_offset(top_k) + top_k->view_offs; + } + GGML_ASSERT(dst_buf != nullptr); + ggml_vk_buffer_copy_async(subctx, dst_buf, dst_off, ctx->prealloc_x, scratch_size, out_size); + } + ctx->prealloc_x_need_sync = true; } @@ -17513,11 +18016,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; case GGML_OP_GET_ROWS: - if (ctx->fused_topk_qsa) { - ggml_vk_topk_qsa(ctx, compute_ctx, cgraph, node_idx); - } else { - ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node); - } + ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node); break; case GGML_OP_GET_ROWS_BACK: @@ -17621,7 +18120,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_OP_CPY: case GGML_OP_CONT: case GGML_OP_DUP: - ggml_vk_cpy(ctx, compute_ctx, src0, node); + if (ctx->fused_topk_qsa) { + ggml_vk_topk_qsa(ctx, compute_ctx, cgraph, node_idx); + } else { + ggml_vk_cpy(ctx, compute_ctx, src0, node); + } break; case GGML_OP_SET_ROWS: @@ -19036,14 +19539,29 @@ static bool ggml_vk_can_fuse_topk_qsa(ggml_backend_vk_context * ctx, const struc } } - const ggml_tensor * get_rows = cgraph->nodes[node_idx + 0]; + const ggml_tensor * pre_cont = cgraph->nodes[node_idx + 0]; + const ggml_tensor * get_rows = cgraph->nodes[node_idx + 1]; const ggml_tensor * add = cgraph->nodes[node_idx + n_ops - 2]; const ggml_tensor * top_k = cgraph->nodes[node_idx + n_ops - 1]; - const ggml_tensor * scores = get_rows->src[0]; // [n_tps, n_blocks, n_stream] + // the kernel reads the score through the transpose it folds away: same values, addressed + // by block instead of by token. Only that exact (1,0,2,3) transpose is transparent. + const ggml_tensor * permute = pre_cont->src[0]; + const ggml_tensor * scores = permute->src[0]; // [n_blocks, n_tps, n_stream] const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream] const ggml_tensor * expanded = add->src[0]; // [n_kv, n_tps, n_stream] + if (permute == nullptr || permute->op != GGML_OP_PERMUTE || scores == nullptr || + ggml_get_op_params_i32(permute, 0) != 1 || ggml_get_op_params_i32(permute, 1) != 0 || + ggml_get_op_params_i32(permute, 2) != 2 || ggml_get_op_params_i32(permute, 3) != 3) { + return false; + } + if (scores->ne[3] != 1 || !ggml_is_contiguous(scores) || + scores->ne[0] != pre_cont->ne[1] || scores->ne[1] != pre_cont->ne[0] || + scores->ne[2] != pre_cont->ne[2]) { + return false; + } + // raw mask: follow the reshape/cpy chain back to the materialized f16 input const ggml_tensor * mask = add->src[1]; while (mask && (mask->op == GGML_OP_RESHAPE || mask->op == GGML_OP_CPY)) { @@ -19061,14 +19579,15 @@ static bool ggml_vk_can_fuse_topk_qsa(ggml_backend_vk_context * ctx, const struc return false; } - const int64_t n_tps = scores->ne[0]; - const int64_t n_blocks = scores->ne[1]; - const int64_t n_stream = scores->ne[2]; + const int64_t n_tps = pre_cont->ne[0]; + const int64_t n_blocks = pre_cont->ne[1]; + const int64_t n_stream = pre_cont->ne[2]; const int64_t n_kv = cell_blk->ne[0]; const int64_t width = top_k->ne[0]; // pin the indexer layout the shader's addressing assumes - if (scores->ne[3] != 1 || cell_blk->ne[1] != n_stream || ggml_nrows(cell_blk) != n_stream || + if (!ggml_is_contiguous(pre_cont) || pre_cont->ne[3] != 1 || + cell_blk->ne[1] != n_stream || ggml_nrows(cell_blk) != n_stream || ggml_nelements(mask) != n_kv * n_tps * n_stream || expanded->ne[0] != n_kv || expanded->ne[1] != n_tps || expanded->ne[2] != n_stream || top_k->ne[1] != n_tps || top_k->ne[2] != n_stream || top_k->ne[3] != 1 || @@ -19649,10 +20168,14 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg // topk_moe often overwrites the source, but for a given row all the src values are // loaded before anything is stored. If there's only one row, this is safe, so treat // this as a special case. - bool is_topk_moe_single_row = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT && - ggml_nrows(cgraph->nodes[i]->src[0]) == 1; - - if (!is_topk_moe_single_row) { + // The fused QSA top-k routes its output into private storage whenever a tensor it + // reads overlaps it, so it never writes memory it still has to read and the reason + // this guard exists does not apply to it. + const bool overlap_safe = (ctx->fused_topk_moe_mode != TOPK_MOE_COUNT && + ggml_nrows(cgraph->nodes[i]->src[0]) == 1) || + ctx->fused_topk_qsa; + + if (!overlap_safe) { for (int j = 0; j < 2; ++j) { ggml_tensor *dst = output_nodes[j]; if (!dst) { @@ -19689,6 +20212,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_topk_qsa = false; + // the nodes run one by one now, so the perf logger must not report them under the + // fused name: a declined fusion used to look like a fused one that got slow + fusion_string = nullptr; } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 7c7d614d88d3..72de451edac0 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -218,12 +218,14 @@ void main() { uint32_t c = (idx + tid) % Bc; uint32_t r = (idx + tid) / Bc; if (idx + tid < Bc * Br) { - if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) { - FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + (j * Bc + c)]); + uint32_t kcol; + bool kv_active = fa_kv_index(j * Bc + c, kcol); + if (kv_active && (!nem1_bounds_check || i * Br + r < p.nem1)) { + FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + kcol]); masksh[c * masksh_stride + r] = m; max_mask = max(max_mask, float(m)); } else { - masksh[c * masksh_stride + r] = FLOAT_TYPE(0); + masksh[c * masksh_stride + r] = USE_SPARSE ? FLOAT_TYPE(NEG_FLT_MAX_OVER_2) : FLOAT_TYPE(0); } } } @@ -258,14 +260,15 @@ void main() { uint32_t c = (idx + tid) / (HSK / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSK / 4 || c < Bc) { FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); - if (!KV_bounds_check || j * Bc + c < KV) { + uint32_t kcol; + if (fa_kv_index(j * Bc + c, kcol)) { if (USE_DECODE_K) { - uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d; + uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * d; uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d]); } } @@ -305,7 +308,9 @@ void main() { } [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { + uint32_t kcol; + bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol); + if (!kv_active) { continue; } @@ -313,12 +318,12 @@ void main() { if (SHMEM_STAGING != 0) { K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; } else if (USE_DECODE_K) { - uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); + uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Sf[r][c] = dot_product(Q_cache[r], K_Tf, Sf[r][c]); @@ -327,7 +332,9 @@ void main() { } } else { [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { + uint32_t kcol; + bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol); + if (!kv_active) { continue; } @@ -336,12 +343,12 @@ void main() { if (SHMEM_STAGING != 0) { K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; } else if (USE_DECODE_K) { - uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); + uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Sf[r][c] = dot_product(Qf[tile_row(r) * qf_stride + d * D_split + d_tid], K_Tf, Sf[r][c]); @@ -493,14 +500,15 @@ void main() { uint32_t c = (idx + tid) / (HSV / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSV / 4 || c < Bc) { FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0); - if (!KV_bounds_check || j * Bc + c < KV) { + uint32_t vcol; + if (fa_kv_index(j * Bc + c, vcol)) { if (USE_DECODE_V) { - uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d; + uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * d; uint ib = coord / BLOCK_SIZE_V; uint iqs = (coord % BLOCK_SIZE_V); V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); } else { - V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]); + V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d]); } } @@ -511,7 +519,9 @@ void main() { } [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { + uint32_t vcol; + bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, vcol); + if (!kv_active) { continue; } @@ -526,12 +536,12 @@ void main() { if (SHMEM_STAGING != 0) { Vf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; } else if (USE_DECODE_V) { - uint coord = (j * Bc + c * cols_per_iter + col_tid) * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid); + uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid); uint ib = coord / BLOCK_SIZE_V; uint iqs = (coord % BLOCK_SIZE_V); Vf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); } else { - Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * v_stride / 4 + d * D_split + d_tid]); + Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Of[r][d] += FLOAT_TYPEV4(Pf[r] * Vf); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 0ff59672debf..1f608a11792a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -28,7 +28,9 @@ const bool OLD_AMD_WINDOWS = (Flags & 8) != 0; // GPU, where the row count is only known after a dedup pass and so cannot be pushed. The // workgroup counts derive from neq1/neq2/neq3 and never from KV, so no indirect dispatch is // needed: only this loop bound changes. Folds away for every other pipeline. -const bool DYNAMIC_KV = (Flags & 16) != 0; +const bool DYNAMIC_KV = (Flags & 32) != 0; +// Sparse: gather binding-7 indices instead of scanning [0,KV); p.split_kv = n_kv_max. +const bool USE_SPARSE = (Flags & 16) != 0; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; @@ -87,6 +89,9 @@ layout (binding = 5) writeonly buffer OV4 {D_TYPEV4 data_ov4[];}; layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; layout (binding = 7) readonly buffer KVB {uint32_t data_kv_dyn[];}; +// The sparse index list shares this binding with the dynamic-KV row count (the two features +// are mutually exclusive), so the same 32-bit words are reinterpreted as signed indices. +#define data_sparse(i) (int(data_kv_dyn[(i)])) #define MASK_OPT_ALL_NEG_INF 1 #define MASK_OPT_ALL_ZERO 2 @@ -199,7 +204,7 @@ ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC uint32_t i, N, KV, split_k_index, Tr, start_j, end_j, gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3, - q_stride, k_stride, v_stride, m_stride, m_row_len, gqa_ratio, split_k_num, output_k_num; + q_stride, k_stride, v_stride, m_stride, m_row_len, gqa_ratio, split_k_num, output_k_num, sparse_base; bool partial_output; void init_indices() @@ -281,6 +286,33 @@ void init_indices() // under the sparse split path, where this dispatch only covers the raw prefix (KV) but // the mask rows span the whole K range. m_row_len = mask_stride_in_split_kv ? p.split_kv : KV; + + // Sparse: the tile shares one mask row (gqa heads, or Br==1). split_k + // partitions the n_kv_max blocks. + if (USE_SPARSE) { + 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; + + uint32_t total_blocks = CEIL_DIV(p.split_kv, Bc); + uint32_t per_blocks = CEIL_DIV(total_blocks, p.k_num); + start_j = min(split_k_index * per_blocks, total_blocks); + end_j = min((split_k_index + 1) * per_blocks, total_blocks); + } +} + +// Resolve a linear KV slot to a real column; false for inactive (sparse padding/-1, or dense OOB). +bool fa_kv_index(uint lin, out uint kv_col) { + if (USE_SPARSE) { + if (lin >= p.split_kv) { + kv_col = 0; + return false; + } + int idx = data_sparse(sparse_base + lin); + kv_col = idx >= 0 ? uint(idx) : 0; + return idx >= 0; + } + kv_col = lin; + return !KV_bounds_check || lin < KV; } // Bias applied to softmax to stay in fp16 range. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 195901a7d90d..f235ebd6632e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -179,9 +179,16 @@ void main() { uint32_t c = (idx + tid) / (Br / 4); uint32_t r = (idx + tid) % (Br / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { - if ((!KV_bounds_check || j * Bc + c < KV)) { + uint32_t kcol; + bool kv_active = fa_kv_index(j * Bc + c, kcol); + if (kv_active) { f16vec4 m; - if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { + if (USE_SPARSE) { + // sparse is gqa-gated (m_stride == 0): all four rows share the value + FLOAT_TYPE mv = FLOAT_TYPE(data_m[m_offset + kcol]); + m = f16vec4(mv); + max_mask = max(max_mask, float(mv)); + } else if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], @@ -209,6 +216,8 @@ void main() { m = f16vec4(0.0); } mask_cache[idx / WorkGroupSize] = m; + } else if (USE_SPARSE) { + mask_cache[idx / WorkGroupSize] = f16vec4(NEG_FLT_MAX_OVER_2); } } } @@ -234,17 +243,19 @@ void main() { uint32_t c = (idx + tid) / (HSK_pad / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSK_pad / 4 || c < Bc) { FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); - if ((!KV_bounds_check || j * Bc + c < KV) && (HSK == HSK_pad || d < HSK / 4)) { + uint32_t kcol; + bool kv_active = fa_kv_index(j * Bc + c, kcol); + if (kv_active && (HSK == HSK_pad || d < HSK / 4)) { #if !defined(BFLOAT16) if (USE_DECODE_K) { - uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d; + uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * d; uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else #endif { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d]); } } @@ -269,7 +280,7 @@ void main() { if (SHMEM_STAGING == 0) { // For quants we always need to dequant into kvsh; for f16/bf16 we can load // directly from global memory when alignment / bounds allow it. - const bool stage_k = USE_DECODE_K || KV_bounds_check || d * 16 + 16 > HSK; + const bool stage_k = USE_DECODE_K || KV_bounds_check || USE_SPARSE || d * 16 + 16 > HSK; if (stage_k) { barrier(); [[unroll]] for (uint32_t idx = 0; idx < Bc * MatBr / 4; idx += gl_WorkGroupSize.x) { @@ -277,17 +288,19 @@ void main() { uint32_t row = (idx + tid) / (MatBr / 4); if (idx + tid < Bc * MatBr / 4) { FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); - if ((!KV_bounds_check || j * Bc + row < KV) && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) { + uint32_t kcol; + bool kv_active = fa_kv_index(j * Bc + row, kcol); + if (kv_active && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) { #if !defined(BFLOAT16) if (USE_DECODE_K) { - uint coord = (j * Bc + row) * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4; + uint coord = kcol * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4; uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else #endif { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + row) * k_stride / 4 + d * 16 / 4 + col_vec]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * 16 / 4 + col_vec]); } } @@ -408,17 +421,19 @@ void main() { uint32_t c = (idx + tid) / (HSV_pad / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSV_pad / 4 || c < Bc) { FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0); - if ((!KV_bounds_check || j * Bc + c < KV) && (HSV == HSV_pad || d < HSV / 4)) { + uint32_t vcol; + bool kv_active = fa_kv_index(j * Bc + c, vcol); + if (kv_active && (HSV == HSV_pad || d < HSV / 4)) { #if !defined(BFLOAT16) if (USE_DECODE_V) { - uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d; + uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * d; uint ib = coord / BLOCK_SIZE_V; uint iqs = (coord % BLOCK_SIZE_V); V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); } else #endif { - V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]); + V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d]); } } @@ -455,21 +470,23 @@ void main() { if (SHMEM_STAGING == 0) { // For quants we always preload via kvsh. For f16/bf16 we only preload when // alignment / bounds force it (otherwise we coopMatLoad direct from data_vv4). - const bool stage_v = USE_DECODE_V || KV_bounds_check; + const bool stage_v = USE_DECODE_V || KV_bounds_check || USE_SPARSE; if (stage_v) { [[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) { const uint idx = i * gl_WorkGroupSize.x + tid; const uint row = idx / v_cols; const uint col = idx % v_cols; - const uint v_row = j * Bc + row; + uint32_t vcol; + bool kv_active = fa_kv_index(j * Bc + row, vcol); + const uint v_row = USE_SPARSE ? vcol : (j * Bc + row); const uint v_col = hsv_tile * MatBc * row_split + col * 4; const uint coord = v_row * v_stride * BLOCK_SIZE_V + v_col; const uint ib = coord / BLOCK_SIZE_V; const uint iqs = coord % BLOCK_SIZE_V; - if (!KV_bounds_check || (v_row < KV && v_col < HSV)) { + if (USE_SPARSE ? (kv_active && v_col < HSV) : (!KV_bounds_check || (v_row < KV && v_col < HSV))) { #if !defined(BFLOAT16) if (USE_DECODE_V) { kvsh[row * vsh_stride + col] = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); @@ -491,7 +508,7 @@ void main() { if (hsv_offset < HSV_pad) { [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { if (SHMEM_STAGING == 0) { - if (!USE_DECODE_V && !KV_bounds_check) { + if (!USE_DECODE_V && !KV_bounds_check && !USE_SPARSE) { // F16/BF16 values can be loaded directly from global memory const uint v_tile_row = j * Bc + bc_chunk * MatBc; const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index 54be1e6daa99..990b4647c755 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -105,6 +105,39 @@ layout (binding = 1) readonly buffer K {uint8_t data_k[];}; layout (binding = 2) readonly buffer V {uint8_t data_v[];}; layout (binding = 3) readonly buffer M {uint8_t data_m[];}; +// f16 aliases for the sparse gather callbacks. +layout (binding = 1) readonly buffer KF16 {float16_t data_kf16[];}; +layout (binding = 2) readonly buffer VF16 {float16_t data_vf16[];}; +layout (binding = 3) readonly buffer MF16 {float16_t data_mf16[];}; + +// K/V/mask f16-element offsets for the current head/batch, set in main(). +uint32_t g_k_off_elem, g_v_off_elem, g_m_off_elem; + +#if !defined(BFLOAT16) +// Gather decode: ignore the pre-resolved block and read the selected KV row via the +// index list. blockCoords[0] = KV slot in [0,n_kv_max), [1] = head dim. +float16_t faGatherK(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) { + if (blockCoords[0] >= p.split_kv) { return float16_t(0); } + const int r = data_sparse(sparse_base + blockCoords[0]); + return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1]]; +} + +float16_t faGatherV(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) { + if (blockCoords[0] >= p.split_kv) { return float16_t(0); } + const int r = data_sparse(sparse_base + blockCoords[0]); + return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1]]; +} +#endif + +// Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk. +ACC_TYPE faAddSparseMask(const uint32_t row, const uint32_t col, const ACC_TYPE elem, const uint32_t jblk) { + const float NEG = uintBitsToFloat(0xFEFFFFFF); + const uint32_t kvslot = jblk * Bc + col; + if (kvslot >= p.split_kv) { return ACC_TYPE(NEG); } + const int r = data_sparse(sparse_base + kvslot); + return r < 0 ? ACC_TYPE(NEG) : elem + ACC_TYPE(data_mf16[g_m_off_elem + row * m_stride + uint(r)]); +} + ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) { return max(x, y); } @@ -189,8 +222,15 @@ void main() { tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v); tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, HSK); - tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, HSK); - tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, HSV); + if (USE_SPARSE) { + // Sparse iterates n_kv_max (in split_kv); the decode callbacks remap each slot. + // Kept behind USE_SPARSE so the dense specialization compiles the original code. + tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, p.split_kv, HSK); + tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, p.split_kv, HSV); + } else { + tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, HSK); + tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, HSV); + } // hint to the compiler that strides are aligned for the aligned variant of the shader if (Clamp != gl_CooperativeMatrixClampModeConstantNV) @@ -248,6 +288,12 @@ void main() { mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } + if (USE_SPARSE) { + g_k_off_elem = (ik2*p.nb12 + ik3*p.nb13) / 2; + g_v_off_elem = (iv2*p.nb22 + iv3*p.nb23) / 2; + g_m_off_elem = m_offset / 2; + } + uint32_t mask_opt = 0; uint32_t mask_opt_idx = ~0; @@ -255,7 +301,7 @@ void main() { for (uint32_t j = start_j; j < end_j; ++j) { coopmat mv = coopmat(0); - if (MASK_ENABLE) { + if (MASK_ENABLE && !USE_SPARSE) { if (USE_MASK_OPT && mask_opt_idx != j / 16) { mask_opt_idx = j / 16; @@ -313,7 +359,9 @@ void main() { coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose); #else const bool k_use_decode = (bs_k > 1u); - if (k_use_decode) { + if (USE_SPARSE) { + coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose, faGatherK); + } else if (k_use_decode) { coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK); } else { coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose); @@ -328,7 +376,9 @@ void main() { } } - if (MASK_ENABLE) { + if (MASK_ENABLE && USE_SPARSE) { + coopMatPerElementNV(S, S, faAddSparseMask, j); + } else if (MASK_ENABLE) { S += slopeMat*coopmat(mv); } @@ -383,7 +433,9 @@ void main() { coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad)); #else const bool v_use_decode = (bs_v > 1u); - if (v_use_decode) { + if (USE_SPARSE) { + coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad), faGatherV); + } else if (v_use_decode) { coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV); } else { coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad)); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp index 306372ba8af6..a088488b3882 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp @@ -37,13 +37,18 @@ layout(push_constant) uniform Parameters { uint nbm1; uint n_batch; uint row_words; // bytes per K row / 4 + uint src_head_stride; // source KV-head stride, WORDS (0 for the single-head MLA row) + uint batch_off; // first mask row of this group (grouped prefill); 0 otherwise } p; const uint LANES = 64; +// Workgroups are {compact rows, KV heads}. The MLA row has one head and no head stride, so +// the y dimension is 1 there and the addressing below reduces to the original single-head form. void main() { - const uint row = gl_WorkGroupID.x; - const uint tid = gl_LocalInvocationIndex; + const uint row = gl_WorkGroupID.x; + const uint head = gl_WorkGroupID.y; + const uint tid = gl_LocalInvocationIndex; const uint kv_c = data_c[0]; // padded compact rows, the FA's runtime KV const uint n_uni = data_c[1]; // unpadded union size @@ -62,11 +67,11 @@ void main() { // Zeroing an unused row writes zero BYTES, which decode to zero for every block-quantised // type here (a zero scale zeroes the block) as well as for f16. The row is -inf in the mask // either way; zeroing only keeps a garbage dot product from reaching the softmax as a NaN. - const uint dst_base = row * p.row_words; + const uint dst_base = head * p.kv_c_max * p.row_words + row * p.row_words; // Stepped 4*LANES for the same reason as flash_attn_gather.comp: row_words is a push // constant, so the copy cannot be [[unroll]]ed, and an f16 row is one iteration this way. if (src < p.n_kv) { - const uint src_base = src * p.nbk1; + const uint src_base = head * p.src_head_stride + src * p.nbk1; uint i = tid; for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { data_kc[dst_base + i] = data_k[src_base + i]; @@ -95,10 +100,11 @@ void main() { // mismatch would step the mask by the wrong amount for every token past the first. The // buffer is allocated for kv_c_max, so a smaller stride simply leaves a tail unused. const float NEG_INF = uintBitsToFloat(0xff800000); - if (tid < p.n_batch) { + // The compact mask is head-independent: one KV head writes it. + if (tid < p.n_batch && head == 0) { float mv = NEG_INF; if (src < p.n_kv) { - mv = float(data_m[tid * p.nbm1 + src]); + mv = float(data_m[(tid + p.batch_off) * p.nbm1 + src]); } data_mc[tid * kv_c + row] = float16_t(mv); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp new file mode 100644 index 000000000000..6269deb0e0e3 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp @@ -0,0 +1,90 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_ballot : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const uint BLOCK_SIZE = 128; + +layout (binding = 0) readonly buffer M {float16_t data_m[];}; +layout (binding = 1) writeonly buffer I {int32_t data_i[];}; + +layout (push_constant) uniform parameter { + uint KV; + uint nem1; + uint nem2; + uint nbm1; + uint nbm2; + uint nbm3; + uint n_kv_max; +} p; + +shared uint count; +shared uint chunk_base; +shared uint sub_tot[32]; // one slot per subgroup: 1024 threads / 32 lanes max = 32 + +// One workgroup per mask row: compact finite-mask KV positions into a per-row index list of +// length n_kv_max, -1 padded. The list must be in ascending position order (it is the FA's +// accumulation order and the run-to-run contract), so the compaction is a deterministic +// ballot scan: per BLOCK_SIZE chunk, each subgroup ballots its finite lanes, and the output +// slots are assigned by subgroup rank and lane rank within the subgroup. No atomics: the +// slot order cannot vary run to run. +void main() { + const uint i1 = gl_WorkGroupID.x; + const uint i2 = gl_WorkGroupID.y; + const uint i3 = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint wave = gl_SubgroupID; + const uint nwave = gl_NumSubgroups; + + if (tid == 0) { + count = 0; + } + barrier(); + + const uint m_base = i3 * p.nbm3 + i2 * p.nbm2 + i1 * p.nbm1; + const uint out_base = ((i3 * p.nem2 + i2) * p.nem1 + i1) * p.n_kv_max; + + for (uint base = 0; base < p.KV; base += BLOCK_SIZE) { + const uint k = base + tid; + bool fin = false; + if (k < p.KV) { + const float v = float(data_m[m_base + k]); + fin = !isinf(v) && !isnan(v); + } + const uvec4 ballot = subgroupBallot(fin); + const uvec4 masked = ballot & gl_SubgroupLtMask; + const uint local = bitCount(masked.x) + bitCount(masked.y) + bitCount(masked.z) + bitCount(masked.w); + if (lane == 0) { + sub_tot[wave] = bitCount(ballot.x) + bitCount(ballot.y) + bitCount(ballot.z) + bitCount(ballot.w); + } + barrier(); + if (tid == 0) { + uint tot = 0; + for (uint w = 0; w < nwave; ++w) { + tot += sub_tot[w]; + } + chunk_base = count; + count += tot; + } + barrier(); + uint wave_prefix = 0; + for (uint w = 0; w < wave; ++w) { + wave_prefix += sub_tot[w]; + } + const uint slot = chunk_base + wave_prefix + local; + if (fin && slot < p.n_kv_max) { + data_i[out_base + slot] = int32_t(k); + } + barrier(); + } + + const uint c = min(count, p.n_kv_max); + for (uint s = c + tid; s < p.n_kv_max; s += gl_WorkGroupSize.x) { + data_i[out_base + s] = int32_t(-1); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp index 513375709300..f5f39b5d9127 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp @@ -46,6 +46,7 @@ layout(push_constant) uniform Parameters { uint max_words; // capacity of the shared bitmap, host-checked uint pad_to; uint count_only; // 1: produce the count, write no index list + uint batch_off; // first top-k row of this group (grouped prefill); 0 otherwise } p; // 12288 words = 393216 compressed rows, ~48 KiB of shared memory. @@ -77,7 +78,7 @@ void main() { for (uint c = tid; c < n_cand; c += gl_WorkGroupSize.x) { const uint t = c / p.n_top_k; const uint j = c - t * p.n_top_k; - const int idx = data_top[t * p.nbt1 + j]; + const int idx = data_top[(t + p.batch_off) * p.nbt1 + j]; if (idx >= 0 && uint(idx) < R) { atomicOr(bitmap[uint(idx) >> 5], 1u << (uint(idx) & 31u)); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp index c8975391b30f..9ebdf7119ef1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp @@ -4,11 +4,13 @@ #extension GL_EXT_shader_16bit_storage : require #extension GL_KHR_shader_subgroup_basic : enable #extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_shuffle : enable #include "types.glsl" layout(constant_id = 0) const int BLOCK_SIZE = 1024; -layout(constant_id = 1) const int QSA = 0; // 1: fuse the qwen4 QSA indexer gather + f16 mask +layout(constant_id = 1) const int QSA = 0; // 1: fuse the qwen4 QSA indexer gather + f16 mask +layout(constant_id = 2) const int EMIT_W = 8; // emit elements per invocation per round; 1 = ballot fallback layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; @@ -16,7 +18,7 @@ layout (binding = 0) readonly buffer A {float data_a[];}; // input values layout (binding = 1) writeonly buffer D {int data_d[];}; // [k, ...] layout (binding = 2) readonly buffer CB {int cell_blk[];}; // QSA: cell->block map [n_kv, n_stream] layout (binding = 3) readonly buffer M {float16_t mask[];}; // QSA: raw f16 kq_mask [n_kv, n_tps, n_stream] -layout (binding = 4) buffer S {float scratch[];}; // QSA: [nrows, n_kv] gathered inputs +layout (binding = 4) buffer S {float scratch[];}; // QSA: [nrows, n_kv] materialized inputs layout (push_constant) uniform parameter { uint ncols; @@ -27,13 +29,19 @@ layout (push_constant) uniform parameter { uint n_stream; // QSA only } p; -#define RADIX_BITS 8 +// 11 + 11 + 10 bits: three passes over the row instead of four, which is one full row of +// traffic less per token. Any digit width is exact for radix select. +#define RADIX_BITS 11 #define RADIX_SIZE (1 << RADIX_BITS) +#define RADIX_PASSES 3 shared uint histo[RADIX_SIZE]; shared uint sh_bucket; shared uint sh_above; -shared uint sg_cnt[64]; // per-subgroup hit counts for the slot scan +shared uint sg_cnt_gt[64]; // per-subgroup strictly-above counts for the slot scan +shared uint sg_cnt_eq[64]; // per-subgroup tie counts for the slot scan + +#define NEG_INF (uintBitsToFloat(0xFF800000u)) // order-preserving float -> uint mapping uint f2ui(float x) { @@ -47,21 +55,27 @@ uint f2ui(float x) { } // QSA element i of row (t,s): score[cell_blk[i,s], t, s] + mask[i,t,s] +// the score is read in its native block-major layout, so consecutive cells of a block (which +// is what cell_blk yields in runs) walk consecutive addresses float gather(uint row, uint i) { const uint t = row % p.n_tps; const uint s = row / p.n_tps; - const uint block = uint(cell_blk[s * p.ncols + i]); - const float a = data_a[(s * p.n_blocks + block) * p.n_tps + t]; const float m = float(mask[(s * p.n_tps + t) * p.ncols + i]); - return a + m; + // a masked cell sums to exactly -inf whatever the block index is, so the score read and + // the block lookup can be skipped: the block may not cover this cell at all + if (m == NEG_INF) { + return NEG_INF; + } + const uint block = uint(cell_blk[s * p.ncols + i]); + return data_a[(s * p.n_tps + t) * p.n_blocks + block] + m; } float load(uint row, uint i, bool first) { if (QSA == 0) { return data_a[row * p.ncols + i]; } - // materialize the scattered gather on the first pass and reuse it after; each - // invocation only touches its own scratch entries, so no barrier is needed + // materialize the scattered gather on the first pass and reuse it after; each invocation + // only touches its own scratch entries, so no barrier is needed const uint off = row * p.ncols + i; if (first) { const float v = gather(row, i); @@ -71,28 +85,158 @@ float load(uint row, uint i, bool first) { return scratch[off]; } +// single emit scan over the row. Values strictly above the threshold fill the slots +// [0, n_above) and the ties at the threshold fill [n_above, k). Each class keeps ascending +// element order, so the result equals the two-pass (above, then ties) output exactly while +// the row is read once. +// +// Slots come from a deterministic exclusive scan, so the output does not depend on +// scheduling. The class split cannot be merged into one ">= threshold" pass: ties must land +// after all strictly larger values or the sparse attention summation order changes. +void emit_rows(const uint row, const uint threshold, const uint n_above) { + const uint tid = gl_LocalInvocationID.x; + const uint ncols = p.ncols; + const uint row_out = row * p.k; + + uint base_gt = 0; + uint base_eq = n_above; + + for (uint rb = 0; rb < ncols; rb += BLOCK_SIZE * EMIT_W) { + const uint e0 = rb + tid * uint(EMIT_W); + uint m_gt = 0u; + uint m_eq = 0u; + if (e0 < ncols) { + [[unroll]] for (uint b = 0; b < EMIT_W; ++b) { + const uint i = e0 + b; + if (i < ncols) { + const uint key = f2ui(load(row, i, false)); + m_gt |= (key > threshold) ? (1u << b) : 0u; + m_eq |= (key == threshold) ? (1u << b) : 0u; + } + } + } + + uint rank_gt = 0; + uint rank_eq = 0; + uint cnt_gt = 0; + uint cnt_eq = 0; + + if (EMIT_W == 1) { + // fallback without subgroup shuffle: one BLOCK_SIZE chunk per round, one ballot + // per class. This is the historical chunked scan, kept bit for bit. + const uvec4 ballot_gt = subgroupBallot(m_gt != 0u); + const uvec4 ballot_eq = subgroupBallot(m_eq != 0u); + rank_gt = subgroupBallotExclusiveBitCount(ballot_gt); + rank_eq = subgroupBallotExclusiveBitCount(ballot_eq); + cnt_gt = subgroupBallotBitCount(ballot_gt); + cnt_eq = subgroupBallotBitCount(ballot_eq); + } else { + // exclusive prefix over the subgroup lanes. The shuffle index is clamped to the + // own lane so that every shuffle reads an executing invocation even in a partial + // last subgroup, and the mask keeps the value only where a lower lane exists. + uvec2 v = uvec2(bitCount(m_gt), bitCount(m_eq)); + [[unroll]] for (uint d = 1; d < gl_SubgroupSize; d <<= 1) { + const uint id = gl_SubgroupInvocationID; + const uint src = (id >= d) ? (id - d) : id; + const uvec2 n = subgroupShuffle(v, src); + v += (id >= d) ? n : uvec2(0u); + } + const uvec2 own = uvec2(bitCount(m_gt), bitCount(m_eq)); + rank_gt = (v - own).x; + rank_eq = (v - own).y; + // the inclusive sum sits on the highest active lane + const uint act = subgroupBallotBitCount(subgroupBallot(true)); + const uvec2 sg = subgroupShuffle(v, act - 1u); + cnt_gt = sg.x; + cnt_eq = sg.y; + } + + if (subgroupElect()) { + sg_cnt_gt[gl_SubgroupID] = cnt_gt; + sg_cnt_eq[gl_SubgroupID] = cnt_eq; + } + barrier(); + + uint sg_gt = 0; + uint sg_eq = 0; + uint tot_gt = 0; + uint tot_eq = 0; + for (uint sg = 0; sg < gl_NumSubgroups; ++sg) { + const uint a = sg_cnt_gt[sg]; + const uint b = sg_cnt_eq[sg]; + if (sg < gl_SubgroupID) { + sg_gt += a; + sg_eq += b; + } + tot_gt += a; + tot_eq += b; + } + const uint slot_gt0 = base_gt + sg_gt + rank_gt; + const uint slot_eq0 = base_eq + sg_eq + rank_eq; + + if (EMIT_W == 1) { + if ((m_gt != 0u) && (slot_gt0 < p.k)) { + data_d[row_out + slot_gt0] = int(e0); + } + if ((m_eq != 0u) && (slot_eq0 < p.k)) { + data_d[row_out + slot_eq0] = int(e0); + } + } else { + if (m_gt != 0u) { + [[unroll]] for (uint b = 0; b < EMIT_W; ++b) { + if ((m_gt & (1u << b)) != 0u) { + const uint slot = slot_gt0 + bitCount(m_gt & ((1u << b) - 1u)); + if (slot < p.k) { + data_d[row_out + slot] = int(e0 + b); + } + } + } + } + if (m_eq != 0u) { + [[unroll]] for (uint b = 0; b < EMIT_W; ++b) { + if ((m_eq & (1u << b)) != 0u) { + const uint slot = slot_eq0 + bitCount(m_eq & ((1u << b) - 1u)); + if (slot < p.k) { + data_d[row_out + slot] = int(e0 + b); + } + } + } + } + } + + base_gt += tot_gt; + base_eq += tot_eq; + barrier(); + } +} + // one workgroup per row: radix-select the K-th largest, then compact it plus enough ties void topk(const uint row) { const uint tid = gl_LocalInvocationID.x; const uint ncols = p.ncols; - const uint row_out = row * p.k; uint prefix = 0; // fixed high bits of the threshold key uint desired = p.k; // count still needed from the candidate range - [[unroll]] for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) { - for (uint i = tid; i < RADIX_SIZE; i += BLOCK_SIZE) { + for (int pass = 0; pass < RADIX_PASSES; ++pass) { + // digits cover bits [shift, shift + dbits): 11 + 11 + 10 = 32, no overlap + const int shift = (pass == 2) ? 0 : (21 - 11 * pass); + const int dbits = (pass == 2) ? 10 : 11; + const uint nbuck = 1u << uint(dbits); + const uint bmask = nbuck - 1u; + const uint hi_mask = ((shift + dbits) >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + dbits)); + const uint prefix_hi = prefix & hi_mask; + const bool first = (pass == 0); + + for (uint i = tid; i < nbuck; i += BLOCK_SIZE) { histo[i] = 0; } barrier(); - const bool first = (shift == 32 - RADIX_BITS); - const uint hi_mask = (shift + RADIX_BITS >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + RADIX_BITS)); - const uint prefix_hi = prefix & hi_mask; for (uint i = tid; i < ncols; i += BLOCK_SIZE) { const uint key = f2ui(load(row, i, first)); if ((key & hi_mask) == prefix_hi) { - atomicAdd(histo[(key >> uint(shift)) & (RADIX_SIZE - 1)], 1u); + atomicAdd(histo[(key >> uint(shift)) & bmask], 1u); } } barrier(); @@ -101,7 +245,7 @@ void topk(const uint row) { if (tid == 0) { uint acc = 0; uint b = 0; - for (int bb = RADIX_SIZE - 1; bb >= 0; --bb) { + for (int bb = int(nbuck) - 1; bb >= 0; --bb) { const uint c = histo[bb]; if (acc + c >= desired) { b = uint(bb); break; } acc += c; @@ -116,46 +260,8 @@ void topk(const uint row) { barrier(); } - - // Emit everything above the threshold, then fill the rest from ties. Slots come from - // an exclusive scan over the candidate flags, one BLOCK_SIZE chunk at a time in ascending - // index order, so the output is identical on every run. The previous atomicAdd slot - // counter made the ORDER scheduling-dependent, and at the tie boundary the SET as well: - // the QSA width is top_k + ratio - 1, so the boundary block's ratio tied cells race for - // ratio - 1 slots and a different cell lost each run (non-repeatable output at depth). - // With the scan the lowest-indexed tied cells win. - const uint threshold = prefix; - uint base = 0; - [[dont_unroll]] for (uint pass = 0; pass < 2; ++pass) { - for (uint c = 0; c < ncols; c += BLOCK_SIZE) { - const uint i = c + tid; - bool hit = false; - if (i < ncols) { - const uint key = f2ui(load(row, i, false)); - hit = (pass == 0) ? (key > threshold) : (key == threshold); - } - const uvec4 ballot = subgroupBallot(hit); - const uint rank_sg = subgroupBallotExclusiveBitCount(ballot); - const uint cnt_sg = subgroupBallotBitCount(ballot); - if (subgroupElect()) { - sg_cnt[gl_SubgroupID] = cnt_sg; - } - barrier(); - uint sg_base = 0; - uint total = 0; - for (uint sg = 0; sg < gl_NumSubgroups; ++sg) { - const uint v = sg_cnt[sg]; - sg_base += (sg < gl_SubgroupID) ? v : 0; - total += v; - } - const uint slot = base + sg_base + rank_sg; - if (hit && slot < p.k) { - data_d[row_out + slot] = int(i); - } - base += total; - barrier(); - } - } + // wanted = K - (number of strictly larger values), so the larger ones take the low slots + emit_rows(row, prefix, p.k - desired); } void main() { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 29dc5265b3fd..002e0d3b5cf9 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -905,6 +905,7 @@ void process_shaders() { string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {}); string_to_spv("fa_mask_opt", "flash_attn_mask_opt.comp", {}); + string_to_spv("fa_sparse_compact", "flash_attn_sparse_compact.comp", {}); string_to_spv("quantize_q8_1", "quantize_q8_1.comp", {}); string_to_spv("quantize_q8_1_subgroup", "quantize_q8_1.comp", {{"USE_SUBGROUPS", "1"}}); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e90a2ccacf88..6607a2e52720 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5463,6 +5463,16 @@ enum ggml_prec ggml_flash_attn_ext_get_prec( return (enum ggml_prec) prec_i32; } +void ggml_flash_attn_ext_set_sparse( + struct ggml_tensor * a, + int32_t n_kv_max) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(n_kv_max >= 0); + + // slot 4 holds the fork's top_k n_kv_raw; the sparse hint lives in slot 5 + ggml_set_op_params_i32(a, 5, n_kv_max); +} + void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 60946a4c1a4c..548e0e68b696 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1079,7 +1079,10 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { const int64_t n_rs = mctx->get_recr()->get_n_rs(); - if (inp_rs->s_copy) { + // ggml-alloc assigns a buffer only to tensors some node reads. A hybrid graph whose + // recurrent side is unused (an empty recurrent set, e.g. an MTP draft context) leaves + // s_copy unallocated, and set_input must then skip it like any dead input. + if (inp_rs->s_copy && inp_rs->s_copy->buffer) { GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer)); int32_t * data = (int32_t *) inp_rs->s_copy->data; @@ -1123,7 +1126,10 @@ void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) { const int64_t n_rs = mctx->get_recr()->get_n_rs(); - if (inp_rs->s_copy) { + // ggml-alloc assigns a buffer only to tensors some node reads. A hybrid graph whose + // recurrent side is unused (an empty recurrent set, e.g. an MTP draft context) leaves + // s_copy unallocated, and set_input must then skip it like any dead input. + if (inp_rs->s_copy && inp_rs->s_copy->buffer) { GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer)); int32_t * data = (int32_t *) inp_rs->s_copy->data; @@ -1197,7 +1203,10 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) { const int64_t n_rs = mctx->get_recr()->get_n_rs(); - if (inp_rs->s_copy) { + // ggml-alloc assigns a buffer only to tensors some node reads. A hybrid graph whose + // recurrent side is unused (an empty recurrent set, e.g. an MTP draft context) leaves + // s_copy unallocated, and set_input must then skip it like any dead input. + if (inp_rs->s_copy && inp_rs->s_copy->buffer) { GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer)); int32_t * data = (int32_t *) inp_rs->s_copy->data; @@ -2526,6 +2535,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * kq_mask, ggml_tensor * sinks, ggml_tensor * v_mla, + int64_t n_kv_max, float kq_scale, int il, ggml_tensor * top_k, @@ -2568,6 +2578,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( if (top_k) { ggml_flash_attn_ext_add_top_k(cur, top_k, n_kv_raw); } + GGML_ASSERT(n_kv_max >= 0 && n_kv_max <= INT32_MAX); + ggml_flash_attn_ext_set_sparse(cur, static_cast(n_kv_max)); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); if (v_mla) { @@ -2717,7 +2729,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = k_cur; ggml_tensor * v = v_cur; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -2816,7 +2828,7 @@ ggml_tensor * llm_graph_context::build_attn( 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, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il); cb(cur, "kqv_out", il); if (inp->self_v_rot) { @@ -2907,7 +2919,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -2992,7 +3004,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -3084,7 +3096,7 @@ ggml_tensor * llm_graph_context::build_attn( v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); } - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il); cb(cur, "kqv_out", il); if (v_rot) { @@ -3157,7 +3169,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = k; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il); cb(cur, "kqv_out", il); if (k_rot) { @@ -3216,7 +3228,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = k_cur; ggml_tensor * v = v_cur; - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 4b9ca0dba350..0475dcf3f7f7 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1138,6 +1138,7 @@ struct llm_graph_context { ggml_tensor * kq_mask, ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + int64_t n_kv_max, float kq_scale, int il, ggml_tensor * top_k = nullptr, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index acb2e798b83a..8865e8b767cc 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2427,30 +2427,18 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (arch == LLM_ARCH_QWEN4EXP && hparams.n_layer_nextn > 0 && params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) { - // A hybrid memory with an empty recurrent layer set fails its buffer - // allocation, so the MTP context gets a PLAIN attention cache over - // the nextn layer(s), dense - the deepseek32 MTP pattern. - llama_kv_cache::layer_filter_cb filter_mtp = - [&](uint32_t il) { return il >= hparams.n_layer(); }; - - res = new llama_kv_cache( - *this, - hparams, - params.type_k, - params.type_v, - !cparams.flash_attn, - cparams.offload_kqv, - cparams.kv_unified, - cparams.n_ctx_seq, - cparams.n_seq_max, - 1, - hparams.n_swa, - hparams.swa_type, - nullptr, - filter_mtp, - nullptr, - nullptr); - break; + // The NextN/MTP block is a full-attention QSA layer, so the draft + // gets the same hybrid-idx memory as the trunk with the filters + // inverted: attention + indexer over the nextn layer(s) only, and + // no recurrent layers at all. An all-false recurrent filter leaves + // the recurrent cache without tensors and without buffers, which is + // exactly what the draft needs (no PLE, no GDN in the MTP block). + filter_attn = [&](uint32_t il) { return il >= hparams.n_layer(); }; + filter_recr = [&](uint32_t) { return false; }; + if (hparams.indexer_head_size > 0 && + hparams.dsv4_compress_ratios[hparams.n_layer()] > 0) { + filter_idx = [&](uint32_t il) { return il >= hparams.n_layer(); }; + } } } diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index d4f7bc69168b..8aa2e6098c26 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -785,7 +785,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il, top_k, raw_k->ne[2]); + // n_kv_max bounds the finite mask entries per row: the SWA raw prefix plus the selection. + // the fork's top_k shader still takes precedence at dispatch; the sparse hint only drives + // the mask-compaction path when that shader declines. + const int64_t n_kv_max = std::min(raw_mask->ne[0], hparams.n_swa) + top_k->ne[0]; + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, n_kv_max, kq_scale, il, top_k, raw_k->ne[2]); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -840,7 +844,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); cb(kq_mask, "hca_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -876,7 +880,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * k = mctx_cur->get_k(ctx0, il); - ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } diff --git a/src/models/models.h b/src/models/models.h index 380860b6e22a..63145cba27b0 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2360,7 +2360,8 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * cur, ggml_tensor * inp_pos, int * sections, - int il); + int il, + bool qsa_allow = true); // dense self-attention restricted to the cells that top_k names ggml_tensor * build_attn_qsa( diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index bf4e44f0634d..bde93a372cff 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -89,6 +89,31 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false); + // The NextN/MTP block is a full-attention QSA layer like every other full-attention + // layer and ships its own trained indexer tensors. A sidecar written through + // llama-model-saver pads compress_ratios to n_layer_all from the trunk's array, whose + // tail stays zero, so a zero there is a padding artifact, not "the draft runs dense". + // Restore the ratio from the trunk's QSA layers when the MTP block has an indexer. + for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) { + if (hparams.dsv4_compress_ratios[il] > 0) { + continue; + } + const std::string probe = "blk." + std::to_string(il) + ".indexer.q_proj.weight"; + if (ml.get_weight(probe.c_str()) == nullptr) { + continue; + } + for (uint32_t t = 0; t < hparams.n_layer(); ++t) { + if (hparams.dsv4_compress_ratios[t] > 0) { + hparams.dsv4_compress_ratios[il] = hparams.dsv4_compress_ratios[t]; + break; + } + } + if (hparams.dsv4_compress_ratios[il] == 0) { + LLAMA_LOG_WARN("%s: MTP layer %u has indexer tensors but the trunk has no QSA ratio, the draft stays dense\n", + __func__, il); + } + } + // PLE n-gram hash embeddings; if the key group is absent every field stays zero hparams.is_ple_impl.reset(); hparams.ple_n_heads = 0; @@ -903,9 +928,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( // the backend attend over the whole cache and merely discard what it read, which is O(n_kv) // per token; with top_k attached, a backend that can compact the active set (the Vulkan // gather-compact path) costs O(n_top_k) instead. n_kv_raw is 0: unlike DeepSeek V4 this - // cache has no dense prefix, every attended cell comes from the selection. Backends without - // that path ignore the extra argument and read the same mask they do today. - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il, top_k, 0); + // cache has no dense prefix, every attended cell comes from the selection. n_kv_max bounds + // the finite mask entries per row (the selection width) and drives the mask-compaction + // sparse path for prefill batches, which the fork's gather-compact path declines (N >= 64). + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, top_k->ne[0], kq_scale, il, top_k, 0); cb(cur, "kqv_out", il); // the rotation is its own inverse, so undo it on the value side of the output @@ -922,12 +948,24 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( ggml_tensor * cur, ggml_tensor * inp_pos, int * sections, - int il) { + int il, + bool qsa_allow) { const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); // indexer reads the same block input as q/k/v; no cache or no ratio means dense - const bool qsa = mctx_hyb != nullptr && mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; + const auto * mctx_idx = mctx_hyb ? mctx_hyb->get_idx() : nullptr; + const bool qsa = mctx_idx != nullptr && hparams.dsv4_compress_ratios[il] > 0 && qsa_allow; + + if (mctx_idx != nullptr && !qsa && model.layers[il].index_k_proj != nullptr) { + // the sparse path is declined for this ubatch, but the pooled-key cache recomputes + // any block above its watermark from the stored indexer keys, so the keys still have + // to be written or a later sparse ubatch pools over garbage. The index cache shares + // the attention cache's slot layout cell for cell, so the k indices are the same. + ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); + k_raw = ggml_reshape_3d(ctx0, k_raw, hparams.indexer_head_size, 1, n_tokens); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->get_k_idxs(), il)); + } ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; @@ -1562,10 +1600,17 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); - // the MTP context holds a plain attention cache over the nextn layer(s) only, the - // deepseek32 pattern: the draft runs dense (no indexer cache, no recurrent state) - auto * inp_attn = build_attn_inp_kv(); - const llama_memory_hybrid_idx_context * mctx_hyb = nullptr; + // the MTP block is a full-attention QSA layer like every trunk full-attention layer, so + // the draft context carries the same hybrid-idx memory as the trunk: attention + indexer + // over the nextn layer(s), and an empty recurrent set (no PLE, no GDN in the draft block) + auto * inp = build_inp_mem_hybrid(); + auto * inp_attn = inp->get_attn(); + // qwen4exp always builds llama_memory_hybrid_idx, so this downcast is safe + const auto * mctx_hyb = static_cast(inp->mctx); + if (mctx_hyb->get_idx()) { + GGML_ASSERT(mctx_hyb->get_idx()->get_n_kv() == inp->mctx->get_attn()->get_n_kv() && + "the indexer cache must track the attention cache cell for cell"); + } // hnorm is the same grouped RMSNorm as every HC norm: rms over one stream, flat gamma ggml_tensor * h_norm = ggml_rms_norm(ctx0, h_state, hparams.f_norm_rms_eps); @@ -1591,7 +1636,10 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_ &inject, il); ggml_build_forward_expand(gf, cur); - cur = build_layer_attn(inp_attn, mctx_hyb, cur, inp_pos, sections, il); + // The draft's sparse path only pays on prefill-sized batches: a decode or verify batch + // of 1-4 rows spends more on the indexer pipeline and its O(n_kv) host scan than the + // sparse FA saves, so those run dense while keeping the indexer keys written. + cur = build_layer_attn(inp_attn, mctx_hyb, cur, inp_pos, sections, il, n_tokens >= 16); res_hc = build_hc_combine(res_hc, cur, inject, il); cur = build_hc_mix(res_hc, diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 75c2fb6449ac..d866b9dade5a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -189,6 +189,33 @@ static void init_tensor_kq_mask(ggml_tensor * tensor, float min = -1.0f, float m ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); } +static void init_tensor_kq_mask_sparse(ggml_tensor * tensor, int64_t n_kv_max) { + GGML_ASSERT(tensor->type == GGML_TYPE_F16); + GGML_ASSERT(n_kv_max > 1 && n_kv_max <= tensor->ne[0]); + + const int64_t ne0 = tensor->ne[0]; + const int64_t nrows = ggml_nrows(tensor); + std::vector data_f32(ggml_nelements(tensor), -INFINITY); + std::vector data_f16(ggml_nelements(tensor)); + std::vector order(ne0); + for (int64_t i = 0; i < ne0; ++i) { + order[i] = i; + } + + std::mt19937 gen(0x5A17); + for (int64_t row = 0; row < nrows; ++row) { + std::shuffle(order.begin(), order.end(), gen); + const int64_t count = n_kv_max - row % std::min(n_kv_max, 17); + std::sort(order.begin(), order.begin() + count); + for (int64_t i = 0; i < count; ++i) { + data_f32[row*ne0 + order[i]] = -0.03125f * (1 + (i + row) % 7); + } + } + + ggml_fp32_to_fp16_row(data_f32.data(), data_f16.data(), data_f16.size()); + ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); +} + // generate a lower triangular matrix static void init_tensor_tril(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { GGML_ASSERT(tensor->type == GGML_TYPE_F32); @@ -433,6 +460,8 @@ static std::string var_to_str(ggml_scale_mode mode) { #define VARS_TO_STR14(a, b, c, d, e, f, g, h, i, j, k, l, m, n) VAR_TO_STR(a) + "," + VARS_TO_STR13(b, c, d, e, f, g, h, i, j, k, l, m, n) #define VARS_TO_STR15(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) VAR_TO_STR(a) + "," + VARS_TO_STR14(b, c, d, e, f, g, h, i, j, k, l, m, n, o) #define VARS_TO_STR16(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) VAR_TO_STR(a) + "," + VARS_TO_STR15(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) +#define VARS_TO_STR17(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q) VAR_TO_STR(a) + "," + VARS_TO_STR16(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q) +#define VARS_TO_STR18(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r) VAR_TO_STR(a) + "," + VARS_TO_STR17(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r) #ifdef GGML_USE_SYCL static bool inline _isinf(float f) { @@ -6208,6 +6237,7 @@ struct test_topk_qsa : public test_case { const int64_t n_tps; const int64_t n_stream; const int width; + const bool degenerate; // equal scores, few distinct mask values and masked cells ggml_tensor * out {}; std::string op_desc(ggml_tensor * t) override { @@ -6216,11 +6246,11 @@ struct test_topk_qsa : public test_case { } std::string vars() override { - return VARS_TO_STR5(n_blocks, n_kv, n_tps, n_stream, width); + return VARS_TO_STR6(n_blocks, n_kv, n_tps, n_stream, width, degenerate); } - test_topk_qsa(int64_t n_blocks = 512, int64_t n_kv = 2048, int64_t n_tps = 2, int64_t n_stream = 1, int width = 1500) - : n_blocks(n_blocks), n_kv(n_kv), n_tps(n_tps), n_stream(n_stream), width(width) {} + test_topk_qsa(int64_t n_blocks = 512, int64_t n_kv = 2048, int64_t n_tps = 2, int64_t n_stream = 1, int width = 1500, bool degenerate = false) + : n_blocks(n_blocks), n_kv(n_kv), n_tps(n_tps), n_stream(n_stream), width(width), degenerate(degenerate) {} double max_err() override { return 0.0; } bool run_whole_graph() override { return true; } @@ -6245,7 +6275,10 @@ struct test_topk_qsa : public test_case { std::vector fusion_test_nodes() override { return { out }; } - // distinct mask ramp + small scores keep every cell value unique, so no top-k ties + // distinct mask ramp + small scores keep every cell value unique, so no top-k ties. + // The degenerate variant instead floods every row with ties and masked cells: the + // threshold then has to be filled from equal values, and the -inf shortcut in the + // fused gather has to agree with the reference while the block index is ignored. void initialize_tensors(ggml_context * ctx) override { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { if (t->op != GGML_OP_NONE) { @@ -6259,12 +6292,18 @@ struct test_topk_qsa : public test_case { std::vector data(ggml_nelements(t)); for (int64_t r = 0; r < ggml_nrows(t); r++) { for (int64_t i = 0; i < n_kv; i++) { - data[r * n_kv + i] = ggml_fp32_to_fp16((float) i); + const bool masked = degenerate && (i % 16 == 0); + const float v = degenerate ? (float) (i % 4) : (float) i; + data[r * n_kv + i] = ggml_fp32_to_fp16(masked ? -INFINITY : v); } } ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(ggml_fp16_t)); } else { init_tensor_uniform(t, 0.0f, 0.5f); + if (degenerate) { + std::vector data(ggml_nelements(t), 1.0f); + ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); + } } } } @@ -6278,8 +6317,35 @@ struct test_topk_qsa : public test_case { ib[i] = (int32_t) b[i]; diff += std::fabs(a[i] - ia[i]) + std::fabs(b[i] - ib[i]); } + if (degenerate) { + // every row holds fewer finite values than the width, so the selection has to be + // filled from the masked -inf cells. Which of the tied cells each backend keeps is + // unspecified (the CPU sort is not stable), so compare the value multiset instead: + // it is unique and it changes if the masked cells lose their -inf value. + std::vector va(n), vb(n); + for (size_t i = 0; i < n; i++) { + va[i] = value_of(ia[i]); + vb[i] = value_of(ib[i]); + } + std::sort(va.begin(), va.end()); + std::sort(vb.begin(), vb.end()); + double miss = 0.0; + for (size_t i = 0; i < n; i++) { + miss += (va[i] == vb[i]) ? 0.0 : 1.0; + } + return diff + miss; + } return diff + jdst(ia.data(), ib.data(), n); } + + // cell value produced by initialize_tensors for the degenerate data: uniform score plus + // a masking value, or exactly -inf inside a masked cell + float value_of(int32_t i) const { + if (i % 16 == 0) { + return -INFINITY; + } + return 1.0f + (float) (i % 4); + } }; enum MoeGatingFunc { @@ -7143,11 +7209,12 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_K; const ggml_type type_V; std::array permute; - const bool kv_view; // K/V as sparse views (default); false = dense permuted like the model KV cache + const bool v_is_view_of_k; + const int64_t n_kv_max; std::string vars() override { - return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view); + return VARS_TO_STR17(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k, n_kv_max); } double max_nmse_err() override { @@ -7164,9 +7231,9 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, - bool kv_view = true) + bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7226,6 +7293,7 @@ struct test_flash_attn_ext : public test_case { ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap); ggml_flash_attn_ext_add_sinks(out, s); + ggml_flash_attn_ext_set_sparse(out, n_kv_max); ggml_flash_attn_ext_set_prec (out, prec); ggml_set_name(out, "out"); @@ -7238,7 +7306,11 @@ struct test_flash_attn_ext : public test_case { // make the sink values more noticeable in order to trigger a test failure when the implementation is wrong init_tensor_uniform(t, -10.0f, 10.0f); } else if (strcmp(t->name, "m") == 0) { - init_tensor_kq_mask(t); + if (n_kv_max > 0) { + init_tensor_kq_mask_sparse(t, n_kv_max); + } else { + init_tensor_kq_mask(t); + } } else { init_tensor_uniform(t); } @@ -7270,6 +7342,8 @@ struct test_flash_attn_ext_top_k : public test_case { const int64_t nh; const int64_t nh_kv; const bool shared_kv; + const int64_t mask_finite; // finite mask cells per row: > 0 overrides (n_kv_raw + n_top_k), + // exercising masks wider than the sparse compaction budget std::string vars() override { return VARS_TO_STR12(kv, nb, n_kv_raw, n_top_k, sinks, ns, ov, type_K, hs, nh, nh_kv, shared_kv); @@ -7288,9 +7362,10 @@ struct test_flash_attn_ext_top_k : public test_case { test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false, int64_t ns = 1, int64_t ov = 0, ggml_type type_K = GGML_TYPE_F16, - int64_t hs = 512, int64_t nh = 64, int64_t nh_kv = 1, bool shared_kv = true) + int64_t hs = 512, int64_t nh = 64, int64_t nh_kv = 1, bool shared_kv = true, + int64_t mask_finite = 0) : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks), ns(ns), ov(ov), type_K(type_K), - hs(hs), nh(nh), nh_kv(nh_kv), shared_kv(shared_kv && nh_kv == 1) {} + hs(hs), nh(nh), nh_kv(nh_kv), shared_kv(shared_kv && nh_kv == 1), mask_finite(mask_finite) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh, ns); @@ -7321,6 +7396,10 @@ struct test_flash_attn_ext_top_k : public test_case { ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hs), 0.0f, 0.0f); ggml_flash_attn_ext_add_sinks(out, s); ggml_flash_attn_ext_add_top_k(out, t, n_kv_raw); + // The fork keeps the sparse hint in op_params[5] (slot 4 is n_kv_raw). Without this + // the backend never takes the sparse path and the test only ever measures the dense + // fallback. n_kv_max is the mask's own finite-cell bound. + ggml_flash_attn_ext_set_sparse(out, (int32_t) (n_kv_raw + (mask_finite > 0 ? mask_finite : n_top_k))); ggml_flash_attn_ext_set_prec (out, GGML_PREC_F32); ggml_set_name(out, "out"); @@ -7356,7 +7435,11 @@ struct test_flash_attn_ext_top_k : public test_case { for (int64_t i = 0; i < kv; ++i) { mask[mrow + i] = i < n_kv_raw ? zero : minus_inf; } - for (int64_t j = 0; j < n_top_k; ++j) { + // top_k always carries n_top_k entries; a wider mask only widens the mask, + // so the selection loop must stay bounded by n_top_k (writing past it would + // run off the end of `top`). + const int64_t finite = mask_finite > 0 ? mask_finite : n_top_k; + for (int64_t j = 0; j < n_top_k && j < finite; ++j) { // offset the selection by the stream too, so a dropped stream stride // reads another sequence's keys and shows up as a mismatch const bool shared = (int64_t) j * 100 < n_top_k * ov; @@ -7370,6 +7453,14 @@ struct test_flash_attn_ext_top_k : public test_case { } top[trow + j] = idx; } + // when the mask is wider than the selection (mask_finite > n_top_k), the + // extra finite cells have no index in top_k: a backend that compacts from + // the mask but bounds the list by n_top_k, or gathers from top_k but + // attends the whole mask, loses the agreement this test checks + for (int64_t j = n_top_k; j < finite; ++j) { + const int32_t idx = (int32_t) ((j * range) / n_top_k + b + s * 7) % (int32_t) range; + mask[mrow + n_kv_raw + idx] = zero; + } } } @@ -9905,6 +9996,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_topk_qsa(512, 2048, 2, 1, 1500)); test_cases.emplace_back(new test_topk_qsa(256, 2048, 4, 2, 2000)); test_cases.emplace_back(new test_topk_qsa(64, 256, 2, 1, 200)); // small k: unfused fallback + // width 2000 over 1920 finite cells: the threshold falls into the masked -inf group + test_cases.emplace_back(new test_topk_qsa(512, 2048, 2, 1, 2000, true)); // ties + masked cells // exhaustive top_k tests //for (int i = 1; i < 9999; ++i) { @@ -10179,6 +10272,36 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q2_0)); test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16)); + // q8_0 KV cases: decode and prompt batches, KV pad, permuted KV, feature flags, and long context + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 2}, 1025, 1, true, true, 8, 30, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + + // MLA shape: the V cache is a sub-view of the K cache, with quantized KV + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + + // Sparse mask hint: supported decode/prefill layouts and dense fallbacks. + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 2}, 4096, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 768)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 768)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512 )); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2304)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + // Qwen QSA: 256/256, gqa 12, budget 2048. + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, 8192, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + + // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes + test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 512, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). for (int64_t kv : { 4096, 16384 }) { @@ -10334,10 +10457,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 17, 1, 1, GGML_TYPE_F16)); test_cases.emplace_back(new test_lightning_indexer(128, 64, 512, 512, 1, 1, GGML_TYPE_F16)); - // sparse top-k FA: (kv, nb, n_kv_raw, n_top_k, sinks). The Vulkan sparse path engages - // when kv >= 3*(n_kv_raw + n_top_k) AND nb >= 64 (prefill-only); the nb < 64 cases - // and the kv=512 case verify dense-fallback parity with the hint attached, the - // nb=64/128 cases exercise the sparse shader itself. + // sparse top-k FA: (kv, nb, n_kv_raw, n_top_k, sinks). Whether the hint is honoured + // depends on the shape, not on nb: the Vulkan paths gate on gqa_ratio > 1 (one selection + // per query, shared by its heads) for per-tile compaction, and on the measured union + // being smaller than the dense cache for the grouped union. The nb < 64 cases and the + // kv=512 case verify dense-fallback parity with the hint attached; the larger shapes + // exercise the sparse shaders themselves. test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 1, 256, 512, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 8, 64, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 17, 64, 128, false)); @@ -10347,6 +10472,30 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k(1024, 64, 65, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 257, 256, 512, false)); + // mask wider than the selection: qwen4exp QSA prefill produces mask rows with more + // finite cells than n_top_k (whole-block selection + the visible tail). A backend + // that bounds the compacted list by the hint instead of the mask drops those cells. + test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false, 1, 0, GGML_TYPE_F16, 512, 64, 1, true, 768)); + // prefill-shaped qwen4exp QSA cases (nb >= 64 is what engages the mask-compaction + // sparse path; the nb=1..4 cells above only cover the gather path) + test_cases.emplace_back(new test_flash_attn_ext_top_k(32768, 64, 0, 2051, false, 1, 0, GGML_TYPE_F16, 256, 24, 2, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(32768, 128, 0, 2051, false, 1, 0, GGML_TYPE_F16, 256, 24, 2, false)); + // The same shapes with realistic adjacent-token overlap. A QSA selection is whole blocks of + // 4 cells, so neighbouring tokens share most of their picks - which is what makes a + // deduplicated UNION (one set per group of query rows, instead of one set per row) worth + // building for prefill. ov=0 above says nothing about it: with disjoint picks the union is + // the whole cache and the gate correctly declines. + // + // These four cells exist to cover the grouped union's host addressing (the union/gather/FA + // dispatch sequence), which dense-fallback cases cannot: a wrong push constant there is + // invisible whenever the union is not taken. That gate is a measurement, so run them with + // GGML_VK_FA_TOPK_UNION_GQA=1 GGML_VK_FA_UNION_FORCE=1 to have the path taken at all - + // without it they cover the fallback, which is still worth asserting but is not this. + for (int nb : { 64, 128 }) { + for (int ov : { 60, 86 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(32768, nb, 0, 2051, false, 1, ov, GGML_TYPE_F16, 256, 24, 2, false)); + } + } // ns > 1: the split-K partial-output path indexes O and L/M by stream, so these cover // the stream stride in both regions (single tile and multi-tile). // small-batch decode (speculative drafts): each token gets its own gathered top-k block, @@ -10696,6 +10845,13 @@ static std::vector> make_test_cases_perf() { // cost-partition probes: no mask; f16 accumulate test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_DEFAULT, GGML_TYPE_F16, GGML_TYPE_F16)); + // Sparse flash attention (n_kv_max hint) decode across KV depths. + // Shapes: 576/512 DeepSeek MLA, 512/512 DeepSeek-V4/GLM-5.2, 256/256 gqa12 Qwen QSA. + for (int64_t kv : {4096, 16384, 32768}) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + } test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); @@ -10897,6 +11053,11 @@ static std::vector> make_test_cases_perf() { for (int kv : { 11008, 19200, 35584, 68352, 133888 }) { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, 2048, 2304, 512, false)); } + // qwen4exp QSA prefill: the shape the sparse mask compaction exists for. 24 query heads + // over 2 KV heads (gqa 12), no dense prefix, 2051 selected cells per row. + for (int nb : { 64, 128, 512 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(32768, nb, 0, 2051, false, 1, 0, GGML_TYPE_F16, 256, 24, 2, false)); + } // small-batch decode at depth: the speculative-draft regime (batch 2-8), where the old // path fell through to dense attention over the whole compressed KV. for (int kv : { 11008, 35584, 133888 }) {