From 4d426d513ef80bc78f1dfb5d546042d7455091e1 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Wed, 23 Sep 2026 15:55:22 +0800 Subject: [PATCH 1/3] CUDA: enable sparse-fa for dsv4 prefill (again) --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 4 ++-- ggml/src/ggml-cuda/fattn.cu | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 84219ca9259d..449a77c5b058 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -2011,7 +2011,7 @@ 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)) } -bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1); +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1, const int ncols2); template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { @@ -2065,7 +2065,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr bool use_logit_softcap = false; #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(cc, dst, ncols1)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, ncols1, ncols2)) { constexpr bool use_sparse_kernel = true; fattn_kernel = flash_attn_ext_f16; use_sparse = true; diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 7098c8b4ca63..b341b341c06c 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -116,9 +116,9 @@ void ggml_cuda_flash_attn_ext_compact_mask( #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) } -bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1) { +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1, const int ncols2) { #if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) - GGML_UNUSED_VARS(cc, dst, ncols1); + GGML_UNUSED_VARS(cc, dst, ncols1, ncols2); return false; #else const ggml_tensor * Q = dst->src[0]; @@ -132,7 +132,8 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_ const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4); - const int64_t n_gather = (ncols1 == 1 ? Q->ne[1] : ncols1) * (int64_t) n_kv_max; + // the dense kernel handles up to 64/ncols2 queries per K/V pass, the single-query gather has to beat that + const int64_t n_gather = (ncols1 == 1 ? std::min(Q->ne[1], 64/ncols2) : ncols1) * (int64_t) n_kv_max; 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 && @@ -148,7 +149,9 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con #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(cc, dst, 1)) { + // a sparse variant at the full tile width gathers the union of its queries once, prefer it for large batches + constexpr bool has_wide_sparse = ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 64/ncols2, ncols2); + if (!(has_wide_sparse && Q->ne[1] > 32/ncols2) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, ncols2)) { ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } @@ -629,7 +632,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // the sparse gather exists only in the MMA kernel: (DKQ, DV, 1, 8) with GQA > 4 const bool sparse_decode = gqa_opt_applies && gqa_ratio > 4 && ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(K->ne[0], V->ne[0], 1, 8) && - ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1); + ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 8); if (!sparse_decode && cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && (Q->ne[0] >= 256 || K->ne[1] >= 8192))) { return BEST_FATTN_KERNEL_VEC; From c5835a883817c1b498800bb1b41cbd9e04dd7364 Mon Sep 17 00:00:00 2001 From: Pascal Date: Wed, 23 Sep 2026 12:04:26 +0200 Subject: [PATCH 2/3] CUDA: unroll the query loop of the sparse mask scan The query loop of flash_attn_mask_to_sparse_indices has a runtime trip count, which keeps the unrolled scan over the values of a lane from issuing its loads together. Template the kernel on ncols1 so the loop is bounded at compile time: batch one decodes compile to straight line code and the scan drops from 46 to 17 us at 49k columns on sparse decode shapes. --- ggml/src/ggml-cuda/fattn.cu | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index b341b341c06c..8f822e726aa0 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -7,10 +7,11 @@ #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // one list per group of ncols1 queries: a column is selected if any query of the group can see it +template __launch_bounds__(256, 1) static __global__ void flash_attn_mask_to_sparse_indices( const half * mask_ptr, int32_t * indices_ptr, int32_t * counts_ptr, const int ne30, const int n_queries, - const int ncols1, const int n_kv_max, const int64_t s31, const int64_t s33) { + const int n_kv_max, const int64_t s31, const int64_t s33) { ggml_cuda_pdl_sync(); constexpr int values_per_lane = 8; @@ -42,7 +43,8 @@ static __global__ void flash_attn_mask_to_sparse_indices( for (int item = 0; item < values_per_lane; ++item) { const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; bool selected = false; - for (int q = 0; q < q1 - q0 && !selected; ++q) { +#pragma unroll + for (int q = 0; q < ncols1 && q < q1 - q0 && !selected; ++q) { selected = i < ne30 && isfinite(__half2float(mask[q*s31 + i])); } selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected); @@ -110,8 +112,10 @@ void ggml_cuda_flash_attn_ext_compact_mask( const dim3 blocks_num((n_queries + ncols1 - 1)/ncols1, 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, counts, int(mask->ne[0]), n_queries, ncols1, n_kv_max, s31, s33); + GGML_ASSERT(ncols1 == 1 || ncols1 == 8); + const auto kernel = ncols1 == 1 ? flash_attn_mask_to_sparse_indices<1> : flash_attn_mask_to_sparse_indices<8>; + ggml_cuda_kernel_launch(kernel, launch_params, + (const half *) mask->data, indices, counts, int(mask->ne[0]), n_queries, n_kv_max, s31, s33); CUDA_CHECK(cudaGetLastError()); #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) } From 1768ddcfaa690e11f82e1a7f4095ae7c8e5ba0b4 Mon Sep 17 00:00:00 2001 From: Pascal Date: Wed, 23 Sep 2026 12:36:04 +0200 Subject: [PATCH 3/3] CUDA: pick the out of bounds check of the sparse mask scan in host code The query loop of the ncols1 == 8 scan keeps a runtime bound and an early exit, so it does not unroll past its first iteration. Template the kernel on whether the last group of queries is partial, decided on the host from n_queries, and hoist the column bound out of the loop: the loop becomes straight line code and the batched sparse op at 49k context drops from 586 to 244 us. --- ggml/src/ggml-cuda/fattn.cu | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 8f822e726aa0..d1fcf58cb17f 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -7,7 +7,7 @@ #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // one list per group of ncols1 queries: a column is selected if any query of the group can see it -template +template __launch_bounds__(256, 1) static __global__ void flash_attn_mask_to_sparse_indices( const half * mask_ptr, int32_t * indices_ptr, int32_t * counts_ptr, const int ne30, const int n_queries, @@ -43,9 +43,11 @@ static __global__ void flash_attn_mask_to_sparse_indices( for (int item = 0; item < values_per_lane; ++item) { const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; bool selected = false; + if (i < ne30) { #pragma unroll - for (int q = 0; q < ncols1 && q < q1 - q0 && !selected; ++q) { - selected = i < ne30 && isfinite(__half2float(mask[q*s31 + i])); + for (int q = 0; q < ncols1; ++q) { + selected |= (!oob || q < q1 - q0) && isfinite(__half2float(mask[q*s31 + i])); + } } selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected); warp_count += __popc(selected_warp[item]); @@ -112,8 +114,11 @@ void ggml_cuda_flash_attn_ext_compact_mask( const dim3 blocks_num((n_queries + ncols1 - 1)/ncols1, 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); + // the last group of queries is partial only if ncols1 does not divide n_queries GGML_ASSERT(ncols1 == 1 || ncols1 == 8); - const auto kernel = ncols1 == 1 ? flash_attn_mask_to_sparse_indices<1> : flash_attn_mask_to_sparse_indices<8>; + const auto kernel = ncols1 == 1 ? flash_attn_mask_to_sparse_indices<1, false> : + n_queries % 8 != 0 ? flash_attn_mask_to_sparse_indices<8, true> : + flash_attn_mask_to_sparse_indices<8, false>; ggml_cuda_kernel_launch(kernel, launch_params, (const half *) mask->data, indices, counts, int(mask->ne[0]), n_queries, n_kv_max, s31, s33); CUDA_CHECK(cudaGetLastError());