@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dc2e8f1157

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

Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>

Matthias Gehre committed Sep 14, 2026 at 11:40 UTC dc2e8f1157173d0f0fc09712f6897e9f6eee52a0

4 files changed +117 -45

csrc/rocm/skinny_gemms_int4.cu

+41 -18

@@ -102,6 +102,18 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) {

102 return min(a, b);

103 }

104

105 +// Reads row `row` of packed zero points [M/8, num_groups] uint32.

106 +//

107 +// The parameters must be signed and the indexing must use shift/mask: with

108 +// unsigned parameters and row/8, row%8, the (A_CHUNK=32, UNRL=8)

109 +// instantiation needs 256 VGPRs and 68 bytes of scratch instead of 239.

110 +__device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed,

111 + const int row, const int group,

112 + const int num_groups) {

113 + const uint32_t w = zp_packed[(row >> 3) * num_groups + group];

114 + return (w >> (4 * (row & 7))) & 0xFu;

115 +}

116 +

117 // W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations

118 // Targets the "sml" case where activations fit in LDS.

119 // A_CHUNK: number of K-elements processed per thread per step.

@@ -116,7 +128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

128 wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By,

129 const uint8_t* B_packed,

130 const scalar_t* __restrict__ A, const scalar_t* scale,

119 - const scalar_t* zero_points,

131 + const uint32_t* zero_points,

132 const scalar_t* __restrict__ BIAS, scalar_t* C,

133 const int _WvPrGrp, const int CuCount) {

134 constexpr int max_lds_len = LDS_SIZE / 2;

@@ -257,7 +269,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

269 if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {

270 if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {

271 uint32_t group_idx = k_ / GROUP_SIZE;

260 - scalar_t zp = zero_points[(m + y) * num_groups + group_idx];

272 + scalar_t zp = __float2s<scalar_t>(static_cast<float>(

273 + zp_nibble(zero_points, m + y, group_idx, num_groups)));

274 #pragma unroll

275 for (uint32_t b = 0; b < A_CHUNK; b++) {

276 cvtB.h[b] = cvtB.h[b] - zp;

@@ -280,8 +293,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

293 }

294 if constexpr (HAS_ZERO_POINTS) {

295 uint32_t group_idx_zp = k_ / GROUP_SIZE;

283 - float zp_f = __s2float(

284 - zero_points[(m + y) * num_groups + group_idx_zp]);

296 + float zp_f = static_cast<float>(

297 + zp_nibble(zero_points, m + y, group_idx_zp, num_groups));

298 partial -= (128.0f + zp_f) * act_sum;

299 } else {

300 partial -= 136.0f * act_sum;

@@ -335,7 +348,7 @@ __global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx,

348 const int By, const uint8_t* B_packed,

349 const scalar_t* __restrict__ A,

350 const scalar_t* scale,

338 - const scalar_t* zero_points,

351 + const uint32_t* zero_points,

352 const scalar_t* __restrict__ BIAS,

353 scalar_t* C, const int _WvPrGrp,

354 const int CuCount) {

@@ -352,7 +365,7 @@ template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,

365 __global__ void __launch_bounds__(WvPrGrp* THRDS)

366 wvSplitK_int4_hf_(const int K, const int M, const int Bx, const int By,

367 const uint8_t* B_packed, const scalar_t* __restrict__ A,

355 - const scalar_t* scale, const scalar_t* zero_points,

368 + const scalar_t* scale, const uint32_t* zero_points,

369 const scalar_t* __restrict__ BIAS, scalar_t* C,

370 const int _WvPrGrp, const int CuCount) {

371 constexpr int max_lds_len = LDS_SIZE / 2;

@@ -508,7 +521,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

521 if constexpr (!std::is_same_v<scalar_t, __hip_bfloat16>) {

522 if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) {

523 uint32_t group_idx = k_ / GROUP_SIZE;

511 - scalar_t zp = zero_points[(m + y) * num_groups + group_idx];

524 + scalar_t zp = __float2s<scalar_t>(static_cast<float>(

525 + zp_nibble(zero_points, m + y, group_idx, num_groups)));

526 #pragma unroll

527 for (uint32_t b = 0; b < A_CHUNK; b++) {

528 cvtB.h[b] = cvtB.h[b] - zp;

@@ -531,8 +545,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

545 }

546 if constexpr (HAS_ZERO_POINTS) {

547 uint32_t group_idx_zp = k_ / GROUP_SIZE;

534 - float zp_f = __s2float(

535 - zero_points[(m + y) * num_groups + group_idx_zp]);

548 + float zp_f = static_cast<float>(

549 + zp_nibble(zero_points, m + y, group_idx_zp, num_groups));

550 partial -= (128.0f + zp_f) * act_sum;

551 } else {

552 partial -= 136.0f * act_sum;

@@ -596,7 +610,7 @@ __global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx,

610 const int By, const uint8_t* B_packed,

611 const scalar_t* __restrict__ A,

612 const scalar_t* scale,

599 - const scalar_t* zero_points,

613 + const uint32_t* zero_points,

614 const scalar_t* __restrict__ BIAS,

615 scalar_t* C, const int _WvPrGrp,

616 const int CuCount) {

@@ -621,7 +635,8 @@ static int mindiv_int4(int N, int div1, int div2) {

635 // in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32)

636 // in_b: activations [N, K] (fp16/bf16)

637 // in_scale: group scales [M, K/group_size] (fp16/bf16)

624 -// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16)

638 +// in_zero_points: optional raw zero points [M/8, K/group_size] (int32,

639 +// 8 uint4 per word along dim 0)

640 // If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric).

641 // If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8).

642 // 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,

679 "Scale must be [M, K/group_size] = [", M_in, ", ", num_groups,

680 "] but got [", in_scale.size(0), ", ", in_scale.size(1), "]");

681 if (in_zero_points.has_value()) {

667 - TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(),

668 - "Zero points dtype must match activation dtype");

682 + // The kernel reads the words as uint32, so either signedness is accepted.

683 + TORCH_CHECK(in_zero_points->dtype() == at::kInt ||

684 + in_zero_points->dtype() == at::kUInt32,

685 + "Zero points must be int32 or uint32 (packed 8x uint4 along "

686 + "dim 0), got ",

687 + in_zero_points->dtype());

688 TORCH_CHECK(in_zero_points->dim() == 2,

670 - "Zero points must be 2D [M, K/group_size], got shape ",

689 + "Zero points must be 2D [M/8, K/group_size], got shape ",

690 in_zero_points->sizes());

672 - TORCH_CHECK(in_zero_points->size(0) == M_in &&

691 + TORCH_CHECK(M_in % 8 == 0,

692 + "M must be divisible by 8 for packed zero points, got ", M_in);

693 + TORCH_CHECK(in_zero_points->is_contiguous(),

694 + "Zero points must be contiguous");

695 + TORCH_CHECK(in_zero_points->size(0) == M_in / 8 &&

696 in_zero_points->size(1) == num_groups,

674 - "Zero points must be [M, K/group_size] = [", M_in, ", ",

697 + "Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ",

698 num_groups, "] but got [", in_zero_points->size(0), ", ",

699 in_zero_points->size(1), "]");

700 }

@@ -769,9 +792,9 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,

792 const fptype* aptr = reinterpret_cast<const fptype*>(in_b.data_ptr());

793 const fptype* sptr =

794 reinterpret_cast<const fptype*>(in_scale.data_ptr());

772 - const fptype* zpptr =

795 + const uint32_t* zpptr =

796 in_zero_points.has_value()

774 - ? reinterpret_cast<const fptype*>(in_zero_points->data_ptr())

797 + ? reinterpret_cast<const uint32_t*>(in_zero_points->data_ptr())

798 : nullptr;

799 const fptype* biasptr =

800 (in_bias.has_value() && in_bias->numel() > 0)

csrc/rocm/torch_bindings.cpp

+2 -1

@@ -31,7 +31,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {

31 rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);

33 // W4A16 grouped skinny GEMM: packed int4 weights, per-group scales,

34 - // optional zero points for asymmetric quantization

34 + // optional zero points [M/8, K/group_size] int32 for asymmetric

35 + // quantization

36 rocm_ops.def(

37 "wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, "

38 "Tensor? in_zero_points, Tensor? in_bias, int CuCount, "

tests/kernels/quantization/test_rdna_hybrid_w4a16.py

+55 -10

@@ -37,6 +37,20 @@ MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE

37 # ---------------------------------------------------------------------------

40 +def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor:

41 + """Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32.

42 +

43 + Row n's nibble lands in word[n//8] at bits 4*(n%8).

44 + """

45 + assert zp_nkg.dtype == torch.int32

46 + N, G = zp_nkg.shape

47 + assert N % 8 == 0

48 + shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None]

49 + return torch.sum(

50 + (zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32

51 + ).contiguous()

52 +

53 +

54 def _rdna_hybrid_w4a16_reference(

55 x_mk: torch.Tensor,

56 w_int4_nk: torch.Tensor,

@@ -50,8 +64,8 @@ def _rdna_hybrid_w4a16_reference(

64 x_mk: [M, K] fp16/bf16

65 w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15]

66 scales_nkg: [N, K//G] fp16/bf16

53 - zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype),

54 - or None for symmetric (uint4b8, dequant subtracts 8)

67 + zp_nkg: [N, K//G] int32 raw zero points in [0, 15], or None for

68 + symmetric (uint4b8, dequant subtracts 8)

69 """

70 G = group_size

71 N, K = w_int4_nk.shape

@@ -110,13 +124,16 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)

124 0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32)

125 ).to(dtype)

126

113 - # Optional raw zero points [N, K//G] in act dtype.

127 + # Optional raw zero points, [N, K//G] int32 for the reference and packed

128 + # along N for the op.

129 if has_zp:

130 zp_nkg = torch.randint(

131 0, 16, (N, K // group_size), device=device, dtype=torch.int32

117 - ).to(dtype)

132 + )

133 + w_zp = _pack_zp_rows_for_kernel(zp_nkg)

134 else:

135 zp_nkg = None

136 + w_zp = None

137

138 from vllm.utils.platform_utils import num_compute_units

139

@@ -124,7 +141,7 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M)

141 x_mk,

142 w_q,

143 scales_nkg,

127 - zp_nkg,

144 + w_zp,

145 None, # bias

146 num_compute_units(),

147 group_size,

@@ -332,7 +349,7 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini

349

350 @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES)

351 def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init):

335 - """uint4 (asymmetric): zero points unpacked to raw values in act dtype."""

352 + """uint4 (asymmetric): zero points are packed [N//8, K//G] int32."""

353 if not torch.cuda.is_available():

354 pytest.skip("CUDA/HIP device not available")

355

@@ -374,10 +391,11 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in

391 )

392 kernel.process_weights_after_loading(layer)

393

377 - # Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15].

378 - assert layer.weight_zero_point.dtype == torch.float16

379 - assert tuple(layer.weight_zero_point.shape) == (N, K // G)

380 - expected_zp = zeros_int4_gn.t().to(torch.float16) # [N, K//G]

394 + # Zero-points: [N//8, K//G] int32, row n's nibble at word[n//8]

395 + # bits 4*(n%8).

396 + assert layer.weight_zero_point.dtype == torch.int32

397 + assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G)

398 + expected_zp = _pack_zp_rows_for_kernel(zeros_int4_gn.t().contiguous())

399 torch.testing.assert_close(layer.weight_zero_point, expected_zp)

400

401 # 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):

536 ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8)

537

538 torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)

539 +

540 +

541 +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")

542 +@pytest.mark.skipif(

543 + not hasattr(torch.ops, "_rocm_C")

544 + or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"),

545 + reason="wvSplitK_int4_g not built",

546 +)

547 +def test_wvsplitk_int4_g_rejects_unpacked_zero_points():

548 + """Act-dtype zero points raise instead of being read as packed words.

549 +

550 + Both layouts are 2D with compatible extents, so only the dtype separates

551 + them; misreading one as the other returns wrong numbers silently.

552 + """

553 + import vllm._custom_ops as ops

554 + from vllm.utils.platform_utils import num_compute_units

555 +

556 + K, N, G, M = 256, 64, 128, 1

557 + a = torch.randn((M, K), device=device, dtype=torch.float16)

558 + w = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view(

559 + torch.int8

560 + )

561 + scales = torch.rand((N, K // G), device=device, dtype=torch.float16)

562 + zp_unpacked = torch.zeros((N, K // G), device=device, dtype=torch.float16)

563 +

564 + with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"):

565 + 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

@@ -77,7 +77,7 @@ def _triton_w4a16_skinny_fmt_kernel(

77 a_ptr, # [M, K] fp16/bf16 activations

78 b_ptr, # [N, K//8] int32 packed (ExLlama shuffle, K is packed dim)

79 scales_ptr, # [N, K//G] fp16/bf16 scales (skinny layout)

80 - zp_ptr, # [N, K//G] fp16/bf16 raw zero-points (when HAS_ZP=True)

80 + zp_ptr, # [N//8, K//G] int32 zero-points (when HAS_ZP=True)

81 c_ptr, # [M, N] fp16/bf16 output

82 # Dimensions

83 M,

@@ -103,8 +103,9 @@ def _triton_w4a16_skinny_fmt_kernel(

103 | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28)

104

105 Scales are [N, K//G] (skinny layout, NOT transposed).

106 - When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G]

107 - and subtracted directly: (nibble - zp_raw) * scale.

106 + When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw

107 + zero-point at word[n//8] bits 4*(n%8), and dequant is

108 + (nibble - zp_raw) * scale.

109 When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric).

110 """

111 pid_m = tl.program_id(0)

@@ -148,9 +149,10 @@ def _triton_w4a16_skinny_fmt_kernel(

149 scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0)

150

151 if HAS_ZP:

151 - zp_ptrs = zp_ptr + offs_n * num_groups + group_idx

152 - zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0)

153 - b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None]

152 + zp_ptrs = zp_ptr + (offs_n // 8) * num_groups + group_idx

153 + zp_word = tl.load(zp_ptrs, mask=scale_mask, other=0)

154 + zp_raw = (zp_word >> (4 * (offs_n % 8))) & 0xF

155 + b_fp = (b - zp_raw[:, None]).to(scales.dtype) * scales[:, None]

156 else:

157 b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None]

158

@@ -192,7 +194,7 @@ def triton_w4a16_skinny_fmt_gemm(

194 scales: torch.Tensor, # [N, K//G] fp16/bf16

195 group_size: int,

196 zp_bias: int = 8,

195 - zp: torch.Tensor | None = None, # [N, K//G] per-group zero-points

197 + zp: torch.Tensor | None = None, # [N//8, K//G] int32 zero-points

198 ) -> torch.Tensor:

199 """

200 Fused W4A16 GEMM reading from skinny weight format [N, K//8].

@@ -203,8 +205,8 @@ def triton_w4a16_skinny_fmt_gemm(

205 scales: Per-group scales [N, K//G], same dtype as a.

206 group_size: Quantization group size (resolved from -1 to K by caller).

207 zp_bias: Constant zero bias (default 8 for unsigned int4).

206 - zp: Raw per-group zero-points [N, K//G] (asymmetric),

207 - stored as zp_raw in activation dtype. When provided,

208 + zp: Raw per-group zero-points [N//8, K//G] int32, row n at

209 + word[n//8] bits 4*(n%8) (asymmetric). When provided,

210 dequant is (nibble - zp_raw) * scale.

211

212 Returns:

@@ -225,8 +227,9 @@ def triton_w4a16_skinny_fmt_gemm(

227 )

228 if zp is not None:

229 assert zp.is_contiguous(), "Zero-points must be contiguous"

228 - assert zp.shape == (N, num_groups), (

229 - f"zp shape mismatch: {zp.shape} vs ({N}, {num_groups})"

230 + assert N % 8 == 0, f"N must be divisible by 8 for packed zp, got {N}"

231 + assert zp.shape == (N // 8, num_groups), (

232 + f"zp shape mismatch: {zp.shape} vs ({N // 8}, {num_groups})"

233 )

234 has_zp = zp is not None

235

@@ -381,7 +384,10 @@ def _rdna_hybrid_w4a16_apply_impl(

384 cu_count: int,

385 group_size: int,

386 ) -> torch.Tensor:

384 - """Dispatch between skinny GEMM and Triton based on batch size M."""

387 + """Dispatch between skinny GEMM and Triton based on batch size M.

388 +

389 + ``w_zp`` is [N//8, K//G] int32 for asymmetric layers, None for symmetric.

390 + """

391 import vllm._custom_ops as ops

392

393 M = x_2d.shape[0]

@@ -520,10 +526,7 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel):

526 assert self.w_zp_name is not None

527 w_zp_raw = getattr(layer, self.w_zp_name)

528 permute_param_layout_(w_zp_raw, input_dim=1, output_dim=0, packed_dim=0)

523 - zp_unpacked = unpack_quantized_values_into_int32(

524 - w_zp_raw.data, c.weight_type, packed_dim=0

525 - )

526 - w_zp = zp_unpacked.to(c.act_type).contiguous()

529 + w_zp = w_zp_raw.data.contiguous()

530 self._transform_param(layer, self.w_zp_name, lambda x: w_zp)

531

532 self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny)