diff --git a/megatron/core/ssm/gated_delta_net/__init__.py b/megatron/core/ssm/gated_delta_net/__init__.py
new file mode 100644
index 00000000000..6514f7b3a87
--- /dev/null
+++ b/megatron/core/ssm/gated_delta_net/__init__.py
@@ -0,0 +1,33 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Gated Delta Net (GDN) family of layers.
+
+This package replaces the former ``megatron/core/ssm/gated_delta_net.py`` module
+at the same import path; the names below preserve that module's public surface.
+"""
+
+from megatron.core.ssm.gated_delta_net.common import (
+    HAVE_FLA,
+    GatedDeltaNetSubmodules,
+    causal_conv1d,
+    chunk_gated_delta_rule,
+    get_parameter_local_cp,
+    l2norm,
+    tensor_a2a_cp2hp,
+    tensor_a2a_hp2cp,
+    torch_chunk_gated_delta_rule,
+)
+from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet
+
+__all__ = [
+    "HAVE_FLA",
+    "GatedDeltaNet",
+    "GatedDeltaNetSubmodules",
+    "causal_conv1d",
+    "chunk_gated_delta_rule",
+    "get_parameter_local_cp",
+    "l2norm",
+    "tensor_a2a_cp2hp",
+    "tensor_a2a_hp2cp",
+    "torch_chunk_gated_delta_rule",
+]
diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net/common.py
similarity index 69%
rename from megatron/core/ssm/gated_delta_net.py
rename to megatron/core/ssm/gated_delta_net/common.py
index 06eb0763e57..7ddaf6c3a1b 100644
--- a/megatron/core/ssm/gated_delta_net.py
+++ b/megatron/core/ssm/gated_delta_net/common.py
@@ -1,21 +1,21 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh.
 
 # Some of this code was adopted from https://github.com/huggingface/transformers
 # This source code is licensed under the Apache license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pylint: disable=unused-import
+
 import logging
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import Optional, Union
+from typing import Callable, Optional, Protocol, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch import Tensor
 
-from megatron.core import tensor_parallel
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.jit import jit_fuser
@@ -38,7 +38,7 @@
     make_sharded_tensors_for_checkpoint,
     sharded_state_dict_default,
 )
-from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push
+from megatron.core.utils import nvtx_range_pop, nvtx_range_push
 
 try:
     from fla.modules.convolution import causal_conv1d
@@ -67,13 +67,45 @@ class GatedDeltaNetSubmodules:
     out_proj: Union[ModuleSpec, type] = IdentityOp
 
 
-class GatedDeltaNet(MegatronModule):
-    """Gated Delta Net (GDN) layer class
+class GatedDeltaRuleInterface(Protocol):
+    """
+    Unified typing protocol for GDN core computation interfaces.
+    """
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        scale: float | None = None,
+        initial_state: torch.Tensor | None = None,
+        output_final_state: bool = False,
+        use_qk_l2norm_in_kernel: bool = False,
+        cu_seqlens: torch.LongTensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]: ...
+
+
+class _GDNBase(MegatronModule):
+    """Common base class for the Gated Delta Net (GDN) family of layers.
 
-    GDN layer takes input with size [s, b, h]
-    and returns output of the same size.
+    Hosts everything the GDN variants share: the fused input projection, causal
+    convolution on q/k/v, the CP all-to-all plumbing, the kernel-input preparation
+    skeleton, the gated output norm + projection, and sharded checkpointing.
     """
 
+    dt_bias_dim: int
+    a_log_dim: int
+    in_proj_qkvg_dim: int
+    in_proj_extra_dim: int
+    in_proj_dim: int
+
+    dt_bias: nn.Parameter
+    A_log: nn.Parameter
+
+    gated_delta_rule: GatedDeltaRuleInterface
+
     def __init__(
         self,
         config: TransformerConfig,
@@ -81,11 +113,13 @@ def __init__(
         layer_number: int = None,
         bias: bool = False,
         conv_bias: bool = False,
-        conv_init: Optional[float] = None,
+        conv_init: float | None = None,
         use_qk_l2norm: bool = True,
         A_init_range: tuple[float, float] = (1, 16),
         pg_collection: ProcessGroupCollection = None,
+        *,
         name: str | None = None,
+        cp_comm_type: str | None = None,
     ):
         """
         Args:
@@ -100,11 +134,14 @@ def __init__(
             pg_collection: The required process groups to use for tensor model parallel and context
                 parallel.
             name (str | None): module instance name passed top-down from its paranet module
+            cp_comm_type (Optional[str]): Accepted for TransformerLayer compatibility and
+                ignored; GDN implements context parallelism with its own all-to-alls rather
+                than the attention CP communication schemes.
         """
-
         if not HAVE_FLA:
             raise ImportError(
-                "FLA is not installed. Please install it with `pip install flash-linear-attention`."
+                "FLA is not installed. Please install it with "
+                "`pip install flash-linear-attention[cuda]`."
             )
 
         super().__init__(config)
@@ -139,10 +176,25 @@ def __init__(
         self.qk_dim_local_tp = self.qk_dim // self.tp_size
         self.v_dim_local_tp = self.v_dim // self.tp_size
 
-        # 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.
-        self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.num_value_heads * 2
+        self.num_v_heads_local_tp = self.num_value_heads // self.tp_size
+        self.num_k_heads_local_tp = self.num_key_heads // self.tp_size
+
+        attrs_to_check = (
+            "dt_bias_dim",
+            "a_log_dim",
+            "in_proj_extra_dim",
+            "in_proj_split_names",
+            "in_proj_split_sections",
+            "feat_dim_split",
+            "gated_delta_rule",
+        )
+        self._setup_variant_attrs()
+        for attr in attrs_to_check:
+            assert getattr(self, attr, None) is not None, f"Attribute {attr} for GDN is not set"
+        # QK, V, gate, shared across all variants
+        self.in_proj_qkvg_dim = self.qk_dim * 2 + self.v_dim * 2
+        self.in_proj_dim = self.in_proj_qkvg_dim + self.in_proj_extra_dim
+
         if self.config.fp8:
             fp8_align_size = get_fp8_align_size(self.config.fp8_recipe)
             assert self.in_proj_dim % fp8_align_size == 0, (
@@ -168,8 +220,6 @@ def __init__(
         self.conv_dim = self.qk_dim * 2 + self.v_dim
         self.conv_dim_local_tp = self.conv_dim // self.tp_size
 
-        # weight shape: [conv_dim, 1, d_conv]
-        # bias shape: [conv_dim]
         self.conv1d = nn.Conv1d(
             in_channels=self.conv_dim_local_tp,
             out_channels=self.conv_dim_local_tp,
@@ -186,34 +236,22 @@ def __init__(
             setattr(self.conv1d.bias, "tensor_model_parallel", True)
             setattr(self.conv1d.bias, "partition_dim", 0)
 
-        # Time step projection (discretization)
-        self.num_v_heads_local_tp = self.num_value_heads // self.tp_size
-        # dt_bias parameter
         self.dt_bias = nn.Parameter(
             torch.empty(
-                self.num_v_heads_local_tp,
-                dtype=config.params_dtype,
-                device=torch.cuda.current_device(),
+                self.dt_bias_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device()
             )
         )
         setattr(self.dt_bias, "tensor_model_parallel", True)
         setattr(self.dt_bias, "partition_dim", 0)
-        # A_log parameter
+
         self.A_log = nn.Parameter(
             torch.empty(
-                self.num_v_heads_local_tp,
-                dtype=config.params_dtype,
-                device=torch.cuda.current_device(),
+                self.a_log_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device()
             )
         )
         setattr(self.A_log, "tensor_model_parallel", True)
         setattr(self.A_log, "partition_dim", 0)
 
-        if self.config.deterministic_mode:
-            self.gated_delta_rule = torch_chunk_gated_delta_rule
-        else:
-            self.gated_delta_rule = chunk_gated_delta_rule
-
         # Output layernorm before projection
         self.out_norm = build_module(
             submodules.out_norm,
@@ -243,23 +281,35 @@ def __init__(
 
         self.reset_parameters()
 
+    def _setup_variant_attrs(self):
+        """Set variant specifics on the module. Called once from ``__init__``.
+
+        Must set:
+        - ``in_proj_dim``
+        - ``in_proj_split_names``
+        - ``in_proj_split_sections``
+        - ``feat_dim_split``
+        - ``dt_bias_dim`` / ``a_log_dim`` (sizes of the gate parameters, which the
+          base class creates after the conv1d module to preserve the original
+          parameter registration order)
+        - ``gated_delta_rule`` (the kernel callable).
+        """
+        raise NotImplementedError
+
     def reset_parameters(self):
         """Reset the parameters."""
         if self.config.perform_initialization:
             with get_cuda_rng_tracker().fork():
-                # conv1d.weight
                 if self.conv_init is not None:
                     nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
-                # dt_bias
                 torch.ones(
-                    self.num_v_heads_local_tp,
-                    out=self.dt_bias.data,
+                    self.dt_bias_dim,
                     dtype=self.config.params_dtype,
                     device=torch.cuda.current_device(),
+                    out=self.dt_bias.data,
                 )
-                # A_log
                 A = torch.empty(
-                    self.num_v_heads_local_tp,
+                    self.A_log.shape[0],
                     dtype=self.config.params_dtype,
                     device=torch.cuda.current_device(),
                 ).uniform_(*self.A_init_range)
@@ -267,265 +317,54 @@ def reset_parameters(self):
 
     def forward(
         self,
-        hidden_states: Tensor,
-        attention_mask: Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[int] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         **kwargs,
-    ):
-        """
-        Perform a forward pass through the GDN module.
-
-        Args:
-            hidden_states (Tensor): Hidden states.
-            attention_mask (Tensor): Attention mask.
-            inference_context (Optional[BaseInferenceContext]): Inference context that manages
-                KV cache.
-            packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format.
-            sequence_len_offset (Optional[int]): Sequence length offset used for
-                inference CUDA graphs.
-
-        Return:
-            (tuple[Tensor, Tensor]) GDN output and bias.
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # pylint: disable=missing-function-docstring
+        raise NotImplementedError
 
-        """
-        # TODO: Deal with attention_mask
-
-        inference_context = deprecate_inference_params(inference_context, inference_params)
+    def _gated_norm_and_a2a(
+        self,
+        core_attn_out: torch.Tensor,
+        gate: torch.Tensor,
+        thd_cp_a2a_inv: torch.Tensor | None,
+        batch: int,
+        seq_len: int,
+        packed_seq_params: PackedSeqParams | None = None,
+    ) -> torch.Tensor:
+        # RMSNorm
+        nvtx_range_push(suffix="gated_norm")
+        norm_out_hp = self._apply_gated_norm(core_attn_out, gate)
+        nvtx_range_pop(suffix="gated_norm")
 
-        seq_len, batch, _ = hidden_states.shape
-        seq_len = seq_len * self.sp_size * self.cp_size
-
-        if inference_context is not None:
-            assert (
-                inference_context.is_static_batching()
-            ), "GDN does not currently support dynamic inference batching."
-            assert not self.config.sequence_parallel
-            # TODO: support inference
-            raise NotImplementedError("GDN does not support inference for now.")
+        # Transpose: b s x --> s b x
+        # From bshd back to sbhd format
+        norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1)
+        norm_out_hp = norm_out_hp.transpose(0, 1).contiguous()
 
+        # CP all to all: HP to CP
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-            assert batch == 1, "Packed sequence expects batch dimension to be 1"
-            assert (
-                not self.config.deterministic_mode
-            ), "Packed sequence does not support deterministic mode."
-
-            # Resolve cu_seqlens with alignment padding handling.
-            cu_seqlens_q = self._resolve_cu_seqlens(
-                packed_seq_params.cu_seqlens_q_padded,
-                packed_seq_params.cu_seqlens_q,
-                seq_len,
-                "cu_seqlens_q",
-                cp_size=self.cp_size,
-            )
-            cu_seqlens_kv = self._resolve_cu_seqlens(
-                packed_seq_params.cu_seqlens_kv_padded,
-                packed_seq_params.cu_seqlens_kv,
-                seq_len,
-                "cu_seqlens_kv",
-                cp_size=self.cp_size,
-            )
-            assert torch.equal(cu_seqlens_q, cu_seqlens_kv), (
-                "Currently only support cu_seqlens_q equals to cu_seqlens_kv, "
-                f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}"
-            )
-            num_packed_seqs = cu_seqlens_q.shape[0] - 1
-            assert num_packed_seqs > 0, (
-                "Number of packed sequences must be greater than 0, "
-                f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}"
-            )
-        else:
-            cu_seqlens_q = None
-            cu_seqlens_kv = None
-
-        # Input projection
-        nvtx_range_push(suffix="in_proj")
-        qkvzba, _ = self.in_proj(hidden_states)
-        nvtx_range_pop(suffix="in_proj")
-
-        # 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.
-            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,
+            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,
-                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
-                )
-                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
-            )
-
-        # Transpose: s b x --> b s x
-        # From sbhd to bshd format
-        qkvzba = qkvzba.transpose(0, 1)
-
-        # Split, reorder, and reshape the tensor into q, k, v, gate, beta, alpha
-        qkv, gate, beta, alpha = torch.split(
-            qkvzba,
-            [
-                (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size,
-                self.v_dim_local_tp // self.cp_size,
-                self.num_value_heads // self.tp_size // self.cp_size,
-                self.num_value_heads // self.tp_size // self.cp_size,
-            ],
-            dim=-1,
-        )
-        gate = gate.reshape(batch, seq_len, -1, self.value_head_dim)
-        beta = beta.reshape(batch, seq_len, -1)
-        alpha = alpha.reshape(batch, seq_len, -1)
-
-        # Convolution on qkv
-        nvtx_range_push(suffix="conv1d")
-        seq_len = qkv.shape[1]
-        qkv_channels_split_sections = [
-            self.qk_dim_local_tp,
-            self.qk_dim_local_tp,
-            self.v_dim_local_tp,
-        ]
-        conv1d_weight = get_parameter_local_cp(
-            self.conv1d.weight,
-            dim=0,
-            cp_group=self.pg_collection.cp,
-            split_sections=qkv_channels_split_sections,
-        )
-        conv1d_bias = (
-            get_parameter_local_cp(
-                self.conv1d.bias,
-                dim=0,
-                cp_group=self.pg_collection.cp,
-                split_sections=qkv_channels_split_sections,
-            )
-            if self.conv_bias
-            else None
-        )
-        if self.config.deterministic_mode:
-            qkv = qkv.transpose(1, 2).contiguous()  # b, s, d -> b, d, s
-            conv_out = F.conv1d(
-                input=qkv,  # Torch-native only accept [b, d, s] format input
-                weight=conv1d_weight,
-                bias=conv1d_bias,
-                stride=self.conv1d.stride,
-                padding=self.conv1d.padding,
-                dilation=self.conv1d.dilation,
-                groups=self.conv_dim_local_tp // self.cp_size,
+                redo_attention_load_balancing=False,
             )
-            qkv = self.act_fn(conv_out[..., :seq_len])
-            qkv = qkv.transpose(1, 2)  # b, d, s -> b, s, d
         else:
-            assert self.activation in ["silu", "swish"]
-            qkv, _ = causal_conv1d(
-                x=qkv,  # FLA conv1d accepts [b, s, d] format input
-                weight=conv1d_weight.squeeze(1),  # d, 1, w -> d, w
-                bias=conv1d_bias,
-                activation=self.activation,
-                initial_state=None,
-                output_final_state=False,
-                cu_seqlens=cu_seqlens_q,
+            norm_out = tensor_a2a_hp2cp(
+                norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp
             )
-        nvtx_range_pop(suffix="conv1d")
-
-        # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous)
-        nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule")
-        query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule(
-            qkv, gate, beta, alpha, batch, seq_len
-        )
-        nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule")
-
-        # Calculate g and beta
-        nvtx_range_push(suffix="g_and_beta")
-        A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp)
-        dt_bias_local_cp = get_parameter_local_cp(
-            self.dt_bias, dim=0, cp_group=self.pg_collection.cp
-        )
-        g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta)
-        nvtx_range_pop(suffix="g_and_beta")
-
-        nvtx_range_push(suffix="gated_delta_rule")
-        core_attn_out, last_recurrent_state = self.gated_delta_rule(
-            query,
-            key,
-            value,
-            g=g,
-            beta=beta,
-            initial_state=None,
-            output_final_state=False,
-            use_qk_l2norm_in_kernel=False,
-            cu_seqlens=cu_seqlens_q,
-        )
-        nvtx_range_pop(suffix="gated_delta_rule")
-
-        def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor):
-            # RMSNorm
-            nvtx_range_push(suffix="gated_norm")
-            norm_out_hp = self._apply_gated_norm(core_attn_out, gate)
-            nvtx_range_pop(suffix="gated_norm")
-
-            # Transpose: b s x --> s b x
-            # From bshd back to sbhd format
-            norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1)
-            norm_out_hp = norm_out_hp.transpose(0, 1).contiguous()
-
-            # CP all to all: HP to CP
-            if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
-                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
-                )
 
-            return norm_out
-
-        if self.recompute_norm_out:
-            self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput()
-            norm_out = self.norm_out_checkpoint.checkpoint(_gated_norm_and_a2a, core_attn_out, gate)
-        else:
-            norm_out = _gated_norm_and_a2a(core_attn_out, gate)
-
-        # Output projection
-        nvtx_range_push(suffix="out_proj")
-        out, out_bias = self.out_proj(norm_out)
-        nvtx_range_pop(suffix="out_proj")
-
-        if self.recompute_norm_out:
-            self.norm_out_checkpoint.discard_output_and_register_recompute(out)
-
-        return out, out_bias
+        return norm_out
 
     @jit_fuser
     def _apply_gated_norm(self, x, gate):
@@ -540,10 +379,21 @@ def _apply_gated_norm(self, x, gate):
         return y
 
     @jit_fuser
-    def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len):
+    def _prepare_input_for_gated_delta_rule(
+        self,
+        qkv: torch.Tensor,
+        gate: torch.Tensor,
+        batch: int,
+        seq_len: int,
+        *gate_feats: tuple[torch.Tensor],
+    ) -> tuple[torch.Tensor, ...]:
         """
-        Prepare query, key, value, gate, beta, alpha tensors for gated delta rule.
+        Prepare the query, key, value, gate, and variant gate-feature tensors for the
+        gated delta rule kernels.
+
         Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations.
+        ``gate_feats`` holds the variant-specific in_proj sections, which are returned
+        contiguous for the decay/gating computation in ``forward``.
         """
         # Split qkv into query_key and value
         query_key, value = torch.split(
@@ -575,13 +425,18 @@ def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_l
         key = key.contiguous()
         value = value.contiguous()
         gate = gate.contiguous()
-        beta = beta.contiguous()
-        alpha = alpha.contiguous()
+        gate_feats = tuple(t.contiguous() for t in gate_feats)
 
-        return query, key, value, gate, beta, alpha
+        return query, key, value, gate, *gate_feats
 
     @jit_fuser
-    def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta):
+    def _compute_g_and_beta(
+        self,
+        A_log_local_cp: torch.Tensor,
+        dt_bias_local_cp: torch.Tensor,
+        alpha: torch.Tensor,
+        beta: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Compute g (decay) and beta (sigmoid) for gated delta rule.
         Fuses exp, softplus, mul, neg, and sigmoid operations.
@@ -669,15 +524,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_gr
 
         sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory(
             sharded_state_dict[f"{prefix}in_proj.weight"],
-            [
-                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,
-            ],
-            ["query", "key", "value", "z", "beta", "alpha"],
+            list(self.in_proj_split_sections),
+            self.in_proj_split_names,
             0,
         )
 
@@ -951,9 +799,9 @@ def tensor_a2a_hp2cp(
 # Torch native gated delta rule
 ####################
 def torch_chunk_gated_delta_rule(
-    query,
-    key,
-    value,
+    q,
+    k,
+    v,
     g,
     beta,
     chunk_size=64,
@@ -961,7 +809,7 @@ def torch_chunk_gated_delta_rule(
     output_final_state=False,
     use_qk_l2norm_in_kernel=False,
     cu_seqlens=None,
-):
+) -> tuple[torch.Tensor, torch.Tensor | None]:
     # pylint: disable=line-too-long
     '''
     Torch-native implementation of chunked gated delta rule for deterministic mode.
@@ -974,6 +822,7 @@ def torch_chunk_gated_delta_rule(
         cu_seqlens is None
     ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now."
 
+    query, key, value = q, k, v
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
         query = l2norm(query, dim=-1, eps=1e-6)
diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py
new file mode 100644
index 00000000000..65d9dc7df0a
--- /dev/null
+++ b/megatron/core/ssm/gated_delta_net/gdn.py
@@ -0,0 +1,278 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh.
+
+# Some of this code was adopted from https://github.com/huggingface/transformers
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+
+from functools import partial
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from megatron.core import tensor_parallel
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.ssm.gated_delta_net.common import (
+    _build_head_perm_for_split_sections,
+    _build_thd_cp_a2a_perm,
+    _GDNBase,
+    causal_conv1d,
+    chunk_gated_delta_rule,
+    get_parameter_local_cp,
+    tensor_a2a_cp2hp,
+    torch_chunk_gated_delta_rule,
+)
+from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push
+
+
+class GatedDeltaNet(_GDNBase):
+    # pylint: disable=missing-class-docstring
+    def _setup_variant_attrs(self):
+        """Set the GDN in_proj sizing, split tables, gate parameter dims, and kernel."""
+        # alpha, beta
+        self.in_proj_extra_dim = self.num_value_heads * 2
+
+        # Per-section sizes (and names) of the in_proj output, local to this TP rank.
+        # Used for the CP head permutation (pre-a2a), for splitting the projection
+        # output (post-a2a), and for the sharded checkpoint split of in_proj.weight.
+        self.in_proj_split_names = ["query", "key", "value", "z", "beta", "alpha"]
+        self.in_proj_split_sections = (
+            self.qk_dim_local_tp,  # q
+            self.qk_dim_local_tp,  # k
+            self.v_dim_local_tp,  # v
+            self.v_dim_local_tp,  # gate (z)
+            self.num_value_heads // self.tp_size,  # beta
+            self.num_value_heads // self.tp_size,  # alpha
+        )
+        self.feat_dim_split = (
+            (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size,  # qkv
+            self.v_dim_local_tp // self.cp_size,  # gate (z)
+            self.num_value_heads // self.tp_size // self.cp_size,  # beta
+            self.num_value_heads // self.tp_size // self.cp_size,  # alpha
+        )
+
+        self.dt_bias_dim = self.num_v_heads_local_tp
+        self.a_log_dim = self.num_v_heads_local_tp
+
+        if self.config.deterministic_mode:
+            self.gated_delta_rule = torch_chunk_gated_delta_rule
+        else:
+            self.gated_delta_rule = chunk_gated_delta_rule
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        inference_context: Optional[BaseInferenceContext] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+        sequence_len_offset: Optional[int] = None,
+        *,
+        inference_params: Optional[BaseInferenceContext] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Perform a forward pass through the GDN module.
+
+        Return:
+            (tuple[torch.Tensor, torch.Tensor]) GDN output and bias.
+        """
+
+        inference_context = deprecate_inference_params(inference_context, inference_params)
+
+        seq_len, batch, _ = hidden_states.shape
+        seq_len = seq_len * self.sp_size * self.cp_size
+
+        if inference_context is not None:
+            assert (
+                inference_context.is_static_batching()
+            ), "GDN does not currently support dynamic inference batching."
+            assert not self.config.sequence_parallel
+            # TODO: support inference
+            raise NotImplementedError("GDN does not support inference for now.")
+
+        if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
+            assert batch == 1, "Packed sequence expects batch dimension to be 1"
+            assert (
+                not self.config.deterministic_mode
+            ), "Packed sequence does not support deterministic mode."
+
+            # Resolve cu_seqlens with alignment padding handling.
+            cu_seqlens_q = self._resolve_cu_seqlens(
+                packed_seq_params.cu_seqlens_q_padded,
+                packed_seq_params.cu_seqlens_q,
+                seq_len,
+                "cu_seqlens_q",
+                cp_size=self.cp_size,
+            )
+            cu_seqlens_kv = self._resolve_cu_seqlens(
+                packed_seq_params.cu_seqlens_kv_padded,
+                packed_seq_params.cu_seqlens_kv,
+                seq_len,
+                "cu_seqlens_kv",
+                cp_size=self.cp_size,
+            )
+            assert torch.equal(cu_seqlens_q, cu_seqlens_kv), (
+                "Currently only support cu_seqlens_q equals to cu_seqlens_kv, "
+                f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}"
+            )
+            num_packed_seqs = cu_seqlens_q.shape[0] - 1
+            assert num_packed_seqs > 0, (
+                "Number of packed sequences must be greater than 0, "
+                f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}"
+            )
+        else:
+            cu_seqlens_q = None
+            cu_seqlens_kv = None
+
+        # Input projection
+        nvtx_range_push(suffix="in_proj")
+        qkvzba, _ = self.in_proj(hidden_states)
+        nvtx_range_pop(suffix="in_proj")
+
+        # 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.
+            head_perm = _build_head_perm_for_split_sections(
+                self.in_proj_split_sections,
+                self.pg_collection.cp.size(),
+                torch.cuda.current_device(),
+            )
+            qkvzba = qkvzba.index_select(-1, head_perm)
+
+        thd_cp_a2a_idx, thd_cp_a2a_inv = None, None
+        if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
+            qkvzba = tensor_a2a_cp2hp(
+                qkvzba,
+                seq_dim=0,
+                head_dim=-1,
+                cp_group=self.pg_collection.cp,
+                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
+                )
+                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
+            )
+
+        # Transpose: s b x --> b s x
+        # From sbhd to bshd format
+        qkvzba = qkvzba.transpose(0, 1)
+
+        # Split the tensor into q, k, v, gate (z), and the variant-specific gate features
+        # (beta, alpha for GDN; f, b, w for GDN2)
+        qkv, gate, beta, alpha = torch.split(qkvzba, self.feat_dim_split, dim=-1)
+        gate = gate.reshape(batch, seq_len, -1, self.value_head_dim)
+
+        # Convolution on qkv
+        nvtx_range_push(suffix="conv1d")
+        seq_len = qkv.shape[1]
+        qkv_channels_split_sections = [
+            self.qk_dim_local_tp,
+            self.qk_dim_local_tp,
+            self.v_dim_local_tp,
+        ]
+        conv1d_weight = get_parameter_local_cp(
+            self.conv1d.weight,
+            dim=0,
+            cp_group=self.pg_collection.cp,
+            split_sections=qkv_channels_split_sections,
+        )
+        conv1d_bias = (
+            get_parameter_local_cp(
+                self.conv1d.bias,
+                dim=0,
+                cp_group=self.pg_collection.cp,
+                split_sections=qkv_channels_split_sections,
+            )
+            if self.conv_bias
+            else None
+        )
+        if self.config.deterministic_mode:
+            qkv = qkv.transpose(1, 2).contiguous()  # b, s, d -> b, d, s
+            conv_out = F.conv1d(
+                input=qkv,  # Torch-native only accept [b, d, s] format input
+                weight=conv1d_weight,
+                bias=conv1d_bias,
+                stride=self.conv1d.stride,
+                padding=self.conv1d.padding,
+                dilation=self.conv1d.dilation,
+                groups=self.conv_dim_local_tp // self.cp_size,
+            )
+            qkv = self.act_fn(conv_out[..., :seq_len])
+            qkv = qkv.transpose(1, 2)  # b, d, s -> b, s, d
+        else:
+            assert self.activation in ["silu", "swish"]
+            qkv, _ = causal_conv1d(
+                x=qkv,  # FLA conv1d accepts [b, s, d] format input
+                weight=conv1d_weight.squeeze(1),  # d, 1, w -> d, w
+                bias=conv1d_bias,
+                activation=self.activation,
+                initial_state=None,
+                output_final_state=False,
+                cu_seqlens=cu_seqlens_q,
+            )
+        nvtx_range_pop(suffix="conv1d")
+
+        A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp)
+        dt_bias_local_cp = get_parameter_local_cp(
+            self.dt_bias, dim=0, cp_group=self.pg_collection.cp
+        )
+
+        # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous)
+        nvtx_range_push(suffix="prepare_input_for_gated_delta_rule")
+        query, key, value, gate, beta, alpha = self._prepare_input_for_gated_delta_rule(
+            qkv, gate, batch, seq_len, beta, alpha
+        )
+        nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule")
+
+        nvtx_range_push(suffix="g_and_beta")
+        g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta)
+        nvtx_range_pop(suffix="g_and_beta")
+
+        nvtx_range_push(suffix="gated_delta_rule")
+        core_attn_out, _ = self.gated_delta_rule(
+            query,
+            key,
+            value,
+            g=g,
+            beta=beta,
+            initial_state=None,
+            output_final_state=False,
+            use_qk_l2norm_in_kernel=False,
+            cu_seqlens=cu_seqlens_q,
+        )
+        nvtx_range_pop(suffix="gated_delta_rule")
+
+        if self.recompute_norm_out:
+            self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            norm_func = partial(
+                self._gated_norm_and_a2a,
+                thd_cp_a2a_inv=thd_cp_a2a_inv,
+                batch=batch,
+                seq_len=seq_len,
+                packed_seq_params=packed_seq_params,
+            )
+            norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate)
+        else:
+            norm_out = self._gated_norm_and_a2a(
+                core_attn_out, gate, thd_cp_a2a_inv, batch, seq_len, packed_seq_params
+            )
+
+        # Output projection
+        nvtx_range_push(suffix="out_proj")
+        out, out_bias = self.out_proj(norm_out)
+        nvtx_range_pop(suffix="out_proj")
+
+        if self.recompute_norm_out:
+            self.norm_out_checkpoint.discard_output_and_register_recompute(out)
+
+        return out, out_bias
diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py
index 074e7740db2..7cd2eb5e104 100644
--- a/tests/unit_tests/ssm/test_gated_delta_net.py
+++ b/tests/unit_tests/ssm/test_gated_delta_net.py
@@ -2,41 +2,27 @@
 
 import copy
 import os
-from unittest import mock
 
 import pytest
 import torch
 import torch.nn.functional as F
 
 from megatron.core import parallel_state
-from megatron.core.models.common.embeddings.rope_utils import (
-    get_pos_emb_on_this_cp_rank as get_tensor_on_this_cp_rank,
-)
 from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
     get_experimental_attention_variant_module_spec,
     get_transformer_block_with_experimental_attention_variant_spec,
 )
-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
+from megatron.core.ssm.gated_delta_net.common import (
     _build_head_perm_for_split_sections,
     _build_thd_cp_a2a_perm,
     tensor_a2a_cp2hp,
     tensor_a2a_hp2cp,
+    torch_chunk_gated_delta_rule,
 )
 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
-from megatron.training.arguments import parse_args
-from megatron.training.checkpointing import load_checkpoint, save_checkpoint
-from megatron.training.global_vars import set_args
-from megatron.training.training import get_model
-from tests.unit_tests.dist_checkpointing import (
-    TempNamedDir,
-    init_basic_mock_args,
-    init_checkpointing_mock_args,
-)
 from tests.unit_tests.test_utilities import Utils
 from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness
 from tests.unit_tests.transformer.test_multi_latent_attention import (
@@ -247,6 +233,78 @@ def run(gdn, hidden_states):
                 rec_grads[name], base_grads[name]
             ), f"Grad not identical for {name} ({rank=})"
 
+    def test_deterministic_mode(self):
+        tp_group = parallel_state.get_tensor_model_parallel_group()
+        cp_group = parallel_state.get_context_parallel_group()
+        pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group)
+
+        det_config = copy.deepcopy(self.transformer_config)
+        det_config.deterministic_mode = True
+
+        gdn_submodules = get_experimental_attention_variant_module_spec(
+            config=det_config
+        ).submodules
+
+        model_parallel_cuda_manual_seed(42)
+        torch.manual_seed(42)
+        gdn = (
+            GatedDeltaNet(
+                det_config,
+                submodules=gdn_submodules,
+                layer_number=1,
+                bias=False,
+                conv_bias=False,
+                conv_init=1.0,
+                use_qk_l2norm=True,
+                A_init_range=(1, 16),
+                pg_collection=pg_collection,
+            )
+            .cuda()
+            .bfloat16()
+        )
+
+        # deterministic_mode must select the torch-native kernel, not FLA.
+        assert gdn.gated_delta_rule is torch_chunk_gated_delta_rule
+
+        micro_batch_size = 2
+        seq_length = 64
+        torch.manual_seed(0)
+        base_input = torch.randn(
+            (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size),
+            device=torch.cuda.current_device(),
+            dtype=torch.bfloat16,
+        )
+
+        def run():
+            hidden_states = base_input.clone().requires_grad_(True)
+            output, _ = gdn(hidden_states, None)
+            output.float().sum().backward()
+            grads = {
+                name: param.grad.detach().clone()
+                for name, param in gdn.named_parameters()
+                if param.grad is not None
+            }
+            gdn.zero_grad(set_to_none=True)
+            return output.detach().clone(), grads, hidden_states.grad.detach().clone()
+
+        out1, grads1, input_grad1 = run()
+        out2, grads2, input_grad2 = run()
+
+        rank = torch.distributed.get_rank()
+        assert torch.equal(out1, out2), f"Output not reproducible ({rank=})"
+        assert torch.equal(input_grad1, input_grad2), f"Input grad not reproducible ({rank=})"
+        assert set(grads1.keys()) == set(grads2.keys())
+        for name in grads1:
+            assert torch.equal(
+                grads1[name], grads2[name]
+            ), f"Grad not reproducible for {name} ({rank=})"
+
+    def test_module_construction(self):
+        gdn = self.gdn
+        assert gdn.in_proj_dim == 2 * gdn.qk_dim + 2 * gdn.v_dim + 2 * gdn.num_value_heads
+        assert gdn.A_log.shape == (gdn.num_value_heads // self.tp_size,)
+        assert gdn.dt_bias.shape == (gdn.num_value_heads // self.tp_size,)
+
     def test_jit_compiled_helpers(self):
         import torch._dynamo
 
@@ -254,62 +312,45 @@ def test_jit_compiled_helpers(self):
         batch = 2
         seq_len = 16
 
+        device = torch.cuda.current_device()
         num_v_heads_local = gdn.num_value_heads // gdn.tp_size // gdn.cp_size
+        num_k_heads_local = gdn.num_key_heads // gdn.tp_size // gdn.cp_size
+        qk_dim_local = gdn.qk_dim_local_tp // gdn.cp_size
+        v_dim_local = gdn.v_dim_local_tp // gdn.cp_size
 
-        qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // gdn.cp_size
         qkv = torch.randn(
-            batch, seq_len, qkv_last_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16
+            batch, seq_len, 2 * qk_dim_local + v_dim_local, device=device, dtype=torch.bfloat16
         )
         gate = torch.randn(
             batch,
             seq_len,
             num_v_heads_local,
             gdn.value_head_dim,
-            device=torch.cuda.current_device(),
-            dtype=torch.bfloat16,
-        )
-        beta = torch.randn(
-            batch,
-            seq_len,
-            num_v_heads_local,
-            device=torch.cuda.current_device(),
-            dtype=torch.bfloat16,
-        )
-        alpha = torch.randn(
-            batch,
-            seq_len,
-            num_v_heads_local,
-            device=torch.cuda.current_device(),
+            device=device,
             dtype=torch.bfloat16,
         )
+        gate_feats = (
+            torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16),
+            torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16),
+        )  # beta, alpha
 
         # Disable dynamo so coverage.py can trace through the method bodies,
         # which are normally wrapped by @jit_fuser (torch.compile).
         with torch._dynamo.config.patch(disable=True):
-            query, key, value, gate_out, beta_out, alpha_out = (
-                gdn._prepare_qkv_for_gated_delta_rule(qkv, gate, beta, alpha, batch, seq_len)
+            query, key, value, gate_out, *gate_feats_out = gdn._prepare_input_for_gated_delta_rule(
+                qkv, gate, batch, seq_len, *gate_feats
             )
 
         assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim)
         assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim)
         assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim)
-        assert query.is_contiguous()
-        assert key.is_contiguous()
-        assert value.is_contiguous()
-
-        A_log_mock = torch.randn(
-            num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16
-        )
-        dt_bias_mock = torch.randn(
-            num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16
-        )
-
-        with torch._dynamo.config.patch(disable=True):
-            g, beta_sig = gdn._compute_g_and_beta(A_log_mock, dt_bias_mock, alpha, beta)
+        for t in (query, key, value, gate_out, *gate_feats_out):
+            assert t.is_contiguous()
 
-        assert g.dtype == torch.float32
-        assert g.shape == alpha.shape
-        assert beta_sig.shape == beta.shape
+        # The variant gate features (beta, alpha) pass through with shapes intact
+        beta_out, alpha_out = gate_feats_out
+        assert beta_out.shape == (batch, seq_len, num_v_heads_local)
+        assert alpha_out.shape == (batch, seq_len, num_v_heads_local)
 
     def test_gpu_forward_thd_correctness(self):
         if self.sp_size > 1:
diff --git a/tests/unit_tests/ssm/test_split_tensor_factory.py b/tests/unit_tests/ssm/test_split_tensor_factory.py
index abb668e16a8..ab9fd434e08 100644
--- a/tests/unit_tests/ssm/test_split_tensor_factory.py
+++ b/tests/unit_tests/ssm/test_split_tensor_factory.py
@@ -7,7 +7,7 @@
 import torch
 
 from megatron.core.dist_checkpointing import ShardedTensor
-from megatron.core.ssm.gated_delta_net import (
+from megatron.core.ssm.gated_delta_net.common import (
     _split_tensor_factory as gated_delta_split_tensor_factory,
 )
 from megatron.core.ssm.mamba_mixer import _split_tensor_factory as mamba_split_tensor_factory