Add bf16 + defer-input-quant support to flashinfer_nvlink_one_sided all2all

The one-sided MoeAlltoAll dispatch workspace was hardcoded for nvfp4
hidden states + fp8 scales, so any other activation dtype overran the
buffer. Parameterize the workspace sizing by bytes-per-elem and whether
an fp8 scale payload is present, then route non-nvfp4 quant configs to
a bf16 dispatch (2 B/elem, no scale) via a new defer_input_quant hint.

trtllm_mxfp4 experts already advertise expects_unquantized_inputs=True
(they call mxfp8_quantize internally). Wire make_mxfp4_moe_kernel to
pass that signal into maybe_make_prepare_finalize, and have the one-
sided prepare() honor the per-call defer_input_quant flag by shipping
a1 as bf16 with no scale payload. Two-sided already handled this.

NOTE: the flashinfer moe_a2a_dispatch C++ kernel only templates top_k
in {1, 2, 4, 8}; models with other top_k (e.g. DeepSeek-V4 top_k=6)
must use flashinfer_nvlink_two_sided instead.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

docs/design/moe_kernel_features.md

@@ -36,7 +36,7 @@ th {
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | 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

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
         )

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:

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
 

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)