@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4282fe52d5

[ROCm][CI] Add More AITER quantization/MoE kernel tests (#49375)

Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

Micah Williamson committed Aug 5, 2026 at 16:27 UTC 4282fe52d5353084a5773e924a508e9346dc2718
5 files changed +2335 -2
.buildkite/test-amd.yaml
+20 -2
@@ -1613,8 +1613,12 @@ steps:
1613 source_file_dependencies:
1614 - csrc/quantization/cutlass_w8a8/moe/
1615 - csrc/moe/
1616 + - csrc/rocm/
1617 - tests/kernels/moe
1618 + - tests/kernels/utils.py
1619 - vllm/model_executor/layers/fused_moe/
1620 + - vllm/model_executor/layers/quantization/mxfp4.py
1621 + - vllm/model_executor/layers/quantization/utils/quant_utils.py
1622 - vllm/distributed/device_communicators/
1623 - vllm/envs.py
1624 - vllm/config
@@ -1634,12 +1638,17 @@ steps:
1638 working_dir: "/vllm-workspace/tests"
1639 source_file_dependencies:
1640 - csrc/quantization/
1641 + - csrc/rocm/
1642 - vllm/model_executor/layers/quantization
1643 - tests/kernels/quantization
1639 - - tests/kernels/quantization/test_rocm_skinny_gemms.py
1644 + - tests/kernels/quant_utils.py
1645 + - tests/kernels/utils.py
1646 - vllm/_aiter_ops.py
1647 + - vllm/_custom_ops.py
1648 + - vllm/envs.py
1649 - vllm/platforms/rocm.py
1650 - vllm/model_executor/kernels/
1651 + - vllm/v1/attention/backends/rocm_aiter_fa.py
1652 commands:
1653 - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1654
@@ -3406,8 +3415,12 @@ steps:
3415 source_file_dependencies:
3416 - csrc/quantization/cutlass_w8a8/moe/
3417 - csrc/moe/
3418 + - csrc/rocm/
3419 - tests/kernels/moe
3420 + - tests/kernels/utils.py
3421 - vllm/model_executor/layers/fused_moe/
3422 + - vllm/model_executor/layers/quantization/mxfp4.py
3423 + - vllm/model_executor/layers/quantization/utils/quant_utils.py
3424 - vllm/distributed/device_communicators/
3425 - vllm/envs.py
3426 - vllm/config
@@ -3426,12 +3439,17 @@ steps:
3439 working_dir: "/vllm-workspace/tests"
3440 source_file_dependencies:
3441 - csrc/quantization/
3442 + - csrc/rocm/
3443 - vllm/model_executor/layers/quantization
3444 - tests/kernels/quantization
3431 - - tests/kernels/quantization/test_rocm_skinny_gemms.py
3445 + - tests/kernels/quant_utils.py
3446 + - tests/kernels/utils.py
3447 - vllm/_aiter_ops.py
3448 + - vllm/_custom_ops.py
3449 + - vllm/envs.py
3450 - vllm/platforms/rocm.py
3451 - vllm/model_executor/kernels/
3452 + - vllm/v1/attention/backends/rocm_aiter_fa.py
3453 commands:
3454 - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3455
tests/kernels/moe/test_rocm_aiter_moe.py new
+1079
@@ -0,0 +1,1079 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Kernel tests for ROCm AITER fused MoE.
4 +
5 +This file owns the ROCm-specific fused-MoE custom-op path:
6 +- custom-op registration and fake-tensor support
7 +- enablement gating for ``VLLM_ROCM_USE_AITER_MOE`` and shared experts
8 +- BF16 accuracy for the fused MoE kernel on representative shapes
9 +- gfx950-only AITER MXFP4 W4A16 MoE support, accuracy, and determinism
10 +- MoE-facing FP8 group-quant activation quality
11 +- deterministic routing and representative gfx942 / gfx950 coverage
12 +
13 +Generic fused-MoE backend selection and non-ROCm kernel coverage live in the
14 +generic MoE test files under ``tests/kernels/moe``.
15 +
16 +Raw ROCm AITER helper-op coverage such as ``group_fp8_quant`` and
17 +``per_tensor_quant`` lives in ``tests/kernels/core/test_rocm_aiter_ops.py``.
18 +This file only keeps the MoE-shaped integration angle for those helpers.
19 +"""
20 +
21 +import importlib
22 +import math
23 +import warnings
24 +from typing import Any, NamedTuple
25 +
26 +import pytest
27 +import torch
28 +import torch.nn.functional as F
29 +
30 +from tests.kernels.utils import _assert_deterministic
31 +from vllm.platforms import current_platform
32 +from vllm.platforms.rocm import on_gfx942, on_gfx950
33 +from vllm.utils.torch_utils import set_random_seed
34 +
35 +pytestmark = pytest.mark.skipif(
36 + not current_platform.is_rocm(), reason="ROCm-specific tests"
37 +)
38 +
39 +
40 +# Helpers -----------------------------------------------------------------
41 +
42 +
43 +def _reload_envs():
44 + import vllm.envs as envs
45 +
46 + return importlib.reload(envs)
47 +
48 +
49 +@pytest.fixture(autouse=True)
50 +def _restore_rocm_env_state():
51 + """Restore global env + AITER flag state after every test."""
52 + yield
53 + _reload_envs()
54 + from vllm._aiter_ops import rocm_aiter_ops
55 +
56 + rocm_aiter_ops.refresh_env_variables()
57 +
58 +
59 +def _assert_aiter_supported() -> None:
60 + from vllm._aiter_ops import is_aiter_found_and_supported
61 +
62 + assert is_aiter_found_and_supported(), (
63 + "aiter is required on supported ROCm hardware for this test"
64 + )
65 +
66 +
67 +def _format_observed_rate(count: int, total: int) -> str:
68 + return f"{count / total:.4%} ({count}/{total})"
69 +
70 +
71 +def _format_allowed_rate(rate: float, total: int) -> str:
72 + allowed_count = int(rate * total)
73 + return f"{rate:.4%} (<= {allowed_count}/{total})"
74 +
75 +
76 +def _quantile(values: torch.Tensor, q: float) -> float:
77 + if values.numel() == 0:
78 + return 0.0
79 + return torch.quantile(values, q).item()
80 +
81 +
82 +def _assert_close_budget(
83 + actual: torch.Tensor,
84 + expected: torch.Tensor,
85 + *,
86 + label: str,
87 + atol: float,
88 + rtol: float = 0.0,
89 + pass_rate: float = 0.99999,
90 + max_violation_factor: float = 3.0,
91 +) -> None:
92 + actual_f = actual.detach().float().flatten()
93 + expected_f = expected.detach().float().flatten()
94 + abs_diff = (actual_f - expected_f).abs()
95 + allowed = atol + rtol * expected_f.abs()
96 +
97 + total = abs_diff.numel()
98 + within = abs_diff <= allowed
99 + passed = int(within.sum().item())
100 + failed = total - passed
101 + allowed_fail_rate = 1.0 - pass_rate
102 +
103 + max_abs = abs_diff.max().item()
104 + mean_abs = abs_diff.mean().item()
105 + p99_abs = _quantile(abs_diff, 0.99)
106 + p999_abs = _quantile(abs_diff, 0.999)
107 + worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item()
108 + max_atol = max_violation_factor * atol
109 + above_max_count = int((abs_diff > max_atol).sum().item())
110 +
111 + msg = (
112 + "[rocm_aiter_moe] "
113 + f"{label}: "
114 + f"pass={passed / total:.4%} ({passed}/{total}) "
115 + f"fail={_format_observed_rate(failed, total)} "
116 + f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
117 + f"atol={atol:g} "
118 + f"rtol={rtol:g} "
119 + f"abs>{max_atol:g}={_format_observed_rate(above_max_count, total)} "
120 + f"allowed_above_max={_format_allowed_rate(0.0, total)} "
121 + f"max_abs={max_abs:.6g} "
122 + f"mean_abs={mean_abs:.6g} "
123 + f"p99_abs={p99_abs:.6g} "
124 + f"p999_abs={p999_abs:.6g} "
125 + f"worst_ratio={worst_ratio:.6g}"
126 + )
127 + print(msg)
128 + if failed > 0:
129 + warnings.warn(msg, stacklevel=2)
130 +
131 + assert passed / total >= pass_rate, msg
132 + assert max_abs <= max_atol, msg
133 + assert mean_abs <= atol * 0.25, msg
134 +
135 +
136 +def _assert_group_quant_quality(
137 + actual: torch.Tensor,
138 + expected: torch.Tensor,
139 + *,
140 + label: str,
141 + preferred_rel: float = 0.06,
142 + pass_rate: float = 0.9999,
143 + max_rel: float = 0.55,
144 + max_fail_rate: float = 0.000005,
145 + mean_limit: float = 0.03,
146 +) -> None:
147 + rel = (
148 + (actual.float() - expected.float()).abs()
149 + / expected.float().abs().clamp_min(1e-5)
150 + ).flatten()
151 + total = rel.numel()
152 + within_preferred_count = int((rel <= preferred_rel).sum().item())
153 + fail_count = total - within_preferred_count
154 + allowed_fail_rate = 1.0 - pass_rate
155 + above_max_count = int((rel > max_rel).sum().item())
156 + mean_rel = rel.mean().item()
157 + max_rel_err = rel.max().item()
158 + p99 = _quantile(rel, 0.99)
159 + p999 = _quantile(rel, 0.999)
160 +
161 + msg = (
162 + "[rocm_aiter_moe] "
163 + f"{label}: "
164 + f"rel<={preferred_rel:g} pass={within_preferred_count / total:.4%} "
165 + f"({within_preferred_count}/{total}) "
166 + f"fail={_format_observed_rate(fail_count, total)} "
167 + f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
168 + f"rel>{max_rel:g}={_format_observed_rate(above_max_count, total)} "
169 + f"allowed_above_max={_format_allowed_rate(max_fail_rate, total)} "
170 + f"mean_rel={mean_rel:.6g} "
171 + f"max_rel={max_rel_err:.6g} "
172 + f"p99={p99:.6g} "
173 + f"p999={p999:.6g}"
174 + )
175 + print(msg)
176 + if fail_count > 0:
177 + warnings.warn(msg, stacklevel=2)
178 +
179 + assert within_preferred_count / total >= pass_rate, msg
180 + assert above_max_count / total <= max_fail_rate, msg
181 + assert mean_rel < mean_limit, msg
182 +
183 +
184 +def ref_moe_forward(
185 + hidden_states: torch.Tensor,
186 + w1: torch.Tensor,
187 + w2: torch.Tensor,
188 + topk_weights: torch.Tensor,
189 + topk_ids: torch.Tensor,
190 + activation: str = "silu",
191 +) -> torch.Tensor:
192 + """Float32 mask-based MoE reference for the ROCm fused-MoE kernel."""
193 + num_tokens, hidden_dim = hidden_states.shape
194 + num_experts = w1.shape[0]
195 + intermediate_dim = w1.shape[1] // 2
196 + topk = topk_ids.shape[1]
197 + device = hidden_states.device
198 +
199 + hidden_states_f = hidden_states.float()
200 + w1_f = w1.float()
201 + w2_f = w2.float()
202 +
203 + expanded = hidden_states_f.view(num_tokens, 1, hidden_dim).expand(
204 + num_tokens, topk, hidden_dim
205 + )
206 + expanded = expanded.reshape(num_tokens * topk, hidden_dim)
207 + output = torch.zeros(
208 + num_tokens * topk, hidden_dim, dtype=torch.float32, device=device
209 + )
210 + flat_topk_ids = topk_ids.view(-1).long()
211 +
212 + for expert_idx in range(num_experts):
213 + expert_mask = flat_topk_ids == expert_idx
214 + if expert_mask.sum() == 0:
215 + continue
216 + gate_up = expanded[expert_mask] @ w1_f[expert_idx].T
217 + gate = gate_up[:, :intermediate_dim]
218 + up = gate_up[:, intermediate_dim:]
219 + if activation == "silu":
220 + act = F.silu(gate) * up
221 + elif activation == "gelu":
222 + act = F.gelu(gate) * up
223 + else:
224 + raise ValueError(f"Unknown activation: {activation}")
225 + output[expert_mask] = act @ w2_f[expert_idx].T
226 +
227 + output = output.view(num_tokens, topk, hidden_dim)
228 + weights = topk_weights.float().view(num_tokens, topk, 1)
229 + return (output * weights).sum(dim=1)
230 +
231 +
232 +class AiterMxfp4MoeCase(NamedTuple):
233 + hidden_states: torch.Tensor
234 + w1_kernel: torch.Tensor
235 + w2_kernel: torch.Tensor
236 + w1_ref: torch.Tensor
237 + w2_ref: torch.Tensor
238 + topk_weights: torch.Tensor
239 + topk_ids: torch.Tensor
240 + moe_config: Any
241 + quant_config: Any
242 +
243 +
244 +def _make_topk_ids(
245 + num_tokens: int,
246 + num_experts: int,
247 + topk: int,
248 + *,
249 + device: str = "cuda",
250 +) -> torch.Tensor:
251 + """Generate distinct expert IDs per token with the same top-k shape as
252 + production routing."""
253 + router_logits = torch.randn(num_tokens, num_experts, device=device)
254 + _, topk_ids = torch.topk(torch.softmax(router_logits, dim=-1), k=topk, dim=-1)
255 + return topk_ids.to(torch.int32)
256 +
257 +
258 +def _shuffle_moe_weights(
259 + w1: torch.Tensor,
260 + w2: torch.Tensor,
261 +) -> tuple[torch.Tensor, torch.Tensor]:
262 + """Shuffle MoE weights into the AITER CK layout used in production."""
263 + from vllm._aiter_ops import rocm_aiter_ops
264 +
265 + w1_shuffled, w2_shuffled = rocm_aiter_ops.shuffle_weights(w1, w2)
266 + w1_shuffled.is_shuffled = True
267 + w2_shuffled.is_shuffled = True
268 + return w1_shuffled, w2_shuffled
269 +
270 +
271 +def _make_moe_case(
272 + *,
273 + num_tokens: int,
274 + hidden_dim: int,
275 + intermediate_dim: int,
276 + num_experts: int,
277 + topk: int,
278 + seed: int,
279 +) -> dict[str, torch.Tensor]:
280 + torch.set_default_device("cuda")
281 + set_random_seed(seed)
282 +
283 + hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
284 + w1 = torch.randn(
285 + num_experts,
286 + intermediate_dim * 2,
287 + hidden_dim,
288 + dtype=torch.bfloat16,
289 + ) / math.sqrt(hidden_dim)
290 + w2 = torch.randn(
291 + num_experts,
292 + hidden_dim,
293 + intermediate_dim,
294 + dtype=torch.bfloat16,
295 + ) / math.sqrt(intermediate_dim)
296 + topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32)
297 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
298 + topk_ids = _make_topk_ids(num_tokens, num_experts, topk)
299 + return {
300 + "hidden_states": hidden_states,
301 + "w1": w1,
302 + "w2": w2,
303 + "topk_weights": topk_weights,
304 + "topk_ids": topk_ids,
305 + }
306 +
307 +
308 +def _interleave_gate_up_rows(t: torch.Tensor) -> torch.Tensor:
309 + """Reorder contiguous ``[gate; up]`` rows into gpt-oss interleaved order."""
310 + e, two_i = t.shape[0], t.shape[1]
311 + i, rest = two_i // 2, t.shape[2:]
312 + perm = (0, 2, 1, *range(3, 3 + len(rest)))
313 + return t.view(e, 2, i, *rest).permute(*perm).contiguous().view(e, two_i, *rest)
314 +
315 +
316 +def _make_aiter_mxfp4_moe_case(
317 + *,
318 + num_tokens: int,
319 + hidden_dim: int,
320 + intermediate_dim: int,
321 + num_experts: int,
322 + topk: int,
323 + seed: int,
324 +) -> AiterMxfp4MoeCase:
325 + from triton_kernels.numerics_details.mxfp import (
326 + downcast_to_mxfp,
327 + upcast_from_mxfp,
328 + )
329 +
330 + from tests.kernels.moe.utils import make_dummy_moe_config
331 + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
332 + Mxfp4MoeBackend,
333 + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
334 + make_mxfp4_moe_quant_config,
335 + )
336 +
337 + torch.set_default_device("cuda")
338 + set_random_seed(seed)
339 +
340 + hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
341 + w1 = torch.randn(
342 + num_experts,
343 + intermediate_dim * 2,
344 + hidden_dim,
345 + dtype=torch.bfloat16,
346 + ) / math.sqrt(hidden_dim)
347 + w2 = torch.randn(
348 + num_experts,
349 + hidden_dim,
350 + intermediate_dim,
351 + dtype=torch.bfloat16,
352 + ) / math.sqrt(intermediate_dim)
353 + topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32)
354 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
355 + topk_ids = _make_topk_ids(num_tokens, num_experts, topk)
356 +
357 + w1_q, w1_scale = downcast_to_mxfp(w1, torch.uint8, axis=-1)
358 + w2_q, w2_scale = downcast_to_mxfp(w2, torch.uint8, axis=-1)
359 +
360 + w1_ref = upcast_from_mxfp(w1_q, w1_scale, torch.bfloat16, axis=-1)
361 + w2_ref = upcast_from_mxfp(w2_q, w2_scale, torch.bfloat16, axis=-1)
362 +
363 + (
364 + w1_kernel,
365 + w2_kernel,
366 + w1_scale_kernel,
367 + w2_scale_kernel,
368 + _,
369 + _,
370 + ) = convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
371 + mxfp4_backend=Mxfp4MoeBackend.AITER,
372 + layer=torch.nn.Module(),
373 + # w13 rows must be gpt-oss-interleaved; the converter de-interleaves them
374 + # back to the contiguous gate/up blocks that ``w1_ref`` / ``ref_moe_forward``
375 + # use. w2 has no gate/up split, so it stays as-is.
376 + w13_weight=_interleave_gate_up_rows(w1_q),
377 + w2_weight=w2_q.clone(),
378 + w13_weight_scale=_interleave_gate_up_rows(w1_scale),
379 + w2_weight_scale=w2_scale.clone(),
380 + )
381 +
382 + moe_config = make_dummy_moe_config(
383 + num_experts=num_experts,
384 + experts_per_token=topk,
385 + hidden_dim=hidden_dim,
386 + intermediate_size=intermediate_dim,
387 + in_dtype=torch.bfloat16,
388 + )
389 + quant_config = make_mxfp4_moe_quant_config(
390 + Mxfp4MoeBackend.AITER,
391 + w1_scale=w1_scale_kernel,
392 + w2_scale=w2_scale_kernel,
393 + )
394 +
395 + return AiterMxfp4MoeCase(
396 + hidden_states=hidden_states,
397 + w1_kernel=w1_kernel,
398 + w2_kernel=w2_kernel,
399 + w1_ref=w1_ref,
400 + w2_ref=w2_ref,
401 + topk_weights=topk_weights,
402 + topk_ids=topk_ids,
403 + moe_config=moe_config,
404 + quant_config=quant_config,
405 + )
406 +
407 +
408 +def _run_fused_moe(
409 + hidden_states: torch.Tensor,
410 + w1: torch.Tensor,
411 + w2: torch.Tensor,
412 + topk_weights: torch.Tensor,
413 + topk_ids: torch.Tensor,
414 + *,
415 + activation_method: int,
416 + quant_method: int,
417 +) -> torch.Tensor:
418 + w1_shuffled, w2_shuffled = _shuffle_moe_weights(w1, w2)
419 + return torch.ops.vllm.rocm_aiter_fused_moe(
420 + hidden_states,
421 + w1_shuffled,
422 + w2_shuffled,
423 + topk_weights,
424 + topk_ids,
425 + expert_mask=None,
426 + activation_method=activation_method,
427 + quant_method=quant_method,
428 + doweight_stage1=False,
429 + )
430 +
431 +
432 +# Custom op tests ---------------------------------------------------------
433 +
434 +
435 +def test_aiter_fused_moe_custom_op_registered():
436 + """The main fused-MoE custom op should stay registered for runtime use."""
437 + _assert_aiter_supported()
438 + import vllm._aiter_ops as aiter_ops # noqa: F401
439 +
440 + assert hasattr(torch.ops.vllm, "rocm_aiter_fused_moe")
441 + assert callable(torch.ops.vllm.rocm_aiter_fused_moe)
442 +
443 +
444 +def test_aiter_asm_moe_tkw1_custom_op_registered():
445 + """The tkw1 custom op should stay registered for FP8 apply-router-weight
446 + paths."""
447 + _assert_aiter_supported()
448 + import vllm._aiter_ops as aiter_ops # noqa: F401
449 +
450 + assert hasattr(torch.ops.vllm, "rocm_aiter_asm_moe_tkw1")
451 + assert callable(torch.ops.vllm.rocm_aiter_asm_moe_tkw1)
452 +
453 +
454 +def test_aiter_fused_moe_fake_tensor_support():
455 + """The fused-MoE op should preserve fake-tensor compatibility for
456 + torch.compile-style tracing."""
457 + _assert_aiter_supported()
458 + import vllm._aiter_ops # noqa: F401
459 +
460 + num_tokens = 16
461 + hidden_dim = 1024
462 + intermediate_dim = 2048
463 + num_experts = 8
464 + topk = 2
465 +
466 + hidden_states = torch.randn(
467 + num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda"
468 + )
469 + w1 = torch.randn(
470 + num_experts,
471 + intermediate_dim * 2,
472 + hidden_dim,
473 + dtype=torch.bfloat16,
474 + device="cuda",
475 + )
476 + w2 = torch.randn(
477 + num_experts,
478 + hidden_dim,
479 + intermediate_dim,
480 + dtype=torch.bfloat16,
481 + device="cuda",
482 + )
483 + topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32, device="cuda")
484 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
485 + topk_ids = _make_topk_ids(num_tokens, num_experts, topk, device="cuda")
486 +
487 + torch.library.opcheck(
488 + torch.ops.vllm.rocm_aiter_fused_moe,
489 + (hidden_states, w1, w2, topk_weights, topk_ids),
490 + kwargs={
491 + "expert_mask": None,
492 + "activation_method": 0,
493 + "quant_method": 0,
494 + "doweight_stage1": False,
495 + },
496 + test_utils=("test_faketensor",),
497 + )
498 +
499 +
500 +# Env gating tests --------------------------------------------------------
501 +
502 +
503 +@pytest.mark.parametrize(
504 + ("use_aiter", "use_moe", "expected"),
505 + [
506 + (True, True, True),
507 + (True, False, False),
508 + (False, True, False),
509 + (False, False, False),
510 + ],
511 +)
512 +def test_aiter_moe_enablement_follows_env(
513 + use_aiter: bool,
514 + use_moe: bool,
515 + expected: bool,
516 + monkeypatch: pytest.MonkeyPatch,
517 +):
518 + """The fused-MoE gate should depend only on the main AITER toggle and the
519 + MoE-specific toggle."""
520 + from vllm._aiter_ops import rocm_aiter_ops
521 +
522 + _assert_aiter_supported()
523 +
524 + with monkeypatch.context() as mp:
525 + mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
526 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_moe else "0")
527 + _reload_envs()
528 + rocm_aiter_ops.refresh_env_variables()
529 +
530 + assert rocm_aiter_ops.is_fused_moe_enabled() is expected
531 +
532 +
533 +@pytest.mark.parametrize(
534 + ("use_aiter", "use_moe", "use_shared", "expected"),
535 + [
536 + (True, True, True, True),
537 + (True, True, False, False),
538 + (True, False, True, False),
539 + (False, True, True, False),
540 + ],
541 +)
542 +def test_aiter_moe_shared_experts_enablement_follows_env(
543 + use_aiter: bool,
544 + use_moe: bool,
545 + use_shared: bool,
546 + expected: bool,
547 + monkeypatch: pytest.MonkeyPatch,
548 +):
549 + """Shared-expert fusion should only be enabled when the fused-MoE path is
550 + enabled too."""
551 + from vllm._aiter_ops import rocm_aiter_ops
552 +
553 + _assert_aiter_supported()
554 +
555 + with monkeypatch.context() as mp:
556 + mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
557 + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_moe else "0")
558 + mp.setenv(
559 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS",
560 + "1" if use_shared else "0",
561 + )
562 + _reload_envs()
563 + rocm_aiter_ops.refresh_env_variables()
564 +
565 + assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected
566 +
567 +
568 +@pytest.mark.parametrize("moe_padding", [True, False])
569 +def test_aiter_moe_padding_env_var(
570 + moe_padding: bool,
571 + monkeypatch: pytest.MonkeyPatch,
572 +):
573 + """The ROCm MoE padding env var should keep its exact parse contract."""
574 + with monkeypatch.context() as mp:
575 + mp.setenv("VLLM_ROCM_MOE_PADDING", "1" if moe_padding else "0")
576 + envs = _reload_envs()
577 + assert envs.VLLM_ROCM_MOE_PADDING is moe_padding
578 +
579 +
580 +# Enum tests --------------------------------------------------------------
581 +
582 +
583 +def test_quant_method_enum_values():
584 + """The AITER quant-method bridge enum should keep its wire values."""
585 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import QuantMethod
586 +
587 + assert QuantMethod.NO == 0
588 + assert QuantMethod.PER_TENSOR == 1
589 + assert QuantMethod.PER_TOKEN == 2
590 + assert QuantMethod.BLOCK_1X32 == 3
591 + assert QuantMethod.BLOCK_1X128 == 4
592 + assert QuantMethod.BLOCK_128x128 == 5
593 +
594 +
595 +def test_activation_method_enum_values():
596 + """The AITER activation-method bridge enum should keep its wire values."""
597 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
598 + ActivationMethod,
599 + )
600 +
601 + assert ActivationMethod.SILU == 0
602 + assert ActivationMethod.GELU == 1
603 +
604 +
605 +# MXFP4 kernel tests ------------------------------------------------------
606 +
607 +
608 +def test_aiter_mxfp4_quant_scheme_support_matches_gfx950():
609 + """AITER MXFP4 MoE support should stay gfx950-only."""
610 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
611 + AiterExperts,
612 + )
613 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
614 + kMxfp4Static,
615 + )
616 +
617 + assert AiterExperts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950()
618 +
619 +
620 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
621 +def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy():
622 + """The gfx950 AITER MXFP4 W4A16 MoE path should match the dequantized
623 + MXFP4 reference."""
624 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
625 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
626 + rocm_aiter_fused_experts,
627 + )
628 +
629 + _assert_aiter_supported()
630 + case = _make_aiter_mxfp4_moe_case(
631 + num_tokens=32,
632 + hidden_dim=512,
633 + intermediate_dim=1024,
634 + num_experts=4,
635 + topk=2,
636 + seed=11,
637 + )
638 + ref_out = ref_moe_forward(
639 + case.hidden_states,
640 + case.w1_ref,
641 + case.w2_ref,
642 + case.topk_weights,
643 + case.topk_ids,
644 + )
645 + out = rocm_aiter_fused_experts(
646 + hidden_states=case.hidden_states,
647 + w1=case.w1_kernel,
648 + w2=case.w2_kernel,
649 + topk_weights=case.topk_weights,
650 + topk_ids=case.topk_ids,
651 + activation=MoEActivation.SWIGLUOAI,
652 + quant_config=case.quant_config,
653 + moe_config=case.moe_config,
654 + expert_map=None,
655 + )
656 +
657 + assert out.shape == case.hidden_states.shape
658 + _assert_close_budget(
659 + out.float(),
660 + ref_out.float(),
661 + label="mi350_mxfp4_w4a16_accuracy",
662 + atol=0.1,
663 + rtol=0.1,
664 + pass_rate=0.99,
665 + max_violation_factor=2.0,
666 + )
667 +
668 +
669 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
670 +def test_aiter_fused_moe_mi350_mxfp4_w4a16_determinism():
671 + """The gfx950 AITER MXFP4 W4A16 MoE path should stay bitwise
672 + deterministic."""
673 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
674 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
675 + rocm_aiter_fused_experts,
676 + )
677 +
678 + _assert_aiter_supported()
679 + case = _make_aiter_mxfp4_moe_case(
680 + num_tokens=8,
681 + hidden_dim=512,
682 + intermediate_dim=1024,
683 + num_experts=4,
684 + topk=2,
685 + seed=13,
686 + )
687 +
688 + def run_mxfp4_moe():
689 + return rocm_aiter_fused_experts(
690 + hidden_states=case.hidden_states,
691 + w1=case.w1_kernel,
692 + w2=case.w2_kernel,
693 + topk_weights=case.topk_weights,
694 + topk_ids=case.topk_ids,
695 + activation=MoEActivation.SWIGLUOAI,
696 + quant_config=case.quant_config,
697 + moe_config=case.moe_config,
698 + expert_map=None,
699 + )
700 +
701 + _assert_deterministic(run_mxfp4_moe, n_runs=4)
702 +
703 +
704 +# FP8 group-quant tests ---------------------------------------------------
705 +
706 +
707 +@pytest.mark.skipif(
708 + not (on_gfx942() or on_gfx950()),
709 + reason="gfx942/gfx950 ROCm only",
710 +)
711 +@pytest.mark.parametrize("num_tokens,hidden_dim", [(16, 2048), (64, 4096), (128, 8192)])
712 +def test_aiter_moe_group_fp8_quant_reconstructs_hidden_states(
713 + num_tokens: int,
714 + hidden_dim: int,
715 +):
716 + """The MoE-facing FP8 group-quant helper should reconstruct hidden states
717 + within the expected FP8 error budget.
718 +
719 + The raw op's shape and standalone roundtrip contracts are covered in
720 + ``tests/kernels/core/test_rocm_aiter_ops.py``. This test keeps the
721 + representative hidden-state sizes that the fused-MoE path actually cares
722 + about.
723 + """
724 + from vllm._aiter_ops import rocm_aiter_ops
725 +
726 + _assert_aiter_supported()
727 + torch.set_default_device("cuda")
728 + set_random_seed(1)
729 +
730 + group_size = 128
731 + hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
732 + x_fp8, scales = rocm_aiter_ops.group_fp8_quant(hidden_states, group_size)
733 +
734 + assert x_fp8.shape == hidden_states.shape
735 + assert scales.shape == (num_tokens, (hidden_dim + group_size - 1) // group_size)
736 + assert scales.dtype == torch.float32
737 +
738 + scales_expanded = scales.repeat_interleave(group_size, dim=1)[:, :hidden_dim]
739 + dequantized = x_fp8.float() * scales_expanded
740 +
741 + _assert_group_quant_quality(
742 + dequantized,
743 + hidden_states,
744 + label=f"group_fp8_quant shape=({num_tokens}, {hidden_dim})",
745 + )
746 +
747 +
748 +# Kernel accuracy tests ---------------------------------------------------
749 +
750 +
751 +@pytest.mark.parametrize(
752 + "num_tokens,hidden_dim,intermediate_dim",
753 + [
754 + (16, 512, 1024),
755 + (128, 2048, 4096),
756 + (2048, 4096, 11008),
757 + ],
758 +)
759 +def test_aiter_fused_moe_bf16_accuracy(
760 + num_tokens: int,
761 + hidden_dim: int,
762 + intermediate_dim: int,
763 +):
764 + """The ROCm AITER fused-MoE BF16 path should match the float32 reference
765 + on representative shapes."""
766 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
767 + ActivationMethod,
768 + QuantMethod,
769 + )
770 +
771 + _assert_aiter_supported()
772 + case = _make_moe_case(
773 + num_tokens=num_tokens,
774 + hidden_dim=hidden_dim,
775 + intermediate_dim=intermediate_dim,
776 + num_experts=8,
777 + topk=2,
778 + seed=0,
779 + )
780 + ref_out = ref_moe_forward(
781 + case["hidden_states"],
782 + case["w1"],
783 + case["w2"],
784 + case["topk_weights"],
785 + case["topk_ids"],
786 + activation="silu",
787 + )
788 + out = _run_fused_moe(
789 + case["hidden_states"],
790 + case["w1"],
791 + case["w2"],
792 + case["topk_weights"],
793 + case["topk_ids"],
794 + activation_method=int(ActivationMethod.SILU),
795 + quant_method=int(QuantMethod.NO),
796 + )
797 +
798 + assert out.shape == (num_tokens, hidden_dim)
799 + _assert_close_budget(
800 + out.float(),
801 + ref_out,
802 + label=f"bf16_accuracy shape=({num_tokens}, {hidden_dim}, {intermediate_dim})",
803 + atol=0.05,
804 + rtol=0.0,
805 + )
806 +
807 +
808 +def test_aiter_fused_moe_gelu_accuracy():
809 + """The GELU activation variant should stay aligned with the float32
810 + reference."""
811 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
812 + ActivationMethod,
813 + QuantMethod,
814 + )
815 +
816 + _assert_aiter_supported()
817 + case = _make_moe_case(
818 + num_tokens=32,
819 + hidden_dim=512,
820 + intermediate_dim=1024,
821 + num_experts=4,
822 + topk=2,
823 + seed=42,
824 + )
825 + ref_out = ref_moe_forward(
826 + case["hidden_states"],
827 + case["w1"],
828 + case["w2"],
829 + case["topk_weights"],
830 + case["topk_ids"],
831 + activation="gelu",
832 + )
833 + out = _run_fused_moe(
834 + case["hidden_states"],
835 + case["w1"],
836 + case["w2"],
837 + case["topk_weights"],
838 + case["topk_ids"],
839 + activation_method=int(ActivationMethod.GELU),
840 + quant_method=int(QuantMethod.NO),
841 + )
842 +
843 + assert out.shape == case["hidden_states"].shape
844 + _assert_close_budget(
845 + out.float(),
846 + ref_out,
847 + label="gelu_accuracy",
848 + atol=0.05,
849 + rtol=0.0,
850 + )
851 +
852 +
853 +def test_aiter_fused_moe_determinism():
854 + """The BF16 fused-MoE kernel should stay bitwise deterministic for the
855 + same inputs."""
856 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
857 + ActivationMethod,
858 + QuantMethod,
859 + )
860 +
861 + _assert_aiter_supported()
862 + case = _make_moe_case(
863 + num_tokens=8,
864 + hidden_dim=256,
865 + intermediate_dim=512,
866 + num_experts=4,
867 + topk=2,
868 + seed=2,
869 + )
870 +
871 + def run_moe():
872 + return _run_fused_moe(
873 + case["hidden_states"],
874 + case["w1"],
875 + case["w2"],
876 + case["topk_weights"],
877 + case["topk_ids"],
878 + activation_method=int(ActivationMethod.SILU),
879 + quant_method=int(QuantMethod.NO),
880 + )
881 +
882 + _assert_deterministic(run_moe, n_runs=4)
883 +
884 +
885 +# Routed end-to-end tests -------------------------------------------------
886 +
887 +
888 +@pytest.mark.parametrize(
889 + "num_tokens,num_experts,topk",
890 + [
891 + (1, 8, 2),
892 + (16, 8, 2),
893 + (64, 16, 4),
894 + ],
895 +)
896 +def test_aiter_fused_moe_end_to_end(
897 + num_tokens: int,
898 + num_experts: int,
899 + topk: int,
900 +):
901 + """The full router-logits to top-k to fused-MoE path should stay accurate."""
902 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
903 + ActivationMethod,
904 + QuantMethod,
905 + )
906 +
907 + _assert_aiter_supported()
908 + torch.set_default_device("cuda")
909 + set_random_seed(7)
910 +
911 + hidden_dim = 512
912 + intermediate_dim = 1024
913 + hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
914 + w1 = torch.randn(
915 + num_experts,
916 + intermediate_dim * 2,
917 + hidden_dim,
918 + dtype=torch.bfloat16,
919 + ) / math.sqrt(hidden_dim)
920 + w2 = torch.randn(
921 + num_experts,
922 + hidden_dim,
923 + intermediate_dim,
924 + dtype=torch.bfloat16,
925 + ) / math.sqrt(intermediate_dim)
926 +
927 + router_logits = torch.randn(num_tokens, num_experts, device="cuda")
928 + router_probs = torch.softmax(router_logits, dim=-1)
929 + topk_weights, topk_ids = torch.topk(router_probs, k=topk, dim=-1)
930 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
931 + topk_weights = topk_weights.float()
932 + topk_ids = topk_ids.to(torch.int32)
933 +
934 + ref_out = ref_moe_forward(hidden_states, w1, w2, topk_weights, topk_ids)
935 + out = _run_fused_moe(
936 + hidden_states,
937 + w1,
938 + w2,
939 + topk_weights,
940 + topk_ids,
941 + activation_method=int(ActivationMethod.SILU),
942 + quant_method=int(QuantMethod.NO),
943 + )
944 +
945 + assert out.shape == (num_tokens, hidden_dim)
946 + assert out.dtype == torch.bfloat16
947 + _assert_close_budget(
948 + out.float(),
949 + ref_out,
950 + label=f"end_to_end tokens={num_tokens} experts={num_experts} topk={topk}",
951 + atol=0.05,
952 + rtol=0.0,
953 + )
954 +
955 +
956 +# Arch-specific tests -----------------------------------------------------
957 +
958 +
959 +@pytest.mark.skipif(
960 + not (on_gfx942() or on_gfx950()),
961 + reason="gfx942/gfx950 ROCm only",
962 +)
963 +def test_aiter_fused_moe_mi3xx_bf16_accuracy():
964 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
965 + ActivationMethod,
966 + QuantMethod,
967 + )
968 +
969 + _assert_aiter_supported()
970 + case = _make_moe_case(
971 + num_tokens=64,
972 + hidden_dim=4096,
973 + intermediate_dim=11008,
974 + num_experts=8,
975 + topk=2,
976 + seed=42,
977 + )
978 + ref_out = ref_moe_forward(
979 + case["hidden_states"],
980 + case["w1"],
981 + case["w2"],
982 + case["topk_weights"],
983 + case["topk_ids"],
984 + )
985 + out = _run_fused_moe(
986 + case["hidden_states"],
987 + case["w1"],
988 + case["w2"],
989 + case["topk_weights"],
990 + case["topk_ids"],
991 + activation_method=int(ActivationMethod.SILU),
992 + quant_method=int(QuantMethod.NO),
993 + )
994 +
995 + _assert_close_budget(
996 + out.float(),
997 + ref_out,
998 + label=(f"mi3xx_bf16_accuracy arch={'gfx942' if on_gfx942() else 'gfx950'}"),
999 + atol=0.05,
1000 + rtol=0.0,
1001 + )
1002 +
1003 +
1004 +@pytest.mark.skipif(
1005 + not (on_gfx942() or on_gfx950()),
1006 + reason="gfx942/gfx950 ROCm only",
1007 +)
1008 +def test_aiter_fused_moe_mi3xx_fp8_accuracy():
1009 + """The MI3xx FP8 per-tensor MoE path should stay within the measured FP8
1010 + error budget."""
1011 + from tests.kernels.moe.utils import make_test_weights
1012 + from vllm._aiter_ops import rocm_aiter_ops
1013 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1014 + ActivationMethod,
1015 + QuantMethod,
1016 + )
1017 +
1018 + _assert_aiter_supported()
1019 + torch.set_default_device("cuda")
1020 + set_random_seed(99)
1021 +
1022 + num_tokens = 32
1023 + hidden_dim = 512
1024 + intermediate_dim = 1024
1025 + num_experts = 4
1026 + topk = 2
1027 +
1028 + hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16) / 10
1029 + fp8_dtype = current_platform.fp8_dtype()
1030 + (w1_bf16, w1_fp8, w1_scale, _), (w2_bf16, w2_fp8, w2_scale, _) = make_test_weights(
1031 + num_experts,
1032 + intermediate_dim * 2,
1033 + hidden_dim,
1034 + torch.bfloat16,
1035 + torch.float8_e4m3fn,
1036 + per_out_ch_quant=False,
1037 + )
1038 + _, a1_scale = rocm_aiter_ops.per_tensor_quant(hidden_states, fp8_dtype)
1039 + router_logits = torch.randn(num_tokens, num_experts, device="cuda")
1040 + router_probs = torch.softmax(router_logits, dim=-1)
1041 + topk_weights, topk_ids = torch.topk(router_probs, k=topk, dim=-1)
1042 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1043 + topk_weights = topk_weights.float()
1044 + topk_ids = topk_ids.to(torch.int32)
1045 +
1046 + ref_out = ref_moe_forward(
1047 + hidden_states,
1048 + w1_bf16,
1049 + w2_bf16,
1050 + topk_weights,
1051 + topk_ids,
1052 + )
1053 +
1054 + w1_shuffled, w2_shuffled = _shuffle_moe_weights(w1_fp8, w2_fp8)
1055 + out = torch.ops.vllm.rocm_aiter_fused_moe(
1056 + hidden_states,
1057 + w1_shuffled,
1058 + w2_shuffled,
1059 + topk_weights,
1060 + topk_ids,
1061 + expert_mask=None,
1062 + activation_method=int(ActivationMethod.SILU),
1063 + quant_method=int(QuantMethod.PER_TENSOR),
1064 + doweight_stage1=False,
1065 + w1_scale=w1_scale,
1066 + w2_scale=w2_scale,
1067 + a1_scale=a1_scale,
1068 + )
1069 +
1070 + assert out.shape == hidden_states.shape
1071 + _assert_close_budget(
1072 + out.float(),
1073 + ref_out,
1074 + label=(f"mi3xx_fp8_accuracy arch={'gfx942' if on_gfx942() else 'gfx950'}"),
1075 + atol=0.02,
1076 + rtol=0.0,
1077 + pass_rate=1.0,
1078 + max_violation_factor=1.5,
1079 + )
tests/kernels/quantization/test_rocm_fp8.py new
+541
@@ -0,0 +1,541 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ROCm-specific tests for FP8 quantization and kernels."""
4 +
5 +import importlib
6 +import math
7 +import warnings
8 +
9 +import pytest
10 +import torch
11 +
12 +import vllm._custom_ops as ops
13 +from tests.kernels.quant_utils import (
14 + ref_dynamic_per_tensor_fp8_quant,
15 + ref_dynamic_per_token_quant,
16 +)
17 +from tests.kernels.utils import _assert_deterministic
18 +from vllm.platforms import current_platform
19 +from vllm.utils.torch_utils import set_random_seed
20 +
21 +pytestmark = pytest.mark.skipif(
22 + not current_platform.is_rocm(), reason="ROCm-specific tests"
23 +)
24 +
25 +DEVICE = "cuda"
26 +FP8_DTYPE = current_platform.fp8_dtype()
27 +fp8_only = pytest.mark.skipif(
28 + not current_platform.supports_fp8(), reason="ROCm FP8 only"
29 +)
30 +
31 +ENV_DEFAULTS = {
32 + "VLLM_ROCM_FP8_PADDING": True,
33 + "VLLM_ROCM_USE_AITER_FP8BMM": True,
34 + "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": False,
35 +}
36 +
37 +SCALED_MM_CASES = [
38 + pytest.param((1, 128, 128), 0.05, 0.15, 0.95, id="1x128x128"),
39 + pytest.param((4, 256, 512), 0.10, 0.30, 0.96, id="4x256x512"),
40 + pytest.param((8, 512, 1024), 0.20, 0.50, 0.99, id="8x512x1024"),
41 +]
42 +
43 +
44 +def _reload_envs():
45 + import vllm.envs as envs
46 +
47 + return importlib.reload(envs)
48 +
49 +
50 +@pytest.fixture(autouse=True)
51 +def _restore_rocm_env_state():
52 + """Restore global env + AITER flag state after every test.
53 +
54 + Tests here reload ``vllm.envs`` (a module singleton) and refresh the
55 + cached ``rocm_aiter_ops`` env flags. Restoring in a fixture teardown
56 + guarantees cleanup even when an assertion fails mid-test, preventing
57 + state from leaking into other tests in the same worker.
58 + """
59 + yield
60 + _reload_envs()
61 + from vllm._aiter_ops import rocm_aiter_ops
62 +
63 + rocm_aiter_ops.refresh_env_variables()
64 +
65 +
66 +def _quantile(flat: torch.Tensor, q: float) -> float:
67 + if flat.numel() == 1:
68 + return flat.item()
69 + return torch.quantile(flat, q).item()
70 +
71 +
72 +def _format_observed_rate(count: int, total: int) -> str:
73 + return f"{count / total:.4%} ({count}/{total})"
74 +
75 +
76 +def _format_allowed_rate(rate: float, total: int) -> str:
77 + allowed_count = math.floor(rate * total)
78 + return f"{rate:.4%} (<= {allowed_count}/{total})"
79 +
80 +
81 +def _assert_abs_error_budget(
82 + actual: torch.Tensor,
83 + expected: torch.Tensor,
84 + *,
85 + label: str,
86 + tight_atol: float,
87 + max_atol: float,
88 + pass_rate: float,
89 +) -> None:
90 + diff = (actual.float() - expected.float()).abs().flatten()
91 + total = diff.numel()
92 + within_tight_count = int((diff <= tight_atol).sum().item())
93 + fail_count = total - within_tight_count
94 + above_max_count = int((diff > max_atol).sum().item())
95 + within_tight = within_tight_count / total
96 + allowed_fail_rate = 1.0 - pass_rate
97 + above_max = above_max_count / total
98 + max_fail_rate = 0.0
99 + max_diff = diff.max().item()
100 + mean_diff = diff.mean().item()
101 + p99 = _quantile(diff, 0.99)
102 + p999 = _quantile(diff, 0.999)
103 + msg = (
104 + f"[rocm_fp8] {label}: abs<={tight_atol} pass={within_tight:.4%} "
105 + f"({within_tight_count}/{total}) "
106 + f"fail={_format_observed_rate(fail_count, total)} "
107 + f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
108 + f"abs>{max_atol}={_format_observed_rate(above_max_count, total)} "
109 + f"allowed_above_max={_format_allowed_rate(max_fail_rate, total)} "
110 + f"max_diff={max_diff:.6g} mean_diff={mean_diff:.6g} "
111 + f"p99={p99:.6g} p999={p999:.6g}"
112 + )
113 + print(msg)
114 + if within_tight < 1.0:
115 + warnings.warn(msg, stacklevel=2)
116 + assert within_tight >= pass_rate, msg
117 + assert above_max <= max_fail_rate, msg
118 +
119 +
120 +def _assert_rel_error_budget(
121 + actual: torch.Tensor,
122 + expected: torch.Tensor,
123 + *,
124 + label: str,
125 + tight_rtol: float,
126 + max_rtol: float,
127 + pass_rate: float,
128 + max_fail_rate: float,
129 +) -> None:
130 + rel = (
131 + (actual.float() - expected.float()).abs()
132 + / expected.float().abs().clamp_min(1e-5)
133 + ).flatten()
134 + total = rel.numel()
135 + within_tight_count = int((rel <= tight_rtol).sum().item())
136 + fail_count = total - within_tight_count
137 + above_max_count = int((rel > max_rtol).sum().item())
138 + within_tight = within_tight_count / total
139 + allowed_fail_rate = 1.0 - pass_rate
140 + above_max = above_max_count / total
141 + max_rel = rel.max().item()
142 + mean_rel = rel.mean().item()
143 + p99 = _quantile(rel, 0.99)
144 + p999 = _quantile(rel, 0.999)
145 + msg = (
146 + f"[rocm_fp8] {label}: rel<={tight_rtol} pass={within_tight:.4%} "
147 + f"({within_tight_count}/{total}) "
148 + f"fail={_format_observed_rate(fail_count, total)} "
149 + f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
150 + f"rel>{max_rtol}={_format_observed_rate(above_max_count, total)} "
151 + f"allowed_above_max={_format_allowed_rate(max_fail_rate, total)} "
152 + f"max_rel={max_rel:.6g} mean_rel={mean_rel:.6g} "
153 + f"p99={p99:.6g} p999={p999:.6g}"
154 + )
155 + print(msg)
156 + if within_tight < 1.0:
157 + warnings.warn(msg, stacklevel=2)
158 + assert within_tight >= pass_rate, msg
159 + assert above_max <= max_fail_rate, msg
160 +
161 +
162 +def _ref_paged_attention(
163 + query: torch.Tensor,
164 + key_cache: torch.Tensor,
165 + value_cache: torch.Tensor,
166 + *,
167 + query_lens: list[int],
168 + kv_lens: list[int],
169 + block_tables: torch.Tensor,
170 + scale: float,
171 +) -> torch.Tensor:
172 + block_tables_np = block_tables.cpu().numpy()
173 + _, block_size, num_kv_heads, head_size = key_cache.shape
174 +
175 + outputs = []
176 + start_idx = 0
177 + for i, query_len in enumerate(query_lens):
178 + kv_len = kv_lens[i]
179 + q = query[start_idx : start_idx + query_len] * scale
180 +
181 + num_kv_blocks = (kv_len + block_size - 1) // block_size
182 + block_indices = block_tables_np[i, :num_kv_blocks]
183 +
184 + k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
185 + v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len]
186 +
187 + if q.shape[1] != k.shape[1]:
188 + repeat = q.shape[1] // k.shape[1]
189 + k = torch.repeat_interleave(k, repeat, dim=1)
190 + v = torch.repeat_interleave(v, repeat, dim=1)
191 +
192 + attn = torch.einsum("qhd,khd->hqk", q, k).float()
193 + mask = torch.triu(
194 + torch.ones(query_len, kv_len, device=q.device),
195 + diagonal=kv_len - query_len + 1,
196 + ).bool()
197 + attn.masked_fill_(mask, float("-inf"))
198 + attn = torch.softmax(attn, dim=-1).to(v.dtype)
199 + outputs.append(torch.einsum("hqk,khd->qhd", attn, v))
200 + start_idx += query_len
201 +
202 + return torch.cat(outputs, dim=0)
203 +
204 +
205 +def test_rocm_fp8_env_defaults(monkeypatch):
206 + with monkeypatch.context() as mp:
207 + for env_name in ENV_DEFAULTS:
208 + mp.delenv(env_name, raising=False)
209 + envs = _reload_envs()
210 + for env_name, expected in ENV_DEFAULTS.items():
211 + assert getattr(envs, env_name) is expected
212 +
213 +
214 +@pytest.mark.parametrize(
215 + ("env_name", "raw_value", "expected"),
216 + [
217 + ("VLLM_ROCM_FP8_PADDING", "0", False),
218 + ("VLLM_ROCM_FP8_PADDING", "1", True),
219 + ("VLLM_ROCM_USE_AITER_FP8BMM", "0", False),
220 + ("VLLM_ROCM_USE_AITER_FP8BMM", "1", True),
221 + ("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0", False),
222 + ("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "1", True),
223 + ],
224 +)
225 +def test_rocm_fp8_env_overrides(env_name, raw_value, expected, monkeypatch):
226 + with monkeypatch.context() as mp:
227 + mp.setenv(env_name, raw_value)
228 + envs = _reload_envs()
229 + assert getattr(envs, env_name) is expected
230 +
231 +
232 +@fp8_only
233 +def test_maybe_pad_fp8_weight_respects_env(monkeypatch):
234 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (
235 + _maybe_pad_fp8_weight,
236 + )
237 +
238 + weight = torch.randn(64, 512, device=DEVICE).to(FP8_DTYPE)
239 +
240 + with monkeypatch.context() as mp:
241 + mp.setenv("VLLM_ROCM_FP8_PADDING", "1")
242 + _reload_envs()
243 + padded = _maybe_pad_fp8_weight(weight)
244 + assert padded.shape == weight.shape
245 + assert padded.dtype == weight.dtype
246 + assert padded.data_ptr() != weight.data_ptr()
247 + torch.testing.assert_close(padded.float(), weight.float())
248 +
249 + with monkeypatch.context() as mp:
250 + mp.setenv("VLLM_ROCM_FP8_PADDING", "0")
251 + _reload_envs()
252 + unpadded = _maybe_pad_fp8_weight(weight)
253 + assert unpadded.data_ptr() == weight.data_ptr()
254 +
255 +
256 +@fp8_only
257 +@pytest.mark.parametrize(
258 + ("use_aiter", "fp8bmm_enabled", "expected"),
259 + [
260 + (True, True, True),
261 + (True, False, False),
262 + (False, True, False),
263 + ],
264 +)
265 +def test_aiter_fp8bmm_enabled_api_respects_env(
266 + use_aiter, fp8bmm_enabled, expected, monkeypatch
267 +):
268 + from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
269 +
270 + assert is_aiter_found_and_supported()
271 +
272 + with monkeypatch.context() as mp:
273 + mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
274 + mp.setenv("VLLM_ROCM_USE_AITER_FP8BMM", "1" if fp8bmm_enabled else "0")
275 + _reload_envs()
276 + rocm_aiter_ops.refresh_env_variables()
277 + assert rocm_aiter_ops.is_fp8bmm_enabled() is expected
278 +
279 +
280 +@fp8_only
281 +@pytest.mark.parametrize(
282 + ("shape", "tight_atol", "max_atol", "pass_rate"), SCALED_MM_CASES
283 +)
284 +def test_fp8_scaled_mm_matches_dequantized_reference(
285 + shape, tight_atol, max_atol, pass_rate
286 +):
287 + torch.manual_seed(0)
288 +
289 + batch, m_dim, k_dim = shape
290 + n_dim = k_dim
291 + a = torch.randn(batch * m_dim, k_dim, device=DEVICE)
292 + b = torch.randn(n_dim, k_dim, device=DEVICE)
293 +
294 + a_fp8, scale_a = ref_dynamic_per_tensor_fp8_quant(a)
295 + b_fp8, scale_b = ref_dynamic_per_tensor_fp8_quant(b)
296 + out = torch._scaled_mm(
297 + a_fp8,
298 + b_fp8.t(),
299 + out_dtype=torch.bfloat16,
300 + scale_a=scale_a,
301 + scale_b=scale_b,
302 + )
303 +
304 + a_dq = a_fp8.float() * scale_a.float()
305 + b_dq = b_fp8.float() * scale_b.float()
306 + ref = torch.mm(a_dq, b_dq.t())
307 +
308 + assert out.shape == (batch * m_dim, n_dim)
309 + assert out.dtype == torch.bfloat16
310 + assert not torch.isnan(out).any()
311 +
312 + _assert_abs_error_budget(
313 + out.float(),
314 + ref,
315 + label=f"scaled_mm shape={shape}",
316 + tight_atol=tight_atol,
317 + max_atol=max_atol,
318 + pass_rate=pass_rate,
319 + )
320 +
321 +
322 +@fp8_only
323 +def test_fp8_scaled_mm_determinism():
324 + torch.manual_seed(5)
325 +
326 + m_dim, k_dim, n_dim = 64, 128, 128
327 + a = torch.randn(m_dim, k_dim, device=DEVICE)
328 + b = torch.randn(n_dim, k_dim, device=DEVICE)
329 + a_fp8, scale_a = ref_dynamic_per_tensor_fp8_quant(a)
330 + b_fp8, scale_b = ref_dynamic_per_tensor_fp8_quant(b)
331 +
332 + def run_gemm():
333 + return torch._scaled_mm(
334 + a_fp8,
335 + b_fp8.t(),
336 + out_dtype=torch.bfloat16,
337 + scale_a=scale_a,
338 + scale_b=scale_b,
339 + )
340 +
341 + _assert_deterministic(run_gemm, n_runs=4)
342 +
343 +
344 +@fp8_only
345 +def test_aiter_fp8_kv_attention_matches_bf16_reference():
346 + import aiter
347 +
348 + from vllm._aiter_ops import is_aiter_found_and_supported
349 + from vllm.v1.attention.backends.rocm_aiter_fa import cp_mha_gather_cache
350 +
351 + assert is_aiter_found_and_supported()
352 +
353 + set_random_seed(10)
354 +
355 + num_q_heads, num_kv_heads = 8, 8
356 + head_size = 128
357 + query_len, kv_len = 4, 128
358 + block_size = 16
359 + num_blocks = 2048
360 + scale = head_size**-0.5
361 +
362 + query = torch.randn(
363 + query_len, num_q_heads, head_size, dtype=torch.bfloat16, device=DEVICE
364 + )
365 + key_cache_fp8 = torch.clamp(
366 + torch.randn(num_blocks, block_size, num_kv_heads, head_size, device=DEVICE),
367 + min=-1.0,
368 + max=1.0,
369 + ).to(FP8_DTYPE)
370 + value_cache_fp8 = torch.clamp(
371 + torch.randn(num_blocks, block_size, num_kv_heads, head_size, device=DEVICE),
372 + min=-1.0,
373 + max=1.0,
374 + ).to(FP8_DTYPE)
375 +
376 + cu_query_lens = torch.tensor([0, query_len], dtype=torch.int32, device=DEVICE)
377 + cu_seq_lens = torch.tensor([0, kv_len], dtype=torch.int32, device=DEVICE)
378 + block_tables = torch.randint(
379 + 0,
380 + num_blocks,
381 + (1, (kv_len + block_size - 1) // block_size),
382 + dtype=torch.int32,
383 + device=DEVICE,
384 + )
385 + token_to_batch = torch.zeros(kv_len, dtype=torch.int32, device=DEVICE)
386 + seq_starts = torch.zeros(1, dtype=torch.int32, device=DEVICE)
387 + gathered_key = torch.empty(
388 + kv_len, num_kv_heads, head_size, dtype=torch.bfloat16, device=DEVICE
389 + )
390 + gathered_value = torch.empty_like(gathered_key)
391 +
392 + cp_mha_gather_cache(
393 + key_cache=key_cache_fp8,
394 + value_cache=value_cache_fp8,
395 + key=gathered_key,
396 + value=gathered_value,
397 + block_tables=block_tables,
398 + k_scales=torch.ones(1, dtype=torch.float32, device=DEVICE),
399 + v_scales=torch.ones(1, dtype=torch.float32, device=DEVICE),
400 + cu_seqlens_kv=cu_seq_lens,
401 + token_to_batch=token_to_batch,
402 + seq_starts=seq_starts,
403 + dequant=True,
404 + kv_cache_layout="NHD",
405 + total_tokens=kv_len,
406 + )
407 +
408 + output = torch.empty_like(query)
409 + aiter.flash_attn_varlen_func(
410 + q=query,
411 + k=gathered_key,
412 + v=gathered_value,
413 + cu_seqlens_q=cu_query_lens,
414 + cu_seqlens_k=cu_seq_lens,
415 + max_seqlen_q=query_len,
416 + max_seqlen_k=kv_len,
417 + min_seqlen_q=1,
418 + dropout_p=0.0,
419 + softmax_scale=scale,
420 + causal=True,
421 + window_size=(-1, -1),
422 + alibi_slopes=None,
423 + return_lse=False,
424 + out=output,
425 + )
426 +
427 + ref = _ref_paged_attention(
428 + query=query,
429 + key_cache=key_cache_fp8.to(torch.bfloat16),
430 + value_cache=value_cache_fp8.to(torch.bfloat16),
431 + query_lens=[query_len],
432 + kv_lens=[kv_len],
433 + block_tables=block_tables,
434 + scale=scale,
435 + )
436 +
437 + _assert_abs_error_budget(
438 + output.float(),
439 + ref.float(),
440 + label="aiter_fp8_kv_attention",
441 + tight_atol=0.15,
442 + max_atol=0.25,
443 + pass_rate=0.98,
444 + )
445 +
446 +
447 +@fp8_only
448 +@pytest.mark.parametrize("shape", [(128, 256), (512, 1024), (1024, 4096)])
449 +def test_fp8_per_tensor_quant_matches_reference(shape):
450 + torch.manual_seed(0)
451 +
452 + x = torch.randn(*shape, dtype=torch.bfloat16, device=DEVICE)
453 +
454 + out, scale = ops.scaled_fp8_quant(x)
455 + ref_out, ref_scale = ref_dynamic_per_tensor_fp8_quant(x)
456 +
457 + assert out.shape == x.shape
458 + assert out.dtype == FP8_DTYPE
459 +
460 + torch.testing.assert_close(scale, ref_scale, rtol=1e-3, atol=1e-4)
461 + code_match_rate = (
462 + (out.view(torch.uint8) == ref_out.view(torch.uint8)).float().mean().item()
463 + )
464 + assert code_match_rate >= 0.999, (
465 + f"FP8 codes match reference only {code_match_rate:.4%} (< 99.9%)"
466 + )
467 +
468 + dequant = out.float() * scale.float()
469 + _assert_rel_error_budget(
470 + dequant,
471 + x.float(),
472 + label=f"per_tensor_quant shape={shape}",
473 + tight_rtol=0.06,
474 + max_rtol=0.60,
475 + pass_rate=0.99,
476 + max_fail_rate=1e-4,
477 + )
478 +
479 +
480 +@fp8_only
481 +def test_aiter_group_fp8_quant_roundtrip():
482 + from aiter import dtypes
483 +
484 + from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
485 +
486 + assert is_aiter_found_and_supported()
487 +
488 + torch.manual_seed(0)
489 +
490 + m_dim, n_dim = 64, 4096
491 + group_size = 128
492 + x = torch.randn(m_dim, n_dim, dtype=torch.bfloat16, device=DEVICE)
493 +
494 + out, scales = rocm_aiter_ops.group_fp8_quant(x, group_size)
495 + assert out.dtype == dtypes.fp8
496 + assert scales.dtype == torch.float32
497 +
498 + scales_expanded = scales.repeat_interleave(group_size, dim=1)[:, :n_dim]
499 + dequant = out.float() * scales_expanded
500 + _assert_rel_error_budget(
501 + dequant,
502 + x.float(),
503 + label="aiter_group_fp8_quant",
504 + tight_rtol=0.06,
505 + max_rtol=0.30,
506 + pass_rate=0.99,
507 + max_fail_rate=1e-4,
508 + )
509 +
510 +
511 +@fp8_only
512 +def test_fp8_per_token_quant_matches_reference():
513 + torch.manual_seed(0)
514 +
515 + x = torch.randn(32, 4096, dtype=torch.bfloat16, device=DEVICE)
516 +
517 + out, scale = ops.scaled_fp8_quant(x, scale=None, use_per_token_if_dynamic=True)
518 + ref_out, ref_scale = ref_dynamic_per_token_quant(x, FP8_DTYPE)
519 +
520 + assert out.shape == x.shape
521 + assert out.dtype == FP8_DTYPE
522 + assert scale.shape == (x.shape[0], 1)
523 +
524 + torch.testing.assert_close(scale, ref_scale, rtol=1e-3, atol=1e-4)
525 + code_match_rate = (
526 + (out.view(torch.uint8) == ref_out.view(torch.uint8)).float().mean().item()
527 + )
528 + assert code_match_rate >= 0.999, (
529 + f"FP8 codes match reference only {code_match_rate:.4%} (< 99.9%)"
530 + )
531 +
532 + dequant = out.float() * scale.float()
533 + _assert_rel_error_budget(
534 + dequant,
535 + x.float(),
536 + label="per_token_quant",
537 + tight_rtol=0.06,
538 + max_rtol=0.30,
539 + pass_rate=0.99,
540 + max_fail_rate=1e-4,
541 + )
tests/kernels/quantization/test_rocm_mxfp4.py new
+638
@@ -0,0 +1,638 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ROCm AITER FP4 and MXFP4 tests.
4 +
5 +This file keeps the ROCm-specific FP4 coverage:
6 +- ROCm AITER FP4 env and enablement gates
7 +- large-shape MXFP4 wrapper checks through vLLM's public helper op
8 +- AITER Triton MXFP4 quant format and determinism
9 +- FP4 GEMM, preshuffled-scale, hardware-quant, and skinny decode
10 + paths
11 +
12 +Generic MXFP4 wrapper/reference checks already live in
13 +``tests/kernels/quantization/test_quark.py``.
14 +"""
15 +
16 +import importlib
17 +import warnings
18 +
19 +import pytest
20 +import torch
21 +
22 +from tests.kernels.utils import _assert_accurate, _assert_deterministic
23 +from tests.quantization.reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
24 +from vllm.platforms import current_platform
25 +from vllm.platforms.rocm import on_gfx950
26 +
27 +pytestmark = pytest.mark.skipif(
28 + not current_platform.is_rocm(), reason="ROCm-specific tests"
29 +)
30 +
31 +LLAMA_SHAPES = [
32 + (64, 4096),
33 + (64, 11008),
34 + (32, 14336),
35 +]
36 +GEMM_ATOL = 0.5
37 +GEMM_RTOL = 0.0
38 +GEMM_MAX_VIOLATION_FACTOR = 2.0
39 +A4W4_GEMM_PASS_RATES = {
40 + (64, 128, 64): 0.9999,
41 + (128, 256, 128): 0.9999,
42 + (128, 4096, 4096): 0.9999,
43 + (256, 4096, 11008): 0.9999,
44 + (64, 8192, 28672): 0.99765,
45 +}
46 +SKINNY_GEMM_PASS_RATES = {
47 + (4096, 4096): 0.9999,
48 + (4096, 11008): 0.99235,
49 + (8192, 8192): 0.99765,
50 +}
51 +PRESHUFFLED_SHAPES = [
52 + (64, 4096, 8192),
53 + (32, 8192, 8192),
54 +]
55 +
56 +
57 +def _reload_envs():
58 + import vllm.envs as envs
59 +
60 + return importlib.reload(envs)
61 +
62 +
63 +@pytest.fixture(autouse=True)
64 +def _restore_rocm_env_state():
65 + """Restore global env + AITER flag state after every test."""
66 + yield
67 + _reload_envs()
68 + from vllm._aiter_ops import rocm_aiter_ops
69 +
70 + rocm_aiter_ops.refresh_env_variables()
71 +
72 +
73 +def _assert_aiter_supported() -> None:
74 + from vllm._aiter_ops import is_aiter_found_and_supported
75 +
76 + assert is_aiter_found_and_supported(), (
77 + "aiter is required on supported ROCm hardware for this test"
78 + )
79 +
80 +
81 +def _format_observed_rate(count: int, total: int) -> str:
82 + return f"{count / total:.4%} ({count}/{total})"
83 +
84 +
85 +def _format_allowed_rate(rate: float, total: int) -> str:
86 + allowed_count = int(rate * total)
87 + return f"{rate:.4%} (<= {allowed_count}/{total})"
88 +
89 +
90 +def _print_close_stats(
91 + label: str,
92 + actual: torch.Tensor,
93 + expected: torch.Tensor,
94 + *,
95 + atol: float,
96 + rtol: float,
97 + pass_rate: float = 1.0,
98 + max_atol: float | None = None,
99 +) -> None:
100 + abs_diff = (actual - expected).abs().float().flatten()
101 + expected_abs = expected.abs().float().flatten()
102 + allowed = atol + rtol * expected_abs
103 + within = abs_diff <= allowed
104 +
105 + total = abs_diff.numel()
106 + passed = int(within.sum().item())
107 + failed = total - passed
108 + allowed_fail_rate = 1.0 - pass_rate
109 +
110 + max_abs = abs_diff.max().item()
111 + mean_abs = abs_diff.mean().item()
112 + p99_abs = torch.quantile(abs_diff, 0.99).item()
113 + p999_abs = torch.quantile(abs_diff, 0.999).item()
114 + worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item()
115 +
116 + msg = (
117 + "[rocm_aiter_fp4] "
118 + f"{label}: "
119 + f"pass={passed / total:.4%} ({passed}/{total}) "
120 + f"fail={_format_observed_rate(failed, total)} "
121 + f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
122 + f"atol={atol:g} "
123 + f"rtol={rtol:g} "
124 + )
125 + if max_atol is not None:
126 + above_max_count = int((abs_diff > max_atol).sum().item())
127 + msg += (
128 + f"abs>{max_atol:g}={_format_observed_rate(above_max_count, total)} "
129 + f"allowed_above_max={_format_allowed_rate(0.0, total)} "
130 + )
131 + msg += (
132 + f"max_abs={max_abs:.6g} "
133 + f"mean_abs={mean_abs:.6g} "
134 + f"p99_abs={p99_abs:.6g} "
135 + f"p999_abs={p999_abs:.6g} "
136 + f"worst_ratio={worst_ratio:.6g}"
137 + )
138 + print(msg)
139 + if failed > 0:
140 + warnings.warn(msg, stacklevel=2)
141 +
142 +
143 +# Env and enablement tests ------------------------------------------------
144 +
145 +
146 +def test_fp4_env_defaults():
147 + """ROCm FP4 env defaults should stay stable for the AITER gates."""
148 + import vllm.envs as envs
149 +
150 + assert envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM is False
151 + assert envs.VLLM_ROCM_USE_AITER_FP4BMM is True
152 +
153 +
154 +@pytest.mark.parametrize(
155 + (
156 + "use_aiter",
157 + "use_fp4_asm_gemm",
158 + "use_fp4bmm",
159 + ),
160 + [
161 + (True, True, True),
162 + (True, True, False),
163 + (True, False, True),
164 + (False, True, True),
165 + ],
166 +)
167 +def test_rocm_aiter_fp4_enablement_follows_env_and_arch(
168 + use_aiter,
169 + use_fp4_asm_gemm,
170 + use_fp4bmm,
171 + monkeypatch,
172 +):
173 + """The ROCm FP4 AITER gates should depend only on the env toggles and the
174 + gfx950 hardware check.
175 +
176 + Expected values are derived from the *actual* platform rather than a patched
177 + arch (gfx950 -> MI355, gfx942 -> MI300). A failure in this group therefore
178 + points straight at the hardware it ran on instead of a mocked value.
179 + """
180 + from vllm._aiter_ops import rocm_aiter_ops
181 +
182 + _assert_aiter_supported()
183 +
184 + on_gfx950_value = on_gfx950()
185 + expected_asm_gemm = use_aiter and use_fp4_asm_gemm and on_gfx950_value
186 + expected_fp4bmm = use_aiter and use_fp4bmm and on_gfx950_value
187 +
188 + with monkeypatch.context() as mp:
189 + mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
190 + mp.setenv(
191 + "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
192 + "1" if use_fp4_asm_gemm else "0",
193 + )
194 + mp.setenv("VLLM_ROCM_USE_AITER_FP4BMM", "1" if use_fp4bmm else "0")
195 + _reload_envs()
196 + rocm_aiter_ops.refresh_env_variables()
197 +
198 + assert (
199 + rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() is expected_asm_gemm
200 + )
201 + assert rocm_aiter_ops.is_fp4bmm_enabled() is expected_fp4bmm
202 +
203 +
204 +# Large-shape MXFP4 wrapper tests -----------------------------------------
205 +
206 +
207 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
208 +@pytest.mark.parametrize("shape", LLAMA_SHAPES)
209 +def test_vllm_quant_dequant_mxfp4_matches_reference_on_large_shapes(shape, dtype):
210 + """The public vLLM MXFP4 QDQ helper should stay exact against the torch
211 + reference even on larger Llama-like shapes."""
212 + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
213 + quant_dequant_mxfp4,
214 + )
215 +
216 + torch.set_default_device("cuda")
217 + torch.manual_seed(0)
218 + rows, cols = shape
219 + x = torch.randn(rows, cols, dtype=dtype)
220 + out = quant_dequant_mxfp4(x)
221 + ref = qdq_mxfp4_torch(x, "even")
222 +
223 + assert out.shape == x.shape
224 + assert out.dtype == dtype
225 + assert torch.isfinite(out).all()
226 + torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0)
227 +
228 +
229 +def test_vllm_quant_dequant_mxfp4_is_deterministic():
230 + """The public vLLM MXFP4 QDQ helper should stay bitwise deterministic."""
231 + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
232 + quant_dequant_mxfp4,
233 + )
234 +
235 + torch.set_default_device("cuda")
236 + torch.manual_seed(3)
237 +
238 + x = torch.randn(64, 128, dtype=torch.bfloat16)
239 +
240 + _assert_deterministic(quant_dequant_mxfp4, x, n_runs=4)
241 +
242 +
243 +# Triton MXFP4 quant tests ------------------------------------------------
244 +
245 +
246 +def test_aiter_dynamic_mxfp4_quant_output_format():
247 + """dynamic_mxfp4_quant returns packed uint8 FP4 values and E8M0 uint8 scales.
248 +
249 + OCP MXFP4: block_size=32, 2 FP4 E2M1 values packed per byte.
250 + Scale shape: (M, K // 32); one E8M0 exponent byte per 32-element block.
251 + """
252 + _assert_aiter_supported()
253 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
254 +
255 + torch.set_default_device("cuda")
256 + M, K = 64, 256
257 + x = torch.randn(M, K, dtype=torch.bfloat16)
258 +
259 + x_fp4, x_scale = dynamic_mxfp4_quant(x)
260 +
261 + # FP4 values packed 2-per-byte; shape (M, K // 2)
262 + assert x_fp4.shape == (M, K // 2), (
263 + f"Expected fp4 shape ({M}, {K // 2}), got {x_fp4.shape}"
264 + )
265 + assert x_fp4.dtype == torch.uint8
266 +
267 + # One E8M0 scale byte per 32-element block; shape (M, K // 32)
268 + assert x_scale.shape == (M, K // 32), (
269 + f"Expected scale shape ({M}, {K // 32}), got {x_scale.shape}"
270 + )
271 + assert x_scale.dtype == torch.uint8
272 +
273 +
274 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
275 +@pytest.mark.parametrize(
276 + "shape",
277 + [
278 + (128, 4096), # Llama-7B hidden
279 + (256, 11008), # Llama-7B FFN
280 + (64, 14336), # Llama-70B FFN
281 + (512, 8192), # Llama-70B hidden
282 + (32, 28672), # DeepSeek-style large FFN
283 + ],
284 +)
285 +def test_aiter_dynamic_mxfp4_quant_llama_shapes(shape, dtype):
286 + """MXFP4 quantization output format is correct for Llama-class weight shapes.
287 +
288 + Targets B200 parity: NVIDIA nvfp4 tests parametrize over Llama shapes
289 + (7168, 14336, 28672...). This covers gfx950 MXFP4 quant format validation.
290 + """
291 + _assert_aiter_supported()
292 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
293 +
294 + torch.set_default_device("cuda")
295 + M, K = shape
296 + x = torch.randn(M, K, dtype=dtype)
297 +
298 + x_fp4, x_scale = dynamic_mxfp4_quant(x)
299 +
300 + assert x_fp4.shape == (M, K // 2)
301 + assert x_fp4.dtype == torch.uint8
302 + assert x_scale.shape == (M, K // 32)
303 + assert x_scale.dtype == torch.uint8
304 + assert not torch.any(torch.isnan(x_fp4.float()))
305 +
306 +
307 +def test_aiter_dynamic_mxfp4_quant_determinism():
308 + """dynamic_mxfp4_quant is bitwise deterministic across repeated runs."""
309 + _assert_aiter_supported()
310 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
311 +
312 + torch.set_default_device("cuda")
313 + torch.manual_seed(7)
314 + x = torch.randn(128, 256, dtype=torch.bfloat16)
315 +
316 + _assert_deterministic(dynamic_mxfp4_quant, x, n_runs=4)
317 +
318 +
319 +# gfx950 hardware FP4 GEMM tests ------------------------------------------
320 +
321 +
322 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
323 +@pytest.mark.parametrize(
324 + "shape",
325 + [
326 + (64, 128, 64), # small square-ish
327 + (128, 256, 128), # medium
328 + (128, 4096, 4096), # Llama-7B hidden square
329 + (256, 4096, 11008), # Llama-7B FFN
330 + (64, 8192, 28672), # Llama-70B FFN
331 + ],
332 +)
333 +def test_aiter_fp4_gemm_a4w4_accuracy(shape):
334 + """AITER A4W4 FP4 GEMM output is close to matmul on dequantized weights.
335 +
336 + Tests gfx950 parity with B200 test_nvfp4_scaled_mm.py: E2M1 x E2M1 GEMM
337 + with block-scaled quantization, multiple Llama-class shapes.
338 + Requires gfx950 hardware; gracefully skips elsewhere.
339 + """
340 + atol = GEMM_ATOL
341 + rtol = GEMM_RTOL
342 + pass_rate = A4W4_GEMM_PASS_RATES[shape]
343 + _assert_aiter_supported()
344 + from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
345 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
346 +
347 + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
348 + quant_dequant_mxfp4,
349 + )
350 +
351 + torch.set_default_device("cuda")
352 + torch.manual_seed(0)
353 + M, K, N = shape
354 +
355 + A = torch.randn(M, K, dtype=torch.bfloat16)
356 + B = torch.randn(N, K, dtype=torch.bfloat16) # weight stored as (N, K)
357 +
358 + # Quantize
359 + A_fp4, A_scale = dynamic_mxfp4_quant(A)
360 + B_fp4, B_scale = dynamic_mxfp4_quant(B)
361 +
362 + # FP4 GEMM computes A_fp4 (M, K) @ B_fp4.T (K, N) into (M, N).
363 + out = gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale)
364 +
365 + assert out.shape == (M, N), f"Expected ({M}, {N}), got {out.shape}"
366 + assert out.dtype == torch.bfloat16
367 + assert not torch.any(torch.isnan(out))
368 +
369 + # Reference: matmul on dequantized FP4 inputs (quark roundtrip)
370 + A_dq = quant_dequant_mxfp4(A)
371 + B_dq = quant_dequant_mxfp4(B)
372 + ref = torch.matmul(A_dq.float(), B_dq.t().float())
373 +
374 + # FP4 has only 3 mantissa bits; large K amplifies the accumulation error.
375 + # These budgets were measured on gfx950 and keep p99/p999 drift visible
376 + # without pretending the kernel is close to BF16 matmul everywhere.
377 + _print_close_stats(
378 + f"a4w4_gemm shape={shape}",
379 + out.float(),
380 + ref,
381 + atol=atol,
382 + rtol=rtol,
383 + pass_rate=pass_rate,
384 + max_atol=GEMM_MAX_VIOLATION_FACTOR * atol,
385 + )
386 + _assert_accurate(
387 + out.float(),
388 + ref,
389 + atol=atol,
390 + rtol=rtol,
391 + pass_rate=pass_rate,
392 + max_violation_factor=GEMM_MAX_VIOLATION_FACTOR,
393 + )
394 +
395 +
396 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
397 +@pytest.mark.parametrize(
398 + "shape",
399 + PRESHUFFLED_SHAPES,
400 +)
401 +def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape):
402 + """The preshuffled FP4 GEMM path should run on tuned shapes with the same
403 + packed inputs vLLM uses in production.
404 +
405 + This kernel only advertises a narrow tuned `(N, K)` set on gfx950. The
406 + important contract here is that those tuned shapes execute with the
407 + production-style shuffled weights and stay deterministic.
408 + """
409 + _assert_aiter_supported()
410 + from aiter import per_1x32_f4_quant_hip
411 + from aiter.ops.shuffle import shuffle_weight
412 + from aiter.ops.triton.gemm_afp4wfp4 import (
413 + gemm_afp4wfp4_preshuffled_weight_scales,
414 + )
415 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
416 +
417 + from vllm._aiter_ops import rocm_aiter_ops
418 +
419 + torch.set_default_device("cuda")
420 + torch.manual_seed(1)
421 + M, K, N = shape
422 +
423 + assert M <= 64
424 + assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K)
425 +
426 + A = torch.randn(M, K, dtype=torch.bfloat16)
427 + B = torch.randn(N, K, dtype=torch.bfloat16)
428 +
429 + B_fp4, B_scale = dynamic_mxfp4_quant(B)
430 + A_q, A_s = per_1x32_f4_quant_hip(A, shuffle=M >= 32)
431 +
432 + scale_rows, scale_cols = B_scale.shape
433 + B_scale = (
434 + B_scale.view(scale_rows // 32, 2, 16, scale_cols // 8, 2, 4, 1)
435 + .permute(0, 3, 5, 2, 4, 1, 6)
436 + .contiguous()
437 + .view(scale_rows, scale_cols)
438 + )
439 + B_fp4 = shuffle_weight(B_fp4, layout=(16, 16))
440 +
441 + if M >= 32:
442 + A_s = A_s.contiguous().view(torch.uint8).reshape(A_s.shape[0] // 32, -1)
443 + else:
444 + A_s = A_s[:M, ...].contiguous().view(torch.uint8)
445 + B_scale = B_scale.contiguous().view(torch.uint8).reshape(B_scale.shape[0] // 32, -1)
446 +
447 + def run_preshuffled() -> torch.Tensor:
448 + y = torch.empty(M, N, device="cuda", dtype=torch.bfloat16)
449 + return gemm_afp4wfp4_preshuffled_weight_scales(
450 + A_q.contiguous().view(torch.uint8),
451 + B_fp4.contiguous().view(torch.uint8).reshape(B_fp4.shape[0] // 16, -1),
452 + A_s,
453 + B_scale,
454 + torch.bfloat16,
455 + y,
456 + )
457 +
458 + out = run_preshuffled()
459 + assert out.shape == (M, N)
460 + assert out.dtype == torch.bfloat16
461 + assert torch.isfinite(out).all()
462 +
463 + _assert_deterministic(run_preshuffled, n_runs=3)
464 +
465 +
466 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
467 +def test_aiter_fp4_gemm_a4w4_determinism():
468 + """AITER FP4 A4W4 GEMM is bitwise deterministic across repeated runs.
469 +
470 + Targets parity with B200 FP4 determinism requirements.
471 + Requires gfx950.
472 + """
473 + _assert_aiter_supported()
474 + from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
475 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
476 +
477 + torch.set_default_device("cuda")
478 + torch.manual_seed(2)
479 +
480 + M, K, N = 128, 256, 128
481 + A = torch.randn(M, K, dtype=torch.bfloat16)
482 + B = torch.randn(N, K, dtype=torch.bfloat16)
483 +
484 + A_fp4, A_scale = dynamic_mxfp4_quant(A)
485 + B_fp4, B_scale = dynamic_mxfp4_quant(B)
486 +
487 + def run_gemm():
488 + return gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale)
489 +
490 + _assert_deterministic(run_gemm, n_runs=4)
491 +
492 +
493 +# gfx950 hardware FP4 dynamic quantization --------------------------------
494 +
495 +
496 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
497 +@pytest.mark.parametrize(
498 + "shape",
499 + [
500 + (128, 256),
501 + (256, 4096),
502 + (64, 14336),
503 + ],
504 +)
505 +def test_aiter_hardware_fp4_dynamic_quant_format(shape):
506 + """aiter hardware FP4 dynamic quant produces correct output format.
507 +
508 + Tests gfx950 hardware-accelerated FP4 quantization (OCP MXFP4 E2M1).
509 + Parity with B200 scaled_fp4_quant: block_size=32, packed uint8 output.
510 + Requires gfx950.
511 + """
512 + _assert_aiter_supported()
513 +
514 + from aiter import dynamic_per_group_scaled_quant_fp4
515 +
516 + torch.set_default_device("cuda")
517 + M, K = shape
518 + group_size = 32
519 +
520 + x = torch.randn(M, K, dtype=torch.bfloat16)
521 + out_fp4 = torch.empty(M, K // 2, dtype=torch.uint8)
522 + scales = torch.empty(M, K // group_size, dtype=torch.uint8)
523 +
524 + dynamic_per_group_scaled_quant_fp4(out_fp4, x, scales, group_size)
525 +
526 + assert out_fp4.shape == (M, K // 2), (
527 + f"Shape {shape}: expected fp4 ({M}, {K // 2}), got {out_fp4.shape}"
528 + )
529 + assert scales.shape == (M, K // group_size), (
530 + f"Shape {shape}: expected scale ({M}, {K // group_size}), got {scales.shape}"
531 + )
532 +
533 +
534 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
535 +def test_aiter_hardware_fp4_quant_vs_triton():
536 + """Hardware FP4 quant should dequantize back to the same MXFP4 values as
537 + the Triton path.
538 +
539 + The packed bytes are not identical on gfx950, but with
540 + ``shuffle_scale=False`` the dequantized results match tightly.
541 + """
542 + _assert_aiter_supported()
543 +
544 + from aiter import dynamic_per_group_scaled_quant_fp4
545 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
546 +
547 + torch.set_default_device("cuda")
548 + torch.manual_seed(5)
549 +
550 + M, K = 128, 256
551 + group_size = 32
552 + x = torch.randn(M, K, dtype=torch.bfloat16)
553 +
554 + # Hardware path
555 + out_hw = torch.empty(M, K // 2, dtype=torch.uint8)
556 + scales_hw = torch.empty(M, K // group_size, dtype=torch.uint8)
557 + dynamic_per_group_scaled_quant_fp4(
558 + out_hw, x, scales_hw, group_size, shuffle_scale=False
559 + )
560 +
561 + # Triton path
562 + out_triton, scales_triton = dynamic_mxfp4_quant(x)
563 +
564 + assert out_hw.shape == out_triton.shape
565 + assert scales_hw.shape == scales_triton.shape
566 +
567 + dq_hw = dq_mxfp4_torch(out_hw, scales_hw, torch.bfloat16)
568 + dq_triton = dq_mxfp4_torch(out_triton, scales_triton, torch.bfloat16)
569 +
570 + _print_close_stats(
571 + "hardware_quant_vs_triton",
572 + dq_hw.float(),
573 + dq_triton.float(),
574 + atol=0.25,
575 + rtol=0.0,
576 + )
577 + torch.testing.assert_close(dq_hw, dq_triton, atol=0.25, rtol=0.0)
578 +
579 +
580 +# ROCm skinny GEMM FP4 tests ----------------------------------------------
581 +
582 +
583 +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
584 +@pytest.mark.parametrize("M", [1, 2, 4, 8]) # decode / skinny batch sizes
585 +@pytest.mark.parametrize("N, K", [(4096, 4096), (4096, 11008), (8192, 8192)])
586 +def test_aiter_fp4_gemm_skinny_shapes(M, N, K):
587 + """FP4 GEMM accuracy for skinny (small-M) shapes (decode phase).
588 +
589 + Skinny GEMMs (M=1..8) are the bottleneck in decode phase.
590 + Tests parity with B200 nvfp4_scaled_mm tests at decode batch sizes.
591 + Requires gfx950.
592 + """
593 + atol = GEMM_ATOL
594 + rtol = GEMM_RTOL
595 + pass_rate = SKINNY_GEMM_PASS_RATES[(N, K)]
596 + _assert_aiter_supported()
597 + from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
598 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
599 +
600 + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
601 + quant_dequant_mxfp4,
602 + )
603 +
604 + torch.set_default_device("cuda")
605 + torch.manual_seed(0)
606 +
607 + A = torch.randn(M, K, dtype=torch.bfloat16)
608 + B = torch.randn(N, K, dtype=torch.bfloat16)
609 +
610 + A_fp4, A_scale = dynamic_mxfp4_quant(A)
611 + B_fp4, B_scale = dynamic_mxfp4_quant(B)
612 +
613 + out = gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale)
614 +
615 + assert out.shape == (M, N)
616 + assert not torch.any(torch.isnan(out))
617 +
618 + A_dq = quant_dequant_mxfp4(A)
619 + B_dq = quant_dequant_mxfp4(B)
620 + ref = torch.matmul(A_dq.float(), B_dq.t().float())
621 +
622 + _print_close_stats(
623 + f"skinny_gemm M={M} N={N} K={K}",
624 + out.float(),
625 + ref,
626 + atol=atol,
627 + rtol=rtol,
628 + pass_rate=pass_rate,
629 + max_atol=GEMM_MAX_VIOLATION_FACTOR * atol,
630 + )
631 + _assert_accurate(
632 + out.float(),
633 + ref,
634 + atol=atol,
635 + rtol=rtol,
636 + pass_rate=pass_rate,
637 + max_violation_factor=GEMM_MAX_VIOLATION_FACTOR,
638 + )
tests/kernels/utils.py
+57
@@ -36,6 +36,63 @@ ALL_OPCHECK_TEST_UTILS: tuple[str, ...] = (
36 )
37
38
39 +def _assert_accurate(
40 + actual: torch.Tensor,
41 + expected: torch.Tensor,
42 + atol: float,
43 + rtol: float = 0.0,
44 + pass_rate: float = 0.99999,
45 + max_violation_factor: float = 3.0,
46 +) -> None:
47 + """Check numeric accuracy with pass-rate, max-error, and mean-error bounds."""
48 + a = actual.detach().float().flatten()
49 + e = expected.detach().float().flatten()
50 +
51 + abs_err = (a - e).abs()
52 + tol = atol + rtol * e.abs()
53 +
54 + rate = (abs_err <= tol).float().mean().item()
55 + assert rate >= pass_rate, (
56 + f"Accuracy pass rate {rate:.6f} < {pass_rate} (atol={atol}, rtol={rtol})"
57 + )
58 +
59 + max_err = abs_err.max().item()
60 + assert max_err <= max_violation_factor * atol, (
61 + f"Max absolute error {max_err:.6f} exceeds {max_violation_factor} * atol={atol}"
62 + )
63 +
64 + mean_err = abs_err.mean().item()
65 + assert mean_err <= atol * 0.25, (
66 + f"Mean absolute error {mean_err:.6f} >= atol * 0.25 = {atol * 0.25:.6f}"
67 + )
68 +
69 +
70 +def _assert_deterministic(
71 + fn,
72 + *args,
73 + n_runs: int = 4,
74 + **kwargs,
75 +) -> None:
76 + """Verify that repeated calls produce bitwise-identical tensor outputs."""
77 +
78 + def _collect(result: Any) -> list[torch.Tensor]:
79 + if isinstance(result, torch.Tensor):
80 + return [result.detach().clone()]
81 + if isinstance(result, (tuple, list)):
82 + return [t.detach().clone() for t in result if isinstance(t, torch.Tensor)]
83 + raise TypeError(f"Unexpected return type {type(result)}")
84 +
85 + reference = _collect(fn(*args, **kwargs))
86 +
87 + for run in range(1, n_runs):
88 + outputs = _collect(fn(*args, **kwargs))
89 + for idx, (ref, out) in enumerate(zip(reference, outputs)):
90 + assert torch.equal(ref, out), (
91 + f"Run {run}: output[{idx}] differs from run 0 "
92 + f"(max diff = {(out.float() - ref.float()).abs().max().item():.2e})"
93 + )
94 +
95 +
96 class QKVInputs(NamedTuple):
97 """
98 Data structure for representing unpacked attention inputs,