Commit dc9879cf6 for llama.cpp
commit dc9879cf66aeb5c2f7c38e9578e6a5f38c497865
Author: Aman Gupta <amangupta052@gmail.com>
Date: Wed Sep 23 23:20:40 2026 +0800
CUDA: enable sparse-fa for dsv4 prefill (again) (#29298)
* CUDA: enable sparse-fa for dsv4 prefill (again)
* 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.
* 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.
---------
Co-authored-by: Pascal <admin@serveurperso.com>
diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index 84219ca92..449a77c5b 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 <int DKQ, int DV, int ncols1, int ncols2>
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<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
use_sparse = true;
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 7098c8b4c..d1fcf58cb 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 <int ncols1, bool oob>
__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,8 +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;
- for (int q = 0; q < q1 - q0 && !selected; ++q) {
- selected = i < ne30 && isfinite(__half2float(mask[q*s31 + i]));
+ if (i < ne30) {
+#pragma unroll
+ 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]);
@@ -110,15 +114,20 @@ 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);
+ // 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, 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());
#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 +141,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<int64_t>(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 +158,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<DKQ, DV, 1, ncols2>(ctx, dst);
return;
}
@@ -629,7 +641,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;