@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 276 lines 9.69 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 """Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100).
5
6 Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with
7 GPTQ-format checkpoint parameters, runs ``process_weights_after_loading``
8 (weight shuffle + zero-point synthesis), then ``apply_weights``, and compares
9 the result against an fp32 reference dequant-and-matmul.
10
11 The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only
12 built for gfx11; tests are skipped elsewhere.
13
14 Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`.
15 """
16
17 import pytest
18 import torch
19
20 from vllm.platforms import current_platform
21
22 if not current_platform.is_rocm():
23 pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
24
25 from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402
26 MPLinearLayerConfig,
27 )
28 from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402
29 RDNA3W4A16LinearKernel,
30 )
31 from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
32 pack_quantized_values_into_int32,
33 )
34 from vllm.model_executor.parameter import ( # noqa: E402
35 GroupQuantScaleParameter,
36 PackedvLLMParameter,
37 )
38 from vllm.platforms.rocm import on_gfx1100 # noqa: E402
39 from vllm.scalar_type import scalar_types # noqa: E402
40 from vllm.utils.torch_utils import set_random_seed # noqa: E402
41
42 device = "cuda"
43
44 WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8
45 PACK_FACTOR = 8 # 8 x 4-bit nibbles per int32
46
47 # Skip everything in this module unless we are on the only architecture the
48 # kernel is built/registered for.
49 gfx1100_only = pytest.mark.skipif(
50 not (
51 on_gfx1100()
52 and hasattr(torch.ops, "_rocm_C")
53 and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
54 ),
55 reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
56 )
57
58
59 # ---------------------------------------------------------------------------
60 # Reference implementation
61 # ---------------------------------------------------------------------------
62
63
64 def _reference(
65 x_mk: torch.Tensor,
66 q_int4_kn: torch.Tensor,
67 scales_gn: torch.Tensor,
68 zeros_gn: torch.Tensor | None,
69 group_size: int,
70 bias: torch.Tensor | None,
71 ) -> torch.Tensor:
72 """fp32 reference for the RDNA3 W4A16 op.
73
74 x_mk: [M, K] fp16/bf16 activations.
75 q_int4_kn: [K, N] int32 raw stored nibbles in [0, 15].
76 scales_gn: [K//G, N] per-group scales (act dtype).
77 zeros_gn: [K//G, N] int32 raw stored zero points in [0, 15], or None
78 for the symmetric path (kernel synthesizes stored zero = 7).
79 group_size: G.
80
81 The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective
82 zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8).
83 """
84 K, N = q_int4_kn.shape
85 s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32) # [K,N]
86 if zeros_gn is None:
87 z_full = torch.full(
88 (K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32
89 )
90 else:
91 z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32)
92 w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full # [K, N]
93 out = x_mk.to(torch.float32) @ w_fp # [M, N]
94 if bias is not None:
95 out = out + bias.to(torch.float32)
96 return out.to(x_mk.dtype)
97
98
99 # ---------------------------------------------------------------------------
100 # Layer construction (GPTQ checkpoint format)
101 # ---------------------------------------------------------------------------
102
103
104 def _build_layer(
105 q_int4_kn: torch.Tensor,
106 scales_gn: torch.Tensor,
107 zeros_gn: torch.Tensor | None,
108 dtype: torch.dtype,
109 ) -> torch.nn.Module:
110 """Build a dummy layer carrying GPTQ-format params, as the loader would."""
111 no_loader = lambda *args, **kwargs: None # noqa: E731
112
113 # qweight: int4 packed along K into int32 -> [K//8, N].
114 qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0)
115
116 class DummyLayer(torch.nn.Module):
117 pass
118
119 layer = DummyLayer()
120 layer.register_parameter(
121 "qweight",
122 PackedvLLMParameter(
123 data=qweight,
124 weight_loader=no_loader,
125 input_dim=0,
126 output_dim=1,
127 packed_dim=0,
128 packed_factor=PACK_FACTOR,
129 ),
130 )
131 layer.register_parameter(
132 "scales",
133 GroupQuantScaleParameter(
134 data=scales_gn.to(dtype),
135 weight_loader=no_loader,
136 input_dim=0,
137 output_dim=1,
138 ),
139 )
140 if zeros_gn is not None:
141 # qzeros: int4 packed along N into int32 -> [K//G, N//8].
142 qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1)
143 layer.register_parameter(
144 "qzeros",
145 PackedvLLMParameter(
146 data=qzeros,
147 weight_loader=no_loader,
148 input_dim=0,
149 output_dim=1,
150 packed_dim=1,
151 packed_factor=PACK_FACTOR,
152 ),
153 )
154 return layer
155
156
157 def _run_kernel(
158 x_mk: torch.Tensor,
159 q_int4_kn: torch.Tensor,
160 scales_gn: torch.Tensor,
161 zeros_gn: torch.Tensor | None,
162 group_size: int,
163 bias: torch.Tensor | None,
164 dtype: torch.dtype,
165 ) -> torch.Tensor:
166 K, N = q_int4_kn.shape
167 has_zp = zeros_gn is not None
168
169 config = MPLinearLayerConfig(
170 full_weight_shape=(K, N),
171 partition_weight_shape=(K, N),
172 weight_type=WEIGHT_TYPE,
173 act_type=dtype,
174 group_size=group_size,
175 zero_points=has_zp,
176 )
177 ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
178 assert ok, f"can_implement rejected a supported config: {reason}"
179
180 layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype)
181 kernel = RDNA3W4A16LinearKernel(
182 config,
183 w_q_param_name="qweight",
184 w_s_param_name="scales",
185 w_zp_param_name="qzeros" if has_zp else None,
186 )
187 kernel.process_weights_after_loading(layer)
188 return kernel.apply_weights(layer, x_mk, bias=bias)
189
190
191 # Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights
192 # to fp32 and accumulates in fp32, so it matches the reference almost exactly
193 # (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2
194 # "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered
195 # as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order
196 # mantissa bits and leaves ~2-3% relative noise that accumulates over K. We
197 # compare on the relative Frobenius norm rather than elementwise, since the
198 # bit-trick noise produces large *relative* errors on individual near-zero
199 # outputs that carry negligible absolute weight.
200 _REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2}
201
202
203 def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype):
204 rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to(
205 torch.float32
206 ).norm()
207 tol = _REL_L2_TOL[dtype]
208 assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}"
209
210
211 # ---------------------------------------------------------------------------
212 # Forward correctness
213 # ---------------------------------------------------------------------------
214
215
216 # (M, K, N, group_size). M spans the scalar decode path (small M) and the
217 # WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's
218 # divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0).
219 MKNG_SHAPES = [
220 (1, 128, 128, 128), # single group, decode
221 (2, 256, 256, 128), # two groups
222 (8, 256, 512, 64), # M=8 scalar, smaller group
223 (16, 512, 256, 128), # M=16 -> WMMA path for bf16
224 (32, 512, 512, 64), # larger prefill
225 ]
226
227
228 @gfx1100_only
229 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
230 @pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"])
231 @pytest.mark.parametrize(
232 "M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES]
233 )
234 def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init):
235 set_random_seed(0)
236 assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0
237
238 groups = K // G
239 x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
240 q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
241 scales_gn = (
242 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
243 ).to(dtype)
244 zeros_gn = (
245 torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32)
246 if has_zp
247 else None
248 )
249
250 out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype)
251 ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None)
252
253 assert out.shape == (M, N) and out.dtype == dtype
254 _assert_close(out, ref, dtype)
255
256
257 @gfx1100_only
258 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
259 @pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"])
260 def test_rdna3_w4a16_bias(dtype, M, dist_init):
261 """Bias is added on both the scalar (M=1) and WMMA (M=32) paths."""
262 set_random_seed(0)
263 K, N, G = 512, 256, 128
264 groups = K // G
265
266 x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
267 q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
268 scales_gn = (
269 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
270 ).to(dtype)
271 bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype)
272
273 out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype)
274 ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias)
275
276 _assert_close(out, ref, dtype)