[MoE] Refine FlashInfer one-sided All2All integration (#51924)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.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,bf16,mxfp8 | 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,mxfp8,fp8 | 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

tests/distributed/test_mnnvl_alltoall.py

@@ -342,6 +342,8 @@ def _one_sided_lifecycle_worker(rank, world_size):
         top_k=2,
         num_experts=world_size * 8,
         hidden_size=4096,
+        x_bytes_per_token=4096 * 2,
+        x_sf_bytes_per_token=0,
     )
 
     # Initialize
@@ -409,12 +411,12 @@ def _one_sided_workspace_grow_worker(rank, world_size):
         hidden_size=4096,
     )
     nvfp4_kwargs = dict(
-        dispatch_dtype_bytes_per_elem=0,
-        dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16,
+        x_bytes_per_token=base_kwargs["hidden_size"] // 2,
+        x_sf_bytes_per_token=base_kwargs["hidden_size"] // 16,
     )
     bf16_kwargs = dict(
-        dispatch_dtype_bytes_per_elem=2,
-        dispatch_scale_bytes_per_token=0,
+        x_bytes_per_token=base_kwargs["hidden_size"] * 2,
+        x_sf_bytes_per_token=0,
     )
 
     # First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16).
@@ -780,11 +782,12 @@ def _one_sided_data_worker(rank, world_size):
         top_k=experts_per_token,
         num_experts=num_experts,
         hidden_size=hidden_size,
-        # Account for the fp8 block-scale payload (a1q_scale: hidden//16 bytes
+        x_bytes_per_token=hidden_size // 2,
+        # Account for the fp8 block-scale payload (x_sf: hidden//16 bytes
         # per token) that is dispatched alongside the nvfp4 hidden states.
         # Without this the dispatch region is under-reserved and the combine
         # payload overflows the per-rank workspace.
-        dispatch_scale_bytes_per_token=hidden_size // 16,
+        x_sf_bytes_per_token=hidden_size // 16,
     )
     assert manager.initialized
     assert manager.moe_alltoall is not None
@@ -798,17 +801,17 @@ def _one_sided_data_worker(rank, world_size):
 
             # Create test data with raw tensors matching the nvfp4 payload
             # sizes the workspace was allocated for:
-            #   a1q: (tokens, hidden_size // 2) — nvfp4 hidden states
-            #   a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors
+            #   x: (tokens, hidden_size // 2) — nvfp4 hidden states
+            #   x_sf: (tokens, hidden_size // 16) — fp8 scaling factors
             torch.manual_seed(rank + 42)
-            a1q = torch.randint(
+            x = torch.randint(
                 0,
                 256,
                 (tokens_per_rank, hidden_size // 2),
                 device=device,
                 dtype=torch.uint8,
             )
-            a1q_scale = torch.randint(
+            x_sf = torch.randint(
                 0,
                 256,
                 (tokens_per_rank, hidden_size // 16),
@@ -830,15 +833,16 @@ def _one_sided_data_worker(rank, world_size):
             )
 
             # --- One-sided dispatch ---
-            payloads = [a1q, a1q_scale, topk_ids, topk_weights]
+            payloads = [x, x_sf, topk_ids, topk_weights]
             recv_payloads = manager.moe_alltoall.dispatch(
                 token_selected_experts=topk_ids,
                 input_payloads=payloads,
                 runtime_max_tokens_per_rank=runtime_max_tokens,
             )
             assert len(recv_payloads) == 4
-            recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads
-            assert recv_a1q.numel() > 0
+            recv_x, recv_x_sf, recv_ids, recv_weights = recv_payloads
+            assert recv_x.numel() > 0
+            assert recv_x_sf.numel() > 0
             assert recv_ids.numel() > 0
 
             # --- Round-trip exact verification ---

tests/kernels/moe/test_moe_layer.py

@@ -130,7 +130,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
     "mori_high_throughput":        {None,         "fp8", "modelopt_fp8"},
     "mori_low_latency":            {None,         "fp8", "modelopt_fp8"},
     "flashinfer_nvlink_two_sided": {None, "fp8_blocked",                 "modelopt_fp4"}, # noqa: E501
-    "flashinfer_nvlink_one_sided": {None,                                "modelopt_fp4"}, # noqa: E501
+    "flashinfer_nvlink_one_sided": {None, "fp8_blocked",                 "modelopt_fp4"}, # noqa: E501
     "deepep_low_latency":          {None, "fp8_blocked",                 "modelopt_fp4"}, # noqa: E501
     "deepep_high_throughput":      {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
     "nixl_ep":                     {None, "fp8_blocked", "modelopt_fp8"},

vllm/config/parallel.py

@@ -677,6 +677,7 @@ class ParallelConfig:
                 "allgather_reducescatter",
                 "deepep_high_throughput",
                 "deepep_low_latency",
+                "flashinfer_nvlink_one_sided",
                 "mori_high_throughput",
                 "mori_low_latency",
                 "nixl_ep",

vllm/distributed/device_communicators/all2all.py

@@ -734,17 +734,13 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
         top_k: int,
         num_experts: int,
         hidden_size: int,
-        dispatch_dtype_bytes_per_elem: int = 0,
-        dispatch_scale_bytes_per_token: int = 0,
+        x_bytes_per_token: int,
+        x_sf_bytes_per_token: int,
     ):
         """Initialize (or grow) the MoeAlltoAll workspace."""
-        if dispatch_dtype_bytes_per_elem == 0:
-            hidden_bytes = hidden_size // 2
-        else:
-            hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
         total_dispatch_payload_size_per_token = (
-            hidden_bytes
-            + dispatch_scale_bytes_per_token
+            x_bytes_per_token
+            + x_sf_bytes_per_token
             + top_k * 4  # int32 topks ids
             + top_k * 4  # float32 topk weights
         )

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -37,6 +38,52 @@ from vllm.utils.import_utils import (
     has_nixl_ep,
 )
 
+
+@dataclass(frozen=True)
+class FlashInferOneSidedDispatchLayout:
+    x_bytes_per_token: int
+    x_sf_bytes_per_token: int
+
+
+def flashinfer_one_sided_dispatch_layout(
+    hidden_dim: int, quant_config: FusedMoEQuantConfig
+) -> FlashInferOneSidedDispatchLayout:
+    """Return the one-sided activation payload layout."""
+    if quant_config.quant_dtype is None:
+        return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0)
+    if quant_config.quant_dtype == "nvfp4":
+        scale_elems = hidden_dim // 16
+        return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems)
+    if quant_config.quant_dtype == "mxfp8":
+        align = quant_config.mx_alignment
+        padded_k = (
+            ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim
+        )
+        scale_elems = padded_k // 32
+        return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems)
+    if (
+        quant_config.use_fp8_w8a8
+        and quant_config.quant_dtype == current_platform.fp8_dtype()
+        and quant_config.block_shape == [128, 128]
+    ):
+        if hidden_dim % 128 != 0:
+            raise NotImplementedError(
+                "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch "
+                f"requires hidden_dim divisible by 128; got {hidden_dim}"
+            )
+        scale_elems = hidden_dim // 128
+        scale_bytes = scale_elems * torch.float32.itemsize
+        return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes)
+    raise NotImplementedError(
+        "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
+        "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 "
+        "(quant_dtype=None) today; got "
+        f"quant_dtype={quant_config.quant_dtype!r}, "
+        f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, "
+        f"block_shape={quant_config.block_shape!r}"
+    )
+
+
 logger = init_logger(__name__)
 
 if current_platform.is_cuda_alike():
@@ -282,34 +329,17 @@ def maybe_make_prepare_finalize(
         max_num_tokens = (
             get_current_vllm_config().scheduler_config.max_num_batched_tokens
         )
-        if quant_config.quant_dtype is None:
-            dispatch_dtype_bytes_per_elem = 2
-            dispatch_scale_bytes_per_token = 0
-        elif quant_config.quant_dtype == "nvfp4":
-            dispatch_dtype_bytes_per_elem = 0
-            dispatch_scale_bytes_per_token = moe.hidden_dim // 16
-        elif quant_config.quant_dtype == "mxfp8":
-            dispatch_dtype_bytes_per_elem = 1
-            align = quant_config.mx_alignment
-            if align > 0:
-                padded_k = ((moe.hidden_dim + align - 1) // align) * align
-            else:
-                padded_k = moe.hidden_dim
-            dispatch_scale_bytes_per_token = padded_k // 32
-        else:
-            raise NotImplementedError(
-                "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
-                "and bf16 (quant_dtype=None) today; got "
-                f"quant_dtype={quant_config.quant_dtype!r}"
-            )
+        dispatch_layout = flashinfer_one_sided_dispatch_layout(
+            moe.hidden_dim, quant_config
+        )
         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_scale_bytes_per_token=dispatch_scale_bytes_per_token,
+            x_bytes_per_token=dispatch_layout.x_bytes_per_token,
+            x_sf_bytes_per_token=dispatch_layout.x_sf_bytes_per_token,
         )
 
     elif moe.use_ag_rs_all2all_kernels and allow_new_interface:

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

@@ -36,6 +36,28 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 logger = init_logger(__name__)
 
 
+def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tensor:
+    """Validate DeepSeek Blockwise FP8 tensors and return TRTLLM layout."""
+    if x.dtype != current_platform.fp8_dtype():
+        raise ValueError(
+            f"DeepSeekFp8 activations must use the platform E4M3 dtype; got {x.dtype}"
+        )
+    if x.ndim != 2 or x.shape[1] % 128 != 0:
+        raise ValueError(
+            "DeepSeekFp8 activations must be [M,K] with K divisible by 128; "
+            f"got {tuple(x.shape)}"
+        )
+    expected_shape = (x.shape[0], x.shape[1] // 128)
+    if x_sf.dtype != torch.float32 or tuple(x_sf.shape) != expected_shape:
+        raise ValueError(
+            "DeepSeekFp8 activation scales must be FP32 [M,K/128]; "
+            f"expected {expected_shape}, got dtype={x_sf.dtype}, "
+            f"shape={tuple(x_sf.shape)}"
+        )
+    # FlashInfer TRTLLM-gen consumes [K/128,M] for DeepSeekFp8/BlockMajorK.
+    return x_sf.t().contiguous()
+
+
 class TrtLlmFp8ExpertsBase:
     """
     Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic
@@ -239,6 +261,10 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
         # Pack topk ids and weights into format expected by the kernel.
         packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
 
+        if a1q_scale is None:
+            raise RuntimeError(
+                "TRT-LLM FP8 experts require precomputed activation scales"
+            )
         assert a1q_scale is not None
 
         is_mxfp8 = self.quant_config.block_shape == [1, 32]
@@ -251,7 +277,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
             use_shuffled_weight = True
             weight_layout = WeightLayout.BlockMajorK
-            hidden_states_scale = a1q_scale.t().contiguous()
+            hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale)
 
         flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
             topk_ids=packed_topk_ids,

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

@@ -30,9 +30,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         top_k: int,
         num_experts: int,
         hidden_size: int,
+        x_bytes_per_token: int,
+        x_sf_bytes_per_token: int,
         num_dispatchers: int = 1,
-        dispatch_dtype_bytes_per_elem: int = 0,
-        dispatch_scale_bytes_per_token: int = 0,
     ):
         super().__init__()
         self.max_num_tokens = max_num_tokens
@@ -40,7 +40,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         self.num_experts = num_experts
         self.hidden_size = hidden_size
         self.num_dispatchers_ = num_dispatchers
-        self.scale_elems_per_token = dispatch_scale_bytes_per_token
 
         device_communicator = get_ep_group().device_communicator
         assert device_communicator is not None
@@ -52,8 +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_scale_bytes_per_token=dispatch_scale_bytes_per_token,
+            x_bytes_per_token=x_bytes_per_token,
+            x_sf_bytes_per_token=x_sf_bytes_per_token,
         )
 
     @property
@@ -98,9 +97,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         )
 
         if defer_input_quant:
-            a1q, a1q_scale = a1, None
+            dispatch_x, dispatch_x_sf = a1, None
         else:
-            a1q, a1q_scale = moe_kernel_quantize_input(
+            dispatch_x, dispatch_x_sf = moe_kernel_quantize_input(
                 a1,
                 quant_config.a1_gscale,
                 quant_config.quant_dtype,
@@ -110,10 +109,9 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
                 mx_alignment=quant_config.mx_alignment,
             )
 
-        payloads = []
-        payloads.append(a1q)
-        if a1q_scale is not None:
-            payloads.append(a1q_scale)
+        payloads = [dispatch_x]
+        if dispatch_x_sf is not None:
+            payloads.append(dispatch_x_sf)
         topk_ids_payload_index = len(payloads)
         payloads.append(topk_ids)
         payloads.append(topk_weights)
@@ -126,23 +124,29 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
             invalid_token_expert_id=-1,  # Follow TRTLLM Pattern
             expert_id_payload_index=topk_ids_payload_index,
         )
-        if a1q_scale is not None:
-            a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
+        if dispatch_x_sf is not None:
+            recv_x, recv_x_sf, topk_ids_recv, topk_weights_recv = recv_payloads
+            x_sf_width = recv_x_sf.shape[-1]
             # Apply scale interleaving only for CUTLASS (not TRT-LLM)
             if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled:
-                a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
-                a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
-                a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
-            assert self.scale_elems_per_token > 0
-            a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token)
+                recv_x_sf = recv_x_sf.view(-1, x_sf_width)
+                recv_x_sf = recv_x_sf.view(torch.uint8)
+                recv_x_sf = nvfp4_block_scale_interleave(recv_x_sf)
+            recv_x_sf = recv_x_sf.view(-1, x_sf_width)
         else:
-            a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
-            a1q_scale_recv = None
-        a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1])
+            recv_x, topk_ids_recv, topk_weights_recv = recv_payloads
+            recv_x_sf = None
+        recv_x = recv_x.view(-1, recv_x.shape[-1])
         topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1])
         topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1])
 
-        return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv
+        return (
+            recv_x,
+            recv_x_sf,
+            None,
+            topk_ids_recv,
+            topk_weights_recv,
+        )
 
     def finalize(
         self,