A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 19:00 UTC
Convert wvSplitKQ to 16x16 MFMA in prep for mi4xx. (#34100)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>
Hashem Hashemi committed
Feb 24, 2026 at 15:35 UTC
a0e50a4260d20d021d6caa137a078ae2d16a8f93
1 file changed
+14
-46
csrc/rocm/skinny_gemms.cu
+14
-46
@@ -1902,7 +1902,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1902
float sB = *s_B;
1903
1904
while (m < M) {
1905
- floatx16 sum[N][YTILE] = {};
1905
+ scalar8 sum[N][YTILE] = {};
1906
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
1907
bigType bigA[N][UNRL] = {};
1908
bigType bigB[YTILE][UNRL];
@@ -1936,7 +1936,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1936
for (uint32_t n = 0; n < N; n++) {
1937
for (int i = 0; i < A_CHUNK; i += 8) {
1938
for (int y = 0; y < YTILE; ++y) {
1939
- sum[n][y] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
1939
+ sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
1940
bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
1941
0);
1942
}
@@ -1949,31 +1949,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1949
for (int n = 0; n < N; n++) {
1950
for (int y = 0; y < YTILE; y++) {
1951
float accm0 = sum[n][y][0];
1952
- float accm16 = sum[n][y][8];
1952
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
1953
1); // row_shl1
1955
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
1954
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
1955
1); // row_shl2
1958
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
1956
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
1957
1); // row_shl3
1961
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
1962
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
1963
- 1); // row_shl8
1964
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
1965
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
1966
- 1); // row_shl9
1967
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
1968
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
1969
- 1); // row_shl10
1970
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
1971
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
1972
- 1); // row_shl11
1973
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
1974
- accm0 += __shfl(accm0, 36);
1975
- accm16 += __shfl(accm16, 52);
1976
- sum[n][y][0] = accm0 + __shfl(accm16, 16);
1958
+ accm0 += __shfl_down(accm0, 20);
1959
+ accm0 += __shfl_down(accm0, 40);
1960
+ sum[n][y][0] = accm0;
1961
}
1962
}
1963
@@ -2064,7 +2048,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2048
float sB = *s_B;
2049
2050
while (m < M) {
2067
- floatx16 sum[N][YTILE] = {};
2051
+ scalar8 sum[N][YTILE] = {};
2052
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
2053
bigType bigA[N][UNRL] = {};
2054
bigType bigB[YTILE][UNRL];
@@ -2100,7 +2084,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2084
for (uint32_t n = 0; n < N; n++) {
2085
for (int i = 0; i < A_CHUNK; i += 8) {
2086
for (int y = 0; y < YTILE; ++y) {
2103
- sum[n][y] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
2087
+ sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
2088
bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
2089
0);
2090
}
@@ -2113,31 +2097,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2097
for (int n = 0; n < N; n++) {
2098
for (int y = 0; y < YTILE; y++) {
2099
float accm0 = sum[n][y][0];
2116
- float accm16 = sum[n][y][8];
2100
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
2101
1); // row_shl1
2119
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
2102
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
2103
1); // row_shl2
2122
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
2104
accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
2105
1); // row_shl3
2125
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
2126
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
2127
- 1); // row_shl8
2128
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
2129
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
2130
- 1); // row_shl9
2131
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
2132
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
2133
- 1); // row_shl10
2134
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
2135
- accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
2136
- 1); // row_shl11
2137
- accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
2138
- accm0 += __shfl(accm0, 36);
2139
- accm16 += __shfl(accm16, 52);
2140
- sum[n][y][0] = accm0 + __shfl(accm16, 16);
2106
+ accm0 += __shfl_down(accm0, 20);
2107
+ accm0 += __shfl_down(accm0, 40);
2108
+ sum[n][y][0] = accm0;
2109
}
2110
}
2111
@@ -2242,16 +2210,16 @@ void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
2210
: nullptr;
2211
switch (N_in) {
2212
case 1:
2245
- WVSPLITKQ(12, 2, 2, 2, 2, 1)
2213
+ WVSPLITKQ(16, 2, 2, 2, 2, 1)
2214
break;
2215
case 2:
2248
- WVSPLITKQ(12, 2, 2, 2, 2, 2)
2216
+ WVSPLITKQ(16, 2, 2, 2, 2, 2)
2217
break;
2218
case 3:
2251
- WVSPLITKQ(8, 2, 2, 1, 1, 3)
2219
+ WVSPLITKQ(16, 2, 2, 2, 2, 3)
2220
break;
2221
case 4:
2254
- WVSPLITKQ(4, 2, 2, 1, 1, 4)
2222
+ WVSPLITKQ(16, 2, 2, 2, 2, 4)
2223
break;
2224
default:
2225
throw std::runtime_error(