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, 23:00 UTC
Improvements to wvSplitKrc skinny GEMM solution (#34304)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>
Hashem Hashemi committed
Mar 10, 2026 at 09:14 UTC
721ae79f50c5f85b301d05f1db71372b1ca85dd6
3 files changed
+168
-97
csrc/rocm/skinny_gemms.cu
+150
-84
@@ -12,6 +12,7 @@
12
#include "../cuda_compat.h"
13
#include "dispatch_utils.h"
14
#include "quantization/w8a8/fp8/common.cuh"
15
+#include "core/batch_invariant.hpp"
16
17
// TODO(rasmith): The kernels in this file are susceptible to integer overflow
18
// issues, do not take strides, and are unable to handle PyTorch tensors that
@@ -1224,17 +1225,14 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1225
#if defined(__gfx950__)
1226
#define WVSPLITKRC_1KPASS
1227
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1227
- int UNRL, int N, int GrpsShrB, int CHUNKK>
1228
+ int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
1229
__global__ void __launch_bounds__(WvPrGrp* THRDS)
1230
__attribute__((amdgpu_waves_per_eu(1, 1)))
1230
- wvSplitKrc_(const int actlN, const int K, const int M, const int Bx,
1231
- const int By, const scalar_t* __restrict__ B,
1232
- const scalar_t* __restrict__ A,
1233
- const scalar_t* __restrict__ BIAS, float* glbl, scalar_t* C,
1234
- const int CuCount) {
1235
- // Use upper half of glbl buffer for atomic reduce counting
1236
- int* cntr = (int*)(&glbl[M * N]);
1237
-
1231
+ wvSplitKrc_(const int actlN, const int K, const int Kap, const int M,
1232
+ const int Bx, const int By, const scalar_t* __restrict__ A,
1233
+ const scalar_t* __restrict__ B,
1234
+ const scalar_t* __restrict__ BIAS, float* glbl, int* cntr,
1235
+ scalar_t* C, const int CuCount) {
1236
constexpr int NTILE = 16;
1237
constexpr int APAD = 1;
1238
constexpr int ASTRD = 64;
@@ -1425,11 +1423,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1423
unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
1424
for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
1425
__builtin_amdgcn_global_load_lds(
1428
- (int*)(&A[min__(
1429
- K * actlN - A_CHUNK,
1430
- kOffcp + K * (n / CHUNKK +
1431
- (N / CHUNKK) * (threadIdx.x / (64 / CHUNKK)) +
1432
- (threadIdx.y % sprdN)))]),
1426
+ (int*)(&A[min__(Kap * actlN - A_CHUNK,
1427
+ kOffcp + Kap * (n / CHUNKK +
1428
+ (N / CHUNKK) * (threadIdx.x /
1429
+ (64 / CHUNKK)) +
1430
+ (threadIdx.y % sprdN)))]),
1431
(int*)(&s[(k +
1432
kFitPdd * ((n / CHUNKK) + (threadIdx.y % sprdN)))]),
1433
16, 0, 0);
@@ -1533,45 +1531,98 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1531
}
1532
}
1533
1534
+ union flt4 {
1535
+ scalar8 s8;
1536
+ float2 f2[2];
1537
+ float4 f4;
1538
+ };
1539
if (m + (threadIdx.x % 16) < M) {
1540
int my_cntr;
1541
int mindx = m + (threadIdx.x % 16);
1542
int g_mindx = m * 4 + (threadIdx.x % 64); // coalesced atomic reduction
1543
scalar_t biases[N / NTILE / GrpsShrB][4] = {};
1544
// Atomic add the output, read biases
1542
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
1543
- for (uint32_t j = 0; j < 4; j++) {
1544
- // int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1545
- // (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1546
- // int adr = mindx + M * nindx;
1547
- int g_nindx =
1548
- j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1549
- int g_adr = g_mindx + M * g_nindx * 4;
1550
- atomicAdd(&glbl[g_adr], sum4[nt][0][j]);
1545
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1546
+ int g_nindx =
1547
+ (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1548
+ int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1549
+ if (DTRMNSTC) {
1550
+ flt4 flt4_ = {.s8 = sum4[nt][0]};
1551
+ __hip_atomic_store((float2*)&glbl[g_adr + M * N * (m0 / Mmod)],
1552
+ flt4_.f2[0], __ATOMIC_RELAXED,
1553
+ __HIP_MEMORY_SCOPE_AGENT);
1554
+ __hip_atomic_store((float2*)&glbl[g_adr + 2 + M * N * (m0 / Mmod)],
1555
+ flt4_.f2[1], __ATOMIC_RELAXED,
1556
+ __HIP_MEMORY_SCOPE_AGENT);
1557
+ } else {
1558
+ for (uint32_t j = 0; j < 4; j++)
1559
+ atomicAdd((&glbl[g_adr + j]), sum4[nt][0][j]);
1560
}
1561
+ }
1562
+
1563
+ __atomic_signal_fence(__ATOMIC_SEQ_CST);
1564
+ asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
1565
+ __atomic_signal_fence(__ATOMIC_SEQ_CST);
1566
+
1567
int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1568
(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1569
int adr_ = mindx + M * nindx_ / 4;
1555
- // Update the complete counter
1570
my_cntr = atomicAdd(&cntr[adr_], 1);
1557
- float vals[N / NTILE / GrpsShrB][4] = {};
1571
+
1572
+ // make sure LDS is free for write out staging
1573
+ if (DTRMNSTC) __syncthreads();
1574
+
1575
+ // Update the complete counter
1576
+ flt4 vals[N / NTILE / GrpsShrB] = {};
1577
// If we're the last k-shard, read back the value and convert...
1578
if (my_cntr + 1 == k_rnd) {
1560
- if (BIAS)
1561
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1562
- for (uint32_t j = 0; j < 4; j++) {
1563
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1564
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1565
- biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1579
+ cntr[adr_] = 0; // clear for next round
1580
+ if constexpr (DTRMNSTC) {
1581
+ #pragma unroll
1582
+ for (int ks = 0; ks < k_rnd; ks++) {
1583
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1584
+ int g_nindx =
1585
+ (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1586
+ int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1587
+ __builtin_amdgcn_global_load_lds(
1588
+ (float4*)(&glbl[g_adr + M * N * ks]),
1589
+ &(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 +
1590
+ nt * THRDS * 4 * k_rnd]),
1591
+ 16, 0, 0);
1592
}
1593
}
1568
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1569
- for (uint32_t j = 0; j < 4; j++) {
1594
+ if (BIAS)
1595
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1596
+ for (uint32_t j = 0; j < 4; j++) {
1597
+ int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1598
+ (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1599
+ biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1600
+ }
1601
+ }
1602
+ asm volatile("s_waitcnt 0");
1603
+ for (int ks = 0; ks < k_rnd; ks++) {
1604
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1605
+ float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) +
1606
+ ks * THRDS * 4 + nt * THRDS * 4 * k_rnd];
1607
+ vals[nt].f4 += eval;
1608
+ }
1609
+ }
1610
+ } else {
1611
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1612
int g_nindx =
1571
- j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1572
- int g_adr = g_mindx + M * g_nindx * 4;
1573
- vals[nt][j] = glbl[g_adr];
1613
+ (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1614
+ int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1615
+ vals[nt].f4 = *(float4*)(&glbl[g_adr]);
1616
+ *(float4*)(&glbl[g_adr]) = {}; // clear out for next round
1617
}
1618
+ if (BIAS)
1619
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1620
+ for (uint32_t j = 0; j < 4; j++) {
1621
+ int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1622
+ (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1623
+ biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1624
+ }
1625
+ }
1626
}
1627
__builtin_amdgcn_sched_barrier(0);
1628
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
@@ -1581,11 +1632,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1632
if (nindx < actlN) {
1633
int adr = mindx + M * nindx;
1634
if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1584
- vals[nt][j] += __bfloat162float(biases[nt][j]);
1585
- C[adr] = __float2bfloat16(vals[nt][j]);
1635
+ vals[nt].s8[j] += __bfloat162float(biases[nt][j]);
1636
+ C[adr] = __float2bfloat16(vals[nt].s8[j]);
1637
} else {
1587
- vals[nt][j] += __half2float(biases[nt][j]);
1588
- C[adr] = __float2half(vals[nt][j]);
1638
+ vals[nt].s8[j] += __half2float(biases[nt][j]);
1639
+ C[adr] = __float2half(vals[nt].s8[j]);
1640
}
1641
}
1642
}
@@ -1604,21 +1655,25 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1655
}
1656
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1657
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1607
- int UNRL, int N, int GrpsShrB, int CHUNKK>
1608
-__global__ void wvSplitKrc_(const int actlN, const int K, const int M,
1609
- const int Bx, const int By, const scalar_t* B,
1610
- const scalar_t* __restrict__ A,
1658
+ int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
1659
+__global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
1660
+ const int M, const int Bx, const int By,
1661
+ const scalar_t* B, const scalar_t* __restrict__ A,
1662
const scalar_t* __restrict__ BIAS, float* glbl,
1612
- // int* cntr,
1613
- scalar_t* C, const int CuCount){UNREACHABLE_CODE}
1663
+ int* cntr, scalar_t* C,
1664
+ const int CuCount){UNREACHABLE_CODE}
1665
#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1666
1667
torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1668
const std::optional<at::Tensor>& in_bias,
1669
const int64_t CuCount) {
1619
- auto M_in = in_a.size(0);
1620
- auto N_in = in_b.size(0);
1621
- auto K_in = in_a.size(1);
1670
+ int _DTRMNSTC = 1; // vllm::vllm_is_batch_invariant();
1671
+
1672
+ auto M_in = in_b.size(0);
1673
+ auto N_in = in_a.size(0);
1674
+ auto K_in = in_b.size(1);
1675
+ auto Kap_in = in_a.stride(0);
1676
+
1677
auto Bx_in =
1678
(in_bias.has_value() && in_bias->numel() > 0)
1679
? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
@@ -1635,13 +1690,9 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1690
1691
auto out_c = torch::empty(
1692
{N_in, M_in},
1638
- torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
1693
+ torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device()));
1694
1695
auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1));
1641
- auto axl_glbl = torch::empty(
1642
- {N_p2 + N_p2 / 4, M_in + M_in / 4},
1643
- torch::TensorOptions().dtype(torch::kFloat32).device(in_b.device()));
1644
- axl_glbl.zero_(); // disable for FAST_UNSAFE_RDC_INIT
1696
1697
dim3 grid(CuCount);
1698
@@ -1649,55 +1700,70 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1700
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1701
// const int max_lds_len = get_lds_size() / 2;
1702
1703
+ // With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
1704
+ // and each working on a 512-shard of K, how many CUs would we need?
1705
+ int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
1706
+
1707
+ // How many of 4 waves in a group can work on same 16 Ms at same time? First
1708
+ // try to maximize this. This reduces the Ms each group works on, i.e.
1709
+ // increasing the number of CUs needed.
1710
+ int GrpsShrB = min(N_p2 / 16, 4);
1711
+
1712
+ // Given the above, how many CUs would we need?
1713
+ int CuNeeded = rndup_cus * GrpsShrB;
1714
+
1715
+ if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size");
1716
+
1717
+ // Can we increase SplitK by shrinking the K-shared to 256?
1718
+ int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
1719
+
1720
+ static torch::Tensor axl_glbl =
1721
+ torch::zeros(
1722
+ 128 * 1024 * (_DTRMNSTC ? 12 : 1),
1723
+ torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()))
1724
+ .detach();
1725
+ static torch::Tensor axl_cntr =
1726
+ torch::zeros(
1727
+ 128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4,
1728
+ torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
1729
+ .detach();
1730
+ auto glbl = axl_glbl.data_ptr<float>();
1731
+ auto cntr = axl_cntr.data_ptr<int>();
1732
+
1733
#define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \
1734
{ \
1735
dim3 block(64, 4); \
1655
- wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK> \
1656
- <<<grid, block, 0, stream>>>(N_in, K_in, M_in, Bx_in, By_in, af4, bf4, \
1657
- biasf4, glbl, c, CuCount); \
1736
+ if (_DTRMNSTC) \
1737
+ wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 1> \
1738
+ <<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
1739
+ af4, bf4, biasf4, glbl, cntr, c, \
1740
+ CuCount); \
1741
+ else \
1742
+ wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 0> \
1743
+ <<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
1744
+ af4, bf4, biasf4, glbl, cntr, c, \
1745
+ CuCount); \
1746
}
1747
1660
- AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitKrc", [&] {
1748
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(in_a.scalar_type(), "wvSplitKrc", [&] {
1749
using fptype = typename scalar<scalar_t>::type;
1662
- fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
1750
+ const fptype* af4 = reinterpret_cast<const fptype*>(in_a.data_ptr());
1751
const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
1752
const fptype* biasf4 =
1753
(in_bias.has_value() && in_bias->numel() > 0)
1754
? reinterpret_cast<const fptype*>(in_bias->data_ptr())
1755
: nullptr;
1756
fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
1669
- auto glbl = axl_glbl.data_ptr<float>();
1670
-
1671
- // With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
1672
- // and each working on a 512-shard of K, how many CUs would we need?
1673
- int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
1674
-
1675
- // How many of 4 waves in a group can work on same 16 Ms at same time? First
1676
- // try to maximize this. This reduces the Ms each group works on, i.e.
1677
- // increasing the number of CUs needed.
1678
- int GrpsShrB = min(N_p2 / 16, 4);
1679
-
1680
- // Given the above, how many CUs would we need?
1681
- int CuNeeded = rndup_cus * GrpsShrB;
1682
-
1683
- if (CuNeeded > CuCount) std::runtime_error("Invalid wvSplitKrc size");
1684
-
1685
- // Can we increase SplitK by shrinking the K-shared to 256?
1686
- int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
1757
1758
switch (N_p2) {
1759
case 16:
1760
WVSPLITKrc(16, 1, 1) break;
1761
case 32:
1692
- if (chunkk == 2)
1693
- WVSPLITKrc(32, 2, 2) else if (chunkk == 1) WVSPLITKrc(32, 2, 1) break;
1762
+ if (chunkk == 2) WVSPLITKrc(32, 2, 2) else WVSPLITKrc(32, 2, 1) break;
1763
case 64:
1695
- if (chunkk == 2)
1696
- WVSPLITKrc(64, 4, 2) else if (chunkk == 1) WVSPLITKrc(64, 4, 1) break;
1764
+ if (chunkk == 2) WVSPLITKrc(64, 4, 2) else WVSPLITKrc(64, 4, 1) break;
1765
case 128:
1698
- if (chunkk == 2)
1699
- WVSPLITKrc(128, 4, 2) else if (chunkk == 1)
1700
- WVSPLITKrc(128, 4, 1) break;
1766
+ if (chunkk == 2) WVSPLITKrc(128, 4, 2) else WVSPLITKrc(128, 4, 1) break;
1767
default:
1768
throw std::runtime_error(
1769
"Unsupported N value: " + std::to_string(M_in) + "," +
tests/kernels/quantization/test_rocm_skinny_gemms.py
+7
-4
@@ -70,7 +70,6 @@ N_FACTORS_WVSPLITKRC = [
70
117,
71
128,
72
]
73
-
73
K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
74
M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
75
@@ -123,10 +122,11 @@ def pad_fp8(weight):
122
@pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC)
123
@pytest.mark.parametrize("dtype", DTYPES)
124
@pytest.mark.parametrize("seed", SEEDS)
125
+@pytest.mark.parametrize("padded_a", [False, True])
126
@pytest.mark.parametrize("bias_mode", BIAS_MODES)
127
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
128
@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
129
-def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
129
+def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode):
130
torch.manual_seed(seed)
131
cu_count = num_compute_units()
132
@@ -141,7 +141,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
141
# Given the above, how many CUs would we need?
142
CuNeeded = rndup_cus * GrpsShrB
143
# candidate for atomic reduce count splitk?
144
- fits_wvsplitkrc = CuNeeded <= cu_count
144
+ fits_wvsplitkrc = (N_p2 * m * ((k + 512 - 1) // 512)) <= 128 * 1024 * 12
145
+ fits_wvsplitkrc &= CuNeeded <= cu_count
146
147
if not fits_wvsplitkrc:
148
pytest.skip("Too large for wvSplitKrc")
@@ -151,6 +152,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
152
) # normalize to avoid large output-bias deltas
153
A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
154
B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
155
+ if padded_a:
156
+ A = pad_fp8(A)
157
158
BIAS = None
159
if bias_mode == 1:
@@ -159,7 +162,7 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
162
BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
163
164
ref_out = torch.nn.functional.linear(A, B, BIAS)
162
- out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
165
+ out = ops.wvSplitKrc(A, B, cu_count, BIAS)
166
167
if xnorm:
168
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
vllm/model_executor/layers/utils.py
+11
-9
@@ -129,10 +129,6 @@ def rocm_unquantized_gemm_impl(
129
k = weight.shape[1]
130
131
cu_count = num_compute_units()
132
- if use_aiter_triton_gemm(n, m, k, x.dtype):
133
- from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
134
-
135
- return gemm_a16w16(x, weight, bias)
132
133
# Next ^2 of n
134
N_p2 = 1 << (n - 1).bit_length()
@@ -145,7 +141,10 @@ def rocm_unquantized_gemm_impl(
141
# Given the above, how many CUs would we need?
142
CuNeeded = rndup_cus * GrpsShrB
143
# candidate for atomic reduce count splitk?
148
- fits_wvsplitkrc = CuNeeded <= cu_count
144
+ fits_wvsplitkrc = (
145
+ N_p2 * m * ((k + 512 - 1) // 512)
146
+ ) <= 128 * 1024 * 12 # deterministic
147
+ fits_wvsplitkrc &= CuNeeded <= cu_count
148
149
use_skinny_reduce_counting = (
150
envs.VLLM_ROCM_USE_SKINNY_GEMM
@@ -157,13 +156,16 @@ def rocm_unquantized_gemm_impl(
156
and k > 512
157
and m % 16 == 0
158
and fits_wvsplitkrc
160
- and x.is_contiguous()
159
+ and weight.is_contiguous()
160
)
161
)
162
if use_skinny_reduce_counting:
164
- x_view = x.reshape(-1, x.size(-1))
165
- out = ops.wvSplitKrc(weight, x_view, cu_count, bias)
166
- return out.reshape(*x.shape[:-1], weight.shape[0])
163
+ return ops.wvSplitKrc(x, weight, cu_count, bias)
164
+
165
+ if use_aiter_triton_gemm(n, m, k, x.dtype):
166
+ from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
167
+
168
+ return gemm_a16w16(x, weight, bias)
169
170
use_skinny = (
171
envs.VLLM_ROCM_USE_SKINNY_GEMM