A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][gpt-oss] Pass GateMode.INTERLEAVE for MXFP4 W4A16 fused MoE (#44893)
Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Rohan Potdar committed
Jun 12, 2026 at 01:02 UTC
fe042382925000e5adfe530a1cc2b91d7a125fd5
2 files changed
+40
vllm/_aiter_ops.py
+24
@@ -167,6 +167,7 @@ def _rocm_aiter_fused_moe_impl(
167
output_dtype: torch.dtype | None = None,
168
hidden_pad: int = 0,
169
intermediate_pad: int = 0,
170
+ gate_mode: str = "",
171
bias1: torch.Tensor | None = None,
172
bias2: torch.Tensor | None = None,
173
moe_sorting_dispatch_policy: int = 0,
@@ -177,6 +178,10 @@ def _rocm_aiter_fused_moe_impl(
178
activation = ActivationType(activation_method)
179
quant_type = QuantType(quant_method)
180
181
+ extra_kwargs: dict = {}
182
+ if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode():
183
+ extra_kwargs["gate_mode"] = gate_mode
184
+
185
return fused_moe(
186
hidden_states,
187
w1,
@@ -198,6 +203,7 @@ def _rocm_aiter_fused_moe_impl(
203
bias1=bias1,
204
bias2=bias2,
205
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
206
+ **extra_kwargs,
207
)
208
209
@@ -219,6 +225,7 @@ def _rocm_aiter_fused_moe_fake(
225
output_dtype: torch.dtype | None = None,
226
hidden_pad: int = 0,
227
intermediate_pad: int = 0,
228
+ gate_mode: str = "",
229
bias1: torch.Tensor | None = None,
230
bias2: torch.Tensor | None = None,
231
moe_sorting_dispatch_policy: int = 0,
@@ -1804,6 +1811,21 @@ class rocm_aiter_ops:
1811
except (ImportError, ModuleNotFoundError):
1812
return False
1813
1814
+ @classmethod
1815
+ @if_aiter_supported
1816
+ @functools.cache
1817
+ def fused_moe_supports_gate_mode(cls) -> bool:
1818
+ """Probe whether the installed aiter.fused_moe accepts `gate_mode`.
1819
+
1820
+ Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14).
1821
+ Builds with older AITER must omit this argument.
1822
+ """
1823
+ import inspect
1824
+
1825
+ from aiter.fused_moe import fused_moe
1826
+
1827
+ return "gate_mode" in inspect.signature(fused_moe).parameters
1828
+
1829
@staticmethod
1830
@if_aiter_supported
1831
def register_ops_once() -> None:
@@ -2172,6 +2194,7 @@ class rocm_aiter_ops:
2194
output_dtype: torch.dtype | None = None,
2195
hidden_pad: int = 0,
2196
intermediate_pad: int = 0,
2197
+ gate_mode: str = "",
2198
bias1: torch.Tensor | None = None,
2199
bias2: torch.Tensor | None = None,
2200
moe_sorting_dispatch_policy: int = 0,
@@ -2194,6 +2217,7 @@ class rocm_aiter_ops:
2217
output_dtype,
2218
hidden_pad,
2219
intermediate_pad,
2220
+ gate_mode,
2221
bias1,
2222
bias2,
2223
moe_sorting_dispatch_policy,
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+16
@@ -351,6 +351,21 @@ def rocm_aiter_fused_experts(
351
intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
352
)
353
354
+ # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs
355
+ # for interleaved vs separated gate and up weights.
356
+ # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by
357
+ # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`,
358
+ # which always sets `is_guinterleave=True`.
359
+ # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling.
360
+ gate_mode = ""
361
+ if quant_config.use_mxfp4_w4a16:
362
+ try:
363
+ from aiter.ops.flydsl.moe_common import GateMode
364
+
365
+ gate_mode = GateMode.INTERLEAVE.value
366
+ except ImportError:
367
+ pass
368
+
369
return rocm_aiter_ops.fused_moe(
370
hidden_states,
371
w1,
@@ -369,6 +384,7 @@ def rocm_aiter_fused_experts(
384
output_dtype=output_dtype,
385
hidden_pad=hidden_pad,
386
intermediate_pad=intermediate_pad,
387
+ gate_mode=gate_mode,
388
bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
389
bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
390
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,