Commit 78febbad9 for llama.cpp
commit 78febbad9d4fcb6d85d65207265c875f5aa1c1af
Author: uvos <carl@uvos.xyz>
Date: Thu Oct 29 09:14:11 2026 +0100
HIP/CUDA: enable sparse attn on CDNA (#31259)
* HIP/CUDA: enable sparse attn on CDNA
* Use __ballot for __ballot_sync on old ROCm versions
* Ignore new large head fattn mma variant
diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index 296aa2b34..31a273919 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -1866,7 +1866,7 @@ static __global__ void flash_attn_ext_f16(
#endif // defined(AMD_WMMA_AVAILABLE)
#if defined(AMD_MFMA_AVAILABLE)
- if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 32)) {
+ if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 32 && !use_sparse)) {
NO_DEVICE_CODE;
return;
}
@@ -2054,7 +2054,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
bool use_sparse = false;
if (logit_softcap == 0.0f) {
constexpr bool use_logit_softcap = false;
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !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, ncols2)) {
constexpr bool use_sparse_kernel = true;
@@ -2077,7 +2077,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
}
}
} else
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
{
constexpr bool use_sparse_kernel = false;
fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 1515572fe..f1ebd4e8b 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -5,7 +5,8 @@
#include "fattn-vec.cuh"
#include "fattn.cuh"
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !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)
@@ -14,10 +15,20 @@ static __global__ void flash_attn_mask_to_sparse_indices(
const int n_kv_max, const int64_t s31, const int64_t s33) {
ggml_cuda_pdl_sync();
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int values_per_lane = 8;
+ using lane_mask_t = std::conditional_t<warp_size == 64, uint64_t, uint32_t>;
+ constexpr uint64_t ballot_mask = ~(lane_mask_t)0;
+ const auto popc = [](const lane_mask_t mask) {
+ if constexpr (warp_size == 64) {
+ return __popcll(mask);
+ } else {
+ return __popc(mask);
+ }
+ };
const int tid = threadIdx.x;
- const int warp = tid / WARP_SIZE;
- const int lane = tid % WARP_SIZE;
+ const int warp = tid / warp_size;
+ const int lane = tid % warp_size;
const int sequence = blockIdx.y;
const int group = blockIdx.x;
@@ -27,7 +38,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
const half * mask = mask_ptr + sequence*s33 + q0*s31;
int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + group)*n_kv_max;
- __shared__ int warp_offsets[256/WARP_SIZE];
+ __shared__ int warp_offsets[256/warp_size];
__shared__ int row_count;
__shared__ int chunk_count;
@@ -37,11 +48,11 @@ static __global__ void flash_attn_mask_to_sparse_indices(
__syncthreads();
for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) {
- uint32_t selected_warp[values_per_lane];
+ lane_mask_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 int i = i0 + (warp*values_per_lane + item)*warp_size + lane;
bool selected = false;
if (i < ne30) {
#pragma unroll
@@ -49,8 +60,8 @@ static __global__ void flash_attn_mask_to_sparse_indices(
selected |= (!oob || q < q1 - q0) && isfinite(__half2float(mask[q*s31 + i]));
}
}
- selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected);
- warp_count += __popc(selected_warp[item]);
+ selected_warp[item] = (lane_mask_t)__ballot_sync(ballot_mask, (int) selected);
+ warp_count += popc(selected_warp[item]);
}
if (lane == 0) {
@@ -61,7 +72,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
if (tid == 0) {
int offset = 0;
#pragma unroll
- for (int iw = 0; iw < 256/WARP_SIZE; ++iw) {
+ for (int iw = 0; iw < 256/warp_size; ++iw) {
const int count = warp_offsets[iw];
warp_offsets[iw] = offset;
offset += count;
@@ -70,16 +81,16 @@ static __global__ void flash_attn_mask_to_sparse_indices(
}
__syncthreads();
- const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1;
+ const lane_mask_t lane_mask = lane == 0 ? 0 : ((lane_mask_t)1 << 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) {
+ 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] & ((lane_mask_t)1 << lane)) && dst < n_kv_max) {
indices[dst] = i;
}
- warp_item_offset += __popc(selected_warp[item]);
+ warp_item_offset += popc(selected_warp[item]);
}
__syncthreads();
@@ -101,13 +112,13 @@ static __global__ void flash_attn_mask_to_sparse_indices(
// the dependent grid reads indices, signal once the row is complete
ggml_cuda_pdl_lc();
}
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
void ggml_cuda_flash_attn_ext_compact_mask(
const ggml_tensor * mask, int32_t * indices, int32_t * counts, int32_t n_queries, int32_t ncols1, int32_t n_kv_max, cudaStream_t stream) {
-#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+#ifdef GGML_USE_MUSA
GGML_UNUSED_VARS(mask, indices, counts, n_queries, ncols1, n_kv_max, stream);
- GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA");
+ GGML_ABORT("sparse flash attention is not supported on MUSA");
#else
const int64_t s31 = mask->nb[1] / sizeof(half);
const int64_t s33 = mask->nb[3] / sizeof(half);
@@ -122,16 +133,17 @@ void ggml_cuda_flash_attn_ext_compact_mask(
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)
+#endif // !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, const int ncols2) {
-#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+#ifdef GGML_USE_MUSA
GGML_UNUSED_VARS(cc, dst, ncols1, ncols2);
return false;
#else
const ggml_tensor * Q = dst->src[0];
const ggml_tensor * K = dst->src[1];
+ const ggml_tensor * V = dst->src[2];
const ggml_tensor * mask = dst->src[3];
float max_bias = 0.0f;
@@ -139,16 +151,21 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_
memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float));
memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
+ // On CDNA only the DKQ=576 kernel variant with ncols2=16 is compiled with sparse support.
+ if (amd_mfma_available(cc) && !ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(int(Q->ne[0]), int(V->ne[0]), 1, 16)) {
+ return false;
+ }
+
const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4);
// 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) &&
+ return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_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<int64_t>(4096, 2*n_gather);
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+ K->ne[1] >= std::max<int64_t>(4096, GGML_CUDA_CC_IS_NVIDIA(cc) ? 2*n_gather : 2LL*n_kv_max);
+#endif // !defined(GGML_USE_MUSA)
}
template <int DKQ, int DV, int ncols2>
@@ -156,7 +173,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
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 !defined(GGML_USE_MUSA)
if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) {
// 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);
@@ -165,7 +182,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
return;
}
}
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
if constexpr (ncols2 <= 8) {
if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) {
@@ -365,7 +382,9 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg
GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
const int gqa_ratio = Q->ne[2] / K->ne[2];
- if (gqa_ratio == 20 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // GLM 4.7 Flash
+ if (amd_mfma_available(cc) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 16)) {
+ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst);
+ } else if (gqa_ratio == 20 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // GLM 4.7 Flash
if (cc >= GGML_CUDA_CC_DGX_SPARK) {
if (Q->ne[1] <= 8) {
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst);
@@ -679,6 +698,11 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
return BEST_FATTN_KERNEL_MMA_F16;
}
+ // On CDNA, the sparse MMA kernel is used for DKQ=576 regardless of the head size heuristics.
+ if (amd_mfma_available(cc) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 16)) {
+ return BEST_FATTN_KERNEL_MMA_F16;
+ }
+
// AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
if (amd_mfma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {
diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h
index 8e8f5a4ee..61769f400 100644
--- a/ggml/src/ggml-cuda/vendors/hip.h
+++ b/ggml/src/ggml-cuda/vendors/hip.h
@@ -254,6 +254,10 @@ typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
#define FP8_AVAILABLE
#endif // HIP_VERSION >= 60300000
+#if HIP_VERSION < 70000000
+#define __ballot_sync(mask, var) __ballot(var)
+#endif //HIP_VERSION < 70000000
+
typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4)));
static __device__ __forceinline__ int __vsubss4(const int a, const int b) {
diff --git a/scripts/hip/gcn-cdna-vgpr-check.py b/scripts/hip/gcn-cdna-vgpr-check.py
index 1f1c2c903..146bc8ff5 100644
--- a/scripts/hip/gcn-cdna-vgpr-check.py
+++ b/scripts/hip/gcn-cdna-vgpr-check.py
@@ -84,6 +84,7 @@ def main():
'_ZL18flash_attn_ext_f16ILi576ELi512ELi8ELi4ELb0ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
'_ZL18flash_attn_ext_vecILi256ELi1EL9ggml_type2ELS0_2ELb0EEvPKcS2_S2_S2_S2_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS6_IjLj3EEiiiiiiiiiiiliiliiiiil',
'_ZL18flash_attn_ext_f16ILi320ELi256ELi2ELi32ELb0ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
+ '_ZL18flash_attn_ext_f16ILi576ELi512ELi1ELi16ELb0ELb1ELb1EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil',
}
functions = parse_log_file(log_file)