iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] W4A16: speed up skinny GEMM decode, including speculative batches (#57199)

Browse Source 
Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
This commit is contained in:
Robert Esclapez
2026-10-13 17:21:30 +00:00
committed by GitHub
co-authored by Claude Opus 4.7 Douglas Lehr
parent 7fad7db330
commit e83222396f
4 changed files with 603 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/rocm/skinny_gemms_int4.cu
+442 -94
View File
@@ -8,6 +8,7 @@
#include <stdexcept>
#include <algorithm>
#include <type_traits>
#include <utility> // 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 <assert.h>
@@ -98,6 +116,14 @@ struct scalar<c10::BFloat16> {
} 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 <typename scalar_t, int THRDS, int A_CHUNK, int N, uint32_t KFIT>
__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<uint32_t>(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<uint32_t>(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 <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
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<uint64_t>(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<uint32_t>(K) + kFit - 1) / kFit;
const uint32_t even = (static_cast<uint32_t>(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<scalar_t, half>) {
#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<scalar_t, half>) {
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<scalar_t, __hip_bfloat16>) {
if constexpr (!FOLD_ZP_INTO_DEQUANT &&
!std::is_same_v<scalar_t, __hip_bfloat16>) {
if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
uint32_t group_idx = k_ / GROUP_SIZE;
scalar_t zp = __float2s<scalar_t>(static_cast<float>(
@@ -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<uint32_t>(K); kb += chunkSpan) {
__syncthreads();
load_act_chunk_into_lds<scalar_t, THRDS, A_CHUNK, N, kFit>(
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<uint32_t>(K) &&
K_whole < static_cast<uint32_t>(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 <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
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 <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GROUP_SIZE = 0, bool HAS_ZERO_POINTS = false>
@@ -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<scalar_t, half>) {
#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<scalar_t, half>) {
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<scalar_t, __hip_bfloat16>) {
if constexpr (!FOLD_ZP_INTO_DEQUANT &&
!std::is_same_v<scalar_t, __hip_bfloat16>) {
if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {
uint32_t group_idx = k_ / GROUP_SIZE;
scalar_t zp = __float2s<scalar_t>(static_cast<float>(
@@ -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<int>(group_stride),
"Scale row stride overflows int");
const int group_stride_i32 = static_cast<int>(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_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
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_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
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_<fptype, _THRDS, _YTILE, _W, _AC, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
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_<fptype, _THRDS, _YTILE, _W, _AC, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
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_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
_HAS_ZP, /*CHUNKED=*/true, lds_elems> \
<<<grid, block, 0, stream>>>(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_<fptype, _THRDS, _YTILE, 16, 16, _UNRL, _N, _GS, \
_HAS_ZP><<<grid, block, 0, stream>>>( \
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
tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+131 -18
View File
@@ -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
vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+25 -5
View File
@@ -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 = (
vllm/platforms/rocm.py
+5
View File
@@ -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