[ROCm][gpt-oss] Pass GateMode.INTERLEAVE for MXFP4 W4A16 fused MoE

vllm/_aiter_ops.py

@@ -167,6 +167,7 @@ def _rocm_aiter_fused_moe_impl(
     output_dtype: torch.dtype | None = None,
     hidden_pad: int = 0,
     intermediate_pad: int = 0,
+    gate_mode: str = "",
     bias1: torch.Tensor | None = None,
     bias2: torch.Tensor | None = None,
     moe_sorting_dispatch_policy: int = 0,
@@ -177,6 +178,10 @@ def _rocm_aiter_fused_moe_impl(
     activation = ActivationType(activation_method)
     quant_type = QuantType(quant_method)
 
+    extra_kwargs: dict = {}
+    if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode():
+        extra_kwargs["gate_mode"] = gate_mode
+
     return fused_moe(
         hidden_states,
         w1,
@@ -198,6 +203,7 @@ def _rocm_aiter_fused_moe_impl(
         bias1=bias1,
         bias2=bias2,
         moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
+        **extra_kwargs,
     )
 
 
@@ -219,6 +225,7 @@ def _rocm_aiter_fused_moe_fake(
     output_dtype: torch.dtype | None = None,
     hidden_pad: int = 0,
     intermediate_pad: int = 0,
+    gate_mode: str = "",
     bias1: torch.Tensor | None = None,
     bias2: torch.Tensor | None = None,
     moe_sorting_dispatch_policy: int = 0,
@@ -1804,6 +1811,21 @@ def are_gdn_triton_kernels_available(cls) -> bool:
         except (ImportError, ModuleNotFoundError):
             return False
 
+    @classmethod
+    @if_aiter_supported
+    @functools.cache
+    def fused_moe_supports_gate_mode(cls) -> bool:
+        """Probe whether the installed aiter.fused_moe accepts `gate_mode`.
+
+        Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14).
+        Builds with older AITER must omit this argument.
+        """
+        import inspect
+
+        from aiter.fused_moe import fused_moe
+
+        return "gate_mode" in inspect.signature(fused_moe).parameters
+
     @staticmethod
     @if_aiter_supported
     def register_ops_once() -> None:
@@ -2172,6 +2194,7 @@ def fused_moe(
         output_dtype: torch.dtype | None = None,
         hidden_pad: int = 0,
         intermediate_pad: int = 0,
+        gate_mode: str = "",
         bias1: torch.Tensor | None = None,
         bias2: torch.Tensor | None = None,
         moe_sorting_dispatch_policy: int = 0,
@@ -2194,6 +2217,7 @@ def fused_moe(
             output_dtype,
             hidden_pad,
             intermediate_pad,
+            gate_mode,
             bias1,
             bias2,
             moe_sorting_dispatch_policy,

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

@@ -351,6 +351,21 @@ def rocm_aiter_fused_experts(
             intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
         )
 
+        # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs
+        # for interleaved vs separated gate and up weights.
+        # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by
+        # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`,
+        # which always sets `is_guinterleave=True`.
+        # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling.
+        gate_mode = ""
+        if quant_config.use_mxfp4_w4a16:
+            try:
+                from aiter.ops.flydsl.moe_common import GateMode
+
+                gate_mode = GateMode.INTERLEAVE.value
+            except ImportError:
+                pass
+
         return rocm_aiter_ops.fused_moe(
             hidden_states,
             w1,
@@ -369,6 +384,7 @@ def rocm_aiter_fused_experts(
             output_dtype=output_dtype,
             hidden_pad=hidden_pad,
             intermediate_pad=intermediate_pad,
+            gate_mode=gate_mode,
             bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
             bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
             moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,