From a14d8d7016a46ea9f4e11cfd0ba4e15e9963a7a7 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Wed, 13 May 2026 20:20:39 -0700
Subject: [PATCH 1/7] Add DeepEP v2 flex dispatcher backend

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 .../core/models/gpt/fine_grained_callables.py |   4 +-
 megatron/core/transformer/moe/fused_a2a.py    | 206 +++++++++++++-
 .../core/transformer/moe/token_dispatcher.py  | 266 +++++++++++++++++-
 .../core/transformer/transformer_config.py    |  12 +-
 .../transformer/moe/test_token_dispatcher.py  |  40 +--
 5 files changed, 499 insertions(+), 29 deletions(-)

diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 0038851ac5a..191786ddca3 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -61,7 +61,7 @@ def should_free_input(name, is_moe, config, num_local_experts):
         return False
     enable_deepep = (
         config.moe_token_dispatcher_type == "flex"
-        and config.moe_flex_dispatcher_backend == "deepep"
+        and config.moe_flex_dispatcher_backend in ("deepep", "deepepv2")
     )
     enable_hybridep = (
         config.moe_token_dispatcher_type == "flex"
@@ -494,7 +494,7 @@ def build_transformer_layer_callables(layer: TransformerLayer):
     is_moe = isinstance(layer.mlp, MoELayer)
     enable_deepep = (
         layer.config.moe_token_dispatcher_type == "flex"
-        and layer.config.moe_flex_dispatcher_backend == "deepep"
+        and layer.config.moe_flex_dispatcher_backend in ("deepep", "deepepv2")
     )
     enable_hybridep = (
         layer.config.moe_token_dispatcher_type == "flex"
diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index e0f7ab9030c..1bb26ec6219 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -8,16 +8,32 @@
 from megatron.core.utils import internal_api
 
 try:
-    from deep_ep import Buffer
     from deep_ep.utils import EventHandle, EventOverlap
+except ImportError:
+    try:
+        from deep_ep import EventHandle, EventOverlap
+    except ImportError:
+        EventHandle = None
+        EventOverlap = None
+
+try:
+    from deep_ep import Buffer
 
     HAVE_DEEP_EP = True
 except ImportError:
     HAVE_DEEP_EP = False
 
+try:
+    from deep_ep import ElasticBuffer
+
+    HAVE_DEEP_EP_V2 = True
+except ImportError:
+    HAVE_DEEP_EP_V2 = False
+
 import torch
 
 _buffer = None
+_elastic_buffer = None
 
 
 def get_hidden_bytes(x: torch.Tensor) -> int:
@@ -71,6 +87,39 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int):
     return _buffer
 
 
+def get_elastic_buffer(
+    group: torch.distributed.ProcessGroup, num_max_tokens_per_rank: int, hidden: int, num_topk: int
+):
+    """Get or create a DeepEP v2 elastic buffer for all-to-all communication."""
+    global _elastic_buffer
+
+    num_bytes = ElasticBuffer.get_buffer_size_hint(
+        group, num_max_tokens_per_rank=num_max_tokens_per_rank, hidden=hidden, num_topk=num_topk
+    )
+
+    if (
+        _elastic_buffer is None
+        or _elastic_buffer.group != group
+        or _elastic_buffer.num_bytes < num_bytes
+        or _elastic_buffer.num_max_tokens_per_rank < num_max_tokens_per_rank
+    ):
+        _elastic_buffer = ElasticBuffer(
+            group,
+            num_bytes=num_bytes,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            hidden=hidden,
+            num_topk=num_topk,
+        )
+    return _elastic_buffer
+
+
+def _capture_elastic_previous_event(buffer, async_finish: bool, allocate_on_comm_stream: bool):
+    """Capture the current stream for DeepEP v2 async communication."""
+    if async_finish and allocate_on_comm_stream:
+        return buffer.capture()
+    return None
+
+
 class FusedDispatch(torch.autograd.Function):
     """Fused dispatch operation for MoE routing combining computation and communication."""
 
@@ -270,6 +319,161 @@ def set_deepep_num_sms(num_sms):
     set_deepep_num_sms = None
 
 
+class DeepepV2Dispatch(torch.autograd.Function):
+    """DeepEP v2 elastic dispatch with autograd support."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        buffer,
+        x,
+        token_indices,
+        token_probs,
+        num_experts,
+        num_max_tokens_per_rank,
+        expert_alignment,
+        num_sms,
+        async_finish=False,
+        allocate_on_comm_stream=False,
+    ):
+        """Forward pass of DeepEP v2 elastic dispatch."""
+        previous_event = _capture_elastic_previous_event(
+            buffer, async_finish, allocate_on_comm_stream
+        )
+        recv_x, recv_token_indices, recv_token_probs, handle, event = buffer.dispatch(
+            x,
+            topk_idx=token_indices,
+            topk_weights=token_probs,
+            num_experts=num_experts,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            expert_alignment=expert_alignment,
+            num_sms=num_sms,
+            previous_event=previous_event,
+            async_with_compute_stream=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+
+        if async_finish:
+            event.current_stream_wait()
+
+        ctx.buffer = buffer
+        ctx.handle = handle
+        ctx.num_sms = num_sms
+        ctx.async_finish = async_finish
+        ctx.allocate_on_comm_stream = allocate_on_comm_stream
+        tokens_per_expert = torch.tensor(handle.num_recv_tokens_per_expert_list)
+
+        return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle)
+
+    @staticmethod
+    def backward(
+        ctx, grad_output, grad_token_indices, grad_token_probs, grad_tokens_per_expert, grad_handle
+    ):
+        """Backward pass of DeepEP v2 elastic dispatch."""
+        previous_event = _capture_elastic_previous_event(
+            ctx.buffer, ctx.async_finish, ctx.allocate_on_comm_stream
+        )
+        grad_x, grad_token_probs, event = ctx.buffer.combine(
+            grad_output.contiguous(),
+            handle=ctx.handle,
+            topk_weights=grad_token_probs.float(),
+            num_sms=ctx.num_sms,
+            previous_event=previous_event,
+            async_with_compute_stream=ctx.async_finish,
+            allocate_on_comm_stream=ctx.allocate_on_comm_stream,
+        )
+        if ctx.async_finish:
+            event.current_stream_wait()
+        return None, grad_x, None, grad_token_probs, None, None, None, None, None, None
+
+
+class DeepepV2Combine(torch.autograd.Function):
+    """DeepEP v2 elastic combine with autograd support."""
+
+    @staticmethod
+    def forward(ctx, buffer, x, handle, num_sms, async_finish=False, allocate_on_comm_stream=False):
+        """Forward pass of DeepEP v2 elastic combine."""
+        previous_event = _capture_elastic_previous_event(
+            buffer, async_finish, allocate_on_comm_stream
+        )
+        combined_x, combined_token_probs, event = buffer.combine(
+            x,
+            handle=handle,
+            num_sms=num_sms,
+            previous_event=previous_event,
+            async_with_compute_stream=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        if async_finish:
+            event.current_stream_wait()
+
+        ctx.buffer = buffer
+        ctx.handle = handle
+        ctx.num_sms = num_sms
+        ctx.async_finish = async_finish
+        ctx.allocate_on_comm_stream = allocate_on_comm_stream
+        return combined_x, combined_token_probs
+
+    @staticmethod
+    def backward(ctx, grad_output, grad_combined_token_probs):
+        """Backward pass of DeepEP v2 elastic combine."""
+        previous_event = _capture_elastic_previous_event(
+            ctx.buffer, ctx.async_finish, ctx.allocate_on_comm_stream
+        )
+        grad_x, _, _, _, event = ctx.buffer.dispatch(
+            grad_output.contiguous(),
+            handle=ctx.handle,
+            num_sms=ctx.num_sms,
+            previous_event=previous_event,
+            async_with_compute_stream=ctx.async_finish,
+            allocate_on_comm_stream=ctx.allocate_on_comm_stream,
+        )
+        if ctx.async_finish:
+            event.current_stream_wait()
+        return None, grad_x, None, None, None, None, None
+
+
+if HAVE_DEEP_EP_V2:
+
+    def deepepv2_dispatch(
+        buffer,
+        x,
+        token_indices,
+        token_probs,
+        num_experts,
+        num_max_tokens_per_rank,
+        expert_alignment=1,
+        num_sms=0,
+        async_finish=False,
+        allocate_on_comm_stream=False,
+    ):
+        """Perform DeepEP v2 elastic dispatch."""
+        return DeepepV2Dispatch.apply(
+            buffer,
+            x.contiguous(),
+            token_indices,
+            token_probs,
+            num_experts,
+            num_max_tokens_per_rank,
+            expert_alignment,
+            num_sms,
+            async_finish,
+            allocate_on_comm_stream,
+        )
+
+    def deepepv2_combine(
+        buffer, x, handle, num_sms=0, async_finish=False, allocate_on_comm_stream=False
+    ):
+        """Perform DeepEP v2 elastic combine."""
+        return DeepepV2Combine.apply(
+            buffer, x.contiguous(), handle, num_sms, async_finish, allocate_on_comm_stream
+        )
+
+else:
+    deepepv2_dispatch = None
+    deepepv2_combine = None
+
+
 try:
     from deep_ep import HybridEPBuffer
 
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index e8f2044650e..9d76751940d 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -19,8 +19,11 @@
 from megatron.core.transformer.enums import CudaGraphModule
 from megatron.core.transformer.moe.fused_a2a import (
     HYBRIDEP_TOKEN_ALIGNMENT,
+    deepepv2_combine,
+    deepepv2_dispatch,
     fused_combine,
     fused_dispatch,
+    get_elastic_buffer,
     hybrid_ep_combine,
     hybrid_ep_dispatch,
     set_deepep_num_sms,
@@ -966,18 +969,31 @@ class _DispatchManager(ABC):
     can be the number of local experts, or the size of sub_group.
     """
 
+    token_probs: Optional[torch.Tensor]
+    dispatched_probs: Optional[torch.Tensor]
+
     @abstractmethod
     def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         """Set up metadata of routing_map and probs."""
         pass
 
     @abstractmethod
-    def dispatch(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        async_finish: bool = True,
+        allocate_on_comm_stream: bool = True,
+    ) -> torch.Tensor:
         """Dispatch the hidden_states according to the routing_map."""
         pass
 
     @abstractmethod
-    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def combine(
+        self,
+        hidden_states: torch.Tensor,
+        async_finish: bool = True,
+        allocate_on_comm_stream: bool = True,
+    ) -> torch.Tensor:
         """Combine the hidden_states after expert processing."""
         pass
 
@@ -986,6 +1002,11 @@ def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
         """Get the permuted hidden states by instances."""
         pass
 
+    @abstractmethod
+    def get_number_of_tokens_per_expert(self) -> torch.Tensor:
+        """Get the number of tokens per expert."""
+        pass
+
     @abstractmethod
     def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Get the restored hidden states by instances."""
@@ -1461,6 +1482,233 @@ def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
         return hidden_states
 
 
+class _DeepepV2Manager(_DispatchManager):
+    """
+    A manager class for the DeepEP v2 ElasticBuffer backend.
+
+    This keeps the original DeepEP backend isolated under "deepep", while "deepepv2"
+    uses the v2 dispatch/combine APIs.
+    """
+
+    def __init__(
+        self,
+        group: torch.distributed.ProcessGroup,
+        num_local_experts: int,
+        router_topk: int,
+        num_experts: int,
+        config: TransformerConfig,
+    ):
+        self.group = group
+        self.num_local_experts = num_local_experts
+        self.config = config
+
+        self.router_topk = router_topk
+        self.num_experts = num_experts
+        self.router_dtype = config.moe_router_dtype
+        self.capacity_factor = config.moe_expert_capacity_factor
+        self.permute_fusion = config.moe_permute_fusion
+        self.num_sms = config.moe_deepep_num_sms
+
+        self.token_indices: Optional[torch.Tensor] = None
+        self.token_probs: Optional[torch.Tensor] = None
+        self.handle = None
+        self.buffer = None
+
+        if deepepv2_dispatch is None:
+            raise ImportError(
+                "DeepEP v2 is not installed. Please install a DeepEP package that provides "
+                "ElasticBuffer."
+            )
+
+    def _get_buffer(self, hidden_states: torch.Tensor):
+        self.buffer = get_elastic_buffer(
+            self.group,
+            num_max_tokens_per_rank=hidden_states.shape[0],
+            hidden=hidden_states.shape[1],
+            num_topk=self.token_indices.shape[1],
+        )
+        return self.buffer
+
+    def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
+        num_tokens = routing_map.shape[0]
+
+        routing_map = routing_map.reshape(num_tokens, self.num_experts)
+        probs = probs.reshape(num_tokens, self.num_experts)
+        # Convert the format of routing map from multihot to indices.
+        self.token_probs, self.token_indices = torch.topk(probs, self.router_topk, dim=-1)
+        # Mask the indices of dropped tokens with -1
+        if self.capacity_factor is not None:
+            mask = self.token_probs == 0
+            self.token_indices = self.token_indices.masked_fill(mask, -1)
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        async_finish: bool = False,
+        allocate_on_comm_stream: bool = False,
+    ) -> torch.Tensor:
+        # DeepEP v2 only supports float32 probs
+        if self.token_probs.dtype != torch.float32:
+            if self.token_probs.dtype in [torch.bfloat16, torch.float16]:
+                logger.warning(
+                    "DeepEP v2 only supports float32 probs, please set --moe-router-dtype=fp32"
+                )
+            self.token_probs = self.token_probs.float()
+        buffer = self._get_buffer(hidden_states)
+        hidden_states, dispatched_indices, dispatched_probs, num_tokens_per_expert, handle = (
+            deepepv2_dispatch(
+                buffer,
+                hidden_states,
+                self.token_indices,
+                self.token_probs,
+                self.num_experts,
+                num_max_tokens_per_rank=hidden_states.shape[0],
+                expert_alignment=1,
+                num_sms=self.num_sms,
+                async_finish=async_finish,
+                allocate_on_comm_stream=allocate_on_comm_stream,
+            )
+        )
+        self.handle = handle
+        self.tokens_per_expert = num_tokens_per_expert
+        self.dispatched_indices = dispatched_indices
+        self.dispatched_probs = dispatched_probs
+
+        return hidden_states
+
+    def combine(
+        self,
+        hidden_states: torch.Tensor,
+        async_finish: bool = False,
+        allocate_on_comm_stream: bool = False,
+    ) -> torch.Tensor:
+        hidden_states, _ = deepepv2_combine(
+            self.buffer,
+            hidden_states,
+            self.handle,
+            num_sms=self.num_sms,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        self.handle = None
+        self.dispatched_indices = None
+        self.dispatched_probs = None
+        return hidden_states
+
+    def _indices_to_multihot(self, indices, probs):
+        """
+        Converts a tensor of indices to a multihot vector.
+
+        Args:
+            indices (torch.Tensor): [num_tokens, topk] token indices, where -1 means masked out.
+            probs (torch.Tensor): [num_tokens, topk] token probabilities.
+
+        Returns:
+            A tuple of (routing_map, probs), where routing_map is the multihot vector
+            and probs is the multihot probabilities.
+        """
+        batch_size = indices.shape[0]
+        multihot_routing_map = torch.zeros(
+            (batch_size, self.num_local_experts), dtype=torch.long, device=indices.device
+        )
+
+        multihot_probs = torch.zeros(
+            (batch_size, self.num_local_experts), dtype=torch.float, device=indices.device
+        )
+
+        mask = indices != -1
+        valid_indices = indices[mask]
+        row_indices = torch.arange(batch_size, device=indices.device).repeat_interleave(
+            mask.sum(dim=1)
+        )
+        multihot_routing_map[row_indices, valid_indices] = 1
+        multihot_probs[row_indices, valid_indices] = probs[mask]
+        return multihot_routing_map.bool(), multihot_probs
+
+    def get_number_of_tokens_per_expert(self) -> torch.Tensor:
+        """
+        Get the number of tokens per expert.
+        """
+        return self.tokens_per_expert
+
+    def _pad_routing_map(
+        self, routing_map: torch.Tensor, tokens_per_expert: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Pad the routing map to the nearest multiple of the pad_multiple.
+        """
+        pad_multiple = get_align_size_for_quantization(self.config)
+
+        num_input_tokens = routing_map.shape[0]
+        target_tokens_per_expert = (
+            torch.ceil(tokens_per_expert / pad_multiple) * pad_multiple
+        ).long()
+
+        # Check if there are enough tokens to pad
+        enough_tokens_to_pad = torch.all(target_tokens_per_expert <= num_input_tokens)
+        if not enough_tokens_to_pad:
+            logger.warning(
+                "Not enough tokens to pad. The total number of tokens received in this rank "
+                "is smaller than the target number of tokens for each expert. "
+                "Falling back to explicit padding within GroupedMLP"
+            )
+        else:
+            if is_experimental_enabled() and self.permute_fusion:
+                from megatron.core.fusions.fused_pad_routing_map import fused_pad_routing_map
+
+                routing_map = fused_pad_routing_map(routing_map, pad_multiple)
+            else:
+                routing_map = pad_routing_map(routing_map, pad_multiple)
+            tokens_per_expert = target_tokens_per_expert
+        return routing_map, tokens_per_expert
+
+    def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if is_experimental_enabled() and self.permute_fusion:
+            self.dispatched_routing_map, self.dispatched_probs = fused_indices_to_multihot(
+                self.dispatched_indices, self.dispatched_probs, self.num_local_experts
+            )
+        else:
+            self.dispatched_routing_map, self.dispatched_probs = self._indices_to_multihot(
+                self.dispatched_indices, self.dispatched_probs
+            )
+        if self.config.moe_router_padding_for_quantization:
+            self.dispatched_routing_map, self.tokens_per_expert = self._pad_routing_map(
+                self.dispatched_routing_map, self.tokens_per_expert
+            )
+
+        self.hidden_shape_before_permute = hidden_states.shape
+        assert self.dispatched_probs.dtype == torch.float32, "DeepEP v2 only supports float32 probs"
+        (
+            hidden_states,
+            permuted_probs,
+            self.reversed_mapping_for_combine,
+            self.pad_offsets,
+            self.tokens_per_expert,
+        ) = permute(
+            hidden_states,
+            self.dispatched_routing_map,
+            probs=self.dispatched_probs,
+            num_out_tokens=self.tokens_per_expert.sum().item(),
+            fused=self.permute_fusion,
+            tokens_per_expert=self.tokens_per_expert,
+            align_size=get_align_size_for_quantization(self.config),
+        )
+        if self.router_dtype == "fp64":
+            permuted_probs = permuted_probs.to(torch.float64)
+        return hidden_states, permuted_probs
+
+    def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = unpermute(
+            hidden_states,
+            self.reversed_mapping_for_combine,
+            restore_shape=self.hidden_shape_before_permute,
+            routing_map=self.dispatched_routing_map,
+            fused=self.permute_fusion,
+            pad_offsets=self.pad_offsets,
+        )
+        return hidden_states
+
+
 class MoEFlexTokenDispatcher(MoETokenDispatcher):
     """A flexible token dispatcher that abstracts the underlying tensor and expert
     parallelism. It uses a single communication group over all TP and EP ranks,
@@ -1487,6 +1735,7 @@ def __init__(
 
         self.num_local_experts = num_local_experts
         self.local_expert_indices = local_expert_indices
+        self._comm_manager: _DispatchManager
         if self.config.moe_flex_dispatcher_backend == "deepep":
             assert self.tp_size * self.ep_size > 1, "DeepEP dispatcher requires TPxEP > 1"
             self._comm_manager = _DeepepManager(
@@ -1497,6 +1746,16 @@ def __init__(
                 config=self.config,
             )
             self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.token_indices']
+        elif self.config.moe_flex_dispatcher_backend == "deepepv2":
+            assert self.tp_size * self.ep_size > 1, "DeepEP v2 dispatcher requires TPxEP > 1"
+            self._comm_manager = _DeepepV2Manager(
+                group=self.tp_ep_group,
+                num_local_experts=self.num_local_experts,
+                router_topk=self.tp_size * self.config.moe_router_topk,
+                num_experts=self.tp_size * self.config.num_moe_experts,
+                config=self.config,
+            )
+            self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.token_indices']
         elif self.config.moe_flex_dispatcher_backend == "hybridep":
             self._comm_manager = _HybridEPManager(
                 group=self.tp_ep_group,
@@ -1508,7 +1767,8 @@ def __init__(
         else:
             raise ValueError(
                 f"Invalid backend: {self.config.moe_flex_dispatcher_backend}"
-                "Please set --moe-flex-dispatcher-backend=deepep or "
+                "Please set --moe-flex-dispatcher-backend=deepep, "
+                "--moe-flex-dispatcher-backend=deepepv2 or "
                 "--moe-flex-dispatcher-backend=hybridep"
             )
 
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 6687d572914..eb9c8af7dd9 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -874,10 +874,10 @@ class TransformerConfig(ModelParallelConfig):
     moe_enable_deepep: bool = False
     """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models."""
 
-    moe_flex_dispatcher_backend: Literal['deepep', 'hybridep'] = "deepep"
+    moe_flex_dispatcher_backend: Literal['deepep', 'deepepv2', 'hybridep'] = "deepep"
     """[Experimental] The backend to use for flex token dispatcher. The default is "deepep".
-    Options are "deepep" and "hybridep". Currently only "hybridep" backend supports 
-    the MNNVL case."""
+    Options are "deepep", "deepepv2" and "hybridep". Currently only "hybridep"
+    backend supports the MNNVL case."""
 
     moe_permute_fusion_into_hybridep: bool = False
     """Fuse token rearrangement ops during token dispatching for HybridEP."""
@@ -1701,7 +1701,7 @@ def __post_init__(self):
         if self.moe_enable_deepep:
             if self.moe_token_dispatcher_type != "flex":
                 raise ValueError("DeepEP backend is only supported with flex token dispatcher.")
-            if self.moe_flex_dispatcher_backend == "hybridep":
+            if self.moe_flex_dispatcher_backend in ("deepepv2", "hybridep"):
                 raise ValueError("Only one backend is supported for flex token dispatcher.")
             self.moe_flex_dispatcher_backend = "deepep"
             warnings.warn(
@@ -1711,10 +1711,10 @@ def __post_init__(self):
 
         if self.moe_token_dispatcher_type == "flex":
             if self.moe_pad_expert_input_to_capacity and (
-                self.moe_enable_deepep or self.moe_flex_dispatcher_backend == "deepep"
+                self.moe_enable_deepep or self.moe_flex_dispatcher_backend in ("deepep", "deepepv2")
             ):
                 raise ValueError(
-                    "Flex token dispatcher with deepep backend does not support "
+                    "Flex token dispatcher with deepep/deepepv2 backend does not support "
                     "moe_pad_expert_input_to_capacity"
                 )
 
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index ddc8b313c19..e769dab664f 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -465,6 +465,12 @@ def is_deep_ep_available():
     return HAVE_DEEP_EP
 
 
+def is_deep_ep_v2_available():
+    from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP_V2
+
+    return HAVE_DEEP_EP_V2
+
+
 def is_hybrid_ep_available():
     from megatron.core.transformer.moe.fused_a2a import HAVE_HYBRIDEP
 
@@ -636,9 +642,18 @@ def test_sequence_packing_thd_e2e_proxy_model(dispatcher):
         Utils.destroy_model_parallel()
 
 
+def skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend):
+    if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available():
+        pytest.skip("Deep EP is not available")
+    if moe_flex_dispatcher_backend == "deepepv2" and not is_deep_ep_v2_available():
+        pytest.skip("Deep EP v2 is not available")
+    if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available():
+        pytest.skip("Hybrid EP is not available")
+
+
 @pytest.mark.skipif(
-    not is_deep_ep_available() and not is_hybrid_ep_available(),
-    reason="Deep EP and Hybrid EP are not available",
+    not is_deep_ep_available() and not is_deep_ep_v2_available() and not is_hybrid_ep_available(),
+    reason="Deep EP, Deep EP v2 and Hybrid EP are not available",
 )
 class TestFlexDispatcher:
     def setup_method(self, method):
@@ -652,7 +667,7 @@ def teardown_method(self, method):
     @pytest.mark.internal
     @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)])
     @pytest.mark.parametrize("permute_fusion", permute_fusion_params)
-    @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "hybridep"])
+    @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "deepepv2", "hybridep"])
     @pytest.mark.parametrize("moe_permute_fusion_into_hybridep", [True, False])
     def test_forward_backward(
         self,
@@ -662,10 +677,7 @@ def test_forward_backward(
         moe_flex_dispatcher_backend,
         moe_permute_fusion_into_hybridep,
     ):
-        if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available():
-            pytest.skip("Deep EP is not available")
-        if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available():
-            pytest.skip("Hybrid EP is not available")
+        skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend)
         if moe_permute_fusion_into_hybridep:
             if permute_fusion or moe_flex_dispatcher_backend != "hybridep":
                 pytest.skip(
@@ -696,7 +708,7 @@ def test_forward_backward(
     @pytest.mark.timeout(120)
     @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)])
     @pytest.mark.parametrize("permute_fusion", permute_fusion_params)
-    @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "hybridep"])
+    @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "deepepv2", "hybridep"])
     @pytest.mark.parametrize("moe_permute_fusion_into_hybridep", [True, False])
     def test_capacity_forward_backward(
         self,
@@ -706,10 +718,7 @@ def test_capacity_forward_backward(
         moe_flex_dispatcher_backend,
         moe_permute_fusion_into_hybridep,
     ):
-        if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available():
-            pytest.skip("Deep EP is not available")
-        if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available():
-            pytest.skip("Hybrid EP is not available")
+        skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend)
         if moe_permute_fusion_into_hybridep:
             if permute_fusion or moe_flex_dispatcher_backend != "hybridep":
                 pytest.skip(
@@ -745,7 +754,7 @@ def test_capacity_forward_backward(
     @pytest.mark.timeout(120)
     @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)])
     @pytest.mark.parametrize("permute_fusion", [True])
-    @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "hybridep"])
+    @pytest.mark.parametrize("moe_flex_dispatcher_backend", ["deepep", "deepepv2", "hybridep"])
     @pytest.mark.parametrize("moe_permute_fusion_into_hybridep", [True, False])
     def test_router_padding_for_fp8_forward_backward(
         self,
@@ -755,10 +764,7 @@ def test_router_padding_for_fp8_forward_backward(
         moe_flex_dispatcher_backend,
         moe_permute_fusion_into_hybridep,
     ):
-        if moe_flex_dispatcher_backend == "deepep" and not is_deep_ep_available():
-            pytest.skip("Deep EP is not available")
-        if moe_flex_dispatcher_backend == "hybridep" and not is_hybrid_ep_available():
-            pytest.skip("Hybrid EP is not available")
+        skip_if_flex_backend_unavailable(moe_flex_dispatcher_backend)
         if moe_permute_fusion_into_hybridep:
             if permute_fusion or moe_flex_dispatcher_backend != "hybridep":
                 pytest.skip(

From 91e7c346ebe7c9cd64d456e742678eb00cb92e79 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Wed, 13 May 2026 21:08:44 -0700
Subject: [PATCH 2/7] Make DeepEP v2 manager inherit DeepEP manager

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 .../core/transformer/moe/token_dispatcher.py  | 128 +-----------------
 1 file changed, 2 insertions(+), 126 deletions(-)

diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index 9d76751940d..f00ead80dc5 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1482,7 +1482,7 @@ def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
         return hidden_states
 
 
-class _DeepepV2Manager(_DispatchManager):
+class _DeepepV2Manager(_DeepepManager):
     """
     A manager class for the DeepEP v2 ElasticBuffer backend.
 
@@ -1498,6 +1498,7 @@ def __init__(
         num_experts: int,
         config: TransformerConfig,
     ):
+        # Do not call _DeepepManager.__init__; v2-only images may not ship the v1 Buffer API.
         self.group = group
         self.num_local_experts = num_local_experts
         self.config = config
@@ -1529,18 +1530,6 @@ def _get_buffer(self, hidden_states: torch.Tensor):
         )
         return self.buffer
 
-    def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
-        num_tokens = routing_map.shape[0]
-
-        routing_map = routing_map.reshape(num_tokens, self.num_experts)
-        probs = probs.reshape(num_tokens, self.num_experts)
-        # Convert the format of routing map from multihot to indices.
-        self.token_probs, self.token_indices = torch.topk(probs, self.router_topk, dim=-1)
-        # Mask the indices of dropped tokens with -1
-        if self.capacity_factor is not None:
-            mask = self.token_probs == 0
-            self.token_indices = self.token_indices.masked_fill(mask, -1)
-
     def dispatch(
         self,
         hidden_states: torch.Tensor,
@@ -1595,119 +1584,6 @@ def combine(
         self.dispatched_probs = None
         return hidden_states
 
-    def _indices_to_multihot(self, indices, probs):
-        """
-        Converts a tensor of indices to a multihot vector.
-
-        Args:
-            indices (torch.Tensor): [num_tokens, topk] token indices, where -1 means masked out.
-            probs (torch.Tensor): [num_tokens, topk] token probabilities.
-
-        Returns:
-            A tuple of (routing_map, probs), where routing_map is the multihot vector
-            and probs is the multihot probabilities.
-        """
-        batch_size = indices.shape[0]
-        multihot_routing_map = torch.zeros(
-            (batch_size, self.num_local_experts), dtype=torch.long, device=indices.device
-        )
-
-        multihot_probs = torch.zeros(
-            (batch_size, self.num_local_experts), dtype=torch.float, device=indices.device
-        )
-
-        mask = indices != -1
-        valid_indices = indices[mask]
-        row_indices = torch.arange(batch_size, device=indices.device).repeat_interleave(
-            mask.sum(dim=1)
-        )
-        multihot_routing_map[row_indices, valid_indices] = 1
-        multihot_probs[row_indices, valid_indices] = probs[mask]
-        return multihot_routing_map.bool(), multihot_probs
-
-    def get_number_of_tokens_per_expert(self) -> torch.Tensor:
-        """
-        Get the number of tokens per expert.
-        """
-        return self.tokens_per_expert
-
-    def _pad_routing_map(
-        self, routing_map: torch.Tensor, tokens_per_expert: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Pad the routing map to the nearest multiple of the pad_multiple.
-        """
-        pad_multiple = get_align_size_for_quantization(self.config)
-
-        num_input_tokens = routing_map.shape[0]
-        target_tokens_per_expert = (
-            torch.ceil(tokens_per_expert / pad_multiple) * pad_multiple
-        ).long()
-
-        # Check if there are enough tokens to pad
-        enough_tokens_to_pad = torch.all(target_tokens_per_expert <= num_input_tokens)
-        if not enough_tokens_to_pad:
-            logger.warning(
-                "Not enough tokens to pad. The total number of tokens received in this rank "
-                "is smaller than the target number of tokens for each expert. "
-                "Falling back to explicit padding within GroupedMLP"
-            )
-        else:
-            if is_experimental_enabled() and self.permute_fusion:
-                from megatron.core.fusions.fused_pad_routing_map import fused_pad_routing_map
-
-                routing_map = fused_pad_routing_map(routing_map, pad_multiple)
-            else:
-                routing_map = pad_routing_map(routing_map, pad_multiple)
-            tokens_per_expert = target_tokens_per_expert
-        return routing_map, tokens_per_expert
-
-    def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if is_experimental_enabled() and self.permute_fusion:
-            self.dispatched_routing_map, self.dispatched_probs = fused_indices_to_multihot(
-                self.dispatched_indices, self.dispatched_probs, self.num_local_experts
-            )
-        else:
-            self.dispatched_routing_map, self.dispatched_probs = self._indices_to_multihot(
-                self.dispatched_indices, self.dispatched_probs
-            )
-        if self.config.moe_router_padding_for_quantization:
-            self.dispatched_routing_map, self.tokens_per_expert = self._pad_routing_map(
-                self.dispatched_routing_map, self.tokens_per_expert
-            )
-
-        self.hidden_shape_before_permute = hidden_states.shape
-        assert self.dispatched_probs.dtype == torch.float32, "DeepEP v2 only supports float32 probs"
-        (
-            hidden_states,
-            permuted_probs,
-            self.reversed_mapping_for_combine,
-            self.pad_offsets,
-            self.tokens_per_expert,
-        ) = permute(
-            hidden_states,
-            self.dispatched_routing_map,
-            probs=self.dispatched_probs,
-            num_out_tokens=self.tokens_per_expert.sum().item(),
-            fused=self.permute_fusion,
-            tokens_per_expert=self.tokens_per_expert,
-            align_size=get_align_size_for_quantization(self.config),
-        )
-        if self.router_dtype == "fp64":
-            permuted_probs = permuted_probs.to(torch.float64)
-        return hidden_states, permuted_probs
-
-    def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = unpermute(
-            hidden_states,
-            self.reversed_mapping_for_combine,
-            restore_shape=self.hidden_shape_before_permute,
-            routing_map=self.dispatched_routing_map,
-            fused=self.permute_fusion,
-            pad_offsets=self.pad_offsets,
-        )
-        return hidden_states
-
 
 class MoEFlexTokenDispatcher(MoETokenDispatcher):
     """A flexible token dispatcher that abstracts the underlying tensor and expert

From 522d5c5f805ee1f059a1d3390cdb4a54fb58d0e8 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Wed, 13 May 2026 21:21:47 -0700
Subject: [PATCH 3/7] Restore DispatchManager base interface

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 .../core/transformer/moe/token_dispatcher.py  | 22 ++-----------------
 1 file changed, 2 insertions(+), 20 deletions(-)

diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index f00ead80dc5..baec260f631 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -969,31 +969,18 @@ class _DispatchManager(ABC):
     can be the number of local experts, or the size of sub_group.
     """
 
-    token_probs: Optional[torch.Tensor]
-    dispatched_probs: Optional[torch.Tensor]
-
     @abstractmethod
     def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         """Set up metadata of routing_map and probs."""
         pass
 
     @abstractmethod
-    def dispatch(
-        self,
-        hidden_states: torch.Tensor,
-        async_finish: bool = True,
-        allocate_on_comm_stream: bool = True,
-    ) -> torch.Tensor:
+    def dispatch(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Dispatch the hidden_states according to the routing_map."""
         pass
 
     @abstractmethod
-    def combine(
-        self,
-        hidden_states: torch.Tensor,
-        async_finish: bool = True,
-        allocate_on_comm_stream: bool = True,
-    ) -> torch.Tensor:
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Combine the hidden_states after expert processing."""
         pass
 
@@ -1002,11 +989,6 @@ def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
         """Get the permuted hidden states by instances."""
         pass
 
-    @abstractmethod
-    def get_number_of_tokens_per_expert(self) -> torch.Tensor:
-        """Get the number of tokens per expert."""
-        pass
-
     @abstractmethod
     def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Get the restored hidden states by instances."""

From 58fe1c55e7aa7ee21b45d9ad9f9a9dc0b23445c1 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Wed, 3 Jun 2026 22:54:17 -0700
Subject: [PATCH 4/7] Clarify DeepEP v2 dispatch comments

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 megatron/core/transformer/moe/fused_a2a.py | 58 ++++++++++++++--------
 1 file changed, 37 insertions(+), 21 deletions(-)

diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index 1bb26ec6219..aeaeb13844d 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -113,13 +113,6 @@ def get_elastic_buffer(
     return _elastic_buffer
 
 
-def _capture_elastic_previous_event(buffer, async_finish: bool, allocate_on_comm_stream: bool):
-    """Capture the current stream for DeepEP v2 async communication."""
-    if async_finish and allocate_on_comm_stream:
-        return buffer.capture()
-    return None
-
-
 class FusedDispatch(torch.autograd.Function):
     """Fused dispatch operation for MoE routing combining computation and communication."""
 
@@ -320,7 +313,7 @@ def set_deepep_num_sms(num_sms):
 
 
 class DeepepV2Dispatch(torch.autograd.Function):
-    """DeepEP v2 elastic dispatch with autograd support."""
+    """Dispatch operation using the DeepEP v2 ElasticBuffer backend."""
 
     @staticmethod
     def forward(
@@ -336,10 +329,11 @@ def forward(
         async_finish=False,
         allocate_on_comm_stream=False,
     ):
-        """Forward pass of DeepEP v2 elastic dispatch."""
-        previous_event = _capture_elastic_previous_event(
-            buffer, async_finish, allocate_on_comm_stream
-        )
+        """Forward pass of dispatch using the DeepEP v2 ElasticBuffer backend."""
+        # Capture the current stream for the communication stream to wait on when
+        # DeepEP v2 allocates output tensors on the communication stream.
+        previous_event = buffer.capture() if async_finish and allocate_on_comm_stream else None
+        # Process the dispatch and keep the handle for the subsequent combine call.
         recv_x, recv_token_indices, recv_token_probs, handle, event = buffer.dispatch(
             x,
             topk_idx=token_indices,
@@ -369,9 +363,10 @@ def forward(
     def backward(
         ctx, grad_output, grad_token_indices, grad_token_probs, grad_tokens_per_expert, grad_handle
     ):
-        """Backward pass of DeepEP v2 elastic dispatch."""
-        previous_event = _capture_elastic_previous_event(
-            ctx.buffer, ctx.async_finish, ctx.allocate_on_comm_stream
+        """Backward pass of dispatch using the DeepEP v2 ElasticBuffer backend."""
+        # The backward pass of dispatch is a combine over the dispatch handle.
+        previous_event = (
+            ctx.buffer.capture() if ctx.async_finish and ctx.allocate_on_comm_stream else None
         )
         grad_x, grad_token_probs, event = ctx.buffer.combine(
             grad_output.contiguous(),
@@ -393,9 +388,7 @@ class DeepepV2Combine(torch.autograd.Function):
     @staticmethod
     def forward(ctx, buffer, x, handle, num_sms, async_finish=False, allocate_on_comm_stream=False):
         """Forward pass of DeepEP v2 elastic combine."""
-        previous_event = _capture_elastic_previous_event(
-            buffer, async_finish, allocate_on_comm_stream
-        )
+        previous_event = buffer.capture() if async_finish and allocate_on_comm_stream else None
         combined_x, combined_token_probs, event = buffer.combine(
             x,
             handle=handle,
@@ -417,8 +410,8 @@ def forward(ctx, buffer, x, handle, num_sms, async_finish=False, allocate_on_com
     @staticmethod
     def backward(ctx, grad_output, grad_combined_token_probs):
         """Backward pass of DeepEP v2 elastic combine."""
-        previous_event = _capture_elastic_previous_event(
-            ctx.buffer, ctx.async_finish, ctx.allocate_on_comm_stream
+        previous_event = (
+            ctx.buffer.capture() if ctx.async_finish and ctx.allocate_on_comm_stream else None
         )
         grad_x, _, _, _, event = ctx.buffer.dispatch(
             grad_output.contiguous(),
@@ -447,7 +440,30 @@ def deepepv2_dispatch(
         async_finish=False,
         allocate_on_comm_stream=False,
     ):
-        """Perform DeepEP v2 elastic dispatch."""
+        """Perform dispatch using the DeepEP v2 ElasticBuffer backend.
+
+        Args:
+            buffer (ElasticBuffer):
+                DeepEP v2 buffer used for all-to-all communication.
+            x (torch.Tensor):
+                Input hidden states to dispatch.
+            token_indices (torch.Tensor):
+                Top-k expert indices for each token.
+            token_probs (torch.Tensor):
+                Top-k routing probabilities for each token.
+            num_experts (int):
+                Total number of experts across the communication group.
+            num_max_tokens_per_rank (int):
+                Maximum number of input tokens on each rank.
+            expert_alignment (int):
+                Alignment applied to per-expert token counts.
+            num_sms (int):
+                Number of SMs used by the dispatch API.
+            async_finish (bool):
+                Whether to use asynchronous communication completion.
+            allocate_on_comm_stream (bool):
+                Whether to allocate DeepEP output buffers on the communication stream.
+        """
         return DeepepV2Dispatch.apply(
             buffer,
             x.contiguous(),

From 9bc31941fdd1653f835acb125d6c832802a16752 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Thu, 11 Jun 2026 02:25:45 -0700
Subject: [PATCH 5/7] Fix DeepEP v2 combine backward defaults

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 megatron/core/transformer/moe/fused_a2a.py        |  2 +-
 megatron/core/transformer/moe/token_dispatcher.py | 12 +++++++++---
 megatron/core/transformer/transformer_config.py   |  4 ++--
 tests/unit_tests/models/test_hybrid_moe_model.py  |  2 +-
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index aeaeb13844d..e493e39876a 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -423,7 +423,7 @@ def backward(ctx, grad_output, grad_combined_token_probs):
         )
         if ctx.async_finish:
             event.current_stream_wait()
-        return None, grad_x, None, None, None, None, None
+        return None, grad_x, None, None, None, None
 
 
 if HAVE_DEEP_EP_V2:
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index baec260f631..b540818a551 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1285,7 +1285,7 @@ def __init__(
                 "DeepEP is not installed. Please install DeepEP package from "
                 "https://github.com/deepseek-ai/deepep."
             )
-        set_deepep_num_sms(config.moe_deepep_num_sms)
+        set_deepep_num_sms(20 if config.moe_deepep_num_sms is None else config.moe_deepep_num_sms)
 
     def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         num_tokens = routing_map.shape[0]
@@ -1512,6 +1512,12 @@ def _get_buffer(self, hidden_states: torch.Tensor):
         )
         return self.buffer
 
+    def _get_num_sms(self):
+        """Resolve DeepEP v2 SM count using ElasticBuffer's analytical default when unset."""
+        if self.num_sms is None or self.num_sms == 0:
+            return self.buffer.get_theoretical_num_sms(self.num_experts, self.router_topk)
+        return self.num_sms
+
     def dispatch(
         self,
         hidden_states: torch.Tensor,
@@ -1535,7 +1541,7 @@ def dispatch(
                 self.num_experts,
                 num_max_tokens_per_rank=hidden_states.shape[0],
                 expert_alignment=1,
-                num_sms=self.num_sms,
+                num_sms=self._get_num_sms(),
                 async_finish=async_finish,
                 allocate_on_comm_stream=allocate_on_comm_stream,
             )
@@ -1557,7 +1563,7 @@ def combine(
             self.buffer,
             hidden_states,
             self.handle,
-            num_sms=self.num_sms,
+            num_sms=self._get_num_sms(),
             async_finish=async_finish,
             allocate_on_comm_stream=allocate_on_comm_stream,
         )
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index eb9c8af7dd9..4259bf30cb9 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -935,8 +935,8 @@ class TransformerConfig(ModelParallelConfig):
     moe_latent_size: Optional[int] = None
     """Latent projection dimension for MoE. If None, MoE latent projections are not used."""
 
-    moe_deepep_num_sms: int = 20
-    """Number of SMs to use for DeepEP."""
+    moe_deepep_num_sms: Optional[int] = None
+    """Number of SMs to use for DeepEP. None uses v1's default or v2's theoretical default."""
 
     moe_hybridep_num_sms: Optional[int] = None
     """Number of SMs to use for HybridEP. None uses the default from DeepEP.
diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py
index f3d5e47a103..4c1837cd900 100644
--- a/tests/unit_tests/models/test_hybrid_moe_model.py
+++ b/tests/unit_tests/models/test_hybrid_moe_model.py
@@ -174,7 +174,7 @@
     "mlp_chunks_for_prefill": 1,
     "moe_apply_probs_on_input": False,
     "moe_aux_loss_coeff": 0.0,
-    "moe_deepep_num_sms": 20,
+    "moe_deepep_num_sms": None,
     "moe_enable_deepep": False,
     "moe_expert_capacity_factor": None,
     "moe_expert_rank_capacity_factor": None,

From b99cede44bdb922d4fe5c4c82e1224c81a11bfc5 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Thu, 11 Jun 2026 02:49:14 -0700
Subject: [PATCH 6/7] Simplify DeepEP v2 SM default handling

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 megatron/core/transformer/moe/token_dispatcher.py | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index b540818a551..d90d9d29287 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1504,19 +1504,16 @@ def __init__(
             )
 
     def _get_buffer(self, hidden_states: torch.Tensor):
-        self.buffer = get_elastic_buffer(
+        buffer = get_elastic_buffer(
             self.group,
             num_max_tokens_per_rank=hidden_states.shape[0],
             hidden=hidden_states.shape[1],
             num_topk=self.token_indices.shape[1],
         )
-        return self.buffer
-
-    def _get_num_sms(self):
-        """Resolve DeepEP v2 SM count using ElasticBuffer's analytical default when unset."""
         if self.num_sms is None or self.num_sms == 0:
-            return self.buffer.get_theoretical_num_sms(self.num_experts, self.router_topk)
-        return self.num_sms
+            self.num_sms = buffer.get_theoretical_num_sms(self.num_experts, self.router_topk)
+        self.buffer = buffer
+        return buffer
 
     def dispatch(
         self,
@@ -1541,7 +1538,7 @@ def dispatch(
                 self.num_experts,
                 num_max_tokens_per_rank=hidden_states.shape[0],
                 expert_alignment=1,
-                num_sms=self._get_num_sms(),
+                num_sms=self.num_sms,
                 async_finish=async_finish,
                 allocate_on_comm_stream=allocate_on_comm_stream,
             )
@@ -1563,7 +1560,7 @@ def combine(
             self.buffer,
             hidden_states,
             self.handle,
-            num_sms=self._get_num_sms(),
+            num_sms=self.num_sms,
             async_finish=async_finish,
             allocate_on_comm_stream=allocate_on_comm_stream,
         )

From a71d0d1c0e260b0a8b91e5258a9c703692f04089 Mon Sep 17 00:00:00 2001
From: tongliu <tongliu@nvidia.com>
Date: Thu, 11 Jun 2026 06:22:39 -0700
Subject: [PATCH 7/7] Use DeepEP v2 automatic SM default

Signed-off-by: tongliu <tongliu@nvidia.com>
---
 megatron/core/transformer/moe/fused_a2a.py      |  4 ++--
 .../core/transformer/moe/token_dispatcher.py    | 17 ++++++++++-------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index e493e39876a..0ce27c7f21f 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -352,7 +352,7 @@ def forward(
 
         ctx.buffer = buffer
         ctx.handle = handle
-        ctx.num_sms = num_sms
+        ctx.num_sms = handle.num_sms
         ctx.async_finish = async_finish
         ctx.allocate_on_comm_stream = allocate_on_comm_stream
         tokens_per_expert = torch.tensor(handle.num_recv_tokens_per_expert_list)
@@ -402,7 +402,7 @@ def forward(ctx, buffer, x, handle, num_sms, async_finish=False, allocate_on_com
 
         ctx.buffer = buffer
         ctx.handle = handle
-        ctx.num_sms = num_sms
+        ctx.num_sms = handle.num_sms if num_sms == 0 else num_sms
         ctx.async_finish = async_finish
         ctx.allocate_on_comm_stream = allocate_on_comm_stream
         return combined_x, combined_token_probs
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index d90d9d29287..ed18b6df511 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1285,7 +1285,10 @@ def __init__(
                 "DeepEP is not installed. Please install DeepEP package from "
                 "https://github.com/deepseek-ai/deepep."
             )
-        set_deepep_num_sms(20 if config.moe_deepep_num_sms is None else config.moe_deepep_num_sms)
+        if config.moe_deepep_num_sms is None:
+            set_deepep_num_sms(20)
+        else:
+            set_deepep_num_sms(config.moe_deepep_num_sms)
 
     def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         num_tokens = routing_map.shape[0]
@@ -1490,7 +1493,10 @@ def __init__(
         self.router_dtype = config.moe_router_dtype
         self.capacity_factor = config.moe_expert_capacity_factor
         self.permute_fusion = config.moe_permute_fusion
-        self.num_sms = config.moe_deepep_num_sms
+        if config.moe_deepep_num_sms is None:
+            self.num_sms = 0
+        else:
+            self.num_sms = config.moe_deepep_num_sms
 
         self.token_indices: Optional[torch.Tensor] = None
         self.token_probs: Optional[torch.Tensor] = None
@@ -1504,16 +1510,13 @@ def __init__(
             )
 
     def _get_buffer(self, hidden_states: torch.Tensor):
-        buffer = get_elastic_buffer(
+        self.buffer = get_elastic_buffer(
             self.group,
             num_max_tokens_per_rank=hidden_states.shape[0],
             hidden=hidden_states.shape[1],
             num_topk=self.token_indices.shape[1],
         )
-        if self.num_sms is None or self.num_sms == 0:
-            self.num_sms = buffer.get_theoretical_num_sms(self.num_experts, self.router_topk)
-        self.buffer = buffer
-        return buffer
+        return self.buffer
 
     def dispatch(
         self,