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

[ROCm][Perf] W4A16: keep skinny GEMM zero-points packed 4-bit (#54965)

Browse Source 
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
This commit is contained in:
Matthias Gehre
2026-09-14 17:40:25 +08:00
committed by GitHub
parent 238cb2b191
commit dc2e8f1157
4 changed files with 117 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/rocm/skinny_gemms_int4.cu
+41 -18
View File
@@ -102,6 +102,18 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) {
return min(a, b);
}
// Reads row `row` of packed zero points [M/8, num_groups] uint32.
//
// The parameters must be signed and the indexing must use shift/mask: with
// unsigned parameters and row/8, row%8, the (A_CHUNK=32, UNRL=8)
// instantiation needs 256 VGPRs and 68 bytes of scratch instead of 239.
__device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed,
const int row, const int group,
const int num_groups) {
const uint32_t w = zp_packed[(row >> 3) * num_groups + group];
return (w >> (4 * (row & 7))) & 0xFu;
}
// W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations
// Targets the "sml" case where activations fit in LDS.
// A_CHUNK: number of K-elements processed per thread per step.
@@ -116,7 +128,7 @@ __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,
const scalar_t* __restrict__ A, const scalar_t* scale,
const scalar_t* zero_points,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C,
const int _WvPrGrp, const int CuCount) {
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -257,7 +269,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
if constexpr (!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 = zero_points[(m + y) * num_groups + group_idx];
scalar_t zp = __float2s<scalar_t>(static_cast<float>(
zp_nibble(zero_points, m + y, group_idx, num_groups)));
#pragma unroll
for (uint32_t b = 0; b < A_CHUNK; b++) {
cvtB.h[b] = cvtB.h[b] - zp;
@@ -280,8 +293,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
if constexpr (HAS_ZERO_POINTS) {
uint32_t group_idx_zp = k_ / GROUP_SIZE;
float zp_f = __s2float(
zero_points[(m + y) * num_groups + group_idx_zp]);
float zp_f = static_cast<float>(
zp_nibble(zero_points, m + y, group_idx_zp, num_groups));
partial -= (128.0f + zp_f) * act_sum;
} else {
partial -= 136.0f * act_sum;
@@ -335,7 +348,7 @@ __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 scalar_t* zero_points,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS,
scalar_t* C, const int _WvPrGrp,
const int CuCount) {
@@ -352,7 +365,7 @@ template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
__global__ void __launch_bounds__(WvPrGrp* THRDS)
wvSplitK_int4_hf_(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 scalar_t* zero_points,
const scalar_t* scale, const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS, scalar_t* C,
const int _WvPrGrp, const int CuCount) {
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -508,7 +521,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
if constexpr (!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 = zero_points[(m + y) * num_groups + group_idx];
scalar_t zp = __float2s<scalar_t>(static_cast<float>(
zp_nibble(zero_points, m + y, group_idx, num_groups)));
#pragma unroll
for (uint32_t b = 0; b < A_CHUNK; b++) {
cvtB.h[b] = cvtB.h[b] - zp;
@@ -531,8 +545,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
if constexpr (HAS_ZERO_POINTS) {
uint32_t group_idx_zp = k_ / GROUP_SIZE;
float zp_f = __s2float(
zero_points[(m + y) * num_groups + group_idx_zp]);
float zp_f = static_cast<float>(
zp_nibble(zero_points, m + y, group_idx_zp, num_groups));
partial -= (128.0f + zp_f) * act_sum;
} else {
partial -= 136.0f * act_sum;
@@ -596,7 +610,7 @@ __global__ void wvSplitK_int4_hf_(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 scalar_t* zero_points,
const uint32_t* zero_points,
const scalar_t* __restrict__ BIAS,
scalar_t* C, const int _WvPrGrp,
const int CuCount) {
@@ -621,7 +635,8 @@ static int mindiv_int4(int N, int div1, int div2) {
// in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32)
// in_b: activations [N, K] (fp16/bf16)
// in_scale: group scales [M, K/group_size] (fp16/bf16)
// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16)
// in_zero_points: optional raw zero points [M/8, K/group_size] (int32,
// 8 uint4 per word along dim 0)
// If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric).
// If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8).
// group_size: 32, 64, or 128
@@ -664,14 +679,22 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
"Scale must be [M, K/group_size] = [", M_in, ", ", num_groups,
"] but got [", in_scale.size(0), ", ", in_scale.size(1), "]");
if (in_zero_points.has_value()) {
TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(),
"Zero points dtype must match activation dtype");
// The kernel reads the words as uint32, so either signedness is accepted.
TORCH_CHECK(in_zero_points->dtype() == at::kInt ||
in_zero_points->dtype() == at::kUInt32,
"Zero points must be int32 or uint32 (packed 8x uint4 along "
"dim 0), got ",
in_zero_points->dtype());
TORCH_CHECK(in_zero_points->dim() == 2,
"Zero points must be 2D [M, K/group_size], got shape ",
"Zero points must be 2D [M/8, K/group_size], got shape ",
in_zero_points->sizes());
TORCH_CHECK(in_zero_points->size(0) == M_in &&
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->size(0) == M_in / 8 &&
in_zero_points->size(1) == num_groups,
"Zero points must be [M, K/group_size] = [", M_in, ", ",
"Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ",
num_groups, "] but got [", in_zero_points->size(0), ", ",
in_zero_points->size(1), "]");
}
@@ -769,9 +792,9 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
const fptype* aptr = reinterpret_cast<const fptype*>(in_b.data_ptr());
const fptype* sptr =
reinterpret_cast<const fptype*>(in_scale.data_ptr());
const fptype* zpptr =
const uint32_t* zpptr =
in_zero_points.has_value()
? reinterpret_cast<const fptype*>(in_zero_points->data_ptr())
? reinterpret_cast<const uint32_t*>(in_zero_points->data_ptr())
: nullptr;
const fptype* biasptr =
(in_bias.has_value() && in_bias->numel() > 0)
csrc/rocm/torch_bindings.cpp
+2 -1
View File
@@ -31,7 +31,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
// W4A16 grouped skinny GEMM: packed int4 weights, per-group scales,
// optional zero points for asymmetric quantization
// optional zero points [M/8, K/group_size] int32 for asymmetric
// quantization
rocm_ops.def(
"wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, "
"Tensor? in_zero_points, Tensor? in_bias, int CuCount, "
tests/kernels/quantization/test_rdna_hybrid_w4a16.py
+55 -10
View File
@@ -37,6 +37,20 @@ MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE
# ---------------------------------------------------------------------------
def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor:
"""Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32.
Row n's nibble lands in word[n//8] at bits 4*(n%8).
"""
assert zp_nkg.dtype == torch.int32
N, G = zp_nkg.shape
assert N % 8 == 0
shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None]
return torch.sum(
(zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32
).contiguous()
def _rdna_hybrid_w4a16_reference(
x_mk: torch.Tensor,
w_int4_nk: torch.Tensor,
@@ -50,8 +64,8 @@ def _rdna_hybrid_w4a16_reference(
x_mk: [M, K] fp16/bf16
w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15]
scales_nkg: [N, K//G] fp16/bf16
zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype),
or None for symmetric (uint4b8, dequant subtracts 8)
zp_nkg: [N, K//G] int32 raw zero points in [0, 15], or None for
symmetric (uint4b8, dequant subtracts 8)
"""
G = group_size
N, K = w_int4_nk.shape
@@ -110,13 +124,16 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)
0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32)
).to(dtype)
# Optional raw zero points [N, K//G] in act dtype.
# Optional raw zero points, [N, K//G] int32 for the reference and packed
# along N for the op.
if has_zp:
zp_nkg = torch.randint(
0, 16, (N, K // group_size), device=device, dtype=torch.int32
).to(dtype)
)
w_zp = _pack_zp_rows_for_kernel(zp_nkg)
else:
zp_nkg = None
w_zp = None
from vllm.utils.platform_utils import num_compute_units
@@ -124,7 +141,7 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)
x_mk,
w_q,
scales_nkg,
zp_nkg,
w_zp,
None, # bias
num_compute_units(),
group_size,
@@ -332,7 +349,7 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini
@pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)
def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init):
"""uint4 (asymmetric): zero points unpacked to raw values in act dtype."""
"""uint4 (asymmetric): zero points are packed [N//8, K//G] int32."""
if not torch.cuda.is_available():
pytest.skip("CUDA/HIP device not available")
@@ -374,10 +391,11 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in
)
kernel.process_weights_after_loading(layer)
# Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15].
assert layer.weight_zero_point.dtype == torch.float16
assert tuple(layer.weight_zero_point.shape) == (N, K // G)
expected_zp = zeros_int4_gn.t().to(torch.float16) # [N, K//G]
# Zero-points: [N//8, K//G] int32, row n's nibble at word[n//8]
# bits 4*(n%8).
assert layer.weight_zero_point.dtype == torch.int32
assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G)
expected_zp = _pack_zp_rows_for_kernel(zeros_int4_gn.t().contiguous())
torch.testing.assert_close(layer.weight_zero_point, expected_zp)
# Quantized weights match symmetric path's layout regardless of zp.
@@ -518,3 +536,30 @@ def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G):
ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)
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.skipif(
not hasattr(torch.ops, "_rocm_C")
or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"),
reason="wvSplitK_int4_g not built",
)
def test_wvsplitk_int4_g_rejects_unpacked_zero_points():
"""Act-dtype zero points raise instead of being read as packed words.
Both layouts are 2D with compatible extents, so only the dtype separates
them; misreading one as the other returns wrong numbers silently.
"""
import vllm._custom_ops as ops
from vllm.utils.platform_utils import num_compute_units
K, N, G, M = 256, 64, 128, 1
a = torch.randn((M, K), device=device, dtype=torch.float16)
w = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view(
torch.int8
)
scales = torch.rand((N, K // G), device=device, dtype=torch.float16)
zp_unpacked = torch.zeros((N, K // G), device=device, dtype=torch.float16)
with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"):
ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None)
vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py
+19 -16
View File
@@ -77,7 +77,7 @@ def _triton_w4a16_skinny_fmt_kernel(
a_ptr, # [M, K] fp16/bf16 activations
b_ptr, # [N, K//8] int32 packed (ExLlama shuffle, K is packed dim)
scales_ptr, # [N, K//G] fp16/bf16 scales (skinny layout)
zp_ptr, # [N, K//G] fp16/bf16 raw zero-points (when HAS_ZP=True)
zp_ptr, # [N//8, K//G] int32 zero-points (when HAS_ZP=True)
c_ptr, # [M, N] fp16/bf16 output
# Dimensions
M,
@@ -103,8 +103,9 @@ def _triton_w4a16_skinny_fmt_kernel(
| (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28)
Scales are [N, K//G] (skinny layout, NOT transposed).
When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G]
and subtracted directly: (nibble - zp_raw) * scale.
When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw
zero-point at word[n//8] bits 4*(n%8), and dequant is
(nibble - zp_raw) * scale.
When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric).
"""
pid_m = tl.program_id(0)
@@ -148,9 +149,10 @@ def _triton_w4a16_skinny_fmt_kernel(
scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0)
if HAS_ZP:
zp_ptrs = zp_ptr + offs_n * num_groups + group_idx
zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0)
b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None]
zp_ptrs = zp_ptr + (offs_n // 8) * num_groups + 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]
else:
b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None]
@@ -192,7 +194,7 @@ def triton_w4a16_skinny_fmt_gemm(
scales: torch.Tensor, # [N, K//G] fp16/bf16
group_size: int,
zp_bias: int = 8,
zp: torch.Tensor | None = None, # [N, K//G] per-group zero-points
zp: torch.Tensor | None = None, # [N//8, K//G] int32 zero-points
) -> torch.Tensor:
"""
Fused W4A16 GEMM reading from skinny weight format [N, K//8].
@@ -203,8 +205,8 @@ def triton_w4a16_skinny_fmt_gemm(
scales: Per-group scales [N, K//G], same dtype as a.
group_size: Quantization group size (resolved from -1 to K by caller).
zp_bias: Constant zero bias (default 8 for unsigned int4).
zp: Raw per-group zero-points [N, K//G] (asymmetric),
stored as zp_raw in activation dtype. When provided,
zp: Raw per-group zero-points [N//8, K//G] int32, row n at
word[n//8] bits 4*(n%8) (asymmetric). When provided,
dequant is (nibble - zp_raw) * scale.
Returns:
@@ -225,8 +227,9 @@ def triton_w4a16_skinny_fmt_gemm(
)
if zp is not None:
assert zp.is_contiguous(), "Zero-points must be contiguous"
assert zp.shape == (N, num_groups), (
f"zp shape mismatch: {zp.shape} vs ({N}, {num_groups})"
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})"
)
has_zp = zp is not None
@@ -381,7 +384,10 @@ def _rdna_hybrid_w4a16_apply_impl(
cu_count: int,
group_size: int,
) -> torch.Tensor:
"""Dispatch between skinny GEMM and Triton based on batch size M."""
"""Dispatch between skinny GEMM and Triton based on batch size M.
``w_zp`` is [N//8, K//G] int32 for asymmetric layers, None for symmetric.
"""
import vllm._custom_ops as ops
M = x_2d.shape[0]
@@ -520,10 +526,7 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel):
assert self.w_zp_name is not None
w_zp_raw = getattr(layer, self.w_zp_name)
permute_param_layout_(w_zp_raw, input_dim=1, output_dim=0, packed_dim=0)
zp_unpacked = unpack_quantized_values_into_int32(
w_zp_raw.data, c.weight_type, packed_dim=0
)
w_zp = zp_unpacked.to(c.act_type).contiguous()
w_zp = w_zp_raw.data.contiguous()
self._transform_param(layer, self.w_zp_name, lambda x: w_zp)
self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny)