[Build] Update DeepEP pin for low-latency fabric support

Bump the repo's DeepEP pin to the revision that adds the use_fabric low-latency buffer option required for cross-node MNNVL initialization.

This keeps the runtime DeepEP low-latency fix deployable instead of relying on a local DeepEP build newer than the repo default.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Roger Wang <hey@rogerw.io>

docker/Dockerfile

@@ -330,7 +330,7 @@ WORKDIR /workspace
 # Build DeepEP wheels
 COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
 # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
-ARG DEEPEP_COMMIT_HASH=73b6ea4
+ARG DEEPEP_COMMIT_HASH=9249c25
 ARG NVSHMEM_VER
 RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p /tmp/ep_kernels_workspace/dist && \

docker/versions.json

@@ -53,7 +53,7 @@
       "default": "cuda"
     },
     "DEEPEP_COMMIT_HASH": {
-      "default": "73b6ea4"
+      "default": "9249c25"
     },
     "GIT_REPO_CHECK": {
       "default": "0"

tests/kernels/moe/test_nvfp4_backend_selection.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import patch
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    NvFp4MoeBackend,
+    select_nvfp4_moe_backend,
+)
+from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import (  # noqa: E501
+    _group_rank_batched_inputs_by_local_expert,
+    _reduce_local_expert_outputs_to_rank_batched_payload,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kNvfp4Dynamic,
+    kNvfp4Static,
+)
+
+
+class _StandardNvFp4Kernel:
+    @staticmethod
+    def is_supported_config(
+        _cls,
+        moe_config,
+        weight_key,
+        activation_key,
+        activation_format,
+    ):
+        if activation_format == mk.FusedMoEActivationFormat.Standard:
+            return True, None
+        return False, f"{activation_format.value} activation format"
+
+
+class _BatchedNvFp4Kernel:
+    @staticmethod
+    def is_supported_config(
+        _cls,
+        moe_config,
+        weight_key,
+        activation_key,
+        activation_format,
+    ):
+        if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
+            return True, None
+        return False, f"{activation_format.value} activation format"
+
+
+class _UnsupportedNvFp4Kernel:
+    @staticmethod
+    def is_supported_config(
+        _cls,
+        moe_config,
+        weight_key,
+        activation_key,
+        activation_format,
+    ):
+        return False, "unsupported"
+
+
+def _make_nvfp4_config(all2all_backend: str):
+    moe_config = make_dummy_moe_config(num_experts=8, hidden_dim=16)
+    moe_config.moe_backend = "flashinfer_cutedsl"
+    moe_config.moe_parallel_config.dp_size = 2
+    moe_config.moe_parallel_config.use_ep = True
+    moe_config.moe_parallel_config.all2all_backend = all2all_backend
+    return moe_config
+
+
+def _fake_backend_to_kernel_cls(backend: NvFp4MoeBackend):
+    if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
+        return [_StandardNvFp4Kernel]
+    if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED:
+        return [_BatchedNvFp4Kernel]
+    return [_UnsupportedNvFp4Kernel]
+
+
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.nvfp4.backend_to_kernel_cls",
+    side_effect=_fake_backend_to_kernel_cls,
+)
+def test_select_nvfp4_backend_uses_standard_cutedsl_for_standard_all2all(
+    mock_backend_to_kernel_cls,
+):
+    moe_config = _make_nvfp4_config("allgather_reducescatter")
+
+    backend, experts_cls = select_nvfp4_moe_backend(
+        moe_config,
+        kNvfp4Static,
+        kNvfp4Dynamic,
+    )
+
+    assert backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
+    assert experts_cls is _StandardNvFp4Kernel
+
+
+@patch(
+    "vllm.model_executor.layers.fused_moe.oracle.nvfp4.backend_to_kernel_cls",
+    side_effect=_fake_backend_to_kernel_cls,
+)
+def test_select_nvfp4_backend_promotes_cutedsl_for_batched_all2all(
+    mock_backend_to_kernel_cls,
+):
+    moe_config = _make_nvfp4_config("flashinfer_nvlink_one_sided")
+
+    backend, experts_cls = select_nvfp4_moe_backend(
+        moe_config,
+        kNvfp4Static,
+        kNvfp4Dynamic,
+    )
+
+    assert moe_config.moe_parallel_config.use_batched_activation_format
+    assert backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
+    assert experts_cls is _BatchedNvFp4Kernel
+
+
+def test_flashinfer_one_sided_rank_batched_regroup_and_reduce():
+    hidden_states = torch.tensor(
+        [
+            [[10.0, 11.0], [20.0, 21.0], [30.0, 31.0], [0.0, 0.0]],
+            [[40.0, 41.0], [50.0, 51.0], [0.0, 0.0], [0.0, 0.0]],
+        ]
+    )
+    hidden_scales = torch.tensor(
+        [
+            [[1.0], [2.0], [3.0], [0.0]],
+            [[4.0], [5.0], [0.0], [0.0]],
+        ]
+    )
+    topk_ids = torch.tensor(
+        [
+            [[2, 0], [3, 2], [1, 0], [0, 0]],
+            [[3, 1], [2, 3], [0, 0], [0, 0]],
+        ],
+        dtype=torch.int32,
+    )
+    topk_weights = torch.tensor(
+        [
+            [[0.70, 0.30], [0.40, 0.60], [0.50, 0.50], [0.0, 0.0]],
+            [[0.80, 0.20], [0.25, 0.75], [0.0, 0.0], [0.0, 0.0]],
+        ],
+        dtype=torch.float32,
+    )
+
+    (
+        batched_hidden_states,
+        batched_hidden_scales,
+        expert_tokens_meta,
+        dispatch_metadata,
+    ) = _group_rank_batched_inputs_by_local_expert(
+        hidden_states,
+        hidden_scales,
+        topk_ids,
+        topk_weights,
+        num_local_experts=2,
+        first_local_expert=2,
+        num_dispatchers=2,
+        runtime_max_tokens_per_rank=4,
+        source_num_tokens=[3, 2],
+    )
+
+    assert expert_tokens_meta.expert_num_tokens.tolist() == [3, 3]
+    torch.testing.assert_close(
+        batched_hidden_states[0, :3],
+        torch.tensor([[10.0, 11.0], [20.0, 21.0], [50.0, 51.0]]),
+    )
+    torch.testing.assert_close(
+        batched_hidden_states[1, :3],
+        torch.tensor([[20.0, 21.0], [40.0, 41.0], [50.0, 51.0]]),
+    )
+    assert batched_hidden_scales is not None
+    torch.testing.assert_close(
+        batched_hidden_scales[0, :3],
+        torch.tensor([[1.0], [2.0], [5.0]]),
+    )
+    torch.testing.assert_close(
+        batched_hidden_scales[1, :3],
+        torch.tensor([[2.0], [4.0], [5.0]]),
+    )
+
+    fused_expert_output = torch.tensor(
+        [
+            [[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [0.0, 0.0]],
+            [[10.0, 10.0], [20.0, 20.0], [30.0, 30.0], [0.0, 0.0]],
+        ]
+    )
+    combine_payload = _reduce_local_expert_outputs_to_rank_batched_payload(
+        fused_expert_output,
+        dispatch_metadata,
+        num_dispatchers=2,
+        runtime_max_tokens_per_rank=4,
+        apply_router_weight_on_input=False,
+    )
+
+    expected = torch.zeros((2, 4, 2))
+    expected[0, 0] = torch.tensor([0.70, 0.70])
+    expected[0, 1] = torch.tensor([5.20, 5.20])
+    expected[1, 0] = torch.tensor([16.0, 16.0])
+    expected[1, 1] = torch.tensor([23.25, 23.25])
+    torch.testing.assert_close(combine_payload, expected)

tools/ep_kernels/install_python_libraries.sh

@@ -8,7 +8,7 @@ set -ex
 #   --nvshmem-ver <ver>  NVSHMEM version 
 
 CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
-DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
+DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"9249c25"}
 NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"}  # Default supports both CUDA 12 and 13
 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
 MODE=${MODE:-install}

vllm/distributed/device_communicators/all2all.py

@@ -414,9 +414,11 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
             num_qps_per_rank=num_qps_per_rank,
         )
         if not current_platform.is_rocm():
+            use_mnnvl = envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL
             kwargs.update(
                 allow_nvlink_for_low_latency_mode=True,
-                allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
+                allow_mnnvl=use_mnnvl,
+                use_fabric=use_mnnvl,
                 explicitly_destroy=True,
             )
         return kwargs

vllm/model_executor/layers/fused_moe/config.py

@@ -990,7 +990,11 @@ class FusedMoEParallelConfig:
 
     @property
     def use_batched_activation_format(self):
-        return self.use_deepep_ll_kernels
+        return (
+            self.use_deepep_ll_kernels
+            or self.use_fi_nvl_one_sided_kernels
+            or self.use_nixl_ep_kernels
+        )
 
     @property
     def use_ag_rs_all2all_kernels(self):

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

@@ -4,7 +4,6 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
@@ -126,7 +125,7 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
 
         # We use global_num_experts due to how moe_align_block_size handles
         # expert_maps.
-        K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
+        K_dim = self.moe_config.hidden_dim if self.moe_config.hidden_dim == K * 2 else K
         output_shape = (local_num_experts, M, K_dim)
         workspace2 = (local_num_experts, M, N)
         workspace1 = output_shape
@@ -163,13 +162,12 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
         assert self.w1_scale.ndim == 3
         assert self.w2_scale.ndim == 3
 
-        input_global_scale = (
-            None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
+        use_prequantized_inputs = (
+            a1q_scale is not None and hidden_states.dtype == torch.uint8
         )
+        input_global_scale = None if use_prequantized_inputs else self.a1_gscale
         flashinfer_hidden_states = (
-            (hidden_states, a1q_scale)
-            if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
-            else hidden_states
+            (hidden_states, a1q_scale) if use_prequantized_inputs else hidden_states
         )
         flashinfer_cutedsl_moe_masked(
             hidden_states=flashinfer_hidden_states,

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

@@ -162,10 +162,9 @@ def select_nvfp4_moe_backend(
     # NOTE(rob): this is kind of a hack. We need to peak into
     # the prepare-finalize selection to determine if we are using
     # the batched or standard expert format.
-    use_batched = config.moe_parallel_config.use_deepep_ll_kernels
     activation_format = (
         mk.FusedMoEActivationFormat.BatchedExperts
-        if use_batched
+        if config.moe_parallel_config.use_batched_activation_format
         else mk.FusedMoEActivationFormat.Standard
     )
 
@@ -205,16 +204,22 @@ def select_nvfp4_moe_backend(
 
         raise ValueError(_make_log_unsupported(backend, reason))
 
+    def _resolve_requested_backend(
+        backend: NvFp4MoeBackend,
+    ) -> NvFp4MoeBackend:
+        if (
+            activation_format == mk.FusedMoEActivationFormat.BatchedExperts
+            and backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
+        ):
+            return NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
+        return backend
+
     # Handle explicit moe_backend from user.
     runner_backend = config.moe_backend
     if runner_backend != "auto":
-        requested_backend = map_nvfp4_backend(runner_backend)
-        # For batched activation format, use batched variant if available.
-        if (
-            activation_format == mk.FusedMoEActivationFormat.BatchedExperts
-            and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
-        ):
-            requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
+        requested_backend = _resolve_requested_backend(
+            map_nvfp4_backend(runner_backend)
+        )
         return _return_or_raise(
             requested_backend, config, weight_key, activation_key, activation_format
         )
@@ -227,7 +232,9 @@ def select_nvfp4_moe_backend(
 
         elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):
             # If user is explicit about backend, validate it.
-            backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]
+            backend = _resolve_requested_backend(
+                fi_2_vllm_backend_map[get_flashinfer_moe_backend()]
+            )
             return _return_or_raise(
                 backend, config, weight_key, activation_key, activation_format
             )

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

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import inspect
 from collections.abc import Callable
 
 import deep_ep
@@ -120,6 +121,20 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         # time. This setting is handled by post_init_setup.
         self.use_ue8m0_dispatch = False
 
+        # Check if DeepEP supports use_nvfp4 in low_latency_dispatch.
+        # This requires the hybrid-ep branch of DeepEP.
+        self.has_nvfp4_support = (
+            "use_nvfp4" in inspect.signature(buffer.low_latency_dispatch).parameters
+        )
+        if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH and not self.has_nvfp4_support:
+            logger.warning_once(
+                "VLLM_DEEPEPLL_NVFP4_DISPATCH=1 but DeepEP does not support "
+                "use_nvfp4 in low_latency_dispatch. Falling back to FP8/BF16 "
+                "dispatch. Install DeepEP from the hybrid-ep branch for "
+                "NvFP4 dispatch support: "
+                "https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep"
+            )
+
     def post_init_setup(self, fused_experts: mk.FusedMoEExperts):
         if not fused_experts.supports_packed_ue8m0_act_scales():
             # Early exit.
@@ -183,27 +198,29 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         assert isinstance(x, (torch.Tensor, tuple))
         q_dtype = quant_config.quant_dtype
 
-        if q_dtype == "nvfp4" and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH:
+        nvfp4_quant_path = (
+            q_dtype == "nvfp4"
+            and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
+            and self.has_nvfp4_support
+        )
+        if nvfp4_quant_path:
             logger.info_once(
-                "Since VLLM_DEEPEPLL_NVFP4_DISPATCH==1, make sure "
-                "using the hybrid-ep branch of DeepEP"
-                "(https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep)"
+                "Quantization is fused with DeepEP nvfp4 dispatch (hybrid-ep branch)"
             )
             assert isinstance(x, tuple)
             x_scales = x[1]
             x = x[0].permute(2, 0, 1)
             num_experts, max_tokens, hidden_dim_by_2 = x.shape
             hidden_dim = hidden_dim_by_2 * 2
-            logger.info_once(
-                "Quantization is fused with DeepEP nvfp4 dispatch for "
-                "FlashInfer CUTEDSL as VLLM_DEEPEPLL_NVFP4_DISPATCH==1"
-            )
         else:
             if q_dtype == "nvfp4":
                 q_dtype = None
                 logger.info_once(
-                    "Using DeepEP bfloat16 dispatch for FlashInfer CUTEDSL as "
-                    "VLLM_DEEPEPLL_NVFP4_DISPATCH==0"
+                    "Using DeepEP bfloat16 dispatch for FlashInfer CUTEDSL "
+                    "(nvfp4 dispatch %s)",
+                    "not supported by this DeepEP build"
+                    if not self.has_nvfp4_support
+                    else "disabled",
                 )
             assert isinstance(x, torch.Tensor)
             num_experts, max_tokens, hidden_dim = x.size()
@@ -260,7 +277,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
 
         use_nvfp4 = False
         nvfp4_dispatch = (
-            quant_config.quant_dtype == "nvfp4" and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
+            quant_config.quant_dtype == "nvfp4"
+            and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
+            and self.has_nvfp4_support
         )
         if nvfp4_dispatch:
             use_nvfp4 = True

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

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Sequence
+from dataclasses import dataclass
+
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -14,6 +17,156 @@ def get_local_sizes():
     return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
 
 
+@dataclass
+class _OneSidedDispatchMetadata:
+    combine_indices: list[torch.Tensor]
+    combine_weights: list[torch.Tensor]
+
+
+def _normalize_source_num_tokens(
+    source_num_tokens: Sequence[int] | None,
+    num_dispatchers: int,
+    runtime_max_tokens_per_rank: int,
+) -> list[int]:
+    if source_num_tokens is None:
+        return [runtime_max_tokens_per_rank] * num_dispatchers
+
+    normalized = [int(x) for x in source_num_tokens[:num_dispatchers]]
+    if len(normalized) < num_dispatchers:
+        normalized.extend(
+            [runtime_max_tokens_per_rank] * (num_dispatchers - len(normalized))
+        )
+    return [min(max(x, 0), runtime_max_tokens_per_rank) for x in normalized]
+
+
+def _group_rank_batched_inputs_by_local_expert(
+    hidden_states: torch.Tensor,
+    hidden_scales: torch.Tensor | None,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    *,
+    num_local_experts: int,
+    first_local_expert: int,
+    num_dispatchers: int,
+    runtime_max_tokens_per_rank: int,
+    source_num_tokens: Sequence[int] | None,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor | None,
+    mk.ExpertTokensMetadata,
+    _OneSidedDispatchMetadata,
+]:
+    hidden_dim = hidden_states.shape[-1]
+    batched_hidden_states = hidden_states.new_empty(
+        (num_local_experts, runtime_max_tokens_per_rank, hidden_dim)
+    )
+    batched_hidden_scales = (
+        None
+        if hidden_scales is None
+        else hidden_scales.new_empty(
+            (num_local_experts, runtime_max_tokens_per_rank, hidden_scales.shape[-1])
+        )
+    )
+    tokens_per_expert = torch.zeros(
+        num_local_experts, dtype=torch.int32, device=hidden_states.device
+    )
+
+    combine_indices: list[torch.Tensor] = []
+    combine_weights: list[torch.Tensor] = []
+    valid_source_num_tokens = _normalize_source_num_tokens(
+        source_num_tokens, num_dispatchers, runtime_max_tokens_per_rank
+    )
+
+    for local_expert in range(num_local_experts):
+        global_expert = first_local_expert + local_expert
+        expert_indices: list[torch.Tensor] = []
+        expert_weights: list[torch.Tensor] = []
+        cursor = 0
+
+        for dispatcher, num_tokens in enumerate(valid_source_num_tokens):
+            if num_tokens == 0:
+                continue
+
+            token_idx, topk_slot_idx = torch.where(
+                topk_ids[dispatcher, :num_tokens] == global_expert
+            )
+            rows = token_idx.numel()
+            if rows == 0:
+                continue
+
+            batched_hidden_states[local_expert, cursor : cursor + rows] = hidden_states[
+                dispatcher, token_idx
+            ]
+            if batched_hidden_scales is not None:
+                assert hidden_scales is not None
+                batched_hidden_scales[local_expert, cursor : cursor + rows] = (
+                    hidden_scales[dispatcher, token_idx]
+                )
+
+            expert_indices.append(
+                dispatcher * runtime_max_tokens_per_rank + token_idx.to(torch.int64)
+            )
+            expert_weights.append(
+                topk_weights[dispatcher, token_idx, topk_slot_idx].contiguous()
+            )
+            cursor += rows
+
+        tokens_per_expert[local_expert] = cursor
+        combine_indices.append(
+            torch.cat(expert_indices)
+            if expert_indices
+            else torch.empty(0, dtype=torch.int64, device=hidden_states.device)
+        )
+        combine_weights.append(
+            torch.cat(expert_weights)
+            if expert_weights
+            else torch.empty(0, dtype=topk_weights.dtype, device=topk_weights.device)
+        )
+
+    expert_tokens_meta = mk.ExpertTokensMetadata(
+        expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None
+    )
+    dispatch_metadata = _OneSidedDispatchMetadata(
+        combine_indices=combine_indices,
+        combine_weights=combine_weights,
+    )
+    return (
+        batched_hidden_states,
+        batched_hidden_scales,
+        expert_tokens_meta,
+        dispatch_metadata,
+    )
+
+
+def _reduce_local_expert_outputs_to_rank_batched_payload(
+    fused_expert_output: torch.Tensor,
+    dispatch_metadata: _OneSidedDispatchMetadata,
+    *,
+    num_dispatchers: int,
+    runtime_max_tokens_per_rank: int,
+    apply_router_weight_on_input: bool,
+) -> torch.Tensor:
+    hidden_dim = fused_expert_output.shape[-1]
+    combine_payload = fused_expert_output.new_zeros(
+        (num_dispatchers, runtime_max_tokens_per_rank, hidden_dim)
+    )
+    flat_payload = combine_payload.view(-1, hidden_dim)
+
+    for local_expert, linear_indices in enumerate(dispatch_metadata.combine_indices):
+        rows = linear_indices.numel()
+        if rows == 0:
+            continue
+
+        expert_output = fused_expert_output[local_expert, :rows]
+        if not apply_router_weight_on_input:
+            expert_output = expert_output * dispatch_metadata.combine_weights[
+                local_expert
+            ].to(expert_output.dtype).unsqueeze(-1)
+        flat_payload.index_add_(0, linear_indices, expert_output)
+
+    return combine_payload
+
+
 class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
     """FlashInfer implementation using the Moe AlltoAll kernel."""
 
@@ -31,8 +184,14 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         self.num_experts = num_experts
         self.hidden_size = hidden_size
         self.num_dispatchers_ = num_dispatchers
-
         self.all2all_manager = get_ep_group().device_communicator.all2all_manager
+        assert self.num_experts % self.num_dispatchers_ == 0, (
+            "flashinfer_nvlink_one_sided requires evenly sharded local experts."
+        )
+        self.num_local_experts = self.num_experts // self.num_dispatchers_
+        self.first_local_expert = self.all2all_manager.rank * self.num_local_experts
+        self.dispatch_metadata: _OneSidedDispatchMetadata | None = None
+
         self.all2all_manager.initialize(
             max_num_tokens=self.max_num_tokens,
             top_k=self.top_k,
@@ -42,10 +201,10 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
 
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
-        return mk.FusedMoEActivationFormat.Standard
+        return mk.FusedMoEActivationFormat.BatchedExperts
 
     def max_num_tokens_per_rank(self) -> int | None:
-        return None
+        return self.max_num_tokens
 
     def num_dispatchers(self) -> int:
         return self.num_dispatchers_
@@ -67,6 +226,10 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         quant_config: FusedMoEQuantConfig,
         defer_input_quant: bool = False,
     ) -> mk.PrepareResultType:
+        if defer_input_quant:
+            raise NotImplementedError(
+                f"{self.__class__.__name__} does not support defer_input_quant=True."
+            )
         if apply_router_weight_on_input:
             topk = topk_ids.size(1)
             assert topk == 1, (
@@ -104,7 +267,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         )
         if a1q_scale is not None:
             a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
-            # Apply scale interleaving only for CUTLASS (not TRT-LLM)
+            # Swizzle after dispatch when the selected MoE kernel expects it.
             if (
                 quant_config.quant_dtype == "nvfp4"
                 and quant_config.is_nvfp4_scale_swizzled
@@ -112,15 +275,32 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
                 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)
-            a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)
+            a1q_scale_recv = a1q_scale_recv.view(
+                self.num_dispatchers_,
+                self.runtime_max_tokens_per_rank,
+                self.hidden_size // 16,
+            )
         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])
-        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])
+        (
+            a1q_recv,
+            a1q_scale_recv,
+            expert_tokens_meta,
+            self.dispatch_metadata,
+        ) = _group_rank_batched_inputs_by_local_expert(
+            a1q_recv,
+            a1q_scale_recv,
+            topk_ids_recv,
+            topk_weights_recv,
+            num_local_experts=self.num_local_experts,
+            first_local_expert=self.first_local_expert,
+            num_dispatchers=self.num_dispatchers_,
+            runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
+            source_num_tokens=global_num_tokens_cpu,
+        )
 
-        return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv
+        return a1q_recv, a1q_scale_recv, expert_tokens_meta, None, None
 
     def finalize(
         self,
@@ -132,15 +312,20 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         weight_and_reduce_impl: mk.TopKWeightAndReduce,
     ) -> None:
         assert self.all2all_manager.moe_alltoall is not None
-
-        ep_size = self.all2all_manager.world_size
-        hidden_size = fused_expert_output.shape[-1]
-        fused_expert_output = fused_expert_output.view(
-            ep_size, self.runtime_max_tokens_per_rank, hidden_size
+        assert self.dispatch_metadata is not None, (
+            "flashinfer_nvlink_one_sided finalize called before prepare"
+        )
+        combine_payload = _reduce_local_expert_outputs_to_rank_batched_payload(
+            fused_expert_output,
+            self.dispatch_metadata,
+            num_dispatchers=self.num_dispatchers_,
+            runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
+            apply_router_weight_on_input=apply_router_weight_on_input,
         )
 
         combined_output = self.all2all_manager.moe_alltoall.combine(
-            payload=fused_expert_output,
+            payload=combine_payload,
             runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
         )
+        self.dispatch_metadata = None
         output.copy_(combined_output)