test_rocm_aiter_moe.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,568 lines 47.9 KB

Raw

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	from unittest import mock
26
27	import pytest
28	import torch
29	import torch.nn.functional as F
30
31	from tests.kernels.utils import _assert_deterministic
32	from vllm.platforms import current_platform
33	from vllm.platforms.rocm import on_cdna, on_gfx942, on_gfx950
34	from vllm.utils.torch_utils import set_random_seed
35
36	pytestmark = pytest.mark.skipif(
37	not current_platform.is_rocm(), reason="ROCm-specific tests"
38	)
39
40
41	# Helpers -----------------------------------------------------------------
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	yield
54	_reload_envs()
55	from vllm._aiter_ops import rocm_aiter_ops
56
57	rocm_aiter_ops.refresh_env_variables()
58
59
60	def _assert_aiter_supported() -> None:
61	from vllm._aiter_ops import is_aiter_found_and_supported
62
63	assert is_aiter_found_and_supported(), (
64	"aiter is required on supported ROCm hardware for this test"
65	)
66
67
68	def _format_observed_rate(count: int, total: int) -> str:
69	return f"{count / total:.4%} ({count}/{total})"
70
71
72	def _format_allowed_rate(rate: float, total: int) -> str:
73	allowed_count = int(rate * total)
74	return f"{rate:.4%} (<= {allowed_count}/{total})"
75
76
77	def _quantile(values: torch.Tensor, q: float) -> float:
78	if values.numel() == 0:
79	return 0.0
80	return torch.quantile(values, q).item()
81
82
83	def _assert_close_budget(
84	actual: torch.Tensor,
85	expected: torch.Tensor,
86	*,
87	label: str,
88	atol: float,
89	rtol: float = 0.0,
90	pass_rate: float = 0.99999,
91	max_violation_factor: float = 3.0,
92	) -> None:
93	actual_f = actual.detach().float().flatten()
94	expected_f = expected.detach().float().flatten()
95	abs_diff = (actual_f - expected_f).abs()
96	allowed = atol + rtol * expected_f.abs()
97
98	total = abs_diff.numel()
99	within = abs_diff <= allowed
100	passed = int(within.sum().item())
101	failed = total - passed
102	allowed_fail_rate = 1.0 - pass_rate
103
104	max_abs = abs_diff.max().item()
105	mean_abs = abs_diff.mean().item()
106	p99_abs = _quantile(abs_diff, 0.99)
107	p999_abs = _quantile(abs_diff, 0.999)
108	worst_ratio = (abs_diff / allowed.clamp_min(1e-12)).max().item()
109	max_atol = max_violation_factor * atol
110	above_max_count = int((abs_diff > max_atol).sum().item())
111
112	msg = (
113	"[rocm_aiter_moe] "
114	f"{label}: "
115	f"pass={passed / total:.4%} ({passed}/{total}) "
116	f"fail={_format_observed_rate(failed, total)} "
117	f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
118	f"atol={atol:g} "
119	f"rtol={rtol:g} "
120	f"abs>{max_atol:g}={_format_observed_rate(above_max_count, total)} "
121	f"allowed_above_max={_format_allowed_rate(0.0, total)} "
122	f"max_abs={max_abs:.6g} "
123	f"mean_abs={mean_abs:.6g} "
124	f"p99_abs={p99_abs:.6g} "
125	f"p999_abs={p999_abs:.6g} "
126	f"worst_ratio={worst_ratio:.6g}"
127	)
128	print(msg)
129	if failed > 0:
130	warnings.warn(msg, stacklevel=2)
131
132	assert passed / total >= pass_rate, msg
133	assert max_abs <= max_atol, msg
134	assert mean_abs <= atol * 0.25, msg
135
136
137	def _assert_group_quant_quality(
138	actual: torch.Tensor,
139	expected: torch.Tensor,
140	*,
141	label: str,
142	preferred_rel: float = 0.06,
143	pass_rate: float = 0.9999,
144	max_rel: float = 0.55,
145	max_fail_rate: float = 0.000005,
146	mean_limit: float = 0.03,
147	) -> None:
148	rel = (
149	(actual.float() - expected.float()).abs()
150	/ expected.float().abs().clamp_min(1e-5)
151	).flatten()
152	total = rel.numel()
153	within_preferred_count = int((rel <= preferred_rel).sum().item())
154	fail_count = total - within_preferred_count
155	allowed_fail_rate = 1.0 - pass_rate
156	above_max_count = int((rel > max_rel).sum().item())
157	mean_rel = rel.mean().item()
158	max_rel_err = rel.max().item()
159	p99 = _quantile(rel, 0.99)
160	p999 = _quantile(rel, 0.999)
161
162	msg = (
163	"[rocm_aiter_moe] "
164	f"{label}: "
165	f"rel<={preferred_rel:g} pass={within_preferred_count / total:.4%} "
166	f"({within_preferred_count}/{total}) "
167	f"fail={_format_observed_rate(fail_count, total)} "
168	f"allowed_fail={_format_allowed_rate(allowed_fail_rate, total)} "
169	f"rel>{max_rel:g}={_format_observed_rate(above_max_count, total)} "
170	f"allowed_above_max={_format_allowed_rate(max_fail_rate, total)} "
171	f"mean_rel={mean_rel:.6g} "
172	f"max_rel={max_rel_err:.6g} "
173	f"p99={p99:.6g} "
174	f"p999={p999:.6g}"
175	)
176	print(msg)
177	if fail_count > 0:
178	warnings.warn(msg, stacklevel=2)
179
180	assert within_preferred_count / total >= pass_rate, msg
181	assert above_max_count / total <= max_fail_rate, msg
182	assert mean_rel < mean_limit, msg
183
184
185	def ref_moe_forward(
186	hidden_states: torch.Tensor,
187	w1: torch.Tensor,
188	w2: torch.Tensor,
189	topk_weights: torch.Tensor,
190	topk_ids: torch.Tensor,
191	activation: str = "silu",
192	) -> torch.Tensor:
193	"""Float32 mask-based MoE reference for the ROCm fused-MoE kernel."""
194	num_tokens, hidden_dim = hidden_states.shape
195	num_experts = w1.shape[0]
196	intermediate_dim = w1.shape[1] // 2
197	topk = topk_ids.shape[1]
198	device = hidden_states.device
199
200	hidden_states_f = hidden_states.float()
201	w1_f = w1.float()
202	w2_f = w2.float()
203
204	expanded = hidden_states_f.view(num_tokens, 1, hidden_dim).expand(
205	num_tokens, topk, hidden_dim
206	)
207	expanded = expanded.reshape(num_tokens * topk, hidden_dim)
208	output = torch.zeros(
209	num_tokens * topk, hidden_dim, dtype=torch.float32, device=device
210	)
211	flat_topk_ids = topk_ids.view(-1).long()
212
213	for expert_idx in range(num_experts):
214	expert_mask = flat_topk_ids == expert_idx
215	if expert_mask.sum() == 0:
216	continue
217	gate_up = expanded[expert_mask] @ w1_f[expert_idx].T
218	gate = gate_up[:, :intermediate_dim]
219	up = gate_up[:, intermediate_dim:]
220	if activation == "silu":
221	act = F.silu(gate) * up
222	elif activation == "gelu":
223	act = F.gelu(gate) * up
224	else:
225	raise ValueError(f"Unknown activation: {activation}")
226	output[expert_mask] = act @ w2_f[expert_idx].T
227
228	output = output.view(num_tokens, topk, hidden_dim)
229	weights = topk_weights.float().view(num_tokens, topk, 1)
230	return (output * weights).sum(dim=1)
231
232
233	class AiterMxfp4MoeCase(NamedTuple):
234	hidden_states: torch.Tensor
235	w1_kernel: torch.Tensor
236	w2_kernel: torch.Tensor
237	w1_ref: torch.Tensor
238	w2_ref: torch.Tensor
239	topk_weights: torch.Tensor
240	topk_ids: torch.Tensor
241	moe_config: Any
242	quant_config: Any
243
244
245	def _make_topk_ids(
246	num_tokens: int,
247	num_experts: int,
248	topk: int,
249	*,
250	device: str = "cuda",
251	) -> torch.Tensor:
252	"""Generate distinct expert IDs per token with the same top-k shape as
253	production routing."""
254	router_logits = torch.randn(num_tokens, num_experts, device=device)
255	_, topk_ids = torch.topk(torch.softmax(router_logits, dim=-1), k=topk, dim=-1)
256	return topk_ids.to(torch.int32)
257
258
259	def _shuffle_moe_weights(
260	w1: torch.Tensor,
261	w2: torch.Tensor,
262	) -> tuple[torch.Tensor, torch.Tensor]:
263	"""Shuffle MoE weights into the AITER CK layout used in production."""
264	from vllm._aiter_ops import rocm_aiter_ops
265
266	w1_shuffled, w2_shuffled = rocm_aiter_ops.shuffle_weights(w1, w2)
267	w1_shuffled.is_shuffled = True
268	w2_shuffled.is_shuffled = True
269	return w1_shuffled, w2_shuffled
270
271
272	def _make_moe_case(
273	*,
274	num_tokens: int,
275	hidden_dim: int,
276	intermediate_dim: int,
277	num_experts: int,
278	topk: int,
279	seed: int,
280	) -> dict[str, torch.Tensor]:
281	torch.set_default_device("cuda")
282	set_random_seed(seed)
283
284	hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
285	w1 = torch.randn(
286	num_experts,
287	intermediate_dim * 2,
288	hidden_dim,
289	dtype=torch.bfloat16,
290	) / math.sqrt(hidden_dim)
291	w2 = torch.randn(
292	num_experts,
293	hidden_dim,
294	intermediate_dim,
295	dtype=torch.bfloat16,
296	) / math.sqrt(intermediate_dim)
297	topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32)
298	topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
299	topk_ids = _make_topk_ids(num_tokens, num_experts, topk)
300	return {
301	"hidden_states": hidden_states,
302	"w1": w1,
303	"w2": w2,
304	"topk_weights": topk_weights,
305	"topk_ids": topk_ids,
306	}
307
308
309	def _interleave_gate_up_rows(t: torch.Tensor) -> torch.Tensor:
310	"""Reorder contiguous ``[gate; up]`` rows into gpt-oss interleaved order."""
311	e, two_i = t.shape[0], t.shape[1]
312	i, rest = two_i // 2, t.shape[2:]
313	perm = (0, 2, 1, *range(3, 3 + len(rest)))
314	return t.view(e, 2, i, rest).permute(perm).contiguous().view(e, two_i, *rest)
315
316
317	def _make_aiter_mxfp4_moe_case(
318	*,
319	num_tokens: int,
320	hidden_dim: int,
321	intermediate_dim: int,
322	num_experts: int,
323	topk: int,
324	seed: int,
325	) -> AiterMxfp4MoeCase:
326	from triton_kernels.numerics_details.mxfp import (
327	downcast_to_mxfp,
328	upcast_from_mxfp,
329	)
330
331	from tests.kernels.moe.utils import make_dummy_moe_config
332	from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
333	Mxfp4MoeBackend,
334	convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
335	make_mxfp4_moe_quant_config,
336	)
337
338	torch.set_default_device("cuda")
339	set_random_seed(seed)
340
341	hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
342	w1 = torch.randn(
343	num_experts,
344	intermediate_dim * 2,
345	hidden_dim,
346	dtype=torch.bfloat16,
347	) / math.sqrt(hidden_dim)
348	w2 = torch.randn(
349	num_experts,
350	hidden_dim,
351	intermediate_dim,
352	dtype=torch.bfloat16,
353	) / math.sqrt(intermediate_dim)
354	topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32)
355	topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
356	topk_ids = _make_topk_ids(num_tokens, num_experts, topk)
357
358	w1_q, w1_scale = downcast_to_mxfp(w1, torch.uint8, axis=-1)
359	w2_q, w2_scale = downcast_to_mxfp(w2, torch.uint8, axis=-1)
360
361	w1_ref = upcast_from_mxfp(w1_q, w1_scale, torch.bfloat16, axis=-1)
362	w2_ref = upcast_from_mxfp(w2_q, w2_scale, torch.bfloat16, axis=-1)
363
364	(
365	w1_kernel,
366	w2_kernel,
367	w1_scale_kernel,
368	w2_scale_kernel,
369	_,
370	_,
371	) = convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
372	mxfp4_backend=Mxfp4MoeBackend.AITER,
373	layer=torch.nn.Module(),
374	# w13 rows must be gpt-oss-interleaved; the converter de-interleaves them
375	# back to the contiguous gate/up blocks that ``w1_ref`` / ``ref_moe_forward``
376	# use. w2 has no gate/up split, so it stays as-is.
377	w13_weight=_interleave_gate_up_rows(w1_q),
378	w2_weight=w2_q.clone(),
379	w13_weight_scale=_interleave_gate_up_rows(w1_scale),
380	w2_weight_scale=w2_scale.clone(),
381	)
382
383	moe_config = make_dummy_moe_config(
384	num_experts=num_experts,
385	experts_per_token=topk,
386	hidden_dim=hidden_dim,
387	intermediate_size=intermediate_dim,
388	in_dtype=torch.bfloat16,
389	)
390	quant_config = make_mxfp4_moe_quant_config(
391	Mxfp4MoeBackend.AITER,
392	w1_scale=w1_scale_kernel,
393	w2_scale=w2_scale_kernel,
394	)
395
396	return AiterMxfp4MoeCase(
397	hidden_states=hidden_states,
398	w1_kernel=w1_kernel,
399	w2_kernel=w2_kernel,
400	w1_ref=w1_ref,
401	w2_ref=w2_ref,
402	topk_weights=topk_weights,
403	topk_ids=topk_ids,
404	moe_config=moe_config,
405	quant_config=quant_config,
406	)
407
408
409	def _run_fused_moe(
410	hidden_states: torch.Tensor,
411	w1: torch.Tensor,
412	w2: torch.Tensor,
413	topk_weights: torch.Tensor,
414	topk_ids: torch.Tensor,
415	*,
416	activation_method: int,
417	quant_method: int,
418	) -> torch.Tensor:
419	w1_shuffled, w2_shuffled = _shuffle_moe_weights(w1, w2)
420	return torch.ops.vllm.rocm_aiter_fused_moe(
421	hidden_states,
422	w1_shuffled,
423	w2_shuffled,
424	topk_weights,
425	topk_ids,
426	expert_mask=None,
427	activation_method=activation_method,
428	quant_method=quant_method,
429	doweight_stage1=False,
430	)
431
432
433	# Custom op tests ---------------------------------------------------------
434
435
436	def test_aiter_fused_moe_custom_op_registered():
437	"""The main fused-MoE custom op should stay registered for runtime use."""
438	_assert_aiter_supported()
439	import vllm._aiter_ops as aiter_ops # noqa: F401
440
441	assert hasattr(torch.ops.vllm, "rocm_aiter_fused_moe")
442	assert callable(torch.ops.vllm.rocm_aiter_fused_moe)
443
444
445	def test_aiter_asm_moe_tkw1_custom_op_registered():
446	"""The tkw1 custom op should stay registered for FP8 apply-router-weight
447	paths."""
448	_assert_aiter_supported()
449	import vllm._aiter_ops as aiter_ops # noqa: F401
450
451	assert hasattr(torch.ops.vllm, "rocm_aiter_asm_moe_tkw1")
452	assert callable(torch.ops.vllm.rocm_aiter_asm_moe_tkw1)
453
454
455	def test_aiter_fused_moe_fake_tensor_support():
456	"""The fused-MoE op should preserve fake-tensor compatibility for
457	torch.compile-style tracing."""
458	_assert_aiter_supported()
459	import vllm._aiter_ops # noqa: F401
460
461	num_tokens = 16
462	hidden_dim = 1024
463	intermediate_dim = 2048
464	num_experts = 8
465	topk = 2
466
467	hidden_states = torch.randn(
468	num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda"
469	)
470	w1 = torch.randn(
471	num_experts,
472	intermediate_dim * 2,
473	hidden_dim,
474	dtype=torch.bfloat16,
475	device="cuda",
476	)
477	w2 = torch.randn(
478	num_experts,
479	hidden_dim,
480	intermediate_dim,
481	dtype=torch.bfloat16,
482	device="cuda",
483	)
484	topk_weights = torch.rand(num_tokens, topk, dtype=torch.float32, device="cuda")
485	topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
486	topk_ids = _make_topk_ids(num_tokens, num_experts, topk, device="cuda")
487
488	torch.library.opcheck(
489	torch.ops.vllm.rocm_aiter_fused_moe,
490	(hidden_states, w1, w2, topk_weights, topk_ids),
491	kwargs={
492	"expert_mask": None,
493	"activation_method": 0,
494	"quant_method": 0,
495	"doweight_stage1": False,
496	},
497	test_utils=("test_faketensor",),
498	)
499
500
501	# Env gating tests --------------------------------------------------------
502
503
504	@pytest.mark.parametrize(
505	("use_aiter", "use_moe", "expected"),
506	[
507	(True, True, True),
508	(True, False, False),
509	(False, True, False),
510	(False, False, False),
511	],
512	)
513	def test_aiter_moe_enablement_follows_env(
514	use_aiter: bool,
515	use_moe: bool,
516	expected: bool,
517	monkeypatch: pytest.MonkeyPatch,
518	):
519	"""The fused-MoE gate should depend only on the main AITER toggle and the
520	MoE-specific toggle."""
521	from vllm._aiter_ops import rocm_aiter_ops
522
523	_assert_aiter_supported()
524
525	with monkeypatch.context() as mp:
526	mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
527	mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_moe else "0")
528	_reload_envs()
529	rocm_aiter_ops.refresh_env_variables()
530
531	assert rocm_aiter_ops.is_fused_moe_enabled() is expected
532
533
534	@pytest.mark.parametrize(
535	("use_aiter", "use_moe", "use_shared", "expected"),
536	[
537	(True, True, True, True),
538	(True, True, False, False),
539	(True, False, True, False),
540	(False, True, True, False),
541	],
542	)
543	def test_aiter_moe_shared_experts_enablement_follows_env(
544	use_aiter: bool,
545	use_moe: bool,
546	use_shared: bool,
547	expected: bool,
548	monkeypatch: pytest.MonkeyPatch,
549	):
550	"""Shared-expert fusion should only be enabled when the fused-MoE path is
551	enabled too."""
552	from vllm._aiter_ops import rocm_aiter_ops
553
554	_assert_aiter_supported()
555
556	with monkeypatch.context() as mp:
557	mp.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
558	mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1" if use_moe else "0")
559	mp.setenv(
560	"VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS",
561	"1" if use_shared else "0",
562	)
563	_reload_envs()
564	rocm_aiter_ops.refresh_env_variables()
565
566	assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected
567
568
569	@pytest.mark.parametrize(
570	"value,expected_act,expected_env",
571	[
572	# Unset / auto / legacy "1" all mean the a4w4 default.
573	(None, "a4w4", {"AITER_SITUV2_A4W4": "1"}),
574	("auto", "a4w4", {"AITER_SITUV2_A4W4": "1"}),
575	("1", "a4w4", {"AITER_SITUV2_A4W4": "1"}),
576	("a4w4", "a4w4", {"AITER_SITUV2_A4W4": "1"}),
577	("0", "a16w4", {}),
578	("a16w4", "a16w4", {}),
579	("A8W4", "a8w4", {"AITER_SITUV2_A8W4": "1"}),
580	],
581	)
582	def test_aiter_moe_situv2_activation_syncs_aiter_env(
583	value: str \| None,
584	expected_act: str,
585	expected_env: dict[str, str],
586	monkeypatch: pytest.MonkeyPatch,
587	):
588	"""Each SiTUv2 activation choice sets exactly its AITER dispatch env."""
589	import os
590
591	from vllm._aiter_ops import rocm_aiter_ops
592
593	_assert_aiter_supported()
594
595	with monkeypatch.context() as mp:
596	# Stale AITER overrides must be replaced by the vLLM selection.
597	mp.setenv("AITER_SITUV2_A8W4", "1")
598	mp.setenv("AITER_SITUV2_A4W4", "1")
599	mp.setenv("VLLM_ROCM_USE_AITER", "1")
600	mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
601	if value is None:
602	mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False)
603	else:
604	mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", value)
605	_reload_envs()
606	rocm_aiter_ops.refresh_env_variables()
607
608	assert rocm_aiter_ops.get_fused_moe_situv2_activation() == expected_act
609	assert rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved() is (
610	expected_act == "a8w4"
611	)
612	assert rocm_aiter_ops.is_fused_moe_situv2_enabled() is (expected_act != "a16w4")
613	for name in ("AITER_SITUV2_A8W4", "AITER_SITUV2_A4W4"):
614	assert os.environ.get(name) == expected_env.get(name)
615
616
617	def test_aiter_moe_situv2_rejects_unknown_activation(
618	monkeypatch: pytest.MonkeyPatch,
619	):
620	with monkeypatch.context() as mp:
621	mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "a8w8")
622	_reload_envs()
623	import vllm.envs as envs
624
625	with pytest.raises(ValueError, match="VLLM_ROCM_USE_AITER_MOE_SITUV2"):
626	_ = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
627
628
629	@pytest.mark.parametrize("moe_padding", [True, False])
630	def test_aiter_moe_padding_env_var(
631	moe_padding: bool,
632	monkeypatch: pytest.MonkeyPatch,
633	):
634	"""The ROCm MoE padding env var should keep its exact parse contract."""
635	with monkeypatch.context() as mp:
636	mp.setenv("VLLM_ROCM_MOE_PADDING", "1" if moe_padding else "0")
637	envs = _reload_envs()
638	assert envs.VLLM_ROCM_MOE_PADDING is moe_padding
639
640
641	# a4w4 (FP4 activation) opt-in gating test ---------------------------------
642
643
644	@pytest.mark.parametrize(
645	("env_value", "model_type", "expected"),
646	[
647	# Unset/"0" is opt-in-only: a4w4 never turns on by itself, even on
648	# a validated model type.
649	(None, "deepseek_v41", False),
650	("0", "deepseek_v41", False),
651	("1", "deepseek_v41", True),
652	("1", "deepseek_v41_text", True),
653	],
654	)
655	def test_aiter_moe_a4w4_dsv4_is_opt_in(
656	env_value: str \| None,
657	model_type: str,
658	expected: bool,
659	monkeypatch: pytest.MonkeyPatch,
660	):
661	"""vllm-project/vllm#58819's a4w4 MoE activation path must stay a plain
662	explicit opt-in (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1), with no
663	default-on behavior, but is still only permitted on the model types
664	this flag has been validated for."""
665	import types
666
667	from tests.kernels.moe.utils import make_dummy_moe_config
668	from vllm._aiter_ops import rocm_aiter_ops
669
670	_assert_aiter_supported()
671
672	fake_vllm_config = types.SimpleNamespace(
673	model_config=types.SimpleNamespace(
674	hf_config=types.SimpleNamespace(model_type=model_type)
675	)
676	)
677
678	with monkeypatch.context() as mp:
679	mp.setenv("VLLM_ROCM_USE_AITER", "1")
680	mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
681	if env_value is None:
682	mp.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False)
683	else:
684	mp.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", env_value)
685	_reload_envs()
686	rocm_aiter_ops.refresh_env_variables()
687
688	with mock.patch(
689	"vllm.config.get_current_vllm_config_or_none",
690	return_value=fake_vllm_config,
691	):
692	moe_config = make_dummy_moe_config(
693	num_experts=4,
694	experts_per_token=2,
695	hidden_dim=256,
696	intermediate_size=512,
697	in_dtype=torch.bfloat16,
698	)
699
700	assert moe_config.use_mxfp4_w4a4_dsv4 is expected
701
702
703	@pytest.mark.parametrize("model_type", ["deepseek_v4", "gpt_oss", None])
704	def test_aiter_moe_a4w4_dsv4_rejects_unvalidated_model_type(
705	model_type: str \| None,
706	monkeypatch: pytest.MonkeyPatch,
707	):
708	"""Forcing a4w4 on a model type this flag hasn't been GSM8K/AgentX
709	validated for must raise: it also flips the MXFP4 weight-shuffle layout
710	(ATOM SEPARATED vs INTERLEAVE), which silently produces garbled output
711	on the wrong model rather than just being a slower path — the same bug
712	class vllm-project/vllm#58819's own second commit fixed for DeepSeek
713	V4.1 itself."""
714	import types
715
716	from tests.kernels.moe.utils import make_dummy_moe_config
717	from vllm._aiter_ops import rocm_aiter_ops
718
719	_assert_aiter_supported()
720
721	fake_vllm_config = types.SimpleNamespace(
722	model_config=types.SimpleNamespace(
723	hf_config=types.SimpleNamespace(model_type=model_type)
724	)
725	)
726
727	with monkeypatch.context() as mp:
728	mp.setenv("VLLM_ROCM_USE_AITER", "1")
729	mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
730	mp.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1")
731	_reload_envs()
732	rocm_aiter_ops.refresh_env_variables()
733
734	with (
735	mock.patch(
736	"vllm.config.get_current_vllm_config_or_none",
737	return_value=fake_vllm_config,
738	),
739	pytest.raises(ValueError, match="VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4"),
740	):
741	make_dummy_moe_config(
742	num_experts=4,
743	experts_per_token=2,
744	hidden_dim=256,
745	intermediate_size=512,
746	in_dtype=torch.bfloat16,
747	)
748
749
750	# Dispatch-policy forwarding test ------------------------------------------
751
752
753	@pytest.mark.parametrize("dispatch_policy", [0, 1, 2])
754	def test_aiter_moe_dispatch_policy_forwarded_to_fused_moe(
755	dispatch_policy: int,
756	monkeypatch: pytest.MonkeyPatch,
757	):
758	"""VLLM_ROCM_AITER_MOE_DISPATCH_POLICY should reach rocm_aiter_ops.fused_moe
759	unchanged, the same forwarding AiterExperts.apply does via
760	rocm_aiter_fused_experts, for every documented policy value (0=auto,
761	1=always single-pass, 2=always multi-pass). See vllm-project/vllm#54966.
762	"""
763	from tests.kernels.moe.utils import make_dummy_moe_config
764	from vllm._aiter_ops import rocm_aiter_ops
765	from vllm.model_executor.layers.fused_moe.activation import MoEActivation
766	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
767	rocm_aiter_fused_experts,
768	)
769
770	_assert_aiter_supported()
771	case = _make_moe_case(
772	num_tokens=8,
773	hidden_dim=256,
774	intermediate_dim=512,
775	num_experts=4,
776	topk=2,
777	seed=5,
778	)
779	w1_shuffled, w2_shuffled = _shuffle_moe_weights(case["w1"], case["w2"])
780	moe_config = make_dummy_moe_config(
781	num_experts=4,
782	experts_per_token=2,
783	hidden_dim=256,
784	intermediate_size=512,
785	in_dtype=torch.bfloat16,
786	)
787
788	with monkeypatch.context() as mp:
789	mp.setenv("VLLM_ROCM_AITER_MOE_DISPATCH_POLICY", str(dispatch_policy))
790	_reload_envs()
791	rocm_aiter_ops.refresh_env_variables()
792
793	assert rocm_aiter_ops.get_moe_dispatch_policy() == dispatch_policy, (
794	"rocm_aiter_ops cached a stale dispatch policy after the env var changed."
795	)
796
797	with mock.patch.object(
798	rocm_aiter_ops, "fused_moe", wraps=rocm_aiter_ops.fused_moe
799	) as fused_moe_mock:
800	rocm_aiter_fused_experts(
801	hidden_states=case["hidden_states"],
802	w1=w1_shuffled,
803	w2=w2_shuffled,
804	topk_weights=case["topk_weights"],
805	topk_ids=case["topk_ids"],
806	moe_config=moe_config,
807	activation=MoEActivation.SILU,
808	moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
809	)
810
811	assert (
812	fused_moe_mock.call_args.kwargs["moe_sorting_dispatch_policy"]
813	== dispatch_policy
814	), (
815	"VLLM_ROCM_AITER_MOE_DISPATCH_POLICY was not forwarded to "
816	"rocm_aiter_ops.fused_moe."
817	)
818
819
820	# Enum tests --------------------------------------------------------------
821
822
823	def test_quant_method_enum_values():
824	"""The AITER quant-method bridge enum should keep its wire values."""
825	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import QuantMethod
826
827	assert QuantMethod.NO == 0
828	assert QuantMethod.PER_TENSOR == 1
829	assert QuantMethod.PER_TOKEN == 2
830	assert QuantMethod.BLOCK_1X32 == 3
831	assert QuantMethod.BLOCK_1X128 == 4
832	assert QuantMethod.BLOCK_128x128 == 5
833
834
835	def test_activation_method_enum_values():
836	"""The AITER activation-method bridge enum should keep its wire values."""
837	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
838	ActivationMethod,
839	)
840
841	assert ActivationMethod.SILU == 0
842	assert ActivationMethod.GELU == 1
843
844
845	# MXFP4 kernel tests ------------------------------------------------------
846
847
848	def test_aiter_mxfp4_quant_scheme_support_matches_gfx950():
849	"""AITER MXFP4 MoE support should stay gfx950-only."""
850	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
851	AiterExperts,
852	)
853	from vllm.model_executor.layers.quantization.utils.quant_utils import (
854	kMxfp4Static,
855	)
856
857	assert AiterExperts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950()
858
859
860	@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
861	def test_aiter_fused_moe_mi350_mxfp4_w4a16_accuracy():
862	"""The gfx950 AITER MXFP4 W4A16 MoE path should match the dequantized
863	MXFP4 reference."""
864	from vllm.model_executor.layers.fused_moe.activation import MoEActivation
865	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
866	rocm_aiter_fused_experts,
867	)
868
869	_assert_aiter_supported()
870	case = _make_aiter_mxfp4_moe_case(
871	num_tokens=32,
872	hidden_dim=512,
873	intermediate_dim=1024,
874	num_experts=4,
875	topk=2,
876	seed=11,
877	)
878	ref_out = ref_moe_forward(
879	case.hidden_states,
880	case.w1_ref,
881	case.w2_ref,
882	case.topk_weights,
883	case.topk_ids,
884	)
885	out = rocm_aiter_fused_experts(
886	hidden_states=case.hidden_states,
887	w1=case.w1_kernel,
888	w2=case.w2_kernel,
889	topk_weights=case.topk_weights,
890	topk_ids=case.topk_ids,
891	activation=MoEActivation.SWIGLUOAI,
892	quant_config=case.quant_config,
893	moe_config=case.moe_config,
894	expert_mask=None,
895	)
896
897	assert out.shape == case.hidden_states.shape
898	_assert_close_budget(
899	out.float(),
900	ref_out.float(),
901	label="mi350_mxfp4_w4a16_accuracy",
902	atol=0.1,
903	rtol=0.1,
904	pass_rate=0.99,
905	max_violation_factor=2.0,
906	)
907
908
909	@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only")
910	def test_aiter_fused_moe_mi350_mxfp4_w4a16_determinism():
911	"""The gfx950 AITER MXFP4 W4A16 MoE path should stay bitwise
912	deterministic."""
913	from vllm.model_executor.layers.fused_moe.activation import MoEActivation
914	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
915	rocm_aiter_fused_experts,
916	)
917
918	_assert_aiter_supported()
919	case = _make_aiter_mxfp4_moe_case(
920	num_tokens=8,
921	hidden_dim=512,
922	intermediate_dim=1024,
923	num_experts=4,
924	topk=2,
925	seed=13,
926	)
927
928	def run_mxfp4_moe():
929	return rocm_aiter_fused_experts(
930	hidden_states=case.hidden_states,
931	w1=case.w1_kernel,
932	w2=case.w2_kernel,
933	topk_weights=case.topk_weights,
934	topk_ids=case.topk_ids,
935	activation=MoEActivation.SWIGLUOAI,
936	quant_config=case.quant_config,
937	moe_config=case.moe_config,
938	expert_mask=None,
939	)
940
941	_assert_deterministic(run_mxfp4_moe, n_runs=4)
942
943
944	# FP8 group-quant tests ---------------------------------------------------
945
946
947	@pytest.mark.skipif(
948	not (on_gfx942() or on_gfx950()),
949	reason="gfx942/gfx950 ROCm only",
950	)
951	@pytest.mark.parametrize("num_tokens,hidden_dim", [(16, 2048), (64, 4096), (128, 8192)])
952	def test_aiter_moe_group_fp8_quant_reconstructs_hidden_states(
953	num_tokens: int,
954	hidden_dim: int,
955	):
956	"""The MoE-facing FP8 group-quant helper should reconstruct hidden states
957	within the expected FP8 error budget.
958
959	The raw op's shape and standalone roundtrip contracts are covered in
960	``tests/kernels/core/test_rocm_aiter_ops.py``. This test keeps the
961	representative hidden-state sizes that the fused-MoE path actually cares
962	about.
963	"""
964	from vllm._aiter_ops import rocm_aiter_ops
965
966	_assert_aiter_supported()
967	torch.set_default_device("cuda")
968	set_random_seed(1)
969
970	group_size = 128
971	hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
972	x_fp8, scales = rocm_aiter_ops.group_fp8_quant(hidden_states, group_size)
973
974	assert x_fp8.shape == hidden_states.shape
975	assert scales.shape == (num_tokens, (hidden_dim + group_size - 1) // group_size)
976	assert scales.dtype == torch.float32
977
978	scales_expanded = scales.repeat_interleave(group_size, dim=1)[:, :hidden_dim]
979	dequantized = x_fp8.float() * scales_expanded
980
981	_assert_group_quant_quality(
982	dequantized,
983	hidden_states,
984	label=f"group_fp8_quant shape=({num_tokens}, {hidden_dim})",
985	)
986
987
988	# Kernel accuracy tests ---------------------------------------------------
989
990
991	@pytest.mark.parametrize(
992	"num_tokens,hidden_dim,intermediate_dim",
993	[
994	(16, 512, 1024),
995	(128, 2048, 4096),
996	(2048, 4096, 11008),
997	],
998	)
999	def test_aiter_fused_moe_bf16_accuracy(
1000	num_tokens: int,
1001	hidden_dim: int,
1002	intermediate_dim: int,
1003	):
1004	"""The ROCm AITER fused-MoE BF16 path should match the float32 reference
1005	on representative shapes."""
1006	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1007	ActivationMethod,
1008	QuantMethod,
1009	)
1010
1011	_assert_aiter_supported()
1012	case = _make_moe_case(
1013	num_tokens=num_tokens,
1014	hidden_dim=hidden_dim,
1015	intermediate_dim=intermediate_dim,
1016	num_experts=8,
1017	topk=2,
1018	seed=0,
1019	)
1020	ref_out = ref_moe_forward(
1021	case["hidden_states"],
1022	case["w1"],
1023	case["w2"],
1024	case["topk_weights"],
1025	case["topk_ids"],
1026	activation="silu",
1027	)
1028	out = _run_fused_moe(
1029	case["hidden_states"],
1030	case["w1"],
1031	case["w2"],
1032	case["topk_weights"],
1033	case["topk_ids"],
1034	activation_method=int(ActivationMethod.SILU),
1035	quant_method=int(QuantMethod.NO),
1036	)
1037
1038	assert out.shape == (num_tokens, hidden_dim)
1039	_assert_close_budget(
1040	out.float(),
1041	ref_out,
1042	label=f"bf16_accuracy shape=({num_tokens}, {hidden_dim}, {intermediate_dim})",
1043	atol=0.05,
1044	rtol=0.0,
1045	)
1046
1047
1048	def test_aiter_fused_moe_gelu_accuracy():
1049	"""The GELU activation variant should stay aligned with the float32
1050	reference."""
1051	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1052	ActivationMethod,
1053	QuantMethod,
1054	)
1055
1056	_assert_aiter_supported()
1057	case = _make_moe_case(
1058	num_tokens=32,
1059	hidden_dim=512,
1060	intermediate_dim=1024,
1061	num_experts=4,
1062	topk=2,
1063	seed=42,
1064	)
1065	ref_out = ref_moe_forward(
1066	case["hidden_states"],
1067	case["w1"],
1068	case["w2"],
1069	case["topk_weights"],
1070	case["topk_ids"],
1071	activation="gelu",
1072	)
1073	out = _run_fused_moe(
1074	case["hidden_states"],
1075	case["w1"],
1076	case["w2"],
1077	case["topk_weights"],
1078	case["topk_ids"],
1079	activation_method=int(ActivationMethod.GELU),
1080	quant_method=int(QuantMethod.NO),
1081	)
1082
1083	assert out.shape == case["hidden_states"].shape
1084	_assert_close_budget(
1085	out.float(),
1086	ref_out,
1087	label="gelu_accuracy",
1088	atol=0.05,
1089	rtol=0.0,
1090	)
1091
1092
1093	def test_aiter_fused_moe_determinism():
1094	"""The BF16 fused-MoE kernel should stay bitwise deterministic for the
1095	same inputs."""
1096	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1097	ActivationMethod,
1098	QuantMethod,
1099	)
1100
1101	_assert_aiter_supported()
1102	case = _make_moe_case(
1103	num_tokens=8,
1104	hidden_dim=256,
1105	intermediate_dim=512,
1106	num_experts=4,
1107	topk=2,
1108	seed=2,
1109	)
1110
1111	def run_moe():
1112	return _run_fused_moe(
1113	case["hidden_states"],
1114	case["w1"],
1115	case["w2"],
1116	case["topk_weights"],
1117	case["topk_ids"],
1118	activation_method=int(ActivationMethod.SILU),
1119	quant_method=int(QuantMethod.NO),
1120	)
1121
1122	_assert_deterministic(run_moe, n_runs=4)
1123
1124
1125	# Routed end-to-end tests -------------------------------------------------
1126
1127
1128	@pytest.mark.parametrize(
1129	"num_tokens,num_experts,topk",
1130	[
1131	(1, 8, 2),
1132	(16, 8, 2),
1133	(64, 16, 4),
1134	],
1135	)
1136	def test_aiter_fused_moe_end_to_end(
1137	num_tokens: int,
1138	num_experts: int,
1139	topk: int,
1140	):
1141	"""The full router-logits to top-k to fused-MoE path should stay accurate."""
1142	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1143	ActivationMethod,
1144	QuantMethod,
1145	)
1146
1147	_assert_aiter_supported()
1148	torch.set_default_device("cuda")
1149	set_random_seed(7)
1150
1151	hidden_dim = 512
1152	intermediate_dim = 1024
1153	hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16)
1154	w1 = torch.randn(
1155	num_experts,
1156	intermediate_dim * 2,
1157	hidden_dim,
1158	dtype=torch.bfloat16,
1159	) / math.sqrt(hidden_dim)
1160	w2 = torch.randn(
1161	num_experts,
1162	hidden_dim,
1163	intermediate_dim,
1164	dtype=torch.bfloat16,
1165	) / math.sqrt(intermediate_dim)
1166
1167	router_logits = torch.randn(num_tokens, num_experts, device="cuda")
1168	router_probs = torch.softmax(router_logits, dim=-1)
1169	topk_weights, topk_ids = torch.topk(router_probs, k=topk, dim=-1)
1170	topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1171	topk_weights = topk_weights.float()
1172	topk_ids = topk_ids.to(torch.int32)
1173
1174	ref_out = ref_moe_forward(hidden_states, w1, w2, topk_weights, topk_ids)
1175	out = _run_fused_moe(
1176	hidden_states,
1177	w1,
1178	w2,
1179	topk_weights,
1180	topk_ids,
1181	activation_method=int(ActivationMethod.SILU),
1182	quant_method=int(QuantMethod.NO),
1183	)
1184
1185	assert out.shape == (num_tokens, hidden_dim)
1186	assert out.dtype == torch.bfloat16
1187	_assert_close_budget(
1188	out.float(),
1189	ref_out,
1190	label=f"end_to_end tokens={num_tokens} experts={num_experts} topk={topk}",
1191	atol=0.05,
1192	rtol=0.0,
1193	)
1194
1195
1196	# Arch-specific tests -----------------------------------------------------
1197
1198
1199	@pytest.mark.skipif(
1200	not (on_gfx942() or on_gfx950()),
1201	reason="gfx942/gfx950 ROCm only",
1202	)
1203	def test_aiter_fused_moe_mi3xx_bf16_accuracy():
1204	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1205	ActivationMethod,
1206	QuantMethod,
1207	)
1208
1209	_assert_aiter_supported()
1210	case = _make_moe_case(
1211	num_tokens=64,
1212	hidden_dim=4096,
1213	intermediate_dim=11008,
1214	num_experts=8,
1215	topk=2,
1216	seed=42,
1217	)
1218	ref_out = ref_moe_forward(
1219	case["hidden_states"],
1220	case["w1"],
1221	case["w2"],
1222	case["topk_weights"],
1223	case["topk_ids"],
1224	)
1225	out = _run_fused_moe(
1226	case["hidden_states"],
1227	case["w1"],
1228	case["w2"],
1229	case["topk_weights"],
1230	case["topk_ids"],
1231	activation_method=int(ActivationMethod.SILU),
1232	quant_method=int(QuantMethod.NO),
1233	)
1234
1235	_assert_close_budget(
1236	out.float(),
1237	ref_out,
1238	label=(f"mi3xx_bf16_accuracy arch={'gfx942' if on_gfx942() else 'gfx950'}"),
1239	atol=0.05,
1240	rtol=0.0,
1241	)
1242
1243
1244	@pytest.mark.skipif(
1245	not (on_gfx942() or on_gfx950()),
1246	reason="gfx942/gfx950 ROCm only",
1247	)
1248	def test_aiter_fused_moe_mi3xx_fp8_accuracy():
1249	"""The MI3xx FP8 per-tensor MoE path should stay within the measured FP8
1250	error budget."""
1251	from tests.kernels.moe.utils import make_test_weights
1252	from vllm._aiter_ops import rocm_aiter_ops
1253	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1254	ActivationMethod,
1255	QuantMethod,
1256	)
1257
1258	_assert_aiter_supported()
1259	torch.set_default_device("cuda")
1260	set_random_seed(99)
1261
1262	num_tokens = 32
1263	hidden_dim = 512
1264	intermediate_dim = 1024
1265	num_experts = 4
1266	topk = 2
1267
1268	hidden_states = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16) / 10
1269	fp8_dtype = current_platform.fp8_dtype()
1270	(w1_bf16, w1_fp8, w1_scale, _), (w2_bf16, w2_fp8, w2_scale, _) = make_test_weights(
1271	num_experts,
1272	intermediate_dim * 2,
1273	hidden_dim,
1274	torch.bfloat16,
1275	torch.float8_e4m3fn,
1276	per_out_ch_quant=False,
1277	)
1278	_, a1_scale = rocm_aiter_ops.per_tensor_quant(hidden_states, fp8_dtype)
1279	router_logits = torch.randn(num_tokens, num_experts, device="cuda")
1280	router_probs = torch.softmax(router_logits, dim=-1)
1281	topk_weights, topk_ids = torch.topk(router_probs, k=topk, dim=-1)
1282	topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1283	topk_weights = topk_weights.float()
1284	topk_ids = topk_ids.to(torch.int32)
1285
1286	ref_out = ref_moe_forward(
1287	hidden_states,
1288	w1_bf16,
1289	w2_bf16,
1290	topk_weights,
1291	topk_ids,
1292	)
1293
1294	w1_shuffled, w2_shuffled = _shuffle_moe_weights(w1_fp8, w2_fp8)
1295	out = torch.ops.vllm.rocm_aiter_fused_moe(
1296	hidden_states,
1297	w1_shuffled,
1298	w2_shuffled,
1299	topk_weights,
1300	topk_ids,
1301	expert_mask=None,
1302	activation_method=int(ActivationMethod.SILU),
1303	quant_method=int(QuantMethod.PER_TENSOR),
1304	doweight_stage1=False,
1305	w1_scale=w1_scale,
1306	w2_scale=w2_scale,
1307	a1_scale=a1_scale,
1308	)
1309
1310	assert out.shape == hidden_states.shape
1311	_assert_close_budget(
1312	out.float(),
1313	ref_out,
1314	label=(f"mi3xx_fp8_accuracy arch={'gfx942' if on_gfx942() else 'gfx950'}"),
1315	atol=0.02,
1316	rtol=0.0,
1317	pass_rate=1.0,
1318	max_violation_factor=1.5,
1319	)
1320
1321
1322	# Weight alignment tests --------------------------------------------------
1323	#
1324	# AITER's CK 2stages MoE kernel rejects an intermediate size not divisible by
1325	# its tile width (64 at or below inter_dim 192, 128 above). Some model + TP
1326	# splits land on an unaligned size (e.g. 1792 / TP=8 = 224), so the AITER path
1327	# rounds the intermediate dim up in ``maybe_roundup_sizes`` and allocates the
1328	# weights at the padded size.
1329
1330	ALIGNMENT_HIDDEN = 64
1331	ALIGNMENT_NUM_EXPERTS = 2
1332
1333
1334	def _make_alignment_moe_config(intermediate: int):
1335	from tests.kernels.moe.utils import make_dummy_moe_config
1336
1337	return make_dummy_moe_config(
1338	num_experts=ALIGNMENT_NUM_EXPERTS,
1339	hidden_dim=ALIGNMENT_HIDDEN,
1340	intermediate_size=intermediate,
1341	)
1342
1343
1344	def _make_aiter_method(moe_config):
1345	"""Build the unquantized method with the backend pinned to AITER.
1346
1347	``select_unquantized_moe_backend`` needs a real ROCm + AITER runtime, so
1348	stub it out and set the backend directly.
1349	"""
1350	from unittest.mock import patch
1351
1352	from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1353	UnquantizedMoeBackend,
1354	)
1355	from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
1356	UnquantizedFusedMoEMethod,
1357	)
1358
1359	with patch(
1360	"vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method"
1361	".select_unquantized_moe_backend",
1362	return_value=(UnquantizedMoeBackend.AITER, None),
1363	):
1364	return UnquantizedFusedMoEMethod(moe_config)
1365
1366
1367	def _roundup(method, moe_config, intermediate):
1368	from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
1369
1370	return method.maybe_roundup_sizes(
1371	hidden_size=ALIGNMENT_HIDDEN,
1372	intermediate_size_per_partition=intermediate,
1373	act_dtype=moe_config.in_dtype,
1374	moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1375	)
1376
1377
1378	@pytest.mark.parametrize(
1379	("intermediate", "expected_padded"),
1380	[
1381	(64, 64), # already aligned, untouched
1382	(96, 128),
1383	(160, 192),
1384	(192, 192), # must stay 192: it is valid and has tuned configs
1385	(224, 256), # K2-Horizon-375B at TP=8
1386	(256, 256),
1387	(448, 512), # K2-Horizon-375B at TP=4
1388	(4096, 4096),
1389	],
1390	)
1391	def test_aiter_moe_roundup_pads_intermediate(
1392	intermediate,
1393	expected_padded,
1394	default_vllm_config,
1395	):
1396	from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1397	aiter_moe_intermediate_alignment,
1398	)
1399
1400	moe_config = _make_alignment_moe_config(intermediate)
1401	method = _make_aiter_method(moe_config)
1402
1403	hidden, padded = _roundup(method, moe_config, intermediate)
1404
1405	assert padded == expected_padded
1406	assert hidden == ALIGNMENT_HIDDEN
1407	# Padding must not move a size across the <= 192 threshold, or the
1408	# alignment we picked is not the one the kernel dispatches the padded
1409	# shape to.
1410	assert aiter_moe_intermediate_alignment(padded) == (
1411	aiter_moe_intermediate_alignment(intermediate)
1412	)
1413
1414
1415	@pytest.mark.parametrize("backend_name", ["TRITON", "FLASHINFER_CUTLASS"])
1416	def test_aiter_moe_roundup_is_not_applied_to_other_backends(
1417	backend_name,
1418	default_vllm_config,
1419	):
1420	"""Backends with no alignment requirement keep the unaligned size."""
1421	from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1422	UnquantizedMoeBackend,
1423	)
1424
1425	intermediate = 224
1426	moe_config = _make_alignment_moe_config(intermediate)
1427	method = _make_aiter_method(moe_config)
1428	method.unquantized_backend = UnquantizedMoeBackend[backend_name]
1429
1430	_, padded = _roundup(method, moe_config, intermediate)
1431
1432	assert padded == intermediate
1433
1434
1435	@pytest.mark.parametrize("intermediate", [224, 448])
1436	def test_aiter_moe_reload_zeroes_intermediate_padding(
1437	intermediate,
1438	monkeypatch,
1439	default_vllm_config,
1440	):
1441	"""A reload must re-zero pad lanes left dirty in the weight storage.
1442
1443	The reload writes back only the logical slices, so the pad lanes hold
1444	garbage (NaN here) that the AITER conversion must zero. Two passes also
1445	cover reloading into storage that already holds a previous conversion.
1446	"""
1447	from unittest.mock import MagicMock
1448
1449	from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1450	aiter_moe_intermediate_alignment,
1451	)
1452
1453	alignment = aiter_moe_intermediate_alignment(intermediate)
1454	padded = -(-intermediate // alignment) * alignment
1455	assert padded != intermediate
1456
1457	moe_config = _make_alignment_moe_config(intermediate)
1458	assert moe_config.is_act_and_mul
1459	method = _make_aiter_method(moe_config)
1460	# RoutedExperts records the rounded-up size after maybe_roundup_sizes.
1461	moe_config.intermediate_size_per_partition = padded
1462	layer = torch.nn.Module()
1463	layer.moe_config = moe_config
1464
1465	method.create_weights(
1466	layer=layer,
1467	num_experts=ALIGNMENT_NUM_EXPERTS,
1468	hidden_size=ALIGNMENT_HIDDEN,
1469	intermediate_size_per_partition=padded,
1470	params_dtype=torch.float32,
1471	)
1472
1473	# Identity shuffle keeps pad lanes sliceable; both stubs need real ROCm.
1474	monkeypatch.setattr(
1475	"vllm.model_executor.layers.fused_moe.oracle.unquantized"
1476	".rocm_aiter_ops.shuffle_weights",
1477	lambda w13, w2: (w13, w2),
1478	)
1479	monkeypatch.setattr(
1480	method,
1481	"_init_moe_kernel",
1482	lambda _layer: setattr(method, "moe_kernel", MagicMock()),
1483	)
1484
1485	up = padded # up projection starts one padded block into the fused rows
1486
1487	for _ in range(2):
1488	gate = torch.randn(ALIGNMENT_NUM_EXPERTS, intermediate, ALIGNMENT_HIDDEN)
1489	up_proj = torch.randn(ALIGNMENT_NUM_EXPERTS, intermediate, ALIGNMENT_HIDDEN)
1490	down = torch.randn(ALIGNMENT_NUM_EXPERTS, ALIGNMENT_HIDDEN, intermediate)
1491
1492	# NaN the pads, then write only the logical slices like the loader.
1493	layer.w13_weight.data.fill_(float("nan"))
1494	layer.w2_weight.data.fill_(float("nan"))
1495	layer.w13_weight.data[:, :intermediate].copy_(gate)
1496	layer.w13_weight.data[:, up : up + intermediate].copy_(up_proj)
1497	layer.w2_weight.data[:, :, :intermediate].copy_(down)
1498
1499	method.process_weights_after_loading(layer)
1500
1501	w13 = layer.w13_weight
1502	w2 = layer.w2_weight
1503	assert torch.equal(w13[:, :intermediate], gate)
1504	assert torch.equal(w13[:, up : up + intermediate], up_proj)
1505	assert torch.equal(w2[:, :, :intermediate], down)
1506	assert torch.all(w13[:, intermediate:up] == 0)
1507	assert torch.all(w13[:, up + intermediate :] == 0)
1508	assert torch.all(w2[:, :, intermediate:] == 0)
1509
1510
1511	def _aiter_accepts_intermediate(intermediate: int, num_tokens: int) -> bool:
1512	"""Run the real AITER CK MoE kernel and report whether it dispatched."""
1513	from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1514	ActivationMethod,
1515	QuantMethod,
1516	)
1517
1518	case = _make_moe_case(
1519	num_tokens=num_tokens,
1520	hidden_dim=1024,
1521	intermediate_dim=intermediate,
1522	num_experts=8,
1523	topk=2,
1524	seed=0,
1525	)
1526	try:
1527	_run_fused_moe(
1528	case["hidden_states"],
1529	case["w1"],
1530	case["w2"],
1531	case["topk_weights"],
1532	case["topk_ids"],
1533	activation_method=int(ActivationMethod.SILU),
1534	quant_method=int(QuantMethod.NO),
1535	)
1536	except RuntimeError:
1537	return False
1538	return True
1539
1540
1541	@pytest.mark.skipif(
1542	not on_cdna(),
1543	reason="CDNA ROCm only",
1544	)
1545	@pytest.mark.parametrize("intermediate", [192, 224, 256])
1546	# With topk=2 and 8 experts, these token counts select block_m 32 / 64 / 128
1547	# respectively (AITER's get_block_size_M, on both 256- and 304-CU parts).
1548	@pytest.mark.parametrize("num_tokens", [1024, 2560, 5120])
1549	def test_aiter_moe_alignment_rule_holds_across_block_m(intermediate, num_tokens):
1550	"""The alignment rule must hold at every reachable ``block_m``.
1551
1552	Stage 1 checks ``NPerBlock`` and stage 2 ``KPerBlock``; both vary with
1553	``block_m``, so a size validated at one tile size says nothing about the
1554	others.
1555	"""
1556	from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1557	aiter_moe_intermediate_alignment,
1558	)
1559
1560	_assert_aiter_supported()
1561
1562	alignment = aiter_moe_intermediate_alignment(intermediate)
1563	padded = -(-intermediate // alignment) * alignment
1564
1565	assert _aiter_accepts_intermediate(intermediate, num_tokens) == (
1566	intermediate % alignment == 0
1567	)
1568	assert _aiter_accepts_intermediate(padded, num_tokens)