@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 361 lines 11 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 fused MoE W4A16 HIP kernel (gfx1100).
5
6 Tests ``moe_gptq_gemm_rdna3`` against the dense ``gptq_gemm_rdna3`` as
7 reference: builds RDNA3-format weights (shuffled int32, synthesized qzeros),
8 runs the fused MoE kernel, and compares per-expert results.
9
10 Model parameters taken from:
11 - cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit
12 (hidden=2048, inter=768, E=128, top_k=8, G=32)
13 - Qwen3.6-35B-A3B-GPTQ-W4A16-G32
14 (hidden=2048, inter=512, E=256, top_k=8, G=32)
15
16 Run `pytest tests/kernels/quantization/test_rdna3_moe_w4a16.py`.
17 """
18
19 import pytest
20 import torch
21
22 from vllm.platforms import current_platform
23
24 if not current_platform.is_rocm():
25 pytest.skip("RDNA3 MoE W4A16 kernel is ROCm-only", allow_module_level=True)
26
27 from vllm import _custom_ops as ops # noqa: E402
28 from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402
29 MoEActivation,
30 apply_moe_activation,
31 )
32 from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402
33 moe_align_block_size,
34 )
35 from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
36 pack_quantized_values_into_int32,
37 )
38 from vllm.platforms.rocm import on_gfx1100 # noqa: E402
39 from vllm.scalar_type import scalar_types # noqa: E402
40
41 device = "cuda"
42
43 gfx1100_only = pytest.mark.skipif(
44 not (
45 on_gfx1100()
46 and hasattr(torch.ops, "_rocm_C")
47 and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
48 ),
49 reason="Requires gfx1100 with moe_gptq_gemm_rdna3 op",
50 )
51
52 # Model configurations: real K/N/top_k/group_size dims, E capped at 16 to
53 # fit in test GPU memory (full E=128/256 would need >20GB for weights alone).
54 # Kernel behavior is E-independent (per-expert tiling), so E=16 is sufficient.
55 MODEL_CONFIGS = [
56 # cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit dims (E capped)
57 pytest.param(16, 2048, 768, 8, 32, id="Qwen3-30B-A3B"),
58 # Qwen3.6-35B-A3B-GPTQ-W4A16-G32 dims (E capped)
59 pytest.param(16, 2048, 512, 8, 32, id="Qwen3.6-35B-A3B"),
60 ]
61
62 # Token counts: decode (1), small batch (4), medium (16), prefill (64)
63 NUM_TOKENS = [1, 4, 16, 64, 256, 512]
64
65
66 def _make_packed_weights(E, K, N):
67 """Create random 4-bit packed weights [E, K/8, N] int32 + shuffle."""
68 w = torch.randint(0, 16, (E, K, N), dtype=torch.int32, device=device)
69 packed = torch.zeros(E, K // 8, N, dtype=torch.int32, device=device)
70 for i in range(8):
71 packed |= (w[:, i::8, :] & 0xF) << (i * 4)
72 for e in range(E):
73 we = packed[e].contiguous()
74 ops.gptq_shuffle(we, 4)
75 packed[e] = we
76 return packed
77
78
79 def _make_scales(E, groups, N, dtype):
80 return torch.rand(E, groups, N, dtype=dtype, device=device) * 0.1
81
82
83 def _make_qzeros(E, groups, N):
84 zeros = torch.full(
85 (groups, N),
86 scalar_types.uint4b8.bias - 1,
87 dtype=torch.int32,
88 device=device,
89 )
90 qz = pack_quantized_values_into_int32(
91 zeros,
92 scalar_types.uint4b8,
93 packed_dim=1,
94 )
95 return qz.unsqueeze(0).expand(E, -1, -1).contiguous()
96
97
98 @gfx1100_only
99 @pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
100 @pytest.mark.parametrize("M", NUM_TOKENS)
101 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
102 @pytest.mark.parametrize("block_size_m", [1, 4])
103 def test_fused_moe_w1_matches_dense(
104 E, K, N_inter, top_k, group_size, M, dtype, block_size_m
105 ):
106 """w1 GEMM via fused kernel matches per-expert dense kernel."""
107 N_gate_up = N_inter * 2
108 groups = K // group_size
109
110 torch.manual_seed(42)
111 x = torch.randn(M, K, dtype=dtype, device=device)
112 w13 = _make_packed_weights(E, K, N_gate_up)
113 w13_s = _make_scales(E, groups, N_gate_up, dtype)
114 w13_z = _make_qzeros(E, groups, N_gate_up)
115
116 topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
117 si, ei, ntp = moe_align_block_size(topk_ids, block_size_m, E)
118
119 # Fused kernel
120 fused_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
121 ops.moe_gptq_gemm_rdna3(
122 x,
123 fused_out,
124 w13,
125 w13_s,
126 w13_z,
127 torch.empty(0, device=device),
128 si,
129 ei,
130 ntp,
131 top_k,
132 block_size_m,
133 False,
134 0,
135 )
136
137 # Per-expert dense reference
138 ref_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
139 for m in range(M):
140 for k in range(top_k):
141 e = topk_ids[m, k].item()
142 flat = m * top_k + k
143 ref = ops.gptq_gemm_rdna3(
144 x[m : m + 1],
145 w13[e],
146 w13_z[e],
147 w13_s[e],
148 False,
149 )
150 ref_out[flat] = ref.squeeze()
151
152 # Split-K atomics can cause minor fp16/bf16 rounding differences
153 # at large K (e.g. K=2048 → 8 K-blocks). Use allclose, not equal.
154 atol = 0.5 if dtype == torch.bfloat16 else 0.1
155 assert torch.allclose(fused_out, ref_out, atol=atol, rtol=0.01), (
156 f"max diff: {(fused_out - ref_out).abs().max().item()}"
157 )
158
159
160 @gfx1100_only
161 @pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
162 @pytest.mark.parametrize("M", NUM_TOKENS)
163 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
164 def test_fused_moe_output_topk_reduces(E, K, N_inter, top_k, group_size, M, dtype):
165 """output_topk fuses moe_sum: multiple experts write to same output row."""
166 groups = K // group_size
167
168 torch.manual_seed(123)
169 x = torch.randn(M * top_k, K, dtype=dtype, device=device)
170 w = _make_packed_weights(E, K, N_inter)
171 ws = _make_scales(E, groups, N_inter, dtype)
172 wz = _make_qzeros(E, groups, N_inter)
173
174 topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
175 topk_w = torch.softmax(
176 torch.randn(M, top_k, device=device),
177 dim=-1,
178 ).float()
179
180 si, ei, ntp = moe_align_block_size(topk_ids, 1, E)
181
182 # Without output_topk: write to [M*top_k, N] then moe_sum
183 flat_out = torch.zeros(M * top_k, N_inter, dtype=dtype, device=device)
184 ops.moe_gptq_gemm_rdna3(
185 x,
186 flat_out,
187 w,
188 ws,
189 wz,
190 topk_w.view(-1),
191 si,
192 ei,
193 ntp,
194 1,
195 1,
196 True,
197 0,
198 )
199 ref = torch.zeros(M, N_inter, dtype=dtype, device=device)
200 ops.moe_sum(flat_out.view(M, top_k, N_inter), ref)
201
202 # With output_topk: write directly to [M, N]
203 fused = torch.zeros(M, N_inter, dtype=dtype, device=device)
204 ops.moe_gptq_gemm_rdna3(
205 x,
206 fused,
207 w,
208 ws,
209 wz,
210 topk_w.view(-1),
211 si,
212 ei,
213 ntp,
214 1,
215 1,
216 True,
217 top_k,
218 )
219
220 atol = 1.0 if dtype == torch.bfloat16 else 0.1
221 assert torch.allclose(fused, ref, atol=atol, rtol=0.01), (
222 f"max diff: {(fused - ref).abs().max().item()}"
223 )
224
225
226 @gfx1100_only
227 @pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
228 @pytest.mark.parametrize("M", NUM_TOKENS)
229 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
230 def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype):
231 """Full MoE forward: w1 + silu_and_mul + w2 with output_topk reduce."""
232 N_gate_up = N_inter * 2
233 hidden = K
234
235 torch.manual_seed(7)
236 x = torch.randn(M, K, dtype=dtype, device=device)
237 w13 = _make_packed_weights(E, K, N_gate_up)
238 w13_s = _make_scales(E, K // group_size, N_gate_up, dtype)
239 w13_z = _make_qzeros(E, K // group_size, N_gate_up)
240 w2 = _make_packed_weights(E, N_inter, hidden)
241 w2_s = _make_scales(E, N_inter // group_size, hidden, dtype)
242 w2_z = _make_qzeros(E, N_inter // group_size, hidden)
243
244 topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
245 topk_w = torch.softmax(
246 torch.randn(M, top_k, device=device),
247 dim=-1,
248 ).float()
249
250 si, ei, ntp = moe_align_block_size(topk_ids, 1, E)
251
252 # Fused path (what apply() does)
253 w1_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
254 ops.moe_gptq_gemm_rdna3(
255 x,
256 w1_out,
257 w13,
258 w13_s,
259 w13_z,
260 torch.empty(0, device=device),
261 si,
262 ei,
263 ntp,
264 top_k,
265 1,
266 False,
267 0,
268 )
269 act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device)
270 apply_moe_activation(MoEActivation.SILU, act_out, w1_out)
271 fused = torch.zeros(M, hidden, dtype=dtype, device=device)
272 ops.moe_gptq_gemm_rdna3(
273 act_out,
274 fused,
275 w2,
276 w2_s,
277 w2_z,
278 topk_w.view(-1),
279 si,
280 ei,
281 ntp,
282 1,
283 1,
284 True,
285 top_k,
286 )
287
288 # Per-expert reference
289 ref = torch.zeros(M, hidden, dtype=dtype, device=device)
290 for m_idx in range(M):
291 for k_idx in range(top_k):
292 e = topk_ids[m_idx, k_idx].item()
293 w = topk_w[m_idx, k_idx].item()
294 r1 = ops.gptq_gemm_rdna3(
295 x[m_idx : m_idx + 1],
296 w13[e],
297 w13_z[e],
298 w13_s[e],
299 False,
300 )
301 a = torch.empty(1, N_inter, dtype=dtype, device=device)
302 apply_moe_activation(MoEActivation.SILU, a, r1)
303 r2 = ops.gptq_gemm_rdna3(
304 a,
305 w2[e],
306 w2_z[e],
307 w2_s[e],
308 False,
309 )
310 ref[m_idx] += r2.squeeze() * w
311
312 # E2E chains w1 + activation + w2 + topk_w + output_topk reduce.
313 # Each step accumulates rounding error (split-K atomics, topk_w
314 # multiply order). Use relative L2 norm like the dense kernel test.
315 diff_l2 = torch.norm(fused.float() - ref.float())
316 ref_l2 = torch.norm(ref.float())
317 rel_l2 = (diff_l2 / ref_l2).item() if ref_l2 > 0 else 0.0
318 threshold = 0.05 if dtype == torch.float16 else 0.10
319 assert rel_l2 < threshold, (
320 f"rel L2 = {rel_l2:.4f} (threshold {threshold}), "
321 f"max abs diff: {(fused - ref).abs().max().item()}"
322 )
323
324
325 @gfx1100_only
326 def test_expert_id_minus_one():
327 """Kernel handles expert_id == -1 (expert parallelism) without crash."""
328 # Qwen3-30B-A3B dims (E capped for memory)
329 E, K, N = 16, 2048, 768
330 groups = K // 32
331
332 w = _make_packed_weights(E, K, N)
333 ws = _make_scales(E, groups, N, torch.bfloat16)
334 wz = _make_qzeros(E, groups, N)
335 x = torch.randn(1, K, dtype=torch.bfloat16, device=device)
336
337 # Manually create sorted_token_ids/expert_ids with -1
338 sorted_ids = torch.tensor([0], dtype=torch.int32, device=device)
339 expert_ids = torch.tensor([-1], dtype=torch.int32, device=device)
340 ntp = torch.tensor([1], dtype=torch.int32, device=device)
341
342 out = torch.zeros(1, N, dtype=torch.bfloat16, device=device)
343 ops.moe_gptq_gemm_rdna3(
344 x,
345 out,
346 w,
347 ws,
348 wz,
349 torch.empty(0, device=device),
350 sorted_ids,
351 expert_ids,
352 ntp,
353 1,
354 1,
355 False,
356 0,
357 )
358 current_platform.synchronize()
359
360 # Output should remain zero (expert skipped)
361 assert torch.equal(out, torch.zeros_like(out))