From c80c3bb260cdb9454c0cf102b61b0d9daaa65972 Mon Sep 17 00:00:00 2001 From: Vikash Loomba Date: Wed, 12 Aug 2026 10:12:33 -0700 Subject: [PATCH] =?UTF-8?q?feat(rocm):=20wvSplitK=20skinny=20GEMM=20for=20?= =?UTF-8?q?decode=20M<=3D4=20=E2=80=94=20RDNA3=20decode=20GEMM=20routing?= =?UTF-8?q?=20(#487)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports vLLM's wvSplitK_hf_sml_ (csrc/rocm/skinny_gemms.cu:351-573) — the split-K, LDS-staged, CU-count-aware skinny GEMM that wins M<=4 shapes — and routes decode-skinny MatmulBT (M 1..4, bf16, K%8==0, activation fits the 64KB LDS stage) to it instead of the 128x128-macro-tile rocBLAS GEMM. gfx1100 (GFX1X/wave32), bf16. VT_ROCM_SKINNY=0 restores the BLAS path for A/B. Measured on 4x RX 7900 XTX (gfx1100), ROCm 7.14, Release: - kernel-level vs the current rocBLAS tile path, same buffers back-to-back: qkv(5120x1024) 2.52x, o_proj(1024x2048) 3.47x, mlp_gateup(3072x1024) 1.78x, lm_head(151936x1024) 3.64x (the issue's worst single case) - in-engine decode, Qwen3-0.6B 128-token steady state: 88.1 vs 70.4 tok/s (+25%) with VT_ROCM_SKINNY on vs off; 0.8B GDN model output unchanged - cross-device: new decode-skinny MatmulBT case green (8/8, NMSE vs CPU) FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: pi:kimi-k3 [pi] --- CMakeLists.txt | 2 + docs/USAGE.md | 11 +- scripts/env-doc-allowlist.txt | 3 + src/vt/rocm/rocm_matmul_hipblaslt.hip | 26 ++++ src/vt/rocm/rocm_skinny_gemm.hip | 175 +++++++++++++++++++++++++ tests/vt/test_backend_cross_device.cpp | 56 ++++++++ 6 files changed, 272 insertions(+), 1 deletion(-) create mode 100644 src/vt/rocm/rocm_skinny_gemm.hip diff --git a/CMakeLists.txt b/CMakeLists.txt index 1357fce58..f915f9240 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/docs/USAGE.md b/docs/USAGE.md index 635e3e55c..643e0934e 100644 --- a/docs/USAGE.md +++ b/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 diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index e66695e24..bf87eca97 100644 --- a/scripts/env-doc-allowlist.txt +++ b/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 diff --git a/src/vt/rocm/rocm_matmul_hipblaslt.hip b/src/vt/rocm/rocm_matmul_hipblaslt.hip index 43919747b..3f8c831ae 100644 --- a/src/vt/rocm/rocm_matmul_hipblaslt.hip +++ b/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(M), static_cast(N), + static_cast(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(N), static_cast(K), 1.f, 0.f); diff --git a/src/vt/rocm/rocm_skinny_gemm.hip b/src/vt/rocm/rocm_skinny_gemm.hip new file mode 100644 index 000000000..411ea672e --- /dev/null +++ b/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 +#include + +#include +#include +#include + +#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 +__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(M)) { + float sum[N][kYtile] = {}; + for (uint32_t k1 = 0; k1 < static_cast(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(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(a); + auto* B = static_cast(b); + switch (M) { + case 1: + wvSplitKSml<1><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + case 2: + wvSplitKSml<2><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + case 3: + wvSplitKSml<3><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + break; + case 4: + wvSplitKSml<4><<>>(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 diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index a56116eb4..71580ec6e 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/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{5120, 1024}, {1024, 2048}}) { + CAPTURE(M); + CAPTURE(N); + CAPTURE(K); + const size_t an = static_cast(M) * K, bn = static_cast(N) * K; + const std::vector a = RandomVec(an, 991); + const std::vector b = RandomVec(bn, 992, -0.5f, 0.5f); + const std::vector a_bf = Bf16Bits(a), b_bf = Bf16Bits(b); + + std::vector ref(static_cast(M) * N, 0); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector 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(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 got(static_cast(M) * N); + dout.Download(got.data()); + // bf16 outputs; compare as f32 NMSE. + std::vector gf(got.size()), rf(got.size()); + for (size_t i = 0; i < got.size(); ++i) { + uint32_t ug = static_cast(got[i]) << 16, ur = static_cast(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;