From f28481db8c798d5563f48c74eba71dfc9dbe5237 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 26 May 2026 02:43:07 -0700
Subject: [PATCH 1/7] add thd sequence packing dispatcher support for main

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 megatron/core/model_parallel_config.py        |   7 +
 megatron/core/transformer/moe/fused_a2a.py    |  12 +-
 .../core/transformer/moe/token_dispatcher.py  |  52 ++++-
 .../core/transformer/transformer_config.py    |  33 ++++
 .../transformer/moe/test_token_dispatcher.py  | 177 +++++++++++++++++-
 5 files changed, 273 insertions(+), 8 deletions(-)

diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
index dabe0d0aced..5212200e7f5 100644
--- a/megatron/core/model_parallel_config.py
+++ b/megatron/core/model_parallel_config.py
@@ -62,6 +62,13 @@ class ModelParallelConfig:
     each rank when using hybrid_context_parallel.
     """
 
+    sequence_packing_scheduler: Optional[Literal['dp_balanced', 'default_dynamic_cp']] = None
+    """
+    Scheduler for sequence packing.
+    dp_balanced: DP-balanced scheduler for sequence packing.
+    default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing.
+    """
+
     hybrid_context_parallel: bool = False
     """
     If true, enables hybrid context parallel. This is used to balance the workload of 
diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index defbe996a72..5722ba111e4 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -53,9 +53,12 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int):
         num_nvl_bytes = max(
             config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes
         )
-        num_rdma_bytes = max(
-            config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
-        )
+        # Local-only EP groups do not need an RDMA buffer, and DeepEP builds
+        # without internode support may not expose RDMA size hints.
+        if group.size() > torch.cuda.device_count():
+            num_rdma_bytes = max(
+                config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
+            )
 
     # Allocate buffer if not existed or not enough buffer
     # NOTES: the adaptive routing configuration of the network **must be off**
@@ -334,6 +337,9 @@ def _validate_hybrid_ep_ib_tx_depth(num_tokens: int, group: torch.distributed.Pr
         f"are at most {max_supported_tokens} for multi-node HybridEP."
     )
 
+# HybridEP dispatch/combine kernels use 64-token chunks for their public APIs.
+HYBRIDEP_TOKEN_ALIGNMENT = 64
+
 
 def init_hybrid_ep_buffer(
     group: torch.distributed.ProcessGroup,
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index d7a1772b61e..90a89eb7584 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -18,6 +18,7 @@
 )
 from megatron.core.transformer.enums import CudaGraphModule
 from megatron.core.transformer.moe.fused_a2a import (
+    HYBRIDEP_TOKEN_ALIGNMENT,
     fused_combine,
     fused_dispatch,
     hybrid_ep_combine,
@@ -1032,11 +1033,41 @@ def __init__(
 
         self.moe_expert_rank_capacity_factor = self.config.moe_expert_rank_capacity_factor
         self.over_budget = torch.zeros(1, dtype=torch.bool, device='cuda')
+        # THD sequence packing can produce different token counts per rank.
+        # HybridEP dispatch expects equal per-rank input sizes, so metadata and
+        # hidden states are padded to the group-wide max and trimmed in combine.
+        self._original_num_tokens: Optional[int] = None
+        self._padded_num_tokens: Optional[int] = None
 
     def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         num_tokens = routing_map.shape[0]
-        self.routing_map = routing_map.reshape(num_tokens, self.num_experts)
-        self.token_probs = probs.reshape(num_tokens, self.num_experts)
+        self._original_num_tokens = num_tokens
+
+        padded_num_tokens = num_tokens
+        if self.config.sequence_packing_scheduler is not None:
+            # Use the actual tp_ep max so all ranks in the MoE communication
+            # group pass the same token count to HybridEP.
+            max_num_tokens_across_ep = torch.tensor(
+                [num_tokens], device=routing_map.device, dtype=torch.long
+            )
+            torch.distributed.all_reduce(
+                max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group
+            )
+            padded_num_tokens = int(max_num_tokens_across_ep.item())
+            padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT
+        self._padded_num_tokens = padded_num_tokens
+
+        routing_map = routing_map.reshape(num_tokens, self.num_experts)
+        probs = probs.reshape(num_tokens, self.num_experts)
+        if self.config.sequence_packing_scheduler is not None and padded_num_tokens > num_tokens:
+            pad_rows = padded_num_tokens - num_tokens
+            routing_map = torch.cat(
+                [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0
+            )
+            probs = torch.cat([probs, probs.new_zeros((pad_rows, self.num_experts))], dim=0)
+
+        self.routing_map = routing_map
+        self.token_probs = probs
 
         if self.moe_expert_rank_capacity_factor is not None:
             pad_multiple = get_align_size_for_quantization(self.config)
@@ -1044,7 +1075,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
             # budget). Tokens above this budget are dropped inside HybridEP; dispatch then
             # sets overflow_flag on the handle (accumulated in over_budget in dispatch()).
             budget = int(
-                routing_map.shape[0]
+                padded_num_tokens
                 * self.config.moe_router_topk
                 * self.moe_expert_rank_capacity_factor
             )
@@ -1055,7 +1086,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         # in dispatch) and does not drop tokens or report overflow.
         # Compute the capacity for each expert at the drop_and_pad mode
         if self.drop_and_pad:
-            num_out_tokens = num_tokens * self.config.moe_router_topk
+            num_out_tokens = padded_num_tokens * self.config.moe_router_topk
             # Drop and pad the input to capacity.
             self.capacity = get_capacity(
                 num_tokens=num_out_tokens,
@@ -1084,6 +1115,11 @@ def dispatch(
             self.token_probs = self.token_probs.float()  # downcast or upcast
         if self.config.fp8 or self.config.fp4:
             self.pad_multiple = get_align_size_for_quantization(self.config)
+        if self._padded_num_tokens is not None and hidden_states.shape[0] < self._padded_num_tokens:
+            pad_rows = self._padded_num_tokens - hidden_states.shape[0]
+            hidden_states = torch.cat(
+                [hidden_states, hidden_states.new_zeros((pad_rows, hidden_states.shape[-1]))], dim=0
+            )
         dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = (
             hybrid_ep_dispatch(
                 x=hidden_states,
@@ -1130,12 +1166,20 @@ def combine(
             pad_multiple=self.pad_multiple,
             fused=self.config.moe_permute_fusion_into_hybridep,
         )
+        if (
+            self._padded_num_tokens is not None
+            and self._original_num_tokens is not None
+            and hidden_states.shape[0] > self._original_num_tokens
+        ):
+            hidden_states = hidden_states[: self._original_num_tokens]
         # Release the used handle/num_permuted_tokens which could change in each iteration.
         # For drop_and_pad mode, we don't need to reset the num_permuted_tokens and
         # num_dispatched_tokens, because their values never change.
         self.handle = None
         if not self.drop_and_pad:
             self.num_permuted_tokens = None
+        self._original_num_tokens = None
+        self._padded_num_tokens = None
         return hidden_states
 
     def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 3e91a2b8042..f4939b09702 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -2557,6 +2557,39 @@ def _scope_to_str(s):
                 self.attention_backend == AttnBackend.flash
             ), "Batch invariant mode only supports FlashAttention"
 
+        if self.sequence_packing_scheduler is not None:
+            # Check TE version.
+            if not HAVE_PACKAGING:
+                raise ImportError(
+                    "packaging is not installed. Please install it with `pip install packaging`."
+                )
+            # TODO: remove this after we fix the convergence issue with TE < 2.9.
+            if not (
+                is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a")
+            ):
+                raise ValueError(
+                    "SFT sequence packing requires Transformer Engine >= 2.9.0 "
+                    f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)."
+                )
+
+            # Needed for passing variable sequences between pp stages.
+            self.variable_seq_lengths = True
+
+            assert self.moe_token_dispatcher_type in ("alltoall", "flex"), (
+                f"sequence_packing only supports moe_token_dispatcher_type in "
+                f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'"
+            )
+
+            supported_schedulers = ['dp_balanced', 'default_dynamic_cp']
+            if (
+                self.sequence_packing_scheduler is not None
+                and self.sequence_packing_scheduler not in supported_schedulers
+            ):
+                raise ValueError(
+                    f"Unsupported scheduler: {self.sequence_packing_scheduler}. "
+                    f"Available schedulers: {supported_schedulers}"
+                )
+
 
 @dataclass
 class MLATransformerConfig(TransformerConfig):
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index 46cf639e059..5e1121d976a 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -1,15 +1,22 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
 import dataclasses
+import math
 
 import pytest
 import torch
 
 from megatron.core import config, parallel_state
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_local_submodules,
+    get_gpt_layer_with_transformer_engine_spec,
+)
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer
 from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
 from megatron.core.transformer.moe.moe_utils import get_capacity
 from megatron.core.transformer.spec_utils import get_submodules
+from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
 from megatron.core.utils import is_te_min_version
@@ -418,6 +425,173 @@ def is_hybrid_ep_available():
     return HAVE_HYBRIDEP
 
 
+def _round_up(value, divisor):
+    return value if divisor <= 1 else (value + divisor - 1) // divisor * divisor
+
+
+def _get_thd_padded_seqlens(seqlens, cp_size, tp_size):
+    # This follows the runtime packed-sequence path used by the Moonlight script:
+    # per-sequence lengths must be CP partitionable, and the packed token count
+    # must be even for TP/SP slicing.
+    cp_divisor = 2 * cp_size if cp_size > 1 else 1
+    padded_seqlens = [_round_up(seqlen, cp_divisor) for seqlen in seqlens]
+    total_seqlen = sum(padded_seqlens)
+    total_alignment = math.lcm(cp_divisor, tp_size)
+    padded_seqlens[-1] += _round_up(total_seqlen, total_alignment) - total_seqlen
+    return padded_seqlens
+
+
+def _to_cu_seqlens(seqlens):
+    cu_seqlens = torch.empty(len(seqlens) + 1, dtype=torch.int32, device="cuda")
+    cu_seqlens[0] = 0
+    cu_seqlens[1:] = torch.cumsum(torch.tensor(seqlens, dtype=torch.int32, device="cuda"), dim=0)
+    return cu_seqlens
+
+
+def _make_thd_packed_seq_params(seqlens, cp_size, tp_size):
+    padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size)
+    cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
+    max_seqlen = max(padded_seqlens)
+    # Match get_batch_on_this_rank_for_sequence_packing(): TE consumes padded
+    # cumulative lengths as both cu_seqlens and cu_seqlens_padded for THD.
+    return PackedSeqParams(
+        qkv_format="thd",
+        cu_seqlens_q=cu_seqlens_padded,
+        cu_seqlens_kv=cu_seqlens_padded,
+        cu_seqlens_q_padded=cu_seqlens_padded,
+        cu_seqlens_kv_padded=cu_seqlens_padded,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_kv=max_seqlen,
+    )
+
+
+def _make_sharded_thd_hidden_states(seqlens, hidden_size, cp_size, tp_size, dtype):
+    padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size)
+    padded_sequences = []
+    for seqlen, padded_seqlen in zip(seqlens, padded_seqlens):
+        sequence = torch.randn(seqlen, hidden_size, device="cuda", dtype=dtype)
+        if padded_seqlen > seqlen:
+            sequence = torch.cat(
+                [
+                    sequence,
+                    torch.zeros(padded_seqlen - seqlen, hidden_size, device="cuda", dtype=dtype),
+                ],
+                dim=0,
+            )
+        padded_sequences.append(sequence)
+
+    hidden_states = torch.cat(padded_sequences, dim=0)
+    if cp_size > 1:
+        import transformer_engine_torch as tex
+
+        cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
+        cp_rank = parallel_state.get_context_parallel_rank()
+        index = tex.thd_get_partitioned_indices(
+            cu_seqlens_padded, hidden_states.shape[0], cp_size, cp_rank
+        )
+        hidden_states = hidden_states.index_select(0, index)
+
+    tp_rank = parallel_state.get_tensor_model_parallel_rank()
+    sequence_parallel_length = hidden_states.shape[0] // tp_size
+    hidden_states = hidden_states[
+        tp_rank * sequence_parallel_length : (tp_rank + 1) * sequence_parallel_length
+    ]
+    return hidden_states.unsqueeze(1).contiguous().requires_grad_(True)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    Utils.world_size % 8 != 0, reason="requires world size divisible by 8 for pp2/cp2/tp2/ep2/etp2"
+)
+@pytest.mark.internal
+@pytest.mark.parametrize("dispatcher", ["alltoall", "deepep", "hybridep"])
+def test_sequence_packing_thd_e2e_proxy_model(dispatcher):
+    """Run packed THD attention + MoE forward/backward with major parallelisms enabled."""
+    if not is_te_min_version("2.9.0"):
+        pytest.skip("SFT sequence packing requires Transformer Engine >= 2.9.0")
+    if dispatcher == "deepep" and not is_deep_ep_available():
+        pytest.skip("Deep EP is not available")
+    if dispatcher == "hybridep" and not is_hybrid_ep_available():
+        pytest.skip("Hybrid EP is not available")
+
+    tp_size, pp_size, cp_size, ep_size, etp_size = 2, 2, 2, 2, 2
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=tp_size,
+        pipeline_model_parallel_size=pp_size,
+        context_parallel_size=cp_size,
+        expert_model_parallel_size=ep_size,
+        expert_tensor_parallel_size=etp_size,
+    )
+    _set_random_seed(seed_=123, data_parallel_random_init=False)
+
+    try:
+        spec = get_gpt_layer_with_transformer_engine_spec(num_experts=4, moe_grouped_gemm=False)
+        transformer_config = TransformerConfig(
+            num_layers=4,
+            hidden_size=1024,
+            ffn_hidden_size=2048,
+            moe_ffn_hidden_size=2048,
+            num_attention_heads=8,
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=pp_size,
+            context_parallel_size=cp_size,
+            expert_model_parallel_size=ep_size,
+            expert_tensor_parallel_size=etp_size,
+            sequence_parallel=True,
+            sequence_packing_scheduler="dp_balanced",
+            max_seqlen_per_dp_cp_rank=1024,
+            cp_comm_type="p2p",
+            num_moe_experts=4,
+            moe_router_topk=2,
+            moe_router_load_balancing_type="aux_loss",
+            moe_token_dispatcher_type=(
+                "flex" if dispatcher in ("deepep", "hybridep") else dispatcher
+            ),
+            moe_flex_dispatcher_backend=(
+                dispatcher if dispatcher in ("deepep", "hybridep") else "deepep"
+            ),
+            moe_grouped_gemm=False,
+            moe_router_dtype="fp32",
+            params_dtype=torch.bfloat16,
+            pipeline_dtype=torch.bfloat16,
+            autocast_dtype=torch.bfloat16,
+            bf16=True,
+            add_bias_linear=False,
+            attention_dropout=0.0,
+            hidden_dropout=0.0,
+            use_cpu_initialization=True,
+        )
+        transformer_block = TransformerBlock(transformer_config, spec).cuda().to(torch.bfloat16)
+
+        torch.manual_seed(1000 + torch.distributed.get_rank())
+        seqlens = [257, 509, 1021]
+        hidden_states = _make_sharded_thd_hidden_states(
+            seqlens, transformer_config.hidden_size, cp_size, tp_size, torch.bfloat16
+        )
+        packed_seq_params = _make_thd_packed_seq_params(seqlens, cp_size, tp_size)
+
+        output = transformer_block(
+            hidden_states=hidden_states, attention_mask=None, packed_seq_params=packed_seq_params
+        )
+        assert output.shape == hidden_states.shape
+        assert torch.isfinite(output).all()
+
+        loss = output.float().square().mean()
+        loss.backward()
+
+        assert hidden_states.grad is not None
+        assert hidden_states.grad.shape == hidden_states.shape
+        assert torch.isfinite(hidden_states.grad).all()
+        assert any(
+            param.grad is not None and torch.isfinite(param.grad).all()
+            for param in transformer_block.parameters()
+            if param.requires_grad
+        )
+    finally:
+        reset_hybrid_ep_buffer()
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.skipif(
     not is_deep_ep_available() and not is_hybrid_ep_available(),
     reason="Deep EP and Hybrid EP are not available",
@@ -427,6 +601,7 @@ def setup_method(self, method):
         pass
 
     def teardown_method(self, method):
+        reset_hybrid_ep_buffer()
         Utils.destroy_model_parallel()
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

From 5c46612678cfb81e45efa90b36483c919c6a63e0 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Mon, 22 Jun 2026 23:40:56 -0700
Subject: [PATCH 2/7] cherry pick from pr 5048

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 .../core/transformer/moe/token_dispatcher.py  |  8 +++-
 .../core/transformer/transformer_config.py    | 39 +++----------------
 .../models/test_hybrid_moe_model.py           |  1 +
 3 files changed, 13 insertions(+), 35 deletions(-)

diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index 90a89eb7584..dfc5963551d 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1044,7 +1044,11 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         self._original_num_tokens = num_tokens
 
         padded_num_tokens = num_tokens
-        if self.config.sequence_packing_scheduler is not None:
+        equalize_thd_token_counts = (
+            self.config.sequence_packing_scheduler is not None
+            or self.config.moe_hybridep_pad_variable_tokens
+        )
+        if equalize_thd_token_counts:
             # Use the actual tp_ep max so all ranks in the MoE communication
             # group pass the same token count to HybridEP.
             max_num_tokens_across_ep = torch.tensor(
@@ -1059,7 +1063,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
 
         routing_map = routing_map.reshape(num_tokens, self.num_experts)
         probs = probs.reshape(num_tokens, self.num_experts)
-        if self.config.sequence_packing_scheduler is not None and padded_num_tokens > num_tokens:
+        if equalize_thd_token_counts and padded_num_tokens > num_tokens:
             pad_rows = padded_num_tokens - num_tokens
             routing_map = torch.cat(
                 [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index f4939b09702..2749e0bbf89 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -799,6 +799,12 @@ class TransformerConfig(ModelParallelConfig):
     moe_permute_fusion_into_hybridep: bool = False
     """Fuse token rearrangement ops during token dispatching for HybridEP."""
 
+    moe_hybridep_pad_variable_tokens: bool = False
+    """Pad uneven local token counts to the HybridEP group maximum before dispatch.
+    This is needed when the frontend supplies locally packed THD inputs whose token counts
+    can differ across ranks, without using Megatron Core's sequence_packing_scheduler.
+    """
+    
     moe_per_layer_logging: bool = False
     """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."""
 
@@ -2557,39 +2563,6 @@ def _scope_to_str(s):
                 self.attention_backend == AttnBackend.flash
             ), "Batch invariant mode only supports FlashAttention"
 
-        if self.sequence_packing_scheduler is not None:
-            # Check TE version.
-            if not HAVE_PACKAGING:
-                raise ImportError(
-                    "packaging is not installed. Please install it with `pip install packaging`."
-                )
-            # TODO: remove this after we fix the convergence issue with TE < 2.9.
-            if not (
-                is_te_min_version("2.9.0") or get_te_version() == PkgVersion("2.9.0.dev0+5b3092a")
-            ):
-                raise ValueError(
-                    "SFT sequence packing requires Transformer Engine >= 2.9.0 "
-                    f"but got {get_te_version()} (TE < 2.9.0 may have convergence issues)."
-                )
-
-            # Needed for passing variable sequences between pp stages.
-            self.variable_seq_lengths = True
-
-            assert self.moe_token_dispatcher_type in ("alltoall", "flex"), (
-                f"sequence_packing only supports moe_token_dispatcher_type in "
-                f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'"
-            )
-
-            supported_schedulers = ['dp_balanced', 'default_dynamic_cp']
-            if (
-                self.sequence_packing_scheduler is not None
-                and self.sequence_packing_scheduler not in supported_schedulers
-            ):
-                raise ValueError(
-                    f"Unsupported scheduler: {self.sequence_packing_scheduler}. "
-                    f"Available schedulers: {supported_schedulers}"
-                )
-
 
 @dataclass
 class MLATransformerConfig(TransformerConfig):
diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py
index 59a6da45a1e..30a69aacc21 100644
--- a/tests/unit_tests/models/test_hybrid_moe_model.py
+++ b/tests/unit_tests/models/test_hybrid_moe_model.py
@@ -313,6 +313,7 @@
     "use_transformer_engine_op_fuser": False,
     "moe_single_grouped_weight": False,
     "moe_single_grouped_bias": False,
+    "moe_hybridep_pad_variable_tokens": False,
 }
 # Fields to ignore entirely (ephemeral, environment-specific, very large).
 SKIP_FIELDS = set()

From 241dc4002182b537044c92569bb3f8f862546172 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Jun 2026 00:17:55 -0700
Subject: [PATCH 3/7] fix linting

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 megatron/core/transformer/moe/fused_a2a.py      | 1 +
 megatron/core/transformer/transformer_config.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index 5722ba111e4..c4b02cb5de2 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -337,6 +337,7 @@ def _validate_hybrid_ep_ib_tx_depth(num_tokens: int, group: torch.distributed.Pr
         f"are at most {max_supported_tokens} for multi-node HybridEP."
     )
 
+
 # HybridEP dispatch/combine kernels use 64-token chunks for their public APIs.
 HYBRIDEP_TOKEN_ALIGNMENT = 64
 
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 2749e0bbf89..61d3313564a 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -804,7 +804,7 @@ class TransformerConfig(ModelParallelConfig):
     This is needed when the frontend supplies locally packed THD inputs whose token counts
     can differ across ranks, without using Megatron Core's sequence_packing_scheduler.
     """
-    
+
     moe_per_layer_logging: bool = False
     """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."""
 

From 81e1ddc7089cfa905930ac014991f98859f7d96c Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 7 Jul 2026 23:59:38 -0700
Subject: [PATCH 4/7] refactor

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 megatron/core/model_parallel_config.py        |   7 -
 .../core/transformer/moe/token_dispatcher.py  |  13 +-
 .../core/transformer/transformer_config.py    |   4 +-
 .../transformer/moe/test_token_dispatcher.py  | 175 +-----------------
 4 files changed, 7 insertions(+), 192 deletions(-)

diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
index 5212200e7f5..dabe0d0aced 100644
--- a/megatron/core/model_parallel_config.py
+++ b/megatron/core/model_parallel_config.py
@@ -62,13 +62,6 @@ class ModelParallelConfig:
     each rank when using hybrid_context_parallel.
     """
 
-    sequence_packing_scheduler: Optional[Literal['dp_balanced', 'default_dynamic_cp']] = None
-    """
-    Scheduler for sequence packing.
-    dp_balanced: DP-balanced scheduler for sequence packing.
-    default_dynamic_cp: Dynamic-CP scheduler for packed sequence balancing.
-    """
-
     hybrid_context_parallel: bool = False
     """
     If true, enables hybrid context parallel. This is used to balance the workload of 
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index dfc5963551d..785da64552c 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1033,9 +1033,8 @@ def __init__(
 
         self.moe_expert_rank_capacity_factor = self.config.moe_expert_rank_capacity_factor
         self.over_budget = torch.zeros(1, dtype=torch.bool, device='cuda')
-        # THD sequence packing can produce different token counts per rank.
-        # HybridEP dispatch expects equal per-rank input sizes, so metadata and
-        # hidden states are padded to the group-wide max and trimmed in combine.
+        # HybridEP dispatch expects equal per-rank input sizes. When requested,
+        # variable token counts are padded to the group-wide max and trimmed in combine.
         self._original_num_tokens: Optional[int] = None
         self._padded_num_tokens: Optional[int] = None
 
@@ -1044,11 +1043,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         self._original_num_tokens = num_tokens
 
         padded_num_tokens = num_tokens
-        equalize_thd_token_counts = (
-            self.config.sequence_packing_scheduler is not None
-            or self.config.moe_hybridep_pad_variable_tokens
-        )
-        if equalize_thd_token_counts:
+        if self.config.moe_hybridep_pad_variable_tokens:
             # Use the actual tp_ep max so all ranks in the MoE communication
             # group pass the same token count to HybridEP.
             max_num_tokens_across_ep = torch.tensor(
@@ -1063,7 +1058,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
 
         routing_map = routing_map.reshape(num_tokens, self.num_experts)
         probs = probs.reshape(num_tokens, self.num_experts)
-        if equalize_thd_token_counts and padded_num_tokens > num_tokens:
+        if padded_num_tokens > num_tokens:
             pad_rows = padded_num_tokens - num_tokens
             routing_map = torch.cat(
                 [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 61d3313564a..9aec528aad3 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -801,8 +801,8 @@ class TransformerConfig(ModelParallelConfig):
 
     moe_hybridep_pad_variable_tokens: bool = False
     """Pad uneven local token counts to the HybridEP group maximum before dispatch.
-    This is needed when the frontend supplies locally packed THD inputs whose token counts
-    can differ across ranks, without using Megatron Core's sequence_packing_scheduler.
+    Enable only when HybridEP inputs can have different token counts across ranks.
+    Leave disabled when inputs are already statically padded to equal token counts.
     """
 
     moe_per_layer_logging: bool = False
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index 5e1121d976a..d0fd4228bd2 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -1,22 +1,16 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
 import dataclasses
-import math
 
 import pytest
 import torch
 
 from megatron.core import config, parallel_state
-from megatron.core.models.gpt.gpt_layer_specs import (
-    get_gpt_layer_local_submodules,
-    get_gpt_layer_with_transformer_engine_spec,
-)
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules
 from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer
 from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
 from megatron.core.transformer.moe.moe_utils import get_capacity
 from megatron.core.transformer.spec_utils import get_submodules
-from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
 from megatron.core.utils import is_te_min_version
@@ -425,173 +419,6 @@ def is_hybrid_ep_available():
     return HAVE_HYBRIDEP
 
 
-def _round_up(value, divisor):
-    return value if divisor <= 1 else (value + divisor - 1) // divisor * divisor
-
-
-def _get_thd_padded_seqlens(seqlens, cp_size, tp_size):
-    # This follows the runtime packed-sequence path used by the Moonlight script:
-    # per-sequence lengths must be CP partitionable, and the packed token count
-    # must be even for TP/SP slicing.
-    cp_divisor = 2 * cp_size if cp_size > 1 else 1
-    padded_seqlens = [_round_up(seqlen, cp_divisor) for seqlen in seqlens]
-    total_seqlen = sum(padded_seqlens)
-    total_alignment = math.lcm(cp_divisor, tp_size)
-    padded_seqlens[-1] += _round_up(total_seqlen, total_alignment) - total_seqlen
-    return padded_seqlens
-
-
-def _to_cu_seqlens(seqlens):
-    cu_seqlens = torch.empty(len(seqlens) + 1, dtype=torch.int32, device="cuda")
-    cu_seqlens[0] = 0
-    cu_seqlens[1:] = torch.cumsum(torch.tensor(seqlens, dtype=torch.int32, device="cuda"), dim=0)
-    return cu_seqlens
-
-
-def _make_thd_packed_seq_params(seqlens, cp_size, tp_size):
-    padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size)
-    cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
-    max_seqlen = max(padded_seqlens)
-    # Match get_batch_on_this_rank_for_sequence_packing(): TE consumes padded
-    # cumulative lengths as both cu_seqlens and cu_seqlens_padded for THD.
-    return PackedSeqParams(
-        qkv_format="thd",
-        cu_seqlens_q=cu_seqlens_padded,
-        cu_seqlens_kv=cu_seqlens_padded,
-        cu_seqlens_q_padded=cu_seqlens_padded,
-        cu_seqlens_kv_padded=cu_seqlens_padded,
-        max_seqlen_q=max_seqlen,
-        max_seqlen_kv=max_seqlen,
-    )
-
-
-def _make_sharded_thd_hidden_states(seqlens, hidden_size, cp_size, tp_size, dtype):
-    padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size)
-    padded_sequences = []
-    for seqlen, padded_seqlen in zip(seqlens, padded_seqlens):
-        sequence = torch.randn(seqlen, hidden_size, device="cuda", dtype=dtype)
-        if padded_seqlen > seqlen:
-            sequence = torch.cat(
-                [
-                    sequence,
-                    torch.zeros(padded_seqlen - seqlen, hidden_size, device="cuda", dtype=dtype),
-                ],
-                dim=0,
-            )
-        padded_sequences.append(sequence)
-
-    hidden_states = torch.cat(padded_sequences, dim=0)
-    if cp_size > 1:
-        import transformer_engine_torch as tex
-
-        cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
-        cp_rank = parallel_state.get_context_parallel_rank()
-        index = tex.thd_get_partitioned_indices(
-            cu_seqlens_padded, hidden_states.shape[0], cp_size, cp_rank
-        )
-        hidden_states = hidden_states.index_select(0, index)
-
-    tp_rank = parallel_state.get_tensor_model_parallel_rank()
-    sequence_parallel_length = hidden_states.shape[0] // tp_size
-    hidden_states = hidden_states[
-        tp_rank * sequence_parallel_length : (tp_rank + 1) * sequence_parallel_length
-    ]
-    return hidden_states.unsqueeze(1).contiguous().requires_grad_(True)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-@pytest.mark.skipif(
-    Utils.world_size % 8 != 0, reason="requires world size divisible by 8 for pp2/cp2/tp2/ep2/etp2"
-)
-@pytest.mark.internal
-@pytest.mark.parametrize("dispatcher", ["alltoall", "deepep", "hybridep"])
-def test_sequence_packing_thd_e2e_proxy_model(dispatcher):
-    """Run packed THD attention + MoE forward/backward with major parallelisms enabled."""
-    if not is_te_min_version("2.9.0"):
-        pytest.skip("SFT sequence packing requires Transformer Engine >= 2.9.0")
-    if dispatcher == "deepep" and not is_deep_ep_available():
-        pytest.skip("Deep EP is not available")
-    if dispatcher == "hybridep" and not is_hybrid_ep_available():
-        pytest.skip("Hybrid EP is not available")
-
-    tp_size, pp_size, cp_size, ep_size, etp_size = 2, 2, 2, 2, 2
-    Utils.initialize_model_parallel(
-        tensor_model_parallel_size=tp_size,
-        pipeline_model_parallel_size=pp_size,
-        context_parallel_size=cp_size,
-        expert_model_parallel_size=ep_size,
-        expert_tensor_parallel_size=etp_size,
-    )
-    _set_random_seed(seed_=123, data_parallel_random_init=False)
-
-    try:
-        spec = get_gpt_layer_with_transformer_engine_spec(num_experts=4, moe_grouped_gemm=False)
-        transformer_config = TransformerConfig(
-            num_layers=4,
-            hidden_size=1024,
-            ffn_hidden_size=2048,
-            moe_ffn_hidden_size=2048,
-            num_attention_heads=8,
-            tensor_model_parallel_size=tp_size,
-            pipeline_model_parallel_size=pp_size,
-            context_parallel_size=cp_size,
-            expert_model_parallel_size=ep_size,
-            expert_tensor_parallel_size=etp_size,
-            sequence_parallel=True,
-            sequence_packing_scheduler="dp_balanced",
-            max_seqlen_per_dp_cp_rank=1024,
-            cp_comm_type="p2p",
-            num_moe_experts=4,
-            moe_router_topk=2,
-            moe_router_load_balancing_type="aux_loss",
-            moe_token_dispatcher_type=(
-                "flex" if dispatcher in ("deepep", "hybridep") else dispatcher
-            ),
-            moe_flex_dispatcher_backend=(
-                dispatcher if dispatcher in ("deepep", "hybridep") else "deepep"
-            ),
-            moe_grouped_gemm=False,
-            moe_router_dtype="fp32",
-            params_dtype=torch.bfloat16,
-            pipeline_dtype=torch.bfloat16,
-            autocast_dtype=torch.bfloat16,
-            bf16=True,
-            add_bias_linear=False,
-            attention_dropout=0.0,
-            hidden_dropout=0.0,
-            use_cpu_initialization=True,
-        )
-        transformer_block = TransformerBlock(transformer_config, spec).cuda().to(torch.bfloat16)
-
-        torch.manual_seed(1000 + torch.distributed.get_rank())
-        seqlens = [257, 509, 1021]
-        hidden_states = _make_sharded_thd_hidden_states(
-            seqlens, transformer_config.hidden_size, cp_size, tp_size, torch.bfloat16
-        )
-        packed_seq_params = _make_thd_packed_seq_params(seqlens, cp_size, tp_size)
-
-        output = transformer_block(
-            hidden_states=hidden_states, attention_mask=None, packed_seq_params=packed_seq_params
-        )
-        assert output.shape == hidden_states.shape
-        assert torch.isfinite(output).all()
-
-        loss = output.float().square().mean()
-        loss.backward()
-
-        assert hidden_states.grad is not None
-        assert hidden_states.grad.shape == hidden_states.shape
-        assert torch.isfinite(hidden_states.grad).all()
-        assert any(
-            param.grad is not None and torch.isfinite(param.grad).all()
-            for param in transformer_block.parameters()
-            if param.requires_grad
-        )
-    finally:
-        reset_hybrid_ep_buffer()
-        Utils.destroy_model_parallel()
-
-
 @pytest.mark.skipif(
     not is_deep_ep_available() and not is_hybrid_ep_available(),
     reason="Deep EP and Hybrid EP are not available",

From 71af71d1e9a639781ff9b1ed197b61b2a5c8d4f6 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Fri, 17 Jul 2026 09:07:46 -0700
Subject: [PATCH 5/7] fix

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 megatron/core/transformer/moe/fused_a2a.py    |  9 ++--
 .../transformer/moe/test_token_dispatcher.py  | 47 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py
index c4b02cb5de2..3df3213dba2 100644
--- a/megatron/core/transformer/moe/fused_a2a.py
+++ b/megatron/core/transformer/moe/fused_a2a.py
@@ -53,12 +53,9 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int):
         num_nvl_bytes = max(
             config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes
         )
-        # Local-only EP groups do not need an RDMA buffer, and DeepEP builds
-        # without internode support may not expose RDMA size hints.
-        if group.size() > torch.cuda.device_count():
-            num_rdma_bytes = max(
-                config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
-            )
+        num_rdma_bytes = max(
+            config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
+        )
 
     # Allocate buffer if not existed or not enough buffer
     # NOTES: the adaptive routing configuration of the network **must be off**
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index d0fd4228bd2..3ad53d6f956 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -7,9 +7,10 @@
 
 from megatron.core import config, parallel_state
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules
-from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer
+from megatron.core.transformer.moe.fused_a2a import HYBRIDEP_TOKEN_ALIGNMENT, reset_hybrid_ep_buffer
 from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
 from megatron.core.transformer.moe.moe_utils import get_capacity
+from megatron.core.transformer.moe.token_dispatcher import _HybridEPManager
 from megatron.core.transformer.spec_utils import get_submodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
@@ -419,6 +420,50 @@ def is_hybrid_ep_available():
     return HAVE_HYBRIDEP
 
 
+def test_hybridep_pad_variable_tokens_metadata(monkeypatch):
+    manager = _HybridEPManager.__new__(_HybridEPManager)
+    manager.group = object()
+    manager.num_local_experts = 2
+    manager.num_experts = 4
+    manager.config = TransformerConfig(
+        num_layers=1,
+        hidden_size=16,
+        num_attention_heads=4,
+        num_moe_experts=4,
+        moe_router_topk=2,
+        moe_hybridep_pad_variable_tokens=True,
+    )
+    manager.moe_expert_rank_capacity_factor = None
+    manager.drop_and_pad = False
+
+    local_num_tokens = 17
+    max_num_tokens_across_ep = 70
+    padded_num_tokens = (
+        max_num_tokens_across_ep
+        + -max_num_tokens_across_ep % HYBRIDEP_TOKEN_ALIGNMENT
+    )
+    routing_map = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.bool)
+    probs = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.float32)
+
+    def fake_all_reduce(tensor, op=None, group=None):
+        assert op == torch.distributed.ReduceOp.MAX
+        assert group is manager.group
+        tensor.fill_(max_num_tokens_across_ep)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce)
+
+    manager.setup_metadata(routing_map, probs)
+
+    assert manager._original_num_tokens == local_num_tokens
+    assert manager._padded_num_tokens == padded_num_tokens
+    assert manager.routing_map.shape == (padded_num_tokens, manager.num_experts)
+    assert manager.token_probs.shape == (padded_num_tokens, manager.num_experts)
+    torch.testing.assert_close(manager.routing_map[:local_num_tokens], routing_map)
+    torch.testing.assert_close(manager.token_probs[:local_num_tokens], probs)
+    assert not manager.routing_map[local_num_tokens:].any()
+    assert not manager.token_probs[local_num_tokens:].any()
+
+
 @pytest.mark.skipif(
     not is_deep_ep_available() and not is_hybrid_ep_available(),
     reason="Deep EP and Hybrid EP are not available",

From 856a8a352f74e3826749f2013616d135c96225b4 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Fri, 17 Jul 2026 09:15:25 -0700
Subject: [PATCH 6/7] fix linting

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 tests/unit_tests/transformer/moe/test_token_dispatcher.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index 3ad53d6f956..3a04f778ad5 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -439,8 +439,7 @@ def test_hybridep_pad_variable_tokens_metadata(monkeypatch):
     local_num_tokens = 17
     max_num_tokens_across_ep = 70
     padded_num_tokens = (
-        max_num_tokens_across_ep
-        + -max_num_tokens_across_ep % HYBRIDEP_TOKEN_ALIGNMENT
+        max_num_tokens_across_ep + -max_num_tokens_across_ep % HYBRIDEP_TOKEN_ALIGNMENT
     )
     routing_map = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.bool)
     probs = torch.ones((local_num_tokens, manager.num_experts), dtype=torch.float32)

From d31e98398977ddb1cab40a897faa9f4ff4ea3b78 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 23 Jul 2026 02:24:07 -0700
Subject: [PATCH 7/7] rename flag

Signed-off-by: HaochenYuan <haocheny@nvidia.com>
---
 megatron/core/transformer/moe/token_dispatcher.py        | 2 +-
 megatron/core/transformer/transformer_config.py          | 9 +++++----
 tests/unit_tests/models/test_hybrid_moe_model.py         | 2 +-
 .../unit_tests/transformer/moe/test_token_dispatcher.py  | 4 ++--
 4 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
index 785da64552c..1067d5e4d25 100644
--- a/megatron/core/transformer/moe/token_dispatcher.py
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -1043,7 +1043,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         self._original_num_tokens = num_tokens
 
         padded_num_tokens = num_tokens
-        if self.config.moe_hybridep_pad_variable_tokens:
+        if self.config.moe_hybridep_pad_uneven_dispatch_inputs:
             # Use the actual tp_ep max so all ranks in the MoE communication
             # group pass the same token count to HybridEP.
             max_num_tokens_across_ep = torch.tensor(
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 9aec528aad3..354c3f5ac3f 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -799,10 +799,11 @@ class TransformerConfig(ModelParallelConfig):
     moe_permute_fusion_into_hybridep: bool = False
     """Fuse token rearrangement ops during token dispatching for HybridEP."""
 
-    moe_hybridep_pad_variable_tokens: bool = False
-    """Pad uneven local token counts to the HybridEP group maximum before dispatch.
-    Enable only when HybridEP inputs can have different token counts across ranks.
-    Leave disabled when inputs are already statically padded to equal token counts.
+    moe_hybridep_pad_uneven_dispatch_inputs: bool = False
+    """Pad uneven HybridEP dispatch inputs to the group maximum before dispatch.
+    Enable when local HybridEP input token counts can differ across ranks, for example
+    with dynamically packed THD inputs. Leave disabled when dispatcher inputs are
+    already padded to equal token counts.
     """
 
     moe_per_layer_logging: bool = False
diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py
index 30a69aacc21..4d59552b78e 100644
--- a/tests/unit_tests/models/test_hybrid_moe_model.py
+++ b/tests/unit_tests/models/test_hybrid_moe_model.py
@@ -313,7 +313,7 @@
     "use_transformer_engine_op_fuser": False,
     "moe_single_grouped_weight": False,
     "moe_single_grouped_bias": False,
-    "moe_hybridep_pad_variable_tokens": False,
+    "moe_hybridep_pad_uneven_dispatch_inputs": False,
 }
 # Fields to ignore entirely (ephemeral, environment-specific, very large).
 SKIP_FIELDS = set()
diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
index 3a04f778ad5..96eb22bc778 100644
--- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py
+++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py
@@ -420,7 +420,7 @@ def is_hybrid_ep_available():
     return HAVE_HYBRIDEP
 
 
-def test_hybridep_pad_variable_tokens_metadata(monkeypatch):
+def test_hybridep_pad_uneven_dispatch_inputs_metadata(monkeypatch):
     manager = _HybridEPManager.__new__(_HybridEPManager)
     manager.group = object()
     manager.num_local_experts = 2
@@ -431,7 +431,7 @@ def test_hybridep_pad_variable_tokens_metadata(monkeypatch):
         num_attention_heads=4,
         num_moe_experts=4,
         moe_router_topk=2,
-        moe_hybridep_pad_variable_tokens=True,
+        moe_hybridep_pad_uneven_dispatch_inputs=True,
     )
     manager.moe_expert_rank_capacity_factor = None
     manager.drop_and_pad = False