diff --git a/tests/kernels/moe/test_mori_moe.py b/tests/kernels/moe/test_mori_moe.py index f06afc8bd6..0bdc2d669d 100644 --- a/tests/kernels/moe/test_mori_moe.py +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 6a0a2b2358..875df4320c 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/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: diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 5c506eae12..bc478b3d0e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/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( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py index 743235bc7e..53427f2deb 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/mori.py +++ b/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