From e83222396f6c9ceb72f80c0d82782071bbd8317c Mon Sep 17 00:00:00 2001 From: Robert Esclapez Date: Tue, 13 Oct 2026 19:21:30 +0200 Subject: [PATCH] [ROCm][Perf] W4A16: speed up skinny GEMM decode, including speculative batches (#57199) Signed-off-by: Robert Esclapez Garcia Co-authored-by: Claude Opus 4.7 Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- csrc/rocm/skinny_gemms_int4.cu | 536 +++++++++++++++--- .../quantization/test_rdna_hybrid_w4a16.py | 149 ++++- .../mixed_precision/rdna_hybrid_w4a16.py | 30 +- vllm/platforms/rocm.py | 5 + 4 files changed, 603 insertions(+), 117 deletions(-) diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index be3cda7be6..b80a71e36d 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -8,6 +8,7 @@ #include #include +#include #include // std::in_range #include "core/cuda_compat.h" @@ -17,10 +18,27 @@ #define __HIP__GFX1X__ #endif +// The decode tuning and the deep-K routing below were swept on gfx115x and are +// kept to it: the rest of gfx11 and gfx12 also compile this kernel, and keep +// the behaviour and the K * N bound they had. +#if defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || \ + defined(__gfx1153__) + #define __HIP__GFX115X__ +#endif + #define LDS_SIZE 64 * 1024 static constexpr int get_lds_size_int4() { return 64 * 1024; } +// Host-side twin of __HIP__GFX115X__, for the tuple choices the dispatch makes. +static bool on_gfx115x_int4() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + return std::string(dprops->gcnArchName).find("gfx115") != std::string::npos; + }(); + return result; +} + #if defined(NDEBUG) #undef NDEBUG #include @@ -98,6 +116,14 @@ struct scalar { } while (0) #endif +// Folding the zero point into the dequant constants was swept on gfx115x; +// elsewhere the original second pass over the unpacked values is kept. +#if defined(__HIP__GFX115X__) +[[maybe_unused]] constexpr bool FOLD_ZP_INTO_DEQUANT = true; +#else +[[maybe_unused]] constexpr bool FOLD_ZP_INTO_DEQUANT = false; +#endif + __device__ inline unsigned int min__(uint32_t a, uint32_t b) { return min(a, b); } @@ -114,17 +140,70 @@ __device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed, return (w >> (4 * (row & 7))) & 0xFu; } +// A_CHUNK halves per lane is two 16-byte LDS accesses. Laid out linearly lane +// i's first piece lands on banks (8i..8i+3) mod 32, so every fourth lane +// collides and LDS read bandwidth halves. The chunked band therefore stores +// piece-major within a block of THRDS lanes -- piece p of lane i at +// p * THRDS + i -- which tiles all 32 banks across eight lanes. The block +// stride is unchanged, so the permutation stays inside one block. +constexpr uint32_t ACT_LDS_PIECE = 8; // halves per 16-byte LDS access +struct alignas(16) lds_piece { + float f[ACT_LDS_PIECE / 2]; +}; + +// Copies A[kBase, kBase + chunkSpan) of every activation row into a +// KFIT-strided LDS band, so the compute body reads every row from LDS even +// when the whole activation does not fit. +// +// Strides by the runtime _WvPrGrp: the waves above it have already returned, +// so the compile-time WvPrGrp rows are not all there to help. The caller owns +// the barriers. Requires K % A_CHUNK == 0, which the host checks. +template +__device__ __forceinline__ void load_act_chunk_into_lds( + scalar_t* s, const scalar_t* __restrict__ A, const int K, + const uint32_t kBase, const uint32_t chunkSpan, const int _WvPrGrp) { + const uint32_t span = min__(chunkSpan, static_cast(K) - kBase); + const uint32_t step = THRDS * A_CHUNK * _WvPrGrp; + const uint32_t off = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + for (uint32_t k = 0; k < span; k += step) { + uint32_t k_in = k + off; + if (k_in >= span) break; + // k is a multiple of THRDS * A_CHUNK, so the block base is + // k + threadIdx.y * THRDS * A_CHUNK and this thread owns lane + // threadIdx.x of it -- no division needed. + const uint32_t blk = k + threadIdx.y * (THRDS * A_CHUNK); +#pragma unroll + for (int n = 0; n < N; n++) { + // Both sides are 16-byte aligned: k_in and KFIT are multiples of + // A_CHUNK, kBase of KFIT, and K of 16. Unstated, the accesses degrade to + // pairs of 2addr_b32. + const auto* src = (const lds_piece*)__builtin_assume_aligned( + &A[static_cast(K) * n + kBase + k_in], 16); + auto* dst = (lds_piece*)__builtin_assume_aligned( + &s[n * KFIT + blk + threadIdx.x * ACT_LDS_PIECE], 16); +#pragma unroll + for (uint32_t p = 0; p < A_CHUNK / ACT_LDS_PIECE; p++) + dst[p * THRDS] = src[p]; + } + } +} + // W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations -// Targets the "sml" case where activations fit in LDS. +// Targets the "sml" case, where the activation fits in LDS, and with CHUNKED +// the deep-K case, where it is staged there a chunk at a time. // A_CHUNK: number of K-elements processed per thread per step. // Weights are packed as 2 int4 values per byte: [M, K/2] bytes. // GROUP_SIZE: 0 = per-channel scale [M], >0 = per-group scale [M, // K/GROUP_SIZE]. // Requires GROUP_SIZE % A_CHUNK == 0 when GROUP_SIZE > 0. // B_row_stride_bytes is the per-row byte stride of the packed weights. +// CHUNKED walks K in LDS-sized chunks so all N activation rows stay in LDS +// even when K * N exceeds it. LDS_LEN is a template parameter so the band +// stride is constant and folds into the LDS address immediate. #if defined(__HIP__GFX1X__) template + int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false, + bool CHUNKED = false, uint32_t LDS_LEN = LDS_SIZE / 2> __global__ void __launch_bounds__(WvPrGrp* THRDS) wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, @@ -132,8 +211,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount, + const int group_stride, const int B_row_stride_bytes) { - constexpr int max_lds_len = LDS_SIZE / 2; + constexpr int max_lds_len = LDS_LEN; union bigTypeA { scalar_t h[A_CHUNK]; @@ -148,39 +228,79 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) __shared__ scalar_t s[max_lds_len]; - for (uint32_t k = 0; k < min__(K * N, max_lds_len); - k += THRDS * WvPrGrp * A_CHUNK) { - uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); + // CHUNKED fills s[] per chunk instead, and needs the waves above _WvPrGrp + // gone before the first barrier: s_barrier is untagged on gfx11, so a wave + // parked at a barrier it does not own would satisfy a reload barrier. + if constexpr (!CHUNKED) { + for (uint32_t k = 0; k < min__(K * N, max_lds_len); + k += THRDS * WvPrGrp * A_CHUNK) { + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); - if (k_in >= min__(K * N, max_lds_len)) break; + if (k_in >= min__(K * N, max_lds_len)) break; - *((bigTypeA*)(&s[k_in])) = *((bigTypeA*)(&A[k_in])); + *((bigTypeA*)(&s[k_in])) = *((bigTypeA*)(&A[k_in])); + } + __syncthreads(); } - __syncthreads(); if (threadIdx.y >= _WvPrGrp) return; uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; - // For per-group, precompute num_groups and scale stride - [[maybe_unused]] const int num_groups = - (GROUP_SIZE > 0) ? (K / GROUP_SIZE) : 0; + // Per-row LDS band stride for CHUNKED, a whole number of K_STEP steps. + constexpr uint32_t K_STEP = THRDS * A_CHUNK * UNRL; + constexpr uint32_t kFit = (LDS_LEN / N) - (LDS_LEN / N) % K_STEP; + static_assert(!CHUNKED || kFit >= K_STEP, + "CHUNKED needs at least one K_STEP of LDS per activation row"); + static_assert(!CHUNKED || static_cast(kFit) * N <= LDS_LEN, + "CHUNKED LDS band overruns s[]"); + [[maybe_unused]] uint32_t kBase = 0; + + // How much is copied per chunk need not be the band stride, and taking kFit + // every time leaves a short final chunk whose k1 loop cannot amortise its + // own setup. Spread K evenly instead, still in whole K_STEP steps so a chunk + // boundary falls between two k1 iterations. + [[maybe_unused]] uint32_t chunkSpan = kFit; + if constexpr (CHUNKED) { + const uint32_t n_chunks = (static_cast(K) + kFit - 1) / kFit; + const uint32_t even = (static_cast(K) + n_chunks - 1) / n_chunks; + chunkSpan = (even + K_STEP - 1) / K_STEP * K_STEP; + } + + // Live waves must run the m loop the same number of times or they stop + // arriving at the reload barriers. Rounding M up to a whole YTILE * _WvPrGrp + // cannot split a workgroup's window, since the m stride is a multiple of it. + uint32_t mEnd = M; + if constexpr (CHUNKED) { + const uint32_t YW = YTILE * _WvPrGrp; + mEnd = (M % YW == 0) ? M : (M - M % YW + YW); + } + + // Scale and zero-point rows carry their own row stride, which the caller + // may pad beyond K / GROUP_SIZE. + [[maybe_unused]] const int num_groups = (GROUP_SIZE > 0) ? group_stride : 0; float sum[N][YTILE]; - while (m < M) { + while (m < mEnd) { for (int i = 0; i < YTILE; i++) for (int n = 0; n < N; n++) sum[n][i] = 0; - bigTypeA bigA[N][UNRL]; bigTypeW bigB[YTILE][UNRL]; - for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + // Load and compute halves, each instantiated twice on a constexpr CHECK + // flag. Whole K_STEP blocks take the unchecked one, where the UNRL weight + // loads clause together instead of being split by a divergent branch that + // drains vmcnt per step; only the ragged tail pays the check. + auto k_load = [&](auto CHECK_T, bigTypeW(&bigB)[YTILE][UNRL], uint32_t k1) { + constexpr bool CHECK = decltype(CHECK_T)::value; #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { uint32_t k = k1 + k2 * THRDS * A_CHUNK; uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; + if constexpr (CHECK) { + if (k_ >= K) break; + } const uint8_t* B_ = &B_packed[(m + 0) * B_row_stride_bytes + k_ / 2]; for (int y = 0; y < YTILE; y++) { @@ -190,15 +310,37 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) bigB[y][k2].f[i] = loadnt((float*)&src[i]); } } + }; + + auto k_compute = [&](auto CHECK_T, const bigTypeW(&bigB)[YTILE][UNRL], + uint32_t k1) { + constexpr bool CHECK = decltype(CHECK_T)::value; + bigTypeA bigA[N][UNRL]; #pragma unroll for (uint32_t k2 = 0; k2 < UNRL; k2++) { uint32_t k = k1 + k2 * THRDS * A_CHUNK; uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; + if constexpr (CHECK) { + if (k_ >= K) break; + } for (int n = 0; n < N; n++) { - bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n]))); + // Only the LDS address is chunk-relative; the scale and zero-point + // lookups below stay on the absolute k_. + if constexpr (CHUNKED) { + // Piece-major band, see ACT_LDS_PIECE. k, not k_, is the block + // base: the lane offset moves in 16-byte pieces instead of + // A_CHUNK, and successive pieces sit a whole lane-plane apart. + const auto* src = (const lds_piece*)__builtin_assume_aligned( + &s[kFit * n + (k - kBase) + threadIdx.x * ACT_LDS_PIECE], 16); + auto* dst = (lds_piece*)&bigA[n][k2]; + #pragma unroll + for (uint32_t p = 0; p < A_CHUNK / ACT_LDS_PIECE; p++) + dst[p] = src[p * THRDS]; + } else { + bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n]))); + } } } @@ -206,7 +348,30 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (uint32_t k2 = 0; k2 < UNRL; k2++) { uint32_t k = k1 + k2 * THRDS * A_CHUNK; uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; + if constexpr (CHECK) { + if (k_ >= K) break; + } + + // The unpack already subtracts a constant, so the zero point rides + // along in it rather than costing a second A_CHUNK-wide pass: it is + // uniform over a group and A_CHUNK never straddles one. Exact in fp16, + // since one ulp is 1 at 1024 and 1/16 at 64. Hoisted out of the n loop + // because it does not depend on the activation row. + [[maybe_unused]] uint32_t BIAS_LO[YTILE], BIAS_HI[YTILE]; + if constexpr (std::is_same_v) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + BIAS_LO[y] = HAS_ZERO_POINTS ? 0x64006400u : 0x64086408u; + BIAS_HI[y] = HAS_ZERO_POINTS ? 0xD400D400u : 0xD480D480u; + if constexpr (FOLD_ZP_INTO_DEQUANT && HAS_ZERO_POINTS && + GROUP_SIZE > 0) { + const uint32_t zp = + zp_nibble(zero_points, m + y, k_ / GROUP_SIZE, num_groups); + BIAS_LO[y] += zp * 0x00010001u; + BIAS_HI[y] += (zp << 4) * 0x00010001u; + } + } + } #pragma unroll for (uint32_t n = 0; n < N; n++) { @@ -216,14 +381,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) if constexpr (std::is_same_v) { constexpr uint32_t FP16_MAGIC = 0x64006400u; - // When HAS_ZERO_POINTS, store raw nibble values; - // the zero-point subtraction below handles the full shift. - // When symmetric, bake -8 into the constants. - constexpr uint32_t BIAS_LO = - HAS_ZERO_POINTS ? 0x64006400u : 0x64086408u; constexpr uint32_t SCALE16 = 0x2C002C00u; - constexpr uint32_t BIAS_HI = - HAS_ZERO_POINTS ? 0xD400D400u : 0xD480D480u; #pragma unroll for (uint32_t w = 0; w < A_CHUNK / 8; w++) { uint32_t qa = bigB[y][k2].u32[w]; @@ -234,15 +392,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t hi1 = (qa & 0x00F000F0u) | FP16_MAGIC; *(half2*)&cvtB.f[w * 4 + 0] = - __hsub2(*(half2*)&lo0, *(const half2*)&BIAS_LO); + __hsub2(*(half2*)&lo0, *(const half2*)&BIAS_LO[y]); *(half2*)&cvtB.f[w * 4 + 1] = __hfma2(*(half2*)&hi0, *(const half2*)&SCALE16, - *(const half2*)&BIAS_HI); + *(const half2*)&BIAS_HI[y]); *(half2*)&cvtB.f[w * 4 + 2] = - __hsub2(*(half2*)&lo1, *(const half2*)&BIAS_LO); + __hsub2(*(half2*)&lo1, *(const half2*)&BIAS_LO[y]); *(half2*)&cvtB.f[w * 4 + 3] = __hfma2(*(half2*)&hi1, *(const half2*)&SCALE16, - *(const half2*)&BIAS_HI); + *(const half2*)&BIAS_HI[y]); } } else { // bf16 path: marlin-style magic-number trick. @@ -267,7 +425,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } } - if constexpr (!std::is_same_v) { + if constexpr (!FOLD_ZP_INTO_DEQUANT && + !std::is_same_v) { if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) { uint32_t group_idx = k_ / GROUP_SIZE; scalar_t zp = __float2s(static_cast( @@ -323,6 +482,70 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } } } + }; + + // Wide-unroll N = 3 performs better with checked K loads. Off gfx115x + // every tuple keeps the checked loop, as before this was tuned; CHUNKED + // launches on gfx115x only. + #if defined(__HIP__GFX115X__) + constexpr bool CHECKED_K_LOOP = !CHUNKED && (N == 3) && (UNRL >= 4); + #else + constexpr bool CHECKED_K_LOOP = !CHUNKED; + #endif + + static_assert(!CHUNKED || !CHECKED_K_LOOP, + "the chunked loop below assumes the unchecked K body"); + + if constexpr (CHUNKED) { + // The refill stays outside the k1 loop: folded in behind a predicate, + // the compiler stays conservative about lgkmcnt/vmcnt across a barrier + // it cannot prove absent, even where it is skipped. + const uint32_t K_whole = K - (K % K_STEP); + for (uint32_t kb = 0; kb < static_cast(K); kb += chunkSpan) { + __syncthreads(); + load_act_chunk_into_lds( + s, A, K, kb, chunkSpan, _WvPrGrp); + __syncthreads(); + kBase = kb; + // Waves past the real M had to reach the barriers above, but have + // no rows to reduce. + if (m >= M) continue; + const uint32_t kStop = min__(K_whole, kb + chunkSpan); + for (uint32_t k1 = kb; k1 < kStop; k1 += K_STEP) { + k_load(std::false_type{}, bigB, k1); + k_compute(std::false_type{}, bigB, k1); + } + // chunkSpan is a multiple of K_STEP, so the ragged tail falls in + // the final chunk. + if (kb + chunkSpan >= static_cast(K) && + K_whole < static_cast(K)) { + k_load(std::true_type{}, bigB, K_whole); + k_compute(std::true_type{}, bigB, K_whole); + } + } + } else if constexpr (CHECKED_K_LOOP) { + for (uint32_t k1 = 0; k1 < K; k1 += K_STEP) { + k_load(std::true_type{}, bigB, k1); + k_compute(std::true_type{}, bigB, k1); + } + } else { + const uint32_t K_whole = K - (K % K_STEP); + for (uint32_t k1 = 0; k1 < K_whole; k1 += K_STEP) { + k_load(std::false_type{}, bigB, k1); + k_compute(std::false_type{}, bigB, k1); + } + if (K_whole < K) { + k_load(std::true_type{}, bigB, K_whole); + k_compute(std::true_type{}, bigB, K_whole); + } + } + + if constexpr (CHUNKED) { + if (m >= M) { + m += CuCount * _WvPrGrp * YTILE; + kBase = 0; + continue; + } } for (int n = 0; n < N; n++) @@ -344,20 +567,21 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } #else // !defined(__HIP__GFX1X__) template + int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false, + bool CHUNKED = false, uint32_t LDS_LEN = LDS_SIZE / 2> __global__ void wvSplitK_int4_hf_sml_( const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, - const int CuCount, const int B_row_stride_bytes) { + const int CuCount, const int group_stride, const int B_row_stride_bytes) { UNREACHABLE_CODE } #endif // defined(__HIP__GFX1X__) -// W4A16 skinny GEMM "medium" kernel: activation matrix marginally exceeds LDS. -// Loads as much of A into LDS as fits; overflowing rows fall back to global -// memory. Also handles M not divisible by YTILE via commitColumn tracking. +// W4A16 skinny GEMM "medium" kernel: the activation exceeds LDS. Loads as +// much of A into LDS as fits; overflowing rows fall back to global memory. +// Also handles M not divisible by YTILE via commitColumn tracking. #if defined(__HIP__GFX1X__) template @@ -367,7 +591,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) const scalar_t* scale, const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount, - const int B_row_stride_bytes) { + const int group_stride, const int B_row_stride_bytes) { constexpr int max_lds_len = LDS_SIZE / 2; union bigTypeA { @@ -410,8 +634,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) if (threadIdx.y >= _WvPrGrp) return; - [[maybe_unused]] const int num_groups = - (GROUP_SIZE > 0) ? (K / GROUP_SIZE) : 0; + // Scale and zero-point rows carry their own row stride, which the caller + // may pad beyond K / GROUP_SIZE. + [[maybe_unused]] const int num_groups = (GROUP_SIZE > 0) ? group_stride : 0; float sum[N][YTILE]; @@ -458,6 +683,24 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t k_ = k + threadIdx.x * A_CHUNK; if (k_ >= K) break; + // Zero point folded into the unpack constants, as in + // wvSplitK_int4_hf_sml_. + [[maybe_unused]] uint32_t BIAS_LO[YTILE], BIAS_HI[YTILE]; + if constexpr (std::is_same_v) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + BIAS_LO[y] = HAS_ZERO_POINTS ? 0x64006400u : 0x64086408u; + BIAS_HI[y] = HAS_ZERO_POINTS ? 0xD400D400u : 0xD480D480u; + if constexpr (FOLD_ZP_INTO_DEQUANT && HAS_ZERO_POINTS && + GROUP_SIZE > 0) { + const uint32_t zp = + zp_nibble(zero_points, m + y, k_ / GROUP_SIZE, num_groups); + BIAS_LO[y] += zp * 0x00010001u; + BIAS_HI[y] += (zp << 4) * 0x00010001u; + } + } + } + #pragma unroll for (uint32_t n = 0; n < N; n++) { #pragma unroll @@ -466,14 +709,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) if constexpr (std::is_same_v) { constexpr uint32_t FP16_MAGIC = 0x64006400u; - // When HAS_ZERO_POINTS, store raw nibble values; - // the zero-point subtraction below handles the full shift. - // When symmetric, bake -8 into the constants. - constexpr uint32_t BIAS_LO = - HAS_ZERO_POINTS ? 0x64006400u : 0x64086408u; constexpr uint32_t SCALE16 = 0x2C002C00u; - constexpr uint32_t BIAS_HI = - HAS_ZERO_POINTS ? 0xD400D400u : 0xD480D480u; #pragma unroll for (uint32_t w = 0; w < A_CHUNK / 8; w++) { uint32_t qa = bigB[y][k2].u32[w]; @@ -484,15 +720,15 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t hi1 = (qa & 0x00F000F0u) | FP16_MAGIC; *(half2*)&cvtB.f[w * 4 + 0] = - __hsub2(*(half2*)&lo0, *(const half2*)&BIAS_LO); + __hsub2(*(half2*)&lo0, *(const half2*)&BIAS_LO[y]); *(half2*)&cvtB.f[w * 4 + 1] = __hfma2(*(half2*)&hi0, *(const half2*)&SCALE16, - *(const half2*)&BIAS_HI); + *(const half2*)&BIAS_HI[y]); *(half2*)&cvtB.f[w * 4 + 2] = - __hsub2(*(half2*)&lo1, *(const half2*)&BIAS_LO); + __hsub2(*(half2*)&lo1, *(const half2*)&BIAS_LO[y]); *(half2*)&cvtB.f[w * 4 + 3] = __hfma2(*(half2*)&hi1, *(const half2*)&SCALE16, - *(const half2*)&BIAS_HI); + *(const half2*)&BIAS_HI[y]); } } else { // bf16 path: marlin-style magic-number trick. @@ -517,7 +753,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } } - if constexpr (!std::is_same_v) { + if constexpr (!FOLD_ZP_INTO_DEQUANT && + !std::is_same_v) { if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) { uint32_t group_idx = k_ / GROUP_SIZE; scalar_t zp = __float2s(static_cast( @@ -610,7 +847,7 @@ __global__ void wvSplitK_int4_hf_( const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, - const int CuCount, const int B_row_stride_bytes) { + const int CuCount, const int group_stride, const int B_row_stride_bytes) { UNREACHABLE_CODE } #endif // defined(__HIP__GFX1X__) @@ -682,6 +919,17 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, TORCH_CHECK(in_scale.size(0) == M_in && in_scale.size(1) == num_groups, "Scale must be [M, K/group_size] = [", M_in, ", ", num_groups, "] but got [", in_scale.size(0), ", ", in_scale.size(1), "]"); + // Scale and zero-point rows are indexed by their own stride, so the caller + // may pad them past K / group_size. Only dim 0 may be padded; dim 1 must + // stay contiguous because the kernel reads consecutive groups. + const int64_t group_stride = in_scale.stride(0); + TORCH_CHECK(in_scale.stride(1) == 1, "Scale must be contiguous in dim 1"); + TORCH_CHECK(group_stride >= num_groups, + "Scale row stride must be at least K/group_size=", num_groups, + ", got ", group_stride); + TORCH_CHECK(std::in_range(group_stride), + "Scale row stride overflows int"); + const int group_stride_i32 = static_cast(group_stride); if (in_zero_points.has_value()) { // The kernel reads the words as uint32, so either signedness is accepted. TORCH_CHECK(in_zero_points->dtype() == at::kInt || @@ -694,8 +942,11 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, in_zero_points->sizes()); TORCH_CHECK(M_in % 8 == 0, "M must be divisible by 8 for packed zero points, got ", M_in); - TORCH_CHECK(in_zero_points->is_contiguous(), - "Zero points must be contiguous"); + TORCH_CHECK(in_zero_points->stride(1) == 1, + "Zero points must be contiguous in dim 1"); + TORCH_CHECK(in_zero_points->stride(0) == group_stride, + "Zero points must share the scale row stride ", group_stride, + ", got ", in_zero_points->stride(0)); TORCH_CHECK(in_zero_points->size(0) == M_in / 8 && in_zero_points->size(1) == num_groups, "Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ", @@ -704,8 +955,18 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, } TORCH_CHECK(K_in % 16 == 0, "K must be divisible by 16"); + // Every kernel indexes A linearly, which a non-contiguous activation would + // break. + TORCH_CHECK(in_b.is_contiguous(), "Activation must be contiguous"); + const int max_lds_len = get_lds_size_int4() / 2; - TORCH_CHECK(K_in * N_in <= (int64_t)(max_lds_len * 1.2), + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // gfx115x walks K in LDS-sized chunks, so K * N is unbounded there. Elsewhere + // only the medium kernel serves shapes past LDS, and within 1.2x of it. + TORCH_CHECK(on_gfx115x_int4() || K_in * N_in <= (int64_t)(max_lds_len * 1.2), "K*N exceeds LDS capacity (medium limit). K=", K_in, " N=", N_in); auto out_c = torch::empty( @@ -714,26 +975,30 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, dim3 grid(CuCount); - const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - -// Dispatch macro: _HAS_ZP selects the HAS_ZERO_POINTS template parameter -#define WVSPLITK_INT4G_LAUNCH(_THRDS, _YTILE, _UNRL, _N, _GS, _HAS_ZP) \ - { \ - dim3 block(_THRDS, 16); \ - int __wvPrGrp = mindiv_int4(M_in, CuCount * _YTILE, 16); \ - if (K_in * N_in <= max_lds_len && M_in % _YTILE == 0) \ - wvSplitK_int4_hf_sml_<<>>( \ - K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ - __wvPrGrp, CuCount, b_row_stride_bytes_i32); \ - else \ - wvSplitK_int4_hf_<<>>( \ - K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ - __wvPrGrp, CuCount, b_row_stride_bytes_i32); \ +// Dispatch macro: _HAS_ZP selects the HAS_ZERO_POINTS template parameter. +// _W is the workgroup's wave count and _AC the per-thread K chunk. Almost every +// tuple takes the 16/16 default through WVSPLITK_INT4G_LAUNCH; only the deep-K +// N = 1 rules below widen _AC. +#define WVSPLITK_INT4G_LAUNCH_W_AC(_THRDS, _YTILE, _W, _AC, _UNRL, _N, _GS, \ + _HAS_ZP) \ + { \ + dim3 block(_THRDS, _W); \ + int __wvPrGrp = mindiv_int4(M_in, CuCount * _YTILE, _W); \ + if (K_in * N_in <= max_lds_len && M_in % _YTILE == 0) \ + wvSplitK_int4_hf_sml_<<>>( \ + K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ + __wvPrGrp, CuCount, group_stride_i32, b_row_stride_bytes_i32); \ + else \ + wvSplitK_int4_hf_<<>>( \ + K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ + __wvPrGrp, CuCount, group_stride_i32, b_row_stride_bytes_i32); \ } +#define WVSPLITK_INT4G_LAUNCH(_THRDS, _YTILE, _UNRL, _N, _GS, _HAS_ZP) \ + WVSPLITK_INT4G_LAUNCH_W_AC(_THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, _HAS_ZP) + #define WVSPLITK_INT4G(_YTILE, _UNRL, _N, _GS, _HAS_ZP) \ WVSPLITK_INT4G_LAUNCH(32, _YTILE, _UNRL, _N, _GS, _HAS_ZP) @@ -745,27 +1010,104 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, else \ WVSPLITK_INT4G(_YTILE, _UNRL, _N, 128, _HAS_ZP) -#define WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \ - { \ - if (K_in * N_in > max_lds_len) { \ - if (_sYT < 30) \ - WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \ - else \ - WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - } else if (__N >= 4 && _sYT >= 480) \ - WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT >= 40) \ - WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \ - WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ - else if (__N >= 3 && _sYT < 40) \ - WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else if (__N >= 2) \ - WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ - else if (_sYT >= 30) \ - WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ - else \ - WVSPLIT_INT4G_GS(1, 4, __N, _HAS_ZP) \ +#define WVSPLITK_INT4G_W_AC(_YTILE, _UNRL, _W, _AC, _N, _GS, _HAS_ZP) \ + WVSPLITK_INT4G_LAUNCH_W_AC(32, _YTILE, _W, _AC, _UNRL, _N, _GS, _HAS_ZP) + +#define WVSPLIT_INT4G_GS_W_AC(_YTILE, _UNRL, _W, _AC, _N, _HAS_ZP) \ + if (group_size == 32) \ + WVSPLITK_INT4G_W_AC(_YTILE, _UNRL, _W, _AC, _N, 32, _HAS_ZP) \ + else if (group_size == 64) \ + WVSPLITK_INT4G_W_AC(_YTILE, _UNRL, _W, _AC, _N, 64, _HAS_ZP) \ + else \ + WVSPLITK_INT4G_W_AC(_YTILE, _UNRL, _W, _AC, _N, 128, _HAS_ZP) + +// Launch for shapes whose activation does not fit LDS in one go, falling back +// to the medium kernel when M rules the chunked path out. Its own macro so the +// CHUNKED instantiations exist only for the tuples this branch selects. +#define WVSPLITK_INT4G_LAUNCH_CHUNKED(_THRDS, _YTILE, _UNRL, _N, _GS, _HAS_ZP) \ + { \ + dim3 block(_THRDS, 16); \ + int __wvPrGrp = mindiv_int4(M_in, CuCount * _YTILE, 16); \ + /* Gate on what the kernel allocates, not the host's max_lds_len. */ \ + constexpr int lds_elems = LDS_SIZE / 2; \ + constexpr int k_step = (_THRDS) * (_UNRL) * 16; \ + constexpr int k_fit = (lds_elems / (_N)) - (lds_elems / (_N)) % k_step; \ + /* K % A_CHUNK == 0 follows from the K % 16 check above. */ \ + if (M_in % (_YTILE) == 0 && k_fit >= k_step) \ + wvSplitK_int4_hf_sml_ \ + <<>>(K_in, M_in, Bx_in, By_in, wptr, aptr, \ + sptr, zpptr, biasptr, cptr, __wvPrGrp, \ + CuCount, group_stride_i32, \ + b_row_stride_bytes_i32); \ + else \ + wvSplitK_int4_hf_<<>>( \ + K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ + __wvPrGrp, CuCount, group_stride_i32, b_row_stride_bytes_i32); \ + } + +#define WVSPLITK_INT4G_CHUNKED(_YTILE, _UNRL, _N, _GS, _HAS_ZP) \ + WVSPLITK_INT4G_LAUNCH_CHUNKED(32, _YTILE, _UNRL, _N, _GS, _HAS_ZP) + +#define WVSPLIT_INT4G_GS_CHUNKED(_YTILE, _UNRL, _N, _HAS_ZP) \ + if (group_size == 32) \ + WVSPLITK_INT4G_CHUNKED(_YTILE, _UNRL, _N, 32, _HAS_ZP) \ + else if (group_size == 64) \ + WVSPLITK_INT4G_CHUNKED(_YTILE, _UNRL, _N, 64, _HAS_ZP) \ + else \ + WVSPLITK_INT4G_CHUNKED(_YTILE, _UNRL, _N, 128, _HAS_ZP) + +#define WVSPLIT_INT4G_TILE(_sYT, __N, _HAS_ZP) \ + { \ + if (K_in * N_in > max_lds_len && !on_gfx115x_int4()) { \ + /* Off gfx115x only the medium kernel serves this, up to 1.2x LDS. */ \ + if (_sYT < 30) \ + WVSPLIT_INT4G_GS(4, 2, __N, _HAS_ZP) \ + else \ + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + } else if (K_in * N_in > max_lds_len) { \ + if (_sYT < 30) \ + WVSPLIT_INT4G_GS_CHUNKED(4, 2, __N, _HAS_ZP) \ + else if (mindiv_int4(M_in, CuCount * 4, 16) < 16) \ + WVSPLIT_INT4G_GS_CHUNKED(2, 2, __N, _HAS_ZP) \ + else \ + WVSPLIT_INT4G_GS_CHUNKED(4, 1, __N, _HAS_ZP) \ + } else if (on_gfx115x_int4() && __N >= 3 && _sYT >= 40 && K_in >= 4096 && \ + mindiv_int4(M_in, CuCount * 4, 16) < 16) \ + /* YTILE = 4 is better here, but mindiv_int4 can trim WvPrGrp \ + harder at YTILE = 4 than at YTILE = 2, and idle wave rows \ + cost more than the tile gains. The K bound keeps this off \ + shallow rows, where the k1 loop cannot amortise the extra \ + m-tiles. */ \ + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if (__N >= 4 && _sYT >= 480) \ + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT >= 40) \ + WVSPLIT_INT4G_GS(4, 1, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT < 40 && (K_in <= 2048 || K_in >= 4096)) \ + WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ + else if (__N >= 3 && _sYT < 40) \ + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if (__N >= 2) \ + WVSPLIT_INT4G_GS(2, 2, __N, _HAS_ZP) \ + else if (on_gfx115x_int4() && __N == 1 && K_in == 4096) \ + /* A_CHUNK = 32 doubles the per-thread load granularity, and a \ + K = 4096 row has the work to amortise it. _W stays at 16 \ + because A_CHUNK = 32 with _W = 32 spills. */ \ + WVSPLIT_INT4G_GS_W_AC(1, 4, 16, 32, __N, _HAS_ZP) \ + else if (on_gfx115x_int4() && __N == 1 && !(_HAS_ZP) && \ + K_in % (32 * 32 * 8) == 0) \ + /* Deep-K rows dividing this tuple's K step exactly, where UNRL = 8 \ + gives each row a 4 KB contiguous run. A K leaving a large ragged \ + tail has too few k1 iterations to hide it. Symmetric only: the \ + tuple sits at the edge of the register budget and the zero-point \ + lookup pushes it over, costing 13.5% asymmetric at K = 8192. */ \ + WVSPLIT_INT4G_GS_W_AC(2, 8, 16, 32, __N, _HAS_ZP) \ + else if (_sYT >= 30) \ + WVSPLIT_INT4G_GS(2, 4, __N, _HAS_ZP) \ + else \ + WVSPLIT_INT4G_GS(1, 4, __N, _HAS_ZP) \ } // Inner dispatch: shared by both symmetric and asymmetric paths @@ -812,9 +1154,15 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, WVSPLIT_INT4G_DISPATCH(false) }); +#undef WVSPLITK_INT4G_LAUNCH_W_AC +#undef WVSPLITK_INT4G_W_AC +#undef WVSPLIT_INT4G_GS_W_AC #undef WVSPLITK_INT4G_LAUNCH #undef WVSPLITK_INT4G #undef WVSPLIT_INT4G_GS +#undef WVSPLITK_INT4G_LAUNCH_CHUNKED +#undef WVSPLITK_INT4G_CHUNKED +#undef WVSPLIT_INT4G_GS_CHUNKED #undef WVSPLIT_INT4G_TILE #undef WVSPLIT_INT4G_DISPATCH diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index da7ef2d437..e88d84c5d9 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -19,7 +19,7 @@ if not current_platform.is_rocm(): pytest.importorskip("triton") -from vllm.platforms.rocm import on_gfx1x # noqa: E402 +from vllm.platforms.rocm import on_gfx1x, on_gfx115x # noqa: E402 device = "cuda" @@ -446,22 +446,28 @@ def _hip_skinny_reference( scales_nkg: torch.Tensor, *, group_size: int, - zp_bias: int, + zp_nkg: torch.Tensor | None, ) -> torch.Tensor: - """Reference for symmetric HIP skinny: C = A @ (W - zp_bias) * S.""" + """Reference for HIP skinny: C = A @ (W - zp) * S. + + ``zp_nkg`` is [N, K//G] int32 raw zero points for the asymmetric case, or + None for symmetric uint4b8, where the kernel subtracts a constant 8. + """ K = a_mk.shape[1] N = w_int4_nk.shape[0] num_groups = K // group_size - w_fp = (w_int4_nk.to(torch.float32) - zp_bias).view(N, num_groups, group_size) + w_g = w_int4_nk.to(torch.float32).view(N, num_groups, group_size) + zp = 8.0 if zp_nkg is None else zp_nkg.to(torch.float32).unsqueeze(-1) s = scales_nkg.to(torch.float32).unsqueeze(-1) - w_dequant = (w_fp * s).view(N, K) + w_dequant = ((w_g - zp) * s).view(N, K) return (a_mk.to(torch.float32) @ w_dequant.t()).to(a_mk.dtype) @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True]) @pytest.mark.parametrize( "M,K,N,G", [ @@ -470,13 +476,50 @@ def _hip_skinny_reference( (1, 512, 256, 128), (2, 512, 256, 64), (3, 256, 512, 64), + # M is the kernel's batch dimension; 4 and 5 reach dispatch tuples that + # 1..3 never take, and 5 is MAX_SKINNY_BATCH_SIZE. + (4, 512, 256, 32), + (4, 256, 512, 64), + (5, 512, 256, 128), + (5, 256, 256, 32), + # N == 8 is exactly one packed zero-point word. + (1, 256, 8, 32), + # K must exceed THRDS * A_CHUNK * UNRL for the K loop to run a whole + # unrolled block; every K above is below that bound, so without these + # the deep-K path goes untested. 2560 is not a multiple of the 1024 K + # step the batch 2 and 4 tuples take, so it also hits the ragged tail. + (1, 4096, 256, 128), + (2, 2560, 256, 32), + (3, 4096, 256, 64), + (4, 2560, 256, 128), + (5, 4096, 256, 32), + # On gfx115x, batch 1 with K == 4096, or symmetric with K a multiple of + # 8192, takes a tuple with A_CHUNK = 32 instead of 16. G = 32 is the + # tight case there: each lane's A_CHUNK spans exactly one scale group. + (1, 4096, 256, 32), + (1, 8192, 256, 128), + (1, 16384, 256, 64), + # K * batch beyond what LDS holds (32768 fp16 elements), so gfx115x + # walks K in chunks; elsewhere the op rejects these, as the layer + # routes them to Triton. The per-row window shrinks as the batch grows, + # so these cover a single reload and several, with and without a short + # final chunk. 20992 is not a multiple of the K step either, so it + # crosses a chunk boundary *and* has a ragged K tail. + (2, 20480, 512, 64), + (3, 16384, 512, 32), + (4, 16384, 512, 128), + (5, 21504, 512, 32), + (5, 20992, 512, 128), ], ) -def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G): +def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G, has_zp): """Test HIP wvSplitK_int4_g kernel directly via _custom_ops.""" import vllm._custom_ops as ops from vllm.utils.platform_utils import num_compute_units + if K * M > LDS_CAPACITY_ELEMENTS and not on_gfx115x(): + pytest.skip("deep-K HIP path is gfx115x only") + set_random_seed(0) a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) @@ -489,10 +532,64 @@ def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G): dtype ) - cu_count = num_compute_units() - out = ops.wvSplitK_int4_g(b_packed_i8, a, scales, cu_count, G) + zp_nkg = None + zp_packed = None + if has_zp: + zp_nkg = torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32) + zp_packed = _pack_zp_rows_for_kernel(zp_nkg) - ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8) + cu_count = num_compute_units() + out = ops.wvSplitK_int4_g(b_packed_i8, a, scales, cu_count, G, zp_packed) + + ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_nkg=zp_nkg) + + torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("G", [32, 128]) +@pytest.mark.parametrize("has_zp", [False, True]) +@pytest.mark.parametrize("M", [2, 64], ids=["M=2_decode", "M=64_prefill"]) +def test_hip_skinny_padded_group_stride(dtype, G, has_zp, M): + """Scale and zero-point rows may be padded past K // group_size. + + Both paths index them by their own row stride, so a padded row is read + correctly rather than bleeding into the next row. M picks the path: the + HIP skinny kernel at 2, the Triton prefill kernel at 64. + """ + from vllm.utils.platform_utils import num_compute_units + + set_random_seed(0) + K, N = 2048, 256 + num_groups = K // G + pad = 7 + + a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) + b_packed_i8 = pack_int4_exllama_shuffle(w_int4_nk).view(torch.int8) + + scales = (0.05 * torch.rand((N, num_groups), device=device)).to(dtype) + scales_padded = torch.zeros((N, num_groups + pad), device=device, dtype=dtype) + scales_padded[:, :num_groups] = scales + scales_view = scales_padded[:, :num_groups] + assert scales_view.stride(0) == num_groups + pad + + zp_nkg = zp_view = None + if has_zp: + zp_nkg = torch.randint(0, 16, (N, num_groups), device=device, dtype=torch.int32) + zp_packed = _pack_zp_rows_for_kernel(zp_nkg) + zp_padded = torch.zeros( + (N // 8, num_groups + pad), device=device, dtype=torch.int32 + ) + zp_padded[:, :num_groups] = zp_packed + zp_view = zp_padded[:, :num_groups] + + cu_count = num_compute_units() + out = torch.ops.vllm.rdna_hybrid_w4a16_apply( + a, b_packed_i8, scales_view, zp_view, None, cu_count, G + ) + ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_nkg=zp_nkg) torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) @@ -510,6 +607,8 @@ def test_hip_skinny_wvSplitK_int4_g(dtype, M, K, N, G): (1, 256, 256, 64), (1, 512, 256, 32), (1, 512, 256, 128), + # Past LDS: HIP on gfx115x, Triton elsewhere. + (2, 20480, 256, 64), (32, 512, 256, 64), (64, 1024, 256, 128), ], @@ -535,7 +634,7 @@ def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G): a, b_packed_i8, scales, None, None, cu_count, G ) - ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8) + ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_nkg=None) torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) @@ -748,17 +847,30 @@ def test_process_weights_pads_cliff_rows(K, expected_weight_stride, dist_init): @pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") @pytest.mark.parametrize( - "M,K,expected_path", + "M,K,gfx115x,expected_path", [ - (4, 8192, "hip"), - (4, 9728, "hip"), - (4, 9856, "triton"), - (MAX_SKINNY_BATCH_SIZE + 1, 1024, "triton"), + (4, 8192, False, "hip"), + (4, 9728, False, "hip"), + (4, 9856, False, "triton"), + (4, 9856, True, "hip"), + (MAX_SKINNY_BATCH_SIZE + 1, 1024, True, "triton"), + ], + ids=[ + "regular_limit", + "medium_range", + "above_medium", + "above_medium_gfx115x", + "batch_too_large", ], - ids=["regular_limit", "medium_range", "above_medium", "batch_too_large"], ) -def test_rdna_hybrid_w4a16_dispatch_boundaries(M, K, expected_path, monkeypatch): - """Route only supported regular and medium skinny shapes to HIP.""" +def test_rdna_hybrid_w4a16_dispatch_boundaries( + M, K, gfx115x, expected_path, monkeypatch +): + """Route only supported regular and medium skinny shapes to HIP. + + gfx115x drops the K*M bound, so a shape past the medium limit still goes to + HIP there; the batch bound holds on every target. + """ import vllm._custom_ops as ops N, G = 16, 128 @@ -777,6 +889,7 @@ def test_rdna_hybrid_w4a16_dispatch_boundaries(M, K, expected_path, monkeypatch) monkeypatch.setattr(ops, "wvSplitK_int4_g", fake_hip) monkeypatch.setattr(hybrid_module, "triton_w4a16_skinny_fmt_gemm", fake_triton) + monkeypatch.setattr(hybrid_module, "_on_gfx115x", lambda: gfx115x) output = hybrid_module._rdna_hybrid_w4a16_apply_impl( x_mk, w_q, scales, None, None, 40, G diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index bed62e89af..20a5f5bb4a 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -57,9 +57,21 @@ def _on_gfx1151() -> bool: return on_gfx1151() +def _on_gfx115x() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx115x + + return on_gfx115x() + + # Maximum activation rows supported by HIP skinny (Python M maps to C++ N_in). # When K*M exceeds regular LDS capacity, the C++ medium kernel caches the # prefix in LDS and reads the remaining activation elements from global memory. +# On gfx115x the K*M term is dropped entirely: deep-K shapes stage every +# activation row in LDS one K chunk at a time, and only when N is not a +# multiple of the kernel's row tile does the medium kernel fall back to reading +# the rows that do not fit from global memory. MAX_SKINNY_BATCH_SIZE = 5 # 64 KiB per-workgroup LDS limit expressed in fp16 elements. # (AMD RDNA has 128 KiB total LDS per CU, but 64 KiB per workgroup.) @@ -86,6 +98,7 @@ def _triton_w4a16_skinny_fmt_kernel( K, K8, # K // 8 num_groups, # K // group_size + stride_sn, # row stride of scales_ptr and zp_ptr, in groups stride_bn, # b_ptr row stride; >= K8, the rows may be padded stride_am, # a_ptr row stride in elements; >= K, rows may be padded # Quantization parameters @@ -146,12 +159,12 @@ def _triton_w4a16_skinny_fmt_kernel( b = (b >> shifts_full) & 0xF group_idx = (k_start * BLOCK_K) // group_size - scale_ptrs = scales_ptr + offs_n * num_groups + group_idx + scale_ptrs = scales_ptr + offs_n * stride_sn + group_idx scale_mask = offs_n < N scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0) if HAS_ZP: - zp_ptrs = zp_ptr + (offs_n // 8) * num_groups + group_idx + zp_ptrs = zp_ptr + (offs_n // 8) * stride_sn + group_idx zp_word = tl.load(zp_ptrs, mask=scale_mask, other=0) zp_raw = (zp_word >> (4 * (offs_n % 8))) & 0xF b_fp = (b - zp_raw[:, None]).to(scales.dtype) * scales[:, None] @@ -216,7 +229,8 @@ def triton_w4a16_skinny_fmt_gemm( """ assert a.stride(1) == 1, "Activation rows must be contiguous" assert b_q.stride(1) == 1, "Weight rows must be contiguous" - assert scales.is_contiguous(), "Scales must be contiguous" + assert scales.stride(1) == 1, "Scale rows must be contiguous" + stride_sn = scales.stride(0) M, K = a.shape N = b_q.shape[0] @@ -230,7 +244,10 @@ def triton_w4a16_skinny_fmt_gemm( f"scales shape mismatch: {scales.shape} vs ({N}, {num_groups})" ) if zp is not None: - assert zp.is_contiguous(), "Zero-points must be contiguous" + assert zp.stride(1) == 1, "Zero-point rows must be contiguous" + assert zp.stride(0) == stride_sn, ( + f"zp row stride {zp.stride(0)} must match scales {stride_sn}" + ) assert N % 8 == 0, f"N must be divisible by 8 for packed zp, got {N}" assert zp.shape == (N // 8, num_groups), ( f"zp shape mismatch: {zp.shape} vs ({N // 8}, {num_groups})" @@ -337,6 +354,7 @@ def triton_w4a16_skinny_fmt_gemm( K, K8, num_groups, + stride_sn, stride_bn, stride_am, group_size=group_size, @@ -466,7 +484,9 @@ def _rdna_hybrid_w4a16_apply_impl( K = x_2d.shape[1] N = w_q.shape[0] - if M <= MAX_SKINNY_BATCH_SIZE and K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS: + if M <= MAX_SKINNY_BATCH_SIZE and ( + K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS or _on_gfx115x() + ): # record_function is not torch.compile-safe; use nullcontext when # compiling to keep the op traceable. ctx = ( diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 807e571002..c34ad0c363 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -218,6 +218,7 @@ _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"]) _ON_GFX11 = "gfx11" in _GCN_ARCH _ON_GFX1100 = "gfx1100" in _GCN_ARCH _ON_GFX1151 = "gfx1151" in _GCN_ARCH +_ON_GFX115X = "gfx115" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) @@ -319,6 +320,10 @@ def on_gfx1151() -> bool: return _ON_GFX1151 +def on_gfx115x() -> bool: + return _ON_GFX115X + + def on_gfx12x() -> bool: return _ON_GFX12X and not _ON_CDNA