@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7a1030431a
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 10, 01:30 UTC

Atomics Reduce Counting Optimization for SplitK Skinny GEMMs. (#29843)

Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>

Hashem Hashemi committed Jan 16, 2026 at 09:45 UTC 7a1030431ade33439902f5b0cf7ff4a69c0fc033
6 files changed +635 -10
csrc/rocm/ops.h
+4
@@ -9,6 +9,10 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
9 const std::optional<at::Tensor>& in_bias,
10 const int64_t CuCount);
11
12 +torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
13 + const std::optional<at::Tensor>& in_bias,
14 + const int64_t CuCount);
15 +
16 void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
17 const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
18 const at::Tensor& scale_a, const at::Tensor& scale_b,
csrc/rocm/skinny_gemms.cu
+545 -9
@@ -287,6 +287,11 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
287 V0 += (s.x + s.y); \
288 }
289
290 +// To avoid LLVM silently upcasting to double
291 +__device__ inline unsigned int min__(uint32_t a, uint32_t b) {
292 + return min(a, b);
293 +}
294 +
295 #if defined(__HIP__GFX9__) // TODO: Add NAVI support
296 // This version targets cases where A[] fits LDS capacity
297 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
@@ -334,11 +339,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
339 // - Then the WG will move to another 8 K elements
340 // TODO: Logic below will only work when K is multiple of 8
341 //----------------------------------------------------
337 - for (uint32_t k = 0; k < min(K * N, max_lds_len);
342 + for (uint32_t k = 0; k < min__(K * N, max_lds_len);
343 k += THRDS * WvPrGrp * A_CHUNK) {
344 uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
345
341 - if (k_in >= min(K * N, max_lds_len)) break;
346 + if (k_in >= min__(K * N, max_lds_len)) break;
347
348 *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
349 }
@@ -633,11 +638,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
638 // - Then the WG will move to another 8 K elements
639 // TODO: Logic below will only work when K is multiple of 8
640 //----------------------------------------------------
636 - for (uint32_t k = 0; k < min(K * N, max_lds_len);
641 + for (uint32_t k = 0; k < min__(K * N, max_lds_len);
642 k += THRDS * WvPrGrp * A_CHUNK) {
643 uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
644
640 - if (k_in >= min(K * N, max_lds_len)) break;
645 + if (k_in >= min__(K * N, max_lds_len)) break;
646
647 *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
648 }
@@ -954,11 +959,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
959 //----------------------------------------------------
960 #define PCML
961 #ifndef PCML
957 - for (uint32_t k = 0; k < min(K * N, max_lds_len);
962 + for (uint32_t k = 0; k < min__(K * N, max_lds_len);
963 k += THRDS * WvPrGrp * A_CHUNK) {
964 uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
965
961 - if (k_in >= min(K * N, max_lds_len)) break;
966 + if (k_in >= min__(K * N, max_lds_len)) break;
967
968 *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
969 }
@@ -975,7 +980,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
980 ? kFit
981 : (kFit - kFit % TUC); // round up to multiple of TUC
982 // if (kFit == 0) kFit = TUC;
978 - kFit = min(kFit, K);
983 + kFit = min__(kFit, K);
984
985 float sum[N][YTILE];
986 scalar8 sum4[N][YTILE];
@@ -1251,6 +1256,7 @@ int mindiv(int N, int div1, int div2) {
1256 }
1257 for (int i = 12; i >= 0; i--)
1258 if (rnds[0] == rnds[i]) return (div2 - i);
1259 + return 0;
1260 }
1261
1262 torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
@@ -1352,6 +1358,536 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1358 return out_c;
1359 }
1360
1361 +#if defined(__gfx950__) // TODO: Add NAVI support
1362 + // This version targets big A[] cases, where it is much larger than LDS
1363 + // capacity
1364 + #define WVSPLITKRC_1KPASS
1365 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1366 + int UNRL, int N, int GrpsShrB>
1367 +
1368 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
1369 + __attribute__((amdgpu_waves_per_eu(1, 1)))
1370 + wvSplitKrc_(const int actlN, const int K, const int M, const int Bx,
1371 + const int By, const scalar_t* __restrict__ B,
1372 + const scalar_t* __restrict__ A,
1373 + const scalar_t* __restrict__ BIAS, float* glbl, scalar_t* C,
1374 + const int CuCount) {
1375 + // Use upper half of glbl buffer for atomic reduce counting
1376 + int* cntr = (int*)(&glbl[M * N]);
1377 +
1378 + constexpr int NTILE = 16;
1379 + constexpr int WVLDS_ = (NTILE * THRDS * A_CHUNK);
1380 + constexpr int APAD = 1;
1381 + constexpr int ASTRD = 64;
1382 + constexpr int BPAD = 1;
1383 + constexpr int BSTRD = 64;
1384 + constexpr int WVLDS = ((WVLDS_ + (WVLDS_ / BSTRD) * 4 * BPAD));
1385 +
1386 + constexpr int max_lds_len = LDS_SIZE / 2;
1387 +
1388 + using scalar16 =
1389 + __attribute__((__vector_size__((A_CHUNK * 2) * sizeof(float)))) float;
1390 + using scalar8 =
1391 + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
1392 + using half4 =
1393 + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
1394 + union bigType {
1395 + scalar_t h[A_CHUNK];
1396 + float f[A_CHUNK / 2];
1397 + unsigned int i[A_CHUNK / 2];
1398 + float2 f2[A_CHUNK / 4];
1399 + unsigned long l[A_CHUNK / 4];
1400 + double d[A_CHUNK / 4];
1401 + half4 h4[A_CHUNK / 4];
1402 + scalar8 h8;
1403 + };
1404 + using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16;
1405 +
1406 + __shared__ scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
1407 + unsigned int* myStg = (unsigned int*)(&stg[WVLDS * (threadIdx.y / GrpsShrB)]);
1408 + __shared__ scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
1409 +
1410 + #ifndef WVSPLITKRC_1KPASS
1411 + constexpr int TUC_ = (THRDS * UNRL * A_CHUNK);
1412 + // find biggest k size that fits padded into LDS
1413 + constexpr uint32_t kFit__ = (max_lds_len - WvPrGrp * WVLDS / GrpsShrB) / N;
1414 + constexpr uint32_t kFit_ = (kFit__ * ASTRD) / (APAD + ASTRD);
1415 + uint32_t kFit = kFit_ - (kFit_ % TUC_);
1416 + uint32_t kfitsPerRdc = (K + kFit - 1) / kFit;
1417 +
1418 + // find best k split to fill the CUs
1419 + if (((K + kfitsPerRdc * kFit - 1) / (kfitsPerRdc * kFit)) * numCuWithFullK <=
1420 + CuCount)
1421 + while (true) {
1422 + while (kFit > TUC_) {
1423 + uint32_t kFit_ = kFit - TUC_;
1424 + if (((K + (kfitsPerRdc * kFit_ - 1)) / (kfitsPerRdc * kFit_)) *
1425 + numCuWithFullK >
1426 + CuCount)
1427 + break;
1428 + kFit = kFit_;
1429 + }
1430 + if (((K + ((kfitsPerRdc - 1) * kFit - 1)) / ((kfitsPerRdc - 1) * kFit)) *
1431 + numCuWithFullK <=
1432 + CuCount)
1433 + kfitsPerRdc--;
1434 + else
1435 + break;
1436 + }
1437 + #else
1438 + int constexpr kFit = 512;
1439 + int constexpr kfitsPerRdc = 1;
1440 + #endif
1441 +
1442 + bool doRdc = (kfitsPerRdc * kFit < K);
1443 + uint32_t numCuWithFullK =
1444 + ((M + (WvPrGrp * YTILE / GrpsShrB) - 1) / (WvPrGrp * YTILE / GrpsShrB));
1445 + uint32_t Mmod = numCuWithFullK * (WvPrGrp * YTILE / GrpsShrB);
1446 +
1447 + // given above k-split, find this wave's position
1448 + uint32_t kFitPdd = kFit + (kFit / ASTRD) * APAD;
1449 + uint32_t m0 = (blockIdx.x * WvPrGrp / GrpsShrB) * YTILE;
1450 + uint32_t m1 = ((threadIdx.y % WvPrGrp) / GrpsShrB) * YTILE;
1451 + uint32_t m = (m0 + m1) % Mmod;
1452 + const uint32_t k_str = (m0 / Mmod) * kFit * kfitsPerRdc;
1453 + uint32_t k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
1454 + const uint32_t k_rnd = (K + kFit * kfitsPerRdc - 1) / (kFit * kfitsPerRdc);
1455 +
1456 + scalar8 sum4[N / NTILE / GrpsShrB][1];
1457 + bigType bigB_[YTILE / GrpsShrB][UNRL];
1458 + const uint32_t bLoader = (threadIdx.y % GrpsShrB);
1459 + uint32_t kBase = 0;
1460 + if (k_str >= K) return;
1461 + if (m >= Mmod) return;
1462 +
1463 + bool noreloada = false;
1464 + constexpr bool FAST_UNSAFE_RDC_INIT = false;
1465 +
1466 + #ifdef WVSPLITKRC_1KPASS
1467 + // Early glbl init, B[] loading, if 1KPASS
1468 + if constexpr (FAST_UNSAFE_RDC_INIT) {
1469 + if (m + (threadIdx.x % 16) < M)
1470 + if (doRdc)
1471 + if (k_str == 0) {
1472 + int mindx = m + (threadIdx.x % 16);
1473 + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1474 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1475 + int adr_ = mindx + M * nindx_ / 4;
1476 + __hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED,
1477 + __HIP_MEMORY_SCOPE_AGENT);
1478 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1479 + for (uint32_t j = 0; j < 4; j++) {
1480 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1481 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1482 + int adr = mindx + M * nindx;
1483 + __hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED,
1484 + __HIP_MEMORY_SCOPE_AGENT);
1485 + }
1486 + }
1487 + }
1488 + }
1489 +
1490 + // Load first B[] chunk
1491 + #pragma unroll
1492 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1493 + uint32_t k = k_str + k2 * THRDS * A_CHUNK;
1494 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1495 + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1496 + #pragma unroll
1497 + for (uint32_t y = 0; y < YTILE / GrpsShrB; y++)
1498 + bigB_[y][k2].h8 = (loadnt(
1499 + (scalar8*)(&B_[min__(y * GrpsShrB + bLoader + m, M - 1) * K])));
1500 + }
1501 + {
1502 + #else
1503 + while (m < Mmod) {
1504 + #endif
1505 +
1506 + #ifndef WVSPLITKRC_1KPASS
1507 + if constexpr (FAST_UNSAFE_RDC_INIT) {
1508 + if (m + (threadIdx.x % 16) < M)
1509 + if (doRdc)
1510 + if (k_str == 0) {
1511 + int mindx = m + (threadIdx.x % 16);
1512 + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1513 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1514 + int adr_ = mindx + M * nindx_ / 4;
1515 + __hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED,
1516 + __HIP_MEMORY_SCOPE_AGENT);
1517 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1518 + for (uint32_t j = 0; j < 4; j++) {
1519 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1520 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1521 + int adr = mindx + M * nindx;
1522 + __hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED,
1523 + __HIP_MEMORY_SCOPE_AGENT);
1524 + }
1525 + }
1526 + }
1527 + }
1528 +
1529 + #endif
1530 +
1531 + #ifndef WVSPLITKRC_1KPASS
1532 + for (uint32_t k1 = k_str; k1 < k_end; k1 += THRDS * A_CHUNK * UNRL) {
1533 + #else
1534 + const uint32_t k1 = k_str;
1535 + {
1536 + #endif
1537 + #ifndef WVSPLITKRC_1KPASS
1538 + const bool reloada = (!noreloada) &&
1539 + ((k1 == k_str) || (k1 == k_str + kBase + kFit)) &&
1540 + (k1 < k_end);
1541 + // load next chunk of A[] to LDS
1542 + if (reloada) {
1543 + if (k1 != k_str) kBase += kFit;
1544 + __syncthreads();
1545 + #else
1546 + const bool reloada = (!noreloada) &&
1547 + ((k1 == k_str) || (k1 == k_str + kBase + kFit)) &&
1548 + (k1 < k_end);
1549 + if (reloada) {
1550 + #endif
1551 + constexpr int sprdN = 4;
1552 + const uint32_t thrd = ((threadIdx.y / sprdN) * THRDS + threadIdx.x);
1553 +
1554 + #ifndef WVSPLITKRC_1KPASS
1555 + #pragma unroll
1556 + for (int k = 0; k < kFit; k += THRDS * (WvPrGrp / sprdN) * A_CHUNK) {
1557 + #else
1558 + const unsigned int k = 0;
1559 + {
1560 + #endif
1561 + unsigned int kOff = k + (thrd * A_CHUNK);
1562 + unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
1563 + const unsigned int k_in = kOffcp + ((threadIdx.y % sprdN)) * K;
1564 + const unsigned int k_ot = kOff + ((threadIdx.y % sprdN)) * kFitPdd;
1565 + for (unsigned int n = 0; n < N / 2; n += sprdN) {
1566 + __builtin_amdgcn_global_load_lds((int*)(&A[k_in + n * K]),
1567 + (int*)(&s[(k_ot + n * kFitPdd)]),
1568 + 16, 0, 0);
1569 + if (((threadIdx.y % sprdN)) + n + N / 2 >= actlN) continue;
1570 + __builtin_amdgcn_global_load_lds(
1571 + (int*)(&A[k_in + (n + N / 2) * K]),
1572 + (int*)(&s[(k_ot + (n + N / 2) * kFitPdd)]), 16, 0, 0);
1573 + }
1574 +
1575 + // Stage loaded B[] to LDS for MFMA swizzling...
1576 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1577 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1578 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1579 + const bool oob_k = (k_ >= K);
1580 + for (uint32_t y = 0; y < YTILE / GrpsShrB; y++) {
1581 + uint32_t idx = threadIdx.x * 4 +
1582 + (y * GrpsShrB + bLoader) * ((THRDS + BPAD) * 4);
1583 + // zero out if oob
1584 + *((scalar8*)&myStg[idx]) =
1585 + (oob_k || (y * GrpsShrB + bLoader + m >= M))
1586 + ? 0
1587 + : bigB_[y][k2].h8;
1588 + }
1589 + }
1590 + }
1591 + }
1592 + }
1593 +
1594 + #ifndef WVSPLITKRC_1KPASS
1595 + // Fire load of next B[] chunk...
1596 + if ((k1 + THRDS * A_CHUNK * UNRL < k_end) &&
1597 + (k1 + THRDS * A_CHUNK * UNRL < K))
1598 + #pragma unroll
1599 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1600 + uint32_t k = k1 + THRDS * A_CHUNK * UNRL + k2 * THRDS * A_CHUNK;
1601 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1602 + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1603 + #pragma unroll
1604 + for (uint32_t y = 0; y < YTILE / GrpsShrB; y++)
1605 + bigB_[y][k2].h8 = (loadnt(
1606 + (scalar8*)(&B_[min__(y * GrpsShrB + bLoader + m, M - 1) * K])));
1607 + }
1608 + #endif
1609 +
1610 + // B[] staging is cooperative across GrpsShrB, so sync here before reading
1611 + // back
1612 + __syncthreads();
1613 +
1614 + // read back B[] swizzled for MFMA...
1615 + bigType bigB[YTILE][UNRL];
1616 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1617 + for (uint32_t y = 0; y < YTILE; y++) {
1618 + unsigned int idx = (threadIdx.x % YTILE) * ((THRDS + BPAD) * 4) +
1619 + (threadIdx.x / YTILE) * 4 + y * 16;
1620 + bigB[y][k2].h8 = *((scalar8*)&myStg[idx]);
1621 + }
1622 + }
1623 +
1624 + // rReadback A[] swizzled for MFMA...
1625 + bigType bigA[N / GrpsShrB][UNRL];
1626 + #pragma unroll
1627 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1628 + uint32_t k = k1 + k2 * THRDS * A_CHUNK - kBase - k_str;
1629 + #pragma unroll
1630 + for (uint32_t nt = 0; nt < N / GrpsShrB; nt += NTILE)
1631 + #pragma unroll
1632 + for (uint32_t n = 0; n < NTILE; n++) {
1633 + uint32_t idxa = (nt + (threadIdx.x % NTILE) +
1634 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) *
1635 + kFitPdd +
1636 + A_CHUNK * ((threadIdx.x / NTILE) + n * 4) + k;
1637 + bigA[nt + n][k2] = *((const bigType*)(&(s[idxa])));
1638 + }
1639 + }
1640 +
1641 + // Do the MFMAs
1642 + #pragma unroll
1643 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1644 + #pragma unroll
1645 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1646 + if constexpr (std::is_same_v<scalar_t, half>) {
1647 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1648 + bigA[nt * NTILE + 0][k2].h4[0], bigB[0][k2].h4[0],
1649 + (k1 == k_str) ? ((scalar8){0}) : sum4[nt][0], 0, 0, 0);
1650 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1651 + bigA[nt * NTILE + 0][k2].h4[1], bigB[0][k2].h4[1], sum4[nt][0], 0,
1652 + 0, 0);
1653 + } else { // bf16
1654 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
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_16x16x16bf16_1k(
1658 + bigA[nt * NTILE + 0][k2].h4[1], bigB[0][k2].h4[1], sum4[nt][0], 0,
1659 + 0, 0);
1660 + }
1661 + #pragma unroll
1662 + for (uint32_t j = 1; j < YTILE; j++) {
1663 + if constexpr (std::is_same_v<scalar_t, half>) {
1664 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1665 + bigA[nt * NTILE + j][k2].h4[0], bigB[j][k2].h4[0], sum4[nt][0],
1666 + 0, 0, 0);
1667 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16f16(
1668 + bigA[nt * NTILE + j][k2].h4[1], bigB[j][k2].h4[1], sum4[nt][0],
1669 + 0, 0, 0);
1670 + } else { // bf16
1671 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(
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_16x16x16bf16_1k(
1675 + bigA[nt * NTILE + j][k2].h4[1], bigB[j][k2].h4[1], sum4[nt][0],
1676 + 0, 0, 0);
1677 + }
1678 + }
1679 + }
1680 + }
1681 + }
1682 +
1683 + if (!doRdc) {
1684 + if (m + (threadIdx.x % 16) < M) {
1685 + scalar_t biases[N / NTILE / GrpsShrB][4] = {0};
1686 + if (BIAS)
1687 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1688 + for (uint32_t j = 0; j < 4; j++) {
1689 + int mindx = m + (threadIdx.x % 16);
1690 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1691 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1692 + biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * M];
1693 + }
1694 + }
1695 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1696 + for (uint32_t j = 0; j < 4; j++) {
1697 + int mindx = m + (threadIdx.x % 16);
1698 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1699 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1700 + int adr = mindx + M * nindx;
1701 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1702 + if (BIAS) sum4[nt][0][j] += __bfloat162float(biases[nt][j]);
1703 + C[adr] = __float2bfloat16(sum4[nt][0][j]);
1704 + } else {
1705 + if (BIAS) sum4[nt][0][j] += __half2float(biases[nt][j]);
1706 + C[adr] = __float2half(sum4[nt][0][j]);
1707 + }
1708 + }
1709 + }
1710 + }
1711 + } else {
1712 + if (m + (threadIdx.x % 16) < M) {
1713 + int my_cntr;
1714 + if (!BIAS) {
1715 + int mindx = m + (threadIdx.x % 16);
1716 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
1717 + for (uint32_t j = 0; j < 4; j++) {
1718 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1719 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1720 + int adr = mindx + M * nindx;
1721 + atomicAdd(&glbl[adr], sum4[nt][0][j]);
1722 + }
1723 + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1724 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1725 + int adr_ = mindx + M * nindx_ / 4;
1726 + my_cntr = atomicAdd(&cntr[adr_], 1);
1727 + float vals[N / NTILE / GrpsShrB][4] = {};
1728 + if (my_cntr + 1 == k_rnd) {
1729 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1730 + for (uint32_t j = 0; j < 4; j++) {
1731 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1732 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1733 + int adr = mindx + M * nindx;
1734 + vals[nt][j] = glbl[adr];
1735 + }
1736 + }
1737 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1738 + for (uint32_t j = 0; j < 4; j++) {
1739 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1740 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1741 + if (nindx >= actlN) break;
1742 + int adr = mindx + M * nindx;
1743 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1744 + C[adr] = __float2bfloat16(vals[nt][j]);
1745 + } else {
1746 + C[adr] = __float2half(vals[nt][j]);
1747 + }
1748 + }
1749 + }
1750 + }
1751 + } else {
1752 + int mindx = m + (threadIdx.x % 16);
1753 + scalar_t biases[N / NTILE / GrpsShrB][4] = {};
1754 + // Atomic add the output, read biases
1755 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
1756 + for (uint32_t j = 0; j < 4; j++) {
1757 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1758 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1759 + int adr = mindx + M * nindx;
1760 + atomicAdd(&glbl[adr], sum4[nt][0][j]);
1761 + biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * M];
1762 + }
1763 + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1764 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1765 + int adr_ = mindx + M * nindx_ / 4;
1766 + // Update the complete counter
1767 + my_cntr = atomicAdd(&cntr[adr_], 1);
1768 + float vals[N / NTILE / GrpsShrB][4] = {};
1769 + // If we're the last k-shard, read back the value and convert...
1770 + if (my_cntr + 1 == k_rnd) {
1771 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1772 + for (uint32_t j = 0; j < 4; j++) {
1773 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1774 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1775 + int adr = mindx + M * nindx;
1776 + vals[nt][j] = glbl[adr];
1777 + }
1778 + }
1779 + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1780 + for (uint32_t j = 0; j < 4; j++) {
1781 + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1782 + (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1783 + if (nindx >= actlN) break;
1784 + int adr = mindx + M * nindx;
1785 + if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1786 + vals[nt][j] += __bfloat162float(biases[nt][j]);
1787 + C[adr] = __float2bfloat16(vals[nt][j]);
1788 + } else {
1789 + vals[nt][j] += __half2float(biases[nt][j]);
1790 + C[adr] = __float2half(vals[nt][j]);
1791 + }
1792 + }
1793 + }
1794 + }
1795 + }
1796 + }
1797 +
1798 + #ifndef WVSPLITKRC_1KPASS
1799 + m0 += CuCount * WvPrGrp * YTILE / GrpsShrB;
1800 + m = (m0 + m1) % Mmod;
1801 + k_str = (m0 / Mmod) * kFit * kfitsPerRdc;
1802 + k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
1803 + if (k_str >= K) break;
1804 + kBase = 0;
1805 + #endif
1806 + }
1807 +}
1808 +#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1809 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1810 + int UNRL, int N, int GrpsShrB>
1811 +__global__ void wvSplitKrc_(const int actlN, const int K, const int M,
1812 + const int Bx, const int By, const scalar_t* B,
1813 + const scalar_t* __restrict__ A,
1814 + const scalar_t* __restrict__ BIAS, float* glbl,
1815 + // int* cntr,
1816 + scalar_t* C, const int CuCount){UNREACHABLE_CODE}
1817 +#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1818 +
1819 +torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1820 + const std::optional<at::Tensor>& in_bias,
1821 + const int64_t CuCount) {
1822 + auto M_in = in_a.size(0);
1823 + auto N_in = in_b.size(0);
1824 + auto K_in = in_a.size(1);
1825 + auto Bx_in =
1826 + (in_bias.has_value() && in_bias->numel() > 0)
1827 + ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
1828 + : 1;
1829 + auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
1830 + in_bias->sizes().size() == 2)
1831 + ? in_bias->size(0)
1832 + : 1;
1833 +
1834 + TORCH_CHECK(in_a.dtype() == in_b.dtype());
1835 + TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
1836 + TORCH_CHECK(in_a.dtype() == torch::kFloat16 ||
1837 + in_a.dtype() == torch::kBFloat16);
1838 +
1839 + auto out_c = torch::empty(
1840 + {N_in, M_in},
1841 + torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
1842 +
1843 + auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1));
1844 + auto axl_glbl = torch::empty(
1845 + {N_p2 + N_p2 / 4, M_in + M_in / 4},
1846 + torch::TensorOptions().dtype(torch::kFloat32).device(in_b.device()));
1847 + axl_glbl.zero_(); // disable for FAST_UNSAFE_RDC_INIT
1848 +
1849 + dim3 grid(CuCount);
1850 +
1851 + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1852 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1853 + // const int max_lds_len = get_lds_size() / 2;
1854 +
1855 +#define WVSPLITKrc(_WvPrGrp, _YTILE, _UNRL, _N, _GrpsShrB) \
1856 + { \
1857 + dim3 block(64, _WvPrGrp); \
1858 + wvSplitKrc_<fptype, 64, _YTILE, _WvPrGrp, 8, _UNRL, _N, _GrpsShrB> \
1859 + <<<grid, block, 0, stream>>>(N_in, K_in, M_in, Bx_in, By_in, af4, bf4, \
1860 + biasf4, glbl, c, CuCount); \
1861 + }
1862 +
1863 + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitKrc", [&] {
1864 + using fptype = typename scalar<scalar_t>::type;
1865 + fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
1866 + const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
1867 + const fptype* biasf4 =
1868 + (in_bias.has_value() && in_bias->numel() > 0)
1869 + ? reinterpret_cast<const fptype*>(in_bias->data_ptr())
1870 + : nullptr;
1871 + fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
1872 + auto glbl = axl_glbl.data_ptr<float>();
1873 + switch (N_p2) {
1874 + case 16:
1875 + WVSPLITKrc(4, 16, 1, 16, 1) break;
1876 + case 32:
1877 + WVSPLITKrc(4, 16, 1, 32, 2) break;
1878 + case 64:
1879 + WVSPLITKrc(4, 16, 1, 64, 2) break;
1880 + case 128:
1881 + WVSPLITKrc(4, 16, 1, 128, 4) break;
1882 + default:
1883 + throw std::runtime_error(
1884 + "Unsupported N value: " + std::to_string(M_in) + "," +
1885 + std::to_string(K_in) + "," + std::to_string(N_in));
1886 + }
1887 + });
1888 + return out_c;
1889 +}
1890 +
1891 #if defined(__HIP__MI3XX__) // TODO: Add NAVI support
1892 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1893 int A_CHUNK, int UNRL, int N>
@@ -1381,7 +1917,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1917 __shared__ fp8_t s[max_lds_len];
1918
1919 for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1384 - k < min(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
1920 + k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
1921 *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
1922 }
1923 __syncthreads();
@@ -1570,7 +2106,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
2106 __shared__ fp8_t s[max_lds_len];
2107
2108 for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1573 - k < min(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
2109 + k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
2110 *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
2111 }
2112 __syncthreads();
csrc/rocm/torch_bindings.cpp
+6
@@ -26,6 +26,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
26 "Tensor");
27 rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
28
29 + // Custom gemm op for skinny matrix-matrix multiplication
30 + rocm_ops.def(
31 + "wvSplitKrc(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
32 + "Tensor");
33 + rocm_ops.impl("wvSplitKrc", torch::kCUDA, &wvSplitKrc);
34 +
35 // wvSplitK for fp8
36 rocm_ops.def(
37 "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor? in_bias, Tensor! out_c, "
tests/kernels/quantization/test_rocm_skinny_gemms.py
+53
@@ -8,9 +8,11 @@ import torch
8 import vllm._custom_ops as ops
9 from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
10 from vllm.platforms import current_platform
11 +from vllm.platforms.rocm import on_gfx950
12 from vllm.utils.platform_utils import get_cu_count
13
14 DTYPES = [torch.bfloat16, torch.float16]
15 +BIAS_MODES = [0, 1, 2]
16 # Specific (N, K, M) combinations for targeted testing
17 NKM_FACTORS_LLMM1 = [
18 # Small, medium, large cases
@@ -43,6 +45,31 @@ 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),
71 +]
72 +
73 NKM_FACTORS_WVSPLITK_FP8 = [
74 # FP8-specific cases with K % 16 == 0
75 (1, 16, 16),
@@ -60,6 +87,32 @@ NKM_FACTORS_WVSPLITK_FP8 = [
87 SEEDS = [0]
88
89
90 +@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC)
91 +@pytest.mark.parametrize("dtype", DTYPES)
92 +@pytest.mark.parametrize("seed", SEEDS)
93 +@pytest.mark.parametrize("bias_mode", BIAS_MODES)
94 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
95 +@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
96 +def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, seed, bias_mode):
97 + torch.manual_seed(seed)
98 + cu_count = get_cu_count()
99 +
100 + xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
101 + A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
102 + B = (torch.rand(m, k, dtype=dtype, device="cuda") - 0.5) * xavier
103 +
104 + BIAS = None
105 + if bias_mode == 1:
106 + BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5
107 + elif bias_mode == 2:
108 + BIAS = torch.rand(n, m, dtype=dtype, device="cuda") - 0.5
109 +
110 + ref_out = torch.nn.functional.linear(A, B, BIAS)
111 + out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
112 +
113 + assert torch.allclose(out, ref_out, rtol=0.01)
114 +
115 +
116 @pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
117 @pytest.mark.parametrize("dtype", DTYPES)
118 @pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16])
vllm/_custom_ops.py
+6
@@ -2072,6 +2072,12 @@ def wvSplitK(
2072 return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
2073
2074
2075 +def wvSplitKrc(
2076 + a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2077 +) -> torch.Tensor:
2078 + return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count)
2079 +
2080 +
2081 def wvSplitKQ(
2082 a: torch.Tensor,
2083 b: torch.Tensor,
vllm/model_executor/layers/utils.py
+21 -1
@@ -129,12 +129,32 @@ def use_aiter_triton_gemm(n, m, k, dtype):
129 def rocm_unquantized_gemm_impl(
130 x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
131 ) -> torch.Tensor:
132 - from vllm.platforms.rocm import on_gfx9
132 + from vllm.platforms.rocm import on_gfx9, on_gfx950
133
134 n = x.numel() / x.size(-1)
135 m = weight.shape[0]
136 k = weight.shape[1]
137
138 + import math
139 +
140 + use_skinny_reduce_counting = (
141 + envs.VLLM_ROCM_USE_SKINNY_GEMM
142 + and on_gfx950()
143 + and x.dtype in [torch.float16, torch.bfloat16]
144 + and (
145 + n >= 16
146 + and n <= 128
147 + and k > 512
148 + and math.ceil(k / 512) * math.ceil(m / 16) < get_cu_count()
149 + )
150 + # k == 2880 and (m == 640 or m == 128))
151 + )
152 + if use_skinny_reduce_counting:
153 + cu_count = get_cu_count()
154 + x_view = x.reshape(-1, x.size(-1))
155 + out = ops.wvSplitKrc(weight, x_view, cu_count, bias)
156 + return out.reshape(*x.shape[:-1], weight.shape[0])
157 +
158 if use_aiter_triton_gemm(n, m, k, x.dtype):
159 from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
160