@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 5.88 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """MXFP8 MoE backend selection for the AITER FlyDSL kernel (gfx950).
4
5 GPU-free: mocks the platform (gfx950) and the aiter fused-MoE enable flag, then
6 exercises the oracle so the FlyDSL backend is auto-picked when usable (including
7 under expert parallelism, since apply() forwards the expert_map as aiter's
8 expert_mask) and skipped (native fallback) when the device is unsupported or the
9 aiter runtime is disabled.
10 """
11
12 import dataclasses
13 from unittest.mock import patch
14
15 import pytest
16
17 from vllm.platforms import current_platform
18
19 if not current_platform.is_rocm():
20 pytest.skip("This test can only run on ROCm.", allow_module_level=True)
21
22 from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402
23 from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402
24 MoEActivation,
25 )
26 from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402
27 _AITER_SWIGLU_ALPHA,
28 _AITER_SWIGLU_BETA,
29 AiterMxfp8Experts,
30 )
31 from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( # noqa: E402
32 Mxfp8EmulationTritonExperts,
33 )
34 from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402
35 FusedMoEActivationFormat,
36 )
37 from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( # noqa: E402
38 Fp8MoeBackend,
39 )
40 from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402
41 _BACKEND_NAME_MAP,
42 _SUPPORTED_BACKENDS,
43 _mxfp8_backend_to_kernel_cls,
44 _select_kernel_cls,
45 select_mxfp8_moe_backend,
46 )
47 from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
48 kMxfp8Dynamic,
49 kMxfp8Static,
50 )
51
52 from vllm._aiter_ops import rocm_aiter_ops # noqa: E402 # isort: skip
53
54 _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe"
55
56
57 def _config(ep_size: int = 1):
58 # AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and
59 # alpha/beta so is_supported_config doesn't reject the config on those grounds.
60 cfg = make_dummy_moe_config(
61 num_experts=128,
62 experts_per_token=4,
63 hidden_dim=6144,
64 activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
65 )
66 cfg = dataclasses.replace(
67 cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA
68 )
69 if ep_size != 1:
70 cfg = dataclasses.replace(
71 cfg,
72 moe_parallel_config=dataclasses.replace(
73 cfg.moe_parallel_config, ep_size=ep_size, use_ep=True
74 ),
75 )
76 return cfg
77
78
79 def _gfx950():
80 """Patch the platform so the device gate (gfx950 / MX) passes off-ROCm."""
81 return patch.multiple(
82 f"{_AITER_MOD}.current_platform",
83 is_rocm=lambda: True,
84 supports_mx=lambda: True,
85 )
86
87
88 def _aiter_moe_enabled(present: bool):
89 return patch.object(rocm_aiter_ops, "is_fused_moe_enabled", return_value=present)
90
91
92 def test_aiter_mxfp8_registered():
93 """The FlyDSL backend is auto-selectable and reachable via --moe-backend aiter."""
94 assert Fp8MoeBackend.AITER_MXFP8 in _SUPPORTED_BACKENDS
95 assert _BACKEND_NAME_MAP["aiter"] is Fp8MoeBackend.AITER_MXFP8
96 assert _mxfp8_backend_to_kernel_cls(Fp8MoeBackend.AITER_MXFP8) == [
97 AiterMxfp8Experts
98 ]
99
100
101 @pytest.mark.parametrize("ep_size", [1, 2])
102 def test_ep_supported(ep_size):
103 """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask."""
104 assert (
105 AiterMxfp8Experts._supports_parallel_config(
106 _config(ep_size).moe_parallel_config
107 )
108 is True
109 )
110
111
112 @pytest.mark.parametrize(
113 "present,ep_size,supported,reason_substr",
114 [
115 (True, 1, True, None), # gfx950 + aiter MoE + TP -> selectable
116 (True, 2, True, None), # gfx950 + aiter MoE + EP -> selectable (expert_mask)
117 (
118 False,
119 1,
120 False,
121 "does not support current device",
122 ), # disabled -> not selected
123 ],
124 )
125 def test_is_supported_config(present, ep_size, supported, reason_substr):
126 with _gfx950(), _aiter_moe_enabled(present):
127 ok, reason = AiterMxfp8Experts.is_supported_config(
128 AiterMxfp8Experts,
129 _config(ep_size),
130 kMxfp8Static,
131 kMxfp8Dynamic,
132 FusedMoEActivationFormat.Standard,
133 )
134 assert ok is supported
135 if reason_substr is not None:
136 assert reason_substr in reason
137
138
139 def test_explicit_moe_backend_aiter():
140 """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear
141 ValueError when the aiter MoE runtime is disabled."""
142 with _gfx950(), _aiter_moe_enabled(True):
143 assert (
144 _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
145 is AiterMxfp8Experts
146 )
147 assert (
148 _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(2))
149 is AiterMxfp8Experts
150 )
151 with (
152 _gfx950(),
153 _aiter_moe_enabled(False),
154 pytest.raises(ValueError, match="does not support current device"),
155 ):
156 _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
157
158
159 def test_gfx950_picks_aiter():
160 """Auto-select on real ROCm hardware with aiter MoE enabled -> FlyDSL wins."""
161 with (
162 patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True),
163 _aiter_moe_enabled(True),
164 ):
165 backend, experts_cls = select_mxfp8_moe_backend(_config())
166 assert backend is Fp8MoeBackend.AITER_MXFP8
167 assert experts_cls is AiterMxfp8Experts
168
169
170 def test_gfx942_picks_emulation():
171 """Flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
172 dot_scaled backend wins instead."""
173 with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False):
174 backend, experts_cls = select_mxfp8_moe_backend(_config())
175 assert backend is Fp8MoeBackend.EMULATION
176 assert experts_cls is Mxfp8EmulationTritonExperts