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 14, 05:30 UTC
[ROCm] Add skinny gemm bias support for dtypes fp16,bf16,fp8 (#24988)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com> Signed-off-by: Hashem Hashemi <159079214+amd-hhashemi@users.noreply.github.com>
Hashem Hashemi committed
Sep 23, 2025 at 11:31 UTC
a3a7828010f4ab8e608fa257d09ab31846e83a9a
7 files changed
+233
-79
csrc/rocm/ops.h
+6
-3
@@ -5,11 +5,14 @@
5
torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
6
const int64_t rows_per_block);
7
8
-torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
8
+torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
9
+ const c10::optional<at::Tensor>& in_bias,
10
const int64_t CuCount);
11
11
-void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
12
- at::Tensor& scale_a, at::Tensor& scale_b, const int64_t CuCount);
12
+void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
13
+ const c10::optional<at::Tensor>& in_bias, at::Tensor& out_c,
14
+ const at::Tensor& scale_a, const at::Tensor& scale_b,
15
+ const int64_t CuCount);
16
17
void paged_attention(
18
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
csrc/rocm/skinny_gemms.cu
+139
-42
@@ -292,8 +292,9 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
292
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
293
int UNRL, int N>
294
__global__ void __launch_bounds__(WvPrGrp* THRDS)
295
- wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
296
- const scalar_t* __restrict__ A, scalar_t* C,
295
+ wvSplitK_hf_sml_(const int K, const int M, const int Bx, const int By,
296
+ const scalar_t* B, const scalar_t* __restrict__ A,
297
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
298
const int _WvPrGrp, const int CuCount) {
299
constexpr int max_lds_len = LDS_SIZE / 2;
300
#if defined(__HIP__MI3XX__)
@@ -484,7 +485,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
485
if (threadIdx.x == 63) {
486
for (int n = 0; n < N; n++) {
487
for (int i = 0; i < YTILE; i++) {
487
- // if (commitColumn[i]) C[m + i + n * M] = __float2half(sum[n][i]);
488
+ if constexpr (std::is_same_v<scalar_t, half>) {
489
+ if (BIAS)
490
+ sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
491
+ } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
492
+ if (BIAS)
493
+ sum[n][i] +=
494
+ __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
495
+ }
496
C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
497
}
498
}
@@ -529,7 +537,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
537
if (threadIdx.x == 63) {
538
for (int n = 0; n < N; n++) {
539
for (int i = 0; i < YTILE; i++) {
532
- // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
540
+ if (BIAS)
541
+ sum4[n][i][0] +=
542
+ __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
543
C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
544
}
545
}
@@ -541,8 +551,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
551
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
552
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
553
int UNRL, int N>
544
-__global__ void wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
545
- const scalar_t* __restrict__ A, scalar_t* C,
554
+__global__ void wvSplitK_hf_sml_(const int K, const int M, const int Bx,
555
+ const int By, const scalar_t* B,
556
+ const scalar_t* __restrict__ A,
557
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
558
const int _WvPrGrp, const int CuCount) {
559
UNREACHABLE_CODE
560
}
@@ -553,8 +565,9 @@ __global__ void wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
565
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
566
int UNRL, int N>
567
__global__ void __launch_bounds__(WvPrGrp* THRDS)
556
- wvSplitK_hf_(const int K, const int M, const scalar_t* B,
557
- const scalar_t* __restrict__ A, scalar_t* C,
568
+ wvSplitK_hf_(const int K, const int M, const int Bx, const int By,
569
+ const scalar_t* B, const scalar_t* __restrict__ A,
570
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
571
const int _WvPrGrp, const int CuCount) {
572
constexpr int max_lds_len = LDS_SIZE / 2;
573
#if defined(__HIP__MI3XX__)
@@ -772,8 +785,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
785
if (threadIdx.x == 63) {
786
for (int n = 0; n < N; n++) {
787
for (int i = 0; i < YTILE; i++) {
775
- if (commitColumn[i])
788
+ if (commitColumn[i]) {
789
+ if constexpr (std::is_same_v<scalar_t, half>) {
790
+ if (BIAS)
791
+ sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
792
+ } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
793
+ if (BIAS)
794
+ sum[n][i] +=
795
+ __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
796
+ }
797
C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
798
+ }
799
}
800
}
801
}
@@ -818,8 +840,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
840
if (threadIdx.x == 63) {
841
for (int n = 0; n < N; n++) {
842
for (int i = 0; i < YTILE; i++) {
821
- // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
822
- C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
843
+ if (commitColumn[i]) {
844
+ if (BIAS)
845
+ sum4[n][i][0] +=
846
+ __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
847
+ C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
848
+ }
849
}
850
}
851
}
@@ -842,8 +868,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
868
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
869
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
870
int UNRL, int N>
845
-__global__ void wvSplitK_hf_(const int K, const int M, const scalar_t* B,
846
- const scalar_t* __restrict__ A, scalar_t* C,
871
+__global__ void wvSplitK_hf_(const int K, const int M, const int Bx,
872
+ const int By, const scalar_t* B,
873
+ const scalar_t* __restrict__ A,
874
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
875
const int _WvPrGrp, const int CuCount) {
876
UNREACHABLE_CODE
877
}
@@ -854,8 +882,9 @@ __global__ void wvSplitK_hf_(const int K, const int M, const scalar_t* B,
882
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
883
int UNRL, int N>
884
__global__ void __launch_bounds__(WvPrGrp* THRDS)
857
- wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
858
- const scalar_t* __restrict__ A, scalar_t* C,
885
+ wvSplitK_hf_big_(const int K, const int M, const int Bx, const int By,
886
+ const scalar_t* B, const scalar_t* __restrict__ A,
887
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
888
const int _WvPrGrp, const int CuCount) {
889
constexpr int max_lds_len = LDS_SIZE / 2;
890
#if defined(__HIP__MI3XX__)
@@ -1124,8 +1153,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1153
if (threadIdx.x == 63) {
1154
for (int n = 0; n < N; n++) {
1155
for (int i = 0; i < YTILE; i++) {
1127
- if (commitColumn[i])
1156
+ if (commitColumn[i]) {
1157
+ if constexpr (std::is_same_v<scalar_t, half>) {
1158
+ if (BIAS)
1159
+ sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
1160
+ } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1161
+ if (BIAS)
1162
+ sum[n][i] +=
1163
+ __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
1164
+ }
1165
C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
1166
+ }
1167
}
1168
}
1169
}
@@ -1166,8 +1204,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1204
if (threadIdx.x == 63) {
1205
for (int n = 0; n < N; n++) {
1206
for (int i = 0; i < YTILE; i++) {
1169
- // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
1170
- C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
1207
+ if (commitColumn[i]) {
1208
+ if (BIAS)
1209
+ sum4[n][i][0] +=
1210
+ __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
1211
+ C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
1212
+ }
1213
}
1214
}
1215
}
@@ -1190,8 +1232,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1232
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1233
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1234
int UNRL, int N>
1193
-__global__ void wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
1194
- const scalar_t* __restrict__ A, scalar_t* C,
1235
+__global__ void wvSplitK_hf_big_(const int K, const int M, const int Bx,
1236
+ const int By, const scalar_t* B,
1237
+ const scalar_t* __restrict__ A,
1238
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
1239
const int _WvPrGrp, const int CuCount) {
1240
UNREACHABLE_CODE
1241
}
@@ -1226,11 +1270,20 @@ int mindiv(int N, int div1, int div2) {
1270
return rtn;
1271
}
1272
1229
-torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1273
+torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1274
+ const c10::optional<at::Tensor>& in_bias,
1275
const int64_t CuCount) {
1276
auto M_in = in_a.size(0);
1277
auto K_in = in_a.size(1);
1278
auto N_in = in_b.size(0);
1279
+ auto Bx_in =
1280
+ (in_bias.has_value() && in_bias->numel() > 0)
1281
+ ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
1282
+ : 1;
1283
+ auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
1284
+ in_bias->sizes().size() == 2)
1285
+ ? in_bias->size(0)
1286
+ : 1;
1287
1288
TORCH_CHECK(in_a.dtype() == in_b.dtype());
1289
TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
@@ -1254,18 +1307,18 @@ torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1307
if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
1308
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
1309
wvSplitK_hf_sml_<fptype, 64, _YTILEs, _WvPrGrp, 8, _UNRLs, _N> \
1257
- <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
1258
- CuCount); \
1310
+ <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4, \
1311
+ biasf4, c, __wvPrGrp, CuCount); \
1312
} else if (K_in * N_in <= max_lds_len * 1.2) { \
1313
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \
1314
wvSplitK_hf_<fptype, 64, _YTILEm, _WvPrGrp, 8, _UNRLm, _N> \
1262
- <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
1263
- CuCount); \
1315
+ <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4, \
1316
+ biasf4, c, __wvPrGrp, CuCount); \
1317
} else { \
1318
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEb, _WvPrGrp); \
1319
wvSplitK_hf_big_<fptype, 64, _YTILEb, _WvPrGrp, 8, _UNRLb, _N> \
1267
- <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
1268
- CuCount); \
1320
+ <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4, \
1321
+ biasf4, c, __wvPrGrp, CuCount); \
1322
} \
1323
}
1324
@@ -1273,6 +1326,10 @@ torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1326
using fptype = typename scalar<scalar_t>::type;
1327
fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
1328
const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
1329
+ const fptype* biasf4 =
1330
+ (in_bias.has_value() && in_bias->numel() > 0)
1331
+ ? reinterpret_cast<const fptype*>(in_bias->data_ptr())
1332
+ : nullptr;
1333
fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
1334
switch (N_in) {
1335
case 1:
@@ -1300,8 +1357,9 @@ torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1357
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1358
int A_CHUNK, int UNRL, int N>
1359
__global__ void __launch_bounds__(WvPrGrp* THRDS)
1303
- wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const fp8_t* B,
1304
- const fp8_t* __restrict__ A, scalar_t* C,
1360
+ wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const int Bx,
1361
+ const int By, const fp8_t* B, const fp8_t* __restrict__ A,
1362
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
1363
const float* __restrict__ s_A,
1364
const float* __restrict__ s_B, const int _WvPrGrp,
1365
const int CuCount) {
@@ -1453,7 +1511,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1511
if (threadIdx.x == 0) {
1512
for (int n = 0; n < N; n++) {
1513
for (int y = 0; y < YTILE; y++) {
1456
- C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0] * sA * sB);
1514
+ if (y + m >= M) break; // To avoid mem access fault.
1515
+ sum[n][y][0] *= sA * sB;
1516
+ if constexpr (std::is_same_v<scalar_t, half>) {
1517
+ if (BIAS)
1518
+ sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
1519
+ } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1520
+ if (BIAS)
1521
+ sum[n][y][0] +=
1522
+ __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
1523
+ }
1524
+ C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]); // * sA * sB);
1525
}
1526
}
1527
}
@@ -1465,7 +1533,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1533
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1534
int A_CHUNK, int UNRL, int N>
1535
__global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
1468
- const fp8_t* B, const fp8_t* __restrict__ A,
1536
+ const int Bx, const int By, const fp8_t* B,
1537
+ const fp8_t* __restrict__ A,
1538
+ const scalar_t* __restrict__ BIAS,
1539
scalar_t* C, const float* __restrict__ s_A,
1540
const float* __restrict__ s_B,
1541
const int _WvPrGrp, const int CuCount) {
@@ -1477,8 +1547,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
1547
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1548
int A_CHUNK, int UNRL, int N>
1549
__global__ void __launch_bounds__(WvPrGrp* THRDS)
1480
- wvSplitKQ_hf_(const int K, const int Kp, const int M, const fp8_t* B,
1481
- const fp8_t* __restrict__ A, scalar_t* C,
1550
+ wvSplitKQ_hf_(const int K, const int Kp, const int M, const int Bx,
1551
+ const int By, const fp8_t* B, const fp8_t* __restrict__ A,
1552
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
1553
const float* __restrict__ s_A, const float* __restrict__ s_B,
1554
const int _WvPrGrp, const int CuCount) {
1555
constexpr int max_lds_len = LDS_SIZE;
@@ -1626,7 +1697,16 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1697
for (int n = 0; n < N; n++) {
1698
for (int y = 0; y < YTILE; y++) {
1699
if (y + m >= M) break; // To avoid mem access fault.
1629
- C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0] * sA * sB);
1700
+ sum[n][y][0] *= sA * sB;
1701
+ if constexpr (std::is_same_v<scalar_t, half>) {
1702
+ if (BIAS)
1703
+ sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
1704
+ } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1705
+ if (BIAS)
1706
+ sum[n][y][0] +=
1707
+ __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
1708
+ }
1709
+ C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
1710
}
1711
}
1712
}
@@ -1638,16 +1718,19 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1718
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1719
int A_CHUNK, int UNRL, int N>
1720
__global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
1641
- const fp8_t* B, const fp8_t* __restrict__ A,
1642
- scalar_t* C, const float* __restrict__ s_A,
1721
+ const int Bx, const int By, const fp8_t* B,
1722
+ const fp8_t* __restrict__ A,
1723
+ const scalar_t* __restrict__ BIAS, scalar_t* C,
1724
+ const float* __restrict__ s_A,
1725
const float* __restrict__ s_B, const int _WvPrGrp,
1726
const int CuCount) {
1727
UNREACHABLE_CODE
1728
}
1729
#endif // defined(__HIP__MI3XX__) TODO: Add NAVI support
1730
1649
-void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1650
- at::Tensor& scale_a, at::Tensor& scale_b,
1731
+void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
1732
+ const c10::optional<at::Tensor>& in_bias, at::Tensor& out_c,
1733
+ const at::Tensor& scale_a, const at::Tensor& scale_b,
1734
const int64_t CuCount) {
1735
static c10::ScalarType kFp8Type = is_fp8_ocp()
1736
? c10::ScalarType::Float8_e4m3fn
@@ -1656,6 +1739,15 @@ void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1739
auto K_in = in_a.size(1);
1740
auto N_in = in_b.size(0);
1741
auto Kp_in = in_a.stride(0);
1742
+ auto Bx_in =
1743
+ (in_bias.has_value() && in_bias->numel() > 0)
1744
+ ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
1745
+ : 1;
1746
+ auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
1747
+ in_bias->sizes().size() == 2)
1748
+ ? in_bias->size(0)
1749
+ : 1;
1750
+
1751
TORCH_CHECK(K_in % 16 == 0, "k % 16 == 0");
1752
TORCH_CHECK(in_a.dtype() == in_b.dtype() && in_a.dtype() == kFp8Type);
1753
TORCH_CHECK(out_c.dtype() == torch::kFloat16 ||
@@ -1673,13 +1765,15 @@ void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1765
if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
1766
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
1767
wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
1676
- <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, a_ptr, b_ptr, c_ptr, \
1677
- s_a, s_b, __wvPrGrp, CuCount); \
1768
+ <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
1769
+ b_ptr, bias_ptr, c_ptr, s_a, s_b, \
1770
+ __wvPrGrp, CuCount); \
1771
} else { \
1772
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \
1773
wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
1681
- <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, a_ptr, b_ptr, c_ptr, \
1682
- s_a, s_b, __wvPrGrp, CuCount); \
1774
+ <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
1775
+ b_ptr, bias_ptr, c_ptr, s_a, s_b, \
1776
+ __wvPrGrp, CuCount); \
1777
} \
1778
}
1779
@@ -1691,6 +1785,9 @@ void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1785
VLLM_DISPATCH_FP8_TYPES(in_a.scalar_type(), "wvSplitKQ", [&] {
1786
auto a_ptr = in_a.data_ptr<fp8_t>();
1787
auto b_ptr = in_b.data_ptr<fp8_t>();
1788
+ auto bias_ptr = (in_bias.has_value() && in_bias->numel() > 0)
1789
+ ? reinterpret_cast<fptype*>(in_bias->data_ptr())
1790
+ : nullptr;
1791
switch (N_in) {
1792
case 1:
1793
WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 1)
csrc/rocm/torch_bindings.cpp
+3
-2
@@ -22,13 +22,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
22
23
// Custom gemm op for skinny matrix-matrix multiplication
24
rocm_ops.def(
25
- "wvSplitK(Tensor in_a, Tensor in_b, int CuCount) -> "
25
+ "wvSplitK(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
26
"Tensor");
27
rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
28
29
// wvSplitK for fp8
30
rocm_ops.def(
31
- "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor! out_c, Tensor scale_a, "
31
+ "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor? in_bias, Tensor! out_c, "
32
+ "Tensor scale_a, "
33
" Tensor scale_b, int CuCount) -> ()");
34
rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
35
tests/kernels/quantization/test_rocm_skinny_gemms.py
+62
-18
@@ -1,12 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+
5
import pytest
6
import torch
7
8
import vllm._custom_ops as ops
9
from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
8
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
9
- rocm_per_tensor_w8a8_scaled_mm_impl)
10
from vllm.platforms import current_platform
11
12
DTYPES = [torch.bfloat16, torch.float16]
@@ -49,6 +49,7 @@ NKM_FACTORS_WVSPLITK_FP8 = [
49
(2, 512, 512),
50
(3, 2048, 2048),
51
(4, 4096, 4096),
52
+ (4, 16400, 2048),
53
# Extended FP8 dimensions not covered by WVSPLITK
54
(1, 14336, 1024),
55
(2, 24576, 2048),
@@ -67,6 +68,9 @@ SEEDS = [0]
68
@torch.inference_mode()
69
def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
70
torch.manual_seed(seed)
71
+ #TODO: Zero-centering the inputs causes errors for LLMM1!
72
+ # Without that the numbers quickly saturate, and may
73
+ # be giving false matches.
74
A = torch.rand(n, k, dtype=dtype, device="cuda")
75
B = torch.rand(m, k, dtype=dtype, device="cuda")
76
@@ -85,11 +89,51 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
89
torch.manual_seed(seed)
90
cu_count = current_platform.get_cu_count()
91
88
- A = torch.rand(n, k, dtype=dtype, device="cuda")
89
- B = torch.rand(m, k, dtype=dtype, device="cuda")
92
+ A = torch.rand(n, k, dtype=dtype, device="cuda") - .5
93
+ B = torch.rand(m, k, dtype=dtype, device="cuda") - .5
94
91
- ref_out = torch.matmul(A, B.t())
92
- out = ops.wvSplitK(B, A, cu_count)
95
+ ref_out = torch.nn.functional.linear(A, B)
96
+ out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)
97
+
98
+ assert torch.allclose(out, ref_out, rtol=0.01)
99
+
100
+
101
+@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
102
+@pytest.mark.parametrize("dtype", DTYPES)
103
+@pytest.mark.parametrize("seed", SEEDS)
104
+@pytest.mark.skipif(not current_platform.is_rocm(),
105
+ reason="only test for rocm")
106
+def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
107
+ torch.manual_seed(seed)
108
+ cu_count = current_platform.get_cu_count()
109
+
110
+ xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
111
+ A = (torch.rand(n, k, dtype=dtype, device="cuda") - .5) * xavier
112
+ B = (torch.rand(m, k, dtype=dtype, device="cuda") - .5) * xavier
113
+ BIAS = torch.rand(m, dtype=dtype, device="cuda") - .5
114
+
115
+ ref_out = torch.nn.functional.linear(A, B, BIAS)
116
+ out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
117
+
118
+ assert torch.allclose(out, ref_out, rtol=0.01)
119
+
120
+
121
+@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
122
+@pytest.mark.parametrize("dtype", DTYPES)
123
+@pytest.mark.parametrize("seed", SEEDS)
124
+@pytest.mark.skipif(not current_platform.is_rocm(),
125
+ reason="only test for rocm")
126
+def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
127
+ torch.manual_seed(seed)
128
+ cu_count = current_platform.get_cu_count()
129
+
130
+ xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
131
+ A = (torch.rand(n, k, dtype=dtype, device="cuda") - .5) * xavier
132
+ B = (torch.rand(m, k, dtype=dtype, device="cuda") - .5) * xavier
133
+ BIAS = torch.rand(n, m, dtype=dtype, device="cuda") - .5
134
+
135
+ ref_out = torch.nn.functional.linear(A, B, BIAS)
136
+ out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
137
138
assert torch.allclose(out, ref_out, rtol=0.01)
139
@@ -103,8 +147,8 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
147
def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
148
torch.manual_seed(seed)
149
106
- A = torch.rand(n, k, device="cuda")
107
- B = torch.rand(m, k, device="cuda")
150
+ A = torch.rand(n, k, device="cuda") - 0.5
151
+ B = torch.rand(m, k, device="cuda") - 0.5
152
153
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
154
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
@@ -123,27 +167,27 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
167
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
168
@pytest.mark.parametrize("dtype", DTYPES)
169
@pytest.mark.parametrize("seed", SEEDS)
126
-@pytest.mark.parametrize("use_bias", [True, False])
170
@pytest.mark.skipif(
171
not (current_platform.is_rocm() and current_platform.supports_fp8()),
172
reason="only test for rocm fp8")
130
-def test_rocm_per_tensor_w8a8_scaled_mm_impl(n, k, m, dtype, seed, use_bias):
173
+def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
174
torch.manual_seed(seed)
175
133
- A = torch.rand(n, k, device="cuda")
134
- B = torch.rand(m, k, device="cuda")
176
+ xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
177
+ A = (torch.rand(n, k, device="cuda") - .5) * xavier
178
+ B = (torch.rand(m, k, device="cuda") - .5) * xavier
179
+ BIAS = torch.rand(m, dtype=dtype, device="cuda") - .5
180
181
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
182
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
183
139
- bias = torch.rand(1, m, dtype=dtype, device="cuda") if use_bias else None
140
-
141
- output = rocm_per_tensor_w8a8_scaled_mm_impl(A, B.t(), dtype, scale_a,
142
- scale_b, bias)
184
ref_out = torch._scaled_mm(A,
185
B.t(),
186
out_dtype=dtype,
187
scale_a=scale_a,
188
scale_b=scale_b,
148
- bias=bias)
149
- assert torch.allclose(output, ref_out, rtol=0.01)
189
+ bias=BIAS)
190
+ out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b,
191
+ current_platform.get_cu_count(), BIAS)
192
+
193
+ assert torch.allclose(out, ref_out, rtol=0.01)
vllm/_custom_ops.py
+15
-8
@@ -1447,17 +1447,24 @@ def LLMM1(a: torch.Tensor, b: torch.Tensor,
1447
return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
1448
1449
1450
-def wvSplitK(a: torch.Tensor, b: torch.Tensor, cu_count: int) -> torch.Tensor:
1451
- return torch.ops._rocm_C.wvSplitK(a, b, cu_count)
1452
-
1453
-
1454
-def wvSplitKQ(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype,
1455
- scale_a: torch.Tensor, scale_b: torch.Tensor,
1456
- cu_count: int) -> torch.Tensor:
1450
+def wvSplitK(a: torch.Tensor,
1451
+ b: torch.Tensor,
1452
+ cu_count: int,
1453
+ bias: torch.Tensor = None) -> torch.Tensor:
1454
+ return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
1455
+
1456
+
1457
+def wvSplitKQ(a: torch.Tensor,
1458
+ b: torch.Tensor,
1459
+ out_dtype: torch.dtype,
1460
+ scale_a: torch.Tensor,
1461
+ scale_b: torch.Tensor,
1462
+ cu_count: int,
1463
+ bias: torch.Tensor = None) -> torch.Tensor:
1464
out = torch.empty((b.shape[0], a.shape[0]),
1465
dtype=out_dtype,
1466
device=b.device)
1460
- torch.ops._rocm_C.wvSplitKQ(a, b, out, scale_a, scale_b, cu_count)
1467
+ torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count)
1468
return out
1469
1470
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+5
-3
@@ -178,10 +178,12 @@ def rocm_per_tensor_w8a8_scaled_mm_impl(qinput: torch.Tensor,
178
scale_b: torch.Tensor,
179
bias: torch.Tensor) -> torch.Tensor:
180
from vllm.platforms.rocm import on_mi3xx
181
- if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi3xx(
182
- ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0 and bias is None:
181
+ if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi3xx() and \
182
+ qinput.shape[0] == 1 and \
183
+ qinput.shape[1] % 16 == 0 and \
184
+ ((bias is None) or (bias.dtype == out_dtype)) :
185
output = ops.wvSplitKQ(weight.t(), qinput, out_dtype, scale_a, scale_b,
184
- current_platform.get_cu_count())
186
+ current_platform.get_cu_count(), bias)
187
else:
188
output = torch._scaled_mm(qinput,
189
weight,
vllm/model_executor/layers/utils.py
+3
-3
@@ -100,7 +100,7 @@ def rocm_unquantized_gemm_impl(
100
k = weight.shape[1]
101
use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and on_gfx9() and \
102
x.dtype in [torch.float16, torch.bfloat16] \
103
- and k % 8 == 0 and bias is None)
103
+ and k % 8 == 0)
104
105
if use_skinny is not True:
106
return torch.nn.functional.linear(x, weight, bias)
@@ -111,9 +111,9 @@ def rocm_unquantized_gemm_impl(
111
cu_count = current_platform.get_cu_count()
112
113
if m > 8 and 0 < n <= 4:
114
- out = ops.wvSplitK(weight, x_view, cu_count)
114
+ out = ops.wvSplitK(weight, x_view, cu_count, bias)
115
return out.view(*x.shape[:-1], weight.shape[0])
116
- elif m % 4 == 0 and n == 1 and k <= 8192:
116
+ elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
117
out = ops.LLMM1(weight, x_view, 4)
118
return out.view(*x.shape[:-1], weight.shape[0])
119
return torch.nn.functional.linear(x, weight, bias)