feat(rocm): wvSplitK skinny GEMM for decode M<=4 — the #487 decode GEMM lever

CMakeLists.txt

@@ -1369,6 +1369,7 @@ if(VLLM_CPP_HIP)
     src/vt/rocm/rocm_gdn_postconv.hip
     src/vt/rocm/rocm_gdn_scan.hip
     src/vt/rocm/rocm_gdn_fused.hip
+    src/vt/rocm/rocm_skinny_gemm.hip
     src/vt/rocm/rocm_ops.hip)
   if(VLLM_CPP_HIP_ARCHITECTURES)
     set_source_files_properties(
@@ -1389,6 +1390,7 @@ if(VLLM_CPP_HIP)
       src/vt/rocm/rocm_gdn_postconv.hip
       src/vt/rocm/rocm_gdn_scan.hip
       src/vt/rocm/rocm_gdn_fused.hip
+      src/vt/rocm/rocm_skinny_gemm.hip
       src/vt/rocm/rocm_ops.hip
       PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}")
   endif()

docs/USAGE.md

@@ -54,7 +54,16 @@ MoE-path coverage is partial: `MoeRouterTopK` (f32/bf16 logits, ungrouped
 softmax, no bias) and `MoeSiluMul` are native; the remaining chain
 (`kSharedExpertGate`, `kMoeCombine`/`kMoeCombineGate`, and the grouped quant
 expert GEMM) is not registered yet, so MoE-bearing models still throw on
-those ops. On a
+those ops.
+
+### ROCm decode GEMM routing (wvSplitK skinny path)
+
+Decode-shaped GEMMs (M<=4, bf16) route to a split-K skinny-GEMM kernel (a port
+of vLLM's `wvSplitK`) instead of the 128x128-tile rocBLAS GEMM that dominates
+decode GPU time ([#487](https://github.com/mudler/vllm.cpp/issues/487)). On by
+default where it fits; `VT_ROCM_SKINNY=0` restores the BLAS path for A/B.
+
+On a
 discrete card there is no CPU fallback tier, so a model whose layers call an op
 that is not registered yet fails loudly with `vt: no kernel for op N on device
 type 5` — that is the memory-safety design working, not a crash. Run with

scripts/env-doc-allowlist.txt

@@ -15,6 +15,7 @@ VT_GEMMA4_PROFILE
 VT_ROCM_GEMM_COMPUTE
 VT_ROCM_GEMV
 VT_ROCM_HIPBLASLT
+VT_ROCM_SKINNY
 VLLM_MM_TOWER_PROFILE
 VT_ARCH_TACTIC_STATS
 VT_ATTN_DECODE_D128
@@ -154,6 +155,7 @@ VT_RMSNORM_GATED_FAST
 VT_ROCM_GEMM_COMPUTE
 VT_ROCM_GEMV
 VT_ROCM_HIPBLASLT
+VT_ROCM_SKINNY
 VT_SILU_FP4_FAST
 VT_SWIZZLE_IN_QUANT
 VT_W4A4_TRUE
@@ -233,3 +235,4 @@ VT_GEMMA4_PROFILE
 VT_ROCM_GEMM_COMPUTE
 VT_ROCM_GEMV
 VT_ROCM_HIPBLASLT
+VT_ROCM_SKINNY

src/vt/rocm/rocm_matmul_hipblaslt.hip

@@ -137,6 +137,17 @@ bool GemvEnabled() {
   return on;
 }
 
+// wvSplitK skinny GEMM (rocm_skinny_gemm.hip, #487): the split-K/LDS-staged
+// decode path that beats the 128x128-tile rocBLAS GEMM at M<=4. Default ON for
+// decode-skinny shapes; VT_ROCM_SKINNY=0 restores the BLAS path for A/B.
+bool SkinnyGemmEnabled() {
+  static const bool on = [] {
+    if (const char* e = std::getenv("VT_ROCM_SKINNY")) return e[0] != '0';
+    return true;
+  }();
+  return on;
+}
+
 // y[n] = alpha * dot(x[0:K], W[n,0:K]) + beta * y[n]
 // grid = N — one block per output row; x cached in LDS; block-reduce over K.
 __global__ void Bf16GemvBTRowKernel(__hip_bfloat16* __restrict__ y,
@@ -437,6 +448,11 @@ void MatmulKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) {
 }
 
 // out[M,N] = a[M,K] @ b[N,K]^T
+// wvSplitK skinny-GEMM host entry (rocm_skinny_gemm.hip, #487). External
+// vt::rocm linkage to match the definition; declared beside its only caller.
+void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N,
+                int K, int device);
+
 // Row-major trick: gemm(OP_T, OP_N, N, M, K, B, K, A, a_rs, C, N)
 //   BLAS: C = op(A)*op(B) with opA=T => A is KxN in col form = row B[N,K]
 //   opB=N, B is KxM col = row A[M,K] with ld=a_rs
@@ -465,6 +481,16 @@ void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b)
     throw std::runtime_error("vt rocm: matmul_bt: bad a stride");
   }
 
+  // Decode-skinny (M<=4) bf16: the wvSplitK port wins over the 128-tile GEMM
+  // (#487). Gate on the LDS-fit the sml variant needs (K*M activation fits the
+  // 64KB staging buffer); everything else stays on the BLAS path.
+  if (bf16 && out.dtype == DType::kBF16 && M >= 1 && M <= 4 && (K % 8) == 0 &&
+      a.stride[0] == K && K * M <= 32768 && SkinnyGemmEnabled()) {
+    WvSplitKBT(s, out.data, a.data, b.data, static_cast<int>(M), static_cast<int>(N),
+               static_cast<int>(K), q.device.index);
+    return;
+  }
+
   // Decode: M=1 BF16 GEMV
   if (M == 1 && bf16 && out.dtype == DType::kBF16 && a.stride[0] == K && GemvEnabled()) {
     Bf16GemvBT(s, out.data, a.data, b.data, static_cast<int>(N), static_cast<int>(K), 1.f, 0.f);

src/vt/rocm/rocm_skinny_gemm.hip

@@ -0,0 +1,175 @@
+// ROCm skinny GEMM for decode (BACKEND-ROCM; issue #487).
+// Port of vLLM's wvSplitK_hf_sml_ (csrc/rocm/skinny_gemms.cu:351-573) — the
+// split-K, LDS-staged, CU-count-aware kernel that wins M<=4 shapes, replacing
+// the 128x128-macro-tile rocBLAS GEMM that dominates decode GPU time. gfx1100
+// (GFX1X/wave32) bf16. De-torched; semantics mirror the donor exactly.
+//
+// Layout mapping to our MatmulBT (out[M,N] = a[M,K] @ b[N,K]^T, decode M<=4):
+//   donor in_a (weight [N,K]) = our b;  donor in_b (x [M,K]) = our a;
+//   donor M_in = our N (output dim);  donor N_in = our M;  C[N_in,M_in] = out[M,N].
+
+#include <hip/hip_bf16.h>
+#include <hip/hip_runtime.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+
+#include "vt/ops.h"
+
+namespace vt::rocm {
+namespace {
+
+constexpr int kLdsSize = 64 * 1024;  // gfx1100 (non-gfx95x)
+constexpr int kThrds = 32;           // wave32
+constexpr int kWvPrGrp = 16;
+constexpr int kYtile = 2;            // proven decode config (donor gfx1x, N small)
+constexpr int kUnrl = 2;
+constexpr int kAChunk = 8;
+
+using scalar8 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
+union bigType {
+  __hip_bfloat16 h[kAChunk];
+  float f[kAChunk / 2];
+  scalar8 h8;
+};
+
+__device__ __forceinline__ unsigned int min__(uint32_t a, uint32_t b) {
+  return min(a, b);
+}
+
+// Donor's persistent-workgroup splitter (host-side in skinny_gemms.cu:1169).
+inline int mindiv(int N, int div1, int div2) {
+  int nPrRnd = div1 * div2;
+  int rnds[13];
+  for (int i = 0; i < 13; i++) {
+    rnds[i] = (N + nPrRnd - 1) / nPrRnd;
+    nPrRnd -= div1;
+  }
+  for (int i = 12; i >= 0; i--)
+    if (rnds[0] == rnds[i]) return (div2 - i);
+  return 0;
+}
+
+// A (activation) fits LDS. N = decode batch (our M). bf16, f32 accum.
+template <int N>
+__global__ void __launch_bounds__(kWvPrGrp * kThrds)
+    wvSplitKSml(const int K, const int Kbp, const int Kap, const int M,
+                const __hip_bfloat16* __restrict__ B, const __hip_bfloat16* __restrict__ A,
+                __hip_bfloat16* C, const int _WvPrGrp, const int CuCount) {
+  constexpr int max_lds_len = kLdsSize / 2;  // bf16 elements
+  __shared__ __hip_bfloat16 s[max_lds_len];
+
+  // Stage the activation row(s) A [N,K] into LDS.
+  for (uint32_t k = (threadIdx.y * kThrds + threadIdx.x) * kAChunk;
+       k < min__(Kap * N, max_lds_len); k += kThrds * kWvPrGrp * kAChunk) {
+    *((bigType*)(&s[k])) = *((const bigType*)(&A[k]));
+  }
+  __syncthreads();
+
+  if (threadIdx.y >= _WvPrGrp) return;
+
+  uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * kYtile;
+
+  while (m < static_cast<uint32_t>(M)) {
+    float sum[N][kYtile] = {};
+    for (uint32_t k1 = 0; k1 < static_cast<uint32_t>(K); k1 += kThrds * kAChunk * kUnrl) {
+      bigType bigA[N][kUnrl] = {};
+      bigType bigB[kYtile][kUnrl];
+#pragma unroll
+      for (uint32_t k2 = 0; k2 < kUnrl; k2++) {
+        uint32_t k = k1 + k2 * kThrds * kAChunk;
+        uint32_t k_ = k + threadIdx.x * kAChunk;
+        const __hip_bfloat16* B_ = &B[min__(k_, K - kAChunk)];
+        for (int y = 0; y < kYtile; y++)
+          bigB[y][k2].h8 = __builtin_nontemporal_load(
+              (const scalar8*)(&B_[min__(y + m, M - 1) * Kbp]));
+      }
+#pragma unroll
+      for (uint32_t k2 = 0; k2 < kUnrl; k2++) {
+        uint32_t k = k1 + k2 * kThrds * kAChunk;
+        uint32_t k_ = k + threadIdx.x * kAChunk;
+        if (k_ >= static_cast<uint32_t>(K)) break;
+        for (int n = 0; n < N; n++) bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
+      }
+      // Interleaved MAC; bf16 pairs unpacked to f32 (donor DOT2C bf16 branch).
+      for (uint32_t k2 = 0; k2 < kUnrl; k2++) {
+        for (int n = 0; n < N; n++) {
+          for (int y = 0; y < kYtile; y++) {
+#pragma unroll
+            for (uint32_t b = 0; b < kAChunk / 2; b++) {
+              float2 a2 = __bfloat1622float2(*((__hip_bfloat162*)(&(bigA[n][k2].h[b * 2]))));
+              float2 b2 = __bfloat1622float2(*((__hip_bfloat162*)(&(bigB[y][k2].h[b * 2]))));
+              sum[n][y] += (a2.x * b2.x) + (a2.y * b2.y);
+            }
+          }
+        }
+      }
+    }
+    __builtin_amdgcn_sched_barrier(0);
+    // Wave32 reduction: DPP row_shr 8/4/2/1 then shfl_xor(16).
+    for (int n = 0; n < N; n++) {
+      for (int y = 0; y < kYtile; y++) {
+        sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, 1);
+        sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, 1);
+        sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, 1);
+        sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, 1);
+        sum[n][y] += __shfl_xor(sum[n][y], 16);
+      }
+    }
+    if (threadIdx.x == (kThrds - 1)) {
+      for (int n = 0; n < N; n++)
+        for (int y = 0; y < kYtile; y++) C[m + y + n * M] = __float2bfloat16(sum[n][y]);
+    }
+    m += CuCount * _WvPrGrp * kYtile;
+  }
+}
+
+int DeviceCuCount(int device) {
+  static const int cache = [] {
+    int dev = 0;
+    if (hipGetDevice(&dev) != hipSuccess) return 96;  // gfx1100 fallback
+    int n = 0;
+    if (hipDeviceGetAttribute(&n, hipDeviceAttributeMultiprocessorCount, dev) != hipSuccess ||
+        n <= 0)
+      return 96;
+    return n;
+  }();
+  (void)device;
+  return cache;
+}
+
+}  // namespace
+
+// out[M,N] = a[M,K] @ b[N,K]^T, bf16 in/out. Only called for the decode-skinny
+// gate below (M in 1..4, K%8==0); the caller keeps every other shape on the
+// BLAS path.
+void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N,
+                int K, int device) {
+  const int cu = DeviceCuCount(device);
+  dim3 grid(cu), block(kThrds, kWvPrGrp);
+  const int wvPrGrp = mindiv(N, cu * kYtile, kWvPrGrp);
+  auto* C = static_cast<__hip_bfloat16*>(out);
+  auto* A = static_cast<const __hip_bfloat16*>(a);
+  auto* B = static_cast<const __hip_bfloat16*>(b);
+  switch (M) {
+    case 1:
+      wvSplitKSml<1><<<grid, block, 0, s>>>(K, K, K, N, B, A, C, wvPrGrp, cu);
+      break;
+    case 2:
+      wvSplitKSml<2><<<grid, block, 0, s>>>(K, K, K, N, B, A, C, wvPrGrp, cu);
+      break;
+    case 3:
+      wvSplitKSml<3><<<grid, block, 0, s>>>(K, K, K, N, B, A, C, wvPrGrp, cu);
+      break;
+    case 4:
+      wvSplitKSml<4><<<grid, block, 0, s>>>(K, K, K, N, B, A, C, wvPrGrp, cu);
+      break;
+    default:
+      throw std::runtime_error("vt rocm: wvSplitK unsupported M=" + std::to_string(M));
+  }
+  if (hipGetLastError() != hipSuccess)
+    throw std::runtime_error("vt rocm: wvSplitK launch failed");
+}
+
+}  // namespace vt::rocm

tests/vt/test_backend_cross_device.cpp

@@ -1979,6 +1979,62 @@ TEST_CASE("MoeRouterTopK matches the CPU oracle (f32 and bf16 logits)") {
   }
 }
 
+TEST_CASE("decode-skinny MatmulBT (wvSplitK path) matches the CPU oracle") {
+  // The M<=4 bf16 decode GEMM routes to the wvSplitK port on ROCm (#487).
+  // Real decode shapes; bf16 in/out. NMSE vs the CPU oracle (split-K reduction
+  // order differs from the CPU sequential sum).
+  for (int64_t M : {1, 4}) {
+    for (auto [N, K] : {std::pair<int64_t,int64_t>{5120, 1024}, {1024, 2048}}) {
+      CAPTURE(M);
+      CAPTURE(N);
+      CAPTURE(K);
+      const size_t an = static_cast<size_t>(M) * K, bn = static_cast<size_t>(N) * K;
+      const std::vector<float> a = RandomVec(an, 991);
+      const std::vector<float> b = RandomVec(bn, 992, -0.5f, 0.5f);
+      const std::vector<uint16_t> a_bf = Bf16Bits(a), b_bf = Bf16Bits(b);
+
+      std::vector<uint16_t> ref(static_cast<size_t>(M) * N, 0);
+      {
+        vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU);
+        Queue cq = cpu.CreateQueue();
+        const Device cd{DeviceType::kCPU, 0};
+        std::vector<uint16_t> ca = a_bf, cb = b_bf;
+        Tensor ta = Tensor::Contiguous(ca.data(), DType::kBF16, cd, {M, K});
+        Tensor tb = Tensor::Contiguous(cb.data(), DType::kBF16, cd, {N, K});
+        Tensor to = Tensor::Contiguous(ref.data(), DType::kBF16, cd, {M, N});
+        vt::MatmulBT(cq, to, ta, tb);
+        cpu.DestroyQueue(cq);
+      }
+      for (DeviceType dt : RegisteredDevices()) {
+        if (!OpAvailable(vt::OpId::kMatmulBT, dt)) continue;
+        CAPTURE(DeviceName(dt));
+        vt::Backend& dev = vt::GetBackend(dt);
+        Queue q = dev.CreateQueue();
+        const Device d{dt, 0};
+        DevBufBytes da(dev, q, an * 2), db(dev, q, bn * 2), dout(dev, q, static_cast<size_t>(M) * N * 2);
+        da.Upload(a_bf.data());
+        db.Upload(b_bf.data());
+        Tensor ta = Tensor::Contiguous(da.ptr(), DType::kBF16, d, {M, K});
+        Tensor tb = Tensor::Contiguous(db.ptr(), DType::kBF16, d, {N, K});
+        Tensor to = Tensor::Contiguous(dout.ptr(), DType::kBF16, d, {M, N});
+        vt::MatmulBT(q, to, ta, tb);
+        std::vector<uint16_t> got(static_cast<size_t>(M) * N);
+        dout.Download(got.data());
+        // bf16 outputs; compare as f32 NMSE.
+        std::vector<float> gf(got.size()), rf(got.size());
+        for (size_t i = 0; i < got.size(); ++i) {
+          uint32_t ug = static_cast<uint32_t>(got[i]) << 16, ur = static_cast<uint32_t>(ref[i]) << 16;
+          std::memcpy(&gf[i], &ug, 4);
+          std::memcpy(&rf[i], &ur, 4);
+        }
+        CHECK(Nmse(rf, gf) <= kNmseTol);
+        dev.DestroyQueue(q);
+      }
+    }
+  }
+}
+
+
 TEST_CASE("reference tier: an op with no native kernel matches the CPU oracle (unified only)") {
   constexpr int64_t kRows = 7, kCols = 48;
   constexpr size_t kN = kRows * kCols;