@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1779c09898
Fork 0

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 6, 08:15 UTC

[ROCm] Enable wvSplitK skinny GEMM kernel for RDNA4/gfx1x decode (#34709)

Signed-off-by: L.B.R. <lbr@mmonad.com> Co-authored-by: L.B.R. <lbr@mmonad.com>

L.B.R. committed Mar 20, 2026 at 15:11 UTC 1779c09898e091167b5d29cb8931ce1f5dea9a47
4 files changed +365 -99
csrc/rocm/skinny_gemms.cu
+268 -92
@@ -26,6 +26,16 @@
26 #define __HIP__GFX9__
27 #endif
28
29 +#if defined(__HIPCC__) && \
30 + (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
31 + defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
32 + #define __HIP__GFX1X__
33 +#endif
34 +
35 +#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
36 + #define __HIP__GFX12__
37 +#endif
38 +
39 #if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
40 #define __HIP__MI3XX__
41 #endif
@@ -37,15 +47,31 @@
47 #endif
48
49 int get_lds_size() {
40 - static bool is_cached = false;
41 - static int result;
42 - if (is_cached == false) {
43 - auto dprops = at::cuda::getCurrentDeviceProperties();
44 - std::string device_arch = dprops->gcnArchName;
45 - size_t substring = device_arch.find("gfx95");
46 - result = (substring == std::string::npos ? 64 * 1024 : 160 * 1024);
47 - is_cached = true;
48 - }
50 + static const int result = [] {
51 + const auto* dprops = at::cuda::getCurrentDeviceProperties();
52 + const std::string device_arch = dprops->gcnArchName;
53 + return device_arch.find("gfx95") == std::string::npos ? 64 * 1024
54 + : 160 * 1024;
55 + }();
56 + return result;
57 +}
58 +
59 +bool on_gfx1x() {
60 + static const bool result = [] {
61 + const auto* dprops = at::cuda::getCurrentDeviceProperties();
62 + const std::string device_arch = dprops->gcnArchName;
63 + return device_arch.find("gfx11") != std::string::npos ||
64 + device_arch.find("gfx12") != std::string::npos;
65 + }();
66 + return result;
67 +}
68 +
69 +bool on_gfx12() {
70 + static const bool result = [] {
71 + const auto* dprops = at::cuda::getCurrentDeviceProperties();
72 + const std::string device_arch = dprops->gcnArchName;
73 + return device_arch.find("gfx12") != std::string::npos;
74 + }();
75 return result;
76 }
77
@@ -286,21 +312,35 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
312 return out_c;
313 }
314
289 -#define DOT2C(V0, V2, V3) \
290 - if constexpr (std::is_same_v<scalar_t, half>) { \
291 - asm("v_dot2c_f32_f16 %0, %2, %3" : "=v"(V0) : "0"(V0), "v"(V2), "v"(V3)); \
292 - } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
293 - float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \
294 - __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \
295 - V0 += (s.x + s.y); \
296 - }
315 +#if defined(__HIP__GFX9__) && !defined(__HIP__GFX1X__)
316 + #define DOT2C(V0, V2, V3) \
317 + if constexpr (std::is_same_v<scalar_t, half>) { \
318 + asm("v_dot2c_f32_f16 %0, %2, %3" \
319 + : "=v"(V0) \
320 + : "0"(V0), "v"(V2), "v"(V3)); \
321 + } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
322 + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \
323 + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \
324 + V0 += (s.x + s.y); \
325 + }
326 +#elif defined(__HIP__GFX1X__)
327 + // gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12)
328 + #define DOT2C(V0, V2, V3) \
329 + if constexpr (std::is_same_v<scalar_t, half>) { \
330 + asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \
331 + } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
332 + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \
333 + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \
334 + V0 += (s.x + s.y); \
335 + }
336 +#endif
337
338 // To avoid LLVM silently upcasting to double
339 __device__ inline unsigned int min__(uint32_t a, uint32_t b) {
340 return min(a, b);
341 }
342
303 -#if defined(__HIP__GFX9__) // TODO: Add NAVI support
343 +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__)
344 // This version targets cases where A[] fits LDS capacity
345 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
346 int UNRL, int N>
@@ -442,14 +482,18 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
482 1); // row_shr2
483 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
484 1); // row_shr1
485 + #if defined(__HIP__GFX9__)
486 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
487 1); // ROW_BCAST15
488 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
489 1); // ROW_BCAST31
490 + #else
491 + sum[n][y] += __shfl_xor(sum[n][y], 16);
492 + #endif
493 }
494 }
495
452 - if (threadIdx.x == 63) {
496 + if (threadIdx.x == (THRDS - 1)) {
497 scalar_t biases[N][YTILE] = {};
498 if (BIAS)
499 for (int n = 0; n < N; n++) {
@@ -469,9 +513,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
513 }
514 }
515 } else {
472 - #pragma unroll
516 + #ifdef __HIP__GFX9__
517 + #pragma unroll
518 for (int n = 0; n < N; n++) {
474 - #pragma unroll
519 + #pragma unroll
520 for (int y = 0; y < YTILE; y++) {
521 /*float accm1 = 0;
522 for (int i=0; i<64; i++)
@@ -498,7 +543,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
543 sum4[n][y][0] = accm;
544 }
545 }
501 - if (threadIdx.x == 63) {
546 + if (threadIdx.x == (THRDS - 1)) {
547 scalar_t biases[N][YTILE] = {};
548 if (BIAS)
549 for (int n = 0; n < N; n++) {
@@ -513,11 +558,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
558 }
559 }
560 }
561 + #endif // __HIP__GFX9__ (MFMA path)
562 }
563 m += CuCount * _WvPrGrp * YTILE;
564 }
565 }
520 -#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
566 +#else
567 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
568 int UNRL, int N>
569 __global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap,
@@ -528,9 +574,9 @@ __global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap,
574 const int _WvPrGrp, const int CuCount) {
575 UNREACHABLE_CODE
576 }
531 -#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
577 +#endif
578
533 -#if defined(__HIP__GFX9__) // TODO: Add NAVI support
579 +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__)
580 // This version targets cases where A[] marginally exceeds LDS capacity
581 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
582 int UNRL, int N>
@@ -657,14 +703,18 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
703 1); // row_shr2
704 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
705 1); // row_shr1
706 + #if defined(__HIP__GFX9__)
707 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
708 1); // ROW_BCAST15
709 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
710 1); // ROW_BCAST31
711 + #else
712 + sum[n][y] += __shfl_xor(sum[n][y], 16);
713 + #endif
714 }
715 }
716
667 - if (threadIdx.x == 63) {
717 + if (threadIdx.x == (THRDS - 1)) {
718 scalar_t biases[N][YTILE] = {};
719 if (BIAS)
720 for (int n = 0; n < N; n++) {
@@ -686,9 +736,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
736 }
737 }
738 } else {
689 - #pragma unroll
739 + #ifdef __HIP__GFX9__
740 + #pragma unroll
741 for (int n = 0; n < N; n++) {
691 - #pragma unroll
742 + #pragma unroll
743 for (int y = 0; y < YTILE; y++) {
744 // float accm1 = 0;
745 // for (int i=0; i<64; i++)
@@ -713,7 +764,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
764 sum4[n][y][0] = accm;
765 }
766 }
716 - if (threadIdx.x == 63) {
767 + if (threadIdx.x == (THRDS - 1)) {
768 scalar_t biases[N][YTILE] = {};
769 if (BIAS)
770 for (int n = 0; n < N; n++) {
@@ -730,6 +781,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
781 }
782 }
783 }
784 + #endif // __HIP__GFX9__ (MFMA path)
785 }
786
787 m += CuCount * _WvPrGrp * YTILE;
@@ -746,7 +798,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
798 }
799 }
800
749 -#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
801 +#else
802 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
803 int UNRL, int N>
804 __global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap,
@@ -756,9 +808,9 @@ __global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap,
808 const int _WvPrGrp, const int CuCount) {
809 UNREACHABLE_CODE
810 }
759 -#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
811 +#endif
812
761 -#if defined(__HIP__GFX9__) // TODO: Add NAVI support
813 +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__)
814 // This version targets big A[] cases, where it is much larger than LDS capacity
815 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
816 int UNRL, int N>
@@ -1004,14 +1056,18 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1056 1); // row_shr2
1057 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
1058 1); // row_shr1
1059 + #if defined(__HIP__GFX9__)
1060 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
1061 1); // ROW_BCAST15
1062 sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
1063 1); // ROW_BCAST31
1064 + #else
1065 + sum[n][y] += __shfl_xor(sum[n][y], 16);
1066 + #endif
1067 }
1068 }
1069
1014 - if (threadIdx.x == 63) {
1070 + if (threadIdx.x == (THRDS - 1)) {
1071 scalar_t biases[N][YTILE] = {};
1072 if (BIAS)
1073 for (int n = 0; n < N; n++) {
@@ -1033,9 +1089,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1089 }
1090 }
1091 } else {
1036 - #pragma unroll
1092 + #ifdef __HIP__GFX9__
1093 + #pragma unroll
1094 for (int n = 0; n < N; n++) {
1038 - #pragma unroll
1095 + #pragma unroll
1096 for (int y = 0; y < YTILE; y++) {
1097 float accm = sum4[n][y][0];
1098 accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
@@ -1057,7 +1114,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1114 sum4[n][y][0] = accm;
1115 }
1116 }
1060 - if (threadIdx.x == 63) {
1117 + if (threadIdx.x == (THRDS - 1)) {
1118 scalar_t biases[N][YTILE] = {};
1119 if (BIAS)
1120 for (int n = 0; n < N; n++) {
@@ -1074,6 +1131,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1131 }
1132 }
1133 }
1134 + #endif // __HIP__GFX9__ (MFMA path)
1135 }
1136
1137 m += CuCount * _WvPrGrp * YTILE;
@@ -1090,7 +1148,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1148 }
1149 }
1150 }
1093 -#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1151 +#else
1152 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1153 int UNRL, int N>
1154 __global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap,
@@ -1101,7 +1159,7 @@ __global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap,
1159 const int _WvPrGrp, const int CuCount) {
1160 UNREACHABLE_CODE
1161 }
1104 -#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1162 +#endif
1163
1164 // Find the min val of div2 that doesn't increase N/(div1*div2)
1165 int mindiv(int N, int div1, int div2) {
@@ -1148,40 +1206,40 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1206 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1207 const int max_lds_len = get_lds_size() / 2;
1208
1151 -#define WVSPLITK(_YTILE, _UNRL, _N) \
1209 +#define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N) \
1210 { \
1153 - dim3 block(64, 16); \
1154 - int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, 16); \
1211 + dim3 block(_THRDS, _WVPRGRP); \
1212 + int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, _WVPRGRP); \
1213 if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) \
1156 - wvSplitK_hf_sml_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1214 + wvSplitK_hf_sml_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N> \
1215 <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1216 By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1217 CuCount); \
1218 else if (Kbp_in * N_in <= max_lds_len * 1.2) \
1161 - wvSplitK_hf_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1219 + wvSplitK_hf_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N> \
1220 <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1221 By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1222 CuCount); \
1223 else \
1166 - wvSplitK_hf_big_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1224 + wvSplitK_hf_big_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N> \
1225 <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1226 By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1227 CuCount); \
1228 }
1229
1172 -#define WVSPLIT_TILE(_sYT, __N) \
1230 +#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \
1231 { \
1232 bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
1233 if (_sYT <= 1) \
1176 - WVSPLITK(1, 4, __N) \
1234 + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \
1235 else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \
1178 - WVSPLITK(2, 2, __N) \
1236 + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \
1237 else if (_sYT <= 4 * 3) \
1180 - WVSPLITK(3, 2, __N) \
1238 + WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \
1239 else if (__N == 4) \
1182 - WVSPLITK(4, 1, __N) \
1240 + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \
1241 else \
1184 - WVSPLITK(4, 2, __N) \
1242 + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \
1243 }
1244
1245 AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
@@ -1198,18 +1256,31 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1256 // then cut the active waves to balance their distribution...
1257 int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4);
1258
1259 + const bool use_wave32 = on_gfx1x();
1260 switch (N_in) {
1261 case 1:
1203 - WVSPLIT_TILE(sYT, 1)
1262 + if (use_wave32)
1263 + WVSPLIT_TILE_CFG(32, 16, sYT, 1)
1264 + else
1265 + WVSPLIT_TILE_CFG(64, 16, sYT, 1)
1266 break;
1267 case 2:
1206 - WVSPLIT_TILE(sYT, 2)
1268 + if (use_wave32)
1269 + WVSPLIT_TILE_CFG(32, 16, sYT, 2)
1270 + else
1271 + WVSPLIT_TILE_CFG(64, 16, sYT, 2)
1272 break;
1273 case 3:
1209 - WVSPLIT_TILE(sYT, 3)
1274 + if (use_wave32)
1275 + WVSPLIT_TILE_CFG(32, 16, sYT, 3)
1276 + else
1277 + WVSPLIT_TILE_CFG(64, 16, sYT, 3)
1278 break;
1279 case 4:
1212 - WVSPLIT_TILE(sYT, 4)
1280 + if (use_wave32)
1281 + WVSPLIT_TILE_CFG(32, 16, sYT, 4)
1282 + else
1283 + WVSPLIT_TILE_CFG(64, 16, sYT, 4)
1284 break;
1285 default:
1286 throw std::runtime_error(
@@ -1653,7 +1724,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1724 #endif
1725 }
1726 }
1656 -#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1727 +#else
1728 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1729 int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
1730 __global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
@@ -1688,6 +1759,8 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1759 TORCH_CHECK(in_a.dtype() == torch::kFloat16 ||
1760 in_a.dtype() == torch::kBFloat16);
1761
1762 + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1763 +
1764 auto out_c = torch::empty(
1765 {N_in, M_in},
1766 torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device()));
@@ -1696,7 +1769,6 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1769
1770 dim3 grid(CuCount);
1771
1699 - const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1772 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1773 // const int max_lds_len = get_lds_size() / 2;
1774
@@ -1773,7 +1845,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1845 return out_c;
1846 }
1847
1776 -#if defined(__HIP__MI3XX__) // TODO: Add NAVI support
1848 +#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
1849 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1850 int A_CHUNK, int UNRL, int N>
1851 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1817,12 +1889,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1889
1890 uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
1891
1820 - using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
1892 float sA = *s_A;
1893 float sB = *s_B;
1894
1895 while (m < M) {
1896 + #ifdef __HIP__GFX12__
1897 + // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
1898 + float sum[N][YTILE] = {};
1899 + #else
1900 + // gfx9: MFMA accumulation
1901 scalar8 sum[N][YTILE] = {};
1902 + #endif
1903 for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
1904 bigType bigA[N][UNRL] = {};
1905 bigType bigB[YTILE][UNRL];
@@ -1854,6 +1931,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1931 #pragma unroll
1932 for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1933 for (uint32_t n = 0; n < N; n++) {
1934 + #ifdef __HIP__GFX12__
1935 + // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
1936 + for (int y = 0; y < YTILE; ++y) {
1937 + #pragma unroll
1938 + for (int i = 0; i < A_CHUNK / 4; i++) {
1939 + sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8(
1940 + bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]);
1941 + }
1942 + }
1943 + #else
1944 + // gfx9: MFMA path
1945 for (int i = 0; i < A_CHUNK; i += 8) {
1946 for (int y = 0; y < YTILE; ++y) {
1947 sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
@@ -1861,11 +1949,33 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1949 0);
1950 }
1951 }
1952 + #endif
1953 }
1954 }
1955 }
1956
1957 // Final reduction
1958 + #ifdef __HIP__GFX12__
1959 + // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
1960 + for (int n = 0; n < N; n++) {
1961 + for (int y = 0; y < YTILE; y++) {
1962 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
1963 + : "=v"(sum[n][y])
1964 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1965 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
1966 + : "=v"(sum[n][y])
1967 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1968 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
1969 + : "=v"(sum[n][y])
1970 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1971 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 "
1972 + : "=v"(sum[n][y])
1973 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1974 + sum[n][y] += __shfl_xor(sum[n][y], 16);
1975 + }
1976 + }
1977 + #else
1978 + // gfx9 MFMA reduction
1979 for (int n = 0; n < N; n++) {
1980 for (int y = 0; y < YTILE; y++) {
1981 float accm0 = sum[n][y][0];
@@ -1880,8 +1990,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1990 sum[n][y][0] = accm0;
1991 }
1992 }
1993 + #endif
1994
1884 - if (threadIdx.x == 0) {
1995 + const bool writeback_lane =
1996 + #ifdef __HIP__GFX12__
1997 + threadIdx.x == (THRDS - 1);
1998 + #else
1999 + threadIdx.x == 0;
2000 + #endif
2001 + if (writeback_lane) {
2002 scalar_t biases[N][YTILE] = {};
2003 if (BIAS)
2004 for (int n = 0; n < N; n++) {
@@ -1892,13 +2009,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2009 for (int n = 0; n < N; n++) {
2010 for (int y = 0; y < YTILE; y++) {
2011 if (y + m >= M) break; // To avoid mem access fault.
1895 - sum[n][y][0] *= sA * sB;
2012 + #ifdef __HIP__GFX12__
2013 + float result = sum[n][y] * sA * sB;
2014 + #else
2015 + float result = sum[n][y][0] * sA * sB;
2016 + #endif
2017 if constexpr (std::is_same_v<scalar_t, half>) {
1897 - sum[n][y][0] += __half2float(biases[n][y]);
2018 + result += __half2float(biases[n][y]);
2019 } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1899 - sum[n][y][0] += __bfloat162float(biases[n][y]);
2020 + result += __bfloat162float(biases[n][y]);
2021 }
1901 - C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
2022 + C[m + y + n * M] = __float2s<scalar_t>(result);
2023 }
2024 }
2025 }
@@ -1906,7 +2027,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2027 m += CuCount * _WvPrGrp * YTILE;
2028 }
2029 }
1909 -#else // !defined(__HIP__MI3XX__) TODO: Add NAVI support
2030 +#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
2031 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2032 int A_CHUNK, int UNRL, int N>
2033 __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
@@ -1918,9 +2039,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
2039 const int _WvPrGrp, const int CuCount) {
2040 UNREACHABLE_CODE
2041 }
1921 -#endif // defined(__HIP__MI3XX__) TODO: Add NAVI support
2042 +#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
2043
1923 -#if defined(__HIP__MI3XX__) // TODO: Add NAVI support
2044 +#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
2045 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2046 int A_CHUNK, int UNRL, int N>
2047 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1963,12 +2084,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2084
2085 uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
2086
1966 - using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
2087 float sA = *s_A;
2088 float sB = *s_B;
2089
2090 while (m < M) {
2091 + #ifdef __HIP__GFX12__
2092 + // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
2093 + float sum[N][YTILE] = {};
2094 + #else
2095 + // gfx9: MFMA accumulation
2096 scalar8 sum[N][YTILE] = {};
2097 + #endif
2098 for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
2099 bigType bigA[N][UNRL] = {};
2100 bigType bigB[YTILE][UNRL];
@@ -2002,6 +2128,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2128 #pragma unroll
2129 for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2130 for (uint32_t n = 0; n < N; n++) {
2131 + #ifdef __HIP__GFX12__
2132 + // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
2133 + for (int y = 0; y < YTILE; ++y) {
2134 + #pragma unroll
2135 + for (int i = 0; i < A_CHUNK / 4; i++) {
2136 + sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8(
2137 + bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]);
2138 + }
2139 + }
2140 + #else
2141 + // gfx9: MFMA path
2142 for (int i = 0; i < A_CHUNK; i += 8) {
2143 for (int y = 0; y < YTILE; ++y) {
2144 sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
@@ -2009,11 +2146,33 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2146 0);
2147 }
2148 }
2149 + #endif
2150 }
2151 }
2152 }
2153
2154 // Final reduction
2155 + #ifdef __HIP__GFX12__
2156 + // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
2157 + for (int n = 0; n < N; n++) {
2158 + for (int y = 0; y < YTILE; y++) {
2159 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
2160 + : "=v"(sum[n][y])
2161 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2162 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
2163 + : "=v"(sum[n][y])
2164 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2165 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
2166 + : "=v"(sum[n][y])
2167 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2168 + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 "
2169 + : "=v"(sum[n][y])
2170 + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2171 + sum[n][y] += __shfl_xor(sum[n][y], 16);
2172 + }
2173 + }
2174 + #else
2175 + // gfx9 MFMA reduction
2176 for (int n = 0; n < N; n++) {
2177 for (int y = 0; y < YTILE; y++) {
2178 float accm0 = sum[n][y][0];
@@ -2028,8 +2187,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2187 sum[n][y][0] = accm0;
2188 }
2189 }
2190 + #endif
2191
2032 - if (threadIdx.x == 0) {
2192 + const bool writeback_lane =
2193 + #ifdef __HIP__GFX12__
2194 + threadIdx.x == (THRDS - 1);
2195 + #else
2196 + threadIdx.x == 0;
2197 + #endif
2198 + if (writeback_lane) {
2199 scalar_t biases[N][YTILE] = {};
2200 if (BIAS)
2201 for (int n = 0; n < N; n++) {
@@ -2040,13 +2206,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2206 for (int n = 0; n < N; n++) {
2207 for (int y = 0; y < YTILE; y++) {
2208 if (y + m >= M) break; // To avoid mem access fault.
2043 - sum[n][y][0] *= sA * sB;
2209 + #ifdef __HIP__GFX12__
2210 + float result = sum[n][y] * sA * sB;
2211 + #else
2212 + float result = sum[n][y][0] * sA * sB;
2213 + #endif
2214 if constexpr (std::is_same_v<scalar_t, half>) {
2045 - sum[n][y][0] += __half2float(biases[n][y]);
2215 + result += __half2float(biases[n][y]);
2216 } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
2047 - sum[n][y][0] += __bfloat162float(biases[n][y]);
2217 + result += __bfloat162float(biases[n][y]);
2218 }
2049 - C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
2219 + C[m + y + n * M] = __float2s<scalar_t>(result);
2220 }
2221 }
2222 }
@@ -2054,7 +2224,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2224 m += CuCount * _WvPrGrp * YTILE;
2225 }
2226 }
2057 -#else // !defined(__HIP__MI3XX__) TODO: Add NAVI support
2227 +#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
2228 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2229 int A_CHUNK, int UNRL, int N>
2230 __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
@@ -2066,7 +2236,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
2236 const int CuCount) {
2237 UNREACHABLE_CODE
2238 }
2069 -#endif // defined(__HIP__MI3XX__) TODO: Add NAVI support
2239 +#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
2240
2241 void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
2242 const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
@@ -2099,24 +2269,30 @@ void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
2269 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
2270 const int max_lds_len = get_lds_size();
2271
2102 -#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2103 - { \
2104 - dim3 block(64, _WvPrGrp); \
2105 - if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
2106 - int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16)); \
2107 - wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
2108 - <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
2109 - By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \
2110 - s_a, s_b, __wvPrGrp, CuCount); \
2111 - } else { \
2112 - int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16)); \
2113 - wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
2114 - <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
2115 - By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \
2116 - s_a, s_b, __wvPrGrp, CuCount); \
2117 - } \
2272 +#define WVSPLITKQ_IMPL(_THRDS, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2273 + { \
2274 + dim3 block(_THRDS, _WvPrGrp); \
2275 + if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
2276 + int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16)); \
2277 + wvSplitKQ_hf_sml_<fptype, fp8_t, _THRDS, _YTILEs, _WvPrGrp, 16, _UNRLs, \
2278 + _N><<<grid, block, 0, stream>>>( \
2279 + K_in, Kap_in, Kbp_in, M_in, Bx_in, By_in, b_ptr, a_ptr, bias_ptr, \
2280 + c_ptr, s_a, s_b, __wvPrGrp, CuCount); \
2281 + } else { \
2282 + int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16)); \
2283 + wvSplitKQ_hf_<fptype, fp8_t, _THRDS, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
2284 + <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
2285 + By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \
2286 + s_a, s_b, __wvPrGrp, CuCount); \
2287 + } \
2288 }
2289
2290 +#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2291 + if (on_gfx12()) \
2292 + WVSPLITKQ_IMPL(32, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2293 + else \
2294 + WVSPLITKQ_IMPL(64, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N)
2295 +
2296 AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] {
2297 using fptype = typename scalar<scalar_t>::type;
2298 auto c_ptr = reinterpret_cast<fptype*>(out_c.data_ptr());
@@ -2136,10 +2312,10 @@ void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
2312 WVSPLITKQ(16, 2, 2, 2, 2, 2)
2313 break;
2314 case 3:
2139 - WVSPLITKQ(16, 2, 2, 2, 2, 3)
2315 + WVSPLITKQ(16, 2, 2, 1, 1, 3)
2316 break;
2317 case 4:
2142 - WVSPLITKQ(16, 2, 2, 2, 2, 4)
2318 + WVSPLITKQ(16, 2, 2, 1, 1, 4)
2319 break;
2320 default:
2321 throw std::runtime_error(
tests/kernels/quantization/test_rocm_skinny_gemms.py
+5 -4
@@ -160,6 +160,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode
160 BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
161 elif bias_mode == 2:
162 BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
163 + elif bias_mode == 3:
164 + BIAS = torch.rand(1, m, dtype=dtype, device="cuda") * 2 - 1
165
166 ref_out = torch.nn.functional.linear(A, B, BIAS)
167 out = ops.wvSplitKrc(A, B, cu_count, BIAS)
@@ -224,10 +226,9 @@ def test_rocm_wvsplitk_kernel(
226 ref_out = torch.nn.functional.linear(A, B, BIAS)
227 out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
228
227 - if xnorm:
228 - assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
229 - else:
230 - assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
229 + # Accumulation error in fp16 GEMM scales with sqrt(K)
230 + atol = torch.finfo(dtype).eps * math.sqrt(k)
231 + torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-2)
232
233
234 @pytest.mark.parametrize("xnorm", [False, True])
tests/model_executor/layers/test_rocm_unquantized_gemm.py new
+89
@@ -0,0 +1,89 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from unittest.mock import MagicMock
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.platforms import current_platform
10 +
11 +if current_platform.is_cuda():
12 + pytest.skip(
13 + "ROCm skinny GEMM tests are not supported on CUDA.",
14 + allow_module_level=True,
15 + )
16 +
17 +from vllm.model_executor.layers import utils
18 +
19 +
20 +def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
21 + x = torch.randn(1, 64, dtype=torch.float16)
22 + weight = torch.randn(128, 64, dtype=torch.float16)
23 +
24 + monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
25 + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
26 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
27 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
28 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
29 + monkeypatch.setattr(utils, "get_cu_count", lambda: 120)
30 +
31 + wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
32 + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
33 + llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
34 + monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
35 +
36 + out = utils.rocm_unquantized_gemm_impl(x, weight, None)
37 + ref = torch.nn.functional.linear(x, weight, None)
38 +
39 + wvsplitk_mock.assert_called_once()
40 + llmm1_mock.assert_not_called()
41 + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
42 +
43 +
44 +def test_rocm_unquantized_gemm_gfx1x_n_gt_4_falls_back(monkeypatch):
45 + x = torch.randn(5, 64, dtype=torch.float16)
46 + weight = torch.randn(128, 64, dtype=torch.float16)
47 +
48 + monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
49 + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
50 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
51 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
52 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
53 + monkeypatch.setattr(utils, "get_cu_count", lambda: 120)
54 +
55 + wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
56 + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
57 + llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
58 + monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
59 +
60 + out = utils.rocm_unquantized_gemm_impl(x, weight, None)
61 + ref = torch.nn.functional.linear(x, weight, None)
62 +
63 + wvsplitk_mock.assert_not_called()
64 + llmm1_mock.assert_not_called()
65 + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
66 +
67 +
68 +def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
69 + x = torch.randn(16, 1024, dtype=torch.float16)
70 + weight = torch.randn(256, 1024, dtype=torch.float16)
71 +
72 + monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
73 + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
74 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
75 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
76 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
77 + monkeypatch.setattr(utils, "get_cu_count", lambda: 120)
78 +
79 + wvsplitkrc_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
80 + monkeypatch.setattr(utils.ops, "wvSplitKrc", wvsplitkrc_mock)
81 + wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
82 + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
83 +
84 + out = utils.rocm_unquantized_gemm_impl(x, weight, None)
85 + ref = torch.nn.functional.linear(x, weight, None)
86 +
87 + wvsplitkrc_mock.assert_called_once()
88 + wvsplitk_mock.assert_not_called()
89 + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
vllm/model_executor/layers/utils.py
+3 -3
@@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype):
122 def rocm_unquantized_gemm_impl(
123 x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
124 ) -> torch.Tensor:
125 - from vllm.platforms.rocm import on_gfx9, on_gfx950
125 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950
126
127 n = x.numel() // x.size(-1)
128 m = weight.shape[0]
@@ -169,12 +169,12 @@ def rocm_unquantized_gemm_impl(
169
170 use_skinny = (
171 envs.VLLM_ROCM_USE_SKINNY_GEMM
172 - and on_gfx9()
172 + and (on_gfx9() or on_gfx1x())
173 and x.dtype in [torch.float16, torch.bfloat16]
174 and k % 8 == 0
175 )
176
177 - if use_skinny is not True:
177 + if not use_skinny:
178 return torch.nn.functional.linear(x, weight, bias)
179
180 x_view = x.reshape(-1, x.size(-1))