From 25ad602ae34aa4a7eb49d1303d9b760a43a77ab3 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Wed, 29 Apr 2026 09:47:56 +0000
Subject: [PATCH 01/17] MTP support with mHC; new mHC contract

---
 megatron/core/models/gpt/gpt_layer_specs.py   |  12 +-
 megatron/core/models/gpt/gpt_model.py         |  12 +-
 megatron/core/transformer/hyper_connection.py |  17 ++
 .../transformer/multi_token_prediction.py     | 204 +++++++++++++-----
 .../core/transformer/transformer_block.py     |  33 ++-
 .../core/transformer/transformer_config.py    |   7 -
 6 files changed, 216 insertions(+), 69 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index 64faa80fce2..1a6e37f1faa 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -786,16 +786,10 @@ def get_gpt_mtp_block_spec_for_backend(
 
     transformer_layer_spec.submodules = copy.copy(transformer_layer_spec.submodules)
 
-    # MTP does not support hyper connections yet; strip HC modules and
-    # downgrade the layer class to plain TransformerLayer.
-    transformer_layer_spec.submodules.self_attention_hyper_connection = IdentityOp
-    transformer_layer_spec.submodules.cross_attention_hyper_connection = IdentityOp
-    transformer_layer_spec.submodules.mlp_hyper_connection = IdentityOp
-    if transformer_layer_spec.module is HyperConnectionTransformerLayer:
-        transformer_layer_spec.module = TransformerLayer
-
     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,
+        enable_hyper_connections=config.enable_hyper_connections,
     )
     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/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index ff63fcab9c2..422a0e71e6a 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -550,7 +550,7 @@ def forward(
             decoder_extra_block_kwargs['input_ids'] = input_ids
 
         # Run decoder.
-        hidden_states = self.decoder(
+        decoder_output = self.decoder(
             hidden_states=decoder_input,
             attention_mask=attention_mask,
             inference_context=inference_context,
@@ -563,6 +563,13 @@ def forward(
             padding_mask=padding_mask,
             **decoder_extra_block_kwargs,
         )
+        # When mHC + MTP, the decoder returns (contracted, multi-stream).
+        # MTP needs multi-stream; lm_head needs contracted.
+        if isinstance(decoder_output, tuple):
+            hidden_states, mhc_multistream = decoder_output
+        else:
+            hidden_states = decoder_output
+            mhc_multistream = None
 
         return self._postprocess(
             hidden_states=hidden_states,
@@ -582,6 +589,7 @@ def forward(
             runtime_gather_output=runtime_gather_output,
             extra_block_kwargs=extra_block_kwargs,
             inference_context=inference_context,
+            mhc_multistream=mhc_multistream,
         )
 
     def _postprocess(
@@ -603,6 +611,7 @@ def _postprocess(
         runtime_gather_output=None,
         extra_block_kwargs=None,
         inference_context=None,
+        mhc_multistream=None,
     ):
         """Postprocesses decoder hidden states to generate logits or compute loss.
 
@@ -631,6 +640,7 @@ def _postprocess(
                 input_ids=input_ids,
                 position_ids=position_ids,
                 hidden_states=hidden_states,
+                mhc_multistream=mhc_multistream,
                 attention_mask=attention_mask,
                 inference_params=None,  # MTP layers don't use KV cache
                 rotary_pos_emb=rotary_pos_emb,
diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py
index 64ec3107213..4d954f37d18 100644
--- a/megatron/core/transformer/hyper_connection.py
+++ b/megatron/core/transformer/hyper_connection.py
@@ -5,6 +5,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch import Tensor
 
 from megatron.core.transformer.module import MegatronModule
@@ -91,6 +92,22 @@ def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tenso
     return proj, r
 
 
+@torch.compile
+def learned_output_contract(
+    hidden_states: Tensor, head_fn: Tensor, base: Tensor, scale: Tensor, n: int, eps: float
+) -> Tensor:
+    """Learned output contraction: n-stream → 1-stream via sigmoid-gated weighted sum."""
+    dtype = hidden_states.dtype
+    hidden_states = hidden_states.to(torch.float32)
+    rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps)
+    mixes = F.linear(hidden_states, head_fn) * rsqrt
+    pre = torch.sigmoid(mixes * scale + base) + 1e-6
+    y = torch.sum(
+        pre.unsqueeze(-1) * hidden_states.view(*hidden_states.shape[:-1], n, -1), dim=-2
+    )
+    return y.to(dtype)
+
+
 # ============================================================================
 # HyperConnectionModule
 # ============================================================================
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index b751ef89cf0..f819aea7ddd 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Callable, List, Optional, Union
 
 import torch
+import torch.nn as nn
 from torch import Tensor
 
 from megatron.core import InferenceParams, parallel_state, tensor_parallel
@@ -27,6 +28,7 @@
     inference_all_gather_from_tensor_model_parallel_region,
 )
 from megatron.core.transformer.enums import AttnMaskType, LayerType
+from megatron.core.transformer.hyper_connection import learned_output_contract
 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
@@ -440,11 +442,15 @@ class MultiTokenPredictionLayerSubmodules:
     layer_norm: LayerNormBuilder
 
     eh_proj: Union[ModuleSpec, type] = None
+    e_proj: Union[ModuleSpec, type] = None
+    h_proj: Union[ModuleSpec, type] = None
     mtp_model_layer: Union[ModuleSpec, type] = None
 
 
 def get_mtp_layer_spec(
-    mtp_model_layer_spec: ModuleSpec, use_transformer_engine: bool
+    mtp_model_layer_spec: ModuleSpec,
+    use_transformer_engine: bool,
+    enable_hyper_connections: bool = False,
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -454,11 +460,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(),
+        enable_hyper_connections=enable_hyper_connections,
     )
 
 
 def get_mtp_layer_spec_for_backend(
-    mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
+    mtp_model_layer_spec: ModuleSpec,
+    backend: BackendSpecProvider,
+    enable_hyper_connections: bool = False,
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -467,15 +476,22 @@ def get_mtp_layer_spec_for_backend(
     """
     column_parallel_linear_impl: type = backend.column_parallel_linear()
     layer_norm_impl = backend.layer_norm()
+
+    submodules_kwargs = dict(
+        enorm=layer_norm_impl,
+        hnorm=layer_norm_impl,
+        mtp_model_layer=mtp_model_layer_spec,
+        layer_norm=layer_norm_impl,
+    )
+    if enable_hyper_connections:
+        submodules_kwargs["e_proj"] = column_parallel_linear_impl
+        submodules_kwargs["h_proj"] = column_parallel_linear_impl
+    else:
+        submodules_kwargs["eh_proj"] = column_parallel_linear_impl
+
     mtp_layer_spec = ModuleSpec(
         module=MultiTokenPredictionLayer,
-        submodules=MultiTokenPredictionLayerSubmodules(
-            enorm=layer_norm_impl,
-            hnorm=layer_norm_impl,
-            eh_proj=column_parallel_linear_impl,
-            mtp_model_layer=mtp_model_layer_spec,
-            layer_norm=layer_norm_impl,
-        ),
+        submodules=MultiTokenPredictionLayerSubmodules(**submodules_kwargs),
     )
     return mtp_layer_spec
 
@@ -816,6 +832,8 @@ def __init__(
                     f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
                 )
 
+        self.mhc_enabled = self.config.enable_hyper_connections
+
         self.enorm = self.submodules.enorm(
             config=self.config,
             hidden_size=self.config.hidden_size,
@@ -828,24 +846,58 @@ def __init__(
             eps=self.config.layernorm_epsilon,
         )
 
-        # For the linear projection at the (k - 1)-th MTP layer, the input is the concatenation
-        # of the i-th token's hidden states and the (i + K)-th token's decoder input,
-        # so the input's shape is [s, b, 2*h].
-        # The output will be send to the following transformer layer,
-        # so the output's shape should be [s, b, h].
-        self.eh_proj = build_module(
-            self.submodules.eh_proj,
-            self.config.hidden_size * 2,
-            self.config.hidden_size,
-            config=self.config,
-            init_method=self.config.init_method,
-            gather_output=False,
-            bias=False,
-            skip_bias_add=False,
-            is_expert=False,
-            tp_comm_buffer_name="mtp_eh_proj",
-            tp_group=pg_collection.tp if pg_collection is not None else None,
-        )
+        if self.mhc_enabled:
+            # mHC mode: separate e_proj and h_proj, operating per-stream.
+            # e_proj: [h] -> [h], applied to embedding then broadcast across streams.
+            # h_proj: [h] -> [h], applied per-stream on hidden states.
+            self.e_proj = build_module(
+                self.submodules.e_proj,
+                self.config.hidden_size,
+                self.config.hidden_size,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=False,
+                skip_bias_add=False,
+                is_expert=False,
+                tp_comm_buffer_name="mtp_e_proj",
+                tp_group=pg_collection.tp if pg_collection is not None else None,
+            )
+            self.h_proj = build_module(
+                self.submodules.h_proj,
+                self.config.hidden_size,
+                self.config.hidden_size,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=False,
+                skip_bias_add=False,
+                is_expert=False,
+                tp_comm_buffer_name="mtp_h_proj",
+                tp_group=pg_collection.tp if pg_collection is not None else None,
+            )
+            self.eh_proj = None
+        else:
+            # For the linear projection at the (k - 1)-th MTP layer, the input is the concatenation
+            # of the i-th token's hidden states and the (i + K)-th token's decoder input,
+            # so the input's shape is [s, b, 2*h].
+            # The output will be send to the following transformer layer,
+            # so the output's shape should be [s, b, h].
+            self.eh_proj = build_module(
+                self.submodules.eh_proj,
+                self.config.hidden_size * 2,
+                self.config.hidden_size,
+                config=self.config,
+                init_method=self.config.init_method,
+                gather_output=False,
+                bias=False,
+                skip_bias_add=False,
+                is_expert=False,
+                tp_comm_buffer_name="mtp_eh_proj",
+                tp_group=pg_collection.tp if pg_collection is not None else None,
+            )
+            self.e_proj = None
+            self.h_proj = None
 
         # Build inner layers: two possible paths
         # 1. Hybrid path: use HybridStack for hybrid pattern support
@@ -884,6 +936,14 @@ def __init__(
             hidden_size=self.config.hidden_size,
             eps=self.config.layernorm_epsilon,
         )
+
+        if self.mhc_enabled:
+            hc_mult = self.config.num_residual_streams
+            hc_dim = self.config.hidden_size * hc_mult
+            self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim))
+            self.hc_head_base = nn.Parameter(torch.zeros(hc_mult))
+            self.hc_head_scale = nn.Parameter(torch.ones(1))
+
         self.offload_context = nullcontext()
 
     def _get_embeddings(
@@ -937,25 +997,49 @@ def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.T
         """
         decoder_input = apply_module(self.enorm)(decoder_input)
         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)
-        # 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)
-        hidden_states, _ = self.eh_proj(hidden_states)
-        # For tensor parallel we need to gather the tensor across the model-parallel
-        # ranks after the linear projection.
-        if not self.training:
-            hidden_states = inference_all_gather_from_tensor_model_parallel_region(
-                hidden_states, self.tp_group, self.config
+
+        if self.mhc_enabled:
+            n = self.config.num_residual_streams
+            h = self.config.hidden_size
+            # hidden_states is [s, b, n*h] (multi-stream).
+            # hnorm operates per-stream on the h dimension.
+            s, b, _ = hidden_states.shape
+            hs_streams = hidden_states.view(s, b, n, h)
+            hs_streams = apply_module(self.hnorm)(hs_streams)
+            hs_streams = make_viewless_tensor(
+                inp=hs_streams, requires_grad=True, keep_graph=True
             )
+            # e_proj: [s, b, h] -> [s, b, h], then broadcast to [s, b, n, h]
+            e_out, _ = self.e_proj(decoder_input)
+            e_out = e_out.unsqueeze(2).expand(s, b, n, h)
+            # h_proj: applied per-stream on the h dimension
+            h_out, _ = self.h_proj(hs_streams)
+            # Combine and flatten back to [s, b, n*h]
+            hidden_states = (e_out + h_out).reshape(s, b, n * h)
         else:
-            hidden_states = gather_from_tensor_model_parallel_region(
-                hidden_states, group=self.tp_group
+            hidden_states = apply_module(self.hnorm)(hidden_states)
+            hidden_states = make_viewless_tensor(
+                inp=hidden_states, requires_grad=True, keep_graph=True
             )
-        # 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)
+            # 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)
+            hidden_states, _ = self.eh_proj(hidden_states)
+            # For tensor parallel we need to gather the tensor across the model-parallel
+            # ranks after the linear projection.
+            if not self.training:
+                hidden_states = inference_all_gather_from_tensor_model_parallel_region(
+                    hidden_states, self.tp_group, self.config
+                )
+            else:
+                hidden_states = gather_from_tensor_model_parallel_region(
+                    hidden_states, group=self.tp_group
+                )
+            # 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
+                )
         return hidden_states
 
     def _proj_and_transformer_layer(
@@ -1023,7 +1107,8 @@ def _proj_and_transformer_layer(
                         sequence_len_offset=sequence_len_offset,
                     )
 
-        hidden_states = self._postprocess(hidden_states)
+        if not self.mhc_enabled:
+            hidden_states = self._postprocess(hidden_states)
 
         return hidden_states
 
@@ -1032,6 +1117,16 @@ def _postprocess(self, hidden_states: torch.Tensor):
         Postprocesses the output of the transformer layers.
         """
 
+        if self.mhc_enabled:
+            hidden_states = learned_output_contract(
+                hidden_states,
+                self.hc_head_fn,
+                self.hc_head_base,
+                self.hc_head_scale,
+                self.config.num_residual_streams,
+                self.config.layernorm_epsilon,
+            )
+
         # Layer norm before shared head layer.
         hidden_states = apply_module(self.final_layernorm)(hidden_states)
         # TENorm produces a "viewed" tensor. This will result in schedule.py's
@@ -1584,6 +1679,7 @@ def forward(
         sequence_len_offset: Optional[Tensor] = None,
         extra_block_kwargs: Optional[dict] = None,
         embedding=None,
+        mhc_multistream: Optional[Tensor] = None,
     ) -> Tensor:
         """
         Perform the forward pass through all of the MTP modules.
@@ -1591,6 +1687,9 @@ def forward(
         Args:
             hidden_states (Tensor): Hidden states for input token with the shape [s, b, h]
                 where s is the sequence length, b is the batch size, and h is the hidden size.
+                Contracted decoder hidden states [s, b, h] when mHC is enabled.
+            mhc_multistream (Tensor, optional): When mHC is enabled, the pre-contraction
+                multi-stream decoder output [s, b, n*h] used as input to MTP depths.
             attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking
                 self-attention.
 
@@ -1600,7 +1699,12 @@ def forward(
         # get hidden states from previous mtp stages
         offset = get_mtp_layer_offset(self.config, self.vp_stage)
         hidden_states_list = list(torch.chunk(hidden_states, 1 + offset, dim=0))
-        hidden_states = hidden_states_list[offset]
+        if mhc_multistream is not None:
+            # mHC mode: use multi-stream for MTP depth input, contracted for loss list.
+            mhc_chunks = list(torch.chunk(mhc_multistream, 1 + offset, dim=0))
+            hidden_states = mhc_chunks[offset]
+        else:
+            hidden_states = hidden_states_list[offset]
         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) = self.layers[layer_idx](
@@ -1618,9 +1722,13 @@ def forward(
                 **(extra_block_kwargs or {}),
             )
 
-            # append the output hidden states of the current mtp layer
-            # to the hidden_states_list
-            hidden_states_list.append(hidden_states)
+            if mhc_multistream is not None:
+                mhc_chunks.append(hidden_states)
+                hidden_states_list.append(self.layers[layer_idx]._postprocess(hidden_states))
+            else:
+                # append the output hidden states of the current mtp layer
+                # to the hidden_states_list
+                hidden_states_list.append(hidden_states)
 
         # concat the hidden states of all mtp layers
         hidden_states = torch.cat(hidden_states_list, dim=0)
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index c991210b431..7a04c07095f 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -6,6 +6,7 @@
 from typing import List, Optional, Set, Tuple, Union, cast
 
 import torch
+import torch.nn as nn
 from torch import Tensor
 
 from megatron.core import parallel_state, tensor_parallel
@@ -22,7 +23,10 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import CheckpointManager
 from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType
-from megatron.core.transformer.hyper_connection import HyperConnectionModule
+from megatron.core.transformer.hyper_connection import (
+    HyperConnectionModule,
+    learned_output_contract,
+)
 from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.torch_norm import LayerNormBuilder
@@ -381,6 +385,12 @@ def build_layer(layer_spec, layer_number):
                 hidden_size=self.config.hidden_size,
                 eps=self.config.layernorm_epsilon,
             )
+            if self.config.enable_hyper_connections:
+                hc_mult = self.config.num_residual_streams
+                hc_dim = self.config.hidden_size * hc_mult
+                self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim))
+                self.hc_head_base = nn.Parameter(torch.zeros(hc_mult))
+                self.hc_head_scale = nn.Parameter(torch.ones(1))
         else:
             self.final_layernorm = None  # Either this or nn.Identity
 
@@ -926,10 +936,20 @@ def forward(
                         intermediate_hidden_states.append(hidden_states)
 
         # Only contract if the final layer norm is in this stage
+        mhc_multistream = None
         if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage():
-            hidden_states = HyperConnectionModule.output_contract(
-                hidden_states, self.num_residual_streams
-            )  # [s, b, n*C] -> [s, b, C]
+            # When MTP is enabled, save pre-contraction multi-stream for MTP input.
+            if self.config.mtp_num_layers is not None:
+                mhc_multistream = hidden_states
+            # [s, b, n*C] -> [s, b, C]
+            hidden_states = learned_output_contract(
+                hidden_states,
+                self.hc_head_fn,
+                self.hc_head_base,
+                self.hc_head_scale,
+                self.config.num_residual_streams,
+                self.config.layernorm_epsilon,
+            )
 
         # Final layer norm.
         if self.final_layernorm is not None:
@@ -949,6 +969,11 @@ def forward(
         if len(extract_layer_indices) > 0:
             return hidden_states, intermediate_hidden_states
 
+        # When mHC + MTP, return both contracted [s,b,h] (for lm_head) and
+        # pre-contraction multi-stream [s,b,n*h] (for MTP input).
+        if mhc_multistream is not None:
+            return hidden_states, mhc_multistream
+
         return hidden_states
 
     def sharded_state_dict(
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 83868dca9a2..1e7ff5a2823 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -1832,13 +1832,6 @@ def __post_init__(self):
                 )
                 self.use_fused_mhc = False
 
-        # Validation for hyper_connections with MTP
-        if self.enable_hyper_connections and self.mtp_num_layers is not None:
-            raise ValueError(
-                "enable_hyper_connections is not compatible with Multi-Token Prediction (MTP). "
-                "Please disable MTP (set mtp_num_layers=None) when using hyper connections."
-            )
-
         if self.fine_grained_activation_offloading:
             assert (
                 not self.cpu_offloading

From 2f363ae4e11032f5a3e9dfeb2538033cedd12dcc Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Wed, 29 Apr 2026 14:45:41 +0000
Subject: [PATCH 02/17] minor fix

---
 .../core/transformer/multi_token_prediction.py    | 15 +++++++++++++++
 megatron/core/transformer/transformer_block.py    |  8 ++++++++
 2 files changed, 23 insertions(+)

diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index f819aea7ddd..52f1e5e86ba 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -943,6 +943,11 @@ def __init__(
             self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim))
             self.hc_head_base = nn.Parameter(torch.zeros(hc_mult))
             self.hc_head_scale = nn.Parameter(torch.ones(1))
+            nn.init.xavier_uniform_(self.hc_head_fn)
+            if self.config.sequence_parallel:
+                setattr(self.hc_head_fn, 'sequence_parallel', True)
+                setattr(self.hc_head_base, 'sequence_parallel', True)
+                setattr(self.hc_head_scale, 'sequence_parallel', True)
 
         self.offload_context = nullcontext()
 
@@ -1011,11 +1016,21 @@ def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.T
             )
             # e_proj: [s, b, h] -> [s, b, h], then broadcast to [s, b, n, h]
             e_out, _ = self.e_proj(decoder_input)
+            e_out = gather_from_tensor_model_parallel_region(
+                e_out, group=self.tp_group
+            )
             e_out = e_out.unsqueeze(2).expand(s, b, n, h)
             # h_proj: applied per-stream on the h dimension
             h_out, _ = self.h_proj(hs_streams)
+            h_out = gather_from_tensor_model_parallel_region(
+                h_out, group=self.tp_group
+            )
             # Combine and flatten back to [s, b, n*h]
             hidden_states = (e_out + h_out).reshape(s, b, n * h)
+            if self.sequence_parallel:
+                hidden_states = scatter_to_sequence_parallel_region(
+                    hidden_states, group=self.tp_group
+                )
         else:
             hidden_states = apply_module(self.hnorm)(hidden_states)
             hidden_states = make_viewless_tensor(
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index 7a04c07095f..d0371596717 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -391,6 +391,11 @@ def build_layer(layer_spec, layer_number):
                 self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim))
                 self.hc_head_base = nn.Parameter(torch.zeros(hc_mult))
                 self.hc_head_scale = nn.Parameter(torch.ones(1))
+                nn.init.xavier_uniform_(self.hc_head_fn)
+                if self.config.sequence_parallel:
+                    setattr(self.hc_head_fn, 'sequence_parallel', True)
+                    setattr(self.hc_head_base, 'sequence_parallel', True)
+                    setattr(self.hc_head_scale, 'sequence_parallel', True)
         else:
             self.final_layernorm = None  # Either this or nn.Identity
 
@@ -940,6 +945,9 @@ def forward(
         if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage():
             # When MTP is enabled, save pre-contraction multi-stream for MTP input.
             if self.config.mtp_num_layers is not None:
+                assert (
+                    len(extract_layer_indices) == 0
+                ), "Feature extraction is not supported with mHC + MTP."
                 mhc_multistream = hidden_states
             # [s, b, n*C] -> [s, b, C]
             hidden_states = learned_output_contract(

From 0c99d96df1b39ddfcd7a96233f0d8f8856e55628 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Thu, 30 Apr 2026 03:34:39 +0000
Subject: [PATCH 03/17] fix new contract dtype

---
 megatron/core/transformer/hyper_connection.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py
index 4d954f37d18..c891ae53fe5 100644
--- a/megatron/core/transformer/hyper_connection.py
+++ b/megatron/core/transformer/hyper_connection.py
@@ -99,6 +99,9 @@ def learned_output_contract(
     """Learned output contraction: n-stream → 1-stream via sigmoid-gated weighted sum."""
     dtype = hidden_states.dtype
     hidden_states = hidden_states.to(torch.float32)
+    head_fn = head_fn.to(torch.float32)
+    base = base.to(torch.float32)
+    scale = scale.to(torch.float32)
     rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps)
     mixes = F.linear(hidden_states, head_fn) * rsqrt
     pre = torch.sigmoid(mixes * scale + base) + 1e-6

From 5b27dc99867e0b9a3ed952fd97122f663180d29a Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Thu, 30 Apr 2026 06:09:10 +0000
Subject: [PATCH 04/17] format and add tests

---
 megatron/core/transformer/hyper_connection.py |   4 +-
 .../transformer/multi_token_prediction.py     |  12 +-
 .../test_multi_token_prediction.py            | 275 ++++++++++++++++++
 3 files changed, 279 insertions(+), 12 deletions(-)

diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py
index c891ae53fe5..900a24ea4a6 100644
--- a/megatron/core/transformer/hyper_connection.py
+++ b/megatron/core/transformer/hyper_connection.py
@@ -105,9 +105,7 @@ def learned_output_contract(
     rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps)
     mixes = F.linear(hidden_states, head_fn) * rsqrt
     pre = torch.sigmoid(mixes * scale + base) + 1e-6
-    y = torch.sum(
-        pre.unsqueeze(-1) * hidden_states.view(*hidden_states.shape[:-1], n, -1), dim=-2
-    )
+    y = torch.sum(pre.unsqueeze(-1) * hidden_states.view(*hidden_states.shape[:-1], n, -1), dim=-2)
     return y.to(dtype)
 
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 52f1e5e86ba..8cbc45a2ebd 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -1011,20 +1011,14 @@ def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.T
             s, b, _ = hidden_states.shape
             hs_streams = hidden_states.view(s, b, n, h)
             hs_streams = apply_module(self.hnorm)(hs_streams)
-            hs_streams = make_viewless_tensor(
-                inp=hs_streams, requires_grad=True, keep_graph=True
-            )
+            hs_streams = make_viewless_tensor(inp=hs_streams, requires_grad=True, keep_graph=True)
             # e_proj: [s, b, h] -> [s, b, h], then broadcast to [s, b, n, h]
             e_out, _ = self.e_proj(decoder_input)
-            e_out = gather_from_tensor_model_parallel_region(
-                e_out, group=self.tp_group
-            )
+            e_out = gather_from_tensor_model_parallel_region(e_out, group=self.tp_group)
             e_out = e_out.unsqueeze(2).expand(s, b, n, h)
             # h_proj: applied per-stream on the h dimension
             h_out, _ = self.h_proj(hs_streams)
-            h_out = gather_from_tensor_model_parallel_region(
-                h_out, group=self.tp_group
-            )
+            h_out = gather_from_tensor_model_parallel_region(h_out, group=self.tp_group)
             # Combine and flatten back to [s, b, n*h]
             hidden_states = (e_out + h_out).reshape(s, b, n * h)
             if self.sequence_parallel:
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index c042545f2bf..f3437292374 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -5,6 +5,7 @@
 
 import pytest
 import torch
+from torch import Tensor
 
 from megatron.core.enums import ModelType
 from megatron.core.extensions.transformer_engine import HAVE_TE
@@ -21,11 +22,13 @@
 from megatron.core.parallel_state import get_context_parallel_group
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.hyper_connection import learned_output_contract
 from megatron.core.transformer.multi_token_prediction import (
     MTPLossLoggingHelper,
     MultiTokenPredictionBlock,
     roll_tensor,
 )
+from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
 from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args
@@ -1071,3 +1074,275 @@ def test_attention_mask_validation_mamba(self):
                 pytest.fail(f"Attention mask validation failed for Mamba hybrid model: {e}")
             else:
                 raise
+
+
+class TestLearnedOutputContract:
+    """Tests for learned_output_contract: shape, dtype, gradient, and numerical correctness."""
+
+    def setup_method(self, method):
+        Utils.initialize_model_parallel(1, 1)
+        model_parallel_cuda_manual_seed(_SEED)
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    def test_shape_and_dtype(self):
+        """Output shape is [*, h] from [*, n*h]; dtype matches input after fp32 round-trip."""
+        seq_len, batch_size, hidden_size, n_streams = 16, 2, 64, 4
+        head_fn = torch.randn(n_streams, n_streams * hidden_size, device='cuda')
+        base = torch.zeros(n_streams, device='cuda')
+        scale = torch.ones(1, device='cuda')
+
+        for dtype in [torch.bfloat16, torch.float16]:
+            hidden_states = torch.randn(
+                seq_len, batch_size, n_streams * hidden_size, device='cuda', dtype=dtype
+            )
+            output = learned_output_contract(
+                hidden_states, head_fn, base, scale, n_streams, eps=1e-6
+            )
+            assert output.shape == (seq_len, batch_size, hidden_size)
+            assert output.dtype == dtype
+
+    def test_gradient_and_numerical_correctness(self):
+        """Gradients flow to all inputs; output matches reference implementation."""
+        torch.manual_seed(_SEED)
+        seq_len, batch_size, hidden_size, n_streams = 2, 1, 8, 2
+        eps = 1e-6
+        hidden_states = torch.randn(
+            seq_len,
+            batch_size,
+            n_streams * hidden_size,
+            device='cuda',
+            dtype=torch.float32,
+            requires_grad=True,
+        )
+        head_fn = torch.randn(n_streams, n_streams * hidden_size, device='cuda', requires_grad=True)
+        base = torch.zeros(n_streams, device='cuda', requires_grad=True)
+        scale = torch.ones(1, device='cuda', requires_grad=True)
+
+        output = learned_output_contract(hidden_states, head_fn, base, scale, n_streams, eps)
+
+        # Numerical reference
+        hs_fp32 = hidden_states.detach().clone()
+        rsqrt_ref = torch.rsqrt(hs_fp32.square().mean(-1, keepdim=True) + eps)
+        mixes_ref = torch.nn.functional.linear(hs_fp32, head_fn.detach()) * rsqrt_ref
+        pre_ref = torch.sigmoid(mixes_ref * scale.detach() + base.detach()) + 1e-6
+        y_ref = torch.sum(
+            pre_ref.unsqueeze(-1) * hs_fp32.view(*hs_fp32.shape[:-1], n_streams, -1), dim=-2
+        )
+        torch.testing.assert_close(output, y_ref, rtol=1e-4, atol=1e-4)
+
+        # Gradient flow
+        output.sum().backward()
+        for name, tensor in [
+            ("hidden_states", hidden_states),
+            ("head_fn", head_fn),
+            ("base", base),
+            ("scale", scale),
+        ]:
+            assert tensor.grad is not None, f"No gradient for {name}"
+            assert not torch.all(tensor.grad == 0), f"Zero gradient for {name}"
+
+
+class TestMHCMTPIntegration:
+    """Integration tests for mHC + MTP: constructor, TransformerBlock output, E2E."""
+
+    def setup_method(self, method):
+        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+        MTPLossLoggingHelper.tracker = {}
+
+    @pytest.mark.parametrize('tp', [1, 2])
+    def test_mtp_constructor_with_mhc(self, tp):
+        """MTP layers have e_proj/h_proj (not eh_proj) and learned contraction params."""
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(tensor_model_parallel_size=tp)
+        config = TransformerConfig(
+            mtp_num_layers=2,
+            num_layers=4,
+            hidden_size=64,
+            num_attention_heads=8,
+            num_residual_streams=4,
+            enable_hyper_connections=True,
+            use_cpu_initialization=True,
+            tensor_model_parallel_size=tp,
+            sequence_parallel=True if tp > 1 else False,
+        )
+        spec = get_gpt_layer_local_spec(enable_hyper_connection=True)
+        mtp_block_spec = get_gpt_mtp_block_spec(
+            config=config, spec=spec, use_transformer_engine=False
+        )
+        mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec)
+
+        n, h = config.num_residual_streams, config.hidden_size
+        for i in range(config.mtp_num_layers):
+            layer = mtp.layers[i]
+            assert layer.e_proj is not None and layer.h_proj is not None
+            assert layer.eh_proj is None
+            assert layer.e_proj.weight.shape == (h // tp, h)
+            assert layer.h_proj.weight.shape == (h // tp, h)
+            assert layer.hc_head_fn.shape == (n, n * h)
+            assert layer.hc_head_base.shape == (n,)
+            assert layer.hc_head_scale.shape == (1,)
+            if tp > 1:
+                assert getattr(layer.hc_head_fn, 'sequence_parallel', False)
+
+    def test_transformer_block_returns_tuple(self):
+        """With mHC+MTP the block returns (contracted, multistream); without MTP just a tensor."""
+        Utils.initialize_model_parallel(1, 1)
+        model_parallel_cuda_manual_seed(_SEED)
+        spec = get_gpt_layer_local_spec(enable_hyper_connection=True)
+
+        seq_len, batch_size, h, n = 16, 2, 64, 4
+
+        # With MTP: should return tuple
+        config_mtp = TransformerConfig(
+            num_layers=2,
+            hidden_size=h,
+            num_attention_heads=4,
+            enable_hyper_connections=True,
+            num_residual_streams=n,
+            use_cpu_initialization=True,
+            mtp_num_layers=2,
+        )
+        block_mtp = TransformerBlock(config_mtp, spec).cuda()
+        hidden_states = torch.randn(seq_len, batch_size, n * h, device='cuda', requires_grad=True)
+        output = block_mtp(hidden_states=hidden_states, attention_mask=None)
+
+        assert isinstance(output, tuple)
+        contracted, multistream = output
+        assert contracted.shape == (seq_len, batch_size, h)
+        assert multistream.shape == (seq_len, batch_size, n * h)
+
+        (contracted.sum() + multistream.sum()).backward()
+        assert hidden_states.grad is not None
+
+        # Without MTP: should return single tensor
+        config_no_mtp = TransformerConfig(
+            num_layers=2,
+            hidden_size=h,
+            num_attention_heads=4,
+            enable_hyper_connections=True,
+            num_residual_streams=n,
+            use_cpu_initialization=True,
+            mtp_num_layers=None,
+        )
+        block_no_mtp = TransformerBlock(config_no_mtp, spec).cuda()
+        hs2 = torch.randn(seq_len, batch_size, n * h, device='cuda')
+        output2 = block_no_mtp(hidden_states=hs2, attention_mask=None)
+        assert isinstance(output2, Tensor)
+        assert output2.shape == (seq_len, batch_size, h)
+
+    @pytest.mark.skipif(
+        not HAVE_TE or not is_te_min_version("1.7.0"), reason="TransformerEngine >= 1.7.0 required"
+    )
+    @pytest.mark.parametrize('tp', [1, 2])
+    def test_e2e_forward_backward(self, tp):
+        """GPTModel E2E with mHC + MTP: finite output, MTP loss logged, gradients on HC params."""
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+
+        seq_length, micro_batch_size = 32, 2
+
+        sys.argv = ['test_multi_token_prediction.py']
+        args = parse_args()
+        args.num_layers = 2
+        args.mtp_num_layers = 2
+        args.mtp_loss_scaling_factor = 0.1
+        args.vocab_size = 128800
+        args.hidden_size = 128
+        args.num_attention_heads = 8
+        args.max_position_embeddings = 256
+        args.micro_batch_size = micro_batch_size
+        args.create_attention_mask_in_dataloader = True
+        args.seq_length = seq_length
+        args.tensor_model_parallel_size = tp
+        args.sequence_parallel = tp > 1
+        args.context_parallel_size = 1
+        args.position_embedding_type = 'rope'
+        args.num_experts = None
+        args.moe_grouped_gemm = False
+        args.train_iters = 1
+        args.lr = 3e-5
+        args.attention_dropout = 0.0
+        args.hidden_dropout = 0.0
+        args.add_bias_linear = False
+        args.swiglu = True
+        args.bf16 = True
+        args.enable_hyper_connections = True
+        args.num_residual_streams = 4
+        args.recompute_granularity = None
+
+        validate_args(args)
+        set_global_variables(args, False)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(tensor_model_parallel_size=tp)
+
+        def model_provider(
+            pre_process=True,
+            post_process=True,
+            layer_spec_fn=get_gpt_layer_with_transformer_engine_spec,
+        ):
+            model_parallel_cuda_manual_seed(_SEED)
+            a = get_args()
+            config = core_transformer_config_from_args(a)
+            layer_spec = layer_spec_fn(
+                a.num_experts,
+                a.moe_grouped_gemm,
+                a.qk_layernorm,
+                enable_hyper_connection=config.enable_hyper_connections,
+            )
+            mtp_spec = get_gpt_mtp_block_spec(
+                config=config, spec=layer_spec, use_transformer_engine=True
+            )
+            return GPTModel(
+                config=config,
+                transformer_layer_spec=layer_spec,
+                mtp_block_spec=mtp_spec,
+                vocab_size=a.vocab_size,
+                max_sequence_length=a.max_position_embeddings,
+                pre_process=pre_process,
+                post_process=post_process,
+                fp16_lm_cross_entropy=a.fp16_lm_cross_entropy,
+                parallel_output=True,
+                share_embeddings_and_output_weights=not a.untie_embeddings_and_output_weights,
+                position_embedding_type=a.position_embedding_type,
+                rotary_percent=a.rotary_percent,
+            )
+
+        gpt_model, _, _ = setup_model_and_optimizer(model_provider, ModelType.encoder_or_decoder)
+
+        data = list(range(seq_length))
+        tokens = 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()
+
+        output = gpt_model[0].forward(
+            input_ids=tokens,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            loss_mask=loss_mask,
+        )
+        assert torch.isfinite(output).all(), f"Non-finite output (TP={tp})"
+
+        tracker = MTPLossLoggingHelper.tracker
+        assert "values" in tracker, f"MTP loss not logged (TP={tp})"
+        assert torch.isfinite(tracker['values']).all()
+        MTPLossLoggingHelper.clean_loss_in_tracker()
+
+        output.mean().backward()
+        hc_param_names = ['hc_head_fn', 'hc_head_base', 'hc_head_scale']
+        for name, param in gpt_model[0].named_parameters():
+            assert param.main_grad is not None, f"No gradient for {name}"
+            if any(n in name for n in hc_param_names):
+                assert not torch.all(param.main_grad == 0), f"Zero gradient for {name}"

From ecfb6f32932aa25dac9ae9d9c1f4c22b11d93081 Mon Sep 17 00:00:00 2001
From: Yuzhong Wang <yuzhongw@nvidia.com>
Date: Thu, 30 Apr 2026 18:17:21 -0700
Subject: [PATCH 05/17] fix mscale

---
 .../transformer/experimental_attention_variant/csa.py     | 4 ++++
 .../deepseek_v4_hybrid_attention.py                       | 8 ++++++++
 2 files changed, 12 insertions(+)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py
index 1c24ecda5c7..86f4cc78f63 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa.py
@@ -120,6 +120,10 @@ def _apply_rope(
             ), "Fused MLA RoPE apply is not imported successfully"
         else:
             rotary_pos_emb, mscale = rotary_pos_emb_module(total_seq_len, packed_seq=False)
+            # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's
+            # concentration factor (mscale) is NOT part of the DSv4 model contract --
+            # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0.
+            mscale = 1.0
     if rotary_pos_emb is not None and ratio > 1:
         rotary_pos_emb = rotary_pos_emb[:total_seq_len:ratio][:rotary_seq_len]
     if rotary_pos_cos is not None and ratio > 1:
diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
index 7aa321a3cd1..ffb9ed33373 100644
--- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
@@ -330,6 +330,10 @@ def forward(
                 ), "Fused MLA RoPE apply is not imported successfully"
             else:
                 rotary_pos_emb, mscale = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq)
+                # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's
+                # concentration factor (mscale) is NOT part of the DSv4 model contract --
+                # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0.
+                mscale = 1.0
         if self.config.apply_rope_fusion:
             core_attn_out = fused_mla_rope_inplace(
                 core_attn_out,
@@ -527,6 +531,10 @@ def get_query_key_value_tensors(
                 ), "Fused MLA RoPE apply is not imported successfully"
             else:
                 rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
+                # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's
+                # concentration factor (mscale) is NOT part of the DSv4 model contract --
+                # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0.
+                mscale = 1.0
 
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
             if packed_seq_params.cu_seqlens_q_padded is not None:

From e1fc6895c06d59660259bc94c3a53359898c64fe Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Fri, 8 May 2026 11:13:38 +0000
Subject: [PATCH 06/17] fix state_dict

---
 .../core/transformer/transformer_block.py     | 26 ++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index d0371596717..abd16d41f55 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -35,7 +35,11 @@
     BaseTransformerLayer,
     get_transformer_layer_offset,
 )
-from megatron.core.transformer.utils import sharded_state_dict_default
+from megatron.core.transformer.utils import (
+    ensure_metadata_has_dp_cp_group,
+    make_sharded_tensors_for_checkpoint,
+    sharded_state_dict_default,
+)
 from megatron.core.typed_torch import apply_module, not_none
 from megatron.core.utils import (
     WrappedTensor,
@@ -949,6 +953,7 @@ def forward(
                     len(extract_layer_indices) == 0
                 ), "Feature extraction is not supported with mHC + MTP."
                 mhc_multistream = hidden_states
+            # DSv4 introduced the new output contraction for mHC.
             # [s, b, n*C] -> [s, b, C]
             hidden_states = learned_output_contract(
                 hidden_states,
@@ -1073,4 +1078,23 @@ def sharded_state_dict(
                     )
                 )
 
+        # Save bare parameters/buffers that are direct attributes of this block
+        # (e.g. hyper-connection learned weights: hc_head_fn, hc_head_base,
+        # hc_head_scale). The named_children loop above would silently drop
+        # these since they are not nn.Module children. Mirrors the handling in
+        # MegatronModule.sharded_state_dict.
+        local_state_dict: dict = {}
+        self._save_to_state_dict(local_state_dict, '', keep_vars=True)
+        if local_state_dict:
+            metadata = ensure_metadata_has_dp_cp_group(metadata)
+            sharded_state_dict.update(
+                make_sharded_tensors_for_checkpoint(
+                    local_state_dict,
+                    prefix,
+                    sharded_offsets=sharded_offsets,
+                    tp_group=self.tp_group,
+                    dp_cp_group=metadata['dp_cp_group'],
+                )
+            )
+
         return sharded_state_dict

From 6d80f37de2e8326e7b309fcf9a49e7b752328bdd Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Fri, 8 May 2026 11:16:57 +0000
Subject: [PATCH 07/17] fix eps in learned_output_contract

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

diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py
index 900a24ea4a6..aa6a4bf7fbb 100644
--- a/megatron/core/transformer/hyper_connection.py
+++ b/megatron/core/transformer/hyper_connection.py
@@ -104,7 +104,7 @@ def learned_output_contract(
     scale = scale.to(torch.float32)
     rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps)
     mixes = F.linear(hidden_states, head_fn) * rsqrt
-    pre = torch.sigmoid(mixes * scale + base) + 1e-6
+    pre = torch.sigmoid(mixes * scale + base) + eps
     y = torch.sum(pre.unsqueeze(-1) * hidden_states.view(*hidden_states.shape[:-1], n, -1), dim=-2)
     return y.to(dtype)
 

From c7adacbe6e3e3df6981362145ab4088398099df2 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Sat, 9 May 2026 01:38:40 +0000
Subject: [PATCH 08/17] fix tests

---
 megatron/core/transformer/multi_token_prediction.py |  4 +++-
 .../transformer/test_multi_token_prediction.py      | 13 ++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 8cbc45a2ebd..3e70f8d8a88 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -1015,10 +1015,12 @@ def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.T
             # e_proj: [s, b, h] -> [s, b, h], then broadcast to [s, b, n, h]
             e_out, _ = self.e_proj(decoder_input)
             e_out = gather_from_tensor_model_parallel_region(e_out, group=self.tp_group)
-            e_out = e_out.unsqueeze(2).expand(s, b, n, h)
             # h_proj: applied per-stream on the h dimension
             h_out, _ = self.h_proj(hs_streams)
             h_out = gather_from_tensor_model_parallel_region(h_out, group=self.tp_group)
+            # Sequence-parallel column projections gather the sequence dimension.
+            s, b, n, h = h_out.shape
+            e_out = e_out.unsqueeze(2).expand(s, b, n, h)
             # Combine and flatten back to [s, b, n*h]
             hidden_states = (e_out + h_out).reshape(s, b, n * h)
             if self.sequence_parallel:
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index f3437292374..580fd23d783 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -1210,7 +1210,7 @@ def test_transformer_block_returns_tuple(self):
             mtp_num_layers=2,
         )
         block_mtp = TransformerBlock(config_mtp, spec).cuda()
-        hidden_states = torch.randn(seq_len, batch_size, n * h, device='cuda', requires_grad=True)
+        hidden_states = torch.randn(seq_len, batch_size, h, device='cuda', requires_grad=True)
         output = block_mtp(hidden_states=hidden_states, attention_mask=None)
 
         assert isinstance(output, tuple)
@@ -1232,7 +1232,7 @@ def test_transformer_block_returns_tuple(self):
             mtp_num_layers=None,
         )
         block_no_mtp = TransformerBlock(config_no_mtp, spec).cuda()
-        hs2 = torch.randn(seq_len, batch_size, n * h, device='cuda')
+        hs2 = torch.randn(seq_len, batch_size, h, device='cuda')
         output2 = block_no_mtp(hidden_states=hs2, attention_mask=None)
         assert isinstance(output2, Tensor)
         assert output2.shape == (seq_len, batch_size, h)
@@ -1287,10 +1287,15 @@ def model_provider(
             pre_process=True,
             post_process=True,
             layer_spec_fn=get_gpt_layer_with_transformer_engine_spec,
+            config=None,
+            pg_collection=None,
+            vp_stage=None,
+            **kwargs,
         ):
             model_parallel_cuda_manual_seed(_SEED)
             a = get_args()
-            config = core_transformer_config_from_args(a)
+            if config is None:
+                config = core_transformer_config_from_args(a)
             layer_spec = layer_spec_fn(
                 a.num_experts,
                 a.moe_grouped_gemm,
@@ -1313,6 +1318,8 @@ def model_provider(
                 share_embeddings_and_output_weights=not a.untie_embeddings_and_output_weights,
                 position_embedding_type=a.position_embedding_type,
                 rotary_percent=a.rotary_percent,
+                pg_collection=pg_collection,
+                vp_stage=vp_stage,
             )
 
         gpt_model, _, _ = setup_model_and_optimizer(model_provider, ModelType.encoder_or_decoder)

From 61f71c8d66797f581c329493aab18c02c92364bf Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Tue, 12 May 2026 00:55:16 +0000
Subject: [PATCH 09/17] add yarn arg original-max-position-embeddings

---
 megatron/training/arguments.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index d09f314c1f9..853973b92cd 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -4530,6 +4530,12 @@ def _add_mla_args(parser):
         default=1.0,
         help="Rotary scaling factor for the rotary embeddings.",
     )
+    group.add_argument(
+        '--original-max-position-embeddings',
+        type=int,
+        default=4096,
+        help="Original maximum position embeddings for the original model, used by yarn.",
+    )
     group.add_argument(
         '--mscale', type=float, default=1.0, help="Mscale for YaRN RoPE in multi-latent attention."
     )

From e9f95e3303b950eba48c7bcd41f3299dc8df887c Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Tue, 12 May 2026 03:44:02 +0000
Subject: [PATCH 10/17] fix mtp spec; add tflops calc

---
 gpt_builders.py               |   7 ++
 megatron/training/training.py | 169 ++++++++++++++++++++++++++++------
 2 files changed, 148 insertions(+), 28 deletions(-)

diff --git a/gpt_builders.py b/gpt_builders.py
index bce11cc252b..f3cf6e6a251 100644
--- a/gpt_builders.py
+++ b/gpt_builders.py
@@ -63,6 +63,13 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
             # Get the decoder layer spec explicitly if no decoder layer in the last stage,
             # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
             transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
+        elif args.experimental_attention_variant is not None:
+            # get_gpt_decoder_layer_specs rejects experimental variants;
+            # build per-layer specs via the experimental entry point.
+            experimental_layer_specs = (
+                get_transformer_layer_with_experimental_attention_variant_spec(config=config)
+            )
+            transformer_layer_spec_for_mtp = experimental_layer_specs[-1]
         else:
             # Define the decoder block spec
             decoder_layer_specs = get_gpt_decoder_layer_specs(
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 5fe7ce224ae..542a01c2862 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -564,42 +564,68 @@ def transformer_flops():
             https://arxiv.org/abs/2305.10403
             https://arxiv.org/abs/2205.05198
             '''
-            ## MLA
-            if args.q_lora_rank is None:
-                q_term = (
-                    args.hidden_size
-                    * args.num_attention_heads
-                    * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                )
-            else:
+            if args.experimental_attention_variant == "dsv4_hybrid":
+                ## DSv4 hybrid MLA projections (per layer, per token).
+                ## In dsv4_hybrid mode, qk_head_dim + qk_pos_emb_head_dim == v_head_dim
+                ## (qk_head_dim is derived as v_head_dim - qk_pos_emb_head_dim), and the
+                ## joint KV is produced by a single hidden -> v_head_dim projection.
+                ## Full core attention is replaced by sparse attention and is accounted
+                ## for in the dsv4_hybrid branch below.
                 q_term = args.q_lora_rank * (
                     args.hidden_size
                     + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                    + 1
+                    + 1  # q norm
                 )
-            standard_self_attn_term = (
-                forward_backward_expansion_factor
-                * fma_expansion_factor
-                * (
-                    ## q lora + rope + q norm
-                    q_term
-                    ## kv lora + rope + kv norm
-                    + args.kv_lora_rank
-                    * (
+                kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim  # kv proj + kv norm
+                ## Grouped low-rank output projection:
+                ##  wo_a:        (n_head * v_head_dim) -> (o_groups * o_lora_rank)
+                ##  linear_proj: (o_groups * o_lora_rank) -> hidden
+                o_term = (
+                    args.num_attention_heads * args.v_head_dim * args.o_lora_rank
+                    + args.o_groups * args.o_lora_rank * args.hidden_size
+                )
+                standard_self_attn_term = (
+                    forward_backward_expansion_factor
+                    * fma_expansion_factor
+                    * (q_term + kv_term + o_term)
+                )
+            else:
+                ## MLA
+                if args.q_lora_rank is None:
+                    q_term = (
+                        args.hidden_size
+                        * args.num_attention_heads
+                        * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                    )
+                else:
+                    q_term = args.q_lora_rank * (
                         args.hidden_size
-                        + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim)
+                        + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
                         + 1
                     )
-                    + args.hidden_size * args.qk_pos_emb_head_dim
-                    ## o proj
-                    + (args.num_attention_heads * args.v_head_dim) * args.hidden_size
-                    ## core attn
-                    + args.seq_length
-                    * (args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim))
-                    / 2  # causal mask (only half of the mask is non-zero)
-                    + args.seq_length * args.num_attention_heads * args.v_head_dim / 2
+                standard_self_attn_term = (
+                    forward_backward_expansion_factor
+                    * fma_expansion_factor
+                    * (
+                        ## q lora + rope + q norm
+                        q_term
+                        ## kv lora + rope + kv norm
+                        + args.kv_lora_rank
+                        * (
+                            args.hidden_size
+                            + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim)
+                            + 1
+                        )
+                        + args.hidden_size * args.qk_pos_emb_head_dim
+                        ## o proj
+                        + (args.num_attention_heads * args.v_head_dim) * args.hidden_size
+                        ## core attn
+                        + args.seq_length
+                        * (args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim))
+                        / 2  # causal mask (only half of the mask is non-zero)
+                        + args.seq_length * args.num_attention_heads * args.v_head_dim / 2
+                    )
                 )
-            )
 
         else:
             ## MHA or GQA
@@ -629,6 +655,7 @@ def transformer_flops():
                 )
             )
 
+        dsv4_hybrid_extra_term = 0
         if is_linear_attention_variant(args.experimental_attention_variant):
             # Calculate number of dense and MoE Transformer MLPs.
             if isinstance(args.linear_attention_freq, int):
@@ -686,6 +713,91 @@ def transformer_flops():
                     "Invalid experimental_attention_variant: "
                     f"{args.experimental_attention_variant}"
                 )
+        elif args.experimental_attention_variant == "dsv4_hybrid":
+            # DSv4 hybrid: full core attention is replaced by sparse attention (CSA),
+            # and selected layers additionally run a learned indexer (DSA).
+            # The MLA-style projection cost per layer is captured in
+            # ``standard_self_attn_term`` above; here we add the extra per-layer FLOPs
+            # for sparse attention, the main compressor, and the indexer.
+            num_linear_attention_layers = 0
+            linear_self_attn_term = 0
+            num_standard_attention_layers = num_layers
+
+            compress_ratios = args.csa_compress_ratios
+            assert compress_ratios is not None, (
+                "csa_compress_ratios must be set for dsv4_hybrid"
+            )
+            assert len(compress_ratios) == num_layers, (
+                f"Invalid length of csa_compress_ratios: {len(compress_ratios)}, "
+                f"expected num_layers + mtp_num_layers ({num_layers})."
+            )
+            # ratio == 0: window-only (no compressor, no indexer)
+            # ratio == 4: window + learned-topk over compressed KV (compressor + indexer)
+            # ratio == 128: window + all compressed KV (compressor only)
+            n_layers_r0 = sum(1 for r in compress_ratios if r == 0)
+            n_layers_r4 = sum(1 for r in compress_ratios if r == 4)
+            n_layers_r128 = sum(1 for r in compress_ratios if r == 128)
+
+            n_head = args.num_attention_heads
+            v_head_dim = args.v_head_dim
+            window = args.csa_window_size
+            seq_len = args.seq_length
+
+            # ---- Sparse attention (replaces full core attention) ----
+            # Per token per layer: n_head * (positions_attended) * v_head_dim, ×2 for
+            # QK^T and softmax @ V. Average valid positions account for the causal mask.
+            sparse_attn_r0 = n_layers_r0 * n_head * window * v_head_dim * 2
+            avg_comp_128 = (seq_len // 128) / 2
+            sparse_attn_r128 = n_layers_r128 * n_head * (window + avg_comp_128) * v_head_dim * 2
+
+            # ---- Main compressor (ratio > 0 layers) ----
+            # Two projections per layer (wkv + wgate): hidden -> coff * v_head_dim.
+            #   ratio == 4:   coff = 2 (overlapping windows)
+            #   ratio == 128: coff = 1 (non-overlapping)
+            main_compressor_term = (
+                n_layers_r4 * args.hidden_size * (2 * v_head_dim) * 2
+                + n_layers_r128 * args.hidden_size * (1 * v_head_dim) * 2
+            )
+
+            # ---- r=4 layers: sparse attention + indexer ----
+            # Indexer parameters are only required when at least one ratio==4 layer exists.
+            if n_layers_r4 > 0:
+                assert args.dsa_indexer_n_heads is not None, (
+                    "dsa_indexer_n_heads must be set for dsv4_hybrid with ratio==4 layers."
+                )
+                assert args.dsa_indexer_head_dim is not None, (
+                    "dsa_indexer_head_dim must be set for dsv4_hybrid with ratio==4 layers."
+                )
+                assert args.dsa_indexer_topk is not None, (
+                    "dsa_indexer_topk must be set for dsv4_hybrid with ratio==4 layers."
+                )
+                idx_n_heads = args.dsa_indexer_n_heads
+                idx_head_dim = args.dsa_indexer_head_dim
+                idx_topk = args.dsa_indexer_topk
+
+                effective_topk_4 = min(idx_topk, seq_len // 4)
+                avg_comp_4 = effective_topk_4 * (1 - effective_topk_4 * 4 / (2 * seq_len))
+                sparse_attn_r4 = n_layers_r4 * n_head * (window + avg_comp_4) * v_head_dim * 2
+
+                # Indexer's own compressor (coff=2, wkv + wgate), Q proj, weights proj,
+                # and scoring each query against the seq_len // 4 compressed positions.
+                indexer_term = (
+                    n_layers_r4 * args.hidden_size * (2 * idx_head_dim) * 2
+                    + n_layers_r4 * args.q_lora_rank * idx_n_heads * idx_head_dim
+                    + n_layers_r4 * args.hidden_size * idx_n_heads
+                    + n_layers_r4 * idx_n_heads * idx_head_dim * (seq_len // 4)
+                )
+            else:
+                sparse_attn_r4 = 0
+                indexer_term = 0
+
+            sparse_attn_term = sparse_attn_r0 + sparse_attn_r4 + sparse_attn_r128
+
+            dsv4_hybrid_extra_term = (
+                forward_backward_expansion_factor
+                * fma_expansion_factor
+                * (sparse_attn_term + main_compressor_term + indexer_term)
+            )
         else:
             num_linear_attention_layers = 0
             linear_self_attn_term = 0
@@ -694,6 +806,7 @@ def transformer_flops():
         self_attn_term = (
             linear_self_attn_term * num_linear_attention_layers
             + standard_self_attn_term * num_standard_attention_layers
+            + dsv4_hybrid_extra_term
         )
 
         total_floating_point_operations = (

From 8664bd9cd3ace9d709b5502112cdfd2665eb5e60 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Tue, 12 May 2026 05:34:03 +0000
Subject: [PATCH 11/17] add functional test

---
 .../model_config.yaml                         | 77 +++++++++++++++++++
 tests/test_utils/recipes/h100/gpt.yaml        |  5 ++
 2 files changed, 82 insertions(+)
 create mode 100644 tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml

diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
new file mode 100644
index 00000000000..a1a34515b85
--- /dev/null
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
@@ -0,0 +1,77 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+  ENABLE_LIGHTWEIGHT_MODE: true
+MODEL_ARGS:
+  --num-layers: 4
+  --hidden-size: 512
+  --num-attention-heads: 8
+  --multi-latent-attention: true
+  --q-lora-rank: 192
+  --kv-lora-rank: 64
+  --qk-head-dim: 16
+  --qk-pos-emb-head-dim: 8
+  --v-head-dim: 16
+  --experimental-attention-variant: dsv4_hybrid
+  --dsa-indexer-n-heads: 64
+  --dsa-indexer-head-dim: 128
+  --dsa-indexer-topk: 512
+  --dsa-indexer-loss-coeff: 0.01
+  --dsa-indexer-use-sparse-loss: true
+  --csa-window-size: 128
+  --csa-compress-ratios: ([0,4,128,4,0])
+  --csa-compress-rotary-base: 40000
+  --attention-backend: fused
+  --enable-hyper-connections: true
+  --num-residual-streams: 4
+  --mhc-sinkhorn-iterations: 20
+  --use-fused-mhc: true
+  --mtp-num-layers: 1
+  --mtp-loss-scaling-factor: 0.1
+  --pipeline-model-parallel-layout: "Et|t|t|tmL"
+  --log-params-norm: true
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --micro-batch-size: 4
+  --global-batch-size: 32
+  --seq-length: 1024
+  --max-position-embeddings: 1024
+  --train-iters: 50
+  --timing-log-level: 0
+  --lr-decay-iters: 320000
+  --save: ${CHECKPOINT_SAVE_PATH}
+  --load: ${CHECKPOINT_LOAD_PATH}
+  --data-path: ${DATA_PATH}/text/the_pile/shard00/my-gpt3_00_text_document
+  --vocab-file: ${DATA_PATH}/text/the_pile/shard00/bpe/vocab.json
+  --merge-file: ${DATA_PATH}/text/the_pile/shard00/bpe/merges.txt
+  --split: 949,50,1
+  --distributed-backend: nccl
+  --lr: 0.00015
+  --lr-decay-style: cosine
+  --min-lr: 1.0e-5
+  --weight-decay: 1e-2
+  --clip-grad: 1.0
+  --lr-warmup-fraction: .01
+  --log-interval: 1
+  --save-interval: 25
+  --eval-interval: 1000
+  --eval-iters: 10
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 2
+  --pipeline-model-parallel-size: 2
+  --sequence-parallel: true
+  --untie-embeddings-and-output-weights: true
+  --deterministic-mode: true
+  --no-gradient-accumulation-fusion: true
+  --attention-softmax-in-fp32: true
+  --use-mcore-models: true
+  --ckpt-format: torch_dist
+  --data-cache-path: ${DATA_CACHE_PATH}
+  --bf16: true
+  --attention-backend: unfused
+  --log-memory-to-tensorboard: true
+TEST_TYPE: ckpt-resume
diff --git a/tests/test_utils/recipes/h100/gpt.yaml b/tests/test_utils/recipes/h100/gpt.yaml
index 5da053b793d..9e74d25a87c 100644
--- a/tests/test_utils/recipes/h100/gpt.yaml
+++ b/tests/test_utils/recipes/h100/gpt.yaml
@@ -367,6 +367,11 @@ products:
       - environment: [dev]
         scope: [mr, mr-github, mr-github-slim]
         platforms: [dgx_h100]
+  - test_case: [gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp]
+    products:
+      - environment: [dev]
+        scope: [mr, mr-github, mr-github-slim]
+        platforms: [dgx_h100]
   - test_case: [gpt3_mcore_te_tp2_pp2_resume_torch_dist_ddp_average_in_collective]
     products:
       - environment: [dev]

From 4ace391c21c01edc1777d854a1905167a8bc3be5 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Tue, 12 May 2026 06:30:05 +0000
Subject: [PATCH 12/17] fix test

---
 .../gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml  | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
index a1a34515b85..98416b69be2 100644
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
@@ -39,6 +39,7 @@ MODEL_ARGS:
   --micro-batch-size: 4
   --global-batch-size: 32
   --seq-length: 1024
+  --position-embedding-type: rope
   --max-position-embeddings: 1024
   --train-iters: 50
   --timing-log-level: 0

From ffb96e9d390225b03d13971bff79650026ddb0d0 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Tue, 12 May 2026 08:51:40 +0000
Subject: [PATCH 13/17] fix test config

---
 .../model_config.yaml                                          | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
index 98416b69be2..6541e9d35cc 100644
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml
@@ -27,7 +27,6 @@ MODEL_ARGS:
   --enable-hyper-connections: true
   --num-residual-streams: 4
   --mhc-sinkhorn-iterations: 20
-  --use-fused-mhc: true
   --mtp-num-layers: 1
   --mtp-loss-scaling-factor: 0.1
   --pipeline-model-parallel-layout: "Et|t|t|tmL"
@@ -62,7 +61,7 @@ MODEL_ARGS:
   --eval-interval: 1000
   --eval-iters: 10
   --transformer-impl: transformer_engine
-  --tensor-model-parallel-size: 2
+  --tensor-model-parallel-size: 1
   --pipeline-model-parallel-size: 2
   --sequence-parallel: true
   --untie-embeddings-and-output-weights: true

From 8e4f78cb92a9c67c4ba9950a4470a3fc4955f1cf Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Tue, 12 May 2026 14:04:58 +0000
Subject: [PATCH 14/17] fix indexer loss logging; fix ckpt

---
 .../experimental_attention_variant/csa.py     |  6 ++-
 .../deepseek_v4_hybrid_attention.py           |  1 +
 .../experimental_attention_variant/dsa.py     | 49 +++++++++++++++----
 .../transformer/multi_latent_attention.py     |  6 +++
 .../core/transformer/transformer_config.py    |  1 +
 megatron/training/training.py                 |  2 +
 6 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py
index 86f4cc78f63..547c1828a95 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa.py
@@ -575,6 +575,7 @@ def __init__(
         pg_collection: Optional[ProcessGroupCollection] = None,
         rotary_pos_emb: nn.Module = None,
         compress_ratio: int = 0,
+        is_mtp_layer: bool = False,
     ):
         super().__init__(config=config)
 
@@ -583,6 +584,8 @@ def __init__(
         self.pg_collection = pg_collection
 
         self.layer_number = layer_number
+        if is_mtp_layer:
+            self.layer_number = self.layer_number + self.config.num_layers
         self.compress_ratio = compress_ratio
         self.window_size = config.csa_window_size
         self.v_head_dim = config.v_head_dim
@@ -732,7 +735,8 @@ def forward(
                             DSAIndexerLossLoggingHelper.save_loss_to_tracker(
                                 loss=indexer_loss,
                                 layer_number=self.layer_number,
-                                num_layers=self.config.num_layers,
+                                num_layers=self.config.num_layers
+                                + (self.config.mtp_num_layers or 0),
                             )
                     else:
                         _, topk_indices_compressed = self.indexer(
diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
index ffb9ed33373..0e0a69cb6e9 100644
--- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
@@ -145,6 +145,7 @@ def __init__(
         core_attn_extra_kwargs = {
             "rotary_pos_emb": self.rotary_pos_emb,
             "compress_ratio": compress_ratio,
+            "is_mtp_layer": is_mtp_layer,
         }
         self.core_attention = build_module(
             submodules.core_attention,
diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index 5d7566b3926..94f0fae781c 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -3,7 +3,7 @@
 import copy
 import math
 from dataclasses import dataclass
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 
@@ -89,11 +89,23 @@ def clean_loss_in_tracker():
         tracker["avg_group"] = None
 
     @staticmethod
-    def reduce_loss_in_tracker():
-        """Collect and reduce the indexer losses across ranks."""
+    def reduce_loss_in_tracker(num_layers: Optional[int] = None):
+        """Collect and reduce the indexer losses across ranks.
+
+        Cross-PP `all_reduce` must be invoked on every rank in the pipeline-parallel group,
+        otherwise ranks without any indexer layer would skip the collective and cause a hang.
+        Pass `num_layers` to lazily initialize the tracker on such ranks so they participate
+        with a zero-filled tensor.
+
+        Args:
+            num_layers: Total number of decoder layers; required to lazily initialize the
+                tracker on ranks where no indexer layer ran.
+        """
         tracker = DSAIndexerLossLoggingHelper.tracker
         if "values" not in tracker:
-            return
+            if num_layers is None:
+                return
+            tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device())
         values = tracker["values"]
 
         torch.distributed.all_reduce(
@@ -120,6 +132,8 @@ def track_indexer_metrics(
         wandb_writer=None,
         total_loss_dict=None,
         per_layer_logging: bool = False,
+        num_layers: Optional[int] = None,
+        csa_compress_ratios: Optional[List[int]] = None,
     ):
         """Track the sparse attention indexer metrics for logging.
 
@@ -130,17 +144,31 @@ def track_indexer_metrics(
             wandb_writer: Weights & Biases writer.
             total_loss_dict: Dictionary to accumulate total losses.
             per_layer_logging: Whether to log per-layer losses.
+            num_layers: Total number of decoder layers (including MTP). Required when running
+                with hybrid attention layouts where some PP ranks may not own any indexer
+                layer; passing it ensures every PP rank participates in the cross-PP
+                `all_reduce`.
+            csa_compress_ratios: Per-layer compress ratios for compressed sparse attention.
+                When provided, the cross-layer average uses the count of layers with
+                ``ratio == 4`` (the only ratio that owns an indexer) as the divisor.
+                Otherwise (legacy DSA path) every layer is assumed to be an indexer layer
+                and the divisor is the tracker tensor size.
         """
-        DSAIndexerLossLoggingHelper.reduce_loss_in_tracker()
+        DSAIndexerLossLoggingHelper.reduce_loss_in_tracker(num_layers=num_layers)
         tracker = DSAIndexerLossLoggingHelper.tracker
         if "values" not in tracker:
             return
 
         indexer_loss_values = tracker["values"] * loss_scale
-        num_layers = indexer_loss_values.shape[0]
 
-        # Average across all layers (assuming all layers have sparse attention)
-        avg_indexer_loss = indexer_loss_values.sum() / num_layers
+        if csa_compress_ratios is not None:
+            num_indexer_layers = sum(1 for r in csa_compress_ratios if r == 4)
+        else:
+            num_indexer_layers = indexer_loss_values.shape[0]
+
+        # Average across layers that actually own an indexer; layers without one
+        # contribute zero in `tracker["values"]` so they must not be in the divisor.
+        avg_indexer_loss = indexer_loss_values.sum() / max(num_indexer_layers, 1)
 
         # Log average loss
         if total_loss_dict is not None:
@@ -1076,10 +1104,13 @@ def __init__(
         v_channels: Optional[int] = None,
         cp_comm_type: str = "p2p",
         pg_collection: ProcessGroupCollection = None,
+        is_mtp_layer: bool = False,
     ):
         super().__init__(config=config)
 
         self.layer_number = layer_number
+        if is_mtp_layer:
+            self.layer_number = self.layer_number + self.config.num_layers
 
         self.indexer = build_module(
             submodules.indexer, config=self.config, pg_collection=pg_collection
@@ -1176,7 +1207,7 @@ def forward(
                 DSAIndexerLossLoggingHelper.save_loss_to_tracker(
                     loss=indexer_loss,
                     layer_number=self.layer_number,
-                    num_layers=self.config.num_layers,
+                    num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0),
                 )
 
             # ===================================
diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
index 6954f39a7fe..ab50f8a9067 100644
--- a/megatron/core/transformer/multi_latent_attention.py
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -198,6 +198,11 @@ def __init__(
                 "'rope' and 'yarn'"
             )
 
+        if self.config.experimental_attention_variant == "dsa":
+            core_attn_extra_kwargs = {"is_mtp_layer": is_mtp_layer}
+        else:
+            core_attn_extra_kwargs = {}
+
         self.core_attention = build_module(
             submodules.core_attention,
             config=self.config,
@@ -209,6 +214,7 @@ def __init__(
             v_channels=self.config.v_head_dim,
             cp_comm_type=cp_comm_type,
             pg_collection=self.pg_collection,
+            **core_attn_extra_kwargs,
         )
 
         # Output.
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 1e7ff5a2823..0059af0643e 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -1379,6 +1379,7 @@ def __post_init__(self):
                 self.tensor_model_parallel_size == 1
             ), "DSv4 Hybrid Attention only supports TP size 1."
             assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention."
+            self.hetereogenous_dist_checkpoint = True
 
         if self.fp8:
             # cannot support first last layer bf16 with delayed scaling
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 542a01c2862..af03a485c69 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -2554,6 +2554,8 @@ def training_log(
             writer=writer,
             wandb_writer=wandb_writer,
             total_loss_dict=total_loss_dict,
+            num_layers=args.num_layers + (args.mtp_num_layers or 0),
+            csa_compress_ratios=args.csa_compress_ratios,
         )
 
     # Dump memory snapshot and print metrics to stdout.

From 9d2fcd63f9f5f93910e7a363b4871395bd6691ff Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Wed, 13 May 2026 08:20:42 +0000
Subject: [PATCH 15/17] add swiglu clamp to shared expert

---
 megatron/core/fusions/fused_bias_swiglu.py    | 61 ++++++++++++--
 megatron/core/transformer/mlp.py              |  1 +
 megatron/core/transformer/moe/experts.py      |  1 +
 .../core/transformer/moe/shared_experts.py    | 10 ++-
 .../unit_tests/fusions/test_swiglu_fusion.py  | 83 +++++++++++++++++++
 .../transformer/moe/test_shared_experts.py    | 78 +++++++++++++++++
 6 files changed, 223 insertions(+), 11 deletions(-)

diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py
index ec195551ffa..4cd678be816 100644
--- a/megatron/core/fusions/fused_bias_swiglu.py
+++ b/megatron/core/fusions/fused_bias_swiglu.py
@@ -58,6 +58,12 @@ def clamped_swiglu(y, clamp_value):
     return res.to(dtype)
 
 
+@jit_fuser
+def bias_clamped_swiglu(y, bias, clamp_value):
+    """SwiGLU with clamping after bias addition."""
+    return clamped_swiglu(y + bias, clamp_value)
+
+
 @jit_fuser
 def clamped_weighted_swiglu(y, weights, clamp_value):
     dtype = y.dtype
@@ -134,6 +140,12 @@ def clamped_swiglu_back(g, y, clamp_value):
     return res.to(dtype)
 
 
+@jit_fuser
+def bias_clamped_swiglu_back(g, y, bias, clamp_value):
+    """Backward of SwiGLU with clamping after bias addition."""
+    return clamped_swiglu_back(g, y + bias, clamp_value)
+
+
 @jit_fuser
 def clamped_weighted_swiglu_back(g, y, weights, clamp_value):
     input_dtype = y.dtype
@@ -149,7 +161,7 @@ class BiasSwiGLUFunction(torch.autograd.Function):
 
     @staticmethod
     @nvtx_decorator()
-    def forward(ctx, input, bias, fp8_input_store, cpu_offload_input):
+    def forward(ctx, input, bias, fp8_input_store, cpu_offload_input, clamp_value):
         """Forward pass of biased SwiGLU activation.
 
         Args:
@@ -157,6 +169,10 @@ def forward(ctx, input, bias, fp8_input_store, cpu_offload_input):
             input (torch.Tensor): Input tensor to apply SwiGLU to.
             bias (torch.Tensor): Bias tensor to be added to input before SwiGLU.
             fp8_input_store (bool): If True, stores intermediate values in FP8 format.
+            cpu_offload_input (bool): If True, mark saved tensors for activation offloading.
+            clamp_value (float | None): If set and positive, clamp the gate input to
+                ``[-inf, clamp_value]`` and the linear input to ``[-clamp_value, clamp_value]``
+                before applying SwiGLU.
 
         Returns:
             torch.Tensor: Result of applying bias addition followed by SwiGLU activation.
@@ -168,6 +184,9 @@ def forward(ctx, input, bias, fp8_input_store, cpu_offload_input):
         ctx.save_for_backward(input_for_backward, bias)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
+        ctx.clamp_value = clamp_value
+        if clamp_value is not None and clamp_value > 0:
+            return bias_clamped_swiglu(input, bias, clamp_value)
         return bias_swiglu(input, bias)
 
     @staticmethod
@@ -184,11 +203,16 @@ def backward(ctx, grad_output):
                 - Gradient with respect to the input tensor
                 - Gradient with respect to the bias tensor
                 - None for fp8_input_store parameter
+                - None for cpu_offload_input parameter
+                - None for clamp_value parameter
         """
         input, bias = ctx.saved_tensors
         input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
-        tmp = bias_swiglu_back(grad_output, input, bias)
-        return tmp, tmp, None, None
+        if ctx.clamp_value is not None and ctx.clamp_value > 0:
+            tmp = bias_clamped_swiglu_back(grad_output, input, bias, ctx.clamp_value)
+        else:
+            tmp = bias_swiglu_back(grad_output, input, bias)
+        return tmp, tmp, None, None, None
 
 
 class SwiGLUFunction(torch.autograd.Function):
@@ -196,13 +220,17 @@ class SwiGLUFunction(torch.autograd.Function):
 
     @staticmethod
     @nvtx_decorator()
-    def forward(ctx, input, fp8_input_store, cpu_offload_input):
+    def forward(ctx, input, fp8_input_store, cpu_offload_input, clamp_value):
         """Forward pass of SwiGLU activation.
 
         Args:
             ctx: Autograd context object for saving tensors for backward pass.
             input (torch.Tensor): Input tensor to apply SwiGLU to.
             fp8_input_store (bool): If True, stores intermediate values in FP8 format.
+            cpu_offload_input (bool): If True, mark saved tensors for activation offloading.
+            clamp_value (float | None): If set and positive, clamp the gate input to
+                ``[-inf, clamp_value]`` and the linear input to ``[-clamp_value, clamp_value]``
+                before applying SwiGLU.
 
         Returns:
             torch.Tensor: Result of applying SwiGLU activation.
@@ -213,6 +241,9 @@ def forward(ctx, input, fp8_input_store, cpu_offload_input):
         ctx.save_for_backward(input_for_backward)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
+        ctx.clamp_value = clamp_value
+        if clamp_value is not None and clamp_value > 0:
+            return clamped_swiglu(input, clamp_value)
         return swiglu(input)
 
     @staticmethod
@@ -228,11 +259,16 @@ def backward(ctx, grad_output):
             tuple: Tuple containing:
                 - Gradient with respect to the input tensor
                 - None for fp8_input_store parameter
+                - None for cpu_offload_input parameter
+                - None for clamp_value parameter
         """
         input = ctx.saved_tensors[0]
         input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
-        tmp = swiglu_back(grad_output, input)
-        return tmp, None, None
+        if ctx.clamp_value is not None and ctx.clamp_value > 0:
+            tmp = clamped_swiglu_back(grad_output, input, ctx.clamp_value)
+        else:
+            tmp = swiglu_back(grad_output, input)
+        return tmp, None, None, None
 
 
 class WeightedSwiGLUFunction(torch.autograd.Function):
@@ -260,7 +296,7 @@ def backward(ctx, grad_output):
         return tmp, wgrad, None, None
 
 
-def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False):
+def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False, clamp_value=None):
     """Implementation of biased SwiGLU that handles different input shapes.
 
     This function reshapes the input if necessary, applies the SwiGLU activation
@@ -272,6 +308,11 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False
             uses the bias-free SwiGLU variant.
         fp8_input_store (bool, optional): Whether to store intermediate values in FP8 format.
             Defaults to False.
+        cpu_offload_input (bool, optional): If True, mark saved tensors for activation
+            offloading. Defaults to False.
+        clamp_value (float | None, optional): If set and positive, clamp the gate input to
+            ``[-inf, clamp_value]`` and the linear input to ``[-clamp_value, clamp_value]``
+            before applying SwiGLU. Defaults to None (no clamping).
 
     Returns:
         torch.Tensor: Result of biased SwiGLU activation.
@@ -283,9 +324,11 @@ def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False
     assert len(ori_shape) in [2, 3]
     input = input.view(-1, ori_shape[-1])
     if bias is not None:
-        output = BiasSwiGLUFunction.apply(input, bias, fp8_input_store, cpu_offload_input)
+        output = BiasSwiGLUFunction.apply(
+            input, bias, fp8_input_store, cpu_offload_input, clamp_value
+        )
     else:
-        output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input)
+        output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input, clamp_value)
 
     return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
 
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index ed288c3e1f4..d55e9daf9f6 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -297,6 +297,7 @@ def forward(
                         self.config.cpu_offloading
                         and self.config.cpu_offloading_activations
                         and HAVE_TE,
+                        self.config.activation_func_clamp_value,
                     )
                 else:
                     raise ValueError("Only support fusion of gelu and swiglu")
diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index b9099068720..6c33e56e70f 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -584,6 +584,7 @@ def bias_act_func(self, intermediate_parallel, bias_parallel, permuted_probs):
                     bias_parallel,
                     permuted_probs,
                     self.config.activation_func_fp8_input_store,
+                    self.config.activation_func_clamp_value,
                 )
             elif self.activation_func == quick_gelu and self.config.gated_linear_unit:
                 intermediate_parallel = weighted_bias_quick_geglu_impl(
diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py
index 61ea47955b8..609b0c16dc4 100644
--- a/megatron/core/transformer/moe/shared_experts.py
+++ b/megatron/core/transformer/moe/shared_experts.py
@@ -269,6 +269,7 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None):
                         intermediate_parallel,
                         bias_parallel,
                         self.config.activation_func_fp8_input_store,
+                        clamp_value=self.config.activation_func_clamp_value,
                     )
                 else:
                     raise ValueError("Only support fusion of gelu and swiglu")
@@ -278,8 +279,13 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None):
                 if self.config.gated_linear_unit:
 
                     def glu(x):
-                        x = torch.chunk(x, 2, dim=-1)
-                        return self.config.activation_func(x[0]) * x[1]
+                        x_glu, x_linear = torch.chunk(x, 2, dim=-1)
+                        if (val := self.config.activation_func_clamp_value) is not None:
+                            x_glu = x_glu.clamp(min=None, max=val)
+                            x_linear = x_linear.clamp(min=-val, max=val)
+                        return self.config.activation_func(x_glu) * (
+                            x_linear + self.config.glu_linear_offset
+                        )
 
                     intermediate_parallel = glu(intermediate_parallel)
                 else:
diff --git a/tests/unit_tests/fusions/test_swiglu_fusion.py b/tests/unit_tests/fusions/test_swiglu_fusion.py
index 58e7069d3f1..645e087d027 100644
--- a/tests/unit_tests/fusions/test_swiglu_fusion.py
+++ b/tests/unit_tests/fusions/test_swiglu_fusion.py
@@ -85,3 +85,86 @@ def test_clamped_weighted_bias_swiglu(input_dtype):
     assert weights_2.grad.dtype == weights.grad.dtype
     if input_dtype == torch.float32:
         assert torch.allclose(weights.grad, weights_2.grad, **tols)
+
+
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("with_bias", [False, True])
+def test_clamped_bias_swiglu_impl(input_dtype, with_bias):
+    """``bias_swiglu_impl`` (with and without bias) must respect clamp_value."""
+    clamp_value = 10.0
+
+    if input_dtype == torch.float32:
+        tols = dict(rtol=1.0e-6, atol=1.0e-6)
+    elif input_dtype == torch.bfloat16:
+        tols = dict(rtol=2.0e-2, atol=1.0e-3)
+    else:
+        raise ValueError(f"Invalid input dtype: {input_dtype}")
+
+    # Use a large input range so the clamp actually triggers in many positions.
+    x = (torch.randn(16, 64, dtype=input_dtype, device="cuda") * 5.0).requires_grad_(True)
+    bias = (
+        torch.randn(64, dtype=input_dtype, device="cuda").requires_grad_(True)
+        if with_bias
+        else None
+    )
+    bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda")
+
+    # Reference: manual clamp + silu in fp32 then cast back, mirroring ``clamped_swiglu``.
+    # Cast BEFORE the bias-add so the addition happens in fp32 (matches the fused
+    # kernel's internal accumulation); a bf16-precision bias-add can flip clamp
+    # saturation near the boundary and yield 0 grad where fp32 sees a finite slope.
+    x_fp32 = x.to(torch.float32)
+    x_eff = x_fp32 + bias.to(torch.float32) if with_bias else x_fp32
+    y_1, y_2 = torch.chunk(x_eff, 2, -1)
+    y_1c = y_1.clamp(min=None, max=clamp_value)
+    y_2c = y_2.clamp(min=-clamp_value, max=clamp_value)
+    y_ref = (F.silu(y_1c) * y_2c).to(input_dtype)
+    y_ref.backward(bwd_input)
+
+    x_2 = x.detach().clone().requires_grad_(True)
+    bias_2 = bias.detach().clone().requires_grad_(True) if with_bias else None
+    bwd_input_2 = bwd_input.detach().clone()
+
+    y_fused = bias_swiglu_impl(x_2, bias_2, clamp_value=clamp_value)
+    y_fused.backward(bwd_input_2)
+
+    assert y_fused.dtype == y_ref.dtype
+    assert torch.allclose(y_ref, y_fused, **tols)
+    assert x_2.grad.dtype == x.grad.dtype
+    assert torch.allclose(x.grad, x_2.grad, **tols)
+    if with_bias:
+        assert bias_2.grad.dtype == bias.grad.dtype
+        if input_dtype == torch.float32:
+            assert torch.allclose(bias.grad, bias_2.grad, **tols)
+
+
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("with_bias", [False, True])
+def test_bias_swiglu_impl_clamp_none_matches_unclamped(input_dtype, with_bias):
+    """``clamp_value=None`` (default) must match the legacy unclamped behavior."""
+    if input_dtype == torch.float32:
+        tols = dict(rtol=1.0e-6, atol=1.0e-6)
+    else:
+        tols = dict(rtol=2.0e-2, atol=1.0e-3)
+
+    x = torch.randn(16, 64, dtype=input_dtype, device="cuda").requires_grad_(True)
+    bias = (
+        torch.randn(64, dtype=input_dtype, device="cuda").requires_grad_(True)
+        if with_bias
+        else None
+    )
+    bwd_input = torch.randn(16, 32, dtype=input_dtype, device="cuda")
+
+    y_unclamped = bias_swiglu_impl(x, bias)
+    y_unclamped.backward(bwd_input)
+
+    x_2 = x.detach().clone().requires_grad_(True)
+    bias_2 = bias.detach().clone().requires_grad_(True) if with_bias else None
+
+    y_default_clamp = bias_swiglu_impl(x_2, bias_2, clamp_value=None)
+    y_default_clamp.backward(bwd_input.detach().clone())
+
+    assert torch.allclose(y_unclamped, y_default_clamp, **tols)
+    assert torch.allclose(x.grad, x_2.grad, **tols)
+    if with_bias:
+        assert torch.allclose(bias.grad, bias_2.grad, **tols)
diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py
index a96c6450af4..fd4ed0306e6 100644
--- a/tests/unit_tests/transformer/moe/test_shared_experts.py
+++ b/tests/unit_tests/transformer/moe/test_shared_experts.py
@@ -114,3 +114,81 @@ def test_shared_expert_forward_backward(self, dispatcher_type: str, tp_size, ep_
             assert torch.allclose(
                 p_overlap.grad, p_no_overlap.grad
             ), f"max diff: {torch.max(torch.abs(p_overlap.grad - p_no_overlap.grad))}"
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("bias_activation_fusion", [True, False])
+    def test_shared_expert_clamped_swiglu(self, bias_activation_fusion):
+        """
+        Verifies that ``activation_func_clamp_value`` is honored for SwiGLU shared
+        experts in both the overlapped and non-overlapped paths, and in both the
+        ``bias_activation_fusion`` and manual-glu code paths.
+        """
+        Utils.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1)
+
+        clamp_value = 1.0
+
+        # Create MoE layer with shared expert overlap enabled.
+        model_parallel_cuda_manual_seed(123)
+        moe_layer_overlap = self.get_moe_layer(
+            moe_shared_expert_overlap=True,
+            moe_token_dispatcher_type="alltoall",
+            activation_func_clamp_value=clamp_value,
+            bias_activation_fusion=bias_activation_fusion,
+        ).to(dtype=torch.bfloat16)
+
+        # Create MoE layer with shared expert overlap disabled, sharing weights.
+        model_parallel_cuda_manual_seed(123)
+        moe_layer_no_overlap = self.get_moe_layer(
+            moe_shared_expert_overlap=False,
+            moe_token_dispatcher_type="alltoall",
+            activation_func_clamp_value=clamp_value,
+            bias_activation_fusion=bias_activation_fusion,
+        ).to(dtype=torch.bfloat16)
+        moe_layer_no_overlap.load_state_dict(moe_layer_overlap.state_dict())
+
+        # Use a large input range to ensure the clamp actually triggers.
+        hidden_states = (
+            torch.randn((32, 2, self.config.hidden_size), device="cuda", dtype=torch.bfloat16) * 5.0
+        )
+        hidden_states = hidden_states.detach().requires_grad_(True)
+        hidden_states_no_overlap = hidden_states.detach().clone().requires_grad_(True)
+
+        output_overlap, _ = moe_layer_overlap(hidden_states)
+        output_no_overlap, _ = moe_layer_no_overlap(hidden_states_no_overlap)
+
+        cos_out = torch.nn.functional.cosine_similarity(
+            output_overlap.flatten().unsqueeze(0).float(),
+            output_no_overlap.flatten().unsqueeze(0).float(),
+        ).item()
+        assert cos_out > 0.999, (
+            f"shared-expert clamp output mismatch (fusion={bias_activation_fusion}): "
+            f"cos sim = {cos_out:.6f}"
+        )
+
+        output_overlap.mean().backward()
+        output_no_overlap.mean().backward()
+
+        for p_overlap, p_no_overlap in zip(
+            moe_layer_overlap.parameters(), moe_layer_no_overlap.parameters()
+        ):
+            assert torch.allclose(p_overlap.grad, p_no_overlap.grad), (
+                f"shared-expert clamp mismatch (fusion={bias_activation_fusion}); "
+                f"max diff: {torch.max(torch.abs(p_overlap.grad - p_no_overlap.grad))}"
+            )
+
+        model_parallel_cuda_manual_seed(123)
+        moe_layer_unclamped = self.get_moe_layer(
+            moe_shared_expert_overlap=False,
+            moe_token_dispatcher_type="alltoall",
+            activation_func_clamp_value=None,
+            bias_activation_fusion=bias_activation_fusion,
+        ).to(dtype=torch.bfloat16)
+        moe_layer_unclamped.load_state_dict(moe_layer_overlap.state_dict())
+
+        hidden_states_unclamped = hidden_states.clone().detach().requires_grad_(True)
+        output_unclamped, _ = moe_layer_unclamped(hidden_states_unclamped)
+        assert not torch.allclose(output_no_overlap, output_unclamped), (
+            "Clamping had no observable effect on shared-expert output; "
+            "activation_func_clamp_value may not be plumbed through."
+        )

From 44e12b14e3bc968f640ec1414a0b66f7c41dd1a1 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Wed, 13 May 2026 12:49:30 +0000
Subject: [PATCH 16/17] fix test

---
 tests/unit_tests/fusions/test_swiglu_fusion.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tests/unit_tests/fusions/test_swiglu_fusion.py b/tests/unit_tests/fusions/test_swiglu_fusion.py
index 645e087d027..283d45d68af 100644
--- a/tests/unit_tests/fusions/test_swiglu_fusion.py
+++ b/tests/unit_tests/fusions/test_swiglu_fusion.py
@@ -134,8 +134,10 @@ def test_clamped_bias_swiglu_impl(input_dtype, with_bias):
     assert torch.allclose(x.grad, x_2.grad, **tols)
     if with_bias:
         assert bias_2.grad.dtype == bias.grad.dtype
-        if input_dtype == torch.float32:
-            assert torch.allclose(bias.grad, bias_2.grad, **tols)
+        bias_grad_cos = torch.nn.functional.cosine_similarity(
+            bias.grad.flatten().float().unsqueeze(0), bias_2.grad.flatten().float().unsqueeze(0)
+        ).item()
+        assert bias_grad_cos > 0.999, f"bias.grad cosine similarity = {bias_grad_cos:.6f}"
 
 
 @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32])
@@ -167,4 +169,7 @@ def test_bias_swiglu_impl_clamp_none_matches_unclamped(input_dtype, with_bias):
     assert torch.allclose(y_unclamped, y_default_clamp, **tols)
     assert torch.allclose(x.grad, x_2.grad, **tols)
     if with_bias:
-        assert torch.allclose(bias.grad, bias_2.grad, **tols)
+        bias_grad_cos = torch.nn.functional.cosine_similarity(
+            bias.grad.flatten().float().unsqueeze(0), bias_2.grad.flatten().float().unsqueeze(0)
+        ).item()
+        assert bias_grad_cos > 0.999, f"bias.grad cosine similarity = {bias_grad_cos:.6f}"

From 60ff377d0eff380231e27332030c3a322edb03d6 Mon Sep 17 00:00:00 2001
From: Hongxiao Bai <hongxiaob@nvidia.com>
Date: Thu, 14 May 2026 04:12:33 +0000
Subject: [PATCH 17/17] update golden values due to new mhc contract

---
 .../golden_values_dev_dgx_h100.json           | 490 +++++++++---------
 1 file changed, 245 insertions(+), 245 deletions(-)

diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json
index a890eb1b600..8a856291495 100644
--- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json
+++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_mhc/golden_values_dev_dgx_h100.json
@@ -4,56 +4,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 10.86149,
-            "2": 10.85467,
-            "3": 10.86695,
-            "4": 10.84622,
-            "5": 10.88467,
-            "6": 10.89675,
-            "7": 10.87274,
-            "8": 10.86587,
-            "9": 10.86993,
-            "10": 10.83755,
-            "11": 10.8946,
-            "12": 10.8795,
-            "13": 10.87683,
-            "14": 10.90365,
+            "1": 10.86153,
+            "2": 10.85471,
+            "3": 10.86694,
+            "4": 10.84628,
+            "5": 10.88468,
+            "6": 10.8968,
+            "7": 10.87276,
+            "8": 10.86586,
+            "9": 10.86987,
+            "10": 10.83767,
+            "11": 10.89459,
+            "12": 10.87954,
+            "13": 10.87684,
+            "14": 10.9036,
             "15": 10.83112,
-            "16": 10.8345,
-            "17": 10.80061,
-            "18": 10.82067,
-            "19": 10.81459,
-            "20": 10.71809,
-            "21": 10.68633,
+            "16": 10.83449,
+            "17": 10.80068,
+            "18": 10.8207,
+            "19": 10.81463,
+            "20": 10.7181,
+            "21": 10.68634,
             "22": 10.53197,
-            "23": 10.70485,
-            "24": 10.58544,
-            "25": 10.51899,
-            "26": 10.58489,
-            "27": 10.60103,
-            "28": 10.53535,
-            "29": 10.57111,
-            "30": 10.33244,
-            "31": 10.05828,
-            "32": 10.42787,
-            "33": 10.42023,
+            "23": 10.70487,
+            "24": 10.58555,
+            "25": 10.51897,
+            "26": 10.5849,
+            "27": 10.60104,
+            "28": 10.53536,
+            "29": 10.57116,
+            "30": 10.33242,
+            "31": 10.05836,
+            "32": 10.42794,
+            "33": 10.42026,
             "34": 10.16983,
-            "35": 10.23073,
-            "36": 10.18747,
-            "37": 10.31252,
-            "38": 10.14214,
-            "39": 10.38141,
-            "40": 10.04843,
-            "41": 10.10327,
+            "35": 10.23071,
+            "36": 10.18758,
+            "37": 10.31243,
+            "38": 10.14212,
+            "39": 10.38137,
+            "40": 10.04847,
+            "41": 10.10333,
             "42": 10.17154,
-            "43": 9.78292,
-            "44": 9.90961,
-            "45": 9.78503,
-            "46": 9.76877,
-            "47": 10.10084,
+            "43": 9.78293,
+            "44": 9.90957,
+            "45": 9.78507,
+            "46": 9.7689,
+            "47": 10.10085,
             "48": 9.80966,
-            "49": 9.48773,
-            "50": 9.86705
+            "49": 9.48775,
+            "50": 9.86712
         }
     },
     "num-zeros": {
@@ -61,56 +61,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 1732.0,
-            "2": 34586.0,
-            "3": 1628.0,
-            "4": 1806.0,
-            "5": 1834.0,
-            "6": 1858.0,
-            "7": 1772.0,
-            "8": 1562.0,
-            "9": 34695.0,
-            "10": 1453.0,
-            "11": 34608.0,
-            "12": 34493.0,
-            "13": 1885.0,
-            "14": 34479.0,
-            "15": 1876.0,
-            "16": 1773.0,
-            "17": 34664.0,
-            "18": 1653.0,
-            "19": 1796.0,
-            "20": 1636.0,
-            "21": 1854.0,
-            "22": 1680.0,
-            "23": 34870.0,
-            "24": 1743.0,
-            "25": 34415.0,
-            "26": 34506.0,
-            "27": 34562.0,
-            "28": 1973.0,
-            "29": 34797.0,
-            "30": 1874.0,
-            "31": 34398.0,
-            "32": 34704.0,
-            "33": 34981.0,
-            "34": 1929.0,
-            "35": 34822.0,
-            "36": 34718.0,
-            "37": 2413.0,
-            "38": 35053.0,
-            "39": 35229.0,
-            "40": 34965.0,
-            "41": 35070.0,
-            "42": 2353.0,
-            "43": 34792.0,
-            "44": 35066.0,
-            "45": 34885.0,
-            "46": 35077.0,
-            "47": 35294.0,
-            "48": 35254.0,
-            "49": 35217.0,
-            "50": 35213.0
+            "1": 1703.0,
+            "2": 1742.0,
+            "3": 1622.0,
+            "4": 1788.0,
+            "5": 1858.0,
+            "6": 1801.0,
+            "7": 1793.0,
+            "8": 1674.0,
+            "9": 1888.0,
+            "10": 1389.0,
+            "11": 1759.0,
+            "12": 1668.0,
+            "13": 1869.0,
+            "14": 1801.0,
+            "15": 1852.0,
+            "16": 1768.0,
+            "17": 1945.0,
+            "18": 1725.0,
+            "19": 1754.0,
+            "20": 1688.0,
+            "21": 1866.0,
+            "22": 1620.0,
+            "23": 2088.0,
+            "24": 1701.0,
+            "25": 1641.0,
+            "26": 1760.0,
+            "27": 1842.0,
+            "28": 2007.0,
+            "29": 1987.0,
+            "30": 1956.0,
+            "31": 1590.0,
+            "32": 1873.0,
+            "33": 2187.0,
+            "34": 1985.0,
+            "35": 1969.0,
+            "36": 1921.0,
+            "37": 2438.0,
+            "38": 2161.0,
+            "39": 2402.0,
+            "40": 2183.0,
+            "41": 2268.0,
+            "42": 2382.0,
+            "43": 2039.0,
+            "44": 2157.0,
+            "45": 2204.0,
+            "46": 2370.0,
+            "47": 2460.0,
+            "48": 2439.0,
+            "49": 2414.0,
+            "50": 2402.0
         }
     },
     "mem-allocated-bytes": {
@@ -118,56 +118,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 555746816.0,
-            "2": 555746816.0,
-            "3": 555746816.0,
-            "4": 555746816.0,
-            "5": 555746816.0,
-            "6": 555746816.0,
-            "7": 555746816.0,
-            "8": 555746816.0,
-            "9": 555746816.0,
-            "10": 555746816.0,
-            "11": 555746816.0,
-            "12": 555746816.0,
-            "13": 555746816.0,
-            "14": 555746816.0,
-            "15": 555746816.0,
-            "16": 555746816.0,
-            "17": 555746816.0,
-            "18": 555746816.0,
-            "19": 555746816.0,
-            "20": 555746816.0,
-            "21": 555746816.0,
-            "22": 555746816.0,
-            "23": 555746816.0,
-            "24": 555746816.0,
-            "25": 555746816.0,
-            "26": 555746816.0,
-            "27": 555746816.0,
-            "28": 555746816.0,
-            "29": 555746816.0,
-            "30": 555746816.0,
-            "31": 555746816.0,
-            "32": 555746816.0,
-            "33": 555746816.0,
-            "34": 555746816.0,
-            "35": 555746816.0,
-            "36": 555746816.0,
-            "37": 555746816.0,
-            "38": 555746816.0,
-            "39": 555746816.0,
-            "40": 555746816.0,
-            "41": 555746816.0,
-            "42": 555746816.0,
-            "43": 555746816.0,
-            "44": 555746816.0,
-            "45": 555746816.0,
-            "46": 555746816.0,
-            "47": 555746816.0,
-            "48": 555746816.0,
-            "49": 555746816.0,
-            "50": 555746816.0
+            "1": 554817024.0,
+            "2": 554817024.0,
+            "3": 555865600.0,
+            "4": 555865600.0,
+            "5": 554817024.0,
+            "6": 554817024.0,
+            "7": 554817024.0,
+            "8": 554817024.0,
+            "9": 554817024.0,
+            "10": 555865600.0,
+            "11": 554817024.0,
+            "12": 554817024.0,
+            "13": 554817024.0,
+            "14": 554817024.0,
+            "15": 554817024.0,
+            "16": 554817024.0,
+            "17": 554817024.0,
+            "18": 554817024.0,
+            "19": 554817024.0,
+            "20": 554817024.0,
+            "21": 554817024.0,
+            "22": 554817024.0,
+            "23": 554817024.0,
+            "24": 554817024.0,
+            "25": 554817024.0,
+            "26": 554817024.0,
+            "27": 554817024.0,
+            "28": 554817024.0,
+            "29": 554817024.0,
+            "30": 554817024.0,
+            "31": 554817024.0,
+            "32": 554817024.0,
+            "33": 554817024.0,
+            "34": 554817024.0,
+            "35": 554817024.0,
+            "36": 554817024.0,
+            "37": 555865600.0,
+            "38": 554817024.0,
+            "39": 554817024.0,
+            "40": 554817024.0,
+            "41": 554817024.0,
+            "42": 554817024.0,
+            "43": 554817024.0,
+            "44": 554817024.0,
+            "45": 554817024.0,
+            "46": 555865600.0,
+            "47": 554817024.0,
+            "48": 554817024.0,
+            "49": 554817024.0,
+            "50": 554817024.0
         }
     },
     "mem-max-allocated-bytes": {
@@ -175,56 +175,56 @@
         "end_step": 50,
         "step_interval": 1,
         "values": {
-            "1": 1728349696.0,
-            "2": 1917909504.0,
-            "3": 1917909504.0,
-            "4": 1917909504.0,
-            "5": 1917909504.0,
-            "6": 1917909504.0,
-            "7": 1917909504.0,
-            "8": 1917909504.0,
-            "9": 1917909504.0,
-            "10": 1917909504.0,
-            "11": 1917909504.0,
-            "12": 1917909504.0,
-            "13": 1917909504.0,
-            "14": 1917909504.0,
-            "15": 1917909504.0,
-            "16": 1917909504.0,
-            "17": 1917909504.0,
-            "18": 1917909504.0,
-            "19": 1917909504.0,
-            "20": 1917909504.0,
-            "21": 1917909504.0,
-            "22": 1917909504.0,
-            "23": 1917909504.0,
-            "24": 1917909504.0,
-            "25": 1917909504.0,
-            "26": 1917909504.0,
-            "27": 1917909504.0,
-            "28": 1917909504.0,
-            "29": 1917909504.0,
-            "30": 1917909504.0,
-            "31": 1917909504.0,
-            "32": 1917909504.0,
-            "33": 1917909504.0,
-            "34": 1917909504.0,
-            "35": 1917909504.0,
-            "36": 1917909504.0,
-            "37": 1917909504.0,
-            "38": 1917909504.0,
-            "39": 1917909504.0,
-            "40": 1917909504.0,
-            "41": 1917909504.0,
-            "42": 1917909504.0,
-            "43": 1917909504.0,
-            "44": 1917909504.0,
-            "45": 1917909504.0,
-            "46": 1917909504.0,
-            "47": 1917909504.0,
-            "48": 1917909504.0,
-            "49": 1917909504.0,
-            "50": 1917909504.0
+            "1": 1746300416.0,
+            "2": 1934879232.0,
+            "3": 1934879232.0,
+            "4": 1934879232.0,
+            "5": 1934879232.0,
+            "6": 1934879232.0,
+            "7": 1934879232.0,
+            "8": 1934879232.0,
+            "9": 1934879232.0,
+            "10": 1934879232.0,
+            "11": 1934879232.0,
+            "12": 1934879232.0,
+            "13": 1934879232.0,
+            "14": 1934879232.0,
+            "15": 1934879232.0,
+            "16": 1934879232.0,
+            "17": 1934879232.0,
+            "18": 1934879232.0,
+            "19": 1934879232.0,
+            "20": 1934879232.0,
+            "21": 1934879232.0,
+            "22": 1934879232.0,
+            "23": 1934879232.0,
+            "24": 1934879232.0,
+            "25": 1934879232.0,
+            "26": 1935927808.0,
+            "27": 1935927808.0,
+            "28": 1935927808.0,
+            "29": 1935927808.0,
+            "30": 1935927808.0,
+            "31": 1935927808.0,
+            "32": 1935927808.0,
+            "33": 1935927808.0,
+            "34": 1935927808.0,
+            "35": 1935927808.0,
+            "36": 1935927808.0,
+            "37": 1935927808.0,
+            "38": 1935927808.0,
+            "39": 1935927808.0,
+            "40": 1935927808.0,
+            "41": 1935927808.0,
+            "42": 1935927808.0,
+            "43": 1935927808.0,
+            "44": 1935927808.0,
+            "45": 1935927808.0,
+            "46": 1935927808.0,
+            "47": 1935927808.0,
+            "48": 1935927808.0,
+            "49": 1935927808.0,
+            "50": 1935927808.0
         }
     },
     "iteration-time": {
@@ -233,55 +233,55 @@
         "step_interval": 1,
         "values": {
             "1": "nan",
-            "2": 30.27287,
-            "3": 0.63036,
-            "4": 0.62463,
-            "5": 0.62389,
-            "6": 0.62241,
-            "7": 0.62274,
-            "8": 0.62116,
-            "9": 0.62223,
-            "10": 0.62501,
-            "11": 0.62222,
-            "12": 0.62201,
-            "13": 0.6223,
-            "14": 0.62539,
-            "15": 0.62434,
-            "16": 0.62424,
-            "17": 0.62735,
-            "18": 0.62325,
-            "19": 0.62244,
-            "20": 0.62506,
-            "21": 0.62317,
-            "22": 0.62235,
-            "23": 0.625,
-            "24": 0.62205,
-            "25": 0.62519,
-            "26": 0.64769,
-            "27": 0.62564,
-            "28": 0.62374,
-            "29": 0.62533,
-            "30": 0.62018,
-            "31": 0.62779,
-            "32": 0.62201,
-            "33": 0.63514,
-            "34": 0.6314,
-            "35": 0.63737,
-            "36": 0.62906,
-            "37": 0.64653,
-            "38": 0.63058,
-            "39": 0.63017,
-            "40": 0.63041,
-            "41": 0.6331,
-            "42": 0.62522,
-            "43": 0.62568,
-            "44": 0.62119,
-            "45": 0.62536,
-            "46": 0.62217,
-            "47": 0.62615,
-            "48": 0.6199,
-            "49": 0.61769,
-            "50": 0.62242
+            "2": 31.19119,
+            "3": 4.64254,
+            "4": 2.06235,
+            "5": 3.10645,
+            "6": 2.09979,
+            "7": 3.00592,
+            "8": 2.84917,
+            "9": 2.47663,
+            "10": 3.44241,
+            "11": 3.16127,
+            "12": 5.01739,
+            "13": 2.65652,
+            "14": 2.48699,
+            "15": 2.65524,
+            "16": 2.11061,
+            "17": 2.67153,
+            "18": 3.8657,
+            "19": 3.07496,
+            "20": 3.26175,
+            "21": 3.6739,
+            "22": 1.84516,
+            "23": 3.5372,
+            "24": 2.81029,
+            "25": 2.85832,
+            "26": 0.7157,
+            "27": 5.60633,
+            "28": 3.04555,
+            "29": 2.37229,
+            "30": 2.82131,
+            "31": 3.33094,
+            "32": 5.55589,
+            "33": 2.19105,
+            "34": 2.07484,
+            "35": 2.60419,
+            "36": 2.92689,
+            "37": 2.97485,
+            "38": 1.62047,
+            "39": 3.13391,
+            "40": 3.59651,
+            "41": 4.22308,
+            "42": 1.86597,
+            "43": 3.16598,
+            "44": 2.03267,
+            "45": 2.76972,
+            "46": 2.09152,
+            "47": 3.69723,
+            "48": 2.47382,
+            "49": 2.18467,
+            "50": 2.99757
         }
     }
-}
\ No newline at end of file
+}