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 8, 19:00 UTC
Adds padding and perf improvements to wvSplitK_fp8 (#33527)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>
Hashem Hashemi committed
Feb 5, 2026 at 14:16 UTC
d5c4800112c12bbcd4955858ef1b415c16ae16e7
3 files changed
+169
-229
csrc/rocm/skinny_gemms.cu
+126
-174
@@ -1899,8 +1899,9 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1899
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1900
int A_CHUNK, int UNRL, int N>
1901
__global__ void __launch_bounds__(WvPrGrp* THRDS)
1902
- wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const int Bx,
1903
- const int By, const fp8_t* B, const fp8_t* __restrict__ A,
1902
+ wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, const int M,
1903
+ const int Bx, const int By, const fp8_t* B,
1904
+ const fp8_t* __restrict__ A,
1905
const scalar_t* __restrict__ BIAS, scalar_t* C,
1906
const float* __restrict__ s_A,
1907
const float* __restrict__ s_B, const int _WvPrGrp,
@@ -1924,9 +1925,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1925
__shared__ fp8_t s[max_lds_len];
1926
1927
for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1927
- k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
1928
+ k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
1929
+ #if defined(__gfx950__)
1930
+ __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
1931
+ #else
1932
*((bigType*)(&s[k])) = *((bigType*)(&A[k]));
1933
+ #endif
1934
}
1935
+ asm volatile("s_waitcnt vmcnt(0)");
1936
__syncthreads();
1937
1938
if (threadIdx.y >= _WvPrGrp) return;
@@ -1934,37 +1940,24 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1940
uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
1941
1942
using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
1937
- floatx16 sum[N][YTILE];
1943
float sA = *s_A;
1944
float sB = *s_B;
1945
1946
while (m < M) {
1942
- for (int i = 0; i < YTILE; i++)
1943
- for (int n = 0; n < N; n++) sum[n][i] = {0.f};
1944
-
1945
- bigType bigA[N][UNRL];
1946
- bigType bigB[YTILE][UNRL];
1947
-
1947
+ floatx16 sum[N][YTILE] = {};
1948
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
1949
- #pragma unroll
1950
- for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1951
- #pragma unroll
1952
- for (uint32_t n = 0; n < N; ++n) bigA[n][k2].h8 = {0.f};
1953
- #pragma unroll
1954
- for (uint32_t y = 0; y < YTILE; ++y) bigB[y][k2].h8 = {0.f};
1955
- }
1949
+ bigType bigA[N][UNRL] = {};
1950
+ bigType bigB[YTILE][UNRL];
1951
1952
// Fetch the weight matrix from memory!
1953
#pragma unroll
1954
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1955
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1956
uint32_t k_ = k + threadIdx.x * A_CHUNK;
1962
- if (k_ >= K) break;
1963
-
1964
- const fp8_t* B_ = &B[(m + 0) * Kp + k_];
1957
+ const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)];
1958
#pragma unroll
1959
for (uint32_t y = 0; y < YTILE; ++y) {
1967
- bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp])));
1960
+ bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
1961
}
1962
}
1963
@@ -1975,16 +1968,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1968
uint32_t k_ = k + threadIdx.x * A_CHUNK;
1969
if (k_ >= K) break;
1970
for (int n = 0; n < N; n++) {
1978
- bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
1971
+ bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
1972
}
1973
}
1974
1975
// Do the matrix multiplication in interleaved manner
1976
#pragma unroll
1977
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1985
- uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1986
- if (k >= K) break;
1987
-
1978
for (uint32_t n = 0; n < N; n++) {
1979
for (int i = 0; i < A_CHUNK; i += 8) {
1980
for (int y = 0; y < YTILE; ++y) {
@@ -2002,48 +1992,27 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1992
for (int y = 0; y < YTILE; y++) {
1993
float accm0 = sum[n][y][0];
1994
float accm16 = sum[n][y][8];
2005
- asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
2006
- : "=v"(accm0)
2007
- : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0));
2008
- asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
2009
- : "=v"(accm16)
2010
- : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16));
2011
- asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
2012
- : "=v"(accm0)
2013
- : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0));
2014
- asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
2015
- : "=v"(accm16)
2016
- : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16));
2017
- asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
2018
- : "=v"(accm0)
2019
- : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0));
2020
- asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
2021
- : "=v"(accm16)
2022
- : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16));
2023
- asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
2024
- : "=v"(accm0)
2025
- : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0));
2026
- asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
2027
- : "=v"(accm16)
2028
- : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16));
2029
- asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
2030
- : "=v"(accm0)
2031
- : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0));
2032
- asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
2033
- : "=v"(accm16)
2034
- : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16));
2035
- asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
2036
- : "=v"(accm0)
2037
- : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0));
2038
- asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
2039
- : "=v"(accm16)
2040
- : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16));
2041
- asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
2042
- : "=v"(accm0)
2043
- : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0));
2044
- asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
2045
- : "=v"(accm16)
2046
- : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16));
1995
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
1996
+ 1); // row_shl1
1997
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
1998
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
1999
+ 1); // row_shl2
2000
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
2001
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
2002
+ 1); // row_shl3
2003
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
2004
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
2005
+ 1); // row_shl8
2006
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
2007
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
2008
+ 1); // row_shl9
2009
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
2010
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
2011
+ 1); // row_shl10
2012
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
2013
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
2014
+ 1); // row_shl11
2015
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
2016
accm0 += __shfl(accm0, 36);
2017
accm16 += __shfl(accm16, 52);
2018
sum[n][y][0] = accm0 + __shfl(accm16, 16);
@@ -2051,19 +2020,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2020
}
2021
2022
if (threadIdx.x == 0) {
2023
+ scalar_t biases[N][YTILE] = {};
2024
+ if (BIAS)
2025
+ for (int n = 0; n < N; n++) {
2026
+ for (int y = 0; y < YTILE; y++) {
2027
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
2028
+ }
2029
+ }
2030
for (int n = 0; n < N; n++) {
2031
for (int y = 0; y < YTILE; y++) {
2032
if (y + m >= M) break; // To avoid mem access fault.
2033
sum[n][y][0] *= sA * sB;
2034
if constexpr (std::is_same_v<scalar_t, half>) {
2059
- if (BIAS)
2060
- sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
2035
+ sum[n][y][0] += __half2float(biases[n][y]);
2036
} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
2062
- if (BIAS)
2063
- sum[n][y][0] +=
2064
- __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
2037
+ sum[n][y][0] += __bfloat162float(biases[n][y]);
2038
}
2066
- C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]); // * sA * sB);
2039
+ C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
2040
}
2041
}
2042
}
@@ -2074,9 +2047,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2047
#else // !defined(__HIP__MI3XX__) TODO: Add NAVI support
2048
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2049
int A_CHUNK, int UNRL, int N>
2077
-__global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
2078
- const int Bx, const int By, const fp8_t* B,
2079
- const fp8_t* __restrict__ A,
2050
+__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
2051
+ const int M, const int Bx, const int By,
2052
+ const fp8_t* B, const fp8_t* __restrict__ A,
2053
const scalar_t* __restrict__ BIAS,
2054
scalar_t* C, const float* __restrict__ s_A,
2055
const float* __restrict__ s_B,
@@ -2089,8 +2062,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
2062
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2063
int A_CHUNK, int UNRL, int N>
2064
__global__ void __launch_bounds__(WvPrGrp* THRDS)
2092
- wvSplitKQ_hf_(const int K, const int Kp, const int M, const int Bx,
2093
- const int By, const fp8_t* B, const fp8_t* __restrict__ A,
2065
+ wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, const int M,
2066
+ const int Bx, const int By, const fp8_t* B,
2067
+ const fp8_t* __restrict__ A,
2068
const scalar_t* __restrict__ BIAS, scalar_t* C,
2069
const float* __restrict__ s_A, const float* __restrict__ s_B,
2070
const int _WvPrGrp, const int CuCount) {
@@ -2113,9 +2087,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2087
__shared__ fp8_t s[max_lds_len];
2088
2089
for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
2116
- k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
2090
+ k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
2091
+ #if defined(__gfx950__)
2092
+ __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
2093
+ #else
2094
*((bigType*)(&s[k])) = *((bigType*)(&A[k]));
2095
+ #endif
2096
}
2097
+ asm volatile("s_waitcnt vmcnt(0)");
2098
__syncthreads();
2099
2100
if (threadIdx.y >= _WvPrGrp) return;
@@ -2123,29 +2102,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2102
uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
2103
2104
using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
2126
- floatx16 sum[N][YTILE];
2105
float sA = *s_A;
2106
float sB = *s_B;
2107
2108
while (m < M) {
2131
- for (int i = 0; i < YTILE; i++)
2132
- for (int n = 0; n < N; n++) sum[n][i] = {0};
2133
-
2134
- bigType bigA[N][UNRL];
2135
- bigType bigB[YTILE][UNRL];
2136
-
2109
+ floatx16 sum[N][YTILE] = {};
2110
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
2111
+ bigType bigA[N][UNRL] = {};
2112
+ bigType bigB[YTILE][UNRL];
2113
+
2114
// Fetch the weight matrix from memory!
2115
#pragma unroll
2116
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2117
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
2118
uint32_t k_ = k + threadIdx.x * A_CHUNK;
2143
- if (k_ >= K) break;
2144
-
2145
- const fp8_t* B_ = &B[(m + 0) * Kp + k_];
2119
+ const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)];
2120
for (int y = 0; y < YTILE; ++y) {
2147
- if (y + m >= M) break; // To avoid mem access fault.
2148
- bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp])));
2121
+ bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
2122
}
2123
}
2124
@@ -2156,20 +2129,16 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2129
uint32_t k_ = k + threadIdx.x * A_CHUNK;
2130
if (k_ >= K) break;
2131
for (int n = 0; n < N; n++) {
2159
- if (k_ + K * n < max_lds_len)
2160
- bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
2132
+ if (k_ + Kap * n < max_lds_len)
2133
+ bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
2134
else
2162
- bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
2135
+ bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n])));
2136
}
2137
}
2138
2139
// Do the matrix multiplication in interleaved manner
2140
#pragma unroll
2141
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2169
- uint32_t k = k1 + k2 * THRDS * A_CHUNK;
2170
- uint32_t k_ = k + threadIdx.x * A_CHUNK;
2171
- if (k_ >= K) break;
2172
-
2142
for (uint32_t n = 0; n < N; n++) {
2143
for (int i = 0; i < A_CHUNK; i += 8) {
2144
for (int y = 0; y < YTILE; ++y) {
@@ -2187,48 +2156,27 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2156
for (int y = 0; y < YTILE; y++) {
2157
float accm0 = sum[n][y][0];
2158
float accm16 = sum[n][y][8];
2190
- asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
2191
- : "=v"(accm0)
2192
- : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0));
2193
- asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
2194
- : "=v"(accm16)
2195
- : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16));
2196
- asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
2197
- : "=v"(accm0)
2198
- : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0));
2199
- asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
2200
- : "=v"(accm16)
2201
- : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16));
2202
- asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
2203
- : "=v"(accm0)
2204
- : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0));
2205
- asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
2206
- : "=v"(accm16)
2207
- : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16));
2208
- asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
2209
- : "=v"(accm0)
2210
- : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0));
2211
- asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
2212
- : "=v"(accm16)
2213
- : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16));
2214
- asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
2215
- : "=v"(accm0)
2216
- : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0));
2217
- asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
2218
- : "=v"(accm16)
2219
- : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16));
2220
- asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
2221
- : "=v"(accm0)
2222
- : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0));
2223
- asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
2224
- : "=v"(accm16)
2225
- : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16));
2226
- asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
2227
- : "=v"(accm0)
2228
- : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0));
2229
- asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
2230
- : "=v"(accm16)
2231
- : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16));
2159
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
2160
+ 1); // row_shl1
2161
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
2162
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
2163
+ 1); // row_shl2
2164
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
2165
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
2166
+ 1); // row_shl3
2167
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
2168
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
2169
+ 1); // row_shl8
2170
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
2171
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
2172
+ 1); // row_shl9
2173
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
2174
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
2175
+ 1); // row_shl10
2176
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
2177
+ accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
2178
+ 1); // row_shl11
2179
+ accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
2180
accm0 += __shfl(accm0, 36);
2181
accm16 += __shfl(accm16, 52);
2182
sum[n][y][0] = accm0 + __shfl(accm16, 16);
@@ -2236,17 +2184,21 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2184
}
2185
2186
if (threadIdx.x == 0) {
2187
+ scalar_t biases[N][YTILE] = {};
2188
+ if (BIAS)
2189
+ for (int n = 0; n < N; n++) {
2190
+ for (int y = 0; y < YTILE; y++) {
2191
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
2192
+ }
2193
+ }
2194
for (int n = 0; n < N; n++) {
2195
for (int y = 0; y < YTILE; y++) {
2196
if (y + m >= M) break; // To avoid mem access fault.
2197
sum[n][y][0] *= sA * sB;
2198
if constexpr (std::is_same_v<scalar_t, half>) {
2244
- if (BIAS)
2245
- sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
2199
+ sum[n][y][0] += __half2float(biases[n][y]);
2200
} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
2247
- if (BIAS)
2248
- sum[n][y][0] +=
2249
- __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
2201
+ sum[n][y][0] += __bfloat162float(biases[n][y]);
2202
}
2203
C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
2204
}
@@ -2259,9 +2211,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2211
#else // !defined(__HIP__MI3XX__) TODO: Add NAVI support
2212
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2213
int A_CHUNK, int UNRL, int N>
2262
-__global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
2263
- const int Bx, const int By, const fp8_t* B,
2264
- const fp8_t* __restrict__ A,
2214
+__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
2215
+ const int M, const int Bx, const int By,
2216
+ const fp8_t* B, const fp8_t* __restrict__ A,
2217
const scalar_t* __restrict__ BIAS, scalar_t* C,
2218
const float* __restrict__ s_A,
2219
const float* __restrict__ s_B, const int _WvPrGrp,
@@ -2270,17 +2222,18 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
2222
}
2223
#endif // defined(__HIP__MI3XX__) TODO: Add NAVI support
2224
2273
-void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
2225
+void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
2226
const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
2227
const at::Tensor& scale_a, const at::Tensor& scale_b,
2228
const int64_t CuCount) {
2229
static c10::ScalarType kFp8Type = is_fp8_ocp()
2230
? c10::ScalarType::Float8_e4m3fn
2231
: c10::ScalarType::Float8_e4m3fnuz;
2280
- auto M_in = in_a.size(0);
2281
- auto K_in = in_a.size(1);
2282
- auto N_in = in_b.size(0);
2283
- auto Kp_in = in_a.stride(0);
2232
+ auto M_in = in_b.size(0);
2233
+ auto K_in = in_b.size(1);
2234
+ auto N_in = in_a.size(0);
2235
+ auto Kap_in = in_a.stride(0);
2236
+ auto Kbp_in = in_b.stride(0);
2237
auto Bx_in =
2238
(in_bias.has_value() && in_bias->numel() > 0)
2239
? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
@@ -2300,23 +2253,22 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
2253
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
2254
const int max_lds_len = get_lds_size();
2255
2303
-#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \
2304
- _N) \
2305
- { \
2306
- dim3 block(64, _WvPrGrp); \
2307
- if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
2308
- int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
2309
- wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
2310
- <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
2311
- b_ptr, bias_ptr, c_ptr, s_a, s_b, \
2312
- __wvPrGrp, CuCount); \
2313
- } else { \
2314
- int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \
2315
- wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
2316
- <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
2317
- b_ptr, bias_ptr, c_ptr, s_a, s_b, \
2318
- __wvPrGrp, CuCount); \
2319
- } \
2256
+#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2257
+ { \
2258
+ dim3 block(64, _WvPrGrp); \
2259
+ if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
2260
+ int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16)); \
2261
+ wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
2262
+ <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
2263
+ By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \
2264
+ s_a, s_b, __wvPrGrp, CuCount); \
2265
+ } else { \
2266
+ int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16)); \
2267
+ wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
2268
+ <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
2269
+ By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \
2270
+ s_a, s_b, __wvPrGrp, CuCount); \
2271
+ } \
2272
}
2273
2274
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] {
@@ -2332,16 +2284,16 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
2284
: nullptr;
2285
switch (N_in) {
2286
case 1:
2335
- WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 1)
2287
+ WVSPLITKQ(12, 2, 2, 2, 2, 1)
2288
break;
2289
case 2:
2338
- WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 2)
2290
+ WVSPLITKQ(12, 2, 2, 2, 2, 2)
2291
break;
2292
case 3:
2341
- WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 3)
2293
+ WVSPLITKQ(8, 2, 2, 1, 1, 3)
2294
break;
2295
case 4:
2344
- WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 4)
2296
+ WVSPLITKQ(4, 2, 2, 1, 1, 4)
2297
break;
2298
default:
2299
throw std::runtime_error(
tests/kernels/quantization/test_rocm_skinny_gemms.py
+40
-52
@@ -73,21 +73,40 @@ NKM_FACTORS_WVSPLITKRC = [
73
NKM_FACTORS_WVSPLITK_FP8 = [
74
# FP8-specific cases with K % 16 == 0
75
(1, 16, 16),
76
+ (1, 32, 16 + 16),
77
(1, 64, 64),
78
+ (1, 64, 64 + 16),
79
+ (1, 64 + 16, 64),
80
+ (1, 64 + 16, 64 + 16),
81
+ (4, 64, 64),
82
+ (4, 64, 64 + 16),
83
+ (4, 64 + 16, 64),
84
+ (4, 64 + 16, 64 + 16),
85
(2, 512, 512),
86
+ (3, 512, 512),
87
+ (3, 512, 512 + 16),
88
+ (4, 512, 512),
89
(3, 2048, 2048),
90
+ (3, 2048, 2048 + 16),
91
+ (4, 2048 + 16, 2048),
92
+ (4, 2048 + 16, 2048 + 16),
93
(4, 4096, 4096),
94
(4, 16400, 2048),
95
+ (4, 16400, 2048 + 16),
96
# Extended FP8 dimensions not covered by WVSPLITK
97
(1, 14336, 1024),
98
(2, 24576, 2048),
99
(4, 32768, 28672),
100
+ (4, 32768 * 2, 28672),
101
+ (4, 32768 * 2, 28672 + 16),
102
+ (4, 32768 * 2 + 16, 28672),
103
+ (4, 32768 * 2 + 16, 28672 + 16),
104
]
105
106
SEEDS = [0]
107
108
90
-def pad_weights_fp8(weight):
109
+def pad_fp8(weight):
110
num_pad = 256 // weight.element_size()
111
import torch.nn.functional as F
112
@@ -195,72 +214,41 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
214
assert torch.allclose(out, ref_out, rtol=0.01)
215
216
217
+@pytest.mark.parametrize("xnorm", [False, True])
218
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
219
@pytest.mark.parametrize("dtype", DTYPES)
220
@pytest.mark.parametrize("seed", SEEDS)
201
-@pytest.mark.parametrize("padded", [False, True])
221
+@pytest.mark.parametrize("padded_a", [False, True])
222
+@pytest.mark.parametrize("padded_b", [False, True])
223
+@pytest.mark.parametrize("biased", [False, True])
224
@pytest.mark.skipif(
225
not (current_platform.is_rocm() and current_platform.supports_fp8()),
226
reason="only test for rocm fp8",
227
)
206
-def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
228
+def test_rocm_wvsplitk_fp8_kernel(
229
+ xnorm, n, k, m, dtype, seed, padded_a, padded_b, biased
230
+):
231
torch.manual_seed(seed)
232
209
- A = torch.rand(n, k, device="cuda") - 0.5
210
- B = torch.rand(m, k, device="cuda") - 0.5
233
+ xavier = math.sqrt(2 / k) if xnorm else 1 # normalize to avoid large deltas
234
+ A = (torch.rand(n, k, device="cuda") * 2 - 1) * xavier
235
+ B = (torch.rand(m, k, device="cuda") * 2 - 1) * xavier
236
237
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
238
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
214
- if padded:
215
- B = pad_weights_fp8(B)
216
-
217
- ref_out = torch._scaled_mm(
218
- A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b
219
- )
220
- out = ops.wvSplitKQ(
221
- B,
222
- A,
223
- dtype,
224
- scale_a,
225
- scale_b,
226
- get_cu_count(),
227
- )
228
-
229
- assert torch.allclose(out, ref_out, rtol=0.01)
239
+ if padded_b:
240
+ B = pad_fp8(B)
241
+ if padded_a:
242
+ A = pad_fp8(A)
243
231
-
232
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
233
-@pytest.mark.parametrize("dtype", DTYPES)
234
-@pytest.mark.parametrize("seed", SEEDS)
235
-@pytest.mark.parametrize("padded", [False, True])
236
-@pytest.mark.skipif(
237
- not (current_platform.is_rocm() and current_platform.supports_fp8()),
238
- reason="only test for rocm fp8",
239
-)
240
-def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded):
241
- torch.manual_seed(seed)
242
-
243
- xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
244
- A = (torch.rand(n, k, device="cuda") - 0.5) * xavier
245
- B = (torch.rand(m, k, device="cuda") - 0.5) * xavier
246
- BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5
247
-
248
- A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
249
- B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
250
- if padded:
251
- B = pad_weights_fp8(B)
244
+ BIAS = None if (not biased) else (torch.rand(m, dtype=dtype, device="cuda") * 2 - 1)
245
246
ref_out = torch._scaled_mm(
247
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
248
)
256
- out = ops.wvSplitKQ(
257
- B,
258
- A,
259
- dtype,
260
- scale_a,
261
- scale_b,
262
- get_cu_count(),
263
- BIAS,
264
- )
249
+ out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, get_cu_count(), BIAS)
250
266
- assert torch.allclose(out, ref_out, rtol=0.01)
251
+ if xnorm:
252
+ assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
253
+ else:
254
+ assert torch.allclose(out, ref_out, 0.01)
vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py
+3
-3
@@ -25,10 +25,10 @@ def rocm_per_tensor_float_w8a8_scaled_mm_impl(
25
bias: torch.Tensor,
26
) -> torch.Tensor:
27
if (
28
- A.shape[0] == 1
29
- and B.shape[1] % 16 == 0
28
+ A.shape[0] <= 4
29
+ and B.shape[0] % 16 == 0 # M TODO: needed?
30
+ and B.shape[1] % 16 == 0 # K
31
and ((bias is None) or (bias.dtype == out_dtype))
31
- and A.is_contiguous()
32
):
33
output = ops.wvSplitKQ(
34
B.t(),