wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 32 Packages Projects Releases Wiki Activity

HIP/CUDA: enable sparse attn on CDNA (#31259)

Browse Source 
* HIP/CUDA: enable sparse attn on CDNA

* Use __ballot for __ballot_sync on old ROCm versions

* Ignore new large head fattn mma variant
This commit is contained in:
uvos authored and GitHub committed 2026-10-29 10:14:11 +02:00
1 parent abb76ab380
commit 78febbad9d
4 files changed
+57 -28
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

ggml/src/ggml-cuda/fattn-mma-f16.cuh
+3 -3
View File
@@ -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>;
ggml/src/ggml-cuda/fattn.cu
+49 -25
View File
@@ -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)) {
ggml/src/ggml-cuda/vendors/hip.h
Vendored
+4
View File
@@ -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) {
scripts/hip/gcn-cdna-vgpr-check.py
+1
View File
@@ -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)