[ROCm][DI] Enable MoRI FP4 dispatch for mxfp4 w4a4 moe (#59596)

Signed-off-by: lcskrishna <lollachaitanya@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/kernels/moe/test_mori_moe.py

@@ -40,11 +40,36 @@ MNKS = [
 ]
 
 
-def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
+def _mxfp4_aiter_weights(w1, w2, start, end):
+    """MXFP4-quantize experts [start, end) into AITER's W4A4 kernel layout."""
+    from aiter.ops.triton.quant import dynamic_mxfp4_quant
+
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+        Mxfp4MoeBackend,
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
+        make_mxfp4_moe_quant_config,
+    )
+
+    def quantize(w):
+        w = w[start:end]
+        q, s = dynamic_mxfp4_quant(w.flatten(0, 1))
+        return q.view(*w.shape[:2], -1), s.view(*w.shape[:2], -1)
+
+    backend = Mxfp4MoeBackend.AITER_MXFP4_MXFP4
+    (q1, s1), (q2, s2) = quantize(w1), quantize(w2)
+    q1, q2, s1, s2, _, _ = convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
+        backend, None, q1, q2, s1, s2
+    )
+    return q1, q2, make_mxfp4_moe_quant_config(backend, s1, s2)
+
+
+def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, scheme, graph):
     from vllm._aiter_ops import rocm_aiter_ops
     from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
         AiterFp8ChannelExperts,
+        AiterMxfp4Experts,
         AiterUnquantizedExperts,
+        rocm_aiter_fused_experts,
     )
     from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
         MoriPrepareAndFinalize,
@@ -53,6 +78,8 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
     rocm_aiter_ops.refresh_env_variables()
     init_workspace_manager(pgi.device)
     set_random_seed(47)
+    fp8 = scheme == "fp8_token_channel"
+    mxfp4 = scheme == "mxfp4"
     dtype = current_platform.fp8_dtype() if fp8 else None
     (_, w1, s1, _), (_, w2, s2, _) = make_test_weights(
         32, n, k, quant_dtype=dtype, per_out_ch_quant=fp8
@@ -86,29 +113,44 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
         moe_parallel_config=parallel,
     )
     start, end = pgi.rank * 16, (pgi.rank + 1) * 16
-    quant = FusedMoEQuantConfig.make(
-        dtype,
-        per_act_token_quant=fp8,
-        per_out_ch_quant=fp8,
-        w1_scale=None if s1 is None else s1[start:end].contiguous(),
-        w2_scale=None if s2 is None else s2[start:end].contiguous(),
-    )
+    if mxfp4:
+        local_w1, local_w2, quant = _mxfp4_aiter_weights(w1, w2, start, end)
+        ref_w1, ref_w2, ref_quant = _mxfp4_aiter_weights(w1, w2, 0, 32)
+    else:
+        quant = FusedMoEQuantConfig.make(
+            dtype,
+            per_act_token_quant=fp8,
+            per_out_ch_quant=fp8,
+            w1_scale=None if s1 is None else s1[start:end].contiguous(),
+            w2_scale=None if s2 is None else s2[start:end].contiguous(),
+        )
+        local_w1, local_w2 = rocm_aiter_ops.shuffle_weights(
+            w1[start:end].contiguous(), w2[start:end].contiguous()
+        )
+        local_w1.is_shuffled = local_w2.is_shuffled = True
     prepare = maybe_make_prepare_finalize(config, quant)
     assert isinstance(prepare, MoriPrepareAndFinalize)
     assert prepare.use_fp8_dispatch == fp8
+    assert prepare.use_fp4_dispatch == mxfp4
     assert prepare.num_dispatchers() == 2
-    experts_cls = AiterFp8ChannelExperts if fp8 else AiterUnquantizedExperts
+    experts_cls = {
+        "bf16": AiterUnquantizedExperts,
+        "fp8_token_channel": AiterFp8ChannelExperts,
+        "mxfp4": AiterMxfp4Experts,
+    }[scheme]
     experts = experts_cls(config, quant)
     assert not experts.expects_unquantized_inputs
     kernel = FusedMoEKernel(prepare, experts)
-    local_w1, local_w2 = rocm_aiter_ops.shuffle_weights(
-        w1[start:end].contiguous(), w2[start:end].contiguous()
-    )
-    local_w1.is_shuffled = local_w2.is_shuffled = True
     expert_mask = torch.zeros(32, dtype=torch.int32, device=pgi.device)
     expert_mask[start:end] = 1
 
     def reference(route_weights):
+        if mxfp4:
+            # AITER's non-EP MXFP4 path quantizes the BF16 input the same way
+            # as FP4 dispatch, so the BF16 tolerances below still apply.
+            return rocm_aiter_fused_experts(
+                x, ref_w1, ref_w2, route_weights, ids, config, quant_config=ref_quant
+            )
         return torch_experts(
             x,
             w1,
@@ -181,9 +223,14 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph):
     not current_platform.is_rocm() or not has_mori(), reason="Requires ROCm MoRI"
 )
 @pytest.mark.parametrize("m,n,k", MNKS)
-@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8_token_channel"])
+@pytest.mark.parametrize("scheme", ["bf16", "fp8_token_channel", "mxfp4"])
 @multi_gpu_test(num_gpus=2)
-def test_mori_moe(m, n, k, fp8):
+def test_mori_moe(m, n, k, scheme):
+    if scheme == "mxfp4":
+        from vllm.platforms.rocm import on_gfx950
+
+        if not on_gfx950():
+            pytest.skip("AITER MXFP4 W4A4 MoE requires gfx950")
     config = VllmConfig(
         parallel_config=ParallelConfig(
             data_parallel_size=2,
@@ -197,4 +244,4 @@ def test_mori_moe(m, n, k, fp8):
         "VLLM_ROCM_USE_AITER_MOE": "1",
         "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS": "0",
     }
-    parallel_launch_with_config(2, _mori_worker, config, env, m, n, k, fp8, m == 32)
+    parallel_launch_with_config(2, _mori_worker, config, env, m, n, k, scheme, m == 32)

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -329,14 +329,24 @@ def maybe_make_prepare_finalize(
     elif moe.use_mori_kernels:
         assert quant_config is not None
 
+        # MXFP4 activations are dispatched as packed FP4 with their e8m0
+        # block scales, which the experts consume without re-quantizing.
+        use_fp4_dispatch = quant_config.use_mxfp4_w4a4
         # Note: We may want to use FP8 dispatch just to reduce
         # data movement.
-        use_fp8_dispatch = (
+        use_fp8_dispatch = not use_fp4_dispatch and (
             quant_config.is_per_act_token
             or quant_config.is_block_quantized
             or quant_config.is_per_tensor
         )
-        if use_fp8_dispatch:
+        if use_fp4_dispatch:
+            # One e8m0 scale per 32-element block. token_hidden_size stays
+            # the unpacked hidden size: mori sizes its token buffers by the
+            # BF16 input type, which also covers the BF16 combine.
+            quant_dtype = torch.float4_e2m1fn_x2
+            scale_dim = moe.hidden_dim // 32
+            scale_type_size = torch.float8_e8m0fnu.itemsize
+        elif use_fp8_dispatch:
             # For PTPC (per token per channel) or per-tensor quant,
             # scale dim is 1. For 1x128 quant, scale dim is
             # hidden_dim // 128
@@ -346,18 +356,20 @@ def maybe_make_prepare_finalize(
                 if (quant_config.is_per_act_token or quant_config.is_per_tensor)
                 else moe.hidden_dim // 128
             )
+            scale_type_size = torch.float32.itemsize
         else:
             # Unquantized dispatch (e.g. AITER with defer_input_quant):
             # dispatch raw BF16/FP16 data, no scales needed.
             quant_dtype = moe.in_dtype
             scale_dim = 0
+            scale_type_size = 0
         all_to_all_args = dict(
             rank=all2all_manager.rank,
             num_ep_ranks=all2all_manager.world_size,
             quant_dtype=quant_dtype,
             token_hidden_size=moe.hidden_dim,
             scale_dim=scale_dim,
-            scale_type_size=0 if scale_dim == 0 else torch.float32.itemsize,
+            scale_type_size=scale_type_size,
             max_num_tokens_per_dp_rank=moe.max_num_tokens,
             input_dtype=moe.in_dtype,
             num_local_experts=moe.num_experts // all2all_manager.world_size,
@@ -370,6 +382,7 @@ def maybe_make_prepare_finalize(
             max_tokens_per_rank=moe.max_num_tokens,
             num_dispatchers=all2all_manager.world_size,
             use_fp8_dispatch=use_fp8_dispatch,
+            use_fp4_dispatch=use_fp4_dispatch,
         )
 
     elif moe.use_fi_nvl_two_sided_kernels:

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

@@ -352,7 +352,11 @@ def rocm_aiter_fused_experts(
         intermediate_pad = 0
         assert moe_config.hidden_dim_unpadded is not None
         assert moe_config.intermediate_size_per_partition_unpadded is not None
-        hidden_pad = hidden_states.shape[1] - moe_config.hidden_dim_unpadded
+        hidden_size = hidden_states.shape[1]
+        if hidden_states.dtype == torch.float4_e2m1fn_x2:
+            # Pre-quantized MXFP4 from MoRI dispatch, packed two per byte.
+            hidden_size *= 2
+        hidden_pad = hidden_size - moe_config.hidden_dim_unpadded
         intermediate_pad = (
             (
                 moe_config.intermediate_size_per_partition
@@ -448,7 +452,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):
 
     @property
     def expects_unquantized_inputs(self) -> bool:
-        # When paired with MoRI, the prepare/finalize handles FP8
+        # When paired with MoRI, the prepare/finalize handles FP8 or MXFP4
         # quantization during dispatch to reduce network traffic,
         # so we should not defer input quantization.
         # Otherwise, AITER fused MoE kernels handle input quantization
@@ -516,7 +520,12 @@ class AiterExperts(mk.FusedMoEExpertsModular):
         # Workspaces are managed internally by AITER.
         workspace1 = (0,)
         workspace2 = (0,)
-        output = (M, K)
+        # MoRI dispatches MXFP4 activations packed two per byte, so K is
+        # half the hidden size there.
+        packed_fp4_input = (
+            self.moe_config.use_mori_kernels and self.quant_config.use_mxfp4_w4a4
+        )
+        output = (M, K * 2 if packed_fp4_input else K)
         return (workspace1, workspace2, output)
 
     def apply(

vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py

@@ -21,12 +21,15 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         max_tokens_per_rank: int,
         num_dispatchers: int,
         use_fp8_dispatch: bool = False,
+        use_fp4_dispatch: bool = False,
     ):
         super().__init__()
+        assert not (use_fp8_dispatch and use_fp4_dispatch)
         self.mori_op = mori_op
         self.num_dispatchers_ = num_dispatchers
         self.max_tokens_per_rank = max_tokens_per_rank
         self.use_fp8_dispatch = use_fp8_dispatch
+        self.use_fp4_dispatch = use_fp4_dispatch
         self._dispatch_topk_ids: torch.Tensor | None = None
 
     @property
@@ -73,7 +76,7 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         )
         scale = None
         # When defer_input_quant is True, the expert kernel handles
-        # quantization internally, so skip FP8 dispatch quantization.
+        # quantization internally, so skip dispatch quantization.
         if self.use_fp8_dispatch and not defer_input_quant:
             from aiter import QuantType, get_hip_quant
 
@@ -92,6 +95,14 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
                 )
                 # mori expects one scale slot per token; broadcast.
                 scale = scale.expand(a1.shape[0], 1).contiguous()
+        elif self.use_fp4_dispatch and not defer_input_quant:
+            from aiter import QuantType, get_hip_quant
+
+            # Packed MXFP4 [M, K // 2] with e8m0 scales [M, K // 32]. The
+            # scales stay unshuffled: AITER fused_moe sorts them itself
+            # when it receives pre-quantized FP4 input.
+            quant_func = get_hip_quant(QuantType.per_1x32)
+            a1, scale = quant_func(a1, shuffle=False)
 
         # mori's combine() reduces over this rank's own [num_tokens, topk]
         # routing. The modular kernel rebinds topk_ids to the dispatched ids