[dev]: faster implementation of mHC fused kernels

megatron/core/models/hybrid/hybrid_block.py

@@ -146,8 +146,7 @@ def forward(
         mhc_recompute_manager=None,
     ) -> Tuple[Tensor, Optional[Tensor]]:
         """Run the wrapped hybrid layer through one layer-boundary mHC update."""
-        residual = hidden_states
-        aggregated, h_res, h_post = self.hyper_connection(
+        aggregated, h_res, h_post, residual = self.hyper_connection(
             hidden_states, mhc_recompute_manager=mhc_recompute_manager
         )
         layer_output, context = self._call_inner_layer(

megatron/core/transformer/hyper_connection.py

@@ -15,14 +15,17 @@
 if TYPE_CHECKING:
     from megatron.core.tensor_parallel.random import CheckpointManager
 
+_MHC_SINKHORN_EPS = 1e-6
+_MHC_COMPUTE_H_EPS = 1e-6
+
 
 @torch.compile
 def _sinkhorn_iterations(input_logits: Tensor, num_iterations: int, eps: float) -> Tensor:
-    row_max = input_logits.max(dim=-1, keepdim=True).values
-    M = torch.exp(input_logits - row_max)
-    for _ in range(num_iterations):
-        M = M / M.sum(dim=-1, keepdim=True).clamp(min=eps)
-        M = M / M.sum(dim=-2, keepdim=True).clamp(min=eps)
+    M = input_logits.softmax(dim=-1) + eps
+    M = M / (M.sum(dim=-2, keepdim=True) + eps)
+    for _ in range(num_iterations - 1):
+        M = M / (M.sum(dim=-1, keepdim=True) + eps)
+        M = M / (M.sum(dim=-2, keepdim=True) + eps)
     return M
 
 
@@ -69,11 +72,11 @@ def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor:
 def native_h_post_bda(
     h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor]
 ) -> Tensor:
-    """Native H_res @ residual + H_post * (x [+ bias])."""
+    """Native H_res.T @ residual + H_post * (x [+ bias])."""
     s, b, n, C = original_residual.shape
     h_res_batched = h_res.view(s * b, n, n)
     residual_batched = original_residual.view(s * b, n, C)
-    mixed = torch.bmm(h_res_batched, residual_batched).view(s, b, n, C)
+    mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched).view(s, b, n, C)
     x_expanded = h_post.unsqueeze(-1) * x.unsqueeze(2)
     if bias is not None:
         bias_expanded = h_post.unsqueeze(-1) * bias.view(1, 1, 1, C)
@@ -92,6 +95,39 @@ def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tenso
     return proj, r
 
 
+@torch.compile
+def native_fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor:
+    """Native 3-way elementwise add (torch.compile fuses into single kernel)."""
+    return a + b + c
+
+
+class BroadcastTensorFused(torch.autograd.Function):
+    """Split one tensor into 3 autograd-graph children sharing the same storage.
+
+    During backward the three incoming gradients are summed with a caller-
+    supplied fused-add function (cuTile or torch.compile fallback) instead of
+    PyTorch's default sequential accumulation.
+    """
+
+    @staticmethod
+    def forward(ctx, x, fused_add_3_fn):
+        """Return three view aliases and save the fused gradient combiner."""
+        ctx.fused_add_3_fn = fused_add_3_fn
+        return x.view_as(x), x.view_as(x), x.view_as(x)
+
+    @staticmethod
+    def backward(ctx, grad1, grad2, grad3):
+        """Combine gradients from the three broadcast aliases."""
+        grads = [g for g in (grad1, grad2, grad3) if g is not None]
+        if len(grads) == 0:
+            return None, None
+        if len(grads) == 1:
+            return grads[0], None
+        if len(grads) == 2:
+            return grads[0] + grads[1], None
+        return ctx.fused_add_3_fn(grad1, grad2, grad3), None
+
+
 @torch.compile
 def learned_output_contract(
     hidden_states: Tensor, head_fn: Tensor, base: Tensor, scale: Tensor, n: int, eps: float
@@ -120,13 +156,13 @@ class HyperConnectionModule(MegatronModule):
     Unified mHC (Manifold-Constrained Hyper-Connections) module.
 
     Implements the complete mHC propagation:
-        x_{l+1} = H_res @ x_l + H_post^T @ F(H_pre @ x_l)
+        x_{l+1} = H_res^T @ x_l + H_post^T @ F(H_pre @ x_l)
 
     This module handles:
     1. Computing learnable mappings: H_pre, H_post, H_res (with Sinkhorn-Knopp projection)
     2. Aggregation: n-stream → 1-stream (H_pre @ x)
     3. Expansion: 1-stream → n-stream (H_post^T @ output)
-    4. Residual merge: H_res @ x + expanded_output
+    4. Residual merge: H_res^T @ x + expanded_output
     5. Block-level expand/contract for TransformerBlock boundaries
 
     Args:
@@ -141,6 +177,8 @@ def __init__(self, config: TransformerConfig, layer_number: int):
         self.n = config.num_residual_streams
         self.hidden_size = config.hidden_size
         self.sinkhorn_iterations = config.mhc_sinkhorn_iterations
+        self.sinkhorn_eps = _MHC_SINKHORN_EPS
+        self.compute_h_eps = _MHC_COMPUTE_H_EPS
 
         # Projection weights for dynamic mappings
         # Input: [s, b, n*C] -> Output: n^2 + 2n values per token
@@ -161,26 +199,35 @@ def __init__(self, config: TransformerConfig, layer_number: int):
         self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n))
         self.norm_eps = 1e-6
 
-        # Choose implementation: fused cuTile kernels vs reference modules.
-        # Both paths expose the same call signatures so the rest of the code
-        # is implementation-agnostic.
+        # Choose implementation: unified fused kernels vs reference modules.
+        # The fused public API selects the backend per operation internally.
+        # fused_add_3 always uses torch.compile (native_fused_add_3) regardless
+        # of the kernel backend. cuTile's register overhead (56 regs/thread for
+        # a trivial a+b+c) is not worth it for a pure memory-bound elementwise op.
+        self._fused_add_3_op = native_fused_add_3
+
         if config.use_fused_mhc:
             from megatron.core.fusions.fused_mhc_kernels import (
                 fused_h_aggregate,
                 fused_h_post_bda,
                 fused_proj_rms,
+                fused_proj_rms_compute_h,
                 fused_sinkhorn,
+                log_fused_mhc_backend_once,
             )
 
+            log_fused_mhc_backend_once()
             self._sinkhorn_op = fused_sinkhorn
             self._h_aggregate_op = fused_h_aggregate
             self._h_post_bda_op = fused_h_post_bda
             self._proj_rms_op = fused_proj_rms
+            self._proj_rms_compute_h_op = fused_proj_rms_compute_h
         else:
             self._sinkhorn_op = native_sinkhorn
             self._h_aggregate_op = native_h_aggregate
             self._h_post_bda_op = native_h_post_bda
             self._proj_rms_op = native_proj_rms
+            self._proj_rms_compute_h_op = None
 
         self._init_weights()
 
@@ -235,10 +282,10 @@ def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
         )
         h = r * proj * alpha_ + self.bias
         # H_pre = σ(α_pre * (θ_pre @ x̃) + b_pre)
-        h_pre = h[..., : self.n].sigmoid()  # [s, b, n]
+        h_pre = h[..., : self.n].sigmoid() + self.compute_h_eps  # [s, b, n]
 
         # H_post = 2σ(α_post * (θ_post @ x̃) + b_post)
-        h_post = h[..., self.n : 2 * self.n].sigmoid() * 2  # [s, b, n]
+        h_post = h[..., self.n : 2 * self.n].sigmoid() * 2
         h_res = h[..., 2 * self.n :]
         return h_pre, h_post, h_res
 
@@ -258,12 +305,35 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
             h_res: [s, b, n, n] - residual mixing matrix (doubly stochastic)
         """
         s, b, _ = x.shape
-        with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"):
-            proj, r = self._projection_and_get_norm(x)
-        with torch.cuda.nvtx.range("HyperConnection::compute_h"):
-            h_pre, h_post, h_res = self._compute_h(proj, r)
+
+        if self._proj_rms_compute_h_op is not None:
+            # Fused path: proj_rms + compute_h in one kernel launch sequence
+            x_2d = x.reshape(s * b, self.n * self.hidden_size)
+            with torch.cuda.nvtx.range("HyperConnection::fused_proj_rms_compute_h"):
+                h_pre, h_post, h_res, _ = self._proj_rms_compute_h_op(
+                    x_2d,
+                    self.mapping_proj.weight,
+                    self.alpha_pre,
+                    self.alpha_post,
+                    self.alpha_res,
+                    self.bias,
+                    self.n,
+                    self.norm_eps,
+                    self.compute_h_eps,
+                )
+            h_pre = h_pre.view(s, b, self.n)
+            h_post = h_post.view(s, b, self.n)
+            h_res = h_res.view(s, b, self.n, self.n)
+        else:
+            # Native path: separate proj_rms + _compute_h
+            with torch.cuda.nvtx.range("HyperConnection::projection_and_get_norm"):
+                proj, r = self._projection_and_get_norm(x)
+            with torch.cuda.nvtx.range("HyperConnection::compute_h"):
+                h_pre, h_post, h_res = self._compute_h(proj, r)
+            h_res = h_res.view(s, b, self.n, self.n)
+
         h_res = self._sinkhorn_op(
-            h_res.view(s, b, self.n, self.n), self.sinkhorn_iterations, self.norm_eps
+            h_res, self.sinkhorn_iterations, self.sinkhorn_eps
         )  # [s, b, n, n]
 
         return h_pre, h_post, h_res
@@ -372,7 +442,7 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor:
         """
         Apply H_res to residual using H_res weights.
 
-        Computes: H_res @ residual
+        Computes: H_res.T @ residual
 
         Args:
             h_res: [s, b, n, n] - residual mixing matrix
@@ -387,33 +457,42 @@ def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor:
         # [s, b, n*C] -> [s, b, n, C] -> [s*b, n, C]
         residual_batched = residual.view(s, b, n, C).view(s * b, n, C)
 
-        # Batch matrix multiply: [s*b, n, n] @ [s*b, n, C] -> [s*b, n, C]
-        mixed = torch.bmm(h_res_batched, residual_batched)
+        # Batch matrix multiply: [s*b, n, n].T @ [s*b, n, C] -> [s*b, n, C]
+        mixed = torch.bmm(h_res_batched.transpose(1, 2), residual_batched)
 
         return mixed.view(s, b, n * C)
 
     def forward(
         self, hidden_states: Tensor, mhc_recompute_manager: Optional['CheckpointManager'] = None
-    ) -> Tuple[Tensor, Tensor, Tensor]:
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         """
         Full mHC forward pass.
 
+        Uses BroadcastTensorFused to split hidden_states into 3 autograd-graph
+        children so that gradient accumulation from the 3 consumers
+        (compute_mappings, aggregate, fused_h_res_h_post_bda) is handled by a
+        single fused add instead of PyTorch's default sequential accumulation.
+
         Args:
             hidden_states: [s, b, n*C] - n-stream hidden states
             mhc_recompute_manager: Optional CheckpointManager for checkpoint management.
                 When provided, uses _forward_with_checkpoint for memory-efficient execution.
 
         Returns:
+            A 4-tuple. This is an intentional breaking change from the older
+            3-tuple API because fused_h_res_h_post_bda consumes the residual
+            branch created by BroadcastTensorFused.
             aggregated: [s, b, C] - aggregated input for layer computation
             h_res: [s, b, n, n] - residual mixing matrix (for fused kernel)
             h_post: [s, b, n] - expansion weights
+            residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda
         """
         if mhc_recompute_manager is not None:
             return self._forward_with_checkpoint(hidden_states, mhc_recompute_manager)
         else:
             return self._forward_normal(hidden_states)
 
-    def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+    def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         """
         Normal forward pass without checkpointing.
 
@@ -424,19 +503,25 @@ def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor
             aggregated: [s, b, C] - aggregated input for layer computation
             h_res: [s, b, n, n] - residual mixing matrix (for fused kernel)
             h_post: [s, b, n] - expansion weights
+            residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda
         """
+        # Split into 3 views to avoid extra grad accumulations in backward
+        hs_for_mappings, hs_for_aggregate, hs_for_residual = BroadcastTensorFused.apply(
+            hidden_states, self._fused_add_3_op
+        )
+
         # Compute mappings
-        h_pre, h_post, h_res = self.compute_mappings(hidden_states)
+        h_pre, h_post, h_res = self.compute_mappings(hs_for_mappings)
 
         # Aggregate for layer input
         with torch.cuda.nvtx.range("HyperConnection::aggregate"):
-            aggregated = self.aggregate(hidden_states, h_pre)
+            aggregated = self.aggregate(hs_for_aggregate, h_pre)
 
-        return aggregated, h_res, h_post
+        return aggregated, h_res, h_post, hs_for_residual
 
     def _forward_with_checkpoint(
         self, hidden_states: Tensor, manager: 'CheckpointManager'
-    ) -> Tuple[Tensor, Tensor, Tensor]:
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
         """
         Forward pass with checkpointing for memory efficiency.
 
@@ -453,17 +538,23 @@ def _forward_with_checkpoint(
             aggregated: [s, b, C] - aggregated input for layer computation
             h_res: [s, b, n, n] - residual mixing matrix (for fused kernel)
             h_post: [s, b, n] - expansion weights
+            residual: [s, b, n*C] - residual view for fused_h_res_h_post_bda
         """
         from megatron.core.tensor_parallel.random import CheckpointWithoutOutput
 
-        h_pre, h_post, h_res = self.compute_mappings(hidden_states)
+        # Split into 3 views to avoid extra grad accumulations in backward
+        hs_for_mappings, hs_for_aggregate, hs_for_residual = BroadcastTensorFused.apply(
+            hidden_states, self._fused_add_3_op
+        )
+
+        h_pre, h_post, h_res = self.compute_mappings(hs_for_mappings)
 
         # Checkpoint aggregate - auto-registers to manager
         aggregated = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(
-            self.aggregate, hidden_states, h_pre
+            self.aggregate, hs_for_aggregate, h_pre
         )
 
-        return aggregated, h_res, h_post
+        return aggregated, h_res, h_post, hs_for_residual
 
     # ==================== Block-level utilities ====================
 
@@ -528,7 +619,7 @@ def fused_h_res_h_post_bda(
         Currently implements the operations sequentially using native PyTorch.
 
         The computation flow is:
-            1. mixed = H_res @ original_residual (apply_h_res)
+            1. mixed = H_res.T @ original_residual (apply_h_res)
             2. expanded = H_post^T @ layer_output (apply_h_post)
             3. output = dropout(expanded + bias) + mixed (bias-dropout-add)
 
@@ -584,7 +675,7 @@ def _fused_h_res_h_post_bda_native(
         h_res, h_post and bda.
 
         When dropout is zero (or inference), uses a single fused/reference kernel
-        for H_res @ residual + H_post * (x + bias). Falls back to unfused
+        for H_res.T @ residual + H_post * (x + bias). Falls back to unfused
         implementation when dropout is needed.
 
         Args:

megatron/core/transformer/transformer_config.py

@@ -1097,13 +1097,17 @@ class TransformerConfig(ModelParallelConfig):
     """Initial value of Gating Factor (alpha in paper)."""
 
     use_fused_mhc: bool = False
-    """Use cuTile fused kernels for mHC operations.
+    """Use unified fused kernels for mHC operations.
 
     When True, attempts to replace the reference mHC modules (SinkhornKnopp,
-    H_aggregate, H_post_bda, ProjRms) with fused cuda.tile (cuTile) autograd
-    functions for better performance on supported GPUs.  Requires cuTile to be
-    installed; if cuTile is unavailable the flag is silently reset to False and
-    a warning is emitted.
+    H_aggregate, H_post_bda, ProjRms) with fused/autograd implementations for
+    better performance on supported GPUs.  Backend selection is internal and
+    op-specific: Triton for Sinkhorn and H_post_bda backward when available,
+    cuTile for the remaining fused kernels when available, then native torch
+    fallback. If every mHC operation uses the native torch fallback,
+    use_fused_mhc remains enabled and a rank-0 warning is emitted. The all-native
+    fallback is functionally equivalent, but may not provide fused backend
+    performance benefits.
     """
 
     mhc_recompute_layer_num: Optional[int] = None
@@ -1970,23 +1974,6 @@ def __post_init__(self):
         if self.use_fused_mhc:
             if not self.enable_hyper_connections:
                 raise ValueError("use_fused_mhc requires enable_hyper_connections=True.")
-            try:
-                from megatron.core.fusions.fused_mhc_kernels import is_cutile_available
-
-                if not is_cutile_available():
-                    warnings.warn(
-                        "use_fused_mhc is enabled but cuda.tile (cuTile) is not installed. "
-                        "Falling back to reference mHC implementations.",
-                        UserWarning,
-                    )
-                    self.use_fused_mhc = False
-            except ImportError:
-                warnings.warn(
-                    "use_fused_mhc is enabled but fused_mhc_kernels module could not be "
-                    "imported. Falling back to reference mHC implementations.",
-                    UserWarning,
-                )
-                self.use_fused_mhc = False
 
         if self.fine_grained_activation_offloading:
             assert (