diff --git a/megatron/core/models/backends.py b/megatron/core/models/backends.py
index a270161ddd6..c543a49e266 100644
--- a/megatron/core/models/backends.py
+++ b/megatron/core/models/backends.py
@@ -10,6 +10,7 @@
     TEColumnParallelGroupedLinear,
     TERowParallelGroupedLinear,
 )
+from megatron.core.post_training.modelopt.layers import Linear
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from megatron.core.transformer.dot_product_attention import DotProductAttention
 from megatron.core.transformer.mlp import MLPSubmodules, TEActivationFunctionBuilder
@@ -99,6 +100,15 @@ def activation_func(self) -> TEActivationFunctionBuilder | None:
 class LocalSpecProvider(BackendSpecProvider):
     """A protocol for providing Local submodules used in Spec building."""
 
+    def linear(self) -> type:
+        """TP-replicated local Linear (modelopt Linear, not TELinear).
+
+        DSA indexer / MLA down-projections call backend.linear(). TESpecProvider
+        still returns TELinear; this method is the TE-off counterpart so a
+        LocalSpecProvider DSA spec does not re-enter Transformer Engine.
+        """
+        return Linear
+
     def column_parallel_linear(self) -> type:
         """Which column parallel linear module the backend uses"""
         return ColumnParallelLinear
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 a76fe6e3a23..a8cc4886717 100644
--- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
+++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
@@ -20,6 +20,7 @@
 )
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.torch_norm import WrappedTorchNorm
 from megatron.core.transformer.transformer_block import (
     TransformerBlockSubmodules,
     get_num_layers_to_build,
@@ -57,6 +58,15 @@
 ##########
 
 
+def _get_standalone_norm(
+    config: TransformerConfig, backend: BackendSpecProvider, *, for_qk=False
+):
+    rms_norm = config.normalization == "RMSNorm"
+    if rms_norm and config.norm_accuracy_compatible:
+        return WrappedTorchNorm
+    return backend.layer_norm(rms_norm=rms_norm, for_qk=for_qk)
+
+
 def get_gated_delta_net_module_spec(
     config: TransformerConfig, backend: BackendSpecProvider = None
 ) -> ModuleSpec:
@@ -65,12 +75,11 @@ def get_gated_delta_net_module_spec(
     if backend is None:
         backend = _get_backend_spec_provider(config=config)
 
-    rms_norm = config.normalization == "RMSNorm"
     attention = ModuleSpec(
         module=GatedDeltaNet,
         submodules=GatedDeltaNetSubmodules(
             in_proj=backend.column_parallel_layer_norm_linear(),
-            out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
+            out_norm=_get_standalone_norm(config, backend),
             out_proj=backend.row_parallel_linear(),
         ),
         metainfo={"fuse_input_layernorm": True},
@@ -82,7 +91,9 @@ def get_dsa_module_spec_for_backend(
     config: TransformerConfig, backend: BackendSpecProvider = None
 ) -> ModuleSpec:
     """Helper function to get module spec for Sparse Attention."""
-    assert config.multi_latent_attention, "Currently only MLA supports sparse attention."
+    assert config.multi_latent_attention, (
+        "Currently only MLA supports sparse attention."
+    )
     assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA."
 
     # Because TransformerEngine does not support sparse attention yet, we use local
@@ -102,12 +113,12 @@ def get_dsa_module_spec_for_backend(
         ),
     )
 
-    # Adjust for RMS norm.
-    rms_norm = config.normalization == "RMSNorm"
     # DSA indexer requires normalized q as input, so here we cannot fuse qk layernorm
     # with linear projection and have to use unfused qk layernorm.
     qk_norm = (
-        backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp
+        _get_standalone_norm(config, backend, for_qk=True)
+        if config.qk_layernorm
+        else IdentityOp
     )
 
     attention = ModuleSpec(
@@ -203,7 +214,9 @@ def get_transformer_layer_with_experimental_attention_variant_spec(
         experimental_attention_spec = None
 
     if 0 in experimental_attention_pattern:
-        standard_attention_spec = _get_self_attention_module_spec(config=config, backend=backend)
+        standard_attention_spec = _get_self_attention_module_spec(
+            config=config, backend=backend
+        )
     else:
         standard_attention_spec = None
 
@@ -228,7 +241,6 @@ def get_transformer_layer_with_experimental_attention_variant_spec(
         dense_mlp_layer_spec, fuse_layernorm_pre_dense = None, False
 
     # Get GPT decoder block layer specs
-    rms_norm = config.normalization == "RMSNorm"
     layer_specs = []
     for layer_number in range(config.num_layers):
         attention = (
@@ -236,7 +248,11 @@ def get_transformer_layer_with_experimental_attention_variant_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
+        mlp = (
+            moe_layer_spec
+            if moe_layer_pattern[layer_number] == 1
+            else dense_mlp_layer_spec
+        )
         fuse_pre_mlp_layernorm = (
             fuse_layernorm_pre_moe
             if moe_layer_pattern[layer_number] == 1
@@ -245,12 +261,12 @@ def get_transformer_layer_with_experimental_attention_variant_spec(
         input_layernorm = (
             IdentityOp
             if attention.metainfo["fuse_input_layernorm"]
-            else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+            else _get_standalone_norm(config, backend)
         )
         pre_mlp_layernorm = (
             IdentityOp
             if fuse_pre_mlp_layernorm
-            else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+            else _get_standalone_norm(config, backend)
         )
 
         layer_specs.append(
@@ -271,7 +287,9 @@ def get_transformer_layer_with_experimental_attention_variant_spec(
 
 
 def get_transformer_block_with_experimental_attention_variant_spec(
-    config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None
+    config: TransformerConfig,
+    vp_stage: Optional[int] = None,
+    pp_rank: Optional[int] = None,
 ) -> TransformerBlockSubmodules:
     """Build transformer block spec with experimental attention variants (e.g., linear attention).
 
@@ -309,17 +327,20 @@ def get_transformer_block_with_experimental_attention_variant_spec(
             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)
+        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)
 
     _validate_dsa_index_share_pipeline_split(config, local_layer_ids)
     layer_specs = [layer_specs[layer_id] for layer_id in local_layer_ids]
 
     # Get GPT decoder block spec
-    rms_norm = config.normalization == "RMSNorm"
     gpt_decoder_block_spec = TransformerBlockSubmodules(
-        layer_specs=layer_specs, layer_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False)
+        layer_specs=layer_specs, layer_norm=_get_standalone_norm(config, backend)
     )
 
     return gpt_decoder_block_spec
@@ -336,7 +357,9 @@ def is_linear_attention_variant(experimental_attention_variant: Optional[str]) -
     return experimental_attention_variant in linear_attention_variants
 
 
-def _validate_dsa_index_share_pipeline_split(config: TransformerConfig, local_layer_ids) -> None:
+def _validate_dsa_index_share_pipeline_split(
+    config: TransformerConfig, local_layer_ids
+) -> None:
     """Ensure DSA top-k sharing does not require top-k indices from another PP stage."""
     if (
         config.experimental_attention_variant != "dsa"
@@ -351,12 +374,16 @@ def _validate_dsa_index_share_pipeline_split(config: TransformerConfig, local_la
     for position, layer_id in enumerate(local_layer_ids):
         layer_number = layer_id + 1
         if not is_dsa_skip_topk_layer(
-            layer_number, config.dsa_indexer_skip_topk_offset, config.dsa_indexer_topk_freq
+            layer_number,
+            config.dsa_indexer_skip_topk_offset,
+            config.dsa_indexer_topk_freq,
         ):
             continue
 
         source_layer_number = source_dsa_compute_layer(
-            layer_number, config.dsa_indexer_skip_topk_offset, config.dsa_indexer_topk_freq
+            layer_number,
+            config.dsa_indexer_skip_topk_offset,
+            config.dsa_indexer_topk_freq,
         )
         source_layer_id = source_layer_number - 1
         if (
@@ -383,7 +410,8 @@ def get_moe_layer_pattern(config: TransformerConfig) -> List[int]:
     if isinstance(config.moe_layer_freq, int):
         # [1,0,0,...,0,1,0,0,...,0,...]
         moe_layer_pattern = [
-            1 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.num_layers)
+            1 if (i % config.moe_layer_freq == 0) else 0
+            for i in range(config.num_layers)
         ]
     elif isinstance(config.moe_layer_freq, list):
         moe_layer_pattern = config.moe_layer_freq
@@ -471,7 +499,9 @@ def _get_self_attention_module_spec(
     if backend is None:
         backend = _get_backend_spec_provider(config=config)
 
-    from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+    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,
@@ -532,7 +562,9 @@ def _get_moe_module_spec(
     if backend is None:
         backend = _get_backend_spec_provider(config=config)
 
-    from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend
+    from megatron.core.models.gpt.moe_module_specs import (
+        get_moe_module_spec_for_backend,
+    )
 
     return (
         get_moe_module_spec_for_backend(
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index 984840b3a87..63a1aa51b02 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -773,7 +773,7 @@ def get_gpt_mtp_block_spec_for_backend(
         raise ValueError(f"Invalid spec: {spec}")
 
     mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        mtp_model_layer_spec=transformer_layer_spec, backend=backend
+        mtp_model_layer_spec=transformer_layer_spec, backend=backend, config=config
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
     if config.mtp_use_repeated_layer:
diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 32a61cf7efc..c7406191dc2 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -556,6 +556,9 @@ def _get_megatron_optimizer_based_on_param_groups(
             # on source of optimizer (Torch or TE/Apex)
             if USING_PYTORCH_OPTIMIZER:
                 adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam
+            elif config.native_unfused_adamw:
+                adam_cls = torch.optim.AdamW if config.decoupled_weight_decay else torch.optim.Adam
+                kwargs.update({"foreach": False, "fused": False})
             else:
                 kwargs["adam_w_mode"] = config.decoupled_weight_decay
                 adam_cls = Adam
diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py
index 24f9a032c47..a11f5c8cc3b 100644
--- a/megatron/core/optimizer/optimizer_config.py
+++ b/megatron/core/optimizer/optimizer_config.py
@@ -247,6 +247,9 @@ class OptimizerConfig:
     adam_eps: float = 1e-08
     """Term added to the denominator to improve numerical stability in Adam optimizer."""
 
+    native_unfused_adamw: bool = False
+    """Use torch.optim.AdamW with foreach=False and fused=False instead of TE/Apex Adam."""
+
     decoupled_weight_decay: bool = True
     """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false,
     original Adam update rule will be used. Defaults to True.
diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index dde238635c2..c7f3a212558 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -64,6 +64,7 @@ def _unfused_absorbed_dsa_fn(
     varlen_starts: Optional[torch.Tensor] = None,
     varlen_ends: Optional[torch.Tensor] = None,
     key_positions: Optional[torch.Tensor] = None,
+    accuracy_compatible: bool = False,
 ) -> torch.Tensor:
     """Unfused absorbed-MLA attention: output stays [sq, b, np, v_channels]."""
     sq, b, np, hn = query.size()
@@ -99,10 +100,15 @@ def _unfused_absorbed_dsa_fn(
     )
 
     attention_scores = attention_scores + index_mask.unsqueeze(1)
-    valid_index_mask = torch.isfinite(index_mask)
-    attention_scores = dsa_masking.masked_softmax(
-        attention_scores.float(), valid_index_mask.unsqueeze(1).expand(b, np, sq, skv), dim=-1
-    )
+    valid_index_mask = torch.isfinite(index_mask).unsqueeze(1).expand(b, np, sq, skv)
+    if accuracy_compatible:
+        attention_scores = _AccuracyCompatibleSoftmax.apply(
+            attention_scores.float(), valid_index_mask
+        )
+    else:
+        attention_scores = dsa_masking.masked_softmax(
+            attention_scores.float(), valid_index_mask, dim=-1
+        )
 
     # Latent value is the first v_channels slice of absorbed key cache.
     value = key[..., :v_channels].permute(1, 2, 0, 3)  # [b,1,skv,v]
@@ -110,6 +116,25 @@ def _unfused_absorbed_dsa_fn(
     return output.permute(2, 0, 1, 3).contiguous()
 
 
+class _AccuracyCompatibleSoftmax(torch.autograd.Function):
+    """Masked softmax with an explicit backward formula for DSA alignment."""
+
+    @staticmethod
+    def forward(ctx, logits: torch.Tensor, valid_mask: torch.Tensor) -> torch.Tensor:
+        probabilities = torch.softmax(logits.masked_fill(~valid_mask, float("-inf")), dim=-1)
+        probabilities = probabilities.masked_fill(~valid_mask, 0.0)
+        ctx.save_for_backward(probabilities, valid_mask)
+        return probabilities
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        probabilities, valid_mask = ctx.saved_tensors
+        grad_logits = probabilities * (
+            grad_output - (grad_output * probabilities).sum(dim=-1, keepdim=True)
+        )
+        return grad_logits.masked_fill(~valid_mask, 0.0), None
+
+
 def _run_sparse_attention(
     *,
     absorbed_mla: bool,
@@ -127,6 +152,7 @@ def _run_sparse_attention(
     topk_length: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Run sparse attention for absorbed and non-absorbed MLA paths."""
+    accuracy_compatible = bool(getattr(config, "dsa_accuracy_compatible", False))
     if absorbed_mla:
         latent_v_channels = int(getattr(config, "kv_lora_rank", 0) or 0)
         if latent_v_channels <= 0:
@@ -143,7 +169,7 @@ def _run_sparse_attention(
                 "Received absorbed layout with explicit value tensor."
             )
         output = None
-        if dsa_kernels.use_fused_dsa_kernels(config):
+        if not accuracy_compatible and dsa_kernels.use_fused_dsa_kernels(config):
             output = dsa_kernels.run_fused_absorbed_sparse_attention(
                 config,
                 query,
@@ -166,6 +192,7 @@ def _run_sparse_attention(
                 varlen_starts=varlen_starts,
                 varlen_ends=varlen_ends,
                 key_positions=key_positions,
+                accuracy_compatible=accuracy_compatible,
             )
         assert output is not None
         output = torch.einsum("sbhc,hdc->sbhd", output, up_v_weight).contiguous()
@@ -182,6 +209,7 @@ def _run_sparse_attention(
         varlen_starts=varlen_starts,
         varlen_ends=varlen_ends,
         key_positions=key_positions,
+        accuracy_compatible=accuracy_compatible,
     )
 
 
@@ -1411,6 +1439,7 @@ def unfused_dsa_fn(
     varlen_starts: Optional[torch.Tensor] = None,
     varlen_ends: Optional[torch.Tensor] = None,
     key_positions: Optional[torch.Tensor] = None,
+    accuracy_compatible: bool = False,
 ):
     """
     Unfused sparse attention implementation.
@@ -1457,6 +1486,27 @@ def unfused_dsa_fn(
         device=query.device,
     )
 
+    if accuracy_compatible:
+        index_mask = torch.full((b, sq, skv), float("-inf"), device=query.device)
+        dsa_masking.scatter_topk_into_index_mask(index_mask, topk_indices)
+        index_mask = dsa_masking.apply_sparse_validity_to_index_mask(
+            index_mask,
+            row_mask=row_mask,
+            varlen_starts=varlen_starts,
+            varlen_ends=varlen_ends,
+            key_positions=key_positions,
+        )
+        valid_index_mask = torch.isfinite(index_mask).unsqueeze(1).expand(b, np, sq, skv)
+        attention_scores = (
+            torch.matmul(query_b.float(), key_b.float().transpose(-1, -2)) * softmax_scale
+        )
+        attention_probs = _AccuracyCompatibleSoftmax.apply(
+            attention_scores + index_mask.unsqueeze(1), valid_index_mask
+        )
+        output = torch.matmul(attention_probs.to(value_b.dtype), value_b)
+        output = output.permute(2, 0, 1, 3).contiguous().view(sq, b, np * hnv)
+        return output.squeeze(1) if query_was_thd else output
+
     seq_chunk_size = 512
     head_chunk_size = 16
     topk_chunk_size = 1024
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 03317b65f1c..e25339b5ff1 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -104,6 +104,14 @@ def gating(self, input: torch.Tensor):
             router_dtype = torch.float32
         elif self.config.moe_router_dtype == 'fp64':
             router_dtype = torch.float64
+        if self.config.router_accuracy_compatible:
+            inp_shape = input.shape
+            logits = torch.mm(
+                input.reshape(-1, inp_shape[-1]).float(), self.weight.float().t()
+            )
+            if self.bias is not None:
+                logits = logits + self.bias.float()
+            return logits.view(*inp_shape[:-1], -1)
         logits = router_gating_linear(input, self.weight, self.bias, router_dtype)
         return logits
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index b20514ce6a4..ee92afe4954 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -11,7 +11,10 @@
 
 from megatron.core import InferenceParams, parallel_state, tensor_parallel
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
-from megatron.core.dist_checkpointing.utils import apply_prefix_mapping, replace_prefix_for_sharding
+from megatron.core.dist_checkpointing.utils import (
+    apply_prefix_mapping,
+    replace_prefix_for_sharding,
+)
 from megatron.core.enums import Fp8Recipe
 from megatron.core.extensions.transformer_engine import HAVE_TE
 from megatron.core.fp8_utils import get_fp8_context
@@ -30,7 +33,7 @@
 from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
-from megatron.core.transformer.torch_norm import LayerNormBuilder
+from megatron.core.transformer.torch_norm import LayerNormBuilder, WrappedTorchNorm
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
@@ -61,7 +64,9 @@
 else:
     TESpecProvider = None
 
-from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
+from megatron.core.transformer.pipeline_parallel_layer_layout import (
+    PipelineParallelLayerLayout,
+)
 
 
 def tie_word_embeddings_state_dict(
@@ -165,7 +170,9 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non
 
     # Handle packed sequences cases
     if packed_seq_params is not None:
-        return _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group)
+        return _roll_tensor_packed_seq(
+            tensor, shifts, dims, packed_seq_params, cp_group
+        )
 
     # Standard rolling behavior when CP is not enabled (cp_group is None or size=1)
     if cp_group is None or cp_group.size() == 1:
@@ -202,17 +209,25 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non
     # Start send and recv ops
     ops = []
     if local_rank != 0:
-        req_send_first_part = torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank)
+        req_send_first_part = torch.distributed.isend(
+            tensor=tensor_send_list[0], dst=prev_rank
+        )
         ops.append(req_send_first_part)
-        req_recv_second_part = torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank)
+        req_recv_second_part = torch.distributed.irecv(
+            tensor=tensor_recv_list[1], src=prev_rank
+        )
         ops.append(req_recv_second_part)
     else:
         # Inserted elements are set to be 0.0.
         tensor_recv_list[1] = 0
     if local_rank != len(global_ranks) - 1:
-        req_recv_first_part = torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank)
+        req_recv_first_part = torch.distributed.irecv(
+            tensor=tensor_recv_list[0], src=next_rank
+        )
         ops.append(req_recv_first_part)
-        req_send_second_part = torch.distributed.isend(tensor=tensor_send_list[1], dst=next_rank)
+        req_send_second_part = torch.distributed.isend(
+            tensor=tensor_send_list[1], dst=next_rank
+        )
         ops.append(req_send_second_part)
     else:
         # For the last CP rank, the removed elements of second part go into the first part
@@ -242,12 +257,14 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No
 
     # Notice: This is a naive implementation to test the correctness,
     # a better solution will only sync the boundary tokens once.
-    assert (
-        dims == -1 or dims == tensor.dim() - 1
-    ), "Packed sequence roll only supports the last dimension."
+    assert dims == -1 or dims == tensor.dim() - 1, (
+        "Packed sequence roll only supports the last dimension."
+    )
     assert shifts == -1, "Packed sequence roll only supports a single-token left shift."
     cu_seqlens = packed_seq_params.cu_seqlens_q
-    assert cu_seqlens is not None, "Packed sequence parameters must provide cu_seqlens_q."
+    assert cu_seqlens is not None, (
+        "Packed sequence parameters must provide cu_seqlens_q."
+    )
 
     rolled_tensor = tensor.clone()
 
@@ -289,7 +306,9 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No
 
         # The following code is very similar as the code in roll_tensor function
         local_chunks = tensor_slice.chunk(2, dim=dims)
-        rolled_chunks = [torch.roll(chunk, shifts=shifts, dims=dims) for chunk in local_chunks]
+        rolled_chunks = [
+            torch.roll(chunk, shifts=shifts, dims=dims) for chunk in local_chunks
+        ]
 
         tensor_send_list = []
         tensor_recv_list = []
@@ -297,10 +316,14 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No
             # Skip empty chunks that can occur when the sequence slice is very small
             if chunk.size(dims) == 0:
                 tensor_send_list.append(
-                    torch.empty(chunk.shape[:-1], dtype=chunk.dtype, device=chunk.device)
+                    torch.empty(
+                        chunk.shape[:-1], dtype=chunk.dtype, device=chunk.device
+                    )
                 )
                 tensor_recv_list.append(
-                    torch.empty(chunk.shape[:-1], dtype=chunk.dtype, device=chunk.device)
+                    torch.empty(
+                        chunk.shape[:-1], dtype=chunk.dtype, device=chunk.device
+                    )
                 )
                 continue
             boundary = chunk.select(dims, shifts).contiguous().clone()
@@ -309,14 +332,22 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No
 
         ops = []
         if local_rank != 0:
-            ops.append(torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank))
-            ops.append(torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank))
+            ops.append(
+                torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank)
+            )
+            ops.append(
+                torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank)
+            )
         else:
             tensor_recv_list[1].zero_()
 
         if local_rank != cp_size - 1:
-            ops.append(torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank))
-            ops.append(torch.distributed.isend(tensor=tensor_send_list[1], dst=next_rank))
+            ops.append(
+                torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank)
+            )
+            ops.append(
+                torch.distributed.isend(tensor=tensor_send_list[1], dst=next_rank)
+            )
         else:
             tensor_recv_list[0].copy_(tensor_send_list[1])
 
@@ -371,11 +402,17 @@ def save_metrics_to_tracker(
 
         tracker = MTPLossLoggingHelper.tracker
         if "loss_values" not in tracker:
-            tracker["loss_values"] = torch.zeros(num_layers, device=torch.cuda.current_device())
+            tracker["loss_values"] = torch.zeros(
+                num_layers, device=torch.cuda.current_device()
+            )
         if "correct_values" not in tracker:
-            tracker["correct_values"] = torch.zeros(num_layers, device=torch.cuda.current_device())
+            tracker["correct_values"] = torch.zeros(
+                num_layers, device=torch.cuda.current_device()
+            )
         if "total_values" not in tracker:
-            tracker["total_values"] = torch.zeros(num_layers, device=torch.cuda.current_device())
+            tracker["total_values"] = torch.zeros(
+                num_layers, device=torch.cuda.current_device()
+            )
 
         tracker["loss_values"][layer_number] += loss.detach()
         tracker["correct_values"][layer_number] += correct.detach()
@@ -404,26 +441,32 @@ def reduce_metrics_in_tracker():
             return
 
         loss_values = tracker["loss_values"]
-        if tracker.get('reduce_group') is not None:
-            torch.distributed.all_reduce(loss_values, group=tracker.get('reduce_group'))
-        if tracker.get('avg_group') is not None:
+        if tracker.get("reduce_group") is not None:
+            torch.distributed.all_reduce(loss_values, group=tracker.get("reduce_group"))
+        if tracker.get("avg_group") is not None:
             torch.distributed.all_reduce(
-                loss_values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG
+                loss_values,
+                group=tracker["avg_group"],
+                op=torch.distributed.ReduceOp.AVG,
             )
 
         for key in ["correct_values", "total_values"]:
             if key not in tracker:
                 continue
             values = tracker[key]
-            if tracker.get('reduce_group') is not None:
-                torch.distributed.all_reduce(values, group=tracker.get('reduce_group'))
-            if tracker.get('avg_group') is not None:
+            if tracker.get("reduce_group") is not None:
+                torch.distributed.all_reduce(values, group=tracker.get("reduce_group"))
+            if tracker.get("avg_group") is not None:
                 torch.distributed.all_reduce(
-                    values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.SUM
+                    values,
+                    group=tracker["avg_group"],
+                    op=torch.distributed.ReduceOp.SUM,
                 )
 
     @staticmethod
-    def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None):
+    def track_mtp_metrics(
+        loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None
+    ):
         """Track the Multi-Token Prediction (MTP) metrics for logging."""
         MTPLossLoggingHelper.reduce_metrics_in_tracker()
         tracker = MTPLossLoggingHelper.tracker
@@ -453,15 +496,16 @@ def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_lo
 
         mtp_num_layers = mtp_losses.shape[0]
         for i in range(mtp_num_layers):
-            loss_name = f"mtp_{i+1} loss"
-            step_acc_name = f"mtp_{i+1}_acceptance_rate"
-            cum_acc_name = f"mtp_{i+1}_cumulative_acceptance_rate"
+            loss_name = f"mtp_{i + 1} loss"
+            step_acc_name = f"mtp_{i + 1}_acceptance_rate"
+            cum_acc_name = f"mtp_{i + 1}_cumulative_acceptance_rate"
 
             loss = mtp_losses[i]
             # Empty masks can leave no valid MTP positions, so clamp denominators to avoid NaNs.
             step_rate = (mtp_corrects[i] / torch.clamp(mtp_totals[i], min=1)) * 100.0
             cum_rate = (
-                mtp_cumulative_corrects[i] / torch.clamp(mtp_cumulative_totals[i], min=1)
+                mtp_cumulative_corrects[i]
+                / torch.clamp(mtp_cumulative_totals[i], min=1)
             ) * 100.0
 
             if total_loss_dict is not None:
@@ -491,7 +535,9 @@ def _mtp_logits_are_vocab_sharded(
 
 
 def _vocab_parallel_argmax(
-    vocab_parallel_logits: Tensor, tp_group: torch.distributed.ProcessGroup, tp_size: int
+    vocab_parallel_logits: Tensor,
+    tp_group: torch.distributed.ProcessGroup,
+    tp_size: int,
 ) -> Tensor:
     """Return global argmax ids from logits sharded across the vocab dimension."""
     vocab_shard_size = vocab_parallel_logits.size(-1)
@@ -505,9 +551,9 @@ def _vocab_parallel_argmax(
     stacked_max_vals = torch.stack(gathered_max_vals, dim=0)
     stacked_argmax = torch.stack(gathered_argmax, dim=0)
     winning_rank = stacked_max_vals.argmax(dim=0)  # [s, b]
-    winning_local_argmax = torch.gather(stacked_argmax, 0, winning_rank.unsqueeze(0)).squeeze(
-        0
-    )  # [s, b]
+    winning_local_argmax = torch.gather(
+        stacked_argmax, 0, winning_rank.unsqueeze(0)
+    ).squeeze(0)  # [s, b]
     return winning_rank * vocab_shard_size + winning_local_argmax  # [s, b]
 
 
@@ -534,7 +580,11 @@ def _compute_mtp_acceptance_counts(
                 "tp_group must be provided when computing MTP acceptance counts "
                 "from vocab-sharded logits under tensor model parallelism."
             )
-        tp_size = torch.distributed.get_world_size(group=tp_group) if tp_group is not None else 1
+        tp_size = (
+            torch.distributed.get_world_size(group=tp_group)
+            if tp_group is not None
+            else 1
+        )
 
         # Apply TP rank offsets only when logits are vocab-sharded; gathered logits already
         # contain global vocab ids in their last dimension.
@@ -577,7 +627,9 @@ class MultiTokenPredictionLayerSubmodules:
 
 
 def get_mtp_layer_spec(
-    mtp_model_layer_spec: ModuleSpec, use_transformer_engine: bool
+    mtp_model_layer_spec: ModuleSpec,
+    use_transformer_engine: bool,
+    config: Optional[TransformerConfig] = None,
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -587,11 +639,14 @@ def get_mtp_layer_spec(
     return get_mtp_layer_spec_for_backend(
         mtp_model_layer_spec,
         backend=TESpecProvider() if use_transformer_engine else LocalSpecProvider(),
+        config=config,
     )
 
 
 def get_mtp_layer_spec_for_backend(
-    mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
+    mtp_model_layer_spec: ModuleSpec,
+    backend: BackendSpecProvider,
+    config: Optional[TransformerConfig] = None,
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -599,7 +654,11 @@ def get_mtp_layer_spec_for_backend(
         ModuleSpec: Module specification with modules from the backend.
     """
     column_parallel_linear_impl: type = backend.column_parallel_linear()
-    layer_norm_impl = backend.layer_norm()
+    layer_norm_impl = (
+        WrappedTorchNorm
+        if config is not None and config.norm_accuracy_compatible
+        else backend.layer_norm()
+    )
     mtp_layer_spec = ModuleSpec(
         module=MultiTokenPredictionLayer,
         submodules=MultiTokenPredictionLayerSubmodules(
@@ -637,14 +696,19 @@ def mtp_on_this_rank(
         # with custom PP layout, we support put MTP layers on any pipeline stage
         if (
             not ignore_virtual
-            and parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None
+            and parallel_state.get_virtual_pipeline_model_parallel_world_size()
+            is not None
         ):
-            assert vp_stage is not None, "vp_stage must be passed if virtual pipeline is enabled"
+            assert vp_stage is not None, (
+                "vp_stage must be passed if virtual pipeline is enabled"
+            )
             num_layers_to_build = layout.layout[pp_rank][vp_stage].count(LayerType.mtp)
             mtp_on_this_rank = num_layers_to_build > 0
         else:
             for vpp_rank in range(len(layout.layout[pp_rank])):
-                num_layers_to_build = layout.layout[pp_rank][vpp_rank].count(LayerType.mtp)
+                num_layers_to_build = layout.layout[pp_rank][vpp_rank].count(
+                    LayerType.mtp
+                )
                 if num_layers_to_build > 0:
                     mtp_on_this_rank = True
                     break
@@ -675,7 +739,9 @@ def get_mtp_ranks(pp_ranks: List[int], config: TransformerConfig) -> List[int]:
     return list(mtp_ranks)
 
 
-def get_mtp_layer_offset(config: TransformerConfig, vp_stage: Optional[int] = None) -> int:
+def get_mtp_layer_offset(
+    config: TransformerConfig, vp_stage: Optional[int] = None
+) -> int:
     """Get the offset of the MTP layer."""
     if config.pipeline_model_parallel_size > 1:
         if config.pipeline_model_parallel_layout:
@@ -690,21 +756,29 @@ def get_mtp_layer_offset(config: TransformerConfig, vp_stage: Optional[int] = No
 
 
 def get_mtp_num_layers_to_build(
-    config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None
+    config: TransformerConfig,
+    vp_stage: Optional[int] = None,
+    pp_rank: Optional[int] = None,
 ) -> int:
     """Get the number of MTP layers to build."""
     if config.pipeline_model_parallel_layout is not None:
         # If we have a custom PP layout, get the number of mtp layers in the layout array.
-        num_layers_to_build = config.pipeline_model_parallel_layout.get_num_layers_to_build(
-            layer_type=LayerType.mtp, vp_stage=vp_stage
+        num_layers_to_build = (
+            config.pipeline_model_parallel_layout.get_num_layers_to_build(
+                layer_type=LayerType.mtp, vp_stage=vp_stage
+            )
         )
-        assert num_layers_to_build == config.mtp_num_layers or num_layers_to_build == 0, (
+        assert (
+            num_layers_to_build == config.mtp_num_layers or num_layers_to_build == 0
+        ), (
             f"Currently, we only support put all of MTP layers on the last pipeline stage, "
             f"so the number of MTP layers to build ({num_layers_to_build}) must match "
             f"mtp_num_layers ({config.mtp_num_layers}) or be 0."
         )
     else:
-        if parallel_state.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage):
+        if parallel_state.is_pipeline_last_stage(
+            ignore_virtual=False, vp_stage=vp_stage
+        ):
             num_layers_to_build = config.mtp_num_layers if config.mtp_num_layers else 0
         else:
             num_layers_to_build = 0
@@ -811,7 +885,11 @@ def process_mtp_loss(
         if input_ids is None:
             return hidden_states
         labels, _ = roll_tensor(
-            input_ids, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+            input_ids,
+            shifts=-1,
+            dims=-1,
+            cp_group=cp_group,
+            packed_seq_params=packed_seq_params,
         )
         derived_labels_from_input_ids = True
 
@@ -829,7 +907,11 @@ def process_mtp_loss(
         # label is fabricated (zeroed). Roll loss_mask in lockstep with the
         # input_ids -> labels shift so that boundary position is masked.
         loss_mask, _ = roll_tensor(
-            loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+            loss_mask,
+            shifts=-1,
+            dims=-1,
+            cp_group=cp_group,
+            packed_seq_params=packed_seq_params,
         )
 
     # Store the original number of tokens before rolling for proper normalization
@@ -846,10 +928,18 @@ def process_mtp_loss(
         if scale_logits_fn is not None:
             mtp_logits = scale_logits_fn(mtp_logits)
         mtp_labels, _ = roll_tensor(
-            mtp_labels, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+            mtp_labels,
+            shifts=-1,
+            dims=-1,
+            cp_group=cp_group,
+            packed_seq_params=packed_seq_params,
         )
         loss_mask, num_tokens = roll_tensor(
-            loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+            loss_mask,
+            shifts=-1,
+            dims=-1,
+            cp_group=cp_group,
+            packed_seq_params=packed_seq_params,
         )
 
         mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
@@ -861,7 +951,12 @@ def process_mtp_loss(
                 torch.sum(mtp_loss) * (num_tokens > 0).to(mtp_loss.dtype)
             ) / num_tokens.clamp(min=1)
             correct, total = _compute_mtp_acceptance_counts(
-                mtp_logits, mtp_labels, loss_mask, output_layer, runtime_gather_output, tp_group
+                mtp_logits,
+                mtp_labels,
+                loss_mask,
+                output_layer,
+                runtime_gather_output,
+                tp_group,
             )
 
             MTPLossLoggingHelper.save_metrics_to_tracker(
@@ -870,7 +965,9 @@ def process_mtp_loss(
                 total,
                 mtp_layer_number,
                 config.mtp_num_layers,
-                avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True),
+                avg_group=parallel_state.get_data_parallel_group(
+                    with_context_parallel=True
+                ),
             )
         mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers
         if config.calculate_per_token_loss:
@@ -954,18 +1051,28 @@ def __init__(
 
         # Validate attention mask type if using transformer-based inner layers
         if self.submodules.mtp_model_layer is not None and hasattr(
-            self.submodules.mtp_model_layer, 'submodules'
+            self.submodules.mtp_model_layer, "submodules"
         ):
             from megatron.core.models.hybrid.hybrid_block import HybridStackSubmodules
-            from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules
+            from megatron.core.transformer.transformer_layer import (
+                TransformerLayerSubmodules,
+            )
 
             layer_submodules = None
-            if isinstance(self.submodules.mtp_model_layer.submodules, HybridStackSubmodules):
-                attention_layer_spec = self.submodules.mtp_model_layer.submodules.attention_layer
-                if hasattr(attention_layer_spec, 'submodules'):
-                    assert isinstance(attention_layer_spec.submodules, TransformerLayerSubmodules)
+            if isinstance(
+                self.submodules.mtp_model_layer.submodules, HybridStackSubmodules
+            ):
+                attention_layer_spec = (
+                    self.submodules.mtp_model_layer.submodules.attention_layer
+                )
+                if hasattr(attention_layer_spec, "submodules"):
+                    assert isinstance(
+                        attention_layer_spec.submodules, TransformerLayerSubmodules
+                    )
                     layer_submodules = attention_layer_spec.submodules
-            elif isinstance(self.submodules.mtp_model_layer.submodules, TransformerLayerSubmodules):
+            elif isinstance(
+                self.submodules.mtp_model_layer.submodules, TransformerLayerSubmodules
+            ):
                 layer_submodules = self.submodules.mtp_model_layer.submodules
             else:
                 raise ValueError(
@@ -973,7 +1080,7 @@ def __init__(
                 )
             if layer_submodules:
                 self_attention_spec = layer_submodules.self_attention
-                attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                attn_mask_type = self_attention_spec.params.get("attn_mask_type", "")
                 assert attn_mask_type in SUPPORTED_ATTN_MASK, (
                     f"Multi-Token Prediction (MTP) is not yet supported with "
                     f"{attn_mask_type} attention mask type. "
@@ -1017,7 +1124,9 @@ def __init__(
         # 2. GPT path: single TransformerLayer
         if mtp_layer_pattern is not None and hybrid_submodules is not None:
             from megatron.core.models.hybrid.hybrid_block import HybridStack
-            from megatron.core.models.hybrid.hybrid_layer_allocation import validate_segment_layers
+            from megatron.core.models.hybrid.hybrid_layer_allocation import (
+                validate_segment_layers,
+            )
 
             self.mtp_model_layer = HybridStack(
                 config=self.config,
@@ -1106,7 +1215,9 @@ def _get_embeddings(
         if self.config.mtp_detach_heads:
             decoder_input = decoder_input.detach()
 
-        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+        hidden_states = make_viewless_tensor(
+            inp=hidden_states, requires_grad=True, keep_graph=True
+        )
         # make_viewless_tensor no-ops when hidden_states is not a view (_base is None),
         # which happens after detach() with mtp_detach_heads. Activation
         # checkpointing (CheckpointFunction.apply) requires at least one input tensor
@@ -1117,14 +1228,20 @@ def _get_embeddings(
 
         return input_ids, position_ids, padding_mask, decoder_input, hidden_states
 
-    def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.Tensor):
+    def _concat_embeddings(
+        self, hidden_states: torch.Tensor, decoder_input: torch.Tensor
+    ):
         """
         Concatenate the tokens before sending to transformer layer.
         """
         decoder_input = apply_module(self.enorm)(decoder_input)
-        decoder_input = make_viewless_tensor(inp=decoder_input, requires_grad=True, keep_graph=True)
+        decoder_input = make_viewless_tensor(
+            inp=decoder_input, requires_grad=True, keep_graph=True
+        )
         hidden_states = apply_module(self.hnorm)(hidden_states)
-        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+        hidden_states = make_viewless_tensor(
+            inp=hidden_states, requires_grad=True, keep_graph=True
+        )
         # At the (k - 1)-th MTP module, concatenates the i-th token's hidden_states
         # and the (i + K)-th token's embedding, and combine them with linear projection.
         hidden_states = torch.cat((decoder_input, hidden_states), -1)
@@ -1141,7 +1258,9 @@ def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.T
             )
         # For sequence parallel, scatter after linear_fc and before transformer layer.
         if self.sequence_parallel:
-            hidden_states = scatter_to_sequence_parallel_region(hidden_states, group=self.tp_group)
+            hidden_states = scatter_to_sequence_parallel_region(
+                hidden_states, group=self.tp_group
+            )
         return hidden_states
 
     def _proj_and_transformer_layer(
@@ -1226,7 +1345,9 @@ def _postprocess(self, hidden_states: torch.Tensor):
         # TENorm produces a "viewed" tensor. This will result in schedule.py's
         # deallocate_output_tensor() throwing an error, so a viewless tensor is
         # created to prevent this.
-        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+        hidden_states = make_viewless_tensor(
+            inp=hidden_states, requires_grad=True, keep_graph=True
+        )
 
         return hidden_states
 
@@ -1404,19 +1525,20 @@ def checkpoint_handler():
                     sequence_len_offset,
                 )
 
-        if self.config.recompute_method == 'uniform':
+        if self.config.recompute_method == "uniform":
             # Uniformly divide the total number of Transformer layers and checkpoint
             # the input activation of each divided chunk.
             # A method to further reduce memory usage reducing checkpoints.
-            assert (
-                self.config.recompute_num_layers == 1
-            ), "recompute_num_layers must be 1 for MTP recompute"
+            assert self.config.recompute_num_layers == 1, (
+                "recompute_num_layers must be 1 for MTP recompute"
+            )
             with outer_quantization_context:
                 outputs = checkpoint_handler()
-        elif self.config.recompute_method == 'block':
+        elif self.config.recompute_method == "block":
             # TODO: implement block-based recompute for MTP
             warnings.warn(
-                "recompute_method == 'block' is not supported for MTP yet." " Skipping recompute."
+                "recompute_method == 'block' is not supported for MTP yet."
+                " Skipping recompute."
             )
             outputs = self._proj_and_transformer_layer(
                 hidden_states=hidden_states,
@@ -1478,17 +1600,21 @@ def forward(
             Union[Tensor, Tuple[Tensor, Tensor]]: The output hidden states tensor of shape
             [s, b, h], and optionally the updated context tensor if cross-attention is used.
         """
-        assert context is None, "multi token prediction + cross attention is not yet supported."
-        input_ids, position_ids, padding_mask, decoder_input, hidden_states = self._get_embeddings(
-            input_ids=input_ids,
-            position_ids=position_ids,
-            padding_mask=padding_mask,
-            embedding=embedding,
-            hidden_states=hidden_states,
-            packed_seq_params=packed_seq_params,
+        assert context is None, (
+            "multi token prediction + cross attention is not yet supported."
+        )
+        input_ids, position_ids, padding_mask, decoder_input, hidden_states = (
+            self._get_embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                padding_mask=padding_mask,
+                embedding=embedding,
+                hidden_states=hidden_states,
+                packed_seq_params=packed_seq_params,
+            )
         )
 
-        if self.config.recompute_granularity == 'full' and self.training:
+        if self.config.recompute_granularity == "full" and self.training:
             hidden_states = self._checkpointed_forward(
                 hidden_states=hidden_states,
                 decoder_input=decoder_input,
@@ -1524,7 +1650,10 @@ def forward(
         return hidden_states, input_ids, position_ids, padding_mask
 
     def sharded_state_dict(
-        self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
+        self,
+        prefix: str = "",
+        sharded_offsets: tuple = (),
+        metadata: Optional[dict] = None,
     ) -> ShardedStateDict:
         """
         Generate a sharded state dictionary for the multi token prediction layer.
@@ -1538,7 +1667,9 @@ def sharded_state_dict(
             ShardedStateDict: A dictionary containing the sharded state of the multi
             token prediction layer.
         """
-        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        sharded_state_dict = super().sharded_state_dict(
+            prefix, sharded_offsets, metadata
+        )
 
         # Backward compatibility: GPT MTP checkpoints were saved with the submodule
         # named 'transformer_layer'. Remap checkpoint keys so old checkpoints load
@@ -1546,7 +1677,8 @@ def sharded_state_dict(
         # since no older checkpoints exist for them.
         if self.mtp_layer_pattern is None:
             apply_prefix_mapping(
-                sharded_state_dict, {f'{prefix}mtp_model_layer.': f'{prefix}transformer_layer.'}
+                sharded_state_dict,
+                {f"{prefix}mtp_model_layer.": f"{prefix}transformer_layer."},
             )
 
         return sharded_state_dict
@@ -1571,7 +1703,8 @@ class MultiTokenPredictionBlockSubmodules:
 
 
 def _get_mtp_block_submodules(
-    config: TransformerConfig, spec: Union[MultiTokenPredictionBlockSubmodules, ModuleSpec]
+    config: TransformerConfig,
+    spec: Union[MultiTokenPredictionBlockSubmodules, ModuleSpec],
 ) -> MultiTokenPredictionBlockSubmodules:
     """
     Retrieve or construct MultiTokenPredictionBlockSubmodules based on the provided specification.
@@ -1672,21 +1805,25 @@ def __init__(
         # to the roll_tensor function for proper boundary communication
         if pg_collection is None:
             # Use default MPU process groups if not provided
-            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp', 'tp'])
+            pg_collection = ProcessGroupCollection.use_mpu_process_groups(
+                required_pgs=["cp", "tp"]
+            )
         else:
             # Ensure the provided process groups include CP
-            assert hasattr(
-                pg_collection, 'cp'
-            ), "MultiTokenPredictionBlock pg_collection must have cp process group"
+            assert hasattr(pg_collection, "cp"), (
+                "MultiTokenPredictionBlock pg_collection must have cp process group"
+            )
 
         self._build_layers(pg_collection)
-        assert len(self.layers) > 0, "MultiTokenPredictionBlock must have at least one layer."
+        assert len(self.layers) > 0, (
+            "MultiTokenPredictionBlock must have at least one layer."
+        )
         self.cp_group = pg_collection.cp
 
         if self.config.mtp_detach_heads:
             # Tag MTP params so the optimizer can clip their gradients separately.
             for param in self.parameters():
-                param.grad_norm_group = 'mtp'
+                param.grad_norm_group = "mtp"
 
     def _build_layers(self, pg_collection):
         # Determine number of depths to build
@@ -1706,7 +1843,9 @@ def build_layer_legacy(layer_spec, layer_number):
                     vp_stage=self.vp_stage,
                     pg_collection=pg_collection,
                     mtp_layer_pattern=self.mtp_layer_pattern,
-                    name=(self.name + f".layers.{layer_number}") if self.name is not None else None,
+                    name=(self.name + f".layers.{layer_number}")
+                    if self.name is not None
+                    else None,
                 )
             return module
 
@@ -1724,7 +1863,9 @@ def build_layer_with_pattern(
                     pg_collection=pg_collection,
                     mtp_layer_pattern=mtp_layer_pattern,
                     hybrid_submodules=hybrid_submodules,
-                    name=(self.name + f".layers.{layer_number}") if self.name is not None else None,
+                    name=(self.name + f".layers.{layer_number}")
+                    if self.name is not None
+                    else None,
                 )
             return module
 
@@ -1816,7 +1957,9 @@ def forward(
 
         for iteration in range(self.config.mtp_num_layers):
             layer_idx = 0 if self.mtp_use_repeated_layer else iteration
-            (hidden_states, input_ids, position_ids, padding_mask) = self.layers[layer_idx](
+            (hidden_states, input_ids, position_ids, padding_mask) = self.layers[
+                layer_idx
+            ](
                 input_ids=input_ids,
                 position_ids=position_ids,
                 hidden_states=hidden_states,
@@ -1841,7 +1984,10 @@ def forward(
         return hidden_states
 
     def sharded_state_dict(
-        self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
+        self,
+        prefix: str = "",
+        sharded_offsets: tuple = (),
+        metadata: Optional[dict] = None,
     ) -> ShardedStateDict:
         """
         Generate a sharded state dictionary for the multi token prediction module.
@@ -1856,16 +2002,18 @@ def sharded_state_dict(
             token prediction module.
         """
         sharded_state_dict = {}
-        layer_prefix = f'{prefix}layers.'
+        layer_prefix = f"{prefix}layers."
         for layer in self.layers:
             offset = get_mtp_layer_offset(self.config, self.vp_stage)
-            sharded_prefix = f'{layer_prefix}{layer.layer_number - 1}.'
+            sharded_prefix = f"{layer_prefix}{layer.layer_number - 1}."
 
-            state_dict_prefix = f'{layer_prefix}{layer.layer_number - 1 - offset}.'
+            state_dict_prefix = f"{layer_prefix}{layer.layer_number - 1 - offset}."
             sharded_pp_offset = []
             layer_sharded_state_dict = layer.sharded_state_dict(
                 state_dict_prefix, sharded_pp_offset, metadata
             )
-            replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix)
+            replace_prefix_for_sharding(
+                layer_sharded_state_dict, state_dict_prefix, sharded_prefix
+            )
             sharded_state_dict.update(layer_sharded_state_dict)
         return sharded_state_dict
diff --git a/megatron/core/transformer/torch_norm.py b/megatron/core/transformer/torch_norm.py
index 5948ae600f9..c75525dcd59 100644
--- a/megatron/core/transformer/torch_norm.py
+++ b/megatron/core/transformer/torch_norm.py
@@ -41,32 +41,41 @@ def __new__(
         zero_centered_gamma: bool = False,
         normalization: str = "LayerNorm",
     ) -> LayerNormInterface:
-        assert (
-            not config.layernorm_zero_centered_gamma
-        ), f"zero_centered_gamma not supported by torch LayerNorm"
+        assert not config.layernorm_zero_centered_gamma, (
+            f"zero_centered_gamma not supported by torch LayerNorm"
+        )
 
-        assert not config.persist_layer_norm, f"persist_layer_norm not supported by torch LayerNorm"
+        assert not config.persist_layer_norm, (
+            f"persist_layer_norm not supported by torch LayerNorm"
+        )
 
-        assert not config.sequence_parallel, f"sequence parallel not supported by torch LayerNorm"
-
-        assert (
-            not config.memory_efficient_layer_norm
-        ), f"memory_efficient_layer_norm not supported by torch LayerNorm"
+        assert not config.memory_efficient_layer_norm, (
+            f"memory_efficient_layer_norm not supported by torch LayerNorm"
+        )
 
         if config.normalization == "LayerNorm":
             norm_cls = torch.nn.LayerNorm
         elif config.normalization == "RMSNorm":
-            assert is_torch_min_version(
-                "2.4.0a0"
-            ), 'Torch RMSNorm requires PyTorch version >= 2.4.0'
+            assert is_torch_min_version("2.4.0a0"), (
+                "Torch RMSNorm requires PyTorch version >= 2.4.0"
+            )
 
             norm_cls = torch.nn.RMSNorm
         elif config.normalization == "L2Norm":
             norm_cls = torch.nn.L2Norm
         else:
-            raise Exception("Only LayerNorm, RMSNorm and L2Norm are currently supported")
+            raise Exception(
+                "Only LayerNorm, RMSNorm and L2Norm are currently supported"
+            )
 
-        return norm_cls(normalized_shape=hidden_size, eps=eps)
+        factory_kwargs = {}
+        if config.normalization == "RMSNorm" and config.norm_accuracy_compatible:
+            factory_kwargs["dtype"] = config.params_dtype
+        norm = norm_cls(normalized_shape=hidden_size, eps=eps, **factory_kwargs)
+        if config.sequence_parallel:
+            for parameter in norm.parameters():
+                parameter.sequence_parallel = True
+        return norm
 
 
 class L2Norm(torch.nn.Module, LayerNormInterface):
@@ -99,7 +108,9 @@ def _norm(self, x: torch.Tensor) -> torch.Tensor:
             torch.Tensor: The L2-normalized tensor.
         """
         x_float = x.float()
-        return (x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)
+        return (
+            x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
+        ).type_as(x)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index bbcf413baee..a8e083341c2 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -25,7 +25,9 @@
     CudaGraphScope,
     InferenceCudaGraphScope,
 )
-from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
+from megatron.core.transformer.pipeline_parallel_layer_layout import (
+    PipelineParallelLayerLayout,
+)
 
 from .._rank_utils import log_single_rank
 from ..fusions.fused_bias_geglu import quick_gelu
@@ -101,7 +103,9 @@ class TransformerConfig(ModelParallelConfig):
     """Number of transformer layers on last pipeline stage.
     None implies equal layer division across PP ranks."""
 
-    pipeline_model_parallel_layout: Optional[Union[str, list, PipelineParallelLayerLayout]] = None
+    pipeline_model_parallel_layout: Optional[
+        Union[str, list, PipelineParallelLayerLayout]
+    ] = None
     """Custom definition of the pipeline parallel partitioning.
     Support type:
     - str: e.g., 'Et*3|(tt|)*29,m|L'. Stages are split by '|', replicated stages or layers
@@ -138,7 +142,9 @@ class TransformerConfig(ModelParallelConfig):
     hidden_size: int = field(default=0, metadata={"argparse_meta": {"default": None}})
     """Transformer hidden size."""
 
-    num_attention_heads: int = field(default=0, metadata={"argparse_meta": {"default": None}})
+    num_attention_heads: int = field(
+        default=0, metadata={"argparse_meta": {"default": None}}
+    )
     """Number of transformer attention heads."""
 
     attention_backend: AttnBackend = AttnBackend.auto
@@ -150,7 +156,7 @@ class TransformerConfig(ModelParallelConfig):
     softmax_scale: Optional[float] = None
     """Softmax scale for attention scaling."""
 
-    softmax_type: Literal['vanilla', 'off-by-one', 'learnable'] = 'vanilla'
+    softmax_type: Literal["vanilla", "off-by-one", "learnable"] = "vanilla"
     """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. 
        Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and 
        and learnable offset."""
@@ -186,14 +192,28 @@ class TransformerConfig(ModelParallelConfig):
     )
     """Epsilon value for any LayerNorm/RMSNorm operations."""
 
+    norm_accuracy_compatible: bool = field(
+        default=False,
+        metadata={"argparse_meta": {"arg_names": ["--norm-accuracy-compatible"]}},
+    )
+    """Use native Torch RMSNorm modules instead of Transformer Engine norm modules for alignment."""
+
+    router_accuracy_compatible: bool = field(
+        default=False,
+        metadata={"argparse_meta": {"arg_names": ["--router-accuracy-compatible"]}},
+    )
+    """Use an explicit fp32 router GEMM instead of the fused Transformer Engine path."""
+
     layernorm_zero_centered_gamma: bool = field(
-        default=False, metadata={"argparse_meta": {"arg_names": ["--apply-layernorm-1p"]}}
+        default=False,
+        metadata={"argparse_meta": {"arg_names": ["--apply-layernorm-1p"]}},
     )
     """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves
     numerical stability."""
 
     add_bias_linear: bool = field(
-        default=True, metadata={"argparse_meta": {"arg_names": ["--disable-bias-linear"]}}
+        default=True,
+        metadata={"argparse_meta": {"arg_names": ["--disable-bias-linear"]}},
     )
     """Include/exclude a bias term in all linear layers (QKV projections, after core attention,
     and two in MLP layer)."""
@@ -236,7 +256,7 @@ class TransformerConfig(ModelParallelConfig):
     - An integer N: Represents a (N-1):1 ratio, one full attention layer after (N-1) SWA layers.
     - A list that defines a custom pattern, e.g.: [1,1,1,1,0,0,0,0], where 1 represents SWA. """
 
-    normalization: Literal['LayerNorm', 'RMSNorm'] = "LayerNorm"
+    normalization: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
     """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`."""
 
     qk_layernorm: bool = False
@@ -281,10 +301,12 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # attention variant
     ####################
-    experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None
+    experimental_attention_variant: Optional[Literal["gated_delta_net", "dsa"]] = None
     """Type of attention variant to use. Currently support gated_delta_net and dsa."""
 
-    experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None
+    experimental_attention_variant_loss_scale_func: Optional[
+        Callable[[torch.Tensor], None]
+    ] = None
     """Optional hook for experimental attention variants to receive the main loss scale."""
 
     ####################
@@ -318,6 +340,12 @@ class TransformerConfig(ModelParallelConfig):
     ``none`` disables fused DSA kernels. Explicit ``tilelang`` or ``cudnn`` enables only that
     backend. Unsupported DSA layouts continue to use the PyTorch fallback."""
 
+    dsa_accuracy_compatible: bool = field(
+        default=False,
+        metadata={"argparse_meta": {"arg_names": ["--dsa-accuracy-compatible"]}},
+    )
+    """Use the full-score DSA fallback with explicit softmax backward for alignment."""
+
     dsa_indexer_rope_interleaved: bool = False
     """Whether DSA indexer RoPE should use MLA-style interleaving."""
 
@@ -505,7 +533,7 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # activation recomputation
     ####################
-    recompute_granularity: Optional[Literal['full', 'selective']] = None
+    recompute_granularity: Optional[Literal["full", "selective"]] = None
     """Determines which type of activation recompute to use.  Megatron-core supports 'selective'
     activation checkpointing where the submodules set in --recompute-modules is checkpointed.
     The default is "core_attn" which is the memory intensive part of attention.
@@ -516,7 +544,7 @@ class TransformerConfig(ModelParallelConfig):
     If set, must be 'selective' or 'full'. 'selective' always uses all layers.
     """
 
-    recompute_method: Optional[Literal['uniform', 'block']] = None
+    recompute_method: Optional[Literal["uniform", "block"]] = None
     """Determines which transformer layers will be recomputed. uniform will uniformly divide the
     total number of transformer layers in a transformer block and recompute the input activation of
     each divided chunk at the specified granularity.  block will recompute the input activations for
@@ -553,16 +581,16 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # fp8 related
     ####################
-    fp8: Optional[Literal['e4m3', 'hybrid']] = field(
+    fp8: Optional[Literal["e4m3", "hybrid"]] = field(
         default=None, metadata={"argparse_meta": {"arg_names": ["--fp8-format"]}}
     )
     """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined
     choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8
     activation and weight tensors and e5m2 for all FP8 output activation gradient tensors."""
 
-    fp8_recipe: Optional[Literal['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom']] = (
-        "delayed"
-    )
+    fp8_recipe: Optional[
+        Literal["tensorwise", "delayed", "mxfp8", "blockwise", "custom"]
+    ] = "delayed"
     """If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined
     choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed'
     uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only,
@@ -590,7 +618,7 @@ class TransformerConfig(ModelParallelConfig):
     fp8_amax_history_len: int = 1
     """The length of the amax history window used for scaling factor computation."""
 
-    fp8_amax_compute_algo: Literal['most_recent', 'max'] = "most_recent"
+    fp8_amax_compute_algo: Literal["most_recent", "max"] = "most_recent"
     """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2
     predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent`
     always chooses the most recently seen value.
@@ -639,13 +667,13 @@ class TransformerConfig(ModelParallelConfig):
     ####################
     # fp4 related
     ####################
-    fp4: Optional[Literal['e2m1']] = field(
+    fp4: Optional[Literal["e2m1"]] = field(
         default=None, metadata={"argparse_meta": {"arg_names": ["--fp4-format"]}}
     )
     """If set, enables the use of FP4 precision through Transformer Engine. Currently only 
     supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0)."""
 
-    fp4_recipe: Optional[Literal['nvfp4', 'custom']] = "nvfp4"
+    fp4_recipe: Optional[Literal["nvfp4", "custom"]] = "nvfp4"
     """If set, enables the use of FP4 precision through Transformer Engine. Currently only
     'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture."""
 
@@ -759,10 +787,10 @@ class TransformerConfig(ModelParallelConfig):
     """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax
     enabled. Defaults to None, which means no scaling."""
 
-    moe_router_score_function: Literal['softmax', 'sigmoid', 'sqrtsoftplus'] = "softmax"
+    moe_router_score_function: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "softmax"
     """Score function for MoE routing. Can be "softmax", "sigmoid" or "sqrtsoftplus"."""
 
-    moe_router_dtype: Optional[Literal['fp32', 'fp64']] = None
+    moe_router_dtype: Optional[Literal["fp32", "fp64"]] = None
     """Data type for routing and expert output weighted averaging. Using fp32 or fp64 can
     improve stability especially when the number of experts is large (e.g. finegrained-moe).
     None means no changes for dtype."""
@@ -816,7 +844,9 @@ class TransformerConfig(ModelParallelConfig):
     If a list of load balancing types is provided for `moe_router_load_balancing_type`,
     a corresponding list of coefficients should be provided here."""
 
-    moe_z_loss_coeff: Optional[float] = None  # 1e-3 would be a good start value for z-loss
+    moe_z_loss_coeff: Optional[float] = (
+        None  # 1e-3 would be a good start value for z-loss
+    )
     """Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended."""
 
     moe_input_jitter_eps: Optional[float] = None
@@ -827,14 +857,14 @@ class TransformerConfig(ModelParallelConfig):
     specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is
     currently unsupported so should remain False."""
 
-    moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex'] = "allgather"
+    moe_token_dispatcher_type: Literal["allgather", "alltoall", "flex"] = "allgather"
     """The type of token dispatcher to use. The default is 'allgather'.
     Options are 'allgather','alltoall' and 'flex'."""
 
     moe_enable_deepep: bool = False
     """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models."""
 
-    moe_flex_dispatcher_backend: Literal['deepep', 'hybridep'] = "deepep"
+    moe_flex_dispatcher_backend: Literal["deepep", "hybridep"] = "deepep"
     """[Experimental] The backend to use for flex token dispatcher. The default is "deepep".
     Options are "deepep" and "hybridep". Currently only "hybridep" backend supports 
     the MNNVL case."""
@@ -860,7 +890,7 @@ class TransformerConfig(ModelParallelConfig):
     max that an expert could see during inference so no tokens are actually dropped. The default
     setting is False."""
 
-    moe_token_drop_policy: Literal['probs', 'position'] = "probs"
+    moe_token_drop_policy: Literal["probs", "position"] = "probs"
     """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with
     the lowest probabilities will be dropped. If "position", tokens at the end of each batch will
     be dropped.
@@ -962,7 +992,9 @@ class TransformerConfig(ModelParallelConfig):
     """DEPRECATED and replaced by cuda_graph_impl.
     When set to true, TransformerLayer layers are swapped with user provided CUDA graphs."""
 
-    cuda_graph_impl: Literal['none', 'local', 'transformer_engine', 'full_iteration'] = "none"
+    cuda_graph_impl: Literal[
+        "none", "local", "transformer_engine", "full_iteration"
+    ] = "none"
     """Determines the CUDA graph capture implementation.
     "none": no CUDA graph.
     "local": MCore CUDA graph implementation. During training, graphable modules own per-layer
@@ -977,7 +1009,9 @@ class TransformerConfig(ModelParallelConfig):
     cuda_graph_modules has no effect when cuda_graph_impl="none" and must be empty when
     cuda_graph_impl="full_iteration"."""
 
-    cuda_graph_modules: Union[str, CudaGraphModule, List[str], List[CudaGraphModule]] = "full"
+    cuda_graph_modules: Union[
+        str, CudaGraphModule, List[str], List[CudaGraphModule]
+    ] = "full"
     """Selects training capture coverage within per-layer CUDA graphs (local and
     transformer_engine implementations).
     Valid values are "attn", "mlp", "moe", "moe_router", "moe_preprocess", and "mamba":
@@ -1018,7 +1052,10 @@ class TransformerConfig(ModelParallelConfig):
 
     cuda_graph_scope: Optional[
         Union[
-            str, CudaGraphModule, CudaGraphScope, List[Union[str, CudaGraphModule, CudaGraphScope]]
+            str,
+            CudaGraphModule,
+            CudaGraphScope,
+            List[Union[str, CudaGraphModule, CudaGraphScope]],
         ]
     ] = None
     """Deprecated: renamed to cuda_graph_modules. Accepted for backward compatibility and
@@ -1061,7 +1098,9 @@ class TransformerConfig(ModelParallelConfig):
     inference_sampling_seed: int = 42
     """ Random seed to use for sampling during inference. """
 
-    symmetric_ar_type: Optional[Literal['two_shot', "one_shot", "multimem_all_reduce"]] = None
+    symmetric_ar_type: Optional[
+        Literal["two_shot", "one_shot", "multimem_all_reduce"]
+    ] = None
     """What type of symmetric all reduce to use. The default is None
     which is no use of symmetric memory.
     """
@@ -1078,7 +1117,7 @@ class TransformerConfig(ModelParallelConfig):
     inference_disable_triton_nvls_kernels: bool = False
     """ If true, disables the use of Triton NVLS kernels during inference. """
 
-    inference_grouped_gemm_backend: Literal['flashinfer', 'torch', 'vllm'] = "vllm"
+    inference_grouped_gemm_backend: Literal["flashinfer", "torch", "vllm"] = "vllm"
     """Specifies the backend to use for grouped GEMM operations during inference.
     Options:
     - 'flashinfer': Uses FlashInfer cutlass_fused_moe. Not compatible with MXFP8.
@@ -1094,7 +1133,7 @@ class TransformerConfig(ModelParallelConfig):
     fp8_recipe='mxfp8'. Set to True to disable fusion and use separate kernel
     launches (useful for debugging)."""
 
-    inference_moe_token_dispatcher_type: Literal['nccl', 'nvls'] = 'nvls'
+    inference_moe_token_dispatcher_type: Literal["nccl", "nvls"] = "nvls"
     """Token dispatcher to use for MoE expert parallelism during inference.
     - 'nccl': AllGather/ReduceScatter via NCCL. Fixed token counts per rank; requires
       decode-only CUDA graphs (forced automatically).
@@ -1127,7 +1166,8 @@ class TransformerConfig(ModelParallelConfig):
     None causes the states to follow the activation dtype."""
 
     use_mamba_mem_eff_path: bool = field(
-        default=True, metadata={"argparse_meta": {"arg_names": ["--disable-mamba-mem-eff-path"]}}
+        default=True,
+        metadata={"argparse_meta": {"arg_names": ["--disable-mamba-mem-eff-path"]}},
     )
     """Controls usage of the memory efficient path for Mamba layers."""
 
@@ -1150,7 +1190,7 @@ class TransformerConfig(ModelParallelConfig):
     quant_recipe: Optional[RecipeConfig] = None
     """Configuration of any per-module quantization settings to be applied to the model"""
 
-    transformer_impl: Literal['local', 'transformer_engine', 'inference_optimized'] = (
+    transformer_impl: Literal["local", "transformer_engine", "inference_optimized"] = (
         "transformer_engine"
     )
     """Transformer implementation to use.
@@ -1278,26 +1318,26 @@ def __post_init__(self):
             )
 
         if self.experimental_attention_variant == "gated_delta_net":
-            assert (
-                self.linear_attention_freq is not None
-            ), f"linear_attention_freq must be set for linear gated_delta_net."
+            assert self.linear_attention_freq is not None, (
+                f"linear_attention_freq must be set for linear gated_delta_net."
+            )
 
             # Check required parameters
-            assert (
-                self.linear_conv_kernel_dim is not None
-            ), "linear_conv_kernel_dim must be set for gated delta net."
-            assert (
-                self.linear_key_head_dim is not None
-            ), "linear_key_head_dim must be set for gated delta net."
-            assert (
-                self.linear_value_head_dim is not None
-            ), "linear_value_head_dim must be set for gated delta net."
-            assert (
-                self.linear_num_key_heads is not None
-            ), "linear_num_key_heads must be set for gated delta net."
-            assert (
-                self.linear_num_value_heads is not None
-            ), "linear_num_value_heads must be set for gated delta net."
+            assert self.linear_conv_kernel_dim is not None, (
+                "linear_conv_kernel_dim must be set for gated delta net."
+            )
+            assert self.linear_key_head_dim is not None, (
+                "linear_key_head_dim must be set for gated delta net."
+            )
+            assert self.linear_value_head_dim is not None, (
+                "linear_value_head_dim must be set for gated delta net."
+            )
+            assert self.linear_num_key_heads is not None, (
+                "linear_num_key_heads must be set for gated delta net."
+            )
+            assert self.linear_num_value_heads is not None, (
+                "linear_num_value_heads must be set for gated delta net."
+            )
             assert self.linear_num_value_heads % self.linear_num_key_heads == 0, (
                 f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of "
                 f"linear_num_key_heads ({self.linear_num_key_heads})."
@@ -1333,7 +1373,9 @@ def __post_init__(self):
         if self.fp8:
             # cannot support first last layer bf16 with delayed scaling
             if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed:
-                raise ValueError("Delayed scaling does not support first / last layer in BF16.")
+                raise ValueError(
+                    "Delayed scaling does not support first / last layer in BF16."
+                )
 
             # max bf16 layers per pipeline stage
             max_bf16_layers_per_pipeline_stage = (
@@ -1344,7 +1386,8 @@ def __post_init__(self):
             if self.first_last_layers_bf16:
                 if (
                     self.num_layers_at_start_in_bf16 < 0
-                    or self.num_layers_at_start_in_bf16 > max_bf16_layers_per_pipeline_stage
+                    or self.num_layers_at_start_in_bf16
+                    > max_bf16_layers_per_pipeline_stage
                 ):
                     raise ValueError(
                         f"num_layers_at_start_in_bf16 ({self.num_layers_at_start_in_bf16}) must be "
@@ -1353,7 +1396,8 @@ def __post_init__(self):
                     )
                 if (
                     self.num_layers_at_end_in_bf16 < 0
-                    or self.num_layers_at_end_in_bf16 > max_bf16_layers_per_pipeline_stage
+                    or self.num_layers_at_end_in_bf16
+                    > max_bf16_layers_per_pipeline_stage
                 ):
                     raise ValueError(
                         f"num_layers_at_end_in_bf16 ({self.num_layers_at_end_in_bf16}) must be "
@@ -1377,7 +1421,8 @@ def __post_init__(self):
                 raise ValueError("fp8_output_proj must be used together with fp8 mode.")
             if self.fp8_recipe != Fp8Recipe.mxfp8:
                 raise ValueError(
-                    f"fp8_output_proj requires fp8_recipe='mxfp8', got " f"'{self.fp8_recipe}'."
+                    f"fp8_output_proj requires fp8_recipe='mxfp8', got "
+                    f"'{self.fp8_recipe}'."
                 )
 
         # FP4 validation
@@ -1385,7 +1430,9 @@ def __post_init__(self):
             raise ValueError("fp4_param must be used together with fp4 mode.")
 
         if self.fp4 and self.fp8:
-            raise ValueError("fp4 and fp8 cannot be used simultaneously. Please choose one.")
+            raise ValueError(
+                "fp4 and fp8 cannot be used simultaneously. Please choose one."
+            )
 
         if self.fp4 and self.fp4_recipe == Fp4Recipe.custom:
             if not self.fp4_quantizer_factory:
@@ -1401,13 +1448,18 @@ def __post_init__(self):
         if self.expert_model_parallel_size > 1 and self.num_moe_experts is None:
             raise ValueError("num_moe_experts must be non None to use expert-parallel.")
 
-        if self.transformer_impl == "inference_optimized" and self.num_moe_experts is not None:
+        if (
+            self.transformer_impl == "inference_optimized"
+            and self.num_moe_experts is not None
+        ):
             if self.expert_tensor_parallel_size > 1:
                 raise ValueError(
                     "Inference-optimized MoE layers does not support expert tensor parallelism."
                 )
             if self.moe_expert_capacity_factor is not None:
-                raise ValueError("Inference-optimized MoE layers only support dropless MoE ")
+                raise ValueError(
+                    "Inference-optimized MoE layers only support dropless MoE "
+                )
             if self.moe_router_padding_for_quantization:
                 raise ValueError(
                     "Inference-optimized MoE layers do not support padded "
@@ -1444,7 +1496,8 @@ def __post_init__(self):
                 )
 
             if (
-                self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.FLASHINFER
+                self.inference_grouped_gemm_backend
+                == InferenceGroupedGemmBackend.FLASHINFER
                 and self.fp8 == "mxfp8"
             ):
                 raise ValueError(
@@ -1466,7 +1519,9 @@ def __post_init__(self):
 
         if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None:
             self.moe_ffn_hidden_size = self.ffn_hidden_size
-            warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.")
+            warnings.warn(
+                "moe_ffn_hidden_size is not set, using ffn_hidden_size instead."
+            )
 
         if self.num_moe_experts is None and self.moe_ffn_hidden_size is not None:
             is_mixed_model = (
@@ -1514,9 +1569,13 @@ def __post_init__(self):
 
         if self.moe_enable_deepep:
             if self.moe_token_dispatcher_type != "flex":
-                raise ValueError("DeepEP backend is only supported with flex token dispatcher.")
+                raise ValueError(
+                    "DeepEP backend is only supported with flex token dispatcher."
+                )
             if self.moe_flex_dispatcher_backend == "hybridep":
-                raise ValueError("Only one backend is supported for flex token dispatcher.")
+                raise ValueError(
+                    "Only one backend is supported for flex token dispatcher."
+                )
             self.moe_flex_dispatcher_backend = "deepep"
             warnings.warn(
                 "moe_enable_deepep is deprecated."
@@ -1539,10 +1598,14 @@ def __post_init__(self):
                     f"num_shared_experts * ffn_size_of_each_shared_expert, "
                     f"but got {self.moe_shared_expert_intermediate_size}"
                 )
-            if self.moe_shared_expert_overlap and self.moe_token_dispatcher_type not in [
-                "alltoall",
-                "flex",
-            ]:
+            if (
+                self.moe_shared_expert_overlap
+                and self.moe_token_dispatcher_type
+                not in [
+                    "alltoall",
+                    "flex",
+                ]
+            ):
                 raise ValueError(
                     f"moe_shared_expert_overlap only works with alltoall or flex token dispatcher."
                 )
@@ -1600,7 +1663,8 @@ def __post_init__(self):
                 )
 
         if self.cpu_offloading and (
-            self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
+            self.cpu_offloading_num_layers < 0
+            or self.cpu_offloading_num_layers >= self.num_layers
         ):
             raise ValueError(
                 f"CPU offloading can be done only for layers less than {self.num_layers}"
@@ -1634,7 +1698,10 @@ def __post_init__(self):
                     'recompute_method must be "block" or "uniform"'
                 )
 
-            if self.recompute_granularity != "selective" and self.recompute_num_layers is None:
+            if (
+                self.recompute_granularity != "selective"
+                and self.recompute_num_layers is None
+            ):
                 raise ValueError(
                     f"When using recompute_granularity: {self.recompute_granularity} "
                     "recompute_num_layers must be between "
@@ -1642,7 +1709,8 @@ def __post_init__(self):
                     f"{self.num_layers // self.pipeline_model_parallel_size}"
                 )
             elif (
-                self.recompute_granularity == "selective" and self.recompute_num_layers is not None
+                self.recompute_granularity == "selective"
+                and self.recompute_num_layers is not None
             ):
                 raise ValueError(
                     f"When using recompute_granularity: {self.recompute_granularity} "
@@ -1681,7 +1749,10 @@ def __post_init__(self):
                     "moe_act in recompute_modules is only supported with moe_grouped_gemm."
                 )
 
-            if "mla_up_proj" in self.recompute_modules and not self.multi_latent_attention:
+            if (
+                "mla_up_proj" in self.recompute_modules
+                and not self.multi_latent_attention
+            ):
                 raise ValueError(
                     "mla_up_proj in recompute_modules is only supported with "
                     "multi_latent_attention."
@@ -1714,8 +1785,11 @@ def __post_init__(self):
                     )
 
             if self.fp8:
-                if "moe_act" in self.recompute_modules or "layernorm" in self.recompute_modules:
-                    if self.fp8_recipe == 'delayed':
+                if (
+                    "moe_act" in self.recompute_modules
+                    or "layernorm" in self.recompute_modules
+                ):
+                    if self.fp8_recipe == "delayed":
                         raise ValueError(
                             "Delayed scaling does not support moe_act and layernorm recompute "
                             "for fp8."
@@ -1741,9 +1815,9 @@ def __post_init__(self):
                 self.recompute_modules.append("moe")
 
         if self.fine_grained_activation_offloading:
-            assert (
-                not self.cpu_offloading
-            ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."
+            assert not self.cpu_offloading, (
+                "fine_grained_activation_offloading cannot be enabled with cpu_offloading."
+            )
             assert self.offload_modules is not None and len(self.offload_modules) > 0
             allowed_modules = {
                 "core_attn",
@@ -1757,16 +1831,22 @@ def __post_init__(self):
             }
             invalid_modules = set(self.offload_modules) - allowed_modules
             assert not invalid_modules, (
-                f'Invalid choices for offload_modules: {invalid_modules}. '
-                f'Allowed modules are: {allowed_modules}'
+                f"Invalid choices for offload_modules: {invalid_modules}. "
+                f"Allowed modules are: {allowed_modules}"
             )
-            if "attn_proj" in self.offload_modules and "core_attn" not in self.offload_modules:
+            if (
+                "attn_proj" in self.offload_modules
+                and "core_attn" not in self.offload_modules
+            ):
                 raise ValueError(
                     "attn_proj cannot be set to offload_modules alone without core_attn "
                     "because the input of attn_proj is the output of core_attn, "
                     "which is needed in core_attn.backward()."
                 )
-            if self.recompute_granularity == "selective" and "moe" in self.recompute_modules:
+            if (
+                self.recompute_granularity == "selective"
+                and "moe" in self.recompute_modules
+            ):
                 offload_inside_moe = {"moe_act", "expert_fc1", "fused_group_mlp"} & set(
                     self.offload_modules
                 )
@@ -1777,20 +1857,25 @@ def __post_init__(self):
                     f"Either remove 'moe' from --recompute-modules or remove "
                     f"{offload_inside_moe} from --offload-modules."
                 )
+            assert self.min_offloaded_tensor_size >= 0, (
+                "min_offloaded_tensor_size must be non-negative."
+            )
             assert (
-                self.min_offloaded_tensor_size >= 0
-            ), "min_offloaded_tensor_size must be non-negative."
-            assert (
-                self.activation_offload_fraction >= 0 and self.activation_offload_fraction <= 1
+                self.activation_offload_fraction >= 0
+                and self.activation_offload_fraction <= 1
             ), "activation_offload_fraction must be in range [0, 1]."
-            assert (
-                self.delta_offload_bytes_across_pp_ranks >= 0
-            ), "delta_offload_bytes_across_pp_ranks must be non-negative."
+            assert self.delta_offload_bytes_across_pp_ranks >= 0, (
+                "delta_offload_bytes_across_pp_ranks must be non-negative."
+            )
 
             if "fused_group_mlp" in self.offload_modules:
                 if not self.use_transformer_engine_op_fuser:
-                    raise ValueError("fused_group_mlp requires use_transformer_engine_op_fuser.")
-                moe_partial_offload = {"expert_fc1", "moe_act"} & set(self.offload_modules)
+                    raise ValueError(
+                        "fused_group_mlp requires use_transformer_engine_op_fuser."
+                    )
+                moe_partial_offload = {"expert_fc1", "moe_act"} & set(
+                    self.offload_modules
+                )
                 if moe_partial_offload:
                     raise ValueError(
                         "fused_group_mlp offloads the whole fused grouped MLP and cannot be "
@@ -1798,7 +1883,9 @@ def __post_init__(self):
                     )
         if self.moe_paged_stash:
             if self.cpu_offloading:
-                raise ValueError("moe_paged_stash cannot be enabled with cpu_offloading.")
+                raise ValueError(
+                    "moe_paged_stash cannot be enabled with cpu_offloading."
+                )
             if self.moe_expert_rank_capacity_factor is None:
                 raise ValueError(
                     "moe_paged_stash requires moe_expert_rank_capacity_factor to be set; "
@@ -1819,7 +1906,8 @@ def __post_init__(self):
             self.num_layers_in_first_pipeline_stage is not None
             or self.num_layers_in_last_pipeline_stage is not None
         ) and (
-            self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split
+            self.account_for_embedding_in_pipeline_split
+            or self.account_for_loss_in_pipeline_split
         ):
             raise ValueError(
                 "num_layers_in_first_pipeline_stage and num_layers_in_last_pipeline_stage cannot be"
@@ -1850,9 +1938,11 @@ def __post_init__(self):
             # Transfer pipeline_model_parallel_layout from str or list to
             # PipelineParallelLayerLayout
             if isinstance(self.pipeline_model_parallel_layout, str):
-                self.pipeline_model_parallel_layout = PipelineParallelLayerLayout.from_str(
-                    layout=self.pipeline_model_parallel_layout,
-                    pipeline_model_parallel_size=self.pipeline_model_parallel_size,
+                self.pipeline_model_parallel_layout = (
+                    PipelineParallelLayerLayout.from_str(
+                        layout=self.pipeline_model_parallel_layout,
+                        pipeline_model_parallel_size=self.pipeline_model_parallel_size,
+                    )
                 )
             elif isinstance(self.pipeline_model_parallel_layout, list):
                 # Since list is not hashable, the initialization will not be cached.
@@ -1876,8 +1966,10 @@ def __post_init__(self):
                 self.virtual_pipeline_model_parallel_size = detected_vpp_size
 
             # Check whether the layout is valid.
-            self.mtp_standalone = self.pipeline_model_parallel_layout.validate_layer_layout(
-                num_layers=self.num_layers, mtp_num_layers=self.mtp_num_layers
+            self.mtp_standalone = (
+                self.pipeline_model_parallel_layout.validate_layer_layout(
+                    num_layers=self.num_layers, mtp_num_layers=self.mtp_num_layers
+                )
             )
 
         # Uneven PP
@@ -1890,7 +1982,9 @@ def __post_init__(self):
 
             if self.num_layers_in_first_pipeline_stage is not None:
                 if self.num_layers_in_first_pipeline_stage <= 0:
-                    raise ValueError("num_layers_in_first_pipeline_stage must be larger than 0")
+                    raise ValueError(
+                        "num_layers_in_first_pipeline_stage must be larger than 0"
+                    )
 
                 if self.virtual_pipeline_model_parallel_size is not None:
                     if (
@@ -1909,7 +2003,9 @@ def __post_init__(self):
 
             if self.num_layers_in_last_pipeline_stage is not None:
                 if self.num_layers_in_last_pipeline_stage <= 0:
-                    raise ValueError("num_layers_in_last_pipeline_stage must be larger than 0")
+                    raise ValueError(
+                        "num_layers_in_last_pipeline_stage must be larger than 0"
+                    )
 
                 if self.virtual_pipeline_model_parallel_size is not None:
                     if (
@@ -1943,8 +2039,13 @@ def __post_init__(self):
 
             # If there are middle PP stages, check number of layers
             # on each middle PP rank is divisible by VPP size.
-            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
-                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
+            if (
+                pipeline_parallel_size
+                and self.virtual_pipeline_model_parallel_size is not None
+            ):
+                num_layers_per_middle_pipeline_rank = (
+                    num_layers // pipeline_parallel_size
+                )
                 if (
                     not num_layers_per_middle_pipeline_rank
                     % self.virtual_pipeline_model_parallel_size
@@ -1957,7 +2058,8 @@ def __post_init__(self):
                     )
 
         elif (
-            self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split
+            self.account_for_embedding_in_pipeline_split
+            or self.account_for_loss_in_pipeline_split
         ):
             if self.virtual_pipeline_model_parallel_size is None:
                 num_layers = self.num_layers
@@ -1990,9 +2092,12 @@ def __post_init__(self):
                         f"{self.pipeline_model_parallel_size}"
                     )
 
-                num_layers_per_pipeline_rank = num_layers // self.pipeline_model_parallel_size
+                num_layers_per_pipeline_rank = (
+                    num_layers // self.pipeline_model_parallel_size
+                )
                 if (
-                    not num_layers_per_pipeline_rank % self.virtual_pipeline_model_parallel_size
+                    not num_layers_per_pipeline_rank
+                    % self.virtual_pipeline_model_parallel_size
                     == 0
                 ):
                     raise ValueError(
@@ -2055,7 +2160,9 @@ def __post_init__(self):
 
         if self.activation_func_fp8_input_store:
             if self.activation_func != F.silu or not self.gated_linear_unit:
-                raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.")
+                raise ValueError(
+                    "Storing activation input in FP8 is supported only for SwiGLU."
+                )
 
         if self.apply_rope_fusion:
             if self.multi_latent_attention:
@@ -2076,33 +2183,46 @@ def __post_init__(self):
                     fused_apply_rotary_pos_emb_thd,
                 )
 
-                if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None:
+                if (
+                    fused_apply_rotary_pos_emb is None
+                    and fused_apply_rotary_pos_emb_thd is None
+                ):
                     raise ValueError(
                         "apply_rope_fusion is not available. Please install TE >= 1.4."
                     )
 
         if self.fused_single_qkv_rope:
             if self.attention_output_gate:
-                raise ValueError("fused_single_qkv_rope does not support gated attention for now.")
+                raise ValueError(
+                    "fused_single_qkv_rope does not support gated attention for now."
+                )
 
         if self.multi_latent_attention and self.rotary_interleaved:
-            raise ValueError("rotary_interleaved does not work with multi_latent_attention.")
+            raise ValueError(
+                "rotary_interleaved does not work with multi_latent_attention."
+            )
 
         # MuP (Maximal Update Parameterization) configuration
         if self.use_mup:
             # Default base_hidden_size to hidden_size (base model case, width_mult=1.0)
             if self.mup_base_hidden_size is None:
                 self.mup_base_hidden_size = self.hidden_size
-            assert self.mup_base_hidden_size > 0, "--mup-base-hidden-size must be positive."
+            assert self.mup_base_hidden_size > 0, (
+                "--mup-base-hidden-size must be positive."
+            )
             # Compute width multiplier
             self.mup_width_mult = self.hidden_size / self.mup_base_hidden_size
 
             # MuP attention scaling: 1/d_head instead of 1/sqrt(d_head).
             if self.softmax_scale is None:
                 base_head_scale = (
-                    1.0 if self.mup_base_head_dim is None else self.mup_base_head_dim**0.5
+                    1.0
+                    if self.mup_base_head_dim is None
+                    else self.mup_base_head_dim**0.5
+                )
+                self.softmax_scale = base_head_scale / (
+                    self.kv_channels**self.mup_attn_scale_power
                 )
-                self.softmax_scale = base_head_scale / (self.kv_channels**self.mup_attn_scale_power)
 
             # MuP output scaling: scale logits by 1/width_mult to keep outputs O(1).
             # Only auto-set if user hasn't explicitly configured it.
@@ -2135,10 +2255,14 @@ def __post_init__(self):
             self.embedding_init_method_std = self.init_method_std
 
         if self.embedding_init_method is None:
-            if self.init_method is None or (self.embedding_init_method_std != self.init_method_std):
+            if self.init_method is None or (
+                self.embedding_init_method_std != self.init_method_std
+            ):
                 # In this case, we set both the init method and the embedding init method to
                 #  whatever std value requested (or defaulted) for the embedding_init_layer
-                self.embedding_init_method = init_method_normal(self.embedding_init_method_std)
+                self.embedding_init_method = init_method_normal(
+                    self.embedding_init_method_std
+                )
             else:
                 # Replicate the current behavior where if you are not changing the std of the
                 #  embedding init differently and the init method is set, we fallback to the
@@ -2172,13 +2296,17 @@ def __post_init__(self):
                 )
 
         if self.num_moe_experts is not None and self.add_bias_linear:
-            assert (
-                self.expert_tensor_parallel_size == 1
-            ), "Bias in Moe is only supported when ETP==1"
+            assert 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 not in (
-            "sigmoid",
-            "sqrtsoftplus",
+        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' and 'sqrtsoftplus' "
@@ -2258,20 +2386,22 @@ def __post_init__(self):
             self.moe_router_num_groups = self.expert_model_parallel_size
 
         if self.enable_cuda_graph or self.external_cuda_graph:
-            assert (
-                self.cuda_graph_impl == "none"
-            ), "Do not use enable_cuda_graph or external_cuda_graph with cuda_graph_impl."
-            assert (
-                not self.enable_cuda_graph or not self.external_cuda_graph
-            ), "enable_cuda_graph and external_cuda_graph cannot be enabled at the same time."
+            assert self.cuda_graph_impl == "none", (
+                "Do not use enable_cuda_graph or external_cuda_graph with cuda_graph_impl."
+            )
+            assert not self.enable_cuda_graph or not self.external_cuda_graph, (
+                "enable_cuda_graph and external_cuda_graph cannot be enabled at the same time."
+            )
 
             if self.enable_cuda_graph:
-                warnings.warn('enable_cuda_graph is deprecated, use cuda_graph_impl=local instead.')
+                warnings.warn(
+                    "enable_cuda_graph is deprecated, use cuda_graph_impl=local instead."
+                )
                 self.cuda_graph_impl = "local"
             if self.external_cuda_graph:
                 warnings.warn(
-                    'external_cuda_graph is deprecated, '
-                    'use cuda_graph_impl=transformer_engine instead.'
+                    "external_cuda_graph is deprecated, "
+                    "use cuda_graph_impl=transformer_engine instead."
                 )
                 self.cuda_graph_impl = "transformer_engine"
 
@@ -2300,8 +2430,8 @@ def _scope_to_str(s):
                 self.cuda_graph_modules = _scope_to_str(scope)
             self.cuda_graph_scope = None
 
-        normalized_scopes, deprecated_scopes, used_full_scope = normalize_cuda_graph_modules(
-            self.cuda_graph_modules
+        normalized_scopes, deprecated_scopes, used_full_scope = (
+            normalize_cuda_graph_modules(self.cuda_graph_modules)
         )
         validate_deprecated_cuda_graph_modules_migration_inputs(
             deprecated_scopes, self.cuda_graph_impl, self.inference_cuda_graph_scope
@@ -2335,7 +2465,9 @@ def _scope_to_str(s):
         self.cuda_graph_modules = normalized_scopes
         assert all(
             isinstance(scope, CudaGraphModule) for scope in self.cuda_graph_modules
-        ), f"cuda_graph_modules must be a list of CudaGraphModule, got {self.cuda_graph_modules}."
+        ), (
+            f"cuda_graph_modules must be a list of CudaGraphModule, got {self.cuda_graph_modules}."
+        )
 
         assert self.cuda_graph_impl in [
             "none",
@@ -2348,7 +2480,10 @@ def _scope_to_str(s):
             self.inference_cuda_graph_scope, self.cuda_graph_impl
         )
 
-        assert self.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[self.cuda_graph_impl], (
+        assert (
+            self.inference_cuda_graph_scope
+            in ALLOWED_INFERENCE_SCOPES[self.cuda_graph_impl]
+        ), (
             "Invalid inference CUDA graph scope "
             f"{self.inference_cuda_graph_scope.name!r} for cuda_graph_impl="
             f"{self.cuda_graph_impl!r}."
@@ -2358,7 +2493,6 @@ def _scope_to_str(s):
         ), 'cuda_graph_modules must be empty when cuda_graph_impl="full_iteration".'
 
         if self.cuda_graph_impl != "none":
-
             if self.cpu_offloading and self.cuda_graph_impl != "full_iteration":
                 raise ValueError("CUDA graphs not supported with CPU offloading.")
 
@@ -2373,51 +2507,60 @@ def _scope_to_str(s):
                     ):
                         if CudaGraphModule.moe_router not in self.cuda_graph_modules:
                             self.cuda_graph_modules.append(CudaGraphModule.moe_router)
-                        if CudaGraphModule.moe_preprocess not in self.cuda_graph_modules:
-                            self.cuda_graph_modules.append(CudaGraphModule.moe_preprocess)
+                        if (
+                            CudaGraphModule.moe_preprocess
+                            not in self.cuda_graph_modules
+                        ):
+                            self.cuda_graph_modules.append(
+                                CudaGraphModule.moe_preprocess
+                            )
 
                 assert (
                     CudaGraphModule.moe not in self.cuda_graph_modules
                     or CudaGraphModule.moe_router not in self.cuda_graph_modules
-                ), 'cuda_graph_modules must not contain both moe and moe_router.'
+                ), "cuda_graph_modules must not contain both moe and moe_router."
                 if CudaGraphModule.moe_preprocess in self.cuda_graph_modules:
-                    assert (
-                        CudaGraphModule.moe_router in self.cuda_graph_modules
-                    ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.'
+                    assert CudaGraphModule.moe_router in self.cuda_graph_modules, (
+                        "moe_preprocess cuda graph is only supported with moe_router cuda graph."
+                    )
                 if self.num_moe_experts is None or self.num_moe_experts <= 1:
                     assert (
                         CudaGraphModule.moe not in self.cuda_graph_modules
                         and CudaGraphModule.moe_router not in self.cuda_graph_modules
-                    ), 'moe cuda graph is only supported for MoE.'
+                    ), "moe cuda graph is only supported for MoE."
                 else:
                     if self.moe_layer_freq == 1 or (
-                        isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq
+                        isinstance(self.moe_layer_freq, list)
+                        and 0 not in self.moe_layer_freq
                     ):
                         assert CudaGraphModule.mlp not in self.cuda_graph_modules, (
-                            'mlp cuda graph is only supported for dense layers, '
-                            'but not found in the model.'
+                            "mlp cuda graph is only supported for dense layers, "
+                            "but not found in the model."
                         )
                     if (
                         self.moe_expert_capacity_factor is None
                         or not self.moe_pad_expert_input_to_capacity
                     ):
-                        assert (
-                            CudaGraphModule.moe not in self.cuda_graph_modules
-                        ), 'moe cuda graph is only supported with drop-padding MoE.'
-                        if self.moe_token_dispatcher_type == 'alltoall' and (
+                        assert CudaGraphModule.moe not in self.cuda_graph_modules, (
+                            "moe cuda graph is only supported with drop-padding MoE."
+                        )
+                        if self.moe_token_dispatcher_type == "alltoall" and (
                             self.moe_expert_capacity_factor is not None
                             or self.moe_router_padding_for_fp8
                         ):
-                            assert CudaGraphModule.moe_preprocess not in self.cuda_graph_modules, (
-                                'moe_preprocess cuda graph is not supported when there are '
-                                'DtoH copies and synchronizations in the preprocess step.'
+                            assert (
+                                CudaGraphModule.moe_preprocess
+                                not in self.cuda_graph_modules
+                            ), (
+                                "moe_preprocess cuda graph is not supported when there are "
+                                "DtoH copies and synchronizations in the preprocess step."
                             )
 
             if self.recompute_granularity:
                 if self.recompute_granularity != "selective":
-                    assert (
-                        self.cuda_graph_impl == "full_iteration"
-                    ), "full recompute is only supported with full iteration CUDA graph."
+                    assert self.cuda_graph_impl == "full_iteration", (
+                        "full recompute is only supported with full iteration CUDA graph."
+                    )
                 else:
                     # The recompute module should be inside or outside of the graph scope.
                     # Recompute module coverring graph scope is not allowed.
@@ -2427,13 +2570,18 @@ def _scope_to_str(s):
                     ):
                         assert (
                             CudaGraphModule.moe_router not in self.cuda_graph_modules
-                        ), "moe recompute is not supported with moe_router CUDA graph with: "
+                        ), (
+                            "moe recompute is not supported with moe_router CUDA graph with: "
+                        )
                         "--cuda-graph-impl transformer_engine."
 
                     # Graphed recompute module doesn't accept random number.
                     # full_cudagraph means either full_iteration impl or an empty per-layer scope
                     # (which captures the whole layer).
-                    if self.cuda_graph_impl == "full_iteration" or not self.cuda_graph_modules:
+                    if (
+                        self.cuda_graph_impl == "full_iteration"
+                        or not self.cuda_graph_modules
+                    ):
                         full_cudagraph = True
                     else:
                         full_cudagraph = False
@@ -2458,7 +2606,9 @@ def _scope_to_str(s):
                                 and CudaGraphModule.moe not in self.cuda_graph_modules
                             )
                             or "moe" not in self.recompute_modules
-                        ), "hidden dropout is not supported with graphed MLP/MoE recomputation."
+                        ), (
+                            "hidden dropout is not supported with graphed MLP/MoE recomputation."
+                        )
                     if self.moe_input_jitter_eps is not None:
                         assert (
                             not full_cudagraph
@@ -2470,8 +2620,14 @@ def _scope_to_str(s):
             if self.fine_grained_activation_offloading:
                 offload_modules = set(self.offload_modules or [])
                 if self.cuda_graph_impl == "local":
-                    local_supported_offload_modules = {"expert_fc1", "moe_act", "fused_group_mlp"}
-                    unsupported_offload_modules = offload_modules - local_supported_offload_modules
+                    local_supported_offload_modules = {
+                        "expert_fc1",
+                        "moe_act",
+                        "fused_group_mlp",
+                    }
+                    unsupported_offload_modules = (
+                        offload_modules - local_supported_offload_modules
+                    )
                     assert not unsupported_offload_modules, (
                         "fine-grained activation offloading with cuda_graph_impl='local' "
                         "only supports offload_modules 'expert_fc1', 'moe_act', and "
@@ -2498,9 +2654,9 @@ def _scope_to_str(s):
                     "are supported only for expert_fc1, moe_act, or fused_group_mlp "
                     "offload when the full MoE module is not captured."
                 )
-                assert (
-                    CudaGraphModule.moe not in self.cuda_graph_modules
-                ), "Token-drop MoE is temporarily not supported with activation offloading."
+                assert CudaGraphModule.moe not in self.cuda_graph_modules, (
+                    "Token-drop MoE is temporarily not supported with activation offloading."
+                )
                 assert self.cuda_graph_warmup_steps > 0, (
                     "cuda_graph_warmup_steps must be greater than 0 when enabling "
                     "fine-grained activation offloading."
@@ -2538,51 +2694,55 @@ def _scope_to_str(s):
                 or fused_sort_chunks_by_index_with_probs is None
                 or fused_unpermute is None
             ):
-                raise ValueError("fused permutation is not available. Please install TE >= 2.1.0.")
+                raise ValueError(
+                    "fused permutation is not available. Please install TE >= 2.1.0."
+                )
 
         if self.overlap_moe_expert_parallel_comm:
             # TODO: remove this after we fix the hang issue with torch version < 2.6.0
-            assert is_torch_min_version(
-                "2.6.0"
-            ), "A2A Overlap encounters hang issue with torch version < 2.6.0"
+            assert is_torch_min_version("2.6.0"), (
+                "A2A Overlap encounters hang issue with torch version < 2.6.0"
+            )
             if self.pipeline_model_parallel_size > 1:
                 assert self.virtual_pipeline_model_parallel_size is not None, (
                     "If enabling EP A2A overlap, virtual_pipeline_model_parallel_size "
                     "must be specified when pipeline_model_parallel_size > 1"
                 )
             # Expert model parallelism requirements
-            assert (
-                self.expert_model_parallel_size > 1
-            ), 'overlap_moe_expert_parallel_comm is only supported with expert model parallelism'
+            assert self.expert_model_parallel_size > 1, (
+                "overlap_moe_expert_parallel_comm is only supported with expert model parallelism"
+            )
             assert self.moe_token_dispatcher_type in [
-                'alltoall',
-                'flex',
-            ], 'overlap_moe_expert_parallel_comm is supported with alltoall/flex token dispatcher'
+                "alltoall",
+                "flex",
+            ], (
+                "overlap_moe_expert_parallel_comm is supported with alltoall/flex token dispatcher"
+            )
 
-            assert (
-                self.recompute_granularity != 'full'
-            ), 'disable full recomputation when enabling overlap_moe_expert_parallel_comm'
-            assert (
-                self.recompute_method is None
-            ), 'disable recomputation method when enabling overlap_moe_expert_parallel_comm'
-            assert (
-                self.recompute_num_layers is None
-            ), 'recompute_num_layers must be None when enabling overlap_moe_expert_parallel_comm'
-            assert (
-                "moe" not in self.recompute_modules
-            ), 'disable moe in recompute_modules when enabling overlap_moe_expert_parallel_comm'
+            assert self.recompute_granularity != "full", (
+                "disable full recomputation when enabling overlap_moe_expert_parallel_comm"
+            )
+            assert self.recompute_method is None, (
+                "disable recomputation method when enabling overlap_moe_expert_parallel_comm"
+            )
+            assert self.recompute_num_layers is None, (
+                "recompute_num_layers must be None when enabling overlap_moe_expert_parallel_comm"
+            )
+            assert "moe" not in self.recompute_modules, (
+                "disable moe in recompute_modules when enabling overlap_moe_expert_parallel_comm"
+            )
 
             # Check if bf16 or fp16 is used
-            assert (
-                self.bf16 or self.fp16
-            ), 'overlap_moe_expert_parallel_comm is only supported with bf16 or fp16 model'
+            assert self.bf16 or self.fp16, (
+                "overlap_moe_expert_parallel_comm is only supported with bf16 or fp16 model"
+            )
 
-            assert (
-                not self.moe_shared_expert_overlap
-            ), 'disable moe_shared_expert_overlap when enabling overlap_moe_expert_parallel_comm'
-            assert (
-                self.mtp_num_layers is None or self.mtp_num_layers == 1
-            ), 'MTP layernum only supports 1 when enabling overlap_moe_expert_parallel_comm.'
+            assert not self.moe_shared_expert_overlap, (
+                "disable moe_shared_expert_overlap when enabling overlap_moe_expert_parallel_comm"
+            )
+            assert self.mtp_num_layers is None or self.mtp_num_layers == 1, (
+                "MTP layernum only supports 1 when enabling overlap_moe_expert_parallel_comm."
+            )
 
             if self.cuda_graph_impl != "none":
                 if self.cuda_graph_impl == "transformer_engine":
@@ -2590,38 +2750,38 @@ def _scope_to_str(s):
                         CudaGraphModule.moe not in self.cuda_graph_modules
                         and CudaGraphModule.mlp not in self.cuda_graph_modules
                     ), (
-                        'CUDA graph scope on moe and mlp is not '
-                        'supported with overlap_moe_expert_parallel_comm'
+                        "CUDA graph scope on moe and mlp is not "
+                        "supported with overlap_moe_expert_parallel_comm"
                     )
 
         # Check delay_wgrad_compute compatibility
         if self.delay_wgrad_compute:
-            assert (
-                self.overlap_moe_expert_parallel_comm
-            ), 'overlap_moe_expert_parallel_comm must be enabled when enabling delay_wgrad_compute'
+            assert self.overlap_moe_expert_parallel_comm, (
+                "overlap_moe_expert_parallel_comm must be enabled when enabling delay_wgrad_compute"
+            )
             if self.cuda_graph_impl == "transformer_engine":
                 assert is_te_min_version("2.10.0"), (
-                    'TE version >= 2.10.0 is required for delay_wgrad_compute with '
-                    'partial cuda graph'
+                    "TE version >= 2.10.0 is required for delay_wgrad_compute with "
+                    "partial cuda graph"
                 )
 
         if self.overlap_dispatch_backward_with_experts_wgrad:
             assert not self.overlap_moe_expert_parallel_comm, (
-                'overlap_moe_expert_parallel_comm must be disabled when enabling '
-                'overlap_dispatch_backward_with_experts_wgrad.'
+                "overlap_moe_expert_parallel_comm must be disabled when enabling "
+                "overlap_dispatch_backward_with_experts_wgrad."
+            )
+            assert is_te_min_version("2.3.0"), (
+                "TE version >= 2.3.0 is required for overlap_dispatch_backward_with_experts_wgrad"
             )
-            assert is_te_min_version(
-                "2.3.0"
-            ), 'TE version >= 2.3.0 is required for overlap_dispatch_backward_with_experts_wgrad'
             assert not self.delay_wgrad_compute, (
-                'delay_wgrad_compute and overlap_dispatch_backward_with_experts_wgrad '
-                'are mutually exclusive; use only one'
+                "delay_wgrad_compute and overlap_dispatch_backward_with_experts_wgrad "
+                "are mutually exclusive; use only one"
             )
 
         if self.ep_overlap_early_attn_memory_release:
             assert self.overlap_moe_expert_parallel_comm, (
-                'overlap_moe_expert_parallel_comm must be enabled when enabling '
-                'ep_overlap_early_attn_memory_release'
+                "overlap_moe_expert_parallel_comm must be enabled when enabling "
+                "ep_overlap_early_attn_memory_release"
             )
 
         if self.context_parallel_size > 1 and self.cp_comm_type is not None:
@@ -2631,14 +2791,14 @@ def _scope_to_str(s):
                     f"the total number of transformer layers ({self.num_layers})!"
                 )
             else:
-                assert isinstance(
-                    self.cp_comm_type, str
-                ), "Unsupported communication type for context parallelism!"
+                assert isinstance(self.cp_comm_type, str), (
+                    "Unsupported communication type for context parallelism!"
+                )
 
-        assert (
-            self.pipeline_model_parallel_size > 0
-        ), f"Pipeline model parallel size must be larger than 0 \
+        assert self.pipeline_model_parallel_size > 0, (
+            f"Pipeline model parallel size must be larger than 0 \
             when enable --standalone-embedding-stage and --standalone-loss-stage"
+        )
 
         if (
             self.num_moe_experts is not None
@@ -2654,10 +2814,14 @@ def _scope_to_str(s):
                 raise ImportError(
                     "packaging is not installed. Please install it with `pip install packaging`."
                 )
-            assert is_torch_min_version("2.7.0a0"), "Must have at least torch version 2.7 or higher"
+            assert is_torch_min_version("2.7.0a0"), (
+                "Must have at least torch version 2.7 or higher"
+            )
             assert is_te_min_version("2.3.0") or get_te_version() == PkgVersion(
                 "2.3.0.dev0+39c0e70"
-            ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce"
+            ), (
+                "Must have at least TE version 2.3 or higher to use symmetric memory all reduce"
+            )
 
         if self.no_rope_freq:
             assert not self.flash_decode, "flash_decode cannot be used with no_rope."
@@ -2684,7 +2848,9 @@ def _scope_to_str(s):
             assert not self.use_kitchen
 
         if self.experimental_attention_variant == "dsa":
-            assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention"
+            assert not self.apply_rope_fusion, (
+                "RoPE fusion is not supported for DSAttention"
+            )
             if self.context_parallel_size > 1:
                 cp_comm_types = (
                     self.cp_comm_type
@@ -2705,9 +2871,9 @@ def _scope_to_str(s):
                 "inference_fuse_tp_communication is only supported "
                 "for inference_optimized transformer implementation."
             )
-            assert (
-                self.num_moe_experts is None
-            ), "--inference-fuse-tp-communication is not supported for MoE models."
+            assert self.num_moe_experts is None, (
+                "--inference-fuse-tp-communication is not supported for MoE models."
+            )
 
         if self.inference_disable_triton_nvls_kernels:
             assert self.transformer_impl == "inference_optimized", (
@@ -2716,9 +2882,9 @@ def _scope_to_str(s):
             )
 
         if self.batch_invariant_mode:
-            assert (
-                self.attention_backend == AttnBackend.flash
-            ), "Batch invariant mode only supports FlashAttention"
+            assert self.attention_backend == AttnBackend.flash, (
+                "Batch invariant mode only supports FlashAttention"
+            )
 
 
 @dataclass
@@ -2789,13 +2955,17 @@ class MLATransformerConfig(TransformerConfig):
 
     def __post_init__(self):
         super().__post_init__()
-        if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn":
+        if (
+            self.multi_latent_attention
+            and self.apply_rope_fusion
+            and self.rope_type != "yarn"
+        ):
             raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.")
 
         if self.attention_output_gate:
             raise NotImplementedError("Output gate is not supported for MLA yet.")
 
         if self.cache_mla_latents:
-            assert (
-                self.apply_rope_fusion is False
-            ), "Rope Fusion is not compatible with caching latents"
+            assert self.apply_rope_fusion is False, (
+                "Rope Fusion is not compatible with caching latents"
+            )
diff --git a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py
index 0a454b5d7ff..d05ab24df38 100644
--- a/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py
+++ b/tests/unit_tests/models/test_experimental_attention_variant_module_specs.py
@@ -49,7 +49,9 @@ def _make_backend(fuse_layernorm=True):
     backend.linear.return_value = _FakeLinear
     backend.column_parallel_linear.return_value = _FakeColumnParallelLinear
     backend.row_parallel_linear.return_value = _FakeRowParallelLinear
-    backend.column_parallel_layer_norm_linear.return_value = _FakeLayerNormColumnParallelLinear
+    backend.column_parallel_layer_norm_linear.return_value = (
+        _FakeLayerNormColumnParallelLinear
+    )
     backend.fuse_layernorm_and_linear.return_value = fuse_layernorm
     backend.core_attention.return_value = _FakeCoreAttention
 
@@ -65,6 +67,7 @@ def _make_config(**overrides):
     defaults = dict(
         num_layers=4,
         normalization="RMSNorm",
+        norm_accuracy_compatible=False,
         qk_layernorm=False,
         multi_latent_attention=False,
         qk_l2_norm=False,
@@ -106,7 +109,12 @@ def _fn(variant):
 
     @pytest.mark.parametrize(
         "variant, expected",
-        [("gated_delta_net", True), ("dsa", False), (None, False), ("some_unknown_variant", False)],
+        [
+            ("gated_delta_net", True),
+            ("dsa", False),
+            (None, False),
+            ("some_unknown_variant", False),
+        ],
     )
     def test_variants(self, variant, expected):
         """Validate linear-attention variant classification across supported and unsupported names."""
@@ -198,7 +206,9 @@ def test_list_freq_wrong_length_raises(self):
     def test_none_for_non_linear_variant(self):
         """Verify non-linear variants default to all-standard attention when freq is None."""
         cfg = _make_config(
-            num_layers=4, linear_attention_freq=None, experimental_attention_variant="dsa"
+            num_layers=4,
+            linear_attention_freq=None,
+            experimental_attention_variant="dsa",
         )
         assert self._fn(cfg) == [0, 0, 0, 0]
 
@@ -294,7 +304,9 @@ def _call(self, cfg=None, backend=None):
         )
 
         if cfg is None:
-            cfg = _make_config(multi_latent_attention=True, qk_l2_norm=False, qk_layernorm=True)
+            cfg = _make_config(
+                multi_latent_attention=True, qk_l2_norm=False, qk_layernorm=True
+            )
         if backend is None:
             backend = _make_backend()
         return get_dsa_module_spec_for_backend(cfg, backend=backend)
@@ -332,7 +344,9 @@ def test_returns_absorbed_mla_self_attention_spec(self):
 
     def test_core_attention_is_dsa(self):
         """Verify MLA core_attention is wrapped with DSAttention."""
-        from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention
+        from megatron.core.transformer.experimental_attention_variant.dsa import (
+            DSAttention,
+        )
 
         spec = self._call()
         core = spec.submodules.core_attention
@@ -340,7 +354,9 @@ def test_core_attention_is_dsa(self):
 
     def test_dsa_indexer_structure(self):
         """Verify DSA indexer wiring uses expected backend linear/norm modules."""
-        from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexer
+        from megatron.core.transformer.experimental_attention_variant.dsa import (
+            DSAIndexer,
+        )
 
         spec = self._call()
         indexer = spec.submodules.core_attention.submodules.indexer
@@ -369,25 +385,48 @@ def test_qk_layernorm_enabled(self, normalization):
         assert spec.submodules.q_layernorm is spec.submodules.kv_layernorm
         backend.layer_norm.assert_any_call(rms_norm=expected_rms, for_qk=True)
 
+    def test_accuracy_compatible_qk_rmsnorm(self):
+        """Verify DSA q/kv norms can use the native Torch RMSNorm builder."""
+        from megatron.core.transformer.torch_norm import WrappedTorchNorm
+
+        backend = _make_backend()
+        cfg = _make_config(
+            multi_latent_attention=True,
+            qk_l2_norm=False,
+            qk_layernorm=True,
+            normalization="RMSNorm",
+            norm_accuracy_compatible=True,
+        )
+        spec = self._call(cfg=cfg, backend=backend)
+
+        assert spec.submodules.q_layernorm is WrappedTorchNorm
+        assert spec.submodules.kv_layernorm is WrappedTorchNorm
+
     def test_qk_layernorm_disabled(self):
         """Verify q/kv layernorm becomes IdentityOp, skipping backend.layer_norm for qk."""
         backend = _make_backend()
-        cfg = _make_config(multi_latent_attention=True, qk_l2_norm=False, qk_layernorm=False)
+        cfg = _make_config(
+            multi_latent_attention=True, qk_l2_norm=False, qk_layernorm=False
+        )
         spec = self._call(cfg=cfg, backend=backend)
         assert spec.submodules.q_layernorm is IdentityOp
         assert spec.submodules.kv_layernorm is IdentityOp
         # backend.layer_norm is still called for the indexer k_norm (for_qk=True at line 94),
         # but NOT for the outer qk_norm (line 105-107 takes the else branch).
         # Exactly one for_qk=True call should exist (from the indexer, not from qk_norm).
-        qk_calls = [c for c in backend.layer_norm.call_args_list if c.kwargs.get("for_qk")]
-        assert (
-            len(qk_calls) == 1
-        ), f"Expected 1 for_qk=True call (indexer only), got {len(qk_calls)}"
+        qk_calls = [
+            c for c in backend.layer_norm.call_args_list if c.kwargs.get("for_qk")
+        ]
+        assert len(qk_calls) == 1, (
+            f"Expected 1 for_qk=True call (indexer only), got {len(qk_calls)}"
+        )
 
     def test_linear_projections(self):
         """Verify Q/KV projection slots and backend.column_parallel_linear call count."""
         backend = _make_backend()
-        cfg = _make_config(multi_latent_attention=True, qk_l2_norm=False, qk_layernorm=True)
+        cfg = _make_config(
+            multi_latent_attention=True, qk_l2_norm=False, qk_layernorm=True
+        )
         spec = self._call(cfg=cfg, backend=backend)
         subs = spec.submodules
         assert subs.linear_q_proj == _FakeColumnParallelLinear
@@ -419,14 +458,18 @@ class TestGetExperimentalAttentionVariantModuleSpec:
     def test_dispatches_to_variant_handler(self, variant, target_fn):
         """Verify dispatcher routes each variant name to its corresponding builder function."""
         backend = _make_backend()
-        cfg = _make_config(experimental_attention_variant=variant, normalization="RMSNorm")
+        cfg = _make_config(
+            experimental_attention_variant=variant, normalization="RMSNorm"
+        )
         with patch(f"{self.MODULE}.{target_fn}") as mock_fn:
             mock_fn.return_value = ModuleSpec(module=MagicMock)
             from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
                 get_experimental_attention_variant_module_spec,
             )
 
-            result = get_experimental_attention_variant_module_spec(cfg, backend=backend)
+            result = get_experimental_attention_variant_module_spec(
+                cfg, backend=backend
+            )
             mock_fn.assert_called_once_with(config=cfg, backend=backend)
             assert result is mock_fn.return_value
 
@@ -451,12 +494,15 @@ class TestGetTransformerLayerWithExperimentalAttentionVariantSpec:
 
     def _make_attention_spec(self, fuse_input_layernorm=True):
         """Construct a mock attention spec with configurable fuse metadata."""
-        return ModuleSpec(module=MagicMock, metainfo={"fuse_input_layernorm": fuse_input_layernorm})
+        return ModuleSpec(
+            module=MagicMock, metainfo={"fuse_input_layernorm": fuse_input_layernorm}
+        )
 
     def _make_mlp_spec(self, fuse_pre_mlp_layernorm=True):
         """Construct a mock MLP spec with configurable fuse metadata."""
         return ModuleSpec(
-            module=MagicMock, metainfo={"fuse_pre_mlp_layernorm": fuse_pre_mlp_layernorm}
+            module=MagicMock,
+            metainfo={"fuse_pre_mlp_layernorm": fuse_pre_mlp_layernorm},
         )
 
     def test_all_experimental_no_moe(self):
@@ -480,7 +526,10 @@ def test_all_experimental_no_moe(self):
                 f"{self.MODULE}.get_experimental_attention_variant_module_spec",
                 return_value=attn_spec,
             ),
-            patch(f"{self.MODULE}._get_dense_mlp_module_spec", return_value=(mlp_spec, True)),
+            patch(
+                f"{self.MODULE}._get_dense_mlp_module_spec",
+                return_value=(mlp_spec, True),
+            ),
         ):
             specs = get_transformer_layer_with_experimental_attention_variant_spec(
                 cfg, backend=backend
@@ -516,8 +565,14 @@ def test_hybrid_attention_pattern(self):
                 f"{self.MODULE}.get_experimental_attention_variant_module_spec",
                 return_value=exp_attn_spec,
             ),
-            patch(f"{self.MODULE}._get_self_attention_module_spec", return_value=std_attn_spec),
-            patch(f"{self.MODULE}._get_dense_mlp_module_spec", return_value=(mlp_spec, True)),
+            patch(
+                f"{self.MODULE}._get_self_attention_module_spec",
+                return_value=std_attn_spec,
+            ),
+            patch(
+                f"{self.MODULE}._get_dense_mlp_module_spec",
+                return_value=(mlp_spec, True),
+            ),
         ):
             specs = get_transformer_layer_with_experimental_attention_variant_spec(
                 cfg, backend=backend
@@ -553,8 +608,13 @@ def test_hybrid_moe_pattern(self):
                 f"{self.MODULE}.get_experimental_attention_variant_module_spec",
                 return_value=attn_spec,
             ),
-            patch(f"{self.MODULE}._get_moe_module_spec", return_value=(moe_spec, False)),
-            patch(f"{self.MODULE}._get_dense_mlp_module_spec", return_value=(dense_spec, True)),
+            patch(
+                f"{self.MODULE}._get_moe_module_spec", return_value=(moe_spec, False)
+            ),
+            patch(
+                f"{self.MODULE}._get_dense_mlp_module_spec",
+                return_value=(dense_spec, True),
+            ),
         ):
             specs = get_transformer_layer_with_experimental_attention_variant_spec(
                 cfg, backend=backend
@@ -618,7 +678,8 @@ def test_get_transformer_block_with_experimental_attention_variant_spec(
         )
         backend = _make_backend()
         fake_layer_specs = [
-            ModuleSpec(module=TransformerLayer, submodules=MagicMock()) for _ in range(num_layers)
+            ModuleSpec(module=TransformerLayer, submodules=MagicMock())
+            for _ in range(num_layers)
         ]
 
         with (
@@ -639,17 +700,25 @@ def test_get_transformer_block_with_experimental_attention_variant_spec(
                 # Without explicit layout, slicing comes from offset + num_layers_to_build.
                 with (
                     patch(
-                        f"{self.MODULE}.get_transformer_layer_offset", return_value=offset
+                        f"{self.MODULE}.get_transformer_layer_offset",
+                        return_value=offset,
                     ) as mock_offset,
                     patch(
-                        f"{self.MODULE}.get_num_layers_to_build", return_value=num_layers_to_build
+                        f"{self.MODULE}.get_num_layers_to_build",
+                        return_value=num_layers_to_build,
                     ) as mock_num_layers,
                 ):
-                    result = get_transformer_block_with_experimental_attention_variant_spec(
-                        cfg, vp_stage=vp_stage, pp_rank=pp_rank
+                    result = (
+                        get_transformer_block_with_experimental_attention_variant_spec(
+                            cfg, vp_stage=vp_stage, pp_rank=pp_rank
+                        )
                     )
-                mock_offset.assert_called_once_with(cfg, vp_stage=vp_stage, pp_rank=pp_rank)
-                mock_num_layers.assert_called_once_with(cfg, vp_stage=vp_stage, pp_rank=pp_rank)
+                mock_offset.assert_called_once_with(
+                    cfg, vp_stage=vp_stage, pp_rank=pp_rank
+                )
+                mock_num_layers.assert_called_once_with(
+                    cfg, vp_stage=vp_stage, pp_rank=pp_rank
+                )
 
         assert isinstance(result, TransformerBlockSubmodules)
         assert result.layer_specs == [fake_layer_specs[i] for i in expected_ids]
diff --git a/tests/unit_tests/models/test_local_spec_provider_linear.py b/tests/unit_tests/models/test_local_spec_provider_linear.py
new file mode 100644
index 00000000000..4564ce0adf8
--- /dev/null
+++ b/tests/unit_tests/models/test_local_spec_provider_linear.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+"""LocalSpecProvider must expose a non-TE backend.linear() for DSA/MLA."""
+
+from megatron.core.extensions.transformer_engine import TELinear
+from megatron.core.models.backends import LocalSpecProvider
+from megatron.core.post_training.modelopt.layers import Linear
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear
+
+
+def test_local_spec_provider_linear_is_replicated_local_linear():
+    backend = LocalSpecProvider()
+    assert backend.linear() is Linear
+    assert backend.linear() is not TELinear
+    assert backend.column_parallel_linear() is ColumnParallelLinear
+    assert backend.linear() is not backend.column_parallel_linear()
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py
index 642aeeb126f..67ae7530216 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py
@@ -27,6 +27,7 @@
     DSAttention,
     DSAttentionSubmodules,
     FusedDSAIndexerLoss,
+    _AccuracyCompatibleSoftmax,
     _run_sparse_attention,
     _validate_nonpacked_cp_uniform_length,
     compute_dsa_indexer_loss,
@@ -68,6 +69,60 @@ def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor
     return x * scale
 
 
+class TestAccuracyCompatibleDSA:
+    """Test the opt-in full-score DSA alignment path."""
+
+    def test_explicit_softmax_backward_matches_formula(self):
+        logits = torch.randn(2, 3, 5, device="cuda", requires_grad=True)
+        valid_mask = torch.ones_like(logits, dtype=torch.bool)
+        valid_mask[..., -1] = False
+        grad_output = torch.randn_like(logits)
+
+        probabilities = _AccuracyCompatibleSoftmax.apply(logits, valid_mask)
+        probabilities.backward(grad_output)
+        expected = probabilities.detach() * (
+            grad_output - (grad_output * probabilities.detach()).sum(dim=-1, keepdim=True)
+        )
+        expected = expected.masked_fill(~valid_mask, 0.0)
+
+        assert torch.equal(logits.grad, expected)
+        assert torch.equal(probabilities[..., -1], torch.zeros_like(probabilities[..., -1]))
+
+    def test_accuracy_compatible_switch_defaults_off(self, monkeypatch):
+        query = torch.randn(8, 1, 2, 8, device="cuda", dtype=torch.bfloat16)
+        key = torch.randn_like(query)
+        value = torch.randn(8, 1, 2, 4, device="cuda", dtype=torch.bfloat16)
+        indices = torch.arange(8, device="cuda").view(1, 8, 1)
+        original = unfused_dsa_fn
+        calls = []
+
+        def capture(*args, **kwargs):
+            calls.append(kwargs.get("accuracy_compatible"))
+            return original(*args, **kwargs)
+
+        monkeypatch.setattr(
+            "megatron.core.transformer.experimental_attention_variant.dsa.unfused_dsa_fn",
+            capture,
+        )
+        common = dict(
+            absorbed_mla=False,
+            query=query,
+            key=key,
+            value=value,
+            up_v_weight=None,
+            topk_indices=indices,
+            softmax_scale=query.size(-1) ** -0.5,
+            mask=None,
+            varlen_starts=None,
+            varlen_ends=None,
+            key_positions=None,
+        )
+        _run_sparse_attention(config=SimpleNamespace(), **common)
+        _run_sparse_attention(config=SimpleNamespace(dsa_accuracy_compatible=True), **common)
+
+        assert calls == [False, True]
+
+
 class TestDSAIndexShareHelpers:
     """Test cross-layer top-k sharing helpers."""
 
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index 9f33dd01920..da05d34b937 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -62,6 +62,42 @@ def test_constructor(self):
         num_weights = sum([p.numel() for p in self.router.parameters()])
         assert num_weights == 12 * 4, num_weights
 
+    @pytest.mark.internal
+    def test_router_accuracy_compatible_gating(self):
+        hidden_states = torch.randn(
+            (3, 1, self.router.config.hidden_size), device="cuda", dtype=torch.bfloat16
+        )
+        self.router.config.router_accuracy_compatible = True
+
+        logits = self.router.gating(hidden_states)
+        expected = torch.mm(
+            hidden_states.reshape(-1, hidden_states.shape[-1]).float(),
+            self.router.weight.float().t(),
+        ).view(3, 1, -1)
+
+        assert logits.dtype == torch.float32
+        assert torch.equal(logits, expected)
+
+    @pytest.mark.internal
+    def test_default_router_gating_stays_native(self, monkeypatch):
+        expected = torch.randn((3, 1, self.router.config.num_moe_experts))
+        called = False
+
+        def fake_router_gating_linear(inp, weight, bias, router_dtype):
+            nonlocal called
+            called = True
+            return expected
+
+        monkeypatch.setattr(
+            "megatron.core.transformer.moe.router.router_gating_linear",
+            fake_router_gating_linear,
+        )
+        hidden_states = torch.randn((3, 1, self.router.config.hidden_size), dtype=torch.bfloat16)
+
+        assert self.router.config.router_accuracy_compatible is False
+        assert self.router.gating(hidden_states) is expected
+        assert called
+
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @pytest.mark.parametrize("moe_router_pre_softmax", [(True), (False)])
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index c3c3944e007..e4c422a5108 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -17,7 +17,9 @@
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec
 from megatron.core.models.hybrid.hybrid_model import HybridModel
-from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
+from megatron.core.num_microbatches_calculator import (
+    destroy_num_microbatches_calculator,
+)
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import get_context_parallel_group
 from megatron.core.process_groups_config import ProcessGroupCollection
@@ -29,10 +31,22 @@
     process_mtp_loss,
     roll_tensor,
 )
+from megatron.core.transformer.torch_norm import WrappedTorchNorm
 from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.utils import get_batch_on_this_cp_rank, is_te_min_version, unwrap_model
-from megatron.training.argument_utils import gpt_config_from_args, hybrid_config_from_args
-from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args
+from megatron.core.utils import (
+    get_batch_on_this_cp_rank,
+    is_te_min_version,
+    unwrap_model,
+)
+from megatron.training.argument_utils import (
+    gpt_config_from_args,
+    hybrid_config_from_args,
+)
+from megatron.training.arguments import (
+    core_transformer_config_from_args,
+    parse_args,
+    validate_args,
+)
 from megatron.training.checkpointing import load_checkpoint, save_checkpoint
 from megatron.training.global_vars import (
     destroy_global_vars,
@@ -45,7 +59,9 @@
 from tests.unit_tests.test_utilities import Utils
 
 if HAVE_TE:
-    from megatron.core.extensions.transformer_engine import TEColumnParallelGroupedLinear
+    from megatron.core.extensions.transformer_engine import (
+        TEColumnParallelGroupedLinear,
+    )
 else:
     TEColumnParallelGroupedLinear = None
 
@@ -54,7 +70,7 @@
 
 class TestMultiTokenPredictionLayer:
     def setup_method(self, method):
-        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+        os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
@@ -62,7 +78,9 @@ def teardown_method(self, method):
         destroy_num_microbatches_calculator()
 
     def _create_config_and_mtp_block_spec(self, tp, cp, use_te=False):
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
         config = TransformerConfig(
             mtp_num_layers=2,
             num_layers=4,
@@ -82,17 +100,52 @@ def _create_config_and_mtp_block_spec(self, tp, cp, use_te=False):
         )
         return config, mtp_block_spec
 
+    def test_accuracy_compatible_norms_override_te_mtp_norms(self):
+        """Accuracy mode routes all MTP-owned norms through native Torch RMSNorm."""
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=1
+        )
+        config = TransformerConfig(
+            mtp_num_layers=1,
+            num_layers=1,
+            hidden_size=64,
+            num_attention_heads=8,
+            normalization="RMSNorm",
+            norm_accuracy_compatible=True,
+            use_cpu_initialization=True,
+        )
+        transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec()
+        mtp_block_spec = get_gpt_mtp_block_spec(
+            config=config,
+            spec=transformer_layer_spec,
+            use_transformer_engine=True,
+        )
+        mtp_layer_spec = mtp_block_spec.layer_specs[0]
+
+        assert mtp_layer_spec.submodules.enorm is WrappedTorchNorm
+        assert mtp_layer_spec.submodules.hnorm is WrappedTorchNorm
+        assert mtp_layer_spec.submodules.layer_norm is WrappedTorchNorm
+        final_norm = mtp_layer_spec.submodules.layer_norm(
+            config=config, hidden_size=config.hidden_size, eps=config.layernorm_epsilon
+        )
+        assert isinstance(final_norm, torch.nn.RMSNorm)
+
     def test_mtp_detach_heads_config(self):
         """Test that mtp_detach_heads config defaults to False."""
         config = TransformerConfig(
-            num_layers=4, hidden_size=64, num_attention_heads=8, use_cpu_initialization=True
+            num_layers=4,
+            hidden_size=64,
+            num_attention_heads=8,
+            use_cpu_initialization=True,
         )
         assert config.mtp_detach_heads is False
 
     def test_constructor_with_detach_heads(self):
         """Test construction of MTP module with mtp_detach_heads=True."""
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=1
+        )
         config = TransformerConfig(
             mtp_num_layers=2,
             num_layers=4,
@@ -112,11 +165,11 @@ def test_constructor_with_detach_heads(self):
 
         # Verify all parameters are tagged for separate MTP grad-norm handling.
         for name, param in mtp.named_parameters():
-            assert (
-                getattr(param, 'grad_norm_group', None) == 'mtp'
-            ), f"Parameter {name} missing grad_norm_group attribute"
+            assert getattr(param, "grad_norm_group", None) == "mtp", (
+                f"Parameter {name} missing grad_norm_group attribute"
+            )
 
-    @pytest.mark.parametrize(('tp'), [(1), (2), (4)])
+    @pytest.mark.parametrize(("tp"), [(1), (2), (4)])
     def test_constructor_local(self, tp):
         """Test basic construction of MTP module."""
 
@@ -142,12 +195,16 @@ def test_constructor_local(self, tp):
             assert num_weights == 15216 * config.mtp_num_layers
 
     @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
-    @pytest.mark.parametrize(('tp', 'cp'), [(1, 1), (1, 2), (2, 1), (2, 2)])
+    @pytest.mark.parametrize(("tp", "cp"), [(1, 1), (1, 2), (2, 1), (2, 2)])
     def test_constructor_ues_te(self, tp, cp):
         """Test basic construction of MTP module."""
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
-        config, mtp_block_spec = self._create_config_and_mtp_block_spec(tp, cp, use_te=True)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
+        config, mtp_block_spec = self._create_config_and_mtp_block_spec(
+            tp, cp, use_te=True
+        )
         mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec)
 
         assert isinstance(mtp, MultiTokenPredictionBlock)
@@ -176,15 +233,22 @@ def test_get_embeddings_rolls_padding_mask(self):
 
         seq_len = 6
         batch_size = 2
-        input_ids = torch.tensor([[1, 2, 3, 4, 0, 0], [5, 6, 7, 0, 0, 0]], dtype=torch.int64)
+        input_ids = torch.tensor(
+            [[1, 2, 3, 4, 0, 0], [5, 6, 7, 0, 0, 0]], dtype=torch.int64
+        )
         position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1)
         padding_mask = torch.tensor(
-            [[True, True, True, True, False, False], [True, True, True, False, False, False]]
+            [
+                [True, True, True, True, False, False],
+                [True, True, True, False, False, False],
+            ]
         )
         hidden_states = torch.randn(seq_len, batch_size, config.hidden_size)
 
         def fake_embedding(input_ids, position_ids):
-            return torch.zeros(seq_len, batch_size, config.hidden_size, dtype=hidden_states.dtype)
+            return torch.zeros(
+                seq_len, batch_size, config.hidden_size, dtype=hidden_states.dtype
+            )
 
         rolled_input_ids, rolled_position_ids, rolled_padding_mask, _, _ = (
             mtp_layer._get_embeddings(
@@ -216,13 +280,17 @@ def test_forward_propagates_rolled_padding_mask(self, monkeypatch):
         batch_size = 2
         input_ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]], dtype=torch.int64)
         position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1)
-        padding_mask = torch.tensor([[True, True, True, False], [True, True, False, False]])
+        padding_mask = torch.tensor(
+            [[True, True, True, False], [True, True, False, False]]
+        )
         hidden_states = torch.randn(seq_len, batch_size, config.hidden_size)
         attention_mask = torch.ones((batch_size, 1, seq_len, seq_len), dtype=torch.bool)
         seen = {}
 
         def fake_embedding(input_ids, position_ids):
-            return torch.zeros(seq_len, batch_size, config.hidden_size, dtype=hidden_states.dtype)
+            return torch.zeros(
+                seq_len, batch_size, config.hidden_size, dtype=hidden_states.dtype
+            )
 
         def fake_proj_and_transformer_layer(
             self,
@@ -278,7 +346,9 @@ def test_get_embeddings_detaches_decoder_input(self):
         position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1)
         # hidden_states arrives without requires_grad (it is detached upstream by the block).
         hidden_states = torch.randn(seq_len, batch_size, config.hidden_size)
-        emb_weight = torch.nn.Parameter(torch.randn(seq_len, batch_size, config.hidden_size))
+        emb_weight = torch.nn.Parameter(
+            torch.randn(seq_len, batch_size, config.hidden_size)
+        )
 
         def fake_embedding(input_ids, position_ids):
             return emb_weight.clone()
@@ -323,8 +393,12 @@ def forward(self, hidden_states, **kwargs):
         seq_len = 4
         batch_size = 2
         input_ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]], dtype=torch.int64).cuda()
-        position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1).cuda()
-        attention_mask = torch.ones((batch_size, 1, seq_len, seq_len), dtype=torch.bool).cuda()
+        position_ids = (
+            torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1).cuda()
+        )
+        attention_mask = torch.ones(
+            (batch_size, 1, seq_len, seq_len), dtype=torch.bool
+        ).cuda()
         hidden_states = torch.randn(
             seq_len, batch_size, config.hidden_size, device="cuda", requires_grad=True
         )
@@ -359,7 +433,9 @@ def fake_embedding(input_ids, position_ids):
             # The returned block output still includes the original hidden-state
             # chunk, so autograd may allocate a zero grad for it through cat().
             if hidden_states.grad is not None:
-                torch.testing.assert_close(hidden_states.grad, torch.zeros_like(hidden_states))
+                torch.testing.assert_close(
+                    hidden_states.grad, torch.zeros_like(hidden_states)
+                )
             assert emb_weight.grad is None
         else:
             assert hidden_states.grad is not None
@@ -370,7 +446,9 @@ def test_process_mtp_loss_detaches_output_weight(self, detach_heads):
         """process_mtp_loss must detach the output-head weight when mtp_detach_heads=True
         so the MTP loss does not update the (shared) output projection weight."""
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=1
+        )
         config = TransformerConfig(
             mtp_num_layers=2,
             num_layers=4,
@@ -426,7 +504,7 @@ class TestMultiTokenPrediction:
     def setup_method(self, method):
         self.seq_length = 32
         self.micro_batch_size = 2
-        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+        os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
@@ -473,7 +551,7 @@ def create_test_args(
         destroy_global_vars()
         destroy_num_microbatches_calculator()
 
-        sys.argv = ['test_multi_token_predictioin.py']
+        sys.argv = ["test_multi_token_predictioin.py"]
         args = parse_args()
         args.num_layers = 2
         args.mtp_num_layers = 2
@@ -488,10 +566,10 @@ def create_test_args(
         args.tensor_model_parallel_size = tp
         args.sequence_parallel = True if tp > 1 else False
         args.context_parallel_size = cp
-        args.position_embedding_type = 'rope'
+        args.position_embedding_type = "rope"
         args.num_experts = 8
         args.train_iters = 1
-        args.ckpt_format = 'torch_dist'
+        args.ckpt_format = "torch_dist"
         args.moe_router_topk = 2
         args.moe_router_pre_softmax = False
         args.lr = 3e-5
@@ -507,10 +585,10 @@ def create_test_args(
             args.moe_grouped_gemm = False
         args.bf16 = True
         if fp8 is not None:
-            args.fp8 = 'e4m3'
+            args.fp8 = "e4m3"
         if full_recompute:
-            args.recompute_granularity = 'full'
-            args.recompute_method = 'uniform'
+            args.recompute_granularity = "full"
+            args.recompute_method = "uniform"
             args.recompute_num_layers = 1
         else:
             args.recompute_granularity = None
@@ -523,19 +601,26 @@ def create_test_args(
 
     def get_batch(self, seq_length, micro_batch_size):
         data = list(range(seq_length))
-        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
-        labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
-        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        input_ids = (
+            torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        )
+        labels = (
+            1
+            + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        )
+        position_ids = (
+            torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        )
         attention_mask = torch.ones(
             (micro_batch_size, 1, seq_length, seq_length), dtype=bool
         ).cuda()
         loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda()
         batch = {
-            'tokens': input_ids,
-            'labels': labels,
-            'loss_mask': loss_mask,
-            'attention_mask': attention_mask,
-            'position_ids': position_ids,
+            "tokens": input_ids,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
         }
         return batch
 
@@ -566,7 +651,9 @@ def get_packed_batch(self, seq_lengths, micro_batch_size):
         # Convert to tensors with shape [batch, total_seq_length]
         input_ids = torch.tensor(input_ids_list, dtype=torch.int64).unsqueeze(0).cuda()
         labels = torch.tensor(labels_list, dtype=torch.int64).unsqueeze(0).cuda()
-        position_ids = torch.tensor(position_ids_list, dtype=torch.int64).unsqueeze(0).cuda()
+        position_ids = (
+            torch.tensor(position_ids_list, dtype=torch.int64).unsqueeze(0).cuda()
+        )
 
         # Create attention mask for packed sequences (all ones for simplicity)
         attention_mask = torch.ones(
@@ -578,7 +665,8 @@ def get_packed_batch(self, seq_lengths, micro_batch_size):
 
         # Create cumulative sequence lengths for PackedSeqParams
         cu_seqlens = torch.tensor(
-            [0] + [sum(seq_lengths[: i + 1]) for i in range(len(seq_lengths))], dtype=torch.int32
+            [0] + [sum(seq_lengths[: i + 1]) for i in range(len(seq_lengths))],
+            dtype=torch.int32,
         ).cuda()
 
         packed_seq_params = PackedSeqParams(
@@ -586,16 +674,16 @@ def get_packed_batch(self, seq_lengths, micro_batch_size):
             cu_seqlens_kv=cu_seqlens,
             max_seqlen_q=max(seq_lengths),
             max_seqlen_kv=max(seq_lengths),
-            qkv_format='thd',
+            qkv_format="thd",
         )
 
         batch = {
-            'tokens': input_ids,
-            'labels': labels,
-            'loss_mask': loss_mask,
-            'attention_mask': attention_mask,
-            'position_ids': position_ids,
-            'packed_seq_params': packed_seq_params,
+            "tokens": input_ids,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "packed_seq_params": packed_seq_params,
         }
         return batch
 
@@ -609,7 +697,9 @@ def test_sharded_state_dict(self, tp, cp):
         args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
         set_args(args)
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
 
         model_parallel_cuda_manual_seed(_SEED)
         pg_collection = ProcessGroupCollection.use_mpu_process_groups()
@@ -638,7 +728,9 @@ def test_forward_backward(self, tmp_path_dist_ckpt, tp, cp, full_recompute):
         """Test MTP forward and backward with gptmodel."""
         tp_ref = 1
         cp_ref = 1
-        args = self.create_test_args(tp_ref, cp_ref, self.seq_length, self.micro_batch_size)
+        args = self.create_test_args(
+            tp_ref, cp_ref, self.seq_length, self.micro_batch_size
+        )
         set_args(args)
         torch.manual_seed(_SEED)
         Utils.initialize_model_parallel(
@@ -659,7 +751,7 @@ def test_forward_backward(self, tmp_path_dist_ckpt, tp, cp, full_recompute):
         tracker = MTPLossLoggingHelper.tracker
         mtp_loss_ref = None
         assert "loss_values" in tracker
-        mtp_loss_ref = tracker['loss_values'].clone()
+        mtp_loss_ref = tracker["loss_values"].clone()
         MTPLossLoggingHelper.clean_metrics_in_tracker()
 
         iteration = 123
@@ -670,7 +762,7 @@ def set_ckpt_path(ckpt_path):
             args.load = ckpt_path
 
         with TempNamedDir(
-            tmp_path_dist_ckpt / 'test_mtp_model_reconfiguration_model_A'
+            tmp_path_dist_ckpt / "test_mtp_model_reconfiguration_model_A"
         ) as ckpt_dir_A:
             set_ckpt_path(ckpt_dir_A)
             save_checkpoint(
@@ -687,12 +779,18 @@ def set_ckpt_path(ckpt_path):
             # Test with different TP/CP configuration
             Utils.destroy_model_parallel()
             args = self.create_test_args(
-                tp, cp, self.seq_length, self.micro_batch_size, full_recompute=full_recompute
+                tp,
+                cp,
+                self.seq_length,
+                self.micro_batch_size,
+                full_recompute=full_recompute,
             )
             set_args(args)
             set_ckpt_path(ckpt_dir_A)
             torch.manual_seed(_SEED)
-            Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+            Utils.initialize_model_parallel(
+                tensor_model_parallel_size=tp, context_parallel_size=cp
+            )
             gpt_model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
                 ModelType.encoder_or_decoder, self.model_provider
             )
@@ -702,7 +800,9 @@ def set_ckpt_path(ckpt_path):
             batch = get_batch_on_this_cp_rank(
                 batch, is_hybrid_cp=False, cp_group=get_context_parallel_group()
             )
-            tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values()
+            tokens, labels, loss_mask, attention_mask, position_ids, output_ref = (
+                batch.values()
+            )
             output = gpt_model[0].forward(
                 input_ids=tokens,
                 position_ids=position_ids,
@@ -712,9 +812,11 @@ def set_ckpt_path(ckpt_path):
             )
             tracker = MTPLossLoggingHelper.tracker
             assert "loss_values" in tracker
-            mtp_loss = tracker['loss_values'].clone()
+            mtp_loss = tracker["loss_values"].clone()
             # Average MTP loss across CP ranks for comparison with reference
-            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp'])
+            pg_collection = ProcessGroupCollection.use_mpu_process_groups(
+                required_pgs=["cp"]
+            )
             torch.distributed.all_reduce(
                 mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
             )
@@ -743,14 +845,21 @@ def test_fp8_support(self, full_recompute):
         """Test MTP with FP8 training enabled."""
         tp = 1
         cp = 1
-        fp8 = 'e4m3'
+        fp8 = "e4m3"
         args = self.create_test_args(
-            tp, cp, self.seq_length, self.micro_batch_size, fp8, full_recompute=full_recompute
+            tp,
+            cp,
+            self.seq_length,
+            self.micro_batch_size,
+            fp8,
+            full_recompute=full_recompute,
         )
         set_args(args)
 
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
         batch = self.get_batch(self.seq_length, self.micro_batch_size)
         tokens, labels, loss_mask, attention_mask, position_ids = batch.values()
         gpt_model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
@@ -765,7 +874,9 @@ def test_fp8_support(self, full_recompute):
             loss_mask=loss_mask,
         )
 
-        assert output.dtype == torch.float32  # Output should be converted back to float32
+        assert (
+            output.dtype == torch.float32
+        )  # Output should be converted back to float32
 
         loss = output.mean()
         loss.backward()
@@ -785,16 +896,18 @@ def test_packed_sequences(self, tp, cp):
         set_args(args)
 
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
 
         # Get packed batch
         batch = self.get_packed_batch(seq_lengths, micro_batch_size=1)
-        tokens = batch['tokens']
-        labels = batch['labels']
-        loss_mask = batch['loss_mask']
-        attention_mask = batch['attention_mask']
-        position_ids = batch['position_ids']
-        packed_seq_params = batch['packed_seq_params']
+        tokens = batch["tokens"]
+        labels = batch["labels"]
+        loss_mask = batch["loss_mask"]
+        attention_mask = batch["attention_mask"]
+        position_ids = batch["position_ids"]
+        packed_seq_params = batch["packed_seq_params"]
 
         # Create model
         model_parallel_cuda_manual_seed(_SEED)
@@ -824,7 +937,7 @@ def test_packed_sequences(self, tp, cp):
         # Verify MTP loss was computed
         tracker = MTPLossLoggingHelper.tracker
         assert "loss_values" in tracker
-        mtp_loss = tracker['loss_values'].clone()
+        mtp_loss = tracker["loss_values"].clone()
         assert mtp_loss.shape[0] == args.mtp_num_layers
         MTPLossLoggingHelper.clean_metrics_in_tracker()
 
@@ -856,12 +969,18 @@ def test_packed_sequences_with_full_recompute(self):
         total_seq_length = sum(seq_lengths)
 
         args = self.create_test_args(
-            tp=1, cp=1, sequence_length=total_seq_length, micro_batch_size=1, full_recompute=True
+            tp=1,
+            cp=1,
+            sequence_length=total_seq_length,
+            micro_batch_size=1,
+            full_recompute=True,
         )
         set_args(args)
 
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=1
+        )
 
         batch = self.get_packed_batch(seq_lengths, micro_batch_size=1)
 
@@ -876,12 +995,12 @@ def test_packed_sequences_with_full_recompute(self):
         )
 
         output = gpt_model[0].forward(
-            input_ids=batch['tokens'],
-            position_ids=batch['position_ids'],
-            attention_mask=batch['attention_mask'],
-            labels=batch['labels'],
-            loss_mask=batch['loss_mask'],
-            packed_seq_params=batch['packed_seq_params'],
+            input_ids=batch["tokens"],
+            position_ids=batch["position_ids"],
+            attention_mask=batch["attention_mask"],
+            labels=batch["labels"],
+            loss_mask=batch["loss_mask"],
+            packed_seq_params=batch["packed_seq_params"],
         )
 
         # Backward must run end-to-end through the recomputed MTP layer.
@@ -893,7 +1012,9 @@ def test_packed_sequences_with_full_recompute(self):
 
     def test_roll_tensor_none_input(self):
         """Test that roll_tensor returns (None, None) when given None input."""
-        Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=1
+        )
         result, sum_val = roll_tensor(None, shifts=-1, dims=-1)
         assert result is None
         assert sum_val is None
@@ -906,7 +1027,9 @@ def test_roll_tensor_shifts_left_and_zeroes_last(self):
         are not provided (RL training): label[i] = input_id[i+1], last position zeroed.
         The end-to-end derivation is covered by process_mtp_loss (see input_ids path).
         """
-        Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=1
+        )
         # Simulate input_ids [batch=2, seq=5]
         input_ids = torch.tensor(
             [[10, 20, 30, 40, 50], [60, 70, 80, 90, 100]], dtype=torch.int64
@@ -926,13 +1049,13 @@ def test_process_mtp_loss_skips_when_no_labels_and_no_input_ids(self):
             hidden_size=8, num_layers=2, num_attention_heads=2, mtp_num_layers=1
         )
         hidden_states = torch.ones(2, 1, 4)
-        called = {'value': False}
+        called = {"value": False}
 
         def output_layer(hidden, weight=None, runtime_gather_output=None):
             return hidden.clone(), None
 
         def compute_language_model_loss(mtp_labels, mtp_logits):
-            called['value'] = True
+            called["value"] = True
             return torch.ones_like(mtp_labels, dtype=mtp_logits.dtype)
 
         out = process_mtp_loss(
@@ -951,7 +1074,7 @@ def compute_language_model_loss(mtp_labels, mtp_logits):
         )
 
         # First chunk is returned unchanged and the loss is never computed.
-        assert not called['value']
+        assert not called["value"]
         assert torch.equal(out, torch.chunk(hidden_states, 2, dim=0)[0])
 
     def test_process_mtp_loss_derives_labels_from_input_ids(self):
@@ -967,13 +1090,13 @@ def test_process_mtp_loss_derives_labels_from_input_ids(self):
         # hidden_states is chunked into (1 + mtp_num_layers) along dim 0.
         hidden_states = torch.ones(2, 1, 5)
         input_ids = torch.tensor([[10, 20, 30, 40, 50]], dtype=torch.long)
-        seen = {'labels': None, 'masked_loss': None}
+        seen = {"labels": None, "masked_loss": None}
 
         def output_layer(hidden, weight=None, runtime_gather_output=None):
             return hidden.clone(), None
 
         def compute_language_model_loss(mtp_labels, mtp_logits):
-            seen['labels'] = mtp_labels.clone()
+            seen["labels"] = mtp_labels.clone()
             # Per-position loss of 1.0 so loss_mask * loss exposes the active mask.
             return torch.ones_like(mtp_labels, dtype=torch.float32)
 
@@ -994,8 +1117,10 @@ def compute_language_model_loss(mtp_labels, mtp_logits):
 
         # input_ids rolled twice (once to SFT format, once in the MTP layer loop):
         # [10,20,30,40,50] -> [20,30,40,50,0] -> [30,40,50,0,0].
-        assert seen['labels'] is not None, "loss should be computed in RL mode"
-        assert torch.equal(seen['labels'], torch.tensor([[30, 40, 50, 0, 0]], dtype=torch.long))
+        assert seen["labels"] is not None, "loss should be computed in RL mode"
+        assert torch.equal(
+            seen["labels"], torch.tensor([[30, 40, 50, 0, 0]], dtype=torch.long)
+        )
 
     @pytest.mark.parametrize("cp", [1, 2])
     def test_roll_tensor_with_packed_sequences(self, cp):
@@ -1004,9 +1129,13 @@ def test_roll_tensor_with_packed_sequences(self, cp):
         For CP=1: Tests standard packed sequence rolling with verified expected values
         For CP=2: Tests CP-enabled rolling executes without errors
         """
-        Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, context_parallel_size=cp
+        )
         cp_group = get_context_parallel_group() if cp > 1 else None
-        cp_rank = torch.distributed.get_rank(group=cp_group) if cp_group is not None else 0
+        cp_rank = (
+            torch.distributed.get_rank(group=cp_group) if cp_group is not None else 0
+        )
 
         if cp == 1:
             # Test case: Simple packed sequences (CP disabled)
@@ -1018,12 +1147,16 @@ def test_roll_tensor_with_packed_sequences(self, cp):
                 cu_seqlens_kv=cu_seqlens,
                 max_seqlen_q=3,
                 max_seqlen_kv=3,
-                qkv_format='thd',
+                qkv_format="thd",
             )
 
             # Roll by -1 (shift left)
             rolled, sum_val = roll_tensor(
-                tensor, shifts=-1, dims=0, cp_group=cp_group, packed_seq_params=packed_seq_params
+                tensor,
+                shifts=-1,
+                dims=0,
+                cp_group=cp_group,
+                packed_seq_params=packed_seq_params,
             )
 
             # Expected: [2, 3, 0, 5, 0] - boundaries at indices 2 and 4 are zeroed
@@ -1059,21 +1192,25 @@ def test_roll_tensor_with_packed_sequences(self, cp):
                 cu_seqlens_kv=cu_seqlens,
                 max_seqlen_q=6,  # max(4, 6) - max local seq length per sequence
                 max_seqlen_kv=6,
-                qkv_format='thd',
+                qkv_format="thd",
             )
 
             # Roll by -1 (shift left) with CP communication
             rolled, sum_val = roll_tensor(
-                tensor, shifts=-1, dims=0, cp_group=cp_group, packed_seq_params=packed_seq_params
+                tensor,
+                shifts=-1,
+                dims=0,
+                cp_group=cp_group,
+                packed_seq_params=packed_seq_params,
             )
 
             # Verify the rolled tensor matches expected values
-            assert (
-                rolled.shape == expected.shape
-            ), f"Shape mismatch: expected {expected.shape}, got {rolled.shape}"
-            assert torch.equal(
-                rolled, expected
-            ), f"CP Rank {cp_rank}: Expected\n{expected}\nbut got\n{rolled}\nDiff:\n{rolled - expected}"
+            assert rolled.shape == expected.shape, (
+                f"Shape mismatch: expected {expected.shape}, got {rolled.shape}"
+            )
+            assert torch.equal(rolled, expected), (
+                f"CP Rank {cp_rank}: Expected\n{expected}\nbut got\n{rolled}\nDiff:\n{rolled - expected}"
+            )
 
             # Verify sum is correct
             assert sum_val.numel() == 1, "Sum should be a scalar"
@@ -1123,10 +1260,22 @@ class DummyOutputLayer:
             def __init__(self, gather_output):
                 self.gather_output = gather_output
 
-        assert _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), None) is False
-        assert _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=False), None) is True
-        assert _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), True) is False
-        assert _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), False) is True
+        assert (
+            _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), None)
+            is False
+        )
+        assert (
+            _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=False), None)
+            is True
+        )
+        assert (
+            _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), True)
+            is False
+        )
+        assert (
+            _mtp_logits_are_vocab_sharded(DummyOutputLayer(gather_output=True), False)
+            is True
+        )
 
     def test_track_mtp_metrics(self):
         """Test tracking MTP metrics including acceptance rate."""
@@ -1137,7 +1286,11 @@ def test_track_mtp_metrics(self):
 
         for i in range(num_layers):
             MTPLossLoggingHelper.save_metrics_to_tracker(
-                loss=loss, correct=correct, total=total, layer_number=i, num_layers=num_layers
+                loss=loss,
+                correct=correct,
+                total=total,
+                layer_number=i,
+                num_layers=num_layers,
             )
 
         class DummyWriter:
@@ -1168,20 +1321,25 @@ def log(self, metrics, iteration):
         # Verify loss uses the legacy normalized MTP loss scaled by loss_scale.
         expected_loss = loss * loss_scale
         for i in range(num_layers):
-            assert f"mtp_{i+1} loss" in writer.scalars
-            assert torch.isclose(torch.as_tensor(writer.scalars[f"mtp_{i+1} loss"]), expected_loss)
-            assert torch.isclose(total_loss_dict[f"mtp_{i+1} loss"], expected_loss)
+            assert f"mtp_{i + 1} loss" in writer.scalars
+            assert torch.isclose(
+                torch.as_tensor(writer.scalars[f"mtp_{i + 1} loss"]), expected_loss
+            )
+            assert torch.isclose(total_loss_dict[f"mtp_{i + 1} loss"], expected_loss)
 
         # Verify acceptance rate is computed as (correct / total) * 100
         expected_rate = (correct / total) * 100.0
         for i in range(num_layers):
-            assert f"mtp_{i+1}_acceptance_rate" in writer.scalars
+            assert f"mtp_{i + 1}_acceptance_rate" in writer.scalars
             assert torch.isclose(
-                torch.as_tensor(writer.scalars[f"mtp_{i+1}_acceptance_rate"]), expected_rate
+                torch.as_tensor(writer.scalars[f"mtp_{i + 1}_acceptance_rate"]),
+                expected_rate,
             )
-            assert f"mtp_{i+1}_cumulative_acceptance_rate" in writer.scalars
+            assert f"mtp_{i + 1}_cumulative_acceptance_rate" in writer.scalars
             assert torch.isclose(
-                torch.as_tensor(writer.scalars[f"mtp_{i+1}_cumulative_acceptance_rate"]),
+                torch.as_tensor(
+                    writer.scalars[f"mtp_{i + 1}_cumulative_acceptance_rate"]
+                ),
                 expected_rate,
             )
 
@@ -1208,16 +1366,23 @@ def log(self, metrics, iteration):
         )
 
         expected_second_rate = (second_correct / second_total) * 100.0
-        expected_cumulative_rate = ((correct + second_correct) / (total + second_total)) * 100.0
+        expected_cumulative_rate = (
+            (correct + second_correct) / (total + second_total)
+        ) * 100.0
         for i in range(num_layers):
             assert torch.isclose(
-                torch.as_tensor(writer.scalars[f"mtp_{i+1}_acceptance_rate"]), expected_second_rate
+                torch.as_tensor(writer.scalars[f"mtp_{i + 1}_acceptance_rate"]),
+                expected_second_rate,
             )
             assert torch.isclose(
-                torch.as_tensor(writer.scalars[f"mtp_{i+1}_cumulative_acceptance_rate"]),
+                torch.as_tensor(
+                    writer.scalars[f"mtp_{i + 1}_cumulative_acceptance_rate"]
+                ),
                 expected_cumulative_rate,
             )
-            assert torch.isclose(total_loss_dict[f"mtp_{i+1} loss"], expected_loss * 2)
+            assert torch.isclose(
+                total_loss_dict[f"mtp_{i + 1} loss"], expected_loss * 2
+            )
 
         # Verify tracker is cleaned
         assert torch.all(MTPLossLoggingHelper.tracker["loss_values"] == 0)
@@ -1234,10 +1399,18 @@ def test_track_mtp_loss_preserves_legacy_normalized_loss_semantics(self):
         layer_number = 0
 
         MTPLossLoggingHelper.save_metrics_to_tracker(
-            loss=first_loss, correct=correct, total=total, layer_number=layer_number, num_layers=1
+            loss=first_loss,
+            correct=correct,
+            total=total,
+            layer_number=layer_number,
+            num_layers=1,
         )
         MTPLossLoggingHelper.save_metrics_to_tracker(
-            loss=second_loss, correct=correct, total=total, layer_number=layer_number, num_layers=1
+            loss=second_loss,
+            correct=correct,
+            total=total,
+            layer_number=layer_number,
+            num_layers=1,
         )
 
         class DummyWriter:
@@ -1265,7 +1438,7 @@ class TestMultiTokenPredictionHybrid:
     def setup_method(self, method):
         self.seq_length = 32
         self.micro_batch_size = 2
-        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+        os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
@@ -1308,7 +1481,7 @@ def create_test_args(
         destroy_global_vars()
         destroy_num_microbatches_calculator()
 
-        sys.argv = ['test_multi_token_prediction_hybrid.py']
+        sys.argv = ["test_multi_token_prediction_hybrid.py"]
         args = parse_args()
         args.mtp_num_layers = 2
         args.mtp_loss_scaling_factor = 0.1
@@ -1324,9 +1497,9 @@ def create_test_args(
         args.tensor_model_parallel_size = tp
         args.sequence_parallel = True if tp > 1 else False
         args.context_parallel_size = cp
-        args.position_embedding_type = 'rope'
+        args.position_embedding_type = "rope"
         args.train_iters = 1
-        args.ckpt_format = 'torch_dist'
+        args.ckpt_format = "torch_dist"
         args.lr = 3e-5
         args.attention_dropout = 0.0
         args.hidden_dropout = 0.0
@@ -1338,10 +1511,10 @@ def create_test_args(
         args.hybrid_layer_pattern = "M*M*/M*/M*"
 
         if fp8 is not None:
-            args.fp8 = 'e4m3'
+            args.fp8 = "e4m3"
         if full_recompute:
-            args.recompute_granularity = 'full'
-            args.recompute_method = 'uniform'
+            args.recompute_granularity = "full"
+            args.recompute_method = "uniform"
             args.recompute_num_layers = 1
         else:
             args.recompute_granularity = None
@@ -1354,19 +1527,26 @@ def create_test_args(
 
     def get_batch(self, seq_length, micro_batch_size):
         data = list(range(seq_length))
-        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
-        labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
-        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        input_ids = (
+            torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        )
+        labels = (
+            1
+            + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        )
+        position_ids = (
+            torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        )
         attention_mask = torch.ones(
             (micro_batch_size, 1, seq_length, seq_length), dtype=bool
         ).cuda()
         loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda()
         batch = {
-            'tokens': input_ids,
-            'labels': labels,
-            'loss_mask': loss_mask,
-            'attention_mask': attention_mask,
-            'position_ids': position_ids,
+            "tokens": input_ids,
+            "labels": labels,
+            "loss_mask": loss_mask,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
         }
         return batch
 
@@ -1377,7 +1557,9 @@ def test_sharded_state_dict_mamba(self, tp, cp):
         args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
         set_args(args)
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
 
         model_parallel_cuda_manual_seed(_SEED)
         pg_collection = ProcessGroupCollection.use_mpu_process_groups()
@@ -1401,7 +1583,9 @@ def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
         """Test MTP forward and backward with Mamba hybrid model."""
         tp_ref = 1
         cp_ref = 1
-        args = self.create_test_args(tp_ref, cp_ref, self.seq_length, self.micro_batch_size)
+        args = self.create_test_args(
+            tp_ref, cp_ref, self.seq_length, self.micro_batch_size
+        )
         set_args(args)
         torch.manual_seed(_SEED)
         Utils.initialize_model_parallel(
@@ -1430,7 +1614,7 @@ def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
         tracker = MTPLossLoggingHelper.tracker
         mtp_loss_ref = None
         assert "loss_values" in tracker
-        mtp_loss_ref = tracker['loss_values'].clone()
+        mtp_loss_ref = tracker["loss_values"].clone()
         MTPLossLoggingHelper.clean_metrics_in_tracker()
 
         iteration = 123
@@ -1440,7 +1624,9 @@ def set_ckpt_path(ckpt_path):
             args.save = ckpt_path
             args.load = ckpt_path
 
-        with TempNamedDir(tmp_path_dist_ckpt / 'test_mtp_mamba_model_reconfiguration') as ckpt_dir:
+        with TempNamedDir(
+            tmp_path_dist_ckpt / "test_mtp_mamba_model_reconfiguration"
+        ) as ckpt_dir:
             set_ckpt_path(ckpt_dir)
             save_checkpoint(
                 iteration,
@@ -1458,7 +1644,9 @@ def set_ckpt_path(ckpt_path):
             set_args(args)
             set_ckpt_path(ckpt_dir)
             torch.manual_seed(_SEED)
-            Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+            Utils.initialize_model_parallel(
+                tensor_model_parallel_size=tp, context_parallel_size=cp
+            )
 
             model_parallel_cuda_manual_seed(_SEED)
             cfg_container = Utils.pretrain_config_from_global_args(args, "hybrid")
@@ -1475,7 +1663,9 @@ def set_ckpt_path(ckpt_path):
             batch = get_batch_on_this_cp_rank(
                 batch, is_hybrid_cp=False, cp_group=get_context_parallel_group()
             )
-            tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values()
+            tokens, labels, loss_mask, attention_mask, position_ids, output_ref = (
+                batch.values()
+            )
             output = mamba_model[0].forward(
                 input_ids=tokens,
                 position_ids=position_ids,
@@ -1485,8 +1675,10 @@ def set_ckpt_path(ckpt_path):
             )
             tracker = MTPLossLoggingHelper.tracker
             assert "loss_values" in tracker
-            mtp_loss = tracker['loss_values'].clone()
-            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp'])
+            mtp_loss = tracker["loss_values"].clone()
+            pg_collection = ProcessGroupCollection.use_mpu_process_groups(
+                required_pgs=["cp"]
+            )
             torch.distributed.all_reduce(
                 mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
             )
@@ -1510,7 +1702,9 @@ def test_attention_mask_validation_mamba(self):
         args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
         set_args(args)
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp, context_parallel_size=cp
+        )
         pg_collection = ProcessGroupCollection.use_mpu_process_groups()
         model_cfg = hybrid_config_from_args(args)
         builder_cls = model_cfg.get_builder_cls()
@@ -1525,6 +1719,8 @@ def test_attention_mask_validation_mamba(self):
             assert mamba_model[0].mtp is not None
         except AssertionError as e:
             if "Multi-Token Prediction (MTP) is not yet supported" in str(e):
-                pytest.fail(f"Attention mask validation failed for Mamba hybrid model: {e}")
+                pytest.fail(
+                    f"Attention mask validation failed for Mamba hybrid model: {e}"
+                )
             else:
                 raise
diff --git a/tests/unit_tests/transformer/test_torch_norm.py b/tests/unit_tests/transformer/test_torch_norm.py
new file mode 100644
index 00000000000..8951eb2f365
--- /dev/null
+++ b/tests/unit_tests/transformer/test_torch_norm.py
@@ -0,0 +1,25 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+
+from megatron.core.transformer.torch_norm import WrappedTorchNorm
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+def _config(**overrides):
+    values = {
+        "num_layers": 1,
+        "hidden_size": 64,
+        "num_attention_heads": 4,
+        "normalization": "RMSNorm",
+    }
+    values.update(overrides)
+    return TransformerConfig(**values)
+
+
+def test_rmsnorm_uses_native_torch_implementation():
+    config = _config(norm_accuracy_compatible=True, params_dtype=torch.bfloat16)
+    norm = WrappedTorchNorm(config=config, hidden_size=64, eps=1e-5)
+
+    assert isinstance(norm, torch.nn.RMSNorm)
+    assert norm.weight.dtype == torch.bfloat16