diff --git a/megatron/core/dist_checkpointing/utils.py b/megatron/core/dist_checkpointing/utils.py
index 161a3477725..87b8eaed64b 100644
--- a/megatron/core/dist_checkpointing/utils.py
+++ b/megatron/core/dist_checkpointing/utils.py
@@ -2,10 +2,13 @@
 
 """ Helpers for manipulating sharded tensors and sharded state dicts. """
 import logging
+import os
 from contextlib import contextmanager
 from time import time
 from typing import Dict, Optional, Tuple
 
+import torch
+
 from .dict_utils import dict_list_map_inplace, extract_matching_values, nested_values
 from .mapping import (
     LocalNonpersistentObject,
@@ -233,6 +236,27 @@ def _replace_prefixes(x):
     dict_list_map_inplace(_replace_prefixes, sharded_state_dict)
 
 
+def _dist_ckpt_fp8_dequant_dtype() -> Optional[torch.dtype]:
+    value = os.getenv("MEGATRON_DIST_CKPT_FP8_DEQUANT_DTYPE", "").strip().lower()
+    if value in ("", "default", "none"):
+        return None
+    dtype_by_name = {
+        "bf16": torch.bfloat16,
+        "bfloat16": torch.bfloat16,
+        "fp16": torch.float16,
+        "float16": torch.float16,
+        "half": torch.float16,
+        "fp32": torch.float32,
+        "float32": torch.float32,
+    }
+    if value not in dtype_by_name:
+        raise ValueError(
+            "MEGATRON_DIST_CKPT_FP8_DEQUANT_DTYPE must be one of "
+            f"{sorted(dtype_by_name)} plus default/none, got {value!r}"
+        )
+    return dtype_by_name[value]
+
+
 def force_all_tensors_to_non_fp8(sharded_state_dict: ShardedStateDict):
     """Force all tensors in state dict to be non-fp8.
 
@@ -241,9 +265,11 @@ def force_all_tensors_to_non_fp8(sharded_state_dict: ShardedStateDict):
     """
     from ..fp8_utils import dequantize_fp8_tensor, is_float8tensor  # Avoid circular import
 
+    dequant_dtype = _dist_ckpt_fp8_dequant_dtype()
+
     for v in nested_values(sharded_state_dict):
         if hasattr(v, "data") and is_float8tensor(v.data):
-            v.data = dequantize_fp8_tensor(v.data)
+            v.data = dequantize_fp8_tensor(v.data, dtype=dequant_dtype)
 
 
 fallback_logger = logging.getLogger(__name__)
diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py
index ddaeb7e8d84..5a7ed553616 100644
--- a/megatron/core/distributed/finalize_model_grads.py
+++ b/megatron/core/distributed/finalize_model_grads.py
@@ -281,7 +281,7 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n
     """
     for model_chunk in model:
         for module in get_attr_wrapped_model(model_chunk, 'modules')():
-            if config.moe_router_enable_expert_bias and hasattr(module, 'expert_bias'):
+            if config.moe_router_enable_expert_bias and getattr(module, 'expert_bias', None) is not None:
                 module.local_tokens_per_expert.zero_()
             if (
                 config.moe_router_load_balancing_type == "global_aux_loss"
@@ -299,7 +299,7 @@ def _update_router_expert_bias(model: List[torch.nn.Module], config: Transformer
     expert_bias_list = []
     for model_chunk in model:
         for module in get_attr_wrapped_model(model_chunk, 'modules')():
-            if hasattr(module, 'expert_bias'):
+            if getattr(module, 'expert_bias', None) is not None:
                 tokens_per_expert_list.append(module.local_tokens_per_expert)
                 expert_bias_list.append(module.expert_bias)
     # For hybrid models with both MoE and Dense layers, this list can be empty.
@@ -472,7 +472,7 @@ def finalize_model_grads(
     if config.timers is not None:
         config.timers('embedding-grads-all-reduce').stop()
 
-    if config.moe_router_enable_expert_bias:
+    if config.moe_router_enable_expert_bias and not config.freeze_e_score_correction_bias:
         _update_router_expert_bias(model, config)
 
     reset_model_temporary_tensors(config, model)
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index d239db4ab0c..0165c7f9b11 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -62,6 +62,25 @@
     te = MagicMock()
     HAVE_TE = False
 
+_TE_QUANTIZED_TENSOR_TYPES = ()
+if HAVE_TE:
+    try:
+        from transformer_engine.pytorch.tensor import QuantizedTensor as _TEQuantizedTensor
+
+        _TE_QUANTIZED_TENSOR_TYPES = _TE_QUANTIZED_TENSOR_TYPES + (_TEQuantizedTensor,)
+    except (ImportError, ModuleNotFoundError, AttributeError):
+        pass
+    try:
+        from transformer_engine.pytorch.quantized_tensor import (
+            QuantizedTensorStorage as _TEQuantizedTensorStorage,
+        )
+
+        _TE_QUANTIZED_TENSOR_TYPES = _TE_QUANTIZED_TENSOR_TYPES + (
+            _TEQuantizedTensorStorage,
+        )
+    except (ImportError, ModuleNotFoundError, AttributeError):
+        pass
+
 
 def _get_extra_te_kwargs(config: TransformerConfig):
     extra_transformer_engine_kwargs = {"params_dtype": config.params_dtype}
@@ -480,6 +499,7 @@ def __init__(
         skip_weight_param_allocation: bool = False,
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        stride: int = 1,
     ):
         if not HAVE_TE:
             raise ImportError(
@@ -515,6 +535,7 @@ def __init__(
         extra_kwargs = _get_extra_te_kwargs(config)
         self.tp_size = get_pg_size(tp_group)
         self.tp_rank = get_pg_rank(tp_group)
+        self.stride = stride
 
         if self.config.delay_wgrad_compute:
             if is_te_min_version("2.3.0"):
@@ -595,6 +616,10 @@ def __init__(
             **extra_kwargs,
         )
 
+        setattr(self.weight, "partition_stride", stride)
+        if bias and hasattr(self, "bias") and self.bias is not None:
+            setattr(self.bias, "partition_stride", stride)
+
         if config.use_cpu_initialization:
             output_size_per_partition = divide(output_size, self.tp_size)
             _ = _initialize_affine_weight_cpu(
@@ -604,7 +629,7 @@ def __init__(
                 output_size_per_partition,
                 0,
                 init_method=condition_init_method(config, init_method),
-                stride=1,
+                stride=stride,
                 return_master_weight=False,
                 rank=self.tp_rank,
                 world_size=self.tp_size,
@@ -614,7 +639,7 @@ def __init__(
                 self.bias = Parameter(
                     torch.empty(output_size_per_partition, dtype=config.params_dtype)
                 )
-                set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+                set_tensor_model_parallel_attributes(self.bias, True, 0, stride)
                 with torch.no_grad():
                     self.bias.zero_()
                 setattr(self.bias, "allreduce", True)
@@ -677,6 +702,7 @@ def __init__(
         skip_weight_param_allocation: bool = False,
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        stride: int = 1,
     ):
         if not HAVE_TE:
             raise ImportError(
@@ -690,6 +716,7 @@ def __init__(
         self._tp_group = tp_group
         world_size = get_pg_size(tp_group)
         rank = get_pg_rank(tp_group)
+        self.stride = stride
 
         super().__init__(
             input_size=input_size,
@@ -710,6 +737,10 @@ def __init__(
             tp_group=tp_group,
         )
 
+        setattr(self.weight, "partition_stride", stride)
+        if bias and hasattr(self, "bias") and self.bias is not None:
+            setattr(self.bias, "partition_stride", stride)
+
         if config.use_cpu_initialization:
             output_size_per_partition = divide(output_size, world_size)
             _ = _initialize_affine_weight_cpu(
@@ -719,7 +750,7 @@ def __init__(
                 output_size_per_partition,
                 0,
                 init_method=condition_init_method(config, init_method),
-                stride=1,
+                stride=stride,
                 return_master_weight=False,
                 rank=rank,
                 world_size=world_size,
@@ -729,7 +760,7 @@ def __init__(
                 self.bias = Parameter(
                     torch.empty(output_size_per_partition, dtype=config.params_dtype)
                 )
-                set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+                set_tensor_model_parallel_attributes(self.bias, True, 0, stride)
                 with torch.no_grad():
                     self.bias.zero_()
                 setattr(self.bias, "allreduce", True)
@@ -1218,6 +1249,188 @@ def fake_int4_quantization_ste(x, group_size):
             
         return x_out
 
+    # ------------------------------------------------------------------
+    # MXFP4 fake-QAT. Matches the FlashInfer MXFP4 rollout kernel:
+    # per-block (1 x 32) absmax -> E8M0 power-of-2 scale -> round to E2M1 grid
+    # -> dequant multiply. Straight-through backward. Gated by
+    # OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG; block size defaults to 32 and is
+    # overridable via OPEN_TRAINING_MXFP4_BLOCK_SIZE for ablations. Mutually
+    # exclusive with OPEN_TRAINING_INT4_FAKE_QAT_FLAG.
+    #
+    # Bit-exact equivalence to the rollout-path quantization is guarded by a
+    # small-card FlashInfer probe against flashinfer.fp4_quantize +
+    # flashinfer.mxfp4_dequantize_host (see
+    # tools/model_convert/dsv4_mxfp4_fake_qat_vs_flashinfer_probe.py).
+    # ------------------------------------------------------------------
+    _MXFP4_E2M1_POS_GRID = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+    _MXFP4_E2M1_MAX = 6.0
+    _MXFP4_BLOCK_SIZE = 32
+
+    # Module-level cached constants for E2M1 rounding. These are tiny
+    # (7-elt boundaries + 8-elt grid) but creating them per-call allocates
+    # temporaries 128+ times per forward when we quantize every expert weight.
+    _MXFP4_E2M1_BOUNDARIES_BF16 = None
+    _MXFP4_E2M1_GRID_BF16 = None
+
+    def _get_e2m1_tables(device, dtype):
+        global _MXFP4_E2M1_BOUNDARIES_BF16, _MXFP4_E2M1_GRID_BF16
+        if (
+            _MXFP4_E2M1_BOUNDARIES_BF16 is None
+            or _MXFP4_E2M1_BOUNDARIES_BF16.device != device
+            or _MXFP4_E2M1_BOUNDARIES_BF16.dtype != dtype
+        ):
+            _MXFP4_E2M1_BOUNDARIES_BF16 = torch.tensor(
+                [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0],
+                dtype=dtype, device=device,
+            )
+            _MXFP4_E2M1_GRID_BF16 = torch.tensor(
+                _MXFP4_E2M1_POS_GRID, dtype=dtype, device=device,
+            )
+        return _MXFP4_E2M1_BOUNDARIES_BF16, _MXFP4_E2M1_GRID_BF16
+
+    def _round_to_e2m1(x: torch.Tensor) -> torch.Tensor:
+        """Round to E2M1 grid.
+
+        Keeps the bit-exact FlashInfer semantics (round-to-nearest,
+        ties-to-even E2M1 code) while minimizing peak memory:
+          - boundaries/grid tables cached at module level (no per-call alloc)
+          - bucket index is int32 (torch.bucketize output is int64 by default
+            but we pass out_int32=True)
+          - sign / grid-lookup / multiply tensors are reused where possible
+
+        Called on x that is already sign-preserved fp32 (or bf16) and clamped
+        to [-E2M1_MAX, E2M1_MAX].
+        """
+        boundaries, grid = _get_e2m1_tables(x.device, x.dtype)
+        mag = x.abs()
+        # int32 bucket index, half the memory of the default int64.
+        idx = torch.bucketize(mag, boundaries, right=False, out_int32=True)
+        # FlashInfer's E2M1 quantizer resolves exact midpoint ties to the
+        # even code index. bucketize(right=False) already selects the lower
+        # even code for boundaries 0, 2, 4, and 6; boundaries 1, 3, and 5
+        # need to advance to the upper even code.
+        tie_up = (mag == boundaries[1]) | (mag == boundaries[3]) | (mag == boundaries[5])
+        idx.add_(tie_up.to(idx.dtype))
+        idx.clamp_(max=len(_MXFP4_E2M1_POS_GRID) - 1)
+        grid_vals = grid[idx]  # same dtype as x
+        del idx
+        # Restore sign via torch.copysign (single tensor output, same dtype).
+        return torch.copysign(grid_vals, x)
+
+    class _FakeMXFP4QuantizationSTE(torch.autograd.Function):
+        """Fake MXFP4 quantization for QAT, shape [M, N] with N as quant axis.
+
+        Backward is straight-through, so nothing from the forward math needs to
+        survive to backward. Run the entire forward under ``torch.no_grad()`` so
+        the autograd engine does not hold references to intermediates.
+
+        Memory strategy (the forward is called for every expert weight every
+        forward pass; DSV4 has up to 128 experts x 2 grouped linears x 4 layers
+        simultaneously live inside one list comprehension, so per-call peak
+        matters a lot):
+          - skip padding when shape is already aligned (DSV4 MoE is aligned)
+          - do the absmax in fp32 (small: numel/32) but keep the bulk division
+            in the original dtype (bf16), not fp32 -- cuts per-call peak 2x
+          - round to E2M1 via cached tables + int32 bucket index (see
+            _round_to_e2m1)
+          - release scale / index / sign tensors as soon as possible
+        """
+
+        @staticmethod
+        def forward(ctx, x, block_size):
+            with torch.no_grad():
+                m, n = x.shape
+                block_size_m, block_size_n = 1, block_size
+
+                # Fast path: no padding required -> view x directly. DSV4 MoE
+                # weights (4096, 4096) and (4096, 2048) always hit this.
+                if m % block_size_m == 0 and n % block_size_n == 0:
+                    x_view = x.view(
+                        m // block_size_m,
+                        block_size_m,
+                        n // block_size_n,
+                        block_size_n,
+                    )
+                    needs_unpad = False
+                else:
+                    m_padded = ceil_div(m, block_size_m) * block_size_m
+                    n_padded = ceil_div(n, block_size_n) * block_size_n
+                    x_padded = torch.zeros(
+                        (m_padded, n_padded),
+                        dtype=x.dtype, device=x.device,
+                    )
+                    x_padded[:m, :n] = x
+                    x_view = x_padded.view(
+                        m_padded // block_size_m,
+                        block_size_m,
+                        n_padded // block_size_n,
+                        block_size_n,
+                    )
+                    needs_unpad = True
+
+                # Per-block absmax in fp32 for numerical stability. Tensor is
+                # 32x smaller than the weight itself, so cost is negligible.
+                x_max = x_view.abs().amax(dim=(1, 3), keepdim=True).float()
+                # scale = ceil(log2(max / E2M1_MAX)), encoded as E8M0 power-of-2.
+                x_max.div_(_MXFP4_E2M1_MAX).clamp_(min=1e-8).log2_().ceil_().clamp_(
+                    min=-127.0, max=127.0
+                )
+                # x_max now holds the log2 exponent; exponentiate into x_scale.
+                x_scale = torch.pow(2.0, x_max).to(x.dtype)
+                del x_max
+
+                # Divide / clamp / round in the input dtype (bf16), not fp32.
+                # Peak memory here is 1x the weight (vs 2x if we went to fp32).
+                x_div = (x_view / x_scale).clamp_(-_MXFP4_E2M1_MAX, _MXFP4_E2M1_MAX)
+                x_q = _round_to_e2m1(x_div)
+                del x_div
+
+                x_dequant = x_q.mul_(x_scale)
+                del x_q, x_scale
+
+                if needs_unpad:
+                    x_dequant_full = x_dequant.reshape(
+                        x_view.size(0) * block_size_m,
+                        x_view.size(2) * block_size_n,
+                    )
+                    x_out = x_dequant_full[:m, :n].contiguous()
+                else:
+                    x_out = x_dequant.view(m, n)
+
+                return x_out
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            return grad_output, None
+
+    def fake_mxfp4_quantization_ste(x, block_size=_MXFP4_BLOCK_SIZE):
+        if _TE_QUANTIZED_TENSOR_TYPES and isinstance(x, _TE_QUANTIZED_TENSOR_TYPES):
+            original_weight = x
+            x = x.dequantize()
+            if original_weight.requires_grad and not x.requires_grad:
+                x = x.detach().requires_grad_(True)
+            # TE's grouped-linear backward writes fused wgrad through these
+            # attributes. Dequantizing primary-FP8 params for forward QAT must
+            # not detach the fake-QAT tensor from the existing main-grad path.
+            for attr in (
+                "main_grad",
+                "grad_added_to_main_grad",
+                "zero_out_wgrad",
+                "overwrite_main_grad",
+            ):
+                if hasattr(original_weight, attr):
+                    setattr(x, attr, getattr(original_weight, attr))
+
+        x_out = _FakeMXFP4QuantizationSTE.apply(x, block_size)
+
+        # Preserve Megatron DDP's ``main_grad`` accumulator: the outer
+        # optimizer reduces into ``param.main_grad``, so callers that look up
+        # ``main_grad`` on the returned tensor must still find it.
+        if hasattr(x, 'main_grad'):
+            x_out.main_grad = x.main_grad
+
+        return x_out
+
     class TEGroupedLinear(te.pytorch.GroupedLinear):
         """
         Wrapper for the Transformer-Engine's `GroupedLinear` layer.
@@ -1311,6 +1524,9 @@ def __init__(
 
             for param in self.parameters():
                 setattr(param, "allreduce", not (is_expert and self.expert_parallel))
+                if is_expert and param.dim() == 2:
+                    setattr(param, "mcore_mxfp4_expert_qat_weight", True)
+                    setattr(param, "mcore_mxfp4_expert_qat_shape", tuple(param.shape))
 
             def merge_extra_states(
                 self,
@@ -1422,14 +1638,35 @@ def _get_weight_tensors(self):
             """Get the weight tensors of the module."""
             weight_tensors = super()._get_weight_tensors()
 
-            if os.getenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") == "1":
+            int4_enabled = os.getenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") == "1"
+            mxfp4_enabled = os.getenv("OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG", "0") == "1"
+
+            if int4_enabled and mxfp4_enabled:
+                raise RuntimeError(
+                    "INT4 and MXFP4 fake QAT are mutually exclusive; set exactly "
+                    "one of OPEN_TRAINING_INT4_FAKE_QAT_FLAG / "
+                    "OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG to 1."
+                )
+
+            if int4_enabled:
                 group_size = int(os.getenv("OPEN_TRAINING_INT4_GROUP_SIZE", "128"))
 
                 weight_tensors = [
-                    fake_int4_quantization_ste(w, group_size) 
+                    fake_int4_quantization_ste(w, group_size)
+                    for w in weight_tensors
+                ]
+            elif mxfp4_enabled:
+                # Standard QAT semantics: fake-MXFP4 is part of the routed
+                # expert forward graph. With fp8_param_gather this acts on the
+                # TE FP8 forward weight tensor after the normal main-param ->
+                # model-param copy, rather than changing the copy boundary.
+                # MXFP4 spec fixes block size at 32; env is overridable for ablations.
+                block_size = int(os.getenv("OPEN_TRAINING_MXFP4_BLOCK_SIZE", "32"))
+                weight_tensors = [
+                    fake_mxfp4_quantization_ste(w, block_size)
                     for w in weight_tensors
                 ]
-                
+
             return weight_tensors
 
         def _encode_extra_state(self, state):
diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py
index 1c52e965cd7..772a3e154ab 100644
--- a/megatron/core/fp8_utils.py
+++ b/megatron/core/fp8_utils.py
@@ -3,6 +3,9 @@
 """Utility functions related to FP8 that are used throughout Megatron core"""
 
 import importlib
+import json
+import math
+import os
 import weakref
 from contextlib import nullcontext
 from functools import wraps
@@ -93,6 +96,97 @@
     te_post_all_gather_processing = None
 
 
+def _mxfp4_qat_min_copy_call() -> Optional[int]:
+    raw_value = os.getenv("OPEN_TRAINING_MXFP4_FAKE_QAT_MIN_COPY_CALL", "")
+    if not raw_value:
+        return None
+    try:
+        value = int(raw_value)
+    except ValueError as exc:
+        raise RuntimeError(
+            "OPEN_TRAINING_MXFP4_FAKE_QAT_MIN_COPY_CALL must be a positive integer."
+        ) from exc
+    if value <= 0:
+        raise RuntimeError(
+            "OPEN_TRAINING_MXFP4_FAKE_QAT_MIN_COPY_CALL must be a positive integer."
+        )
+    return value
+
+
+def _mxfp4_qat_copy_boundary_enabled() -> bool:
+    """Return whether the old fake-MXFP4 -> TE-FP8 copy-boundary probe is enabled."""
+    return os.getenv("OPEN_TRAINING_MXFP4_FAKE_QAT_COPY_BOUNDARY", "0") == "1"
+
+
+def _mxfp4_qat_applies_to_copy_call(copy_call_index: Optional[int]) -> tuple[bool, Optional[int]]:
+    if not _mxfp4_qat_copy_boundary_enabled():
+        return False, _mxfp4_qat_min_copy_call()
+
+    min_copy_call = _mxfp4_qat_min_copy_call()
+    if min_copy_call is None:
+        return True, None
+    if copy_call_index is None:
+        raise RuntimeError(
+            "OPEN_TRAINING_MXFP4_FAKE_QAT_MIN_COPY_CALL requires the "
+            "fp8_param_gather copy-call index."
+        )
+    return copy_call_index >= min_copy_call, min_copy_call
+
+
+def _maybe_fake_mxfp4_expert_qat_main_param_shard(
+    model_param: torch.Tensor,
+    main_param: Optional[torch.Tensor],
+    start_offset: Optional[int],
+    *,
+    copy_call_index: Optional[int] = None,
+) -> Optional[torch.Tensor]:
+    if main_param is None:
+        return None
+    if os.getenv("OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG", "0") != "1":
+        return main_param
+    if not getattr(model_param, "mcore_mxfp4_expert_qat_weight", False):
+        return main_param
+
+    fake_qat_applied, _ = _mxfp4_qat_applies_to_copy_call(copy_call_index)
+    if not fake_qat_applied:
+        return main_param
+
+    block_size = int(os.getenv("OPEN_TRAINING_MXFP4_BLOCK_SIZE", "32"))
+    if block_size <= 0:
+        raise RuntimeError(
+            "OPEN_TRAINING_MXFP4_BLOCK_SIZE must be a positive integer when "
+            "OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG=1."
+        )
+
+    weight_shape = getattr(
+        model_param, "mcore_mxfp4_expert_qat_shape", tuple(model_param.shape)
+    )
+    if len(weight_shape) != 2 or weight_shape[-1] % block_size != 0:
+        raise RuntimeError(
+            "MXFP4 expert QAT with fp8_param_gather requires a 2D expert weight "
+            f"whose K dimension is divisible by block size {block_size}; got {weight_shape}."
+        )
+    if start_offset is None:
+        raise RuntimeError(
+            "MXFP4 expert QAT with fp8_param_gather requires shard start offsets."
+        )
+    if start_offset % block_size != 0 or main_param.numel() % block_size != 0:
+        raise RuntimeError(
+            "MXFP4 expert QAT with fp8_param_gather requires optimizer shards "
+            f"aligned to group-{block_size} boundaries; got start_offset={start_offset}, "
+            f"numel={main_param.numel()}."
+        )
+
+    from megatron.core.extensions.transformer_engine import fake_mxfp4_quantization_ste
+
+    fake_main_param = fake_mxfp4_quantization_ste(
+        main_param.view(-1, block_size), block_size
+    ).view_as(main_param)
+    if hasattr(main_param, "main_grad"):
+        fake_main_param.main_grad = main_param.main_grad
+    return fake_main_param
+
+
 def is_float8tensor(tensor: torch.Tensor) -> bool:
     """Check if a tensor is a Transformer Engine Float8Tensor.
 
@@ -110,12 +204,293 @@ def is_mxfp8tensor(tensor: torch.Tensor) -> bool:
     return HAVE_TE_MXFP8TENSOR and isinstance(tensor, MXFP8Tensor)
 
 
-def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor:
+def dequantize_fp8_tensor(fp8_tensor: torch.Tensor, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
     """Dequantize a fp8 tensor to a higher precision tensor."""
     if is_te_min_version("2.0"):
-        return fp8_tensor.dequantize()
+        return fp8_tensor.dequantize(dtype=dtype) if dtype is not None else fp8_tensor.dequantize()
+
+    tensor = fp8_tensor.from_float8()
+    return tensor.to(dtype=dtype) if dtype is not None else tensor
+
+
+_MXFP4_QAT_FP8_COPY_TRACE_REMAINING = None
+_MXFP4_QAT_FP8_COPY_TRACE_CALL_INDEX = 0
+_MXFP4_QAT_FP8_COPY_TRACE_RECORD_INDEX = 0
+_MXFP4_QAT_FP8_COPY_TRACE_COUNTS_BY_CALL: dict[int, int] = {}
+_MXFP4_E2M1_POS_GRID = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+_MXFP4_E2M1_BOUNDARIES = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
+
+
+def _mxfp4_qat_fp8_copy_trace_remaining() -> int:
+    global _MXFP4_QAT_FP8_COPY_TRACE_REMAINING
+    if _MXFP4_QAT_FP8_COPY_TRACE_REMAINING is None:
+        raw_limit = os.getenv("SIRL_DSV4_QAT_FP8_COPY_TRACE_LIMIT", "")
+        _MXFP4_QAT_FP8_COPY_TRACE_REMAINING = int(raw_limit) if raw_limit else 0
+    return _MXFP4_QAT_FP8_COPY_TRACE_REMAINING
+
+
+def _decrement_mxfp4_qat_fp8_copy_trace_remaining() -> None:
+    global _MXFP4_QAT_FP8_COPY_TRACE_REMAINING
+    _MXFP4_QAT_FP8_COPY_TRACE_REMAINING = max(
+        0, _mxfp4_qat_fp8_copy_trace_remaining() - 1
+    )
+
+
+def _next_mxfp4_qat_fp8_copy_trace_call_index() -> int:
+    global _MXFP4_QAT_FP8_COPY_TRACE_CALL_INDEX
+    _MXFP4_QAT_FP8_COPY_TRACE_CALL_INDEX += 1
+    return _MXFP4_QAT_FP8_COPY_TRACE_CALL_INDEX
+
+
+def _mxfp4_qat_fp8_copy_trace_limit_per_call() -> Optional[int]:
+    raw_limit = os.getenv("SIRL_DSV4_QAT_FP8_COPY_TRACE_LIMIT_PER_CALL", "")
+    return int(raw_limit) if raw_limit else None
+
+
+def _mxfp4_qat_fp8_copy_trace_sample_elems() -> int:
+    raw_limit = os.getenv("SIRL_DSV4_QAT_FP8_COPY_TRACE_SAMPLE_ELEMS", "")
+    return int(raw_limit) if raw_limit else 0
+
+
+def _evenly_spaced_indices(numel: int, limit: int) -> List[int]:
+    if numel <= 0 or limit <= 0:
+        return []
+    if limit >= numel:
+        return list(range(numel))
+    if limit == 1:
+        return [0]
+    return sorted({round(index * (numel - 1) / (limit - 1)) for index in range(limit)})
+
+
+def _mxfp4_e2m1_code(value: float) -> int:
+    mag = abs(value)
+    code = 0
+    for boundary in _MXFP4_E2M1_BOUNDARIES:
+        if mag > boundary:
+            code += 1
+        elif mag == boundary and code % 2 == 1:
+            code += 1
+        else:
+            break
+    return min(code, len(_MXFP4_E2M1_POS_GRID) - 1)
+
+
+def _nearest_mxfp4_boundary_distance(value: float) -> dict[str, float]:
+    mag = abs(value)
+    boundary = min(_MXFP4_E2M1_BOUNDARIES, key=lambda candidate: abs(mag - candidate))
+    return {
+        "nearest_e2m1_boundary": boundary,
+        "nearest_e2m1_boundary_distance": abs(mag - boundary),
+    }
+
+
+def _sample_mxfp4_qat_fp8_copy_boundary(
+    *,
+    raw_main_param: torch.Tensor,
+    fake_main_param: torch.Tensor,
+    actual: torch.Tensor,
+    block_size: int,
+    sample_elems: int,
+) -> List[dict[str, object]]:
+    raw_flat = raw_main_param.detach().float().reshape(-1)
+    fake_flat = fake_main_param.detach().float().reshape(-1)
+    actual_flat = actual.detach().float().reshape(-1)
+    if raw_flat.numel() != fake_flat.numel() or raw_flat.numel() != actual_flat.numel():
+        raise RuntimeError(
+            "MXFP4 QAT FP8-copy sample trace requires raw/fake/actual numel equality: "
+            f"raw={raw_flat.numel()} fake={fake_flat.numel()} actual={actual_flat.numel()}"
+        )
+    if raw_flat.numel() % block_size != 0:
+        raise RuntimeError(
+            "MXFP4 QAT FP8-copy sample trace requires shard numel aligned to block size: "
+            f"numel={raw_flat.numel()} block_size={block_size}"
+        )
+
+    samples: List[dict[str, object]] = []
+    for flat_index in _evenly_spaced_indices(raw_flat.numel(), sample_elems):
+        group_start = (flat_index // block_size) * block_size
+        group_end = group_start + block_size
+        raw_value = float(raw_flat[flat_index].cpu())
+        fake_value = float(fake_flat[flat_index].cpu())
+        actual_value = float(actual_flat[flat_index].cpu())
+        group_absmax = float(raw_flat[group_start:group_end].abs().max().cpu())
+        scale_exponent = math.ceil(math.log2(max(group_absmax / 6.0, 1e-8)))
+        scale_exponent = min(127, max(-127, scale_exponent))
+        scale_value = float(2.0 ** scale_exponent)
+        raw_over_scale = raw_value / scale_value
+        fake_over_scale = fake_value / scale_value
+        actual_over_scale = actual_value / scale_value
+        samples.append(
+            {
+                "flat_index": int(flat_index),
+                "group_start": int(group_start),
+                "group_end": int(group_end),
+                "block_size": int(block_size),
+                "scale_exponent": int(scale_exponent),
+                "scale_value": scale_value,
+                "group_absmax": group_absmax,
+                "raw_value": raw_value,
+                "fake_value": fake_value,
+                "actual_value": actual_value,
+                "raw_over_scale": raw_over_scale,
+                "fake_over_scale": fake_over_scale,
+                "actual_over_scale": actual_over_scale,
+                "raw_e2m1_code": _mxfp4_e2m1_code(raw_over_scale),
+                "fake_e2m1_code": _mxfp4_e2m1_code(fake_over_scale),
+                **_nearest_mxfp4_boundary_distance(raw_over_scale),
+            }
+        )
+    return samples
+
+
+def _should_record_mxfp4_qat_fp8_copy_trace(copy_call_index: int) -> bool:
+    per_call_limit = _mxfp4_qat_fp8_copy_trace_limit_per_call()
+    if per_call_limit is None:
+        return _mxfp4_qat_fp8_copy_trace_remaining() > 0
+
+    count = _MXFP4_QAT_FP8_COPY_TRACE_COUNTS_BY_CALL.get(copy_call_index, 0)
+    return count < per_call_limit
+
+
+def _mark_recorded_mxfp4_qat_fp8_copy_trace(copy_call_index: int) -> int:
+    global _MXFP4_QAT_FP8_COPY_TRACE_RECORD_INDEX
+    record_index = _MXFP4_QAT_FP8_COPY_TRACE_RECORD_INDEX
+    _MXFP4_QAT_FP8_COPY_TRACE_RECORD_INDEX += 1
+
+    per_call_limit = _mxfp4_qat_fp8_copy_trace_limit_per_call()
+    if per_call_limit is None:
+        _decrement_mxfp4_qat_fp8_copy_trace_remaining()
     else:
-        return fp8_tensor.from_float8()
+        _MXFP4_QAT_FP8_COPY_TRACE_COUNTS_BY_CALL[copy_call_index] = (
+            _MXFP4_QAT_FP8_COPY_TRACE_COUNTS_BY_CALL.get(copy_call_index, 0) + 1
+        )
+    return record_index
+
+
+class _TraceFormatDict(dict):
+    def __missing__(self, key):
+        return "unknown"
+
+
+def _trace_rank() -> str:
+    try:
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            return str(torch.distributed.get_rank())
+    except RuntimeError:
+        pass
+    return os.getenv("RANK", "unknown")
+
+
+def _format_mxfp4_qat_fp8_copy_trace_path(
+    path_template: str,
+    *,
+    copy_call_index: int,
+    record_index: int,
+) -> str:
+    return path_template.format_map(
+        _TraceFormatDict(
+            rank=_trace_rank(),
+            pid=str(os.getpid()),
+            copy_call=str(copy_call_index),
+            trace_index=str(record_index),
+        )
+    )
+
+
+def _maybe_record_mxfp4_qat_fp8_copy_trace(
+    model_param: torch.Tensor,
+    fake_main_param: Optional[torch.Tensor],
+    raw_main_param: Optional[torch.Tensor],
+    start_offset: Optional[int],
+    copy_call_index: int,
+) -> None:
+    """Trace main-param -> TE FP8 compute-param copy loss.
+
+    This is disabled by default. It is intentionally train-side: rollout export
+    tracing proves TE-FP8-export -> native-FP4 pack/dequant. Standard QAT keeps
+    fake-MXFP4 in the TEGroupedLinear forward graph; this hook only applies
+    fake-MXFP4 at the copy boundary when
+    OPEN_TRAINING_MXFP4_FAKE_QAT_COPY_BOUNDARY=1 for old diagnostics.
+    """
+    path_template = os.getenv("SIRL_DSV4_QAT_FP8_COPY_TRACE_PATH", "")
+    if not path_template or not _should_record_mxfp4_qat_fp8_copy_trace(copy_call_index):
+        return
+    if fake_main_param is None:
+        return
+    if os.getenv("OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG", "0") != "1":
+        return
+    if not getattr(model_param, "mcore_mxfp4_expert_qat_weight", False):
+        return
+    if start_offset is None:
+        return
+    if raw_main_param is None:
+        return
+
+    with torch.no_grad():
+        fake_qat_applied, min_copy_call = _mxfp4_qat_applies_to_copy_call(copy_call_index)
+        expected = fake_main_param.detach().float().reshape(-1)
+        raw = raw_main_param.detach().float().reshape(-1)
+        if raw.numel() != expected.numel():
+            raise RuntimeError(
+                "MXFP4 QAT FP8-copy trace requires raw and fake main-param shards "
+                f"with matching numel, got raw={raw.numel()} fake={expected.numel()}."
+            )
+        dequantized_model = dequantize_fp8_tensor(model_param).detach().float().reshape(-1)
+        actual = dequantized_model.narrow(0, int(start_offset), expected.numel())
+        diff = actual - expected
+        max_abs_error = diff.abs().max()
+        rmse = diff.pow(2).mean().sqrt()
+        expected_rmse = expected.pow(2).mean().sqrt().clamp_min(1e-12)
+        actual_rmse = actual.pow(2).mean().sqrt()
+        record_index = _mark_recorded_mxfp4_qat_fp8_copy_trace(copy_call_index)
+        record = {
+            "schema": "dsv4_mxfp4_qat_fp8_copy_trace.v1",
+            "trace_index": record_index,
+            "copy_call_index": copy_call_index,
+            "fake_qat_copy_boundary_enabled": _mxfp4_qat_copy_boundary_enabled(),
+            "fake_qat_applied_to_fp8_copy": fake_qat_applied,
+            "fake_qat_min_copy_call": min_copy_call,
+            "rank": _trace_rank(),
+            "pid": os.getpid(),
+            "param_name": getattr(model_param, "mcore_mxfp4_expert_qat_name", None),
+            "start_offset": int(start_offset),
+            "numel": int(expected.numel()),
+            "model_shape": list(getattr(model_param, "shape", ())),
+            "main_param_shape": list(fake_main_param.shape),
+            "model_param_type": f"{type(model_param).__module__}.{type(model_param).__qualname__}",
+            "model_param_dtype": str(getattr(model_param, "dtype", None)),
+            "fake_main_param_dtype": str(fake_main_param.dtype),
+            "max_abs_error": float(max_abs_error.cpu()),
+            "rmse": float(rmse.cpu()),
+            "relative_rmse": float((rmse / expected_rmse).cpu()),
+            "expected_absmax": float(expected.abs().max().cpu()),
+            "actual_absmax": float(actual.abs().max().cpu()),
+            "expected_mean": float(expected.mean().cpu()),
+            "actual_mean": float(actual.mean().cpu()),
+            "expected_rms": float(expected_rmse.cpu()),
+            "actual_rms": float(actual_rmse.cpu()),
+        }
+        sample_elems = _mxfp4_qat_fp8_copy_trace_sample_elems()
+        if sample_elems:
+            block_size = int(os.getenv("OPEN_TRAINING_MXFP4_BLOCK_SIZE", "32"))
+            record["e2m1_positive_grid"] = list(_MXFP4_E2M1_POS_GRID)
+            record["e2m1_boundaries"] = list(_MXFP4_E2M1_BOUNDARIES)
+            record["sampled_main_param_values"] = _sample_mxfp4_qat_fp8_copy_boundary(
+                raw_main_param=raw,
+                fake_main_param=expected,
+                actual=actual,
+                block_size=block_size,
+                sample_elems=sample_elems,
+            )
+
+    path = _format_mxfp4_qat_fp8_copy_trace_path(
+        path_template,
+        copy_call_index=copy_call_index,
+        record_index=record_index,
+    )
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    with open(path, "a", encoding="utf-8") as handle:
+        handle.write(json.dumps(record, sort_keys=True) + "\n")
 
 
 def _resolve_callable_from_python_import_path(dotted_path: str):
@@ -241,16 +616,17 @@ def _quantize_param_shard_impl(
             return
 
         from transformer_engine.pytorch.tensor.utils import cast_master_weights_to_fp8
+        copy_call_index = _next_mxfp4_qat_fp8_copy_trace_call_index()
 
-        args = [model_params, main_params, start_offsets, data_parallel_group]
         if fsdp_shard_model_params is not None:
             if not HAVE_PACKAGING:
                 raise ImportError(
                     "packaging not found, please install it with `pip install packaging`"
                 )
-            if get_te_version() == PkgVersion("2.3.0.dev0+5fdd7bb") or is_te_min_version("2.3.0"):
-                args.append(fsdp_shard_model_params)
-            else:
+            if not (
+                get_te_version() == PkgVersion("2.3.0.dev0+5fdd7bb")
+                or is_te_min_version("2.3.0")
+            ):
                 raise NotImplementedError(
                     f"FSDP with --fp8-param-gather is not supported in TE v{get_te_version()}"
                 )
@@ -262,7 +638,28 @@ def _quantize_param_shard_impl(
         if te_post_all_gather_processing is not None:
             kwargs["manual_post_all_gather_processing"] = True
 
+        raw_main_params = main_params
+        main_params = [
+            _maybe_fake_mxfp4_expert_qat_main_param_shard(
+                model_param,
+                main_param,
+                start_offset,
+                copy_call_index=copy_call_index,
+            )
+            for model_param, main_param, start_offset in zip(
+                model_params, main_params, start_offsets
+            )
+        ]
+        args = [model_params, main_params, start_offsets, data_parallel_group]
+        if fsdp_shard_model_params is not None:
+            args.append(fsdp_shard_model_params)
         cast_master_weights_to_fp8(*args, **kwargs)
+        for model_param, main_param, raw_main_param, start_offset in zip(
+            model_params, main_params, raw_main_params, start_offsets
+        ):
+            _maybe_record_mxfp4_qat_fp8_copy_trace(
+                model_param, main_param, raw_main_param, start_offset, copy_call_index
+            )
 
     def _correct_amax_history_if_needed_impl(model: List[torch.nn.Module]) -> None:
         pass
@@ -297,11 +694,13 @@ def _quantize_param_shard_impl(
         if fsdp_shard_model_params is None:
             fsdp_shard_model_params = [None] * len(model_params)
 
+        copy_call_index = _next_mxfp4_qat_fp8_copy_trace_call_index()
         for model_param, main_param, start_offset, fsdp_shard_model_param in zip(
             model_params, main_params, start_offsets, fsdp_shard_model_params
         ):
             if main_param is None:
                 continue
+            raw_main_param = main_param
 
             if fsdp_shard_model_param is not None:
                 shard_model_param = fsdp_shard_model_param
@@ -311,6 +710,13 @@ def _quantize_param_shard_impl(
                 ]
 
             quantizer = model_param._quantizer
+            main_param = _maybe_fake_mxfp4_expert_qat_main_param_shard(
+                model_param,
+                main_param,
+                start_offset,
+                copy_call_index=copy_call_index,
+            )
+            trace_main_param = main_param
             # When not using --fp8-param-gather, the main_param (fp32) is first cast to bf16/fp16,
             # and then cast to fp8 during forward.
             # Although it's not necessary when --fp8-param-gather is enabled, we still keep this
@@ -326,6 +732,9 @@ def _quantize_param_shard_impl(
                 quantizer=quantizer,
             )
             quantizer.update_quantized(main_param, out)
+            _maybe_record_mxfp4_qat_fp8_copy_trace(
+                model_param, trace_main_param, raw_main_param, start_offset, copy_call_index
+            )
 
         amaxes = []
         scales = []
@@ -387,11 +796,13 @@ def _quantize_param_shard_impl(
         if fsdp_shard_model_params is None:
             fsdp_shard_model_params = [None] * len(model_params)
 
+        copy_call_index = _next_mxfp4_qat_fp8_copy_trace_call_index()
         for model_param, main_param, start_offset, fsdp_shard_model_param in zip(
             model_params, main_params, start_offsets, fsdp_shard_model_params
         ):
             if main_param is None:
                 continue
+            raw_main_param = main_param
 
             if fsdp_shard_model_param is not None:
                 shard_model_param = fsdp_shard_model_param
@@ -404,6 +815,13 @@ def _quantize_param_shard_impl(
             # and then cast to fp8 during forward.
             # Although it's not necessary when --fp8-param-gather is enabled, we still keep this
             # logic to keep numerical consistency. So here cast the main_param to model_param.dtype.
+            main_param = _maybe_fake_mxfp4_expert_qat_main_param_shard(
+                model_param,
+                main_param,
+                start_offset,
+                copy_call_index=copy_call_index,
+            )
+            trace_main_param = main_param
             main_param = main_param.to(model_param.dtype)
             cast_to_fp8(
                 main_param.view(1, -1),
@@ -412,6 +830,9 @@ def _quantize_param_shard_impl(
                 model_param._fp8_dtype,
                 out=shard_model_param.view(1, -1),
             )
+            _maybe_record_mxfp4_qat_fp8_copy_trace(
+                model_param, trace_main_param, raw_main_param, start_offset, copy_call_index
+            )
 
         amaxes = []
         scales = []
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 486a498dd73..c4d69ee2758 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -305,6 +305,7 @@ def __init__(
         extra_block_kwargs=None,
         runtime_gather_output: Optional[bool] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Initialize the schedule plan of all Transformer layers' sub-modules.
 
@@ -322,6 +323,7 @@ def __init__(
             extra_block_kwargs: Additional keyword arguments for blocks.
             runtime_gather_output: Whether to gather output at runtime.
             loss_mask (torch.Tensor): Used to mask out some portions of the loss
+            padding_mask (torch.Tensor): Used to exclude padding tokens from MoE routing losses.
 
         Returns:
             The model chunk schedule plan.
@@ -346,6 +348,7 @@ def __init__(
         self._model_chunk_state.labels = labels
         self._model_chunk_state.mtp_hidden_states = None
         self._model_chunk_state.loss_mask = loss_mask
+        self._model_chunk_state.padding_mask = padding_mask
         self._model_chunk_state.packed_seq_params = packed_seq_params
         self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
         self._model_chunk_state.runtime_gather_output = runtime_gather_output
diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
index e6d6fa03ce7..36b095613b8 100644
--- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
+++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
@@ -1,10 +1,11 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
-from typing import Optional
+from typing import List, Optional
 
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.backends import BackendSpecProvider
 from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
-from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.experimental_attention_variant.dsa import (
     DSAIndexer,
     DSAIndexerSubmodules,
@@ -17,6 +18,34 @@
     MLASelfAttentionSubmodules,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import (
+    TransformerBlockSubmodules,
+    get_num_layers_to_build,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import (
+    TransformerLayer,
+    TransformerLayerSubmodules,
+    get_transformer_layer_offset,
+)
+
+try:
+    import transformer_engine as te  # type: ignore[import-untyped]  # pylint: disable=unused-import
+
+    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+
+    HAVE_TE = True
+except ImportError:
+    HAVE_TE = False
+
+try:
+    import nvidia_kitchen  # type: ignore[import-not-found]  # pylint: disable=unused-import
+
+    from megatron.core.extensions.kitchen import KitchenSpecProvider
+
+    HAVE_KITCHEN = True
+except ImportError:
+    HAVE_KITCHEN = False
 
 
 def is_linear_attention_variant(experimental_attention_variant: str) -> bool:
@@ -136,3 +165,258 @@ def get_experimental_attention_variant_module_spec_for_backend(
         raise ValueError(
             f"Invalid experimental attention variant: {experimental_attention_variant}"
         )
+
+
+def get_experimental_attention_variant_module_spec(
+    config: TransformerConfig, backend: Optional[BackendSpecProvider] = None
+) -> ModuleSpec:
+    """Build an attention module spec from ``config.experimental_attention_variant``.
+
+    This companion API is used by custom model specs that need to patch the
+    experimental attention spec while reusing Megatron's block construction.
+    """
+    if backend is None:
+        backend = _get_backend_spec_provider(config)
+
+    return get_experimental_attention_variant_module_spec_for_backend(
+        backend=backend,
+        sharded_state_dict_keys_map={},
+        experimental_attention_variant=config.experimental_attention_variant,
+        qk_layernorm=config.qk_layernorm,
+        qk_l2_norm=config.qk_l2_norm,
+        multi_latent_attention=config.multi_latent_attention,
+        mla_down_proj_use_column_parallel=False,
+        normalization=config.normalization,
+        fallback_to_eager_attn=config.fallback_to_eager_attn,
+    )
+
+
+def get_transformer_layer_with_experimental_attention_variant_spec(
+    config: TransformerConfig, backend: Optional[BackendSpecProvider] = None
+) -> List[ModuleSpec]:
+    """Build per-layer Transformer specs for experimental-attention GPT blocks."""
+    if backend is None:
+        backend = _get_backend_spec_provider(config)
+
+    experimental_attention_pattern = [0] * config.num_layers
+    if is_linear_attention_variant(config.experimental_attention_variant):
+        experimental_attention_pattern = _get_linear_attention_pattern(config)
+    elif config.experimental_attention_variant is not None:
+        experimental_attention_pattern = [1] * config.num_layers
+
+    experimental_attention_spec = (
+        get_experimental_attention_variant_module_spec(config=config, backend=backend)
+        if 1 in experimental_attention_pattern
+        else None
+    )
+    standard_attention_spec = (
+        _get_self_attention_module_spec(config=config, backend=backend)
+        if 0 in experimental_attention_pattern
+        else None
+    )
+
+    moe_layer_pattern = (
+        _get_moe_layer_pattern(config)
+        if config.num_moe_experts is not None
+        else [0] * config.num_layers
+    )
+
+    moe_layer_spec = (
+        _get_moe_module_spec(config=config, backend=backend)
+        if 1 in moe_layer_pattern
+        else None
+    )
+    dense_mlp_layer_spec = (
+        _get_dense_mlp_module_spec(config=config, backend=backend)
+        if 0 in moe_layer_pattern
+        else None
+    )
+
+    rms_norm = config.normalization == "RMSNorm"
+    layer_specs = []
+    for layer_number in range(config.num_layers):
+        attention = (
+            experimental_attention_spec
+            if experimental_attention_pattern[layer_number] == 1
+            else standard_attention_spec
+        )
+        mlp = moe_layer_spec if moe_layer_pattern[layer_number] == 1 else dense_mlp_layer_spec
+        input_layernorm = (
+            IdentityOp
+            if attention.metainfo["fuse_input_layernorm"]
+            else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+        )
+        pre_mlp_layernorm = (
+            IdentityOp
+            if mlp.metainfo["fuse_pre_mlp_layernorm"]
+            else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+        )
+
+        layer_specs.append(
+            ModuleSpec(
+                module=TransformerLayer,
+                submodules=TransformerLayerSubmodules(
+                    input_layernorm=input_layernorm,
+                    self_attention=attention,
+                    self_attn_bda=get_bias_dropout_add,
+                    pre_mlp_layernorm=pre_mlp_layernorm,
+                    mlp=mlp,
+                    mlp_bda=get_bias_dropout_add,
+                ),
+            )
+        )
+
+    return layer_specs
+
+
+def get_transformer_block_with_experimental_attention_variant_spec(
+    config: TransformerConfig,
+    vp_stage: Optional[int] = None,
+    pp_rank: Optional[int] = None,
+) -> TransformerBlockSubmodules:
+    """Build a TransformerBlock spec from experimental-attention layer specs."""
+    backend = _get_backend_spec_provider(config)
+    layer_specs = get_transformer_layer_with_experimental_attention_variant_spec(
+        config=config,
+        backend=backend,
+    )
+
+    if config.pipeline_model_parallel_layout is not None:
+        local_layer_ids = config.pipeline_model_parallel_layout.get_layer_id_list(
+            layer_type=LayerType.decoder,
+            vp_stage=vp_stage,
+            pp_rank=pp_rank,
+        )
+    else:
+        offset = get_transformer_layer_offset(
+            config,
+            vp_stage=vp_stage,
+            pp_rank=pp_rank,
+        )
+        num_layers_to_build = get_num_layers_to_build(
+            config,
+            vp_stage=vp_stage,
+            pp_rank=pp_rank,
+        )
+        local_layer_ids = range(offset, offset + num_layers_to_build)
+
+    rms_norm = config.normalization == "RMSNorm"
+    return TransformerBlockSubmodules(
+        layer_specs=[layer_specs[layer_id] for layer_id in local_layer_ids],
+        layer_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
+    )
+
+
+def _get_backend_spec_provider(config: TransformerConfig) -> BackendSpecProvider:
+    assert config.transformer_impl == "transformer_engine", (
+        "Experimental GPT decoder block spec only supports transformer_engine."
+    )
+    if config.use_kitchen:
+        assert HAVE_KITCHEN
+        return KitchenSpecProvider(
+            fallback=TESpecProvider(fallback_to_eager_attn=config.fallback_to_eager_attn)
+        )
+    assert HAVE_TE
+    return TESpecProvider(fallback_to_eager_attn=config.fallback_to_eager_attn)
+
+
+def _get_moe_layer_pattern(config: TransformerConfig) -> List[int]:
+    if isinstance(config.moe_layer_freq, int):
+        return [1 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.num_layers)]
+    if isinstance(config.moe_layer_freq, list):
+        assert len(config.moe_layer_freq) == config.num_layers, (
+            f"Invalid length of moe_layer_freq: {len(config.moe_layer_freq)}, "
+            f"expected {config.num_layers}."
+        )
+        return config.moe_layer_freq
+    raise ValueError(f"Invalid moe_layer_freq: {type(config.moe_layer_freq)}, {config.moe_layer_freq}")
+
+
+def _get_linear_attention_pattern(config: TransformerConfig) -> List[int]:
+    if isinstance(config.linear_attention_freq, int):
+        return [
+            0 if ((i + 1) % config.linear_attention_freq == 0) else 1
+            for i in range(config.num_layers)
+        ]
+    if isinstance(config.linear_attention_freq, list):
+        assert len(config.linear_attention_freq) == config.num_layers, (
+            f"Invalid length of linear_attention_freq: {len(config.linear_attention_freq)}, "
+            f"expected {config.num_layers}."
+        )
+        return config.linear_attention_freq
+    if config.linear_attention_freq is None:
+        if is_linear_attention_variant(config.experimental_attention_variant):
+            return [1] * config.num_layers
+        return [0] * config.num_layers
+    raise ValueError(
+        f"Invalid linear_attention_freq: {type(config.linear_attention_freq)}, "
+        f"{config.linear_attention_freq}"
+    )
+
+
+def _get_self_attention_module_spec(
+    config: TransformerConfig,
+    backend: Optional[BackendSpecProvider] = None,
+) -> ModuleSpec:
+    if backend is None:
+        backend = _get_backend_spec_provider(config)
+
+    from megatron.core.models.gpt.gpt_layer_specs import (
+        get_gpt_layer_with_transformer_engine_spec,
+    )
+
+    layer_spec = get_gpt_layer_with_transformer_engine_spec(
+        num_experts=config.num_moe_experts,
+        moe_grouped_gemm=config.moe_grouped_gemm,
+        qk_layernorm=config.qk_layernorm,
+        multi_latent_attention=config.multi_latent_attention,
+        moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
+        qk_l2_norm=config.qk_l2_norm,
+        use_kitchen=config.use_kitchen,
+        use_te_activation_func=config.use_te_activation_func,
+        fallback_to_eager_attn=config.fallback_to_eager_attn,
+    )
+    attn_spec = layer_spec.submodules.self_attention
+    if config.multi_latent_attention:
+        attn_spec.metainfo["fuse_input_layernorm"] = False
+    else:
+        attn_spec.metainfo["fuse_input_layernorm"] = backend.fuse_layernorm_and_linear()
+    return attn_spec
+
+
+def _get_dense_mlp_module_spec(
+    config: TransformerConfig,
+    backend: Optional[BackendSpecProvider] = None,
+) -> ModuleSpec:
+    if backend is None:
+        backend = _get_backend_spec_provider(config)
+
+    from megatron.core.models.gpt.gpt_layer_specs import get_mlp_module_spec_for_backend
+
+    mlp_spec = get_mlp_module_spec_for_backend(
+        backend=backend,
+        num_experts=None,
+        use_te_activation_func=config.use_te_activation_func,
+    )
+    mlp_spec.metainfo["fuse_pre_mlp_layernorm"] = backend.fuse_layernorm_and_linear()
+    return mlp_spec
+
+
+def _get_moe_module_spec(
+    config: TransformerConfig,
+    backend: Optional[BackendSpecProvider] = None,
+) -> ModuleSpec:
+    if backend is None:
+        backend = _get_backend_spec_provider(config)
+
+    from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend
+
+    mlp_spec = get_moe_module_spec_for_backend(
+        backend=backend,
+        num_experts=config.num_moe_experts,
+        moe_grouped_gemm=config.moe_grouped_gemm,
+        moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
+        use_te_activation_func=config.use_te_activation_func,
+    )
+    mlp_spec.metainfo["fuse_pre_mlp_layernorm"] = False
+    return mlp_spec
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 60094976a9a..3a8ab781a6c 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -120,12 +120,20 @@ def forward_impl(self):
         if not self.gpt_model.pre_process:
             self.chunk_state.decoder_input = self.gpt_model.decoder.input_tensor
         # Run GPTModel._preprocess
-        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = (
             self.gpt_model._preprocess(
                 input_ids=self.chunk_state.input_ids,
                 position_ids=self.chunk_state.position_ids,
                 decoder_input=self.chunk_state.decoder_input,
                 packed_seq_params=self.chunk_state.packed_seq_params,
+                padding_mask=self.chunk_state.padding_mask,
             )
         )
 
@@ -135,6 +143,7 @@ def forward_impl(self):
         self.chunk_state.rotary_pos_cos = rotary_pos_cos
         self.chunk_state.rotary_pos_sin = rotary_pos_sin
         self.chunk_state.sequence_len_offset = sequence_len_offset
+        self.chunk_state.padding_mask = padding_mask
         return decoder_input
 
 
@@ -383,7 +392,10 @@ def submodule_post_attn_forward(node: ScheduleNode, hidden_states: torch.Tensor)
             with get_fine_grained_offloading_context(layer.offload_mlp_norm):
                 pre_mlp_layernorm_output = layer.pre_mlp_layernorm(hidden_states)
 
-        probs, routing_map = layer.mlp.route(pre_mlp_layernorm_output)
+        probs, routing_map = layer.mlp.route(
+            pre_mlp_layernorm_output,
+            padding_mask=node.chunk_state.padding_mask,
+        )
         local_tokens, probs, _ = layer.mlp.preprocess(pre_mlp_layernorm_output, probs, routing_map)
 
         # Detach here for mlp_bda residual connection
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 1fd52f65a5d..b1263bfd7c7 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -284,6 +284,7 @@ def _preprocess(
         decoder_input: Tensor = None,
         inference_context: BaseInferenceContext = None,
         packed_seq_params: PackedSeqParams = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Preprocesses inputs for the transformer decoder.
 
@@ -297,6 +298,12 @@ def _preprocess(
         in_inference_mode = inference_context is not None and not self.training
 
         # Decoder embedding.
+        if padding_mask is not None:
+            assert padding_mask.shape == input_ids.shape, (
+                f"padding_mask shape {padding_mask.shape} does not match "
+                f"input_ids shape {input_ids.shape}"
+            )
+
         if decoder_input is not None:
             pass
         elif self.pre_process:
@@ -306,6 +313,15 @@ def _preprocess(
             # decoder will get hidden_states from encoder.input_tensor
             decoder_input = None
 
+        if padding_mask is not None and self.config.sequence_parallel:
+            padding_mask = (
+                tensor_parallel.scatter_to_sequence_parallel_region(
+                    padding_mask.transpose(0, 1).contiguous()
+                )
+                .transpose(0, 1)
+                .contiguous()
+            )
+
         # Rotary positional embeddings (embedding is None for PP intermediate devices)
         rotary_pos_emb = None
         rotary_pos_cos = None
@@ -403,13 +419,14 @@ def _preprocess(
             rotary_pos_cos,
             rotary_pos_sin,
             sequence_len_offset,
+            padding_mask,
         )
         if rotary_pos_cos_sin is not None:
             # only in the case of flashinfer fused rope will we
             # return this extra tensor
             # this is for backwards compatibility with
             # legacy unit tests, which break if you
-            # return a 6 tuple instead of 5.
+            # return a 7 tuple instead of 6.
             preproc_output += (rotary_pos_cos_sin,)
 
         return preproc_output
@@ -446,6 +463,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         mtp_kwargs: Optional[dict] = {},
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
@@ -469,13 +487,19 @@ def forward(
             decoder_input=decoder_input,
             inference_context=inference_context,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
-        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset) = (
-            preproc_output[:5]
-        )
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = preproc_output[:6]
 
-        rotary_pos_cos_sin = preproc_output[5] if len(preproc_output) == 6 else None
+        rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
 
         # Run decoder.
         hidden_states = self.decoder(
@@ -488,6 +512,8 @@ def forward(
             rotary_pos_cos_sin=rotary_pos_cos_sin,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
+            padding_mask=padding_mask,
+            input_ids=input_ids,
             **(extra_block_kwargs or {}),
         )
 
@@ -733,6 +759,7 @@ def build_schedule_plan(
         runtime_gather_output: Optional[bool] = None,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Builds a computation schedule plan for the model.
 
@@ -758,6 +785,7 @@ def build_schedule_plan(
             inference_params (InferenceParams, optional):
                 Parameters for inference. Defaults to None.
             loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
+            padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
 
         Returns:
             TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -779,6 +807,7 @@ def build_schedule_plan(
             extra_block_kwargs,
             runtime_gather_output,
             loss_mask,
+            padding_mask,
         )
 
     def sharded_state_dict(
diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
index 28487c3b367..24fd1c7d51e 100644
--- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
+++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION and Alibaba PAI. All rights reserved.
 from collections import defaultdict
-from typing import Dict
+from typing import Any, Dict
 
 import torch
 
@@ -93,6 +93,7 @@ def _set_sub_optimizer_grads(self):
                     fp32_param.grad = grad.to(fp32_param.dtype)
                     fp32_param.requires_grad = True
                 else:
+                    fp32_param.grad = None
                     fp32_param.requires_grad = False
 
         # Sync the grads from GPU to CPU.
@@ -101,6 +102,7 @@ def _set_sub_optimizer_grads(self):
                 gpu_param = self.cpu_copys_map_gpu_param[param]
                 grad = getattr(gpu_param, "decoupled_grad", gpu_param.grad)
                 if grad is None:
+                    param.grad = None
                     param.requires_grad = False
                     continue
 
@@ -122,7 +124,7 @@ def param_copy_back_gpu_hook(optimizer, args, kwargs):
                     for param in _param_generator(optimizer):
                         gpu_param = self.cpu_copys_map_gpu_param[param]
                         gpu_param.data.copy_(param.data, non_blocking=True)
-                self._d2h_stream.record_event().wait(torch.cuda.current_stream())
+                self._h2d_stream.record_event().wait(torch.cuda.current_stream())
 
             return param_copy_back_gpu_hook
 
@@ -370,15 +372,20 @@ def _update_fp32_params_by_new_state(self):
         if not self.param_update_in_fp32:
             return
         for param, v in self.state.items():
-            fp32_param = self.param_to_fp32_param[param]
-            fp32_param.data.copy_(v["master_param"])
+            inner_param = self.param_to_inner_param.get(param, param)
+            if inner_param is param:
+                continue
+            inner_param.data.copy_(v["master_param"].detach(), non_blocking=False)
 
     def update_fp32_param_by_new_param(self):
         """
-        Update the fp32 parameters by the new parameters.
+        Refresh optimizer-side parameter copies after model weights are loaded
+        or otherwise changed outside the optimizer.
         """
-        for param, fp32_param in self.param_to_fp32_param.items():
-            fp32_param.data.copy_(param)
+        for param, inner_param in self.param_to_inner_param.items():
+            if inner_param is param:
+                continue
+            inner_param.data.copy_(param.detach(), non_blocking=False)
 
     def _register_load_state_dict_hooks(self):
         def pre_load_state_dict_hook(self, state_dict):
@@ -442,6 +449,8 @@ def zero_grad(self, set_to_none: bool = True):
         Zero or zero to none the gradients of all the parameters in the model.
         """
         super(HybridDeviceOptimizer, self).zero_grad(set_to_none)
+        for optimizer in self.sub_optimizers:
+            optimizer.zero_grad(set_to_none)
         for group in self.param_groups:
             for param in group["params"]:
                 if hasattr(param, "decoupled_grad"):
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index eac21a3ea8e..3a8f5f6bec3 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -330,9 +330,10 @@ def _build_model_and_main_param_groups(
         shard_float16_groups = []
         shard_fp32_groups = []
         shard_fp32_from_float16_groups = []
+        model_param_group_index_map = {}
 
         # Allocate (or slice) each group's param shard.
-        for group_range in opt_group_ranges:
+        for group_index, group_range in enumerate(opt_group_ranges):
 
             # Params of this group.
             model_float16_params_this_group = []
@@ -415,7 +416,13 @@ def _build_model_and_main_param_groups(
 
                 # fp32 params.
                 elif model_param.type() == 'torch.cuda.FloatTensor':
-                    shard_model_param = model_param.view(-1)[param_range.start : param_range.end]
+                    # HybridDeviceOptimizer rebuilds torch optimizer param groups from these
+                    # shards. Wrap the shard view as a leaf Parameter so torch.optim accepts
+                    # trainable FP32 params kept by model-specific _keep_fp32 handling.
+                    shard_model_param = torch.nn.Parameter(
+                        model_param.detach().view(-1)[param_range.start : param_range.end],
+                        requires_grad=True,
+                    )
                     model_fp32_params_this_group.append(model_param)
                     shard_fp32_params_this_group.append(shard_model_param)
                     tensor_parallel.copy_tensor_model_parallel_attributes(
@@ -445,12 +452,19 @@ def _build_model_and_main_param_groups(
                     *shard_float16_params_this_group,
                 ]
 
+            for group_order, model_param in enumerate(model_fp32_params_this_group):
+                model_param_group_index_map[model_param] = (group_index, group_order)
+            offset = len(model_fp32_params_this_group)
+            for i, model_param in enumerate(model_float16_params_this_group):
+                model_param_group_index_map[model_param] = (group_index, offset + i)
+
         return (
             model_float16_groups,
             model_fp32_groups,
             shard_float16_groups,
             shard_fp32_groups,
             shard_fp32_from_float16_groups,
+            model_param_group_index_map,
         )
 
     def __init__(
@@ -581,7 +595,7 @@ def __init__(
                         param.main_param_sharded = True
 
         # Optimizer ranges.
-        (self.model_param_group_index_map, self.opt_group_ranges) = (
+        (_, self.opt_group_ranges) = (
             self._build_optimizer_group_ranges(self.optimizer.param_groups, self.gbuf_ranges)
         )
 
@@ -592,6 +606,7 @@ def __init__(
             self.shard_float16_groups,
             self.shard_fp32_groups,
             self.shard_fp32_from_float16_groups,
+            self.model_param_group_index_map,
         ) = self._build_model_and_main_param_groups(
             self.gbuf_ranges, self.model_param_gbuf_map, self.opt_group_ranges, config
         )
diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py
index f18309217c3..6c7561af548 100644
--- a/megatron/core/pipeline_parallel/p2p_communication.py
+++ b/megatron/core/pipeline_parallel/p2p_communication.py
@@ -181,17 +181,18 @@ def _communicate_shapes(self, tensor_send_next, tensor_send_prev, recv_prev, rec
             (recv_prev_shape, recv_next_shape)
         """
         config = self.config
+        num_dims = 4 if config.dsv4_mode else 3
         recv_prev_shape_tensor = None
         recv_next_shape_tensor = None
         send_prev_shape_tensor = None
         send_next_shape_tensor = None
         if recv_prev:
             recv_prev_shape_tensor = torch.empty(
-                (3,), device=torch.cuda.current_device(), dtype=torch.int64
+                (num_dims,), device=torch.cuda.current_device(), dtype=torch.int64
             )
         if recv_next:
             recv_next_shape_tensor = torch.empty(
-                (3,), device=torch.cuda.current_device(), dtype=torch.int64
+                (num_dims,), device=torch.cuda.current_device(), dtype=torch.int64
             )
         if tensor_send_prev is not None:
             send_prev_shape_tensor = torch.tensor(
@@ -241,11 +242,11 @@ def _communicate_shapes(self, tensor_send_next, tensor_send_prev, recv_prev, rec
             # should take this out once the bug with batch_isend_irecv is resolved.
             torch.cuda.synchronize()
 
-        recv_prev_shape = [0, 0, 0]
+        recv_prev_shape = [0] * num_dims
         if recv_prev_shape_tensor is not None:
             recv_prev_shape = recv_prev_shape_tensor.tolist()
 
-        recv_next_shape = [0, 0, 0]
+        recv_next_shape = [0] * num_dims
         if recv_next_shape_tensor is not None:
             recv_next_shape = recv_next_shape_tensor.tolist()
 
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index a8fdf2324f2..9adb8fe483c 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -1036,7 +1036,10 @@ def enable_grad_sync():
 
     model_type = get_model_type(model[0])
 
-    tensor_shape = [seq_length, micro_batch_size, config.hidden_size]
+    if config.dsv4_mode:
+        tensor_shape = [seq_length, micro_batch_size, config.dsv4_hc_mult, config.hidden_size]
+    else:
+        tensor_shape = [seq_length, micro_batch_size, config.hidden_size]
     tensor_shape[0] = tensor_shape[0] // cp_group.size()
     if config.sequence_parallel:
         tensor_shape[0] = tensor_shape[0] // tp_group.size()
@@ -1982,7 +1985,12 @@ def get_tensor_shapes(
     if config.sequence_parallel:
         effective_seq_length = effective_seq_length // tp_group.size()
 
-    tensor_shapes.append((effective_seq_length, micro_batch_size, config.hidden_size))
+    if config.dsv4_mode:
+        tensor_shapes.append(
+            (effective_seq_length, micro_batch_size, config.dsv4_hc_mult, config.hidden_size)
+        )
+    else:
+        tensor_shapes.append((effective_seq_length, micro_batch_size, config.hidden_size))
     return tensor_shapes
 
 
diff --git a/megatron/core/tensor_parallel/mappings.py b/megatron/core/tensor_parallel/mappings.py
index 9ff69c9dc31..ed21bedc98c 100644
--- a/megatron/core/tensor_parallel/mappings.py
+++ b/megatron/core/tensor_parallel/mappings.py
@@ -19,7 +19,7 @@
     dist_reduce_scatter_func = torch.distributed._reduce_scatter_base
 
 
-def _reduce(input_, group):
+def _reduce(input_, group, fp32=False):
     """All-reduce the input tensor across model parallel group."""
     assert group is not None, "group should not be None"
 
@@ -28,7 +28,13 @@ def _reduce(input_, group):
         return input_
 
     # All-reduce.
-    torch.distributed.all_reduce(input_.contiguous(), group=group)
+    if fp32:
+        orig_dtype = input_.dtype
+        input_fp32 = input_.float().contiguous()
+        torch.distributed.all_reduce(input_fp32, group=group)
+        input_.copy_(input_fp32.to(orig_dtype))
+    else:
+        torch.distributed.all_reduce(input_.contiguous(), group=group)
 
     return input_
 
@@ -77,6 +83,24 @@ def _split_along_first_dim(input_, group):
     return output
 
 
+def split_along_nth_dim(input_, dim, group):
+    """Split a tensor along an arbitrary dimension and keep this rank's slice."""
+    assert group is not None, "group should not be None"
+
+    world_size = group.size()
+    if world_size == 1:
+        return input_
+
+    dim_size = input_.size(dim)
+    assert (
+        dim_size % world_size == 0
+    ), f"Dimension {dim} of the tensor should be divisible by tensor parallel size"
+    local_dim_size = dim_size // world_size
+    rank = group.rank()
+    dim_offset = rank * local_dim_size
+    return input_.narrow(dim, dim_offset, local_dim_size).contiguous()
+
+
 def _gather_along_last_dim(input_, group):
     """Gather tensors and concatinate along the last dimension."""
 
@@ -198,20 +222,21 @@ class _CopyToModelParallelRegion(torch.autograd.Function):
     """Pass the input to the model parallel region."""
 
     @staticmethod
-    def symbolic(graph, input_, group):
+    def symbolic(graph, input_, group, all_reduce_grad_fp32):
         """Symbolic function for tracing."""
         return input_
 
     @staticmethod
-    def forward(ctx, input_, group):
+    def forward(ctx, input_, group, all_reduce_grad_fp32):
         """Forward function."""
         ctx.group = group
+        ctx.all_reduce_grad_fp32 = all_reduce_grad_fp32
         return input_
 
     @staticmethod
     def backward(ctx, grad_output):
         """Backward function."""
-        return _reduce(grad_output, ctx.group), None
+        return _reduce(grad_output, ctx.group, fp32=ctx.all_reduce_grad_fp32), None, None
 
 
 class _ReduceFromModelParallelRegion(torch.autograd.Function):
@@ -466,10 +491,10 @@ def backward(ctx, *grad_output):
 # -----------------
 
 
-def copy_to_tensor_model_parallel_region(input_, group=None):
+def copy_to_tensor_model_parallel_region(input_, group=None, all_reduce_grad_fp32=False):
     """Wrapper for autograd function: forward: copy, backward allreduce"""
     group = get_tensor_model_parallel_group_if_none(group)
-    return _CopyToModelParallelRegion.apply(input_, group)
+    return _CopyToModelParallelRegion.apply(input_, group, all_reduce_grad_fp32)
 
 
 def reduce_from_tensor_model_parallel_region(input_, group=None):
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index 8dcf196da94..afd030711c8 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -101,12 +101,23 @@ def __init__(
 
         # If this is a gated linear unit we double the output width
         # see https://arxiv.org/pdf/2002.05202.pdf
+        # For GLU/SwiGLU, TP shards store interleaved [gate, up] portions.
+        # Preserving stride=2 is required for correct checkpoint resharding.
         if self.config.gated_linear_unit:
             ffn_hidden_size *= 2
+            fc1_stride = 2
+            if self.config.use_kitchen:
+                # Kitchen Linear doesn't support stride != 1.
+                fc1_stride = 1
+        else:
+            fc1_stride = 1
+
+        # Use moe_latent_size only for routed experts. Shared experts stay on hidden_size.
+        use_latent_size = (self.config.moe_latent_size is not None) and is_expert
 
         self.linear_fc1 = build_module(
             submodules.linear_fc1,
-            self.input_size,
+            self.input_size if not use_latent_size else self.config.moe_latent_size,
             ffn_hidden_size,
             config=self.config,
             init_method=self.config.init_method,
@@ -116,6 +127,7 @@ def __init__(
             is_expert=is_expert,
             tp_comm_buffer_name="fc1",
             tp_group=tp_group,
+            stride=fc1_stride,
         )
 
         if self.config.use_te_activation_func and not (submodules.activation_func is None):
@@ -126,7 +138,7 @@ def __init__(
         self.linear_fc2 = build_module(
             submodules.linear_fc2,
             self.config.ffn_hidden_size,
-            self.config.hidden_size,
+            self.config.hidden_size if not use_latent_size else self.config.moe_latent_size,
             config=self.config,
             init_method=self.config.output_layer_init_method,
             bias=self.config.add_bias_linear,
diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py
index 2330df91b52..9ed43105569 100644
--- a/megatron/core/transformer/module.py
+++ b/megatron/core/transformer/module.py
@@ -376,6 +376,45 @@ def float_conversion(val):
     return conversion_helper(val, float_conversion)
 
 
+def _detach_keep_fp32_tensors(module: torch.nn.Module):
+    """Temporarily remove tensors marked _keep_fp32 from recursive dtype conversion."""
+    preserved_parameters = []
+    preserved_buffers = []
+    for child_module in module.modules():
+        for name, param in child_module._parameters.items():
+            if param is not None and getattr(param, '_keep_fp32', False):
+                assert param.dtype == torch.float32, (
+                    f"Parameter {name} is marked _keep_fp32 but has dtype {param.dtype}"
+                )
+                preserved_parameters.append((child_module, name, param))
+                child_module._parameters[name] = None
+        for name, buffer in child_module._buffers.items():
+            if buffer is not None and getattr(buffer, '_keep_fp32', False):
+                assert buffer.dtype == torch.float32, (
+                    f"Buffer {name} is marked _keep_fp32 but has dtype {buffer.dtype}"
+                )
+                preserved_buffers.append((child_module, name, buffer))
+                child_module._buffers[name] = None
+    return preserved_parameters, preserved_buffers
+
+
+def _restore_keep_fp32_tensors(preserved_tensors):
+    preserved_parameters, preserved_buffers = preserved_tensors
+    for child_module, name, param in preserved_parameters:
+        child_module._parameters[name] = param
+    for child_module, name, buffer in preserved_buffers:
+        child_module._buffers[name] = buffer
+
+
+def _convert_module_preserving_fp32_tensors(module: torch.nn.Module, convertor):
+    """Convert a module to low precision while preserving explicit fp32 tensors."""
+    preserved_tensors = _detach_keep_fp32_tensors(module)
+    try:
+        return convertor(module)
+    finally:
+        _restore_keep_fp32_tensors(preserved_tensors)
+
+
 class Float16Module(MegatronModule):
     """Float 16 Module.
 
@@ -398,13 +437,17 @@ def __init__(self, config: TransformerConfig, module: torch.nn.Module):
         self.pg_collection = getattr(module, 'pg_collection', None)
 
         if self.fp16:
-            self.add_module('module', module.half())
+            self.add_module(
+                'module', _convert_module_preserving_fp32_tensors(module, lambda m: m.half())
+            )
 
             def float16_convertor(val):
                 return val.half()
 
         elif self.bf16:
-            self.add_module('module', module.bfloat16())
+            self.add_module(
+                'module', _convert_module_preserving_fp32_tensors(module, lambda m: m.bfloat16())
+            )
 
             def float16_convertor(val):
                 return val.bfloat16()
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 10d10f667fe..8b8b30dcb86 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -93,6 +93,13 @@ def set_layer_number(self, layer_number: int):
         self.layer_number = layer_number
         self.router.set_layer_number(layer_number)
 
+    def set_is_mtp(self):
+        """Mark this MoE layer as an MTP layer."""
+        if hasattr(self.router, 'set_is_mtp'):
+            self.router.set_is_mtp()
+        else:
+            self.router.is_mtp = True
+
 
 class MoELayer(BaseMoELayer):
     """Mixture of Experts layer.
@@ -177,14 +184,13 @@ def __init__(
         # Cudagraph tensor store for resuming the forward pass from the end of the cudagraph.
         self.cudagraph_tensor_store = MoECudaGraphTensorStore()
 
-    @maybe_skip_or_early_return_by_cudagraph("route")
-    def route(self, hidden_states: torch.Tensor):
+    def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None):
         """Compute token routing for preprocessing.
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
         """
-        probs, routing_map = self.router(hidden_states)
+        probs, routing_map = self.router(hidden_states, padding_mask, input_ids=input_ids)
         return probs, routing_map
 
     @maybe_skip_or_early_return_by_cudagraph("preprocess")
@@ -270,7 +276,7 @@ def combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Ten
             output = output + shared_expert_output
         return output
 
-    def forward(self, hidden_states: torch.Tensor):
+    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None):
         """Forward pass for the MoE layer.
 
         The forward pass comprises four main steps:
@@ -281,6 +287,8 @@ def forward(self, hidden_states: torch.Tensor):
 
         Args:
             hidden_states (torch.Tensor): The input tensor to the MoE layer.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+            input_ids (torch.Tensor, optional): Input token IDs for DSV4 hash routing.
 
         Returns:
             A tuple containing the output tensor and the MLP bias, if any.
@@ -290,12 +298,14 @@ def forward(self, hidden_states: torch.Tensor):
                 "During training, performance may degrade if MoE and tensor parallelism"
                 "are enabled without also enabling sequence parallelism."
             )
+        if padding_mask is not None:
+            padding_mask = padding_mask.transpose(0, 1).bool()
 
         # MoE forward: route -> dispatch -> compute -> combine
-        def custom_forward(hidden_states):
+        def custom_forward(hidden_states, padding_mask=None, input_ids=None):
             try:
                 shared_expert_output = self.shared_experts_compute(hidden_states)
-                probs, routing_map = self.route(hidden_states)
+                probs, routing_map = self.route(hidden_states, padding_mask, input_ids=input_ids)
                 hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map)
             except MoECudaGraphPartialCaptureSignal as e:
                 # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator.
@@ -318,11 +328,13 @@ def custom_forward(hidden_states):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     parallel_state.get_tensor_model_parallel_group(),
                     hidden_states,
+                    padding_mask,
+                    input_ids,
                 )
             else:
-                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states, padding_mask, input_ids)
         else:
-            outputs = custom_forward(hidden_states)
+            outputs = custom_forward(hidden_states, padding_mask, input_ids)
 
         return outputs
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 71912617f12..44de396ff0f 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -2,7 +2,7 @@
 
 import math
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 
@@ -11,6 +11,7 @@
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -48,6 +49,7 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ):
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
@@ -98,9 +100,18 @@ def switch_load_balancing_loss_func(
         topk (int): The number of experts selected for each token.
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
+    # Apply padding mask to probs if provided
+    if padding_mask is not None:
+        # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
+        mask_expanded = padding_mask.unsqueeze(-1)
+        probs = probs * mask_expanded
+
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -120,21 +131,60 @@ def switch_load_balancing_loss_func(
     return aux_loss
 
 
-def z_loss_func(logits, z_loss_coeff):
+def z_loss_func(
+    logits: torch.Tensor, z_loss_coeff: float, padding_mask: Optional[torch.Tensor] = None
+) -> torch.Tensor:
     """Encourages the router's logits to remain small to enhance stability.
     Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
     Args:
         logits (torch.Tensor): The logits of the router.
+        z_loss_coeff (float): The coefficient for the z-loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                               Shape [num_tokens]. True = padding (exclude),
+                                               False = valid (include). Defaults to None.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
     """
-
-    z_loss = torch.mean(torch.square(torch.logsumexp(logits, dim=-1))) * z_loss_coeff
+    logsum = torch.logsumexp(logits, dim=-1)
+    z_loss_values = torch.square(logsum)
+
+    if padding_mask is not None:
+        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
+        valid_mask = ~padding_mask
+        # Only compute z_loss for valid (non-padding) tokens
+        z_loss_values = z_loss_values * valid_mask
+        # Compute mean over valid tokens only
+        num_valid_tokens = valid_mask.sum()
+        z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff
+    else:
+        z_loss = torch.mean(z_loss_values) * z_loss_coeff
     return z_loss
 
 
+def get_tokens_per_expert_and_token_count(
+    routing_map: torch.Tensor,
+    reduce_group: torch.distributed.ProcessGroup,
+    topk: int = None,
+    with_padding_mask: bool = False,
+) -> torch.Tensor:
+    """
+    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with padding mask.
+    """
+    local_tokens_per_expert = routing_map.sum(dim=0)
+    global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
+        local_tokens_per_expert, reduce_group
+    )
+    if with_padding_mask:
+        local_num_tokens = local_tokens_per_expert.sum() / topk
+        total_num_tokens = global_tokens_per_expert.sum() / topk
+    else:
+        local_num_tokens = routing_map.shape[0]
+        total_num_tokens = local_num_tokens * reduce_group.size()
+    return global_tokens_per_expert, local_num_tokens, total_num_tokens
+
+
 def sinkhorn(cost: torch.Tensor, tol: float = 0.0001):
     """Sinkhorn based MoE routing function"""
     cost = torch.exp(cost)
@@ -537,6 +587,9 @@ def topk_routing_with_score_function(
     score_function: str = "softmax",
     expert_bias: Optional[torch.Tensor] = None,
     fused: bool = False,
+    is_mtp: bool = False,
+    tid2eid: Optional[torch.Tensor] = None,
+    input_ids: Optional[torch.Tensor] = None,
 ):
     """Compute the routing probabilities and map for top-k selection with score function.
     Args:
@@ -546,8 +599,12 @@ def topk_routing_with_score_function(
         num_groups (int): Number of groups for routed experts.
         group_topk (int): Number of selected groups for each token.
         scaling_factor (float): Scaling factor of routing score in top-k selection.
-        score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
+        score_function (str): The score function to use. Can be "softmax", "sigmoid",
+            or "sqrtsoftplus".
         expert_bias (torch.Tensor): The bias added to logits for expert routing.
+        is_mtp (bool, optional): Whether this is an MTP layer. MTP layers bypass routing replay.
+        tid2eid (torch.Tensor, optional): Token-to-expert-id mapping for DSV4 hash routing.
+        input_ids (torch.Tensor, optional): Flat input token IDs for hash routing lookup.
     Returns:
         Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             - routing_probs (torch.Tensor): A tensor of shape [num_tokens, num_experts] containing
@@ -587,8 +644,11 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
         else:
             return torch.topk(scores, k=topk, dim=1)
 
-    from sirl.utils.routing_replay import get_routing_replay_compute_topk
-    compute_topk = get_routing_replay_compute_topk(compute_topk)
+    from sirl.utils.replay_base import routing_replay_manager
+    # MTP layers and hash-routed layers (tid2eid is not None) bypass replay
+    # since MTP routing is non-standard and hash routing is deterministic.
+    if not is_mtp and tid2eid is None:
+        compute_topk = routing_replay_manager.get_topk_fn(compute_topk, return_probs=True)
 
     if score_function == "softmax":
         if use_pre_softmax:
@@ -606,6 +666,25 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
         else:
             scores, top_indices = compute_topk(scores, topk, num_groups, group_topk)
         probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores
+    elif score_function == "sqrtsoftplus":
+        assert num_groups is None
+        assert group_topk is None
+        scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits)
+        if tid2eid is not None:
+            assert not tid2eid.requires_grad
+            assert input_ids is not None and not input_ids.requires_grad
+            assert input_ids.numel() == logits.shape[0], (
+                f"input_ids token count {input_ids.numel()} does not match router logits "
+                f"token count {logits.shape[0]}"
+            )
+            top_indices = tid2eid[input_ids].long()
+            assert torch.all(top_indices >= 0)
+        else:
+            assert expert_bias is not None
+            scores_for_routing = scores + expert_bias
+            _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)
+        scores = torch.gather(scores, dim=1, index=top_indices).type_as(logits)
+        probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
     else:
         raise ValueError(f"Invalid score_function: {score_function}")
 
@@ -632,12 +711,22 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
 
 
 def compute_routing_scores_for_aux_loss(
-    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False
+    logits: torch.Tensor,
+    topk: int,
+    score_function: str,
+    fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ):
     """Compute routing scores based on the score function.
 
     Args:
         logits (torch.Tensor): The logits tensor after gating, shape: [num_tokens, num_experts].
+        topk (int): The number of top-k indices to compute.
+        score_function (str): The score function to use. Can be either "softmax" or "sigmoid".
+        fused (bool, optional): Whether to use the fused version. Defaults to False.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
 
     Returns:
         torch.Tensor: The normalized routing scores.
@@ -647,20 +736,30 @@ def compute_routing_scores_for_aux_loss(
             raise ValueError(
                 "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
             )
-        return fused_compute_score_for_moe_aux_loss(
+        routing_map, scores = fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
-
-    if score_function == "softmax":
-        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
-    elif score_function == "sigmoid":
-        scores = torch.sigmoid(logits)
-        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
     else:
-        raise ValueError(f"Invalid score_function: {score_function}")
+        if score_function == "softmax":
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        elif score_function == "sigmoid":
+            scores = torch.sigmoid(logits)
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        elif score_function == "sqrtsoftplus":
+            scores = torch.nn.functional.softplus(logits.float()).sqrt()
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        else:
+            raise ValueError(f"Invalid score_function: {score_function}")
+
+        _, top_indices = torch.topk(scores, k=topk, dim=1)
+        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
 
-    _, top_indices = torch.topk(scores, k=topk, dim=1)
-    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    # Apply padding mask to scores if provided
+    if padding_mask is not None:
+        # Invert padding_mask and make True indicates valid tokens
+        valid_mask = (~padding_mask).unsqueeze(-1)
+        routing_map = routing_map * valid_mask
+        scores = scores * valid_mask
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index ae238f93119..141aeee28fb 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -6,7 +6,6 @@
 import torch
 
 from megatron.core.jit import jit_fuser
-from megatron.core.tensor_parallel import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,
@@ -14,6 +13,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
+    get_tokens_per_expert_and_token_count,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -28,7 +28,10 @@ class Router(ABC, MegatronModule):
     """Base Router class"""
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None
+        self,
+        config: TransformerConfig,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        layer_number: Optional[int] = None,
     ) -> None:
         """
         Initialize the Router module.
@@ -42,6 +45,7 @@ def __init__(
         self.num_experts = self.config.num_moe_experts
         self.moe_aux_loss_func = None
         self.layer_number = None
+        self.is_mtp = False
         self.tp_group = pg_collection.tp
         self.cp_group = pg_collection.cp
         self.tp_cp_group = pg_collection.tp_cp
@@ -62,6 +66,10 @@ def __init__(
         # So we need to know if the model is configured to calculate per token loss.
         self.calculate_per_token_loss = self.config.calculate_per_token_loss
         self.reset_parameters()
+        if self.config.moe_router_freeze_gate:
+            self.weight.requires_grad = False
+            if self.bias is not None:
+                self.bias.requires_grad = False
 
     def reset_parameters(self):
         """Reset the router parameters."""
@@ -90,6 +98,11 @@ def gating(self, input: torch.Tensor):
         if self.bias is not None and self.bias.device.type == 'cpu':
             self.bias.data = self.bias.data.to(device=torch.cuda.current_device())
 
+        if self.config.moe_router_freeze_gate:
+            assert not self.weight.requires_grad
+            if self.bias is not None:
+                assert not self.bias.requires_grad
+
         # Convert to specified datatype for routing computation if enabled
         router_dtype = input.dtype
         if self.config.moe_router_dtype == 'fp32':
@@ -100,7 +113,7 @@ def gating(self, input: torch.Tensor):
         return logits
 
     @abstractmethod
-    def routing(self, logits: torch.Tensor):
+    def routing(self, logits: torch.Tensor, input_ids: Optional[torch.Tensor] = None):
         """Routing function.
 
         Args:
@@ -113,12 +126,14 @@ def routing(self, logits: torch.Tensor):
         raise NotImplementedError("Routing function not implemented.")
 
     @abstractmethod
-    def forward(self, input: torch.Tensor):
+    def forward(self, input: torch.Tensor, input_ids: Optional[torch.Tensor] = None):
         """
         Forward pass of the router.
 
         Args:
             input (torch.Tensor): Input tensor.
+            input_ids (torch.Tensor, optional): Input token IDs for routing modes that
+                depend on token identity.
         """
         raise NotImplementedError("Forward function not implemented.")
 
@@ -144,42 +159,34 @@ class TopKRouter(Router):
     """
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None
+        self,
+        config: TransformerConfig,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        layer_number: Optional[int] = None,
     ) -> None:
         """Initialize the zero token dropping router.
 
         Args:
             config (TransformerConfig): The configuration for the transformer model.
             pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations.
+            layer_number (int, optional): Layer number for DeepSeek-V4 hash routing.
         """
         super().__init__(config=config, pg_collection=pg_collection)
+        self.layer_number = layer_number
         self.topk = self.config.moe_router_topk
         self.routing_type = self.config.moe_router_load_balancing_type
         self.score_function = self.config.moe_router_score_function
         self.input_jitter = None
 
-        self.enable_expert_bias = self.config.moe_router_enable_expert_bias
-        if self.enable_expert_bias:
-            self.register_buffer(
-                'local_tokens_per_expert',
-                torch.zeros(
-                    self.config.num_moe_experts,
-                    dtype=torch.float32,
-                    device=torch.cuda.current_device(),
-                ),
-                persistent=False,
-            )
-            self.register_buffer(
-                'expert_bias',
-                torch.zeros(
-                    self.config.num_moe_experts,
-                    dtype=torch.float32,
-                    device=torch.cuda.current_device(),
-                ),
-            )
-        else:
-            self.local_tokens_per_expert = None
-            self.expert_bias = None
+        self._routing_mode_initialized = False
+        self.enable_expert_bias = False
+        self.tid2eid = None
+        self._frozen_expert_bias_snapshot = None
+        self._routing_replay_registered = False
+        if layer_number is not None:
+            self._init_routing_mode(layer_number)
+        elif not self.config.dsv4_mode:
+            self._init_routing_mode(0)
 
         # Initialize global tokens per expert for global aux loss
         if self.get_aux_loss_coeff("global_aux_loss") > 0:
@@ -201,8 +208,97 @@ def __init__(
             self.global_tokens_per_expert = None
             self.ga_steps = None
 
-        from sirl.utils.routing_replay import register_routing_replay
-        register_routing_replay(self)
+        self._register_routing_replay_if_needed()
+
+    def _register_routing_replay_if_needed(self):
+        from sirl.utils.replay_base import routing_replay_manager
+
+        if self._routing_replay_registered:
+            if hasattr(self, "routing_replay"):
+                self.routing_replay.metadata.update(
+                    routing_replay_manager._module_metadata(self, "routing_replay")
+                )
+            return
+        routing_replay_manager.register_to_module(self, "routing_replay")
+        self._routing_replay_registered = hasattr(self, "routing_replay")
+
+    def _set_expert_bias_buffers(self, enabled: bool):
+        if enabled:
+            local_tokens_per_expert = torch.zeros(
+                self.config.num_moe_experts,
+                dtype=torch.float32,
+                device=torch.cuda.current_device(),
+            )
+            expert_bias = torch.zeros(
+                self.config.num_moe_experts,
+                dtype=torch.float32,
+                device=torch.cuda.current_device(),
+            )
+        else:
+            local_tokens_per_expert = None
+            expert_bias = None
+
+        if 'local_tokens_per_expert' in self._buffers:
+            self.local_tokens_per_expert = local_tokens_per_expert
+        else:
+            self.register_buffer(
+                'local_tokens_per_expert',
+                local_tokens_per_expert,
+                persistent=False,
+            )
+
+        if 'expert_bias' in self._buffers:
+            self.expert_bias = expert_bias
+        else:
+            self.register_buffer('expert_bias', expert_bias)
+
+    def _init_routing_mode(self, layer_number):
+        assert not self._routing_mode_initialized
+        self._routing_mode_initialized = True
+
+        mode_hash = (
+            self.config.dsv4_mode
+            and layer_number <= self.config.dsv4_n_hash_layers
+            and not self.is_mtp
+        )
+
+        self.enable_expert_bias = (
+            self.config.moe_router_enable_expert_bias and not mode_hash
+        )
+        self._set_expert_bias_buffers(self.enable_expert_bias)
+
+        if mode_hash:
+            self.tid2eid = torch.nn.Parameter(
+                torch.full(
+                    (self.config.vocab_size, self.topk),
+                    fill_value=-1,
+                    dtype=torch.int32,
+                ),
+                requires_grad=False,
+            )
+
+    def set_is_mtp(self):
+        """Mark this router as belonging to an MTP layer."""
+        self.is_mtp = True
+        if not self.config.dsv4_mode or not self._routing_mode_initialized:
+            return
+        if self.tid2eid is None:
+            return
+
+        # DeepSeek V4 MTP checkpoint stores gate bias, not tid2eid. Reduced
+        # layer-count smoke tests can otherwise number MTP inside the hash-routed
+        # prefix and leave all MTP tid2eid entries at -1.
+        del self.tid2eid
+        self.tid2eid = None
+        self.enable_expert_bias = self.config.moe_router_enable_expert_bias
+        self._set_expert_bias_buffers(self.enable_expert_bias)
+
+    def set_layer_number(self, layer_number: int):
+        """Set the layer number and initialize DSV4 routing mode."""
+        self.layer_number = layer_number
+        if not self._routing_mode_initialized:
+            self._init_routing_mode(layer_number)
+        self._register_routing_replay_if_needed()
 
     def _maintain_float32_expert_bias(self):
         """
@@ -271,22 +367,29 @@ def is_aux_loss_enabled(self) -> bool:
         return False
 
     def _apply_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        with_padding_mask: bool = False,
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
-        tokens_per_expert = routing_map.sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                topk=self.topk,
+                with_padding_mask=with_padding_mask,
+            )
         )
-        num_tokens = routing_map.shape[0]
-        total_num_tokens = num_tokens * self.tp_cp_group.size()
 
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
-            tokens_per_expert=tokens_per_expert,
+            tokens_per_expert=global_tokens_per_expert,
             total_num_tokens=total_num_tokens,
             topk=self.topk,
             num_experts=self.config.num_moe_experts,
@@ -294,7 +397,12 @@ def _apply_aux_loss(
             fused=self.config.moe_router_fusion,
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group
+            probs,
+            aux_loss_coeff,
+            aux_loss,
+            "load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -305,6 +413,7 @@ def _apply_seq_aux_loss(
         routing_map: torch.Tensor,
         seq_length: int,
         bsz: int,
+        with_padding_mask: bool = False,
     ):
         """Apply the sequence-level auxiliary loss for the given scores and routing map.
 
@@ -318,17 +427,21 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
-        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+        routing_map = routing_map.reshape(seq_length, -1)
+
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk * bsz,
+            )
         )
 
-        total_num_tokens = seq_length * self.tp_cp_group.size()
-
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
-                tokens_per_expert=tokens_per_expert,
+                tokens_per_expert=global_tokens_per_expert,
                 total_num_tokens=total_num_tokens,
                 topk=self.topk,
                 num_experts=self.config.num_moe_experts,
@@ -338,30 +451,39 @@ def _apply_seq_aux_loss(
             / bsz
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group
+            probs,
+            seq_aux_loss_coeff,
+            aux_loss,
+            "seq_load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
     def _apply_global_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        with_padding_mask: bool = False,
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
         if global_aux_loss_coeff == 0:
             return probs
 
-        tokens_per_expert = routing_map.sum(dim=0)
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_dp_cp_group
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_dp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk,
+            )
         )
-
-        self.global_tokens_per_expert += tokens_per_expert
+        self.global_tokens_per_expert += global_tokens_per_expert
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        num_tokens = scores_for_aux_loss.shape[0]
-        total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
-
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
             tokens_per_expert=averated_tokens_per_expert,
@@ -377,6 +499,7 @@ def _apply_global_aux_loss(
             global_aux_loss,
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -387,6 +510,7 @@ def attach_and_log_load_balancing_loss(
         aux_loss: torch.Tensor,
         aux_loss_name: str,
         reduce_group: torch.distributed.ProcessGroup,
+        valid_token_count: Optional[torch.Tensor] = None,
     ):
         """Attach aux loss function to activation and add to logging."""
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
@@ -411,12 +535,13 @@ def attach_and_log_load_balancing_loss(
             # which scales both the main_loss gradient and aux_loss gradient by
             # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads function.
             # To correct this scaling, we need to scale the aux_loss by num_local_tokens here.
-            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0])
+            num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0]
+            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens)
         else:
             activation = MoEAuxLossAutoScaler.apply(activation, aux_loss)
         return activation
 
-    def apply_z_loss(self, logits):
+    def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
         """Encourages the router's logits to remain small to enhance stability.
         Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
@@ -429,8 +554,7 @@ def apply_z_loss(self, logits):
         if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled():
             # Skip Z loss calculations when using torch.no_grad() or checkpointing.
             moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()
-            z_loss = z_loss_func(logits, moe_z_loss_coeff)
-            scale_up = 1.0
+            z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask)
             if self.calculate_per_token_loss:
                 # The expected final scaling for z_loss gradients is
                 # 1/(num_micro_batches * dp_size).
@@ -439,7 +563,8 @@ def apply_z_loss(self, logits):
                 # which scales both the main_loss gradient and z_loss gradient by
                 # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads().
                 # To correct this scaling, we need to scale the z_loss by num_local_tokens here.
-                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * logits.shape[0])
+                num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0]
+                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens)
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
@@ -473,31 +598,47 @@ def apply_input_jitter(self, input: torch.Tensor):
             return input
 
     @jit_fuser
-    def _apply_expert_bias(self, routing_map: torch.Tensor):
+    def _apply_expert_bias(
+        self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
+    ):
         """
         Update expert bias and tokens_per_expert
         Prevent extra local tokens accumulation on evaluation or activation recomputation
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
+                if padding_mask is not None:
+                    routing_map = routing_map & (~padding_mask).unsqueeze(-1)
                 self.local_tokens_per_expert += routing_map.sum(dim=0)
 
-    def routing(self, logits: torch.Tensor):
+    def routing(
+        self,
+        logits: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor] = None,
+    ):
         """Top-k routing function
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
+            input_ids (torch.Tensor, optional): Input token IDs for hash routing (DSV4).
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
             routing_map (torch.Tensor): The mapping of token to experts assignment,
                 with shape [num_tokens, num_experts].
         """
+        if self.config.dsv4_mode:
+            assert self._routing_mode_initialized
+
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
 
+        if padding_mask is not None:
+            padding_mask = padding_mask.reshape(-1)
+
         # Apply Z-Loss
-        logits = self.apply_z_loss(logits)
+        logits = self.apply_z_loss(logits, padding_mask=padding_mask)
 
         # Calculate probs and routing_map for token dispatching
         if self.routing_type == "sinkhorn":
@@ -513,6 +654,9 @@ def routing(self, logits: torch.Tensor):
                 score_function=self.score_function,
                 expert_bias=self.expert_bias,
                 fused=self.config.moe_router_fusion,
+                is_mtp=self.is_mtp,
+                tid2eid=self.tid2eid,
+                input_ids=input_ids.view(-1) if self.tid2eid is not None and input_ids is not None else None,
             )
 
         # Apply token dropping to probs and routing_map.
@@ -530,18 +674,35 @@ def routing(self, logits: torch.Tensor):
         if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
-                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion
+                logits,
+                self.topk,
+                self.score_function,
+                fused=self.config.moe_router_fusion,
+                padding_mask=padding_mask,
+            )
+            probs = self._apply_aux_loss(
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
-            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)
             probs = self._apply_seq_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                seq_length,
+                bsz,
+                with_padding_mask=padding_mask is not None,
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
 
         # Optionally apply expert bias
-        self._apply_expert_bias(routing_map)
+        self._apply_expert_bias(routing_map, padding_mask=padding_mask)
 
         return probs, routing_map
 
@@ -551,15 +712,25 @@ def reset_global_aux_loss_tracker(self):
             self.global_tokens_per_expert.zero_()
             self.ga_steps.zero_()
 
-    def forward(self, input: torch.Tensor):
+    def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None):
         """
         Forward pass of the router.
 
         Args:
             input (torch.Tensor): Input tensor.
+            padding_mask (torch.Tensor, optional): Padding mask for MoE routing losses.
+            input_ids (torch.Tensor, optional): Input token IDs for DSV4 hash routing.
         """
         self._maintain_float32_expert_bias()
 
+        if self.config.freeze_e_score_correction_bias and self.enable_expert_bias:
+            if self._frozen_expert_bias_snapshot is None:
+                self._frozen_expert_bias_snapshot = self.expert_bias.clone()
+            else:
+                assert torch.equal(
+                    self.expert_bias, self._frozen_expert_bias_snapshot
+                ), "expert_bias was modified but freeze_e_score_correction_bias is enabled"
+
         # Apply input jitter
         input = self.apply_input_jitter(input)
         logits = self.gating(input)
@@ -568,7 +739,7 @@ def forward(self, input: torch.Tensor):
             # Apply force load balancing with random logits for benchmark
             logits = apply_random_logits(logits)
 
-        probs, routing_map = self.routing(logits)
+        probs, routing_map = self.routing(logits, padding_mask=padding_mask, input_ids=input_ids)
 
         return probs, routing_map
 
diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py
index ab075d94e52..0a821818472 100644
--- a/megatron/core/transformer/moe/shared_experts.py
+++ b/megatron/core/transformer/moe/shared_experts.py
@@ -47,6 +47,9 @@ def __init__(
         assert config.add_bias_linear == False, "bias is not supported in the shared experts, "
         "please set '--disable-bias-linear' instead."
 
+        if not config.activation_func_clamp_shared_expert:
+            config.activation_func_clamp_value = None
+
         config.ffn_hidden_size = config.moe_shared_expert_intermediate_size
         # TODO(Hepteract): pass pg_collection to MLP after refactoring MLP
         super().__init__(config=config, submodules=submodules, tp_group=pg_collection.tp)
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index f33f6f05e14..21ac8782efe 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -675,6 +675,8 @@ def __init__(
             vp_stage=vp_stage,
             layer_number=self.layer_number + diff_transformer_layer_offset,
         )
+        if hasattr(self.transformer_layer.mlp, 'set_is_mtp'):
+            self.transformer_layer.mlp.set_is_mtp()
 
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index 023db1fe75a..580d3c4d904 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -387,6 +387,19 @@ def build_layer(layer_spec, layer_number):
         else:
             self.final_layernorm = None  # Either this or nn.Identity
 
+        if self.config.dsv4_mode:
+            from sirl.plugins.models.deepseek_v4.ops.hyper_connection import (
+                DeepSeekV4HyperConnectionUtil,
+                HCHeadParams,
+            )
+
+            self.hc_util = DeepSeekV4HyperConnectionUtil(self.config)
+            if self.has_final_layernorm_in_this_stage():
+                self.hc_head_params = HCHeadParams(self.config)
+                for param in self.hc_head_params.parameters():
+                    # The DSV4 HC head uses these parameters under torch.no_grad().
+                    param.requires_grad_(False)
+
     def has_final_layernorm_in_this_stage(self):
         """
         Check if this vpp stage contains the final layernorm.
@@ -419,6 +432,15 @@ def has_final_layernorm_in_this_stage(self):
     def _get_layer(self, layer_number: int):
         return self.layers[layer_number]
 
+    def _prepare_router_input_ids(self, input_ids: Tensor) -> Tensor:
+        assert input_ids.dim() == 2
+        input_ids = input_ids.transpose(0, 1).contiguous()
+        if self.config.sequence_parallel:
+            input_ids = tensor_parallel.scatter_to_sequence_parallel_region(
+                input_ids, group=self.tp_group
+            )
+        return input_ids
+
     def _checkpointed_forward(
         self,
         hidden_states: Tensor,
@@ -429,12 +451,20 @@ def _checkpointed_forward(
         attention_bias: Tensor,
         packed_seq_params: PackedSeqParams,
         use_inner_quantization_context: bool,
+        padding_mask: Optional[Tensor] = None,
+        input_ids: Optional[Tensor] = None,
     ):
         """Forward method with activation checkpointing."""
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+                hidden_states,
+                attention_mask,
+                context,
+                context_mask,
+                rotary_pos_emb,
+                padding_mask=None,
+                input_ids=None,
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
@@ -465,6 +495,8 @@ def custom_forward(
                             attention_bias=attention_bias,
                             inference_context=None,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
+                            input_ids=input_ids,
                         )
                 return hidden_states, context
 
@@ -484,6 +516,8 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
+                    input_ids,
                 )
             else:
                 return tensor_parallel.checkpoint(
@@ -494,6 +528,8 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
+                    input_ids,
                 )
 
         if self.config.recompute_method == 'uniform':
@@ -527,7 +563,13 @@ def checkpoint_handler(forward_func):
                     hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1))
                 else:
                     hidden_states, context = custom(layer_idx, layer_idx + 1)(
-                        hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+                        hidden_states,
+                        attention_mask,
+                        context,
+                        context_mask,
+                        rotary_pos_emb,
+                        padding_mask,
+                        input_ids,
                     )
         else:
             raise ValueError("Invalid activation recompute method.")
@@ -599,6 +641,8 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
+        input_ids: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         dynamic_inference_decode_only: Optional[bool] = None,
@@ -669,6 +713,13 @@ def forward(
         #   is called here to be future-proof and corner-case-proof.
         hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
 
+        if self.config.dsv4_mode and self.pre_process:
+            hidden_states = self.hc_util.block_expand(hidden_states)
+        if self.config.dsv4_mode and input_ids is not None:
+            input_ids = self._prepare_router_input_ids(input_ids)
+            assert input_ids.size(0) == hidden_states.size(0)
+            assert input_ids.size(1) == hidden_states.size(1)
+
         if self.config.sequence_parallel:
             rng_context = tensor_parallel.get_cuda_rng_tracker().fork()
         else:
@@ -708,6 +759,8 @@ def forward(
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,
                     use_inner_quantization_context=use_inner_quantization_context,
+                    padding_mask=padding_mask,
+                    input_ids=input_ids,
                 )
             else:
                 for l_no, layer in enumerate(self.layers):
@@ -745,6 +798,8 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
+                            padding_mask=padding_mask,
+                            input_ids=input_ids,
                         )
 
                     if (
@@ -754,6 +809,14 @@ def forward(
                     ):
                         hidden_states = self.group_prefetch_offload_commit_async(hidden_states)
 
+        if self.config.dsv4_mode and self.post_process and hasattr(self, "hc_head_params"):
+            hidden_states = self.hc_util.block_head(
+                hidden_states,
+                self.hc_head_params.hc_head_fn,
+                self.hc_head_params.hc_head_scale,
+                self.hc_head_params.hc_head_base,
+            )
+
         # Final layer norm.
         if self.final_layernorm is not None:
             hidden_states = self.final_layernorm(hidden_states)
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index a0aa109b586..7f8853b2bfc 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -171,6 +171,9 @@ class TransformerConfig(ModelParallelConfig):
     """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func
     is quick_gelu."""
 
+    activation_func_clamp_shared_expert: bool = True
+    """If False, shared experts ignore activation_func_clamp_value while routed experts keep it."""
+
     num_moe_experts: Optional[int] = None
     """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None
     for no MoE."""
@@ -241,7 +244,7 @@ class TransformerConfig(ModelParallelConfig):
     # attention variant
     ####################
     experimental_attention_variant: Optional[str] = None
-    """Type of attention variant to use. Currently support gated_delta_net and dsa."""
+    """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4."""
 
     ####################
     # attention variant: gated_delta_net
@@ -291,6 +294,48 @@ class TransformerConfig(ModelParallelConfig):
     """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the
     top-k indices."""
 
+    ####################
+    # attention variant: DeepSeek-V4
+    ####################
+    dsv4_mode: bool = False
+    """Enable DeepSeek-V4 hyper-connection, routing, and sparse-attention shape semantics."""
+
+    dsv4_hc_mult: Optional[int] = None
+    """DeepSeek-V4 Hyper-Connection stream multiplier."""
+
+    dsv4_hc_sinkhorn_iters: int = 20
+    """DeepSeek-V4 Hyper-Connection Sinkhorn iterations."""
+
+    dsv4_hc_eps: float = 1e-6
+    """DeepSeek-V4 Hyper-Connection epsilon."""
+
+    dsv4_compress_ratios: Optional[List[int]] = None
+    """DeepSeek-V4 per-layer compression ratios."""
+
+    dsv4_compress_rope_theta: float = 40000.0
+    """DeepSeek-V4 compressor RoPE theta."""
+
+    dsv4_o_groups: Optional[int] = None
+    """DeepSeek-V4 grouped output projection group count."""
+
+    dsv4_o_lora_rank: Optional[int] = None
+    """DeepSeek-V4 output projection LoRA rank."""
+
+    dsv4_n_hash_layers: int = 0
+    """Number of initial DeepSeek-V4 layers using hash routing."""
+
+    dsv4_window_size: int = 4096
+    """DeepSeek-V4 local attention window size."""
+
+    freeze_e_score_correction_bias: bool = False
+    """Freeze MoE expert score correction bias during training."""
+
+    moe_router_freeze_gate: bool = False
+    """Freeze MoE router gate weights during training."""
+
+    vocab_size: Optional[int] = None
+    """Vocabulary size used to initialize DeepSeek-V4 hash-routing tables."""
+
     ####################
     # initialization
     ####################
@@ -687,6 +732,9 @@ class TransformerConfig(ModelParallelConfig):
     moe_apply_probs_on_input: bool = False
     """Apply probs on input of experts instead of applying after activation and glu."""
 
+    moe_latent_size: Optional[int] = None
+    """Latent projection dimension for MoE. If None, MoE latent projections are not used."""
+
     ##################
     # Context Parallel
     ##################
@@ -897,6 +945,9 @@ def __post_init__(self):
             self.experimental_attention_variant = self.linear_attention_type
             self.linear_attention_type = None
 
+        if self.experimental_attention_variant == "dsv4":
+            self.dsv4_mode = True
+
         if self.experimental_attention_variant in ["gated_delta_net"]:
             assert (
                 self.linear_attention_freq is not None
@@ -1541,10 +1592,13 @@ def __post_init__(self):
                 self.expert_tensor_parallel_size == 1
             ), "Bias in Moe is only supported when ETP==1"
 
-        if self.moe_router_enable_expert_bias and self.moe_router_score_function != "sigmoid":
+        if self.moe_router_enable_expert_bias and self.moe_router_score_function not in (
+            "sigmoid",
+            "sqrtsoftplus",
+        ):
             raise ValueError(
-                "Expert bias for aux-loss-free routing only supports sigmoid score function."
-                "Please set --moe-router-score-function sigmoid for sigmoid score function."
+                "Expert bias for aux-loss-free routing only supports sigmoid or sqrtsoftplus score function. "
+                "Please set --moe-router-score-function to sigmoid or sqrtsoftplus."
             )
 
         if self.num_moe_experts and self.fp8:
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 5a42001b9a8..4ddbfe4540c 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import functools
 import logging
 import warnings
 from abc import ABC
@@ -308,6 +309,10 @@ def __init__(
             layer_number=self.layer_number,
             **attention_optional_kwargs,
         )
+        if self.config.dsv4_mode and getattr(self.self_attention, "indexer", None) is not None:
+            # DSV4 indexer parameters produce integer top-k indices, so no grad flows to them.
+            for param in self.self_attention.indexer.parameters():
+                param.requires_grad_(False)
 
         # [Module 3: BiasDropoutFusion]
         self.self_attn_bda = build_module(submodules.self_attn_bda)
@@ -391,6 +396,28 @@ def __init__(
             eps=self.config.layernorm_epsilon
         )
 
+        if self.config.dsv4_mode:
+            hc_mult = self.config.dsv4_hc_mult
+            hc_dim = hc_mult * self.config.hidden_size
+            mix_size = (2 + hc_mult) * hc_mult
+            self.hc_attn_fn = torch.nn.Parameter(torch.empty(mix_size, hc_dim, dtype=torch.float32))
+            self.hc_attn_base = torch.nn.Parameter(torch.empty(mix_size, dtype=torch.float32))
+            self.hc_attn_scale = torch.nn.Parameter(torch.empty(3, dtype=torch.float32))
+            self.hc_ffn_fn = torch.nn.Parameter(torch.empty(mix_size, hc_dim, dtype=torch.float32))
+            self.hc_ffn_base = torch.nn.Parameter(torch.empty(mix_size, dtype=torch.float32))
+            self.hc_ffn_scale = torch.nn.Parameter(torch.empty(3, dtype=torch.float32))
+            for param in (
+                self.hc_attn_fn,
+                self.hc_attn_base,
+                self.hc_attn_scale,
+                self.hc_ffn_fn,
+                self.hc_ffn_base,
+                self.hc_ffn_scale,
+            ):
+                param._keep_fp32 = True
+                # The DSV4 HC mixer uses these parameters under torch.no_grad().
+                param.requires_grad_(False)
+
         self.recompute_input_layernorm = False
         self.recompute_pre_mlp_layernorm = False
         self.recompute_mlp = False
@@ -472,8 +499,14 @@ def forward(self, *args, **kwargs):
         # this is only used to uniquely identify decode and non-decode cuda graph
         # runners in the cuda graph manager
         kwargs.pop("dynamic_inference_decode_only", None)
+        input_ids = kwargs.pop("input_ids", None)
         hidden_states, context = self._forward_attention(*args, **kwargs)
-        output = self._forward_mlp(hidden_states, kwargs.get("inference_context", None))
+        output = self._forward_mlp(
+            hidden_states,
+            kwargs.get("inference_context", None),
+            padding_mask=kwargs.get("padding_mask", None),
+            input_ids=input_ids,
+        )
         return output, context
 
     def _forward_attention(
@@ -490,6 +523,7 @@ def _forward_attention(
         inference_context: Optional[Any] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[Any] = None,
     ):
@@ -531,6 +565,19 @@ def _forward_attention(
         # Residual connection.
         residual = hidden_states
 
+        if self.config.dsv4_mode:
+            from sirl.plugins.models.deepseek_v4.ops.hyper_connection import (
+                DeepSeekV4HyperConnectionUtil,
+            )
+
+            hc_util = DeepSeekV4HyperConnectionUtil(self.config)
+            hidden_states, hc_attn_post, hc_attn_comb = hc_util.layer_pre(
+                hidden_states,
+                self.hc_attn_fn,
+                self.hc_attn_scale,
+                self.hc_attn_base,
+            )
+
         if self.offload_attn_norm:
             hidden_states = fine_grained_offloading_group_start(hidden_states, name="attn_norm")
         # Optional Input Layer norm
@@ -567,17 +614,22 @@ def _forward_attention(
                 attention_output_with_bias[0]
             )
 
-        attention_output, attention_output_bias = attention_output_with_bias
-        attention_output = self.post_self_attn_layernorm(attention_output)
-        attention_output_with_bias = (attention_output, attention_output_bias)
-
-        # TODO: could we move `bias_dropout_add_exec_handler` itself
-        # inside the module provided in the `bias_dropout_add_spec` module?
         nvtx_range_push(suffix="self_attn_bda")
-        with self.bias_dropout_add_exec_handler():
-            hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)(
-                attention_output_with_bias, residual, self.hidden_dropout
+        if self.config.dsv4_mode:
+            hidden_states = hc_util.layer_post(
+                attention_output_with_bias, residual, hc_attn_post, hc_attn_comb
             )
+        else:
+            attention_output, attention_output_bias = attention_output_with_bias
+            attention_output = self.post_self_attn_layernorm(attention_output)
+            attention_output_with_bias = (attention_output, attention_output_bias)
+
+            # TODO: could we move `bias_dropout_add_exec_handler` itself
+            # inside the module provided in the `bias_dropout_add_spec` module?
+            with self.bias_dropout_add_exec_handler():
+                hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)(
+                    attention_output_with_bias, residual, self.hidden_dropout
+                )
         nvtx_range_pop(suffix="self_attn_bda")
 
         if self.offload_attn_norm:
@@ -611,7 +663,7 @@ def _forward_attention(
 
         return hidden_states, context
 
-    def _forward_mlp(self, hidden_states, inference_context=None):
+    def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None, input_ids=None):
         """
         Perform a forward pass through the feed-forward layer.
 
@@ -630,6 +682,19 @@ def _forward_mlp(self, hidden_states, inference_context=None):
         # Residual connection.
         residual = hidden_states
 
+        if self.config.dsv4_mode:
+            from sirl.plugins.models.deepseek_v4.ops.hyper_connection import (
+                DeepSeekV4HyperConnectionUtil,
+            )
+
+            hc_util = DeepSeekV4HyperConnectionUtil(self.config)
+            hidden_states, hc_ffn_post, hc_ffn_comb = hc_util.layer_pre(
+                hidden_states,
+                self.hc_ffn_fn,
+                self.hc_ffn_scale,
+                self.hc_ffn_base,
+            )
+
         if self.offload_mlp_norm:
             hidden_states = fine_grained_offloading_group_start(hidden_states, name="mlp_norm")
         # Optional Layer norm post the cross-attention.
@@ -662,7 +727,14 @@ def _forward_mlp(self, hidden_states, inference_context=None):
             assert (
                 not self.recompute_pre_mlp_layernorm
             ), "Recomputation is not supported for CUDA graph."
-            cudagraph_outputs = self.mlp(pre_mlp_layernorm_output)
+            cudagraph_outputs = self.mlp(
+                pre_mlp_layernorm_output,
+                **(
+                    {"padding_mask": padding_mask, "input_ids": input_ids}
+                    if self.is_moe_layer
+                    else {}
+                ),
+            )
             nvtx_range_pop(suffix="mlp")
             return cudagraph_outputs + [residual]
         elif self.recompute_mlp:
@@ -671,7 +743,9 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 from megatron.core.extensions.transformer_engine import te_checkpoint
 
                 mlp_output_with_bias = te_checkpoint(
-                    self.mlp,
+                    functools.partial(self.mlp, padding_mask=padding_mask, input_ids=input_ids)
+                    if self.is_moe_layer
+                    else self.mlp,
                     False,
                     tensor_parallel.random.get_cuda_rng_tracker,
                     self.pg_collection.tp,
@@ -679,15 +753,44 @@ def _forward_mlp(self, hidden_states, inference_context=None):
                 )
             else:
                 mlp_output_with_bias = tensor_parallel.checkpoint(
-                    self.mlp, False, pre_mlp_layernorm_output
+                    functools.partial(self.mlp, padding_mask=padding_mask, input_ids=input_ids)
+                    if self.is_moe_layer
+                    else self.mlp,
+                    False,
+                    pre_mlp_layernorm_output,
                 )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
             num_chunks = min(self.config.mlp_chunks_for_prefill, pre_mlp_layernorm_output.shape[0])
             chunks = pre_mlp_layernorm_output.chunk(num_chunks, dim=0)
+            input_id_chunks = (
+                input_ids.chunk(num_chunks, dim=0)
+                if self.is_moe_layer and input_ids is not None
+                else [None] * len(chunks)
+            )
+            padding_mask_chunks = (
+                padding_mask.chunk(num_chunks, dim=1)
+                if self.is_moe_layer and padding_mask is not None
+                else [None] * len(chunks)
+            )
 
             # Compute outputs for each chunk
-            outputs = [self.mlp(chunk) for chunk in chunks]
+            outputs = [
+                self.mlp(
+                    chunk,
+                    **(
+                        {"padding_mask": padding_mask_chunk, "input_ids": input_id_chunk}
+                        if self.is_moe_layer
+                        else {}
+                    ),
+                )
+                for chunk, padding_mask_chunk, input_id_chunk in zip(
+                    chunks,
+                    padding_mask_chunks,
+                    input_id_chunks,
+                    strict=True,
+                )
+            ]
 
             # Aggregate chunk outputs
             mlp_output = torch.cat([out for out, _ in outputs], dim=0)
@@ -695,7 +798,14 @@ def _forward_mlp(self, hidden_states, inference_context=None):
             bias_output = torch.stack(bias_chunks, dim=0).sum(dim=0) if bias_chunks else None
             mlp_output_with_bias = (mlp_output, bias_output)
         else:
-            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+            mlp_output_with_bias = self.mlp(
+                pre_mlp_layernorm_output,
+                **(
+                    {"padding_mask": padding_mask, "input_ids": input_ids}
+                    if self.is_moe_layer
+                    else {}
+                ),
+            )
 
         mlp_output, mlp_output_bias = mlp_output_with_bias
         mlp_output = self.post_mlp_layernorm(mlp_output)
@@ -709,9 +819,14 @@ def _forward_mlp(self, hidden_states, inference_context=None):
             )
         nvtx_range_pop(suffix="mlp")
 
-        return self._forward_post_mlp(mlp_output_with_bias, residual)
+        return self._forward_post_mlp(
+            mlp_output_with_bias,
+            residual,
+            hc_ffn_post=hc_ffn_post if self.config.dsv4_mode else None,
+            hc_ffn_comb=hc_ffn_comb if self.config.dsv4_mode else None,
+        )
 
-    def _forward_post_mlp(self, mlp_output_with_bias, residual):
+    def _forward_post_mlp(self, mlp_output_with_bias, residual, *, hc_ffn_post=None, hc_ffn_comb=None):
         """
         Perform operations after the MLP computation.
 
@@ -730,10 +845,20 @@ def _forward_post_mlp(self, mlp_output_with_bias, residual):
         # TODO: could we move `bias_dropout_add_exec_handler` itself
         # inside the module provided in the `bias_dropout_add_spec` module?
         nvtx_range_push(suffix="mlp_bda")
-        with self.bias_dropout_add_exec_handler():
-            hidden_states = self.mlp_bda(self.training, self.config.bias_dropout_fusion)(
-                mlp_output_with_bias, residual, self.hidden_dropout
+        if self.config.dsv4_mode:
+            from sirl.plugins.models.deepseek_v4.ops.hyper_connection import (
+                DeepSeekV4HyperConnectionUtil,
+            )
+
+            hc_util = DeepSeekV4HyperConnectionUtil(self.config)
+            hidden_states = hc_util.layer_post(
+                mlp_output_with_bias, residual, hc_ffn_post, hc_ffn_comb
             )
+        else:
+            with self.bias_dropout_add_exec_handler():
+                hidden_states = self.mlp_bda(self.training, self.config.bias_dropout_fusion)(
+                    mlp_output_with_bias, residual, self.hidden_dropout
+                )
         nvtx_range_pop(suffix="mlp_bda")
         if self.offload_mlp_norm:
             (hidden_states,) = fine_grained_offloading_group_commit(
@@ -853,7 +978,11 @@ def _te_cuda_graph_capture(self, *args, **kwargs):
                 )
             )
         ):
-            hidden_states = self._forward_mlp(hidden_states)
+            hidden_states = self._forward_mlp(
+                hidden_states,
+                padding_mask=kwargs.get("padding_mask", None),
+                input_ids=kwargs.get("input_ids", None),
+            )
         if not isinstance(hidden_states, list) and not isinstance(hidden_states, tuple):
             cuda_graph_outputs = [hidden_states]
         else:
@@ -945,7 +1074,11 @@ def _te_cuda_graph_replay(self, *args, **kwargs):
             output = self._forward_post_mlp(mlp_output_with_bias, mlp_residual)
         else:
             # CUDA Graph does not capture the MLP/MoE part at all.
-            output = self._forward_mlp(*cuda_graph_output)
+            output = self._forward_mlp(
+                *cuda_graph_output,
+                padding_mask=kwargs.get("padding_mask", None),
+                input_ids=kwargs.get("input_ids", None),
+            )
         return output, context
 
     def _get_te_cuda_graph_replay_args(self, *args, **kwargs):
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 83736acdc4b..1432da6e30e 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1786,6 +1786,9 @@ def _add_network_size_args(parser):
     group.add_argument('--activation-func-clamp-value', type=float, default=None,
                        help='Clamp the output of the linear_fc1 in the activation function. Only used when '
                             'activation_func is quick_gelu.')
+    group.add_argument('--no-activation-func-clamp-shared-expert', action='store_false',
+                       dest='activation_func_clamp_shared_expert',
+                       help='Do not apply activation_func_clamp_value inside shared expert MLPs.')
     group.add_argument('--glu-linear-offset', type=float, default=0.0,
                        help='Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). '
                             'Only used when gated_linear_unit is True')
@@ -3238,9 +3241,9 @@ def _add_moe_args(parser):
     group.add_argument('--moe-router-fusion', action='store_true',
                        help='Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in TransformerEngine 2.7.0 and above.')
     group.add_argument('--moe-router-score-function', type=str,
-                       choices=['softmax', 'sigmoid'],
+                       choices=['softmax', 'sigmoid', 'sqrtsoftplus'],
                        default='softmax',
-                       help='Score function for MoE TopK routing. Can be "softmax" or "sigmoid".')
+                       help='Score function for MoE TopK routing. Can be "softmax", "sigmoid", or "sqrtsoftplus".')
     group.add_argument('--moe-router-topk', type=int, default=2,
                        help='Number of experts to route to for each token. The default is 2.')
     group.add_argument('--moe-router-pre-softmax', action='store_true',
@@ -3256,6 +3259,10 @@ def _add_moe_args(parser):
                        help='TopK routing with dynamic expert bias in the aux-loss-free load balancing strategy. '
                        'The routing decision is based on the sum of the routing scores and the expert bias. '
                        'See https://arxiv.org/abs/2408.15664 for details.')
+    group.add_argument('--moe-router-freeze-gate', action='store_true',
+                       help='Freeze MoE router gate weights during training.')
+    group.add_argument('--freeze-e-score-correction-bias', action='store_true',
+                       help='Freeze MoE expert score correction bias during training.')
     group.add_argument('--moe-router-bias-update-rate', type=float, default=1e-3,
                        help='Expert bias update rate in the aux-loss-free load balancing strategy. '
                        'The expert bias is updated based on the number of assigned tokens to each expert in a global batch, '
@@ -3337,6 +3344,12 @@ def _add_mla_args(parser):
                        help="Dimension of the head in the V projection.")
     group.add_argument('--rotary-scaling-factor', type=float, default=1.0,
                        help="Rotary scaling factor for the rotary embeddings.")
+    group.add_argument('--original-max-position-embeddings', type=int, default=4096,
+                       help='Original maximum position embeddings for YaRN RoPE.')
+    group.add_argument('--beta-fast', type=float, default=32,
+                       help='YaRN beta fast.')
+    group.add_argument('--beta-slow', type=float, default=1,
+                       help='YaRN beta slow.')
     group.add_argument('--mscale', type=float, default=1.0,
                        help="Mscale for YaRN RoPE in multi-latent attention.")
     group.add_argument('--mscale-all-dim', type=float, default=0.0,
@@ -3348,8 +3361,8 @@ def _add_mla_args(parser):
 
 def _add_experimental_attention_variant_args(parser):
     group = parser.add_argument_group(title="experimental_attention_variant")
-    group.add_argument('--experimental-attention-variant', default=None, choices=['gated_delta_net', 'dsa'], type=str,
-                       help='Type of attention variant to use. Currently support gated_delta_net and dsa.')
+    group.add_argument('--experimental-attention-variant', default=None, choices=['gated_delta_net', 'dsa', 'dsv4'], type=str,
+                       help='Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4.')
 
     # Linear attention
     group.add_argument('--linear-attention-type', default=None, choices=['gated_delta_net'], type=str,
@@ -3386,6 +3399,26 @@ def _add_experimental_attention_variant_args(parser):
     group.add_argument('--dsa-indexer-use-sparse-loss', action='store_true',
                        help='Use sparse indexer loss. If set, the indexer loss will be computed using the top-k indices.')
 
+    # DeepSeek-V4
+    group.add_argument('--dsv4-hc-mult', default=None, type=int,
+                       help='DeepSeek-V4 Hyper-Connection stream multiplier.')
+    group.add_argument('--dsv4-hc-sinkhorn-iters', default=20, type=int,
+                       help='DeepSeek-V4 Hyper-Connection Sinkhorn iterations.')
+    group.add_argument('--dsv4-hc-eps', default=1e-6, type=float,
+                       help='DeepSeek-V4 Hyper-Connection epsilon.')
+    group.add_argument('--dsv4-compress-ratios', nargs='+', default=None, type=int,
+                       help='DeepSeek-V4 per-layer compression ratios.')
+    group.add_argument('--dsv4-compress-rope-theta', default=40000.0, type=float,
+                       help='DeepSeek-V4 compressor RoPE theta.')
+    group.add_argument('--dsv4-o-groups', default=None, type=int,
+                       help='DeepSeek-V4 grouped output projection group count.')
+    group.add_argument('--dsv4-o-lora-rank', default=None, type=int,
+                       help='DeepSeek-V4 output projection LoRA rank.')
+    group.add_argument('--dsv4-n-hash-layers', default=0, type=int,
+                       help='Number of initial DeepSeek-V4 layers using hash routing.')
+    group.add_argument('--dsv4-window-size', default=4096, type=int,
+                       help='DeepSeek-V4 local attention window size.')
+
     return parser
 
 def _add_heterogeneous_args(parser):