A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 16:30 UTC
Perf tuning and expansion of cases covered for wvSplitKrc (#33493)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>
Hashem Hashemi committed
Feb 7, 2026 at 05:33 UTC
ed17f54c8bccc2f1a19f83f65b36d837698b8792
3 files changed
+214
-223
csrc/rocm/skinny_gemms.cu
+143
-184
@@ -1365,13 +1365,12 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1365
return out_c;
1366
}
1367
1368
-#if defined(__gfx950__) // TODO: Add NAVI support
1369
- // This version targets big A[] cases, where it is much larger than LDS
1370
- // capacity
1368
+// This version targets cases skinny where CUs are not filled
1369
+// Wave-SplitK is used with reduction done via atomics.
1370
+#if defined(__gfx950__)
1371
#define WVSPLITKRC_1KPASS
1372
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1373
- int UNRL, int N, int GrpsShrB>
1374
-
1373
+ int UNRL, int N, int GrpsShrB, int CHUNKK>
1374
__global__ void __launch_bounds__(WvPrGrp* THRDS)
1375
__attribute__((amdgpu_waves_per_eu(1, 1)))
1376
wvSplitKrc_(const int actlN, const int K, const int M, const int Bx,
@@ -1383,12 +1382,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1382
int* cntr = (int*)(&glbl[M * N]);
1383
1384
constexpr int NTILE = 16;
1386
- constexpr int WVLDS_ = (NTILE * THRDS * A_CHUNK);
1385
constexpr int APAD = 1;
1386
constexpr int ASTRD = 64;
1387
constexpr int BPAD = 1;
1390
- constexpr int BSTRD = 64;
1391
- constexpr int WVLDS = ((WVLDS_ + (WVLDS_ / BSTRD) * 4 * BPAD));
1388
+ constexpr int WVLDS_ = THRDS * A_CHUNK / CHUNKK;
1389
+ constexpr int WVLDS = ((WVLDS_ + A_CHUNK * BPAD)) * YTILE;
1390
1391
constexpr int max_lds_len = LDS_SIZE / 2;
1392
@@ -1442,17 +1440,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1440
break;
1441
}
1442
#else
1445
- int constexpr kFit = 512;
1443
+ int constexpr kFit = 512 / CHUNKK;
1444
int constexpr kfitsPerRdc = 1;
1445
#endif
1446
1449
- bool doRdc = (kfitsPerRdc * kFit < K);
1447
+ bool doRdc = true; // Assuming (kfitsPerRdc * kFit < K) is always true
1448
uint32_t numCuWithFullK =
1449
((M + (WvPrGrp * YTILE / GrpsShrB) - 1) / (WvPrGrp * YTILE / GrpsShrB));
1450
uint32_t Mmod = numCuWithFullK * (WvPrGrp * YTILE / GrpsShrB);
1451
1452
// given above k-split, find this wave's position
1455
- uint32_t kFitPdd = kFit + (kFit / ASTRD) * APAD;
1453
+ uint32_t kFitPdd = kFit * CHUNKK + ((kFit * CHUNKK) / ASTRD) * APAD;
1454
uint32_t m0 = (blockIdx.x * WvPrGrp / GrpsShrB) * YTILE;
1455
uint32_t m1 = ((threadIdx.y % WvPrGrp) / GrpsShrB) * YTILE;
1456
uint32_t m = (m0 + m1) % Mmod;
@@ -1460,8 +1458,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1458
uint32_t k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
1459
const uint32_t k_rnd = (K + kFit * kfitsPerRdc - 1) / (kFit * kfitsPerRdc);
1460
1463
- scalar8 sum4[N / NTILE / GrpsShrB][1];
1464
- bigType bigB_[YTILE / GrpsShrB][UNRL];
1461
+ scalar8 sum4[N / NTILE / GrpsShrB][1] = {0};
1462
+ bigType bigB_[YTILE / GrpsShrB / CHUNKK][UNRL];
1463
const uint32_t bLoader = (threadIdx.y % GrpsShrB);
1464
uint32_t kBase = 0;
1465
if (k_str >= K) return;
@@ -1498,12 +1496,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1496
#pragma unroll
1497
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1498
uint32_t k = k_str + k2 * THRDS * A_CHUNK;
1501
- uint32_t k_ = k + threadIdx.x * A_CHUNK;
1499
+ uint32_t k_ = k + (threadIdx.x % (THRDS / CHUNKK)) * A_CHUNK;
1500
const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1501
#pragma unroll
1504
- for (uint32_t y = 0; y < YTILE / GrpsShrB; y++)
1505
- bigB_[y][k2].h8 = (loadnt(
1506
- (scalar8*)(&B_[min__(y * GrpsShrB + bLoader + m, M - 1) * K])));
1502
+ for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK)
1503
+ bigB_[y / CHUNKK][k2].h8 = (loadnt(
1504
+ (scalar8*)(&B_[min__((y + threadIdx.x / (THRDS / CHUNKK)) * GrpsShrB +
1505
+ bLoader + m,
1506
+ M - 1) *
1507
+ K])));
1508
}
1509
{
1510
#else
@@ -1556,48 +1557,51 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1557
if (reloada) {
1558
#endif
1559
constexpr int sprdN = 4;
1559
- const uint32_t thrd = ((threadIdx.y / sprdN) * THRDS + threadIdx.x);
1560
+ const uint32_t thrd = threadIdx.x % (THRDS / CHUNKK);
1561
1562
#ifndef WVSPLITKRC_1KPASS
1563
#pragma unroll
1563
- for (int k = 0; k < kFit; k += THRDS * (WvPrGrp / sprdN) * A_CHUNK) {
1564
+ for (int k = 0; k < kFit;
1565
+ k += (THRDS * (WvPrGrp / sprdN) * A_CHUNK) / CHUNKK) {
1566
#else
1567
const unsigned int k = 0;
1568
{
1569
#endif
1570
unsigned int kOff = k + (thrd * A_CHUNK);
1569
- unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
1570
- const unsigned int k_in = kOffcp + ((threadIdx.y % sprdN)) * K;
1571
- const unsigned int k_ot = kOff + ((threadIdx.y % sprdN)) * kFitPdd;
1572
- for (unsigned int n = 0; n < N / 2; n += sprdN) {
1573
- __builtin_amdgcn_global_load_lds((int*)(&A[k_in + n * K]),
1574
- (int*)(&s[(k_ot + n * kFitPdd)]),
1575
- 16, 0, 0);
1576
- if (((threadIdx.y % sprdN)) + n + N / 2 >= actlN) continue;
1571
+ unsigned int kOffcp =
1572
+ k_str + kOff; // min__(K - A_CHUNK, k_str + kOff);
1573
+ for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
1574
__builtin_amdgcn_global_load_lds(
1578
- (int*)(&A[k_in + (n + N / 2) * K]),
1579
- (int*)(&s[(k_ot + (n + N / 2) * kFitPdd)]), 16, 0, 0);
1575
+ (int*)(&A[min__(
1576
+ K * actlN - A_CHUNK,
1577
+ kOffcp + K * (n / CHUNKK +
1578
+ (N / CHUNKK) * (threadIdx.x / (64 / CHUNKK)) +
1579
+ (threadIdx.y % sprdN)))]),
1580
+ (int*)(&s[(k +
1581
+ kFitPdd * ((n / CHUNKK) + (threadIdx.y % sprdN)))]),
1582
+ 16, 0, 0);
1583
}
1584
1585
// Stage loaded B[] to LDS for MFMA swizzling...
1586
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1587
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1585
- uint32_t k_ = k + threadIdx.x * A_CHUNK;
1588
+ uint32_t k_ = k + (threadIdx.x % (THRDS / CHUNKK)) * A_CHUNK;
1589
const bool oob_k = (k_ >= K);
1587
- for (uint32_t y = 0; y < YTILE / GrpsShrB; y++) {
1588
- uint32_t idx = threadIdx.x * 4 +
1589
- (y * GrpsShrB + bLoader) * ((THRDS + BPAD) * 4);
1590
+ for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK) {
1591
+ uint32_t idx =
1592
+ (threadIdx.x % (THRDS / CHUNKK)) * 4 +
1593
+ ((y + threadIdx.x / (THRDS / CHUNKK)) * GrpsShrB + bLoader) *
1594
+ ((THRDS / CHUNKK + BPAD) * 4);
1595
// zero out if oob
1596
*((scalar8*)&myStg[idx]) =
1592
- (oob_k || (y * GrpsShrB + bLoader + m >= M))
1597
+ (oob_k) // TODO: ever necessary (y*GrpsShrB+bLoader+m>=M) ?
1598
? 0
1594
- : bigB_[y][k2].h8;
1599
+ : bigB_[y / CHUNKK][k2].h8;
1600
}
1601
}
1602
}
1603
}
1604
}
1600
-
1605
#ifndef WVSPLITKRC_1KPASS
1606
// Fire load of next B[] chunk...
1607
if ((k1 + THRDS * A_CHUNK * UNRL < k_end) &&
@@ -1608,40 +1612,50 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1612
uint32_t k_ = k + threadIdx.x * A_CHUNK;
1613
const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1614
#pragma unroll
1611
- for (uint32_t y = 0; y < YTILE / GrpsShrB; y++)
1612
- bigB_[y][k2].h8 = (loadnt(
1613
- (scalar8*)(&B_[min__(y * GrpsShrB + bLoader + m, M - 1) * K])));
1615
+ for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK)
1616
+ bigB_[y / CHUNKK][k2].h8 = (loadnt(
1617
+ (scalar8*)(&B_[min__((y + threadIdx.x / (THRDS / CHUNKK)) *
1618
+ GrpsShrB +
1619
+ bLoader + m,
1620
+ M - 1) *
1621
+ K])));
1622
}
1623
#endif
1624
1625
// B[] staging is cooperative across GrpsShrB, so sync here before reading
1618
- // back
1626
+ // back. This wait is currently inserted by compiler, but not gauranteed.
1627
+ asm volatile("s_waitcnt 0");
1628
__syncthreads();
1629
1630
// read back B[] swizzled for MFMA...
1622
- bigType bigB[YTILE][UNRL];
1631
+ bigType bigB[YTILE / CHUNKK][UNRL];
1632
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1624
- for (uint32_t y = 0; y < YTILE; y++) {
1625
- unsigned int idx = (threadIdx.x % YTILE) * ((THRDS + BPAD) * 4) +
1626
- (threadIdx.x / YTILE) * 4 + y * 16;
1633
+ for (uint32_t y = 0; y < YTILE / CHUNKK; y++) {
1634
+ unsigned int idx =
1635
+ (threadIdx.x % YTILE) * ((THRDS / CHUNKK + BPAD) * 4) +
1636
+ (threadIdx.x / YTILE) * 4 + y * 16;
1637
bigB[y][k2].h8 = *((scalar8*)&myStg[idx]);
1638
}
1639
}
1640
1641
// rReadback A[] swizzled for MFMA...
1632
- bigType bigA[N / GrpsShrB][UNRL];
1642
+ bigType bigA[N / GrpsShrB / CHUNKK][UNRL];
1643
#pragma unroll
1644
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1645
uint32_t k = k1 + k2 * THRDS * A_CHUNK - kBase - k_str;
1646
#pragma unroll
1647
for (uint32_t nt = 0; nt < N / GrpsShrB; nt += NTILE)
1648
#pragma unroll
1639
- for (uint32_t n = 0; n < NTILE; n++) {
1640
- uint32_t idxa = (nt + (threadIdx.x % NTILE) +
1641
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) *
1642
- kFitPdd +
1643
- A_CHUNK * ((threadIdx.x / NTILE) + n * 4) + k;
1644
- bigA[nt + n][k2] = *((const bigType*)(&(s[idxa])));
1649
+ for (uint32_t n = 0; n < NTILE / CHUNKK; n++) {
1650
+ uint32_t idxa =
1651
+ ((nt + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) % (N / CHUNKK) +
1652
+ (threadIdx.x % NTILE)) *
1653
+ kFitPdd +
1654
+ ((nt + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) /
1655
+ (N / CHUNKK)) *
1656
+ A_CHUNK * (64 / CHUNKK) +
1657
+ A_CHUNK * ((threadIdx.x / NTILE) + n * 4) + k;
1658
+ bigA[nt / CHUNKK + n][k2] = *((const bigType*)(&(s[idxa])));
1659
}
1660
}
1661
@@ -1650,152 +1664,75 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1664
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1665
#pragma unroll
1666
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1653
- if constexpr (std::is_same_v<scalar_t, half>) {
1654
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1655
- bigA[nt * NTILE + 0][k2].h4[0], bigB[0][k2].h4[0],
1656
- (k1 == k_str) ? ((scalar8){0}) : sum4[nt][0], 0, 0, 0);
1657
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1658
- bigA[nt * NTILE + 0][k2].h4[1], bigB[0][k2].h4[1], sum4[nt][0], 0,
1659
- 0, 0);
1660
- } else { // bf16
1661
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
1662
- bigA[nt * NTILE + 0][k2].h4[0], bigB[0][k2].h4[0],
1663
- (k1 == k_str) ? ((scalar8){0}) : sum4[nt][0], 0, 0, 0);
1664
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
1665
- bigA[nt * NTILE + 0][k2].h4[1], bigB[0][k2].h4[1], sum4[nt][0], 0,
1666
- 0, 0);
1667
- }
1667
#pragma unroll
1669
- for (uint32_t j = 1; j < YTILE; j++) {
1668
+ for (uint32_t j = 0; j < YTILE / CHUNKK; j++) {
1669
if constexpr (std::is_same_v<scalar_t, half>) {
1671
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1672
- bigA[nt * NTILE + j][k2].h4[0], bigB[j][k2].h4[0], sum4[nt][0],
1673
- 0, 0, 0);
1674
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1675
- bigA[nt * NTILE + j][k2].h4[1], bigB[j][k2].h4[1], sum4[nt][0],
1676
- 0, 0, 0);
1670
+ sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x32_f16(
1671
+ bigA[nt * (YTILE / CHUNKK) + j][k2].h8, bigB[j][k2].h8,
1672
+ sum4[nt][0], 0, 0, 0);
1673
} else { // bf16
1678
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
1679
- bigA[nt * NTILE + j][k2].h4[0], bigB[j][k2].h4[0], sum4[nt][0],
1680
- 0, 0, 0);
1681
- sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
1682
- bigA[nt * NTILE + j][k2].h4[1], bigB[j][k2].h4[1], sum4[nt][0],
1683
- 0, 0, 0);
1674
+ sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x32_bf16(
1675
+ bigA[nt * (YTILE / CHUNKK) + j][k2].h8, bigB[j][k2].h8,
1676
+ sum4[nt][0], 0, 0, 0);
1677
}
1678
}
1679
}
1680
}
1681
}
1682
1690
- if (!doRdc) {
1691
- if (m + (threadIdx.x % 16) < M) {
1692
- scalar_t biases[N / NTILE / GrpsShrB][4] = {0};
1683
+ if (m + (threadIdx.x % 16) < M) {
1684
+ int my_cntr;
1685
+ int mindx = m + (threadIdx.x % 16);
1686
+ int g_mindx = m * 4 + (threadIdx.x % 64); // coalesced atomic reduction
1687
+ scalar_t biases[N / NTILE / GrpsShrB][4] = {};
1688
+ // Atomic add the output, read biases
1689
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
1690
+ for (uint32_t j = 0; j < 4; j++) {
1691
+ // int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1692
+ // (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1693
+ // int adr = mindx + M * nindx;
1694
+ int g_nindx =
1695
+ j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1696
+ int g_adr = g_mindx + M * g_nindx * 4;
1697
+ atomicAdd(&glbl[g_adr], sum4[nt][0][j]);
1698
+ }
1699
+ int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1700
+ (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1701
+ int adr_ = mindx + M * nindx_ / 4;
1702
+ // Update the complete counter
1703
+ my_cntr = atomicAdd(&cntr[adr_], 1);
1704
+ float vals[N / NTILE / GrpsShrB][4] = {};
1705
+ // If we're the last k-shard, read back the value and convert...
1706
+ if (my_cntr + 1 == k_rnd) {
1707
if (BIAS)
1708
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1709
for (uint32_t j = 0; j < 4; j++) {
1696
- int mindx = m + (threadIdx.x % 16);
1710
int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1711
(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1699
- biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * M];
1712
+ biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1713
}
1714
}
1715
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1716
for (uint32_t j = 0; j < 4; j++) {
1704
- int mindx = m + (threadIdx.x % 16);
1705
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1706
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1707
- int adr = mindx + M * nindx;
1708
- if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1709
- if (BIAS) sum4[nt][0][j] += __bfloat162float(biases[nt][j]);
1710
- C[adr] = __float2bfloat16(sum4[nt][0][j]);
1711
- } else {
1712
- if (BIAS) sum4[nt][0][j] += __half2float(biases[nt][j]);
1713
- C[adr] = __float2half(sum4[nt][0][j]);
1714
- }
1717
+ int g_nindx =
1718
+ j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1719
+ int g_adr = g_mindx + M * g_nindx * 4;
1720
+ vals[nt][j] = glbl[g_adr];
1721
}
1722
}
1717
- }
1718
- } else {
1719
- if (m + (threadIdx.x % 16) < M) {
1720
- int my_cntr;
1721
- if (!BIAS) {
1722
- int mindx = m + (threadIdx.x % 16);
1723
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
1724
- for (uint32_t j = 0; j < 4; j++) {
1725
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1726
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1727
- int adr = mindx + M * nindx;
1728
- atomicAdd(&glbl[adr], sum4[nt][0][j]);
1729
- }
1730
- int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1731
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1732
- int adr_ = mindx + M * nindx_ / 4;
1733
- my_cntr = atomicAdd(&cntr[adr_], 1);
1734
- float vals[N / NTILE / GrpsShrB][4] = {};
1735
- if (my_cntr + 1 == k_rnd) {
1736
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1737
- for (uint32_t j = 0; j < 4; j++) {
1738
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1739
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1740
- int adr = mindx + M * nindx;
1741
- vals[nt][j] = glbl[adr];
1742
- }
1743
- }
1744
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1745
- for (uint32_t j = 0; j < 4; j++) {
1746
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1747
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1748
- if (nindx >= actlN) break;
1749
- int adr = mindx + M * nindx;
1750
- if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1751
- C[adr] = __float2bfloat16(vals[nt][j]);
1752
- } else {
1753
- C[adr] = __float2half(vals[nt][j]);
1754
- }
1755
- }
1756
- }
1757
- }
1758
- } else {
1759
- int mindx = m + (threadIdx.x % 16);
1760
- scalar_t biases[N / NTILE / GrpsShrB][4] = {};
1761
- // Atomic add the output, read biases
1762
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
1763
- for (uint32_t j = 0; j < 4; j++) {
1764
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1765
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1723
+ __builtin_amdgcn_sched_barrier(0);
1724
+ for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1725
+ for (uint32_t j = 0; j < 4; j++) {
1726
+ int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1727
+ (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1728
+ if (nindx < actlN) {
1729
int adr = mindx + M * nindx;
1767
- atomicAdd(&glbl[adr], sum4[nt][0][j]);
1768
- biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * M];
1769
- }
1770
- int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1771
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1772
- int adr_ = mindx + M * nindx_ / 4;
1773
- // Update the complete counter
1774
- my_cntr = atomicAdd(&cntr[adr_], 1);
1775
- float vals[N / NTILE / GrpsShrB][4] = {};
1776
- // If we're the last k-shard, read back the value and convert...
1777
- if (my_cntr + 1 == k_rnd) {
1778
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1779
- for (uint32_t j = 0; j < 4; j++) {
1780
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1781
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1782
- int adr = mindx + M * nindx;
1783
- vals[nt][j] = glbl[adr];
1784
- }
1785
- }
1786
- for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1787
- for (uint32_t j = 0; j < 4; j++) {
1788
- int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1789
- (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1790
- if (nindx >= actlN) break;
1791
- int adr = mindx + M * nindx;
1792
- if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1793
- vals[nt][j] += __bfloat162float(biases[nt][j]);
1794
- C[adr] = __float2bfloat16(vals[nt][j]);
1795
- } else {
1796
- vals[nt][j] += __half2float(biases[nt][j]);
1797
- C[adr] = __float2half(vals[nt][j]);
1798
- }
1730
+ if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1731
+ vals[nt][j] += __bfloat162float(biases[nt][j]);
1732
+ C[adr] = __float2bfloat16(vals[nt][j]);
1733
+ } else {
1734
+ vals[nt][j] += __half2float(biases[nt][j]);
1735
+ C[adr] = __float2half(vals[nt][j]);
1736
}
1737
}
1738
}
@@ -1814,7 +1751,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1751
}
1752
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1753
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1817
- int UNRL, int N, int GrpsShrB>
1754
+ int UNRL, int N, int GrpsShrB, int CHUNKK>
1755
__global__ void wvSplitKrc_(const int actlN, const int K, const int M,
1756
const int Bx, const int By, const scalar_t* B,
1757
const scalar_t* __restrict__ A,
@@ -1859,10 +1796,10 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1796
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1797
// const int max_lds_len = get_lds_size() / 2;
1798
1862
-#define WVSPLITKrc(_WvPrGrp, _YTILE, _UNRL, _N, _GrpsShrB) \
1799
+#define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \
1800
{ \
1864
- dim3 block(64, _WvPrGrp); \
1865
- wvSplitKrc_<fptype, 64, _YTILE, _WvPrGrp, 8, _UNRL, _N, _GrpsShrB> \
1801
+ dim3 block(64, 4); \
1802
+ wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK> \
1803
<<<grid, block, 0, stream>>>(N_in, K_in, M_in, Bx_in, By_in, af4, bf4, \
1804
biasf4, glbl, c, CuCount); \
1805
}
@@ -1877,15 +1814,37 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1814
: nullptr;
1815
fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
1816
auto glbl = axl_glbl.data_ptr<float>();
1817
+
1818
+ // With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
1819
+ // and each working on a 512-shard of K, how many CUs would we need?
1820
+ int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
1821
+
1822
+ // How many of 4 waves in a group can work on same 16 Ms at same time? First
1823
+ // try to maximize this. This reduces the Ms each group works on, i.e.
1824
+ // increasing the number of CUs needed.
1825
+ int GrpsShrB = min(N_p2 / 16, 4);
1826
+
1827
+ // Given the above, how many CUs would we need?
1828
+ int CuNeeded = rndup_cus * GrpsShrB;
1829
+
1830
+ if (CuNeeded > CuCount) std::runtime_error("Invalid wvSplitKrc size");
1831
+
1832
+ // Can we increase SplitK by shrinking the K-shared to 256?
1833
+ int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
1834
+
1835
switch (N_p2) {
1836
case 16:
1882
- WVSPLITKrc(4, 16, 1, 16, 1) break;
1837
+ WVSPLITKrc(16, 1, 1) break;
1838
case 32:
1884
- WVSPLITKrc(4, 16, 1, 32, 2) break;
1839
+ if (chunkk == 2)
1840
+ WVSPLITKrc(32, 2, 2) else if (chunkk == 1) WVSPLITKrc(32, 2, 1) break;
1841
case 64:
1886
- WVSPLITKrc(4, 16, 1, 64, 2) break;
1842
+ if (chunkk == 2)
1843
+ WVSPLITKrc(64, 4, 2) else if (chunkk == 1) WVSPLITKrc(64, 4, 1) break;
1844
case 128:
1888
- WVSPLITKrc(4, 16, 1, 128, 4) break;
1845
+ if (chunkk == 2)
1846
+ WVSPLITKrc(128, 4, 2) else if (chunkk == 1)
1847
+ WVSPLITKrc(128, 4, 1) break;
1848
default:
1849
throw std::runtime_error(
1850
"Unsupported N value: " + std::to_string(M_in) + "," +
tests/kernels/quantization/test_rocm_skinny_gemms.py
+52
-31
@@ -45,31 +45,28 @@ NKM_FACTORS_WVSPLITK = [
45
(4, 256, 8),
46
]
47
48
-NKM_FACTORS_WVSPLITKRC = [
49
- (16, 2880, 128),
50
- (16, 2880, 640),
51
- (17, 2880, 128),
52
- (17, 2880, 640),
53
- (25, 2880, 128),
54
- (25, 2880, 640),
55
- (31, 2880, 128),
56
- (31, 2880, 640),
57
- (32, 2880, 128),
58
- (32, 2880, 640),
59
- (40, 2880, 128),
60
- (40, 2880, 640),
61
- (60, 2880, 128),
62
- (60, 2880, 640),
63
- (64, 2880, 128),
64
- (64, 2880, 640),
65
- (81, 2880, 128),
66
- (81, 2880, 640),
67
- (98, 2880, 128),
68
- (98, 2880, 640),
69
- (128, 2880, 128),
70
- (128, 2880, 640),
48
+N_FACTORS_WVSPLITKRC = [
49
+ 13,
50
+ 16,
51
+ 17,
52
+ 25,
53
+ 29,
54
+ 31,
55
+ 32,
56
+ 41,
57
+ 51,
58
+ 64,
59
+ 71,
60
+ 81,
61
+ 91,
62
+ 103,
63
+ 117,
64
+ 128,
65
]
66
67
+K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
68
+M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
69
+
70
NKM_FACTORS_WVSPLITK_FP8 = [
71
# FP8-specific cases with K % 16 == 0
72
(1, 16, 16),
@@ -113,30 +110,54 @@ def pad_fp8(weight):
110
return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
111
112
116
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC)
113
+@pytest.mark.parametrize("xnorm", [False, True])
114
+@pytest.mark.parametrize("n", N_FACTORS_WVSPLITKRC)
115
+@pytest.mark.parametrize("k", K_FACTORS_WVSPLITKRC)
116
+@pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC)
117
@pytest.mark.parametrize("dtype", DTYPES)
118
@pytest.mark.parametrize("seed", SEEDS)
119
@pytest.mark.parametrize("bias_mode", BIAS_MODES)
120
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
121
@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
122
-def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, seed, bias_mode):
122
+def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
123
torch.manual_seed(seed)
124
cu_count = get_cu_count()
125
126
- xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
127
- A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
128
- B = (torch.rand(m, k, dtype=dtype, device="cuda") - 0.5) * xavier
126
+ # Next ^2 of n
127
+ N_p2 = 1 << (n - 1).bit_length()
128
+ # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
129
+ # and each working on a 512-shard of K, how many CUs would we need?
130
+ rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512)
131
+ # How many of 4 waves in a group can work on same 16 Ms at same time?
132
+ # This reduces the Ms each group works on, i.e. increasing the number of CUs needed.
133
+ GrpsShrB = min(N_p2 // 16, 4)
134
+ # Given the above, how many CUs would we need?
135
+ CuNeeded = rndup_cus * GrpsShrB
136
+ # candidate for atomic reduce count splitk?
137
+ fits_wvsplitkrc = CuNeeded <= cu_count
138
+
139
+ if not fits_wvsplitkrc:
140
+ pytest.skip("Too large for wvSplitKrc")
141
+
142
+ xavier = (
143
+ math.sqrt(2 / k) if xnorm else 1
144
+ ) # normalize to avoid large output-bias deltas
145
+ A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
146
+ B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
147
148
BIAS = None
149
if bias_mode == 1:
132
- BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5
150
+ BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
151
elif bias_mode == 2:
134
- BIAS = torch.rand(n, m, dtype=dtype, device="cuda") - 0.5
152
+ BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
153
154
ref_out = torch.nn.functional.linear(A, B, BIAS)
155
out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
156
139
- assert torch.allclose(out, ref_out, rtol=0.01)
157
+ if xnorm:
158
+ assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
159
+ else:
160
+ assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
161
162
163
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
vllm/model_executor/layers/utils.py
+19
-8
@@ -145,32 +145,43 @@ def rocm_unquantized_gemm_impl(
145
) -> torch.Tensor:
146
from vllm.platforms.rocm import on_gfx9, on_gfx950
147
148
- n = x.numel() / x.size(-1)
148
+ n = x.numel() // x.size(-1)
149
m = weight.shape[0]
150
k = weight.shape[1]
151
152
- import math
153
-
152
+ cu_count = get_cu_count()
153
if use_aiter_triton_gemm(n, m, k, x.dtype):
154
from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
155
156
return gemm_a16w16(x, weight, bias)
157
158
+ # Next ^2 of n
159
+ N_p2 = 1 << (n - 1).bit_length()
160
+ # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
161
+ # and each working on a 512-shard of K, how many CUs would we need?
162
+ rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512)
163
+ # How many of 4 waves in a group can work on same 16 Ms at same time?
164
+ # This reduces the Ms each group works on, i.e. increasing the number of CUs needed.
165
+ GrpsShrB = min(N_p2 // 16, 4)
166
+ # Given the above, how many CUs would we need?
167
+ CuNeeded = rndup_cus * GrpsShrB
168
+ # candidate for atomic reduce count splitk?
169
+ fits_wvsplitkrc = CuNeeded <= cu_count
170
+
171
use_skinny_reduce_counting = (
172
envs.VLLM_ROCM_USE_SKINNY_GEMM
173
and on_gfx950()
174
and x.dtype in [torch.float16, torch.bfloat16]
175
and (
164
- n >= 16
165
- and n <= 128
176
+ 10 <= n <= 128
177
+ and k % 8 == 0
178
and k > 512
167
- and math.ceil(k / 512) * math.ceil(m / 16) < get_cu_count()
179
+ and m % 16 == 0
180
+ and fits_wvsplitkrc
181
and x.is_contiguous()
182
)
170
- # k == 2880 and (m == 640 or m == 128))
183
)
184
if use_skinny_reduce_counting:
173
- cu_count = get_cu_count()
185
x_view = x.reshape(-1, x.size(-1))
186
out = ops.wvSplitKrc(weight, x_view, cu_count, bias)
187
return out.reshape(*x.shape[:-1], weight.shape[0])