Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

HIP: RDNA3 mma FA, faster AMD transpose, tune AMD (#22880)

Browse Source 
Adds RDNA3 support to the CUDA mma FA kernel. To make the RDNA3 tensor cores work with the FP16 accumulation for VKQ the tiles they need to be 32 logical units long in direction of the attention head; for head sizes 80 and 112 that are not exactly divided by 32 the regular length of 16 with FP32 accumulation is used instead. The longer tiles also enable more efficient transposition for a warp size of 32 which is why it's also used for RDNA4. However, this scrambles the data layout of the accumulators along the attention head dimension. To prevent accidental misuse I added another entry to ggml_cuda_mma::data_layout.

I also tuned the kernel parameters for RDNA3, RDNA4, and CDNA1 in general, during which I discovered that the kernel can be made to work for head sizes up to 256 for CDNA. For RDNA3/4 I was not able to get better performance that the tile kernel for head sizes > 128.
This commit is contained in:
Johannes Gäßler
2026-05-14 22:58:58 +02:00
committed by GitHub
parent d81e63dcfd
commit 3e037f313c
3 changed files with 400 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/mma.cuh
+136 -13
View File
@@ -80,6 +80,7 @@ namespace ggml_cuda_mma {
DATA_LAYOUT_J_MAJOR = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3.
DATA_LAYOUT_I_MAJOR_MIRRORED = 20, // Volta, matrix A&B for RDNA3.
DATA_LAYOUT_J_MAJOR_MIRRORED = 30,
DATA_LAYOUT_I_MAJOR_SCRAMBLED = 40, // Scrambled matrix C for faster transposition (RDNA4/CDNA), convert to float to unscramble.
};
// Implemented mma combinations are:
// - (I_MAJOR, I_MAJOR) -> I_MAJOR
@@ -312,13 +313,19 @@ namespace ggml_cuda_mma {
half2 x[ne] = {{0.0f, 0.0f}};
static constexpr __device__ bool supported() {
if (I == 16 && J == 8) return true;
if (I == 16 && J == 8) return true;
if (I == 16 && J == 16) return true;
if (I == 32 && J == 8) return true;
return false;
}
static __device__ __forceinline__ int get_i(const int l) {
if constexpr (I == 16 && J == 8) {
return threadIdx.x % 16;
} else if constexpr (I == 16 && J == 16) {
return threadIdx.x % 16;
} else if constexpr (I == 32 && J == 8) {
return (threadIdx.x % 16) * 2 + l / (ne/2);
} else {
NO_DEVICE_CODE;
return -1;
@@ -327,7 +334,15 @@ namespace ggml_cuda_mma {
static __device__ __forceinline__ int get_j(const int l) {
if constexpr (I == 16 && J == 8) {
return ne * (threadIdx.x / 16) + l;
return (threadIdx.x / 16) * ne + l;
} else if constexpr (I == 16 && J == 16) {
#ifdef RDNA3
return l*2 + (threadIdx.x / 16);
#else
return (threadIdx.x / 16) * ne + l;
#endif // RDNA3
} else if constexpr (I == 32 && J == 8) {
return (threadIdx.x / 16) * (ne/2) + l % (ne/2);
} else {
NO_DEVICE_CODE;
return -1;
@@ -338,13 +353,19 @@ namespace ggml_cuda_mma {
half2 x[ne] = {{0.0f, 0.0f}};
static constexpr __device__ bool supported() {
if (I == 16 && J == 8) return true;
if (I == 16 && J == 8) return true;
if (I == 16 && J == 16) return true;
if (I == 32 && J == 8) return true;
return false;
}
static __device__ __forceinline__ int get_i(const int l) {
if constexpr (I == 16 && J == 8) {
return threadIdx.x % 16;
} else if constexpr (I == 16 && J == 16) {
return threadIdx.x % 16;
} else if constexpr (I == 32 && J == 8) {
return (threadIdx.x % 16) * 2 + l / (ne/2);
} else {
NO_DEVICE_CODE;
return -1;
@@ -353,7 +374,11 @@ namespace ggml_cuda_mma {
static __device__ __forceinline__ int get_j(const int l) {
if constexpr (I == 16 && J == 8) {
return ne * (threadIdx.x / 16) + l;
return (threadIdx.x / 16) * ne + l;
} else if constexpr (I == 16 && J == 16) {
return (threadIdx.x / 16) * ne + l;
} else if constexpr (I == 32 && J == 8) {
return (threadIdx.x / 16) * (ne/2) + l % (ne/2);
} else {
NO_DEVICE_CODE;
return -1;
@@ -516,12 +541,15 @@ namespace ggml_cuda_mma {
if (I == 16 && J == 16) return true;
if (I == 16 && J == 8) return true;
if (I == 16 && J == 4) return true;
if (I == 32 && J == 8) return true;
return false;
}
static __device__ __forceinline__ int get_i(const int /*l*/) {
if constexpr (supported()) {
static __device__ __forceinline__ int get_i(const int l) {
if constexpr (I == 16) {
return threadIdx.x % 16;
} else if constexpr (I == 32) {
return (threadIdx.x % 16) * 2 + l / (ne/2);
} else {
NO_DEVICE_CODE;
return -1;
@@ -529,8 +557,10 @@ namespace ggml_cuda_mma {
}
static __device__ __forceinline__ int get_j(const int l) {
if constexpr (supported()) {
if constexpr (I == 16) {
return l;
} else if constexpr (I == 32) {
return l % (ne/2);
} else {
NO_DEVICE_CODE;
return -1;
@@ -644,6 +674,40 @@ namespace ggml_cuda_mma {
}
};
template <int I_, int J_>
struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> {
static constexpr int I = I_;
static constexpr int J = J_;
static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_SCRAMBLED;
static constexpr int ne = I * J / ggml_cuda_get_physical_warp_size();
half2 x[ne] = {{0.0f, 0.0f}};
static constexpr __device__ bool supported() {
if (I == 16 && J == 16) return true;
return false;
}
static __device__ __forceinline__ int get_i(const int l) {
return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l);
}
};
static __device__ __forceinline__ tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> unscramble(const tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & t) {
#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> ret;
#pragma unroll
for (int l0 = 0; l0 < t.ne/2; ++l0) {
ret.x[2*l0 + 0] = __lows2half2(t.x[l0], t.x[l0 + t.ne/2]);
ret.x[2*l0 + 1] = __highs2half2(t.x[l0], t.x[l0 + t.ne/2]);
}
return ret;
#else
NO_DEVICE_CODE;
GGML_UNUSED(t);
#endif // defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
}
#if defined(TURING_MMA_AVAILABLE)
template <int I, int J>
static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) {
@@ -660,6 +724,21 @@ namespace ggml_cuda_mma {
ret.x[0] = ggml_cuda_movmatrix(t.x[0]);
ret.x[1] = ggml_cuda_movmatrix(t.x[1]);
return ret;
}
#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
static __device__ __forceinline__ tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> get_half2(
const tile<16, 16, float, DATA_LAYOUT_I_MAJOR> & tile_float) {
tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> ret;
#pragma unroll
for (int l = 0; l < tile_float.ne; ++l) {
float tmp[2];
int i = threadIdx.x / 16;
tmp[i] = tile_float.x[l];
i ^= 1;
tmp[i] = __shfl_xor_sync(0xFFFFFFFF, tile_float.x[l], 16, WARP_SIZE);
ret.x[l] = make_half2(tmp[0], tmp[1]);
}
return ret;
}
#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
@@ -802,21 +881,35 @@ namespace ggml_cuda_mma {
#endif // defined(VOLTA_MMA_AVAILABLE)
}
template <typename T>
template <int I, typename T, data_layout dl>
static __device__ __forceinline__ void load_ldmatrix_trans(
tile<16, 8, T> & t, const T * __restrict__ xs0, const int stride) {
tile<I, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) {
#ifdef TURING_MMA_AVAILABLE
static_assert(I == 16, "bad tile width");
static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
int * xi = (int *) t.x;
const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2);
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
: "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
: "l"(xs));
#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
half * xh = (half *) t.x;
static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
if constexpr (I == 32) {
#pragma unroll
for (int l = 0; l < t.ne; ++l) {
xh[2*l + 0] = ((const half *) xs0)[(2*t.get_j(l) + 0)*(2*stride) + t.get_i(l)];
xh[2*l + 1] = ((const half *) xs0)[(2*t.get_j(l) + 1)*(2*stride) + t.get_i(l)];
for (int l0 = 0; l0 < t.ne/2; ++l0) {
const half2 tmp0 = xs0[(2*t.get_j(l0) + 0)*stride + t.get_i(l0)/2];
const half2 tmp1 = xs0[(2*t.get_j(l0) + 1)*stride + t.get_i(l0)/2];
t.x[l0] = __lows2half2(tmp0, tmp1);
t.x[l0 + t.ne/2] = __highs2half2(tmp0, tmp1);
}
} else {
half * xh = (half *) t.x;
#pragma unroll
for (int l = 0; l < t.ne; ++l) {
xh[2*l + 0] = ((const half *) xs0)[(2*t.get_j(l) + 0)*(2*stride) + t.get_i(l)];
xh[2*l + 1] = ((const half *) xs0)[(2*t.get_j(l) + 1)*(2*stride) + t.get_i(l)];
}
}
#else
GGML_UNUSED_VARS(t, xs0, stride);
@@ -972,6 +1065,20 @@ namespace ggml_cuda_mma {
#endif // TURING_MMA_AVAILABLE
}
static __device__ __forceinline__ void mma(
tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR> & A,
const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> & B) {
#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
tile<16, 8, half2> * D16 = (tile<16, 8, half2> *) &D;
const tile<16, 8, half2> * A16 = (const tile<16, 8, half2> *) &A;
mma(D16[0], A16[0], B);
mma(D16[1], A16[1], B);
#else
GGML_UNUSED_VARS(D, A, B);
NO_DEVICE_CODE;
#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)
}
template <data_layout dl_ab, data_layout dl_d>
static __device__ __forceinline__ void mma(
tile<16, 8, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<8, 8, float, dl_ab> & B) {
@@ -1296,6 +1403,22 @@ namespace ggml_cuda_mma {
#endif // defined(VOLTA_MMA_AVAILABLE)
}
static __device__ __forceinline__ void mma(
tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & A,
const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) {
#if defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
using halfx16_t = __attribute__((ext_vector_type(16))) _Float16;
halfx16_t * xD = (halfx16_t *) D.x;
const halfx16_t * xA = (const halfx16_t *) A.x;
const halfx16_t * xB = (const halfx16_t *) B.x;
xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[0], xB[0], xD[0], /*opsel =*/ 0);
xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[1], xB[0], xD[0], /*opsel =*/ 1);
#else
GGML_UNUSED_VARS(D, A, B);
NO_DEVICE_CODE;
#endif // TURING_MMA_AVAILABLE
}
template <data_layout dl_d, data_layout dl_ab>
static __device__ __forceinline__ void mma(
tile<16, 16, int, dl_d> & D, const tile<16, 4, int, dl_ab> & A, const tile<16, 4, int, dl_ab> & B) {