From cc26d00966ac9ff24baa1230ff1c4db9fe9eaa60 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Mon, 26 Oct 2026 12:16:50 -0700 Subject: [PATCH 1/4] [ROCm][RDNA] skinny GEMMs: no VGPR spills in wvSplitK on gfx11 and gfx12 wvSplitK_hf_sml_, wvSplitK_hf_ and wvSplitK_hf_big_ held every activation chunk of an unrolled step in registers next to the weights, and bf16 widened the weights to f32. With 1024-thread workgroups a wave gets 128 VGPRs on gfx1150 (192 on gfx1151), and 102 of the 388 kernels spilled on gfx1150, 78 on gfx1151. A spilling GEMM is slow, and its speed depends on the largest scratch an earlier kernel of the process used (1.3x slower after a 256 B/lane kernel, 2.1x after 4 KiB). On gfx1x: - bf16 multiplies with v_dot2_f32_bf16, as fp16 with v_dot2_f32_f16. - A tile holds the activation chunks while (N + YTILE) * UNRL * A_CHUNK / 2 operand VGPRs plus 64 fit the wave's VGPRs; otherwise it reads each chunk from LDS where it is used. gfx9 code is unchanged. No VGPR spills in the 388 kernels on gfx1150, gfx1151 and gfx1201. bench_rocm_skinny_gemm.py, N 1..4, against gfx11: gfx1150 bf16 1.050x, fp16 1.008x; gfx1151 1.00x. The 4096x2048 bf16 GEMM keeps 186-210 us whatever scratch an earlier kernel used (242-651 us before). test_rocm_skinny_gemms.py: 172 passed on gfx1150 and gfx1151. AI assistance was used for this patch. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- csrc/rocm/skinny_gemms.cu | 268 +++++++++++++++++++++++++++----------- 1 file changed, 194 insertions(+), 74 deletions(-) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index b214cfa804a9..6598ff9d8fdf 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -324,17 +324,36 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b, V0 += (s.x + s.y); \ } #elif defined(__HIP__GFX1X__) - // gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12) - #define DOT2C(V0, V2, V3) \ - if constexpr (std::is_same_v) { \ - asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ - } else if constexpr (std::is_same_v) { \ - float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ - __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ - V0 += (s.x + s.y); \ + // gfx1x: v_dot2_f32_f16 and v_dot2_f32_bf16 (VOP3-P, gfx11 and gfx12). The + // bf16 dot keeps the weights packed; widening them to f32 first doubles + // their registers and makes the large tiles spill. + #define DOT2C(V0, V2, V3) \ + if constexpr (std::is_same_v) { \ + asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ + } else if constexpr (std::is_same_v) { \ + asm("v_dot2_f32_bf16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ } #endif +#if defined(__HIP__GFX1X__) + // VGPRs a lane holds: 1024 on gfx1102, gfx1103 and gfx1150..gfx1153, else + // 1536. + #if defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || \ + defined(__gfx1152__) || defined(__gfx1153__) +inline constexpr int kLaneVgprs = 1024; + #else +inline constexpr int kLaneVgprs = 1536; + #endif +// The VGPRs a wave of a THREADS-thread workgroup gets: its wave32 waves share +// a WGP's four SIMDs. +__device__ constexpr int wave_vgprs(int threads) { + const int vgprs = kLaneVgprs / ((threads / 32 + 3) / 4); + return vgprs < 256 ? vgprs : 256; +} +// Registers beside the operand tiles: sums, addresses, loop state. +inline constexpr int kOperandSlackVgprs = 64; +#endif + // To avoid LLVM silently upcasting to double __device__ inline unsigned int min__(uint32_t a, uint32_t b) { return min(a, b); @@ -483,10 +502,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) // are being worked on by each wave. //---------------------------------------------------- float sum[N][YTILE] = {}; - scalar8 sum4[N][YTILE] = {}; + [[maybe_unused]] scalar8 sum4[N][YTILE] = {}; for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - bigType bigA[N][UNRL] = {}; bigType bigB[YTILE][UNRL]; // Fetch the weight matrix from memory! #pragma unroll @@ -498,29 +516,59 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } - // Fetch activation matrix from either just LDS or from both LDS / memory + #if defined(__HIP__GFX1X__) + // Hold the activation chunks too while they fit; else read each + // where it is used, so only the weights' loads stay in flight. + constexpr bool hold_a = + (N + YTILE) * UNRL * (A_CHUNK / 2) + kOperandSlackVgprs <= + wave_vgprs(WvPrGrp * THRDS); + #else + constexpr bool hold_a = true; + #endif + if constexpr (hold_a) { + bigType bigA[N][UNRL] = {}; + // Fetch activation matrix from either just LDS or from both LDS / + // memory #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + } } - } - // Do the matrix multiplication in interleaved manner - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (uint32_t n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) + // Do the matrix multiplication in interleaved manner + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, + 0); + } + } + } + } else { + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k_ = k1 + k2 * THRDS * A_CHUNK + threadIdx.x * A_CHUNK; + #pragma unroll + for (uint32_t n = 0; n < N; n++) { + bigType a = {}; + if (k_ < K) a = *((const bigType*)(&(s[k_ + Kap * n]))); + #pragma unroll + for (int y = 0; y < YTILE; y++) + #pragma unroll for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + DOT2C(sum[n][y], a.f[b], bigB[y][k2].f[b]) } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); } } } @@ -713,10 +761,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) while (m < M) { float sum[N][YTILE] = {}; - scalar8 sum4[N][YTILE] = {}; + [[maybe_unused]] scalar8 sum4[N][YTILE] = {}; for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - bigType bigA[N][UNRL] = {}; bigType bigB[YTILE][UNRL]; // Fetch the weight matrix from memory! #pragma unroll @@ -728,32 +775,67 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } - // Fetch activation matrix from either just LDS or from both LDS / memory + #if defined(__HIP__GFX1X__) + // Hold the activation chunks too while they fit; else read each + // where it is used, so only the weights' loads stay in flight. + constexpr bool hold_a = + (N + YTILE) * UNRL * (A_CHUNK / 2) + kOperandSlackVgprs <= + wave_vgprs(WvPrGrp * THRDS); + #else + constexpr bool hold_a = true; + #endif + if constexpr (hold_a) { + bigType bigA[N][UNRL] = {}; + // Fetch activation matrix from either just LDS or from both LDS / + // memory #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - if (k_ + Kap * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - else - bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + } } - } - // Do the matrix multiplication in interleaved manner - for (uint32_t n = 0; n < N; n++) { + // Do the matrix multiplication in interleaved manner + for (uint32_t n = 0; n < N; n++) { + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, + 0); + } + } + } + } else { + #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) + uint32_t k_ = k1 + k2 * THRDS * A_CHUNK + threadIdx.x * A_CHUNK; + #pragma unroll + for (uint32_t n = 0; n < N; n++) { + bigType a = {}; + if (k_ < K) { + if (k_ + Kap * n < max_lds_len) + a = *((const bigType*)(&(s[k_ + Kap * n]))); + else + a = *((const bigType*)(&(A[k_ + Kap * n]))); + } + #pragma unroll + for (int y = 0; y < YTILE; y++) + #pragma unroll for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + DOT2C(sum[n][y], a.f[b], bigB[y][k2].f[b]) } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); } } } @@ -1010,7 +1092,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) // are being worked on by each wave. //---------------------------------------------------- float sum[N][YTILE] = {}; - scalar8 sum4[N][YTILE] = {}; + [[maybe_unused]] scalar8 sum4[N][YTILE] = {}; //---------------------------------------------------- // Fetch weight matrix B in interleaved K-split! @@ -1030,7 +1112,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) // TODO: Logic below will only work when K is multiple of 8 //---------------------------------------------------- for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - bigType bigA[N][UNRL] = {}; bigType bigB[YTILE][UNRL]; #ifdef PCML @@ -1067,37 +1148,76 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } - // Fetch activation matrix from either just LDS or from both LDS / memory + #if defined(__HIP__GFX1X__) + // Hold the activation chunks too while they fit; else read each + // where it is used, so only the weights' loads stay in flight. + constexpr bool hold_a = + (N + YTILE) * UNRL * (A_CHUNK / 2) + kOperandSlackVgprs <= + wave_vgprs(WvPrGrp * THRDS); + #else + constexpr bool hold_a = true; + #endif + if constexpr (hold_a) { + bigType bigA[N][UNRL] = {}; + // Fetch activation matrix from either just LDS or from both LDS / + // memory #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { #ifdef PCML - bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); + bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); #else - if (k_ + Kap * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - else - bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); #endif + } } - } - // Do the matrix multiplication in interleaved manner + // Do the matrix multiplication in interleaved manner + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, + 0); + } + } + } + } else { + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k_ = k1 + k2 * THRDS * A_CHUNK + threadIdx.x * A_CHUNK; + #pragma unroll + for (uint32_t n = 0; n < N; n++) { + bigType a = {}; + if (k_ < K) { + #ifdef PCML + a = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); + #else + if (k_ + Kap * n < max_lds_len) + a = *((const bigType*)(&(s[k_ + Kap * n]))); + else + a = *((const bigType*)(&(A[k_ + Kap * n]))); + #endif + } + #pragma unroll + for (int y = 0; y < YTILE; y++) #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (uint32_t n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + DOT2C(sum[n][y], a.f[b], bigB[y][k2].f[b]) } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); } } } From 9c15ea7ffd74fd1fc974da562f2fce7e5c588f08 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Mon, 26 Oct 2026 12:16:50 -0700 Subject: [PATCH 2/4] [ROCm][RDNA] int4 skinny GEMMs: no spilling tile on 1024-VGPR archs, only reachable tiles built The int4 dispatch macros expand every tile for N = 1..5, so tiles only picked for N = 1 were also built for N = 2..5, where they are never launched: 480 of the 1684 kernels, and most of the spilling ones. Those arms are now `if constexpr (__N == 1)`; the dispatch is unchanged. The MoE (YTILE 4, W 32, AC 32, UNRL 4) tile for N = 1, K = 2048 spills with the 128 VGPRs a wave gets on a 1024-VGPR arch: gfx1102, gfx1103 and gfx1150..gfx1153 take (4, 4, 16, 16) there; gfx1151 keeps it. Kernel code is unchanged. Spilling kernels: 334 -> 24 on gfx1150 (the tile above, not launched there), 293 -> 8 on gfx1151 (at the VGPR limit, as before). W4A16 group 128, M = 1..4, against gfx11: gfx1150 MoE 1.19x, dense 0.99x-1.00x; gfx1151 1.00x. test_rdna_hybrid_w4a16.py (102 passed), test_hybrid_w4a16_moe.py (70 passed) and test_rocm_skinny_gemms.py (172 passed) on gfx1150 and gfx1151. AI assistance was used for this patch. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- csrc/rocm/skinny_gemms_int4_kernels.cuh | 278 +++++++++++++----------- 1 file changed, 148 insertions(+), 130 deletions(-) diff --git a/csrc/rocm/skinny_gemms_int4_kernels.cuh b/csrc/rocm/skinny_gemms_int4_kernels.cuh index 0653a9bb8c41..c8b27db136cc 100644 --- a/csrc/rocm/skinny_gemms_int4_kernels.cuh +++ b/csrc/rocm/skinny_gemms_int4_kernels.cuh @@ -55,6 +55,20 @@ static bool on_gfx_int4() { return result; } +// Archs with 1024 VGPRs a SIMD lane, where a wave of a 1024-thread workgroup +// gets 128 VGPRs (192 with 1536). +static bool on_1024_vgpr_int4() { + static const bool result = [] { + const std::string_view arch = + at::cuda::getCurrentDeviceProperties()->gcnArchName; + for (std::string_view a : + {"gfx1102", "gfx1103", "gfx1150", "gfx1152", "gfx1153"}) + if (arch.starts_with(a)) return true; + return false; + }(); + return result; +} + #if defined(NDEBUG) #undef NDEBUG #include @@ -1309,64 +1323,67 @@ static int mindiv_int4(int N, int div1, int div2) { TORCH_CHECK(false, "Unsupported group_size=", group_size); \ } -#define WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \ - { \ - if (K_in * N_in > max_lds_len) { \ - /* Chunked branch. Same wave-packing rule as the fitting branch \ - below, for the same reason: it is a property of M, not of \ - chunking. */ \ - if (_sYT < 30) \ - WVSPLIT_INT4G_GS_CHUNKED(4, 2, __N, _HAS_ZP) \ - else if (mindiv_int4(M_in, CuCount * 4, 16) < 16) \ - WVSPLIT_INT4G_GS_CHUNKED(2, 2, __N, _HAS_ZP) \ - else \ - WVSPLIT_INT4G_GS_CHUNKED(4, 1, __N, _HAS_ZP) \ - } else if (__N >= 3 && _sYT >= 40 && K_in >= 4096 && \ - mindiv_int4(M_in, CuCount * 4, 16) < 16) \ - /* YTILE=4 is the better tuple in this range, but only when it \ - does not cost wave slots. mindiv_int4 trims WvPrGrp to even \ - out the m-rounds, and at YTILE=4 it can trim harder than at \ - YTILE=2: M=5376 gives 14 (2 of 16 wave-rows idle) against 15 at \ - YTILE=2. The K bound keeps this off shallow rows, where the k1 \ - loop is too short to amortise the extra m-tiles. The chunked \ - branch above needs no such bound: it only runs when K*N exceeds \ - LDS, so K is always past it. */ \ - WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else if (__N >= 4 && _sYT >= 480) \ - WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT >= 40) \ - WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \ - WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT < 40) \ - WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else if (__N >= 2) \ - WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else if (on_gfx_int4<11, 12>() && __N == 1 && K_in == 4096) \ - /* Tuned for gfx1151 (Qwen3.5 W4A16 decode: GDN out_proj at \ - M=2048, K=4096, N=1). AC=32 doubles per-thread global load \ - granularity; the K=4096 row has enough work to amortize the \ - wider load. W stays at 16 (vs 32 in the bf16 K=2048 branch) \ - because int4 dequant inflates VGPR pressure -- AC=32 + W=32 \ - spills. Lifts kernel ~70% -> ~84% of LPDDR5X peak post- \ - overhead. K<=2048 already at ~87% by default and untouched. \ - Verify per shape with \ - benchmarks/kernels/sweep_int4g_kernel.py. */ \ - WVSPLITK_INT4G_GS_W_AC(1, 4, 16, 32, __N, _HAS_ZP) \ - else if (on_gfx_int4<11, 12>() && __N == 1 && K_in % (32 * 32 * 8) == 0) \ - /* Deep-K decode rows that divide this tuple's K step exactly \ - (THRDS*A_CHUNK*UNRL = 8192), e.g. gemma-4-31B down_proj at \ - K=8192 and K=16384, M=5376. AC=32 issues global_load_b128 \ - and UNRL=8 makes each row's contiguous run 4 KB. \ - \ \ - The exact-division guard is load-bearing: a K that leaves a \ - large ragged tail has too few k1 iterations to hide it behind \ - and prefers the default tuple. Geometry and row stride have \ - to be tuned together; re-verify both with \ - benchmarks/kernels/sweep_int4g_kernel.py. */ \ - WVSPLITK_INT4G_GS_W_AC(2, 8, 16, 32, __N, _HAS_ZP) \ - else /* N=1: YTILE=2 beats YTILE=1 across all CuCount values */ \ - WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ +#define WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \ + { \ + if (K_in * N_in > max_lds_len) { \ + /* Chunked branch. Same wave-packing rule as the fitting branch \ + below, for the same reason: it is a property of M, not of \ + chunking. */ \ + if (_sYT < 30) \ + WVSPLIT_INT4G_GS_CHUNKED(4, 2, __N, _HAS_ZP) \ + else if (mindiv_int4(M_in, CuCount * 4, 16) < 16) \ + WVSPLIT_INT4G_GS_CHUNKED(2, 2, __N, _HAS_ZP) \ + else \ + WVSPLIT_INT4G_GS_CHUNKED(4, 1, __N, _HAS_ZP) \ + } else if (__N >= 3 && _sYT >= 40 && K_in >= 4096 && \ + mindiv_int4(M_in, CuCount * 4, 16) < 16) \ + /* YTILE=4 is the better tuple in this range, but only when it \ + does not cost wave slots. mindiv_int4 trims WvPrGrp to even \ + out the m-rounds, and at YTILE=4 it can trim harder than at \ + YTILE=2: M=5376 gives 14 (2 of 16 wave-rows idle) against 15 at \ + YTILE=2. The K bound keeps this off shallow rows, where the k1 \ + loop is too short to amortise the extra m-tiles. The chunked \ + branch above needs no such bound: it only runs when K*N exceeds \ + LDS, so K is always past it. */ \ + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if (__N >= 4 && _sYT >= 480) \ + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT >= 40) \ + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \ + WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT < 40) \ + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if (__N >= 2) \ + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if constexpr (__N == 1) { \ + /* N>=2 never gets here: the N=1 tuples are not built for it. */ \ + if (on_gfx_int4<11, 12>() && K_in == 4096) \ + /* Tuned for gfx1151 (Qwen3.5 W4A16 decode: GDN out_proj at \ + M=2048, K=4096, N=1). AC=32 doubles per-thread global load \ + granularity; the K=4096 row has enough work to amortize the \ + wider load. W stays at 16 (vs 32 in the bf16 K=2048 branch) \ + because int4 dequant inflates VGPR pressure -- AC=32 + W=32 \ + spills. Lifts kernel ~70% -> ~84% of LPDDR5X peak post- \ + overhead. K<=2048 already at ~87% by default and untouched. \ + Verify per shape with \ + benchmarks/kernels/sweep_int4g_kernel.py. */ \ + WVSPLITK_INT4G_GS_W_AC(1, 4, 16, 32, __N, _HAS_ZP) \ + else if (on_gfx_int4<11, 12>() && K_in % (32 * 32 * 8) == 0) \ + /* Deep-K decode rows that divide this tuple's K step exactly \ + (THRDS*A_CHUNK*UNRL = 8192), e.g. gemma-4-31B down_proj at \ + K=8192 and K=16384, M=5376. AC=32 issues global_load_b128 \ + and UNRL=8 makes each row's contiguous run 4 KB. \ + \ \ + The exact-division guard is load-bearing: a K that leaves a \ + large ragged tail has too few k1 iterations to hide it behind \ + and prefers the default tuple. Geometry and row stride have \ + to be tuned together; re-verify both with \ + benchmarks/kernels/sweep_int4g_kernel.py. */ \ + WVSPLITK_INT4G_GS_W_AC(2, 8, 16, 32, __N, _HAS_ZP) \ + else /* N=1: YTILE=2 beats YTILE=1 across all CuCount values */ \ + WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ + } \ } // Inner dispatch: shared by both symmetric and asymmetric paths @@ -1699,78 +1716,79 @@ __global__ void moe_wvSplitK_int4_hf_( } \ } -#define MOE_WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \ - { \ - if (K_in * __N > max_lds_len) { \ - if (_sYT < 30) \ - MOE_WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \ - else \ - MOE_WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - } else if (__N >= 4 && _sYT >= 480) \ - MOE_WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT >= 40) \ - MOE_WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \ - MOE_WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT < 40) \ - MOE_WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else if (__N >= 2) \ - MOE_WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else /* N=1 split per K, picked from a full 4-axis (Y, U, W, AC) \ - Cartesian sweep on gfx1151. Re-verified in fresh procs \ - against the local kernel; the two K ranges land on \ - different optima: \ - K >= 1024 (gemm1 K=2048): Y=4, U=4, W=16, AC=16 \ - 64.17 us/call @ 182.6 GiB/s vs Y2U4's 66.40 us; \ - wider YTILE keeps each CU's m-tile fully fed at \ - gfx1151's 20-CU width when M is comparable to \ - CuCount * 16 * Y. \ - K < 1024 (gemm2 K=768): Y=4, U=2, W=32, AC=32 \ - ~10% faster than the previous (Y8 U2 W=AC=16) heur: \ - 2 wide-load batches per k-step instead of 4 narrow \ - ones, fewer waitcnt boundaries amortize HBM page-mode \ - costs. GFX9 wave64 stays on the previous heuristic \ - via MOE_WVSPLIT_INT4G_GS_W_AC's fallback. \ - LOW-VGPR override (gfx1x_int4 only): \ - On Qwen3.5-35B-A3B (gemm2 K=512), the (W=32, AC=32) \ - instantiation compiles to 157 VGPRs/wave; combined \ - with WG=1024 threads (32 wave32) this caps occupancy \ - at 32 active wave32/CU = 50% of peak. Switching to \ - (W=16, AC=16, U=2) drops VGPRs to 113 and WG=512 \ - threads (16 wave32) so 3 WGs fit per CU = 48 active \ - wave32/CU = 75% of peak. Better latency hiding for \ - memory-bound kernels. */ \ - { \ - if (on_gfx_int4<11, 12>() && K_in == 2048) \ - /* gfx1151 K=2048 N=1 (Qwen3.5 MoE gate_up at M=1024, \ - * E=256): 4-axis sweep via \ - * benchmarks/kernels/sweep_int4g_moe_kernel.py finds \ - * (Y=4, U=4, W=32, AC=32) at 44.6 us vs the (Y=4, U=4, \ - * W=16, AC=16) default at 46.5 us -- ~1.04x. Wider per- \ - * thread loads (AC=32) + max WG threads (W=32) amortize the \ - * per-row work the K=2048 case has plenty of. Narrow guard: \ - * only swept at M~1024; other M at K=2048 may benefit too \ - * but were not measured. */ \ - MOE_WVSPLIT_INT4G_GS_W_AC(4, 32, 32, 4, __N, _HAS_ZP) \ - else if (K_in >= 1024) \ - MOE_WVSPLIT_INT4G_GS(4, 4, __N, _HAS_ZP) \ - else if (on_gfx_int4<11, 12>() && K_in == 512) \ - /* gfx1151 K=512 N=1 (Qwen3.5 MoE down-proj). 4-axis sweep \ - * via benchmarks/kernels/sweep_int4g_moe_kernel.py finds \ - * (Y=4, U=1, W=32, AC=8) at 28.1 us vs the LOW-VGPR pick \ - * (Y=4, U=2, W=16, AC=16) at 31.4 us -- ~1.12x at \ - * fuse_silu_mul=False; ~1.18x at fuse_silu_mul=True (the \ - * down+silu trace shape). Pure-work %DRAM lifts ~62% \ - * -> ~78% post per-launch floor subtraction. W=32 + AC=8 \ - * keeps WG threads at 1024 (32 wv32) but the narrow per- \ - * iter loads avoid the (W=32, AC=32) VGPR-spill penalty the \ - * (W=32, AC=32) default below describes. */ \ - MOE_WVSPLIT_INT4G_GS_W_AC(4, 32, 8, 1, __N, _HAS_ZP) \ - else if (on_gfx_int4<11, 12>()) \ - MOE_WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \ - else \ - MOE_WVSPLIT_INT4G_GS_W_AC(4, 32, 32, 2, __N, _HAS_ZP) \ - } \ +#define MOE_WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \ + { \ + if (K_in * __N > max_lds_len) { \ + if (_sYT < 30) \ + MOE_WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \ + else \ + MOE_WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + } else if (__N >= 4 && _sYT >= 480) \ + MOE_WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT >= 40) \ + MOE_WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \ + MOE_WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT < 40) \ + MOE_WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if (__N >= 2) \ + MOE_WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if constexpr (__N == 1) { \ + /* N=1 split per K, picked from a full 4-axis (Y, U, W, AC) \ + Cartesian sweep on gfx1151. Re-verified in fresh procs \ + against the local kernel; the two K ranges land on \ + different optima: \ + K >= 1024 (gemm1 K=2048): Y=4, U=4, W=16, AC=16 \ + 64.17 us/call @ 182.6 GiB/s vs Y2U4's 66.40 us; \ + wider YTILE keeps each CU's m-tile fully fed at \ + gfx1151's 20-CU width when M is comparable to \ + CuCount * 16 * Y. \ + K < 1024 (gemm2 K=768): Y=4, U=2, W=32, AC=32 \ + ~10% faster than the previous (Y8 U2 W=AC=16) heur: \ + 2 wide-load batches per k-step instead of 4 narrow \ + ones, fewer waitcnt boundaries amortize HBM page-mode \ + costs. GFX9 wave64 stays on the previous heuristic \ + via MOE_WVSPLIT_INT4G_GS_W_AC's fallback. \ + LOW-VGPR override (gfx1x_int4 only): \ + On Qwen3.5-35B-A3B (gemm2 K=512), the (W=32, AC=32) \ + instantiation compiles to 157 VGPRs/wave; combined \ + with WG=1024 threads (32 wave32) this caps occupancy \ + at 32 active wave32/CU = 50% of peak. Switching to \ + (W=16, AC=16, U=2) drops VGPRs to 113 and WG=512 \ + threads (16 wave32) so 3 WGs fit per CU = 48 active \ + wave32/CU = 75% of peak. Better latency hiding for \ + memory-bound kernels. */ \ + if (on_gfx_int4<11, 12>() && !on_1024_vgpr_int4() && K_in == 2048) \ + /* gfx1151 K=2048 N=1 (Qwen3.5 MoE gate_up at M=1024, \ + * E=256): 4-axis sweep via \ + * benchmarks/kernels/sweep_int4g_moe_kernel.py finds \ + * (Y=4, U=4, W=32, AC=32) at 44.6 us vs the (Y=4, U=4, \ + * W=16, AC=16) default at 46.5 us -- ~1.04x. Wider per- \ + * thread loads (AC=32) + max WG threads (W=32) amortize the \ + * per-row work the K=2048 case has plenty of. Narrow guard: \ + * only swept at M~1024; other M at K=2048 may benefit too \ + * but were not measured. Not on 1024-VGPR archs: its 128 \ + * VGPRs a wave there spill; they take (4, 4, 16, 16). */ \ + MOE_WVSPLIT_INT4G_GS_W_AC(4, 32, 32, 4, __N, _HAS_ZP) \ + else if (K_in >= 1024) \ + MOE_WVSPLIT_INT4G_GS(4, 4, __N, _HAS_ZP) \ + else if (on_gfx_int4<11, 12>() && K_in == 512) \ + /* gfx1151 K=512 N=1 (Qwen3.5 MoE down-proj). 4-axis sweep \ + * via benchmarks/kernels/sweep_int4g_moe_kernel.py finds \ + * (Y=4, U=1, W=32, AC=8) at 28.1 us vs the LOW-VGPR pick \ + * (Y=4, U=2, W=16, AC=16) at 31.4 us -- ~1.12x at \ + * fuse_silu_mul=False; ~1.18x at fuse_silu_mul=True (the \ + * down+silu trace shape). Pure-work %DRAM lifts ~62% \ + * -> ~78% post per-launch floor subtraction. W=32 + AC=8 \ + * keeps WG threads at 1024 (32 wv32) but the narrow per- \ + * iter loads avoid the (W=32, AC=32) VGPR-spill penalty the \ + * (W=32, AC=32) default below describes. */ \ + MOE_WVSPLIT_INT4G_GS_W_AC(4, 32, 8, 1, __N, _HAS_ZP) \ + else if (on_gfx_int4<11, 12>()) \ + MOE_WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \ + else \ + MOE_WVSPLIT_INT4G_GS_W_AC(4, 32, 32, 2, __N, _HAS_ZP) \ + } \ } #define MOE_WVSPLIT_INT4G_DISPATCH(_HAS_ZP) \ From f993e6b90b0bf2ebf9f1ce1c742f9f240c25ac7d Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Mon, 26 Oct 2026 12:16:50 -0700 Subject: [PATCH 3/4] [ROCm][RDNA] skinny GEMMs: wvSplitK launch rows of their own for gfx1151 and gfx1150 The wave32 wvSplitK dispatch was one chain for every gfx11 arch, tuned shape by shape on gfx1151. A sweep of every non-spilling launch configuration (YTILE 1-4, UNRL 1/2/4/8, A_CHUNK 8/16/32, WvPrGrp 16/32, N = 1..5, bf16 and fp16) finds one tile per arch within ~1% of the best for nearly every shape: - gfx1151: (YTILE 1, UNRL 4, A_CHUNK 16); (2, 4, 8) for K < 1024. - gfx1150, gfx1152, gfx1153: (2, 1, 16); (1, 2, 16) when M <= 32 * CUs. - A_CHUNK 8 when K % 16 != 0: wider chunks give wrong results there. Other archs take the generic rows. The sweep op covers every knob the dispatcher uses, and sweep_bf16_kernel.py pads the weights as they are padded at load. Interleaved A/B against the previous rows, 43 shapes (13 not used for tuning, odd M, K 768/1000/2880), N = 1..5, bf16/fp16 and bias: gfx1151 +0.2% to +1.3%, gfx1150 +0.6% to +1.1%; against upstream's rows +1.8% to +4.2% and +0.8% to +3.9%. 99.3-99.8% of the per-shape best. test_rocm_skinny_gemms.py: 190 passed on gfx1150 and gfx1151, with new K % 16 != 0 and N = 5 cases. AI assistance was used for this patch. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- benchmarks/kernels/sweep_bf16_kernel.py | 27 +- csrc/rocm/skinny_gemms.cu | 234 ++++++++---------- .../quantization/test_rocm_skinny_gemms.py | 5 + 3 files changed, 117 insertions(+), 149 deletions(-) diff --git a/benchmarks/kernels/sweep_bf16_kernel.py b/benchmarks/kernels/sweep_bf16_kernel.py index 7c2c861cc85a..3de7651eb6a9 100644 --- a/benchmarks/kernels/sweep_bf16_kernel.py +++ b/benchmarks/kernels/sweep_bf16_kernel.py @@ -7,13 +7,14 @@ (unquantized) wvSplitK path used by every Linear that hasn't been replaced by an int4/int8/AWQ wrapper. -Limitation: the in-tree ``wvSplitK_sweep`` op (csrc/rocm/skinny_gemms.cu) -currently exposes only ``(ytile, unrl)`` as runtime args; ``WvPrGrp`` and -``A_CHUNK`` are macro-baked at 16 and 8 respectively. The script therefore -sweeps only YTILE x UNRL. For 4-axis sweeps (the (W, AC) corner of the -search space) use the out-of-tree HIP harness in -``rdna35-asm-expert/microbench/wv_splitk/`` until ``wvSplitK_sweep`` is -extended to match ``wvSplitK_int4g_sweep`` (which already takes all 4 axes). +The in-tree ``wvSplitK_sweep`` op (csrc/rocm/skinny_gemms.cu, built with +``VLLM_SKINNY_GEMM_SWEEP_BF16=1``) takes every launch knob of the dispatcher: +YTILE x UNRL x A_CHUNK x WvPrGrp, for N = 1..5. Combinations that spill on +the current arch are worth excluding from any conclusion (their speed depends +on the process's scratch history), and A_CHUNK > 8 needs K % A_CHUNK == 0. +Weights are padded as vLLM pads them at load on gfx11 (a row stride that is a +multiple of 2048 B gets one more 128 B line); unpadded, a weight that lands on +an aliasing address can run 2x slower. For each (M, K, N) cell the script also calls the production ``wvSplitK`` op once to record the dispatcher's chosen wall-clock — useful both as a @@ -69,9 +70,9 @@ (1024, 4096, "K=4096 small-M sanity"), ] -YTILES = [1, 2] -UNRLS = [1, 2, 4] -ACHUNKS = [8, 16] +YTILES = [1, 2, 3, 4] +UNRLS = [1, 2, 4, 8] +ACHUNKS = [8, 16, 32] WVPRGRPS = [16, 32] # Working-set guards used by the dispatcher's variant selection @@ -141,7 +142,9 @@ def run_sweep(shapes, batch_sizes, warmup, rep, dtype, csv_path): # three variants under the hood. Just noting it for the user. pass - A = torch.randn(M, K, dtype=dtype, device="cuda") * 0.01 + pad = 64 if (K * 2) % 2048 == 0 else 0 + A = torch.empty(M, K + pad, dtype=dtype, device="cuda")[:, :K] + A.copy_(torch.randn(M, K, dtype=dtype, device="cuda") * 0.01) B = torch.randn(N, K, dtype=dtype, device="cuda") * 0.01 weight_bytes = M * K * A.element_size() @@ -157,7 +160,7 @@ def run_sweep(shapes, batch_sizes, warmup, rep, dtype, csv_path): for ytile, unrl, ac, w in itertools.product( YTILES, UNRLS, ACHUNKS, WVPRGRPS ): - if M % ytile != 0: + if K % ac != 0: skipped += 1 continue try: diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 6598ff9d8fdf..4a827de2107d 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -67,6 +67,27 @@ static bool on_gfx() { return result; } +bool on_gfx1151() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx1151") != std::string::npos; + }(); + return result; +} + +// gfx1150, gfx1152 and gfx1153: 1024 VGPRs a lane and a 128-bit memory bus. +bool on_gfx1150_class() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx1150") != std::string::npos || + device_arch.find("gfx1152") != std::string::npos || + device_arch.find("gfx1153") != std::string::npos; + }(); + return result; +} + #if defined(NDEBUG) #undef NDEBUG #include @@ -325,8 +346,7 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b, } #elif defined(__HIP__GFX1X__) // gfx1x: v_dot2_f32_f16 and v_dot2_f32_bf16 (VOP3-P, gfx11 and gfx12). The - // bf16 dot keeps the weights packed; widening them to f32 first doubles - // their registers and makes the large tiles spill. + // bf16 dot keeps the weights packed instead of widening them to f32. #define DOT2C(V0, V2, V3) \ if constexpr (std::is_same_v) { \ asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ @@ -1395,15 +1415,13 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const int max_lds_len = get_lds_size() / 2; -// Like WVSPLITK_CFG but lets the caller pick A_CHUNK. The original macro -// hard-coded A_CHUNK=8; this is a compatibility-preserving expansion that -// lets a single dispatch branch select wider/narrower per-thread global -// loads. Used today by the gfx11 K=2048 N=1 fast path (Qwen3.5 decode): -// AC=16 with WvPrGrp=32 and UNRL=8 beats the (AC=8, W=16, UN=4) default -// by ~30% on M ∈ {256, 1024} and ~4% on M=248320. Reproduce the YT/UN -// half of the sweep via benchmarks/kernels/sweep_bf16_kernel.py; the -// A_CHUNK and WvPrGrp axes are not yet exposed by wvSplitK_sweep (TODO: -// extend the op to 4 axes like wvSplitK_int4g_sweep already does). +// WVSPLITK_CFG_AC arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N, A_CHUNK). +// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9) +// WVPRGRP = waves per group +// YTILE = output rows per thread tile +// UNRL = K-loop unroll factor +// N = batch size (passed through from the switch in wvSplitK) +// A_CHUNK = elements per thread per load; above 8 it needs K % A_CHUNK == 0 #define WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N, _AC) \ { \ dim3 block(_THRDS, _WVPRGRP); \ @@ -1425,113 +1443,50 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, CuCount); \ } -// Backwards-compatible wrapper: existing call sites get A_CHUNK=8. #define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N) \ WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N, 8) -#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \ - { \ - bool fit_lds = (Kbp_in * N_in <= max_lds_len); \ - constexpr bool any11 = std::ranges::any_of(kRocmArchs, is_gfx<11>); \ - constexpr bool all11 = std::ranges::all_of(kRocmArchs, is_gfx<11>); \ - if constexpr (any11) { \ - if (all11 || on_gfx<11>()) { \ - if (_sYT <= 1) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ - else if (K_in < 1024) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 4, __N) \ - else if ((K_in % 1024 == 512) && (_sYT >= 40 || K_in >= 4096)) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \ - else if ((K_in == 2048) && (__N == 1)) \ - /* Tuned for gfx1151 (Qwen3.5 decode shapes M ∈ {256, 1024, \ - 248320}, K=2048, N=1): beats the (AC=8, W=16, UN=4) baseline of \ - the K_in<=2048 branch below by 1.31-1.37x on small/mid M and 3.9% \ - on the lm_head-sized M. Compiles to 145 VGPRs / occupancy 9 (vs \ - 46 / 8 for the default); zero spills. VGPR-bound but DRAM- \ - saturated past 92% of LPDDR5X peak after subtracting the per- \ - launch dispatch floor. Verify per shape with \ - benchmarks/kernels/sweep_bf16_kernel.py. */ \ - WVSPLITK_CFG_AC(32, 32, 1, 8, __N, 16) \ - /* gfx1151 AC=32 fast paths. Each cell beats AC=16 by >=2% with \ - z>1.96 in a 10-rep do_bench A/B (stderr 0.1-1.0us per cell, mean \ - delta 1.5-3 us per cell). Other K%2048==0 cells stay on the AC=16 \ - fallbacks below where AC=32 was a tie or lost (notably 4096x4096 \ - N=4 was -2.7%, do not extrapolate to untested cells). Re-verify \ - per shape with benchmarks/kernels/sweep_bf16_kernel.py (extend the \ - ACHUNKS list to include 32 and rebuild with \ - VLLM_SKINNY_GEMM_SWEEP_BF16=1). */ \ - else if ((K_in == 2048) && (__N == 2 || __N == 3)) \ - /* M=2560 K=2048 N=2: 1.057x (z=21.5); N=3: 1.049x (z=12.8) */ \ - WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 32) \ - else if ((K_in == 4096) && (__N == 1)) \ - /* M=2560 K=4096 N=1: 1.041x (z=3.7); UR=4 not 2 for N=1 */ \ - WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 4, __N, 32) \ - else if ((K_in == 4096) && (__N == 2) && (M_in < 4096)) \ - /* M<4096 K=4096 N=2: 1.057x (z=13.1), W=16 wins at this M */ \ - WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 32) \ - else if ((K_in == 4096) && (__N == 2) && (M_in >= 4096)) \ - /* M>=4096 K=4096 N=2: 1.028x (z=6.2), W=32 wins at larger M */ \ - WVSPLITK_CFG_AC(_THRDS, 32, 1, 2, __N, 32) \ - else if ((K_in == 4096) && (__N == 3)) \ - /* M=2560 K=4096 N=3: 1.031x (z=4.9) */ \ - WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 32) \ - else if ((K_in == 8192) && (__N == 2)) \ - /* M=2560 K=8192 N=2: 1.040x (z=9.3), W=32 wins at this K */ \ - WVSPLITK_CFG_AC(_THRDS, 32, 1, 2, __N, 32) \ - else if ((K_in % 2048 == 0) && (__N == 2)) \ - /* gfx1151 K%2048==0, N=2 only: YT=2 + W=32 + AC=16 + UR=4. \ - sweep_bf16_kernel.py 4-axis sweep showed this is the best \ - N=2 config across K in {2048, 4096, 8192} and 4096x4096, \ - 1.06x (K=8192) to 1.60x (K=2048) over the prior AC=8 default. */ \ - WVSPLITK_CFG_AC(_THRDS, 32, 2, 4, __N, 16) \ - else if ((K_in % 2048 == 0) && (__N != 2)) \ - /* gfx1151 K%2048==0, N in {1, 3, 4} (K=2048 N=1 handled above): \ - YT=1 + W=16 + AC=16 + UR=4. N=3/4 want YT=1 not YT=2 (LDS/VGPR \ - pressure from W=32 hurts them); same config also wins for N=1. \ - sweep showed 1.11x-1.19x (N=1), 1.23x-1.98x (N=3), \ - 1.59x-2.73x (N=4) over the prior AC=8 defaults. */ \ - WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 4, __N, 16) \ - else if (K_in <= 2048) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ - else if (__N >= 2 && !fit_lds) { \ - if (K_in % 1024 == 0 && Kbp_in < max_lds_len / 2) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 4, __N) \ - else \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ - } else if (__N == 1) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 2, __N) \ - else \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 1, __N) \ - } \ - } \ - if constexpr (!all11) { \ - if (!any11 || !on_gfx<11>()) { \ - if (_sYT <= 1) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ - else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \ - else if (_sYT <= 4 * 3) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \ - else if (__N == 4) \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \ - else \ - WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \ - } \ - } \ +#define WVSPLIT_TILE_GENERIC(_THRDS, _WVPRGRP, _sYT, __N) \ + { \ + bool fit_lds = (Kbp_in * N_in <= max_lds_len); \ + if (_sYT <= 1) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \ + else if (_sYT <= 4 * 3) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \ + else if (__N == 4) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \ + else \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \ + } + +// gfx1151 and the gfx1150 class take the tile that a sweep of every +// non-spilling (YTILE, UNRL, A_CHUNK, WVPRGRP) found within ~1% of the best +// one for nearly every (M, K, N) on that arch: 99.3-99.5% of the per-shape +// best on average (benchmarks/kernels/sweep_bf16_kernel.py). The other archs +// keep the generic rows. +#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \ + { \ + if constexpr (_THRDS == 32) { \ + if (on_gfx1151()) { \ + if (K_in < 1024 || K_in % 16 != 0) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 4, __N) \ + else \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 4, __N, 16) \ + } else if (on_gfx1150_class()) { \ + if (K_in % 16 != 0) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + else if (_sYT <= 8) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 16) \ + else \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 2, 1, __N, 16) \ + } else \ + WVSPLIT_TILE_GENERIC(_THRDS, _WVPRGRP, _sYT, __N) \ + } else \ + WVSPLIT_TILE_GENERIC(_THRDS, _WVPRGRP, _sYT, __N) \ } -// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N). -// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9) -// WVPRGRP= waves per group (always 16) -// YTILE = output rows per thread tile -// UNRL = K-loop unroll factor -// N = batch size (passed through from the switch in wvSplitK) -// All gfx11xx boards, gfx1151 included, route through WVSPLIT_TILE_CFG. -// Giving gfx1151 a chain of its own would restrict it to AC=8 configs and drop -// the nine AC=16/AC=32 branches WVSPLIT_TILE_CFG carries for K in {2048, 4096, -// 8192} and K%2048==0, which costs 1.0-1.4% decode on Qwen3-30B-A3B (K=2048), -// Qwen3.x-35B and Qwen3-Omni-30B. Dense models are unaffected: Gemma-3-4B's -// K=2560 hits K%1024==512, a rule both paths share. #define WVSPLIT_TILE(_sYT, __N) \ { \ constexpr bool any32 = std::ranges::any_of(kRocmArchs, is_gfx<11, 12>); \ @@ -1824,15 +1779,10 @@ torch::Tensor wvSplitK_fused_silu_gate_mul( // Sweep function disabled by default to reduce compile time. // Build with -DVLLM_SKINNY_GEMM_SWEEP_BF16 (or the umbrella -// -DVLLM_SKINNY_GEMM_SWEEP) to enable. Compared to the dispatcher's -// (YTILE, UNRL) heuristic this lets the caller also pick A_CHUNK and -// WvPrGrp at runtime, which is necessary to verify whether the (W=32, -// AC=16) trick used by the K=2048 N=1 fast path also recovers other -// K%2048 == 0 shapes that currently route through (W=16, AC=8). -// The YTILE grid is restricted to {1, 2} -- the production dispatcher -// never picks YTILE > 2 for the slow K%2048 shapes (YT=1 for N=1, -// YT=2 for N>=2 + !fit_lds) -- which keeps the template-instantiation -// count to 24 combos x 4 N x 3 kernel variants = 288. +// -DVLLM_SKINNY_GEMM_SWEEP) to enable. The caller picks every launch knob +// the dispatcher can: YTILE {1..4}, UNRL {1, 2, 4, 8}, A_CHUNK {8, 16, 32}, +// WvPrGrp {16, 32}, for N = 1..5, and the kernel variant is chosen as in +// WVSPLITK_CFG_AC. 96 combos x 5 N x 3 kernel variants per dtype. #ifdef VLLM_SKINNY_GEMM_SWEEP_BF16 torch::Tensor wvSplitK_sweep(const at::Tensor& in_a, const at::Tensor& in_b, const std::optional& in_bias, @@ -1857,7 +1807,6 @@ torch::Tensor wvSplitK_sweep(const at::Tensor& in_a, const at::Tensor& in_b, TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0"); TORCH_CHECK(in_a.dtype() == torch::kFloat16 || in_a.dtype() == torch::kBFloat16); - TORCH_CHECK(M_in % ytile == 0, "M must be divisible by ytile=", ytile); auto out_c = torch::empty( {N_in, M_in}, @@ -1900,6 +1849,8 @@ torch::Tensor wvSplitK_sweep(const at::Tensor& in_a, const at::Tensor& in_b, WVSPLITK_SWEEP_LAUNCH(_THRDS, _YTILE, _UNRL, 3, _AC, _WVPRGRP) break; \ case 4: \ WVSPLITK_SWEEP_LAUNCH(_THRDS, _YTILE, _UNRL, 4, _AC, _WVPRGRP) break; \ + case 5: \ + WVSPLITK_SWEEP_LAUNCH(_THRDS, _YTILE, _UNRL, 5, _AC, _WVPRGRP) break; \ default: \ TORCH_CHECK(false, "Unsupported N=", N_in); \ } @@ -1914,13 +1865,16 @@ torch::Tensor wvSplitK_sweep(const at::Tensor& in_a, const at::Tensor& in_b, "; allowed: 16, 32"); \ } - #define WVSPLITK_SWEEP_AC(_THRDS, _YTILE, _UNRL) \ - if (achunk == 8) { \ - WVSPLITK_SWEEP_WVPRGRP(_THRDS, _YTILE, _UNRL, 8) \ - } else if (achunk == 16) { \ - WVSPLITK_SWEEP_WVPRGRP(_THRDS, _YTILE, _UNRL, 16) \ - } else { \ - TORCH_CHECK(false, "Unsupported achunk=", achunk, "; allowed: 8, 16"); \ + #define WVSPLITK_SWEEP_AC(_THRDS, _YTILE, _UNRL) \ + if (achunk == 8) { \ + WVSPLITK_SWEEP_WVPRGRP(_THRDS, _YTILE, _UNRL, 8) \ + } else if (achunk == 16) { \ + WVSPLITK_SWEEP_WVPRGRP(_THRDS, _YTILE, _UNRL, 16) \ + } else if (achunk == 32) { \ + WVSPLITK_SWEEP_WVPRGRP(_THRDS, _YTILE, _UNRL, 32) \ + } else { \ + TORCH_CHECK(false, "Unsupported achunk=", achunk, \ + "; allowed: 8, 16, 32"); \ } #define WVSPLITK_SWEEP_UNRL(_THRDS, _YTILE) \ @@ -1930,17 +1884,23 @@ torch::Tensor wvSplitK_sweep(const at::Tensor& in_a, const at::Tensor& in_b, WVSPLITK_SWEEP_AC(_THRDS, _YTILE, 2) \ } else if (unrl == 4) { \ WVSPLITK_SWEEP_AC(_THRDS, _YTILE, 4) \ + } else if (unrl == 8) { \ + WVSPLITK_SWEEP_AC(_THRDS, _YTILE, 8) \ } else { \ TORCH_CHECK(false, "Unsupported unrl=", unrl); \ } - #define WVSPLITK_SWEEP_YTILE(_THRDS) \ - if (ytile == 1) { \ - WVSPLITK_SWEEP_UNRL(_THRDS, 1) \ - } else if (ytile == 2) { \ - WVSPLITK_SWEEP_UNRL(_THRDS, 2) \ - } else { \ - TORCH_CHECK(false, "Unsupported ytile=", ytile, "; allowed: 1, 2"); \ + #define WVSPLITK_SWEEP_YTILE(_THRDS) \ + if (ytile == 1) { \ + WVSPLITK_SWEEP_UNRL(_THRDS, 1) \ + } else if (ytile == 2) { \ + WVSPLITK_SWEEP_UNRL(_THRDS, 2) \ + } else if (ytile == 3) { \ + WVSPLITK_SWEEP_UNRL(_THRDS, 3) \ + } else if (ytile == 4) { \ + WVSPLITK_SWEEP_UNRL(_THRDS, 4) \ + } else { \ + TORCH_CHECK(false, "Unsupported ytile=", ytile, "; allowed: 1-4"); \ } AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK_sweep", [&] { diff --git a/tests/kernels/quantization/test_rocm_skinny_gemms.py b/tests/kernels/quantization/test_rocm_skinny_gemms.py index 51b08461d3fe..a434c1d46202 100644 --- a/tests/kernels/quantization/test_rocm_skinny_gemms.py +++ b/tests/kernels/quantization/test_rocm_skinny_gemms.py @@ -101,6 +101,11 @@ (1, 64, 8), (2, 128, 8), (4, 256, 8), + # K % 16 != 0: the gfx1150/gfx1151 rows fall back to A_CHUNK = 8 + (1, 4096 + 8, 4096 + 1), + (5, 1000, 2048), + # Speculative-decode batch + (5, 2048, 2560), ] # N is bucketed up to its next ^2 (16/32/64/128) and the remainder is masked From 79aa663ba6679d3d1c7093023204bdefa7850238 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Mon, 26 Oct 2026 12:43:44 -0700 Subject: [PATCH 4/4] [ROCm][RDNA] skinny GEMMs: drop the wvSplitK VGPR-budget fallback The fallback that reads each activation chunk from LDS when a tile's operands exceed the wave's VGPRs only served the large tiles of the previous rows. No launched tile needs it now (the tuned rows, the generic rows and the fused epilogues all fit), so the kernels go back to holding the activation chunks, as before. 0 VGPR spills on gfx1150 and gfx1151; the launched kernels compile to the same instructions as the same tiles upstream. AI assistance was used for this patch. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- csrc/rocm/skinny_gemms.cu | 252 ++++++++++---------------------------- 1 file changed, 66 insertions(+), 186 deletions(-) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 4a827de2107d..d2aef69828df 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -355,25 +355,6 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b, } #endif -#if defined(__HIP__GFX1X__) - // VGPRs a lane holds: 1024 on gfx1102, gfx1103 and gfx1150..gfx1153, else - // 1536. - #if defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || \ - defined(__gfx1152__) || defined(__gfx1153__) -inline constexpr int kLaneVgprs = 1024; - #else -inline constexpr int kLaneVgprs = 1536; - #endif -// The VGPRs a wave of a THREADS-thread workgroup gets: its wave32 waves share -// a WGP's four SIMDs. -__device__ constexpr int wave_vgprs(int threads) { - const int vgprs = kLaneVgprs / ((threads / 32 + 3) / 4); - return vgprs < 256 ? vgprs : 256; -} -// Registers beside the operand tiles: sums, addresses, loop state. -inline constexpr int kOperandSlackVgprs = 64; -#endif - // To avoid LLVM silently upcasting to double __device__ inline unsigned int min__(uint32_t a, uint32_t b) { return min(a, b); @@ -522,9 +503,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) // are being worked on by each wave. //---------------------------------------------------- float sum[N][YTILE] = {}; - [[maybe_unused]] scalar8 sum4[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; bigType bigB[YTILE][UNRL]; // Fetch the weight matrix from memory! #pragma unroll @@ -536,59 +518,29 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } - #if defined(__HIP__GFX1X__) - // Hold the activation chunks too while they fit; else read each - // where it is used, so only the weights' loads stay in flight. - constexpr bool hold_a = - (N + YTILE) * UNRL * (A_CHUNK / 2) + kOperandSlackVgprs <= - wave_vgprs(WvPrGrp * THRDS); - #else - constexpr bool hold_a = true; - #endif - if constexpr (hold_a) { - bigType bigA[N][UNRL] = {}; - // Fetch activation matrix from either just LDS or from both LDS / - // memory + // Fetch activation matrix from either just LDS or from both LDS / memory #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - } + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); } + } - // Do the matrix multiplication in interleaved manner - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (uint32_t n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) - } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, - 0); - } - } - } - } else { - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k_ = k1 + k2 * THRDS * A_CHUNK + threadIdx.x * A_CHUNK; - #pragma unroll - for (uint32_t n = 0; n < N; n++) { - bigType a = {}; - if (k_ < K) a = *((const bigType*)(&(s[k_ + Kap * n]))); - #pragma unroll - for (int y = 0; y < YTILE; y++) - #pragma unroll + // Do the matrix multiplication in interleaved manner + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], a.f[b], bigB[y][k2].f[b]) + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); } } } @@ -781,9 +733,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) while (m < M) { float sum[N][YTILE] = {}; - [[maybe_unused]] scalar8 sum4[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; bigType bigB[YTILE][UNRL]; // Fetch the weight matrix from memory! #pragma unroll @@ -795,67 +748,32 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } - #if defined(__HIP__GFX1X__) - // Hold the activation chunks too while they fit; else read each - // where it is used, so only the weights' loads stay in flight. - constexpr bool hold_a = - (N + YTILE) * UNRL * (A_CHUNK / 2) + kOperandSlackVgprs <= - wave_vgprs(WvPrGrp * THRDS); - #else - constexpr bool hold_a = true; - #endif - if constexpr (hold_a) { - bigType bigA[N][UNRL] = {}; - // Fetch activation matrix from either just LDS or from both LDS / - // memory + // Fetch activation matrix from either just LDS or from both LDS / memory #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - if (k_ + Kap * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - else - bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); - } + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); } + } - // Do the matrix multiplication in interleaved manner - for (uint32_t n = 0; n < N; n++) { - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) - } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, - 0); - } - } - } - } else { - #pragma unroll + // Do the matrix multiplication in interleaved manner + for (uint32_t n = 0; n < N; n++) { for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k_ = k1 + k2 * THRDS * A_CHUNK + threadIdx.x * A_CHUNK; - #pragma unroll - for (uint32_t n = 0; n < N; n++) { - bigType a = {}; - if (k_ < K) { - if (k_ + Kap * n < max_lds_len) - a = *((const bigType*)(&(s[k_ + Kap * n]))); - else - a = *((const bigType*)(&(A[k_ + Kap * n]))); - } - #pragma unroll - for (int y = 0; y < YTILE; y++) - #pragma unroll + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], a.f[b], bigB[y][k2].f[b]) + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); } } } @@ -1112,7 +1030,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) // are being worked on by each wave. //---------------------------------------------------- float sum[N][YTILE] = {}; - [[maybe_unused]] scalar8 sum4[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; //---------------------------------------------------- // Fetch weight matrix B in interleaved K-split! @@ -1132,6 +1050,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) // TODO: Logic below will only work when K is multiple of 8 //---------------------------------------------------- for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; bigType bigB[YTILE][UNRL]; #ifdef PCML @@ -1168,76 +1087,37 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); } - #if defined(__HIP__GFX1X__) - // Hold the activation chunks too while they fit; else read each - // where it is used, so only the weights' loads stay in flight. - constexpr bool hold_a = - (N + YTILE) * UNRL * (A_CHUNK / 2) + kOperandSlackVgprs <= - wave_vgprs(WvPrGrp * THRDS); - #else - constexpr bool hold_a = true; - #endif - if constexpr (hold_a) { - bigType bigA[N][UNRL] = {}; - // Fetch activation matrix from either just LDS or from both LDS / - // memory + // Fetch activation matrix from either just LDS or from both LDS / memory #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { #ifdef PCML - bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); + bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); #else - if (k_ + Kap * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - else - bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); #endif - } } + } - // Do the matrix multiplication in interleaved manner - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (uint32_t n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) - } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, - 0); - } - } - } - } else { - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k_ = k1 + k2 * THRDS * A_CHUNK + threadIdx.x * A_CHUNK; - #pragma unroll - for (uint32_t n = 0; n < N; n++) { - bigType a = {}; - if (k_ < K) { - #ifdef PCML - a = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); - #else - if (k_ + Kap * n < max_lds_len) - a = *((const bigType*)(&(s[k_ + Kap * n]))); - else - a = *((const bigType*)(&(A[k_ + Kap * n]))); - #endif - } - #pragma unroll - for (int y = 0; y < YTILE; y++) + // Do the matrix multiplication in interleaved manner #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], a.f[b], bigB[y][k2].f[b]) + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); } } }