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 15, 07:15 UTC
[ROCm][Kernel] Add gfx950 support for skinny gemms (#18010)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
May 31, 2025 at 09:40 UTC
306d60401dbd066f64298e02ca73d4f2075d7bf6
5 files changed
+91
-54
csrc/rocm/skinny_gemms.cu
+70
-43
@@ -13,14 +13,34 @@
13
#include "dispatch_utils.h"
14
#include "quantization/fp8/common.cuh"
15
16
-#if defined(__HIPCC__) && (defined(__gfx90a__) || defined(__gfx942__))
17
- #define __HIP__MI300_MI250__
16
+#if defined(__HIPCC__) && \
17
+ (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
18
+ #define __HIP__GFX9__
19
#endif
20
20
-#if defined(__HIPCC__) && defined(__gfx942__)
21
- #define __HIP__MI300__
21
+#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
22
+ #define __HIP__MI3XX__
23
#endif
24
25
+#if defined(__gfx950__)
26
+ #define LDS_SIZE 160 * 1024
27
+#else
28
+ #define LDS_SIZE 64 * 1024
29
+#endif
30
+
31
+int get_lds_size() {
32
+ static bool is_cached = false;
33
+ static int result;
34
+ if (is_cached == false) {
35
+ auto dprops = at::cuda::getCurrentDeviceProperties();
36
+ std::string device_arch = dprops->gcnArchName;
37
+ size_t substring = device_arch.find("gfx95");
38
+ result = (substring == std::string::npos ? 64 * 1024 : 160 * 1024);
39
+ is_cached = true;
40
+ }
41
+ return result;
42
+}
43
+
44
#if defined(NDEBUG)
45
#undef NDEBUG
46
#include <assert.h>
@@ -267,7 +287,7 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
287
V0 += (s.x + s.y); \
288
}
289
270
-#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
290
+#if defined(__HIP__GFX9__) // TODO: Add NAVI support
291
// This version targets cases where A[] fits LDS capacity
292
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
293
int UNRL, int N>
@@ -275,7 +295,8 @@ __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,
297
const int _WvPrGrp, const int CuCount) {
278
- #if defined(__HIP__MI300__)
298
+ constexpr int max_lds_len = LDS_SIZE / 2;
299
+ #if defined(__HIP__MI3XX__)
300
constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
301
#else
302
constexpr bool use_mfma = false;
@@ -295,13 +316,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
316
};
317
318
//----------------------------------------------------
298
- // Reserving 64 KB of LDS to have 1 WG / CU
319
+ // Reserving 64/160 KB of LDS to have 1 WG / CU
320
// Goal is to bring the activation matrix A to the LDS
321
// and use it across the lifetime of the work group
322
// TODO: When activation matrix is larger than 64 KB
323
// then this is not goint to work!
324
//----------------------------------------------------
304
- __shared__ scalar_t s[1024 * 32];
325
+ __shared__ scalar_t s[max_lds_len];
326
327
//----------------------------------------------------
328
// Fetch the activation matrix to LDS
@@ -312,11 +333,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
333
// - Then the WG will move to another 8 K elements
334
// TODO: Logic below will only work when K is multiple of 8
335
//----------------------------------------------------
315
- for (uint32_t k = 0; k < min(K * N, 32 * 1024);
336
+ for (uint32_t k = 0; k < min(K * N, max_lds_len);
337
k += THRDS * WvPrGrp * A_CHUNK) {
338
uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
339
319
- if (k_in >= min(K * N, 32 * 1024)) break;
340
+ if (k_in >= min(K * N, max_lds_len)) break;
341
342
*((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
343
}
@@ -517,7 +538,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
538
m += CuCount * _WvPrGrp * YTILE;
539
}
540
}
520
-#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
541
+#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
542
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
543
int UNRL, int N>
544
__global__ void wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
@@ -525,9 +546,9 @@ __global__ void wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
546
const int _WvPrGrp, const int CuCount) {
547
UNREACHABLE_CODE
548
}
528
-#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
549
+#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
550
530
-#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
551
+#if defined(__HIP__GFX9__) // TODO: Add NAVI support
552
// This version targets cases where A[] marginally exceeds LDS capacity
553
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
554
int UNRL, int N>
@@ -535,7 +556,8 @@ __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,
558
const int _WvPrGrp, const int CuCount) {
538
- #if defined(__HIP__MI300__)
559
+ constexpr int max_lds_len = LDS_SIZE / 2;
560
+ #if defined(__HIP__MI3XX__)
561
constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
562
#else
563
constexpr bool use_mfma = false;
@@ -561,7 +583,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
583
// TODO: When activation matrix is larger than 64 KB
584
// then this is not goint to work!
585
//----------------------------------------------------
564
- __shared__ scalar_t s[1024 * 32];
586
+ __shared__ scalar_t s[max_lds_len];
587
588
//----------------------------------------------------
589
// Computation of columns that need to be committed to memory!
@@ -598,11 +620,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
620
// - Then the WG will move to another 8 K elements
621
// TODO: Logic below will only work when K is multiple of 8
622
//----------------------------------------------------
601
- for (uint32_t k = 0; k < min(K * N, 32 * 1024);
623
+ for (uint32_t k = 0; k < min(K * N, max_lds_len);
624
k += THRDS * WvPrGrp * A_CHUNK) {
625
uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
626
605
- if (k_in >= min(K * N, 32 * 1024)) break;
627
+ if (k_in >= min(K * N, max_lds_len)) break;
628
629
*((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
630
}
@@ -686,7 +708,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
708
// Fetch A activation matrix in interleaved fashion from LDS or memory
709
710
for (int n = 0; n < N; n++) {
689
- if (k_ + K * n < 32 * 1024)
711
+ if (k_ + K * n < max_lds_len)
712
bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
713
else
714
bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
@@ -817,7 +839,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
839
}
840
}
841
820
-#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
842
+#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
843
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
844
int UNRL, int N>
845
__global__ void wvSplitK_hf_(const int K, const int M, const scalar_t* B,
@@ -825,9 +847,9 @@ __global__ void wvSplitK_hf_(const int K, const int M, const scalar_t* B,
847
const int _WvPrGrp, const int CuCount) {
848
UNREACHABLE_CODE
849
}
828
-#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
850
+#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
851
830
-#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
852
+#if defined(__HIP__GFX9__) // TODO: Add NAVI support
853
// This version targets big A[] cases, where it is much larger than LDS capacity
854
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
855
int UNRL, int N>
@@ -835,7 +857,8 @@ __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,
859
const int _WvPrGrp, const int CuCount) {
838
- #if defined(__HIP__MI300__)
860
+ constexpr int max_lds_len = LDS_SIZE / 2;
861
+ #if defined(__HIP__MI3XX__)
862
constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
863
#else
864
constexpr bool use_mfma = false;
@@ -855,13 +878,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
878
};
879
880
//----------------------------------------------------
858
- // Reserving 64 KB of LDS to have 1 WG / CU
881
+ // Reserving 64/160 KB of LDS to have 1 WG / CU
882
// Goal is to bring the activation matrix A to the LDS
883
// and use it across the lifetime of the work group
884
// TODO: When activation matrix is larger than 64 KB
885
// then this is not goint to work!
886
//----------------------------------------------------
864
- __shared__ scalar_t s[1024 * 32];
887
+ __shared__ scalar_t s[max_lds_len];
888
889
//----------------------------------------------------
890
// Computation of columns that need to be committed to memory!
@@ -902,11 +925,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
925
//----------------------------------------------------
926
#define PCML
927
#ifndef PCML
905
- for (uint32_t k = 0; k < min(K * N, 32 * 1024);
928
+ for (uint32_t k = 0; k < min(K * N, max_lds_len);
929
k += THRDS * WvPrGrp * A_CHUNK) {
930
uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
931
909
- if (k_in >= min(K * N, 32 * 1024)) break;
932
+ if (k_in >= min(K * N, max_lds_len)) break;
933
934
*((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
935
}
@@ -916,7 +939,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
939
#define TUC (THRDS * UNRL * A_CHUNK)
940
uint32_t kBase = 0;
941
// find biggest k size that fits in LDS
919
- uint32_t kFit = (32 * 1024) / N;
942
+ uint32_t kFit = (max_lds_len) / N;
943
// kFit = (kFit%TWC==0) ? kFit : (kFit-kFit%TWC+TWC); //round up to multiple
944
// of TUC
945
kFit = (kFit % TUC == 0)
@@ -1164,7 +1187,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1187
}
1188
}
1189
}
1167
-#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1190
+#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1191
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1192
int UNRL, int N>
1193
__global__ void wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
@@ -1172,7 +1195,7 @@ __global__ void wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
1195
const int _WvPrGrp, const int CuCount) {
1196
UNREACHABLE_CODE
1197
}
1175
-#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1198
+#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1199
1200
int mindiv(int N, int div1, int div2) {
1201
int nPrRnd = div1 * div2;
@@ -1222,17 +1245,18 @@ torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1245
1246
const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1247
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1248
+ const int max_lds_len = get_lds_size() / 2;
1249
1250
#define WVSPLITK(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \
1251
_N) \
1252
{ \
1253
dim3 block(64, _WvPrGrp); \
1230
- if ((K_in * N_in <= 32 * 1024) && (M_in % _YTILEs == 0)) { \
1254
+ if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
1255
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
1256
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); \
1235
- } else if (K_in * N_in <= 32 * 1024 * 1.2) { \
1259
+ } else if (K_in * N_in <= max_lds_len * 1.2) { \
1260
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \
1261
wvSplitK_hf_<fptype, 64, _YTILEm, _WvPrGrp, 8, _UNRLm, _N> \
1262
<<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
@@ -1272,7 +1296,7 @@ torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1296
return out_c;
1297
}
1298
1275
-#if defined(__HIP__MI300__) // TODO: Add NAVI support
1299
+#if defined(__HIP__MI3XX__) // TODO: Add NAVI support
1300
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1301
int A_CHUNK, int UNRL, int N>
1302
__global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1281,6 +1305,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1305
const float* __restrict__ s_A,
1306
const float* __restrict__ s_B, const int _WvPrGrp,
1307
const int CuCount) {
1308
+ constexpr int max_lds_len = LDS_SIZE;
1309
using scalar8 =
1310
__attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float;
1311
using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int;
@@ -1296,10 +1321,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1321
scalar8 h8;
1322
};
1323
1299
- __shared__ fp8_t s[1024 * 64];
1324
+ __shared__ fp8_t s[max_lds_len];
1325
1326
for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1302
- k < min(K * N, 64 * 1024); k += THRDS * WvPrGrp * A_CHUNK) {
1327
+ k < min(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
1328
*((bigType*)(&s[k])) = *((bigType*)(&A[k]));
1329
}
1330
__syncthreads();
@@ -1436,7 +1461,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1461
m += CuCount * _WvPrGrp * YTILE;
1462
}
1463
}
1439
-#else // !defined(__HIP__MI300__) TODO: Add NAVI support
1464
+#else // !defined(__HIP__MI3XX__) TODO: Add NAVI support
1465
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1466
int A_CHUNK, int UNRL, int N>
1467
__global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
@@ -1446,9 +1471,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
1471
const int _WvPrGrp, const int CuCount) {
1472
UNREACHABLE_CODE
1473
}
1449
-#endif // defined(__HIP__MI300__) TODO: Add NAVI support
1474
+#endif // defined(__HIP__MI3XX__) TODO: Add NAVI support
1475
1451
-#if defined(__HIP__MI300__) // TODO: Add NAVI support
1476
+#if defined(__HIP__MI3XX__) // TODO: Add NAVI support
1477
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1478
int A_CHUNK, int UNRL, int N>
1479
__global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1456,6 +1481,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1481
const fp8_t* __restrict__ A, scalar_t* C,
1482
const float* __restrict__ s_A, const float* __restrict__ s_B,
1483
const int _WvPrGrp, const int CuCount) {
1484
+ constexpr int max_lds_len = LDS_SIZE;
1485
using scalar8 =
1486
__attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float;
1487
using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int;
@@ -1471,10 +1497,10 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1497
scalar8 h8;
1498
};
1499
1474
- __shared__ fp8_t s[1024 * 64];
1500
+ __shared__ fp8_t s[max_lds_len];
1501
1502
for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1477
- k < min(K * N, 64 * 1024); k += THRDS * WvPrGrp * A_CHUNK) {
1503
+ k < min(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
1504
*((bigType*)(&s[k])) = *((bigType*)(&A[k]));
1505
}
1506
__syncthreads();
@@ -1517,7 +1543,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1543
uint32_t k_ = k + threadIdx.x * A_CHUNK;
1544
if (k_ >= K) break;
1545
for (int n = 0; n < N; n++) {
1520
- if (k_ + K * n < 64 * 1024)
1546
+ if (k_ + K * n < max_lds_len)
1547
bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
1548
else
1549
bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
@@ -1608,7 +1634,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1634
m += CuCount * _WvPrGrp * YTILE;
1635
}
1636
}
1611
-#else // !defined(__HIP__MI300__) TODO: Add NAVI support
1637
+#else // !defined(__HIP__MI3XX__) TODO: Add NAVI support
1638
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1639
int A_CHUNK, int UNRL, int N>
1640
__global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
@@ -1618,7 +1644,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
1644
const int CuCount) {
1645
UNREACHABLE_CODE
1646
}
1621
-#endif // defined(__HIP__MI300__) TODO: Add NAVI support
1647
+#endif // defined(__HIP__MI3XX__) TODO: Add NAVI support
1648
1649
void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1650
at::Tensor& scale_a, at::Tensor& scale_b,
@@ -1638,12 +1664,13 @@ void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1664
dim3 grid(CuCount);
1665
const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1666
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1667
+ const int max_lds_len = get_lds_size();
1668
1669
#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \
1670
_N) \
1671
{ \
1672
dim3 block(64, _WvPrGrp); \
1646
- if ((K_in * N_in <= 64 * 1024) && (M_in % _YTILEs == 0)) { \
1673
+ if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
1674
int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
1675
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, \
tests/kernels/quant_utils.py
+9
-5
@@ -8,7 +8,7 @@ from vllm.platforms import current_platform
8
9
# Using the default value (240.0) from pytorch will cause accuracy
10
# issue on dynamic quantization models. Here use 224.0 for rocm.
11
-ROCM_FP8_MAX = 224.0
11
+ROCM_FP8FNUZ_MAX = 224.0
12
FP8_DTYPE = current_platform.fp8_dtype()
13
14
@@ -26,9 +26,11 @@ def ref_dynamic_per_token_quant(x: torch.tensor,
26
27
qtype_traits = torch.iinfo(quant_dtype) if quant_dtype == torch.int8 \
28
else torch.finfo(quant_dtype)
29
- qtype_traits_max = ROCM_FP8_MAX if current_platform.is_rocm() \
29
+ qtype_traits_max = ROCM_FP8FNUZ_MAX if current_platform.is_rocm() \
30
+ and current_platform.is_fp8_fnuz() \
31
else qtype_traits.max
31
- qtype_traits_min = -ROCM_FP8_MAX if current_platform.is_rocm() \
32
+ qtype_traits_min = -ROCM_FP8FNUZ_MAX if current_platform.is_rocm() \
33
+ and current_platform.is_fp8_fnuz() \
34
else qtype_traits.min
35
qtype_max = as_float32_tensor(qtype_traits_max)
36
s_1 = as_float32_tensor(1.0)
@@ -70,9 +72,11 @@ def ref_dynamic_per_tensor_fp8_quant(x: torch.tensor) \
72
-> tuple[torch.tensor, torch.tensor]:
73
74
fp8_traits = torch.finfo(FP8_DTYPE)
73
- fp8_traits_max = ROCM_FP8_MAX if current_platform.is_rocm() \
75
+ fp8_traits_max = ROCM_FP8FNUZ_MAX if current_platform.is_rocm() \
76
+ and current_platform.is_fp8_fnuz() \
77
else fp8_traits.max
75
- fp8_traits_min = -ROCM_FP8_MAX if current_platform.is_rocm() \
78
+ fp8_traits_min = -ROCM_FP8FNUZ_MAX if current_platform.is_rocm() \
79
+ and current_platform.is_fp8_fnuz() \
80
else fp8_traits.min
81
fp8_max = as_float32_tensor(fp8_traits_max)
82
one = as_float32_tensor(1.0)
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+2
-2
@@ -155,8 +155,8 @@ def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
155
scale_b: torch.Tensor, bias: torch.Tensor,
156
input_2d: torch.Tensor,
157
output_shape: list) -> torch.Tensor:
158
- from vllm.platforms.rocm import on_mi250_mi300
159
- if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi250_mi300(
158
+ from vllm.platforms.rocm import on_mi3xx
159
+ if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi3xx(
160
) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0:
161
output = ops.wvSplitKQ(weight.t(), qinput, out_dtype, scale_a, scale_b,
162
current_platform.get_cu_count())
vllm/model_executor/layers/utils.py
+2
-2
@@ -70,9 +70,9 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
70
def rocm_unquantized_gemm(x: torch.Tensor,
71
weight: torch.Tensor,
72
bias: Optional[torch.Tensor] = None):
73
- from vllm.platforms.rocm import on_mi250_mi300
73
+ from vllm.platforms.rocm import on_gfx9
74
k = weight.shape[1]
75
- use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi250_mi300() and \
75
+ use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and on_gfx9() and \
76
x.dtype in [torch.float16, torch.bfloat16] \
77
and k % 8 == 0 and bias is None)
78
vllm/platforms/rocm.py
+8
-2
@@ -105,9 +105,15 @@ def on_gfx1x() -> bool:
105
106
107
@cache
108
-def on_mi250_mi300() -> bool:
108
+def on_mi3xx() -> bool:
109
GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
110
- return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])
110
+ return any(arch in GPU_ARCH for arch in ["gfx942", "gfx950"])
111
+
112
+
113
+@cache
114
+def on_gfx9() -> bool:
115
+ GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
116
+ return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
117
118
119
@cache