From 20edd90d618a18cfc4c4a60e3673b9890bc1bb4f Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Tue, 8 Sep 2026 02:59:59 +0300 Subject: [PATCH 01/15] vulkan: sparse flash attention for qwen4exp top-k masks Port PR #28105's sparse flash-attention compaction and wire it to the qwen4exp QSA prefill mask. The mask of a QSA layer is exactly the top-k selection intersected with causality, so only n_kv_max (= top-k width) cells per row are finite; the backend now compacts those positions per mask row and flash attention reads K/V/mask through the per-row index list instead of scanning the whole cache. - ggml_flash_attn_ext_set_sparse stores the per-row finite bound in op_params[5] (op_params[4] stays the fork's n_kv_raw); CUDA fattn passes the hint through for reference. - flash_attn_sparse_compact.comp builds the per-row index list with a deterministic subgroup-ballot scan (ascending position order, -1 padded). Upstream's atomic slot assignment is a race: the list order is the softmax accumulation order, so the run-to-run bits differ; the scan makes the sparse path bit-stable and identical between the cache on/off arms, which the A/B harness requires. - vulkan FA pipelines gain USE_SPARSE (bit 16) and the fork's DYNAMIC_KV moves to bit 32; cm2's sparse-only tensor-layout updates and gather offsets stay behind USE_SPARSE so the dense specialization keeps its codegen (an unguarded runtime KV select halved dense throughput on gfx1151). - The sparse gate follows the tiling contract of the shader: one index list and one mask row are resolved per DISPATCH TILE, so every row of a tile has to be the same query. That holds when the rows are the gqa heads of one token (gqa_ratio > 1, i.e. decode); large-N shapes run with gqa_ratio == 1 and correctly decline to dense. It also declines when the cache is under max(4096, min_ratio * n_kv_max) cells. FA_SPARSE_DISABLE reverts to dense for A/B. - Extend flash_attn_union/gather_union with a KV-head dimension and a batch offset, plus a grouped prefill driver (64-row groups, opt-in via GGML_VK_FA_TOPK_UNION_GQA): one compact set per group with the scratch reused per group. Inert by default; the per-row sparse path measured ahead of any shared-set compaction. That pp512 measurement was the broken configuration: it took the shared-tile path, which is fast and wrong. With the tiling fixed (see the following commit), the sparse path serves decode (gqa_ratio > 1) and prefill declines to dense. Micro model pp512 and A/B figures above therefore do not describe this commit as merged; re-measure before quoting them. --- ggml/include/ggml.h | 6 + ggml/src/ggml-cuda/fattn-common.cuh | 24 +- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 197 +++++--- ggml/src/ggml-cuda/fattn-tile.cuh | 12 +- ggml/src/ggml-cuda/fattn-vec.cuh | 2 +- ggml/src/ggml-cuda/fattn.cu | 136 ++++++ ggml/src/ggml-vulkan/ggml-vulkan.cpp | 426 ++++++++++++++++-- .../vulkan-shaders/flash_attn.comp | 46 +- .../vulkan-shaders/flash_attn_base.glsl | 36 +- .../vulkan-shaders/flash_attn_cm1.comp | 49 +- .../vulkan-shaders/flash_attn_cm2.comp | 64 ++- .../flash_attn_gather_union.comp | 18 +- .../flash_attn_sparse_compact.comp | 90 ++++ .../vulkan-shaders/flash_attn_union.comp | 3 +- .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + ggml/src/ggml.c | 10 + src/llama-graph.cpp | 17 +- src/llama-graph.h | 1 + src/models/deepseek4.cpp | 10 +- src/models/qwen4exp.cpp | 7 +- tests/test-backend-ops.cpp | 120 ++++- 21 files changed, 1092 insertions(+), 183 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp 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..f4d128517dcd 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1141,6 +1141,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 +1973,19 @@ 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; + 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 +2182,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 +2551,7 @@ 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 + 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 +4137,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 +4146,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 +6145,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 { @@ -11959,7 +11984,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 +12103,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: @@ -12143,7 +12175,7 @@ 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; @@ -12163,6 +12195,172 @@ 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; + + 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; + + ggml_vk_sync_buffers(ctx, subctx); + 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), st.ul_buf, st.kv_buf }, 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), st.ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), + st.kc_buf, st.mc_buf, st.kv_buf }, 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), st.ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), + st.vc_buf, st.mc_buf, st.kv_buf }, 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]; + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + dst_buf.offset += batch_off * (uint64_t) dst->nb[1]; + vk_subbuffer sinks_buf = q_buf; + const vk_subbuffer k_buf = st.kc_buf, v_buf = st.vc_buf, mask_buf = st.mc_buf; + + const vk_flash_attn_push_constants pc = { N, st.kv_c, + (uint32_t) N, neq2, 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, st.kv_buf}, + pc, { dispatch_x, workgroups_y, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t) N, neq2, neq2, 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, { (uint32_t) N, HSV, neq2 }); + } else { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, q_buf, st.kv_buf}, + 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 +12391,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 +12445,41 @@ 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; + static const bool union_gqa_env = getenv("GGML_VK_FA_TOPK_UNION_GQA") != nullptr; + // The grouped union for a GQA cache stays opt-in: the per-row sparse compaction above + // attends each row's own selection (n_kv_max rows) and measured ahead of any shared-set + // compaction at depth. + if ((union_gqa_env && (!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)) && 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. + const uint32_t union_batch = gqa_form ? 64u : 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. + const 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_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 +12497,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 +12516,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 +12534,51 @@ 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 = 64; + 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); + const size_t gneed = gkc_sz + gvc_sz + gmc_sz + gul_sz; + 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 +12611,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 +12623,7 @@ 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, }; ggml_vk_dispatch_pipeline(ctx, subctx, gather_pipe, { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), @@ -12386,10 +12645,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 +12842,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 +12929,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 +12997,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 +13007,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 +13047,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 +13068,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 +13115,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 +13147,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 +13215,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 +13265,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 +13280,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) { 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..016923e829b9 --- /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[8]; + +// 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/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..23fc3aa1f59a 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2526,6 +2526,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 +2569,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 +2720,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 +2819,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 +2910,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 +2995,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 +3087,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 +3160,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 +3219,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/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/qwen4exp.cpp b/src/models/qwen4exp.cpp index bf4e44f0634d..b7ec7ea02496 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -903,9 +903,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 diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 75c2fb6449ac..b7390098ae34 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) { @@ -7143,11 +7172,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 +7194,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 +7256,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 +7269,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 +7305,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 +7325,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 +7359,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 +7398,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 +7416,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; + } } } @@ -10179,6 +10233,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 }) { @@ -10347,6 +10431,14 @@ 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)); // 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 +10788,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 +10996,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 }) { From 7e56d92e8659beaeadec1fdc32bdeab7f7c51be9 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Thu, 10 Sep 2026 21:39:16 +0300 Subject: [PATCH 02/15] vulkan: size the sparse-FA compaction per-subgroup tally to the workgroup The compaction shader tallied each chunk's per-subgroup finite counts in a shared array of 8, but the pipeline runs 1024 threads = 16 subgroups on wave-64 devices. Waves 8..15 wrote past the array: their counts were lost, the slot assignment shifted, and each mask row kept only ~930 of its ~2051 finite positions - a pseudo-random subset of the selection. Attention then read the wrong half of the selected cells: text stayed locally coherent but the model lost global structure (long-range analysis hallucinated non-existent issues). On the micro model at pf12288/c16384 the compacted list now matches the mask row exactly (2051/2051, ascending, in-bounds) and greedy decode tokens are identical to the dense-mask arm; residual logits drift is ~1e-5..1e-4 relative from online-softmax reblocking, same class as any summation reordering. --- .../ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 016923e829b9..6269deb0e0e3 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp @@ -24,7 +24,7 @@ layout (push_constant) uniform parameter { shared uint count; shared uint chunk_base; -shared uint sub_tot[8]; +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 From c46ecbcd38063d427d5890abc8cf79a0fdc22e39 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Thu, 10 Sep 2026 22:06:14 +0300 Subject: [PATCH 03/15] docs: recipe - shared per-subgroup tally sizing bug in the sparse FA compaction --- docs/fix-sparse-compact-subtot-overflow.md | 37 ++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 docs/fix-sparse-compact-subtot-overflow.md 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. From 949cc2cd4f134e26b2f16eca4c5f84fe76690930 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Fri, 11 Sep 2026 18:27:24 +0300 Subject: [PATCH 04/15] docs: recipe - the sparse-FA shared-tile tiling contract Write down how the "one index list and one mask row per dispatch tile" bug was localized, since the same shape of mistake is easy to re-introduce: a tile resolves one row for all its rows, so the gate has to encode the invariant that makes the rows interchangeable (gqa_ratio > 1), not a proxy for it (N >= 64). Includes the diagnostic method that found it -- sweep how much the query rows share and watch the error collapse monotonically -- and the two measurement traps hit on the way: a test that never set the sparse hint, so it measured dense while claiming to test sparse, and a generated shader header whose DEPENDS did not list the .comp sources, so shader edits were not compiled at all. --- docs/fix-sparse-fa-shared-tile-list.md | 85 ++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 docs/fix-sparse-fa-shared-tile-list.md 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. From ec2265962af96e62922be1ea238a1670c9eb9ed9 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Fri, 11 Sep 2026 22:44:05 +0300 Subject: [PATCH 05/15] vulkan: fix grouped union prefill host addressing for non-square head/batch shapes The grouped union prefill path (per-group union of the query rows' top-k selections, then dense FA over the compact set) was correct only when the query-head count happened to equal the rows in a group. Three host-side shape mismatches hid behind that coincidence, all of them passing the group's dimensions where the shaders expect the destination tensor's: - the FA push constant ne1 carried the group's row count, but the shader uses ne1 as the head-to-head stride of dst ([HSV, n_head_q, n_batch, ns]): o_offset + iq2*HSV + row*ne1*HSV. It must be q->ne[2]. With nh != nb every head but the first was written to another head's rows, and at nb=128 the write ran past the tensor (DEVICE_LOST). - the group's slice of dst advanced by dst->nb[1], the head stride, instead of nb[2], the batch-row stride. - the split-K reduce was dispatched with x enumerating rows and with ne1/ne2 swapped against its own convention (x enumerates heads, z the rows of the split buffer), so every shape small enough to engage split_k was wrong. Found by bisecting the shape: the path failed for nh < nb and passed for nh == nb, which pointed at the host rather than at the shader. Coverage: the union gate is a measurement of the actual overlap, so the call that produces the estimate must itself decline, and test-backend-ops computes a case once - leaving the path with no deterministic coverage. GGML_VK_FA_UNION_FORCE=1 admits it without the estimate, for tests and for A/B runs; it can only cost a slow step, never correctness. The qwen4exp prefill cases with realistic adjacent-token overlap document the two variables they need. test-backend-ops on Vulkan: 13369/13369 with the union forced, with the gate alone (dense fallback) and by default. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 36 +++++++++++++++++++++++----- tests/test-backend-ops.cpp | 26 ++++++++++++++++---- 2 files changed, 52 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f4d128517dcd..08950927abb1 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -12322,13 +12322,22 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co 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[1]; + 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, - (uint32_t) N, neq2, 1, + neq2, rows, 1, neq2, 1, st.n_head_kv, 1, st.n_head_kv, 1, @@ -12348,9 +12357,13 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, q_buf, st.kv_buf}, pc, { dispatch_x, workgroups_y, 1 }); ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t) N, neq2, neq2, 1, split_k, false }; + // 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, { (uint32_t) N, HSV, neq2 }); + {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, st.kv_buf}, @@ -12478,8 +12491,18 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // roughly dense cost for the one step it takes the estimate to catch up. // 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. - const bool worth_it = kv_c_est != 0 && kv_c_est < kv_c_bound && + 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 @@ -12623,7 +12646,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, + 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), diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b7390098ae34..8d4951313e56 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10418,10 +10418,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)); @@ -10439,6 +10441,22 @@ static std::vector> make_test_cases_eval() { // 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, From 161ab59765f6123be75d2816deb7cf574ec65380 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Fri, 11 Sep 2026 22:44:17 +0300 Subject: [PATCH 06/15] docs: recipe - push constants carry the destination shape, not the dispatch tile's Record the grouped-union host-addressing bug as a recipe: the three shape mismatches (ne1 as the head stride, the group's dst slice stepping by nb[1], the split-K reduce's inverted convention), the pass/fail symmetry that located them, and the two traps that delayed it - a PASS verdict from a dense fallback when the gate is a measurement, and a loose -p regex claiming a verdict for a case that never ran. --- ...ix-fa-push-constants-group-vs-dst-shape.md | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 docs/fix-fa-push-constants-group-vs-dst-shape.md 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. From 0b08831cf7531735a9e7463fb2f40846cc8b9fe4 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Sat, 12 Sep 2026 00:52:07 +0300 Subject: [PATCH 07/15] feat: enable the grouped union prefill for GQA caches by default It was opt-in (GGML_VK_FA_TOPK_UNION_GQA=1) pending full-model quality verification; the probes passed, so it becomes the default with =0 as the opt-out. The overlap gate is unchanged and still declines wherever compaction would not pay. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 08950927abb1..22963a556e8c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -12467,11 +12467,13 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // 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; - static const bool union_gqa_env = getenv("GGML_VK_FA_TOPK_UNION_GQA") != nullptr; - // The grouped union for a GQA cache stays opt-in: the per-row sparse compaction above - // attends each row's own selection (n_kv_max rows) and measured ahead of any shared-set - // compaction at depth. - if ((union_gqa_env && (!union_env || union_env[0] != '0')) && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && + // 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)) && ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16 && ggml_vk_fa_union_stat_init(ctx)) { From 0209047a362ffdfa3c9b07b543ed4b1023105990 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Sat, 12 Sep 2026 12:46:12 +0300 Subject: [PATCH 08/15] vulkan: overlap the union scan of group g+1 with flash attention of group g The union scan (one workgroup, ~4 ms per group at depth) ran serialized behind the previous group's FA behind full barriers. Per-group slots for the union index list and the kv-count word make scan(g+1) data-independent of FA(g), so it is issued right after it with no barrier and overlaps it on the GPU. Allocation: the prealloc_y sizing gains one union-list slot per group (gul_all = gul_sz * n_groups); the fa_union_stat slot stride is 16 bytes. The estimator still reads slot 0 (group 0's last count). test-backend-ops FLASH_ATTN_EXT: 13369/13369. pp4096 full model (Q4_K_S, RADV STRIX_HALO): d2048 494.5 (was 500, noise band) d32768 392.2 (was 343, +14%) d65536 352.1 (was 287, +23%) d131072 301.7 (was 248, +22%) --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 64 +++++++++++++++++++++------- 1 file changed, 49 insertions(+), 15 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 22963a556e8c..99d0c7aa98a3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -12220,17 +12220,47 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co const uint32_t k_row_words = st.row_bytes / 4; const uint32_t v_row_words = st.v_row_bytes / 4; - 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; - - ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_union_push_constants upc = { + // 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 = 16; + 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, batch_off, + (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), st.ul_buf, st.kv_buf }, upc, { 1, 1, 1 }); + { 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 = { @@ -12242,8 +12272,8 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co batch_off, }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, - { ggml_vk_tensor_subbuffer(ctx, k), st.ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), - st.kc_buf, st.mc_buf, st.kv_buf }, gkpc, { st.kv_c, st.n_head_kv, 1 }); + { 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), @@ -12253,8 +12283,8 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co batch_off, }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, - { ggml_vk_tensor_subbuffer(ctx, v), st.ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), - st.vc_buf, st.mc_buf, st.kv_buf }, gvpc, { st.kv_c, st.n_head_kv, 1 }); + { 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 @@ -12354,7 +12384,7 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co 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, st.kv_buf}, + {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, @@ -12366,7 +12396,7 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co {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, st.kv_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, q_buf, kv_g}, pc, { workgroups_x, workgroups_y, 1 }); } } @@ -12575,7 +12605,11 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ 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); - const size_t gneed = gkc_sz + gvc_sz + gmc_sz + gul_sz; + // 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); From 2b8bfbfee4ccdf4f6b1e2012b9d1445c1f5c88c8 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Sun, 13 Sep 2026 20:23:59 +0300 Subject: [PATCH 09/15] vulkan: key the union estimate and stat buffer by the real group size The grouped union prices a GQA batch per group, so the estimate slot must be keyed by the group actually measured, min(64, n_batch): pricing a speculative 2-4 row batch as a 64-row group read past the end of the top-k tensor and poisoned the prefill slot, which made prefill decline on every depth for as long as decoding ran. The stat buffer now has one slot per group of the largest supported batch, and the estimator read is clamped. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 39 +++++++++++++++++++++++----- 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 99d0c7aa98a3..8381e3c8ec79 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1976,6 +1976,17 @@ 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 @@ -2551,6 +2562,8 @@ 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; + // 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; @@ -12124,7 +12137,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 &) { @@ -12133,7 +12146,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; } @@ -12181,7 +12194,10 @@ static uint32_t ggml_vk_fa_union_estimate(ggml_backend_vk_context * ctx, uint32_ 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; } @@ -12229,7 +12245,7 @@ static void ggml_vk_flash_attn_union_groups(ggml_backend_vk_context * ctx, vk_co // 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 = 16; + 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. @@ -12505,13 +12521,24 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ 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 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. - const uint32_t union_batch = gqa_form ? 64u : n_batch; + // + // 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); @@ -12597,7 +12624,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // 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 = 64; + 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]); From a391db92a109bb316468a211687a9819d1d5375e Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Sun, 13 Sep 2026 20:27:06 +0300 Subject: [PATCH 10/15] graph: skip an s_copy input that no node reads 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. a qwen4exp MTP draft context with an all-false recurrent filter) leaves s_copy unallocated, and set_input then aborted on the null buffer. Skip it like any dead input. --- src/llama-graph.cpp | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 23fc3aa1f59a..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; From 29b5b842f98b77c70e4f613f5b70ff0c6e57e786 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Sun, 13 Sep 2026 20:27:06 +0300 Subject: [PATCH 11/15] qwen4exp: run the MTP draft block with its own indexer (sparse QSA) The NextN/MTP block is a full-attention QSA layer: it ships its own trained indexer tensors and its norms sit far from their zero initialization, so the block's attention used the indexer during training. The draft was running that layer dense, which made server pp512 degrade with depth (a dense full-context FA over the whole ubatch, growing linearly): 6936/4494/2235/790 t/s at d0/8k/32k/131k on the micro model against 10631/6445/3943/1983 without the draft. The sidecar's compress_ratios[n_layer] == 0 is a padding artifact of llama-model-saver (the trunk's array padded to n_layer_all), not a statement about the block, so the loader restores the ratio from the trunk's QSA layers when the nextn layer really carries indexer tensors. The MTP context now 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. An empty recurrent set allocates nothing; the old comment claiming its buffer allocation fails was wrong. The draft's sparse path only pays on prefill-sized batches, so graph_mtp enables QSA for n_tokens >= 16 and decode/verify batches run dense while still writing the raw indexer keys (pooled keys are recomputed above the pool watermark from the stored keys). The hybrid graph inputs also skip an s_copy no node reads. Micro model, llama-server, pp512/tg128, one session, dense draft vs QSA draft: pp 6258/4118/2120/734 vs 6157/4579/2350/1019, tg 124/114/94/53 vs 124/112/87/49. Recipe: docs/qwen4exp-mtp-sparse-draft.md --- src/llama-model.cpp | 36 ++++++++---------------- src/models/models.h | 3 +- src/models/qwen4exp.cpp | 61 ++++++++++++++++++++++++++++++++++++----- 3 files changed, 68 insertions(+), 32 deletions(-) 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/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 b7ec7ea02496..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; @@ -923,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; @@ -1563,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); @@ -1592,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, From c00d9b99f2b39193c83d9453273ebd95821fe263 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Sun, 13 Sep 2026 20:27:22 +0300 Subject: [PATCH 12/15] docs: recipe - sparse QSA for the qwen4exp MTP draft block --- docs/qwen4exp-mtp-sparse-draft.md | 87 +++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 docs/qwen4exp-mtp-sparse-draft.md 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. From 188988f2d7801fe2841dc3907865cb66f1b3daca Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Mon, 14 Sep 2026 15:25:35 +0300 Subject: [PATCH 13/15] vulkan: widen the top-k radix pass and batch the emit scan The radix selection used 8-bit digits, so a 32-bit key took four passes over the row. 11+11+10 bits does it in three, which is one full row of traffic less per token. Any digit width is exact for radix select, so the result is unchanged. The emit scan read the row once per class and once more per chunk. Both classes now share one pass: values strictly above the threshold fill [0, n_above) and the ties at the threshold fill [n_above, k). Each class keeps ascending element order, so the output equals the two-pass (above, then ties) result bit for bit while the row is read once. The two classes cannot be merged into a single ">= threshold" pass: the ties must land after all strictly larger values or the sparse attention summation order changes. Slots still come from a deterministic exclusive scan, so the output never depends on scheduling. The emit scan batches EMIT_W elements per invocation and prefix-sums them with subgroup shuffles; a per-chunk ballot path (EMIT_W=1) is kept for devices without subgroup shuffle. Spec constant 2 selects between them. Measured on the qwen4exp prefill shape: plain TOP_K 678 -> 532 ms per window. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 7 +- .../vulkan-shaders/topk_radix_select.comp | 206 +++++++++++++----- tests/test-backend-ops.cpp | 48 +++- 3 files changed, 200 insertions(+), 61 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8381e3c8ec79..23e4f94cb892 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -6441,10 +6441,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); 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..36e50a4bc6dd 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,17 @@ 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 // order-preserving float -> uint mapping uint f2ui(float x) { @@ -60,8 +66,8 @@ 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 +77,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 +237,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 +252,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/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8d4951313e56..8ef205bc2a9c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6237,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 { @@ -6245,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; } @@ -6274,7 +6275,9 @@ 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. 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) { @@ -6288,12 +6291,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)); + } } } } @@ -6307,8 +6316,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 { @@ -9959,6 +9995,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) { From 380b00a2d2cbe16ce45429510ac0473e760fa3fb Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Mon, 14 Sep 2026 15:31:57 +0300 Subject: [PATCH 14/15] vulkan: fold the QSA score transpose into the top-k gather The QSA indexer built its cell-major score by materializing the (1,0,2,3) transpose of the block score and then gathering cells from that copy. At 135k context the copy is 1.09 GB and runs at ~10 GB/s, ~20x below DRAM: each workgroup reads a strided column of the block-major source. The fused kernel wants values, not a layout, so it now reads the score in its native block-major layout, where the run of consecutive blocks that cell_blk yields walks consecutive addresses. The transpose is folded away and never executed; the pattern anchor moves from the gather to the transpose in front of it. A masked cell sums to exactly -inf whatever the block index is, so the kernel skips the block lookup and the score read for it. That is not just a shortcut: the block index of a masked cell is unconstrained, so the shortcut and the general path agree by construction rather than by luck. Every node of the pattern still exists for the unfused fallback (small k) and for the other backends. The gather is value-identical, so the selection and the fingerprints are unchanged. Measured on the qwen4exp prefill shape: fused TOPK_QSA 23.3 -> 14.1 ms per window, the standalone transpose 2.70 s -> 0.52 s, prefill wall 45.2 -> 39.5 s. The fused kernel reads the block score while it writes the cell indices, so it can only write them into the destination when nothing it reads shares storage with it. ggml-alloc does hand the output the memory of the block score, which made the fusion guard decline the whole site: at 131k, 5 of the 12 QSA layers fell back to the unfused transpose + top-k, ~125 ms per site per window. The guard is right, so the output moves instead: when an input overlaps the destination the indices go to the tail of the shared scratch and a buffer copy fills the output after a barrier. The test is on the kernel's read set, which is the whole hazard, rather than on the pattern's elided intermediates - so the guard can be skipped for this fusion without hiding a read-after-write. GGML_VK_QSA_PRIV_FORCE=1 admits the private route without an overlap. Which route a site takes is a property of the allocator's layout, so no test case reaches the private route on its own; the knob is its A/B arm. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 167 +++++++++++++----- .../vulkan-shaders/topk_radix_select.comp | 14 +- tests/test-backend-ops.cpp | 3 +- 3 files changed, 140 insertions(+), 44 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 23e4f94cb892..2151bf90f775 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 ] @@ -16336,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]; @@ -16350,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; @@ -16360,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) { @@ -16376,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; } @@ -17949,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: @@ -18057,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: @@ -19472,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)) { @@ -19497,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 || @@ -20085,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) { 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 36e50a4bc6dd..9ebdf7119ef1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp @@ -41,6 +41,8 @@ shared uint sh_above; 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) { uint y = floatBitsToUint(x); @@ -53,13 +55,19 @@ 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) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8ef205bc2a9c..d866b9dade5a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6277,7 +6277,8 @@ struct test_topk_qsa : public test_case { // 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. + // 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) { From 929aa8d5b89084d07965cb431aceda219d9becf8 Mon Sep 17 00:00:00 2001 From: Aleksey Kutashov Date: Mon, 14 Sep 2026 19:46:22 +0300 Subject: [PATCH 15/15] vulkan: clear the fusion label when a fusion is declined The perf logger names a node's interval after the fusion it was chosen for, and the name is written after the fusion guard has already had its say. A declined fusion therefore kept the fused name while its nodes ran one by one, so the logger billed the fallback under the fused name. That is not cosmetic: it is how a profile showed "12 x TOPK_QSA" at 131k when 7 sites were fused and 5 had declined and were paying for an unfused transpose + top-k. The label now follows the decision. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 2151bf90f775..5d2daa1f827f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -20212,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; } }