diff --git a/hybrid_builders.py b/hybrid_builders.py index 7e1c58682ac..ae337237cb7 100644 --- a/hybrid_builders.py +++ b/hybrid_builders.py @@ -1,12 +1,12 @@ # Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. -from model_provider import count_parameters_in_layer +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module +from megatron.core.transformer.spec_utils import ModuleSpec, import_module from megatron.training import print_rank_0 from megatron.training.arguments import core_transformer_config_from_args -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec +from model_provider import count_parameters_in_layer def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None): @@ -21,6 +21,8 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, ), "inference_fuse_tp_communication is not supported for HybridModel" elif args.spec is not None: hybrid_stack_spec = import_module(args.spec) + if callable(hybrid_stack_spec) and not isinstance(hybrid_stack_spec, ModuleSpec): + hybrid_stack_spec = hybrid_stack_spec(config) else: raise ValueError("You must provide a valid hybrid layer spec via --spec") diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 895d46e9b3d..fc79cc9a7db 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -871,6 +871,29 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool return fp8_context + def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False): + """Return a context manager that disables TE quantization. + + Use this around submodule construction or execution that must stay in a higher + precision while its enclosing module uses an FP8 or FP4 context. + + Args: + config: Transformer configuration that controls quantization. + is_init: Whether to disable the parameter-initialization context instead of + the forward autocast context. + + Returns: + A disabled TE quantization context when quantization is active, otherwise a + no-op context. + """ + if is_init: + if not (config.fp8_param or config.fp4_param): + return nullcontext() + return transformer_engine.pytorch.fp8_model_init(enabled=False) + if not (config.fp8 or config.fp4): + return nullcontext() + return transformer_engine.pytorch.fp8_autocast(enabled=False) + else: def get_fp8_recipe(config: TransformerConfig): @@ -881,6 +904,10 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool """Returns dummy fp8 context manager since TE is not available.""" return nullcontext() + def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False): + """Return a no-op context manager since TE is not available.""" + return nullcontext() + if HAVE_TE: from transformer_engine.pytorch.fp8 import FP8GlobalStateManager diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py index 2eb4007f75c..92466ba7391 100644 --- a/megatron/core/fusions/fused_bias_dropout.py +++ b/megatron/core/fusions/fused_bias_dropout.py @@ -1,10 +1,13 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. -from typing import Optional, Tuple +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from typing import TYPE_CHECKING, Optional, Tuple import torch from megatron.core.jit import jit_fuser +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + # pylint: disable=missing-function-docstring @@ -80,7 +83,26 @@ def bias_dropout_add_fused_inference( return _bias_dropout_add_func(x_with_bias, residual, prob, False) -def get_bias_dropout_add(training, fused): +def get_bias_dropout_add( + training, fused, mhc_recompute_manager: Optional['CheckpointWithoutOutputManager'] = None +): + """ + Get the bias-dropout-add function. + + Args: + training: Whether in training mode. + fused: Whether to use fused implementation. + mhc_recompute_manager: Optional CheckpointWithoutOutputManager for checkpoint management. + When provided, the returned function will wrap the BDA operation with + CheckpointWithoutOutput for memory-efficient recomputation. + + Returns: + A callable that performs bias-dropout-add operation. + """ + if mhc_recompute_manager is not None: + # Return a checkpointed version that handles tuple unpacking internally + return _get_checkpointed_bda(training, fused, mhc_recompute_manager) + if fused: # jit scripting for a nn.module (with dropout) is not # triggering the fusion kernel. For now, we use two @@ -92,3 +114,68 @@ def get_bias_dropout_add(training, fused): return bias_dropout_add_fused_inference else: return bias_dropout_add_unfused(training) + + +def _get_checkpointed_bda(training, fused, mhc_recompute_manager: 'CheckpointWithoutOutputManager'): + """ + Create a checkpointed bias-dropout-add function. + + This function handles: + 1. Tuple unpacking for x_with_bias (required because save_for_backward can't save tuples) + 2. Non-tensor arguments like dropout probability (handled by CheckpointWithoutOutput) + 3. Auto-registration to the CheckpointWithoutOutputManager + + Args: + training: Whether in training mode. + fused: Whether to use fused implementation. + mhc_recompute_manager: CheckpointWithoutOutputManager for checkpoint management. + + Returns: + A callable that performs checkpointed bias-dropout-add operation. + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Get the underlying BDA function + if fused: + if training: + bda_func = bias_dropout_add_fused_train + else: + bda_func = bias_dropout_add_fused_inference + else: + bda_func = bias_dropout_add_unfused(training) + + def _checkpointed_bda(x_with_bias, residual, prob): + """ + Checkpointed BDA that handles tuple unpacking internally. + + Args: + x_with_bias: Either a tuple (x, bias) or a single tensor x. + residual: Residual tensor. + prob: Dropout probability. + + Returns: + Output tensor after bias-dropout-add. + """ + # Create checkpoint with manager + ckpt = CheckpointWithoutOutput(ckpt_manager=mhc_recompute_manager) + + # Handle case where x_with_bias might be a single tensor (e.g., from IdentityOp) + if isinstance(x_with_bias, tuple): + x, bias = x_with_bias + else: + x = x_with_bias + bias = None + + # Wrapper function that re-packs the tuple for the actual BDA function + def _bda_wrapper(output, bias, res, dropout): + return bda_func((output, bias), res, dropout) + + # Call checkpoint with unpacked arguments + result = ckpt.checkpoint(_bda_wrapper, x, bias, residual, prob) + + # No-op when manager is set - manager handles all discarding uniformly + ckpt.discard_output_and_register_recompute(result) + + return result + + return _checkpointed_bda diff --git a/megatron/core/fusions/fused_mhc_kernels.py b/megatron/core/fusions/fused_mhc_kernels.py new file mode 100644 index 00000000000..f94d0cafbf2 --- /dev/null +++ b/megatron/core/fusions/fused_mhc_kernels.py @@ -0,0 +1,3028 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Fused kernels for mHC (Manifold-Constrained Hyper-Connections). + +Uses Triton and cuda.tile (cuTile) kernels when available, with PyTorch +reference implementations as fallback. Reference (non-fused) implementations +live in ``megatron.core.transformer.hyper_connection`` and are used when fused +kernels are unavailable or when the ``use_fused_mhc`` config flag is False. + +Four fused operations: + - sinkhorn: Sinkhorn-Knopp projection to doubly stochastic matrix + - h_aggregate: weighted n-stream -> 1-stream aggregation + - h_post_bda: fused H_res.T @ residual + H_post * (x + bias) + - proj_rms_compute_h: fused projection + RMS normalization + compute_h +""" + +import logging +import math +import os +import shutil +import subprocess +import warnings +from typing import Optional, Tuple + +import torch +from torch import Tensor + +from megatron.core._rank_utils import log_single_rank, safe_get_rank + +logger = logging.getLogger(__name__) +LOG2E = math.log2(math.e) + + +def _env_flag(name: str) -> bool: + return os.getenv(name, "0").lower() in ("1", "true", "yes", "on") + + +def _forced_backend() -> Tuple[str, Optional[Exception]]: + value = os.getenv("MHC_FORCE_BACKEND", "auto").strip().lower() + value = value.replace("-", "_").replace("+", "_") + aliases = { + "auto": "auto", + "mixed": "auto", + "default": "auto", + "native": "native", + "torch": "native", + "pytorch": "native", + "none": "native", + "triton": "triton", + "triton_native": "triton", + "cutile": "cutile", + "cu_tile": "cutile", + "cuda_tile": "cutile", + } + if value not in aliases: + valid = ", ".join(sorted(aliases)) + return "auto", ValueError( + f"Unsupported MHC_FORCE_BACKEND={value!r}; expected one of: {valid}" + ) + return aliases[value], None + + +# --------------------------------------------------------------------------- +# Check cuTile availability +# --------------------------------------------------------------------------- +_CUTILE_AVAILABLE = False +_CUTILE_EXPERIMENTAL_AVAILABLE = False +_CUTILE_DEVICE_SUPPORT_CACHE: Optional[bool] = None +_CUTILE_DEVICE_SUPPORT_ERROR: Optional[str] = None +try: + import cuda.tile as ct + + _CUTILE_AVAILABLE = True + try: + import cuda.tile_experimental as ct_experimental + + _CUTILE_EXPERIMENTAL_AVAILABLE = True + except ImportError: + pass +except ImportError: + pass + + +# --------------------------------------------------------------------------- +# Check Triton availability +# --------------------------------------------------------------------------- +_TRITON_AVAILABLE = False +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + pass + + +_MHC_FORCED_BACKEND, _MHC_BACKEND_VALIDATION_ERROR = _forced_backend() + + +def _record_mhc_backend_validation_error(error: Exception) -> None: + global _MHC_BACKEND_VALIDATION_ERROR + if _MHC_BACKEND_VALIDATION_ERROR is None: + _MHC_BACKEND_VALIDATION_ERROR = error + + +def _raise_mhc_backend_validation_error() -> None: + if _MHC_BACKEND_VALIDATION_ERROR is not None: + raise _MHC_BACKEND_VALIDATION_ERROR + if _MHC_FORCED_BACKEND == "cutile" and not is_cutile_available(): + raise RuntimeError( + "MHC_FORCE_BACKEND=cutile was requested, but cuTile does not support " + f"the current device: {_CUTILE_DEVICE_SUPPORT_ERROR}" + ) + + +if _MHC_FORCED_BACKEND == "native": + _TRITON_AVAILABLE = False + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False +elif _MHC_FORCED_BACKEND == "triton": + if not _TRITON_AVAILABLE: + _record_mhc_backend_validation_error( + RuntimeError("MHC_FORCE_BACKEND=triton was requested, but Triton is not available") + ) + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False +elif _MHC_FORCED_BACKEND == "cutile": + if not _CUTILE_AVAILABLE: + _record_mhc_backend_validation_error( + RuntimeError("MHC_FORCE_BACKEND=cutile was requested, but cuTile is not available") + ) + _TRITON_AVAILABLE = False + +if _env_flag("MHC_DISABLE_TRITON"): + if _MHC_FORCED_BACKEND == "triton": + _record_mhc_backend_validation_error( + ValueError("MHC_FORCE_BACKEND=triton conflicts with MHC_DISABLE_TRITON=1") + ) + _TRITON_AVAILABLE = False + +if _env_flag("MHC_DISABLE_CUTILE"): + if _MHC_FORCED_BACKEND == "cutile": + _record_mhc_backend_validation_error( + ValueError("MHC_FORCE_BACKEND=cutile conflicts with MHC_DISABLE_CUTILE=1") + ) + _CUTILE_AVAILABLE = False + _CUTILE_EXPERIMENTAL_AVAILABLE = False + + +def is_cutile_available() -> bool: + """Return True if cuTile fused kernels are enabled.""" + return _CUTILE_AVAILABLE and _cutile_supports_current_device() + + +def _get_tileiras_path() -> Optional[str]: + """Return the tileiras compiler path if it can be found.""" + tileiras = shutil.which("tileiras") + if tileiras is not None: + return tileiras + + cuda_home = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or "/usr/local/cuda" + candidate = os.path.join(cuda_home, "bin", "tileiras") + if os.path.exists(candidate): + return candidate + return None + + +def _cutile_supports_current_device() -> bool: + """Return whether cuTile can compile for the current CUDA device.""" + global _CUTILE_DEVICE_SUPPORT_CACHE, _CUTILE_DEVICE_SUPPORT_ERROR + + if not _CUTILE_AVAILABLE: + return False + if _CUTILE_DEVICE_SUPPORT_CACHE is not None: + return _CUTILE_DEVICE_SUPPORT_CACHE + + if not torch.cuda.is_available(): + _CUTILE_DEVICE_SUPPORT_ERROR = "CUDA is not available" + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + major, minor = torch.cuda.get_device_capability() + arch = f"sm_{major}{minor}" + tileiras = _get_tileiras_path() + if tileiras is None: + _CUTILE_DEVICE_SUPPORT_ERROR = "tileiras compiler was not found" + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + try: + result = subprocess.run( + [tileiras, "--gpu-name", arch], capture_output=True, check=False, text=True, timeout=10 + ) + except (OSError, subprocess.SubprocessError) as exc: + _CUTILE_DEVICE_SUPPORT_ERROR = str(exc) + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + output = f"{result.stdout}\n{result.stderr}" + if "Cannot find option named" in output and arch in output: + _CUTILE_DEVICE_SUPPORT_ERROR = output.strip().splitlines()[0] + _CUTILE_DEVICE_SUPPORT_CACHE = False + return False + + _CUTILE_DEVICE_SUPPORT_CACHE = True + return True + + +def is_triton_available() -> bool: + """Return True if Triton is enabled for supported mHC kernels.""" + return _TRITON_AVAILABLE + + +# ============================================================================ +# Triton implementations (only defined when triton is available) +# ============================================================================ + +if _TRITON_AVAILABLE: + TLOG2E = tl.constexpr(LOG2E) + + # ============================================================================ + # Sinkhorn-Knopp + # ============================================================================ + + @triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], key=["HC", "NUM_ITERS"] + ) + @triton.jit + def _triton_sinkhorn_fwd_kernel( + inp_ptr, out_ptr, M_init_ptr, N_batch, eps, HC: tl.constexpr, NUM_ITERS: tl.constexpr + ): + """Grid: (N_batch,). Each program handles one [HC, HC] matrix.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + logits = tl.load(inp_ptr + mat_ptrs).to(tl.float32) + row_max = tl.max(logits, axis=1) + # Subtract row_max before exp to keep the exponent numerically stable. + M = tl.exp2((logits - row_max[:, None]) * TLOG2E) + tl.store(M_init_ptr + mat_ptrs, M.to(M_init_ptr.dtype.element_ty)) + + row_sum = tl.sum(M, axis=1) + M = M / row_sum[:, None] + eps + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + for _ in range(NUM_ITERS - 1): + row_sum = tl.sum(M, axis=1) + M = M / (row_sum[:, None] + eps) + col_sum = tl.sum(M, axis=0) + M = M / (col_sum[None, :] + eps) + + tl.store(out_ptr + mat_ptrs, M.to(out_ptr.dtype.element_ty)) + + @triton.autotune( + configs=[triton.Config({}, num_warps=nw) for nw in (1, 2, 4, 8)], key=["HC", "NUM_ITERS"] + ) + @triton.jit + def _triton_sinkhorn_bwd_kernel( + grad_out_ptr, + M_init_ptr, + grad_inp_ptr, + ws_M_ptr, + ws_rs_ptr, + ws_cs_ptr, + N_batch, + eps, + HC: tl.constexpr, + NUM_ITERS: tl.constexpr, + ): + """Grid: (N_batch,). Each program handles one [HC, HC] backward.""" + pid = tl.program_id(0) + if pid >= N_batch: + return + + base = pid * HC * HC + M_ws_base = pid * 2 * NUM_ITERS * HC * HC + v_ws_base = pid * NUM_ITERS + offs_r = tl.arange(0, HC) + offs_c = tl.arange(0, HC) + mat_ptrs = base + offs_r[:, None] * HC + offs_c[None, :] + + M = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + for t in range(NUM_ITERS): + ws_off = M_ws_base + (2 * t) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + row_sum = tl.sum(M, axis=1) + tl.store(ws_rs_ptr + (v_ws_base + t) * HC + offs_r, row_sum) + if t == 0: + M = M / row_sum[:, None] + eps + else: + M = M / (row_sum[:, None] + eps) + + ws_off = M_ws_base + (2 * t + 1) * HC * HC + tl.store(ws_M_ptr + ws_off + offs_r[:, None] * HC + offs_c[None, :], M) + + col_sum = tl.sum(M, axis=0) + tl.store(ws_cs_ptr + (v_ws_base + t) * HC + offs_c, col_sum) + M = M / (col_sum[None, :] + eps) + + # M is the final forward output. It is the right value for the first VJP + # through the last column-normalization step. + grad = tl.load(grad_out_ptr + mat_ptrs).to(tl.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + + col_s = tl.load(ws_cs_ptr + (v_ws_base + t) * HC + offs_c).to(tl.float32) + grad = grad / (col_s[None, :] + eps) + col_corr = tl.sum(grad * M, axis=0) + grad = grad - col_corr[None, :] + M = tl.load( + ws_M_ptr + + M_ws_base + + (2 * t + 1) * HC * HC + + offs_r[:, None] * HC + + offs_c[None, :] + ).to(tl.float32) + + row_s = tl.load(ws_rs_ptr + (v_ws_base + t) * HC + offs_r).to(tl.float32) + if t == 0: + grad = grad / row_s[:, None] + row_corr = tl.sum(grad * (M - eps), axis=1) + else: + grad = grad / (row_s[:, None] + eps) + row_corr = tl.sum(grad * M, axis=1) + grad = grad - row_corr[:, None] + M = tl.load( + ws_M_ptr + M_ws_base + (2 * t) * HC * HC + offs_r[:, None] * HC + offs_c[None, :] + ).to(tl.float32) + + M_init = tl.load(M_init_ptr + mat_ptrs).to(tl.float32) + grad = grad * M_init + tl.store(grad_inp_ptr + mat_ptrs, grad.to(grad_inp_ptr.dtype.element_ty)) + + def _triton_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + dev = input_logits.device + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + inp = input_logits.contiguous().view(N_batch, hc, hc) + _triton_sinkhorn_fwd_kernel[(N_batch,)](inp, out, M_init, N_batch, eps, hc, num_iterations) + return out.view(original_shape), M_init.view(original_shape) + + def _triton_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + dev = grad_output.device + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + go = grad_output.contiguous().view(N_batch, hc, hc) + mi = M_init.contiguous().view(N_batch, hc, hc) + ws_M = torch.empty(N_batch * 2 * num_iterations * hc * hc, dtype=torch.float32, device=dev) + ws_rs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations * hc, dtype=torch.float32, device=dev) + _triton_sinkhorn_bwd_kernel[(N_batch,)]( + go, mi, grad_input, ws_M, ws_rs, ws_cs, N_batch, eps, hc, num_iterations + ) + return grad_input.view(original_shape) + + class TritonFusedSinkhorn(torch.autograd.Function): + """Autograd wrapper for Triton fused Sinkhorn.""" + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """Run Triton Sinkhorn forward and save initial matrix for backward.""" + out, M_init = _triton_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, grad_output: Tensor): + """Run Triton Sinkhorn backward.""" + (M_init,) = ctx.saved_tensors + grad_input = _triton_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + def triton_fused_sinkhorn( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + """Apply Triton fused Sinkhorn with autograd support.""" + return TritonFusedSinkhorn.apply(input_logits, num_iterations, eps) + + # ============================================================================ + # H_aggregate forward + # ============================================================================ + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_h_agg_fwd_kernel( + x_ptr, + h_ptr, + out_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_x_s, + stride_x_n, + stride_x_c, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """out[s, c] = sum_i x[s, i, c] * h[s, i].""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for i in tl.static_range(N): + x_i = tl.load( + x_ptr + offs_s[:, None] * stride_x_s + i * stride_x_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + h_i = tl.load(h_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + acc += h_i[:, None] * x_i + tl.store( + out_ptr + offs_s[:, None] * C + offs_c[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + def _triton_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + x_flat = x.contiguous().view(sb, n, C) + h_flat = h_pre.contiguous().view(sb, n) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_h_agg_fwd_kernel[grid]( + x_flat, h_flat, out, sb, C, n, x_flat.stride(0), x_flat.stride(1), x_flat.stride(2) + ) + return out.view(s, b, C) + + # ============================================================================ + # H_post BDA + # ============================================================================ + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_fwd_kernel( + hr_ptr, + orig_ptr, + hp_ptr, + x_ptr, + bias_ptr, + out_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_hr_s, + stride_hr_i, + stride_hr_j, + stride_orig_s, + stride_orig_n, + stride_orig_c, + stride_out_s, + stride_out_n, + stride_out_c, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """out = hr.T @ orig + hp * (x + bias).""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load(x_ptr + offs_s[:, None] * C + offs_c[None, :], mask=mask_2d, other=0.0).to( + tl.float32 + ) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + hp_i = tl.load(hp_ptr + offs_s * N + i, mask=mask_s, other=0.0).to(tl.float32) + out_i = hp_i[:, None] * x_tile + + for j in tl.static_range(N): + hr_ji = tl.load( + hr_ptr + offs_s * stride_hr_s + j * stride_hr_i + i * stride_hr_j, + mask=mask_s, + other=0.0, + ).to(tl.float32) + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + out_i += hr_ji[:, None] * orig_j + + tl.store( + out_ptr + offs_s[:, None] * stride_out_s + i * stride_out_n + offs_c[None, :], + out_i.to(out_ptr.dtype.element_ty), + mask=mask_2d, + ) + + def _triton_h_post_bda_fwd( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] + ) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=dev) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_fwd_kernel[grid]( + hr_flat, + orig_flat, + hp_flat, + x_flat, + bias if bias is not None else x_flat, + out, + sb, + C, + n, + hr_flat.stride(0), + hr_flat.stride(1), + hr_flat.stride(2), + orig_flat.stride(0), + orig_flat.stride(1), + orig_flat.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + HAS_BIAS=(bias is not None), + ) + return out.view(s, b, n, C) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_bwd_g_x_orig_kernel( + go_ptr, + hr_ptr, + hp_ptr, + g_orig_ptr, + g_x_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_go_s, + stride_go_n, + stride_go_c, + stride_hr_s, + stride_hr_i, + stride_hr_j, + stride_orig_s, + stride_orig_n, + stride_orig_c, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """g_x = hp @ go, g_orig = hr @ go.""" + pid_s = tl.program_id(0) + pid_c = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + mask_s = offs_s < sb + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + g_x_acc = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + hp_j = tl.load(hp_ptr + offs_s * N + j, mask=mask_s, other=0.0).to(tl.float32) + g_x_acc += hp_j[:, None] * go_j + tl.store( + g_x_ptr + offs_s[:, None] * C + offs_c[None, :], + g_x_acc.to(g_x_ptr.dtype.element_ty), + mask=mask_2d, + ) + + for i in tl.static_range(N): + g_orig_i = tl.zeros((BLOCK_S, BLOCK_C), dtype=tl.float32) + for j in tl.static_range(N): + go_j = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + j * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + hr_ij = tl.load( + hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + mask=mask_s, + other=0.0, + ).to(tl.float32) + g_orig_i += hr_ij[:, None] * go_j + tl.store( + g_orig_ptr + offs_s[:, None] * stride_orig_s + i * stride_orig_n + offs_c[None, :], + g_orig_i.to(g_orig_ptr.dtype.element_ty), + mask=mask_2d, + ) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_C": bc, "BLOCK_S": bs}, num_warps=nw) + for bc in (64, 128, 256, 512) + for bs in (1, 2, 4, 8) + for nw in (2, 4, 8) + ], + key=["C", "N"], + ) + @triton.jit + def _triton_hpb_bwd_g_hp_hr_kernel( + go_ptr, + orig_ptr, + x_ptr, + bias_ptr, + g_hr_ptr, + g_hp_ptr, + sb, + C: tl.constexpr, + N: tl.constexpr, + stride_go_s, + stride_go_n, + stride_go_c, + stride_orig_s, + stride_orig_n, + stride_orig_c, + stride_hr_s, + stride_hr_i, + stride_hr_j, + HAS_BIAS: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_S: tl.constexpr, + ): + """g_hp = sum_c go*(x+bias), g_hr = orig @ go.T.""" + pid_s = tl.program_id(0) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + mask_s = offs_s < sb + + g_hp_acc = tl.zeros((BLOCK_S, N), dtype=tl.float32) + g_hr_acc = tl.zeros((BLOCK_S, N * N), dtype=tl.float32) + + for c_start in range(0, C, BLOCK_C): + offs_c = c_start + tl.arange(0, BLOCK_C) + mask_c = offs_c < C + mask_2d = mask_s[:, None] & mask_c[None, :] + + x_tile = tl.load( + x_ptr + offs_s[:, None] * C + offs_c[None, :], mask=mask_2d, other=0.0 + ).to(tl.float32) + if HAS_BIAS: + bias_tile = tl.load(bias_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + x_tile += bias_tile[None, :] + + for i in tl.static_range(N): + go_i = tl.load( + go_ptr + offs_s[:, None] * stride_go_s + i * stride_go_n + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hp = tl.sum(go_i * x_tile, axis=1) + g_hp_acc += tl.where( + tl.arange(0, N)[None, :] == i, + dot_hp[:, None], + tl.zeros((BLOCK_S, N), dtype=tl.float32), + ) + for j in tl.static_range(N): + orig_j = tl.load( + orig_ptr + + offs_s[:, None] * stride_orig_s + + j * stride_orig_n + + offs_c[None, :], + mask=mask_2d, + other=0.0, + ).to(tl.float32) + dot_hr = tl.sum(go_i * orig_j, axis=1) + g_hr_acc += tl.where( + tl.arange(0, N * N)[None, :] == j * N + i, + dot_hr[:, None], + tl.zeros((BLOCK_S, N * N), dtype=tl.float32), + ) + + offs_n = tl.arange(0, N) + tl.store( + g_hp_ptr + offs_s[:, None] * N + offs_n[None, :], + g_hp_acc.to(g_hp_ptr.dtype.element_ty), + mask=mask_s[:, None], + ) + + # N is expected to stay small for mHC, so this simple extraction is acceptable. + nn_offs = tl.arange(0, N * N) + for i in tl.static_range(N): + for j in tl.static_range(N): + col_mask = (nn_offs == (i * N + j)).to(tl.float32) + val = tl.sum(g_hr_acc * col_mask[None, :], axis=1) + tl.store( + g_hr_ptr + offs_s * stride_hr_s + i * stride_hr_i + j * stride_hr_j, + val.to(g_hr_ptr.dtype.element_ty), + mask=mask_s, + ) + + def _triton_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + dev = h_res.device + + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=dev) + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=dev) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=dev) + g_x = torch.empty(sb, C, dtype=x.dtype, device=dev) + + go_flat = grad_output.contiguous().view(sb, n, C) + hr_flat = h_res.contiguous().view(sb, n, n) + orig_flat = original_residual.contiguous().view(sb, n, C) + hp_flat = h_post.contiguous().view(sb, n) + x_flat = x.contiguous().view(sb, C) + + grid_a = lambda META: (triton.cdiv(sb, META["BLOCK_S"]), triton.cdiv(C, META["BLOCK_C"])) + _triton_hpb_bwd_g_x_orig_kernel[grid_a]( + go_flat, + hr_flat, + hp_flat, + g_res, + g_x, + sb, + C, + n, + go_flat.stride(0), + go_flat.stride(1), + go_flat.stride(2), + hr_flat.stride(0), + hr_flat.stride(1), + hr_flat.stride(2), + g_res.stride(0), + g_res.stride(1), + g_res.stride(2), + ) + + grid_b = lambda META: (triton.cdiv(sb, META["BLOCK_S"]),) + _triton_hpb_bwd_g_hp_hr_kernel[grid_b]( + go_flat, + orig_flat, + x_flat, + bias if bias is not None else x_flat, + g_hr, + g_hp, + sb, + C, + n, + go_flat.stride(0), + go_flat.stride(1), + go_flat.stride(2), + orig_flat.stride(0), + orig_flat.stride(1), + orig_flat.stride(2), + g_hr.stride(0), + g_hr.stride(1), + g_hr.stride(2), + HAS_BIAS=(bias is not None), + ) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + +_TRITON_IMPLS = ( + { + "sinkhorn": triton_fused_sinkhorn, + "h_aggregate_fwd": _triton_h_aggregate_fwd, + "h_post_bda_fwd": _triton_h_post_bda_fwd, + "h_post_bda_bwd": _triton_h_post_bda_bwd, + } + if _TRITON_AVAILABLE + else {"sinkhorn": None, "h_aggregate_fwd": None, "h_post_bda_fwd": None, "h_post_bda_bwd": None} +) + + +# ============================================================================ +# CuTile implementations (only defined when cuda.tile is available) +# ============================================================================ + +if _CUTILE_AVAILABLE: + ConstInt = ct.Constant[int] + PAD_ZERO = ct.PaddingMode.ZERO + + # -- Sinkhorn kernels ---------------------------------------------------- + + @ct.kernel + def _ct_sinkhorn_fwd_kernel( + inp, out, M_init_out, eps, HC: ConstInt, NUM_ITERS: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + logits = ct.load(inp, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + row_max = ct.max(logits, axis=2, keepdims=True) + M = ct.exp2((logits - row_max) * LOG2E) + ct.store( + M_init_out, + index=(pid, 0, 0), + tile=ct.reshape(M.astype(M_init_out.dtype), (TILE_SIZE, HC, HC)), + ) + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / row_sum + eps + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) + for _ in range(NUM_ITERS - 1): + row_sum = ct.sum(M, axis=2, keepdims=True) + M = M / (row_sum + eps) + col_sum = ct.sum(M, axis=1, keepdims=True) + M = M / (col_sum + eps) + ct.store(out, index=(pid, 0, 0), tile=ct.reshape(M.astype(out.dtype), (TILE_SIZE, HC, HC))) + + @ct.kernel + def _ct_sinkhorn_bwd_kernel( + grad_out, + M_init, + grad_inp, + ws_M, + ws_rs, + ws_cs, + eps, + HC: ConstInt, + NUM_ITERS: ConstInt, + TILE_SIZE: ConstInt, + ): + pid = ct.bid(0) + M_base = pid * (2 * NUM_ITERS) + v_base = pid * NUM_ITERS + + M = ct.load(M_init, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + for t in range(NUM_ITERS): + ct.store(ws_M, index=(M_base + 2 * t, 0, 0), tile=M) + row_sum = ct.sum(M, axis=2, keepdims=True) + ct.store(ws_rs, index=(v_base + t, 0, 0), tile=row_sum) + if t == 0: + M = M / row_sum + eps + else: + M = M / (row_sum + eps) + ct.store(ws_M, index=(M_base + 2 * t + 1, 0, 0), tile=M) + col_sum = ct.sum(M, axis=1, keepdims=True) + ct.store(ws_cs, index=(v_base + t, 0, 0), tile=col_sum) + M = M / (col_sum + eps) + + grad = ct.load(grad_out, index=(pid, 0, 0), shape=(TILE_SIZE, HC, HC)).astype(ct.float32) + for t_rev in range(NUM_ITERS): + t = NUM_ITERS - 1 - t_rev + col_s = ct.load(ws_cs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, 1, HC)) + grad = grad / (col_s + eps) + col_corr = ct.sum(grad * M, axis=1, keepdims=True) + grad = grad - col_corr + M = ct.load(ws_M, index=(M_base + 2 * t + 1, 0, 0), shape=(TILE_SIZE, HC, HC)) + row_s = ct.load(ws_rs, index=(v_base + t, 0, 0), shape=(TILE_SIZE, HC, 1)) + if t == 0: + grad = grad / row_s + row_corr = ct.sum(grad * (M - eps), axis=2, keepdims=True) + else: + grad = grad / (row_s + eps) + row_corr = ct.sum(grad * M, axis=2, keepdims=True) + grad = grad - row_corr + M = ct.load(ws_M, index=(M_base + 2 * t, 0, 0), shape=(TILE_SIZE, HC, HC)) + grad = grad * M + ct.store(grad_inp, index=(pid, 0, 0), tile=grad.astype(grad_inp.dtype)) + + def _sinkhorn_autotune_tile_sizes(N_batch): + """Generate autotune search space for sinkhorn kernels.""" + for ts in (1, 2, 4, 8, 16, 32, 64, 128): + if ts <= N_batch: + yield ts + + _sinkhorn_fwd_best_cfg: dict = {} + _sinkhorn_bwd_best_cfg: dict = {} + + def _cutile_sinkhorn_fwd( + input_logits: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tuple[Tensor, Tensor]: + original_shape = input_logits.shape + hc = original_shape[-1] + N_batch = input_logits.numel() // (hc * hc) + dev = input_logits.device + stream = torch.cuda.current_stream() + out = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + M_init = torch.empty(N_batch, hc, hc, dtype=input_logits.dtype, device=dev) + inp = input_logits.view(N_batch, hc, hc) + + cache_key = (N_batch, hc, num_iterations) + cached = _sinkhorn_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + ts = cached if cached is not None else math.gcd(N_batch, 128) + ct.launch( + stream, + (math.ceil(N_batch / ts), 1, 1), + _ct_sinkhorn_fwd_kernel, + (inp, out, M_init, eps, hc, num_iterations, ts), + ) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_SIZE=ts) for ts in _sinkhorn_autotune_tile_sizes(N_batch) + ] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), + kernel=_ct_sinkhorn_fwd_kernel, + args_fn=lambda cfg: (inp, out, M_init, eps, hc, num_iterations, cfg.TILE_SIZE), + search_space=configs, + ) + best_ts = tuned.tuned_config.TILE_SIZE + _sinkhorn_fwd_best_cfg[cache_key] = best_ts + ct.launch( + stream, + (math.ceil(N_batch / best_ts), 1, 1), + _ct_sinkhorn_fwd_kernel, + (inp, out, M_init, eps, hc, num_iterations, best_ts), + ) + + return out.view(original_shape), M_init.view(original_shape) + + def _cutile_sinkhorn_bwd( + grad_output: Tensor, M_init: Tensor, num_iterations: int, eps: float = 1e-6 + ) -> Tensor: + original_shape = grad_output.shape + hc = original_shape[-1] + N_batch = grad_output.numel() // (hc * hc) + dev = grad_output.device + stream = torch.cuda.current_stream() + grad_input = torch.empty(N_batch, hc, hc, dtype=grad_output.dtype, device=dev) + go = grad_output.view(N_batch, hc, hc) + mi = M_init.view(N_batch, hc, hc) + + cache_key = (N_batch, hc, num_iterations) + cached = _sinkhorn_bwd_best_cfg.get(cache_key) + + def _alloc_and_launch(ts): + ws_M = torch.empty( + N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev + ) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + ct.launch( + stream, + (math.ceil(N_batch / ts), 1, 1), + _ct_sinkhorn_bwd_kernel, + (go, mi, grad_input, ws_M, ws_rs, ws_cs, eps, hc, num_iterations, ts), + ) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + ts = cached if cached is not None else math.gcd(N_batch, 128) + _alloc_and_launch(ts) + else: + from types import SimpleNamespace + + configs = [ + SimpleNamespace(TILE_SIZE=ts) for ts in _sinkhorn_autotune_tile_sizes(N_batch) + ] + # Allocate workspace for largest tile size (all configs share same workspace shape). + ws_M = torch.empty( + N_batch * 2 * num_iterations, hc, hc, dtype=torch.float32, device=dev + ) + ws_rs = torch.empty(N_batch * num_iterations, hc, 1, dtype=torch.float32, device=dev) + ws_cs = torch.empty(N_batch * num_iterations, 1, hc, dtype=torch.float32, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(N_batch / cfg.TILE_SIZE), 1, 1), + kernel=_ct_sinkhorn_bwd_kernel, + args_fn=lambda cfg: ( + go, + mi, + grad_input, + ws_M, + ws_rs, + ws_cs, + eps, + hc, + num_iterations, + cfg.TILE_SIZE, + ), + search_space=configs, + ) + best_ts = tuned.tuned_config.TILE_SIZE + _sinkhorn_bwd_best_cfg[cache_key] = best_ts + # Re-launch with best config. + _alloc_and_launch(best_ts) + + return grad_input.view(original_shape) + + # -- H_aggregate kernels ------------------------------------------------- + + @ct.kernel + def _ct_h_agg_fwd_kernel(x, h_pre, out, N: ConstInt, TILE_M: ConstInt, TILE_C: ConstInt): + pid = ct.bid(0) + num_tiles = ct.num_tiles(x, axis=2, shape=(TILE_M, N, TILE_C)) + h_tile = ct.load(h_pre, index=(pid, 0), shape=(TILE_M, N), padding_mode=PAD_ZERO) + h_tile = ct.expand_dims(h_tile, axis=2) + for j in range(num_tiles): + x_tile = ct.load(x, index=(pid, 0, j), shape=(TILE_M, N, TILE_C), padding_mode=PAD_ZERO) + acc = ct.sum(x_tile * h_tile, axis=1).astype(ct.float32) + ct.store(out, index=(pid, j), tile=acc.astype(out.dtype)) + + @ct.kernel + def _ct_h_agg_bwd_kernel(go, x, h_pre, gx, gh, N: ConstInt, TILE_M: ConstInt, TILE_C: ConstInt): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(go, axis=1, shape=(TILE_M, TILE_C)) + h_tile = ct.load(h_pre, index=(pid, 0), shape=(TILE_M, N), padding_mode=PAD_ZERO) + h_expanded = ct.expand_dims(h_tile, axis=2) + gh_acc = ct.full((TILE_M, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + go_tile = ct.load( + go, index=(pid, ct_idx), shape=(TILE_M, TILE_C), padding_mode=PAD_ZERO + ) + go_expanded = ct.expand_dims(go_tile, axis=1) + x_tile = ct.load( + x, index=(pid, 0, ct_idx), shape=(TILE_M, N, TILE_C), padding_mode=PAD_ZERO + ) + gx_tile = go_expanded * h_expanded + ct.store(gx, index=(pid, 0, ct_idx), tile=gx_tile.astype(gx.dtype)) + # Reduce in fp32: the torch reference evaluates this product in fp32 + # under torch.compile, and a bf16 product makes grad_h ~3x noisier. + gh_acc += ct.sum(go_expanded.astype(ct.float32) * x_tile.astype(ct.float32), axis=2) + ct.store(gh, index=(pid, 0), tile=gh_acc.astype(gh.dtype)) + + def _cutile_h_aggregate_fwd(x: Tensor, h_pre: Tensor) -> Tensor: + s, b, n, C = x.shape + sb = s * b + stream = torch.cuda.current_stream() + out = torch.empty(sb, C, dtype=x.dtype, device=x.device) + x_flat = x.view(sb, n, C) + h_flat = h_pre.view(sb, n) + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + ct.launch( + stream, (math.ceil(sb / tm),), _ct_h_agg_fwd_kernel, (x_flat, h_flat, out, n, tm, tc) + ) + + return out.view(s, b, C) + + def _cutile_h_aggregate_bwd( + grad_output: Tensor, x: Tensor, h_pre: Tensor + ) -> Tuple[Tensor, Tensor]: + s, b, n, C = x.shape + sb = s * b + stream = torch.cuda.current_stream() + gx = torch.empty(sb, n, C, dtype=x.dtype, device=x.device) + gh = torch.empty(sb, n, dtype=h_pre.dtype, device=x.device) + go_flat = grad_output.view(sb, C) + x_flat = x.view(sb, n, C) + h_flat = h_pre.view(sb, n) + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + tm, tc = math.gcd(sb, 4), math.gcd(C, 1024) + ct.launch( + stream, + (math.ceil(sb / tm),), + _ct_h_agg_bwd_kernel, + (go_flat, x_flat, h_flat, gx, gh, n, tm, tc), + ) + + return gx.view(s, b, n, C), gh.view(s, b, n) + + # -- H_post BDA kernels -------------------------------------------------- + + @ct.kernel + def _ct_hpb_fwd_kernel( + hr, orig, hp, x, out, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hp_exp = ct.expand_dims(hp_tile, axis=2) # (TILE_SIZE, N, 1) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + for ct_idx in range(num_c_tiles): + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + x_exp = ct.expand_dims(x_tile, axis=1) # (TILE_SIZE, 1, TILE_C) + out_tile = hp_exp * x_exp # (TILE_SIZE, N, TILE_C) + for j in range(N): + hr_row = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + hr_col = ct.reshape(hr_row, (TILE_SIZE, N, 1)) + orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + out_tile = out_tile + hr_col * orig_row + ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) + + @ct.kernel + def _ct_hpb_fwd_bias_kernel( + hr, orig, hp, x, bias, out, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + pid = ct.bid(0) + num_c_tiles = ct.num_tiles(x, axis=1, shape=(TILE_SIZE, TILE_C)) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hp_exp = ct.expand_dims(hp_tile, axis=2) # (TILE_SIZE, N, 1) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + for ct_idx in range(num_c_tiles): + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) + xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # (TILE_SIZE, 1, TILE_C) + out_tile = hp_exp * xb_exp # (TILE_SIZE, N, TILE_C) + for j in range(N): + hr_row = ct.extract(hr_tile, (0, j, 0), shape=(TILE_SIZE, 1, N)) + hr_col = ct.reshape(hr_row, (TILE_SIZE, N, 1)) + orig_row = ct.extract(orig_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + out_tile = out_tile + hr_col * orig_row + ct.store(out, index=(pid, 0, ct_idx), tile=out_tile.astype(out.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_x_orig_kernel( + go, hr, hp, g_orig, g_x, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + """Compute g_x = hp @ go and g_orig = hr @ go. + + Grid: (ceil(sb / TILE_SIZE), ceil(C / TILE_C)). + 2D grid — no loop, no accumulators. + """ + pid = ct.bid(0) + ct_idx = ct.bid(1) + hp_tile = ct.load(hp, index=(pid, 0), shape=(TILE_SIZE, N), padding_mode=PAD_ZERO) + hr_tile = ct.load(hr, index=(pid, 0, 0), shape=(TILE_SIZE, N, N), padding_mode=PAD_ZERO) + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + g_x_tile = ct.full((TILE_SIZE, 1, TILE_C), 0, dtype=ct.float32) + g_orig_tile = ct.full((TILE_SIZE, N, TILE_C), 0, dtype=ct.float32) + for j in range(N): + hp_j = ct.extract(hp_tile, (0, j), shape=(TILE_SIZE, 1)) + hp_j_exp = ct.expand_dims(hp_j, axis=2) # [TS, 1, 1] + go_j = ct.extract(go_tile, (0, j, 0), shape=(TILE_SIZE, 1, TILE_C)) + g_x_tile = g_x_tile + hp_j_exp * go_j + hr_col_j = ct.extract(hr_tile, (0, 0, j), shape=(TILE_SIZE, N, 1)) + g_orig_tile = g_orig_tile + hr_col_j * go_j + ct.store( + g_x, + index=(pid, ct_idx), + tile=ct.reshape(g_x_tile, (TILE_SIZE, TILE_C)).astype(g_x.dtype), + ) + ct.store(g_orig, index=(pid, 0, ct_idx), tile=g_orig_tile.astype(g_orig.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_hp_hr_kernel( + go, orig, x, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + """Compute g_hp = sum(go * x) and g_hr = orig @ go.T (no bias). + + Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. + """ + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) + acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + x_exp = ct.expand_dims(x_tile, axis=1) # [TS, 1, TC] + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + acc_g_hp = acc_g_hp + ct.sum(go_tile * x_exp, axis=2, keepdims=True) + acc_g_hr = acc_g_hr + ct.sum( + ct.expand_dims(orig_tile, axis=2) * ct.expand_dims(go_tile, axis=1), axis=3 + ) + ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) + ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) + + @ct.kernel + def _ct_hpb_bwd_g_hp_hr_bias_kernel( + go, orig, x, bias, g_hr, g_hp, N: ConstInt, TILE_C: ConstInt, TILE_SIZE: ConstInt + ): + """Compute g_hp = sum(go * (x+bias)) and g_hr = orig @ go.T (with bias). + + Grid: (ceil(sb / TILE_SIZE),). Loops over C-tiles. + """ + pid = ct.bid(0) + num_c_tiles = ct.cdiv(go.shape[2], TILE_C) + acc_g_hp = ct.full((TILE_SIZE, N, 1), 0, dtype=ct.float32) + acc_g_hr = ct.full((TILE_SIZE, N, N), 0, dtype=ct.float32) + for ct_idx in range(num_c_tiles): + x_tile = ct.load( + x, index=(pid, ct_idx), shape=(TILE_SIZE, TILE_C), padding_mode=PAD_ZERO + ) + bias_tile = ct.load(bias, index=(ct_idx,), shape=(TILE_C,), padding_mode=PAD_ZERO) + xb_exp = ct.expand_dims(x_tile + bias_tile, axis=1) # [TS, 1, TC] + go_tile = ct.load( + go, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + orig_tile = ct.load( + orig, index=(pid, 0, ct_idx), shape=(TILE_SIZE, N, TILE_C), padding_mode=PAD_ZERO + ) + acc_g_hp = acc_g_hp + ct.sum(go_tile * xb_exp, axis=2, keepdims=True) + acc_g_hr = acc_g_hr + ct.sum( + ct.expand_dims(orig_tile, axis=2) * ct.expand_dims(go_tile, axis=1), axis=3 + ) + ct.store(g_hp, index=(pid, 0), tile=ct.reshape(acc_g_hp, (TILE_SIZE, N)).astype(g_hp.dtype)) + ct.store(g_hr, index=(pid, 0, 0), tile=acc_g_hr.astype(g_hr.dtype)) + + # -- H_post BDA autotune configs & caches -------------------------------- + + def _hpb_autotune_configs(sb, C): + """Generate TILE_SIZE × TILE_C search space for h_post_bda kernels.""" + for tile_size in (1, 2, 4, 8): + for tile_c in (32, 64, 128, 256, 512, 1024): + if tile_c <= C and tile_size <= sb: + yield {"TILE_SIZE": tile_size, "TILE_C": tile_c} + + _hpb_fwd_best_cfg: dict = {} + _hpb_bwd_g_x_orig_best_cfg: dict = {} + _hpb_bwd_g_hp_hr_best_cfg: dict = {} + + def _cutile_h_post_bda_fwd( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] + ) -> Tensor: + s, b, n, C = original_residual.shape + sb = s * b + stream = torch.cuda.current_stream() + out = torch.empty(sb, n, C, dtype=h_res.dtype, device=h_res.device) + hr_flat = h_res.view(sb, n, n) + orig_flat = original_residual.view(sb, n, C) + hp_flat = h_post.view(sb, n) + x_flat = x.view(sb, C) + + cache_key = (sb, n, C, bias is not None) + cached = _hpb_fwd_best_cfg.get(cache_key) + kernel = _ct_hpb_fwd_bias_kernel if bias is not None else _ct_hpb_fwd_kernel + + # Autotune disabled — causes cudaErrorLaunchFailure during training. + if cached is not None: + ts, tc = cached + else: + ts, tc = 1, math.gcd(C, 1024) + args = (hr_flat, orig_flat, hp_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (out, n, tc, ts) + ct.launch(stream, (math.ceil(sb / ts),), kernel, args) + + return out.view(s, b, n, C) + + def _cutile_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + stream = torch.cuda.current_stream() + g_hr = torch.empty(sb, n, n, dtype=h_res.dtype, device=h_res.device) + g_res = torch.empty(sb, n, C, dtype=original_residual.dtype, device=h_res.device) + g_hp = torch.empty(sb, n, dtype=h_post.dtype, device=h_res.device) + g_x = torch.empty(sb, C, dtype=x.dtype, device=h_res.device) + go_flat = grad_output.view(sb, n, C) + hr_flat = h_res.view(sb, n, n) + orig_flat = original_residual.view(sb, n, C) + hp_flat = h_post.view(sb, n) + x_flat = x.view(sb, C) + + # --- Kernel A: g_x, g_orig (2D grid, no loop) --- + cache_key_a = ('hpb_bwd_g_x_orig', sb, n, C) + cached_a = _hpb_bwd_g_x_orig_best_cfg.get(cache_key_a) + + if cached_a is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached_a is not None: + ts, tc = cached_a + else: + ts, tc = 1, math.gcd(C, 1024) + ct.launch( + stream, + (math.ceil(sb / ts), math.ceil(C / tc)), + _ct_hpb_bwd_g_x_orig_kernel, + (go_flat, hr_flat, hp_flat, g_res, g_x, n, tc, ts), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _hpb_autotune_configs(sb, C)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE), math.ceil(C / cfg.TILE_C)), + kernel=_ct_hpb_bwd_g_x_orig_kernel, + args_fn=lambda cfg: ( + go_flat, + hr_flat, + hp_flat, + g_res, + g_x, + n, + cfg.TILE_C, + cfg.TILE_SIZE, + ), + search_space=configs, + ) + best = tuned.tuned_config + _hpb_bwd_g_x_orig_best_cfg[cache_key_a] = (best.TILE_SIZE, best.TILE_C) + ct.launch( + stream, + (math.ceil(sb / best.TILE_SIZE), math.ceil(C / best.TILE_C)), + _ct_hpb_bwd_g_x_orig_kernel, + (go_flat, hr_flat, hp_flat, g_res, g_x, n, best.TILE_C, best.TILE_SIZE), + ) + + # --- Kernel B: g_hp, g_hr (1D grid, loops C-tiles) --- + cache_key_b = ('hpb_bwd_g_hp_hr', sb, n, C, bias is not None) + cached_b = _hpb_bwd_g_hp_hr_best_cfg.get(cache_key_b) + hp_hr_kernel = ( + _ct_hpb_bwd_g_hp_hr_bias_kernel if bias is not None else _ct_hpb_bwd_g_hp_hr_kernel + ) + + if cached_b is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached_b is not None: + ts, tc = cached_b + else: + ts, tc = 1, math.gcd(C, 1024) + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (g_hr, g_hp, n, tc, ts) + ct.launch(stream, (math.ceil(sb / ts),), hp_hr_kernel, args) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _hpb_autotune_configs(sb, C)] + + def _hp_hr_args_fn(cfg): + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + return args + (g_hr, g_hp, n, cfg.TILE_C, cfg.TILE_SIZE) + + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(sb / cfg.TILE_SIZE),), + kernel=hp_hr_kernel, + args_fn=_hp_hr_args_fn, + search_space=configs, + ) + best = tuned.tuned_config + _hpb_bwd_g_hp_hr_best_cfg[cache_key_b] = (best.TILE_SIZE, best.TILE_C) + args = (go_flat, orig_flat, x_flat) + if bias is not None: + args = args + (bias,) + args = args + (g_hr, g_hp, n, best.TILE_C, best.TILE_SIZE) + ct.launch(stream, (math.ceil(sb / best.TILE_SIZE),), hp_hr_kernel, args) + + g_bias = g_x.sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.view(s, b, n, n), + g_res.view(s, b, n, C), + g_hp.view(s, b, n), + g_x.view(s, b, C), + g_bias, + ) + + # -- Proj RMS kernels ---------------------------------------------------- + + @ct.kernel + def _ct_proj_rms_fwd_kernel( + A, + B, + PROJ, + NORM, + R, + M: int, + N: int, + K: int, + eps: float, + TILE_M: ConstInt, + TILE_N: ConstInt, + TILE_K: ConstInt, + SPLIT_K: ConstInt, + ): + ''' + Grid: (num_tiles_m, num_tiles_k). + Fused matmul + norm + r: proj, norm, r in one pass over K. + R is a retained signature placeholder; r is computed after split-K + reduction from NORM. + ''' + tile_m_id = ct.bid(0) + split_k_id = ct.bid(1) + num_m_tiles = ct.cdiv(M, TILE_M) + num_k_tiles = ct.cdiv(K, TILE_K) + num_k_tiles_per_split = ct.cdiv(num_k_tiles, SPLIT_K) + tile_k_id_start = split_k_id * num_k_tiles_per_split + tile_k_id_end = ct.minimum(tile_k_id_start + num_k_tiles_per_split, num_k_tiles) + acc = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32) + sum_sq = ct.full((TILE_M, 1), 0.0, dtype=ct.float32) + for tile_k_id in range(tile_k_id_start, tile_k_id_end): + a_tile = ct.load( + A, index=(tile_m_id, tile_k_id), shape=(TILE_M, TILE_K), padding_mode=PAD_ZERO + ) + b_tile = ct.load(B, index=(0, tile_k_id), shape=(TILE_N, TILE_K), padding_mode=PAD_ZERO) + acc = ct.mma( + a_tile.astype(ct.tfloat32), b_tile.transpose().astype(ct.tfloat32), acc=acc + ) + # Square in fp32: a bf16 square/reduction loses ~2e-3 relative on the + # RMS scale, which native (fp32) does not. + a_tile_f32 = a_tile.astype(ct.float32) + sum_sq += ct.sum(a_tile_f32 * a_tile_f32, axis=1, keepdims=True) + + bid_m_k = tile_m_id + split_k_id * num_m_tiles + ct.store(PROJ, index=(bid_m_k, 0), tile=acc.astype(PROJ.dtype)) + ct.store(NORM, index=(bid_m_k, 0), tile=sum_sq.astype(NORM.dtype)) + + # -- Sigmoid helper for cuTile kernels ------------------------------------ + + @ct.function + def _ct_sigmoid(x): + """Sigmoid via exp2: σ(x) = 1 / (1 + 2^(-x * log2(e))).""" + return 1.0 / (1.0 + ct.exp2(-x * LOG2E)) + + # -- Reduce split-K + compute_h kernel ------------------------------------ + + @ct.kernel + def _ct_reduce_compute_h_kernel( + Y_acc, + R_acc, + Bias, + Alpha_pre, + Alpha_post, + Alpha_res, + H_PRE, + H_POST, + H_RES, + R, + PROJ_OUT, + M: int, + N: int, + K: int, + n: ConstInt, + eps: float, + compute_h_eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + SPLIT_K: ConstInt, + ): + """Reduce split-K partial proj/norm, compute r, and apply compute_h activations. + + Grid: (ceil(M / TILE_SIZE_M),). + TILE_SIZE_N = next_power_of_2(N) so one tile covers the full N dimension. + Alpha_{pre,post,res} are [1] tensors (scalar parameters). + """ + bid_m = ct.bid(0) + num_bid_m = ct.cdiv(M, TILE_SIZE_M) + + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + # 1. Reduce split-K partials for each logical output segment. + pre_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + post_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + r_accum = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + + for split_idx in ct.static_iter(range(SPLIT_K)): + bid_m_k = bid_m + split_idx * num_bid_m + pre_tile = ct.load( + Y_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + post_tile = ct.load( + Y_acc, index=(bid_m_k, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + pre_accum = pre_accum + ct.astype(pre_tile, ct.float32) + post_accum = post_accum + ct.astype(post_tile, ct.float32) + + r_tile = ct.load( + R_acc, index=(bid_m_k, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO + ) + r_accum = r_accum + ct.astype(r_tile, ct.float32) + + # Store reduced projection segments for backward. + ct.store(PROJ_OUT, index=(bid_m, 0), tile=pre_accum.astype(PROJ_OUT.dtype)) + ct.store(PROJ_OUT, index=(bid_m, 1), tile=post_accum.astype(PROJ_OUT.dtype)) + + # 2. Compute r = norm / sqrt(K) + denom = ct.full((TILE_SIZE_M, 1), K * 1.0, dtype=ct.float32) + r_val = ct.sqrt(ct.truediv(r_accum, denom)) + + ct.store(R, index=(bid_m, 0), tile=r_val.astype(R.dtype)) + + # 3. Apply compute_h directly into split outputs. + inv_r_eps = 1.0 / (r_val + eps) + bias_pre = ct.load(Bias, index=(0, 0), shape=(1, n), padding_mode=PAD_ZERO) + bias_post = ct.load(Bias, index=(0, 1), shape=(1, n), padding_mode=PAD_ZERO) + bias_pre = ct.astype(bias_pre, ct.float32) + bias_post = ct.astype(bias_post, ct.float32) + + h_pre_linear = pre_accum * alpha_pre * inv_r_eps + bias_pre + h_post_linear = post_accum * alpha_post * inv_r_eps + bias_post + h_pre = _ct_sigmoid(h_pre_linear) + compute_h_eps + h_post = _ct_sigmoid(h_post_linear) * 2.0 + + ct.store(H_PRE, index=(bid_m, 0), tile=h_pre.astype(H_PRE.dtype)) + ct.store(H_POST, index=(bid_m, 0), tile=h_post.astype(H_POST.dtype)) + + for res_chunk in ct.static_iter(range(n)): + res_accum = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + for split_idx in ct.static_iter(range(SPLIT_K)): + bid_m_k = bid_m + split_idx * num_bid_m + res_tile = ct.load( + Y_acc, + index=(bid_m_k, 2 + res_chunk), + shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + res_accum = res_accum + ct.astype(res_tile, ct.float32) + + bias_res = ct.load(Bias, index=(0, 2 + res_chunk), shape=(1, n), padding_mode=PAD_ZERO) + bias_res = ct.astype(bias_res, ct.float32) + h_res = res_accum * alpha_res * inv_r_eps + bias_res + ct.store(PROJ_OUT, index=(bid_m, 2 + res_chunk), tile=res_accum.astype(PROJ_OUT.dtype)) + ct.store(H_RES, index=(bid_m, res_chunk), tile=h_res.astype(H_RES.dtype)) + + def _next_power_of_2(n: int) -> int: + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n + + def _proj_rms_fwd_autotune_configs(N): + """Generate autotune search space for proj_rms forward kernel.""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128) + split_ks = (1, 2, 4, 8, 16) + for tile_m in tile_ms: + for tile_k in tile_ks: + for split_k in split_ks: + yield {"TILE_M": tile_m, "TILE_N": TILE_N, "TILE_K": tile_k, 'SPLIT_K': split_k} + + def _default_tile_m(M: int) -> int: + """Pick a tile size that avoids unmasked stores past the M dimension.""" + for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m <= M and M % tile_m == 0: + return tile_m + return 1 + + def _default_proj_rms_fwd_config(M: int, K: int, TILE_N: int): + """Static fallback for skinny MHC projection when autotune cache is absent.""" + split_k = 16 if K >= 16384 else 8 if K >= 8192 else 1 + return _default_tile_m(M), TILE_N, min(128, K), split_k + + # Cache the best config across calls (keyed by M, N, K). + _proj_rms_fwd_best_cfg: dict = {} + + # -- Reduce + compute_h launcher ------------------------------------------ + + def _reduce_compute_h_autotune_configs(M): + """Generate autotune search space for reduce_compute_h kernel.""" + min_tile_m = 16 if M >= 16 else 1 + for tile_m in (128, 64, 32, 16, 8, 4, 2, 1): + if tile_m < min_tile_m: + continue + if tile_m <= M and M % tile_m == 0: + yield tile_m + + def _default_reduce_compute_h_tile_m(M: int) -> int: + """Pick a reduce tile size with enough blocks to cover the GPU.""" + try: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + except Exception: + num_sms = 128 + + valid = [tm for tm in _reduce_compute_h_autotune_configs(M)] + for tm in valid: + if math.ceil(M / tm) >= num_sms: + return tm + return valid[-1] if valid else 1 + + _reduce_compute_h_best_cfg: dict = {} + + def _cutile_reduce_compute_h( + proj_acc: Tensor, + norm_acc: Tensor, + bias: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + n: int, + M: int, + N: int, + K: int, + eps: float, + compute_h_eps: float, + _proj_tile_m: int, + tile_n: int, + split_k: int, + out_dtype: torch.dtype, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Launch reduce split-K + compute_h kernel. + + Returns: + h_pre: [M, n] sigmoid-activated pre weights + h_post: [M, n] 2*sigmoid-activated post weights + h_res: [M, n*n] residual logits + r: [M, 1] r = norm / sqrt(K) + proj_reduced: [M, N] reduced projection (for backward) + """ + dev = proj_acc.device + stream = torch.cuda.current_stream() + + bias_2d = bias.unsqueeze(0).contiguous() # [1, N] + + # Mapping outputs follow the promoted input/parameter dtype (fp32 for + # mHC's keep_in_fp32 parameters); the reduced projection is kept in the + # fp32 accumulator dtype because the backward consumes it. + h_pre_out = torch.empty(M, n, dtype=out_dtype, device=dev) + h_post_out = torch.empty(M, n, dtype=out_dtype, device=dev) + h_res_out = torch.empty(M, N - 2 * n, dtype=out_dtype, device=dev) + r_out = torch.empty(M, 1, dtype=out_dtype, device=dev) + proj_out = torch.empty(M, N, dtype=proj_acc.dtype, device=dev) + + default_tm = _default_reduce_compute_h_tile_m(M) + cache_key = (M, N, K, n, split_k) + cached = _reduce_compute_h_best_cfg.get(cache_key) + + def _make_args(tm): + return ( + proj_acc, + norm_acc, + bias_2d, + alpha_pre, + alpha_post, + alpha_res, + h_pre_out, + h_post_out, + h_res_out, + r_out, + proj_out, + M, + N, + K, + n, + eps, + compute_h_eps, + tm, + tile_n, + split_k, + ) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + tm = cached if cached is not None else default_tm + ct.launch(stream, (math.ceil(M / tm),), _ct_reduce_compute_h_kernel, _make_args(tm)) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(TILE_M=tm) for tm in _reduce_compute_h_autotune_configs(M)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M),), + kernel=_ct_reduce_compute_h_kernel, + args_fn=lambda cfg: _make_args(cfg.TILE_M), + search_space=configs, + ) + best_tm = tuned.tuned_config.TILE_M + _reduce_compute_h_best_cfg[cache_key] = best_tm + ct.launch( + stream, (math.ceil(M / best_tm),), _ct_reduce_compute_h_kernel, _make_args(best_tm) + ) + + return h_pre_out, h_post_out, h_res_out, r_out, proj_out + + # -- Combined proj_rms + compute_h forward -------------------------------- + + def _cutile_proj_rms_compute_h_fwd( + x: Tensor, + weight: Tensor, + bias: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + n: int, + eps: float, + compute_h_eps: float, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Fused proj_rms + compute_h forward. + + Launches the existing _ct_proj_rms_fwd_kernel (split-K matmul + partial norm), + then _ct_reduce_compute_h_kernel (reduce + r + activations). + + Returns: + h_pre: [M, n] activated pre weights + h_post: [M, n] activated post weights + h_res: [M, n*n] residual logits + r: [M, 1] r = norm / sqrt(K) + proj_reduced: [M, N] reduced projection (for backward) + """ + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + dev = x.device + # The mHC mapping is a keep_in_fp32 computation (see + # HyperConnectionModule._projection_and_get_norm): keep the split-K + # partials and the mapping outputs in fp32 even when the activations + # arrive in bf16, otherwise this path is ~170x less accurate than native. + acc_dtype = torch.float32 + stream = torch.cuda.current_stream() + + cache_key = (M, N, K) + cached = _proj_rms_fwd_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk, split_k = cached + else: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + + proj_acc = torch.empty(split_k * M, N, dtype=acc_dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=acc_dtype, device=dev) + # _ct_proj_rms_fwd_kernel keeps R in its signature; reduce_compute_h + # computes r from norm_acc. + r_placeholder = torch.empty(split_k * M, 1, dtype=acc_dtype, device=dev) + + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + (x, weight, proj_acc, norm_acc, r_placeholder, M, N, K, eps, tm, tn, tk, split_k), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _proj_rms_fwd_autotune_configs(N)] + configs = [cfg for cfg in configs if cfg.TILE_K <= K and M % cfg.TILE_M == 0] + if len(configs) == 0: + tm, tn, tk, split_k = _default_proj_rms_fwd_config(M, K, TILE_N) + proj_acc = torch.empty(split_k * M, N, dtype=acc_dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=acc_dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. + r_placeholder = torch.empty(split_k * M, 1, dtype=acc_dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + ( + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + tm, + tn, + tk, + split_k, + ), + ) + else: + mx_split_k = max(cfg.SPLIT_K for cfg in configs) + proj_acc = torch.empty(mx_split_k * M, N, dtype=acc_dtype, device=dev) + norm_acc = torch.empty(mx_split_k * M, 1, dtype=acc_dtype, device=dev) + # Signature placeholder for autotune launches; not read. + r_placeholder = torch.empty(mx_split_k * M, 1, dtype=acc_dtype, device=dev) + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(M / cfg.TILE_M), cfg.SPLIT_K), + kernel=_ct_proj_rms_fwd_kernel, + args_fn=lambda cfg: ( + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + cfg.TILE_M, + cfg.TILE_N, + cfg.TILE_K, + cfg.SPLIT_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _proj_rms_fwd_best_cfg[cache_key] = ( + best.TILE_M, + best.TILE_N, + best.TILE_K, + best.SPLIT_K, + ) + tm, tn, tk, split_k = best.TILE_M, best.TILE_N, best.TILE_K, best.SPLIT_K + + proj_acc = torch.empty(split_k * M, N, dtype=acc_dtype, device=dev) + norm_acc = torch.empty(split_k * M, 1, dtype=acc_dtype, device=dev) + # Signature placeholder for the cuTile kernel; not read. + r_placeholder = torch.empty(split_k * M, 1, dtype=acc_dtype, device=dev) + ct.launch( + stream, + (math.ceil(M / tm), split_k), + _ct_proj_rms_fwd_kernel, + ( + x, + weight, + proj_acc, + norm_acc, + r_placeholder, + M, + N, + K, + eps, + tm, + tn, + tk, + split_k, + ), + ) + + # Launch reduce + compute_h kernel + h_pre, h_post, h_res, r, proj_reduced = _cutile_reduce_compute_h( + proj_acc, + norm_acc, + bias, + alpha_pre, + alpha_post, + alpha_res, + n, + M, + N, + K, + eps, + compute_h_eps, + tm, + TILE_N, + split_k, + torch.promote_types(x.dtype, weight.dtype), + ) + return h_pre, h_post, h_res, r, proj_reduced + + # -- Fused compute_h + proj_rms backward kernels ---------------------------- + + @ct.kernel + def _ct_fused_grad_h_proj_kernel( + GRAD_H_PRE, # [M, n] + GRAD_H_POST, # [M, n] + GRAD_H_RES, # [M, n*n] + H_PRE, # [M, n] + H_POST, # [M, n] + PROJ, # [M, N] + R, # [M, 1] + GRAD_R_EXT, # [M, 1] + Alpha_pre, # [1] + Alpha_post, # [1] + Alpha_res, # [1] + GRAD_H, # [M, TILE_SIZE_N] output + GRAD_PROJ, # [M, TILE_SIZE_N] output + GRAD_R_TOTAL, # [M, 1] output + M: int, + N: int, + n: ConstInt, + eps: float, + compute_h_eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + HAS_GRAD_H_PRE: ConstInt, + HAS_GRAD_H_POST: ConstInt, + HAS_GRAD_H_RES: ConstInt, + HAS_GRAD_R_EXT: ConstInt, + ): + """Precompute grad_h, grad_proj, and grad_r_total for downstream backward kernels. + + Grid: (ceil(M / TILE_SIZE_M),). + """ + tile_m_id = ct.bid(0) + + alpha_pre = ct.load(Alpha_pre, index=(0,), shape=(1,)).item() + alpha_post = ct.load(Alpha_post, index=(0,), shape=(1,)).item() + alpha_res = ct.load(Alpha_res, index=(0,), shape=(1,)).item() + + r_tile = ct.load(R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) + r_tile = ct.astype(r_tile, ct.float32) + + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + grad_r_from_h = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + + # Clear the padded columns once inside this kernel. Valid columns are + # overwritten by the segment stores below. + zero_full = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0.0, dtype=ct.float32) + ct.store(GRAD_H, index=(tile_m_id, 0), tile=zero_full.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=zero_full.astype(GRAD_PROJ.dtype)) + + if HAS_GRAD_H_PRE: + gy_pre = ct.load( + GRAD_H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + gy_pre = ct.astype(gy_pre, ct.float32) + else: + gy_pre = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + h_pre = ct.load(H_PRE, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO) + h_pre = ct.astype(h_pre, ct.float32) + proj_pre = ct.load( + PROJ, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + proj_pre = ct.astype(proj_pre, ct.float32) + sigmoid_pre = h_pre - compute_h_eps + grad_h_pre = gy_pre * sigmoid_pre * (1.0 - sigmoid_pre) + grad_proj_pre = grad_h_pre * alpha_pre * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_pre * proj_pre * alpha_pre * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True + ) + ct.store(GRAD_H, index=(tile_m_id, 0), tile=grad_h_pre.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 0), tile=grad_proj_pre.astype(GRAD_PROJ.dtype)) + + if HAS_GRAD_H_POST: + gy_post = ct.load( + GRAD_H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + gy_post = ct.astype(gy_post, ct.float32) + else: + gy_post = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + h_post = ct.load( + H_POST, index=(tile_m_id, 0), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + h_post = ct.astype(h_post, ct.float32) + proj_post = ct.load( + PROJ, index=(tile_m_id, 1), shape=(TILE_SIZE_M, n), padding_mode=PAD_ZERO + ) + proj_post = ct.astype(proj_post, ct.float32) + sigmoid_post = h_post * 0.5 + grad_h_post = gy_post * sigmoid_post * (1.0 - sigmoid_post) * 2.0 + grad_proj_post = grad_h_post * alpha_post * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_post * proj_post * alpha_post * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True + ) + ct.store(GRAD_H, index=(tile_m_id, 1), tile=grad_h_post.astype(GRAD_H.dtype)) + ct.store(GRAD_PROJ, index=(tile_m_id, 1), tile=grad_proj_post.astype(GRAD_PROJ.dtype)) + + for res_chunk in ct.static_iter(range(n)): + if HAS_GRAD_H_RES: + grad_h_res = ct.load( + GRAD_H_RES, + index=(tile_m_id, res_chunk), + shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + grad_h_res = ct.astype(grad_h_res, ct.float32) + else: + grad_h_res = ct.full((TILE_SIZE_M, n), 0.0, dtype=ct.float32) + proj_res = ct.load( + PROJ, + index=(tile_m_id, 2 + res_chunk), + shape=(TILE_SIZE_M, n), + padding_mode=PAD_ZERO, + ) + proj_res = ct.astype(proj_res, ct.float32) + grad_proj_res = grad_h_res * alpha_res * inv_r_eps + grad_r_from_h += ct.sum( + grad_h_res * proj_res * alpha_res * (-inv_r_eps * inv_r_eps), axis=1, keepdims=True + ) + ct.store(GRAD_H, index=(tile_m_id, 2 + res_chunk), tile=grad_h_res.astype(GRAD_H.dtype)) + ct.store( + GRAD_PROJ, + index=(tile_m_id, 2 + res_chunk), + tile=grad_proj_res.astype(GRAD_PROJ.dtype), + ) + + if HAS_GRAD_R_EXT: + grad_r_ext_tile = ct.load( + GRAD_R_EXT, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO + ) + grad_r_ext_tile = ct.astype(grad_r_ext_tile, ct.float32) + else: + grad_r_ext_tile = ct.full((TILE_SIZE_M, 1), 0.0, dtype=ct.float32) + grad_r_total = grad_r_from_h + grad_r_ext_tile + + ct.store(GRAD_R_TOTAL, index=(tile_m_id, 0), tile=grad_r_total.astype(GRAD_R_TOTAL.dtype)) + + @ct.kernel + def _ct_fused_grad_x_weight_kernel( + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_SIZE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output + M: int, + N: int, + K: int, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + TILE_SIZE_K: ConstInt, + ): + """Compute grad_x and grad_weight simultaneously. + + Grid: (ceil(K / TILE_SIZE_K),). + Each block handles one K-tile and loops over all M-tiles. + Per M-tile: computes and stores grad_x, accumulates grad_weight. + """ + tile_k_id = ct.bid(0) + NUM_M_TILES = ct.cdiv(M, TILE_SIZE_M) + + # Load weight tile once — only depends on K-tile + weight_tile = ct.load( + WEIGHT, index=(0, tile_k_id), shape=(TILE_SIZE_N, TILE_SIZE_K), padding_mode=PAD_ZERO + ) + + acc_grad_weight = ct.full((TILE_SIZE_K, TILE_SIZE_N), 0.0, dtype=ct.float32) + + for tile_m_id in range(NUM_M_TILES): + grad_proj_tile = ct.load( + GRAD_PROJ, + index=(tile_m_id, 0), + shape=(TILE_SIZE_M, TILE_SIZE_N), + padding_mode=PAD_ZERO, + ) + x_tile = ct.load( + X, + index=(tile_m_id, tile_k_id), + shape=(TILE_SIZE_M, TILE_SIZE_K), + padding_mode=PAD_ZERO, + ) + grad_r_total = ct.load( + GRAD_R_TOTAL, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO + ) + r_tile = ct.load(R, index=(tile_m_id, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) + r_tile = ct.astype(r_tile, ct.float32) + + # grad_x = grad_proj @ weight + grad_r_total * x / (r * K) + inv_rK = 1.0 / (r_tile * K) + acc_grad_x = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) + acc_grad_x = ct.mma( + grad_proj_tile.astype(ct.tfloat32), weight_tile.astype(ct.tfloat32), acc=acc_grad_x + ) + ct.store(GRAD_X, index=(tile_m_id, tile_k_id), tile=acc_grad_x.astype(GRAD_X.dtype)) + + # Accumulate grad_weight += x.T @ grad_proj + acc_grad_weight = ct.mma( + x_tile.transpose().astype(ct.tfloat32), + grad_proj_tile.astype(ct.tfloat32), + acc=acc_grad_weight, + ) + + ct.store( + GRAD_WEIGHT, + index=(0, tile_k_id), + tile=acc_grad_weight.transpose().astype(GRAD_WEIGHT.dtype), + ) + + @ct.kernel + def _ct_scalar_grads_partials_kernel( + GRAD_H, # [M, TILE_SIZE_N] precomputed + PROJ, # [M, N] + R, # [M, 1] + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] output + GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] output + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] output + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] output + M: int, + N: int, + n: int, + eps: float, + TILE_SIZE_M: ConstInt, + TILE_SIZE_N: ConstInt, + ): + """Compute per-M-tile scalar-gradient partials. + + Grid: (ceil(M / TILE_SIZE_M),). Each block processes one M-tile. + """ + bid_m = ct.bid(0) + + offsets = ct.arange(TILE_SIZE_N, dtype=ct.int32) + one = ct.full((TILE_SIZE_N,), 1.0, dtype=ct.float32) + zero = ct.full((TILE_SIZE_N,), 0.0, dtype=ct.float32) + mask_pre = ct.where(ct.less(offsets, n), one, zero) + mask_post = ct.where(ct.less(offsets, 2 * n), one, zero) - mask_pre + mask_res = one - mask_pre - mask_post + + mask_pre_2d = ct.reshape(mask_pre, (1, TILE_SIZE_N)) + mask_post_2d = ct.reshape(mask_post, (1, TILE_SIZE_N)) + mask_res_2d = ct.reshape(mask_res, (1, TILE_SIZE_N)) + + grad_h = ct.load( + GRAD_H, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO + ) + proj_tile = ct.load( + PROJ, index=(bid_m, 0), shape=(TILE_SIZE_M, TILE_SIZE_N), padding_mode=PAD_ZERO + ) + proj_tile = ct.astype(proj_tile, ct.float32) + r_tile = ct.load(R, index=(bid_m, 0), shape=(TILE_SIZE_M, 1), padding_mode=PAD_ZERO) + r_tile = ct.astype(r_tile, ct.float32) + + r_eps = r_tile + eps + inv_r_eps = 1.0 / r_eps + + ga_all = grad_h * proj_tile * inv_r_eps + ga_pre = ct.reshape(ct.sum(ga_all * mask_pre_2d), (1, 1)) + ga_post = ct.reshape(ct.sum(ga_all * mask_post_2d), (1, 1)) + ga_res = ct.reshape(ct.sum(ga_all * mask_res_2d), (1, 1)) + partial_gb = ct.sum(grad_h, axis=0, keepdims=False) + ct.store( + GRAD_ALPHA_PRE_PARTIALS, + index=(bid_m, 0), + tile=ga_pre.astype(GRAD_ALPHA_PRE_PARTIALS.dtype), + ) + ct.store( + GRAD_ALPHA_POST_PARTIALS, + index=(bid_m, 0), + tile=ga_post.astype(GRAD_ALPHA_POST_PARTIALS.dtype), + ) + ct.store( + GRAD_ALPHA_RES_PARTIALS, + index=(bid_m, 0), + tile=ga_res.astype(GRAD_ALPHA_RES_PARTIALS.dtype), + ) + ct.store( + GRAD_BIAS_PARTIALS, + index=(bid_m, 0), + tile=ct.reshape(partial_gb, (1, TILE_SIZE_N)).astype(GRAD_BIAS_PARTIALS.dtype), + ) + + @ct.kernel + def _ct_scalar_grads_reduce_kernel( + GRAD_ALPHA_PRE_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_POST_PARTIALS, # [num_m_blocks, 1] + GRAD_ALPHA_RES_PARTIALS, # [num_m_blocks, 1] + GRAD_BIAS_PARTIALS, # [num_m_blocks, TILE_SIZE_N] + GRAD_ALPHA_PRE, # [1, 1] output + GRAD_ALPHA_POST, # [1, 1] output + GRAD_ALPHA_RES, # [1, 1] output + GRAD_BIAS, # [1, TILE_SIZE_N] output + NUM_M_BLOCKS: int, + TILE_SIZE_N: ConstInt, + ): + """Reduce scalar-gradient partials and write final dtype outputs.""" + acc_pre = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_post = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_res = ct.full((1, 1), 0.0, dtype=ct.float32) + acc_bias = ct.full((1, TILE_SIZE_N), 0.0, dtype=ct.float32) + + for bid_m in range(NUM_M_BLOCKS): + acc_pre += ct.load( + GRAD_ALPHA_PRE_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO + ).astype(ct.float32) + acc_post += ct.load( + GRAD_ALPHA_POST_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO + ).astype(ct.float32) + acc_res += ct.load( + GRAD_ALPHA_RES_PARTIALS, index=(bid_m, 0), shape=(1, 1), padding_mode=PAD_ZERO + ).astype(ct.float32) + acc_bias += ct.load( + GRAD_BIAS_PARTIALS, index=(bid_m, 0), shape=(1, TILE_SIZE_N), padding_mode=PAD_ZERO + ).astype(ct.float32) + + ct.store(GRAD_ALPHA_PRE, index=(0, 0), tile=acc_pre.astype(GRAD_ALPHA_PRE.dtype)) + ct.store(GRAD_ALPHA_POST, index=(0, 0), tile=acc_post.astype(GRAD_ALPHA_POST.dtype)) + ct.store(GRAD_ALPHA_RES, index=(0, 0), tile=acc_res.astype(GRAD_ALPHA_RES.dtype)) + ct.store(GRAD_BIAS, index=(0, 0), tile=acc_bias.astype(GRAD_BIAS.dtype)) + + @ct.kernel + def _ct_fused_compute_h_proj_rms_bwd_small_k_kernel( + X, # [M, K] + WEIGHT, # [N, K] + GRAD_PROJ, # [M, TILE_N] precomputed + GRAD_R_TOTAL, # [M, 1] precomputed + R, # [M, 1] + GRAD_X, # [M, K] output + GRAD_WEIGHT, # [N, K] output + M: int, + N: int, + K: int, + TILE_N_SIZE: ConstInt, + ): + """Fused backward (small K path) with work-stealing. + + Grid: (num_sms, 2). + bid(1)==0: grad_weight via work-stealing over K-tiles, loops M. + bid(1)==1: grad_x via work-stealing over (M×K) tiles. + Scalar gradients are computed by the separate partial/reduce kernels. + """ + zero_pad = ct.PaddingMode.ZERO + + TILE_DB_SIZE_M = 128 + TILE_DB_SIZE_K = 64 + NUM_M_TILES = ct.cdiv(M, TILE_DB_SIZE_M) + NUM_K_TILES = ct.cdiv(K, TILE_DB_SIZE_K) + + if ct.bid(1) == 0: + # --- grad_weight path --- + for tile_id in range(ct.bid(0), NUM_K_TILES, ct.num_blocks(0)): + accumulator_db = ct.full((TILE_DB_SIZE_K, TILE_N_SIZE), 0.0, dtype=ct.float32) + for m_tile in range(NUM_M_TILES): + x_tile = ct.load( + X, + index=(m_tile, tile_id), + shape=(TILE_DB_SIZE_M, TILE_DB_SIZE_K), + padding_mode=zero_pad, + ) + grad_proj_tile = ct.load( + GRAD_PROJ, + index=(m_tile, 0), + shape=(TILE_DB_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + + accumulator_db = ct.mma( + x_tile.transpose().astype(ct.tfloat32), + grad_proj_tile.astype(ct.tfloat32), + acc=accumulator_db, + ) + + ct.store( + GRAD_WEIGHT, + index=(0, tile_id), + tile=accumulator_db.transpose().astype(GRAD_WEIGHT.dtype), + allow_tma=False, + ) + + TILE_DA_SIZE_M = 128 + TILE_DA_SIZE_K = 256 + NUM_DA_TILES = ct.cdiv(M, TILE_DA_SIZE_M) * ct.cdiv(K, TILE_DA_SIZE_K) + NUM_DA_K_TILES = ct.cdiv(K, TILE_DA_SIZE_K) + + if ct.bid(1) == 1: + # --- grad_x path --- + for tile_id in range(ct.bid(0), NUM_DA_TILES, ct.num_blocks(0)): + b_tile_idx = tile_id % NUM_DA_K_TILES + dd_tile_idx = tile_id // NUM_DA_K_TILES + + grad_proj_tile = ct.load( + GRAD_PROJ, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, TILE_N_SIZE), + padding_mode=zero_pad, + ) + grad_r_total = ct.load( + GRAD_R_TOTAL, + index=(dd_tile_idx, 0), + shape=(TILE_DA_SIZE_M, 1), + padding_mode=zero_pad, + ) + r_tile = ct.load( + R, index=(dd_tile_idx, 0), shape=(TILE_DA_SIZE_M, 1), padding_mode=zero_pad + ) + r_tile = ct.astype(r_tile, ct.float32) + + x_tile = ct.load( + X, + index=(dd_tile_idx, b_tile_idx), + shape=(TILE_DA_SIZE_M, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + inv_rK = 1.0 / (r_tile * K) + accumulator_da = (grad_r_total * inv_rK) * ct.astype(x_tile, ct.float32) + + weight_tile = ct.load( + WEIGHT, + index=(0, b_tile_idx), + shape=(TILE_N_SIZE, TILE_DA_SIZE_K), + padding_mode=zero_pad, + ) + accumulator_da = ct.mma( + grad_proj_tile.astype(ct.tfloat32), + weight_tile.astype(ct.tfloat32), + acc=accumulator_da, + ) + ct.store( + GRAD_X, + index=(dd_tile_idx, b_tile_idx), + tile=accumulator_da.astype(GRAD_X.dtype), + ) + + def _fused_grad_x_weight_autotune_configs(N): + """Autotune search space for fused grad_x + grad_weight kernel.""" + TILE_N = _next_power_of_2(N) + tile_ms = (32, 64, 128) + tile_ks = (32, 64, 128, 256) + for tile_m in tile_ms: + for tile_k in tile_ks: + yield {"TILE_SIZE_M": tile_m, "TILE_SIZE_N": TILE_N, "TILE_SIZE_K": tile_k} + + _fused_grad_x_weight_best_cfg: dict = {} + + def _cutile_fused_compute_h_proj_rms_bwd( + x: Tensor, + weight: Tensor, + grad_h_pre: Tensor, + grad_h_post: Tensor, + grad_h_res: Tensor, + h_pre: Tensor, + h_post: Tensor, + h_res: Tensor, + proj: Tensor, + r: Tensor, + grad_r_ext: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float, + compute_h_eps: float, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Fused compute_h + proj_rms backward. + + Returns: + grad_x: [M, K] + grad_weight: [N, K] + grad_alpha_pre: [1] + grad_alpha_post: [1] + grad_alpha_res: [1] + grad_bias: [N] + """ + M, K = x.shape + N = weight.shape[0] + TILE_N = _next_power_of_2(N) + assert TILE_N <= 256, f"TILE_SIZE_N too large: {TILE_N}" + dev = x.device + stream = torch.cuda.current_stream() + + grad_x = torch.empty_like(x) + grad_weight = torch.empty_like(weight) + has_grad_r_ext = grad_r_ext is not None + has_grad_r_ext_flag = int(has_grad_r_ext) + grad_r_ext_arg = grad_r_ext if has_grad_r_ext else r + has_grad_h_pre = grad_h_pre is not None + has_grad_h_post = grad_h_post is not None + has_grad_h_res = grad_h_res is not None + grad_h_pre_arg = grad_h_pre if has_grad_h_pre else h_pre + grad_h_post_arg = grad_h_post if has_grad_h_post else h_post + grad_h_res_arg = grad_h_res if has_grad_h_res else h_res + + # 0. Precompute grad_h, grad_proj, grad_r_total + grad_h_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) + grad_proj_buf = torch.empty(M, TILE_N, dtype=torch.float32, device=dev) + grad_r_total_buf = torch.empty(M, 1, dtype=torch.float32, device=dev) + + tile_m_precomp = _default_tile_m(M) + ct.launch( + stream, + (math.ceil(M / tile_m_precomp),), + _ct_fused_grad_h_proj_kernel, + ( + grad_h_pre_arg, + grad_h_post_arg, + grad_h_res_arg, + h_pre, + h_post, + proj, + r, + grad_r_ext_arg, + alpha_pre, + alpha_post, + alpha_res, + grad_h_buf, + grad_proj_buf, + grad_r_total_buf, + M, + N, + n, + eps, + compute_h_eps, + tile_m_precomp, + TILE_N, + int(has_grad_h_pre), + int(has_grad_h_post), + int(has_grad_h_res), + has_grad_r_ext_flag, + ), + ) + + if K >= 8192: + # 1. Fused grad_x + grad_weight kernel — 1D grid (K-tiles), loops M + cache_key = ('grad_x_weight', M, N, K) + cached = _fused_grad_x_weight_best_cfg.get(cache_key) + + if cached is not None or not _CUTILE_EXPERIMENTAL_AVAILABLE: + if cached is not None: + tm, tn, tk = cached + else: + tm, tn, tk = 128, TILE_N, 128 + ct.launch( + stream, + (math.ceil(K / tk),), + _ct_fused_grad_x_weight_kernel, + ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + tm, + tn, + tk, + ), + ) + else: + from types import SimpleNamespace + + configs = [SimpleNamespace(**c) for c in _fused_grad_x_weight_autotune_configs(N)] + tuned = ct_experimental.autotune_launch( + stream, + grid_fn=lambda cfg: (math.ceil(K / cfg.TILE_SIZE_K),), + kernel=_ct_fused_grad_x_weight_kernel, + args_fn=lambda cfg: ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + cfg.TILE_SIZE_M, + cfg.TILE_SIZE_N, + cfg.TILE_SIZE_K, + ), + search_space=configs, + ) + best = tuned.tuned_config + _fused_grad_x_weight_best_cfg[cache_key] = ( + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ) + ct.launch( + stream, + (math.ceil(K / best.TILE_SIZE_K),), + _ct_fused_grad_x_weight_kernel, + ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + best.TILE_SIZE_M, + best.TILE_SIZE_N, + best.TILE_SIZE_K, + ), + ) + else: + num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count + ct.launch( + stream, + (num_sms, 2, 1), + _ct_fused_compute_h_proj_rms_bwd_small_k_kernel, + ( + x, + weight, + grad_proj_buf, + grad_r_total_buf, + r, + grad_x, + grad_weight, + M, + N, + K, + TILE_N, + ), + ) + + # 2. Separate lightweight kernel for scalar gradients (grad_alpha, grad_bias) + tile_m_scalar = min(128, M) + num_m_blocks = math.ceil(M / tile_m_scalar) + grad_alpha_pre_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_alpha_post_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_alpha_res_partials = torch.empty(num_m_blocks, 1, dtype=torch.float32, device=dev) + grad_bias_partials = torch.empty(num_m_blocks, TILE_N, dtype=torch.float32, device=dev) + grad_alpha_pre = torch.empty(1, 1, dtype=alpha_pre.dtype, device=dev) + grad_alpha_post = torch.empty(1, 1, dtype=alpha_post.dtype, device=dev) + grad_alpha_res = torch.empty(1, 1, dtype=alpha_res.dtype, device=dev) + grad_bias = torch.empty(1, TILE_N, dtype=bias.dtype, device=dev) + + ct.launch( + stream, + (num_m_blocks,), + _ct_scalar_grads_partials_kernel, + ( + grad_h_buf, + proj, + r, + grad_alpha_pre_partials, + grad_alpha_post_partials, + grad_alpha_res_partials, + grad_bias_partials, + M, + N, + n, + eps, + tile_m_scalar, + TILE_N, + ), + ) + ct.launch( + stream, + (1,), + _ct_scalar_grads_reduce_kernel, + ( + grad_alpha_pre_partials, + grad_alpha_post_partials, + grad_alpha_res_partials, + grad_bias_partials, + grad_alpha_pre, + grad_alpha_post, + grad_alpha_res, + grad_bias, + num_m_blocks, + TILE_N, + ), + ) + + return ( + grad_x, + grad_weight, + grad_alpha_pre.view_as(alpha_pre), + grad_alpha_post.view_as(alpha_post), + grad_alpha_res.view_as(alpha_res), + grad_bias.view(-1)[:N], + ) + + +# ============================================================================ +# Unified public dispatch +# ============================================================================ +# The public fused API chooses the fastest validated backend per operation: +# +# sinkhorn fwd/bwd: Triton -> cuTile -> torch +# h_post_bda fwd/bwd: Triton -> cuTile -> torch +# h_aggregate fwd: Triton -> cuTile -> torch +# h_aggregate bwd: cuTile -> torch +# proj_rms/proj_rms_compute_h: cuTile -> torch +# +# Runtime CUDA launch failures are intentionally not swallowed; after such an +# error the CUDA context may not be safely reusable for fallback work. +# ============================================================================ + +from megatron.core.transformer.hyper_connection import ( + native_fused_add_3, + native_h_aggregate, + native_h_post_bda, + native_sinkhorn, +) + +_BACKEND_INFO_LOGGED = False + + +def _select_triton_cutile_native(triton_impl) -> str: + if triton_impl is not None: + return "triton" + if is_cutile_available(): + return "cutile" + return "native" + + +def _mhc_backend_status() -> Tuple[str, bool]: + """Return backend description and whether every backend is native.""" + sinkhorn = _select_triton_cutile_native(_get_triton_sinkhorn()) + h_aggregate_fwd = _select_triton_cutile_native(_get_triton_h_aggregate_fwd()) + h_aggregate_bwd = "cutile" if is_cutile_available() else "native" + h_post_bda_fwd = _select_triton_cutile_native(_get_triton_h_post_bda_fwd()) + h_post_bda_bwd = _select_triton_cutile_native(_get_triton_h_post_bda_bwd()) + proj_rms_compute_h = "cutile" if is_cutile_available() else "native" + selected = ( + sinkhorn, + h_aggregate_fwd, + h_aggregate_bwd, + h_post_bda_fwd, + h_post_bda_bwd, + proj_rms_compute_h, + ) + message = ( + f"MHC_FORCE_BACKEND={_MHC_FORCED_BACKEND}; " + f"sinkhorn={sinkhorn}; " + f"h_aggregate=fwd:{h_aggregate_fwd},bwd:{h_aggregate_bwd}; " + f"h_post_bda=fwd:{h_post_bda_fwd},bwd:{h_post_bda_bwd}; " + f"proj_rms_compute_h={proj_rms_compute_h}" + ) + return message, all(backend == "native" for backend in selected) + + +def _mhc_backend_selection() -> str: + """Return a concise description of the selected mHC fused backends.""" + message, _ = _mhc_backend_status() + return message + + +def log_fused_mhc_backend_once() -> None: + """Log the fused mHC backend selection once per process.""" + _raise_mhc_backend_validation_error() + global _BACKEND_INFO_LOGGED + if _BACKEND_INFO_LOGGED: + return + _BACKEND_INFO_LOGGED = True + backend_selection, all_native = _mhc_backend_status() + log_single_rank( + logger, + logging.WARNING if all_native else logging.INFO, + f"[mHC] fused backend selection: {backend_selection}", + ) + if all_native and safe_get_rank() == 0: + warnings.warn( + "[mHC] No accelerated mHC backend is available; falling back to native torch " + "implementations. The fallback is functionally equivalent, but may not provide " + "the performance benefits of fused mHC backends.", + UserWarning, + stacklevel=2, + ) + + +def fused_add_3(a: Tensor, b: Tensor, c: Tensor) -> Tensor: + """Add three tensors using the native torch.compile-backed implementation.""" + return native_fused_add_3(a, b, c) + + +def _get_triton_sinkhorn(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["sinkhorn"] + + +def _get_triton_h_aggregate_fwd(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["h_aggregate_fwd"] + + +def _get_triton_h_post_bda_fwd(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["h_post_bda_fwd"] + + +def _get_triton_h_post_bda_bwd(): + if not _TRITON_AVAILABLE: + return None + return _TRITON_IMPLS["h_post_bda_bwd"] + + +def _torch_h_aggregate_bwd(grad_output: Tensor, x: Tensor, h_pre: Tensor) -> Tuple[Tensor, Tensor]: + grad_output_expanded = grad_output.unsqueeze(2) + grad_x = grad_output_expanded * h_pre.unsqueeze(-1) + grad_h = torch.sum(grad_output_expanded * x, dim=-1) + return grad_x.to(dtype=x.dtype), grad_h.to(dtype=h_pre.dtype) + + +@torch.compile +def _torch_h_post_bda_bwd( + grad_output: Tensor, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: + s, b, n, C = original_residual.shape + sb = s * b + go = grad_output.reshape(sb, n, C) + hr = h_res.reshape(sb, n, n) + orig = original_residual.reshape(sb, n, C) + hp = h_post.reshape(sb, n) + x_flat = x.reshape(sb, C) + + g_hr = torch.bmm(orig, go.transpose(1, 2)).view(s, b, n, n) + g_res = torch.bmm(hr, go).view(s, b, n, C) + g_x = torch.sum(go * hp.unsqueeze(-1), dim=1).view(s, b, C) + xb = x_flat if bias is None else x_flat + bias.view(1, C) + g_hp = torch.sum(go * xb.unsqueeze(1), dim=2).view(s, b, n) + g_bias = g_x.reshape(sb, C).sum(dim=0).to(dtype=bias.dtype) if bias is not None else None + return ( + g_hr.to(dtype=h_res.dtype), + g_res.to(dtype=original_residual.dtype), + g_hp.to(dtype=h_post.dtype), + g_x.to(dtype=x.dtype), + g_bias, + ) + + +@torch.compile +def _torch_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float, + compute_h_eps: float = 1e-6, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + # compute_mappings() hands us activations in the activation dtype while the + # mapping parameters are keep_in_fp32, so matmul would reject the pair. + # Compute in the wider of the two, matching the unfused path's fp32 upcast + # without letting a lower-precision parameter downcast the activations. + x = x.to(torch.promote_types(x.dtype, weight.dtype)) + proj = torch.matmul(x, weight.t()) + r = x.norm(dim=-1, keepdim=True) / math.sqrt(x.shape[-1]) + alpha = torch.cat( + [alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(weight.shape[0] - 2 * n)], + dim=-1, + ) + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + compute_h_eps + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res = h[..., 2 * n :] + return h_pre, h_post, h_res, r + + +if _CUTILE_AVAILABLE: + + class CutileSinkhornKnopp(torch.autograd.Function): + """cuTile Sinkhorn-Knopp projection fallback.""" + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6): + """Run cuTile Sinkhorn forward and save initial matrix for backward.""" + output, M_init = _cutile_sinkhorn_fwd(input_logits, num_iterations, eps) + ctx.save_for_backward(M_init) + ctx.num_iterations = num_iterations + ctx.eps = eps + return output + + @staticmethod + def backward(ctx, grad_output): + """Run cuTile Sinkhorn backward.""" + (M_init,) = ctx.saved_tensors + grad_input = _cutile_sinkhorn_bwd(grad_output, M_init, ctx.num_iterations, ctx.eps) + return grad_input, None, None + + class CutileHAggregate(torch.autograd.Function): + """cuTile n-stream weighted aggregation.""" + + @staticmethod + def forward(ctx, x: Tensor, h_pre: Tensor): + """Run cuTile h_aggregate forward.""" + output = _cutile_h_aggregate_fwd(x, h_pre) + ctx.save_for_backward(x, h_pre) + return output + + @staticmethod + def backward(ctx, grad_output): + """Run cuTile h_aggregate backward.""" + x, h_pre = ctx.saved_tensors + return _cutile_h_aggregate_bwd(grad_output, x, h_pre) + + class CutileProjRmsComputeH(torch.autograd.Function): + """cuTile projection + RMS norm + compute_h activations.""" + + @staticmethod + def forward( + ctx, + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + compute_h_eps: float = 1e-6, + ): + """Run fused cuTile projection, RMS normalization, and compute_h forward.""" + h_pre, h_post, h_res, r, proj_reduced = _cutile_proj_rms_compute_h_fwd( + x, weight, bias, alpha_pre, alpha_post, alpha_res, n, eps, compute_h_eps + ) + ctx.save_for_backward( + x, + weight, + h_pre, + h_post, + h_res, + proj_reduced, + r, + alpha_pre, + alpha_post, + alpha_res, + bias, + ) + ctx.n = n + ctx.eps = eps + ctx.compute_h_eps = compute_h_eps + return h_pre, h_post, h_res, r + + @staticmethod + def backward(ctx, grad_h_pre, grad_h_post, grad_h_res, grad_r_ext): + """Run fused cuTile projection, RMS normalization, and compute_h backward.""" + ( + x, + weight, + h_pre, + h_post, + h_res, + proj, + r, + alpha_pre, + alpha_post, + alpha_res, + bias_param, + ) = ctx.saved_tensors + + grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias = ( + _cutile_fused_compute_h_proj_rms_bwd( + x, + weight, + grad_h_pre, + grad_h_post, + grad_h_res, + h_pre, + h_post, + h_res, + proj, + r, + grad_r_ext, + alpha_pre, + alpha_post, + alpha_res, + bias_param, + ctx.n, + ctx.eps, + ctx.compute_h_eps, + ) + ) + + return (grad_x, grad_weight, grad_ap, grad_apo, grad_ar, grad_bias, None, None, None) + + +class FusedHAggregate(torch.autograd.Function): + """H_aggregate with Triton/cuTile/torch forward and cuTile/torch backward.""" + + @staticmethod + def forward(ctx, x: Tensor, h_pre: Tensor): + """Run h_aggregate forward using the best available backend.""" + triton_fwd = _get_triton_h_aggregate_fwd() + if triton_fwd is not None: + output = triton_fwd(x, h_pre) + elif is_cutile_available(): + output = _cutile_h_aggregate_fwd(x, h_pre) + else: + output = native_h_aggregate(x, h_pre) + ctx.save_for_backward(x, h_pre) + return output + + @staticmethod + def backward(ctx, grad_output): + """Run h_aggregate backward using the best available backend.""" + x, h_pre = ctx.saved_tensors + if is_cutile_available(): + return _cutile_h_aggregate_bwd(grad_output, x, h_pre) + return _torch_h_aggregate_bwd(grad_output, x, h_pre) + + +class FusedHPostBDA(torch.autograd.Function): + """H_post_bda with Triton/cuTile/torch forward and backward.""" + + @staticmethod + def forward( + ctx, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + x: Tensor, + bias: Optional[Tensor], + ): + """Run h_post_bda forward using the best available backend.""" + triton_fwd = _get_triton_h_post_bda_fwd() + if triton_fwd is not None: + output = triton_fwd(h_res, original_residual, h_post, x, bias) + elif is_cutile_available(): + output = _cutile_h_post_bda_fwd(h_res, original_residual, h_post, x, bias) + else: + output = native_h_post_bda(h_res, original_residual, h_post, x, bias) + if bias is not None: + ctx.save_for_backward(h_res, original_residual, h_post, x, bias) + ctx.has_bias = True + else: + ctx.save_for_backward(h_res, original_residual, h_post, x) + ctx.has_bias = False + return output + + @staticmethod + def backward(ctx, grad_output): + """Run h_post_bda backward using the best available backend.""" + if ctx.has_bias: + h_res, orig_res, h_post, x, bias = ctx.saved_tensors + else: + h_res, orig_res, h_post, x = ctx.saved_tensors + bias = None + + triton_bwd = _get_triton_h_post_bda_bwd() + if triton_bwd is not None: + return triton_bwd(grad_output, h_res, orig_res, h_post, x, bias) + if is_cutile_available(): + return _cutile_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) + return _torch_h_post_bda_bwd(grad_output, h_res, orig_res, h_post, x, bias) + + +def fused_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Project logits to a doubly stochastic matrix using Triton, cuTile, then torch.""" + _raise_mhc_backend_validation_error() + triton_sinkhorn = _get_triton_sinkhorn() + if triton_sinkhorn is not None: + return triton_sinkhorn(input_logits, num_iterations, eps) + if is_cutile_available(): + return CutileSinkhornKnopp.apply(input_logits, num_iterations, eps) + return native_sinkhorn(input_logits, num_iterations, eps) + + +def fused_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + """Weighted n-stream to 1-stream aggregation using Triton/cuTile/torch.""" + _raise_mhc_backend_validation_error() + if _TRITON_AVAILABLE or is_cutile_available(): + return FusedHAggregate.apply(x, h_pre) + return native_h_aggregate(x, h_pre) + + +def fused_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + """Fused H_res.T @ residual + H_post * (x + bias).""" + _raise_mhc_backend_validation_error() + if _TRITON_AVAILABLE or is_cutile_available(): + return FusedHPostBDA.apply(h_res, original_residual, h_post, x, bias) + return native_h_post_bda(h_res, original_residual, h_post, x, bias) + + +def fused_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + compute_h_eps: float = 1e-6, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Projection + RMS norm + compute_h split outputs using cuTile, then torch.""" + _raise_mhc_backend_validation_error() + if is_cutile_available(): + return CutileProjRmsComputeH.apply( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps + ) + return _torch_proj_rms_compute_h( + x, weight, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps + ) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 1fd5dcfae37..cf1c4a31fb0 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -29,17 +29,28 @@ @triton.jit def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): - token_idx = -1 - this_seq_len = 0 + # Cast ``pid_m`` and ``cu_seqlens`` loads to a single shared dtype so + # the loop-body reassignments don't surface as + # "initial value is int32 but redefined as int64" in newer Triton + # versions (which promote ``// Python_int`` to int64). + pid_m = pid_m.to(tl.int64) + token_idx = tl.full((), -1, dtype=tl.int64) + this_seq_len = tl.full((), 0, dtype=tl.int64) seq_idx = 0 - last_cum_seqlen = tl.load(cu_seqlens) // cp_size + last_cum_seqlen = tl.load(cu_seqlens).to(tl.int64) // cp_size while seq_idx < seq_num: - cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1).to(tl.int64) // cp_size if token_idx == -1 and cur_cum_seqlen > pid_m: token_idx = pid_m - last_cum_seqlen this_seq_len = cur_cum_seqlen - last_cum_seqlen last_cum_seqlen = cur_cum_seqlen seq_idx += 1 + # Padding tokens beyond cu_seqlens[-1] (from THD CUDA-graph padding) + # never match any sequence, leaving token_idx == -1. Clamp to 0 so + # the cos/sin table loads stay in-bounds; the wrong RoPE result is + # harmless because padding positions are excluded by loss_mask. + if token_idx == -1: + token_idx = tl.full((), 0, dtype=tl.int64) if cp_size > 1: if token_idx < this_seq_len // 2: token_idx = token_idx + cp_rank * this_seq_len // 2 @@ -65,29 +76,34 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): restore_value=["Q"], ) @triton.jit -def rotary_fwd_q_kernel( +def _mla_rope_fwd_inplace_kernel( Q, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor Q. + Forward pass: apply RoPE inplace to the trailing emb_dim elements. + Reads from interleaved layout, writes back to interleaved layout. Input: - Q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + Q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size: batch size for sbhd format, not used for thd format @@ -97,15 +113,24 @@ def rotary_fwd_q_kernel( pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: + if position_ids is not None: + token_idx = tl.load(position_ids + pid_m) + elif cu_seqlens_q is None: token_idx = pid_m // batch_size else: token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -113,7 +138,7 @@ def rotary_fwd_q_kernel( Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads # x1 = t[..., 0::2], x2 = t[..., 1::2] x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 @@ -124,10 +149,14 @@ def rotary_fwd_q_kernel( x_left = x_1 * cos_left - x_2 * sin_left x_right = x_2 * cos_right + x_1 * sin_right - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - tl.store(Q + x_left_off, x_left, mask=mask) - tl.store(Q + x_right_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + tl.store(Q + x_1_off, x_left, mask=mask) + tl.store(Q + x_2_off, x_right, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) @triton.autotune( @@ -145,29 +174,34 @@ def rotary_fwd_q_kernel( restore_value=["DO"], ) @triton.jit -def rotary_bwd_q_kernel( +def _mla_rope_bwd_inplace_kernel( DO, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor DO. + Backward pass: inverse RoPE inplace on the trailing emb_dim elements. + Reads from interleaved layout, writes to interleaved layout. Input: - DO: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + DO: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size, seq_num, and cu_seqlens_q are the same as in the forward pass @@ -175,15 +209,24 @@ def rotary_bwd_q_kernel( pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: + if position_ids is not None: + token_idx = tl.load(position_ids + pid_m) + elif cu_seqlens_q is None: token_idx = pid_m // batch_size else: token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -191,25 +234,32 @@ def rotary_bwd_q_kernel( DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(DO + x_left_off, mask=mask) - x_right = tl.load(DO + x_right_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(DO + x_1_off, mask=mask) + x_right = tl.load(DO + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 x_1 = x_left * cos_left + x_right * sin_right x_2 = -x_left * sin_left + x_right * cos_right - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 tl.store(DO + x_1_off, x_1, mask=mask) tl.store(DO + x_2_off, x_2, mask=mask) -class ApplyMLARotaryEmbQ(torch.autograd.Function): +class _FusedMLARoPEInplace(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's query. + Autograd function for applying RoPE inplace to the trailing emb_dim + elements of a multi-head tensor (leaving the first nope_dim elements unchanged). """ @staticmethod @@ -218,22 +268,26 @@ def forward( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved=False, + inverse=False, + remove_interleaving=False, + position_ids=None, ): """ - Forward function for ApplyMLARotaryEmbQ. + Forward function for _FusedMLARoPEInplace. Args: - q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, negate sin inside the kernel to apply the inverse rotation """ assert not rotary_interleaved max_seqlen = None @@ -241,6 +295,7 @@ def forward( seq_num = None if cu_seqlens_q is None: # sbhd + assert position_ids is None max_seqlen, batch_size, nheads, headdim = q.shape q = q.view(-1, nheads, headdim) total_seqlen = q.shape[0] @@ -248,33 +303,43 @@ def forward( # thd total_seqlen, nheads, headdim = q.shape seq_num = len(cu_seqlens_q) - 1 + if position_ids is not None: + assert position_ids.shape == (total_seqlen,) assert q.stride(-1) == 1 - assert cos.is_contiguous() - assert sin.is_contiguous() - assert headdim == qk_head_dim + emb_dim + assert cos.stride(-1) == 1 + assert sin.stride(-1) == 1 + assert headdim == nope_dim + emb_dim assert emb_dim % 4 == 0 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid]( + _mla_rope_fwd_inplace_kernel[grid]( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, nheads, batch_size, seq_num, cu_seqlens_q, + position_ids, q.stride(0), q.stride(1), + cos.stride(0), + sin.stride(0), cp_rank, cp_size, + INVERSE=inverse, + REMOVE_INTERLEAVING=remove_interleaving, ) - ctx.save_for_backward(cos, sin) - ctx.qk_head_dim = qk_head_dim + ctx.save_for_backward(cos, sin, *(() if position_ids is None else (position_ids,))) + ctx.has_position_ids = position_ids is not None + ctx.nope_dim = nope_dim ctx.emb_dim = emb_dim ctx.cu_seqlens_q = cu_seqlens_q ctx.rotary_interleaved = rotary_interleaved + ctx.inverse = inverse + ctx.remove_interleaving = remove_interleaving ctx.cp_rank = cp_rank ctx.cp_size = cp_size if cu_seqlens_q is None: @@ -284,13 +349,17 @@ def forward( @staticmethod def backward(ctx, grad): """ - Backward function for ApplyMLARotaryEmbQ. + Backward function for _FusedMLARoPEInplace. Args: - grad: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + grad: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] """ - cos, sin = ctx.saved_tensors + if ctx.has_position_ids: + cos, sin, position_ids = ctx.saved_tensors + else: + cos, sin = ctx.saved_tensors + position_ids = None max_seqlen = None batch_size = None seq_num = None @@ -300,65 +369,126 @@ def backward(ctx, grad): total_seqlen = grad.shape[0] else: seq_num = len(ctx.cu_seqlens_q) - 1 + if ctx.has_position_ids: + grad = grad.contiguous() total_seqlen, nheads, headdim = grad.shape assert grad.stride(-1) == 1 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid]( + _mla_rope_bwd_inplace_kernel[grid]( grad, cos, sin, - ctx.qk_head_dim, + ctx.nope_dim, ctx.emb_dim, nheads, batch_size, seq_num, ctx.cu_seqlens_q, + position_ids, grad.stride(0), grad.stride(1), + cos.stride(0), + sin.stride(0), ctx.cp_rank, ctx.cp_size, + INVERSE=ctx.inverse, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_q is None: grad = grad.view(max_seqlen, batch_size, nheads, headdim) - return grad, None, None, None, None, None, None, None, None + return grad, None, None, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_q( +def fused_mla_rope_inplace( t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - qk_head_dim: int, + nope_dim: int, emb_dim: int, cu_seqlens_q: Optional[torch.Tensor] = None, cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, -): + inverse: bool = False, + remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: """ - Fused function for applying YARN RoPE to MLA's query. - This function inplace modifies the input tensor t. - Along the last dimension of t, the last emb_dim elements are applied with RoPE. - The first qk_head_dim elements are not modified. - It is an experimental feature and may change in future versions. + Fused RoPE applied inplace to the trailing emb_dim elements of a tensor, + leaving the first nope_dim elements unchanged. It supports both sbhd and thd input formats. + When ``inverse=True`` the rotation is reversed, which is useful for + undoing RoPE on the attention output. + For the notations below, seq_len is the length of the sequence per batch for sbhd format, total_seq_len is the total length of the sequences for thd format. max_seq_len is the maximum length of the sequences in the input tensor. Args: - t: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + t: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, apply the inverse rotation + remove_interleaving: if True, output RoPE dims in non-interleaved layout + position_ids: optional THD row positions. When supplied, these positions + replace the built-in CP row-to-position mapping. Returns: t: inplace modified input tensor """ - return ApplyMLARotaryEmbQ.apply( - t, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved + return _FusedMLARoPEInplace.apply( + t, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q, + cp_rank, + cp_size, + rotary_interleaved, + inverse, + remove_interleaving, + position_ids, + ) + + +def fused_mla_rope_out_of_place( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + nope_dim: int, + emb_dim: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, + inverse: bool = False, + remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Apply the fused RoPE kernel without modifying the input tensor. + + Use this wrapper when an upstream autograd function may have retained its + output for backward. The underlying kernel remains in-place, so a private + copy is required to keep the retained tensor unchanged. + """ + return fused_mla_rope_inplace( + t.clone(), + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens_q, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, + inverse=inverse, + remove_interleaving=remove_interleaving, + position_ids=position_ids, ) @@ -376,7 +506,7 @@ def fused_apply_mla_rope_for_q( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_fwd_kv_kernel( +def _mla_rope_fwd_kv_split_kernel( KV, K_POS_EMB, O_KEY, @@ -399,12 +529,12 @@ def rotary_fwd_kv_kernel( stride_v_nheads, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's key and value. - It splits the input tensor KV into key and value, - and concatenates the processed RoPE to the key. + Forward pass: split KV into key and value, apply RoPE to k_pos_emb, + and concatenate the result onto key. Input: KV: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -460,14 +590,24 @@ def rotary_fwd_kv_kernel( x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_left_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads - + k_dim - + tl.arange(0, emb_dim // 2)[None, :] - ) - x_right_off = x_left_off + emb_dim // 2 - tl.store(K_ptr + x_left_off, x_left, mask=mask) - tl.store(K_ptr + x_right_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] * 2 + ) + x_2_off = x_1_off + 1 + tl.store(K_ptr + x_1_off, x_left, mask=mask) + tl.store(K_ptr + x_2_off, x_right, mask=mask) + else: + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) @triton.autotune( @@ -484,7 +624,7 @@ def rotary_fwd_kv_kernel( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_bwd_kv_kernel( +def _mla_rope_bwd_kv_split_kernel( dK, dV, dKV, @@ -507,10 +647,11 @@ def rotary_bwd_kv_kernel( stride_demb_seq, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's key and value. + Backward pass for the KV-split RoPE. Input: dK: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -555,10 +696,16 @@ def rotary_bwd_kv_kernel( dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim mask = x_off < head_num * stride_dk_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(dK_ptr + x_left_off, mask=mask) - x_right = tl.load(dK_ptr + x_right_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(dK_ptr + x_1_off, mask=mask) + x_right = tl.load(dK_ptr + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(dK_ptr + x_left_off, mask=mask) + x_right = tl.load(dK_ptr + x_right_off, mask=mask) x_left_accum += x_left x_right_accum += x_right x_left_accum = tl.sum(x_left_accum, axis=0) @@ -578,9 +725,10 @@ def rotary_bwd_kv_kernel( tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) -class ApplyMLARotaryEmbKV(torch.autograd.Function): +class _FusedMLARoPEKVSplit(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's key and value. + Autograd function for applying RoPE to MLA's key and value. + Splits KV, applies RoPE to k_pos_emb, concatenates onto key. """ @staticmethod @@ -597,9 +745,10 @@ def forward( cp_rank, cp_size, rotary_interleaved=False, + remove_interleaving=False, ): """ - Forward function for ApplyMLARotaryEmbKV. + Forward function for _FusedMLARoPEKVSplit. Args: kv: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -634,7 +783,7 @@ def forward( o_value = kv.new_empty(total_seqlen, nheads, v_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid]( + _mla_rope_fwd_kv_split_kernel[grid]( kv, k_pos_emb, o_key, @@ -657,8 +806,10 @@ def forward( o_value.stride(1), cp_rank, cp_size, + REMOVE_INTERLEAVING=remove_interleaving, ) ctx.save_for_backward(cos, sin) + ctx.remove_interleaving = remove_interleaving ctx.rotary_interleaved = rotary_interleaved ctx.emb_dim = emb_dim ctx.k_dim = k_dim @@ -674,7 +825,7 @@ def forward( @staticmethod def backward(ctx, dk, dv): """ - Backward function for ApplyMLARotaryEmbKV. + Backward function for _FusedMLARoPEKVSplit. Args: dk: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -702,7 +853,7 @@ def backward(ctx, dk, dv): d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid]( + _mla_rope_bwd_kv_split_kernel[grid]( dk, dv, d_kv, @@ -725,14 +876,15 @@ def backward(ctx, dk, dv): d_emb.stride(0), ctx.cp_rank, ctx.cp_size, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_kv is None: d_kv = d_kv.view(max_seqlen, batch_size, nheads, ctx.k_dim + ctx.v_dim) d_emb = d_emb.view(max_seqlen, batch_size, 1, ctx.emb_dim) - return d_kv, d_emb, None, None, None, None, None, None, None, None, None + return d_kv, d_emb, None, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_kv( +def fused_mla_rope_kv_split( kv: torch.Tensor, k_pos_emb: torch.Tensor, cos: torch.Tensor, @@ -744,9 +896,10 @@ def fused_apply_mla_rope_for_kv( cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, -): + remove_interleaving: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: """ - Fused function for applying YARN RoPE to MLA's key and value. + Fused function for applying RoPE to MLA's key and value. It splits the input tensor kv into key and value, and concatenates the processed RoPE to the key. @@ -761,13 +914,14 @@ def fused_apply_mla_rope_for_kv( cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_kv: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + remove_interleaving: if True, output RoPE dims in non-interleaved layout Returns: key: [seq_len, batch_size, head_num, emb_dim + k_dim] or [total_seq_len, head_num, emb_dim + k_dim] value: [seq_len, batch_size, head_num, v_dim] or [total_seq_len, head_num, v_dim] """ - return ApplyMLARotaryEmbKV.apply( + return _FusedMLARoPEKVSplit.apply( kv, k_pos_emb, cos, @@ -779,4 +933,64 @@ def fused_apply_mla_rope_for_kv( cp_rank, cp_size, rotary_interleaved, + remove_interleaving, + ) + + +def fused_apply_mla_rope_for_q( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + qk_head_dim: int, + emb_dim: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Backward-compatible in-place MLA query RoPE API. + + New callers should choose :func:`fused_mla_rope_inplace` or + :func:`fused_mla_rope_out_of_place` explicitly. This legacy name keeps + its original mutation behavior and does not add a clone to the hot path. + """ + return fused_mla_rope_inplace( + t, + cos, + sin, + qk_head_dim, + emb_dim, + cu_seqlens_q=cu_seqlens_q, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, + ) + + +def fused_apply_mla_rope_for_kv( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + emb_dim: int, + k_dim: int, + v_dim: int, + cu_seqlens_kv: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Backward-compatible name for the MLA key/value split RoPE API.""" + return fused_mla_rope_kv_split( + kv, + k_pos_emb, + cos, + sin, + emb_dim, + k_dim, + v_dim, + cu_seqlens_kv=cu_seqlens_kv, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, ) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index 9fab25a3fae..0468ddd14ae 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -95,6 +95,8 @@ def _apply_rotary_pos_emb_bshd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, multi_latent_attention: Optional[bool] = None, ) -> Tensor: """Apply rotary positional embedding to input tensor T. @@ -118,6 +120,13 @@ def _apply_rotary_pos_emb_bshd( ) mla_rotary_interleaved = multi_latent_attention + # Some callers may pass freqs with an extra singleton axis, e.g. + # t: [s, b, d] and freqs: [s, 1, 1, d]. In that case, broadcasting would + # accidentally expand to [s, s, b, d]. Squeeze the extra singleton axis to + # keep freqs rank aligned with t. + if freqs.dim() == t.dim() + 1 and freqs.size(-2) == 1: + freqs = freqs.squeeze(-2) + rot_dim = freqs.shape[-1] # ideally t_pass is empty so rotary pos embedding is applied to all tensor t @@ -132,8 +141,18 @@ def _apply_rotary_pos_emb_bshd( # second part is sine component, need to change signs with _rotate_half method cos_ = (torch.cos(freqs) * mscale).to(t.dtype) sin_ = (torch.sin(freqs) * mscale).to(t.dtype) + if inverse: + sin_ = -sin_ t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_) + + # Fallback to original permutation + # DSv4 applies rope on V and O, so we need to uninterleave the tensor. + # The existing MLA code is safe because the dot product is permutation-invariant. + if mla_rotary_interleaved and mla_output_remove_interleaving: + x1, x2 = torch.chunk(t, 2, dim=-1) + t = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) + return torch.cat((t, t_pass), dim=-1) @@ -193,20 +212,28 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = None, + max_seqlen: Optional[int] = None, ) -> Tensor: - """A baseline implementation of applying RoPE for `thd` format. + """Apply RoPE for `thd` format using vectorized CUDA operations. + + When ``max_seqlen`` is supplied, this path performs no GPU-to-CPU sync and is + compatible with CUDA Graph capture. The compatibility path for legacy callers + that omit ``max_seqlen`` retains one GPU-to-CPU sync. Args: - t (Tensor): Input tensor T is of shape [t, h, d] - cu_seqlens(Tensor): Cumulative sum of sequence lengths in a batch for `t`, - with shape [b + 1] and dtype torch.int32. - freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d] - cp_group (torch.distributed.ProcessGroup): The context parallel group + t (Tensor): Input tensor of shape [total_tokens, h, d] + cu_seqlens (Tensor): Cumulative sequence lengths, shape [num_seqs + 1], int32. + freqs (Tensor): RoPE frequencies, shape [max_s, 1, 1, d] or [total_tokens, 1, 1, d] + cp_group: Context parallel group + max_seqlen: Global max sequence length for this packed batch when known. Supplying it + avoids the compatibility-path host sync used by legacy callers. Returns: - Tensor: Shape [t, h, d]. The input tensor after applying RoPE. + Tensor: Shape [total_tokens, h, d]. Input with RoPE applied. """ if multi_latent_attention is not None: warnings.warn( @@ -219,53 +246,71 @@ def _apply_rotary_pos_emb_thd( raise ValueError("cp_group must be provided for THD format RoPE") cp_size = cp_group.size() cp_rank = cp_group.rank() - seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist() - sequence_splits = torch.split(t, seqlens) - total_seqlen = int(cu_seqlens[-1].item()) - has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen - - # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains positions for the whole packed - # batch. Each sequence must therefore use its cu_seqlens offset when selecting the local CP - # front/back slices. For example, with cu_seqlens=[0, 4, 8], cp_size=2, rank 0 should use - # positions [0, 3, 4, 7], not [0, 3, 0, 3]. - # 2. Otherwise: freqs contains only max sequence length positions. Each packed sequence should - # reuse positions starting from 0, preserving the legacy THD behavior. - if has_packed_freqs: - # CASE 1: Exact mapping with offsets - local_freqs = [] - for i, x in enumerate(sequence_splits): - # cu_seqlens[i] is the starting offset of this sequence in the original batch - seq_start_offset = cu_seqlens[i].item() - local_freqs.append( - _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) - ) - freqs = torch.cat(local_freqs, dim=0) - return _apply_rotary_pos_emb_bshd( - t.unsqueeze(1), - freqs, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - - # CASE 2: Traditional mapping without offsets. Apply RoPE one sequence at a time so the second - # and later packed sequences do not look like continuations of the first sequence. - output = torch.empty_like(t) - output_offset = 0 - for x in sequence_splits: - freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) - output_slice = _apply_rotary_pos_emb_bshd( - x.unsqueeze(1), - freq_slice, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - output.narrow(0, output_offset, x.size(0)).copy_(output_slice) - output_offset += x.size(0) + total_tokens = t.shape[0] + device = t.device + + token_pos = torch.arange(total_tokens, device=device, dtype=torch.int64) + + # `cu_seqlens` describes the global packed sequence. With CP, `t` is already + # CP-partitioned, so build a local cumulative-length view before assigning + # local tokens to packed sequences. + cu_seqlens_i64 = cu_seqlens.to(torch.int64) + global_seq_lens = cu_seqlens_i64[1:] - cu_seqlens_i64[:-1] + local_seq_lens = global_seq_lens // cp_size if cp_size > 1 else global_seq_lens + local_cu_seqlens = torch.zeros_like(cu_seqlens_i64) + local_cu_seqlens[1:] = torch.cumsum(local_seq_lens, dim=0) + + # `searchsorted(..., right=True) - 1` returns the local sequence index. The + # clamp guards padded tokens that sit beyond the final real local token; they + # get a harmless frequency and are later masked out. + seq_idx = torch.searchsorted(local_cu_seqlens, token_pos, right=True) - 1 + seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2) + + local_seq_start = local_cu_seqlens[seq_idx] + local_pos = token_pos - local_seq_start + local_seq_len = local_seq_lens[seq_idx] + global_seq_start = cu_seqlens_i64[seq_idx] - return output + if cp_size > 1: + cp_seg = local_seq_len // 2 + full_seqlen = local_seq_len * cp_size + is_first_half = local_pos < cp_seg + freq_pos = torch.where( + is_first_half, + cp_rank * cp_seg + local_pos, + full_seqlen - (cp_rank + 1) * cp_seg + (local_pos - cp_seg), + ) + else: + freq_pos = local_pos.to(torch.int64) + + if max_seqlen is None: + # Backward compatibility for callers that predate ``max_seqlen``. This retains + # the old packed-frequency semantics at the cost of a GPU-to-CPU sync. Updated + # training paths pass ``max_seqlen`` and stay CUDA-graph safe. + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == int(cu_seqlens[-1].item()) + else: + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) > max_seqlen + if exact_packed_freqs: + # `freqs` covers all positions across all sequences (used for non-1D + # RoPE / VLMs); shift by the per-sequence start offset so each token + # samples its absolute position. When `freqs` only spans one max-len + # sequence, no shift is needed. + freq_pos = freq_pos + global_seq_start + + # Padded positions can sit outside the frequency table. Clamp them into + # range; downstream padding masks exclude those positions from the result. + freq_pos = freq_pos.clamp(min=0, max=freqs.shape[0] - 1) + freqs_packed = freqs[freq_pos] + + return _apply_rotary_pos_emb_bshd( + t.unsqueeze(1), + freqs_packed, + rotary_interleaved=rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ).squeeze(1) def apply_rotary_pos_emb( @@ -276,6 +321,9 @@ def apply_rotary_pos_emb( mscale: float = 1.0, cp_group: torch.distributed.ProcessGroup = None, mla_rotary_interleaved: bool = False, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, + max_seqlen: Optional[int] = None, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -312,6 +360,12 @@ def apply_rotary_pos_emb( "Using unfused implementation." ) use_unfused = True + if inverse: + warnings.warn( + "inverse RoPE is not supported by TE's fused RoPE. " + "Using unfused implementation." + ) + use_unfused = True if not use_unfused: assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) @@ -333,6 +387,8 @@ def apply_rotary_pos_emb( rotary_interleaved=config.rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) else: return _apply_rotary_pos_emb_thd( @@ -343,6 +399,9 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + max_seqlen=max_seqlen, ) diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 0e560f939f2..e591e4ff90d 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -205,6 +205,47 @@ def forward( return emb + def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False, cp_group=None): + """Materialize cached cos/sin tensors for ``[seq_len, ..., dim]``.""" + self.max_seq_len_cached = seq_len + self.offset_cached = offset + self.dtype_cached = dtype + self.packed_seq_cached = packed_seq + + emb = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group) + self.register_buffer("cos_cached", emb.cos().to(dtype).contiguous(), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype).contiguous(), persistent=False) + + def get_cached_cos_sin( + self, + seq_len, + offset=0, + dtype=torch.get_default_dtype(), + packed_seq=False, + cp_group=None, + mscale=None, + ): + """Get cached cos and sin values. + + The cache is rebuilt on first use or whenever ``seq_len`` grows + beyond the cached length, or any of ``offset`` / ``dtype`` / + ``packed_seq`` changes from the previous call. + ``YarnRotaryEmbedding`` overrides this to also bake its + concentration factor into the cached cos/sin (controlled by + ``mscale``); for the base class without a concentration + factor the argument is accepted-and-ignored for API uniformity. + """ + del mscale # base class has no concentration factor + if ( + not hasattr(self, "max_seq_len_cached") + or seq_len > self.max_seq_len_cached + or offset != self.offset_cached + or dtype != self.dtype_cached + or packed_seq != self.packed_seq_cached + ): + self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group) + return (self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]) + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): state_dict.pop(f'{prefix}inv_freq', None) return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py index 166ef9b41e7..cb8a03d0b2b 100644 --- a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py @@ -186,13 +186,18 @@ def forward( emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) return emb, _mscale - def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False): + def _set_cos_sin_cache( + self, seq_len, offset, dtype, packed_seq=False, cp_group=None, mscale=None + ): self.max_seq_len_cached = seq_len self.offset_cached = offset self.dtype_cached = dtype self.packed_seq_cached = packed_seq + self.mscale_cached = mscale - emb, _mscale = self.forward(seq_len, offset, packed_seq) + emb, _mscale = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group) + if mscale is not None: + _mscale = mscale self.register_buffer( "cos_cached", (emb.cos() * _mscale).to(dtype).contiguous(), persistent=False ) @@ -201,16 +206,34 @@ def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False): ) def get_cached_cos_sin( - self, seq_len, offset=0, dtype=torch.get_default_dtype(), packed_seq=False + self, + seq_len, + offset=0, + dtype=torch.get_default_dtype(), + packed_seq=False, + cp_group=None, + mscale=None, ): - """Get cached cos and sin values.""" + """Get cached cos and sin values. + + Args: + mscale: when ``None`` (default), the cached cos/sin are + multiplied by yarn's internal concentration factor (the + normal long-context behaviour). When a float is supplied, + that value is used in place of the internal factor — e.g. + the DSv4 hybrid model passes ``mscale=1.0`` to enforce + its "pure rotation" contract and keep the fused / + unfused rope paths bit-equivalent. + """ if ( - seq_len > self.max_seq_len_cached + not hasattr(self, "max_seq_len_cached") + or seq_len > self.max_seq_len_cached or offset != self.offset_cached or dtype != self.dtype_cached or packed_seq != self.packed_seq_cached + or mscale != getattr(self, "mscale_cached", None) ): - self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq) + self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group, mscale) return (self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index a76fe6e3a23..5189d264e59 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -10,6 +10,18 @@ AbsorbedMLASelfAttention, AbsorbedMLASelfAttentionSubmodules, ) +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerSubmodules, @@ -130,6 +142,57 @@ def get_dsa_module_spec_for_backend( return attention +def get_dsv4_hybrid_module_spec_for_backend( + config: TransformerConfig, backend: BackendSpecProvider = None +) -> ModuleSpec: + """Build the native SBHD DSv4 hybrid-attention module spec.""" + assert config.multi_latent_attention, "Currently only MLA supports sparse attention." + assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA." + + rms_norm = config.normalization == "RMSNorm" + qk_norm = ( + backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp + ) + + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=backend.linear(), + linear_wgate=backend.linear(), + norm=backend.layer_norm(rms_norm=True, for_qk=False), + ), + ) + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=backend.linear(), + linear_weights_proj=backend.linear(), + compressor=compressor_spec, + ), + ) + core_attention = ModuleSpec( + module=CompressedSparseAttention, + submodules=CompressedSparseAttentionSubmodules( + compressor=compressor_spec, indexer=indexer_spec + ), + ) + + return ModuleSpec( + module=DSv4HybridSelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=DSv4HybridSelfAttentionSubmodules( + linear_q_down_proj=backend.linear(), + linear_q_up_proj=backend.column_parallel_linear(), + linear_kv_proj=backend.column_parallel_linear(), + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=qk_norm, + kv_layernorm=qk_norm, + ), + metainfo={"fuse_input_layernorm": False}, + ) + + def get_experimental_attention_variant_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: @@ -142,6 +205,8 @@ def get_experimental_attention_variant_module_spec( return get_gated_delta_net_module_spec(config=config, backend=backend) elif config.experimental_attention_variant == "dsa": return get_dsa_module_spec_for_backend(config=config, backend=backend) + elif config.experimental_attention_variant == "dsv4_hybrid": + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=backend) else: raise ValueError( f"Invalid experimental attention variant: {config.experimental_attention_variant}" diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 0042cbea010..93d0c596622 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -8,7 +8,7 @@ import copy from contextlib import nullcontext from dataclasses import dataclass -from typing import Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import torch from torch import Tensor, nn @@ -25,14 +25,27 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward +from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.cuda_graphs import annotate_first_last_layer +from megatron.core.transformer.hyper_connection import ( + HyperConnectionModule, + learned_output_contract, +) from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.module import ( + MegatronModule, + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, +) from megatron.core.transformer.multi_latent_attention import FusedMLASelfAttention from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_layer import TransformerLayer -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.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor @@ -46,12 +59,211 @@ class HybridStackSubmodules: gdn_layer: Union[ModuleSpec, type] = IdentityOp attention_layer: Union[ModuleSpec, type] = IdentityOp dsa_layer: Union[ModuleSpec, type] = IdentityOp + csa_layer: Union[ModuleSpec, type] = IdentityOp + hca_layer: Union[ModuleSpec, type] = IdentityOp mla_layer: Union[ModuleSpec, type] = IdentityOp + window_layer: Union[ModuleSpec, type] = IdentityOp mlp_layer: Union[ModuleSpec, type] = IdentityOp moe_layer: Union[ModuleSpec, type] = IdentityOp mtp_block_spec: Optional[ModuleSpec] = None +class HyperConnectionHybridLayer(MegatronModule): + """Layer-boundary mHC wrapper for HybridStack layers. + + Hybrid layers already own their local residual paths. Each wrapped layer is + treated as one function by aggregating n streams to its input, running the + existing layer, and feeding only the layer delta back through mHC expansion. + + This wrapper nests the inner layer under inner_layer. Checkpoints cannot + switch between mHC-enabled and ordinary HybridStacks without key migration. + """ + + def __init__(self, config: TransformerConfig, layer: MegatronModule) -> None: + super().__init__(config=config) + self.inner_layer = layer + self.layer_number = layer.layer_number + self.hyper_connection = HyperConnectionModule(config=config, layer_number=self.layer_number) + if config.params_dtype is not None: + convert_module_to_dtype_except_fp32_marked(self.hyper_connection, config.params_dtype) + if hasattr(layer, 'tp_group'): + self.tp_group = layer.tp_group + + def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: + """Delegate Mamba inference state shape requests to the wrapped layer.""" + if not hasattr(self.inner_layer, 'mamba_state_shapes_per_request'): + return None + return self.inner_layer.mamba_state_shapes_per_request() + + def _call_inner_layer( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + ) -> Tuple[Tensor, Optional[Tensor]]: + if isinstance(self.inner_layer, TransformerLayer): + output = self.inner_layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + _called_from_hybrid_mhc_wrapper=True, + ) + else: + # Mamba-like layers only consume the common HybridStack arguments. + output = self.inner_layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + ) + + if isinstance(output, tuple): + context = output[1] if len(output) > 1 else None + return output[0], context + return output, None + + def _call_inner_transformer_layer_without_local_bda( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext], + rotary_pos_emb: Optional[Tensor], + sequence_len_offset: Optional[Tensor], + packed_seq_params: Optional[PackedSeqParams], + padding_mask: Optional[Tensor], + ) -> Optional[Tuple[Tuple[Tensor, Optional[Tensor]], Optional[Tensor], float, bool]]: + """Return a raw branch output for split Hybrid TransformerLayer instances. + + Hybrid layers are normally attention-only or MLP/MoE-only. For those + layers, bypass the inner layer's local residual/BDA and let the mHC BDA + own that operation directly. + """ + if not isinstance(self.inner_layer, TransformerLayer): + return None + + layer = self.inner_layer + if InferenceMode.is_active() and layer.config.inference_fuse_tp_communication: + return None + + has_attention = not isinstance(layer.self_attention, IdentityOp) + has_cross_attention = not isinstance(layer.cross_attention, IdentityOp) + has_mlp = not isinstance(layer.mlp, IdentityOp) + + if has_cross_attention or has_attention == has_mlp: + return None + + if has_attention: + output_with_bias, attn_norm_manager, residual = ( + layer._forward_self_attention_output_with_bias( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + ) + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, attn_norm_manager, forced_released_tensors=[residual] + ) + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + output_with_bias, residual = layer._forward_mlp_output_with_bias( + hidden_states, + inference_context=inference_context, + padding_mask=padding_mask, + packed_seq_params=packed_seq_params, + ) + if layer.mlp_norm_manager is not None: + output_with_bias = layer._group_offload_output_with_bias( + output_with_bias, layer.mlp_norm_manager, forced_released_tensors=[residual] + ) + layer.mlp_norm_manager = None + return output_with_bias, None, layer.hidden_dropout, layer.config.bias_dropout_fusion + + def forward( + self, + hidden_states: Tensor, + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Tensor] = None, + sequence_len_offset: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[Tensor] = None, + mhc_recompute_manager=None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """Run the wrapped hybrid layer through one layer-boundary mHC update.""" + aggregated, h_res, h_post, residual = self.hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager, return_residual=True + ) + fast_path_result = self._call_inner_transformer_layer_without_local_bda( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + ) + + if fast_path_result is None: + layer_output, context = self._call_inner_layer( + aggregated, + attention_mask, + inference_context, + rotary_pos_emb, + sequence_len_offset, + packed_seq_params, + padding_mask, + ) + if self.config.fp32_residual_connection and aggregated.dtype != layer_output.dtype: + aggregated = aggregated.to(layer_output.dtype) + layer_output_with_bias = (layer_output - aggregated, None) + dropout_prob = 0.0 + bias_dropout_fusion = False + else: + layer_output_with_bias, context, dropout_prob, bias_dropout_fusion = fast_path_result + + layer_output = layer_output_with_bias[0] + if layer_output.shape != aggregated.shape: + raise RuntimeError( + "HyperConnectionHybridLayer requires wrapped branches to preserve " + f"hidden-state shape. Got {tuple(layer_output.shape)} from wrapped branch " + f"vs {tuple(aggregated.shape)} input." + ) + + is_last_in_recompute_block = bool( + mhc_recompute_manager is not None + and getattr(mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_bda_manager = None if is_last_in_recompute_block else mhc_recompute_manager + hidden_states = self.hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + layer_output_with_bias, + dropout_prob=dropout_prob, + training=self.training, + fused=bias_dropout_fusion, + manager=mhc_bda_manager, + ) + if ( + self.config.fp32_residual_connection + and self.config.params_dtype is not None + and hidden_states.dtype != self.config.params_dtype + ): + hidden_states = hidden_states.to(self.config.params_dtype) + return hidden_states, context + + class HybridStack(MegatronModule): """ Constructor for the HybridStack class. @@ -111,6 +323,8 @@ def __init__( self.input_tensor = None self.pg_collection = pg_collection + self._mhc_block_end_plan: Optional[List[bool]] = None + assert layer_type_list is not None, ( "layer_type_list must be provided. It should be pre-computed from " "--hybrid-layer-pattern by HybridModel." @@ -162,6 +376,28 @@ def __init__( pp_layer_offset=pp_layer_offset, name=(name + f".layers.{i}") if name is not None else None, ) + elif layer_type == LayerSymbols.CSA: + layer = build_module( + submodules.csa_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) + elif layer_type == LayerSymbols.HCA: + layer = build_module( + submodules.hca_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) elif layer_type == LayerSymbols.MLA: layer = build_module( submodules.mla_layer, @@ -172,6 +408,17 @@ def __init__( add_layer_offset=False, pp_layer_offset=pp_layer_offset, ) + elif layer_type == LayerSymbols.WINDOW: + layer = build_module( + submodules.window_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + name=(name + f".layers.{i}") if name is not None else None, + ) elif layer_type == LayerSymbols.MLP: layer = build_module( submodules.mlp_layer, @@ -202,6 +449,8 @@ def __init__( ) else: raise ValueError("unexpected layer_type") + if self.config.enable_hyper_connections: + layer = HyperConnectionHybridLayer(config=self.config, layer=layer) self.layers.append(layer) if self.config.cuda_graph_impl == "local": @@ -218,6 +467,18 @@ def __init__( eps=self.config.layernorm_epsilon, ) + if self.config.enable_hyper_connections and self.post_process and not self.is_mtp_layer: + hc_mult = self.config.num_residual_streams + hc_dim = self.config.hidden_size * hc_mult + self.hc_head_fn = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) + self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) + self.hc_head_scale = mark_keep_in_fp32(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) + def _fuse_mla_down_proj(self, submodules: HybridStackSubmodules) -> HybridStackSubmodules: # Avoid modifying the original object so users don't get surprised about their `submodules` # being modified underneath them. @@ -258,6 +519,49 @@ def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int return layer.mamba_state_shapes_per_request() return None + def _compute_mhc_block_end_plan(self) -> List[bool]: + """Compute deterministic per-layer mHC recompute block boundaries.""" + num_layers = len(self.layers) + block_ends: List[bool] = [False] * num_layers + if num_layers == 0: + return block_ends + + layers_per_block = self.config.mhc_recompute_layer_num + for layer_idx in range(num_layers): + is_last_in_stack = layer_idx == num_layers - 1 + block_ends[layer_idx] = is_last_in_stack or ( + layers_per_block is not None and (layer_idx + 1) % layers_per_block == 0 + ) + return block_ends + + def _build_mhc_recompute_layer_plan( + self, use_mhc_recompute: bool + ) -> Tuple[List[Optional[CheckpointManager]], List[bool]]: + """Build single-use recompute managers for this forward pass.""" + num_layers = len(self.layers) + if not use_mhc_recompute or num_layers == 0: + return [None] * num_layers, [False] * num_layers + + if self._mhc_block_end_plan is None: + self._mhc_block_end_plan = self._compute_mhc_block_end_plan() + block_ends = self._mhc_block_end_plan + + layer_managers: List[Optional[CheckpointManager]] = [None] * num_layers + manager = CheckpointManager() + for layer_idx in range(num_layers): + layer_managers[layer_idx] = manager + if block_ends[layer_idx] and layer_idx != num_layers - 1: + manager = CheckpointManager() + return layer_managers, block_ends + + @staticmethod + def _finalize_mhc_recompute_layer( + manager: Optional[CheckpointManager], hidden_states: Tensor, is_block_end: bool + ) -> None: + """Finalize the current mHC recompute block when its last layer finishes.""" + if manager is not None and is_block_end: + manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -297,6 +601,11 @@ def forward( if isinstance(hidden_states, WrappedTensor): hidden_states = hidden_states.unwrap() + if self.config.enable_hyper_connections and self.pre_process and not self.is_mtp_layer: + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.config.num_residual_streams + ) + if inference_context and inference_context.is_static_batching(): # NOTE(bnorick): match BaseInferenceContext attributes for # mamba_ssm.utils.generation.BaseInferenceContext, @@ -344,6 +653,14 @@ def get_inner_quant_context(config, layer_number): def get_inner_quant_context(config, layer_number): return nullcontext() + use_mhc_recompute = ( + self.training + and self.config.enable_hyper_connections + and self.config.recompute_granularity == 'selective' + and "mhc" in self.config.recompute_modules + ) + mhc_layer_managers, mhc_block_ends = self._build_mhc_recompute_layer_plan(use_mhc_recompute) + with outer_fp8_context: if self.config.recompute_granularity == 'full' and self.training: hidden_states = checkpointed_forward( @@ -359,14 +676,18 @@ def get_inner_quant_context(config, layer_number): use_inner_quantization_context=(use_inner_fp8_context or use_fp4_context), ) else: - for layer in self.layers: + for layer_idx, layer in enumerate(self.layers): # Layers have 1-indexed layer numbers attribute. inner_quant_context = get_inner_quant_context( self.config, layer.layer_number - 1 ) + mhc_manager = mhc_layer_managers[layer_idx] + if mhc_manager is not None: + mhc_manager.is_last_layer_in_recompute_block = mhc_block_ends[layer_idx] + with inner_quant_context: - if isinstance(layer, TransformerLayer): - hidden_states, _ = layer( + if isinstance(layer, (TransformerLayer, HyperConnectionHybridLayer)): + layer_kwargs = dict( hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, @@ -375,6 +696,11 @@ def get_inner_quant_context(config, layer_number): packed_seq_params=packed_seq_params, padding_mask=padding_mask, ) + if mhc_manager is not None and isinstance( + layer, HyperConnectionHybridLayer + ): + layer_kwargs["mhc_recompute_manager"] = mhc_manager + hidden_states, _ = layer(**layer_kwargs) else: # MambaLayer, Expert, or MLP hidden_states = layer( hidden_states=hidden_states, @@ -389,6 +715,25 @@ def get_inner_quant_context(config, layer_number): if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] + self._finalize_mhc_recompute_layer( + manager=mhc_manager, + hidden_states=hidden_states, + is_block_end=mhc_block_ends[layer_idx], + ) + + mhc_multistream = None + if self.config.enable_hyper_connections and self.post_process and not self.is_mtp_layer: + if (self.config.mtp_num_layers or 0) > 0: + mhc_multistream = hidden_states + 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.post_process and self.post_layer_norm: hidden_states = self.final_norm(hidden_states) @@ -399,6 +744,8 @@ def get_inner_quant_context(config, layer_number): inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True ) + if mhc_multistream is not None: + return hidden_states, mhc_multistream return hidden_states def sharded_state_dict( @@ -423,6 +770,7 @@ def sharded_state_dict( dict: The sharded state dictionary for the current object. """ + sharded_offsets = sharded_offsets or () sharded_state_dict = {} layer_prefix = f'{prefix}layers.' @@ -457,6 +805,20 @@ def 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 diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index 83a6163b88d..948feb5aaa7 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -18,12 +18,16 @@ class Symbols: GDN = 'G' ATTENTION = "*" DS_ATTENTION = "D" + CSA = "C" # DSv4 Compressed Sparse Attention (compress_ratio=4) + HCA = "H" # DSv4 Heavily Compressed Attention (compress_ratio=128) MLA = "+" + WINDOW = "W" # DSv4 sliding-window-only attention (compress_ratio=0) MLP = "-" MOE = 'E' PIPE = '|' MTP_SEPARATOR = "/" - VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLA, MLP, MOE} + VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, CSA, HCA, MLA, WINDOW, MLP, MOE} + MLA_ATTENTION = {MLA, DS_ATTENTION, CSA, HCA, WINDOW} @classmethod def name_sorted_valid_layer_symbols(cls) -> list[str]: @@ -174,10 +178,10 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: Examples: >>> get_hybrid_layer_counts("M*M*") - {'*': 2, 'G': 0, 'D': 0, 'M': 2, '-': 0, 'E': 0} + {'*': 2, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, 'W': 0} >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") - {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0} + {'*': 1, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'M': 8, '+': 0, '-': 4, 'E': 0, 'W': 0} """ parsed = parse_hybrid_pattern(pattern) counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()} @@ -293,9 +297,11 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) f"Valid symbols are: {valid_chars}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in pattern and (Symbols.DS_ATTENTION in pattern or Symbols.MLA in pattern): - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # MLA variants may coexist, but standard attention cannot share a model with them. + if Symbols.ATTENTION in pattern and any(symbol in pattern for symbol in Symbols.MLA_ATTENTION): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) def validate_segment_layers(segment: str) -> List[str]: @@ -321,9 +327,11 @@ def validate_segment_layers(segment: str) -> List[str]: f"one of {Symbols.VALID_LAYERS}" ) - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in segment and (Symbols.DS_ATTENTION in segment or Symbols.MLA in segment): - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + # MLA variants may coexist, but standard attention cannot share a model with them. + if Symbols.ATTENTION in segment and any(symbol in segment for symbol in Symbols.MLA_ATTENTION): + raise ValueError( + "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model" + ) return layer_type_list diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 03fef58159f..fe5c1dbb6ec 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -1,4 +1,5 @@ # Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. +from dataclasses import replace from functools import partial from megatron.core.extensions.transformer_engine import ( @@ -10,6 +11,9 @@ TERowParallelLinear, ) from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, +) from megatron.core.models.gpt.moe_module_specs import ( get_inference_optimized_moe_spec, get_moe_module_spec, @@ -49,6 +53,7 @@ MultiTokenPredictionLayerSubmodules, ) from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( MoETransformerLayer, TransformerLayer, @@ -77,7 +82,11 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, + # Hybrid MTP selects the combined projection normally and + # per-stream projections when mHC is enabled. eh_proj=TEColumnParallelLinear, + e_proj=TEColumnParallelLinear, + h_proj=TEColumnParallelLinear, mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), @@ -340,7 +349,10 @@ submodules=MultiTokenPredictionLayerSubmodules( enorm=TENorm, hnorm=TENorm, + # Keep both projection forms available for Hybrid MTP. eh_proj=InferenceColumnParallelLinear, + e_proj=InferenceColumnParallelLinear, + h_proj=InferenceColumnParallelLinear, mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), @@ -355,3 +367,39 @@ # Backward-compatible aliases mamba_stack_spec = hybrid_stack_spec mamba_inference_stack_spec = hybrid_inference_stack_spec + + +def hybrid_dsv4_stack_spec(config: TransformerConfig) -> ModuleSpec: + """Build a HybridStack with fixed-ratio DSv4 C/H/W attention layers. + + The cloned stack deliberately preserves the ordinary ``D`` DSA layer and + the ``+`` MLA layer from :data:`hybrid_stack_spec`. + """ + assert ( + config.transformer_impl == "transformer_engine" + ), "DSv4 HybridModel currently supports only the transformer-engine implementation." + + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + + dsv4_attention = get_dsv4_hybrid_module_spec_for_backend( + config=config, backend=TESpecProvider() + ) + + def wrap_dsv4_layer(compress_ratio: int) -> ModuleSpec: + attention = replace( + dsv4_attention, params={**dsv4_attention.params, "compress_ratio": compress_ratio} + ) + return ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, self_attention=attention, self_attn_bda=get_bias_dropout_add + ), + ) + + submodules = replace( + hybrid_stack_spec.submodules, + csa_layer=wrap_dsv4_layer(compress_ratio=4), + hca_layer=wrap_dsv4_layer(compress_ratio=128), + window_layer=wrap_dsv4_layer(compress_ratio=0), + ) + return replace(hybrid_stack_spec, submodules=submodules) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f0358de57b9..93c31653f83 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -514,7 +514,7 @@ def forward( # assert attention_mask is None, "The attention mask is ignored and should be set to None" # Run decoder. - hidden_states = self.decoder( + decoder_output = self.decoder( hidden_states=decoder_input, attention_mask=attention_mask, inference_context=inference_context, @@ -522,6 +522,11 @@ def forward( packed_seq_params=packed_seq_params, padding_mask=padding_mask, ) + if isinstance(decoder_output, tuple): + hidden_states, mhc_multistream = decoder_output + else: + hidden_states = decoder_output + mhc_multistream = None output_weight = None if self.share_embeddings_and_output_weights: @@ -543,6 +548,7 @@ def forward( input_ids=input_ids, position_ids=position_ids, hidden_states=hidden_states, + mhc_multistream=mhc_multistream, attention_mask=attention_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py index c87ccd5ff31..45bf910f84d 100644 --- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py @@ -370,8 +370,11 @@ def _update_fp32_params_by_new_state(self): if not self.param_update_in_fp32: return for param, v in self.state.items(): - fp32_param = self.param_to_fp32_param[param] - fp32_param.data.copy_(v["master_param"]) + # Native FP32 params do not need a separate master parameter and are + # intentionally absent from param_to_fp32_param. + fp32_param = self.param_to_fp32_param.get(param) + if fp32_param is not None: + fp32_param.data.copy_(v["master_param"]) def update_fp32_param_by_new_param(self): """ diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index dac16f4a2ee..fd8c1864902 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -7,7 +7,6 @@ import math import warnings from abc import ABC, abstractmethod -from itertools import chain from logging import getLogger from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -1147,8 +1146,25 @@ def sharded_state_dict( state_dict = self.state_dict() + # Optimizer state ids enumerate the inner optimizer params: the fp32 main + # copies of float16 params, native fp32 params, and any frozen params, + # interleaved in the original param-group order. Map each fp32 main copy + # back to its model-side param; all other params already are model params. + main_param_id_to_model_param = { + id(main_param): model_param + for model_group, main_group in zip( + self.float16_groups, self.fp32_from_float16_groups, strict=True + ) + for model_param, main_param in zip(model_group, main_group, strict=True) + } + + def model_params_in_optimizer_order(): + for inner_group in self.optimizer.param_groups: + for param in inner_group['params']: + yield main_param_id_to_model_param.get(id(param), param) + id_to_sharded_param_map = get_param_id_to_sharded_param_map( - model_sharded_state_dict, chain.from_iterable(g for g in self.float16_groups) + model_sharded_state_dict, model_params_in_optimizer_order() ) _backfill_gtp_sharded_param_map( @@ -1159,6 +1175,20 @@ def sharded_state_dict( assert len(state_dict['fp32_from_fp16_params']) == len( state_dict['optimizer']['param_groups'] ) + # State ids of the fp32 main copies only, skipping native fp32 and frozen params. + float16_param_ids_per_group = [] + for state_group, inner_group in zip( + state_dict['optimizer']['param_groups'], self.optimizer.param_groups, strict=True + ): + float16_param_ids_per_group.append( + [ + param_id + for param_id, param in zip( + state_group['params'], inner_group['params'], strict=True + ) + if id(param) in main_param_id_to_model_param + ] + ) state_dict['fp32_from_fp16_params'] = [ [ make_sharded_optimizer_tensor( @@ -1166,10 +1196,10 @@ def sharded_state_dict( fp32_param, prefix=f'optimizer.state.fp32_param', ) - for param_id, fp32_param in zip(state_group['params'], fp32_group) + for param_id, fp32_param in zip(param_ids, fp32_group, strict=True) ] - for fp32_group, state_group in zip( - state_dict['fp32_from_fp16_params'], state_dict['optimizer']['param_groups'] + for fp32_group, param_ids in zip( + state_dict['fp32_from_fp16_params'], float16_param_ids_per_group, strict=True ) ] diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index 1daeacc9027..42b350ebd00 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -277,7 +277,7 @@ def _get_experimental_attention_variant_loss_scale_func(config): if loss_scale_func is not None: return loss_scale_func - if getattr(config, 'experimental_attention_variant', None) == 'dsa': + if getattr(config, 'experimental_attention_variant', None) in ('dsa', 'dsv4_hybrid'): from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, ) diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index a9619ea4819..57041441a47 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Parts of the code here are adapted from PyTorch # repo: https://github.com/pytorch/pytorch @@ -628,7 +628,9 @@ def forward( @staticmethod def backward(ctx, *args): """Backward pass.""" - if not torch.autograd._is_checkpoint_valid(): + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + if not torch.autograd._is_checkpoint_valid() and not is_graph_capturing(): raise RuntimeError( "Checkpointing is not compatible with .grad(), " "please use .backward() if possible" @@ -679,10 +681,67 @@ def checkpoint( return CheckpointFunction.apply(function, distribute_saved_activations, *args) +def _save_args_to_ctx(ctx, args): + """Save mixed tensor/non-tensor arguments into autograd ctx. + + Since save_for_backward only supports tensors, this function separates + tensor and non-tensor arguments, saving tensors via save_for_backward + and storing non-tensor metadata (indices and values) as ctx attributes. + + Use _load_args_from_ctx to reconstruct the original args. + """ + tensor_args = [] + non_tensor_entries = [] + + for index, arg in enumerate(args): + if isinstance(arg, torch.Tensor): + tensor_args.append(arg) + continue + non_tensor_entries.append((index, arg)) + + ctx.save_for_backward(*detach_variable(tuple(tensor_args))) + ctx._non_tensor_entries = tuple(non_tensor_entries) + ctx._total_args_count = len(args) + + +def _load_args_from_ctx(ctx): + """Load and reconstruct mixed tensor/non-tensor arguments from autograd ctx. + + This is the inverse of _save_args_to_ctx. It retrieves tensors from + ctx.saved_tensors and merges them with stored non-tensor arguments + to reconstruct the original args in their original order. + + Returns: + tuple of reconstructed arguments in their original order. + """ + + def _detach_with_grad(tensor): + detached = tensor.detach() + detached.requires_grad_(tensor.requires_grad) + return detached + + tensor_iter = iter(_detach_with_grad(t) for t in ctx.saved_tensors) + total_args_count = ctx._total_args_count + non_tensor_map = dict(ctx._non_tensor_entries) + + reconstructed_args = [] + for index in range(total_args_count): + if index in non_tensor_map: + reconstructed_args.append(non_tensor_map[index]) + else: + reconstructed_args.append(next(tensor_iter)) + return tuple(reconstructed_args) + + class CheckpointWithoutOutputFunction(torch.autograd.Function): """ Checkpoint Function Helper for CheckpointWithoutOutput. Save context for recompute. + + Handles both tensor and non-tensor arguments: + - Tensor arguments are saved via save_for_backward + - Non-tensor arguments (int, float, bool, None, etc.) are stored separately + in ctx attributes and reconstructed during recomputation """ @staticmethod @@ -705,7 +764,10 @@ def forward( with torch.no_grad(), fwd_ctx: outputs = run_function(*args) - ctx.save_for_backward(*detach_variable(args)) + + # Save tensor and non-tensor arguments into ctx for recomputation + _save_args_to_ctx(ctx, args) + # the CheckpointWithoutOutput object is passed in, then it can access the saved input # tensors later for recomputation checkpoint_without_output_obj.ctx = ctx @@ -722,10 +784,60 @@ def backward(ctx, *args): torch.autograd.backward(outputs, args) ctx.outputs = None ctx.inputs = None - grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in inputs) + grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in inputs) return (None, None) + grads +class CheckpointManager: + """ + Manages multiple CheckpointWithoutOutput objects within a TransformerBlock + cross layer recomputations, enabling unified recomputation during backward pass. + This is particularly useful for scenarios where multiple checkpoint operations have + sequential dependencies (i.e., the output of one checkpoint is the input of the next). + + Usage: + ckptManager = CheckpointManager() + ckpt_function = CheckpointWithoutOutput(ckpt_manager=ckptManager) + ckpt_function.checkpoint(run_function, *args) + # other checkpointed operations + ckpt_manager.discard_all_outputs_and_register_unified_recompute(final_output) + """ + + def __init__(self): + self.checkpoints = [] + # Set by TransformerBlock before each layer forward. + # When True, the layer should keep block-boundary output uncheckpointed. + self.is_last_layer_in_recompute_block = False + + def add_checkpoint(self, ckpt): + """Add a checkpoint to the manager.""" + if not isinstance(ckpt, CheckpointWithoutOutput): + raise TypeError("Expected CheckpointWithoutOutput object") + if ckpt.outputs is None: + raise ValueError("CheckpointWithoutOutput must call checkpoint() before adding") + self.checkpoints.append(ckpt) + + def discard_all_outputs_and_register_unified_recompute(self, hook_tensor): + """Discard all checkpoint outputs to save memory and register unified recompute hook.""" + for ckpt in self.checkpoints: + for output in ckpt.outputs: + output.untyped_storage().resize_(0) + + # Register unified recompute hook + if hook_tensor.requires_grad: + hook_tensor.register_hook(self._unified_recompute_hook) + + def _unified_recompute_hook(self, grad_output): + for ckpt in self.checkpoints: + # Call _recompute for each checkpoint in forward order + # The _recompute method will restore the output tensor storage + ckpt._recompute(None) + + +# Compatibility for the already-reviewed mHC prerequisite API. +CheckpointWithoutOutputManager = CheckpointManager + + class CheckpointWithoutOutput(object): """ Checkpoint a model or part of the model and release the output. @@ -740,8 +852,19 @@ class CheckpointWithoutOutput(object): discarded output tensors are directly saved in the following modules for backward computation. """ - def __init__(self, fp8=False): - self.fp8 = fp8 is not None + def __init__(self, fp8=False, ckpt_manager=None): + """ + Initialize CheckpointWithoutOutput. + + Args: + fp8: Whether to use FP8 mode. Defaults to False. + ckpt_manager: Optional CheckpointManager instance. When provided, + checkpoint() will auto-register to the manager, and + discard_output_and_register_recompute() will only discard + output without registering individual hooks. + """ + self.fp8 = bool(fp8) + self.ckpt_manager = ckpt_manager self.run_function = None self.fwd_cpu_rng_state = None self.fwd_cuda_rng_state = None @@ -750,7 +873,12 @@ def __init__(self, fp8=False): self.outputs = None def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_Ts]) -> _R: - """Checkpoint function.""" + """ + Checkpoint function. + + If ckpt_manager was provided during initialization, this checkpoint + will be automatically registered to the manager after execution. + """ # If in cuda graph warmup, disable checkpointing, as 'discard_output_and_register_recompute' # may be called in a separate graph warmup. @@ -767,6 +895,11 @@ def checkpoint(self, run_function: Callable[[Unpack[_Ts]], _R], *args: Unpack[_T self.outputs = outputs if isinstance(self.outputs, torch.Tensor): self.outputs = (self.outputs,) + + # Auto-register to manager if provided + if self.ckpt_manager is not None: + self.ckpt_manager.add_checkpoint(self) + return outputs def _recompute(self, _): @@ -775,7 +908,7 @@ def _recompute(self, _): from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup # The recomputation has been triggered already. Just return. - # Handle cudagraphs, do nothing if currently in graph warmup + # Handle cudagraphs: do nothing if currently in graph warmup if self.ctx is None or is_graph_warmup(): return @@ -797,17 +930,8 @@ def _recompute(self, _): recompute_ctx = contextlib.nullcontext() fp8_ctx = contextlib.nullcontext() - # Store the inputs for backward pass - inputs = self.ctx.saved_tensors - - def detach(t): - if isinstance(t, torch.Tensor): - requires_grad = t.requires_grad - t = t.detach() - t.requires_grad_(requires_grad) - return t - - inputs = tuple(detach(t) for t in inputs) + # Reconstruct full args list from saved ctx + inputs = _load_args_from_ctx(self.ctx) with torch.enable_grad(), fp8_ctx, recompute_ctx: outputs = self.run_function(*inputs) @@ -840,10 +964,11 @@ def discard_output_and_register_recompute(self, hook_tensor): in the forward pass and the gradient of the hook_tensor is computed before the recomputed tensors are used. """ - + # When ckpt_manager is set, this is a no-op. + # Manager handles all discarding and hook registration uniformly. from megatron.core.transformer.cuda_graphs import is_graph_warmup - if is_graph_warmup(): + if self.ckpt_manager is not None or is_graph_warmup(): return # use resize to release the output tensor memory and still keep the metadata in the tensors. diff --git a/megatron/core/transformer/__init__.py b/megatron/core/transformer/__init__.py index 0e3cdcfa57e..75e3b485c4f 100644 --- a/megatron/core/transformer/__init__.py +++ b/megatron/core/transformer/__init__.py @@ -1,6 +1,10 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from .module import MegatronModule from .spec_utils import ModuleSpec, build_module from .transformer_config import MLATransformerConfig, TransformerConfig -from .transformer_layer import TransformerLayer, TransformerLayerSubmodules +from .transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, + TransformerLayerSubmodules, +) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 682b75fb701..1f94b853abe 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -303,6 +303,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -314,6 +315,7 @@ def __init__( self.config = config self.layer_number = layer_number self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.attn_mask_type = attn_mask_type self.attention_type = attention_type @@ -1490,8 +1492,16 @@ def forward( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + rope_freqs_max_seqlen = ( + max(rope_max_seqlen_q, rope_max_seqlen_kv) + if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None + else None + ) else: cu_seqlens_q = cu_seqlens_kv = None + rope_freqs_max_seqlen = None if split_qkv: if q_pos_emb is not None: @@ -1504,6 +1514,7 @@ def forward( cu_seqlens=cu_seqlens_q, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_freqs_max_seqlen, ) else: query = inference_context.apply_rotary_emb_query( @@ -1522,6 +1533,7 @@ def forward( cu_seqlens=cu_seqlens_kv, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_freqs_max_seqlen, ) else: query, key, value = apply_fused_qkv_rotary_pos_emb( @@ -1658,6 +1670,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -1673,6 +1686,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -2075,6 +2089,7 @@ def __init__( attn_mask_type: AttnMaskType = AttnMaskType.padding, cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -2089,6 +2104,7 @@ def __init__( attention_type="cross", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, name=name, ) diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index e0b6af7aa7f..9991e6828d1 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -147,6 +147,7 @@ def __init__( pg_collection: ProcessGroupCollection = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -161,6 +162,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) # Resolve which classes to use for Q and KV linear up projections and norms, based on @@ -447,8 +449,16 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + rope_freqs_max_seqlen = ( + max(rope_max_seqlen_q, rope_max_seqlen_kv) + if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None + else None + ) else: cu_seqlens_q = cu_seqlens_kv = None + rope_freqs_max_seqlen = None # ========================================= # Q down projection @@ -636,6 +646,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -646,6 +657,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # query: [num_tokens, n, (kv_lora_rank + qk_pos_emb_head_dim)] diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py new file mode 100644 index 00000000000..9f32588a864 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,909 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import copy +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +from megatron.core.fp8_utils import get_fp8_disabled_context +from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +from megatron.core.models.common.embeddings import RotaryEmbedding, apply_rotary_pos_emb +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, + fused_qk_topk_naive, + rotate_activation, +) +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_range_pop, nvtx_range_push + +# --------------------------------------------------------------------------- +# Helper functions for index computation +# --------------------------------------------------------------------------- + + +@lru_cache(maxsize=8) +def _get_window_topk_idxs_cached(window_size: int, seqlen: int, device_str: str) -> torch.Tensor: + """Compute sliding-window indices for a single sequence (cached). + + Returns: + indices: [seqlen, window_size] int tensor, -1 for invalid positions. + """ + base = torch.arange(seqlen, device=device_str).unsqueeze(1) + offsets = torch.arange(window_size, device=device_str) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix + + +def get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + """Sliding-window indices [batch, seqlen, window_size].""" + matrix = _get_window_topk_idxs_cached(window_size, seqlen, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +@lru_cache(maxsize=8) +def _get_compress_topk_idxs_cached( + ratio: int, seqlen: int, offset: int, device_str: str +) -> torch.Tensor: + """Compute all-compressed-positions indices for a single sequence (cached). + + Returns: + indices: [seqlen, seqlen // ratio] int tensor, -1 for future positions. + """ + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device_str).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix + + +def get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + """All-compressed-position indices [batch, seqlen, seqlen // ratio].""" + matrix = _get_compress_topk_idxs_cached(ratio, seqlen, offset, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +# --------------------------------------------------------------------------- +# Helper functions for RoPE +# --------------------------------------------------------------------------- + + +def _apply_rope( + x: torch.Tensor, + nope_dim: int, + pos_dim: int, + rotary_pos_emb_module: RotaryEmbedding, + config: TransformerConfig, + rotary_seq_len: int, + ratio: int = 1, + cp_group: torch.distributed.ProcessGroup = None, +) -> torch.Tensor: + """Apply RoPE to the last ``qk_pos_emb_head_dim`` dims, leaving the rest unchanged. + + Accepts both 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs. When the input is 3-D a temporary head dimension is inserted for + ``apply_rotary_pos_emb`` and removed before returning. + """ + if ratio == 1: + total_seq_len = rotary_seq_len + else: + total_seq_len = rotary_seq_len * ratio + # 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 + # regardless of which rotary class is in use. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if config.apply_rope_fusion: + # ``mscale=1.0`` keeps the cached cos/sin free of yarn's + # concentration factor so the fused kernel sees the same + # rotation as the unfused split-rotate path (DSv4 "pure + # rotation" contract). + rotary_pos_cos, rotary_pos_sin = rotary_pos_emb_module.get_cached_cos_sin( + total_seq_len, dtype=x.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: + # Compressed-attention callers instantiate ``YarnRotaryEmbedding`` + # whenever ``compress_ratio > 1`` (regardless of ``config.rope_type``); + # its ``forward`` returns ``(emb, mscale)``. Base ``RotaryEmbedding`` + # returns a single tensor. Unpack either form uniformly; the + # caller-side ``mscale=1.0`` keeps the yarn concentration factor + # out of the rotation. + result = rotary_pos_emb_module(total_seq_len, packed_seq=False) + if isinstance(result, tuple): + rotary_pos_emb = result[0] + else: + rotary_pos_emb = result + 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: + rotary_pos_cos = rotary_pos_cos[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_sin is not None and ratio > 1: + rotary_pos_sin = rotary_pos_sin[:total_seq_len:ratio][:rotary_seq_len] + + squeeze_head = x.dim() == 3 + if squeeze_head: + x = x.unsqueeze(-2) + if config.apply_rope_fusion: + out = fused_mla_rope_inplace( + x, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + cp_group.rank(), + cp_group.size(), + remove_interleaving=True, + ) + else: + x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) + x_pe = apply_rotary_pos_emb( + x_pe, + rotary_pos_emb, + config=config, + cu_seqlens=None, + mscale=mscale, + cp_group=cp_group, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + out = torch.cat([x_nope, x_pe], dim=-1) + if squeeze_head: + out = out.squeeze(-2) + return out + + +# --------------------------------------------------------------------------- +# Sparse attention kernel (unfused, differentiable) +# --------------------------------------------------------------------------- + + +def unfused_compressed_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Differentiable sparse attention with MQA and attention sink. + + Args: + query: [sq, b, np, hn] multi-head query. + kv_full: [n_kv, b, hn] single-head KV (original + compressed). + attn_sink: [np] per-head learnable bias. + topk_indices: [b, sq, topk] indices into kv_full (int32, -1 = invalid). + softmax_scale: float + + Returns: + output: [sq, b, np * hn] + """ + sq, b, np_, hn = query.size() + + # --- Gather KV at topk positions --- + # kv_full: [n_kv, b, hn] -> [b, n_kv, hn] + kv_t = kv_full.permute(1, 0, 2) + + safe_indices = topk_indices.clamp(min=0).long() # [b, sq, topk] + safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, -1, hn) # [b, sq, topk, hn] + # [b, n_kv, hn] -> [b, 1, n_kv, hn] -> gather -> [b, sq, topk, hn] + kv_gathered = torch.gather( + kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=safe_indices_exp + ) + + # --- Attention scores --- + # query: [sq, b, np, hn] -> [b, np, sq, hn] + q = query.permute(1, 2, 0, 3).float() + kv_g = kv_gathered.float() # [b, sq, topk, hn] + + # [b, np, sq, topk] + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_g) * softmax_scale + + # Mask invalid + invalid_mask = (topk_indices < 0).unsqueeze(1) # [b, 1, sq, topk] + scores = scores.masked_fill(invalid_mask, float("-inf")) + + # --- Softmax with attention sink --- + sink = attn_sink.view(1, np_, 1, 1).float() + scores_max = scores.max(dim=-1, keepdim=True).values # [b, np, sq, 1] + scores_max = torch.max(scores_max, sink) + + exp_scores = torch.exp(scores - scores_max) # [b, np, sq, topk] + exp_sink = torch.exp(sink - scores_max) # [1, np, 1, 1] + + sum_exp = exp_scores.sum(dim=-1, keepdim=True) + exp_sink + attn_weights = exp_scores / sum_exp # [b, np, sq, topk] + + # --- Weighted sum --- + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_g) + output = output.to(query.dtype) + + # [b, np, sq, hn] -> [sq, b, np, hn] -> [sq, b, np * hn] + output = output.permute(2, 0, 1, 3).contiguous() + output = output.reshape(sq, b, np_ * hn) + return output + + +@torch.no_grad() +def _compute_unfused_csa_non_compressed_lse( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + window_indices: torch.Tensor, + softmax_scale: float, + chunk_size: int = 512, +) -> torch.Tensor: + """Return the detached sliding-window-plus-sink log mass for the CSA teacher. + + Args: + query: Query tensor in ``[sq, batch, heads, head_dim]`` layout. + kv_full: Original (non-compressed) KV in ``[sk, batch, head_dim]`` layout. + attn_sink: Per-head sink logits in ``[heads]`` layout. + window_indices: Local per-batch window indices in ``[batch, sq, window]`` layout. + softmax_scale: Scale applied to query-key logits. + chunk_size: Maximum number of flattened query rows processed at once. + + Returns: + Detached FP32 log-sum-exp values in ``[batch, heads, sq]`` layout. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if query.ndim != 4: + raise ValueError(f"query must have shape [sq, batch, heads, dim], got {query.shape}") + if attn_sink.ndim != 1: + raise ValueError(f"attn_sink must be 1D, got shape {tuple(attn_sink.shape)}") + + seqlen_q, batch_size, num_heads, head_dim = query.shape + if kv_full.ndim != 3 or kv_full.shape[1:] != (batch_size, head_dim): + raise ValueError( + "non-compressed KV must have shape " + f"[sk, {batch_size}, {head_dim}], got {tuple(kv_full.shape)}" + ) + if window_indices.ndim != 3 or window_indices.shape[:2] != (batch_size, seqlen_q): + raise ValueError( + "window_indices must have shape " + f"[{batch_size}, {seqlen_q}, window], got {tuple(window_indices.shape)}" + ) + if attn_sink.numel() != num_heads: + raise ValueError(f"attn_sink must contain {num_heads} values, got {attn_sink.numel()}") + if not (query.device == kv_full.device == attn_sink.device == window_indices.device): + raise ValueError("query, kv_full, attn_sink, and window_indices must share a device") + + n_kv = kv_full.shape[0] + q_flat = query.detach().permute(1, 0, 2, 3).reshape(-1, num_heads, head_dim) + kv_flat = kv_full.detach().permute(1, 0, 2).reshape(-1, head_dim) + batch_offsets = ( + torch.arange(batch_size, device=window_indices.device, dtype=torch.int64) * n_kv + ).view(batch_size, 1, 1) + window_indices_i64 = window_indices.to(dtype=torch.int64) + global_indices = torch.where( + window_indices_i64 >= 0, window_indices_i64 + batch_offsets, window_indices_i64 + ).reshape(batch_size * seqlen_q, -1) + + sink = attn_sink.detach().to(dtype=torch.float32).view(1, num_heads) + lse_chunks = [] + for start in range(0, q_flat.shape[0], chunk_size): + end = min(start + chunk_size, q_flat.shape[0]) + indices = global_indices[start:end] + gathered_kv = kv_flat.index_select(0, indices.clamp(min=0).reshape(-1)).reshape( + end - start, indices.shape[-1], head_dim + ) + window_logits = torch.einsum("rhd,rkd->rhk", q_flat[start:end].float(), gathered_kv.float()) + window_logits = (window_logits * softmax_scale).masked_fill( + (indices < 0).unsqueeze(1), float("-inf") + ) + lse_chunks.append(torch.logaddexp(torch.logsumexp(window_logits, dim=-1), sink)) + + if lse_chunks: + lse_flat = torch.cat(lse_chunks, dim=0) + else: + lse_flat = torch.empty((0, num_heads), dtype=torch.float32, device=query.device) + return lse_flat.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() + + +# --------------------------------------------------------------------------- +# Compressor +# --------------------------------------------------------------------------- + + +@dataclass +class CompressorSubmodules: + """Submodule specs for CSA and HCA Compressor.""" + + linear_wkv: Union[ModuleSpec, type] = None + linear_wgate: Union[ModuleSpec, type] = None + norm: Union[ModuleSpec, type] = None + + +class Compressor(MegatronModule): + """Gated pooling compressor for CSA and HCA sparse attention. + + Compresses a sequence of tokens into a shorter sequence by pooling groups of + ``compress_ratio`` tokens using learned gated weights. + + For ``compress_ratio == 4``, overlapping compression is used (``coff = 2``). + For ``compress_ratio == 128``, non-overlapping compression is used (``coff = 1``). + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressorSubmodules, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError("Compressor requires an explicit ProcessGroupCollection") + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + + self.rotary_pos_emb = rotary_pos_emb + + proj_out_dim = self.coff * head_dim + + with get_fp8_disabled_context(config, is_init=True): + self.linear_wkv = build_module( + submodules.linear_wkv, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wkv") if name is not None else None, + ) + + self.linear_wgate = build_module( + submodules.linear_wgate, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wgate") if name is not None else None, + ) + + # keep to high precision (FP32 in the reference DeepSeek V4 checkpoint) + _ape = torch.empty( + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = mark_keep_in_fp32(nn.Parameter(_ape)) + + norm_config = copy.copy(config) + norm_config.normalization = "RMSNorm" + self.norm = build_module( + submodules.norm, config=norm_config, hidden_size=head_dim, eps=config.layernorm_epsilon + ) + + def backward_dw(self): + """Compute deferred weight gradients for the compressor projections.""" + self.linear_wkv.backward_dw() + self.linear_wgate.backward_dw() + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + """Apply overlapping window transform for 4x compression. + + Input shape: [n_groups, ratio, b, coff * head_dim] + Output shape: [n_groups, 2 * ratio, b, head_dim] + """ + n_groups, ratio, b_dim, _ = tensor.size() + d = self.head_dim + new_tensor = tensor.new_full((n_groups, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + new_tensor[1:, :ratio] = tensor[:-1, :, :, :d] + return new_tensor + + def _project(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Project compressor values and gates outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + kv, _ = self.linear_wkv(x) + score, _ = self.linear_wgate(x) + return kv, score + + def forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """Compress hidden states into shorter KV sequence. + + Args: + x: [sq, b, hidden_size] + + Returns: + compressed_kv [sq // ratio, b, head_dim] or None if too short. + """ + nvtx_range_push("compressor") + + sq, b, _ = x.size() + ratio = self.compress_ratio + + if sq < ratio: + nvtx_range_pop("compressor") + return None + + kv, score = self._project(x) # [sq, b, coff * head_dim] + + cutoff = (sq // ratio) * ratio + if cutoff < sq: + kv = kv[:cutoff] + score = score[:cutoff] + + n_compressed = cutoff // ratio + + # Reshape: [n_compressed, ratio, b, coff * head_dim] + kv = kv.view(n_compressed, ratio, b, -1) + score = score.view(n_compressed, ratio, b, -1) + + # APE: [ratio, coff * head_dim] -> [1, ratio, 1, coff * head_dim] + score = score + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + + weights = torch.softmax(score, dim=1, dtype=torch.float32).to(kv.dtype) + kv = (kv * weights).sum(dim=1) # [n_compressed, b, head_dim] + + kv = self.norm(kv.to(x.dtype)) + + kv = _apply_rope( + kv, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + n_compressed, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + if self.rotate: + kv = rotate_activation(kv) + + nvtx_range_pop("compressor") + return kv # [n_compressed, b, head_dim] + + +# --------------------------------------------------------------------------- +# CSAIndexer +# --------------------------------------------------------------------------- + + +@dataclass +class CSAIndexerSubmodules: + """Submodule specs for CSAIndexer.""" + + linear_wq_b: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + compressor: Union[ModuleSpec, type] = None + + +class CSAIndexer(MegatronModule): + """Learned top-k retrieval over compressed positions for CSA sparse attention. + + Computes index scores to select the most relevant compressed KV positions for each + query. Reuses the scoring logic from ``DSAIndexer`` (einsum -> relu -> weight -> sum + -> topk) and ``rotate_activation`` (Hadamard transform) from ``dsa.py``. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CSAIndexerSubmodules, + compress_ratio: int, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError("CSAIndexer requires an explicit ProcessGroupCollection") + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.hidden_size = config.hidden_size + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.q_lora_rank = ( + config.q_lora_rank if config.q_lora_rank is not None else config.hidden_size + ) + + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim**-0.5 + + self.rotary_pos_emb = rotary_pos_emb + + # Q projection + self.linear_wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.index_n_heads * self.index_head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wq_b") if name is not None else None, + ) + + # The reference DeepSeek V4 checkpoint keeps this projection in BF16. + with get_fp8_disabled_context(config, is_init=True): + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_weights_proj") if name is not None else None, + ) + + # Own compressor (smaller head_dim, with Hadamard rotation) + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=compress_ratio, + head_dim=self.index_head_dim, + rotate=True, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".compressor") if name is not None else None, + ) + + def backward_dw(self): + """Compute deferred weight gradients for the indexer projections.""" + self.linear_wq_b.backward_dw() + self.linear_weights_proj.backward_dw() + self.compressor.backward_dw() + + def _project_weights(self, x: torch.Tensor) -> torch.Tensor: + """Project indexer weights outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + weights, _ = self.linear_weights_proj(x) + return weights + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute Q, compressed K, and weights before top-k selection.""" + nvtx_range_push("indexer_before_topk") + + sq, bsz, _ = x.size() + + # Q path + q, _ = self.linear_wq_b(qr) # [sq, b, n_heads * head_dim] + q = q.reshape(sq, bsz, self.index_n_heads, self.index_head_dim) + q = _apply_rope( + q, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + sq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + q = rotate_activation(q) + + # K path: own compressor + k = self.compressor(x) # [sq//ratio, b, index_head_dim] + + weights = self._project_weights(x) # [sq, b, n_heads] + weights = weights * (self.index_n_heads**-0.5) + + nvtx_range_pop("indexer_before_topk") + return q, k, weights + + def forward( + self, x: torch.Tensor, qr: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (index_scores, topk_indices).""" + nvtx_range_push("indexer") + q, k, weights = self.forward_before_topk(x, qr) + nvtx_range_push("indexer_qk_topk") + effective_topk = min(self.index_topk, k.size(0)) + index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, effective_topk, mask) + nvtx_range_pop("indexer_qk_topk") + nvtx_range_pop("indexer") + return index_scores, topk_indices + + +# --------------------------------------------------------------------------- +# CompressedSparseAttention (core attention) +# --------------------------------------------------------------------------- + + +@dataclass +class CompressedSparseAttentionSubmodules: + """Submodule specs for CompressedSparseAttention.""" + + compressor: Union[ModuleSpec, type] = None + indexer: Union[ModuleSpec, type] = None + + +class CompressedSparseAttention(MegatronModule): + """Sparse core attention for CompressedSparseAttention. + + Combines sliding window attention with compressed KV attention. The spec always + provides compressor and indexer submodule specs; this ``__init__`` inspects + ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: + + * ``ratio == 0``: window-only (compressor and indexer NOT built) + * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) + * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressedSparseAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + k_channels: Optional[int] = None, + v_channels: Optional[int] = None, + cp_comm_type: str = "p2p", + pg_collection: Optional[ProcessGroupCollection] = None, + rotary_pos_emb: nn.Module = None, + compress_ratio: int = 0, + is_mtp_layer: bool = False, + name: str | None = None, + ): + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError( + "CompressedSparseAttention requires an explicit ProcessGroupCollection" + ) + self.pg_collection = pg_collection + + self.layer_number = layer_number + self.config.num_layers if is_mtp_layer else layer_number + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.v_head_dim = config.v_head_dim + + self.n_local_heads = config.num_attention_heads + + if softmax_scale is None: + softmax_scale = config.v_head_dim**-0.5 + self.softmax_scale = softmax_scale + + # Learnable attention sink per head, kept in reference-checkpoint FP32. + self.attn_sink = mark_keep_in_fp32( + nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) + ) + + # Conditionally build Compressor (ratio > 1) + if self.compress_ratio > 1 and submodules.compressor is not None: + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=self.compress_ratio, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".compressor") if name is not None else None, + ) + else: + self.compressor = None + + # Conditionally build Indexer (ratio == 4) + if ( + self.compress_ratio == 4 + and not config.csa_dense_mode + and submodules.indexer is not None + ): + self.indexer = build_module( + submodules.indexer, + config=config, + compress_ratio=self.compress_ratio, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".indexer") if name is not None else None, + ) + else: + self.indexer = None + + def backward_dw(self): + """Compute deferred gradients for the optional compressor and indexer projections.""" + if self.compressor is not None: + self.compressor.backward_dw() + if self.indexer is not None: + self.indexer.backward_dw() + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + x: torch.Tensor = None, + qr: torch.Tensor = None, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params=None, + ) -> torch.Tensor: + """Forward pass for CompressedSparseAttention. + + Args: + query: [sq, b, np, v_head_dim] + key: [sq, b, 1, v_head_dim] (single-head MQA; head dim squeezed internally) + value: unused (key == value in MQA) + attention_mask: attention mask (may be None for causal). + x: [sq, b, hidden_size] original hidden states. + qr: [sq, b, q_lora_rank] compressed query representation. + + Returns: + output: [sq, b, np * v_head_dim] + """ + nvtx_range_push("compressed_sparse_attn") + assert ( + packed_seq_params is None + ), "Packed sequence not supported for CompressedSparseAttention" + + sq, b, np, hn = query.size() + + # --- Step 1: Prepare single-head KV (squeeze singleton head dim) --- + kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] + + # --- Step 2: Compression --- + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) # [n_compressed, b, v_head_dim] + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + n_compressed = 0 + else: + kv_full = kv + n_compressed = 0 + + offset = sq # compressed indices start after original positions + + # --- Step 3: Window indices --- + window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) + + # --- Step 4: Compressed indices --- + indexer_loss = None + + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(b, -1, -1) + ) # [b, sq, n_compressed] + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det + ) + indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) + weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, kv, self.attn_sink, window_idxs, self.softmax_scale + ) + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer, + weights_for_unfused, + k_indexer, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed), + indexer_loss_coeff, + causal_mask, + self.config.dsa_indexer_use_sparse_loss, + self.indexer.pg_collection, + None, + None, + None, + None, + self.config.calculate_per_token_loss, + True, + non_compressed_lse, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + else: + _, topk_indices_compressed = self.indexer(x_det, qr_det, mask=causal_mask) + + n_valid_per_pos = positions // self.compress_ratio # [sq, 1] + valid = (topk_indices_compressed >= 0) & (topk_indices_compressed < n_valid_per_pos) + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + nvtx_range_pop("compressed_indices") + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + # --- Step 5: Sparse attention --- + nvtx_range_push("sparse_attn_kernel") + output = unfused_compressed_sparse_attn( + query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + + # --- Step 6: Attach indexer loss --- + if indexer_loss is not None and self.training and torch.is_grad_enabled(): + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + nvtx_range_pop("compressed_sparse_attn") + return output 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 new file mode 100644 index 00000000000..df16ba09086 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -0,0 +1,696 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +from dataclasses import dataclass +from typing import NoReturn, Optional, Union + +import torch + +from megatron.core import tensor_parallel +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.attention import Attention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.torch_norm import LayerNormBuilder +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.typed_torch import apply_module +from megatron.core.utils import get_pg_size, is_te_min_version + +try: + from megatron.core.fusions.fused_mla_yarn_rope_apply import ( + fused_mla_rope_inplace, + fused_mla_rope_out_of_place, + ) +except Exception: + fused_mla_rope_inplace = None + fused_mla_rope_out_of_place = None + + +if HAVE_TE: + from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input +else: + TEColumnParallelLinear, TELinear, set_save_original_input = (None, None, None) + + +@torch.compile +def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Fused RMS normalization for query tensor (no learnable weight).""" + return q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + + +@dataclass +class DSv4HybridSelfAttentionSubmodules: + """Submodules for the DSv4HybridAttention layer.""" + + q_layernorm: LayerNormBuilder + kv_layernorm: LayerNormBuilder + + linear_q_down_proj: Union[ModuleSpec, type] = None + linear_q_up_proj: Union[ModuleSpec, type] = None + linear_kv_proj: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + + +class DSv4HybridAttention(Attention): + """DeepSeek-v4 Hybrid Attention layer.""" + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, + name: str | None = None, + ) -> None: + + if pg_collection is None: + raise ValueError("DSv4 hybrid attention requires an explicit ProcessGroupCollection.") + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attention_type=attention_type, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + name=name, + ) + self.config: MLATransformerConfig + + assert ( + get_pg_size(self.pg_collection.tp) == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + + assert ( + not self.checkpoint_core_attention + ), "Checkpoint core attention is not supported in DSv4 Hybrid Attention." + assert ( + not self.offload_qkv_linear + ), "Offload qkv linear is not supported in DSv4 Hybrid Attention." + + self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + + self.q_head_dim = self.config.v_head_dim + + self.key_hidden_size = self.q_head_dim + self.val_hidden_size = self.config.v_head_dim + + self.recompute_up_proj = ( + self.config.recompute_granularity == 'selective' + and "mla_up_proj" in self.config.recompute_modules + ) + self.qkv_up_checkpoint = None + + self.softmax_scale = None + + ratio_idx = self.config.num_layers + layer_number - 1 if is_mtp_layer else layer_number - 1 + if compress_ratio is None: + compress_ratio = self.config.csa_compress_ratios[ratio_idx] + use_compressed_yarn = compress_ratio > 1 + rope_base = ( + self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base + ) + self._dsv4_compress_ratio = compress_ratio + self._dsv4_rope_base = rope_base + self._dsv4_uses_yarn_rope = use_compressed_yarn + if not use_compressed_yarn: + self.rotary_pos_emb = RotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_percent=self.config.rotary_percent, + rotary_base=rope_base, + cp_group=self.pg_collection.cp, + ) + else: + self.rotary_pos_emb = YarnRotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_base=rope_base, + scaling_factor=self.config.rotary_scaling_factor, + original_max_position_embeddings=self.config.original_max_position_embeddings, + beta_fast=self.config.beta_fast, + beta_slow=self.config.beta_slow, + mscale=self.config.mscale, + mscale_all_dim=self.config.mscale_all_dim, + cp_group=self.pg_collection.cp, + ) + + core_attn_extra_kwargs = { + "rotary_pos_emb": self.rotary_pos_emb, + "compress_ratio": compress_ratio, + "is_mtp_layer": is_mtp_layer, + "name": (name + ".core_attention") if name is not None else None, + } + self.core_attention = build_module( + submodules.core_attention, + config=self.config, + layer_number=self.layer_number, + attn_mask_type=self.attn_mask_type, + attention_type=self.attention_type, + softmax_scale=self.softmax_scale, + k_channels=self.q_head_dim, + v_channels=self.config.v_head_dim, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + **core_attn_extra_kwargs, + ) + + # Output. + self.o_local_groups = self.config.o_groups + assert ( + self.query_projection_size % self.config.o_groups == 0 + ), "num_attention_heads * v_head_dim must be divisible by o_groups" + group_proj_in_size = self.query_projection_size // self.config.o_groups + group_proj_out_size = self.config.o_groups * self.config.o_lora_rank + + _linear_o_group_proj = torch.empty( + group_proj_out_size, + group_proj_in_size, + device=torch.cuda.current_device(), + dtype=self.config.params_dtype, + ) + self.config.init_method(_linear_o_group_proj) + self.linear_o_group_proj = torch.nn.Parameter(_linear_o_group_proj) + + linear_proj_in_size = self.config.o_groups * self.config.o_lora_rank + + self.linear_proj = build_module( + submodules.linear_proj, + linear_proj_in_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name='proj', + tp_group=self.pg_collection.tp, + ) + + if ( + HAVE_TE + and isinstance(self.linear_proj, TELinear) + and ( + ( + self.config.fp8 + and self.config.fp8_recipe != 'delayed' + and is_te_min_version("2.6.0dev0") + ) + or (self.config.fp4 and is_te_min_version("2.7.0.dev0")) + ) + ): + # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and + # linear_proj also saves the quantized tensor of this output. Here we set the + # linear_proj to save the original input tensors to avoid the extra memory usage of + # the quantized tensor. + set_save_original_input(self.linear_proj) + + def forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + position_ids=None, + sequence_len_offset=None, + *, + inference_params=None, + ): + """Forward pass for DeepSeek-v4 Hybrid Attention""" + assert ( + rotary_pos_emb is None + ), "Rotary position embeddings should not be passed into DSv4HybridAttention." + assert ( + attention_bias is None + ), "Attention bias should not be passed into DSv4HybridAttention." + assert ( + rotary_pos_cos is None and rotary_pos_sin is None + ), "DSv4HybridAttention does not support Flash Decoding" + assert ( + not rotary_pos_cos_sin + ), "Flash-infer rope has not been tested with DSv4HybridAttention." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridAttention." + assert ( + packed_seq_params is None + ), "Packed sequence is not supported for DSv4HybridAttention." + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention - + # self or cross attn. + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + hidden_states, key_value_states, position_ids, None, inference_context=inference_context + ) + + # TODO: Currently, TE can only accept contiguous tensors for MLA + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + # ================================== + # core attention computation + # ================================== + # Need corresponding TE change + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=None, + x=hidden_states, + qr=q_compressed, + ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) + + if self.recompute_up_proj: + assert self.qkv_up_checkpoint is not None + self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) + self.qkv_up_checkpoint = None + + # inverse RoPE on last qk_pos_emb_head_dim of each head + seq_len = core_attn_out.size(0) + n_heads = self.num_attention_heads_per_partition + pos_dim = self.config.qk_pos_emb_head_dim + nope_dim = self.config.v_head_dim - pos_dim + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), n_heads, -1) + rope_seqlen = seq_len + # 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 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rope_seqlen, dtype=hidden_states.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rope_seqlen, packed_seq=False) + else: + rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=False) + if self.config.apply_rope_fusion: + # Fused DSA backward retains the raw attention output O. Applying + # inverse RoPE to its view in-place corrupts the retained O used by + # the softmax backward, so this call needs private storage. + assert fused_mla_rope_out_of_place is not None + core_attn_out = fused_mla_rope_out_of_place( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + inverse=True, + remove_interleaving=True, + ) + else: + content_part, rot_part = torch.split( + core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 + ) + rot_part = apply_rotary_pos_emb( + rot_part, + rotary_pos_emb, + self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + inverse=True, + mla_output_remove_interleaving=True, + ) + core_attn_out = torch.cat([content_part, rot_part], dim=-1) + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) + + # Grouped output + core_attn_out = core_attn_out.view( + core_attn_out.size(0), core_attn_out.size(1), self.o_local_groups, -1 + ) + wo_a_weight = self.linear_o_group_proj.view( + self.o_local_groups, self.config.o_lora_rank, -1 + ) + core_attn_out = torch.einsum("...gd,grd->...gr", core_attn_out, wo_a_weight) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + + # ================= + # Output. [sq, b, h] + # ================= + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: + output, bias = self.linear_proj(core_attn_out) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + + return output, bias + + +class DSv4HybridSelfAttention(DSv4HybridAttention): + """DSv4Hybrid Self-attention layer class + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type=AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, + name: str | None = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type="self", + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + compress_ratio=compress_ratio, + name=name, + ) + + q_down_proj_kwargs = {} + if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' + else: + raise ValueError(f"Unsupported linear_q_down_proj: {submodules.linear_q_down_proj}") + + self.linear_q_down_proj = build_module( + submodules.linear_q_down_proj, + self.config.hidden_size, + self.config.q_lora_rank, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_down_proj', + skip_weight_param_allocation=False, + tp_group=None, + name=(name + ".linear_q_down_proj") if name is not None else None, + **q_down_proj_kwargs, + ) + + self.linear_q_up_proj = build_module( + submodules.linear_q_up_proj, + self.config.q_lora_rank, + self.config.num_attention_heads * self.q_head_dim, + 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='q_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_q_up_proj") if name is not None else None, + ) + + self.linear_kv_proj = build_module( + submodules.linear_kv_proj, + self.config.hidden_size, + self.config.v_head_dim, + 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='kv_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_kv_proj") if name is not None else None, + ) + self.kv_layernorm = submodules.kv_layernorm( + hidden_size=self.config.v_head_dim, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + self.q_layernorm = submodules.q_layernorm( + hidden_size=self.config.q_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + def get_query_key_value_tensors( + self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_context=None, + *, + inference_params=None, + ): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + # s = sequence length, b = batch size, h = hidden size, n = num attention heads + # Attention heads [s, b, n*h] + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" + assert ( + packed_seq_params is None + ), "Packed sequence is not supported for DSv4HybridAttention." + + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridSelfAttention." + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, None, hidden_states, self.config, None + ) + + # rotary_pos_emb:[s, b, 1, 64] + # 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 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + else: + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + + # ========================================= + # QKV down projection and layernorm + # ========================================= + # q_compressed: [s, b, q_lora_rank] + q_compressed, _ = self.linear_q_down_proj(hidden_states) + + kv_compressed = hidden_states + k_pos_emb = None + + # ========================================= + # Apply norm + # ========================================= + + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + q_compressed = apply_module(self.q_layernorm)(q_compressed) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + + def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + """Apply the up projection and RoPE to the SBHD query and key.""" + # q_compressed: [s, b, q_lora_rank] + # q: [s, b, n * (qk_head_dim + qk_pos_emb_head_dim)] + q, _ = self.linear_q_up_proj(q_compressed) + + # q: [num_tokens, n, q_head_dim] + q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) + q = _q_rms_norm(q, self.config.layernorm_epsilon) + + kv, _ = self.linear_kv_proj(kv_compressed) + kv = self.kv_layernorm(kv) + + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + if k_pos_emb is not None: + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + + if self.config.apply_rope_fusion: + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() + query = fused_mla_rope_inplace( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + None, + cp_rank, + cp_size, + remove_interleaving=True, + ) + kv = kv.unsqueeze(-2) + kv = fused_mla_rope_inplace( + kv, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + None, + cp_rank, + cp_size, + remove_interleaving=True, + ) + key = kv + value = kv + else: + q_len = q.size()[0] + # Keep direct forward calls with shorter sequences aligned to their inputs. + rotary_pos_emb = rotary_pos_emb[0:q_len] + + # q_no_pe: [num_tokens, n, qk_head_dim] + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_no_pe, q_pos_emb = torch.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + # RoPE and query (shared for wkv and latent) + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_pos_emb = apply_rotary_pos_emb( + q_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + # query: [num_tokens, n, (qk_head_dim + v_head_dim)] + query = torch.cat([q_no_pe, q_pos_emb], dim=-1) + + pos_dim = self.config.qk_pos_emb_head_dim + kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) + + # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = apply_rotary_pos_emb( + k_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + + # Single head: key = value = [num_tokens, 1, v_head_dim] + kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) + key = kv + value = kv + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + return query, key, value + + if self.recompute_up_proj: + quantization = self.config.fp8 or self.config.fp4 + self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) + query, key, value = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + else: + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + + return query, key, value, q_compressed, kv_compressed + + def backward_dw(self) -> NoReturn: + """Execute weight gradient computation""" + self._backward_kv_proj() + self._backward_q_proj() + self.core_attention.backward_dw() + self._backward_output_proj() + + def _backward_kv_proj(self): + """Computes weight gradients of KV projection layers""" + self.linear_kv_proj.backward_dw() + + def _backward_q_proj(self): + """Computes weight gradients of Q projection layers""" + self.linear_q_down_proj.backward_dw() + self.linear_q_up_proj.backward_dw() + + def _backward_output_proj(self): + """Computes weight gradients of output projection layer""" + self.linear_proj.backward_dw() + + def set_for_recompute_input_layernorm(self): + """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" + set_save_original_input(self.linear_q_down_proj) + set_save_original_input(self.linear_kv_proj) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index bfead2a25c5..fe8a1456208 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -7,7 +7,6 @@ import torch -from megatron.core import parallel_state from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -296,32 +295,72 @@ def save_loss_to_tracker( return tracker = DSAIndexerLossLoggingHelper.tracker + # Hybrid MTP layer numbers can exceed ``num_layers + mtp_num_layers`` + # because every prediction depth can contain multiple hybrid layers. + needed = max(num_layers, layer_number) if "values" not in tracker: - tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["values"] = torch.zeros(needed, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < needed: + grown = torch.zeros( + needed, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown tracker["values"][layer_number - 1] += loss.detach() tracker["reduce_group"] = reduce_group tracker["avg_group"] = avg_group @staticmethod - def clean_loss_in_tracker(): + def clean_loss_in_tracker(preserve_groups: bool = False): """Clear the indexer losses.""" tracker = DSAIndexerLossLoggingHelper.tracker + reduce_group = tracker.get("reduce_group") if preserve_groups else None + avg_group = tracker.get("avg_group") if preserve_groups else None if "values" in tracker: tracker["values"].zero_() - tracker["reduce_group"] = None - tracker["avg_group"] = None + tracker["reduce_group"] = reduce_group + tracker["avg_group"] = avg_group @staticmethod - def reduce_loss_in_tracker(): - """Collect and reduce the indexer losses across ranks.""" + def reduce_loss_in_tracker( + pg_collection: ProcessGroupCollection, num_layers: Optional[int] = None + ): + """Collect and reduce indexer losses across every pipeline rank. + + Args: + pg_collection: Process groups used for pipeline and data-parallel reductions. + num_layers: Total number of decoder and MTP layers. When provided, ranks without + local indexer losses contribute zeros to the pipeline-wide reduction. + """ tracker = DSAIndexerLossLoggingHelper.tracker - if "values" not in tracker: + pp_group = pg_collection.pp + + # Pipeline ranks can own different attention variants, so first agree on + # a common tracker size. Cache the result because layer allocation is + # static and the negotiation requires a device-to-host synchronization. + if tracker.get("agreed_size") is not None: + size = tracker["agreed_size"] + else: + local_size = tracker["values"].shape[0] if "values" in tracker else (num_layers or 0) + size_t = torch.tensor( + [local_size], device=torch.cuda.current_device(), dtype=torch.long + ) + torch.distributed.all_reduce(size_t, op=torch.distributed.ReduceOp.MAX, group=pp_group) + size = int(size_t.item()) + tracker["agreed_size"] = size + if size == 0: return + if "values" not in tracker: + tracker["values"] = torch.zeros(size, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < size: + grown = torch.zeros( + size, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown values = tracker["values"] - torch.distributed.all_reduce( - values, group=parallel_state.get_pipeline_model_parallel_group() - ) + torch.distributed.all_reduce(values, group=pp_group) # Reduce indexer losses across ranks. if tracker.get('reduce_group') is not None: torch.distributed.all_reduce(values, group=tracker.get('reduce_group')) @@ -330,9 +369,7 @@ def reduce_loss_in_tracker(): values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG ) torch.distributed.all_reduce( - values, - group=parallel_state.get_data_parallel_group(with_context_parallel=False), - op=torch.distributed.ReduceOp.AVG, + values, group=pg_collection.dp, op=torch.distributed.ReduceOp.AVG ) @staticmethod @@ -340,9 +377,13 @@ def track_indexer_metrics( loss_scale: float, iteration: int, writer, + pg_collection: ProcessGroupCollection, wandb_writer=None, total_loss_dict=None, per_layer_logging: bool = False, + num_layers: Optional[int] = None, + num_indexer_layers: Optional[int] = None, + preserve_groups: bool = False, ): """Track the sparse attention indexer metrics for logging. @@ -350,20 +391,27 @@ def track_indexer_metrics( loss_scale: Scale factor for the loss. iteration: Current training iteration. writer: TensorBoard writer. + pg_collection: Process groups used for pipeline and data-parallel reductions. 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 and MTP layers. Passing it makes ranks + without a local indexer participate in the pipeline reduction. + num_indexer_layers: Number of layers that own an indexer. Defaults to the + tracker size when every tracked layer owns one. + preserve_groups: Keep the saved reduction groups for CUDA Graph replays. """ - DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() + DSAIndexerLossLoggingHelper.reduce_loss_in_tracker( + pg_collection=pg_collection, 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 num_indexer_layers is None: + num_indexer_layers = indexer_loss_values.shape[0] + avg_indexer_loss = indexer_loss_values.sum() / max(num_indexer_layers, 1) # Log average loss if total_loss_dict is not None: @@ -378,7 +426,7 @@ def track_indexer_metrics( if wandb_writer is not None: wandb_writer.log({"indexer loss": avg_indexer_loss}, iteration) - DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + DSAIndexerLossLoggingHelper.clean_loss_in_tracker(preserve_groups=preserve_groups) def compute_dsa_indexer_loss( @@ -396,6 +444,7 @@ def compute_dsa_indexer_loss( key_positions: Optional[torch.Tensor] = None, query_valid_rows: Optional[torch.Tensor] = None, calculate_per_token_loss: bool = False, + non_compressed_lse: torch.Tensor | None = None, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -421,6 +470,10 @@ def compute_dsa_indexer_loss( varlen_starts: Optional row-wise key start bounds [sq] for packed THD. varlen_ends: Optional row-wise key end bounds [sq] for packed THD. key_positions: Optional global key positions [sk] for packed THD. + non_compressed_lse: Optional detached FP32 log-sum-exp contribution + [batch, heads, seqlen_q] from teacher keys that are intentionally + omitted from ``key``. When provided, the selected ``key`` logits + are normalized with this external mass before heads are summed. Returns: index_loss: KL divergence loss (scalar). @@ -489,8 +542,8 @@ def compute_dsa_indexer_loss( attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask # [b, np, sq, sk] -> [b, np, sq, sk] - attention_scores = dsa_masking.masked_softmax( - attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1 + attention_scores = _compute_indexer_teacher_probabilities( + attention_scores, attention_valid_mask, non_compressed_lse=non_compressed_lse ) # [b, sq, sk] -> [b, sq, sk] index_log_scores = dsa_masking.masked_log_softmax( @@ -504,7 +557,7 @@ def compute_dsa_indexer_loss( # attention scores are scattered to TP ranks in head dimension. torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) # The target is already non-negative because it is a sum of softmax probabilities. - attention_scores = dsa_indexer_loss.normalize_indexer_target(attention_scores) + attention_scores = _normalize_indexer_teacher_target(attention_scores, non_compressed_lse) return dsa_indexer_loss.indexer_loss_from_target( attention_scores, index_log_scores, @@ -514,6 +567,52 @@ def compute_dsa_indexer_loss( ) +def _compute_indexer_teacher_probabilities( + attention_scores: torch.Tensor, + attention_valid_mask: torch.Tensor, + non_compressed_lse: torch.Tensor | None = None, +) -> torch.Tensor: + """Normalize selected teacher logits, optionally with omitted attention mass. + + ``non_compressed_lse`` is a sufficient statistic for teacher logits that + must participate in the softmax denominator but must not appear in the + compressed-key target returned by this helper. + """ + b, np, sq, sk = attention_scores.shape + expanded_valid_mask = attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk) + if non_compressed_lse is None: + return dsa_masking.masked_softmax(attention_scores.float(), expanded_valid_mask, dim=-1) + + expected_shape = (b, np, sq) + if tuple(non_compressed_lse.shape) != expected_shape: + raise ValueError( + "non_compressed_lse must have shape [batch, heads, seqlen_q], " + f"got {tuple(non_compressed_lse.shape)}, expected {expected_shape}" + ) + if non_compressed_lse.device != attention_scores.device: + raise ValueError( + "non_compressed_lse and attention_scores must be on the same device, " + f"got {non_compressed_lse.device} and {attention_scores.device}" + ) + if non_compressed_lse.requires_grad: + raise ValueError("non_compressed_lse must be detached") + + masked_scores = attention_scores.float().masked_fill(~expanded_valid_mask, float("-inf")) + compressed_lse = torch.logsumexp(masked_scores, dim=-1) + full_lse = torch.logaddexp(non_compressed_lse.float(), compressed_lse) + probabilities = torch.exp(masked_scores - full_lse.unsqueeze(-1)) + return torch.where(expanded_valid_mask, probabilities, torch.zeros_like(probabilities)) + + +def _normalize_indexer_teacher_target( + target: torch.Tensor, non_compressed_lse: torch.Tensor | None +) -> torch.Tensor: + """L1-normalize teacher mass without changing the legacy DSA path.""" + if non_compressed_lse is None: + return dsa_indexer_loss.normalize_indexer_target(target) + return target / target.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) + + def _compute_index_scores( q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor, use_relu: bool = True ) -> torch.Tensor: @@ -627,6 +726,7 @@ def fwd_fused_indexer_loss_naive( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """Naive implementation of forward pass for indexer loss.""" index_scores, topk_indices = fused_qk_topk_naive( @@ -656,6 +756,7 @@ def fwd_fused_indexer_loss_naive( key_positions=key_positions, query_valid_rows=query_valid_rows, calculate_per_token_loss=calculate_per_token_loss, + non_compressed_lse=non_compressed_lse, ) return topk_indices, indexer_loss @@ -680,6 +781,7 @@ def bwd_fused_indexer_loss_naive( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """Naive implementation of backward pass for indexer loss.""" query, _ = dsa_layout.ensure_sbhd(query, "query") @@ -752,8 +854,8 @@ def bwd_fused_indexer_loss_naive( else: index_valid_mask = base_valid_mask attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask - attention_scores_softmax = dsa_masking.masked_softmax( - attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1 + attention_scores_softmax = _compute_indexer_teacher_probabilities( + attention_scores, attention_valid_mask, non_compressed_lse=non_compressed_lse ) # Free attention_scores immediately del attention_scores @@ -776,7 +878,9 @@ def bwd_fused_indexer_loss_naive( # L1 normalize. Fully masked packed/varlen rows can have zero summed # attention mass; clamp the denominator so those rows stay finite and are # later zeroed by the row-valid loss mask. - attention_scores_normalized = dsa_indexer_loss.normalize_indexer_target(attention_scores_sum) + attention_scores_normalized = _normalize_indexer_teacher_target( + attention_scores_sum, non_compressed_lse + ) # Free attention_scores_sum - no longer needed after normalization del attention_scores_sum @@ -890,6 +994,7 @@ def bwd_fused_indexer_loss_naive( "query_valid_rows", "calculate_per_token_loss", "use_relu", + "non_compressed_lse", ) @@ -916,6 +1021,7 @@ def forward( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """ Fused forward: index_scores never materialized in full. @@ -938,10 +1044,17 @@ def forward( query_valid_rows=query_valid_rows, calculate_per_token_loss=calculate_per_token_loss, use_relu=use_relu, + non_compressed_lse=non_compressed_lse, ) # Save for backward (recomputation strategy) - ctx.save_for_backward(q, weights, k, query, key, topk_indices) + saved_non_compressed_lse = ( + non_compressed_lse + if non_compressed_lse is not None + else q.new_empty(0, dtype=torch.float32) + ) + ctx.save_for_backward(q, weights, k, query, key, topk_indices, saved_non_compressed_lse) + ctx.has_non_compressed_lse = non_compressed_lse is not None ctx.softmax_scale = softmax_scale ctx.loss_coeff = loss_coeff ctx.sparse_loss = sparse_loss @@ -953,6 +1066,7 @@ def forward( ctx.query_valid_rows = query_valid_rows ctx.calculate_per_token_loss = calculate_per_token_loss ctx.use_relu = use_relu + ctx.num_inputs = len(ctx.needs_input_grad) return topk_indices, loss @@ -961,7 +1075,8 @@ def backward(ctx, grad_topk_indices, grad_loss): """ Backward: Recompute what we need. """ - q, weights, k, query, key, topk_indices = ctx.saved_tensors + q, weights, k, query, key, topk_indices, saved_non_compressed_lse = ctx.saved_tensors + non_compressed_lse = saved_non_compressed_lse if ctx.has_non_compressed_lse else None grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( q, @@ -982,6 +1097,7 @@ def backward(ctx, grad_topk_indices, grad_loss): query_valid_rows=ctx.query_valid_rows, calculate_per_token_loss=ctx.calculate_per_token_loss, use_relu=ctx.use_relu, + non_compressed_lse=non_compressed_lse, ) grad_by_name = { @@ -991,8 +1107,10 @@ def backward(ctx, grad_topk_indices, grad_loss): # query and key are detached in forward, so return None for their gradients. "query": None, "key": None, + "non_compressed_lse": None, } - return tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) + gradients = tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) + return gradients[: ctx.num_inputs] class DSAIndexerLossAutoScaler(torch.autograd.Function): diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py new file mode 100644 index 00000000000..21932c74a1a --- /dev/null +++ b/megatron/core/transformer/hyper_connection.py @@ -0,0 +1,845 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import math +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_decorator + +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: + 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 + + +class SinkhornKnopp(torch.autograd.Function): + """Sinkhorn-Knopp projection to doubly stochastic matrix. + + This is an autograd.Function because the iterative forward is re-executed + during backward (under torch.enable_grad) so that PyTorch's autograd can + differentiate through it without storing all intermediate iteration states. + """ + + @staticmethod + def forward(ctx, input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Run Sinkhorn iterations and save inputs for backward recomputation.""" + M = _sinkhorn_iterations(input_logits, num_iterations, eps) + ctx.save_for_backward(input_logits) + ctx.num_iterations = num_iterations + ctx.eps = eps + return M + + @staticmethod + def backward(ctx, grad_output: Tensor): + """Recompute forward under enable_grad and back-propagate.""" + (input_logits,) = ctx.saved_tensors + with torch.enable_grad(): + logits = input_logits.detach().requires_grad_(True) + M = _sinkhorn_iterations(logits, ctx.num_iterations, ctx.eps) + M.backward(grad_output) + return logits.grad, None, None + + +def native_sinkhorn(input_logits: Tensor, num_iterations: int, eps: float = 1e-6) -> Tensor: + """Native Sinkhorn-Knopp (autograd.Function wrapper).""" + return SinkhornKnopp.apply(input_logits, num_iterations, eps) + + +@torch.compile +def native_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + """Native n-stream weighted aggregation: out = sum_j(h_pre_j * x_j).""" + return (x * h_pre.unsqueeze(-1)).sum(dim=2) + + +@torch.compile +def native_h_post_bda( + h_res: Tensor, original_residual: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + """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.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) + return x_expanded + bias_expanded + mixed + return x_expanded + mixed + + +@torch.compile +def native_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6) -> Tuple[Tensor, Tensor]: + """Native fused projection + RMS normalization.""" + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + v = norm / math.sqrt(K) + eps + r = 1.0 / v + 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 +) -> Tensor: + """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) + eps + y = torch.sum(pre.unsqueeze(-1) * hidden_states.view(*hidden_states.shape[:-1], n, -1), dim=-2) + return y.to(dtype) + + +# ============================================================================ +# HyperConnectionModule +# ============================================================================ + + +# TODO: keep hyper connection in fp32 computation +class HyperConnectionModule(MegatronModule): + """ + Unified mHC (Manifold-Constrained Hyper-Connections) module. + + Implements the complete mHC propagation: + 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^T @ x + expanded_output + 5. Block-level expand/contract for TransformerBlock boundaries + + Args: + config: TransformerConfig with hyper-connection fields + layer_number: Current layer index for initialization + """ + + def __init__(self, config: TransformerConfig, layer_number: int): + super().__init__(config) + self.config = config + self.layer_number = layer_number + 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 + # - H_pre: n values + # - H_post: n values + # - H_res: n^2 values (before Sinkhorn projection) + self.mapping_proj = nn.Linear( + self.n * self.hidden_size, self.n * self.n + 2 * self.n, bias=False + ) + + init_alpha = config.mhc_init_gating_factor + # Learnable scaling factors (Eq. 5 in paper) + self.alpha_pre = nn.Parameter(torch.full((1,), init_alpha)) + self.alpha_post = nn.Parameter(torch.full((1,), init_alpha)) + self.alpha_res = nn.Parameter(torch.full((1,), init_alpha)) + + # Static bias terms + self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n)) + mark_keep_in_fp32(self.mapping_proj.weight) + mark_keep_in_fp32(self.alpha_pre) + mark_keep_in_fp32(self.alpha_post) + mark_keep_in_fp32(self.alpha_res) + mark_keep_in_fp32(self.bias) + self.norm_eps = 1e-6 + + # 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 + + # The fused path computes the projection and compute_h in one op, so + # _projection_and_get_norm — and therefore _proj_rms_op — is only ever + # reached on the unfused path. + self._proj_rms_op = native_proj_rms + + if config.use_fused_mhc: + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + 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_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_compute_h_op = None + + self._init_weights() + + def _init_weights(self) -> None: + """Initialize weights for stable training.""" + nn.init.xavier_uniform_(self.mapping_proj.weight) + + # Set sequence_parallel attribute on parameters for gradient synchronization + # across TP ranks when sequence_parallel is enabled. + # This is required because HyperConnectionModule uses non-TP-aware layers + # (nn.Linear, nn.RMSNorm) whose gradients need to be all-reduced. + if self.config.sequence_parallel: + setattr(self.mapping_proj.weight, 'sequence_parallel', True) + setattr(self.alpha_pre, 'sequence_parallel', True) + setattr(self.alpha_post, 'sequence_parallel', True) + setattr(self.alpha_res, 'sequence_parallel', True) + setattr(self.bias, 'sequence_parallel', True) + + def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]: + """ + Projection + RMS normalization. + + Args: + x: [s, b, n*C] - n-stream hidden states + """ + s, b, nC = x.shape + # The mHC mapping computation runs in FP32: the parameters are kept in + # FP32 and the activations are upcast here, then compute_mappings casts + # the bounded mixing weights back to the activation dtype. + x_2d = x.reshape(s * b, nC).to(torch.float32) + weight = self.mapping_proj.weight.to(torch.float32) + proj, r = self._proj_rms_op(x_2d, weight, self.norm_eps) + return proj.view(s, b, -1), r.view(s, b, 1) + + @torch.compile + def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Compute h from projected hidden states and scaling factors. + + Args: + proj: [s, b, n^2 + 2n] - projected hidden states + r: [s, b, 1] - scaling factors + + Returns: + h_pre: [s, b, n] - aggregation weights + h_post: [s, b, n] - expansion weights + h_res: [s, b, n^2] - residual mixing logits + """ + alpha_ = torch.cat( + [ + self.alpha_pre.expand(self.n), + self.alpha_post.expand(self.n), + self.alpha_res.expand(self.n * self.n), + ], + dim=-1, + ) + + h = r * proj * alpha_ + self.bias + # H_pre = σ(α_pre * (θ_pre @ x̃) + b_pre) + 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 + h_res = h[..., 2 * self.n :] + return h_pre, h_post, h_res + + @nvtx_decorator(message="HyperConnection::compute_mappings") + def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Compute mHC mappings from input hidden states. + + Reference: Eq. (5) and (8) in mHC paper + + Args: + x: [s, b, n*C] - n-stream hidden states + + Returns: + h_pre: [s, b, n] - aggregation weights (sigmoid activated) + h_post: [s, b, n] - expansion weights (2*sigmoid activated) + h_res: [s, b, n, n] - residual mixing matrix (doubly stochastic) + """ + s, b, _ = x.shape + + 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, self.sinkhorn_iterations, self.sinkhorn_eps + ) # [s, b, n, n] + + # The mixing weights are bounded (sigmoid outputs / doubly stochastic + # matrix), so after the FP32 computation they are safe to apply to the + # streams in the activation dtype. + dtype = x.dtype + return h_pre.to(dtype), h_post.to(dtype), h_res.to(dtype) + + @torch.compile + def _apply_h_post(self, x: Tensor, h_post: Tensor) -> Tensor: + """ + Core implementation of H_post application to a single tensor. + + Computes: H_post^T @ x + + Args: + x: Input tensor, can be either: + - [s, b, C] - standard hidden states + - [C] - bias tensor (will be broadcast) + h_post: [s, b, n] - expansion weights + + Returns: + output: [s, b, n*C] - expanded tensor + """ + n = self.n + s, b, _ = h_post.shape + + if x.dim() == 1: + # x is bias with shape [C], need to broadcast to [s, b, 1, C] + C = x.shape[0] + x_expanded = x.unsqueeze(0).unsqueeze(0).unsqueeze(0).expand(s, b, 1, C) + else: + # x is [s, b, C] + C = x.shape[-1] + x_expanded = x.unsqueeze(2) # [s, b, 1, C] + + # h_post^T @ x : [s, b, n, 1] * [s, b, 1, C] -> [s, b, n, C] + # Using broadcast multiply instead of einsum + result = h_post.unsqueeze(-1) * x_expanded + return result.view(s, b, n * C) + + @nvtx_decorator(message="HyperConnection::apply_h_post") + def apply_h_post( + self, + x_with_bias: Tuple[Tensor, Optional[Tensor]], + h_post: Tensor, + manager: Optional['CheckpointManager'] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Apply H_post to x and optionally bias, with optional checkpointing. + + This is the unified entry point that handles both normal execution + and checkpoint-based execution for memory efficiency. + + Args: + x_with_bias: Tuple of (x, bias) where: + - x: [s, b, C] - hidden states + - bias: [C] or None - optional bias tensor + h_post: [s, b, n] - expansion weights + manager: Optional CheckpointManager for checkpoint management. + When provided, wraps _apply_h_post with CheckpointWithoutOutput. + + Returns: + Tuple of (x_out, bias_out) where: + - x_out: [s, b, n*C] - expanded hidden states + - bias_out: [s, b, n*C] or None - expanded bias if input bias was not None + """ + x, bias = x_with_bias + + if manager is not None: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + # Checkpoint _apply_h_post to discard the output + x_out = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self._apply_h_post, x, h_post + ) + + # Checkpoint _apply_h_post for bias if not None + if bias is not None: + bias_out = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + self._apply_h_post, bias, h_post + ) + else: + bias_out = None + else: + # Normal execution without checkpoint + x_out = self._apply_h_post(x, h_post) + bias_out = self._apply_h_post(bias, h_post) if bias is not None else None + + return x_out, bias_out + + def aggregate(self, x: Tensor, h_pre: Tensor) -> Tensor: + """ + Aggregate n-stream to 1-stream. + + Args: + x: [s, b, n*C] - n-stream hidden states + h_pre: [s, b, n] - aggregation weights + + Returns: + aggregated: [s, b, C] - single stream hidden states + """ + s, b, _ = x.shape + C = self.hidden_size + x_streams = x.view(s, b, self.n, C) + return self._h_aggregate_op(x_streams, h_pre) + + @torch.compile + def apply_h_res(self, h_res: Tensor, residual: Tensor) -> Tensor: + """ + Apply H_res to residual using H_res weights. + + Computes: H_res.T @ residual + + Args: + h_res: [s, b, n, n] - residual mixing matrix + residual: [s, b, n*C] - n-stream hidden states + """ + s, b, _ = residual.shape + n = self.n + C = self.hidden_size + + # Reshape for bmm: [s, b, n, n] -> [s*b, n, n] + h_res_batched = h_res.view(s * b, n, n) + # [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].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, + return_residual: bool = False, + ) -> Tuple[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: + The compatible 3-tuple ``(aggregated, h_res, h_post)`` by default. + HybridModel callers set ``return_residual=True`` to also receive 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: + result = self._forward_with_checkpoint(hidden_states, mhc_recompute_manager) + else: + result = self._forward_normal(hidden_states) + return result if return_residual else result[:3] + + def _forward_normal(self, hidden_states: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """ + Normal forward pass without checkpointing. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + + Returns: + 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(hs_for_mappings) + + # Aggregate for layer input + with torch.cuda.nvtx.range("HyperConnection::aggregate"): + aggregated = self.aggregate(hs_for_aggregate, h_pre) + + return aggregated, h_res, h_post, hs_for_residual + + def _forward_with_checkpoint( + self, hidden_states: Tensor, manager: 'CheckpointManager' + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """ + Forward pass with checkpointing for memory efficiency. + + compute_mappings is called directly (not checkpointed) since its outputs + (h_pre, h_post, h_res) are needed downstream. Only aggregate is wrapped with + CheckpointWithoutOutput and auto-registered to the manager. + apply_h_res is deferred to fused_h_res_h_post_bda for kernel fusion. + + Args: + hidden_states: [s, b, n*C] - n-stream hidden states + manager: CheckpointManager for unified recomputation + + Returns: + 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 + + # 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, hs_for_aggregate, h_pre + ) + + return aggregated, h_res, h_post, hs_for_residual + + # ==================== Block-level utilities ==================== + + @staticmethod + def input_expand(x: Tensor, n: int) -> Tensor: + """ + Expand 1-stream to n-stream at TransformerBlock entry. + + Simple replication strategy: each stream initialized as a copy of input. + + Args: + x: [s, b, C] - single stream hidden states + n: Number of residual streams + + Returns: + expanded: [s, b, n*C] - n-stream hidden states + """ + s, b, C = x.shape + # Replicate input to n streams + expanded = x.unsqueeze(2).expand(s, b, n, C).contiguous() + return expanded.view(s, b, n * C) + + @staticmethod + def output_contract(x: Tensor, n: int) -> Tensor: + """ + Contract n-stream to 1-stream at TransformerBlock exit. + + Simple averaging strategy: average all streams. + + Args: + x: [s, b, n*C] - n-stream hidden states + n: Number of residual streams + + Returns: + contracted: [s, b, C] - single stream hidden states + """ + s, b, nC = x.shape + C = nC // n + # Average all streams + x_streams = x.view(s, b, n, C) + contracted = x_streams.mean(dim=2) + return contracted + + # ==================== Fused kernel placeholder ==================== + + @nvtx_decorator(message="HyperConnection::fused_h_res_h_post_bda") + def fused_h_res_h_post_bda( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + manager: Optional['CheckpointManager'] = None, + ) -> Tensor: + """ + Fused kernel combining apply_h_res, apply_h_post and bias-dropout-add. + + This is a placeholder for future kernel fusion optimization. + Currently implements the operations sequentially using native PyTorch. + + The computation flow is: + 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) + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states (before H_res applied) + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) where: + - x: [s, b, C] - layer output (attention or MLP output) + - bias: [C] or None - optional bias tensor + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + manager: Optional CheckpointManager for checkpoint management. + When provided, each operation is wrapped with CheckpointWithoutOutput. + + Returns: + output: [s, b, n*C] - final output after all operations + """ + if manager is not None: + return self._fused_h_res_h_post_bda_with_checkpoint( + h_res, + original_residual, + h_post, + layer_output_with_bias, + dropout_prob, + training, + fused, + manager, + ) + else: + return self._fused_h_res_h_post_bda_native( + h_res, + original_residual, + h_post, + layer_output_with_bias, + dropout_prob, + training, + fused, + ) + + def _fused_h_res_h_post_bda_native( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + ) -> Tensor: + """ + h_res, h_post and bda. + + When dropout is zero (or inference), uses a single fused/reference kernel + for H_res.T @ residual + H_post * (x + bias). Falls back to unfused + implementation when dropout is needed. + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + + Returns: + output: [s, b, n*C] - final output + """ + x, bias = layer_output_with_bias + + if dropout_prob == 0.0 or not training: + s, b, _ = original_residual.shape + n = self.n + C = self.hidden_size + orig_reshaped = original_residual.view(s, b, n, C) + output = self._h_post_bda_op(h_res, orig_reshaped, h_post, x, bias) + return output.view(s, b, n * C) + + from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add + + with torch.cuda.nvtx.range("HyperConnection::apply_h_res"): + mixed = self.apply_h_res(h_res, original_residual) + with torch.cuda.nvtx.range("HyperConnection::apply_h_post"): + x_expanded = self._apply_h_post(x, h_post) + bias_expanded = self._apply_h_post(bias, h_post) if bias is not None else None + bda_func = get_bias_dropout_add(training, fused) + with torch.cuda.nvtx.range("HyperConnection::bda"): + output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) + return output + + @nvtx_decorator(message="HyperConnection::fused_h_res_h_post_bda_with_checkpoint") + def _fused_h_res_h_post_bda_with_checkpoint( + self, + h_res: Tensor, + original_residual: Tensor, + h_post: Tensor, + layer_output_with_bias: Tuple[Tensor, Optional[Tensor]], + dropout_prob: float, + training: bool, + fused: bool, + manager: 'CheckpointManager', + ) -> Tensor: + """ + Checkpointed variant of _fused_h_res_h_post_bda_native. + + Wraps compute in CheckpointWithoutOutput for activation memory savings. + Cannot reuse _native directly because checkpoint requires all args to be + positional Tensors; tuple/Optional/scalar args are unpacked or captured + via closure instead. + + Args: + h_res: [s, b, n, n] - residual mixing matrix + original_residual: [s, b, n*C] - n-stream hidden states + h_post: [s, b, n] - expansion weights + layer_output_with_bias: Tuple of (x, bias) + dropout_prob: Dropout probability + training: Whether in training mode + fused: Whether to use fused BDA implementation + manager: CheckpointManager for checkpoint management + + Returns: + output: [s, b, n*C] - final output + """ + from megatron.core.tensor_parallel.random import CheckpointWithoutOutput + + x, bias = layer_output_with_bias + n = self.n + C = self.hidden_size + + # Fast path: no dropout — use fused/reference h_post_bda kernel (same as _native) + if dropout_prob == 0.0 or not training: + + def _fused_wrapper(h_res, original_residual, h_post, x, *optional_bias): + s, b, _ = original_residual.shape + orig_reshaped = original_residual.view(s, b, n, C) + b_arg = optional_bias[0] if optional_bias else None + return self._h_post_bda_op(h_res, orig_reshaped, h_post, x, b_arg).view(s, b, n * C) + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + if bias is not None: + output = ckpt.checkpoint(_fused_wrapper, h_res, original_residual, h_post, x, bias) + else: + output = ckpt.checkpoint(_fused_wrapper, h_res, original_residual, h_post, x) + + # Slow path: dropout required — fused kernel does not support dropout, + # fall back to sequential apply_h_res + apply_h_post + bda + else: + from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add + + bda_func = get_bias_dropout_add(training, fused) + has_bias = bias is not None + + def _native_wrapper(h_res, original_residual, h_post, x, *optional_bias): + with torch.cuda.nvtx.range("HyperConnection::apply_h_res"): + mixed = self.apply_h_res(h_res, original_residual) + with torch.cuda.nvtx.range("HyperConnection::apply_h_post"): + x_expanded = self._apply_h_post(x, h_post) + if has_bias: + bias_expanded = self._apply_h_post(optional_bias[0], h_post) + else: + bias_expanded = None + with torch.cuda.nvtx.range("HyperConnection::bda"): + output = bda_func((x_expanded, bias_expanded), mixed, dropout_prob) + return output + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + if has_bias: + output = ckpt.checkpoint(_native_wrapper, h_res, original_residual, h_post, x, bias) + else: + output = ckpt.checkpoint(_native_wrapper, h_res, original_residual, h_post, x) + + return output + + +# ==================== Checkpoint utilities for mHC ==================== + + +class HyperConnectionCheckpoint: + """ + Checkpoint utility for mHC intermediate activations. + + Implements the paper's "recomputing strategy" to reduce memory footprint + by discarding intermediate n-stream activations and recomputing on-the-fly. + """ + + @staticmethod + def compute_optimal_block_size(num_layers: int, num_streams: int) -> int: + """ + Compute optimal recomputation block size. + + From paper Eq. (20): L_r^* ≈ sqrt(nL/(n+2)) + + Args: + num_layers: Total number of transformer layers + num_streams: Number of residual streams (n) + + Returns: + block_size: Optimal block size for checkpointing + """ + block_size = int(math.sqrt(num_streams * num_layers / (num_streams + 2))) + return max(1, block_size) diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 558b1b07a15..bf28600a1aa 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -433,6 +433,40 @@ def float_conversion(val): return conversion_helper(val, float_conversion) +def mark_keep_in_fp32(tensor: torch.Tensor) -> torch.Tensor: + """Mark a parameter or buffer so that ``Float16Module`` keeps it in FP32. + + Args: + tensor: The parameter or buffer to mark. + + Returns: + The same tensor, for call-site convenience. + """ + tensor.keep_in_fp32 = True + return tensor + + +def convert_module_to_dtype_except_fp32_marked( + module: torch.nn.Module, dtype: torch.dtype +) -> torch.nn.Module: + """Cast floating-point parameters and buffers except those marked to stay in FP32. + + Args: + module: The module to convert in place. + dtype: The target floating-point dtype. + + Returns: + The converted module. + """ + return module._apply( + lambda tensor: ( + tensor.to(dtype) + if tensor.is_floating_point() and not getattr(tensor, 'keep_in_fp32', False) + else tensor + ) + ) + + class Float16Module(MegatronModule): """Float 16 Module. @@ -455,13 +489,17 @@ def __init__(self, config: TransformerConfig, module: torch.nn.Module): self.pg_collection = getattr(module, 'pg_collection', None) if self.fp16: - self.add_module('module', module.half()) + self.add_module( + 'module', convert_module_to_dtype_except_fp32_marked(module, torch.half) + ) def float16_convertor(val): return val.half() elif self.bf16: - self.add_module('module', module.bfloat16()) + self.add_module( + 'module', convert_module_to_dtype_except_fp32_marked(module, torch.bfloat16) + ) def float16_convertor(val): return val.bfloat16() diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index aa21e78ce86..f42dacb9bb1 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -144,6 +144,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific @@ -156,6 +157,7 @@ def __init__( attn_mask_type=attn_mask_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) self.config: MLATransformerConfig @@ -484,6 +486,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): if pg_collection is None: @@ -498,6 +501,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -728,8 +732,16 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + rope_freqs_max_seqlen = ( + max(rope_max_seqlen_q, rope_max_seqlen_kv) + if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None + else None + ) else: cu_seqlens_q = cu_seqlens_kv = None + rope_freqs_max_seqlen = None # ========================================= # QKV down projection and layernorm @@ -941,6 +953,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -951,6 +964,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] @@ -1239,6 +1253,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, ): @@ -1254,6 +1269,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, pp_layer_offset=pp_layer_offset, name=name, ) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index b37c4c9d0f4..51893748649 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 @@ -28,7 +29,8 @@ inference_all_gather_from_tensor_model_parallel_region, ) from megatron.core.transformer.enums import AttnMaskType, LayerType -from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.hyper_connection import learned_output_contract +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_block import TransformerBlockSubmodules @@ -573,6 +575,8 @@ 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 @@ -944,6 +948,13 @@ def __init__( stacklevel=2, ) hybrid_submodules = mamba_submodules + if self.config.enable_hyper_connections and ( + mtp_layer_pattern is None or hybrid_submodules is None + ): + raise ValueError( + "Multi-token prediction with hyper connections requires the HybridModel " + "MTP contract: both mtp_layer_pattern and hybrid_submodules must be provided." + ) self.sequence_parallel = config.sequence_parallel self.submodules = submodules self.layer_number = layer_number + get_mtp_layer_offset(self.config, vp_stage) @@ -951,6 +962,7 @@ def __init__( self.cp_group = pg_collection.cp self.tp_group = pg_collection.tp if pg_collection is not None else None self.mtp_layer_pattern = mtp_layer_pattern + self.mhc_enabled = self.config.enable_hyper_connections # Validate attention mask type if using transformer-based inner layers if self.submodules.mtp_model_layer is not None and hasattr( @@ -992,25 +1004,51 @@ 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, - name=(name + ".eh_proj") if name is not None else None, - ) + if self.mhc_enabled: + projection_kwargs = { + "config": self.config, + "init_method": self.config.init_method, + "gather_output": False, + "bias": False, + "skip_bias_add": False, + "is_expert": False, + "tp_group": pg_collection.tp if pg_collection is not None else None, + } + self.e_proj = build_module( + self.submodules.e_proj, + self.config.hidden_size, + self.config.hidden_size, + tp_comm_buffer_name="mtp_e_proj", + name=(name + ".e_proj") if name is not None else None, + **projection_kwargs, + ) + self.h_proj = build_module( + self.submodules.h_proj, + self.config.hidden_size, + self.config.hidden_size, + tp_comm_buffer_name="mtp_h_proj", + name=(name + ".h_proj") if name is not None else None, + **projection_kwargs, + ) + self.eh_proj = None + else: + # Combine each hidden state with the corresponding future-token embedding. + 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, + name=(name + ".eh_proj") if name 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 @@ -1051,6 +1089,17 @@ 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 = mark_keep_in_fp32(nn.Parameter(torch.randn(hc_mult, hc_dim))) + self.hc_head_base = mark_keep_in_fp32(nn.Parameter(torch.zeros(hc_mult))) + self.hc_head_scale = mark_keep_in_fp32(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() def _get_embeddings( @@ -1136,25 +1185,48 @@ 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 InferenceMode.is_active(): - 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 + seq_len, batch_size, _ = hidden_states.shape + hidden_streams = hidden_states.view(seq_len, batch_size, n, h) + hidden_streams = apply_module(self.hnorm)(hidden_streams) + hidden_streams = make_viewless_tensor( + inp=hidden_streams, requires_grad=True, keep_graph=True + ) + embedded, _ = self.e_proj(decoder_input) + embedded = gather_from_tensor_model_parallel_region(embedded, group=self.tp_group) + projected_hidden, _ = self.h_proj(hidden_streams) + projected_hidden = gather_from_tensor_model_parallel_region( + projected_hidden, group=self.tp_group ) + seq_len, batch_size, n, h = projected_hidden.shape + embedded = embedded.unsqueeze(2).expand(seq_len, batch_size, n, h) + hidden_states = (embedded + projected_hidden).reshape(seq_len, batch_size, n * h) + if self.sequence_parallel: + hidden_states = scatter_to_sequence_parallel_region( + hidden_states, group=self.tp_group + ) 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) + hidden_states = torch.cat((decoder_input, hidden_states), -1) + hidden_states, _ = self.eh_proj(hidden_states) + if InferenceMode.is_active(): + 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 + ) + 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( @@ -1225,7 +1297,8 @@ def _proj_and_transformer_layer( padding_mask=padding_mask, ) - hidden_states = self._postprocess(hidden_states) + if not self.mhc_enabled: + hidden_states = self._postprocess(hidden_states) return hidden_states @@ -1234,6 +1307,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 @@ -1806,6 +1889,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. @@ -1813,6 +1897,8 @@ 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. + mhc_multistream (Tensor, optional): Pre-contraction decoder output [s, b, n*h] + used as the input to MTP depths when hyper connections are enabled. attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking self-attention. @@ -1822,7 +1908,11 @@ 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_chunks = list(torch.chunk(mhc_multistream, 1 + offset, dim=0)) + hidden_states = mhc_chunks[offset] + else: + hidden_states = hidden_states_list[offset] if self.config.mtp_detach_heads: hidden_states = hidden_states.detach() @@ -1845,9 +1935,11 @@ 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: + 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 0415035ffbe..732adbbebd1 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -1,8 +1,9 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import logging from contextlib import nullcontext from dataclasses import dataclass -from typing import List, Optional, Set, Union, cast +from typing import List, Optional, Set, Tuple, Union, cast import torch from torch import Tensor @@ -21,8 +22,10 @@ from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward +from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager from megatron.core.transformer.cuda_graphs import annotate_first_last_layer from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType +from megatron.core.transformer.hyper_connection import HyperConnectionModule 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 @@ -318,6 +321,7 @@ def __init__( self.offload_context, self.group_prefetch_offload_commit_async = nullcontext(), None self.config._cpu_offloading_context = None + self.num_residual_streams = config.num_residual_streams self._build_layers() self.num_layers_per_pipeline_rank = len(self.layers) @@ -483,6 +487,46 @@ def __call__(self, *args, **kwargs): return super().__call__(*args, **kwargs)[0] return super().__call__(*args, **kwargs) + def _build_mhc_recompute_layer_plan( + self, use_mhc_recompute: bool + ) -> Tuple[List[Optional[CheckpointWithoutOutputManager]], List[bool]]: + """Pre-build per-layer MHC recompute managers and block-end markers.""" + num_layers = len(self.layers) + layer_managers: List[Optional[CheckpointWithoutOutputManager]] = [None] * num_layers + is_recompute_block_end: List[bool] = [False] * num_layers + + if not use_mhc_recompute or num_layers == 0: + return layer_managers, is_recompute_block_end + + mhc_recompute_layer_num = self.config.mhc_recompute_layer_num + mhc_manager = CheckpointWithoutOutputManager() + + for l_no in range(num_layers): + is_last_in_transformer_block = l_no == num_layers - 1 + is_last_in_recompute_block = is_last_in_transformer_block + if mhc_recompute_layer_num is not None: + is_last_in_recompute_block = is_last_in_transformer_block or ( + (l_no + 1) % mhc_recompute_layer_num == 0 + ) + + layer_managers[l_no] = mhc_manager + is_recompute_block_end[l_no] = is_last_in_recompute_block + + if is_last_in_recompute_block and not is_last_in_transformer_block: + mhc_manager = CheckpointWithoutOutputManager() + + return layer_managers, is_recompute_block_end + + @staticmethod + def _finalize_mhc_recompute_layer( + mhc_manager: Optional[CheckpointWithoutOutputManager], + hidden_states: Tensor, + is_last_in_recompute_block: bool, + ) -> None: + """Finalize MHC recompute state for the current layer when block ends.""" + if mhc_manager is not None and is_last_in_recompute_block: + mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -592,6 +636,13 @@ def forward( # is called here to be future-proof and corner-case-proof. hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + # Expand hidden states for hyper connections at the start of the block + # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage + if self.config.enable_hyper_connections and self.pre_process: + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.num_residual_streams + ) # [s, b, C] -> [s, b, n*C] + if self.config.sequence_parallel: rng_context = tensor_parallel.get_cuda_rng_tracker().fork() else: @@ -619,6 +670,18 @@ def forward( use_inner_quantization_context = False outer_quantization_context = nullcontext() + # Determine if MHC recompute should be used + # Only enable when: training mode AND hyper connections AND 'mhc' in recompute_modules + use_mhc_recompute = ( + self.training + and self.config.enable_hyper_connections + and self.config.recompute_granularity == 'selective' + and "mhc" in self.config.recompute_modules + ) + mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( + use_mhc_recompute + ) + with rng_context, outer_quantization_context: # Forward pass. if self.config.recompute_granularity == 'full' and self.training: @@ -660,6 +723,19 @@ def forward( else: inner_quantization_context = nullcontext() + mhc_manager = mhc_layer_managers[l_no] + if mhc_manager is not None: + mhc_manager.is_last_layer_in_recompute_block = ( + mhc_is_last_in_recompute_block[l_no] + ) + + # Only thread mhc_recompute_manager when the layer is mHC and a + # manager actually exists. Plain TransformerLayer (and its + # MoETransformerLayer subclass) doesn't accept this kwarg, and + # its CUDA-graph machinery rejects unrecognized non-tensor kwargs. + extra_layer_kwargs = ( + {"mhc_recompute_manager": mhc_manager} if mhc_manager is not None else {} + ) with self.offload_context, inner_quantization_context: hidden_states, context = layer( hidden_states=hidden_states, @@ -675,7 +751,13 @@ def forward( packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, padding_mask=padding_mask, + **extra_layer_kwargs, ) + self._finalize_mhc_recompute_layer( + mhc_manager=mhc_manager, + hidden_states=hidden_states, + is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], + ) if ( torch.is_grad_enabled() @@ -688,6 +770,12 @@ def forward( if (l_no + layer_offset) in extract_layer_indices: intermediate_hidden_states.append(hidden_states) + # Only contract if the final layer norm is in this stage + 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] + # Final layer norm. if self.final_layernorm is not None: hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..37eb907bb25 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -291,8 +291,10 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None - """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( + None + ) + """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" @@ -343,6 +345,22 @@ class TransformerConfig(ModelParallelConfig): dsa_indexer_k_norm_fp32: bool = False """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" + #################### + # Compressed sparse attention + #################### + csa_window_size: int = 128 + """Sliding window size for compressed sparse attention.""" + + csa_compress_ratios: Optional[List[int]] = None + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + + csa_compress_rotary_base: float = 40000.0 + """RoPE base for compressed KV positions in compressed sparse attention.""" + + csa_dense_mode: bool = False + """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be + disabled.""" + #################### # linear attention #################### @@ -546,7 +564,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "gdn_norm_out". + "shared_experts", "gdn_norm_out", "mhc". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -556,7 +574,11 @@ class TransformerConfig(ModelParallelConfig): "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. - "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use output-discarding checkpointing, + "mhc": recompute HyperConnection intermediate activations via + CheckpointWithoutOutput + CheckpointWithoutOutputManager. Requires + enable_hyper_connections=True. Cannot be used with "mlp". + "moe_act", "layernorm", "mla_up_proj", "gdn_norm_out", and "mhc" use + output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ @@ -1093,6 +1115,42 @@ class TransformerConfig(ModelParallelConfig): CudaGraphScope instances deserialized from pre-refactor checkpoints are converted to their string names before normalization so existing CUDA_GRAPH_MODULES_DEPRECATIONS handles them.""" + #################### + # Hyper-Connection Configuration + #################### + enable_hyper_connections: bool = False + """Enable mHC residual connections.""" + + num_residual_streams: int = 4 + """Number of residual streams (n in paper).""" + + mhc_sinkhorn_iterations: int = 20 + """Number of Sinkhorn-Knopp iterations for doubly stochastic projection.""" + + mhc_init_gating_factor: float = 0.01 + """Initial value of Gating Factor (alpha in paper).""" + + use_fused_mhc: bool = False + """Use fused kernels for mHC operations when supported. + + Backend selection is operation-specific, with native torch fallbacks that + preserve the same public behavior when Triton or cuTile is unavailable. + """ + + mhc_recompute_layer_num: Optional[int] = None + """Number of layers per MHC recompute block. + + When set, every `mhc_recompute_layer_num` layers form a recompute block. The last layer + in each recompute block (i.e., layer_number % mhc_recompute_layer_num == 0 or the final + layer in the transformer block) will: + - NOT checkpoint its final MLP BDA + - Register the unified recompute hook on its MLP BDA output + - A new CheckpointWithoutOutputManager is created for subsequent layers + + If None, all layers in the transformer block share a single recompute block. + + Must be a positive integer when set.""" + #################### # miscellaneous #################### @@ -1396,6 +1454,31 @@ def __post_init__(self): "dsa_indexer_skip_topk_offset must be non-negative, got " f"{self.dsa_indexer_skip_topk_offset}." ) + elif self.experimental_attention_variant == "dsv4_hybrid": + assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention." + assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set" + mtp_layers = self.mtp_num_layers or 0 + expected_len = self.num_layers + mtp_layers + assert len(self.csa_compress_ratios) >= expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must be at least " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_len})" + ) + assert all( + ratio in [0, 4, 128] for ratio in self.csa_compress_ratios + ), "csa_compress_ratios must be 0, 4, or 128" + assert ( + self.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + assert ( + self.context_parallel_size == 1 + ), "DSv4 Hybrid Attention does not support context parallelism yet." + assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." + if self.dsa_kernel_backend != "none": + raise ValueError( + "The native SBHD DSv4 slice requires dsa_kernel_backend='none'; " + "fused DSv4 backends are added by the follow-up kernel integration." + ) + self.hetereogenous_dist_checkpoint = True if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -1800,6 +1883,7 @@ def __post_init__(self): "moe", "shared_experts", "gdn_norm_out", + "mhc", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1871,6 +1955,46 @@ def __post_init__(self): if "moe" not in self.recompute_modules: self.recompute_modules.append("moe") + # Validation for "mhc" in recompute_modules + if self.recompute_granularity == "selective" and "mhc" in self.recompute_modules: + if not self.enable_hyper_connections: + raise ValueError( + "'mhc' in recompute_modules requires enable_hyper_connections=True." + ) + if "mlp" in self.recompute_modules: + raise ValueError( + "'mhc' and 'mlp' in recompute_modules cannot be used together. " + "They use different checkpoint mechanisms that may conflict." + ) + if self.mhc_recompute_layer_num is not None and ( + isinstance(self.mhc_recompute_layer_num, bool) + or not isinstance(self.mhc_recompute_layer_num, int) + or self.mhc_recompute_layer_num < 1 + ): + raise ValueError( + "mhc_recompute_layer_num must be a positive integer when " + "'mhc' is in recompute_modules." + ) + if self.fine_grained_activation_offloading: + raise NotImplementedError( + "'mhc' in recompute_modules + fine_grained_activation_offloading is " + "not yet supported. The mHC recompute hook currently fires before " + "the offloading backward chunk is initialized, causing tensor_pop " + "on a None chunk. Disable one of them." + ) + + if self.enable_hyper_connections and not ( + self.recompute_granularity == "selective" and "mhc" in self.recompute_modules + ): + warnings.warn( + "HyperConnections are enabled but 'mhc' is not in " + "recompute_modules with selective recompute. Consider adding 'mhc' to " + "recompute_modules with selective recompute to reduce activation memory." + ) + + if self.use_fused_mhc and not self.enable_hyper_connections: + raise ValueError("use_fused_mhc requires enable_hyper_connections=True.") + if self.fine_grained_activation_offloading: assert ( not self.cpu_offloading @@ -2969,10 +3093,12 @@ class MLATransformerConfig(TransformerConfig): """Rank of Query tensor's low rank representation.""" kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" + """Rank of Key and Value tensors' low rank representation. + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_pos_emb_head_dim: int = 64 """Dimension of the position embedding in the QK projection.""" @@ -3010,6 +3136,12 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + o_groups: int = 8 + """Number of groups for grouped low-rank output projection (wo_a).""" + + o_lora_rank: int = 1024 + """Low-rank dimension per group for grouped output (wo_a). Used when o_groups > 0.""" + cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. This is only for the dynamic inference backend and requires that @@ -3022,12 +3154,41 @@ class MLATransformerConfig(TransformerConfig): def __post_init__(self): super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + if ( + self.multi_latent_attention + and self.apply_rope_fusion + and self.rope_type != "yarn" + and self.experimental_attention_variant != "dsv4_hybrid" + ): raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") if self.attention_output_gate: raise NotImplementedError("Output gate is not supported for MLA yet.") + # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim. + if self.experimental_attention_variant == "dsv4_hybrid": + assert ( + not self.mla_down_proj_fusion + ), "MLA down projection fusion must be disabled for DSv4 hybrid mode." + assert self.q_lora_rank is not None, "DSv4 hybrid mode requires q_lora_rank." + assert self.o_groups > 0, "DSv4 hybrid mode requires o_groups to be positive." + assert self.o_lora_rank > 0, "DSv4 hybrid mode requires o_lora_rank to be positive." + assert ( + self.num_attention_heads * self.v_head_dim + ) % self.o_groups == 0, ( + "num_attention_heads * v_head_dim must be divisible by o_groups." + ) + log_single_rank( + logger, + logging.WARNING, + "DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + "v_head_dim and qk_pos_emb_head_dim", + ) + derived = self.v_head_dim - self.qk_pos_emb_head_dim + assert derived > 0, "v_head_dim must be greater than qk_pos_emb_head_dim." + self.qk_head_dim = derived + self.kv_lora_rank = derived + if self.cache_mla_latents: assert ( self.apply_rope_fusion is False diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 5c55f2abe6c..a3b176fa463 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -8,6 +8,9 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Union +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + import torch import torch.distributed from torch import Tensor @@ -268,14 +271,17 @@ class TransformerLayerSubmodules: """ input_layernorm: LayerNormBuilder = IdentityOp + self_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp self_attention: Union[ModuleSpec, type] = IdentityOp self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp pre_cross_attn_layernorm: LayerNormBuilder = IdentityOp + cross_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp cross_attention: Union[ModuleSpec, type] = IdentityOp cross_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp pre_mlp_layernorm: LayerNormBuilder = IdentityOp + mlp_hyper_connection: Union[ModuleSpec, type] = IdentityOp mlp: MlpBuilder | type[IdentityOp] = IdentityOp mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp @@ -368,6 +374,8 @@ def __init__( attention_optional_kwargs["pg_collection"] = pg_collection if pp_layer_offset is not None: attention_optional_kwargs["pp_layer_offset"] = pp_layer_offset + if is_mtp_layer: + attention_optional_kwargs["is_mtp_layer"] = True # [Module 2: SelfAttention] self.self_attention = build_module( @@ -566,6 +574,91 @@ def _get_layer_offset(config: TransformerConfig): ) return get_transformer_layer_offset(config) + @staticmethod + def _group_offload_output_with_bias( + output_with_bias, offload_manager, forced_released_tensors: Optional[list[Tensor]] = None + ): + """Commit a fine-grained offload group for a raw branch output tuple.""" + if isinstance(output_with_bias, tuple): + output = offload_manager.group_offload( + output_with_bias[0], forced_released_tensors=forced_released_tensors + ) + return (output, *output_with_bias[1:]) + return offload_manager.group_offload( + output_with_bias, forced_released_tensors=forced_released_tensors + ) + + def _forward_self_attention_output_with_bias( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[Tensor] = None, + *, + inference_params: Optional[Any] = None, + ): + """Run input norm and self-attention, returning the raw output before BDA.""" + inference_context = deprecate_inference_params(inference_context, inference_params) + + attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") + if self.recompute_input_layernorm: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + apply_module(self.input_layernorm), hidden_states + ) + else: + with attn_norm_manager as hidden_states: + input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) + + if isinstance(input_layernorm_output, tuple): + if len(input_layernorm_output) != 2: + raise ValueError( + "When the output of input_layernorm is a tuple, it is expected " + f"to have 2 elements (output, residual), but got " + f"{len(input_layernorm_output)}" + ) + input_layernorm_output, residual = input_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + using_fused_tp_inference_kernel = ( + InferenceMode.is_active() and self.config.inference_fuse_tp_communication + ) + if using_fused_tp_inference_kernel: + self._set_proj_residual(residual) + + nvtx_range_push(suffix="self_attention") + attention_output_with_bias = self.self_attention( + input_layernorm_output, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + nvtx_range_pop(suffix="self_attention") + + if self.recompute_input_layernorm: + self.input_layernorm_checkpoint.discard_output_and_register_recompute( + attention_output_with_bias[0] + ) + + return attention_output_with_bias, attn_norm_manager, residual + def _forward_attention( self, hidden_states: Tensor, @@ -612,43 +705,16 @@ def _forward_attention( otherwise None. """ inference_context = deprecate_inference_params(inference_context, inference_params) - - # Optional Input Layer norm - attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") - if self.recompute_input_layernorm: - self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() - with attn_norm_manager as hidden_states: - input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( - apply_module(self.input_layernorm), hidden_states - ) - else: - with attn_norm_manager as hidden_states: - input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) - - if isinstance(input_layernorm_output, tuple): - if len(input_layernorm_output) != 2: - raise ValueError( - f"When the output of input_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(input_layernorm_output)}" - ) - input_layernorm_output, residual = input_layernorm_output - else: - residual = hidden_states - - if self.config.fp32_residual_connection: - residual = residual.float() + input_layernorm_output, residual, attn_state = self._run_input_layernorm(hidden_states) using_fused_tp_inference_kernel = ( InferenceMode.is_active() and self.config.inference_fuse_tp_communication ) - if using_fused_tp_inference_kernel: # Set the residual for fused reduce-scatter + add + layer-norm + all-gather # operation in attention's out_proj (linear_proj) self._set_proj_residual(residual) - # Self attention. nvtx_range_push(suffix="self_attention") attention_output_with_bias = self.self_attention( input_layernorm_output, @@ -664,13 +730,71 @@ def _forward_attention( ) nvtx_range_pop(suffix="self_attention") - if self.recompute_input_layernorm: + if self._input_layernorm_checkpoint_active: # discard the output of the input layernorm and register the recompute # as a gradient hook of attention_output_with_bias[0] self.input_layernorm_checkpoint.discard_output_and_register_recompute( attention_output_with_bias[0] ) + hidden_states = self._apply_self_attn_bda_step( + attention_output_with_bias, residual, attn_state + ) + return self._run_cross_attention(hidden_states, context, context_mask, inference_context) + + def _run_input_layernorm(self, hidden_states): + """Run input layernorm with optional output-discarding checkpoint and + fine-grained activation offloading. + + Sets ``self._input_layernorm_checkpoint_active`` so the caller can gate + the post-attention discard-and-register hook on the same condition. The + flag is consumed by the next ``self._apply_self_attn_bda_step`` step. + + Returns: + Tuple ``(input_layernorm_output, residual, attn_state)`` where + ``attn_state`` is an opaque payload subclasses can use to thread + extra intermediates (e.g. mHC ``h_res``/``h_post``) through to + ``_apply_self_attn_bda_step``. Base returns ``()``. + """ + self.attn_norm_manager = self.off_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) + self._input_layernorm_checkpoint_active = self.recompute_input_layernorm + if self._input_layernorm_checkpoint_active: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput() + with self.attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + apply_module(self.input_layernorm), hidden_states + ) + else: + with self.attn_norm_manager as hidden_states: + input_layernorm_output = apply_module(self.input_layernorm)(hidden_states) + + if isinstance(input_layernorm_output, tuple): + if len(input_layernorm_output) != 2: + raise ValueError( + f"When the output of input_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(input_layernorm_output)}" + ) + input_layernorm_output, residual = input_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + return input_layernorm_output, residual, () + + def _apply_self_attn_bda_step(self, attention_output_with_bias, residual, attn_state=()): + """bias-dropout-add for self-attention output + post-step offload commit. + + Subclasses override this to swap in a fused kernel that consumes extra + intermediates threaded via ``attn_state`` (the third element returned + by ``_run_input_layernorm``). Base ignores ``attn_state``. + """ + using_fused_tp_inference_kernel = ( + InferenceMode.is_active() and self.config.inference_fuse_tp_communication + ) # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="self_attn_bda") @@ -688,11 +812,14 @@ def _forward_attention( # Delay the offload of the attention norm until after the self_attn_bda has been computed # because the residual is needed in the self_attn_bda. - hidden_states = attn_norm_manager.group_offload( + hidden_states = self.attn_norm_manager.group_offload( hidden_states, forced_released_tensors=[residual] ) + self.attn_norm_manager = None + return hidden_states - # Optional Layer norm after self-attention + def _run_cross_attention(self, hidden_states, context, context_mask, inference_context): + """Optional pre-cross-attn layernorm + cross-attention + bda block.""" pre_cross_attn_layernorm_output = apply_module(self.pre_cross_attn_layernorm)(hidden_states) if isinstance(pre_cross_attn_layernorm_output, tuple): @@ -709,7 +836,6 @@ def _forward_attention( if self.config.fp32_residual_connection: residual = residual.float() - # Cross attention. attention_output_with_bias = self.cross_attention( pre_cross_attn_layernorm_output, attention_mask=context_mask, @@ -737,6 +863,13 @@ def forward(self, *args, **kwargs): This method calls the core computation of a transformer layer, including self-attention, cross-attention (if applicable), and feed-forward operations. """ + called_from_hybrid_mhc_wrapper = kwargs.pop("_called_from_hybrid_mhc_wrapper", False) + if self.config.enable_hyper_connections and not called_from_hybrid_mhc_wrapper: + raise RuntimeError( + "TransformerLayer.forward() must not be called directly when " + "enable_hyper_connections=True. HyperConnectionHybridLayer must drive " + "the wrapped TransformerLayer through this path." + ) hidden_states, context = self._forward_attention(*args, **kwargs) output = self._forward_mlp( hidden_states, @@ -799,6 +932,76 @@ def _maybe_reflatten_from_moe(self, output, packed_seq_params, mbs): return output return output.transpose(0, 1).reshape(mbs * packed_seq_params.tokens_per_sample, 1, -1) + def _run_pre_mlp_layernorm(self, hidden_states): + """Run pre-MLP layernorm (with optional recompute and offload), unpack a + tuple-output layernorm, and apply the fp32-residual cast. + + Returns: + Tuple ``(pre_mlp_layernorm_output, residual, mlp_state)`` where + ``mlp_state`` is an opaque payload subclasses can use to thread + extra intermediates (e.g. mHC ``mlp_h_res`` / ``mlp_hc_h_post``) + through to ``_apply_mlp_bda_step``. Base returns ``()``. + """ + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + + if isinstance(pre_mlp_layernorm_output, tuple): + if len(pre_mlp_layernorm_output) != 2: + raise ValueError( + f"When the output of pre_mlp_layernorm is a tuple, it is " + f"expected to have 2 elements (output, residual), but " + f"got {len(pre_mlp_layernorm_output)}" + ) + pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + else: + # Residual connection. + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + return pre_mlp_layernorm_output, residual, () + + def _forward_mlp_output_with_bias( + self, + hidden_states: Tensor, + inference_context: BaseInferenceContext | None = None, + padding_mask: Tensor | None = None, + packed_seq_params=None, + ) -> tuple[tuple[Tensor, Tensor | None], Tensor]: + """Run pre-MLP norm and MLP/MoE, returning the raw output before BDA.""" + pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + + if isinstance(pre_mlp_layernorm_output, tuple): + if len(pre_mlp_layernorm_output) != 2: + raise ValueError( + "When the output of pre_mlp_layernorm is a tuple, it is expected " + f"to have 2 elements (output, residual), but got " + f"{len(pre_mlp_layernorm_output)}" + ) + pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + else: + residual = hidden_states + + if self.config.fp32_residual_connection: + residual = residual.float() + + pre_mlp_layernorm_output, padding_mask, moe_unflatten_mbs = self._maybe_unflatten_for_moe( + pre_mlp_layernorm_output, padding_mask, packed_seq_params + ) + + mlp_output_with_bias = self._run_mlp( + pre_mlp_layernorm_output, residual, padding_mask, inference_context + ) + + if moe_unflatten_mbs is not None: + mlp_output, mlp_bias = mlp_output_with_bias + mlp_output = self._maybe_reflatten_from_moe( + mlp_output, packed_seq_params, moe_unflatten_mbs + ) + mlp_output_with_bias = (mlp_output, mlp_bias) + + return mlp_output_with_bias, residual + def _forward_mlp( self, hidden_states: Tensor, @@ -822,29 +1025,58 @@ def _forward_mlp( Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ + pre_mlp_layernorm_output, residual, mlp_state = self._run_pre_mlp_layernorm(hidden_states) - # Optional Layer norm post the cross-attention. - pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states) + pre_mlp_layernorm_output, padding_mask, moe_unflatten_mbs = self._maybe_unflatten_for_moe( + pre_mlp_layernorm_output, padding_mask, packed_seq_params + ) - if isinstance(pre_mlp_layernorm_output, tuple): - if len(pre_mlp_layernorm_output) != 2: - raise ValueError( - f"When the output of pre_mlp_layernorm is a tuple, it is " - f"expected to have 2 elements (output, residual), but " - f"got {len(pre_mlp_layernorm_output)}" - ) - pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + mlp_output_with_bias = self._run_mlp( + pre_mlp_layernorm_output, residual, padding_mask, inference_context + ) + + if moe_unflatten_mbs is not None: + mlp_output, mlp_bias = mlp_output_with_bias + mlp_output = self._maybe_reflatten_from_moe( + mlp_output, packed_seq_params, moe_unflatten_mbs + ) + mlp_output_with_bias = (mlp_output, mlp_bias) + + if ( + self.is_moe_layer + and self.config.cuda_graph_impl == "transformer_engine" + and self.training + and is_graph_capturing() + and CudaGraphModule.moe_router in self.config.cuda_graph_modules + ): + if self.recompute_pre_mlp_layernorm: + # Register the recompute hooks to all the cudagraph output tensors, because some + # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be + # recomputed in backward pass. For example, the router path and the shared expert + # path. So only register in one path is risky. + for tensor in mlp_output_with_bias: + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) + return list(mlp_output_with_bias) + [residual] else: - # Residual connection. - residual = hidden_states + return self._apply_mlp_bda_step(mlp_output_with_bias, residual, mlp_state) - if self.config.fp32_residual_connection: - residual = residual.float() + def _run_mlp( + self, + pre_mlp_layernorm_output: Tensor, + residual: Tensor, + padding_mask: Tensor | None, + inference_context: BaseInferenceContext | None, + ): + """Execute the MLP submodule with the appropriate variant. - pre_mlp_layernorm_output, padding_mask, moe_unflatten_mbs = self._maybe_unflatten_for_moe( - pre_mlp_layernorm_output, padding_mask, packed_seq_params - ) + Picks between the recompute (te_checkpoint / tensor_parallel.checkpoint), + chunked-prefill, and direct-call paths. Shared by both + :class:`TransformerLayer` and :class:`HyperConnectionTransformerLayer` so + the MLP-call branching stays in one place. + Returns: + ``mlp_output_with_bias``: tuple of (mlp_output, mlp_bias). + """ nvtx_range_push(suffix="mlp") # Potentially chunk the MLP computation during prefill to minimize the peak activation size should_chunk_mlp_for_prefill = ( @@ -915,46 +1147,49 @@ def _forward_mlp( pre_mlp_layernorm_output, padding_mask=padding_mask ) - if moe_unflatten_mbs is not None: - mlp_output, mlp_bias = mlp_output_with_bias - mlp_output = self._maybe_reflatten_from_moe( - mlp_output, packed_seq_params, moe_unflatten_mbs - ) - mlp_output_with_bias = (mlp_output, mlp_bias) - nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias - if ( - self.is_moe_layer - and self.config.cuda_graph_impl == "transformer_engine" - and self.training - and is_graph_capturing() - and CudaGraphModule.moe_router in self.config.cuda_graph_modules - ): - if self.recompute_pre_mlp_layernorm: - # Register the recompute hooks to all the cudagraph output tensors, because some - # tensors are in parallel execution paths and they all need pre_mlp_layernorm to be - # recomputed in backward pass. For example, the router path and the shared expert - # path. So only register in one path is risky. - for tensor in mlp_output_with_bias: - self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) - return list(mlp_output_with_bias) + [residual] - else: - return self._forward_post_mlp(mlp_output_with_bias, residual) - - def _forward_post_mlp( - self, mlp_output_with_bias: tuple[Tensor, Tensor | None], residual: Tensor + def _apply_mlp_bda_step( + self, + mlp_output_with_bias: tuple[Tensor, Tensor | None], + residual: Tensor, + mlp_state: tuple = (), ) -> Tensor: """ - Perform operations after the MLP computation. + Perform operations after the MLP computation: bias-dropout-add for + the MLP output + post-step offload commit + viewless-tensor wrap. + + Subclasses override this to swap in a fused kernel that consumes extra + intermediates threaded via ``mlp_state`` (the third element returned + by ``_run_pre_mlp_layernorm``). Base ignores ``mlp_state``. Args: mlp_output_with_bias (Tensor): Output tensor of the MLP layer with bias. residual (Tensor): Residual tensor. + mlp_state: Opaque payload from ``_run_pre_mlp_layernorm``. Default ``()``. Returns: output (Tensor): Transformed hidden states of shape [s, b, h]. """ + # Back-compat shim: prior to the MLP-hook refactor this method was named + # `_forward_post_mlp` and took only (mlp_output_with_bias, residual). If a + # subclass still overrides the legacy name, route through it and emit a + # DeprecationWarning. `mlp_state` is dropped — the legacy contract didn't + # have it. To be removed in a future release. + for klass in type(self).__mro__: + if klass is TransformerLayer: + break + if "_forward_post_mlp" in vars(klass): + warnings.warn( + "TransformerLayer._forward_post_mlp has been renamed to " + "_apply_mlp_bda_step and gained an `mlp_state` parameter. " + "Override `_apply_mlp_bda_step` instead; the legacy hook " + "will be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + return klass._forward_post_mlp(self, mlp_output_with_bias, residual) using_fused_tp_inference_kernel = ( InferenceMode.is_active() and self.config.inference_fuse_tp_communication @@ -1304,10 +1539,10 @@ def _te_cuda_graph_replay_impl(self, args, kwargs, context): nvtx_range_pop(suffix="mlp") # If we early returned, layernorm recompute hooks were attached to the output buffer - # of the cudagraph, so disable the recompute hooks inside _forward_post_mlp + # of the cudagraph, so disable the recompute hooks inside _apply_mlp_bda_step recompute_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm self.recompute_pre_mlp_layernorm = False - output = self._forward_post_mlp(mlp_output_with_bias, residual) + output = self._apply_mlp_bda_step(mlp_output_with_bias, residual) self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm else: # If EP overlap is enabled, needs to return same outputs as submodule.attn @@ -1506,6 +1741,316 @@ def get_layer_norm_weights(self): return +class HyperConnectionTransformerLayer(TransformerLayer): + """A transformer layer with Manifold-Constrained Hyper-Connections (mHC). + + Extends TransformerLayer by adding hyper connection modules around self-attention + and MLP. The n-stream hidden states are aggregated before each sub-layer and + expanded back afterwards using learned mappings (H_pre, H_post, H_res). + + Cross-attention hyper connection is not supported. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + hidden_dropout: Optional[float] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + vp_stage: Optional[int] = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + hidden_dropout=hidden_dropout, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + + if submodules.cross_attention_hyper_connection is not IdentityOp: + raise ValueError( + "HyperConnectionTransformerLayer does not support cross-attention " + "hyper connections. Use IdentityOp for cross_attention_hyper_connection." + ) + + assert submodules.self_attention_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires self_attention_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + assert submodules.mlp_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires mlp_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + + # mHC over a single MoE-MLP layer is not supported in this implementation; + # compose mHC with MoE by wrapping MoE inside a HyperConnectionHybridLayer + # (HybridStack path) instead. This guard fires at setup so misconfigured + # specs fail fast rather than producing silently-wrong shapes at runtime. + if self.is_moe_layer: + raise NotImplementedError( + "HyperConnectionTransformerLayer does not support MoE MLP submodules. " + "To combine mHC with MoE, wrap the MoE block as a HybridStack layer " + "via HyperConnectionHybridLayer instead." + ) + + self.self_attention_hyper_connection = build_module( + submodules.self_attention_hyper_connection, + config=self.config, + layer_number=self.layer_number, + ) + + self.mlp_hyper_connection = build_module( + submodules.mlp_hyper_connection, config=self.config, layer_number=self.layer_number + ) + + # When mHC recompute is active, skip checkpointing if the layernorm + # is IdentityOp (fused into TE linear) — there is nothing to recompute. + self.mhc_checkpoint_input_layernorm = not isinstance(self.input_layernorm, IdentityOp) + self.mhc_checkpoint_pre_mlp_layernorm = not isinstance(self.pre_mlp_layernorm, IdentityOp) + + # Set per-call by __call__ from kwargs so forward can read it without re-piping + # the manager through the CUDA-graph kwarg path (CheckpointWithoutOutputManager + # is not a CUDA-graph-supported type and gets stripped during capture). Read by + # _run_input_layernorm, _apply_self_attn_bda_step, _run_pre_mlp_layernorm, and + # _apply_mlp_bda_step — do not delete; appears unused only at the class level. + self._mhc_recompute_manager: Optional['CheckpointWithoutOutputManager'] = None + + def __call__(self, *args, **kwargs): + # Pull the manager off kwargs before super().__call__ hands them to the + # CUDA-graph machinery (which can't handle a CheckpointWithoutOutputManager). + # forward() reads the value back from self. + self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None) + return super().__call__(*args, **kwargs) + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Override to produce n-stream hidden_states of shape [s, b, n*C]. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. The base class returns [s, b, C], but mHC layers operate on + n-stream hidden states of shape [s, b, n*C]. + """ + static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + hs = static_inputs["hidden_states"] + n = self.config.num_residual_streams + static_inputs["hidden_states"] = torch.ones( + (hs.shape[0], hs.shape[1], n * self.config.hidden_size), + dtype=hs.dtype, + requires_grad=hs.requires_grad, + device=hs.device, + ) + return static_inputs + + def _get_submodules_under_cudagraphs(self): + """Override to include hyper connection modules. + + The base TransformerLayer._get_submodules_under_cudagraphs does not include + self_attention_hyper_connection / mlp_hyper_connection. Their learnable + parameters (mapping_proj, alpha_*, bias) need manual pre-forward hooks + during CUDA graph replay so that parameter all-gathers are triggered. + """ + submodules = super()._get_submodules_under_cudagraphs() + + if not self.config.cuda_graph_scope: + return submodules + + if CudaGraphScope.attn in self.config.cuda_graph_scope: + submodules.append(self.self_attention_hyper_connection) + # HC layer rejects MoE MLPs in __init__, so only the dense (mlp) scope applies. + if CudaGraphScope.mlp in self.config.cuda_graph_scope: + submodules.append(self.mlp_hyper_connection) + return submodules + + def forward(self, *args, **kwargs): + """Forward pass with MHC recompute manager support. + + Inherits ``_forward_attention`` and ``_forward_mlp`` from base; the + mHC-specific behavior is contained in the ``_run_input_layernorm``, + ``_apply_self_attn_bda_step``, ``_run_pre_mlp_layernorm``, and + ``_apply_mlp_bda_step`` overrides, which read the manager off + ``self`` and thread per-call intermediates through the + ``attn_state`` / ``mlp_state`` slots. + + Override exists only to skip the ``enable_hyper_connections`` assert + on base ``TransformerLayer.forward``. + """ + hidden_states, context = self._forward_attention(*args, **kwargs) + output = self._forward_mlp( + hidden_states, + kwargs.get("inference_context", None), + padding_mask=kwargs.get("padding_mask", None), + packed_seq_params=kwargs.get("packed_seq_params", None), + ) + return output, context + + def _run_input_layernorm(self, hidden_states): + """HC input layernorm: hyper-connection pre-wrap + mHC-aware checkpoint. + + Threads ``h_res`` and ``h_post`` (produced by the hyper-connection + pre-wrap) to ``_apply_self_attn_bda_step`` via the ``attn_state`` slot + in the return tuple. Also sets + ``self._input_layernorm_checkpoint_active`` for the post-self-attn + discard hook. + + Returns ``(input_layernorm_output, residual, (h_res, h_post))`` where + ``residual`` is the n-stream hidden state captured before + aggregation — it flows to ``_apply_self_attn_bda_step`` via the base + skeleton's ``residual`` argument, and ``(h_res, h_post)`` flows via + ``attn_state``. + """ + # Capture the n-stream residual BEFORE self_attention_hyper_connection + # aggregates n-stream -> single-stream. The fused bda kernel needs the + # original n-stream tensor. + residual = hidden_states + + nvtx_range_push(suffix="self_attention_hyper_connection") + hidden_states, h_res, h_post = self.self_attention_hyper_connection( + hidden_states, mhc_recompute_manager=self._mhc_recompute_manager + ) + nvtx_range_pop(suffix="self_attention_hyper_connection") + + self.attn_norm_manager = self.off_interface( + self.offload_attn_norm, hidden_states, "attn_norm" + ) + self._input_layernorm_checkpoint_active = self.recompute_input_layernorm or ( + self._mhc_recompute_manager is not None and self.mhc_checkpoint_input_layernorm + ) + if self._input_layernorm_checkpoint_active: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self._mhc_recompute_manager + ) + with self.attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + self.input_layernorm, hidden_states + ) + else: + with self.attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm(hidden_states) + + return input_layernorm_output, residual, (h_res, h_post) + + def _apply_self_attn_bda_step(self, attention_output_with_bias, residual, attn_state): + """HC fused bias-dropout-add: combines apply_h_res + apply_h_post + bda. + + Unpacks ``h_res`` and ``h_post`` from ``attn_state`` (threaded by + ``_run_input_layernorm`` via the base skeleton). + """ + h_res, h_post = attn_state + nvtx_range_push(suffix="self_attention_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.self_attention_hyper_connection.fused_h_res_h_post_bda( + h_res, + residual, + h_post, + attention_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + self._mhc_recompute_manager, + ) + nvtx_range_pop(suffix="self_attention_fused_h_res_h_post_bda") + # HC omits forced_released_tensors — the n-stream residual is consumed + # by the fused kernel above, so the base class's "release residual after + # commit" trick doesn't apply. + hidden_states = self.attn_norm_manager.group_offload(hidden_states) + self.attn_norm_manager = None + return hidden_states + + def _run_pre_mlp_layernorm(self, hidden_states): + """HC pre-mlp layernorm: hyper-connection pre-wrap + mHC-aware checkpoint. + + Threads ``mlp_h_res`` and ``mlp_hc_h_post`` (produced by the + hyper-connection pre-wrap) to ``_apply_mlp_bda_step`` via the + ``mlp_state`` slot in the return tuple. + + Returns ``(pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post))`` + where ``residual`` is the n-stream hidden state captured before + aggregation — it flows to ``_apply_mlp_bda_step`` via the base + skeleton's ``residual`` argument, and ``(mlp_h_res, mlp_hc_h_post)`` + flows via ``mlp_state``. + """ + # Capture the n-stream residual BEFORE mlp_hyper_connection + # aggregates n-stream -> single-stream. The fused bda kernel needs the + # original n-stream tensor. + residual = hidden_states + + nvtx_range_push(suffix="mlp_hyper_connection") + hidden_states, mlp_h_res, mlp_hc_h_post = self.mlp_hyper_connection( + hidden_states, mhc_recompute_manager=self._mhc_recompute_manager + ) + nvtx_range_pop(suffix="mlp_hyper_connection") + + self.mlp_norm_manager = self.off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") + checkpoint_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm or ( + self._mhc_recompute_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ) + if checkpoint_pre_mlp_layernorm: + self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self._mhc_recompute_manager + ) + with self.mlp_norm_manager as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( + self.pre_mlp_layernorm, hidden_states + ) + else: + with self.mlp_norm_manager as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + + return pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post) + + def _apply_mlp_bda_step(self, mlp_output_with_bias, residual, mlp_state): + """HC fused bias-dropout-add for MLP: combines apply_h_res + apply_h_post + bda. + + Unpacks ``mlp_h_res`` and ``mlp_hc_h_post`` from ``mlp_state`` (threaded + by ``_run_pre_mlp_layernorm`` via the base skeleton). Computes the + per-call ``mhc_mlp_bda_manager`` from ``self._mhc_recompute_manager``: + the last layer of a recompute block does NOT pass the manager into the + fused-bda checkpoint — the block-end finalize hook handles its output + discard. + """ + mlp_h_res, mlp_hc_h_post = mlp_state + + is_last_in_recompute_block = bool( + self._mhc_recompute_manager is not None + and getattr(self._mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_mlp_bda_manager = None if is_last_in_recompute_block else self._mhc_recompute_manager + + if self.recompute_pre_mlp_layernorm or ( + mhc_mlp_bda_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ): + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + + nvtx_range_push(suffix="mlp_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.mlp_hyper_connection.fused_h_res_h_post_bda( + mlp_h_res, + residual, + mlp_hc_h_post, + mlp_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + mhc_mlp_bda_manager, + ) + nvtx_range_pop(suffix="mlp_fused_h_res_h_post_bda") + + # HC omits forced_released_tensors — the n-stream residual is consumed + # by the fused kernel above, so the base class's "release residual after + # commit" trick doesn't apply. + if self.mlp_norm_manager is not None: + hidden_states = self.mlp_norm_manager.group_offload(hidden_states) + self.mlp_norm_manager = None + + output = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + return output + + class MoETransformerLayer(TransformerLayer): """ A Transformer layer specialized for Mixture-of-Experts (MoE) architectures. @@ -1704,7 +2249,16 @@ def _forward_mlp_postprocess(self, residual, output, shared_expert_output, mlp_b self.mlp.fwd_execution_map = "postprocess" output = apply_module(self.mlp)(None, intermediate_tensors=(output, shared_expert_output)) - return self._forward_post_mlp((output, mlp_bias), residual) + out = self._apply_mlp_bda_step((output, mlp_bias), residual) + + if is_graph_capturing() and not is_graph_warmup(): + for attr_name, attr in self.token_dispatcher_attrs.items(): + weak_ref = make_weakref(attr, inplace=False) + self.token_dispatcher_attrs[attr_name] = weak_ref + obj, name = self._resolve_token_dispatcher_attr(attr_name) + setattr(obj, name, weak_ref) + + return out def _forward_mlp( self, hidden_states, inference_context=None, padding_mask=None, packed_seq_params=None diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 70f26c64d56..d59aafb2d84 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module +from megatron.core.transformer.spec_utils import ModuleSpec, import_module from megatron.training.config import ( CheckpointConfig, DistributedInitConfig, @@ -274,6 +274,81 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: return field_docstrings +def _normalize_dsv4_hybrid_csa_compress_ratios( + args: Namespace, kw_args: dict, pattern: str +) -> None: + """Normalize compact DSv4 HybridModel ratios into a per-layer config list.""" + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + variant = kw_args.get( + 'experimental_attention_variant', getattr(args, 'experimental_attention_variant', None) + ) + if variant != 'dsv4_hybrid': + return + + fixed_ratio_map = {Symbols.WINDOW: 0, Symbols.CSA: 4, Symbols.HCA: 128} + ratio_symbols = set(fixed_ratio_map) + sections = pattern.split(Symbols.MTP_SEPARATOR) + layers = ''.join(section.replace(Symbols.PIPE, '') for section in sections) + attention_symbols = [symbol for symbol in layers if symbol in ratio_symbols] + compact_len = len(attention_symbols) + full_len = len(layers) + + def compact_and_full(provided: list[int]) -> tuple[list[int], list[int]]: + compact = [] + full = [] + compact_iter = iter(provided) + for symbol in layers: + if symbol in ratio_symbols: + ratio = next(compact_iter) + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + compact.append(ratio) + full.append(ratio) + else: + full.append(0) + return compact, full + + provided_ratios = getattr(args, 'csa_compress_ratios', None) + if provided_ratios is None: + compact_ratios = [fixed_ratio_map[symbol] for symbol in attention_symbols] + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = compact_and_full( + compact_ratios + ) + return + + provided = list(provided_ratios) + if len(provided) == compact_len: + args.csa_compress_ratios, kw_args['csa_compress_ratios'] = compact_and_full(provided) + elif len(provided) == full_len: + compact = [] + for ratio, symbol in zip(provided, layers): + if symbol in ratio_symbols: + expected = fixed_ratio_map[symbol] + assert ratio == expected, ( + f"csa_compress_ratios has ratio {ratio} for hybrid symbol " + f"'{symbol}', expected {expected}." + ) + compact.append(ratio) + else: + assert ratio == 0, ( + "csa_compress_ratios should not pad non-DSv4 hybrid symbol " + f"'{symbol}' with non-zero ratio {ratio}." + ) + args.csa_compress_ratios = compact + kw_args['csa_compress_ratios'] = provided + else: + raise AssertionError( + f"csa_compress_ratios length ({len(provided)}) must equal either the " + f"number of W/C/H attention symbols ({compact_len}) or the legacy " + f"number of all layers in the hybrid pattern ({full_len}) for pattern " + f"'{pattern}'." + ) + + def core_transformer_config_from_args(args, config_class=None): from megatron.core.activations import squared_relu from megatron.core.fusions.fused_bias_geglu import quick_gelu @@ -345,8 +420,19 @@ def core_transformer_config_from_args(args, config_class=None): if args.hybrid_layer_pattern is not None: kw_args['is_hybrid_model'] = True from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols - if Symbols.DS_ATTENTION in args.hybrid_layer_pattern: - kw_args['experimental_attention_variant'] = 'dsa' + + pattern = args.hybrid_layer_pattern + has_dsv4_attention = any( + symbol in pattern for symbol in (Symbols.WINDOW, Symbols.CSA, Symbols.HCA) + ) + has_dsa = Symbols.DS_ATTENTION in pattern + if getattr(args, 'experimental_attention_variant', None) is None: + if has_dsv4_attention: + kw_args['experimental_attention_variant'] = 'dsv4_hybrid' + elif has_dsa: + kw_args['experimental_attention_variant'] = 'dsa' + + _normalize_dsv4_hybrid_csa_compress_ratios(args, kw_args, pattern) kw_args['inference_sampling_seed'] = args.seed @@ -504,7 +590,10 @@ def hybrid_config_from_args( not transformer_cfg.inference_fuse_tp_communication ), "inference_fuse_tp_communication is not supported for HybridModel" elif args.spec is not None: - kwargs["hybrid_stack_spec"] = import_module(args.spec) + hybrid_stack_spec = import_module(args.spec) + if callable(hybrid_stack_spec) and not isinstance(hybrid_stack_spec, ModuleSpec): + hybrid_stack_spec = hybrid_stack_spec(transformer_cfg) + kwargs["hybrid_stack_spec"] = hybrid_stack_spec kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy kwargs["hybrid_layer_pattern"] = args.hybrid_layer_pattern diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d4f9eb9c0de..0beee5ce913 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -323,6 +323,15 @@ def no_rope_freq_type(x): # it's a single int but in str return int(x) + +def compress_ratios_type(x): + """Parse per-layer compression ratios for compressed sparse attention.""" + if isinstance(x, list): + return x + assert isinstance(x, str) + return _eval_pattern(x) + + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -881,10 +890,9 @@ def validate_args(args, defaults={}): args.rank ) - # Infer use of MLA from unified pattern - if args.hybrid_layer_pattern and ( - Symbols.MLA in args.hybrid_layer_pattern - or Symbols.DS_ATTENTION in args.hybrid_layer_pattern + # All MLA-based hybrid attention symbols use MLA projections. + if args.hybrid_layer_pattern and any( + symbol in args.hybrid_layer_pattern for symbol in Symbols.MLA_ATTENTION ): args.multi_latent_attention = True @@ -2213,6 +2221,7 @@ def _add_network_size_args(parser): "no_rope_freq", "moe_layer_freq", "linear_attention_freq", + "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "cp_comm_type", @@ -3399,6 +3408,11 @@ def _add_mla_args(parser): help="Mscale for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") + group.add_argument('--o-groups', type=int, default=8, + help="Number of groups for grouped low-rank output projection (wo_a).") + group.add_argument('--o-lora-rank', type=int, default=1024, + help="Low-rank dimension per group for grouped output (wo_a). " + "Used when o-groups > 0.") group.add_argument('--cache-mla-latents', action='store_true', default=False, help="If set caches the mla down projected latents with mla flash decode.") group.add_argument( @@ -3423,6 +3437,15 @@ def _add_experimental_attention_variant_args(parser): 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + '--csa-compress-ratios', + type=compress_ratios_type, + default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a Python list expression such as "[0,0,4,128,4,128]" or ' + '"([0]+[4,128]*2)*3". Valid values are 0, 4, and 128, and the ' + 'list length must be at least num_layers plus mtp_num_layers.', + ) return parser def _add_heterogeneous_args(parser): diff --git a/megatron/training/training.py b/megatron/training/training.py index acdd727b82d..ec1954256a2 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2842,12 +2842,23 @@ def training_log( # Track sparse attention indexer loss. if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0: indexer_loss_scale = 1 / get_num_microbatches() + assert isinstance( + pg_collection, ProcessGroupCollection + ), "DSA indexer logging requires a ProcessGroupCollection" DSAIndexerLossLoggingHelper.track_indexer_metrics( loss_scale=indexer_loss_scale, iteration=iteration, writer=writer, + pg_collection=pg_collection, wandb_writer=wandb_writer, total_loss_dict=total_loss_dict, + num_layers=args.num_layers + (args.mtp_num_layers or 0), + num_indexer_layers=( + sum(ratio == 4 for ratio in args.csa_compress_ratios) + if args.csa_compress_ratios is not None + else None + ), + preserve_groups=args.cuda_graph_impl != "none", ) # Dump memory snapshot and print metrics to stdout. diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index f93e09a43b7..7c319e0a14a 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_optimizer.py @@ -79,6 +79,27 @@ def sharded_state_dict(self): return sharded_state_dict +class NativeFp32Model(torch.nn.Module): + """Parameters for an interleaved trainable/frozen BF16 and FP32 group.""" + + def __init__(self): + super().__init__() + self.pre = torch.nn.Linear(8, 8, bias=False) + self.frozen = torch.nn.Linear(8, 8, bias=False) + self.frozen.weight.requires_grad_(False) + self.gate = torch.nn.Parameter(torch.zeros(24, dtype=torch.float32)) + self.post = torch.nn.Linear(8, 8, bias=False) + self.config = TransformerConfig( + hidden_size=8, num_attention_heads=1, num_layers=1, bf16=True + ) + + def sharded_state_dict(self): + return { + key: ShardedTensor.from_rank_offsets(key, value) + for key, value in self.state_dict(keep_vars=True).items() + } + + class SwigluFactoryModel(torch.nn.Module): def __init__(self, pp_separate_model: bool = False): super().__init__() @@ -238,6 +259,65 @@ def test_optimizer_params(self, tmp_path_dist_ckpt): ] ) + def test_float16_optimizer_with_native_fp32_and_frozen_params(self): + """Native FP32 and frozen param ids must not shift BF16 checkpoint state.""" + from megatron.core.optimizer import OptimizerConfig + from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params + from megatron.core.transformer.module import ( + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, + ) + + Utils.initialize_model_parallel(1, 1) + model = NativeFp32Model().cuda() + model.gate = mark_keep_in_fp32(model.gate) + convert_module_to_dtype_except_fp32_marked(model, torch.bfloat16) + assert model.pre.weight.dtype == torch.bfloat16 + assert model.frozen.weight.dtype == torch.bfloat16 + assert not model.frozen.weight.requires_grad + assert model.gate.dtype == torch.float32 + assert model.post.weight.dtype == torch.bfloat16 + + # Use an explicit trainable BF16/frozen BF16/FP32/trainable BF16 order. + # Module.parameters() would yield the root gate before child parameters. + ordered_params = [model.pre.weight, model.frozen.weight, model.gate, model.post.weight] + for param in ordered_params: + if param.requires_grad: + param.grad = torch.zeros_like(param) + inner_optim = Adam(ordered_params) + inner_optim.step() + + optim = Float16OptimizerWithFloat16Params( + inner_optim, + OptimizerConfig(optimizer='adam', lr=1e-4, bf16=True), + None, + lambda opt, cfg: None, + ) + sharded_state_dict = optim.sharded_state_dict(model.sharded_state_dict()) + + # FP32 main copies pair with the BF16 params only, in optimizer order. + fp32_params = sharded_state_dict['fp32_from_fp16_params'][0] + assert [(sharded.key, tuple(sharded.data.shape)) for sharded in fp32_params] == [ + ('optimizer.state.fp32_param.pre.weight', (8, 8)), + ('optimizer.state.fp32_param.post.weight', (8, 8)), + ] + + # The frozen parameter has neither optimizer state nor an fp32 main copy. + state = sharded_state_dict['optimizer']['state'] + assert 1 not in state + + # Per-param state maps every trainable param, including native FP32, to the right key. + expected = {0: ('pre.weight', (8, 8)), 2: ('gate', (24,)), 3: ('post.weight', (8, 8))} + for param_id, (model_key, shape) in expected.items(): + for state_key in ('exp_avg', 'exp_avg_sq'): + sharded = state[param_id][state_key] + assert sharded.key == f'optimizer.state.{state_key}.{model_key}', sharded.key + assert tuple(sharded.data.shape) == shape, ( + param_id, + sharded.key, + sharded.data.shape, + ) + def initialize_pp_agnostic_model(pre_process=True, post_process=True, seed=0, **config_kwargs): torch.manual_seed(seed) diff --git a/tests/unit_tests/fusions/test_bias_dropout_fusion.py b/tests/unit_tests/fusions/test_bias_dropout_fusion.py index f8b23900543..a7c63626e93 100644 --- a/tests/unit_tests/fusions/test_bias_dropout_fusion.py +++ b/tests/unit_tests/fusions/test_bias_dropout_fusion.py @@ -319,3 +319,86 @@ def test_fp32_residual_precision_advantage(self): f"fp32 residual error ({err_fp32:.6e}) should be less than " f"bf16 residual error ({err_bf16:.6e})" ) + + +# ============================================================================ +# Tests for the mHC recompute path of get_bias_dropout_add +# ============================================================================ +# +# When ``mhc_recompute_manager`` is provided, ``get_bias_dropout_add`` returns +# a closure that wraps the underlying BDA in ``CheckpointWithoutOutput`` and +# auto-registers with the supplied ``CheckpointWithoutOutputManager``. These tests cover +# that branch (which is otherwise only invoked indirectly from the mHC layer +# forward path). + + +class TestBiasDropoutAddMhcRecompute: + """Direct coverage for ``_get_checkpointed_bda``.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import initialize_rng_tracker + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + from tests.unit_tests.test_utilities import Utils + + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("fused", [False, True]) + @pytest.mark.parametrize("with_bias", [True, False]) + def test_checkpointed_bda_forward_backward(self, fused, with_bias): + """Closure runs forward+backward and registers with the manager.""" + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + + torch.manual_seed(0) + manager = CheckpointWithoutOutputManager() + bda = get_bias_dropout_add(training=True, fused=fused, mhc_recompute_manager=manager) + + x = torch.randn(8, 4, 16, device="cuda", requires_grad=True) + residual = torch.randn_like(x, requires_grad=True) + bias = torch.zeros(16, device="cuda") if with_bias else None + x_with_bias = (x, bias) if with_bias else x + + out = bda(x_with_bias, residual, 0.0) + assert out.shape == x.shape + assert out.dtype == x.dtype + assert len(manager.checkpoints) == 1, "checkpoint should auto-register with manager" + + loss = out.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + + assert x.grad is not None and torch.isfinite(x.grad).all() + assert residual.grad is not None and torch.isfinite(residual.grad).all() + + def test_checkpointed_bda_chained_managers(self): + """Two checkpointed BDAs chained on one manager both register.""" + from megatron.core.tensor_parallel.random import CheckpointWithoutOutputManager + + torch.manual_seed(0) + manager = CheckpointWithoutOutputManager() + bda = get_bias_dropout_add(training=True, fused=False, mhc_recompute_manager=manager) + + x = torch.randn(4, 2, 8, device="cuda", requires_grad=True) + residual = torch.randn_like(x, requires_grad=True) + + y1 = bda((x, None), residual, 0.0) + y2 = bda((y1, None), residual, 0.0) + + assert len(manager.checkpoints) == 2, "each call should register a new checkpoint" + loss = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + assert x.grad is not None + + def test_get_bda_without_manager_unchanged(self): + """The default (manager=None) path returns the regular BDA, not a closure.""" + unfused = get_bias_dropout_add(training=True, fused=False) + fused = get_bias_dropout_add(training=False, fused=True) + # Both must be callable; neither should be the mHC closure (which has __closure__ over manager). + assert callable(unfused) and callable(fused) + assert getattr(unfused, "__name__", "") != "_checkpointed_bda" + assert getattr(fused, "__name__", "") != "_checkpointed_bda" diff --git a/tests/unit_tests/fusions/test_fused_mhc_kernels.py b/tests/unit_tests/fusions/test_fused_mhc_kernels.py new file mode 100644 index 00000000000..d916162c338 --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_mhc_kernels.py @@ -0,0 +1,1557 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Unit tests for unified fused mHC kernels and native implementations. + +Each test compares the fused kernel's forward output AND backward gradients +against a pure-PyTorch differentiable reference to catch numerical drift +introduced by kernel fusion. The public fused API tests also exercise backend +dispatch and fallback selection through the same entry points used by mHC. +""" + +import math +from typing import Optional + +import pytest +import torch +from torch import Tensor + +from megatron.core.fusions.fused_mhc_kernels import is_cutile_available, is_triton_available +from megatron.core.transformer.hyper_connection import ( + native_h_aggregate, + native_h_post_bda, + native_proj_rms, + native_sinkhorn, +) + +_require_cutile = pytest.mark.skipif( + not is_cutile_available(), reason="cuTile unavailable for current device" +) +_require_triton = pytest.mark.skipif(not is_triton_available(), reason="Triton not installed") + + +@pytest.fixture(autouse=True) +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + +DTYPE = torch.bfloat16 +DEVICE = "cuda" +FWD_ATOL, FWD_RTOL = 2e-2, 2e-2 +BWD_ATOL, BWD_RTOL = 5e-2, 5e-2 +RAND_LO, RAND_HI = -0.1, 0.1 +COSINE_SIM_THRESH = 0.999 + + +def _assert_cosine_similar(a: Tensor, b: Tensor, threshold: float, msg: str = ""): + """Assert that flattened tensors have cosine similarity >= threshold.""" + a_flat = a.flatten().float() + b_flat = b.flatten().float() + sim = torch.nn.functional.cosine_similarity(a_flat.unsqueeze(0), b_flat.unsqueeze(0)).item() + assert sim >= threshold, ( + f"{msg}: cosine similarity {sim:.6f} < {threshold} " + f"(max_abs_diff={torch.max(torch.abs(a_flat - b_flat)):.6e})" + ) + + +def _rand(*shape, **kwargs): + """Uniform in [RAND_LO, RAND_HI] to keep magnitudes small for bf16 stability.""" + return torch.empty(*shape, dtype=DTYPE, device=DEVICE, **kwargs).uniform_(RAND_LO, RAND_HI) + + +def _info(): + if is_triton_available() and is_cutile_available(): + backend = "triton+cuTile" + elif is_triton_available(): + backend = "triton+native" + elif is_cutile_available(): + backend = "cuTile" + else: + backend = "native" + print(f"\n [backend: {backend}]") + + +# ============================================================================ +# Pure-PyTorch differentiable references (used by both fwd AND bwd tests) +# ============================================================================ + + +def _ref_sinkhorn(logits: Tensor, num_iters: int, eps: float = 1e-6) -> Tensor: + M = logits.softmax(dim=-1) + eps + M = M / (M.sum(dim=-2, keepdim=True) + eps) + for _ in range(num_iters - 1): + M = M / (M.sum(dim=-1, keepdim=True) + eps) + M = M / (M.sum(dim=-2, keepdim=True) + eps) + return M + + +def _ref_h_aggregate(x: Tensor, h_pre: Tensor) -> Tensor: + return (x * h_pre.unsqueeze(-1)).sum(dim=2) + + +def _ref_h_post_bda( + h_res: Tensor, orig_res: Tensor, h_post: Tensor, x: Tensor, bias: Optional[Tensor] +) -> Tensor: + s, b, n, C = orig_res.shape + h_res_batched = h_res.view(s * b, n, n) + orig_batched = orig_res.view(s * b, n, C) + mixed = torch.bmm(h_res_batched.transpose(1, 2), orig_batched).view(s, b, n, C) + x_exp = h_post.unsqueeze(-1) * x.unsqueeze(2) + out = x_exp + mixed + if bias is not None: + out = out + h_post.unsqueeze(-1) * bias.view(1, 1, 1, C) + return out + + +def _h_post_bda_transpose_case(): + h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], dtype=DTYPE, device=DEVICE) + orig = torch.tensor([[[[10.0, 100.0], [1.0, 2.0]]]], dtype=DTYPE, device=DEVICE) + h_post = torch.zeros(1, 1, 2, dtype=DTYPE, device=DEVICE) + x = torch.zeros(1, 1, 2, dtype=DTYPE, device=DEVICE) + expected = torch.tensor([[[[13.0, 106.0], [24.0, 208.0]]]], dtype=DTYPE, device=DEVICE) + return h_res, orig, h_post, x, expected + + +def _ref_proj_rms(x: Tensor, weight: Tensor, eps: float = 1e-6): + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = 1.0 / (norm / math.sqrt(K) + eps) + return proj, r + + +def _ref_proj_rms_compute_h( + x: Tensor, + weight: Tensor, + alpha_pre: Tensor, + alpha_post: Tensor, + alpha_res: Tensor, + bias: Tensor, + n: int, + eps: float = 1e-6, + compute_h_eps: float = 1e-6, +): + """Reference: fused proj_rms + compute_h.""" + proj = torch.matmul(x, weight.t()) + norm = x.norm(dim=-1, keepdim=True) + K = x.shape[-1] + r = norm / math.sqrt(K) # [M, 1] + N = proj.shape[-1] + alpha = torch.cat([alpha_pre.expand(n), alpha_post.expand(n), alpha_res.expand(N - 2 * n)]) + h = proj * alpha.unsqueeze(0) / (r + eps) + bias.unsqueeze(0) + h_pre = h[..., :n].sigmoid() + compute_h_eps + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res = h[..., 2 * n :] + return h_pre, h_post, h_res, r + + +# ============================================================================ +# Sinkhorn +# ============================================================================ + + +class TestNativeSinkhorn: + """Tests for the native SinkhornKnopp implementation.""" + + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, iters): + """native_sinkhorn fwd output and bwd grad must match the inline PyTorch reference.""" + _info() + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + # -- native_sinkhorn path (autograd.Function) -- + inp_f = data.clone().requires_grad_(True) + out_f = native_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + grad_f = inp_f.grad.clone() + + # -- inline torch reference (fully differentiable) -- + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestFusedSinkhorn: + """Public fused sinkhorn dispatch/fallback plus numerical correctness.""" + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10)]) + def test_fwd_bwd_vs_reference(self, s, b, n, iters): + """E2E: public fused fwd output and bwd grad must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_sinkhorn + + _info() + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + # -- fused path -- + inp_f = data.clone().requires_grad_(True) + out_f = fused_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + grad_f = inp_f.grad.clone() + + # -- reference path (fully differentiable) -- + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + grad_r = inp_r.grad.clone() + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(grad_f, grad_r, atol=BWD_ATOL, rtol=BWD_RTOL) + + +# ============================================================================ +# H_aggregate +# ============================================================================ + + +class TestNativeHAggregate: + """Tests for native_h_aggregate.""" + + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, C): + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + grad_out = _rand(s, b, C) + + xf = x_data.clone().requires_grad_(True) + hf = h_data.clone().requires_grad_(True) + of = native_h_aggregate(xf, hf) + of.backward(grad_out) + + xr = x_data.clone().requires_grad_(True) + hr = h_data.clone().requires_grad_(True) + oref = _ref_h_aggregate(xr, hr) + oref.backward(grad_out) + + torch.testing.assert_close(of, oref, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestFusedHAggregate: + """Public fused h_aggregate dispatch/fallback plus numerical correctness.""" + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256)]) + def test_fwd_bwd_vs_reference(self, s, b, n, C): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_h_aggregate + + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + grad_out = _rand(s, b, C) + + # -- fused path -- + xf = x_data.clone().requires_grad_(True) + hf = h_data.clone().requires_grad_(True) + of = fused_h_aggregate(xf, hf) + of.backward(grad_out) + + # -- reference path -- + xr = x_data.clone().requires_grad_(True) + hr = h_data.clone().requires_grad_(True) + oref = _ref_h_aggregate(xr, hr) + oref.backward(grad_out) + + torch.testing.assert_close(of, oref, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + torch.testing.assert_close(hf.grad, hr.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + +class TestTritonHAggregate: + """Tests for Triton h_aggregate forward against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 1, 2, 256), (64, 8, 4, 4096)]) + def test_fwd_vs_reference(self, s, b, n, C): + from megatron.core.fusions.fused_mhc_kernels import _triton_h_aggregate_fwd + + _info() + x_data = _rand(s, b, n, C) + h_data = _rand(s, b, n) + + out_t = _triton_h_aggregate_fwd(x_data, h_data) + out_r = _ref_h_aggregate(x_data, h_data) + + torch.testing.assert_close(out_t, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + + +# ============================================================================ +# H_post BDA +# ============================================================================ + + +class TestNativeHPostBDA: + """Tests for native_h_post_bda.""" + + def test_forward_uses_h_res_transpose(self): + h_res, orig, h_post, x, expected = _h_post_bda_transpose_case() + out = native_h_post_bda(h_res, orig, h_post, x, bias=None) + + torch.testing.assert_close(out, expected, atol=0.0, rtol=0.0) + + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) + def test_fwd_bwd_vs_torch_reference(self, s, b, n, C, with_bias): + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + def _make_inputs(): + hr = hr_data.clone().requires_grad_(True) + orig = orig_data.clone().requires_grad_(True) + hp = hp_data.clone().requires_grad_(True) + x = x_data.clone().requires_grad_(True) + bi = bias_data.clone().requires_grad_(True) if with_bias else None + return hr, orig, hp, x, bi + + hr_f, orig_f, hp_f, x_f, bi_f = _make_inputs() + out_f = native_h_post_bda(hr_f, orig_f, hp_f, x_f, bi_f) + out_f.backward(grad_out) + + hr_r, orig_r, hp_r, x_r, bi_r = _make_inputs() + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + for name, gf, gr in [ + ("h_res", hr_f.grad, hr_r.grad), + ("orig_res", orig_f.grad, orig_r.grad), + ("h_post", hp_f.grad, hp_r.grad), + ("x", x_f.grad, x_r.grad), + ]: + torch.testing.assert_close( + gf, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + bi_f.grad, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on bias" + ) + + +class TestFusedHPostBDA: + """Public fused h_post_bda dispatch/fallback plus numerical correctness.""" + + def test_forward_uses_h_res_transpose(self): + from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda + + h_res, orig, h_post, x, expected = _h_post_bda_transpose_case() + out = fused_h_post_bda(h_res, orig, h_post, x, bias=None) + + torch.testing.assert_close(out, expected, atol=0.0, rtol=0.0) + + @pytest.mark.flaky_in_dev + @_require_cutile + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256)]) + def test_fwd_bwd_vs_reference(self, s, b, n, C, with_bias): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_h_post_bda + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + def _make_inputs(): + hr = hr_data.clone().requires_grad_(True) + orig = orig_data.clone().requires_grad_(True) + hp = hp_data.clone().requires_grad_(True) + x = x_data.clone().requires_grad_(True) + bi = bias_data.clone().requires_grad_(True) if with_bias else None + return hr, orig, hp, x, bi + + # -- fused path -- + hr_f, orig_f, hp_f, x_f, bi_f = _make_inputs() + out_f = fused_h_post_bda(hr_f, orig_f, hp_f, x_f, bi_f) + out_f.backward(grad_out) + + # -- reference path -- + hr_r, orig_r, hp_r, x_r, bi_r = _make_inputs() + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + for name, gf, gr in [ + ("h_res", hr_f.grad, hr_r.grad), + ("orig_res", orig_f.grad, orig_r.grad), + ("h_post", hp_f.grad, hp_r.grad), + ("x", x_f.grad, x_r.grad), + ]: + torch.testing.assert_close( + gf, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + bi_f.grad, bi_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on bias" + ) + + +class TestTritonHPostBDA: + """Tests for Triton h_post_bda kernels against PyTorch reference.""" + + @_require_triton + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) + def test_fwd_vs_reference(self, s, b, n, C, with_bias): + """Triton hpb forward output must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_fwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + + out_t = _triton_h_post_bda_fwd(hr_data, orig_data, hp_data, x_data, bias_data) + out_r = _ref_h_post_bda(hr_data, orig_data, hp_data, x_data, bias_data) + + torch.testing.assert_close(out_t, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + + @_require_triton + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize( + "s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096), (128, 1, 8, 7168)] + ) + def test_bwd_vs_reference(self, s, b, n, C, with_bias): + """Triton hpb backward grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import _triton_h_post_bda_bwd + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + # -- Triton backward -- + g_hr_t, g_res_t, g_hp_t, g_x_t, g_bias_t = _triton_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data + ) + + # -- Reference backward via autograd -- + hr_r = hr_data.clone().requires_grad_(True) + orig_r = orig_data.clone().requires_grad_(True) + hp_r = hp_data.clone().requires_grad_(True) + x_r = x_data.clone().requires_grad_(True) + bi_r = bias_data.clone().requires_grad_(True) if with_bias else None + out_r = _ref_h_post_bda(hr_r, orig_r, hp_r, x_r, bi_r) + out_r.backward(grad_out) + + for name, gt, gr in [ + ("h_res", g_hr_t, hr_r.grad), + ("orig_res", g_res_t, orig_r.grad), + ("h_post", g_hp_t, hp_r.grad), + ("x", g_x_t, x_r.grad), + ]: + torch.testing.assert_close( + gt, gr, atol=BWD_ATOL, rtol=BWD_RTOL, msg=f"Triton backward mismatch on {name}" + ) + if with_bias: + torch.testing.assert_close( + g_bias_t, + bi_r.grad, + atol=BWD_ATOL, + rtol=BWD_RTOL, + msg="Triton backward mismatch on bias", + ) + + @_require_triton + @_require_cutile + @pytest.mark.parametrize("with_bias", [True, False]) + @pytest.mark.parametrize("s,b,n,C", [(2, 4, 4, 1024), (1, 2, 2, 256), (64, 8, 4, 4096)]) + def test_triton_vs_cutile(self, s, b, n, C, with_bias): + """Triton and cuTile backward must produce identical results.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_post_bda_bwd, + _triton_h_post_bda_bwd, + ) + + _info() + hr_data = _rand(s, b, n, n) + orig_data = _rand(s, b, n, C) + hp_data = _rand(s, b, n) + x_data = _rand(s, b, C) + bias_data = _rand(C) if with_bias else None + grad_out = _rand(s, b, n, C) + + triton_out = _triton_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data + ) + cutile_out = _cutile_h_post_bda_bwd( + grad_out, hr_data, orig_data, hp_data, x_data, bias_data + ) + + for i, name in enumerate(["h_res", "orig_res", "h_post", "x", "bias"]): + if triton_out[i] is None: + continue + torch.testing.assert_close( + triton_out[i], + cutile_out[i], + atol=BWD_ATOL, + rtol=BWD_RTOL, + msg=f"Triton vs cuTile mismatch on {name}", + ) + + +class TestTritonHPostBDABwdE2EDebug: + """Debug: run E2E forward, then compare cuTile vs Triton backward per-output.""" + + @_require_triton + @_require_cutile + def test_e2e_inputs_no_nan(self): + """Feed actual E2E backward inputs to Triton kernel and check for NaN.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_post_bda_bwd, + _triton_h_post_bda_bwd, + fused_h_aggregate, + fused_h_post_bda, + fused_sinkhorn, + ) + + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + # Run E2E forward to produce realistic backward inputs + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + x_2d = hs.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post_val = h[..., n : 2 * n].sigmoid() * 2 + h_res = fused_sinkhorn(h[..., 2 * n :].view(s, b, n, n), sinkhorn_iters, eps) + _ = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post_val, layer_out_data, layer_bias_data + ) + go = torch.ones_like(output) + + # Capture inputs (detach from graph) + hr = h_res.detach() + orig = hs.view(s, b, n, C).detach() + hp = h_post_val.detach() + x = layer_out_data.detach() + bias = layer_bias_data.detach() + + # Compare cuTile vs Triton per-output + ct_out = _cutile_h_post_bda_bwd(go, hr, orig, hp, x, bias) + tr_out = _triton_h_post_bda_bwd(go, hr, orig, hp, x, bias) + + names = ["g_hr", "g_res", "g_hp", "g_x", "g_bias"] + for name, ct_t, tr_t in zip(names, ct_out, tr_out): + if tr_t is None: + continue + assert not tr_t.isnan().any(), f"Triton {name} has NaN" + assert not tr_t.isinf().any(), f"Triton {name} has Inf" + torch.testing.assert_close( + tr_t, + ct_t, + atol=BWD_ATOL, + rtol=BWD_RTOL, + msg=f"Triton vs cuTile mismatch on {name} (E2E inputs)", + ) + + +# ============================================================================ +# Triton: Sinkhorn +# ============================================================================ + + +class TestTritonSinkhorn: + @_require_triton + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5), (1, 1, 2, 10), (8, 4, 4, 20)]) + def test_fwd_bwd_vs_reference(self, s, b, n, iters): + from megatron.core.fusions.fused_mhc_kernels import triton_fused_sinkhorn + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + inp_f = data.clone().requires_grad_(True) + out_f = triton_fused_sinkhorn(inp_f, iters, eps) + out_f.backward(grad_out) + + inp_r = data.clone().requires_grad_(True) + out_r = _ref_sinkhorn(inp_r, iters, eps) + out_r.backward(grad_out) + + torch.testing.assert_close(out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(inp_f.grad, inp_r.grad, atol=BWD_ATOL, rtol=BWD_RTOL) + + @_require_triton + @_require_cutile + @pytest.mark.parametrize("s,b,n,iters", [(2, 4, 4, 5)]) + def test_triton_vs_cutile(self, s, b, n, iters): + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_sinkhorn_bwd, + _cutile_sinkhorn_fwd, + triton_fused_sinkhorn, + ) + + eps = 1e-6 + data = _rand(s, b, n, n) + grad_out = _rand(s, b, n, n) + + inp_t = data.clone().requires_grad_(True) + out_t = triton_fused_sinkhorn(inp_t, iters, eps) + out_t.backward(grad_out) + + out_c, M_init = _cutile_sinkhorn_fwd(data.clone(), iters, eps) + grad_c = _cutile_sinkhorn_bwd(grad_out, M_init, iters, eps) + + torch.testing.assert_close(out_t, out_c, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(inp_t.grad, grad_c, atol=BWD_ATOL, rtol=BWD_RTOL) + + +# ============================================================================ +# Proj RMS +# ============================================================================ + + +class TestNativeProjRms: + """Tests for native_proj_rms.""" + + @pytest.mark.parametrize("M,N,K", [(256, 20, 4096), (64, 8, 512)]) + def test_fwd_bwd_vs_torch_reference(self, M, N, K): + _info() + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + grad_proj = _rand(M, N) + grad_r = _rand(M, 1) + + xf = x_data.clone().requires_grad_(True) + wf = w_data.clone().requires_grad_(True) + proj_f, r_f = native_proj_rms(xf, wf, eps) + (proj_f * grad_proj + r_f * grad_r).sum().backward() + + xr = x_data.clone().requires_grad_(True) + wr = w_data.clone().requires_grad_(True) + proj_r, r_r = _ref_proj_rms(xr, wr, eps) + (proj_r * grad_proj + r_r * grad_r).sum().backward() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + + +# ============================================================================ +# Proj RMS + Compute H (fused) +# ============================================================================ + + +class TestFusedProjRmsComputeH: + """Public fused proj_rms_compute_h dispatch/fallback plus numerical correctness.""" + + @pytest.mark.parametrize("M,n,K", [(256, 4, 4096), (64, 2, 512), (128, 4, 2048)]) + def test_fwd_bwd_vs_reference(self, M, n, K): + """E2E: public fused fwd output and bwd grads must match the PyTorch reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms_compute_h + + _info() + N = n * n + 2 * n + eps = 1e-6 + x_data = _rand(M, K) + w_data = _rand(N, K) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + grad_y = _rand(M, N) + grad_h_pre = grad_y[:, :n] + grad_h_post = grad_y[:, n : 2 * n] + grad_h_res = grad_y[:, 2 * n :] + grad_r = _rand(M, 1) + + def _make_inputs(): + return ( + x_data.clone().requires_grad_(True), + w_data.clone().requires_grad_(True), + ap_data.clone().requires_grad_(True), + apo_data.clone().requires_grad_(True), + ar_data.clone().requires_grad_(True), + bias_data.clone().requires_grad_(True), + ) + + # -- fused path -- + xf, wf, apf, apof, arf, bf = _make_inputs() + h_pre_f, h_post_f, h_res_f, r_f = fused_proj_rms_compute_h( + xf, wf, apf, apof, arf, bf, n, eps + ) + loss_f = ( + (h_pre_f * grad_h_pre).sum() + + (h_post_f * grad_h_post).sum() + + (h_res_f * grad_h_res).sum() + + (r_f * grad_r).sum() + ) + loss_f.backward() + + # -- reference path -- + xr, wr, apr, apor, arr, br = _make_inputs() + h_pre_r, h_post_r, h_res_r, r_r = _ref_proj_rms_compute_h( + xr, wr, apr, apor, arr, br, n, eps + ) + loss_r = ( + (h_pre_r * grad_h_pre).sum() + + (h_post_r * grad_h_post).sum() + + (h_res_r * grad_h_res).sum() + + (r_r * grad_r).sum() + ) + loss_r.backward() + + torch.testing.assert_close( + h_pre_f, h_pre_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_pre mismatch" + ) + torch.testing.assert_close( + h_post_f, h_post_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post mismatch" + ) + torch.testing.assert_close( + h_res_f, h_res_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_res mismatch" + ) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="r mismatch") + torch.testing.assert_close( + xf.grad, xr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on x" + ) + torch.testing.assert_close( + wf.grad, wr.grad, atol=BWD_ATOL, rtol=BWD_RTOL, msg="backward mismatch on weight" + ) + _assert_cosine_similar( + apf.grad, apr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_pre" + ) + _assert_cosine_similar( + apof.grad, apor.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_post" + ) + _assert_cosine_similar( + arf.grad, arr.grad, COSINE_SIM_THRESH, msg="backward mismatch on alpha_res" + ) + _assert_cosine_similar(bf.grad, br.grad, COSINE_SIM_THRESH, msg="backward mismatch on bias") + + +# ============================================================================ +# End-to-end pipeline (all four kernels chained) +# ============================================================================ + + +class TestEndToEndNative: + """Full mHC pipeline using native modules. + + proj_rms -> compute_h -> sinkhorn -> aggregate -> h_post_bda. + Compares the native modules against inline PyTorch reference. + """ + + def test_full_pipeline_fwd_bwd(self): + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_native_modules(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = native_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = native_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = native_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_inline_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_m, agg_m, grad_m = _run_native_modules() + out_r, agg_r, grad_r = _run_inline_ref() + + torch.testing.assert_close( + agg_m, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated output mismatch" + ) + torch.testing.assert_close( + out_m, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_m, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) + + +class TestEndToEndFused: + """Full mHC pipeline using the public fused API.""" + + @_require_cutile + def test_full_pipeline_fwd_bwd(self): + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_sinkhorn, + ) + + _info() + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = fused_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return proj.detach(), r.detach(), output.detach(), aggregated.detach(), hs.grad.clone() + + proj_f, r_f, out_f, agg_f, grad_f = _run_fused() + proj_r, r_r, out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close(proj_f, proj_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(r_f, r_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close(agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_f, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward)" + ) + + def test_full_pipeline_fused_compute_h(self): + """E2E: fused proj_rms_compute_h replaces separate proj_rms + compute_h.""" + from megatron.core.fusions.fused_mhc_kernels import ( + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms_compute_h, + fused_sinkhorn, + ) + + _info() + s, b, n, C = 8, 4, 4, 1024 + N = n * n + 2 * n + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(N, n * C) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_compute_h(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_compute_h() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch" + ) + torch.testing.assert_close( + out_f, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch" + ) + _assert_cosine_similar( + grad_f, + grad_r, + COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused compute_h)", + ) + + +# ============================================================================ +# fused_add_3 kernel tests +# ============================================================================ + + +class TestFusedAdd3: + """Tests for fused_add_3 (torch.compile backend, no cuTile dependency).""" + + def test_fused_add_3_forward_bf16(self): + """fused_add_3 matches a + b + c for bf16 tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + b = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + c = torch.randn(128, 256, dtype=DTYPE, device=DEVICE) + result = fused_add_3(a, b, c) + expected = (a.float() + b.float() + c.float()).to(DTYPE) + torch.testing.assert_close(result, expected, atol=FWD_ATOL, rtol=FWD_RTOL) + + def test_fused_add_3_forward_fp32(self): + """fused_add_3 matches a + b + c for fp32 tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + b = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + c = torch.randn(128, 256, dtype=torch.float32, device=DEVICE) + result = fused_add_3(a, b, c) + expected = a + b + c + torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5) + + def test_fused_add_3_large_tensor(self): + """fused_add_3 handles large tensors.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + b = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + c = torch.randn(8192, 4096, dtype=DTYPE, device=DEVICE) + result = fused_add_3(a, b, c) + expected = (a.float() + b.float() + c.float()).to(DTYPE) + torch.testing.assert_close(result, expected, atol=FWD_ATOL, rtol=FWD_RTOL) + + def test_fused_add_3_gradient(self): + """fused_add_3 produces correct gradients.""" + from megatron.core.fusions.fused_mhc_kernels import fused_add_3 + + a = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + b = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + c = torch.randn(64, 128, dtype=torch.float32, device=DEVICE, requires_grad=True) + result = fused_add_3(a, b, c) + result.sum().backward() + torch.testing.assert_close(a.grad, torch.ones_like(a)) + torch.testing.assert_close(b.grad, torch.ones_like(b)) + torch.testing.assert_close(c.grad, torch.ones_like(c)) + + +# ============================================================================ +# End-to-end pipeline with BroadcastTensorFused +# ============================================================================ + + +class TestEndToEndNativeBroadcast: + """Full mHC pipeline using native modules + BroadcastTensorFused. + + Same pipeline as TestEndToEndNative but hidden_states is split via + BroadcastTensorFused so each consumer (proj_rms/compute_h, aggregate, + h_post_bda) gets a distinct autograd graph node. Verifies gradient + correctness versus the inline reference that uses the tensor directly. + """ + + def test_full_pipeline_fwd_bwd(self): + from megatron.core.transformer.hyper_connection import ( + BroadcastTensorFused, + native_fused_add_3, + ) + + _info() + s, b, n, C = 2, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + # Split via BroadcastTensorFused + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, native_fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = native_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = native_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = native_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_inline_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_b, agg_b, grad_b = _run_broadcast() + out_r, agg_r, grad_r = _run_inline_ref() + + torch.testing.assert_close( + agg_b, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (broadcast)" + ) + torch.testing.assert_close( + out_b, out_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="h_post_bda output mismatch (broadcast)" + ) + _assert_cosine_similar( + grad_b, grad_r, COSINE_SIM_THRESH, msg="hidden_states grad (E2E backward, broadcast)" + ) + + +class TestEndToEndFusedBroadcast: + """Full mHC pipeline through public fused dispatch + BroadcastTensorFused.""" + + def test_full_pipeline_fwd_bwd(self): + from megatron.core.fusions.fused_mhc_kernels import ( + fused_add_3, + fused_h_aggregate, + fused_h_post_bda, + fused_sinkhorn, + ) + from megatron.core.transformer.hyper_connection import BroadcastTensorFused + + _info() + s, b, n, C = 8, 4, 4, 1024 + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(n * n + 2 * n, n * C) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + proj, r = native_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = fused_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + proj, r = _ref_proj_rms(x_2d, w, eps) + proj = proj.view(s, b, -1) + r = r.view(s, b, 1) + + h = r * proj + h_pre = h[..., :n].sigmoid() + h_post = h[..., n : 2 * n].sigmoid() * 2 + h_res_logits = h[..., 2 * n :] + h_res = _ref_sinkhorn(h_res_logits.view(s, b, n, n), sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_broadcast() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, agg_r, atol=FWD_ATOL, rtol=FWD_RTOL, msg="aggregated mismatch (fused broadcast)" + ) + torch.testing.assert_close( + out_f, + out_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, + msg="h_post_bda output mismatch (fused broadcast)", + ) + _assert_cosine_similar( + grad_f, + grad_r, + COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused broadcast)", + ) + + def test_full_pipeline_fused_compute_h_broadcast(self): + """E2E: fused proj_rms_compute_h + BroadcastTensorFused.""" + from megatron.core.fusions.fused_mhc_kernels import ( + fused_add_3, + fused_h_aggregate, + fused_h_post_bda, + fused_proj_rms_compute_h, + fused_sinkhorn, + ) + from megatron.core.transformer.hyper_connection import BroadcastTensorFused + + _info() + s, b, n, C = 8, 4, 4, 1024 + N = n * n + 2 * n + eps = 1e-6 + sinkhorn_iters = 5 + + hs_data = _rand(s, b, n * C) + w_data = _rand(N, n * C) + ap_data = _rand(1) + apo_data = _rand(1) + ar_data = _rand(1) + bias_data = _rand(N) + layer_out_data = _rand(s, b, C) + layer_bias_data = _rand(C) + + def _run_fused_compute_h_broadcast(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + hs_map, hs_agg, hs_res = BroadcastTensorFused.apply(hs, fused_add_3) + + x_2d = hs_map.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = fused_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = fused_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = fused_h_aggregate(hs_agg.view(s, b, n, C), h_pre) + + output = fused_h_post_bda( + h_res, hs_res.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + def _run_ref(): + hs = hs_data.clone().requires_grad_(True) + w = w_data.clone().requires_grad_(True) + ap = ap_data.clone().requires_grad_(True) + apo = apo_data.clone().requires_grad_(True) + ar = ar_data.clone().requires_grad_(True) + bias_p = bias_data.clone().requires_grad_(True) + + x_2d = hs.reshape(s * b, n * C) + h_pre, h_post, h_res_logits, _ = _ref_proj_rms_compute_h( + x_2d, w, ap, apo, ar, bias_p, n, eps + ) + + h_pre = h_pre.view(s, b, n) + h_post = h_post.view(s, b, n) + h_res_logits = h_res_logits.view(s, b, n, n) + h_res = _ref_sinkhorn(h_res_logits, sinkhorn_iters, eps) + + aggregated = _ref_h_aggregate(hs.view(s, b, n, C), h_pre) + + output = _ref_h_post_bda( + h_res, hs.view(s, b, n, C), h_post, layer_out_data, layer_bias_data + ) + + loss = output.sum() + aggregated.sum() + loss.backward() + return output.detach(), aggregated.detach(), hs.grad.clone() + + out_f, agg_f, grad_f = _run_fused_compute_h_broadcast() + out_r, agg_r, grad_r = _run_ref() + + torch.testing.assert_close( + agg_f, + agg_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, + msg="aggregated mismatch (fused compute_h broadcast)", + ) + torch.testing.assert_close( + out_f, + out_r, + atol=FWD_ATOL, + rtol=FWD_RTOL, + msg="h_post_bda output mismatch (fused compute_h broadcast)", + ) + _assert_cosine_similar( + grad_f, + grad_r, + COSINE_SIM_THRESH, + msg="hidden_states grad (E2E backward, fused compute_h broadcast)", + ) + + +class TestFusedProjRmsComputeHKeepFp32: + """Regression: the mHC mapping must stay fp32-accurate with bf16 activations. + + HyperConnectionModule marks mapping_proj.weight / alpha_* / bias as + keep_in_fp32 and the unfused path upcasts the activations to fp32, so the + fused path must not silently degrade the mapping to the activation dtype. + Tolerances here are ~15x tighter than the module-level FWD_ATOL/RTOL: the + bug this guards against (bf16 split-K partials and bf16 mapping outputs) + cost 170x accuracy while staying well inside the loose tolerances. + """ + + # Native fp32 reference achieves ~8e-6 rms; the bf16-output bug gave ~1.5e-3. + MAPPING_RMS_TOL = 1e-4 + # r is a pure reduction: native reaches ~5e-8, the bug gave ~2.7e-3. + R_RMS_TOL = 1e-5 + + @staticmethod + def _rms_rel(actual: Tensor, ref64: Tensor) -> float: + a = actual.detach().to(torch.float64) + r = ref64.detach().to(torch.float64) + return ((a - r).pow(2).mean().sqrt() / r.abs().max().clamp_min(1e-30)).item() + + @pytest.mark.parametrize("M,n,hidden", [(256, 4, 4096), (128, 4, 1024)]) + def test_bf16_activations_fp32_params(self, M, n, hidden): + """Fused mapping with production dtypes must match an fp64 reference.""" + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms_compute_h + + _info() + K = n * hidden + N = n * n + 2 * n + eps = compute_h_eps = 1e-6 + + # Activations arrive in bf16; the mapping parameters are keep_in_fp32. + x = torch.randn(M, K, device=DEVICE, dtype=torch.float32).to(torch.bfloat16) + w = torch.randn(N, K, device=DEVICE, dtype=torch.float32) / math.sqrt(K) + alpha_pre = torch.full((1,), 0.1, device=DEVICE, dtype=torch.float32) + alpha_post = torch.full((1,), 0.1, device=DEVICE, dtype=torch.float32) + alpha_res = torch.full((1,), 0.1, device=DEVICE, dtype=torch.float32) + bias = torch.randn(N, device=DEVICE, dtype=torch.float32) * 0.1 + + h_pre, h_post, h_res, r = fused_proj_rms_compute_h( + x, w, alpha_pre, alpha_post, alpha_res, bias, n, eps, compute_h_eps + ) + + # fp64 reference from the same input values. + x64, w64 = x.to(torch.float64), w.to(torch.float64) + proj64 = x64 @ w64.t() + r64 = x64.norm(dim=-1, keepdim=True) / math.sqrt(K) + alpha64 = torch.cat( + [ + alpha_pre.to(torch.float64).expand(n), + alpha_post.to(torch.float64).expand(n), + alpha_res.to(torch.float64).expand(N - 2 * n), + ], + dim=-1, + ) + h64 = proj64 * alpha64.unsqueeze(0) / (r64 + eps) + bias.to(torch.float64).unsqueeze(0) + + assert self._rms_rel(h_pre, h64[..., :n].sigmoid() + compute_h_eps) < self.MAPPING_RMS_TOL + assert self._rms_rel(h_post, h64[..., n : 2 * n].sigmoid() * 2) < self.MAPPING_RMS_TOL + assert self._rms_rel(h_res, h64[..., 2 * n :]) < self.MAPPING_RMS_TOL + assert self._rms_rel(r, r64) < self.R_RMS_TOL + + def test_native_fallback_accepts_mixed_dtypes(self): + """The native fallback must run with bf16 activations x fp32 parameters.""" + from megatron.core.fusions.fused_mhc_kernels import _torch_proj_rms_compute_h + + M, n, K = 64, 4, 512 + N = n * n + 2 * n + x = torch.randn(M, K, device=DEVICE, dtype=torch.bfloat16) + w = torch.randn(N, K, device=DEVICE, dtype=torch.float32) / math.sqrt(K) + one = torch.full((1,), 0.1, device=DEVICE, dtype=torch.float32) + bias = torch.zeros(N, device=DEVICE, dtype=torch.float32) + + h_pre, h_post, h_res, r = _torch_proj_rms_compute_h(x, w, one, one, one, bias, n, 1e-6) + for t in (h_pre, h_post, h_res, r): + assert torch.isfinite(t).all() + + @_require_cutile + def test_no_fp32_activation_copy(self): + """cuTile must consume the bf16 activations, not an fp32 copy of them. + + The kernels load and cast each tile independently, so normalizing the + activation dtype ahead of the launch would only cost memory. + """ + from megatron.core.fusions.fused_mhc_kernels import fused_proj_rms_compute_h + + _info() + M, n, hidden = 4096, 4, 4096 + K, N = n * hidden, n * n + 2 * n + x = torch.randn(M, K, device=DEVICE, dtype=torch.bfloat16) + w = torch.randn(N, K, device=DEVICE, dtype=torch.float32) / math.sqrt(K) + one = torch.full((1,), 0.1, device=DEVICE, dtype=torch.float32) + bias = torch.zeros(N, device=DEVICE, dtype=torch.float32) + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_allocated() + fused_proj_rms_compute_h(x, w, one, one, one, bias, n, 1e-6) + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - before + + fp32_activation_copy = M * K * 4 + assert peak < fp32_activation_copy // 2, ( + f"peak allocation {peak} suggests the activations were upcast " + f"(an fp32 copy of x is {fp32_activation_copy} bytes)" + ) + + def test_public_entry_point_native_branch(self, monkeypatch): + """The public op must also run natively — this is the forced-native path. + + MHC_FORCE_BACKEND=native and an auto run on a cuTile-less container both + land here, and the module always calls the public op when + use_fused_mhc=True, so this branch is what a real training job hits. + """ + from megatron.core.fusions import fused_mhc_kernels as fused_mod + + _info() + monkeypatch.setattr(fused_mod, "is_cutile_available", lambda: False) + + M, n, K = 64, 4, 512 + N = n * n + 2 * n + x = torch.randn(M, K, device=DEVICE, dtype=torch.bfloat16) + w = torch.randn(N, K, device=DEVICE, dtype=torch.float32) / math.sqrt(K) + one = torch.full((1,), 0.1, device=DEVICE, dtype=torch.float32) + bias = torch.zeros(N, device=DEVICE, dtype=torch.float32) + + outs = fused_mod.fused_proj_rms_compute_h(x, w, one, one, one, bias, n, 1e-6) + for t in outs: + assert torch.isfinite(t).all() + + +class TestCutileHAggregateBackwardReduction: + """Regression: h_aggregate backward is the only other cuTile-only op. + + `grad_h` is a reduction over the hidden dimension; evaluating the product + in bf16 before the fp32 accumulate made it ~2.5x noisier than the torch + reference at production widths. + """ + + @_require_cutile + @pytest.mark.parametrize("tokens,n,C", [(8192, 4, 1536), (2048, 4, 4096)]) + def test_grad_h_reduction_not_worse_than_torch(self, tokens, n, C): + """cuTile grad_h must be at least as accurate as the torch reference.""" + from megatron.core.fusions.fused_mhc_kernels import ( + _cutile_h_aggregate_bwd, + _torch_h_aggregate_bwd, + ) + + _info() + s, b = 1, tokens + x = torch.randn(s, b, n, C, device=DEVICE, dtype=torch.float32).to(DTYPE) + h_pre = (torch.rand(s, b, n, device=DEVICE, dtype=torch.float32) + 0.5).to(DTYPE) + go = torch.randn(s, b, C, device=DEVICE, dtype=torch.float32).to(DTYPE) + + goe = go.double().unsqueeze(2) + ref_gx = goe * h_pre.double().unsqueeze(-1) + ref_gh = (goe * x.double()).sum(-1) + + def rel(a: Tensor, r: Tensor) -> float: + a, r = a.double().flatten(), r.double().flatten() + return ((a - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() + + t_gx, t_gh = _torch_h_aggregate_bwd(go, x, h_pre) + c_gx, c_gh = _cutile_h_aggregate_bwd(go, x, h_pre) + + # grad_x is a pure elementwise product — it must stay bit-identical. + assert torch.equal(c_gx, t_gx) + # A bf16 product before the fp32 accumulate showed up here as ~2.5x. + assert rel(c_gh, ref_gh) <= rel(t_gh, ref_gh) * 1.1 diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 810f48092ee..a63edaba6a3 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -16,12 +16,16 @@ try: from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_apply_mla_rope_for_kv, fused_apply_mla_rope_for_q, + fused_mla_rope_inplace, + fused_mla_rope_kv_split, + fused_mla_rope_out_of_place, ) -except: - fused_apply_mla_rope_for_kv = None +except Exception: fused_apply_mla_rope_for_q = None + fused_mla_rope_inplace = None + fused_mla_rope_kv_split = None + fused_mla_rope_out_of_place = None def dtype_tols(dtype): @@ -54,7 +58,9 @@ def test_packed_freqs_returns_offset_mapped_output_for_context_parallel(self): t = torch.randn(4, 2, 8) freqs = torch.randn(8, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[0:1], freqs[3:4], freqs[4:5], freqs[7:8]], dim=0) expected = rope_utils_module._apply_rotary_pos_emb_bshd( @@ -69,7 +75,9 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se t = torch.randn(4, 2, 8) freqs = torch.randn(4, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[1:2], freqs[2:3]], dim=0) expected_slices = [] @@ -83,9 +91,64 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se torch.testing.assert_close(out, expected) + def test_missing_max_seqlen_preserves_legacy_packed_freq_mapping(self): + cp_group = FakeCPGroup(size=2, rank=0) + cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32) + t = torch.randn(4, 2, 8) + freqs = torch.randn(8, 1, 1, 8) -def _test_fused_apply_mla_rope_for_q(input_format): - assert fused_apply_mla_rope_for_q is not None + legacy_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group + ) + explicit_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) + + torch.testing.assert_close(legacy_out, explicit_out) + + def test_shared_max_seqlen_maps_asymmetric_query_sequences_from_zero(self): + cp_group = FakeCPGroup(size=1, rank=0) + cu_seqlens_q = torch.tensor([0, 3, 6], dtype=torch.int32) + t = torch.randn(6, 2, 8) + freqs = torch.randn(4, 1, 1, 8) + + max_seqlen_q = 3 + max_seqlen_kv = freqs.size(0) + assert max_seqlen_q < max_seqlen_kv < t.size(0) + combined_max_seqlen = max(max_seqlen_q, max_seqlen_kv) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens_q, freqs, cp_group=cp_group, max_seqlen=combined_max_seqlen + ) + compatibility_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens_q, freqs, cp_group=cp_group + ) + + expected_freqs = freqs[torch.tensor([0, 1, 2, 0, 1, 2])] + expected = rope_utils_module._apply_rotary_pos_emb_bshd( + t.unsqueeze(1), expected_freqs + ).squeeze(1) + + torch.testing.assert_close(out, expected) + torch.testing.assert_close(out, compatibility_out) + + +class _SaveOutputForBackward(torch.autograd.Function): + """Minimal stand-in for a kernel whose backward consumes its output.""" + + @staticmethod + def forward(ctx, tensor): + output = tensor.clone() + ctx.save_for_backward(output) + return output + + @staticmethod + def backward(ctx, _grad_output): + (saved_output,) = ctx.saved_tensors + return saved_output + + +def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleaving=False): + assert fused_mla_rope_inplace is not None num_heads = 32 q_dim = 128 emb_dim = 64 @@ -97,6 +160,7 @@ def _test_fused_apply_mla_rope_for_q(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -142,15 +206,25 @@ def _test_fused_apply_mla_rope_for_q(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + inverse=inverse, + mla_output_remove_interleaving=remove_interleaving, ) pytorch_output = torch.concat([no_pe, pe_output], dim=-1) pytorch_output.backward(pytorch_bwd_input, retain_graph=True) - fused_output = fused_apply_mla_rope_for_q( - fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens + fused_output = fused_mla_rope_inplace( + fused_fwd_input, + cos, + sin, + q_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=inverse, + remove_interleaving=remove_interleaving, ) fused_output.backward(fused_bwd_input, retain_graph=True) @@ -169,8 +243,8 @@ def _test_fused_apply_mla_rope_for_q(input_format): ) -def _test_fused_apply_mla_rope_for_kv(input_format): - assert fused_apply_mla_rope_for_kv is not None +def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): + assert fused_mla_rope_kv_split is not None num_heads = 32 k_dim = 128 v_dim = 128 @@ -183,6 +257,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -241,9 +316,11 @@ def _test_fused_apply_mla_rope_for_kv(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + mla_output_remove_interleaving=remove_interleaving, ) if input_format == "sbhd": pe_output = pe_output.expand(-1, -1, num_heads, -1) @@ -255,7 +332,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): (pytorch_k_output, pytorch_v_output), (pytorch_bwd_k_input, pytorch_bwd_v_input) ) - fused_k_output, fused_v_output = fused_apply_mla_rope_for_kv( + fused_k_output, fused_v_output = fused_mla_rope_kv_split( fused_fwd_kv_input, fused_fwd_emb_input, cos, @@ -264,6 +341,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): k_dim, v_dim, cu_seqlens_kv=cu_seqlens, + remove_interleaving=remove_interleaving, ) torch.autograd.backward( (fused_k_output, fused_v_output), (fused_bwd_k_input, fused_bwd_v_input) @@ -301,13 +379,136 @@ def _test_fused_apply_mla_rope_for_kv(input_format): @pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) -class TestFusedApplyMLARope: +class TestFusedMLARope: @pytest.mark.flaky_in_dev - def test_forward_backward_for_q(self, input_format): - _test_fused_apply_mla_rope_for_q(input_format) + @pytest.mark.parametrize("inverse", [False, True]) + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_inplace_forward_backward(self, input_format, inverse, remove_interleaving): + _test_fused_mla_rope_inplace( + input_format, inverse=inverse, remove_interleaving=remove_interleaving + ) + + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_kv_split_forward_backward(self, input_format, remove_interleaving): + _test_fused_mla_rope_kv_split(input_format, remove_interleaving=remove_interleaving) + + +@pytest.mark.experimental +@pytest.mark.internal +@pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("input_format", ["sbhd", "thd"]) +def test_out_of_place_inverse_rope_preserves_upstream_saved_output(input_format): + """Post-attention inverse RoPE must not overwrite an output saved for backward.""" + assert fused_mla_rope_out_of_place is not None + seqlen = 32 + batch_size = 1 + num_heads = 2 + nope_dim = 16 + emb_dim = 64 + dtype = torch.bfloat16 + + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=seqlen) + freqs, mscale = yarn_rope(seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + + if input_format == "sbhd": + shape = (seqlen, batch_size, num_heads, nope_dim + emb_dim) + cu_seqlens = None + else: + shape = (2 * seqlen, num_heads, nope_dim + emb_dim) + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device="cuda") + + unsafe_source = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + unsafe_attention_output = _SaveOutputForBackward.apply(unsafe_source) + unsafe_reference = unsafe_attention_output.detach().clone() + unsafe_inverse_output = fused_mla_rope_inplace( + unsafe_attention_output, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + + assert unsafe_inverse_output.data_ptr() == unsafe_attention_output.data_ptr() + assert not torch.equal(unsafe_attention_output, unsafe_reference) + + source = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + attention_output = _SaveOutputForBackward.apply(source) + saved_reference = attention_output.detach().clone() + + inverse_output = fused_mla_rope_out_of_place( + attention_output, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + expected_inverse_output = fused_mla_rope_inplace( + saved_reference.clone(), + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + + assert inverse_output.data_ptr() != attention_output.data_ptr() + torch.testing.assert_close(attention_output, saved_reference, rtol=0, atol=0) + torch.testing.assert_close(inverse_output, expected_inverse_output, rtol=0, atol=0) + + inverse_output.backward(torch.randn_like(inverse_output).contiguous()) + torch.testing.assert_close(source.grad, saved_reference, rtol=0, atol=0) + + +@pytest.mark.experimental +@pytest.mark.internal +@pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("input_format", ["sbhd", "thd"]) +def test_legacy_query_api_remains_in_place(input_format): + """The legacy API keeps its original mutation behavior and allocation profile.""" + assert fused_apply_mla_rope_for_q is not None + seqlen = 32 + batch_size = 1 + num_heads = 2 + nope_dim = 16 + emb_dim = 64 + dtype = torch.bfloat16 + + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=seqlen) + freqs, mscale = yarn_rope(seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + + if input_format == "sbhd": + shape = (seqlen, batch_size, num_heads, nope_dim + emb_dim) + cu_seqlens = None + else: + shape = (2 * seqlen, num_heads, nope_dim + emb_dim) + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device="cuda") + + query = torch.randn(shape, dtype=dtype, device="cuda") + reference = query.clone() + expected = fused_mla_rope_inplace( + reference.clone(), cos, sin, nope_dim, emb_dim, cu_seqlens_q=cu_seqlens + ) + output = fused_apply_mla_rope_for_q( + query, cos, sin, qk_head_dim=nope_dim, emb_dim=emb_dim, cu_seqlens_q=cu_seqlens + ) - def test_forward_backward_for_kv(self, input_format): - _test_fused_apply_mla_rope_for_kv(input_format) + assert output.data_ptr() == query.data_ptr() + assert not torch.equal(query, reference) + torch.testing.assert_close(output, expected, rtol=0, atol=0) class TestApplyRotaryPosEmbMlaFusionConflict: diff --git a/tests/unit_tests/models/test_hybrid_mhc.py b/tests/unit_tests/models/test_hybrid_mhc.py new file mode 100644 index 00000000000..d805581bf25 --- /dev/null +++ b/tests/unit_tests/models/test_hybrid_mhc.py @@ -0,0 +1,395 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.models.hybrid.hybrid_block import ( + HybridStack, + HybridStackSubmodules, + HyperConnectionHybridLayer, +) +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.moe_layer import MoELayer +from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import TransformerLayer +from tests.unit_tests.test_utilities import Utils + + +class _DummyHybridLayer(MegatronModule): + """Same-shape residual layer used to isolate HybridStack mHC plumbing.""" + + def __init__(self, config: TransformerConfig, layer_number: int, **_kwargs): + super().__init__(config=config) + self.layer_number = layer_number + self.proj = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.seen_hidden_shapes = [] + + def forward( + self, + hidden_states, + attention_mask=None, + inference_context=None, + packed_seq_params=None, + **_kwargs, + ): + self.seen_hidden_shapes.append(tuple(hidden_states.shape)) + return hidden_states + 0.125 * self.proj(hidden_states) + + +class _StubTransformerLayer(TransformerLayer): + """Minimal TransformerLayer that exercises only the mHC wrapper guard.""" + + def __init__(self, config: TransformerConfig): + torch.nn.Module.__init__(self) + self.config = config + self.layer_number = 1 + + def _forward_attention(self, *args, **kwargs): + hidden_states = kwargs.get("hidden_states", args[0] if args else None) + return hidden_states, None + + def _forward_mlp(self, hidden_states, *_args, **_kwargs): + return hidden_states + + +def _get_pg_collection() -> ProcessGroupCollection: + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'pp', 'cp']) + + +def _get_dummy_submodules() -> HybridStackSubmodules: + layer_spec = ModuleSpec(module=_DummyHybridLayer) + return HybridStackSubmodules( + mamba_layer=layer_spec, + gdn_layer=layer_spec, + attention_layer=layer_spec, + dsa_layer=layer_spec, + mlp_layer=layer_spec, + moe_layer=layer_spec, + ) + + +def _get_dummy_stack_spec() -> ModuleSpec: + return ModuleSpec( + module=HybridStack, params={"post_layer_norm": False}, submodules=_get_dummy_submodules() + ) + + +def _get_config(num_layers: int, **kwargs) -> TransformerConfig: + return TransformerConfig( + num_layers=num_layers, + hidden_size=32, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=2, + hidden_dropout=0.0, + mhc_sinkhorn_iterations=3, + **kwargs, + ) + + +def _get_stack( + config: TransformerConfig, + num_local_layers: int, + *, + pre_process: bool = True, + post_process: bool = True, + pp_layer_offset: int = 0, +) -> HybridStack: + return HybridStack( + config=config, + submodules=_get_dummy_submodules(), + pre_process=pre_process, + post_process=post_process, + post_layer_norm=False, + layer_type_list=[Symbols.MAMBA] * num_local_layers, + pp_layer_offset=pp_layer_offset, + pg_collection=_get_pg_collection(), + ) + + +def test_mhc_mtp_requires_hybrid_contract(): + config = _get_config(num_layers=1, mtp_num_layers=1) + + with pytest.raises(ValueError, match="requires the HybridModel MTP contract"): + MultiTokenPredictionLayer( + config=config, submodules=object(), layer_number=1, pg_collection=None + ) + + +@pytest.mark.internal +class TestHybridStackMHC: + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_constructor_and_sharded_state(self): + config = _get_config(num_layers=3) + stack = _get_stack(config, num_local_layers=3) + + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in stack.layers) + assert stack.hc_head_fn.shape == ( + config.num_residual_streams, + config.hidden_size * config.num_residual_streams, + ) + state = stack.sharded_state_dict(prefix="decoder.", metadata={}) + for name in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + assert f"decoder.{name}" in state + + @pytest.mark.parametrize( + "recompute_kwargs", + [ + {}, + { + "recompute_granularity": "selective", + "recompute_modules": ["core_attn", "mhc"], + "mhc_recompute_layer_num": 2, + }, + { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + }, + ], + ids=["none", "selective_mhc", "full_uniform"], + ) + def test_forward_backward(self, recompute_kwargs): + config = _get_config(num_layers=3, **recompute_kwargs) + stack = _get_stack(config, num_local_layers=3).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + + output = stack(hidden_states, attention_mask=None) + + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + output.float().sum().backward() + assert hidden_states.grad is not None + for layer in stack.layers: + assert layer.inner_layer.proj.weight.grad is not None + assert layer.hyper_connection.mapping_proj.weight.grad is not None + assert all( + shape == (8, 2, config.hidden_size) + for shape in layer.inner_layer.seen_hidden_shapes + ) + for name in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + assert getattr(stack, name).grad is not None + + def test_fused_bf16_forward_backward(self): + config = _get_config( + num_layers=2, bf16=True, params_dtype=torch.bfloat16, use_fused_mhc=True + ) + stack = _get_stack(config, num_local_layers=2).cuda().bfloat16() + hidden_states = torch.randn( + 8, 2, config.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + + output = stack(hidden_states, attention_mask=None) + + assert output.shape == hidden_states.shape + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output).all() + output.float().sum().backward() + assert hidden_states.grad is not None + assert all( + layer.hyper_connection.mapping_proj.weight.grad is not None for layer in stack.layers + ) + + def test_pipeline_boundary_shapes(self): + config = _get_config(num_layers=2) + first_stage = _get_stack( + config, num_local_layers=1, pre_process=True, post_process=False + ).cuda() + last_stage = _get_stack( + config, num_local_layers=1, pre_process=False, post_process=True, pp_layer_offset=1 + ).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda") + + pipeline_hidden = first_stage(hidden_states, attention_mask=None) + assert pipeline_hidden.shape == (8, 2, config.hidden_size * config.num_residual_streams) + + last_stage.set_input_tensor(pipeline_hidden.detach()) + output = last_stage(hidden_states, attention_mask=None) + assert output.shape == hidden_states.shape + + def test_real_attention_mlp_forward_backward(self): + config = _get_config(num_layers=2) + stack = HybridStack( + config=config, + submodules=hybrid_stack_spec.submodules, + post_layer_norm=False, + layer_type_list=[Symbols.ATTENTION, Symbols.MLP], + pp_layer_offset=0, + pg_collection=_get_pg_collection(), + ).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + + output = stack(hidden_states, attention_mask=None) + + assert output.shape == hidden_states.shape + assert all(isinstance(layer, HyperConnectionHybridLayer) for layer in stack.layers) + assert all(isinstance(layer.inner_layer, TransformerLayer) for layer in stack.layers) + output.float().sum().backward() + assert hidden_states.grad is not None + assert all( + layer.hyper_connection.mapping_proj.weight.grad is not None for layer in stack.layers + ) + + def test_real_moe_raw_branch_forward_backward(self): + config = _get_config( + num_layers=1, + num_moe_experts=2, + moe_ffn_hidden_size=64, + moe_grouped_gemm=True, + add_bias_linear=False, + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'pp', 'cp', 'ep', 'expt_tp', 'tp_ep', 'expt_dp'] + ) + stack = HybridStack( + config=config, + submodules=hybrid_stack_spec.submodules, + post_layer_norm=False, + layer_type_list=[Symbols.MOE], + pp_layer_offset=0, + pg_collection=pg_collection, + ).cuda() + hidden_states = torch.randn(8, 2, config.hidden_size, device="cuda", requires_grad=True) + + output = stack(hidden_states, attention_mask=None) + + wrapped_layer = stack.layers[0] + assert isinstance(wrapped_layer, HyperConnectionHybridLayer) + assert isinstance(wrapped_layer.inner_layer.mlp, MoELayer) + assert output.shape == hidden_states.shape + output.float().sum().backward() + assert hidden_states.grad is not None + assert wrapped_layer.hyper_connection.mapping_proj.weight.grad is not None + assert any(param.grad is not None for param in wrapped_layer.inner_layer.mlp.parameters()) + + def test_hybrid_model_forward_backward(self): + config = _get_config(num_layers=3) + model = HybridModel( + config=config, + hybrid_stack_spec=_get_dummy_stack_spec(), + vocab_size=64, + max_sequence_length=8, + hybrid_layer_pattern="M*-", + parallel_output=False, + ).cuda() + input_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + position_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + + logits = model(input_ids=input_ids, position_ids=position_ids, attention_mask=None) + + assert logits.shape == (2, 8, model.vocab_size) + assert torch.isfinite(logits).all() + logits.float().mean().backward() + assert all(layer.inner_layer.proj.weight.grad is not None for layer in model.decoder.layers) + assert all( + layer.hyper_connection.mapping_proj.weight.grad is not None + for layer in model.decoder.layers + ) + + def test_hybrid_model_mtp_forward_backward(self): + config = _get_config(num_layers=1, mtp_num_layers=1, mtp_loss_scaling_factor=0.1) + model = HybridModel( + config=config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=64, + max_sequence_length=8, + hybrid_layer_pattern="-/-", + parallel_output=True, + ).cuda() + input_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + position_ids = torch.arange(8, dtype=torch.int64, device="cuda").repeat((2, 1)) + + logits = model(input_ids=input_ids, position_ids=position_ids, attention_mask=None) + + assert logits.shape == (2, 8, model.vocab_size) + assert torch.isfinite(logits).all() + assert not any("mtp_model_layer.hc_head_" in name for name, _ in model.named_parameters()) + logits.float().mean().backward() + mtp_params = [param for name, param in model.named_parameters() if name.startswith("mtp.")] + assert mtp_params + assert all(param.grad is not None for param in mtp_params) + + def test_recompute_plan(self): + config = _get_config( + num_layers=3, + recompute_granularity="selective", + recompute_modules=["core_attn", "mhc"], + mhc_recompute_layer_num=2, + ) + stack = _get_stack(config, num_local_layers=3) + + managers, block_ends = stack._build_mhc_recompute_layer_plan(True) + + assert block_ends == [False, True, True] + assert managers[0] is managers[1] + assert managers[1] is not managers[2] + + def test_boundary_bda_skips_recompute_manager(self, monkeypatch): + config = _get_config(num_layers=1) + layer = HyperConnectionHybridLayer( + config=config, layer=_DummyHybridLayer(config, layer_number=1) + ) + hidden_states = torch.randn( + 4, 2, config.hidden_size * config.num_residual_streams, requires_grad=True + ) + manager = type("_FakeManager", (), {})() + manager.is_last_layer_in_recompute_block = True + seen_managers = [] + + def fake_hyper_connection_forward( + hidden_states, mhc_recompute_manager=None, return_residual=False + ): + assert mhc_recompute_manager is manager + assert return_residual + sequence_length, batch_size, _ = hidden_states.shape + n = config.num_residual_streams + hidden_size = config.hidden_size + aggregated = hidden_states.view(sequence_length, batch_size, n, hidden_size).mean(dim=2) + h_res = torch.empty(sequence_length, batch_size, n, n) + h_post = torch.empty(sequence_length, batch_size, n) + return aggregated, h_res, h_post, hidden_states + + def fake_bda( + h_res, residual, h_post, output_with_bias, dropout_prob, training, fused, manager=None + ): + seen_managers.append(manager) + return residual + + monkeypatch.setattr(layer.hyper_connection, "forward", fake_hyper_connection_forward) + monkeypatch.setattr(layer.hyper_connection, "fused_h_res_h_post_bda", fake_bda) + + output, _ = layer(hidden_states, attention_mask=None, mhc_recompute_manager=manager) + assert output is hidden_states + assert seen_managers == [None] + + manager.is_last_layer_in_recompute_block = False + layer(hidden_states, attention_mask=None, mhc_recompute_manager=manager) + assert seen_managers[-1] is manager + + def test_transformer_layer_wrapper_escape_hatch(self): + config = _get_config(num_layers=1) + layer = _StubTransformerLayer(config) + hidden_states = torch.randn(4, 2, config.hidden_size) + + with pytest.raises(RuntimeError, match="must not be called directly"): + layer.forward(hidden_states=hidden_states, attention_mask=None) + + output, context = layer.forward( + hidden_states=hidden_states, attention_mask=None, _called_from_hybrid_mhc_wrapper=True + ) + assert output is hidden_states + assert context is None diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index eb871568046..ee0fa726ec8 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -68,6 +68,10 @@ "cpu_offloading_weights": False, "cross_entropy_fusion_impl": "native", "cross_entropy_loss_fusion": True, + "csa_compress_ratios": None, + "csa_compress_rotary_base": 40000.0, + "csa_dense_mode": False, + "csa_window_size": 128, "cuda_graph_impl": "none", "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], @@ -100,6 +104,7 @@ "embedding_init_method_std": 0.014, "enable_autocast": False, "enable_cuda_graph": False, + "enable_hyper_connections": False, "ep_overlap_early_attn_memory_release": False, "experimental_attention_variant": None, "experimental_attention_variant_loss_scale_func": None, @@ -171,6 +176,9 @@ "mamba_training_ssm_states_dtype": None, "masked_softmax_fusion": True, "memory_efficient_layer_norm": False, + "mhc_init_gating_factor": 0.01, + "mhc_recompute_layer_num": None, + "mhc_sinkhorn_iterations": 20, "microbatch_group_size_per_vp_stage": 1, "mlp_chunks_for_prefill": 1, "mlp_chunks_for_training": 1, @@ -257,6 +265,7 @@ "num_microbatches_with_partial_activation_checkpoints": None, "num_moe_experts": 128, "num_query_groups": 2, + "num_residual_streams": 4, "output_layer_init_method": {}, "overlap_moe_expert_parallel_comm": False, "overlap_p2p_comm": False, diff --git a/tests/unit_tests/pipeline_parallel/test_schedules.py b/tests/unit_tests/pipeline_parallel/test_schedules.py index 92db675d193..3d50763ea61 100644 --- a/tests/unit_tests/pipeline_parallel/test_schedules.py +++ b/tests/unit_tests/pipeline_parallel/test_schedules.py @@ -377,7 +377,8 @@ def test_dsa_indexer_loss_scale_accepts_dict_output_tensor(): ) -def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config(): +@pytest.mark.parametrize("variant", ["dsa", "dsv4_hybrid"]) +def test_indexer_loss_scale_defaults_from_variant_without_mutating_config(variant): from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, ) @@ -385,7 +386,7 @@ def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config(): config = SimpleNamespace( calculate_per_token_loss=True, experimental_attention_variant_loss_scale_func=None, - experimental_attention_variant='dsa', + experimental_attention_variant=variant, grad_scale_func=lambda tensor: tensor * 7.0, num_moe_experts=None, mtp_num_layers=None, diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index 8b4c181ee30..83a1b517d4d 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -19,6 +19,13 @@ ) +def expected_layer_counts(nonzero=None): + """Build an exact count map while keeping zero-count symbols explicit.""" + counts = {symbol: 0 for symbol in Symbols.VALID_LAYERS} + counts.update(nonzero or {}) + return counts + + @pytest.mark.internal class TestPatternFromRatios: @@ -79,6 +86,7 @@ def test_valid_patterns(self): ("GEGEGE*E", ['G', 'E', 'G', 'E', 'G', 'E', '*', 'E']), ("MDMD", ['M', 'D', 'M', 'D']), ("M+M+", ['M', '+', 'M', '+']), + ("WECEH+", ['W', 'E', 'C', 'E', 'H', '+']), ] for pattern, expected in test_cases: result = validate_segment_layers(pattern) @@ -108,6 +116,12 @@ def test_invalid_symbols_cause_failure(self): # own decoupled RoPE). validate_segment_layers("M+M*-") + def test_dsv4_attention_symbols(self): + assert {Symbols.WINDOW, Symbols.CSA, Symbols.HCA, Symbols.MLA} <= Symbols.MLA_ATTENTION + assert validate_segment_layers("WDCH+") == ["W", "D", "C", "H", "+"] + with pytest.raises(ValueError): + validate_segment_layers("W*C") + @pytest.mark.internal class TestGetHybridTotalLayerCount: @@ -315,170 +329,60 @@ def test_dataclass_equality(self): class TestGetHybridLayerCounts: def test_simple_pattern(self): - assert get_hybrid_layer_counts("M*M*") == { - '*': 2, - 'D': 0, - 'G': 0, - 'M': 2, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("M*M*") == expected_layer_counts({'*': 2, 'M': 2}) def test_all_layer_types(self): # Not allowed to have both standard Attention and MLA/DSA, so we do separate asserts. - assert get_hybrid_layer_counts("MG*-E") == { - '*': 1, - 'D': 0, - 'G': 1, - 'M': 1, - '+': 0, - '-': 1, - 'E': 1, - } - assert get_hybrid_layer_counts("MGD-E") == { - '*': 0, - 'D': 1, - 'G': 1, - 'M': 1, - '+': 0, - '-': 1, - 'E': 1, - } - assert get_hybrid_layer_counts("MG+-E") == { - '*': 0, - 'D': 0, - 'G': 1, - 'M': 1, - '+': 1, - '-': 1, - 'E': 1, - } + assert get_hybrid_layer_counts("MG*-E") == expected_layer_counts( + {'*': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + ) + assert get_hybrid_layer_counts("MGD-E") == expected_layer_counts( + {'D': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1} + ) + assert get_hybrid_layer_counts("MGDCHW+-E") == expected_layer_counts( + {'D': 1, 'C': 1, 'H': 1, 'W': 1, 'G': 1, 'M': 1, '+': 1, '-': 1, 'E': 1} + ) def test_with_pipes(self): # Pipes should be skipped in counting - assert get_hybrid_layer_counts("M*|M*") == { - '*': 2, - 'D': 0, - 'G': 0, - 'M': 2, - '+': 0, - '-': 0, - 'E': 0, - } - assert get_hybrid_layer_counts("M-M-|M-M*-") == { - '*': 1, - 'D': 0, - 'G': 0, - 'M': 4, - '+': 0, - '-': 4, - 'E': 0, - } + assert get_hybrid_layer_counts("M*|M*") == expected_layer_counts({'*': 2, 'M': 2}) + assert get_hybrid_layer_counts("M-M-|M-M*-") == expected_layer_counts( + {'*': 1, 'M': 4, '-': 4} + ) def test_with_mtp(self): # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers - assert get_hybrid_layer_counts("M*M*/MM/MM") == { - '*': 2, - 'D': 0, - 'G': 0, - 'M': 6, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("M*M*/MM/MM") == expected_layer_counts({'*': 2, 'M': 6}) def test_with_pipes_and_mtp(self): # Main: M-M-|M-M*- -> 1 attn, 4 mamba, 4 mlp # MTP: MM x 2 depths -> +4 mamba - assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == { - '*': 1, - 'D': 0, - 'G': 0, - 'M': 8, - '+': 0, - '-': 4, - 'E': 0, - } + assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == expected_layer_counts( + {'*': 1, 'M': 8, '-': 4} + ) def test_moe_pattern(self): - assert get_hybrid_layer_counts("MEME") == { - '*': 0, - 'D': 0, - 'G': 0, - 'M': 2, - '+': 0, - '-': 0, - 'E': 2, - } + assert get_hybrid_layer_counts("MEME") == expected_layer_counts({'M': 2, 'E': 2}) def test_mtp_with_attention(self): # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP - assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == { - '*': 3, - 'D': 0, - 'G': 0, - 'M': 7, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == expected_layer_counts({'*': 3, 'M': 7}) def test_gdn_pattern(self): - assert get_hybrid_layer_counts("GMGM") == { - '*': 0, - 'D': 0, - 'G': 2, - 'M': 2, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("GMGM") == expected_layer_counts({'G': 2, 'M': 2}) def test_gdn_hybrid_pattern(self): # GDN + Mamba + Attention - assert get_hybrid_layer_counts("G*GM*") == { - '*': 2, - 'D': 0, - 'G': 2, - 'M': 1, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("G*GM*") == expected_layer_counts({'*': 2, 'G': 2, 'M': 1}) def test_dsa_pattern(self): - assert get_hybrid_layer_counts("DMDM") == { - '*': 0, - 'D': 2, - 'G': 0, - 'M': 2, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("DMDM") == expected_layer_counts({'D': 2, 'M': 2}) def test_mla_pattern(self): - assert get_hybrid_layer_counts("+M+M") == { - '*': 0, - 'D': 0, - 'G': 0, - 'M': 2, - '+': 2, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("+M+M") == expected_layer_counts({'+': 2, 'M': 2}) def test_empty_pattern(self): - assert get_hybrid_layer_counts("") == { - '*': 0, - 'D': 0, - 'G': 0, - 'M': 0, - '+': 0, - '-': 0, - 'E': 0, - } + assert get_hybrid_layer_counts("") == expected_layer_counts() @pytest.mark.internal @@ -761,7 +665,7 @@ def test_standard_layer_types(self): """Standard symbols each produce a single-entry map at local index 0.""" maps = get_layer_maps_from_layer_type_list(["*", "M", "-", "E"]) # We always get all symbols returned, not only those contained in the pattern. - assert len(maps) == 7 + assert len(maps) == len(Symbols.VALID_LAYERS) attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter( Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE )(maps) diff --git a/tests/unit_tests/test_argument_utils.py b/tests/unit_tests/test_argument_utils.py index 7c0b30d3d56..19a242c2c42 100644 --- a/tests/unit_tests/test_argument_utils.py +++ b/tests/unit_tests/test_argument_utils.py @@ -1,6 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import signal +import types from argparse import ArgumentError, ArgumentParser, Namespace from dataclasses import dataclass, field from typing import Callable, Literal, Optional, Union @@ -10,9 +11,12 @@ from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.optimizer import OptimizerConfig +from megatron.core.transformer.spec_utils import ModuleSpec from megatron.training.argument_utils import ( ArgumentGroupFactory, TypeInferenceError, + _normalize_dsv4_hybrid_csa_compress_ratios, + hybrid_config_from_args, pretrain_cfg_container_from_args, ) from megatron.training.config import PretrainConfigContainer @@ -83,6 +87,97 @@ class ConfigWithLiteral: """Precision level""" +class TestDsv4HybridCsaCompressRatioNormalization: + """Test the DSv4 HybridModel compression-ratio normalization contract.""" + + @pytest.mark.parametrize( + ("provided", "expected_config_ratios"), + [ + (None, [0, 0, 0, 4, 128, 0]), + ([0, 4, 128], [0, 0, 0, 4, 128, 0]), + ([0, 0, 0, 4, 128, 0], [0, 0, 0, 4, 128, 0]), + ], + ) + def test_normalizes_default_compact_and_padded_ratios(self, provided, expected_config_ratios): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=provided) + config_kwargs = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, config_kwargs, "-W|EC/H-") + + assert args.csa_compress_ratios == [0, 4, 128] + assert config_kwargs['csa_compress_ratios'] == expected_config_ratios + + @pytest.mark.parametrize( + ("provided", "message"), + [ + ([0, 8, 128], "ratio 8.*symbol 'C'.*expected 4"), + ([1, 0, 0, 4, 128, 0], "non-DSv4 hybrid symbol '-'.*non-zero ratio 1"), + ([0, 4], r"length \(2\).*W/C/H attention symbols \(3\)"), + ], + ) + def test_rejects_invalid_ratios(self, provided, message): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=provided) + + with pytest.raises(AssertionError, match=message): + _normalize_dsv4_hybrid_csa_compress_ratios(args, {}, "-W|EC/H-") + + def test_ordinary_d_does_not_consume_a_dsv4_ratio(self): + args = Namespace(experimental_attention_variant='dsv4_hybrid', csa_compress_ratios=[4]) + config_kwargs = {} + + _normalize_dsv4_hybrid_csa_compress_ratios(args, config_kwargs, "D-C/D") + + assert args.csa_compress_ratios == [4] + assert config_kwargs['csa_compress_ratios'] == [0, 0, 4, 0] + + +class TestHybridConfigFromArgs: + """Test static and config-aware hybrid stack spec resolution.""" + + @staticmethod + def _args(): + return Namespace( + spec=["test_module", "test_spec"], + fp16_lm_cross_entropy=False, + hybrid_layer_pattern="M", + position_embedding_type="none", + rotary_percent=1.0, + rotary_base=10000, + make_vocab_size_divisible_by=128, + rotary_seq_len_interpolation_factor=None, + max_position_embeddings=1024, + untie_embeddings_and_output_weights=False, + padded_vocab_size=128, + ) + + @staticmethod + def _transformer_config(): + return types.SimpleNamespace( + transformer_impl="transformer_engine", inference_fuse_tp_communication=False + ) + + @patch("megatron.training.argument_utils.import_module") + def test_preserves_static_module_spec(self, mock_import_module): + static_spec = ModuleSpec(module=object) + mock_import_module.return_value = static_spec + + config = hybrid_config_from_args(self._args(), config=self._transformer_config()) + + assert config.hybrid_stack_spec is static_spec + + @patch("megatron.training.argument_utils.import_module") + def test_resolves_config_aware_spec_factory(self, mock_import_module): + static_spec = ModuleSpec(module=object) + spec_factory = MagicMock(return_value=static_spec) + mock_import_module.return_value = spec_factory + transformer_config = self._transformer_config() + + config = hybrid_config_from_args(self._args(), config=transformer_config) + + spec_factory.assert_called_once_with(transformer_config) + assert config.hybrid_stack_spec is static_spec + + class TestArgumentGroupFactoryBasic: """Test basic functionality of ArgumentGroupFactory.""" diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index dc65d541455..2fc55962ae8 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -23,6 +23,31 @@ reason_for_no_mxfp8 = "MXFP8 requires Transformer Engine and device arch >= 10" +@pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine is not installed") +@pytest.mark.parametrize( + ("is_init", "config_values", "te_helper"), + [ + ( + False, + {"fp8": "hybrid", "fp4": None, "fp8_param": False, "fp4_param": False}, + "fp8_autocast", + ), + (True, {"fp8": None, "fp4": None, "fp8_param": True, "fp4_param": False}, "fp8_model_init"), + ], +) +def test_get_fp8_disabled_context_uses_disabled_te_context(is_init, config_values, te_helper): + config = Mock(**config_values) + disabled_context = Mock() + + with patch.object( + fp8_utils.transformer_engine.pytorch, te_helper, return_value=disabled_context + ) as te_context: + result = fp8_utils.get_fp8_disabled_context(config, is_init=is_init) + + assert result is disabled_context + te_context.assert_called_once_with(enabled=False) + + class MockTELinear(nn.Module): """Mock TE Linear module for testing.""" diff --git a/tests/unit_tests/test_optimizer_cpu_offloading.py b/tests/unit_tests/test_optimizer_cpu_offloading.py index 33febbb3eb0..379acc9dbda 100644 --- a/tests/unit_tests/test_optimizer_cpu_offloading.py +++ b/tests/unit_tests/test_optimizer_cpu_offloading.py @@ -17,6 +17,20 @@ from torch.optim import Adam as GPUAdam from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer +from megatron.core.transformer.module import ( + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, +) + + +class Fp32MarkedToyNet(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4, bias=False) + self.scale = mark_keep_in_fp32(nn.Parameter(torch.ones(4))) + + def forward(self, x): + return self.proj(x) * self.scale class Net(nn.Module): @@ -71,6 +85,52 @@ def setup_seed(seed): torch.backends.cudnn.benchmark = False # Disable auto-tuner for reproducibility +def test_load_state_dict_with_native_fp32_param(): + """Round-trip state for a BF16 toy net with a parameter marked to stay in FP32.""" + model = Fp32MarkedToyNet().cuda() + convert_module_to_dtype_except_fp32_marked(model, torch.bfloat16) + assert model.proj.weight.dtype == torch.bfloat16 + assert model.scale.dtype == torch.float32 + + optimizer = HybridDeviceOptimizer( + model.parameters(), + offload_fraction=1.0, + cpu_optimizer_cls=Adam, + gpu_optimizer_cls=GPUAdam, + param_update_in_fp32=True, + overlap_cpu_optimizer_d2h_h2d=False, + lr=1e-3, + ) + inputs = torch.ones(2, 4, device="cuda", dtype=torch.bfloat16) + model(inputs).sum().backward() + optimizer.step() + + restored_model = Fp32MarkedToyNet().cuda() + convert_module_to_dtype_except_fp32_marked(restored_model, torch.bfloat16) + restored_model.load_state_dict(model.state_dict()) + restored_optimizer = HybridDeviceOptimizer( + restored_model.parameters(), + offload_fraction=1.0, + cpu_optimizer_cls=Adam, + gpu_optimizer_cls=GPUAdam, + param_update_in_fp32=True, + overlap_cpu_optimizer_d2h_h2d=False, + lr=1e-3, + ) + restored_optimizer.load_state_dict(optimizer.state_dict()) + + assert set(restored_optimizer.state) == set(restored_model.parameters()) + assert restored_model.proj.weight in restored_optimizer.param_to_fp32_param + assert restored_model.scale not in restored_optimizer.param_to_fp32_param + assert torch.equal( + restored_optimizer.param_to_fp32_param[restored_model.proj.weight], + optimizer.param_to_fp32_param[model.proj.weight], + ) + + restored_model(inputs).sum().backward() + restored_optimizer.step() + + @pytest.mark.skipif( torch.__version__ < '2.3.0', reason=( diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py new file mode 100644 index 00000000000..e7f68369431 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,1418 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + _apply_rope, + _compute_unfused_csa_non_compressed_lse, + get_compress_topk_idxs, + get_window_topk_idxs, + unfused_compressed_sparse_attn, +) +from megatron.core.transformer.experimental_attention_variant.dsa import ( + FusedDSAIndexerLoss, + compute_dsa_indexer_loss, + fused_qk_topk_naive, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + + +def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + """Mock implementation of hadamard_transform for testing without the library installed.""" + return x * scale + + +class _DisabledContextTracker: + """Track whether a projection runs inside the FP8-disabled context.""" + + def __init__(self): + self.depth = 0 + self.entries = 0 + + def __call__(self, _config, is_init=False): + assert not is_init + return self + + def __enter__(self): + self.depth += 1 + self.entries += 1 + return self + + def __exit__(self, _exc_type, _exc_value, _traceback): + self.depth -= 1 + return False + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Automatically patch hadamard_transform in both dsa and csa modules if not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# =========================================================================== +# Helper function tests +# =========================================================================== + + +class _SingleRankTP: + @staticmethod + def size(): + return 1 + + +class _SingleRankPG: + tp = _SingleRankTP() + + +def test_unfused_csa_non_compressed_lse_matches_window_and_sink_oracle(): + torch.manual_seed(17) + seqlen_q, batch_size, n_kv = 3, 2, 5 + num_heads, head_dim = 2, 4 + query = torch.randn(seqlen_q, batch_size, num_heads, head_dim, requires_grad=True) + kv_full = torch.randn(n_kv, batch_size, head_dim, requires_grad=True) + sink = torch.randn(num_heads, requires_grad=True) + window_indices = torch.tensor([[[-1, 0], [0, 1], [1, 2]], [[-1, 0], [0, 2], [2, 4]]]) + + expected = torch.empty(batch_size, num_heads, seqlen_q) + with torch.no_grad(): + for batch in range(batch_size): + for row in range(seqlen_q): + for head in range(num_heads): + logits = [sink[head]] + for key_index in window_indices[batch, row]: + if key_index >= 0: + logits.append( + torch.dot(query[row, batch, head], kv_full[key_index, batch]) + ) + expected[batch, head, row] = torch.logsumexp(torch.stack(logits), dim=0) + + actual = _compute_unfused_csa_non_compressed_lse( + query, kv_full, sink, window_indices, softmax_scale=1.0 + ) + + assert actual.shape == (batch_size, num_heads, seqlen_q) + assert actual.dtype == torch.float32 + assert not actual.requires_grad + torch.testing.assert_close(actual, expected) + for teacher_tensor in (query, kv_full, sink): + assert teacher_tensor.grad is None + + +def _independent_csa_indexer_loss( + index_scores, + topk_indices, + query, + compressed_kv, + window_kv, + window_indices, + sink, + *, + sparse_loss, + loss_coeff, +): + """Compute a small-loop CSA teacher oracle with the complete denominator.""" + batch_size, seqlen_q, n_compressed = index_scores.shape + num_heads = query.shape[2] + losses = [] + for batch in range(batch_size): + for row in range(seqlen_q): + selected = ( + topk_indices[batch, row].tolist() if sparse_loss else list(range(n_compressed)) + ) + target = [] + for compressed_index in selected: + head_mass = 0.0 + for head in range(num_heads): + non_compressed_logits = [sink[head]] + for window_index in window_indices[batch, row]: + if window_index >= 0: + non_compressed_logits.append( + torch.dot(query[row, batch, head], window_kv[window_index, batch]) + ) + compressed_logits = [ + torch.dot(query[row, batch, head], compressed_kv[key_index, batch]) + for key_index in selected + ] + denominator = torch.logsumexp( + torch.stack(non_compressed_logits + compressed_logits), dim=0 + ) + selected_position = selected.index(compressed_index) + head_mass = head_mass + torch.exp( + compressed_logits[selected_position] - denominator + ) + target.append(head_mass) + target = torch.stack(target) + target = target / target.sum() + predict_log = torch.log_softmax(index_scores[batch, row, selected], dim=-1) + losses.append((target * (torch.log(target) - predict_log)).sum()) + return torch.stack(losses).mean() * loss_coeff + + +@pytest.mark.parametrize("sparse_loss", [False, True], ids=["dense", "sparse"]) +def test_csa_indexer_loss_uses_full_attention_denominator(sparse_loss): + torch.manual_seed(29) + seqlen_q, batch_size, num_heads, head_dim = 4, 1, 2, 3 + n_compressed, index_heads, index_dim = 3, 2, 2 + index_topk, loss_coeff = 2, 0.7 + + q = torch.randn(seqlen_q, batch_size, index_heads, index_dim, requires_grad=True) + weights = torch.randn(seqlen_q, batch_size, index_heads, requires_grad=True) + k = torch.randn(n_compressed, batch_size, index_dim, requires_grad=True) + query = torch.randn(seqlen_q, batch_size, num_heads, head_dim, requires_grad=True) + window_kv = torch.randn(seqlen_q, batch_size, head_dim, requires_grad=True) + compressed_kv = torch.randn(n_compressed, batch_size, head_dim, requires_grad=True) + sink = torch.randn(num_heads, requires_grad=True) + window_indices = torch.tensor([[[-1, 0], [0, 1], [1, 2], [2, 3]]]) + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, window_kv, sink, window_indices, softmax_scale=1.0 + ) + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, num_heads, -1) + compressed_mask = torch.zeros(seqlen_q, n_compressed) + + q_reference = q.detach().clone().requires_grad_(True) + weights_reference = weights.detach().clone().requires_grad_(True) + k_reference = k.detach().clone().requires_grad_(True) + index_scores_reference, topk_reference = fused_qk_topk_naive( + q_reference, k_reference, weights_reference, index_topk + ) + loss_reference = compute_dsa_indexer_loss( + index_scores_reference, + topk_reference, + query.detach(), + key_for_loss.detach(), + 1.0, + loss_coeff, + sparse_loss, + _SingleRankPG(), + mask=compressed_mask, + non_compressed_lse=non_compressed_lse, + ) + loss_reference.backward() + + topk_actual, loss_actual = FusedDSAIndexerLoss.apply( + q, + weights, + k, + query, + key_for_loss, + 1.0, + index_topk, + loss_coeff, + compressed_mask, + sparse_loss, + _SingleRankPG(), + None, + None, + None, + None, + False, + True, + non_compressed_lse, + ) + loss_actual.backward() + + independent_loss = _independent_csa_indexer_loss( + index_scores_reference.detach(), + topk_reference, + query.detach(), + compressed_kv.detach(), + window_kv.detach(), + window_indices, + sink.detach(), + sparse_loss=sparse_loss, + loss_coeff=loss_coeff, + ) + + torch.testing.assert_close(loss_actual, independent_loss) + torch.testing.assert_close(loss_actual, loss_reference) + torch.testing.assert_close(topk_actual, topk_reference) + torch.testing.assert_close(q.grad, q_reference.grad) + torch.testing.assert_close(weights.grad, weights_reference.grad) + torch.testing.assert_close(k.grad, k_reference.grad) + for teacher_tensor in (query, window_kv, compressed_kv, sink): + assert teacher_tensor.grad is None + + +class TestGetWindowTopkIdxs: + """Test get_window_topk_idxs helper.""" + + def test_basic_shape(self): + batch_size, seqlen, window_size = 2, 16, 4 + idxs = get_window_topk_idxs(window_size, batch_size, seqlen, torch.device("cpu")) + assert idxs.shape == (batch_size, seqlen, window_size) + + def test_causal_no_future(self): + """Indices should never exceed the query position.""" + seqlen, window_size = 32, 8 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + for i in range(seqlen): + valid = idxs[0, i][idxs[0, i] >= 0] + assert torch.all(valid <= i), f"Position {i} has future indices" + + def test_invalid_marked_minus_one(self): + """Early positions that cannot fill the window should use -1.""" + seqlen, window_size = 8, 4 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs[0, 0, 0] == -1 or idxs[0, 0, 0] == 0 + for pos in range(window_size, seqlen): + assert torch.all(idxs[0, pos] >= 0), f"Position {pos} has invalid -1" + + def test_window_larger_than_seqlen(self): + """Window larger than sequence length should still work.""" + seqlen, window_size = 4, 16 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs.shape == (1, seqlen, window_size) + + +class TestGetCompressTopkIdxs: + """Test get_compress_topk_idxs helper.""" + + def test_basic_shape(self): + ratio, batch_size, seqlen, offset = 4, 2, 32, 32 + idxs = get_compress_topk_idxs(ratio, batch_size, seqlen, offset, torch.device("cpu")) + n_compressed = seqlen // ratio + assert idxs.shape == (batch_size, seqlen, n_compressed) + + def test_offset_applied(self): + """Valid indices should be >= offset.""" + ratio, seqlen, offset = 4, 32, 100 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + valid = idxs[idxs >= 0] + if valid.numel() > 0: + assert torch.all(valid >= offset), "Valid indices should be offset" + + def test_causal_no_future(self): + """Compressed indices should respect causality.""" + ratio, seqlen, offset = 4, 32, 32 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + for i in range(seqlen): + n_valid = (i + 1) // ratio + valid = idxs[0, i][idxs[0, i] >= 0] + assert valid.numel() <= n_valid, f"Position {i} has too many valid compressed indices" + + def test_ratio_128(self): + """Test with large compression ratio.""" + ratio, seqlen, offset = 128, 256, 256 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + assert idxs.shape == (1, seqlen, seqlen // ratio) + + +# =========================================================================== +# unfused_compressed_sparse_attn tests +# =========================================================================== + + +class TestUnfusedCompressedSparseAttn: + """Test the unfused compressed sparse attention kernel.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_output_shape(self): + """Test output shape of unfused compressed sparse attention.""" + sq, b, np_, hn = 16, 2, 4, 64 + n_kv = sq + sq // 4 + topk = 8 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + + assert output.shape == (sq, b, np_ * hn) + assert output.dtype == query.dtype + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_invalid_indices_masked(self): + """Test that -1 indices are properly masked.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + topk_indices = torch.full((b, sq, topk), -1, dtype=torch.int32).cuda() + topk_indices[:, :, 0] = 0 + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + assert not torch.isnan(output).any(), "Output should not contain NaN" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_gradient_flow(self): + """Test that gradients flow through sparse attention.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + kv_full = torch.randn(n_kv, b, hn, dtype=torch.float32).cuda().requires_grad_(True) + attn_sink = torch.nn.Parameter(torch.zeros(np_, dtype=torch.float32).cuda()) + + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert kv_full.grad is not None + assert attn_sink.grad is not None + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +def _make_mla_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + csa_compress_ratios=None, + csa_window_size=8, + csa_dense_mode=False, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + dsa_indexer_use_sparse_loss=False, +): + """Helper to create MLATransformerConfig for CSA tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0] * num_layers + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + csa_dense_mode=csa_dense_mode, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + + +def _make_compressor_submodules(): + """Create Compressor submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressorSubmodules( + linear_wkv=ModuleSpec(module=TELinear), + linear_wgate=ModuleSpec(module=TELinear), + norm=ModuleSpec(module=TENorm), + ) + + +def _make_csa_indexer_submodules(): + """Create CSAIndexer submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_weights_proj=ModuleSpec(module=TELinear), + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + ) + + +def _make_csa_submodules(): + """Create CompressedSparseAttention submodules spec.""" + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressedSparseAttentionSubmodules( + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + indexer=ModuleSpec(module=CSAIndexer, submodules=_make_csa_indexer_submodules()), + ) + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressor: + """Test Compressor module.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_output_shape(self, compress_ratio): + """Test that compressor produces correct output shape.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + + expected_len = seq_len // compress_ratio + assert output is not None + assert output.shape == (expected_len, batch_size, head_dim) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_too_short_input(self, compress_ratio): + """Test that compressor returns None when input is shorter than compress_ratio.""" + short_len = compress_ratio - 1 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + assert output is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_gradient_flow(self, compress_ratio): + """Test that gradients flow through the compressor.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + output = compressor(x) + loss = output.sum() + loss.backward() + + assert x.grad is not None + for name, param in compressor.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_projection_disables_fp8(self, compress_ratio, monkeypatch): + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=self.config.v_head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + tracker = _DisabledContextTracker() + calls = [] + + for name, projection in ( + ('linear_wkv', compressor.linear_wkv), + ('linear_wgate', compressor.linear_wgate), + ): + original_forward = projection.forward + + def checked_forward(*args, _name=name, _forward=original_forward, **kwargs): + assert tracker.depth > 0, f"{_name} ran outside the FP8-disabled context" + calls.append(_name) + return _forward(*args, **kwargs) + + monkeypatch.setattr(projection, 'forward', checked_forward) + + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn( + compress_ratio * 2, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + compressor(x) + + assert calls == ['linear_wkv', 'linear_wgate'] + assert tracker.entries == 1 + + +# =========================================================================== +# CSAIndexer tests +# =========================================================================== + + +@pytest.mark.parametrize("seqlen", [32, 128]) +class TestCSAIndexer: + """Test CSAIndexer module basic functionality.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.compress_ratio = 4 + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4], dsa_indexer_topk=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.indexer = CSAIndexer( + config=cls.config, + submodules=_make_csa_indexer_submodules(), + compress_ratio=cls.compress_ratio, + rotary_pos_emb=cls.rotary_pos_emb, + pg_collection=cls.pg_collection, + ) + + yield + Utils.destroy_model_parallel() + + def test_csa_indexer_constructor(self, seqlen): + """Test CSAIndexer initialization.""" + assert isinstance(self.indexer, CSAIndexer) + assert self.indexer.compress_ratio == self.compress_ratio + assert self.indexer.index_n_heads == self.config.dsa_indexer_n_heads + assert self.indexer.index_head_dim == self.config.dsa_indexer_head_dim + assert self.indexer.index_topk == self.config.dsa_indexer_topk + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward(self, seqlen): + """Test CSAIndexer forward pass.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + index_scores, topk_indices = self.indexer(x, qr) + n_compressed = seqlen // self.compress_ratio + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward_before_topk(self, seqlen): + """Test CSAIndexer forward_before_topk.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + q, k, weights = self.indexer.forward_before_topk(x, qr) + + assert q.shape == ( + seqlen, + batch_size, + self.config.dsa_indexer_n_heads, + self.config.dsa_indexer_head_dim, + ) + n_compressed = seqlen // self.compress_ratio + assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) + assert weights.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_weights_projection_disables_fp8(self, seqlen, monkeypatch): + tracker = _DisabledContextTracker() + self.indexer.cuda() + original_forward = self.indexer.linear_weights_proj.forward + + def checked_forward(*args, **kwargs): + assert tracker.depth > 0, "indexer weights projection ran under FP8" + return original_forward(*args, **kwargs) + + monkeypatch.setattr(self.indexer.linear_weights_proj, 'forward', checked_forward) + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn(seqlen, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + weights = self.indexer._project_weights(x) + + assert weights.shape == (seqlen, 1, self.config.dsa_indexer_n_heads) + assert tracker.entries == 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_with_mask(self, seqlen): + """Test CSAIndexer with causal mask.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + n_compressed = seqlen // self.compress_ratio + causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(seqlen, -1) + positions = torch.arange(1, seqlen + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) + + index_scores, topk_indices = self.indexer(x, qr, mask=causal_mask) + + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + + +# =========================================================================== +# CompressedSparseAttention tests +# =========================================================================== + + +class TestCompressedSparseAttentionRatio1: + """Test CompressedSparseAttention with compress_ratio=1 (window-only).""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[0, 0, 0, 0], csa_window_size=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.csa = CompressedSparseAttention( + config=cls.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=cls.pg_collection, + rotary_pos_emb=rotary_pos_emb, + compress_ratio=0, + ) + + yield + Utils.destroy_model_parallel() + + def test_ratio1_no_compressor(self): + """With ratio=1, compressor and indexer should not be built.""" + assert self.csa.compressor is None + assert self.csa.indexer is None + + def test_mtp_layer_number_is_offset(self): + """MTP attention layers are numbered after all decoder layers.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + compress_ratio=0, + is_mtp_layer=True, + ) + + assert csa.layer_number == self.config.num_layers + 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_forward(self): + """Test forward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_backward(self): + """Test backward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.train() + self.csa.cuda() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressedSparseAttentionCompressed: + """Test CompressedSparseAttention with compress_ratio > 1.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=1.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _get_layer_number(self, compress_ratio): + """Return a layer_number (1-indexed) whose compress_ratio matches.""" + for i, r in enumerate(self.config.csa_compress_ratios): + if r == compress_ratio: + return i + 1 + raise ValueError(f"No layer with compress_ratio={compress_ratio}") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_constructor(self, compress_ratio): + """Test that compressor/indexer are conditionally built.""" + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + assert csa.compressor is not None + if compress_ratio == 4: + assert csa.indexer is not None + elif compress_ratio == 128: + assert csa.indexer is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_forward(self, compress_ratio): + """Test forward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward(self, compress_ratio): + """Test backward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.train() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + for name, param in csa.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_eval_mode(self, compress_ratio): + """Test forward pass in eval mode.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.eval() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + with torch.no_grad(): + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +# =========================================================================== +# _apply_rope tests +# =========================================================================== + + +class TestApplyRope: + """Test ``_apply_rope`` — the layout-aware RoPE wrapper used by + Compressor / CSAIndexer / hybrid-attention callers. + + Behaviours covered: + + * 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs both work (3-D gets a temporary head-dim unsqueeze). + * Only the trailing ``pos_dim`` components are rotated; the leading + ``nope_dim`` slice is bit-exact unchanged. + * Both ``RotaryEmbedding`` (returns ``Tensor``) and + ``YarnRotaryEmbedding`` (returns ``(emb, mscale)`` tuple) — DSv4 + hybrid silently swaps the class based on ``compress_ratio``. + * Both unfused and fused (``config.apply_rope_fusion=True``) paths + produce the same output (within bf16 precision). + * For ``ratio > 1`` the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(0) + model_parallel_cuda_manual_seed(0) + cls = request.cls + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + # head_dim 32 = nope 24 + pos 8 + cls.config = _make_mla_config(v_head_dim=32, qk_pos_emb_head_dim=8) + yield + Utils.destroy_model_parallel() + + def _make_rotary(self, kind: str): + from megatron.core.models.common.embeddings import RotaryEmbedding, YarnRotaryEmbedding + + pos_dim = self.config.qk_pos_emb_head_dim + if kind == 'rope': + return RotaryEmbedding( + pos_dim, rotary_percent=1.0, rotary_base=10000, cp_group=self.pg_collection.cp + ) + if kind == 'yarn': + return YarnRotaryEmbedding( + pos_dim, + rotary_base=40000, + scaling_factor=40, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1.0, + mscale_all_dim=0.0, + cp_group=self.pg_collection.cp, + ) + raise ValueError(kind) + + def _config_with(self, *, apply_rope_fusion: bool): + # Reuse the class-level config; only flip the fusion flag. + cfg = self.config + cfg.apply_rope_fusion = apply_rope_fusion + return cfg + + _ROTARY_FUSION_COMBOS = [ + pytest.param('rope', False, id='rope-unfused'), + pytest.param('rope', True, id='rope-fused'), + pytest.param('yarn', False, id='yarn-unfused'), + pytest.param('yarn', True, id='yarn-fused'), + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize(("rotary_kind", "apply_rope_fusion"), _ROTARY_FUSION_COMBOS) + @pytest.mark.parametrize("input_ndim", [3, 4], ids=['3d', '4d']) + @pytest.mark.parametrize("ratio", [1, 4], ids=['ratio_1', 'ratio_4']) + def test_apply_rope(self, rotary_kind, apply_rope_fusion, input_ndim, ratio): + """Output shape == input shape; no NaN; nope-dim slice is + bit-exact unchanged. Sweeps the valid combinations of rotary + class × apply_rope_fusion × input rank × ratio. Yarn's + tuple-return is covered by the ``'yarn-*'`` combos. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads = 8, 2, 4 + cfg = self._config_with(apply_rope_fusion=apply_rope_fusion) + + shape = (seq, batch, head_dim) if input_ndim == 3 else (seq, batch, heads, head_dim) + x = torch.randn(*shape, dtype=torch.bfloat16, device='cuda') + # ``fused_mla_rope_inplace`` mutates the input — give it a copy so + # the nope-dim equality check below still has the original. + out = _apply_rope( + x.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + assert out.shape == x.shape + assert out.dtype == x.dtype + assert not torch.isnan(out).any() + # The leading nope_dim slice is the identity portion of RoPE. + assert torch.equal( + out[..., :nope], x[..., :nope] + ), "RoPE must not touch the first nope_dim components" + # Trailing pos_dim should rotate at non-zero positions. + pe_changed = (out[..., nope:] != x[..., nope:]).any(dim=-1).flatten() + assert pe_changed[ + 1: + ].any(), "RoPE should rotate the trailing pos_dim components for seq > 0" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_3d_input_matches_4d_with_single_head(self, rotary_kind): + """For a single-head input, the 3-D ``(s, b, d)`` and 4-D + ``(s, b, 1, d)`` invocations must produce numerically identical + output (3-D path just inserts a temporary head dim). + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch = 8, 2 + cfg = self._config_with(apply_rope_fusion=False) + + x_3d = torch.randn(seq, batch, head_dim, dtype=torch.bfloat16, device='cuda') + x_4d = x_3d.unsqueeze(-2) + + out_3d = _apply_rope( + x_3d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_4d = _apply_rope( + x_4d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + + assert out_3d.shape == x_3d.shape + assert out_4d.shape == x_4d.shape + assert torch.equal(out_3d, out_4d.squeeze(-2)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_ratio_strides_rotary_table(self, rotary_kind): + """For ``ratio > 1``, the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. The result + with ``ratio=k`` must equal an ``apply_rope`` call on the same + positions of a length-``rotary_seq_len * k`` table. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads, ratio = 4, 1, 2, 4 + cfg = self._config_with(apply_rope_fusion=False) + + x_comp = torch.randn(seq, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda') + out_comp = _apply_rope( + x_comp.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + x_full = torch.zeros( + seq * ratio, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda' + ) + x_full[::ratio][:seq] = x_comp + out_full = _apply_rope( + x_full, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq * ratio, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_ref = out_full[::ratio][:seq] + + assert torch.allclose(out_comp, out_ref, rtol=1e-3, atol=1e-3), ( + f"ratio={ratio} stride mismatch: " + f"max abs diff = {(out_comp - out_ref).abs().max().item():.3e}" + ) + + +# =========================================================================== +# csa_dense_mode tests +# =========================================================================== + + +class TestCompressedSparseAttentionDenseMode: + """Test that csa_dense_mode=True disables the indexer for ratio=4 layers.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], csa_window_size=8, csa_dense_mode=True + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_disables_indexer_for_ratio4(self): + """With csa_dense_mode=True, ratio=4 layers should NOT build an indexer.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + assert csa.compress_ratio == 4 + assert csa.compressor is not None, "Compressor should still be built" + assert csa.indexer is None, "Indexer should be disabled in dense mode" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_forward_ratio4(self): + """Forward pass should work for ratio=4 in dense mode (uses all compressed positions).""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +class TestCSAHighPrecisionParams: + """Reference-checkpoint FP32 parameters survive BF16 model conversion.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ape_and_attn_sink_stay_fp32_after_bf16_conversion(self): + from megatron.core.transformer.module import Float16Module + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + name="decoder.layers.0.self_attention.core_attention", + ) + + assert csa.attn_sink.dtype == torch.float32 + assert csa.compressor.ape.dtype == torch.float32 + assert csa.indexer.compressor.ape.dtype == torch.float32 + + bf16_module = Float16Module(config=self.config, module=csa) + + assert bf16_module.module.attn_sink.dtype == torch.float32 + assert bf16_module.module.compressor.ape.dtype == torch.float32 + assert bf16_module.module.indexer.compressor.ape.dtype == torch.float32 + assert bf16_module.module.compressor.linear_wkv.weight.dtype == torch.bfloat16 + assert bf16_module.module.compressor.linear_wgate.weight.dtype == torch.bfloat16 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py index 135c4802dd3..1881ffaaa40 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py @@ -28,6 +28,8 @@ DSAttention, DSAttentionSubmodules, FusedDSAIndexerLoss, + _compute_indexer_teacher_probabilities, + _normalize_indexer_teacher_target, _run_sparse_attention, _validate_nonpacked_cp_uniform_length, compute_dsa_indexer_loss, @@ -1658,6 +1660,25 @@ def test_rotate_activation_dtype_check(self): rotate_activation(x) +def test_indexer_teacher_probability_accepts_detached_external_mass(): + """An omitted-key LSE participates in the denominator without entering the target support.""" + attention_scores = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + valid_mask = torch.ones((1, 2, 2), dtype=torch.bool) + non_compressed_lse = torch.tensor([[[0.5, 1.5]]]) + + actual = _compute_indexer_teacher_probabilities( + attention_scores, valid_mask, non_compressed_lse + ) + expected_denominator = torch.logaddexp( + torch.logsumexp(attention_scores, dim=-1), non_compressed_lse + ) + expected = torch.exp(attention_scores - expected_denominator.unsqueeze(-1)) + torch.testing.assert_close(actual, expected) + + normalized = _normalize_indexer_teacher_target(actual.sum(dim=1), non_compressed_lse) + torch.testing.assert_close(normalized.sum(dim=-1), torch.ones((1, 2))) + + @pytest.mark.parametrize("seqlen_and_topk", [[16, 32], [64, 32]]) class TestComputeDSAIndexerLoss: """Test compute_dsa_indexer_loss function.""" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py index 8d63a9bee11..749fb43d3fb 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py @@ -681,3 +681,101 @@ def test_packed_cp_tp2_sequence_parallel_shared_skip_backend_matches_unfused_ref finally: DSAIndexerLossLoggingHelper.clean_loss_in_tracker() Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_indexer_loss_tracker_grows_for_hybrid_mtp_layer_numbers(): + """Hybrid MTP layers can have a layer number beyond the nominal layer count.""" + DSAIndexerLossLoggingHelper.tracker = {} + try: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(2.0, device="cuda"), layer_number=7, num_layers=5 + ) + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(3.0, device="cuda"), layer_number=9, num_layers=5 + ) + + values = DSAIndexerLossLoggingHelper.tracker["values"] + assert values.shape == (9,) + torch.testing.assert_close(values[6], torch.tensor(2.0, device="cuda")) + torch.testing.assert_close(values[8], torch.tensor(3.0, device="cuda")) + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_indexer_metrics_average_only_ratio4_layers(monkeypatch: pytest.MonkeyPatch): + """Window and compressed-only layers must not dilute the indexer-loss average.""" + DSAIndexerLossLoggingHelper.tracker = { + "values": torch.tensor([0.0, 3.0, 0.0, 0.0, 0.0], device="cuda") + } + monkeypatch.setattr(DSAIndexerLossLoggingHelper, "reduce_loss_in_tracker", lambda **_: None) + total_loss_dict = {} + try: + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=0.5, + iteration=1, + writer=None, + pg_collection=object(), + total_loss_dict=total_loss_dict, + num_layers=5, + num_indexer_layers=1, + ) + + torch.testing.assert_close( + total_loss_dict["indexer loss"], torch.tensor(1.5, device="cuda") + ) + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +def test_indexer_loss_cleanup_can_preserve_graph_reduction_groups(): + """CUDA Graph replays keep stable group objects while clearing accumulated loss.""" + reduce_group = object() + avg_group = object() + DSAIndexerLossLoggingHelper.tracker = { + "values": torch.ones(2), + "reduce_group": reduce_group, + "avg_group": avg_group, + } + try: + DSAIndexerLossLoggingHelper.clean_loss_in_tracker(preserve_groups=True) + assert torch.count_nonzero(DSAIndexerLossLoggingHelper.tracker["values"]) == 0 + assert DSAIndexerLossLoggingHelper.tracker["reduce_group"] is reduce_group + assert DSAIndexerLossLoggingHelper.tracker["avg_group"] is avg_group + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +def test_indexer_metrics_reduce_across_pipeline_rank_without_indexer(): + """Every pipeline rank must join indexer loss reduction, even without a local indexer.""" + if Utils.world_size < 2: + pytest.skip("Cross-pipeline indexer reduction requires at least two distributed ranks") + + Utils.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=2) + DSAIndexerLossLoggingHelper.tracker = {} + try: + num_layers = 5 + if parallel_state.get_pipeline_model_parallel_rank() == 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(3.0, device="cuda"), layer_number=2, num_layers=num_layers + ) + + total_loss_dict = {} + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=0.5, + iteration=1, + writer=None, + pg_collection=ProcessGroupCollection.use_mpu_process_groups(required_pgs=['pp', 'dp']), + total_loss_dict=total_loss_dict, + num_layers=num_layers, + num_indexer_layers=1, + ) + + torch.testing.assert_close( + total_loss_dict["indexer loss"], torch.tensor(1.5, device="cuda") + ) + assert torch.count_nonzero(DSAIndexerLossLoggingHelper.tracker["values"]) == 0 + finally: + DSAIndexerLossLoggingHelper.tracker = {} + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py new file mode 100644 index 00000000000..ca35e040aab --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -0,0 +1,617 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + +_SEED = 42 + + +def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Patch hadamard_transform in dsa/csa modules if the library is not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + _mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# --------------------------------------------------------------------------- +# Config / spec helpers +# --------------------------------------------------------------------------- + + +def _make_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + q_lora_rank=64, + o_groups=8, + o_lora_rank=64, + csa_compress_ratios=None, + csa_window_size=8, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + **extra_config_kwargs, +): + """Create an MLATransformerConfig for DSv4 hybrid attention tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0, 4, 128, 4] + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=q_lora_rank, + kv_lora_rank=v_head_dim - qk_pos_emb_head_dim, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + **extra_config_kwargs, + ) + + +def _make_attention_spec(config): + """Build the full DSv4HybridSelfAttention ModuleSpec using the canonical spec builder.""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + +def test_module_spec_is_built_from_explicit_backend(): + """The neutral spec builder should use only its explicitly supplied backend.""" + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + Compressor, + CSAIndexer, + ) + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + class Linear: + pass + + class ColumnParallelLinear: + pass + + class RowParallelLinear: + pass + + class Norm: + pass + + class Backend: + def linear(self): + return Linear + + def column_parallel_linear(self): + return ColumnParallelLinear + + def row_parallel_linear(self): + return RowParallelLinear + + def layer_norm(self, rms_norm=False, for_qk=False, has_residual=False): + return Norm + + spec = get_dsv4_hybrid_module_spec_for_backend(_make_config(), Backend()) + + assert spec.module is DSv4HybridSelfAttention + assert spec.submodules.linear_q_down_proj is Linear + assert spec.submodules.linear_q_up_proj is ColumnParallelLinear + assert spec.submodules.linear_kv_proj is ColumnParallelLinear + assert spec.submodules.linear_proj is RowParallelLinear + assert spec.submodules.core_attention.module is CompressedSparseAttention + assert spec.submodules.core_attention.submodules.compressor.module is Compressor + assert spec.submodules.core_attention.submodules.indexer.module is CSAIndexer + + +def test_config_includes_mtp_ratio_and_derives_dimensions(): + """DSv4 config should account for MTP and derive its shared Q/KV content width.""" + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128]) + + expected_content_dim = config.v_head_dim - config.qk_pos_emb_head_dim + assert config.qk_head_dim == expected_content_dim + assert config.kv_lora_rank == expected_content_dim + assert config.hetereogenous_dist_checkpoint is True + + +def test_config_rejects_context_parallelism(): + """The SBHD slice should fail early instead of silently accepting unsupported CP.""" + with pytest.raises(AssertionError, match="does not support context parallelism"): + _make_config(context_parallel_size=2) + + +def test_config_rejects_fused_backend_in_native_slice(): + """Fused DSv4 backends belong to the follow-up kernel-integration slice.""" + with pytest.raises(ValueError, match="requires dsa_kernel_backend='none'"): + _make_config(dsa_kernel_backend="cudnn") + + +def test_config_accepts_hybrid_model_ratio_tail(): + """HybridModel may expand each MTP depth into multiple attention layers.""" + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128, 4]) + assert config.csa_compress_ratios == [0, 4, 128, 4] + + +def test_hybrid_dsv4_stack_spec_assigns_fixed_ratios_and_preserves_main_layers(): + """C/H/W use DSv4 while ordinary D DSA and + MLA remain unchanged.""" + from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_dsv4_stack_spec, + hybrid_stack_spec, + ) + + stack_spec = hybrid_dsv4_stack_spec(_make_config()) + submodules = stack_spec.submodules + baseline = hybrid_stack_spec.submodules + + assert submodules.dsa_layer is baseline.dsa_layer + assert submodules.mla_layer is baseline.mla_layer + assert submodules.csa_layer.submodules.self_attention.params["compress_ratio"] == 4 + assert submodules.hca_layer.submodules.self_attention.params["compress_ratio"] == 128 + assert submodules.window_layer.submodules.self_attention.params["compress_ratio"] == 0 + + +def test_constructor_requires_explicit_process_groups(): + """Production DSv4 construction must not read process groups from global MPU state.""" + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + with pytest.raises(ValueError, match="explicit ProcessGroupCollection"): + DSv4HybridSelfAttention(config=None, submodules=None, layer_number=1) + + +def _build_attention(config, layer_number, pg_collection, **kwargs): + """Instantiate a DSv4HybridSelfAttention from config.""" + from megatron.core.transformer.spec_utils import build_module + + spec = _make_attention_spec(config) + return build_module( + spec, config=config, layer_number=layer_number, pg_collection=pg_collection, **kwargs + ) + + +# =========================================================================== +# Constructor tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionConstructor: + """Test construction of DSv4HybridSelfAttention in the supported TP=1 configuration.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_basic_construction(self): + """Verify the layer builds and has the expected sub-modules.""" + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert isinstance(attn, DSv4HybridSelfAttention) + assert hasattr(attn, 'linear_q_down_proj') + assert hasattr(attn, 'linear_q_up_proj') + assert hasattr(attn, 'linear_kv_proj') + assert hasattr(attn, 'linear_proj') + assert hasattr(attn, 'linear_o_group_proj') + assert hasattr(attn, 'core_attention') + assert hasattr(attn, 'q_layernorm') + assert hasattr(attn, 'kv_layernorm') + + def test_q_head_dim_equals_v_head_dim(self): + """q_head_dim must equal v_head_dim for DSv4 hybrid.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert attn.q_head_dim == config.v_head_dim + + def test_current_main_constructor_kwargs(self): + """Current TransformerLayer forwards module names and pipeline offsets.""" + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention( + config, + layer_number=1, + pg_collection=pg, + pp_layer_offset=0, + name="decoder.layers.0.self_attention", + ) + + assert attn._pp_layer_offset == 0 + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_rope_base_varies_with_compress_ratio(self, layer_number): + """Layers with compress_ratio > 1 should use csa_compress_rotary_base.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + ratios = [0, 4, 128, 4] + config = _make_config(csa_compress_ratios=ratios) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=layer_number, pg_collection=pg) + + ratio = ratios[layer_number - 1] + if ratio > 1: + expected_base = config.csa_compress_rotary_base + else: + expected_base = config.rotary_base + + # inv_freq is derived from rotary_base; verify the correct base was used + dim = config.qk_pos_emb_head_dim + recomputed_inv_freq = 1.0 / ( + expected_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + assert torch.allclose( + attn.rotary_pos_emb.inv_freq.cpu(), recomputed_inv_freq, rtol=1e-5, atol=1e-5 + ) + + +# =========================================================================== +# Forward / backward tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionForwardBackward: + """Test forward and backward passes of DSv4HybridSelfAttention.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_forward_output_shape(self, layer_number): + """Forward should produce [sq, b, hidden_size] output.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + + @pytest.mark.parametrize("layer_number", [1, 2]) + def test_backward_gradient_flow(self, layer_number): + """Backward should produce gradients for all trainable parameters.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.train() + + hidden = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + output, bias = attn(hidden_states=hidden, attention_mask=None) + loss = output.sum() + loss.backward() + + assert hidden.grad is not None, "No gradient on hidden_states" + for name, param in attn.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" + + def test_eval_mode(self): + """Forward should work in eval mode.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + attn.eval() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + with torch.no_grad(): + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert not torch.isnan(output).any() + + def test_different_seq_lengths(self): + """Forward should handle various sequence lengths.""" + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=2, pg_collection=self.pg).cuda() + + for seq_len in [64, 128, 256]: + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + + +# =========================================================================== +# get_query_key_value_tensors tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridQKV: + """Test get_query_key_value_tensors internals.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config() + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_qkv_shapes(self): + """Query, key, value should have correct shapes.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) + + n_heads = self.config.num_attention_heads + v_dim = self.config.v_head_dim + + assert q.shape == (seq_len, batch_size, n_heads, v_dim) + # key and value are single-head (MQA-style) with an extra head dim + assert k.shape[-1] == v_dim + assert v.shape[-1] == v_dim + + def test_key_equals_value(self): + """In the wkv path, key and value should be the same tensor.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + assert torch.equal(k, v), "key and value should be identical in wkv path" + + +# =========================================================================== +# Grouped output projection tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridGroupedOutput: + """Test that grouped output projection (wo_a) parameters are created.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_o_group_proj_shape(self): + """linear_o_group_proj should have the correct shape.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + o_groups = 8 + o_lora_rank = 64 + config = _make_config(o_groups=o_groups, o_lora_rank=o_lora_rank) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + expected_out = o_groups * o_lora_rank + expected_in = (config.v_head_dim * config.num_attention_heads) // o_groups + assert attn.linear_o_group_proj.shape == (expected_out, expected_in) + assert attn.linear_o_group_proj.requires_grad + + +# =========================================================================== +# apply_rope_fusion tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridRopeFusion: + """Test that apply_rope_fusion=True works for both yarn and non-yarn layers. + + DSv4 Hybrid uses YarnRotaryEmbedding for layers with compress_ratio > 1 + and standard RotaryEmbedding for layers with compress_ratio <= 1. The + fused RoPE path must obtain cos/sin from both embedding classes via + get_cached_cos_sin. + + compress_ratios=[0, 4, 128, 4]: layer 1 has ratio 0 (standard + RotaryEmbedding), layers 2-4 have ratio > 1 (YarnRotaryEmbedding). + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_rope_fusion_forward_backward_parity(self): + """Fused RoPE forward/backward succeeds and matches the unfused path.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + fused_config = _make_config(apply_rope_fusion=True) + attn_fused = _build_attention(fused_config, layer_number=4, pg_collection=self.pg).cuda() + attn_fused.train() + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + unfused_config = _make_config(apply_rope_fusion=False) + attn_unfused = _build_attention( + unfused_config, layer_number=4, pg_collection=self.pg + ).cuda() + attn_unfused.train() + + hidden = torch.randn( + seq_len, batch_size, fused_config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + out_fused, _ = attn_fused(hidden_states=hidden, attention_mask=None) + out_unfused, _ = attn_unfused(hidden_states=hidden, attention_mask=None) + + assert out_fused.shape == (seq_len, batch_size, fused_config.hidden_size) + assert torch.isfinite(out_fused).all() + # The remaining difference is bf16 accumulation order between the fused + # Triton kernel and eager PyTorch operations. + torch.testing.assert_close(out_fused, out_unfused, atol=3e-2, rtol=3e-2) + + hidden_fused = hidden.detach().clone().requires_grad_(True) + hidden_unfused = hidden.detach().clone().requires_grad_(True) + + attn_fused(hidden_states=hidden_fused, attention_mask=None)[0].sum().backward() + attn_unfused(hidden_states=hidden_unfused, attention_mask=None)[0].sum().backward() + + assert hidden_fused.grad is not None + for name, param in attn_fused.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" diff --git a/tests/unit_tests/transformer/test_hyper_connection_recompute.py b/tests/unit_tests/transformer/test_hyper_connection_recompute.py new file mode 100644 index 00000000000..d9ed2cfa7e0 --- /dev/null +++ b/tests/unit_tests/transformer/test_hyper_connection_recompute.py @@ -0,0 +1,434 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +""" +Unit tests for HyperConnection block-level recomputation. + +Tests the following functionality: +1. HyperConnectionModule._forward_with_checkpoint correctness +2. HyperConnectionModule.apply_h_post with CheckpointManager +3. Multiple HyperConnectionModules chained with a single CheckpointManager +4. Partial checkpoint (last layer not checkpointed) +5. TransformerConfig 'mhc' in recompute_modules option +""" + +import pytest +import torch + +from megatron.core.tensor_parallel.random import CheckpointManager, model_parallel_cuda_manual_seed +from megatron.core.transformer.hyper_connection import HyperConnectionModule +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +class TestHyperConnectionCheckpoint: + """Test HyperConnectionModule checkpoint functionality.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_hyper_connection_module(self, hidden_size=64, num_residual_streams=4): + """Create a HyperConnectionModule for testing.""" + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_residual_streams, + mhc_sinkhorn_iterations=5, # Fewer iterations for faster tests + mhc_init_gating_factor=0.01, + ) + module = HyperConnectionModule(config=config, layer_number=1) + module.cuda() + return module + + def test_apply_h_res_uses_h_res_transpose(self): + """apply_h_res should compute H_res.T @ residual.""" + module = self._create_hyper_connection_module(hidden_size=4, num_residual_streams=2) + h_res = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], device='cuda') + residual = torch.tensor([[[10.0, 100.0, 3.0, 4.0, 1.0, 2.0, 5.0, 6.0]]], device='cuda') + expected = torch.tensor( + [[[13.0, 106.0, 18.0, 22.0, 24.0, 208.0, 26.0, 32.0]]], device='cuda' + ) + + mixed = module.apply_h_res(h_res, residual) + + torch.testing.assert_close(mixed, expected, atol=0.0, rtol=0.0) + + def test_forward_preserves_three_tuple_api_and_hybrid_can_request_residual(self): + module = self._create_hyper_connection_module(hidden_size=8, num_residual_streams=2) + hidden_states = torch.randn(4, 1, 16, device='cuda', requires_grad=True) + + compatible_output = module(hidden_states) + hybrid_output = module(hidden_states, return_residual=True) + + assert len(compatible_output) == 3 + assert len(hybrid_output) == 4 + for compatible, hybrid in zip(compatible_output, hybrid_output[:3]): + torch.testing.assert_close(compatible, hybrid) + assert hybrid_output[3].shape == hidden_states.shape + + def test_forward_normal_vs_checkpoint_correctness(self): + """ + Test that _forward_with_checkpoint produces the same outputs as _forward_normal. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + residual = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + # Clone inputs for comparison + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + residual_ckpt = residual.detach().clone().requires_grad_(True) + + # Forward without checkpoint (reference) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref, residual_ref = module._forward_normal(hidden_states) + mixed_ref = module.apply_h_res(h_res_ref, residual) + loss_ref = aggregated_ref.sum() + mixed_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states.grad.clone() + grad_residual_ref = residual.grad.clone() + + # Forward with checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt, residual_ckpt_out = ( + module._forward_with_checkpoint(hidden_states_ckpt, manager) + ) + mixed_ckpt = module.apply_h_res(h_res_ckpt, residual_ckpt) + # Calculate loss before discarding outputs + loss_ckpt = aggregated_ckpt.sum() + mixed_ckpt.sum() + h_post_ckpt.sum() + + # Register unified recompute hook + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + # Backward pass + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5), ( + f"Hidden states gradients mismatch:\n" + f"Checkpoint: {grad_hidden_ckpt}\n" + f"Reference: {grad_hidden_ref}" + ) + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-5), ( + f"Residual gradients mismatch:\n" + f"Checkpoint: {grad_residual_ckpt}\n" + f"Reference: {grad_residual_ref}" + ) + + def test_apply_h_post_with_checkpoint(self): + """ + Test that apply_h_post with manager produces correct gradients. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + x = torch.randn(seq_len, batch_size, hidden_size, device='cuda', requires_grad=True) + bias = torch.randn(hidden_size, device='cuda') + h_post = torch.randn(seq_len, batch_size, num_streams, device='cuda', requires_grad=True) + + # Clone inputs + x_ckpt = x.detach().clone().requires_grad_(True) + h_post_ckpt = h_post.detach().clone().requires_grad_(True) + + # Reference: without checkpoint (manager=None) + torch.manual_seed(42) + x_out_ref, bias_out_ref = module.apply_h_post((x, bias), h_post, manager=None) + loss_ref = x_out_ref.sum() + if bias_out_ref is not None: + loss_ref = loss_ref + bias_out_ref.sum() + loss_ref.backward() + grad_x_ref = x.grad.clone() + grad_h_post_ref = h_post.grad.clone() + + # With checkpoint (manager provided) + torch.manual_seed(42) + manager = CheckpointManager() + x_out_ckpt, bias_out_ckpt = module.apply_h_post( + (x_ckpt, bias), h_post_ckpt, manager=manager + ) + loss_ckpt = x_out_ckpt.sum() + if bias_out_ckpt is not None: + loss_ckpt = loss_ckpt + bias_out_ckpt.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_h_post_ckpt = h_post_ckpt.grad.clone() + + # Verify gradients + assert torch.allclose(grad_x_ckpt, grad_x_ref, atol=1e-5) + assert torch.allclose(grad_h_post_ckpt, grad_h_post_ref, atol=1e-5) + + def test_forward_with_manager_parameter(self): + """ + Test forward() method with mhc_recompute_manager parameter. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + module = self._create_hyper_connection_module(hidden_size, num_streams) + + # Create input tensors + hidden_states = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + # Clone inputs + hidden_states_ckpt = hidden_states.detach().clone().requires_grad_(True) + + # Reference: forward without manager (uses _forward_normal) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module.forward( + hidden_states, mhc_recompute_manager=None + ) + loss_ref = aggregated_ref.sum() + h_res_ref.sum() + h_post_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states.grad.clone() + + # With manager (uses _forward_with_checkpoint) + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + loss_ckpt = aggregated_ckpt.sum() + h_res_ckpt.sum() + h_post_ckpt.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5) + + +class TestMHCBlockRecomputeIntegration: + """Test CheckpointManager integration with HyperConnection.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_multiple_hyper_connections_in_chain(self): + """ + Test that multiple HyperConnectionModules can be chained together + with a single CheckpointManager. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + # Create multiple HyperConnection modules (simulating multiple layers) + config = TransformerConfig( + num_layers=4, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + + modules = [ + HyperConnectionModule(config=config, layer_number=i + 1).cuda() for i in range(3) + ] + + # Create input tensors + hidden_states_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + residual_ref = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + # Reference: forward without checkpoint + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + h = hidden_states_ref + r = residual_ref + for module in modules: + agg, h_res, h_post = module.forward(h, mhc_recompute_manager=None) + agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=None) + mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] + h = agg + mixed + r = h + + loss_ref = h.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + # With checkpoint using single manager + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointManager() + + h = hidden_states_ckpt + r = residual_ckpt + for module in modules: + agg, h_res, h_post = module.forward(h, mhc_recompute_manager=manager) + agg, _ = module.apply_h_post((0.1 * agg, None), h_post, manager=manager) + mixed = module.apply_h_res(h_res, r) # Apply h_res to get mixed [s, b, n*C] + h = agg + mixed + r = h + + loss_ckpt = h.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + loss_ckpt.backward() + + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + # Verify gradients + assert torch.allclose( + grad_hidden_ckpt, grad_hidden_ref, atol=1e-4 + ), f"Chained HyperConnection hidden gradients mismatch" + assert torch.allclose( + grad_residual_ckpt, grad_residual_ref, atol=1e-4 + ), f"Chained HyperConnection residual gradients mismatch" + + def test_partial_checkpoint_last_layer_not_checkpointed(self): + """ + Test that when is_last_layer_in_block=True, the final output is NOT checkpointed. + This simulates the TransformerBlock behavior where the last layer's MLP BDA + serves as the hook_tensor for unified recompute. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + ) + + module = HyperConnectionModule(config=config, layer_number=1).cuda() + + hidden_states_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + residual_ref = torch.randn( + seq_len, batch_size, num_streams * hidden_size, device='cuda', requires_grad=True + ) + + hidden_states_ckpt = hidden_states_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + # Reference + torch.manual_seed(42) + torch.cuda.manual_seed(42) + aggregated_ref, h_res_ref, h_post_ref = module.forward( + hidden_states_ref, mhc_recompute_manager=None + ) + aggregated_ref, _ = module.apply_h_post( + (0.1 * aggregated_ref, None), h_post_ref, manager=None + ) + mixed_ref = module.apply_h_res( + h_res_ref, residual_ref + ) # Apply h_res to get mixed [s, b, n*C] + # Simulate BDA that is NOT checkpointed (last layer) + output_ref = aggregated_ref + 0.5 * mixed_ref + loss_ref = output_ref.sum() + loss_ref.backward() + grad_hidden_ref = hidden_states_ref.grad.clone() + + # With manager - checkpoint everything except final output + torch.manual_seed(42) + torch.cuda.manual_seed(42) + manager = CheckpointManager() + aggregated_ckpt, h_res_ckpt, h_post_ckpt = module.forward( + hidden_states_ckpt, mhc_recompute_manager=manager + ) + + aggregated_ckpt, _ = module.apply_h_post( + (0.1 * aggregated_ckpt, None), h_post_ckpt, manager=manager + ) + mixed_ckpt = module.apply_h_res( + h_res_ckpt, residual_ckpt + ) # Apply h_res to get mixed [s, b, n*C] + # Simulate BDA that is NOT checkpointed (last layer) - this is the hook_tensor + output_ckpt = aggregated_ckpt + 0.5 * mixed_ckpt + + # Register unified recompute on the output (which is not checkpointed) + manager.discard_all_outputs_and_register_unified_recompute(output_ckpt) + + loss_ckpt = output_ckpt.sum() + loss_ckpt.backward() + grad_hidden_ckpt = hidden_states_ckpt.grad.clone() + + # Verify gradients match + assert torch.allclose(grad_hidden_ckpt, grad_hidden_ref, atol=1e-5) + + +class TestTransformerConfigRecomputeMhc: + """Test 'mhc' in recompute_modules configuration.""" + + def test_config_default_value(self): + """Test that 'mhc' is not in recompute_modules by default.""" + config = TransformerConfig(num_layers=2, hidden_size=64, num_attention_heads=4) + assert "mhc" not in config.recompute_modules + + def test_config_enable_mhc_recompute(self): + """Test enabling 'mhc' in recompute_modules.""" + config = TransformerConfig( + num_layers=2, + hidden_size=64, + num_attention_heads=4, + enable_hyper_connections=True, + num_residual_streams=4, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + ) + assert "mhc" in config.recompute_modules + assert config.enable_hyper_connections is True + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/transformer/test_mhc_block_manager.py b/tests/unit_tests/transformer/test_mhc_block_manager.py new file mode 100644 index 00000000000..0d4f40bba7d --- /dev/null +++ b/tests/unit_tests/transformer/test_mhc_block_manager.py @@ -0,0 +1,522 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core.tensor_parallel.random import ( + CheckpointManager, + CheckpointWithoutOutput, + CheckpointWithoutOutputManager, + initialize_rng_tracker, +) +from tests.unit_tests.test_utilities import Utils + + +class TestCheckpointWithoutOutputManagerAPI: + """Test CheckpointWithoutOutput integration with CheckpointWithoutOutputManager.""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_reviewed_manager_name_is_compatible_alias(self): + """The #4531 manager name remains a compatible public alias.""" + assert CheckpointWithoutOutputManager is CheckpointManager + assert isinstance(CheckpointWithoutOutputManager(), CheckpointManager) + + def test_auto_register(self): + """CheckpointWithoutOutput auto-registers to manager when ckpt_manager is provided.""" + manager = CheckpointWithoutOutputManager() + + def func(x): + return x * 2 + 1 + + input_t = torch.randn(4, 4, device='cuda', requires_grad=True) + + ckpt = CheckpointWithoutOutput(ckpt_manager=manager) + y = ckpt.checkpoint(func, input_t) + + assert len(manager.checkpoints) == 1 + assert manager.checkpoints[0] is ckpt + + ckpt2 = CheckpointWithoutOutput(ckpt_manager=manager) + y2 = ckpt2.checkpoint(torch.nn.functional.gelu, y) + + assert len(manager.checkpoints) == 2 + assert manager.checkpoints[1] is ckpt2 + + loss = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss) + loss.backward() + + assert input_t.grad is not None + + def test_discard_is_noop_with_manager(self): + """discard_output_and_register_recompute is a NO-OP when ckpt_manager is set.""" + manager = CheckpointWithoutOutputManager() + + def func1(x): + return x * 2 + + def func2(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + y1_ref = func1(input_ref) + y2_ref = func2(y1_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + ckpt1 = CheckpointWithoutOutput(ckpt_manager=manager) + y1 = ckpt1.checkpoint(func1, input_ckpt) + ckpt1.discard_output_and_register_recompute(y1) + + ckpt2 = CheckpointWithoutOutput(ckpt_manager=manager) + y2 = ckpt2.checkpoint(func2, y1) + ckpt2.discard_output_and_register_recompute(y2) + + assert y1.untyped_storage().size() > 0, "y1 should NOT be discarded yet" + assert y2.untyped_storage().size() > 0, "y2 should NOT be discarded yet" + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert y1.untyped_storage().size() == 0, "y1 should be discarded after manager call" + assert y2.untyped_storage().size() == 0, "y2 should be discarded after manager call" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6) + + def test_backward_compat_without_manager(self): + """CheckpointWithoutOutput without ckpt_manager should work exactly as before.""" + + def func(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + y_ref = func(input_ref) + z_ref = y_ref * 2 + loss_ref = z_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + ckpt = CheckpointWithoutOutput() + y = ckpt.checkpoint(func, input_ckpt) + z = y * 2 + ckpt.discard_output_and_register_recompute(z) + + assert y.untyped_storage().size() == 0 + + loss_ckpt = z.sum() + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6) + + def test_error_handling(self): + """CheckpointWithoutOutputManager rejects invalid add_checkpoint calls.""" + manager = CheckpointWithoutOutputManager() + + with pytest.raises(TypeError): + manager.add_checkpoint("not a checkpoint") + + ckpt = CheckpointWithoutOutput() + with pytest.raises(ValueError): + manager.add_checkpoint(ckpt) + + +class TestCheckpointManagerSequentialChain: + """Test CheckpointWithoutOutputManager with sequential checkpoint chains.""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_basic_sequential_chain(self): + """Three sequential checkpoints: gradients match non-checkpointed version.""" + + def func1(x): + return x * 2 + 1 + + def func2(x): + return torch.nn.functional.gelu(x) + + def func3(x): + return x * x + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + y1_ref = func1(input_ref) + y2_ref = func2(y1_ref) + y3_ref = func3(y2_ref) + loss_ref = y3_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + manager = CheckpointWithoutOutputManager() + + y1 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func1, input_ckpt) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func2, y1) + y3 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func3, y2) + + loss_ckpt = y3.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert y1.untyped_storage().size() == 0, "y1 storage should be released" + assert y2.untyped_storage().size() == 0, "y2 storage should be released" + assert y3.untyped_storage().size() == 0, "y3 storage should be released" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose( + grad_ckpt, grad_ref, atol=1e-6 + ), f"Gradients mismatch!\nWith manager: {grad_ckpt}\nReference: {grad_ref}" + + def test_sequential_chain_with_dropout(self): + """RNG state is restored during recompute so dropout gradients match.""" + + def func_with_dropout(x): + return torch.nn.functional.dropout(x, p=0.3, training=True) + + def func2(x): + return torch.nn.functional.gelu(x) + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + y1_ref = func_with_dropout(input_ref) + y2_ref = func2(y1_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + manager = CheckpointWithoutOutputManager() + + y1 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_with_dropout, input_ckpt) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func2, y1) + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose( + grad_ckpt, grad_ref, atol=1e-6 + ), f"Gradients with dropout mismatch!\nWith manager: {grad_ckpt}\nReference: {grad_ref}" + + def test_multiple_outputs(self): + """CheckpointWithoutOutputManager handles functions that return multiple outputs.""" + + def func_multi_output(x): + return x * 2, x + 1 + + def func_combine(a, b): + return a + b + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + y1a_ref, y1b_ref = func_multi_output(input_ref) + y2_ref = func_combine(y1a_ref, y1b_ref) + loss_ref = y2_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + manager = CheckpointWithoutOutputManager() + + y1a, y1b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + func_multi_output, input_ckpt + ) + y2 = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_combine, y1a, y1b) + + loss_ckpt = y2.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with multiple outputs!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + +class TestCheckpointManagerPartialCheckpoint: + """Test CheckpointWithoutOutputManager with partial checkpointing (some ops not checkpointed).""" + + def setup_method(self, method): + Utils.initialize_model_parallel() + initialize_rng_tracker(force_reset=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_partial_checkpoint(self): + """ + Only f and h are checkpointed; g is a regular operation. + + Computation chain: + a --[f]--> b --[g]--> c --[h]--> d --[sum]--> loss + """ + + def func_f(x): + return torch.nn.functional.gelu(x * 2 + 1) + + def func_g(x): + return x * 3 - 2 + + def func_h(x): + return torch.sigmoid(x) + x + + input_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + b_ref = func_f(input_ref) + c_ref = func_g(b_ref) + d_ref = func_h(c_ref) + loss_ref = d_ref.sum() + loss_ref.backward() + grad_ref = input_ref.grad.clone() + + input_ckpt = input_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + b = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_f, input_ckpt) + c = func_g(b) + d = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(func_h, c) + + loss_ckpt = d.sum() + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert b.untyped_storage().size() == 0, "b storage should be released" + assert d.untyped_storage().size() == 0, "d storage should be released" + assert c.untyped_storage().size() > 0, "c storage should NOT be released (not checkpointed)" + + loss_ckpt.backward() + grad_ckpt = input_ckpt.grad.clone() + + assert torch.allclose(grad_ckpt, grad_ref, atol=1e-6), ( + f"Gradients mismatch with partial checkpoint!\n" + f"With manager: {grad_ckpt}\nReference: {grad_ref}" + ) + + def test_partial_checkpoint_with_tuple_output(self): + """ + Mimics HyperConnection's computation pattern with tuple outputs. + + - compute_mappings: checkpointed, returns tuple (h_pre, h_post, h_res) + - aggregate: NOT checkpointed + - apply_h_res: checkpointed + - apply_h_post: checkpointed + """ + + def compute_mappings(x): + h_pre = torch.sigmoid(x.mean(dim=-1, keepdim=True).expand_as(x)) + h_post = torch.tanh(x.sum(dim=-1, keepdim=True).expand_as(x)) + h_res = torch.relu(x) + return h_pre, h_post, h_res + + def aggregate(x, h_pre): + return x * h_pre + + def apply_h_res(h_res, residual): + return h_res + residual * 0.5 + + def apply_h_post(y, h_post): + return y * h_post + y + + x_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + residual_ref = torch.randn(4, 4, device='cuda', requires_grad=True) + + h_pre_ref, h_post_ref, h_res_ref = compute_mappings(x_ref) + agg_ref = aggregate(x_ref, h_pre_ref) + y_ref = torch.nn.functional.gelu(agg_ref) + mixed_ref = apply_h_res(h_res_ref, residual_ref) + output_ref = apply_h_post(y_ref, h_post_ref) + final_ref = output_ref + mixed_ref + loss_ref = final_ref.sum() + loss_ref.backward() + grad_x_ref = x_ref.grad.clone() + grad_residual_ref = residual_ref.grad.clone() + + x_ckpt = x_ref.detach().clone().requires_grad_(True) + residual_ckpt = residual_ref.detach().clone().requires_grad_(True) + + manager = CheckpointWithoutOutputManager() + + h_pre, h_post, h_res = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + compute_mappings, x_ckpt + ) + agg = aggregate(x_ckpt, h_pre) + y = torch.nn.functional.gelu(agg) + mixed = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint( + apply_h_res, h_res, residual_ckpt + ) + output = CheckpointWithoutOutput(ckpt_manager=manager).checkpoint(apply_h_post, y, h_post) + + final = output + mixed + loss_ckpt = final.sum() + + manager.discard_all_outputs_and_register_unified_recompute(loss_ckpt) + + assert h_pre.untyped_storage().size() == 0, "h_pre storage should be released" + assert h_post.untyped_storage().size() == 0, "h_post storage should be released" + assert h_res.untyped_storage().size() == 0, "h_res storage should be released" + assert mixed.untyped_storage().size() == 0, "mixed storage should be released" + assert output.untyped_storage().size() == 0, "output storage should be released" + + assert agg.untyped_storage().size() > 0, "agg storage should NOT be released" + assert y.untyped_storage().size() > 0, "y storage should NOT be released" + + loss_ckpt.backward() + grad_x_ckpt = x_ckpt.grad.clone() + grad_residual_ckpt = residual_ckpt.grad.clone() + + assert torch.allclose( + grad_x_ckpt, grad_x_ref, atol=1e-6 + ), f"Gradients for x mismatch!\nWith manager: {grad_x_ckpt}\nReference: {grad_x_ref}" + assert torch.allclose(grad_residual_ckpt, grad_residual_ref, atol=1e-6), ( + f"Gradients for residual mismatch!\n" + f"With manager: {grad_residual_ckpt}\nReference: {grad_residual_ref}" + ) + + +# ============================================================================ +# Block-level mHC recompute coverage +# ============================================================================ +# +# These tests instantiate a full ``TransformerBlock`` with mHC enabled to +# exercise: +# * ``_build_mhc_recompute_layer_plan`` (per-layer ``CheckpointWithoutOutputManager`` +# allocation, including the ``mhc_recompute_layer_num`` boundary case), +# * ``_finalize_mhc_recompute_layer`` (manager finalization at block end), +# * the ``HyperConnectionModule.input_expand`` / ``output_contract`` calls +# in ``TransformerBlock.forward`` for ``pre_process`` / ``post_process`` +# stages. +# +# Single-process (no PP) so they can run on a single-GPU CI lane. + + +class TestTransformerBlockMHCRecompute: + """End-to-end ``TransformerBlock`` forward with mHC selective recompute.""" + + def setup_method(self, method): + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _make_mhc_block(num_layers, num_streams=4, mhc_recompute_layer_num=None): + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + ) + from megatron.core.transformer.hyper_connection import HyperConnectionModule + from megatron.core.transformer.transformer_block import TransformerBlock + from megatron.core.transformer.transformer_config import TransformerConfig + from megatron.core.transformer.transformer_layer import HyperConnectionTransformerLayer + + config = TransformerConfig( + num_layers=num_layers, + hidden_size=64, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + mhc_recompute_layer_num=mhc_recompute_layer_num, + recompute_granularity='selective', + recompute_modules=['mhc'], + hidden_dropout=0.0, + attention_dropout=0.0, + ) + spec = get_gpt_layer_with_transformer_engine_spec() + spec.module = HyperConnectionTransformerLayer + spec.submodules.self_attention_hyper_connection = HyperConnectionModule + spec.submodules.mlp_hyper_connection = HyperConnectionModule + return TransformerBlock(config, spec, pre_process=True, post_process=True).cuda(), config + + def _check_recompute_plan(self, block, expected_block_ends): + """Drive ``_build_mhc_recompute_layer_plan`` directly and check the boundary list.""" + block.train() + managers, ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=True) + assert len(managers) == len(block.layers) + assert ends == expected_block_ends, f"got {ends}, expected {expected_block_ends}" + # Layers in the same recompute block share a manager; new block → new manager. + last_was_end = True + last_mgr = None + for mgr, end in zip(managers, ends): + assert mgr is not None + if last_was_end: + assert mgr is not last_mgr, "new recompute block should get a new manager" + else: + assert mgr is last_mgr, "layers within a recompute block share a manager" + last_was_end = end + last_mgr = mgr + + def test_recompute_plan_no_layer_num(self): + """Without ``mhc_recompute_layer_num`` only the final layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4) + self._check_recompute_plan(block, expected_block_ends=[False, False, False, True]) + + def test_recompute_plan_with_layer_num(self): + """With ``mhc_recompute_layer_num=2`` every other layer ends a recompute block.""" + block, _ = self._make_mhc_block(num_layers=4, mhc_recompute_layer_num=2) + self._check_recompute_plan(block, expected_block_ends=[False, True, False, True]) + + def test_recompute_plan_disabled(self): + """``use_mhc_recompute=False`` returns an all-None / all-False plan.""" + block, _ = self._make_mhc_block(num_layers=3) + managers, ends = block._build_mhc_recompute_layer_plan(use_mhc_recompute=False) + assert managers == [None, None, None] + assert ends == [False, False, False] + + def test_block_forward_input_expand_output_contract(self): + """Forward exercises ``input_expand`` (pre) and ``output_contract`` (post).""" + block, config = self._make_mhc_block(num_layers=2, mhc_recompute_layer_num=2) + block.train() + + seq_len = 8 + batch_size = 2 + # Input is [s, b, hidden_size]; the block must expand to [s, b, n*hidden_size] + # internally, then contract back to [s, b, hidden_size] before final layernorm. + hidden_states = torch.randn( + seq_len, batch_size, config.hidden_size, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=torch.bool, device='cuda') + + out = block(hidden_states=hidden_states, attention_mask=attention_mask) + assert out.shape == hidden_states.shape, ( + f"output_contract should restore original shape, got {tuple(out.shape)} " + f"vs expected {tuple(hidden_states.shape)}" + ) + # Backward should flow through the recompute path without error. + out.sum().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() diff --git a/tests/unit_tests/transformer/test_module.py b/tests/unit_tests/transformer/test_module.py index 92f15b2f46d..5faf6c81ef1 100644 --- a/tests/unit_tests/transformer/test_module.py +++ b/tests/unit_tests/transformer/test_module.py @@ -4,7 +4,7 @@ import torch from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.module import Float16Module, MegatronModule +from megatron.core.transformer.module import Float16Module, MegatronModule, mark_keep_in_fp32 from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -163,3 +163,18 @@ def test_bf16_module(self): x = torch.ones((2, 2)).cuda() # inputs are converted to bf16 then outputs are converted to fp32 assert bf16_module(x).dtype == torch.float32 + + @pytest.mark.parametrize( + ('precision', 'dtype'), [('fp16', torch.float16), ('bf16', torch.bfloat16)] + ) + def test_keep_in_fp32_params(self, precision, dtype): + transformer_config = self.transformer_config + megatron_module = self.megatron_module + megatron_module.fp32_param = mark_keep_in_fp32( + torch.nn.Parameter(torch.zeros(4, dtype=torch.float32, device='cuda')) + ) + setattr(transformer_config, precision, True) + float16_module = Float16Module(config=transformer_config, module=megatron_module) + + assert float16_module.module.linear.weight.dtype == dtype + assert float16_module.module.fp32_param.dtype == torch.float32 diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index ee1eb02267f..6702b3c6992 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -22,6 +22,7 @@ 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, @@ -1625,3 +1626,52 @@ 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 the learned n-stream to one-stream mHC contraction.""" + + 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() + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_shape_and_dtype(self, dtype): + hidden_size, n_streams = 32, 4 + hidden_states = torch.randn(8, 2, n_streams * hidden_size, device="cuda", dtype=dtype) + 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") + + output = learned_output_contract(hidden_states, head_fn, base, scale, n_streams, eps=1e-6) + + assert output.shape == (8, 2, hidden_size) + assert output.dtype == dtype + + def test_gradient_and_reference(self): + hidden_size, n_streams, eps = 8, 2, 1e-6 + hidden_states = torch.randn( + 2, 1, 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) + rsqrt = torch.rsqrt(hidden_states.square().mean(-1, keepdim=True) + eps) + mixes = torch.nn.functional.linear(hidden_states, head_fn) * rsqrt + weights = torch.sigmoid(mixes * scale + base) + eps + expected = torch.sum( + weights.unsqueeze(-1) + * hidden_states.view(*hidden_states.shape[:-1], n_streams, hidden_size), + dim=-2, + ) + torch.testing.assert_close(output, expected) + + output.sum().backward() + for tensor in (hidden_states, head_fn, base, scale): + assert tensor.grad is not None + assert torch.count_nonzero(tensor.grad) > 0 diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 93650cf13b0..033f36597c1 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc @@ -16,6 +16,7 @@ ) from megatron.core.tensor_parallel.random import ( HAVE_TE, + CheckpointWithoutOutputManager, initialize_rng_tracker, model_parallel_cuda_manual_seed, ) @@ -23,6 +24,7 @@ from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, TransformerLayer, get_transformer_layer_offset, ) @@ -30,6 +32,44 @@ from tests.unit_tests.test_utilities import Utils +def _make_mhc_layer_spec(**kwargs): + """Build a layer spec with HyperConnectionModule submodules. + + The ``enable_hyper_connection`` kwarg on ``gpt_layer_specs`` is added by + the GPT-wiring follow-up split, so this helper patches the mHC submodules + directly to keep the unit tests self-contained for this split. + """ + from megatron.core.transformer.hyper_connection import HyperConnectionModule + + layer_spec = get_gpt_layer_with_transformer_engine_spec(**kwargs) + layer_spec.module = HyperConnectionTransformerLayer + layer_spec.submodules.self_attention_hyper_connection = HyperConnectionModule + layer_spec.submodules.mlp_hyper_connection = HyperConnectionModule + return layer_spec + + +def _make_mhc_config(hidden_size=64, num_streams=4, **extra): + """Build a TransformerConfig with common MHC defaults. + + Any default can be overridden via **extra + (e.g. ``_make_mhc_config(num_layers=8, recompute_modules=["core_attn", "mhc"])``). + """ + base = dict( + num_layers=2, + hidden_size=hidden_size, + num_attention_heads=4, + use_cpu_initialization=True, + enable_hyper_connections=True, + num_residual_streams=num_streams, + mhc_sinkhorn_iterations=5, + mhc_init_gating_factor=0.01, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + base.update(extra) + return TransformerConfig(**base) + + class TestParallelTransformerLayer: def setup_method(self, method): @@ -418,3 +458,767 @@ def test_deprecated_full_iteration_inference_scope_string_matches_new_granularit assert block.config.cuda_graph_modules == [] assert _no_layers_have_manager(block) _reset_cudagraph_state() + + +class TestTransformerLayerWithHyperConnectionRecompute: + """Test TransformerLayer with HyperConnection and MHC block recomputation.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_layer_with_hyper_connection( + self, hidden_size=64, num_streams=4, layer_number=1, **extra + ): + """Create a HyperConnectionTransformerLayer with hyper connection enabled.""" + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + recompute_modules=["core_attn", "mhc"], + recompute_granularity='selective', + **extra, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=layer_number + ) + layer.cuda() + return layer, config + + def test_forward_with_hyper_connection_recompute(self): + """ + Test that TransformerLayer forward works correctly with HyperConnection + and MHC block recomputation enabled. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() # Enable training mode for recomputation + + # Input shape: [seq_len, batch_size, n * hidden_size] for hyper connections + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Create manager for MHC block recomputation + manager = CheckpointWithoutOutputManager() + + # Forward pass with recompute manager + manager.is_last_layer_in_recompute_block = True + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Expected output shape {(seq_len, batch_size, n_channels)}, got {output.shape}" + + # Register unified recompute hook at block boundary. + manager.discard_all_outputs_and_register_unified_recompute(output) + + # Backward pass should work without error + loss = output.sum() + loss.backward() + + # Verify gradients exist + assert hidden_states.grad is not None, "Gradients should be computed for hidden_states" + assert hidden_states.grad.shape == hidden_states.shape + + def test_intermediate_layer_with_recompute(self): + """ + Test TransformerLayer as an intermediate layer (not last in block). + In this case, MLP BDA should also be checkpointed. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + + layer, config = self._create_layer_with_hyper_connection(hidden_size, num_streams) + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + manager = CheckpointWithoutOutputManager() + + # Forward pass - NOT the last layer in block + manager.is_last_layer_in_recompute_block = False + output, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + mhc_recompute_manager=manager, + ) + + # Verify output shape + assert output.shape == (seq_len, batch_size, n_channels) + + # Backward pass should work + loss = output.sum() + # For intermediate layers, we need to pass output to next layer + # Here we just register the recompute hook on output for testing + manager.discard_all_outputs_and_register_unified_recompute(loss) + + loss.backward() + + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + + def test_multiple_layers_chain_with_recompute(self): + """ + Test multiple TransformerLayers chained together with a single + CheckpointWithoutOutputManager, simulating TransformerBlock behavior. + """ + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + num_layers = 3 + + layers = [ + self._create_layer_with_hyper_connection( + hidden_size, num_streams, layer_number=i + 1, num_layers=num_layers + )[0] + for i in range(num_layers) + ] + + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Single manager for all layers (like TransformerBlock) + manager = CheckpointWithoutOutputManager() + + # Forward through all layers + h = hidden_states + for i, layer in enumerate(layers): + is_last = i == num_layers - 1 + manager.is_last_layer_in_recompute_block = is_last + h, _ = layer( + hidden_states=h, attention_mask=attention_mask, mhc_recompute_manager=manager + ) + if is_last: + manager.discard_all_outputs_and_register_unified_recompute(h) + + # Backward pass + loss = h.sum() + loss.backward() + + # Verify gradients + assert hidden_states.grad is not None + assert hidden_states.grad.shape == hidden_states.shape + # Check that gradient is non-trivial (not all zeros) + assert hidden_states.grad.abs().sum() > 0 + + +class TestMHCRecomputeMemorySaving: + """Verify that 'mhc' in recompute_modules actually reduces peak GPU memory.""" + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @staticmethod + def _run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute, + recompute_block_size=2, + ): + """Run a full forward + backward pass and return (peak memory, output grad). + + When use_recompute=True, a new CheckpointWithoutOutputManager is created every + `recompute_block_size` layers, mirroring TransformerBlock's + _build_mhc_recompute_layer_plan logic. + """ + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + num_layers=num_layers, + recompute_modules=["core_attn", "mhc"] if use_recompute else None, + recompute_granularity='selective' if use_recompute else None, + ) + layer_spec = _make_mhc_layer_spec() + layers = [ + HyperConnectionTransformerLayer( + config, layer_spec.submodules, layer_number=i + 1 + ).cuda() + for i in range(num_layers) + ] + for layer in layers: + layer.train() + + n_channels = num_streams * hidden_size + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + + manager = CheckpointWithoutOutputManager() if use_recompute else None + + h = hidden_states + for i, layer in enumerate(layers): + is_last_in_block = (i == num_layers - 1) or ((i + 1) % recompute_block_size == 0) + kwargs = dict(hidden_states=h, attention_mask=attention_mask) + if manager is not None: + manager.is_last_layer_in_recompute_block = is_last_in_block + kwargs['mhc_recompute_manager'] = manager + h, _ = layer(**kwargs) + if manager is not None and is_last_in_block: + manager.discard_all_outputs_and_register_unified_recompute(h) + if i < num_layers - 1: + manager = CheckpointWithoutOutputManager() + + loss = h.sum() + loss.backward() + torch.cuda.synchronize() + + peak_mem = torch.cuda.max_memory_allocated() + grad = hidden_states.grad.clone() + + del layers, hidden_states, h, loss, manager + torch.cuda.empty_cache() + + return peak_mem, grad + + def test_recompute_reduces_peak_memory(self): + """Peak memory with recompute (block_size=2) should be lower than without.""" + num_layers = 8 + hidden_size = 128 + num_streams = 4 + seq_len = 64 + batch_size = 4 + + peak_no_recompute, _ = self._run_forward_backward( + num_layers, hidden_size, num_streams, seq_len, batch_size, use_recompute=False + ) + peak_recompute, _ = self._run_forward_backward( + num_layers, + hidden_size, + num_streams, + seq_len, + batch_size, + use_recompute=True, + recompute_block_size=2, + ) + + saving_pct = (peak_no_recompute - peak_recompute) / peak_no_recompute * 100 + + assert peak_recompute < peak_no_recompute, ( + f"Recompute should reduce peak memory, but got " + f"no_recompute={peak_no_recompute / 1e6:.1f}MB vs " + f"recompute={peak_recompute / 1e6:.1f}MB " + f"(saving={saving_pct:.1f}%)" + ) + + +class TestMHCWithCudaGraph: + """Test HyperConnectionTransformerLayer compatibility with CUDA graphs. + + CUDA graph capture requires static computation graphs and fixed tensor shapes. + These tests verify that the mHC layer properly supports the CUDA graph interface + defined in GraphableMegatronModule and TransformerLayer. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123, use_cudagraphable_rng=True, force_reset_rng=True) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer(self, hidden_size=64, num_streams=4, **extra_config): + config = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams, **extra_config) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_get_layer_static_inputs_shape_for_mhc(self): + """get_layer_static_inputs must return [s, b, n*C] for mHC layers. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. If the shape is [s, b, C] instead of [s, b, n*C], the graph + capture will produce a shape mismatch at the first hyper connection module. + """ + layer, config = self._create_mhc_layer() + seq_length = 32 + micro_batch_size = 2 + + static_inputs = layer.get_layer_static_inputs(seq_length, micro_batch_size) + hidden_states = static_inputs["hidden_states"] + + expected_hidden_dim = config.num_residual_streams * config.hidden_size + assert hidden_states.shape[-1] == expected_hidden_dim, ( + f"get_layer_static_inputs returns hidden dim {hidden_states.shape[-1]} " + f"but mHC expects {expected_hidden_dim} (n={config.num_residual_streams} * " + f"C={config.hidden_size}). " + f"HyperConnectionTransformerLayer must override get_layer_static_inputs." + ) + + def test_submodules_under_cudagraphs_includes_hyper_connection(self): + """_get_submodules_under_cudagraphs must include hyper connection modules. + + CUDA graph manual hooks are set up for parameters of submodules returned + by this method. Missing hyper connection modules means their parameters + (mapping_proj, alpha_*, bias) will not get proper pre-forward hooks during + graph replay, leading to stale parameter values. + """ + layer, config = self._create_mhc_layer() + + submodules = layer._get_submodules_under_cudagraphs() + + hc_modules_found = any( + hasattr(m, 'mapping_proj') for submod in submodules for m in submod.modules() + ) + assert hc_modules_found, ( + "_get_submodules_under_cudagraphs does not include HyperConnectionModule. " + "Parameters like mapping_proj, alpha_pre/post/res will not be updated " + "during CUDA graph replay." + ) + + def test_forward_through_te_cuda_graph_capture_path(self): + """_te_cuda_graph_capture must produce correct output shapes for mHC. + + TE CUDA graph capture calls _te_cuda_graph_capture() during warmup. + For mHC layers, the input must be n-stream [s, b, n*C] and output must + also be [s, b, n*C]. + """ + layer, config = self._create_mhc_layer() + layer.eval() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + with torch.no_grad(): + outputs = layer._te_cuda_graph_capture( + hidden_states=hidden_states, attention_mask=attention_mask + ) + + if isinstance(outputs, tuple): + output = outputs[0] + else: + output = outputs + + assert output.shape == (seq_len, batch_size, n_channels), ( + f"_te_cuda_graph_capture output shape {output.shape} != " + f"expected {(seq_len, batch_size, n_channels)}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection(self): + """End-to-end CUDA graph capture and replay for forward+backward with mHC. + + Captures both the forward and backward pass of HyperConnectionTransformerLayer + into a torch.cuda.CUDAGraph and replays it with fresh input data, verifying + that the computation graph is fully static (capturable) and produces correct + output shapes and non-trivial gradients. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream to trigger lazy allocations + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + out, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + # Set .grad to None so backward allocates fresh gradient tensors in the + # graph's private memory pool during capture. + layer.zero_grad(set_to_none=True) + static_input.grad = None + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer(hidden_states=static_input, attention_mask=attention_mask) + output.sum().backward() + + # Replay with new input data. + # Use no_grad because backward inside the captured graph already + # bumped the autograd version counter on static_input, making + # in-place copy_ illegal without disabling grad tracking. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Verify numerical consistency: graph replay should match eager execution + # with the same input and weights. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer(hidden_states=eager_input, attention_mask=attention_mask) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_cuda_graph_fwd_bwd_with_hyper_connection_and_recompute(self): + """CUDA graph capture+replay for fwd+bwd with mHC and CheckpointWithoutOutputManager. + + When a CheckpointWithoutOutputManager is used, additional CheckpointWithoutOutput + objects are created for layernorm and hyper-connection operations. The + manager discards intermediate activations during forward (storage.resize_(0)) + and recomputes them during backward via a unified gradient hook. + This test verifies the full capture+replay still works correctly. + """ + layer, config = self._create_mhc_layer() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + static_input = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Warmup on side stream; fresh manager per iteration to avoid stale state. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + out, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=mgr, + ) + mgr.discard_all_outputs_and_register_unified_recompute(out) + out.sum().backward() + torch.cuda.current_stream().wait_stream(s) + + layer.zero_grad(set_to_none=True) + static_input.grad = None + + capture_mgr = CheckpointWithoutOutputManager() + capture_mgr.is_last_layer_in_recompute_block = True + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + output, _ = layer( + hidden_states=static_input, + attention_mask=attention_mask, + mhc_recompute_manager=capture_mgr, + ) + capture_mgr.discard_all_outputs_and_register_unified_recompute(output) + output.sum().backward() + + # Replay with new input data. + with torch.no_grad(): + static_input.copy_(torch.randn_like(static_input)) + g.replay() + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + assert ( + static_input.grad is not None + ), "Gradients should be computed for static_input after graph replay" + assert static_input.grad.shape == static_input.shape + assert static_input.grad.abs().sum() > 0, "Gradients should be non-trivial" + + # Numerical consistency: graph replay vs eager with the same input. + test_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + + with torch.no_grad(): + static_input.copy_(test_data) + g.replay() + graph_out = output.detach().clone() + graph_grad = static_input.grad.detach().clone() + + eager_mgr = CheckpointWithoutOutputManager() + eager_mgr.is_last_layer_in_recompute_block = True + eager_input = test_data.clone().requires_grad_(True) + eager_output, _ = layer( + hidden_states=eager_input, + attention_mask=attention_mask, + mhc_recompute_manager=eager_mgr, + ) + eager_mgr.discard_all_outputs_and_register_unified_recompute(eager_output) + eager_output.sum().backward() + + assert torch.allclose(graph_out, eager_output.detach(), atol=1e-5), ( + f"Graph vs eager output mismatch: " + f"max diff = {(graph_out - eager_output.detach()).abs().max().item()}" + ) + assert torch.allclose(graph_grad, eager_input.grad, atol=1e-5), ( + f"Graph vs eager gradient mismatch: " + f"max diff = {(graph_grad - eager_input.grad).abs().max().item()}" + ) + + def test_mcore_cudagraph_manager_with_mhc_recompute_manager(self): + """MCore CudaGraphManager must not crash on mhc_recompute_manager kwarg. + + When cuda_graph_impl="local" is set, HyperConnectionTransformerLayer.__call__ + runs first and pops mhc_recompute_manager off kwargs before + super().__call__ → MegatronModule.__call__ → CudaGraphManager.__call__, + which iterates over all kwargs to check supported types. + CheckpointWithoutOutputManager (used by mhc_recompute_manager) is not a + CUDA-graph-supported type. + + This test verifies that mhc_recompute_manager is properly extracted + from kwargs before the CudaGraphManager sees them, preventing the + AssertionError that would otherwise occur. + """ + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + assert hasattr( + layer, 'cudagraph_manager' + ), "Layer should have cudagraph_manager with cuda_graph_impl='local'" + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = CheckpointWithoutOutputManager() + mgr.is_last_layer_in_recompute_block = True + + output, context = layer( + hidden_states=hidden_states, attention_mask=attention_mask, mhc_recompute_manager=mgr + ) + + assert output.shape == (seq_len, batch_size, n_channels) + + def test_mcore_cudagraph_manager_without_mhc_recompute_manager(self): + """MCore CudaGraphManager path works when mhc_recompute_manager is None.""" + layer, config = self._create_mhc_layer(cuda_graph_impl="local", cuda_graph_scope="attn") + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == (seq_len, batch_size, n_channels) + + +class TestMHCWithOffloading: + """Test HyperConnectionTransformerLayer with fine-grained activation offloading. + + Fine-grained activation offloading transfers specific activations (e.g., layernorm + inputs) to CPU during forward and reloads them during backward. These tests verify + that the mHC layer's multi-stream architecture works correctly with offloading. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _create_mhc_layer_with_offloading( + self, hidden_size=64, num_streams=4, offload_modules=None + ): + if offload_modules is None: + offload_modules = ["attn_norm", "mlp_norm"] + + config = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=offload_modules, + ) + layer_spec = _make_mhc_layer_spec() + layer = HyperConnectionTransformerLayer(config, layer_spec.submodules) + layer.cuda() + return layer, config + + def test_forward_backward_with_offloading(self): + """Forward+backward should work with activation offloading enabled. + + This exercises the off_interface context manager around layernorms in + the mHC forward path, including the group_commit that commits the + offloading group for the aggregated 1-stream layernorm input. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + layer, config = self._create_mhc_layer_with_offloading() + layer.train() + + seq_len = 8 + batch_size = 2 + n_channels = config.num_residual_streams * config.hidden_size + + hidden_states = torch.randn( + seq_len, batch_size, n_channels, device='cuda', requires_grad=True + ) + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler( + pp_rank=0, vp_size=1, vp_stage=0, min_offloaded_tensor_size=0 + ) + + output, context = layer(hidden_states=hidden_states, attention_mask=attention_mask) + + assert output.shape == ( + seq_len, + batch_size, + n_channels, + ), f"Output shape {output.shape} != expected {(seq_len, batch_size, n_channels)}" + + loss = output.sum() + loss.backward() + + assert hidden_states.grad is not None, "Gradients should flow through offloaded path" + assert hidden_states.grad.shape == hidden_states.shape + assert hidden_states.grad.abs().sum() > 0, "Gradients should be non-trivial" + + PipelineOffloadManager.reset_instance() + + def test_offloading_numerical_equivalence(self): + """Offloaded forward+backward must produce the same result as non-offloaded. + + Compares outputs and gradients between a layer with offloading disabled + vs enabled to ensure the offloading path does not corrupt activations. + """ + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + PipelineOffloadManager.reset_instance() + + hidden_size = 64 + num_streams = 4 + seq_len = 8 + batch_size = 2 + n_channels = num_streams * hidden_size + + torch.manual_seed(42) + input_data = torch.randn(seq_len, batch_size, n_channels, device='cuda') + attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool, device='cuda') + + # Run without offloading + config_no_offload = _make_mhc_config(hidden_size=hidden_size, num_streams=num_streams) + layer_spec = _make_mhc_layer_spec() + layer_no_offload = HyperConnectionTransformerLayer( + config_no_offload, layer_spec.submodules + ).cuda() + layer_no_offload.train() + + h1 = input_data.clone().detach().requires_grad_(True) + out1, _ = layer_no_offload(hidden_states=h1, attention_mask=attention_mask) + out1.sum().backward() + grad_no_offload = h1.grad.clone() + out1_detached = out1.detach().clone() + + # Run with offloading using the same weights + config_offload = _make_mhc_config( + hidden_size=hidden_size, + num_streams=num_streams, + fine_grained_activation_offloading=True, + offload_modules=["attn_norm", "mlp_norm"], + ) + layer_offload = HyperConnectionTransformerLayer( + config_offload, layer_spec.submodules + ).cuda() + layer_offload.load_state_dict(layer_no_offload.state_dict()) + layer_offload.train() + + mgr = PipelineOffloadManager.get_instance() + mgr.init_model_chunk_offload_handler( + pp_rank=0, vp_size=1, vp_stage=0, min_offloaded_tensor_size=0 + ) + + h2 = input_data.clone().detach().requires_grad_(True) + out2, _ = layer_offload(hidden_states=h2, attention_mask=attention_mask) + out2.sum().backward() + grad_offload = h2.grad.clone() + + PipelineOffloadManager.reset_instance() + + assert torch.allclose(out1_detached, out2.detach(), atol=1e-5), ( + f"Forward outputs differ: max diff = " + f"{(out1_detached - out2.detach()).abs().max().item()}" + ) + assert torch.allclose(grad_no_offload, grad_offload, atol=1e-5), ( + f"Gradients differ: max diff = " + f"{(grad_no_offload - grad_offload).abs().max().item()}" + )