From 2d66c4e298b293f570e9bf28cb4b8e39c6382e72 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Wed, 20 May 2026 03:01:36 -0700
Subject: [PATCH 01/12] refactor a2a op to reuse existing code from mapping.py

---
 megatron/core/ssm/mamba_context_parallel.py | 36 ++++-----------------
 1 file changed, 7 insertions(+), 29 deletions(-)

diff --git a/megatron/core/ssm/mamba_context_parallel.py b/megatron/core/ssm/mamba_context_parallel.py
index 3297728d5fe..5c040716069 100644
--- a/megatron/core/ssm/mamba_context_parallel.py
+++ b/megatron/core/ssm/mamba_context_parallel.py
@@ -6,7 +6,7 @@
 import torch.nn.functional as F
 
 from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.tensor_parallel import all_to_all
+from megatron.core.tensor_parallel.mappings import all_to_all_hp2sp, all_to_all_sp2hp
 from megatron.core.utils import is_te_min_version
 
 try:
@@ -302,7 +302,6 @@ def _slice_vector_param(self, param: torch.Tensor, has_hdim: bool = False) -> to
         return param[start:end]
 
 
-# TODO(duncan): Consider combining with all_to_all_sp2hp in mappings.py and using einops.rearrange
 def _all_to_all_cp2hp(
     input_: torch.Tensor, cp_group: torch.distributed.ProcessGroup
 ) -> torch.Tensor:
@@ -324,24 +323,12 @@ def _all_to_all_cp2hp(
     """
     assert input_.dim() == 3, "all_to_all_cp2hp assumes 3-d input shape."
     s_in, b_in, h_in = input_.shape
-    # Squash the first two dimensions -> [s*b, h]
-    input_ = input_.reshape(-1, h_in)
-    # Split into world_size chunks along the h dimension
-    world_size = cp_group.size()
-    h_out = h_in // world_size
-    split_tensors = torch.split(input_, split_size_or_sections=h_out, dim=1)
-    # Concat the chunks along the s*b dimension
-    concat_tensor = torch.cat(split_tensors, dim=0)
-    # TODO(duncan): Can the following be optimized by using the non-single (tensor list) version of
-    # all-to-all?
-    # Swap chunks of dim0 across the cp ranks
-    output = all_to_all(cp_group, concat_tensor)
-    # Recover the s and b dimensions
-    output = output.reshape(s_in * world_size, b_in, h_out)
+    s_out, h_out = s_in * cp_group.size(), h_in // cp_group.size()
+    output = all_to_all_sp2hp(input_, group=cp_group)
+    output = output.reshape(s_out, b_in, h_out)
     return output
 
 
-# TODO(duncan): Consider combining with all_to_all_hp2sp in mappings.py and using einops.rearrange
 def _all_to_all_hp2cp(
     input_: torch.Tensor, cp_group: torch.distributed.ProcessGroup
 ) -> torch.Tensor:
@@ -363,18 +350,9 @@ def _all_to_all_hp2cp(
     """
     assert input_.dim() == 3, "all_to_all_hp2cp assumes 3-d input shape."
     s_in, b_in, h_in = input_.shape
-    # Squash the first two dimensions -> [s*b, h]
-    input_ = input_.reshape(-1, h_in)
-    # Swap chunks of dim0 across the cp ranks
-    input_exchanged = all_to_all(cp_group, input_)
-    # Split into world_size chunks along the s*b dimension
-    world_size = cp_group.size()
-    s_out = s_in // world_size
-    split_tensors = torch.split(input_exchanged, split_size_or_sections=s_out * b_in, dim=0)
-    # Concat the chunks along the h dimension
-    output = torch.cat(split_tensors, dim=-1)
-    # Recover the s and b dimensions
-    output = output.reshape(s_out, b_in, h_in * world_size)
+    s_out, h_out = s_in // cp_group.size(), h_in * cp_group.size()
+    output = all_to_all_hp2sp(input_, group=cp_group)
+    output = output.reshape(s_out, b_in, h_out)
     return output
 
 

From 7d8f381636435b424bab30d1cf66c57dd1416bda Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Wed, 20 May 2026 23:05:19 -0700
Subject: [PATCH 02/12] fuse per-seq a2a into a unified one

---
 megatron/core/ssm/gated_delta_net.py         | 104 +++++++---
 tests/unit_tests/ssm/test_gated_delta_net.py | 194 ++++++++++++++++++-
 2 files changed, 269 insertions(+), 29 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index 6c2c3d9fea9..f3a2b7c88b2 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -348,25 +348,31 @@ def forward(
 
         # CP All to All: CP to HP
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-            unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // self.cp_size, dim=0)
-            outputs = []
-            for qkvzba_i in unpacked_qkvzba:
-                qkvzba_i = tensor_a2a_cp2hp(
-                    qkvzba_i,
-                    seq_dim=0,
-                    head_dim=-1,
-                    cp_group=self.pg_collection.cp,
-                    split_sections=[
-                        self.qk_dim_local_tp,
-                        self.qk_dim_local_tp,
-                        self.v_dim_local_tp,
-                        self.v_dim_local_tp,
-                        self.num_value_heads // self.tp_size,
-                        self.num_value_heads // self.tp_size,
-                    ],
+            # Batched: one a2a on the full local THD tensor, then one local
+            # permutation that reorders rank-grouped output into per-sequence
+            # natural order. The permutation also folds in the per-sequence
+            # `_undo_attention_load_balancing`, so it's disabled inside the
+            # a2a call.
+            qkvzba = tensor_a2a_cp2hp(
+                qkvzba,
+                seq_dim=0,
+                head_dim=-1,
+                cp_group=self.pg_collection.cp,
+                split_sections=[
+                    self.qk_dim_local_tp,
+                    self.qk_dim_local_tp,
+                    self.v_dim_local_tp,
+                    self.v_dim_local_tp,
+                    self.num_value_heads // self.tp_size,
+                    self.num_value_heads // self.tp_size,
+                ],
+                undo_attention_load_balancing=False,
+            )
+            if self.cp_size > 1:
+                thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm(
+                    cu_seqlens_q, self.cp_size, seq_len
                 )
-                outputs.append(qkvzba_i)
-            qkvzba = torch.cat(outputs, dim=0)
+                qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx)
         else:
             qkvzba = tensor_a2a_cp2hp(
                 qkvzba,
@@ -495,14 +501,15 @@ def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor):
 
             # CP all to all: HP to CP
             if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-                unpacked_norm_out = _unpack_sequence(norm_out_hp, cu_seqlens_q, dim=0)
-                outputs = []
-                for norm_out_i in unpacked_norm_out:
-                    norm_out_i = tensor_a2a_hp2cp(
-                        norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp
-                    )
-                    outputs.append(norm_out_i)
-                norm_out = torch.cat(outputs, dim=0)
+                if self.cp_size > 1:
+                    norm_out_hp = norm_out_hp.index_select(0, thd_cp_a2a_inv)
+                norm_out = tensor_a2a_hp2cp(
+                    norm_out_hp,
+                    seq_dim=0,
+                    head_dim=-1,
+                    cp_group=self.pg_collection.cp,
+                    redo_attention_load_balancing=False,
+                )
             else:
                 norm_out = tensor_a2a_hp2cp(
                     norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp
@@ -591,7 +598,7 @@ def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta):
 
     def _resolve_cu_seqlens(
         self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1
-    ):
+    ) -> torch.Tensor:
         """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding."""
         if cu_seqlens_padded is not None:
             cu_seqlens = cu_seqlens_padded
@@ -713,7 +720,8 @@ def _backward_out_proj(self):
         self.out_proj.backward_dw()
 
 
-def _unpack_sequence(x, cu_seqlens, dim=1):
+# Used by tests/unit_tests/ssm/test_gated_delta_net.py
+def _unpack_sequence(x, cu_seqlens, dim=1) -> list[torch.Tensor]:
     unpacked_x = []
     cu_seqlens_list = cu_seqlens.tolist()
     num_seqs = len(cu_seqlens_list) - 1
@@ -725,6 +733,46 @@ def _unpack_sequence(x, cu_seqlens, dim=1):
     return unpacked_x
 
 
+def _build_thd_cp_a2a_perm(
+    cu_seqlens: torch.Tensor, cp_size: int, t_global: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    cu = cu_seqlens.to(dtype=torch.long)
+    t_local = t_global // cp_size
+
+    positions = torch.arange(t_global, device=cu.device)
+    seq_idx = torch.bucketize(positions, cu[1:], right=True)
+    seq_lens = torch.diff(cu)
+    halves = seq_lens // (2 * cp_size)  # per-sequence half-chunk size
+    local_starts = cu[:-1] // cp_size
+    global_starts = cu[:-1]
+
+    half_i = halves[seq_idx]
+    pos_in_seq = positions - global_starts[seq_idx]
+
+    natural_chunk = pos_in_seq // half_i  # in [0, 2*cp)
+    offset = pos_in_seq - natural_chunk * half_i
+
+    # Invert the ordering produced by `_undo_attention_load_balancing`:
+    #   natural_chunk < cp:   load_balanced = 2 * natural_chunk
+    #   natural_chunk >= cp:  load_balanced = 4*cp - 2*natural_chunk - 1
+    lb_chunk = torch.where(
+        natural_chunk < cp_size, 2 * natural_chunk, 4 * cp_size - 2 * natural_chunk - 1
+    )
+
+    # In the per-sequence load-balanced layout each rank owns load-balanced
+    # chunks (2r) and (2r+1), in that order, of every sequence.
+    rank = lb_chunk // 2
+    half_within_rank = lb_chunk - 2 * rank
+    k = half_within_rank * half_i + offset
+
+    idx = rank * t_local + local_starts[seq_idx] + k
+
+    inv = torch.empty_like(idx)
+    inv[idx] = positions
+
+    return idx, inv
+
+
 ####################
 # Sharded state dict utilities
 ####################
diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 6af2b3ccd42..f103549dc54 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -17,7 +17,13 @@
 )
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.ssm.gated_delta_net import GatedDeltaNet
+from megatron.core.ssm.gated_delta_net import (
+    GatedDeltaNet,
+    _build_thd_cp_a2a_perm,
+    _unpack_sequence,
+    tensor_a2a_cp2hp,
+    tensor_a2a_hp2cp,
+)
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.core.utils import unwrap_model
@@ -488,3 +494,189 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packi
         micro_batch_size=4,
         sequence_packing=sequence_packing,
     )
+
+
+@pytest.mark.parametrize("cp_size", [2, 4])
+@pytest.mark.internal
+class TestBatchedThdAllToAll:
+    """Verify batched-a2a + permute matches the per-sequence loop in GDN."""
+
+    @pytest.fixture(scope='function', autouse=True)
+    def setup_method(self, cp_size):
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            context_parallel_size=cp_size,
+        )
+        model_parallel_cuda_manual_seed(123)
+        self.cp_size = cp_size
+        self.cp_group = parallel_state.get_context_parallel_group()
+
+    def teardown_method(self):
+        Utils.destroy_model_parallel()
+
+    @staticmethod
+    def _per_seq_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None):
+        cp_size = cp_group.size()
+        unpacked = _unpack_sequence(local_t, cu_seqlens // cp_size, dim=0)
+        outputs = []
+        for x in unpacked:
+            outputs.append(
+                tensor_a2a_cp2hp(
+                    x,
+                    seq_dim=0,
+                    head_dim=-1,
+                    cp_group=cp_group,
+                    split_sections=split_sections,
+                    undo_attention_load_balancing=True,
+                )
+            )
+        return torch.cat(outputs, dim=0)
+
+    @staticmethod
+    def _per_seq_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
+        unpacked = _unpack_sequence(global_t, cu_seqlens, dim=0)
+        outputs = []
+        for x in unpacked:
+            outputs.append(
+                tensor_a2a_hp2cp(
+                    x,
+                    seq_dim=0,
+                    head_dim=-1,
+                    cp_group=cp_group,
+                    split_sections=split_sections,
+                    redo_attention_load_balancing=True,
+                )
+            )
+        return torch.cat(outputs, dim=0)
+
+    # ---- Optimized: single a2a + production permutation helper ----
+
+    @staticmethod
+    def _batched_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None):
+        cp_size = cp_group.size()
+        t_global = int(cu_seqlens[-1].item())
+        naive = tensor_a2a_cp2hp(
+            local_t,
+            seq_dim=0,
+            head_dim=-1,
+            cp_group=cp_group,
+            split_sections=split_sections,
+            undo_attention_load_balancing=False,
+        )
+        idx, _ = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global)
+        return naive.index_select(0, idx)
+
+    @staticmethod
+    def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
+        cp_size = cp_group.size()
+        t_global = int(cu_seqlens[-1].item())
+        _, inv = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global)
+        permuted = global_t.index_select(0, inv)
+        return tensor_a2a_hp2cp(
+            permuted,
+            seq_dim=0,
+            head_dim=-1,
+            cp_group=cp_group,
+            split_sections=split_sections,
+            redo_attention_load_balancing=False,
+        )
+
+    # ---- Tests ----
+
+    @pytest.mark.parametrize(
+        "cu_seqlens",
+        [
+            (0, 32, 64),  # 2 equal sequences
+            (0, 32, 64, 96, 128),  # 4 equal sequences (matches existing THD test)
+            (0, 16, 48, 80),  # 3 unequal sequences
+        ],
+    )
+    def test_cp2hp_batched_matches_per_seq(self, cu_seqlens):
+        cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
+        if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any():
+            pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}")
+
+        T_global = cu_seqlens[-1]
+        T_local = T_global // self.cp_size
+        hidden = 32
+        torch.manual_seed(42 + self.cp_size)
+        local_t = (
+            torch.rand(T_local, 1, hidden, device=torch.cuda.current_device())
+            .bfloat16()
+            .contiguous()
+        )
+
+        out_ref = self._per_seq_a2a_cp2hp(local_t, cu, self.cp_group)
+        out_opt = self._batched_a2a_cp2hp(local_t, cu, self.cp_group)
+
+        rank = torch.distributed.get_rank()
+        assert out_opt.shape == out_ref.shape, (out_opt.shape, out_ref.shape)
+        # Both paths apply the same a2a kernel; only the surrounding pack/cat
+        # differs. Equality should be bitwise.
+        torch.testing.assert_close(
+            out_opt,
+            out_ref,
+            atol=0.0,
+            rtol=0.0,
+            msg=lambda m: f"Batched CP->HP mismatch on rank={rank}: {m}",
+        )
+
+    @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64), (0, 32, 64, 96, 128), (0, 16, 48, 80)])
+    def test_hp2cp_batched_matches_per_seq(self, cu_seqlens):
+        cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
+        if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any():
+            pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}")
+
+        T_global = cu_seqlens[-1]
+        hidden = 32
+        # Hidden must be divisible by cp_size for the HP-sharded input layout.
+        assert hidden % self.cp_size == 0
+        h_local = hidden // self.cp_size
+        torch.manual_seed(42 + self.cp_size)
+        global_t = (
+            torch.rand(T_global, 1, h_local, device=torch.cuda.current_device())
+            .bfloat16()
+            .contiguous()
+        )
+
+        out_ref = self._per_seq_a2a_hp2cp(global_t, cu, self.cp_group)
+        out_opt = self._batched_a2a_hp2cp(global_t, cu, self.cp_group)
+
+        rank = torch.distributed.get_rank()
+        assert out_opt.shape == out_ref.shape, (out_opt.shape, out_ref.shape)
+        torch.testing.assert_close(
+            out_opt,
+            out_ref,
+            atol=0.0,
+            rtol=0.0,
+            msg=lambda m: f"Batched HP->CP mismatch on rank={rank}: {m}",
+        )
+
+    @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64, 96, 128)])
+    def test_cp2hp_hp2cp_round_trip(self, cu_seqlens):
+        """cp2hp followed by hp2cp on the batched path should be the identity."""
+        cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
+        if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any():
+            pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}")
+
+        T_global = cu_seqlens[-1]
+        T_local = T_global // self.cp_size
+        hidden = 32
+        torch.manual_seed(7)
+        local_t = (
+            torch.rand(T_local, 1, hidden, device=torch.cuda.current_device())
+            .bfloat16()
+            .contiguous()
+        )
+
+        mid = self._batched_a2a_cp2hp(local_t, cu, self.cp_group)
+        back = self._batched_a2a_hp2cp(mid, cu, self.cp_group)
+
+        torch.testing.assert_close(
+            back,
+            local_t,
+            atol=0.0,
+            rtol=0.0,
+            msg=lambda m: f"Batched cp2hp -> hp2cp not identity: {m}",
+        )

From 0dd7928d4730b2bc60956a9248bde287c0313540 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Thu, 21 May 2026 02:10:07 -0700
Subject: [PATCH 03/12] use qwen3 model config for testing

---
 tests/unit_tests/ssm/test_gated_delta_net.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index f103549dc54..575219e0f70 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -79,17 +79,18 @@ def setup_method(self, tp_size, sp, cp_size):
 
         # Initialize model
         self.transformer_config = TransformerConfig(
-            hidden_size=256,
-            linear_conv_kernel_dim=2,
-            linear_key_head_dim=64,
-            linear_value_head_dim=64,
-            linear_num_key_heads=4,
-            linear_num_value_heads=8,
+            hidden_size=2048,
+            linear_conv_kernel_dim=4,
+            linear_key_head_dim=128,
+            linear_value_head_dim=128,
+            linear_num_key_heads=16,
+            linear_num_value_heads=32,
             num_layers=1,
             normalization="RMSNorm",
             use_cpu_initialization=True,
             layernorm_zero_centered_gamma=True,
-            num_attention_heads=8,
+            num_attention_heads=16,
+            num_query_groups=2,
             activation_func=F.silu,
             bf16=True,
             tensor_model_parallel_size=tp_size,

From 5ad0a4493c5ee4adf564cdc06be062b90fb8fcf7 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Fri, 22 May 2026 01:55:08 -0700
Subject: [PATCH 04/12] add head perm

---
 megatron/core/ssm/gated_delta_net.py         | 68 +++++++++++++-------
 tests/unit_tests/ssm/test_gated_delta_net.py | 67 +++++++++----------
 2 files changed, 73 insertions(+), 62 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index f3a2b7c88b2..151ac85c9ff 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -139,6 +139,24 @@ def __init__(
         self.qk_dim_local_tp = self.qk_dim // self.tp_size
         self.v_dim_local_tp = self.v_dim // self.tp_size
 
+        if self.cp_size > 1:
+            head_perm = _build_head_perm_for_split_sections(
+                [
+                    self.qk_dim_local_tp,
+                    self.qk_dim_local_tp,
+                    self.v_dim_local_tp,
+                    self.v_dim_local_tp,
+                    self.num_value_heads // self.tp_size,
+                    self.num_value_heads // self.tp_size,
+                ],
+                self.cp_size,
+                torch.cuda.current_device(),
+            )
+        else:
+            head_perm = None
+        # Registered as a non-persistent buffer to exclude it from state_dict
+        self.register_buffer("_thd_head_perm", head_perm, persistent=False)
+
         # Input projection (hidden_states -> q, k, v, gate, beta, alpha)
         # TODO: for now, output gate is forced for GDN.
         # We may remove this restriction in the future.
@@ -347,25 +365,19 @@ def forward(
         nvtx_range_pop(suffix="in_proj")
 
         # CP All to All: CP to HP
+        if self.cp_size > 1:
+            qkvzba = qkvzba.index_select(-1, self._thd_head_perm)
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-            # Batched: one a2a on the full local THD tensor, then one local
-            # permutation that reorders rank-grouped output into per-sequence
-            # natural order. The permutation also folds in the per-sequence
-            # `_undo_attention_load_balancing`, so it's disabled inside the
-            # a2a call.
+            # Batched: one a2a on the full local THD tensor, then two local
+            # permutations -- one on the head dim (so a single fused no-split
+            # a2a still produces the per-channel scatter layout) and one on
+            # the seq dim (rank-grouped -> per-seq natural order, also folds
+            # in `_undo_attention_load_balancing`).
             qkvzba = tensor_a2a_cp2hp(
                 qkvzba,
                 seq_dim=0,
                 head_dim=-1,
                 cp_group=self.pg_collection.cp,
-                split_sections=[
-                    self.qk_dim_local_tp,
-                    self.qk_dim_local_tp,
-                    self.v_dim_local_tp,
-                    self.v_dim_local_tp,
-                    self.num_value_heads // self.tp_size,
-                    self.num_value_heads // self.tp_size,
-                ],
                 undo_attention_load_balancing=False,
             )
             if self.cp_size > 1:
@@ -375,18 +387,7 @@ def forward(
                 qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx)
         else:
             qkvzba = tensor_a2a_cp2hp(
-                qkvzba,
-                seq_dim=0,
-                head_dim=-1,
-                cp_group=self.pg_collection.cp,
-                split_sections=[
-                    self.qk_dim_local_tp,
-                    self.qk_dim_local_tp,
-                    self.v_dim_local_tp,
-                    self.v_dim_local_tp,
-                    self.num_value_heads // self.tp_size,
-                    self.num_value_heads // self.tp_size,
-                ],
+                qkvzba, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp
             )
 
         # Transpose: s b x --> b s x
@@ -773,6 +774,23 @@ def _build_thd_cp_a2a_perm(
     return idx, inv
 
 
+def _build_head_perm_for_split_sections(
+    split_sections: List[int], cp_size: int, device: torch.device
+) -> torch.Tensor:
+    assert all(
+        s % cp_size == 0 for s in split_sections
+    ), f"split_sections {split_sections} must be divisible by cp_size {cp_size} for GDN"
+    offset = 0
+    parts = []
+    for s in split_sections:
+        parts.append(
+            torch.arange(offset, offset + s, device=device, dtype=torch.long).view(cp_size, -1)
+        )
+        offset += s
+
+    return torch.cat(parts, dim=-1).view(-1)
+
+
 ####################
 # Sharded state dict utilities
 ####################
diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 575219e0f70..4b6abbd2a82 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -19,6 +19,7 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.ssm.gated_delta_net import (
     GatedDeltaNet,
+    _build_head_perm_for_split_sections,
     _build_thd_cp_a2a_perm,
     _unpack_sequence,
     tensor_a2a_cp2hp,
@@ -77,7 +78,7 @@ def setup_method(self, tp_size, sp, cp_size):
         cp_group = parallel_state.get_context_parallel_group()
         pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group)
 
-        # Initialize model
+        # Initialize model, with the same config as Qwen Next except `num_layers`
         self.transformer_config = TransformerConfig(
             hidden_size=2048,
             linear_conv_kernel_dim=4,
@@ -499,8 +500,8 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packi
 
 @pytest.mark.parametrize("cp_size", [2, 4])
 @pytest.mark.internal
-class TestBatchedThdAllToAll:
-    """Verify batched-a2a + permute matches the per-sequence loop in GDN."""
+class TestFusedThdAllToAll:
+    """Verify fused 1 AllToAll + permute matches the per-sequence, per-channel loop in GDN."""
 
     @pytest.fixture(scope='function', autouse=True)
     def setup_method(self, cp_size):
@@ -557,12 +558,17 @@ def _per_seq_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
     def _batched_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None):
         cp_size = cp_group.size()
         t_global = int(cu_seqlens[-1].item())
+        if split_sections is not None and cp_size > 1:
+            head_perm = _build_head_perm_for_split_sections(
+                list(split_sections), cp_size, local_t.device
+            )
+            local_t = local_t.index_select(-1, head_perm)
         naive = tensor_a2a_cp2hp(
             local_t,
             seq_dim=0,
             head_dim=-1,
             cp_group=cp_group,
-            split_sections=split_sections,
+            split_sections=None,  # always single fused a2a
             undo_attention_load_balancing=False,
         )
         idx, _ = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global)
@@ -583,8 +589,6 @@ def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
             redo_attention_load_balancing=False,
         )
 
-    # ---- Tests ----
-
     @pytest.mark.parametrize(
         "cu_seqlens",
         [
@@ -593,34 +597,35 @@ def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
             (0, 16, 48, 80),  # 3 unequal sequences
         ],
     )
-    def test_cp2hp_batched_matches_per_seq(self, cu_seqlens):
+    @pytest.mark.parametrize("split_sections", [(8, 8, 4, 4, 4, 4)])
+    @pytest.mark.skip
+    def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections):
         cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
-        if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any():
+        if (torch.diff(cu) % self.cp_size != 0).any():
             pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}")
 
         T_global = cu_seqlens[-1]
         T_local = T_global // self.cp_size
         hidden = 32
-        torch.manual_seed(42 + self.cp_size)
+        if split_sections is not None:
+            assert sum(split_sections) == hidden, (split_sections, hidden)
+        torch.manual_seed(42)
         local_t = (
             torch.rand(T_local, 1, hidden, device=torch.cuda.current_device())
             .bfloat16()
             .contiguous()
         )
 
-        out_ref = self._per_seq_a2a_cp2hp(local_t, cu, self.cp_group)
-        out_opt = self._batched_a2a_cp2hp(local_t, cu, self.cp_group)
+        out_ref = self._per_seq_a2a_cp2hp(
+            local_t, cu, self.cp_group, split_sections=list(split_sections)
+        )
+        out_fused = self._batched_a2a_cp2hp(
+            local_t, cu, self.cp_group, split_sections=list(split_sections)
+        )
 
         rank = torch.distributed.get_rank()
-        assert out_opt.shape == out_ref.shape, (out_opt.shape, out_ref.shape)
-        # Both paths apply the same a2a kernel; only the surrounding pack/cat
-        # differs. Equality should be bitwise.
-        torch.testing.assert_close(
-            out_opt,
-            out_ref,
-            atol=0.0,
-            rtol=0.0,
-            msg=lambda m: f"Batched CP->HP mismatch on rank={rank}: {m}",
+        assert torch.equal(out_fused, out_ref), (
+            f"Batched CP->HP mismatch on rank={rank} " f"(split_sections={split_sections})"
         )
 
     @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64), (0, 32, 64, 96, 128), (0, 16, 48, 80)])
@@ -634,7 +639,7 @@ def test_hp2cp_batched_matches_per_seq(self, cu_seqlens):
         # Hidden must be divisible by cp_size for the HP-sharded input layout.
         assert hidden % self.cp_size == 0
         h_local = hidden // self.cp_size
-        torch.manual_seed(42 + self.cp_size)
+        torch.manual_seed(42)
         global_t = (
             torch.rand(T_global, 1, h_local, device=torch.cuda.current_device())
             .bfloat16()
@@ -642,19 +647,13 @@ def test_hp2cp_batched_matches_per_seq(self, cu_seqlens):
         )
 
         out_ref = self._per_seq_a2a_hp2cp(global_t, cu, self.cp_group)
-        out_opt = self._batched_a2a_hp2cp(global_t, cu, self.cp_group)
+        out_fused = self._batched_a2a_hp2cp(global_t, cu, self.cp_group)
 
         rank = torch.distributed.get_rank()
-        assert out_opt.shape == out_ref.shape, (out_opt.shape, out_ref.shape)
-        torch.testing.assert_close(
-            out_opt,
-            out_ref,
-            atol=0.0,
-            rtol=0.0,
-            msg=lambda m: f"Batched HP->CP mismatch on rank={rank}: {m}",
-        )
+        assert torch.equal(out_fused, out_ref), f"Batched HP->CP mismatch on rank={rank}"
 
     @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64, 96, 128)])
+    @pytest.mark.skip
     def test_cp2hp_hp2cp_round_trip(self, cu_seqlens):
         """cp2hp followed by hp2cp on the batched path should be the identity."""
         cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
@@ -674,10 +673,4 @@ def test_cp2hp_hp2cp_round_trip(self, cu_seqlens):
         mid = self._batched_a2a_cp2hp(local_t, cu, self.cp_group)
         back = self._batched_a2a_hp2cp(mid, cu, self.cp_group)
 
-        torch.testing.assert_close(
-            back,
-            local_t,
-            atol=0.0,
-            rtol=0.0,
-            msg=lambda m: f"Batched cp2hp -> hp2cp not identity: {m}",
-        )
+        assert torch.equal(back, local_t), "Batched cp2hp -> hp2cp not identity"

From a188b247a070fcd6ff8a5e1c6202f74d8bbfc3bf Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Fri, 22 May 2026 20:53:52 -0700
Subject: [PATCH 05/12] fix test

---
 tests/unit_tests/ssm/test_gated_delta_net.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 4b6abbd2a82..35d49e22982 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -498,23 +498,23 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packi
     )
 
 
-@pytest.mark.parametrize("cp_size", [2, 4])
+@pytest.mark.parametrize("cp_size", [2, 4], scope="class")
 @pytest.mark.internal
 class TestFusedThdAllToAll:
     """Verify fused 1 AllToAll + permute matches the per-sequence, per-channel loop in GDN."""
 
-    @pytest.fixture(scope='function', autouse=True)
-    def setup_method(self, cp_size):
+    @pytest.fixture(scope='class', autouse=True)
+    def setup_method(self, request, cp_size):
         Utils.initialize_model_parallel(
             tensor_model_parallel_size=1,
             pipeline_model_parallel_size=1,
             context_parallel_size=cp_size,
         )
         model_parallel_cuda_manual_seed(123)
-        self.cp_size = cp_size
-        self.cp_group = parallel_state.get_context_parallel_group()
-
-    def teardown_method(self):
+        # Attach on the class so every test method can read self.cp_*.
+        request.cls.cp_size = cp_size
+        request.cls.cp_group = parallel_state.get_context_parallel_group()
+        yield
         Utils.destroy_model_parallel()
 
     @staticmethod
@@ -598,7 +598,6 @@ def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
         ],
     )
     @pytest.mark.parametrize("split_sections", [(8, 8, 4, 4, 4, 4)])
-    @pytest.mark.skip
     def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections):
         cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
         if (torch.diff(cu) % self.cp_size != 0).any():
@@ -653,7 +652,6 @@ def test_hp2cp_batched_matches_per_seq(self, cu_seqlens):
         assert torch.equal(out_fused, out_ref), f"Batched HP->CP mismatch on rank={rank}"
 
     @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64, 96, 128)])
-    @pytest.mark.skip
     def test_cp2hp_hp2cp_round_trip(self, cu_seqlens):
         """cp2hp followed by hp2cp on the batched path should be the identity."""
         cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())

From aa3fbf9817a76db2e82a3dad7efe7579ce7460b1 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Tue, 26 May 2026 21:02:56 -0700
Subject: [PATCH 06/12] move unused function to test file

---
 megatron/core/ssm/gated_delta_net.py         | 13 -------------
 tests/unit_tests/ssm/test_gated_delta_net.py | 13 ++++++++++++-
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index 151ac85c9ff..01c37531cfb 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -721,19 +721,6 @@ def _backward_out_proj(self):
         self.out_proj.backward_dw()
 
 
-# Used by tests/unit_tests/ssm/test_gated_delta_net.py
-def _unpack_sequence(x, cu_seqlens, dim=1) -> list[torch.Tensor]:
-    unpacked_x = []
-    cu_seqlens_list = cu_seqlens.tolist()
-    num_seqs = len(cu_seqlens_list) - 1
-    for i in range(num_seqs):
-        idx_start = cu_seqlens_list[i]
-        idx_end = cu_seqlens_list[i + 1]
-        chunked_index = [slice(None)] * dim + [slice(idx_start, idx_end)]
-        unpacked_x.append(x[tuple(chunked_index)])
-    return unpacked_x
-
-
 def _build_thd_cp_a2a_perm(
     cu_seqlens: torch.Tensor, cp_size: int, t_global: int
 ) -> Tuple[torch.Tensor, torch.Tensor]:
diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 35d49e22982..8fa8e56833a 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -21,7 +21,6 @@
     GatedDeltaNet,
     _build_head_perm_for_split_sections,
     _build_thd_cp_a2a_perm,
-    _unpack_sequence,
     tensor_a2a_cp2hp,
     tensor_a2a_hp2cp,
 )
@@ -52,6 +51,18 @@
     HAVE_FLA = False
 
 
+def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[torch.Tensor]:
+    unpacked_x = []
+    cu_seqlens_list = cu_seqlens.tolist()
+    num_seqs = len(cu_seqlens_list) - 1
+    for i in range(num_seqs):
+        idx_start = cu_seqlens_list[i]
+        idx_end = cu_seqlens_list[i + 1]
+        chunked_index = [slice(None)] * dim + [slice(idx_start, idx_end)]
+        unpacked_x.append(x[tuple(chunked_index)])
+    return unpacked_x
+
+
 @pytest.mark.parametrize(
     ("tp_size", "sp", "cp_size"),
     [(1, False, 1), (2, False, 1), (2, True, 1), (1, False, 2), (2, False, 2), (2, True, 2)],

From 757612e1fdcc060d720d2f17a3d8c3cfe3419a3e Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Tue, 26 May 2026 21:43:18 -0700
Subject: [PATCH 07/12] update comments

---
 megatron/core/ssm/gated_delta_net.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index 01c37531cfb..79152eab3a1 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -366,13 +366,9 @@ def forward(
 
         # CP All to All: CP to HP
         if self.cp_size > 1:
+            # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a.
             qkvzba = qkvzba.index_select(-1, self._thd_head_perm)
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-            # Batched: one a2a on the full local THD tensor, then two local
-            # permutations -- one on the head dim (so a single fused no-split
-            # a2a still produces the per-channel scatter layout) and one on
-            # the seq dim (rank-grouped -> per-seq natural order, also folds
-            # in `_undo_attention_load_balancing`).
             qkvzba = tensor_a2a_cp2hp(
                 qkvzba,
                 seq_dim=0,
@@ -381,6 +377,9 @@ def forward(
                 undo_attention_load_balancing=False,
             )
             if self.cp_size > 1:
+                # Permute at the seq dim so that a single unsectioned a2a
+                # is equivalent to per-sequence a2a.
+                # This also folds the ``_undo_attention_load_balancing`` step.
                 thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm(
                     cu_seqlens_q, self.cp_size, seq_len
                 )

From 2a80dcf5fd7ea41677384ec40e2a663bafe21671 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Tue, 2 Jun 2026 17:33:00 -0700
Subject: [PATCH 08/12] move head_perm to forward

---
 megatron/core/ssm/gated_delta_net.py | 32 +++++++++++-----------------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index 79152eab3a1..b4d7b32681a 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -139,24 +139,6 @@ def __init__(
         self.qk_dim_local_tp = self.qk_dim // self.tp_size
         self.v_dim_local_tp = self.v_dim // self.tp_size
 
-        if self.cp_size > 1:
-            head_perm = _build_head_perm_for_split_sections(
-                [
-                    self.qk_dim_local_tp,
-                    self.qk_dim_local_tp,
-                    self.v_dim_local_tp,
-                    self.v_dim_local_tp,
-                    self.num_value_heads // self.tp_size,
-                    self.num_value_heads // self.tp_size,
-                ],
-                self.cp_size,
-                torch.cuda.current_device(),
-            )
-        else:
-            head_perm = None
-        # Registered as a non-persistent buffer to exclude it from state_dict
-        self.register_buffer("_thd_head_perm", head_perm, persistent=False)
-
         # Input projection (hidden_states -> q, k, v, gate, beta, alpha)
         # TODO: for now, output gate is forced for GDN.
         # We may remove this restriction in the future.
@@ -367,7 +349,19 @@ def forward(
         # CP All to All: CP to HP
         if self.cp_size > 1:
             # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a.
-            qkvzba = qkvzba.index_select(-1, self._thd_head_perm)
+            head_perm = _build_head_perm_for_split_sections(
+                [
+                    self.qk_dim_local_tp,
+                    self.qk_dim_local_tp,
+                    self.v_dim_local_tp,
+                    self.v_dim_local_tp,
+                    self.num_value_heads // self.tp_size,
+                    self.num_value_heads // self.tp_size,
+                ],
+                self.pg_collection.cp.size(),
+                torch.cuda.current_device(),
+            )
+            qkvzba = qkvzba.index_select(-1, head_perm)
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
             qkvzba = tensor_a2a_cp2hp(
                 qkvzba,

From a2bf80452890de99d5bee62df3660e664c7f2824 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Tue, 2 Jun 2026 22:50:17 -0700
Subject: [PATCH 09/12] add lru cache

---
 megatron/core/ssm/gated_delta_net.py         |  8 +++++---
 tests/unit_tests/ssm/test_gated_delta_net.py | 16 +++++-----------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index b4d7b32681a..5dae20ed9b0 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -7,6 +7,7 @@
 
 import logging
 from dataclasses import dataclass, replace
+from functools import lru_cache
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -350,14 +351,14 @@ def forward(
         if self.cp_size > 1:
             # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a.
             head_perm = _build_head_perm_for_split_sections(
-                [
+                (
                     self.qk_dim_local_tp,
                     self.qk_dim_local_tp,
                     self.v_dim_local_tp,
                     self.v_dim_local_tp,
                     self.num_value_heads // self.tp_size,
                     self.num_value_heads // self.tp_size,
-                ],
+                ),
                 self.pg_collection.cp.size(),
                 torch.cuda.current_device(),
             )
@@ -754,8 +755,9 @@ def _build_thd_cp_a2a_perm(
     return idx, inv
 
 
+@lru_cache(maxsize=8)
 def _build_head_perm_for_split_sections(
-    split_sections: List[int], cp_size: int, device: torch.device
+    split_sections: Tuple[int], cp_size: int, device: torch.device
 ) -> torch.Tensor:
     assert all(
         s % cp_size == 0 for s in split_sections
diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 8fa8e56833a..f8342753fea 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -570,9 +570,7 @@ def _batched_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None):
         cp_size = cp_group.size()
         t_global = int(cu_seqlens[-1].item())
         if split_sections is not None and cp_size > 1:
-            head_perm = _build_head_perm_for_split_sections(
-                list(split_sections), cp_size, local_t.device
-            )
+            head_perm = _build_head_perm_for_split_sections(split_sections, cp_size, local_t.device)
             local_t = local_t.index_select(-1, head_perm)
         naive = tensor_a2a_cp2hp(
             local_t,
@@ -608,7 +606,7 @@ def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None):
             (0, 16, 48, 80),  # 3 unequal sequences
         ],
     )
-    @pytest.mark.parametrize("split_sections", [(8, 8, 4, 4, 4, 4)])
+    @pytest.mark.parametrize("split_sections", [(8, 8, 4, 16, 32, 4)])
     def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections):
         cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device())
         if (torch.diff(cu) % self.cp_size != 0).any():
@@ -616,9 +614,7 @@ def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections):
 
         T_global = cu_seqlens[-1]
         T_local = T_global // self.cp_size
-        hidden = 32
-        if split_sections is not None:
-            assert sum(split_sections) == hidden, (split_sections, hidden)
+        hidden = sum(split_sections)
         torch.manual_seed(42)
         local_t = (
             torch.rand(T_local, 1, hidden, device=torch.cuda.current_device())
@@ -626,11 +622,9 @@ def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections):
             .contiguous()
         )
 
-        out_ref = self._per_seq_a2a_cp2hp(
-            local_t, cu, self.cp_group, split_sections=list(split_sections)
-        )
+        out_ref = self._per_seq_a2a_cp2hp(local_t, cu, self.cp_group, split_sections=split_sections)
         out_fused = self._batched_a2a_cp2hp(
-            local_t, cu, self.cp_group, split_sections=list(split_sections)
+            local_t, cu, self.cp_group, split_sections=split_sections
         )
 
         rank = torch.distributed.get_rank()

From 41a848fecdb13bccae1da20b0a4dfdafd966c7f4 Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Thu, 4 Jun 2026 17:48:52 -0700
Subject: [PATCH 10/12] type annotation

---
 megatron/core/ssm/gated_delta_net.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index 5dae20ed9b0..33ae38b8715 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -8,7 +8,7 @@
 import logging
 from dataclasses import dataclass, replace
 from functools import lru_cache
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -85,7 +85,7 @@ def __init__(
         conv_bias: bool = False,
         conv_init: Optional[float] = None,
         use_qk_l2norm: bool = True,
-        A_init_range: Tuple[float, float] = (1, 16),
+        A_init_range: tuple[float, float] = (1, 16),
         pg_collection: ProcessGroupCollection = None,
         name: str | None = None,
     ):
@@ -290,7 +290,7 @@ def forward(
                 inference CUDA graphs.
 
         Return:
-            (Tuple[Tensor, Tensor]) GDN output and bias.
+            (tuple[Tensor, Tensor]) GDN output and bias.
 
         """
         # TODO: Deal with attention_mask
@@ -717,7 +717,7 @@ def _backward_out_proj(self):
 
 def _build_thd_cp_a2a_perm(
     cu_seqlens: torch.Tensor, cp_size: int, t_global: int
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     cu = cu_seqlens.to(dtype=torch.long)
     t_local = t_global // cp_size
 
@@ -757,7 +757,7 @@ def _build_thd_cp_a2a_perm(
 
 @lru_cache(maxsize=8)
 def _build_head_perm_for_split_sections(
-    split_sections: Tuple[int], cp_size: int, device: torch.device
+    split_sections: tuple[int], cp_size: int, device: torch.device
 ) -> torch.Tensor:
     assert all(
         s % cp_size == 0 for s in split_sections
@@ -777,7 +777,7 @@ def _build_head_perm_for_split_sections(
 # Sharded state dict utilities
 ####################
 def _split_tensor_factory(
-    orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int
+    orig_sh_ten: ShardedTensor, split_sections: list[int], split_names: list[str], split_dim: int
 ) -> ShardedTensorFactory:
     """Builds a factory that splits a given ShardedTensor into several independent chunks."""
     assert isinstance(orig_sh_ten, ShardedTensor), type(orig_sh_ten)
@@ -843,7 +843,7 @@ def get_parameter_local_cp(
     param: torch.Tensor,
     dim: int,
     cp_group: torch.distributed.ProcessGroup,
-    split_sections: Optional[List[int]] = None,
+    split_sections: Optional[list[int]] = None,
 ) -> torch.Tensor:
     """Get the local parameter for the current context parallel rank.
 
@@ -851,7 +851,7 @@ def get_parameter_local_cp(
         param (torch.Tensor): The entire parameter to get the local parameter for.
         dim (int): The dimension to split the parameter along. Usually the dimension of head.
         cp_group (torch.distributed.ProcessGroup): The context parallel group.
-        split_sections (Optional[List[int]]): If not None,
+        split_sections (Optional[list[int]]): If not None,
             first split the parameter along the dimension dim into sections,
             then get the local hidden parallel weights separately,
             finally concatenate the local hidden parallel weights along the dimension dim.
@@ -889,7 +889,7 @@ def tensor_a2a_cp2hp(
     seq_dim: int,
     head_dim: int,
     cp_group: torch.distributed.ProcessGroup,
-    split_sections: Optional[List[int]] = None,
+    split_sections: Optional[list[int]] = None,
     undo_attention_load_balancing: bool = True,
 ):
     """All-to-all context parallel to hidden parallel.
@@ -900,7 +900,7 @@ def tensor_a2a_cp2hp(
         seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0.
         head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2.
         cp_group (torch.distributed.ProcessGroup): The context parallel group.
-        split_sections (Optional[List[int]]): If not None, split the tensor along the dimension
+        split_sections (Optional[list[int]]): If not None, split the tensor along the dimension
             head_dim into sections first, then do all-to-all for each section separately,
             finally concatenate the separated tensors along the dimension head_dim.
         undo_attention_load_balancing (bool): Whether to undo the attention load balancing of CP.
@@ -952,7 +952,7 @@ def tensor_a2a_hp2cp(
     seq_dim: int,
     head_dim: int,
     cp_group: torch.distributed.ProcessGroup,
-    split_sections: Optional[List[int]] = None,
+    split_sections: Optional[list[int]] = None,
     redo_attention_load_balancing: bool = True,
 ):
     """All-to-all hidden parallel to context parallel.
@@ -963,7 +963,7 @@ def tensor_a2a_hp2cp(
         seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0.
         head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2.
         cp_group (torch.distributed.ProcessGroup): The context parallel group.
-        split_sections (Optional[List[int]]): If not None, first split the tensor along the
+        split_sections (Optional[list[int]]): If not None, first split the tensor along the
             dimension head_dim into sections, then do all-to-all for each section separately,
             finally concatenate the separated tensors along the dimension head_dim.
         redo_attention_load_balancing (bool): Whether to redo the attention load balancing of HP.

From 859a6dc1a8cfb69a634507cc70289a411d7b87cf Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Fri, 5 Jun 2026 07:18:47 -0700
Subject: [PATCH 11/12] fix tuple

---
 megatron/core/ssm/gated_delta_net.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
index 33ae38b8715..2521145c467 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net.py
@@ -757,7 +757,7 @@ def _build_thd_cp_a2a_perm(
 
 @lru_cache(maxsize=8)
 def _build_head_perm_for_split_sections(
-    split_sections: tuple[int], cp_size: int, device: torch.device
+    split_sections: tuple[int, ...], cp_size: int, device: torch.device
 ) -> torch.Tensor:
     assert all(
         s % cp_size == 0 for s in split_sections

From 5d261d15eb0b2111ed96bba24127ffa82687831c Mon Sep 17 00:00:00 2001
From: Xuanteng Huang <xuantengh@nvidia.com>
Date: Mon, 8 Jun 2026 21:31:51 -0700
Subject: [PATCH 12/12] fix test nvls

---
 tests/unit_tests/ssm/test_gated_delta_net.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index f8342753fea..074e7740db2 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
+import os
 from unittest import mock
 
 import pytest
@@ -50,6 +51,10 @@
 except ImportError:
     HAVE_FLA = False
 
+# https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-multi-rank-gpu-enable
+# NVLS doesn't support one single GPU to be shared by multiple ranks, so disable this in test
+os.environ.update({"NCCL_NVLS_ENABLE": "0"})
+
 
 def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[torch.Tensor]:
     unpacked_x = []