diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 4e3706645ef..2fdcdceae83 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -36,7 +36,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | -| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| flashinfer_nvlink_one_sided | standard | nvfp4,bf16 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 6a15d3f6168..dde60bb2c2e 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -584,8 +584,18 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): top_k: int, num_experts: int, hidden_size: int, + dispatch_dtype_bytes_per_elem: int = 0, + dispatch_has_fp8_scale: bool = True, ): - """Initialize the MoeAlltoAll workspace.""" + """Initialize the MoeAlltoAll workspace. + + dispatch_dtype_bytes_per_elem: bytes/elem for the dispatched hidden + states. Use 0 as a sentinel for sub-byte nvfp4 (0.5 B/elem); use + 1 for fp8, 2 for bf16/fp16. + dispatch_has_fp8_scale: whether a per-16-elem fp8 scale tensor is + dispatched alongside the hidden states (true for nvfp4/fp8, + false for bf16 passthrough). + """ if self.initialized: return @@ -614,9 +624,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): ep_config = MnnvlConfig( comm_backend=CustomCommunicator(self.cpu_group), ) + if dispatch_dtype_bytes_per_elem == 0: + hidden_bytes = hidden_size // 2 # nvfp4 + else: + hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem + scale_bytes = hidden_size // 16 if dispatch_has_fp8_scale else 0 total_dispatch_payload_size_per_token = ( - hidden_size // 2 # nvfp4 hidden states - + hidden_size // 16 # fp8 scaling factors + hidden_bytes + + scale_bytes + top_k * 4 # int32 topks ids + top_k * 4 # float32 topk weights ) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index fba1d4c692a..aeef2c32fd1 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -92,6 +92,7 @@ def maybe_make_prepare_finalize( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, allow_new_interface: bool = False, use_monolithic: bool = False, + defer_input_quant: bool = False, ) -> FusedMoEPrepareAndFinalize | None: # NOTE(rob): we are migrating each quant_method to hold the MK # in all cases. The allow_new_interface=False flag allow us to fall @@ -239,12 +240,26 @@ def maybe_make_prepare_finalize( max_num_tokens = ( get_current_vllm_config().scheduler_config.max_num_batched_tokens ) + if defer_input_quant or quant_config.quant_dtype is None: + # Experts (e.g. trtllm_mxfp4 with mxfp8 activations) quantize + # post-dispatch; ship bf16 tokens with no per-token scale payload. + dispatch_dtype_bytes_per_elem, dispatch_has_fp8_scale = 2, False + elif quant_config.quant_dtype == "nvfp4": + dispatch_dtype_bytes_per_elem, dispatch_has_fp8_scale = 0, True + else: + raise NotImplementedError( + "flashinfer_nvlink_one_sided dispatch only supports nvfp4, " + "bf16, and defer_input_quant paths today; got " + f"quant_dtype={quant_config.quant_dtype!r}" + ) prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize( max_num_tokens=max_num_tokens, top_k=moe.experts_per_token, num_experts=moe.num_experts, hidden_size=moe.hidden_dim, num_dispatchers=all2all_manager.world_size, + dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, + dispatch_has_fp8_scale=dispatch_has_fp8_scale, ) elif moe.use_ag_rs_all2all_kernels and allow_new_interface: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index f476d980d55..55b1f1185a8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1250,6 +1250,15 @@ def make_mxfp4_moe_kernel( """Create a FusedMoEKernel for the given MXFP4 backend.""" is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) + # Some experts (trtllm_mxfp4 with mxfp8 activations) prefer bf16 tokens + # on dispatch and quantize internally; signal this to the prepare/finalize + # so workspace + prepare path ship bf16 instead of the quant_config dtype. + from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( + TrtLlmMxfp4ExpertsBase, + ) + + defer_input_quant = issubclass(experts_cls, TrtLlmMxfp4ExpertsBase) + # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, @@ -1257,6 +1266,7 @@ def make_mxfp4_moe_kernel( routing_tables=routing_tables, allow_new_interface=True, use_monolithic=is_monolithic, + defer_input_quant=defer_input_quant, ) assert prepare_finalize is not None diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index a04ff3b8b68..40b94cfbec6 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -31,6 +31,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo num_experts: int, hidden_size: int, num_dispatchers: int = 1, + dispatch_dtype_bytes_per_elem: int = 0, + dispatch_has_fp8_scale: bool = True, ): super().__init__() self.max_num_tokens = max_num_tokens @@ -49,6 +51,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo top_k=self.top_k, num_experts=self.num_experts, hidden_size=self.hidden_size, + dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, + dispatch_has_fp8_scale=dispatch_has_fp8_scale, ) @property @@ -92,14 +96,19 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo else a1.shape[0] ) - a1q, a1q_scale = moe_kernel_quantize_input( - a1, - quant_config.a1_gscale, - quant_config.quant_dtype, - quant_config.per_act_token_quant, - quant_config.block_shape, - is_fp4_scale_swizzled=False, # delay swizzle to after comm - ) + if defer_input_quant: + # Experts (e.g. trtllm_mxfp4_moe with mxfp8 activations) will + # quantize post-dispatch. Ship bf16 tokens and skip scales. + a1q, a1q_scale = a1, None + else: + a1q, a1q_scale = moe_kernel_quantize_input( + a1, + quant_config.a1_gscale, + quant_config.quant_dtype, + quant_config.per_act_token_quant, + quant_config.block_shape, + is_fp4_scale_swizzled=False, # delay swizzle to after comm + ) payloads = [] payloads.append(a1q)