From 78febbad9d4fcb6d85d65207265c875f5aa1c1af Mon Sep 17 00:00:00 2001 From: uvos Date: Thu, 29 Oct 2026 09:14:11 +0100 Subject: [PATCH] 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 --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 6 +-- ggml/src/ggml-cuda/fattn.cu | 74 ++++++++++++++++++---------- ggml/src/ggml-cuda/vendors/hip.h | 4 ++ scripts/hip/gcn-cdna-vgpr-check.py | 1 + 4 files changed, 57 insertions(+), 28 deletions(-) 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; 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 __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; + 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(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(4096, 2*n_gather); -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + K->ne[1] >= std::max(4096, GGML_CUDA_CC_IS_NVIDIA(cc) ? 2*n_gather : 2LL*n_kv_max); +#endif // !defined(GGML_USE_MUSA) } template @@ -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)