diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 8b7077d14e..a20ca41d41 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -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) { 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(static_cast( + 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( + 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 ) { 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(static_cast( + 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( + 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(in_b.data_ptr()); const fptype* sptr = reinterpret_cast(in_scale.data_ptr()); - const fptype* zpptr = + const uint32_t* zpptr = in_zero_points.has_value() - ? reinterpret_cast(in_zero_points->data_ptr()) + ? reinterpret_cast(in_zero_points->data_ptr()) : nullptr; const fptype* biasptr = (in_bias.has_value() && in_bias->numel() > 0) diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index a8114ce196..14d1822707 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -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, " diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index 562d1ef2ef..fe3db8dab3 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -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) 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 9964c891b6..eb36308bcc 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 @@ -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)