[ROCm][CI] Add MI355 dense NVFP4 and MoRI kernel mirrors (#58281)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 23, 2026 at 03:06 UTC
153242a314153637999eb6ebe8dd830e63433bb6
4 files changed
+405
-2
.buildkite/test-amd.yaml
+54
-2
@@ -1886,7 +1886,7 @@ steps:
1886
- vllm/config/
1887
- tests/kernels/quantization/test_rocm_skinny_gemms.py
1888
commands:
1889
- - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1889
+ - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
1890
1891
- label: ":amd: (MI300) DeepEP FP8 MoE Kernels" # TBD
1892
timeout_in_minutes: 55
@@ -4307,6 +4307,35 @@ steps:
4307
commands:
4308
- pytest -v -s kernels/moe/test_moe_layer.py
4309
4310
+- label: ":amd: (MI355) MoRI EP Numerics"
4311
+ timeout_in_minutes: 20
4312
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
4313
+ dind: false
4314
+ agent_pool: mi355_2
4315
+ num_gpus: 2
4316
+ working_dir: "/vllm-workspace/tests"
4317
+ source_file_dependencies:
4318
+ - tests/kernels/moe/test_mori_moe.py
4319
+ - tests/kernels/moe/utils.py
4320
+ - tests/kernels/moe/modular_kernel_tools/parallel_utils.py
4321
+ - tests/kernels/utils.py
4322
+ - tests/utils.py
4323
+ - tests/conftest.py
4324
+ - vllm/utils/torch_utils.py
4325
+ - vllm/model_executor/layers/fused_moe/
4326
+ - vllm/model_executor/layers/quantization/
4327
+ - vllm/distributed/
4328
+ - vllm/config/
4329
+ - vllm/forward_context.py
4330
+ - vllm/v1/worker/workspace.py
4331
+ - vllm/platforms/
4332
+ - vllm/_aiter_ops.py
4333
+ - requirements/rocm-test.txt
4334
+ optional: true
4335
+ commands:
4336
+ - amd-smi
4337
+ - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 pytest -v -s kernels/moe/test_mori_moe.py
4338
+
4339
- label: ":amd: (MI355) Quantization Kernels Shard %N"
4340
timeout_in_minutes: 120
4341
parallelism: 6
@@ -4332,7 +4361,30 @@ steps:
4361
- tests/kernels/quantization/test_rocm_skinny_gemms.py
4362
optional: true
4363
commands:
4335
- - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
4364
+ - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
4365
+
4366
+- label: ":amd: (MI355) NVFP4 Dense Emulation Shard %N"
4367
+ timeout_in_minutes: 10
4368
+ parallelism: 3
4369
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
4370
+ dind: false
4371
+ agent_pool: mi355_1
4372
+ working_dir: "/vllm-workspace/tests"
4373
+ source_file_dependencies:
4374
+ - tests/kernels/quantization/test_nvfp4_emulation.py
4375
+ - tests/kernels/quantization/nvfp4_utils.py
4376
+ - tests/conftest.py
4377
+ - vllm/model_executor/kernels/linear/
4378
+ - vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
4379
+ - vllm/config/kernel.py
4380
+ - vllm/platforms/
4381
+ - vllm/envs.py
4382
+ - vllm/utils/torch_utils.py
4383
+ - requirements/rocm-test.txt
4384
+ optional: true
4385
+ commands:
4386
+ - amd-smi
4387
+ - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
4388
4389
- label: ":amd: (MI355) FP8 MoE Kernels" # TBD
4390
timeout_in_minutes: 40
.buildkite/test_areas/kernels.yaml
+60
@@ -330,6 +330,8 @@ steps:
330
- vllm/envs.py
331
depends_on:
332
- image-build-amd
333
+ commands:
334
+ - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
335
336
- label: ":nvidia: (L4) MoE Kernels Shard %N"
337
key: kernels-moe-test
@@ -522,6 +524,32 @@ steps:
524
exit 1
525
;;
526
esac
527
+ mirror:
528
+ amd:
529
+ label: ":amd: (MI355) NVFP4 Dense Emulation Shard %N"
530
+ dind: false
531
+ device: mi355_1
532
+ num_devices: 1
533
+ timeout_in_minutes: 10
534
+ working_dir: "/vllm-workspace/tests"
535
+ depends_on:
536
+ - image-build-amd
537
+ source_file_dependencies:
538
+ - tests/kernels/quantization/test_nvfp4_emulation.py
539
+ - tests/kernels/quantization/nvfp4_utils.py
540
+ - tests/conftest.py
541
+ - vllm/model_executor/kernels/linear/
542
+ - vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
543
+ - vllm/config/kernel.py
544
+ - vllm/platforms/
545
+ - vllm/envs.py
546
+ - vllm/utils/torch_utils.py
547
+ - requirements/rocm-test.txt
548
+ # W4A4 format numerics and replay via emulation. TODO: native packed
549
+ # activation/KV-cache layouts and CUDA-only kernels need AMD backends.
550
+ commands:
551
+ - amd-smi
552
+ - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
553
554
- label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly"
555
key: b12x-linear-kernels-dgx-spark-nightly
@@ -701,6 +729,38 @@ steps:
729
commands:
730
- pytest -v -s kernels/moe/test_moe_layer.py
731
- pytest -v -s kernels/moe/test_deepep_v2_moe.py
732
+ mirror:
733
+ amd:
734
+ label: ":amd: (MI355) MoRI EP Numerics"
735
+ dind: false
736
+ device: mi355_2
737
+ num_devices: 2
738
+ timeout_in_minutes: 20
739
+ working_dir: "/vllm-workspace/tests"
740
+ depends_on:
741
+ - image-build-amd
742
+ source_file_dependencies:
743
+ - tests/kernels/moe/test_mori_moe.py
744
+ - tests/kernels/moe/utils.py
745
+ - tests/kernels/moe/modular_kernel_tools/parallel_utils.py
746
+ - tests/kernels/utils.py
747
+ - tests/utils.py
748
+ - tests/conftest.py
749
+ - vllm/utils/torch_utils.py
750
+ - vllm/model_executor/layers/fused_moe/
751
+ - vllm/model_executor/layers/quantization/
752
+ - vllm/distributed/
753
+ - vllm/config/
754
+ - vllm/forward_context.py
755
+ - vllm/v1/worker/workspace.py
756
+ - vllm/platforms/
757
+ - vllm/_aiter_ops.py
758
+ - requirements/rocm-test.txt
759
+ # Same-node MoRI/AITER BF16 and per-token/channel FP8. TODO: the
760
+ # DeepEP v2 FlashInfer backend matrix has no direct AMD implementation.
761
+ commands:
762
+ - amd-smi
763
+ - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 pytest -v -s kernels/moe/test_mori_moe.py
764
765
- label: ":amd: (MI355) FP8 MoE Kernels"
766
key: amd-fp8-moe-kernels-mi355
tests/kernels/moe/test_mori_moe.py
new
+196
@@ -0,0 +1,196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Two-rank MoRI/AITER numerical counterpart to the DeepEP v2 MoE matrix."""
4
+
5
+from dataclasses import replace
6
+
7
+import pytest
8
+import torch
9
+
10
+from tests.kernels.moe.modular_kernel_tools.parallel_utils import (
11
+ parallel_launch_with_config,
12
+)
13
+from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights
14
+from tests.kernels.utils import torch_experts
15
+from tests.utils import multi_gpu_test
16
+from vllm.config import ParallelConfig, SchedulerConfig, VllmConfig
17
+from vllm.forward_context import set_forward_context
18
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
20
+ maybe_make_prepare_finalize,
21
+)
22
+from vllm.model_executor.layers.fused_moe.config import (
23
+ FusedMoEParallelConfig,
24
+ FusedMoEQuantConfig,
25
+)
26
+from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
27
+from vllm.platforms import current_platform
28
+from vllm.utils.import_utils import has_mori
29
+from vllm.utils.torch_utils import set_random_seed
30
+from vllm.v1.worker.workspace import init_workspace_manager
31
+
32
+MNKS = [
33
+ (1, 256, 256),
34
+ (2, 256, 512),
35
+ (3, 1024, 2048),
36
+ (32, 256, 1024),
37
+ (45, 512, 2048),
38
+ (64, 1024, 1024),
39
+ (222, 1024, 2048),
40
+]
41
+
42
+
43
+def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
44
+ from vllm._aiter_ops import rocm_aiter_ops
45
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts
46
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
47
+ MoriPrepareAndFinalize,
48
+ )
49
+
50
+ rocm_aiter_ops.refresh_env_variables()
51
+ init_workspace_manager(pgi.device)
52
+ set_random_seed(47)
53
+ dtype = current_platform.fp8_dtype() if fp8 else None
54
+ (_, w1, s1, _), (_, w2, s2, _) = make_test_weights(
55
+ 32, n, k, quant_dtype=dtype, per_out_ch_quant=fp8
56
+ )
57
+ # Share global expert weights, but give each rank different token inputs.
58
+ set_random_seed(101 + pgi.rank)
59
+ x = torch.randn(m, k, device=pgi.device, dtype=torch.bfloat16) / 10
60
+ routing = torch.randn(m, 32, device=pgi.device)
61
+ # Guarantee a route to both ranks, including the one-token case.
62
+ routing[:, 0] = routing[:, 16] = float("inf")
63
+ ids = routing.topk(6, dim=-1).indices.to(torch.int32)
64
+ weights = torch.randn(m, 6, device=pgi.device, dtype=torch.float32)
65
+ parallel = replace(
66
+ FusedMoEParallelConfig.make_no_parallel(),
67
+ dp_size=2,
68
+ dp_rank=pgi.rank,
69
+ ep_size=2,
70
+ ep_rank=pgi.rank,
71
+ use_ep=True,
72
+ all2all_backend="mori_high_throughput",
73
+ )
74
+ config = replace(
75
+ make_dummy_moe_config(
76
+ num_experts=32,
77
+ num_local_experts=16,
78
+ experts_per_token=6,
79
+ hidden_dim=k,
80
+ intermediate_size=n,
81
+ max_num_tokens=256,
82
+ ),
83
+ moe_parallel_config=parallel,
84
+ )
85
+ start, end = pgi.rank * 16, (pgi.rank + 1) * 16
86
+ quant = FusedMoEQuantConfig.make(
87
+ dtype,
88
+ per_act_token_quant=fp8,
89
+ per_out_ch_quant=fp8,
90
+ w1_scale=None if s1 is None else s1[start:end].contiguous(),
91
+ w2_scale=None if s2 is None else s2[start:end].contiguous(),
92
+ )
93
+ prepare = maybe_make_prepare_finalize(config, quant)
94
+ assert isinstance(prepare, MoriPrepareAndFinalize)
95
+ assert prepare.use_fp8_dispatch == fp8
96
+ assert prepare.num_dispatchers() == 2
97
+ experts = AiterExperts(config, quant)
98
+ assert not experts.expects_unquantized_inputs
99
+ kernel = FusedMoEKernel(prepare, experts)
100
+ local_w1, local_w2 = rocm_aiter_ops.shuffle_weights(
101
+ w1[start:end].contiguous(), w2[start:end].contiguous()
102
+ )
103
+ local_w1.is_shuffled = local_w2.is_shuffled = True
104
+ expert_mask = torch.zeros(32, dtype=torch.int32, device=pgi.device)
105
+ expert_mask[start:end] = 1
106
+
107
+ def reference(route_weights):
108
+ return torch_experts(
109
+ x,
110
+ w1,
111
+ w2,
112
+ route_weights,
113
+ ids,
114
+ w1_scale=s1,
115
+ w2_scale=s2,
116
+ quant_dtype=dtype,
117
+ per_act_token_quant=fp8,
118
+ )
119
+
120
+ def run():
121
+ with set_forward_context(None, vllm_config):
122
+ return kernel.apply(
123
+ x,
124
+ local_w1,
125
+ local_w2,
126
+ weights,
127
+ ids,
128
+ activation=MoEActivation.SILU,
129
+ global_num_experts=32,
130
+ expert_map=expert_mask,
131
+ apply_router_weight_on_input=False,
132
+ )
133
+
134
+ def check(actual):
135
+ expected = reference(weights)
136
+ assert actual.shape == expected.shape
137
+ assert actual.dtype == expected.dtype
138
+ local_weights = weights * ((ids >= start) & (ids < end))
139
+ missing_remote = reference(local_weights)
140
+ norm = expected.float().norm()
141
+ assert norm > 0
142
+ norm_budget = 0.2 if fp8 else 0.06
143
+ assert (actual.float() - expected.float()).norm() / norm < norm_budget
144
+ # The norm check also rejects low-signal zero or rank-local results.
145
+ for incorrect in (torch.zeros_like(expected), missing_remote):
146
+ assert (incorrect.float() - expected.float()).norm() / norm > norm_budget
147
+ if fp8:
148
+ assert (
149
+ torch.isclose(actual, expected, atol=0.2, rtol=0.2)
150
+ .float()
151
+ .mean()
152
+ .item()
153
+ > 0.99
154
+ )
155
+ else:
156
+ torch.testing.assert_close(actual, expected, atol=0.06, rtol=0.06)
157
+ assert torch.isfinite(actual).all()
158
+
159
+ actual = run()
160
+ check(actual)
161
+ if graph:
162
+ saved = actual.clone()
163
+ run()
164
+ capture = torch.cuda.CUDAGraph()
165
+ with torch.cuda.graph(capture):
166
+ captured = run()
167
+ for factor in (-0.5, 1.5):
168
+ x.mul_(factor)
169
+ weights.mul_(-0.75)
170
+ capture.replay()
171
+ torch.accelerator.synchronize()
172
+ check(captured)
173
+ torch.testing.assert_close(actual, saved, atol=0, rtol=0)
174
+
175
+
176
+@pytest.mark.skipif(
177
+ not current_platform.is_rocm() or not has_mori(), reason="Requires ROCm MoRI"
178
+)
179
+@pytest.mark.parametrize("m,n,k", MNKS)
180
+@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8_token_channel"])
181
+@multi_gpu_test(num_gpus=2)
182
+def test_mori_moe(m, n, k, fp8):
183
+ config = VllmConfig(
184
+ parallel_config=ParallelConfig(
185
+ data_parallel_size=2,
186
+ enable_expert_parallel=True,
187
+ all2all_backend="mori_high_throughput",
188
+ ),
189
+ scheduler_config=SchedulerConfig.default_factory(max_num_batched_tokens=256),
190
+ )
191
+ env = {
192
+ "VLLM_ROCM_USE_AITER": "1",
193
+ "VLLM_ROCM_USE_AITER_MOE": "1",
194
+ "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": "0",
195
+ }
196
+ parallel_launch_with_config(2, _mori_worker, config, env, m, n, k, fp8, m == 32)
tests/kernels/quantization/test_nvfp4_emulation.py
+95
@@ -4,9 +4,21 @@ from typing import cast
4
5
import pytest
6
import torch
7
+from nvfp4_utils import (
8
+ FLOAT4_E2M1_MAX,
9
+ FLOAT8_E4M3_MAX,
10
+ break_fp4_bytes,
11
+ kE2M1ToFloat,
12
+)
13
from safetensors import safe_open
14
15
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16
+from vllm.config import VllmConfig, set_current_vllm_config
17
+from vllm.config.kernel import KernelConfig
18
+from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
19
+from vllm.model_executor.kernels.linear.nvfp4.emulation import (
20
+ EmulationNvFp4LinearKernel,
21
+)
22
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
23
from vllm.model_executor.layers.fused_moe.config import (
24
FusedMoEConfig,
@@ -34,6 +46,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
46
from vllm.platforms import current_platform
47
from vllm.transformers_utils.repo_utils import hf_api
48
from vllm.triton_utils import triton
49
+from vllm.utils.torch_utils import set_random_seed
50
51
if current_platform.is_rocm():
52
from vllm.platforms.rocm import on_gfx950
@@ -769,3 +782,85 @@ def test_nvfp4_moe_correctness(
782
atol=0.0 if on_gfx950() else 0.02,
783
rtol=0,
784
)
785
+
786
+
787
+def _dense_nvfp4_reference(
788
+ layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
789
+) -> torch.Tensor:
790
+ """Decode E2M1 nibbles and both scale levels without production helpers."""
791
+ levels = kE2M1ToFloat.cpu()
792
+ weight = break_fp4_bytes(layer.weight.cpu(), torch.float32)
793
+ weight *= (
794
+ layer.weight_scale.cpu().float() * layer.weight_global_scale.cpu()
795
+ ).repeat_interleave(16, dim=-1)
796
+ blocks = x.cpu().float().reshape(x.shape[0], -1, 16)
797
+ global_scale = layer.input_global_scale_inv.cpu()
798
+ scales = (
799
+ (blocks.abs().amax(-1, keepdim=True) * (1.0 / FLOAT4_E2M1_MAX) * global_scale)
800
+ .clamp(max=FLOAT8_E4M3_MAX)
801
+ .to(torch.float8_e4m3fn)
802
+ .float()
803
+ )
804
+ inverse = torch.where(scales == 0, 0, global_scale / scales)
805
+ scaled = (blocks * inverse).abs().clamp(max=FLOAT4_E2M1_MAX)
806
+ distances = (scaled.unsqueeze(-1) - levels).abs()
807
+ # E2M1 ties round to the even significand (even encoding).
808
+ nearest_mask = distances == distances.amin(-1, keepdim=True)
809
+ even_mask = nearest_mask & (torch.arange(8, device="cpu") % 2 == 0)
810
+ candidates = torch.where(even_mask.any(-1, keepdim=True), even_mask, nearest_mask)
811
+ nearest = candidates.to(torch.int8).argmax(-1)
812
+ activation = levels[nearest] * blocks.sign() * (scales / global_scale)
813
+ expected = activation.reshape_as(x).to(x.dtype) @ weight.to(x.dtype).T
814
+ if bias is not None:
815
+ expected += bias.cpu()
816
+ return expected
817
+
818
+
819
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm NVFP4 fallback")
820
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
821
+@pytest.mark.parametrize("m,n,k", [(1, 32, 64), (17, 128, 256)])
822
+@pytest.mark.parametrize("with_bias", [False, True])
823
+def test_nvfp4_dense_emulation(
824
+ dtype: torch.dtype, m: int, n: int, k: int, with_bias: bool
825
+) -> None:
826
+ """W4A4 fallback preserves packed scales, bias and outputs across replay."""
827
+ set_random_seed(47)
828
+ layer = torch.nn.Module()
829
+ # Every signed E2M1 code appears, including both zero encodings.
830
+ codes = (torch.randperm(n * k, device="cuda") % 16).reshape(n, k)
831
+ layer.weight = (codes[:, ::2] | (codes[:, 1::2] << 4)).to(torch.uint8)
832
+ layer.weight_scale = (
833
+ torch.arange(n * k // 16, device="cuda").reshape(n, k // 16) % 7 / 8
834
+ ).to(torch.float8_e4m3fn)
835
+ layer.weight_global_scale = torch.tensor(0.25, device="cuda")
836
+ layer.input_global_scale_inv = torch.tensor(3.0, device="cuda")
837
+ x = torch.randn(m, k, device="cuda", dtype=dtype)
838
+ x[:, :16] = 0
839
+ bias = torch.randn(n, device="cuda", dtype=dtype) if with_bias else None
840
+ with set_current_vllm_config(
841
+ VllmConfig(kernel_config=KernelConfig(linear_backend="emulation"))
842
+ ):
843
+ kernel = init_nvfp4_linear_kernel()
844
+ assert isinstance(kernel, EmulationNvFp4LinearKernel)
845
+ kernel.process_weights_after_loading(layer)
846
+ original_weight = layer.weight.clone()
847
+ original_scale = layer.weight_scale.view(torch.uint8).clone()
848
+ actual = kernel.apply_weights(layer, x, bias)
849
+ torch.testing.assert_close(
850
+ actual.cpu(), _dense_nvfp4_reference(layer, x, bias), atol=0.03, rtol=0.01
851
+ )
852
+ assert actual.abs().max() > 0.1
853
+ saved = actual.clone()
854
+ graph = torch.cuda.CUDAGraph()
855
+ with torch.cuda.graph(graph):
856
+ captured = kernel.apply_weights(layer, x, bias)
857
+ for sign in (-1, 1):
858
+ x.mul_(sign * 0.75)
859
+ graph.replay()
860
+ torch.accelerator.synchronize()
861
+ torch.testing.assert_close(
862
+ captured.cpu(), _dense_nvfp4_reference(layer, x, bias), atol=0.03, rtol=0.01
863
+ )
864
+ torch.testing.assert_close(actual, saved, atol=0, rtol=0)
865
+ assert torch.equal(layer.weight, original_weight)
866
+ assert torch.equal(layer.weight_scale.view(torch.uint8), original_scale)