Fix dtype mismatch in topk_softplus_sqrt for DeepEP backends

DeepEP requires int64 topk indices, but DeepSeek-V4's hash MoE creates
input_ids and hash_indices_table as int32. The CUDA kernel dispatches on
topk_ids dtype and assumes all index tensors match, causing a crash:
"expected scalar type Long but found Int".

Cast input_ids and hash_indices_table to match indices_type before
calling the kernel. The hash table cast is cached since it's static.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

tests/distributed/test_mnnvl_alltoall.py

@@ -19,6 +19,7 @@ from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_one_sided,
     has_flashinfer_nvlink_two_sided,
 )
+from vllm.utils.import_utils import has_deep_ep_v2
 from vllm.utils.network_utils import get_open_port
 
 from ..utils import init_test_distributed_environment
@@ -194,6 +195,10 @@ requires_ptrace = pytest.mark.skipif(
     not _has_sys_ptrace(),
     reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
 )
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
+)
 
 # NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
 # their own @requires_two_sided / @requires_one_sided decorators, and
@@ -772,3 +777,76 @@ def _one_sided_data_worker(rank, world_size):
 def test_one_sided_dispatch_combine(world_size):
     """Test FlashInfer one-sided dispatch/combine with actual data flow."""
     _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
+
+
+# ---------------------------------------------------------------------------
+# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
+# ---------------------------------------------------------------------------
+#
+# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
+# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
+#
+# Uses EP group because the DeepEP v2 manager is constructed with an
+# EP-scoped communicator in production. With tp=world_size the EP group
+# spans all ranks.
+# ---------------------------------------------------------------------------
+
+
+def _deepep_v2_lifecycle_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        DeepEPV2All2AllManager,
+    )
+
+    cpu_group = get_ep_group().cpu_group
+    manager = DeepEPV2All2AllManager(cpu_group)
+
+    assert manager.rank == rank
+    assert manager.world_size == world_size
+    assert manager._num_sms is None
+
+    hidden_size = 7168
+    num_experts = world_size * 32
+    num_topk = 8
+    max_tokens = 256
+
+    handle_kwargs = dict(
+        num_max_tokens_per_rank=max_tokens,
+        hidden=hidden_size,
+        num_topk=num_topk,
+        num_experts=num_experts,
+        use_fp8_dispatch=False,
+    )
+
+    handle = manager.get_handle(handle_kwargs)
+    assert handle is not None
+    assert manager._num_sms is not None
+    assert manager._num_sms > 0
+
+    torch.distributed.barrier()
+
+    # get_handle again with same args should return cached handle
+    handle2 = manager.get_handle(dict(handle_kwargs))
+    assert handle2 is handle
+
+    torch.distributed.barrier()
+
+    # Destroy clears the cache
+    manager.destroy()
+    assert len(manager.handle_cache._cache) == 0
+
+    torch.distributed.barrier()
+
+    # Re-create after destroy
+    handle3 = manager.get_handle(dict(handle_kwargs))
+    assert handle3 is not None
+
+    torch.distributed.barrier()
+    manager.destroy()
+
+
+@requires_multi_gpu
+@requires_deep_ep_v2
+@pytest.mark.parametrize("world_size", [2])
+def test_deepep_v2_manager_lifecycle(world_size):
+    """Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
+    _spawn_workers(_deepep_v2_lifecycle_worker, world_size)

tests/kernels/moe/modular_kernel_tools/common.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from types import SimpleNamespace
 from typing import Any
 
 import torch
@@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.utils.import_utils import (
     has_aiter,
     has_deep_ep,
+    has_deep_ep_v2,
     has_deep_gemm,
     has_mori,
 )
@@ -150,6 +152,10 @@ class Config:
         make env data for vllm launch.
         """
         vllm_config = VllmConfig()
+        vllm_config.model_config = SimpleNamespace(
+            enforce_eager=True,
+            is_moe=True,
+        )
         vllm_config.parallel_config.data_parallel_size = self.world_size
         vllm_config.parallel_config.enable_expert_parallel = True
 
@@ -243,6 +249,10 @@ class Config:
             or info.backend == "deepep_low_latency"
         )
 
+    def needs_deep_ep_v2(self):
+        info = prepare_finalize_info(self.prepare_finalize_type)
+        return info.backend == "deepep_v2"
+
     def needs_aiter(self):
         info = expert_info(self.fused_experts_type)
         return info.needs_aiter
@@ -319,6 +329,8 @@ class Config:
         # Check dependencies (turn into asserts?)
         if self.needs_deep_ep() and not has_deep_ep():
             return False, "Needs DeepEP, but DeepEP not available."
+        if self.needs_deep_ep_v2() and not has_deep_ep_v2():
+            return False, "Needs DeepEP v2, but DeepEP v2 not available."
         if self.needs_deep_gemm() and not has_deep_gemm():
             return False, "Needs DeepGEMM, but DeepGEMM not available."
         if self.needs_aiter() and not has_aiter():  # noqa: SIM103
@@ -653,6 +665,54 @@ def make_modular_kernel(
     return modular_kernel
 
 
+def _maybe_convert_weights_for_experts(
+    config: Config,
+    rank_weights: WeightTensors,
+) -> WeightTensors:
+    """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK)."""
+    from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+        Fp8MoeBackend,
+        convert_to_fp8_moe_kernel_format,
+    )
+
+    fe_type = config.fused_experts_type
+    fe_name = getattr(fe_type, "__name__", "")
+
+    backend: Fp8MoeBackend | None = None
+    if fe_name == "TrtLlmFp8ExpertsModular":
+        backend = Fp8MoeBackend.FLASHINFER_TRTLLM
+    elif fe_name == "FlashInferExperts":
+        backend = Fp8MoeBackend.FLASHINFER_CUTLASS
+
+    if backend is None or not rank_weights.is_quantized():
+        return rank_weights
+
+    mock_layer = SimpleNamespace(
+        weight_block_size=config.quant_block_shape,
+        moe_config=SimpleNamespace(
+            is_act_and_mul=True,
+            intermediate_size_per_partition=config.N,
+        ),
+        activation=SimpleNamespace(is_gated=True),
+    )
+
+    w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
+        fp8_backend=backend,
+        layer=mock_layer,
+        w13=rank_weights.w1,
+        w2=rank_weights.w2,
+        w13_scale=rank_weights.w1_scale,
+        w2_scale=rank_weights.w2_scale,
+        w13_input_scale=None,
+        w2_input_scale=None,
+    )
+
+    return WeightTensors(
+        w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale,
+        w1_gs=rank_weights.w1_gs, w2_gs=rank_weights.w2_gs,
+    )
+
+
 def run_modular_kernel(
     pgi: ProcessGroupInfo,
     vllm_config: VllmConfig,
@@ -665,6 +725,7 @@ def run_modular_kernel(
 
     # weights for rank
     rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
+    rank_weights = _maybe_convert_weights_for_experts(config, rank_weights)
 
     if config.quant_dtype == "nvfp4":
         gscale = _make_gscale(config.num_local_experts)
@@ -710,6 +771,8 @@ def run_modular_kernel(
         [num_tokens] * config.world_size, device="cuda", dtype=torch.int
     )
 
+    torch.distributed.barrier()
+
     with set_forward_context(
         None,
         vllm_config,

tests/kernels/moe/modular_kernel_tools/mk_objects.py

@@ -37,9 +37,11 @@ from vllm.utils.flashinfer import (
     has_flashinfer_cutlass_fused_moe,
     has_flashinfer_nvlink_one_sided,
 )
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 from vllm.utils.import_utils import (
     has_aiter,
     has_deep_ep,
+    has_deep_ep_v2,
     has_deep_gemm,
     has_mori,
 )
@@ -222,6 +224,19 @@ if has_deep_ep() and not current_platform.has_device_capability(100):
         backend="deepep_low_latency",
     )
 
+if has_deep_ep_v2() and current_platform.has_device_capability(100):
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
+        DeepEPV2PrepareAndFinalize,
+    )
+
+    register_prepare_and_finalize(
+        DeepEPV2PrepareAndFinalize,
+        standard_format,
+        common_float_types,
+        blocked_quantization_support=True,
+        backend="deepep_v2",
+    )
+
 if has_mori():
     from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
         MoriPrepareAndFinalize,
@@ -297,6 +312,19 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
         supports_expert_map=True,
     )
 
+if has_flashinfer_trtllm_fused_moe() and current_platform.has_device_capability(100):
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
+        TrtLlmFp8ExpertsModular,
+    )
+
+    register_experts(
+        TrtLlmFp8ExpertsModular,
+        standard_format,
+        fp8_types,
+        blocked_quantization_support=True,
+        supports_expert_map=True,
+    )
+
 if has_aiter():
     from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
         AiterExperts,

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

@@ -96,7 +96,7 @@ def _worker_parallel_launch(
     if vllm_config is not None:
         cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank)
 
-    try:
+    def _run_worker():
         worker(
             ProcessGroupInfo(
                 world_size=world_size,
@@ -111,6 +111,13 @@ def _worker_parallel_launch(
             *args,
             **worker_kwargs,
         )
+
+    try:
+        if vllm_config is not None:
+            with set_current_vllm_config(vllm_config):
+                _run_worker()
+        else:
+            _run_worker()
     except Exception as ex:
         print(ex)
         traceback.print_exc()

tests/kernels/moe/parallel_utils.py

@@ -15,7 +15,7 @@ from torch.distributed import ProcessGroup
 from torch.multiprocessing import spawn  # pyright: ignore[reportPrivateImportUsage]
 from typing_extensions import ParamSpec
 
-from vllm.utils.import_utils import has_deep_ep
+from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
 from vllm.utils.network_utils import get_open_port
 
 if has_deep_ep():
@@ -26,6 +26,11 @@ if has_deep_ep():
         DeepEPLLPrepareAndFinalize,
     )
 
+if has_deep_ep_v2():
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
+        DeepEPV2PrepareAndFinalize,
+    )
+
 ## Parallel Processes Utils
 
 P = ParamSpec("P")
@@ -55,11 +60,10 @@ def _worker_parallel_launch(
     torch.accelerator.set_device_index(local_rank)
     device = torch.device("cuda", local_rank)
     torch.distributed.init_process_group(
-        backend="cpu:gloo,cuda:nccl",
+        backend="nccl",
         init_method=init_method,
         rank=rank,
         world_size=world_size,
-        device_id=device,
     )
     barrier = torch.tensor([rank], device=device)
     torch.distributed.all_reduce(barrier)
@@ -200,3 +204,42 @@ def make_deepep_a2a(
 
     assert deepep_ll_args is not None
     return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape)
+
+
+@dataclasses.dataclass
+class DeepEPV2Args:
+    num_local_experts: int
+    num_experts: int
+    num_topk: int
+    hidden_size: int
+    max_tokens_per_rank: int
+    use_fp8_dispatch: bool
+
+
+def make_deepep_v2_a2a(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    v2_args: DeepEPV2Args,
+    use_cudagraph: bool = False,
+):
+    import deep_ep
+
+    buffer = deep_ep.ElasticBuffer(
+        group=pg,
+        num_max_tokens_per_rank=v2_args.max_tokens_per_rank,
+        hidden=v2_args.hidden_size,
+        num_topk=v2_args.num_topk,
+        use_fp8_dispatch=v2_args.use_fp8_dispatch,
+        explicitly_destroy=True,
+    )
+    return DeepEPV2PrepareAndFinalize(
+        buffer=buffer,
+        num_dispatchers=pgi.world_size,
+        dp_size=dp_size,
+        rank_expert_offset=pgi.rank * v2_args.num_local_experts,
+        num_experts=v2_args.num_experts,
+        num_topk=v2_args.num_topk,
+        use_fp8_dispatch=v2_args.use_fp8_dispatch,
+        use_cudagraph=use_cudagraph,
+    )

tests/kernels/moe/test_deepep_v2_moe.py

@@ -0,0 +1,580 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test DeepEP v2 (ElasticBuffer) dispatch-combine logic.
+Compares against a pure-PyTorch reference MoE implementation.
+"""
+
+import dataclasses
+
+import pytest
+import torch.distributed
+from torch.distributed import ProcessGroup
+
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm import _custom_ops as ops
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import TritonExperts
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.utils.import_utils import has_deep_ep_v2
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import init_workspace_manager
+
+from ...utils import multi_gpu_test
+from .parallel_utils import ProcessGroupInfo, parallel_launch
+
+if has_deep_ep_v2():
+    from .parallel_utils import DeepEPV2Args, make_deepep_v2_a2a
+
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="Requires DeepEP v2 (ElasticBuffer)",
+)
+
+
+def make_weights(
+    e, n, k, dtype
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if dtype in [torch.float16, torch.bfloat16]:
+        w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
+        w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+        return w1, w2, None, None
+
+    assert dtype == torch.float8_e4m3fn
+    w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16)
+    w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16)
+
+    n_b_scales = 2 * n
+    k_b_scales = k
+    w1_q = torch.empty_like(w1, dtype=dtype)
+    w2_q = torch.empty_like(w2, dtype=dtype)
+    w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32)
+    w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32)
+    for expert in range(e):
+        w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
+            w1[expert], use_per_token_if_dynamic=True
+        )
+        w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
+            w2[expert], use_per_token_if_dynamic=True
+        )
+    return w1_q, w2_q, w1_scale, w2_scale
+
+
+@dataclasses.dataclass
+class TestConfig:
+    dtype: torch.dtype
+    topk: int
+    m: int
+    k: int
+    n: int
+    num_experts: int
+
+
+@dataclasses.dataclass
+class TestTensors:
+    rank_tokens: torch.Tensor
+    rank_token_scales: torch.Tensor | None
+    topk: torch.Tensor
+    topk_weights: torch.Tensor
+    config: TestConfig
+
+    @staticmethod
+    def make(config: TestConfig) -> "TestTensors":
+        assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
+        token_dtype = (
+            torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype
+        )
+        rank_tokens = (
+            torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10
+        )
+
+        topk = torch.stack(
+            [
+                torch.randperm(config.num_experts, device="cuda")[: config.topk]
+                for _ in range(config.m)
+            ]
+        ).to(dtype=torch.int64)
+        topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda")
+        return TestTensors(
+            rank_tokens=rank_tokens,
+            rank_token_scales=None,
+            topk=topk,
+            topk_weights=topk_weights,
+            config=config,
+        )
+
+
+def make_modular_kernel(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    hidden_size: int,
+    num_experts: int,
+    num_local_experts: int,
+    topk: int,
+    q_dtype: torch.dtype | None,
+    use_fp8_dispatch: bool,
+    quant_config: FusedMoEQuantConfig,
+    use_cudagraph: bool = False,
+) -> FusedMoEKernel:
+    v2_args = DeepEPV2Args(
+        num_local_experts=num_local_experts,
+        num_experts=num_experts,
+        num_topk=topk,
+        hidden_size=hidden_size,
+        max_tokens_per_rank=8192,
+        use_fp8_dispatch=use_fp8_dispatch,
+    )
+
+    a2a = make_deepep_v2_a2a(
+        pg=pg,
+        pgi=pgi,
+        dp_size=dp_size,
+        v2_args=v2_args,
+        use_cudagraph=use_cudagraph,
+    )
+
+    moe_config = make_dummy_moe_config(
+        num_experts=num_local_experts,
+        experts_per_token=topk,
+        hidden_dim=hidden_size,
+    )
+
+    fused_experts = TritonExperts(
+        moe_config=moe_config,
+        quant_config=quant_config,
+    )
+
+    mk = FusedMoEKernel(
+        prepare_finalize=a2a,
+        fused_experts=fused_experts,
+        inplace=False,
+    )
+    return mk
+
+
+def deepep_v2_moe_impl(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    num_experts: int,
+    topk: int,
+    use_fp8_dispatch: bool,
+    per_act_token_quant: bool,
+) -> torch.Tensor:
+    num_local_experts = w1.size(0)
+
+    def build_expert_map():
+        expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32)
+        s = pgi.rank * num_local_experts
+        e = s + num_local_experts
+        expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
+        device = torch.accelerator.current_device_index()
+        return expert_map.to(device=device, dtype=torch.int32)
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    q_dtype = torch.float8_e4m3fn if is_quantized else None
+
+    quant_config = FusedMoEQuantConfig.make(
+        q_dtype,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        per_act_token_quant=per_act_token_quant,
+        a1_scale=test_tensors.rank_token_scales,
+    )
+
+    hidden_size = test_tensors.rank_tokens.size(1)
+
+    mk: FusedMoEKernel = make_modular_kernel(
+        pg,
+        pgi,
+        dp_size,
+        hidden_size,
+        num_experts,
+        num_local_experts,
+        topk,
+        q_dtype,
+        use_fp8_dispatch,
+        quant_config,
+    )
+
+    out = mk.apply(
+        hidden_states=test_tensors.rank_tokens,
+        w1=w1,
+        w2=w2,
+        topk_weights=test_tensors.topk_weights,
+        topk_ids=test_tensors.topk,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=build_expert_map(),
+        apply_router_weight_on_input=False,
+    )
+
+    return out
+
+
+def torch_moe_impl(
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+):
+    a, topk_ids, topk_weights = (
+        test_tensors.rank_tokens,
+        test_tensors.topk,
+        test_tensors.topk_weights,
+    )
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    a_dtype = a.dtype
+    if is_quantized:
+        w1 = w1.to(dtype=torch.float32) * w1_scale
+        w2 = w2.to(dtype=torch.float32) * w2_scale
+        a = a.to(dtype=torch.float32)
+
+    m, _ = a.shape
+    topk = topk_ids.size(1)
+    out = torch.zeros_like(a)
+
+    for i in range(m):
+        a_i = a[i]
+        o_i = out[i]
+        for j in range(topk):
+            e = topk_ids[i][j]
+            e_w = topk_weights[i][j]
+            w1_e = w1[e]
+            w2_e = w2[e]
+            o_i += (
+                SiluAndMul()(a_i @ w1_e.transpose(0, 1)) @ w2_e.transpose(0, 1)
+            ) * e_w
+
+    if is_quantized:
+        out = out.to(dtype=a_dtype)
+
+    return out
+
+
+def _deep_ep_v2_moe(
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    use_fp8_dispatch: bool,
+    per_act_token_quant: bool,
+):
+    device = torch.device(f"cuda:{pgi.local_rank}")
+    init_workspace_manager(device)
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    device_idx = torch.accelerator.current_device_index()
+    w1 = w1.to(device=device_idx)
+    w2 = w2.to(device=device_idx)
+    if is_quantized:
+        assert w1_scale is not None and w2_scale is not None
+        w1_scale = w1_scale.to(device=device_idx)
+        w2_scale = w2_scale.to(device=device_idx)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    test_tensors = TestTensors.make(config)
+
+    with set_current_vllm_config(VllmConfig()):
+        # Reference
+        torch_combined = torch_moe_impl(
+            test_tensors,
+            w1,
+            w2,
+            w1_scale,
+            w2_scale,
+            per_act_token_quant,
+        )
+
+        # Splice experts for this rank
+        num_local_experts = config.num_experts // pgi.world_size
+        e_start = num_local_experts * pgi.rank
+        e_end = e_start + num_local_experts
+        w1_ep = w1[e_start:e_end]
+        w2_ep = w2[e_start:e_end]
+
+        w1_scale_ep, w2_scale_ep = None, None
+        if is_quantized:
+            w1_scale_ep = w1_scale[e_start:e_end]  # type: ignore
+            w2_scale_ep = w2_scale[e_start:e_end]  # type: ignore
+
+        deepep_combined = deepep_v2_moe_impl(
+            pg,
+            pgi,
+            dp_size,
+            test_tensors,
+            w1_ep,
+            w2_ep,
+            w1_scale_ep,
+            w2_scale_ep,
+            config.num_experts,
+            config.topk,
+            use_fp8_dispatch,
+            per_act_token_quant,
+        )
+
+    torch.testing.assert_close(
+        torch_combined,
+        deepep_combined,
+        atol=6e-2,
+        rtol=6e-2,
+    )
+
+
+MNKs = [
+    (1, 256, 256),
+    (2, 256, 512),
+    (3, 1024, 2048),
+    (32, 256, 1024),
+    (45, 512, 2048),
+    (64, 1024, 1024),
+    (222, 1024, 2048),
+]
+
+DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("m,n,k", MNKs)
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [6])
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
+@multi_gpu_test(num_gpus=2)
+@requires_deep_ep_v2
+def test_deep_ep_v2_moe(
+    dtype: torch.dtype,
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    world_dp_size: tuple[int, int],
+    workspace_init,
+):
+    per_act_token_quant = False
+    use_fp8_dispatch = False
+
+    set_random_seed(7)
+    world_size, dp_size = world_dp_size
+    config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts)
+
+    w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
+
+    parallel_launch(
+        world_size,
+        _deep_ep_v2_moe,
+        dp_size,
+        config,
+        w1,
+        w2,
+        w1_scale,
+        w2_scale,
+        use_fp8_dispatch,
+        per_act_token_quant,
+    )
+
+
+def _deep_ep_v2_moe_cudagraph(
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+):
+    """Worker function: verify DeepEP v2 + TrtLLM FP8 with do_expand=False."""
+    import tempfile
+
+    from tests.kernels.moe.test_moe_layer import make_fused_moe_layer
+    from vllm.distributed import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    device = torch.device(f"cuda:{pgi.local_rank}")
+    init_workspace_manager(device)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    test_tensors = TestTensors.make(config)
+    num_local_experts = config.num_experts // pgi.world_size
+    hidden_size = config.k
+
+    # Create FP8 weights directly, then dequantize for bf16 reference.
+    w1_fp8 = torch.randn(
+        (config.num_experts, 2 * config.n, config.k),
+        device="cuda", dtype=torch.bfloat16,
+    ).to(torch.float8_e4m3fn)
+    w2_fp8 = torch.randn(
+        (config.num_experts, config.k, config.n),
+        device="cuda", dtype=torch.bfloat16,
+    ).to(torch.float8_e4m3fn)
+    w1_ref = w1_fp8.to(torch.bfloat16)
+    w2_ref = w2_fp8.to(torch.bfloat16)
+
+    from vllm.config import KernelConfig
+    vllm_cfg = VllmConfig()
+    vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm")
+
+    with set_current_vllm_config(vllm_cfg):
+        # Initialize vLLM parallel state (needed by FusedMoE layer)
+        temp_file = tempfile.mktemp()
+        init_distributed_environment(
+            world_size=pgi.world_size,
+            rank=pgi.rank,
+            distributed_init_method=f"file://{temp_file}",
+            local_rank=pgi.local_rank,
+            backend="nccl",
+        )
+        initialize_model_parallel(tensor_model_parallel_size=1)
+        # Reference MoE using dequantized bf16 weights
+        torch_combined = torch_moe_impl(
+            test_tensors, w1_ref, w2_ref, None, None, False,
+        )
+
+        # Use the production pipeline: make_fused_moe_layer creates
+        # a FusedMoE layer, quantizes weights, runs
+        # process_weights_after_loading (TrtLLM W31 swap + BlockMajorK
+        # shuffle), and selects the kernel.
+        # Quantize weights using production helper, EP-slice, then
+        # convert to TrtLLM format.
+        from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves
+        from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
+            TrtLlmFp8ExpertsModular,
+        )
+        from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+            Fp8MoeBackend,
+            convert_to_fp8_moe_kernel_format,
+        )
+
+        block_shape = [128, 128]
+        qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape)
+
+        # EP-slice before format conversion
+        e_start = num_local_experts * pgi.rank
+        e_end = e_start + num_local_experts
+        w1_ep = qw.w13_weight[e_start:e_end]
+        w2_ep = qw.w2_weight[e_start:e_end]
+        w1_scale_ep = qw.w13_weight_scale[e_start:e_end]
+        w2_scale_ep = qw.w2_weight_scale[e_start:e_end]
+
+        # Convert to TrtLLM format (W31 swap + BlockMajorK shuffle)
+        class _MockLayer:
+            weight_block_size = block_shape
+            class moe_config:
+                is_act_and_mul = True
+                intermediate_size_per_partition = config.n
+            class activation:
+                is_gated = True
+
+        w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = \
+            convert_to_fp8_moe_kernel_format(
+                fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM,
+                layer=_MockLayer(),
+                w13=w1_ep, w2=w2_ep,
+                w13_scale=w1_scale_ep, w2_scale=w2_scale_ep,
+                w13_input_scale=None, w2_input_scale=None,
+            )
+
+        # Build TrtLLM expert with correct EP params
+        quant_config = FusedMoEQuantConfig.make(
+            torch.float8_e4m3fn,
+            block_shape=block_shape,
+            w1_scale=w1_scale_ep,
+            w2_scale=w2_scale_ep,
+        )
+        moe_config = make_dummy_moe_config(
+            num_experts=num_local_experts,
+            experts_per_token=config.topk,
+            hidden_dim=hidden_size,
+            intermediate_size_per_partition=config.n,
+        )
+        fused_experts = TrtLlmFp8ExpertsModular(
+            moe_config=moe_config,
+            quant_config=quant_config,
+        )
+
+        v2_args = DeepEPV2Args(
+            num_local_experts=num_local_experts,
+            num_experts=config.num_experts,
+            num_topk=config.topk,
+            hidden_size=hidden_size,
+            max_tokens_per_rank=8192,
+            use_fp8_dispatch=False,
+        )
+        a2a = make_deepep_v2_a2a(
+            pg=pg, pgi=pgi, dp_size=dp_size,
+            v2_args=v2_args, use_cudagraph=True,
+        )
+        mk_kernel = FusedMoEKernel(
+            prepare_finalize=a2a,
+            fused_experts=fused_experts,
+            inplace=False,
+        )
+
+        for _ in range(3):
+            out = mk_kernel.apply(
+                hidden_states=test_tensors.rank_tokens,
+                w1=w1_ep, w2=w2_ep,
+                topk_weights=test_tensors.topk_weights,
+                topk_ids=test_tensors.topk,
+                activation=MoEActivation.SILU,
+                global_num_experts=config.num_experts,
+                expert_map=None,
+                apply_router_weight_on_input=False,
+            )
+
+        torch.testing.assert_close(
+            torch_combined, out, atol=6e-2, rtol=6e-2,
+        )
+
+
+@pytest.mark.parametrize("m,n,k", [(32, 256, 1024)])
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [6])
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
+@multi_gpu_test(num_gpus=2)
+@requires_deep_ep_v2
+def test_deep_ep_v2_moe_cudagraph(
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    world_dp_size: tuple[int, int],
+    workspace_init,
+):
+    set_random_seed(7)
+    world_size, dp_size = world_dp_size
+    config = TestConfig(
+        dtype=torch.float8_e4m3fn, topk=topk, m=m, k=k, n=n,
+        num_experts=num_experts,
+    )
+
+    parallel_launch(
+        world_size,
+        _deep_ep_v2_moe_cudagraph,
+        dp_size,
+        config,
+        None,  # weights created inside worker
+        None,
+        None,
+        None,
+    )

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:-"b306af06af"}
 NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"}  # Default supports both CUDA 12 and 13
 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
 MODE=${MODE:-install}

vllm/config/parallel.py

@@ -42,6 +42,7 @@ All2AllBackend = Literal[
     "pplx",
     "deepep_high_throughput",
     "deepep_low_latency",
+    "deepep_v2",
     "mori",
     "nixl_ep",
     "allgather_reducescatter",

vllm/config/vllm.py

@@ -1227,12 +1227,14 @@ class VllmConfig:
             assert a2a_backend in [
                 "deepep_low_latency",
                 "deepep_high_throughput",
+                "deepep_v2",
             ], (
-                "Microbatching currently only supports the deepep_low_latency and "
-                f"deepep_high_throughput all2all backend. {a2a_backend} is not "
-                "supported. To fix use --all2all-backend=deepep_low_latency or "
-                "--all2all-backend=deepep_high_throughput and install the DeepEP"
-                " kernels."
+                "Microbatching currently only supports the deepep_low_latency, "
+                "deepep_high_throughput, and deepep_v2 all2all backends. "
+                f"{a2a_backend} is not supported. To fix use "
+                "--all2all-backend=deepep_low_latency, "
+                "--all2all-backend=deepep_high_throughput, or "
+                "--all2all-backend=deepep_v2 and install the DeepEP kernels."
             )
 
             if not self.model_config.disable_cascade_attn:

vllm/distributed/device_communicators/all2all.py

@@ -15,7 +15,7 @@ from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_one_sided,
     has_flashinfer_nvlink_two_sided,
 )
-from vllm.utils.import_utils import has_deep_ep, has_mori
+from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori
 
 from .base_device_communicator import All2AllManagerBase, Cache
 
@@ -760,3 +760,65 @@ class MoriAll2AllManager(All2AllManagerBase):
             mori_kwargs, self._make_handle
         )
         return handle
+
+
+class DeepEPV2All2AllManager(All2AllManagerBase):
+    """
+    All2All communication based on DeepEP v2 ElasticBuffer (unified API).
+    Uses NCCL Gin backend with analytical SM calculation.
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
+        assert has_deep_ep_v2(), (
+            "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 "
+            "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4."
+        )
+        super().__init__(cpu_group, tcp_store_group)
+        self._device_group = device_group
+        self.handle_cache = Cache()
+        self._num_sms: int | None = None
+
+    def _make_all2all_kwargs(
+        self,
+        num_max_tokens_per_rank: int,
+        hidden: int,
+        num_topk: int,
+        use_fp8_dispatch: bool,
+    ) -> dict:
+        return dict(
+            group=self._device_group if self._device_group is not None
+            else self.cpu_group,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            hidden=hidden,
+            num_topk=num_topk,
+            use_fp8_dispatch=use_fp8_dispatch,
+            allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE,
+            prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP,
+            allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION),
+            explicitly_destroy=True,
+        )
+
+    def get_handle(self, kwargs):
+        import deep_ep  # type: ignore[import-not-found]
+
+        num_experts = kwargs.pop("num_experts", 256)
+        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("DeepEP v2 all2all args %s", buffer_kwargs)
+        handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create(
+            buffer_kwargs, deep_ep.ElasticBuffer
+        )
+        if self._num_sms is None:
+            self._num_sms = handle.get_theoretical_num_sms(
+                num_experts=num_experts,
+                num_topk=kwargs["num_topk"],
+            )
+        return handle
+
+    def max_sms_used(self) -> int | None:
+        return self._num_sms
+
+    def destroy(self):
+        with self.handle_cache._lock:
+            for _, handle in self.handle_cache._cache.items():
+                handle.destroy()
+            self.handle_cache._cache.clear()

vllm/distributed/device_communicators/cuda_communicator.py

@@ -137,6 +137,13 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 from .all2all import MoriAll2AllManager
 
                 self.all2all_manager = MoriAll2AllManager(self.cpu_group)
+            elif self.all2all_backend == "deepep_v2":
+                from .all2all import DeepEPV2All2AllManager
+
+                self.all2all_manager = DeepEPV2All2AllManager(
+                    self.cpu_group, tcp_store_group,
+                    device_group=self.device_group,
+                )
             elif self.all2all_backend == "nixl_ep":
                 from .all2all import NixlEPAll2AllManager
 

vllm/envs.py

@@ -233,6 +233,9 @@ if TYPE_CHECKING:
     VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024
     VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False
     VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False
+    VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True
+    VLLM_DEEPEP_V2_PREFER_OVERLAP: bool = False
+    VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION: bool = False
     VLLM_DBO_COMM_SMS: int = 20
     VLLM_PATTERN_MATCH_DEBUG: str | None = None
     VLLM_DEBUG_DUMP_PATH: str | None = None
@@ -1623,6 +1626,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool(
         int(os.getenv("VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL", "0"))
     ),
+    # DeepEP v2 ElasticBuffer: enable two-tier NVLink+RDMA hybrid mode
+    "VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE": lambda: bool(
+        int(os.getenv("VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE", "1"))
+    ),
+    # DeepEP v2 ElasticBuffer: use fewer SMs at slight throughput cost
+    "VLLM_DEEPEP_V2_PREFER_OVERLAP": lambda: bool(
+        int(os.getenv("VLLM_DEEPEP_V2_PREFER_OVERLAP", "0"))
+    ),
+    # DeepEP v2 ElasticBuffer: trade precision for transfer size in combine
+    "VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION": lambda: bool(
+        int(os.getenv("VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION", "0"))
+    ),
     # The number of SMs to allocate for communication kernels when running DBO
     # the rest of the SMs on the device will be allocated to compute
     "VLLM_DBO_COMM_SMS": lambda: int(os.getenv("VLLM_DBO_COMM_SMS", "20")),

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -29,7 +29,12 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two
     FlashInferNVLinkTwoSidedPrepareAndFinalize,
 )
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
+from vllm.utils.import_utils import (
+    has_deep_ep,
+    has_deep_ep_v2,
+    has_mori,
+    has_nixl_ep,
+)
 
 logger = init_logger(__name__)
 
@@ -40,6 +45,8 @@ if current_platform.is_cuda_alike():
             DEEPEP_QUANT_BLOCK_SHAPE,
             DeepEPLLPrepareAndFinalize,
         )
+    if has_deep_ep_v2():
+        from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
     if has_mori():
         from .prepare_finalize.mori import MoriPrepareAndFinalize
     if has_nixl_ep():
@@ -78,6 +85,11 @@ def maybe_roundup_layer_hidden_size(
             hidden_size
         )
 
+    if moe_parallel_config.use_deepep_v2_kernels:
+        hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
+            hidden_size, act_dtype
+        )
+
     if moe_parallel_config.use_nixl_ep_kernels:
         hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
             hidden_size
@@ -181,6 +193,36 @@ def maybe_make_prepare_finalize(
             physical_to_global=physical_to_global,
             local_expert_global_ids=local_expert_global_ids,
         )
+    elif moe.use_deepep_v2_kernels:
+        assert moe.dp_size == all2all_manager.dp_world_size
+
+        use_fp8_dispatch = (
+            quant_config is not None
+            and quant_config.quant_dtype == current_platform.fp8_dtype()
+            and quant_config.is_block_quantized
+        )
+        all_to_all_args = dict(
+            num_max_tokens_per_rank=moe.max_num_tokens,
+            hidden=moe.hidden_dim,
+            num_topk=moe.experts_per_token,
+            num_experts=moe.num_experts,
+            use_fp8_dispatch=use_fp8_dispatch,
+        )
+        handle = all2all_manager.get_handle(all_to_all_args)
+        vllm_config = get_current_vllm_config()
+        use_cudagraph = not vllm_config.model_config.enforce_eager
+
+        prepare_finalize = DeepEPV2PrepareAndFinalize(
+            buffer=handle,
+            num_dispatchers=all2all_manager.world_size,
+            dp_size=all2all_manager.dp_world_size,
+            rank_expert_offset=all2all_manager.rank * moe.num_local_experts,
+            num_experts=moe.num_experts,
+            num_topk=moe.experts_per_token,
+            use_fp8_dispatch=use_fp8_dispatch,
+            use_cudagraph=use_cudagraph,
+        )
+
     elif moe.use_mori_kernels:
         assert quant_config is not None
 

vllm/model_executor/layers/fused_moe/config.py

@@ -1063,6 +1063,10 @@ class FusedMoEParallelConfig:
     def use_nixl_ep_kernels(self):
         return self.use_all2all_kernels and self.all2all_backend == "nixl_ep"
 
+    @property
+    def use_deepep_v2_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend == "deepep_v2"
+
     @staticmethod
     def flatten_tp_across_dp_and_pcp(
         tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1350,6 +1354,10 @@ class FusedMoEConfig:
     def use_nixl_ep_kernels(self):
         return self.moe_parallel_config.use_nixl_ep_kernels
 
+    @property
+    def use_deepep_v2_kernels(self):
+        return self.moe_parallel_config.use_deepep_v2_kernels
+
     @property
     def needs_round_robin_routing_tables(self):
         return self.moe_parallel_config.needs_round_robin_routing_tables

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

@@ -100,6 +100,14 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
     Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.
     """
 
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return (
+            not moe_parallel_config.use_all2all_kernels
+            or moe_parallel_config.use_ag_rs_all2all_kernels
+            or moe_parallel_config.use_deepep_v2_kernels
+        ) and not moe_parallel_config.enable_eplb
+
     @staticmethod
     def _supports_quant_scheme(
         weight_key: QuantKey | None,
@@ -199,7 +207,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale,
             num_experts=global_num_experts,
-            top_k=self.topk,
+            top_k=topk_ids.size(1),
             n_group=None,
             topk_group=None,
             intermediate_size=self.intermediate_size_per_partition,

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

@@ -85,6 +85,18 @@ def _get_priority_backends(
     def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None:
         backends.insert(0, backends.pop(backends.index(backend)))
 
+    # With DeepEP v2 contiguous layout (do_expand=False), tensors are
+    # worst-case allocated with padding. TrtLLM's tile-level skipping
+    # avoids wasted compute on padding rows; other backends process all rows.
+    if (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+        and moe_config.moe_parallel_config.use_deepep_v2_kernels
+        and activation_key == kFp8Dynamic128Sym
+        and weight_key == kFp8Static128BlockSym
+    ):
+        _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_TRTLLM)
+
     # On Hopper for Block Fp8, prefer Triton for TP and FI CUTLASS for EP.
     if (
         current_platform.is_cuda()

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

@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
+
+import deep_ep
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceContiguous,
+    TopKWeightAndReduceDelegate,
+)
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm.utils.math_utils import round_up
+from vllm.v1.worker.ubatching import (
+    dbo_current_ubatch_id,
+)
+
+
+class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
+    """
+    Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
+
+    Supports two modes controlled by the `use_cudagraph` constructor arg:
+
+    **Decode mode (use_cudagraph=True):**
+      - do_expand=False, do_cpu_sync=False
+      - Tokens returned in original order with recv_topk_idx (global IDs)
+      - Worst-case tensor allocation; padding rows zeroed via
+        handle.psum_num_recv_tokens_per_scaleup_rank
+      - Fully cudagraph-capturable
+      - Expert kernel sorts internally (expert_tokens_meta=None)
+
+    **Prefill mode (use_cudagraph=False):**
+      - do_expand=True, do_cpu_sync=True
+      - Per-expert-contiguous layout; exact memory allocation
+      - Saves GPU memory (no worst-case allocation)
+      - Not cudagraph-capturable (CPU polling), but prefill doesn't
+        use cudagraphs anyway
+      - Provides expert_tokens_meta for efficient batched expert kernels
+
+    Both modes use async_with_compute_stream=False (synchronous from
+    caller's perspective). The ElasticBuffer handles comm internally.
+    """
+
+    @staticmethod
+    def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype) -> int:
+        hidden_size_bytes = hidden_size * dtype.itemsize
+        xfer_atom_size = 512  # 32 * 16 (size(int4))
+        if hidden_size_bytes % xfer_atom_size == 0:
+            return hidden_size
+
+        hidden_size_bytes = round_up(hidden_size_bytes, xfer_atom_size)
+        return hidden_size_bytes // dtype.itemsize
+
+    def __init__(
+        self,
+        buffer: deep_ep.ElasticBuffer,
+        num_dispatchers: int,
+        dp_size: int,
+        rank_expert_offset: int,
+        num_experts: int,
+        num_topk: int,
+        use_fp8_dispatch: bool = False,
+        use_cudagraph: bool = False,
+    ):
+        super().__init__()
+        self.buffer = buffer
+        self.num_dispatchers_ = num_dispatchers
+        self.dp_size = dp_size
+        self.rank_expert_offset = rank_expert_offset
+        self.num_experts = num_experts
+        self.num_topk = num_topk
+        self.use_fp8_dispatch = use_fp8_dispatch
+        self.use_cudagraph = use_cudagraph
+
+        # DBO microbatching: one handle slot per micro-batch.
+        self.handles: list[deep_ep.EPHandle | None] = [None, None]
+
+    def num_dispatchers(self) -> int:
+        return self.num_dispatchers_
+
+    def output_is_reduced(self) -> bool:
+        return True
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return None
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return torch.int64
+
+    def _do_dispatch(
+        self,
+        tokens: torch.Tensor,
+        token_scales: torch.Tensor | None,
+        rank_topk_ids: torch.Tensor,
+        rank_topk_weights: torch.Tensor,
+        num_experts: int,
+        a1_scale: torch.Tensor | None,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool,
+    ) -> Callable:
+        has_scales = token_scales is not None
+
+        token_data = tokens
+        if has_scales:
+            token_data = (tokens, token_scales)
+
+        # Decode: do_expand=False + do_cpu_sync=False (cudagraph-safe)
+        # Prefill: do_expand=True + do_cpu_sync=True (memory-efficient)
+        do_expand = not self.use_cudagraph
+        do_cpu_sync = not self.use_cudagraph
+
+        (
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            handle,
+            event,
+        ) = self.buffer.dispatch(
+            x=token_data,
+            topk_idx=rank_topk_ids,
+            topk_weights=rank_topk_weights,
+            num_experts=num_experts,
+            do_expand=do_expand,
+            do_cpu_sync=do_cpu_sync,
+            async_with_compute_stream=False,
+        )
+
+        a2a_idx = dbo_current_ubatch_id()
+        self.handles[a2a_idx] = handle
+
+        return lambda: self._receiver(
+            event,
+            has_scales,
+            recv_x,
+            recv_topk_idx,
+            num_experts,
+            handle.num_recv_tokens_per_expert_list,
+            recv_topk_weights,
+            handle.psum_num_recv_tokens_per_scaleup_rank,
+            a1_scale,
+            quant_config,
+            defer_input_quant=defer_input_quant,
+        )
+
+    def _receiver(
+        self,
+        event: deep_ep.EventOverlap,
+        has_scales: bool,
+        recv_x: tuple[torch.Tensor, torch.Tensor] | torch.Tensor,
+        recv_topk_idx: torch.Tensor | None,
+        num_experts: int,
+        recv_expert_num_tokens: list[int],
+        recv_topk_weights: torch.Tensor | None,
+        psum_recv_per_rank: torch.Tensor,
+        a1_scale: torch.Tensor | None,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool,
+    ) -> mk.PrepareResultType:
+        if event.event is not None:
+            event.current_stream_wait()
+
+        if isinstance(recv_x, tuple):
+            expert_x, expert_x_scale = recv_x
+        else:
+            expert_x, expert_x_scale = recv_x, None
+
+        if recv_topk_idx is None:
+            # do_expand=True (prefill mode): build topk_ids from
+            # per-expert token counts.
+            total_tokens = sum(recv_expert_num_tokens)
+            if total_tokens > 0:
+                recv_topk_idx = torch.empty(
+                    total_tokens,
+                    dtype=torch.int64,
+                    device=expert_x.device,
+                )
+                offset = 0
+                for i, count in enumerate(recv_expert_num_tokens):
+                    if count > 0:
+                        recv_topk_idx[offset:offset + count].fill_(
+                            i + self.rank_expert_offset)
+                        offset += count
+            else:
+                recv_topk_idx = torch.empty(
+                    0, dtype=torch.int64, device=expert_x.device,
+                )
+            recv_topk_idx = recv_topk_idx.unsqueeze(1)
+        else:
+            # do_expand=False (decode/cudagraph mode): recv_topk_idx has
+            # LOCAL expert IDs (-1 for non-local and padding rows).
+            # Convert valid local IDs to global. Rows with -1 are
+            # skipped by expert kernels (TrtLLM tile-level skipping,
+            # DeepGemm is_computation_valid), so no need to zero
+            # hidden states, scales, or weights for padding rows.
+            valid_mask = recv_topk_idx >= 0
+            recv_topk_idx = torch.where(
+                valid_mask,
+                recv_topk_idx + self.rank_expert_offset,
+                recv_topk_idx,
+            )
+
+        # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1]
+        if recv_topk_weights is not None and recv_topk_weights.ndim == 1:
+            recv_topk_weights = recv_topk_weights.unsqueeze(1)
+
+        expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
+            recv_expert_num_tokens, device=expert_x.device,
+        )
+
+        if not quant_config.is_block_quantized and not defer_input_quant:
+            expert_x_scale = None
+            if expert_x.numel() != 0:
+                expert_x, expert_x_scale = moe_kernel_quantize_input(
+                    expert_x,
+                    a1_scale,
+                    quant_dtype=quant_config.quant_dtype,
+                    per_act_token_quant=False,
+                    block_shape=quant_config.block_shape,
+                    is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,
+                )
+
+        return (
+            expert_x,
+            expert_x_scale,
+            expert_tokens_meta,
+            recv_topk_idx,
+            recv_topk_weights,
+        )
+
+    def supports_async(self) -> bool:
+        return True
+
+    def prepare_async(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.ReceiverType:
+        if apply_router_weight_on_input:
+            topk = topk_ids.size(1)
+            assert topk == 1, (
+                "apply_router_weight_on_input is only implemented for topk=1"
+            )
+            a1 = a1 * topk_weights.to(a1.dtype)
+
+        if quant_config.is_block_quantized and not defer_input_quant:
+            a1q, a1q_scale = moe_kernel_quantize_input(
+                a1,
+                quant_config.a1_scale,
+                quant_dtype=quant_config.quant_dtype,
+                per_act_token_quant=quant_config.per_act_token_quant,
+                block_shape=quant_config.block_shape,
+            )
+            if a1q_scale is not None and a1q_scale.numel() == 1:
+                a1q_scale = a1q_scale.view(1, 1)
+            a1_post_scale = None
+        else:
+            a1q = a1
+            a1q_scale = None
+            a1_post_scale = (
+                quant_config.a1_gscale
+                if quant_config.quant_dtype == "nvfp4"
+                else quant_config.a1_scale
+            )
+
+        return self._do_dispatch(
+            tokens=a1q,
+            token_scales=a1q_scale,
+            rank_topk_ids=topk_ids,
+            rank_topk_weights=topk_weights,
+            num_experts=num_experts,
+            a1_scale=a1_post_scale,
+            quant_config=quant_config,
+            defer_input_quant=defer_input_quant,
+        )
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareResultType:
+        receiver = self.prepare_async(
+            a1,
+            topk_weights,
+            topk_ids,
+            num_experts,
+            expert_map,
+            apply_router_weight_on_input,
+            quant_config,
+            defer_input_quant,
+        )
+        return receiver()
+
+    def _finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+        do_async: bool,
+    ) -> Callable | None:
+        a2a_idx = dbo_current_ubatch_id()
+        handle = self.handles[a2a_idx]
+        assert handle is not None
+
+        if fused_expert_output.numel() != 0:
+            if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
+                weight_and_reduce_impl = TopKWeightAndReduceContiguous()
+            fused_expert_output = weight_and_reduce_impl.apply(
+                output=None,
+                fused_expert_output=fused_expert_output,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+            )
+
+        if fused_expert_output.dtype != torch.bfloat16:
+            raise ValueError(
+                f"DeepEP v2 combine requires bfloat16 input, "
+                f"got {fused_expert_output.dtype}"
+            )
+
+        combined_x, _, event = self.buffer.combine(
+            x=fused_expert_output,
+            handle=handle,
+            topk_weights=None,
+            async_with_compute_stream=False,
+        )
+
+        output.copy_(combined_x, non_blocking=True)
+        return None
+
+    def finalize_async(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> Callable:
+        self._finalize(
+            output,
+            fused_expert_output,
+            topk_weights,
+            topk_ids,
+            apply_router_weight_on_input,
+            weight_and_reduce_impl,
+            False,
+        )
+        return lambda: None
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        self._finalize(
+            output,
+            fused_expert_output,
+            topk_weights,
+            topk_ids,
+            apply_router_weight_on_input,
+            weight_and_reduce_impl,
+            False,
+        )

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

@@ -278,6 +278,20 @@ class FusedTopKBiasRouter(BaseRouter):
         input_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute routing using fused top-k with bias."""
+        # The topk kernel dispatches dtype based on topk_ids (set by
+        # indices_type) and assumes input_tokens/hash_indices_table match.
+        # Cast them here so backends like DeepEP that require int64 indices
+        # don't hit a dtype mismatch against the model's int32 buffers.
+        hash_table = self._hash_indices_table
+        if indices_type is not None:
+            if input_ids is not None:
+                input_ids = input_ids.to(dtype=indices_type)
+            if (hash_table is not None
+                    and hash_table.dtype != indices_type):
+                self._hash_indices_table = hash_table.to(
+                    dtype=indices_type)
+                hash_table = self._hash_indices_table
+
         topk_weights, topk_ids = fused_topk_bias(
             hidden_states=hidden_states,
             gating_output=router_logits,
@@ -289,7 +303,7 @@ class FusedTopKBiasRouter(BaseRouter):
             renormalize=self.renormalize,
             indices_type=indices_type,
             input_tokens=input_ids,
-            hash_indices_table=self._hash_indices_table,
+            hash_indices_table=hash_table,
             routed_scaling_factor=self.routed_scaling_factor,
         )
 

vllm/utils/import_utils.py

@@ -405,6 +405,61 @@ def has_deep_ep() -> bool:
     return _has_module("deep_ep")
 
 
+DEEPEP_V2_MIN_NCCL_VERSION_RAW = 23004  # 2.30.4
+
+
+def _get_runtime_nccl_version() -> int | None:
+    """Get the runtime NCCL version by loading the actual library.
+
+    Returns the raw version int (e.g. 23004 for 2.30.4), or None on failure.
+    torch.cuda.nccl.version() is a compile-time constant from the PyTorch
+    wheel and does not reflect a separately installed NCCL.
+    """
+    import ctypes
+
+    try:
+        from vllm.utils.nccl import find_nccl_library
+
+        lib = ctypes.CDLL(find_nccl_library())
+        version = ctypes.c_int()
+        lib.ncclGetVersion(ctypes.byref(version))
+        return version.value
+    except Exception:
+        return None
+
+
+def _format_nccl_raw_version(raw: int) -> str:
+    s = str(raw)
+    return f"{s[0]}.{s[1:3].lstrip('0') or '0'}.{s[3:].lstrip('0') or '0'}"
+
+
+def has_deep_ep_v2() -> bool:
+    """Whether deep_ep with ElasticBuffer (v2 API) is available.
+
+    Requires both the ElasticBuffer class in the deep_ep module and
+    NCCL >= 2.30.4 (GIN backend), checked against the runtime library.
+    """
+    if not _has_module("deep_ep"):
+        return False
+    import deep_ep  # type: ignore[import-not-found]
+
+    if not hasattr(deep_ep, "ElasticBuffer"):
+        return False
+    try:
+        nccl_ver = _get_runtime_nccl_version()
+        if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW:
+            logger.info_once(
+                "DeepEP v2 requires NCCL >= %s but found %s. "
+                "deepep_v2 backend will not be available.",
+                _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW),
+                _format_nccl_raw_version(nccl_ver) if nccl_ver else "unknown",
+            )
+            return False
+    except Exception:
+        return False
+    return True
+
+
 def has_deep_gemm() -> bool:
     """Whether the optional `deep_gemm` package is available.