From d025a7f095a142800c76afed069a3c88a956924e Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Mon, 15 Jun 2026 23:49:21 -0700 Subject: [PATCH 01/22] Add gated delta product (GDP) mixer and spec (#45) Signed-off-by: Deepak Narayanan Co-authored-by: Deepak Narayanan Co-authored-by: Claude Opus 4.6 Signed-off-by: Keshav Santhanam --- .../core/models/hybrid/hybrid_layer_specs.py | 51 ++ megatron/core/ssm/gated_delta_product.py | 687 ++++++++++++++++++ megatron/core/ssm/gdp_context_parallel.py | 325 +++++++++ 3 files changed, 1063 insertions(+) create mode 100644 megatron/core/ssm/gated_delta_product.py create mode 100644 megatron/core/ssm/gdp_context_parallel.py diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 03fef58159f..0ee36323c95 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -16,6 +16,10 @@ ) from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules +from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer +from megatron.core.ssm.gated_delta_product import ( + MambaMixerSubmodules as GatedDeltaProductMixerSubmodules, +) from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules from megatron.core.ssm.mlp_layer import MLPLayer @@ -87,6 +91,19 @@ ) +def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): + return ModuleSpec( + module=MambaLayer, + submodules=MambaLayerSubmodules( + mixer=ModuleSpec( + module=GatedDeltaProductMixer, + submodules=GatedDeltaProductMixerSubmodules(in_proj=in_proj, out_proj=out_proj), + ), + mamba_bda=get_bias_dropout_add, + ), + ) + + hybrid_stack_spec = ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( @@ -217,6 +234,22 @@ ) +gated_delta_product_stack_spec = ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + mamba_layer=_get_gated_delta_product_mamba_layer_spec( + TELayerNormColumnParallelLinear, TERowParallelLinear + ), + gdn_layer=hybrid_stack_spec.submodules.gdn_layer, + attention_layer=hybrid_stack_spec.submodules.attention_layer, + dsa_layer=hybrid_stack_spec.submodules.dsa_layer, + mlp_layer=hybrid_stack_spec.submodules.mlp_layer, + moe_layer=hybrid_stack_spec.submodules.moe_layer, + mtp_block_spec=hybrid_stack_spec.submodules.mtp_block_spec, + ), +) + + hybrid_inference_stack_spec = ModuleSpec( module=HybridStack, submodules=HybridStackSubmodules( @@ -352,6 +385,24 @@ ) +gated_delta_product_inference_stack_spec = ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + mamba_layer=_get_gated_delta_product_mamba_layer_spec( + InferenceLayerNormColumnParallelLinear, InferenceRowParallelLinear + ), + gdn_layer=hybrid_inference_stack_spec.submodules.gdn_layer, + attention_layer=hybrid_inference_stack_spec.submodules.attention_layer, + dsa_layer=hybrid_inference_stack_spec.submodules.dsa_layer, + mlp_layer=hybrid_inference_stack_spec.submodules.mlp_layer, + moe_layer=hybrid_inference_stack_spec.submodules.moe_layer, + mtp_block_spec=hybrid_inference_stack_spec.submodules.mtp_block_spec, + ), +) + + # Backward-compatible aliases mamba_stack_spec = hybrid_stack_spec mamba_inference_stack_spec = hybrid_inference_stack_spec +gdp_stack_spec = gated_delta_product_stack_spec +gdp_inference_stack_spec = gated_delta_product_inference_stack_spec diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py new file mode 100644 index 00000000000..3236fa000cf --- /dev/null +++ b/megatron/core/ssm/gated_delta_product.py @@ -0,0 +1,687 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +# Some of this code was adopted from https://github.com/state-spaces/mamba/ +# This source code is licensed under the Apache license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math +from dataclasses import dataclass, replace +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from megatron.core.dist_checkpointing import ShardedTensor +from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gdp_context_parallel import GDPContextParallel +from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.utils import ( + make_sharded_tensors_for_checkpoint, + sharded_state_dict_default, +) +from megatron.core.utils import deprecate_inference_params + +try: + from causal_conv1d import causal_conv1d_fn, causal_conv1d_update +except ImportError: + causal_conv1d_fn = None + causal_conv1d_update = None + +try: + from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated + + HAVE_MAMBA_SSM = True +except ImportError: + from unittest.mock import MagicMock + + RMSNormGated = MagicMock() + HAVE_MAMBA_SSM = False + +try: + from einops import rearrange + + HAVE_EINOPS = True +except ImportError: + HAVE_EINOPS = False + +try: + from fla.ops.gated_delta_product import chunk_gated_delta_product + from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + + +logger = logging.getLogger(__name__) + + +class ExtendedRMSNorm(RMSNormGated): + """ + RMSNormGated with sharded state dict. + """ + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + """Sharding along axis 0, bias not sharded""" + state_dict = self.state_dict(prefix="", keep_vars=True) + return make_sharded_tensors_for_checkpoint( + state_dict, prefix, {"weight": 0}, sharded_offsets + ) + + +@dataclass +class MambaMixerSubmodules: + """ + Contains the module specs for the input and output linear layers. + """ + + in_proj: Union[ModuleSpec, type] = None + out_proj: Union[ModuleSpec, type] = None + + +class GatedDeltaProductMixer(MegatronModule): + """ + Args: + config: The config of the model. + submodules: Contains the module specs for the input and output linear layers. + d_model: The hidden size of the model. + d_state: The state size of the SSM. + d_conv: The number of channels in the causal convolution. + conv_init: The initialization range for the causal convolution weights. + expand: The expansion factor for the SSM. + headdim: The hidden size of each attention head. + ngroups: The number of attention heads. + A_init_range: The initialization range for the attention weights. + D_has_hdim: Whether the D parameter has the same number of dimensions as the hidden + state. + rmsnorm: Whether to use root mean square normalization. + norm_before_gate: Whether to apply normalization before the gating mechanism. + dt_min: The minimum value of the dt parameter. + dt_max: The maximum value of the dt parameter. + dt_init: The initialization value of the dt parameter. + dt_scale: The scaling factor for the dt parameter. + dt_init_floor: The minimum value of the dt parameter after initialization. + bias: Whether to use bias in the linear layers. + conv_bias: Whether to use bias in the causal convolution. + chunk_size: The chunk size for the fused kernel. + use_mem_eff_path: Whether to use the memory-efficient path for the Mamba model. + layer_number: The layer number of this Mamba layer. + pg_collection: The required process groups to use for tensor model parallel and context + parallel. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: MambaMixerSubmodules, + d_model, + d_conv=4, + conv_init=None, + expand=2, + A_init_range=(0, 16), + D_has_hdim=False, + rmsnorm=True, + norm_before_gate=False, + dt_min=0.001, + dt_max=0.1, + dt_init="random", + dt_scale=1.0, + dt_init_floor=1e-4, + bias=False, + conv_bias=False, + # Fused kernel and sharding options + chunk_size=128, + layer_number=None, + use_mem_eff_path=None, + d_state=None, + headdim=None, + ngroups=None, + pg_collection: ProcessGroupCollection = None, + pp_layer_offset: int = 0, + ): + if not HAVE_MAMBA_SSM: + raise ImportError( + "MambaSSM is not installed. Please install it with `pip install mamba-ssm`." + ) + + if not HAVE_FLA: + raise ImportError("FLA is not installed") + + super().__init__(config) + + self.num_householder = 3 + + self.config = config + self.d_model = d_model + self.d_conv = d_conv + self.conv_init = conv_init + self.D_has_hdim = D_has_hdim + self.rmsnorm = rmsnorm + self.norm_before_gate = norm_before_gate + assert pg_collection is not None, "pg_collection must be provided for MambaMixer" + self.pg_collection = pg_collection + + self.d_state = self.config.mamba_state_dim + self.headdim = self.config.mamba_head_dim + self.ngroups = self.config.mamba_num_groups + self.nheads = self.config.mamba_num_heads + self.d_inner = self.nheads * self.headdim + + self.layer_number = layer_number + self.pp_layer_offset = pp_layer_offset + self.cached_batch_size = None + + tp_size = self.pg_collection.tp.size() + + self.nheads_local_tp = self.nheads // tp_size + self.d_inner_local_tp = self.d_inner // tp_size + self.ngroups_local_tp = self.ngroups // tp_size + + # Assume sequence parallelism: input is already partitioned along the sequence dimension + self.in_proj = build_module( + submodules.in_proj, + self.d_model, + self.d_inner * (1 + self.num_householder) + + self.ngroups * self.d_state * (self.num_householder + 1) + + self.nheads * (self.num_householder + 1), # zVKQba + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="fc1", + tp_group=self.pg_collection.tp, + ) + + conv_dim = ( + self.d_inner_local_tp * self.num_householder + + (self.num_householder + 1) * self.ngroups_local_tp * self.d_state + ) # V K Q + with get_cuda_rng_tracker().fork(): + # weight shape: [conv_dim, 1, d_conv] + # bias shape: [conv_dim] + self.conv1d = nn.Conv1d( + in_channels=conv_dim, + out_channels=conv_dim, + bias=conv_bias, + kernel_size=d_conv, + groups=conv_dim, + padding=d_conv - 1, + device=torch.cuda.current_device(), + dtype=config.params_dtype, + ) + setattr(self.conv1d.weight, "tensor_model_parallel", True) + if conv_bias: + setattr(self.conv1d.bias, "tensor_model_parallel", True) + + if self.conv_init is not None: + nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) + + self.activation = "silu" + self.act = nn.SiLU() + + with get_cuda_rng_tracker().fork(): + # MCore Mamba2 initialization + # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max + dt = torch.exp( + torch.rand( + self.nheads_local_tp, + device=torch.cuda.current_device(), + dtype=config.params_dtype, + ) + * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ).clamp(min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + + # Our initialization would set all Linear.bias to zero, + # need to mark this one as _no_reinit + self.dt_bias._no_reinit = True + # Just to be explicit. Without this we already don't + # put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + setattr(self.dt_bias, "tensor_model_parallel", True) + + # A parameter + assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0] + A = torch.empty( + self.nheads_local_tp, dtype=torch.float32, device=torch.cuda.current_device() + ).uniform_(*A_init_range) + A_log = torch.log(A) # Keep A_log in fp32 + self.A_log = nn.Parameter(A_log) + self.A_log._no_weight_decay = True + setattr(self.A_log, "tensor_model_parallel", True) + + # D "skip", in Mamba2 but not in GDN or GDP + self.D = None + + if self.rmsnorm: + assert RMSNormGated is not None + self.norm = ExtendedRMSNorm( + self.d_inner_local_tp, + eps=1e-5, + group_size=self.d_inner_local_tp // self.ngroups_local_tp, + norm_before_gate=self.norm_before_gate, + device=torch.cuda.current_device(), + dtype=config.params_dtype, + ) + setattr(self.norm.weight, 'tensor_model_parallel', True) + + # Assume sequence parallelism: input is partitioned along d_inner and + # output is partitioned along the sequence dimension + self.out_proj = build_module( + submodules.out_proj, + self.d_inner, + self.d_model, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=bias, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="fc2", + tp_group=self.pg_collection.tp, + ) + + # Regarding `conv1d`.{`weight`, `bias`}, `dt_bias`, `A_log`, and `D`: these are the + # trainable variables for the current tensor parallel rank, with each tensor parallel rank + # having indepdendent trainable variables. All context parallel ranks in a tensor parallel + # rank store the same trainable variables, but only use and update their unique/independent + # slice of them. + self.cp = GDPContextParallel( + cp_group=self.pg_collection.cp, + d_inner_local_tp=self.d_inner_local_tp, + nheads_local_tp=self.nheads_local_tp, + ngroups_local_tp=self.ngroups_local_tp, + d_state=self.d_state, + num_householder=self.num_householder, + headdim=self.headdim, + conv1d_cp1=self.conv1d, + dt_bias_cp1=self.dt_bias, + A_log_cp1=self.A_log, + D_cp1=self.D, + D_has_hdim=self.D_has_hdim, + ) + + def forward( + self, + hidden_states, + inference_context=None, + *, + inference_params: Optional[BaseInferenceContext] = None, + packed_seq_params=None, + ): + """Run the gated delta product mixer on hidden states.""" + if packed_seq_params is not None: + raise NotImplementedError( + "GatedDeltaProductMixer does not support packed sequences yet." + ) + + seq_len, batch_size, dim = hidden_states.shape + + conv_state, ssm_state = None, None + if inference_context is not None: + assert ( + inference_context.is_static_batching() + ), "Mamba does not currently support dynamic inference batching." + assert not self.config.sequence_parallel + conv_state, ssm_state = self._get_states_from_cache(inference_context, batch_size) + + zVKQba, _ = self.in_proj(hidden_states) + + zVKQba = self.cp.pre_conv_ssm(zVKQba) + + zVKQba = rearrange(zVKQba, "l b d -> b l d").contiguous() + + z, VKQ, ba = torch.split( + zVKQba, + [ + self.cp.d_inner_local_tpcp, + self.cp.d_inner_local_tpcp * self.num_householder + + (self.num_householder + 1) * self.cp.ngroups_local_tpcp * self.d_state, + self.cp.nheads_local_tpcp * (self.num_householder + 1), + ], + dim=-1, + ) + + VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + + # Decode + if inference_context is not None and inference_context.seqlen_offset > 0: + VKQ = causal_conv1d_update( + VKQ, + conv_state, + rearrange(self.conv1d.weight, "d 1 w -> d w"), + self.conv1d.bias, + self.activation, + ) + else: + # Prefill + if conv_state is not None: + # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv + # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. + conv_state.copy_( + F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0)) + ) # Update state (B D W) + # Train + VKQ = causal_conv1d_fn( + x=VKQ, + weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + bias=self.cp.get_conv1d_bias(), + activation=self.activation, + ) + + VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() + + value, key, query = torch.split( + VKQ, + [ + self.cp.d_inner_local_tpcp * self.num_householder, + self.cp.ngroups_local_tpcp * self.d_state * self.num_householder, + self.cp.ngroups_local_tpcp * self.d_state, + ], + dim=-1, + ) + + b, a = torch.split( + ba, + [self.cp.nheads_local_tpcp * self.num_householder, self.cp.nheads_local_tpcp], + dim=-1, + ) + + z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).contiguous() + value = rearrange( + value, "b l (m h p) -> b (l m) h p", m=self.num_householder, p=self.headdim + ).contiguous() + key = rearrange( + key, "b l (m g n) -> b (l m) g n", m=self.num_householder, n=self.d_state + ).contiguous() + query = rearrange(query, "b l (g n) -> b l g n", n=self.d_state).contiguous() + + b, a = b.contiguous(), a.contiguous() + beta = b.sigmoid() + beta = rearrange(beta, "b l (m h) -> b (l m) h", m=self.num_householder).contiguous() + + # If the model is loaded in fp16, without the .float() here, A might be -inf + g = -self.cp.get_A_log().float().exp() * F.softplus(a.float() + self.cp.get_dt_bias()) + + if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: + query = query.repeat_interleave( + self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=2 + ) + key = key.repeat_interleave( + self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=2 + ) + + # Decode + if inference_context is not None and inference_context.seqlen_offset > 0: + + g_new = g.new_zeros(g.shape[0], g.shape[1], self.num_householder, g.shape[2]) + g_new[:, :, 0] = g + g = rearrange(g_new, '... t n h -> ... (t n) h') + + query_new = query.new_zeros( + query.shape[0], query.shape[1], self.num_householder, query.shape[2], query.shape[3] + ) + query_new[:, :, -1] = query + query = rearrange(query_new, '... t n h d-> ... (t n) h d') + + core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=ssm_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + core_attn_out = rearrange( + core_attn_out, '... (t n) h d -> ... t n h d', n=self.num_householder + )[..., -1, :, :].contiguous() + # Train or Prefill + else: + core_attn_out, last_recurrent_state = chunk_gated_delta_product( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=(ssm_state is not None), + num_householder=self.num_householder, + use_qk_l2norm_in_kernel=True, + ) + + if ssm_state is not None: + ssm_state.copy_(last_recurrent_state) + + y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() + y = self.cp.post_conv_ssm(y) + if self.rmsnorm: + z = rearrange(z, "b l h p -> l b (h p)").contiguous() + z = self.cp.post_conv_ssm(z) + y = self.norm(y, z) + + out, out_bias = self.out_proj(y) + + return out, out_bias + + def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None): + """ + allocate inference cache + """ + device = self.out_proj.weight.device + conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype + conv_state = torch.zeros( + batch_size, self.conv1d.weight.shape[0], self.d_conv, device=device, dtype=conv_dtype + ) + ssm_dtype = self.in_proj.weight.dtype if dtype is None else dtype + # ssm_dtype = torch.float32 + ssm_state = torch.zeros( + batch_size, + self.nheads_local_tp, + self.d_state, + self.headdim, + device=device, + dtype=ssm_dtype, + ) + return conv_state, ssm_state + + def mamba_state_shapes_per_request(self) -> Tuple[Tuple[int], Tuple[int]]: + """Returns the Mamba conv and SSM state shapes per request.""" + conv_states_shape = (self.conv1d.weight.shape[0], self.d_conv) + ssm_states_shape = (self.nheads_local_tp, self.d_state, self.headdim) + return (conv_states_shape, ssm_states_shape) + + def _get_states_from_cache(self, inference_context, batch_size, *, inference_params=None): + """Initializes or retrieves the SSM state tensors from the cache. + + At the start of any inference (at the prefill step), if there is no cache or if the + cached batch size has changed, then new tensors are initialized and stored in the cache. + Otherwise the existing tensors are retrieved from the cache and zeroed out. + """ + + inference_context = deprecate_inference_params(inference_context, inference_params) + + assert inference_context is not None + assert self.layer_number is not None + if ( + self.layer_number not in inference_context.key_value_memory_dict + or batch_size != self.cached_batch_size + ): + conv_state = torch.zeros( + batch_size, + self.conv1d.weight.shape[0], + self.d_conv, + device=self.conv1d.weight.device, + dtype=self.conv1d.weight.dtype, + ) + ssm_state = torch.zeros( + batch_size, + self.nheads_local_tp, + self.d_state, + self.headdim, + device=self.in_proj.weight.device, + dtype=self.in_proj.weight.dtype, + ) + inference_context.key_value_memory_dict[self.layer_number] = (conv_state, ssm_state) + self.cached_batch_size = batch_size + else: + conv_state, ssm_state = inference_context.key_value_memory_dict[self.layer_number] + # TODO: Remove reference to `inference_context.sequence_len_offset` for dynamic batching + if inference_context.sequence_len_offset == 0: + conv_state.zero_() + ssm_state.zero_() + return conv_state, ssm_state + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + """Provide a sharded state dictionary for distributed checkpointing.""" + sharded_state_dict = {} + # Parameters + self._save_to_state_dict(sharded_state_dict, "", keep_vars=True) + sharded_state_dict = make_sharded_tensors_for_checkpoint( + sharded_state_dict, + prefix, + tensor_parallel_layers_axis_map={ + "A_log": 0, + "dt_bias": 0, + "D": 0, + }, # parameters sharded across TP + sharded_offsets=sharded_offsets, + ) + # Submodules + for name, module in self.named_children(): + if name == "conv1d": + # Add TP sharding for Conv1d + module_sd = module.state_dict(prefix="", keep_vars=True) + module_sharded_sd = make_sharded_tensors_for_checkpoint( + module_sd, f"{prefix}{name}.", {f"weight": 0, f"bias": 0}, sharded_offsets + ) + + else: + module_sharded_sd = sharded_state_dict_default( + module, f"{prefix}{name}.", sharded_offsets, metadata + ) + + sharded_state_dict.update(module_sharded_sd) + + # At this point the TP sharding is correctly defined for each tensor, but some of the + # tensors must be additionally split into separate parts + in_proj_dim = ( + self.d_inner_local_tp * (1 + self.num_householder) + + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state + + self.nheads_local_tp * (1 + self.num_householder) + ) + assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim, ( + in_proj_dim, + sharded_state_dict[f"{prefix}in_proj.weight"], + ) + + sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory( + sharded_state_dict[f"{prefix}in_proj.weight"], + [ + self.d_inner_local_tp, + self.d_inner_local_tp * self.num_householder, + self.ngroups_local_tp * self.d_state * self.num_householder, + self.ngroups_local_tp * self.d_state, + self.nheads_local_tp * self.num_householder, + self.nheads_local_tp, + ], + ["z", "V", "K", "Q", "b", "a"], + 0, + ) + + conv_dim = ( + self.d_inner_local_tp * self.num_householder + + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state + ) + assert sharded_state_dict[f"{prefix}conv1d.weight"].data.size(0) == conv_dim, ( + conv_dim, + sharded_state_dict[f"{prefix}conv1d.weight"], + ) + + for conv_layer_name in ["conv1d.weight"]: + sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory( + sharded_state_dict[f"{prefix}{conv_layer_name}"], + [ + self.d_inner_local_tp * self.num_householder, + self.ngroups_local_tp * self.d_state * self.num_householder, + self.ngroups_local_tp * self.d_state, + ], + ["V", "K", "Q"], + 0, + ) + + return sharded_state_dict + + +def _split_tensor_factory( + orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int +) -> ShardedTensorFactory: + """Builds a factory that splits a given ShardedTensor into several independent chunks.""" + assert isinstance(orig_sh_ten, ShardedTensor), type(orig_sh_ten) + orig_sh_ten_no_data = orig_sh_ten.without_data() # remove `data` reference + + if sum(split_sections) != orig_sh_ten_no_data.local_shape[split_dim]: + raise ValueError( + f"Split sections must cover the whole dimension size, " + f"got {split_sections=} vs dimensions size " + f"{orig_sh_ten_no_data.local_shape[split_dim]}" + ) + + assert not isinstance( + split_sections, int + ), "Splitting into predefined section sizes is supported (`split_sections` must be a list)" + assert len(split_sections) == len(split_names), (len(split_sections), len(split_names)) + + @torch.no_grad() + def sh_ten_build_fn( + key: str, t: torch.Tensor, replica_id: ReplicaId, flattened_range: Optional[slice] + ): + factory_sh_ten = replace( + orig_sh_ten_no_data, + key=key, + data=t, + dtype=t.dtype, + replica_id=replica_id, + flattened_range=flattened_range, + ) + + chunk_sh_tens = [] + split_start = 0 + for split_size, split_name in zip(split_sections, split_names): + split_chunks = factory_sh_ten.narrow(split_dim, split_start, split_size) + for sh_ten in split_chunks: + sh_ten.key = f"{sh_ten.key}.{split_name}" + chunk_sh_tens.extend(split_chunks) + split_start += split_size + + assert split_start == orig_sh_ten_no_data.local_shape[split_dim], ( + split_start, + orig_sh_ten_no_data.local_shape[split_dim], + ) + assert sum(sh_ten.data.numel() for sh_ten in chunk_sh_tens) == t.numel(), ( + chunk_sh_tens, + t.shape, + ) + return chunk_sh_tens + + @torch.no_grad() + def sh_ten_merge_fn(sub_state_dict): + return torch.cat(sub_state_dict) + + return ShardedTensorFactory( + orig_sh_ten.key, orig_sh_ten.data, sh_ten_build_fn, sh_ten_merge_fn, orig_sh_ten.replica_id + ) diff --git a/megatron/core/ssm/gdp_context_parallel.py b/megatron/core/ssm/gdp_context_parallel.py new file mode 100644 index 00000000000..ff118ceee5f --- /dev/null +++ b/megatron/core/ssm/gdp_context_parallel.py @@ -0,0 +1,325 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +""" +Context parallel support for Gated Delta Product (GDP) with num_householder > 1. + +The key difference from GDNContextParallel (which assumes a single copy of V/K/b) +is that GDP has `num_householder` copies of V, K, and b (beta). The in_proj output +layout is: + + [z(d_inner), V(d_inner*M), K(ngroups*d_state*M), Q(ngroups*d_state), b(nheads*M), a(nheads)] + +where M = num_householder. Similarly, the conv1d operates on: + + [V(d_inner*M), K(ngroups*d_state*M), Q(ngroups*d_state)] + +The all-to-all communication and parameter slicing must account for this. + +Strategy for householder-multiplied tensors (V, K, b): + We fold the M (householder) dimension into the batch dimension before calling + the standard all-to-all, then unfold afterward. This ensures each householder + copy is independently partitioned by heads across CP ranks. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from einops import repeat + + HAVE_EINOPS = True +except ImportError: + HAVE_EINOPS = False + +# Re-use the load balancing and all-to-all helpers from the existing module. +from megatron.core.ssm.mamba_context_parallel import ( + _all_to_all_cp2hp, + _all_to_all_hp2cp, + _redo_attention_load_balancing, + _undo_attention_load_balancing, +) + + +class GDPContextParallel: + """ + Context parallel support for Gated Delta Product (GDP) models with num_householder >= 1. + + Handles the "all-to-all" CP strategy where heads are partitioned across CP ranks + and each rank processes the full sequence for its head partition. Correctly handles + the num_householder multiplier on V, K, and beta projections. + + Args: + cp_group: The process group for context parallel. + d_inner_local_tp: d_inner on the current TP rank. + nheads_local_tp: nheads on the current TP rank. + ngroups_local_tp: ngroups on the current TP rank. + d_state: SSM state dimension. + num_householder: Number of householder reflections (M). + headdim: Dimension per head. + conv1d_cp1: The conv1d module for cp_size=1. + dt_bias_cp1: The dt_bias parameter for cp_size=1. + A_log_cp1: The A_log parameter for cp_size=1. + D_cp1: The D parameter for cp_size=1 (can be None). + D_has_hdim: Whether D is sized to the hidden dimension. + """ + + def __init__( + self, + cp_group: torch.distributed.ProcessGroup, + d_inner_local_tp: int, + nheads_local_tp: int, + ngroups_local_tp: int, + d_state: int, + num_householder: int, + headdim: int, + conv1d_cp1: nn.Conv1d, + dt_bias_cp1: torch.Tensor, + A_log_cp1: torch.Tensor, + D_cp1: torch.Tensor, + D_has_hdim: bool, + ) -> None: + if not HAVE_EINOPS: + raise ImportError("einops is required but cannot be imported") + + self.cp_group = cp_group + self.d_inner_local_tp = d_inner_local_tp + self.nheads_local_tp = nheads_local_tp + self.ngroups_local_tp = ngroups_local_tp + self.d_state = d_state + self.num_householder = num_householder + self.headdim = headdim + self.conv1d_cp1 = conv1d_cp1 + self.dt_bias_cp1 = dt_bias_cp1 + self.A_log_cp1 = A_log_cp1 + self.D_cp1 = D_cp1 + self.D_has_hdim = D_has_hdim + + self.cp_size = self.cp_group.size() + + M = self.num_householder + + if self.cp_size == 1: + self.d_inner_local_tpcp = self.d_inner_local_tp + self.nheads_local_tpcp = self.nheads_local_tp + self.ngroups_local_tpcp = self.ngroups_local_tp + return + + self.cp_rank = self.cp_group.rank() + + assert ( + self.nheads_local_tp % self.cp_size == 0 + ), "nheads must be evenly divisible by tp_size * cp_size" + self.nheads_local_tpcp = self.nheads_local_tp // self.cp_size + + self.d_inner_local_tpcp = self.d_inner_local_tp // self.cp_size + + # Group repeat logic (same as GDNContextParallel) + if self.ngroups_local_tp < self.cp_size: + assert ( + self.cp_size % self.ngroups_local_tp == 0 + ), "cp_size must be evenly divisible by ngroups/tp_size" + self.group_repeat_count = self.cp_size // self.ngroups_local_tp + self.ngroups_local_tpcp = 1 + else: + assert ( + self.ngroups_local_tp % self.cp_size == 0 + ), "ngroups must be evenly divisible by tp_size * cp_size" + self.group_repeat_count = 1 + self.ngroups_local_tpcp = self.ngroups_local_tp // self.cp_size + + def pre_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor: + """ + All-to-all from sequence-partitioned to head-partitioned layout, before conv + SSM. + + Input layout (last dim): + [z, V, K, Q, b, a] with sizes + [d_inner, d_inner*M, ngroups*d_state*M, ngroups*d_state, nheads*M, nheads] + + Output layout (last dim, after head partitioning): + [z, V, K, Q, b, a] with sizes + [d_inner/cp, d_inner/cp*M, ngroups_cp*d_state*M, + ngroups_cp*d_state, nheads/cp*M, nheads/cp] + """ + if self.cp_size == 1: + return input_ + + M = self.num_householder + l, b, _ = input_.shape + + z, V, K, Q, b_proj, a = torch.split( + input_, + [ + self.d_inner_local_tp, # z + self.d_inner_local_tp * M, # V (M copies) + self.ngroups_local_tp * self.d_state * M, # K (M copies) + self.ngroups_local_tp * self.d_state, # Q (single) + self.nheads_local_tp * M, # beta (M copies) + self.nheads_local_tp, # a (single) + ], + dim=-1, + ) + + # z: [l, b, d_inner] -> [l*cp, b, d_inner/cp] + z = _all_to_all_cp2hp(z, self.cp_group) + + # V: [l, b, M * d_inner] -> fold M into batch -> all-to-all -> unfold + # Layout within last dim is (m, h, p). Folding M into batch ensures each + # householder copy is independently split by heads. + V = V.view(l, b, M, self.d_inner_local_tp) + V = V.reshape(l, b * M, self.d_inner_local_tp) + V = _all_to_all_cp2hp(V, self.cp_group) # (l*cp, b*M, d_inner/cp) + V = V.reshape(l * self.cp_size, b, M * self.d_inner_local_tpcp) + + # K: [l, b, M * ngroups * d_state] -> group repeat each copy -> fold M -> all-to-all + K = K.view(l, b, M, self.ngroups_local_tp * self.d_state) + K_parts = [] + for i in range(M): + Ki = K[:, :, i, :] # (l, b, ngroups_tp * d_state) + Ki = repeat( + Ki, + "l b (g n) -> l b (g r n)", + g=self.ngroups_local_tp, + n=self.d_state, + r=self.group_repeat_count, + ) + K_parts.append(Ki) + K = torch.stack(K_parts, dim=2) # (l, b, M, ngroups_tp * r * d_state) + K = K.reshape(l, b * M, -1) + K = _all_to_all_cp2hp(K, self.cp_group) # (l*cp, b*M, ngroups_tpcp * d_state) + K = K.reshape(l * self.cp_size, b, M * self.ngroups_local_tpcp * self.d_state) + + # Q: [l, b, ngroups * d_state] -> group repeat -> all-to-all (single copy, no M) + Q = repeat( + Q, + "l b (g n) -> l b (g r n)", + g=self.ngroups_local_tp, + n=self.d_state, + r=self.group_repeat_count, + ) + Q = _all_to_all_cp2hp(Q, self.cp_group) # (l*cp, b, ngroups_tpcp * d_state) + + # b_proj (beta): [l, b, M * nheads] -> fold M -> all-to-all -> unfold + b_proj = b_proj.view(l, b, M, self.nheads_local_tp) + b_proj = b_proj.reshape(l, b * M, self.nheads_local_tp) + b_proj = _all_to_all_cp2hp(b_proj, self.cp_group) # (l*cp, b*M, nheads/cp) + b_proj = b_proj.reshape(l * self.cp_size, b, M * self.nheads_local_tpcp) + + # a: [l, b, nheads] -> [l*cp, b, nheads/cp] + a = _all_to_all_cp2hp(a, self.cp_group) + + output = torch.cat([z, V, K, Q, b_proj, a], dim=-1) + output = _undo_attention_load_balancing(output, self.cp_size) + + return output + + def post_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor: + """Method to be applied after the conv + SSM (on y and z, which have no M dim).""" + if self.cp_size == 1: + return input_ + else: + return _all_to_all_hp2cp( + _redo_attention_load_balancing(input_, self.cp_size), self.cp_group + ) + + def conv1d(self, input_: torch.Tensor) -> torch.Tensor: + """Performs conv1d using sliced weights for the current CP rank.""" + if self.cp_size == 1: + return self.conv1d_cp1(input_) + else: + return F.conv1d( + input=input_, + weight=self.get_conv1d_weight(), + bias=self.get_conv1d_bias(), + stride=self.conv1d_cp1.stride, + padding=self.conv1d_cp1.padding, + dilation=self.conv1d_cp1.dilation, + groups=self.conv1d_channels(), + ) + + def conv1d_channels(self): + """Number of conv channels on the current CP rank.""" + M = self.num_householder + return ( + self.d_inner_local_tpcp * M + + self.ngroups_local_tpcp * self.d_state * M + + self.ngroups_local_tpcp * self.d_state + ) + + def get_conv1d_weight(self) -> torch.Tensor: + """Returns sliced conv1d weight for the current CP rank.""" + return self._slice_conv_param(self.conv1d_cp1.weight) + + def get_conv1d_bias(self) -> torch.Tensor: + """Returns sliced conv1d bias for the current CP rank.""" + return self._slice_conv_param(self.conv1d_cp1.bias) + + def get_dt_bias(self) -> torch.Tensor: + """Returns sliced dt_bias for the current CP rank.""" + return self._slice_vector_param(self.dt_bias_cp1) + + def get_A_log(self) -> torch.Tensor: + """Returns sliced A_log for the current CP rank.""" + return self._slice_vector_param(self.A_log_cp1) + + def get_D(self) -> torch.Tensor: + """Returns sliced D for the current CP rank.""" + return self._slice_vector_param(self.D_cp1, has_hdim=self.D_has_hdim) + + def _slice_conv_param(self, param: torch.Tensor) -> torch.Tensor: + """ + Slices a cp_size=1 conv1d parameter along the channel dimension, + returning the channels needed on the current CP rank. + + Conv param layout (dim 0): + [V(d_inner * M), K(ngroups * d_state * M), Q(ngroups * d_state)] + + For V and K (which have M copies), we reshape to (M, per_copy_channels, ...), + slice the per-copy channels for this CP rank, then flatten back. + """ + if self.cp_size == 1 or param is None: + return param + + M = self.num_householder + extra_dims = param.shape[1:] # (1, d_conv) for weight, () for bias + + V, K, Q = torch.split( + param, + [ + self.d_inner_local_tp * M, + self.ngroups_local_tp * self.d_state * M, + self.ngroups_local_tp * self.d_state, + ], + dim=0, + ) + + # V: (M * d_inner_tp, ...) -> slice heads for this CP rank + V = V.view(M, self.d_inner_local_tp, *extra_dims) + v_size = self.d_inner_local_tpcp + v_start = self.cp_rank * v_size + V_sliced = V[:, v_start : v_start + v_size, ...].reshape(M * v_size, *extra_dims) + + # K: (M * ngroups_tp * d_state, ...) -> slice groups for this CP rank + K = K.view(M, self.ngroups_local_tp * self.d_state, *extra_dims) + k_size = self.ngroups_local_tpcp * self.d_state + k_start = (self.cp_rank // self.group_repeat_count) * k_size + K_sliced = K[:, k_start : k_start + k_size, ...].reshape(M * k_size, *extra_dims) + + # Q: (ngroups_tp * d_state, ...) -> slice groups (single copy, no M) + q_size = self.ngroups_local_tpcp * self.d_state + q_start = (self.cp_rank // self.group_repeat_count) * q_size + Q_sliced = Q[q_start : q_start + q_size, ...] + + return torch.cat([V_sliced, K_sliced, Q_sliced], dim=0).contiguous() + + def _slice_vector_param(self, param: torch.Tensor, has_hdim: bool = False) -> torch.Tensor: + """ + Slices a per-head vector parameter (dt_bias, A_log, D) for the current CP rank. + These are single-copy (no householder dimension). + """ + if self.cp_size == 1: + return param + + size = self.d_inner_local_tpcp if has_hdim else self.nheads_local_tpcp + start = self.cp_rank * size + return param[start : start + size] From 299c118de6c9385123629ceb4c9f983066b567f7 Mon Sep 17 00:00:00 2001 From: "Mikail Khona (NVIDIA)" Date: Thu, 23 Jul 2026 20:56:18 -0700 Subject: [PATCH 02/22] Skip Muon on GDP's inproj matrix (#119) Signed-off-by: Mikail Khona Signed-off-by: Mikail Khona (NVIDIA) Co-authored-by: Mikail Khona Signed-off-by: Keshav Santhanam --- .../core/optimizer/emerging_optimizers.py | 21 ++++++++++--------- megatron/core/ssm/gated_delta_product.py | 3 +++ 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 53ac956b35c..a42b6238e06 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -130,6 +130,11 @@ def _is_nonlinear_or_embedding(param): return getattr(param, 'is_embedding_or_output_parameter', False) or len(param.shape) != 2 +def _is_muon_excluded(param): + """True for parameters that should use the scalar optimizer instead of Muon.""" + return not getattr(param, 'use_muon', True) or _is_nonlinear_or_embedding(param) + + def _get_qkv_split_shapes(model_cfg) -> list[int]: """Compute QKV split shapes from model config.""" query_projection_size = ( @@ -467,11 +472,9 @@ def _default_adam_based_eopt_config_to_kwargs( init_state_fn=_eopt_init_state_fn, config_to_kwargs=_muon_config_to_kwargs, default_param_overrides={ - ParamKey( - predicate=ParamPredicate( - name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding - ) - ): {'optimizer': 'adam'} + ParamKey(predicate=ParamPredicate(name="muon_excluded", fn=_is_muon_excluded)): { + 'optimizer': 'adam' + } }, ), "adaptive_muon": EmergingOptimizerEntry( @@ -479,11 +482,9 @@ def _default_adam_based_eopt_config_to_kwargs( init_state_fn=_eopt_init_state_fn, config_to_kwargs=_adaptive_muon_config_to_kwargs, default_param_overrides={ - ParamKey( - predicate=ParamPredicate( - name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding - ) - ): {'optimizer': 'adam'} + ParamKey(predicate=ParamPredicate(name="muon_excluded", fn=_is_muon_excluded)): { + 'optimizer': 'adam' + } }, ), } diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 3236fa000cf..afbb63b6824 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -200,6 +200,9 @@ def __init__( tp_comm_buffer_name="fc1", tp_group=self.pg_collection.tp, ) + setattr(self.in_proj.weight, "use_muon", False) + if self.in_proj.bias is not None: + setattr(self.in_proj.bias, "use_muon", False) conv_dim = ( self.d_inner_local_tp * self.num_householder From a3feea26af24e53cbdb01650b4ceccabe0991352 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Fri, 24 Jul 2026 11:10:23 +0200 Subject: [PATCH 03/22] Fix GDP checkpoint resharding across TP sizes Give each householder copy its own checkpoint key in GatedDeltaProductMixer's sharded_state_dict so distributed checkpointing concatenates TP shards within a copy before the copies are merged. This keeps resharded (e.g. TP=2 -> TP=1) in_proj/conv1d tensors in the semantic [M0-all-ranks, M1-all-ranks, ...] order that the forward rearranges expect. Add _get_in_proj_checkpoint_split_layout and _get_conv_checkpoint_split_layout helpers (applied to both weight and bias) and a resharding unit test. Squashed from two commits on kezhik/dev-arch-mar2026: 465077d86347 Fix GDP checkpoint resharding across TP sizes 2759392384fc Fix GDP checkpoint resharding unit test Adapted to the renamed gated_delta_product.py module (was gated_delta_product_original_v4.py on the source branch); the test import was updated to match. Signed-off-by: Deepak Narayanan Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 96 ++++++++++++++----- .../unit_tests/ssm/test_gdp_tp_checkpoint.py | 95 ++++++++++++++++++ 2 files changed, 168 insertions(+), 23 deletions(-) create mode 100644 tests/unit_tests/ssm/test_gdp_tp_checkpoint.py diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index afbb63b6824..edbaaf92d89 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -592,19 +592,31 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_state_dict[f"{prefix}in_proj.weight"], ) - sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory( - sharded_state_dict[f"{prefix}in_proj.weight"], - [ - self.d_inner_local_tp, - self.d_inner_local_tp * self.num_householder, - self.ngroups_local_tp * self.d_state * self.num_householder, - self.ngroups_local_tp * self.d_state, - self.nheads_local_tp * self.num_householder, - self.nheads_local_tp, - ], - ["z", "V", "K", "Q", "b", "a"], - 0, + # V, K, and b are laid out householder-major on every TP rank: + # + # rank r: [M0-local-r, M1-local-r, ..., M(M-1)-local-r] + # + # Treating the entire M-expanded block as one TP shard would make a + # resharded TP=1 tensor rank-major instead: + # + # [rank0-all-M, rank1-all-M, ...] + # + # That ordering is incompatible with the forward rearranges, which + # expect [M0-all-ranks, M1-all-ranks, ...]. Give every householder + # copy its own checkpoint key so DCP concatenates TP shards within a + # copy before the copies are concatenated by the factory merge. + in_proj_split_sections, in_proj_split_names = _get_in_proj_checkpoint_split_layout( + self.d_inner_local_tp, + self.ngroups_local_tp * self.d_state, + self.nheads_local_tp, + self.num_householder, ) + for in_proj_param in ["in_proj.weight", "in_proj.bias"]: + key = f"{prefix}{in_proj_param}" + if key in sharded_state_dict: + sharded_state_dict[key] = _split_tensor_factory( + sharded_state_dict[key], in_proj_split_sections, in_proj_split_names, 0 + ) conv_dim = ( self.d_inner_local_tp * self.num_householder @@ -615,21 +627,59 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_state_dict[f"{prefix}conv1d.weight"], ) - for conv_layer_name in ["conv1d.weight"]: - sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory( - sharded_state_dict[f"{prefix}{conv_layer_name}"], - [ - self.d_inner_local_tp * self.num_householder, - self.ngroups_local_tp * self.d_state * self.num_householder, - self.ngroups_local_tp * self.d_state, - ], - ["V", "K", "Q"], - 0, - ) + conv_split_sections, conv_split_names = _get_conv_checkpoint_split_layout( + self.d_inner_local_tp, self.ngroups_local_tp * self.d_state, self.num_householder + ) + for conv_param in ["conv1d.weight", "conv1d.bias"]: + key = f"{prefix}{conv_param}" + if key in sharded_state_dict: + sharded_state_dict[key] = _split_tensor_factory( + sharded_state_dict[key], conv_split_sections, conv_split_names, 0 + ) return sharded_state_dict +def _get_in_proj_checkpoint_split_layout( + d_inner_local_tp: int, group_state_local_tp: int, nheads_local_tp: int, num_householder: int +) -> Tuple[List[int], List[str]]: + """Return TP-reshardable splits for the packed ``[z,V,K,Q,b,a]`` projection.""" + sections = ( + [d_inner_local_tp] + + [d_inner_local_tp] * num_householder + + [group_state_local_tp] * num_householder + + [group_state_local_tp] + + [nheads_local_tp] * num_householder + + [nheads_local_tp] + ) + names = ( + ["z"] + + [f"V{i}" for i in range(num_householder)] + + [f"K{i}" for i in range(num_householder)] + + ["Q"] + + [f"b{i}" for i in range(num_householder)] + + ["a"] + ) + return sections, names + + +def _get_conv_checkpoint_split_layout( + d_inner_local_tp: int, group_state_local_tp: int, num_householder: int +) -> Tuple[List[int], List[str]]: + """Return TP-reshardable splits for the packed ``[V,K,Q]`` convolution.""" + sections = ( + [d_inner_local_tp] * num_householder + + [group_state_local_tp] * num_householder + + [group_state_local_tp] + ) + names = ( + [f"V{i}" for i in range(num_householder)] + + [f"K{i}" for i in range(num_householder)] + + ["Q"] + ) + return sections, names + + def _split_tensor_factory( orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int ) -> ShardedTensorFactory: diff --git a/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py b/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py new file mode 100644 index 00000000000..740600f03a3 --- /dev/null +++ b/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py @@ -0,0 +1,95 @@ +"""Regression tests for GDP tensor-parallel checkpoint resharding.""" + +from collections import defaultdict + +import torch + +from megatron.core.dist_checkpointing import ShardedTensor +from megatron.core.ssm.gated_delta_product import ( + _get_in_proj_checkpoint_split_layout, + _split_tensor_factory, +) + + +def test_householder_components_reshard_tp2_to_tp1_in_semantic_order(): + """Each householder copy must gather across TP ranks before copies are concatenated.""" + num_householder = 3 + local_sections, names = _get_in_proj_checkpoint_split_layout( + d_inner_local_tp=2, + group_state_local_tp=1, + nheads_local_tp=1, + num_householder=num_householder, + ) + + # Local layout is [z, V0, V1, V2, K0, K1, K2, Q, b0, b1, b2, a]. + rank_data = [ + torch.tensor([0, 1, 10, 11, 20, 21, 30, 31, 40, 50, 60, 70, 80, 90, 100, 110]), + torch.tensor([2, 3, 12, 13, 22, 23, 32, 33, 41, 51, 61, 71, 81, 91, 101, 111]), + ] + + checkpoint_chunks = defaultdict(list) + for tp_rank, local_data in enumerate(rank_data): + sharded_tensor = ShardedTensor.from_rank_offsets( + "in_proj.weight", local_data, (0, tp_rank, 2) + ) + factory = _split_tensor_factory(sharded_tensor, local_sections, names, split_dim=0) + for chunk in factory.build(): + checkpoint_chunks[chunk.key].append(chunk.data) + + # Simulate DCP assembling every semantic checkpoint key for a TP=1 load. + assembled_checkpoint = { + key: torch.cat(chunks, dim=0) for key, chunks in checkpoint_chunks.items() + } + + global_sections, global_names = _get_in_proj_checkpoint_split_layout( + d_inner_local_tp=4, + group_state_local_tp=2, + nheads_local_tp=2, + num_householder=num_householder, + ) + target_tensor = ShardedTensor.from_rank_offsets( + "in_proj.weight", torch.empty(sum(global_sections), dtype=torch.int64), (0, 0, 1) + ) + target_factory = _split_tensor_factory( + target_tensor, global_sections, global_names, split_dim=0 + ) + loaded_chunks = [assembled_checkpoint[chunk.key] for chunk in target_factory.build()] + reloaded = target_factory.merge_fn(loaded_chunks) + + expected = torch.tensor( + [ + 0, + 1, + 2, + 3, + 10, + 11, + 12, + 13, + 20, + 21, + 22, + 23, + 30, + 31, + 32, + 33, + 40, + 41, + 50, + 51, + 60, + 61, + 70, + 71, + 80, + 81, + 90, + 91, + 100, + 101, + 110, + 111, + ] + ) + torch.testing.assert_close(reloaded, expected) From d669b4ac267e0c8d523d4154d2c7150d3e17c683 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Fri, 24 Jul 2026 11:40:22 +0200 Subject: [PATCH 04/22] Add GDP dynamic inference support Add a dynamic-batching inference path to GatedDeltaProductMixer that separates decode and prefill requests, runs each through the GDP kernels, and merges the results back into packed token order. Decode uses causal_conv1d_update plus the fused recurrent gated-delta-rule kernel with per-request state gathered and scattered through DynamicInferenceContext's slot-indexed caches; prefill runs a single variable-length chunk_gated_delta_product call. MVP scope excludes context parallelism, speculative decoding, chunked prefill, prefix caching, and CUDA-graph capture. Cherry-picked from kezhik/dev-arch-mar2026 commit 296b04884456 and adapted to this fork: - Applied onto the renamed gated_delta_product.py (was gated_delta_product_original_v4.py on the source branch). - Ported megatron/core/ssm/_packed_seq_helpers.py, a dependency the source branch already had but this fork lacked. Only check_fla_sequence_packing_support is used here, so the import is narrowed to that symbol to avoid unused imports. Signed-off-by: Deepak Narayanan Signed-off-by: Keshav Santhanam --- megatron/core/ssm/_packed_seq_helpers.py | 87 +++++++ megatron/core/ssm/gated_delta_product.py | 317 ++++++++++++++++++++++- 2 files changed, 401 insertions(+), 3 deletions(-) create mode 100644 megatron/core/ssm/_packed_seq_helpers.py diff --git a/megatron/core/ssm/_packed_seq_helpers.py b/megatron/core/ssm/_packed_seq_helpers.py new file mode 100644 index 00000000000..d12e87d8ad7 --- /dev/null +++ b/megatron/core/ssm/_packed_seq_helpers.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared helpers for SSM mixers handling packed (THD-format) sequences. + +Lifted from `MambaMixer._create_packed_seq_idx` so GDP, KDA, DPv2, GDN can +share a single reference implementation (avoids drift across mixers). +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.utils import is_causal_conv1d_min_version + + +def get_cu_seqlens(packed_seq_params: PackedSeqParams) -> torch.Tensor: + """Pick the right cu_seqlens tensor (padded if available).""" + if packed_seq_params.cu_seqlens_q_padded is not None: + return packed_seq_params.cu_seqlens_q_padded + return packed_seq_params.cu_seqlens_q + + +def build_packed_seq_idx(packed_seq_params: PackedSeqParams, total_tokens: int) -> torch.Tensor: + """Build the per-token sequence index tensor used by varlen kernels. + + For ``packed_seq_params.cu_seqlens_q[_padded]`` of the form + ``[0, 5, 7, 11]`` and ``total_tokens=16`` returns + ``[0,0,0,0,0, 1,1, 2,2,2,2, 3,3,3,3,3]`` (shape ``[1, total_tokens]``, + int32). The trailing chunk after ``cu_seqlens[-1]`` is treated as one + extra sequence so the output covers every token in the pack. If + ``cu_seqlens[-1] == total_tokens`` no extra index is added. + + This is the per-token tensor consumed by ``causal_conv1d_fn(seq_idx=...)`` + and by Mamba's fused conv+SSM kernel as ``seq_idx``. + + ``total_tokens`` must equal the *post-parallelism-gather* sequence length + that the kernel will actually consume — not the caller's + ``hidden_states.shape[0]`` which may be sequence-parallel-sharded and/or + context-parallel-sliced. The robust pattern (mirrors ``mamba_mixer.py``) + is to call this *after* ``in_proj`` (SP all-gather) and ``pre_conv_ssm`` + (CP all-to-all), passing the post-gather tensor's seq dim — that way + the helper is agnostic to TP/SP/CP shapes upstream. + """ + cu_seqlens = get_cu_seqlens(packed_seq_params) + # Guard against a caller passing an upstream-sliced ``total_tokens`` + # (e.g. ``hidden_states.shape[0]`` from before the SP all-gather / CP + # all-to-all). Without this check, the trailing-chunk diff below goes + # negative and ``repeat_interleave`` fails with the unhelpful message + # ``repeats can not be negative``. + last_cu = int(cu_seqlens[-1].item()) + assert total_tokens >= last_cu, ( + f"build_packed_seq_idx: total_tokens={total_tokens} is smaller than " + f"cu_seqlens[-1]={last_cu}. This usually means the caller passed an " + f"upstream-sliced seq_len (SP-sharded or pre-CP-all-to-all). Pass the " + f"post-gather length instead — e.g. ``zVKQba.shape[0]`` taken after " + f"``self.cp.pre_conv_ssm(...)``." + ) + total_tokens_tensor = torch.tensor( + [total_tokens], dtype=cu_seqlens.dtype, device=cu_seqlens.device + ) + cu_seqlens_with_max = torch.cat([cu_seqlens, total_tokens_tensor]) + seq_lengths = cu_seqlens_with_max[1:] - cu_seqlens_with_max[:-1] + seq_idx = torch.repeat_interleave( + torch.arange(seq_lengths.numel(), device=cu_seqlens.device), seq_lengths + ) + return seq_idx.to(torch.int32).unsqueeze(0) + + +def check_fla_sequence_packing_support() -> Tuple[bool, Optional[str]]: + """Lighter sibling of `_check_mamba_sequence_packing_support` for FLA-backed mixers. + + GDP/KDA/DPv2/GDN reach into FLA's chunk_kda / chunk_gated_delta_product / + chunk_gated_delta_rule, all of which manage their own variable-length + state internally. The only shared external dependency is the causal + conv1d kernel — `causal_conv1d_fn(seq_idx=...)` was added in 1.4.0 and + is required to reset the conv state at packed-document boundaries. + + Mamba2's stricter `mamba_ssm` minimums (used by `mamba_split_conv1d_scan_combined`) + do not apply. + """ + conv1d_min = "1.4.0" + if not is_causal_conv1d_min_version(conv1d_min): + return False, f"causal_conv1d >= {conv1d_min} is required for packed sequences" + return True, None diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index edbaaf92d89..610911758d8 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -15,8 +15,14 @@ from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory -from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext +from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( + tensor_get_slice_after, + tensor_masked_update, + tensor_merge, +) from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm._packed_seq_helpers import check_fla_sequence_packing_support from megatron.core.ssm.gdp_context_parallel import GDPContextParallel from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig @@ -26,13 +32,15 @@ make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import deprecate_inference_params +from megatron.core.utils import deprecate_inference_params, is_using_quantization_scales try: from causal_conv1d import causal_conv1d_fn, causal_conv1d_update + from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states except ImportError: causal_conv1d_fn = None causal_conv1d_update = None + causal_conv1d_varlen_states = None try: from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated @@ -335,9 +343,11 @@ def forward( conv_state, ssm_state = None, None if inference_context is not None: + if inference_context.is_dynamic_batching(): + return self._dynamic_inference(hidden_states, inference_context) assert ( inference_context.is_static_batching() - ), "Mamba does not currently support dynamic inference batching." + ), "GDP inference must be either static or dynamic batching." assert not self.config.sequence_parallel conv_state, ssm_state = self._get_states_from_cache(inference_context, batch_size) @@ -481,6 +491,307 @@ def forward( return out, out_bias + # ------------------------------------------------------------------ + # Dynamic-batching inference. + # + # Mirrors ``MambaMixer._dynamic_inference`` / ``_ssm_decode`` / ``_ssm_prefill`` + # (same ``_ssm_`` naming and the same request-level control flow), but runs + # the Gated Delta Product kernels instead of the Mamba2 scan. The per-request + # recurrent state (short-conv state + matrix-valued SSM state) is read/written + # through the slot-indexed caches owned by ``DynamicInferenceContext``. + # + # MVP scope: this path does not yet support context parallelism (cp_size > 1), + # speculative decoding, chunked prefill, Mamba prefix caching, or CUDA-graph + # capture. The reshapes mirror the static ``forward`` math with batch/seq + # repurposed for the packed dynamic layout. + # ------------------------------------------------------------------ + def _dynamic_inference( + self, hidden_states: torch.Tensor, context: DynamicInferenceContext + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Execute one dynamic inference step by separating decode and prefill + requests, running each through the GDP kernels independently, and merging + the results back into packed token order.""" + ok, reason = check_fla_sequence_packing_support() + assert ok, reason + assert self.cp.cp_size == 1, "Context parallel is not supported for GDP dynamic inference" + assert ( + not context.is_chunked_prefill_enabled() + ), "GDP dynamic inference does not support chunked prefill yet." + + # GDP-style layers register as Mamba layers, so the same (conv_state, + # ssm_state) accessor and per-layer slab layout apply. + conv_state, ssm_state = context.mamba_states_cache(self.layer_number - self.pp_layer_offset) + + padded_dims = context.padded_batch_dimensions + token_count = padded_dims.token_count + decode_req_count = padded_dims.decode_req_count + prefill_req_count = padded_dims.prefill_req_count + metadata = context.mamba_metadata + + # Input projection over the full packed batch. + zVKQba, _ = self.in_proj(hidden_states) + + y_decode = None + y_prefill = None + + # --- Decode partition (placed first in the packed batch) --------- + if decode_req_count > 0: + # MVP: exactly one token per decode request (no speculative tokens). + zVKQba_decode = zVKQba[:decode_req_count] if prefill_req_count > 0 else zVKQba + y_decode = self._ssm_decode( + zVKQba_decode.transpose(0, 1), conv_state, ssm_state, metadata.batch_indices_decode + ).transpose(0, 1) + + # --- Prefill partition ------------------------------------------- + if prefill_req_count > 0: + if decode_req_count > 0: + # Mixed batch: gather the prefill tokens out of the packed tensor. + zVKQba_prefill = torch.empty_like(zVKQba) + tensor_get_slice_after( + zVKQba, zVKQba_prefill, metadata.device_decode_prefill, check_bounds=False + ) + else: + zVKQba_prefill = zVKQba + y_prefill = self._ssm_prefill( + zVKQba_prefill, + conv_state=conv_state, + ssm_state=ssm_state, + seq_idx=metadata.seq_idx, + cu_seqlens=metadata.cu_seqlens, + batch_indices=metadata.batch_indices_prefill, + ) + + # --- Merge back into packed token order -------------------------- + if y_decode is not None and y_prefill is not None: + y = torch.empty( + [token_count, 1, y_prefill.shape[-1]], + dtype=y_prefill.dtype, + device=y_prefill.device, + ) + tensor_merge(y_decode, y_prefill, metadata.device_decode_prefill, output_tensor=y) + elif y_decode is not None: + y = y_decode + elif y_prefill is not None: + y = y_prefill + else: + raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests") + + # Zero padding positions to avoid corrupting quantization amax calculations. + if is_using_quantization_scales(self.config): + y[context.padding_slice] = 0.0 + + out, out_bias = self.out_proj(y) + return out, out_bias + + def _ssm_decode( + self, + zVKQba: torch.Tensor, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + batch_indices: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Single-token-per-request decode. ``zVKQba`` is ``[1, decode_req_count, + proj_dim]``; returns ``[1, decode_req_count, d_inner]``. The conv and SSM + states are read/written in place at the slots named by ``batch_indices`` + (``-1`` marks padding slots).""" + seq_len, _, _ = zVKQba.shape + assert seq_len == 1, "GDP decode supports one token per request" + zVKQba = zVKQba.squeeze(0) # [n, proj_dim] + M = self.num_householder + + z, VKQ, ba = torch.split( + zVKQba, + [ + self.cp.d_inner_local_tpcp, + self.cp.d_inner_local_tpcp * M + + (M + 1) * self.cp.ngroups_local_tpcp * self.d_state, + self.cp.nheads_local_tpcp * (M + 1), + ], + dim=-1, + ) + + # Indexed conv update: reads/writes the per-request conv state rows + # selected by ``batch_indices``, in place. ``self.activation`` must be the + # activation *string* so the kernel enables SiLU (a bool would disable it). + VKQ = causal_conv1d_update( + VKQ, + conv_state, + rearrange(self.conv1d.weight, "d 1 w -> d w"), + self.conv1d.bias, + self.activation, + conv_state_indices=batch_indices, + ) + + value, key, query = torch.split( + VKQ, + [ + self.cp.d_inner_local_tpcp * M, + self.cp.ngroups_local_tpcp * self.d_state * M, + self.cp.ngroups_local_tpcp * self.d_state, + ], + dim=-1, + ) + b, a = torch.split(ba, [self.cp.nheads_local_tpcp * M, self.cp.nheads_local_tpcp], dim=-1) + + # Reshape to the fla layout with batch=n requests, seq length 1, and the + # householder copies folded into the sequence dimension (static path, l=1). + value = rearrange(value, "n (m h p) -> n m h p", m=M, p=self.headdim).contiguous() + key = rearrange(key, "n (m g s) -> n m g s", m=M, s=self.d_state).contiguous() + query = rearrange(query, "n (g s) -> n 1 g s", s=self.d_state).contiguous() + z = rearrange(z, "n (h p) -> n 1 h p", p=self.headdim).contiguous() + beta = rearrange(b.sigmoid(), "n (m h) -> n m h", m=M).contiguous() + g = -self.cp.get_A_log().float().exp() * F.softplus(a.float() + self.cp.get_dt_bias()) + g = rearrange(g, "n h -> n 1 h").contiguous() + + if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: + rep = self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp + query = query.repeat_interleave(rep, dim=2) + key = key.repeat_interleave(rep, dim=2) + + # Interleave the (length-1) query / decay with householder zeros so the + # recurrent kernel sees an (1 * M)-length sequence (matches static decode). + g_new = g.new_zeros(g.shape[0], g.shape[1], M, g.shape[2]) + g_new[:, :, 0] = g + g = rearrange(g_new, "n t m h -> n (t m) h") + query_new = query.new_zeros( + query.shape[0], query.shape[1], M, query.shape[2], query.shape[3] + ) + query_new[:, :, -1] = query + query = rearrange(query_new, "n t m h d -> n (t m) h d") + + # Gather this step's per-request initial states. ``.clamp`` (NOT in-place) + # returns a new tensor, so ``batch_indices`` keeps its -1 padding sentinels + # for the scatter below; the padding rows' outputs are never scattered back. + gather_idx = batch_indices.clamp(min=0) + initial_state = ssm_state[gather_idx] + + core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + core_attn_out = rearrange(core_attn_out, "n (t m) h d -> n t m h d", m=M)[ + ..., -1, :, : + ].contiguous() # [n, 1, h, d] + + # Scatter updated states back into the cache (skips -1 padding slots). + tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) + + y = rearrange(core_attn_out, "n t h p -> t n (h p)").contiguous() # [1, n, d_inner] + if self.rmsnorm: + z = rearrange(z, "n t h p -> t n (h p)").contiguous() + y = self.norm(y, z) + return y + + def _ssm_prefill( + self, + zVKQba: torch.Tensor, + conv_state: Optional[torch.Tensor] = None, + ssm_state: Optional[torch.Tensor] = None, + seq_idx: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + batch_indices: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Variable-length prefill over all prefill requests in one varlen call. + ``zVKQba`` is ``[l, 1, proj_dim]``; returns ``[l, 1, d_inner]``. Fresh + requests start from a zero recurrent state (no prefix caching in the MVP); + the resulting final conv/SSM states are written back into the caches.""" + is_dynamic_batching = seq_idx is not None + M = self.num_householder + + # l b d -> b l d + zVKQba = rearrange(zVKQba, "l b d -> b l d").contiguous() + + z, VKQ, ba = torch.split( + zVKQba, + [ + self.cp.d_inner_local_tpcp, + self.cp.d_inner_local_tpcp * M + + (M + 1) * self.cp.ngroups_local_tpcp * self.d_state, + self.cp.nheads_local_tpcp * (M + 1), + ], + dim=-1, + ) + + if conv_state is not None and is_dynamic_batching: + assert batch_indices is not None + # Capture per-request final conv states (before the conv consumes the + # inputs) and write them into the prefill requests' cache rows. + conv_varlen_states = causal_conv1d_varlen_states( + VKQ.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] + ) + tensor_masked_update(conv_state, batch_indices, conv_varlen_states) + # Maintain channels-last memory layout so causal_conv1d_fn can use seq_idx. + VKQ = VKQ.transpose(1, 2) + else: + VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + + seqlen = VKQ.size(2) + if causal_conv1d_fn is None: + VKQ = self.act(self.cp.conv1d(VKQ)[..., :seqlen]) + else: + VKQ = causal_conv1d_fn( + x=VKQ, + weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + bias=self.cp.get_conv1d_bias(), + activation=self.activation, + seq_idx=seq_idx, + ) + VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() + + value, key, query = torch.split( + VKQ, + [ + self.cp.d_inner_local_tpcp * M, + self.cp.ngroups_local_tpcp * self.d_state * M, + self.cp.ngroups_local_tpcp * self.d_state, + ], + dim=-1, + ) + b, a = torch.split(ba, [self.cp.nheads_local_tpcp * M, self.cp.nheads_local_tpcp], dim=-1) + + # batch = 1 packed sequence of length T; householder folded into seq. + value = rearrange(value, "b l (m h p) -> b (l m) h p", m=M, p=self.headdim).contiguous() + key = rearrange(key, "b l (m g s) -> b (l m) g s", m=M, s=self.d_state).contiguous() + query = rearrange(query, "b l (g s) -> b l g s", s=self.d_state).contiguous() + z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).contiguous() + beta = rearrange(b.sigmoid(), "b l (m h) -> b (l m) h", m=M).contiguous() + g = -self.cp.get_A_log().float().exp() * F.softplus(a.float() + self.cp.get_dt_bias()) + g = g.contiguous() + + if self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp > 1: + rep = self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp + query = query.repeat_interleave(rep, dim=2) + key = key.repeat_interleave(rep, dim=2) + + core_attn_out, last_recurrent_state = chunk_gated_delta_product( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=ssm_state is not None, + num_householder=M, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + + # Write per-request final SSM states into the cache for subsequent decode. + if ssm_state is not None and is_dynamic_batching: + tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) + + y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() + if self.rmsnorm: + z = rearrange(z, "b l h p -> l b (h p)").contiguous() + y = self.norm(y, z) + return y + def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None): """ allocate inference cache From 210ae414beae9aa7ad6046e55de29764d92d049f Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Fri, 24 Jul 2026 18:09:45 +0200 Subject: [PATCH 05/22] Add inter-document masking (packed sequence) support to GDP training path Wire THD/packed-sequence support into GatedDeltaProductMixer's training and prefill forward path: build cu_seqlens and a per-token seq_idx from packed_seq_params (after the in_proj sequence-parallel all-gather and the context-parallel all-to-all), thread seq_idx through causal_conv1d_fn to reset convolution boundaries at document edges, and pass cu_seqlens to chunk_gated_delta_product. GDPContextParallel.pre_conv_ssm/post_conv_ssm now take packed_seq_params so the load-balancing undo/redo uses the packed layout. Add self.chunk_size and a causal_conv1d version check in __init__. Integrates NVIDIA-NeMo/nv-mistralai-megatron PR #122 (commit b8726f2edf4b, "IDM for GDP from internal gitlab") onto this branch: - Reconciled the forward() inference branch with the dynamic-inference path already on staging; #122's "no packed sequences during inference" assert now sits alongside the dynamic-batching dispatch. - Kept the _packed_seq_helpers.py port already added here (content-identical to #122's; only formatting differed) and dropped #122's duplicate copy. Signed-off-by: Deepak Narayanan Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 61 +++- megatron/core/ssm/gdp_context_parallel.py | 23 +- tests/unit_tests/ssm/test_gdp_packed_seq.py | 315 ++++++++++++++++++++ 3 files changed, 385 insertions(+), 14 deletions(-) create mode 100644 tests/unit_tests/ssm/test_gdp_packed_seq.py diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 610911758d8..ec1663f772d 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -22,7 +22,11 @@ tensor_merge, ) from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm._packed_seq_helpers import check_fla_sequence_packing_support +from megatron.core.ssm._packed_seq_helpers import ( + build_packed_seq_idx, + check_fla_sequence_packing_support, + get_cu_seqlens, +) from megatron.core.ssm.gdp_context_parallel import GDPContextParallel from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig @@ -164,6 +168,15 @@ def __init__( super().__init__(config) + # Inference-time contract: ``MambaInferenceStateConfig.from_model`` in + # megatron/core/inference/config.py reads ``layer.mixer.chunk_size`` to + # size the SSM scan blocks. + self.chunk_size = chunk_size + + # Check that the causal_conv1d version is new enough or fail + ok, reason = check_fla_sequence_packing_support() + assert ok, reason + self.num_householder = 3 self.config = config @@ -334,11 +347,6 @@ def forward( packed_seq_params=None, ): """Run the gated delta product mixer on hidden states.""" - if packed_seq_params is not None: - raise NotImplementedError( - "GatedDeltaProductMixer does not support packed sequences yet." - ) - seq_len, batch_size, dim = hidden_states.shape conv_state, ssm_state = None, None @@ -349,11 +357,32 @@ def forward( inference_context.is_static_batching() ), "GDP inference must be either static or dynamic batching." assert not self.config.sequence_parallel + assert packed_seq_params is None, ( + "GDP does not currently support packed sequences during inference. " + "Packing is only wired through the training/prefill (chunk) path." + ) conv_state, ssm_state = self._get_states_from_cache(inference_context, batch_size) + # Build cu_seqlens for the chunked recurrence (FLA) when running with + # packed (THD) sequences on the training/prefill path. + cu_seqlens_packed = None + if packed_seq_params is not None: + # ``hidden_states`` is [seq_len, batch, dim]; THD requires batch=1. + assert batch_size == 1, "Packed sequences require batch=1 (THD/varlen format)." + cu_seqlens_packed = get_cu_seqlens(packed_seq_params) + zVKQba, _ = self.in_proj(hidden_states) - zVKQba = self.cp.pre_conv_ssm(zVKQba) + zVKQba = self.cp.pre_conv_ssm(zVKQba, packed_seq_params=packed_seq_params) + + # Build seq_idx *after* in_proj's SP all-gather and pre_conv_ssm's CP + # all-to-all. ``zVKQba.shape[0]`` is now the true pack_length, so the + # helper produces a seq_idx matching the conv1d's input length + # regardless of SP/CP/TP upstream-slicing. Mirrors mamba_mixer.py + # which calls _create_packed_seq_idx after the same gather points. + seq_idx_packed = None + if packed_seq_params is not None: + seq_idx_packed = build_packed_seq_idx(packed_seq_params, zVKQba.shape[0]) zVKQba = rearrange(zVKQba, "l b d -> b l d").contiguous() @@ -368,7 +397,16 @@ def forward( dim=-1, ) - VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + # ``causal_conv1d_fn`` expects a ``[B, D, L]`` tensor. + # But the expected memory layout varies depending on whether seq_idx is set. + if seq_idx_packed is None: + # Default path: channels-first contiguous, stride(2) == 1. + VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + else: + # ``causal_conv1d_fn(seq_idx=...)`` requires channels-last memory but [B, D, L] + # logical shape. This keeps the channels contiguous in memory. + VKQ = VKQ.contiguous() + VKQ = rearrange(VKQ, "b l d -> b d l") # Decode if inference_context is not None and inference_context.seqlen_offset > 0: @@ -388,11 +426,13 @@ def forward( F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0)) ) # Update state (B D W) # Train + # causal_conv1d uses seq_idx_packed to reset the convolution boundaries VKQ = causal_conv1d_fn( x=VKQ, weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), bias=self.cp.get_conv1d_bias(), activation=self.activation, + seq_idx=seq_idx_packed, ) VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() @@ -475,16 +515,17 @@ def forward( output_final_state=(ssm_state is not None), num_householder=self.num_householder, use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens_packed, ) if ssm_state is not None: ssm_state.copy_(last_recurrent_state) y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() - y = self.cp.post_conv_ssm(y) + y = self.cp.post_conv_ssm(y, packed_seq_params=packed_seq_params) if self.rmsnorm: z = rearrange(z, "b l h p -> l b (h p)").contiguous() - z = self.cp.post_conv_ssm(z) + z = self.cp.post_conv_ssm(z, packed_seq_params=packed_seq_params) y = self.norm(y, z) out, out_bias = self.out_proj(y) diff --git a/megatron/core/ssm/gdp_context_parallel.py b/megatron/core/ssm/gdp_context_parallel.py index ff118ceee5f..8447ecdd87f 100644 --- a/megatron/core/ssm/gdp_context_parallel.py +++ b/megatron/core/ssm/gdp_context_parallel.py @@ -21,10 +21,14 @@ copy is independently partitioned by heads across CP ranks. """ +from typing import Optional + import torch import torch.nn as nn import torch.nn.functional as F +from megatron.core.packed_seq_params import PackedSeqParams + try: from einops import repeat @@ -33,6 +37,8 @@ HAVE_EINOPS = False # Re-use the load balancing and all-to-all helpers from the existing module. +# The load-balancing helpers already handle packed (THD) input via their +# ``packed_seq_params`` argument, so GDP just threads it through below. from megatron.core.ssm.mamba_context_parallel import ( _all_to_all_cp2hp, _all_to_all_hp2cp, @@ -128,7 +134,9 @@ def __init__( self.group_repeat_count = 1 self.ngroups_local_tpcp = self.ngroups_local_tp // self.cp_size - def pre_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor: + def pre_conv_ssm( + self, input_: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None + ) -> torch.Tensor: """ All-to-all from sequence-partitioned to head-partitioned layout, before conv + SSM. @@ -140,6 +148,10 @@ def pre_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor: [z, V, K, Q, b, a] with sizes [d_inner/cp, d_inner/cp*M, ngroups_cp*d_state*M, ngroups_cp*d_state, nheads/cp*M, nheads/cp] + + ``packed_seq_params`` must be passed for THD/SFT input — without it + the post-all-to-all undo uses the non-packed zigzag pattern, which + scrambles token order across pack boundaries. """ if self.cp_size == 1: return input_ @@ -209,17 +221,20 @@ def pre_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor: a = _all_to_all_cp2hp(a, self.cp_group) output = torch.cat([z, V, K, Q, b_proj, a], dim=-1) - output = _undo_attention_load_balancing(output, self.cp_size) + output = _undo_attention_load_balancing(output, self.cp_size, packed_seq_params) return output - def post_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor: + def post_conv_ssm( + self, input_: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None + ) -> torch.Tensor: """Method to be applied after the conv + SSM (on y and z, which have no M dim).""" if self.cp_size == 1: return input_ else: return _all_to_all_hp2cp( - _redo_attention_load_balancing(input_, self.cp_size), self.cp_group + _redo_attention_load_balancing(input_, self.cp_size, packed_seq_params), + self.cp_group, ) def conv1d(self, input_: torch.Tensor) -> torch.Tensor: diff --git a/tests/unit_tests/ssm/test_gdp_packed_seq.py b/tests/unit_tests/ssm/test_gdp_packed_seq.py new file mode 100644 index 00000000000..10be89a70c7 --- /dev/null +++ b/tests/unit_tests/ssm/test_gdp_packed_seq.py @@ -0,0 +1,315 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GDP v4 packed-sequence + context-parallel equivalence tests. + +Verifies that ``GatedDeltaProductMixer`` (v4) produces forward outputs and +parameter gradients under ``cp_size=2`` that match a ``cp_size=1`` reference +run on the same packed (THD/SFT-format) input. + +Style mirrors ``test_mamba_context_parallel.py``: ``Utils.initialize_model_parallel``, +``@pytest.mark.internal``, fixed-seed bf16 tensors, tolerance via +``torch.testing.assert_close``. + +Run with:: + + torchrun --nproc_per_node=2 -m pytest \\ + tests/unit_tests/ssm/test_gdp_packed_seq.py -m internal -v +""" +from __future__ import annotations + +import os +from typing import List + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.extensions.transformer_engine import ( + TELayerNormColumnParallelLinear, + TERowParallelLinear, +) +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer, MambaMixerSubmodules +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + import einops # noqa: F401 + import mamba_ssm # noqa: F401 + + HAVE_MAMBA_DEPS = True +except ImportError: + HAVE_MAMBA_DEPS = False + +try: + import fla # noqa: F401 + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + + +# Skip the whole file when bare ``pytest`` is invoked outside torchrun. The +# CP=2 reference setup needs a multi-rank world to be meaningful; without a +# distributed launcher the fixture would either fail loudly or worse, hang. +# Reported as SKIPPED with a clear ``reason`` so contributors running the +# unit suite locally see why and how to invoke it properly. The torchrun +# command sets WORLD_SIZE in the environment for every worker process. +_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + _WORLD_SIZE < 2, + reason=( + "CP=2 equivalence test requires a multi-rank world; run via " + "``torchrun --nproc_per_node=2 -m pytest ... -m internal``." + ), + ), +] + + +# Pack shapes used to parametrize both forward and backward equivalence tests. +# Each segment length must be a multiple of ``2 * cp_size = 4`` so that +# ``tex.thd_get_partitioned_indices`` can split the pack evenly across CP +# ranks (mirrors ``sft_dataset.py``'s pad_granularity). +PACK_SHAPES = [ + pytest.param([16, 8, 24], id="headline"), + pytest.param([48], id="single-long"), + pytest.param([8, 8, 8, 8, 8, 8], id="many-short"), + pytest.param([4, 20, 12, 12], id="mixed-short-long"), + pytest.param([40, 4, 4], id="head-heavy"), + pytest.param([4, 4, 40], id="tail-heavy"), +] + + +def _make_packed_seq_params(seq_lens: List[int]) -> PackedSeqParams: + """Build a PackedSeqParams for a single THD pack with these segment lengths.""" + cu = torch.tensor( + [0] + list(torch.cumsum(torch.tensor(seq_lens), 0).tolist()), + dtype=torch.int32, + device="cuda", + ) + total = int(cu[-1].item()) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=None, + cu_seqlens_kv_padded=None, + max_seqlen_q=total, + max_seqlen_kv=total, + ) + + +def _make_config(cp_size: int) -> TransformerConfig: + """Small-but-shape-valid TransformerConfig for the v4 GDP mixer.""" + return TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + num_query_groups=4, + ffn_hidden_size=128, + normalization="RMSNorm", + bf16=True, + mamba_num_heads=4, + mamba_head_dim=16, + mamba_num_groups=4, + mamba_state_dim=16, + tensor_model_parallel_size=1, + sequence_parallel=False, + context_parallel_size=cp_size, + ) + + +def _build_mixer(cp_group): + """Construct a v4 GDP mixer wired to the given CP group.""" + config = _make_config(cp_group.size()) + pg = ProcessGroupCollection(tp=parallel_state.get_tensor_model_parallel_group(), cp=cp_group) + submodules = MambaMixerSubmodules( + in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear + ) + mixer = GatedDeltaProductMixer( + config=config, + submodules=submodules, + d_model=config.hidden_size, + layer_number=1, + pg_collection=pg, + ) + return mixer.cuda().bfloat16(), config + + +def _sync_weights_from_rank0(mixer): + """Broadcast every parameter from global rank 0 so all ranks share weights.""" + for p in mixer.parameters(): + torch.distributed.broadcast(p.data, src=0) + + +def _build_cp_pair(): + """Build CP=2 + per-rank CP=1 reference mixers with identical weights. + + Returns ``(mixer_cp2, mixer_cp1, config, cp_group, cp_rank)``. The cp=1 + instance lives in a 1-rank subgroup (containing only this rank), so the + same mixer code path runs in cp=1 mode and provides a numerical reference. + """ + cp_group = parallel_state.get_context_parallel_group() + cp_rank = parallel_state.get_context_parallel_rank() + global_rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + + cp1_groups = [torch.distributed.new_group(ranks=[r]) for r in range(world_size)] + cp1_group = cp1_groups[global_rank] + + mixer_cp2, _ = _build_mixer(cp_group) + mixer_cp1, config = _build_mixer(cp1_group) + _sync_weights_from_rank0(mixer_cp2) + mixer_cp1.load_state_dict(mixer_cp2.state_dict()) + return mixer_cp2, mixer_cp1, config, cp_group, cp_rank + + +def _make_hidden_packed(seq_lens, hidden_size): + """Build a packed [total_tokens, 1, hidden] input + matching PackedSeqParams. + + The tensor is broadcast from rank 0 so cp=1 reference and cp=2 sliced + paths see bit-identical input. + """ + psp = _make_packed_seq_params(seq_lens) + total_tokens = sum(seq_lens) + torch.manual_seed(0) + hidden_full = torch.randn(total_tokens, 1, hidden_size, device="cuda", dtype=torch.bfloat16) + torch.distributed.broadcast(hidden_full, src=0) + return hidden_full, psp + + +@pytest.mark.internal +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA + NCCL") +@pytest.mark.skipif( + torch.cuda.device_count() < 2 if torch.cuda.is_available() else True, + reason="CP=2 test requires at least 2 GPUs", +) +@pytest.mark.skipif(not HAVE_MAMBA_DEPS, reason="GDP mixer requires mamba_ssm + einops") +@pytest.mark.skipif(not HAVE_FLA, reason="GDP mixer requires fla") +class TestGDPPackedSequence: + """v4 GDP forward + backward equivalence under CP=2 with packed (THD) input.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + """Initialize TP=1 PP=1 CP=2 model parallel state for every test.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=2 + ) + model_parallel_cuda_manual_seed(123) + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("seq_lens", PACK_SHAPES) + def test_forward_equivalence(self, seq_lens): + """CP=2 forward output (sliced+gathered) matches CP=1 reference on the + same packed input to bf16 tolerance. + """ + import transformer_engine_torch as tex + + mixer_cp2, mixer_cp1, config, cp_group, cp_rank = _build_cp_pair() + mixer_cp2.eval() + mixer_cp1.eval() + + hidden_full, psp = _make_hidden_packed(seq_lens, config.hidden_size) + total_tokens = hidden_full.shape[0] + + with torch.no_grad(): + ref_out = mixer_cp1(hidden_full, packed_seq_params=psp) + ref_out = ref_out[0] if isinstance(ref_out, tuple) else ref_out + assert ref_out.shape == hidden_full.shape + + idx = tex.thd_get_partitioned_indices( + psp.cu_seqlens_q, total_tokens, cp_group.size(), cp_rank + ) + hidden_local = hidden_full.index_select(0, idx) + + with torch.no_grad(): + cp2_out_local = mixer_cp2(hidden_local, packed_seq_params=psp) + cp2_out_local = cp2_out_local[0] if isinstance(cp2_out_local, tuple) else cp2_out_local + + # Scatter the per-rank slice back to its original positions, then + # all-reduce(SUM) across the CP group to reconstruct the full output. + cp2_out_full = torch.zeros_like(ref_out) + scatter_index = idx.long().view(-1, 1, 1).expand_as(cp2_out_local) + cp2_out_full.scatter_(0, scatter_index, cp2_out_local) + torch.distributed.all_reduce(cp2_out_full, group=cp_group) + + torch.testing.assert_close(cp2_out_full, ref_out, atol=5e-2, rtol=5e-2) + + @pytest.mark.parametrize("seq_lens", PACK_SHAPES) + def test_backward_equivalence(self, seq_lens): + """CP=2 parameter gradients (after all-reduce across CP) match CP=1 + reference gradients on the same packed input. + + Mechanics: weights are replicated across CP ranks, so for any scalar + loss ``L_full`` computed over the full output, + ``dL_full/dW = sum_{r in cp_ranks} dL_local/dW`` where ``L_local`` is + the same loss restricted to the rank's local output slice. + + Loss = ``out.float().pow(2).sum()`` so every output element + contributes — exercises every weight that feeds the output. + Tolerance is looser than the forward test: bf16 backward accumulates + error through the chain of matmuls and the all-to-all forward+backward + in CP. + """ + import transformer_engine_torch as tex + + mixer_cp2, mixer_cp1, config, cp_group, cp_rank = _build_cp_pair() + mixer_cp2.eval() + mixer_cp1.eval() + for p in mixer_cp1.parameters(): + p.grad = None + for p in mixer_cp2.parameters(): + p.grad = None + + hidden_full, psp = _make_hidden_packed(seq_lens, config.hidden_size) + total_tokens = hidden_full.shape[0] + + # CP=1 reference: full-sequence forward + backward. + ref_out = mixer_cp1(hidden_full, packed_seq_params=psp) + ref_out = ref_out[0] if isinstance(ref_out, tuple) else ref_out + ref_loss = ref_out.float().pow(2).sum() + ref_loss.backward() + + # CP=2: local slice forward + backward. ``L_full = sum_{r in cp_ranks} L_local`` + # by construction (sum of squares is additive across token partitions), + # so the all-reduced grad equals the reference grad up to bf16 noise. + idx = tex.thd_get_partitioned_indices( + psp.cu_seqlens_q, total_tokens, cp_group.size(), cp_rank + ) + hidden_local = hidden_full.index_select(0, idx) + out_local = mixer_cp2(hidden_local, packed_seq_params=psp) + out_local = out_local[0] if isinstance(out_local, tuple) else out_local + loss_local = out_local.float().pow(2).sum() + loss_local.backward() + + cp1_params = dict(mixer_cp1.named_parameters()) + cp2_params = dict(mixer_cp2.named_parameters()) + assert set(cp1_params) == set(cp2_params) + + mismatches = [] + n_compared = 0 + for name in sorted(cp1_params): + g1 = cp1_params[name].grad + g2 = cp2_params[name].grad + if g1 is None and g2 is None: + continue + assert g1 is not None, f"cp=1 has no grad for {name} but cp=2 does" + assert g2 is not None, f"cp=2 has no grad for {name} but cp=1 does" + g2_reduced = g2.clone().contiguous() + torch.distributed.all_reduce(g2_reduced, group=cp_group) + try: + torch.testing.assert_close(g2_reduced, g1, atol=8e-2, rtol=8e-2) + n_compared += 1 + except AssertionError as e: + mismatches.append((name, tuple(g1.shape), str(e).splitlines()[0])) + assert not mismatches, ( + f"{len(mismatches)} parameter(s) mismatched out of " + f"{n_compared + len(mismatches)} compared:\n" + + "\n".join(f" {n} {s}: {m}" for n, s, m in mismatches) + ) + assert n_compared > 0, "no parameters received a gradient — test setup is wrong" From 02a7cf9b176cb54e0bc502d7cff21c8e47f32fc4 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 29 Jul 2026 18:40:28 -0700 Subject: [PATCH 06/22] Tag GDP packed parameters for refit Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index ec1663f772d..70e257759ad 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -225,6 +225,18 @@ def __init__( if self.in_proj.bias is not None: setattr(self.in_proj.bias, "use_muon", False) + # The fused projection packs independently TP-sharded components. Refit + # uses their local sizes to preserve semantic order when TP size changes. + in_proj_partition_sizes, _ = _get_in_proj_checkpoint_split_layout( + self.d_inner_local_tp, + self.ngroups_local_tp * self.d_state, + self.nheads_local_tp, + self.num_householder, + ) + setattr(self.in_proj.weight, "partition_sizes", in_proj_partition_sizes) + if self.in_proj.bias is not None: + setattr(self.in_proj.bias, "partition_sizes", in_proj_partition_sizes) + conv_dim = ( self.d_inner_local_tp * self.num_householder + (self.num_householder + 1) * self.ngroups_local_tp * self.d_state @@ -243,8 +255,17 @@ def __init__( dtype=config.params_dtype, ) setattr(self.conv1d.weight, "tensor_model_parallel", True) + setattr(self.conv1d.weight, "partition_dim", 0) if conv_bias: setattr(self.conv1d.bias, "tensor_model_parallel", True) + setattr(self.conv1d.bias, "partition_dim", 0) + + conv_partition_sizes, _ = _get_conv_checkpoint_split_layout( + self.d_inner_local_tp, self.ngroups_local_tp * self.d_state, self.num_householder + ) + setattr(self.conv1d.weight, "partition_sizes", conv_partition_sizes) + if conv_bias: + setattr(self.conv1d.bias, "partition_sizes", conv_partition_sizes) if self.conv_init is not None: nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) From c3af950aef3290a5bd9fa5623f6df3aca0c5dca6 Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Mon, 3 Aug 2026 13:36:28 -0700 Subject: [PATCH 07/22] Fail fast on invalid module specs Signed-off-by: Keshav Santhanam --- megatron/core/transformer/spec_utils.py | 17 +++++------ .../unit_tests/transformer/test_spec_utils.py | 28 ++++++++++++++++++- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/megatron/core/transformer/spec_utils.py b/megatron/core/transformer/spec_utils.py index ba9c22f01b6..0a50e2900b2 100644 --- a/megatron/core/transformer/spec_utils.py +++ b/megatron/core/transformer/spec_utils.py @@ -1,14 +1,10 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. import functools -import logging import types from collections.abc import Callable from dataclasses import dataclass, field from typing import Any, Tuple, Union -logger = logging.getLogger(__name__) - - @dataclass class ModuleSpec: """This is a Module Specification dataclass. @@ -50,10 +46,15 @@ def import_module(module_path: Tuple[str]): base_path, name = module_path try: module = __import__(base_path, globals(), locals(), [name]) - except ImportError as e: - logger.error(f"couldn't import module due to {e}") - return None - return vars(module)[name] + except ImportError as exc: + raise ImportError( + f"Could not import module '{base_path}' for spec '{name}': {exc}" + ) from exc + + try: + return vars(module)[name] + except KeyError as exc: + raise ImportError(f"Could not find spec '{name}' in module '{base_path}'") from exc # pylint: disable=missing-function-docstring diff --git a/tests/unit_tests/transformer/test_spec_utils.py b/tests/unit_tests/transformer/test_spec_utils.py index e464b09380e..78b0fb4ab02 100644 --- a/tests/unit_tests/transformer/test_spec_utils.py +++ b/tests/unit_tests/transformer/test_spec_utils.py @@ -5,7 +5,12 @@ import pytest -from megatron.core.transformer.spec_utils import ModuleSpec, build_module, get_submodules +from megatron.core.transformer.spec_utils import ( + ModuleSpec, + build_module, + get_submodules, + import_module, +) def dummy_method(x: int, y: str) -> dict: @@ -79,6 +84,27 @@ def test_build_module_by_call(self): assert mixed.y == 'ghi' +class TestImportModule: + """Unit tests for dynamic spec imports.""" + + def test_missing_module_raises(self): + with pytest.raises( + ImportError, + match="Could not import module 'megatron.core.models.does_not_exist'", + ): + import_module(('megatron.core.models.does_not_exist', 'missing_spec')) + + def test_missing_spec_raises(self): + with pytest.raises( + ImportError, + match=( + "Could not find spec 'does_not_exist' in module " + "'megatron.core.transformer.identity_op'" + ), + ): + import_module(('megatron.core.transformer.identity_op', 'does_not_exist')) + + class OtherChild: def __init__(self, x: int): self.x = x From e2c8855d0671e8a255acbfa9a6b65d90a7048d6f Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Mon, 3 Aug 2026 14:10:41 -0700 Subject: [PATCH 08/22] Fix GDP mixer module naming Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 4 ++++ tests/unit_tests/ssm/test_gdp_packed_seq.py | 1 + 2 files changed, 5 insertions(+) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 70e257759ad..e4af4e3818b 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -127,6 +127,7 @@ class GatedDeltaProductMixer(MegatronModule): layer_number: The layer number of this Mamba layer. pg_collection: The required process groups to use for tensor model parallel and context parallel. + name: Module instance name passed top-down from its parent module. """ def __init__( @@ -157,6 +158,7 @@ def __init__( ngroups=None, pg_collection: ProcessGroupCollection = None, pp_layer_offset: int = 0, + name: str | None = None, ): if not HAVE_MAMBA_SSM: raise ImportError( @@ -220,6 +222,7 @@ def __init__( is_expert=False, tp_comm_buffer_name="fc1", tp_group=self.pg_collection.tp, + name=(name + f".in_proj") if name is not None else None, ) setattr(self.in_proj.weight, "use_muon", False) if self.in_proj.bias is not None: @@ -337,6 +340,7 @@ def __init__( is_expert=False, tp_comm_buffer_name="fc2", tp_group=self.pg_collection.tp, + name=(name + f".out_proj") if name is not None else None, ) # Regarding `conv1d`.{`weight`, `bias`}, `dt_bias`, `A_log`, and `D`: these are the diff --git a/tests/unit_tests/ssm/test_gdp_packed_seq.py b/tests/unit_tests/ssm/test_gdp_packed_seq.py index 10be89a70c7..f38b7c5b67b 100644 --- a/tests/unit_tests/ssm/test_gdp_packed_seq.py +++ b/tests/unit_tests/ssm/test_gdp_packed_seq.py @@ -136,6 +136,7 @@ def _build_mixer(cp_group): d_model=config.hidden_size, layer_number=1, pg_collection=pg, + name="decoder.layers.0.mixer", ) return mixer.cuda().bfloat16(), config From 9c4f6deda2bf153bbeafb40f59c6a10601af9920 Mon Sep 17 00:00:00 2001 From: Roger Waleffe Date: Wed, 15 Jul 2026 22:32:56 -0700 Subject: [PATCH 09/22] Update FLOPs calculation to account for GDP Signed-off-by: Keshav Santhanam --- megatron/training/training.py | 65 +++++++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 3 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index acdd727b82d..0178cb6ab78 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -492,6 +492,38 @@ def mamba_layer_flops(total_tokens, hidden_size, state_dim=16, + (2 * total_tokens * d_in * hidden_size) # out_proj ) + def gated_delta_product_layer_flops( + total_tokens, + hidden_size, + state_dim=128, + head_dim=64, + num_groups=8, + num_heads=None, + conv_kernel_dim=4, + num_householder=3, + ): + """Calculate FLOPs for a Gated Delta Product (GDP) layer.""" + if num_heads is None: + d_inner = 2 * hidden_size + num_heads = d_inner // head_dim + else: + d_inner = num_heads * head_dim + in_proj_dim = ( + d_inner * (1 + num_householder) + + num_groups * state_dim * (1 + num_householder) + + num_heads * (1 + num_householder) + ) + conv_dim = d_inner * num_householder + num_groups * state_dim * (1 + num_householder) + non_core_flops = 2 * total_tokens * ( + hidden_size * in_proj_dim + + conv_kernel_dim * conv_dim + + d_inner * hidden_size + ) + # Best-case recurrent GDP core estimate. The FLA chunk kernel may do additional + # score/solve/WY work, but this keeps the implementation-agnostic lower bound explicit. + core_flops = (4 * num_householder + 3) * total_tokens * d_inner * state_dim + return non_core_flops + core_flops + def gdn_layer_flops(total_tokens, hidden_size, qk_head_dim=128, v_head_dim=128, num_qk_heads=16, num_v_heads=32, @@ -522,20 +554,28 @@ def hybrid_flops(total_tokens, seqlen_squared_sum, hidden_size, mlp_expansion=4.0, swiglu=False, moe_latent_size=None, moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1, + use_gated_delta_product=False, gdn_qk_head_dim=128, gdn_v_head_dim=128, gdn_num_qk_heads=16, gdn_num_v_heads=32, gdn_conv_kernel_dim=4, vocab_size=256000, mtp_num_layers=0): """Calculate total FLOPs for the hybrid model.""" + mamba_flops = ( + gated_delta_product_layer_flops(total_tokens, hidden_size, + mamba_state_dim, mamba_head_dim, + mamba_num_groups, mamba_num_heads) + if use_gated_delta_product + else mamba_layer_flops(total_tokens, hidden_size, + mamba_state_dim, mamba_head_dim, + mamba_num_groups, mamba_num_heads) + ) flops_fwd = ( num_attn_layers * attn_layer_flops(total_tokens, seqlen_squared_sum, hidden_size, num_attn_heads, gqa, gqa_groups, kv_channels) + num_mlp_layers * mlp_layer_flops(total_tokens, hidden_size, mlp_expansion, swiglu) + - num_mamba_layers * mamba_layer_flops(total_tokens, hidden_size, - mamba_state_dim, mamba_head_dim, - mamba_num_groups, mamba_num_heads) + + num_mamba_layers * mamba_flops + num_moe_layers * moe_layer_flops(total_tokens, hidden_size, moe_ffn_hidden_size, shared_expert_ffn_hidden_size, num_experts_routed_to, moe_latent_size, swiglu) + @@ -846,6 +886,24 @@ def transformer_flops(): ) return total_floating_point_operations + def _uses_gated_delta_product_spec(args): + """Return True when the selected hybrid stack spec swaps Mamba layers to GDP.""" + def _split_spec_part(part): + return str(part).replace('[', ' ').replace(']', ' ').replace(',', ' ').split() + + spec = getattr(args, 'spec', None) + if spec is None: + return False + if isinstance(spec, str): + spec_parts = _split_spec_part(spec) + else: + spec_parts = [] + for part in spec: + spec_parts.extend(_split_spec_part(part)) + if not spec_parts: + return False + return spec_parts[-1] in {'gdp_stack_spec', 'gated_delta_product_stack_spec'} + # Main entrypoint for FLOPs calculation. if is_hybrid_model(args): # Calculate the number of each type of layer. @@ -884,6 +942,7 @@ def transformer_flops(): kv_channels=args.kv_channels, mlp_expansion=args.ffn_hidden_size / args.hidden_size, swiglu=args.swiglu, + use_gated_delta_product=_uses_gated_delta_product_spec(args), moe_latent_size=args.moe_latent_size, moe_ffn_hidden_size=(args.moe_ffn_hidden_size if args.moe_ffn_hidden_size is not None else args.ffn_hidden_size), From 0e952670cf401b11c53f720929bc9bd18e8912dd Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 11:06:45 -0700 Subject: [PATCH 10/22] fix(optimizer): skip Muon for Mamba input projection Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/ssm/mamba_mixer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index e374592a125..80c8c894c55 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -305,6 +305,7 @@ def __init__( self.nheads_local_tp, # dt ] setattr(self.in_proj.weight, "partition_sizes", in_proj_partition_sizes) + setattr(self.in_proj.weight, "use_muon", False) if not self.use_mem_eff_path: log_single_rank( From f286a4cf25cbd099cbff0299d8209cfb41b0a88d Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 11:12:51 -0700 Subject: [PATCH 11/22] chore: add missing GDP test copyright header Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- tests/unit_tests/ssm/test_gdp_tp_checkpoint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py b/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py index 740600f03a3..8c4b2f1e456 100644 --- a/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py +++ b/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py @@ -1,3 +1,5 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + """Regression tests for GDP tensor-parallel checkpoint resharding.""" from collections import defaultdict From 2feb52b7b02b575900b9926b75451202901021b2 Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 11:33:02 -0700 Subject: [PATCH 12/22] chore: apply PR formatting fixes Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/transformer/spec_utils.py | 1 + tests/unit_tests/transformer/test_spec_utils.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/spec_utils.py b/megatron/core/transformer/spec_utils.py index 0a50e2900b2..36c7001988b 100644 --- a/megatron/core/transformer/spec_utils.py +++ b/megatron/core/transformer/spec_utils.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, field from typing import Any, Tuple, Union + @dataclass class ModuleSpec: """This is a Module Specification dataclass. diff --git a/tests/unit_tests/transformer/test_spec_utils.py b/tests/unit_tests/transformer/test_spec_utils.py index 78b0fb4ab02..6e86e4d7d2f 100644 --- a/tests/unit_tests/transformer/test_spec_utils.py +++ b/tests/unit_tests/transformer/test_spec_utils.py @@ -89,8 +89,7 @@ class TestImportModule: def test_missing_module_raises(self): with pytest.raises( - ImportError, - match="Could not import module 'megatron.core.models.does_not_exist'", + ImportError, match="Could not import module 'megatron.core.models.does_not_exist'" ): import_module(('megatron.core.models.does_not_exist', 'missing_spec')) From 2fe04d629436e2ba0beb13fdc1a83b867ab47560 Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 12:27:22 -0700 Subject: [PATCH 13/22] refactor(ssm): rename packed sequence helper module Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 4 ++-- .../ssm/{_packed_seq_helpers.py => packed_seq_helpers.py} | 0 2 files changed, 2 insertions(+), 2 deletions(-) rename megatron/core/ssm/{_packed_seq_helpers.py => packed_seq_helpers.py} (100%) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index e4af4e3818b..93420904773 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -22,12 +22,12 @@ tensor_merge, ) from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm._packed_seq_helpers import ( +from megatron.core.ssm.gdp_context_parallel import GDPContextParallel +from megatron.core.ssm.packed_seq_helpers import ( build_packed_seq_idx, check_fla_sequence_packing_support, get_cu_seqlens, ) -from megatron.core.ssm.gdp_context_parallel import GDPContextParallel from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig from megatron.core.transformer.module import MegatronModule diff --git a/megatron/core/ssm/_packed_seq_helpers.py b/megatron/core/ssm/packed_seq_helpers.py similarity index 100% rename from megatron/core/ssm/_packed_seq_helpers.py rename to megatron/core/ssm/packed_seq_helpers.py From 20d9d5b9b60e5efcf88bc5324e95cc0a424ea2ee Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 12:46:07 -0700 Subject: [PATCH 14/22] docs(ssm): describe gated delta product mixer Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 93420904773..9f926740b47 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -99,7 +99,24 @@ class MambaMixerSubmodules: class GatedDeltaProductMixer(MegatronModule): - """ + """Gated Delta Product (GDP) sequence mixer for hybrid models. + + The mixer accepts hidden states with shape ``[sequence, batch, hidden]`` and returns + a projected tensor with the same shape plus the optional output-projection bias. It + serves as the mixer inside ``MambaLayer``, allowing a hybrid stack to select GDP layers + without changing the surrounding layer interface. + + GDP projects each token into an output gate and the ``V``, ``K``, ``Q``, beta, and decay + terms used by a sequence of Householder updates. A depthwise causal convolution mixes + local context in ``V/K/Q``; the FLA GDP recurrence then updates a matrix-valued state, + and gated RMS normalization plus the output projection map the result back to the + model hidden size. + + The module shards projections and recurrent parameters across tensor-parallel ranks, + redistributes sequence and head dimensions for context parallelism, handles packed + THD training sequences, manages static and dynamic inference state, and exposes + semantic sharded-state-dict partitions for checkpoint resharding across TP sizes. + Args: config: The config of the model. submodules: Contains the module specs for the input and output linear layers. From fe926f08d6a769e664b89b6ffe1871d203d9193d Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 16:06:51 -0700 Subject: [PATCH 15/22] refactor(ssm): remove unused GDP constructor arguments Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 14 -------------- tests/unit_tests/ssm/test_gdp_tp_checkpoint.py | 17 +++++++++++++++++ 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 9f926740b47..101e00bfd9f 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -121,12 +121,8 @@ class GatedDeltaProductMixer(MegatronModule): config: The config of the model. submodules: Contains the module specs for the input and output linear layers. d_model: The hidden size of the model. - d_state: The state size of the SSM. d_conv: The number of channels in the causal convolution. conv_init: The initialization range for the causal convolution weights. - expand: The expansion factor for the SSM. - headdim: The hidden size of each attention head. - ngroups: The number of attention heads. A_init_range: The initialization range for the attention weights. D_has_hdim: Whether the D parameter has the same number of dimensions as the hidden state. @@ -134,13 +130,10 @@ class GatedDeltaProductMixer(MegatronModule): norm_before_gate: Whether to apply normalization before the gating mechanism. dt_min: The minimum value of the dt parameter. dt_max: The maximum value of the dt parameter. - dt_init: The initialization value of the dt parameter. - dt_scale: The scaling factor for the dt parameter. dt_init_floor: The minimum value of the dt parameter after initialization. bias: Whether to use bias in the linear layers. conv_bias: Whether to use bias in the causal convolution. chunk_size: The chunk size for the fused kernel. - use_mem_eff_path: Whether to use the memory-efficient path for the Mamba model. layer_number: The layer number of this Mamba layer. pg_collection: The required process groups to use for tensor model parallel and context parallel. @@ -154,25 +147,18 @@ def __init__( d_model, d_conv=4, conv_init=None, - expand=2, A_init_range=(0, 16), D_has_hdim=False, rmsnorm=True, norm_before_gate=False, dt_min=0.001, dt_max=0.1, - dt_init="random", - dt_scale=1.0, dt_init_floor=1e-4, bias=False, conv_bias=False, # Fused kernel and sharding options chunk_size=128, layer_number=None, - use_mem_eff_path=None, - d_state=None, - headdim=None, - ngroups=None, pg_collection: ProcessGroupCollection = None, pp_layer_offset: int = 0, name: str | None = None, diff --git a/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py b/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py index 8c4b2f1e456..2636e95ab8b 100644 --- a/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py +++ b/tests/unit_tests/ssm/test_gdp_tp_checkpoint.py @@ -2,17 +2,34 @@ """Regression tests for GDP tensor-parallel checkpoint resharding.""" +import inspect from collections import defaultdict import torch from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.ssm.gated_delta_product import ( + GatedDeltaProductMixer, _get_in_proj_checkpoint_split_layout, _split_tensor_factory, ) +def test_constructor_drops_unused_mamba_compatibility_arguments(): + """GDP dimensions come from TransformerConfig, not ignored constructor overrides.""" + parameters = inspect.signature(GatedDeltaProductMixer.__init__).parameters + removed_parameters = { + "expand", + "dt_init", + "dt_scale", + "use_mem_eff_path", + "d_state", + "headdim", + "ngroups", + } + assert parameters.keys().isdisjoint(removed_parameters) + + def test_householder_components_reshard_tp2_to_tp1_in_semantic_order(): """Each householder copy must gather across TP ranks before copies are concatenated.""" num_householder = 3 From 0d03d60ea4341ed687c20cbe29bb2b5713a3952c Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 17:04:41 -0700 Subject: [PATCH 16/22] fix(ssm): configure GDP Householder reflections Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 3 +- .../core/transformer/transformer_config.py | 8 +++++ megatron/training/checkpointing.py | 6 ++++ megatron/training/training.py | 5 ++- tests/unit_tests/test_checkpointing.py | 26 ++++++++++++++++ .../test_num_floating_point_operations.py | 31 +++++++++++++++++++ .../transformer/test_transformer_config.py | 25 +++++++++++++++ 7 files changed, 102 insertions(+), 2 deletions(-) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 101e00bfd9f..366b8b43a2b 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -182,7 +182,7 @@ def __init__( ok, reason = check_fla_sequence_packing_support() assert ok, reason - self.num_householder = 3 + self.num_householder = config.gdp_num_householder self.config = config self.d_model = d_model @@ -198,6 +198,7 @@ def __init__( self.headdim = self.config.mamba_head_dim self.ngroups = self.config.mamba_num_groups self.nheads = self.config.mamba_num_heads + assert self.nheads is not None, "mamba_num_heads must be set for GatedDeltaProductMixer" self.d_inner = self.nheads * self.headdim self.layer_number = layer_number diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..35199b78d26 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1189,6 +1189,9 @@ class TransformerConfig(ModelParallelConfig): """The number of heads used in Mamba layers. If None, the number of heads will be hidden_size * expand // mamba_head_dim.""" + gdp_num_householder: int = 3 + """The number of Householder reflections used in Gated Delta Product layers.""" + mamba_training_ssm_states_dtype: Optional[torch.dtype] = None """dtype of the materialized inter-chunk SSM states in Mamba training forwards and backwards. None causes the states to follow the activation dtype.""" @@ -1316,6 +1319,11 @@ def __post_init__(self): f"Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True." ) + if self.gdp_num_householder < 1: + raise ValueError( + f"gdp_num_householder must be positive, got {self.gdp_num_householder}." + ) + # Apply BF16 matmul precision setting if needed if self.bf16 and self.disable_bf16_reduced_precision_matmul: torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index ded7fc4d70f..4783af3e390 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -168,6 +168,8 @@ def _compare(arg_name, old_arg_name=None, default=None): _compare('num_layers') _compare('hidden_size') _compare('num_attention_heads') + if hasattr(args, 'gdp_num_householder'): + _compare('gdp_num_householder', default=3) _compare('add_position_embedding', default=True) if args.vocab_file: _compare('max_position_embeddings') @@ -2010,6 +2012,10 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('mamba_head_dim', force=True) _set_arg('mamba_num_groups', force=True) _set_arg('mamba_num_heads', force=True) + # GDP checkpoints created before this argument existed always used three reflections. + if not hasattr(checkpoint_args, 'gdp_num_householder'): + setattr(checkpoint_args, 'gdp_num_householder', 3) + _set_arg('gdp_num_householder', force=True) # We need to be able to override hybrid_layer_pattern from the command-line so that different # pipelining can be specified when re-loading a model (e.g. for inference or post-training). _set_arg('hybrid_layer_pattern') diff --git a/megatron/training/training.py b/megatron/training/training.py index 0178cb6ab78..5c82dc8ae7f 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -495,12 +495,12 @@ def mamba_layer_flops(total_tokens, hidden_size, state_dim=16, def gated_delta_product_layer_flops( total_tokens, hidden_size, + num_householder, state_dim=128, head_dim=64, num_groups=8, num_heads=None, conv_kernel_dim=4, - num_householder=3, ): """Calculate FLOPs for a Gated Delta Product (GDP) layer.""" if num_heads is None: @@ -546,6 +546,7 @@ def gdn_layer_flops(total_tokens, hidden_size, def hybrid_flops(total_tokens, seqlen_squared_sum, hidden_size, num_attn_layers, num_mamba_layers, num_mlp_layers, num_moe_layers, + gdp_num_householder, num_gdn_layers=0, mamba_state_dim=128, mamba_head_dim=64, mamba_num_groups=8, mamba_num_heads=128, @@ -562,6 +563,7 @@ def hybrid_flops(total_tokens, seqlen_squared_sum, hidden_size, """Calculate total FLOPs for the hybrid model.""" mamba_flops = ( gated_delta_product_layer_flops(total_tokens, hidden_size, + gdp_num_householder, mamba_state_dim, mamba_head_dim, mamba_num_groups, mamba_num_heads) if use_gated_delta_product @@ -936,6 +938,7 @@ def _split_spec_part(part): mamba_head_dim=args.mamba_head_dim, mamba_num_groups=args.mamba_num_groups, mamba_num_heads=args.mamba_num_heads, + gdp_num_householder=args.gdp_num_householder, num_attn_heads=args.num_attention_heads, gqa=args.group_query_attention, gqa_groups=args.num_query_groups, diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index 2e717e4424f..15ab11458ae 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -24,6 +24,7 @@ _build_sharded_state_dict_metadata, _load_base_checkpoint, get_checkpoint_tracker_filename, + load_args_from_checkpoint, load_checkpoint, maybe_save_dataloader_state, read_metadata, @@ -149,6 +150,31 @@ def test_maybe_save_dataloader_state_skips_empty_state_after_barriers(tmp_path): save.assert_not_called() +@pytest.mark.parametrize( + ("checkpoint_args", "configured_num_householder", "expected_num_householder"), + [(SimpleNamespace(gdp_num_householder=5), 3, 5), (SimpleNamespace(), 5, 3)], +) +def test_load_args_restores_gdp_num_householder_from_checkpoint( + checkpoint_args, configured_num_householder, expected_num_householder +): + args = SimpleNamespace( + load="checkpoint", + iteration=0, + gdp_num_householder=configured_num_householder, + use_tokenizer_model_from_checkpoint_args=False, + use_mp_args_from_checkpoint_args=False, + ) + state_dict = {"args": checkpoint_args, "iteration": 12} + + with mock.patch( + "megatron.training.checkpointing._load_base_checkpoint", + return_value=(state_dict, "checkpoint", False, CheckpointType.LEGACY), + ): + restored_args, _ = load_args_from_checkpoint(args) + + assert restored_args.gdp_num_householder == expected_num_householder + + def create_checkpoint(load_path, ckpt_format): """Setup a dummy checkpoint directory.""" iteration = 123 diff --git a/tests/unit_tests/test_num_floating_point_operations.py b/tests/unit_tests/test_num_floating_point_operations.py index f358ada8fc1..df5e4191843 100644 --- a/tests/unit_tests/test_num_floating_point_operations.py +++ b/tests/unit_tests/test_num_floating_point_operations.py @@ -96,6 +96,7 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s args.mamba_head_dim = 64 args.mamba_num_groups = 8 args.mamba_num_heads = 128 + args.gdp_num_householder = 3 return args @@ -255,6 +256,36 @@ def test_hybrid_attention_layers_count(self): assert flops_doubled - flops_bshd == expected_delta +class TestGatedDeltaProductFlops: + """GDP FLOPs must use the Householder count from the model configuration.""" + + def test_householder_count_changes_flops(self): + args = _make_hybrid_args() + args.spec = ["megatron.core.models.hybrid.hybrid_layer_specs", "gdp_stack_spec"] + batch_size = 4 + + flops_m3 = num_floating_point_operations(args, batch_size) + args.gdp_num_householder = 4 + flops_m4 = num_floating_point_operations(args, batch_size) + + total_tokens = batch_size * args.seq_length + d_inner = args.mamba_num_heads * args.mamba_head_dim + group_state_dim = args.mamba_num_groups * args.mamba_state_dim + forward_delta_per_layer = ( + 2 + * total_tokens + * ( + args.hidden_size * (d_inner + group_state_dim + args.mamba_num_heads) + + 4 * (d_inner + group_state_dim) + ) + + 4 * total_tokens * d_inner * args.mamba_state_dim + ) + num_gdp_layers = 2 + expected_delta = 3 * num_gdp_layers * forward_delta_per_layer + + assert flops_m4 - flops_m3 == expected_delta + + class TestPaddingRemoval: """``total_real_tokens_in_batch`` removes padding from token-linear FLOPs. diff --git a/tests/unit_tests/transformer/test_transformer_config.py b/tests/unit_tests/transformer/test_transformer_config.py index febb3842789..24339c12b5a 100644 --- a/tests/unit_tests/transformer/test_transformer_config.py +++ b/tests/unit_tests/transformer/test_transformer_config.py @@ -30,3 +30,28 @@ def test_ep_a2a_overlap_accepts_supported_mtp_layer_counts(mtp_num_layers: int | def test_ep_a2a_overlap_rejects_unsupported_mtp_layer_counts(mtp_num_layers: int): with pytest.raises(AssertionError, match="MTP supports at most one layer"): _make_overlap_config(mtp_num_layers) + + +def test_gdp_num_householder_defaults_to_three(): + config = TransformerConfig(num_layers=1, hidden_size=128, num_attention_heads=4) + + assert config.gdp_num_householder == 3 + + +def test_gdp_num_householder_accepts_positive_values(): + config = TransformerConfig( + num_layers=1, hidden_size=128, num_attention_heads=4, gdp_num_householder=5 + ) + + assert config.gdp_num_householder == 5 + + +@pytest.mark.parametrize("num_householder", [0, -1]) +def test_gdp_num_householder_rejects_non_positive_values(num_householder: int): + with pytest.raises(ValueError, match="gdp_num_householder must be positive"): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + gdp_num_householder=num_householder, + ) From 0ceca771fc26ec99023d583525eddb22d9b4e441 Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 17:37:27 -0700 Subject: [PATCH 17/22] perf(ssm): avoid packed sequence host sync Packed sequence indices are built on every SSM layer forward. Reading cu_seqlens[-1] with .item() synchronizes the GPU and CPU even though the value is used only for validation. Remove the eager host read and let repeat_interleave reject invalid negative lengths. The number of output indices is already known from total_tokens, so pass it as output_size. This also avoids synchronizing to infer the output shape and keeps index construction on-device. Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/ssm/packed_seq_helpers.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/megatron/core/ssm/packed_seq_helpers.py b/megatron/core/ssm/packed_seq_helpers.py index d12e87d8ad7..c959f4b2793 100644 --- a/megatron/core/ssm/packed_seq_helpers.py +++ b/megatron/core/ssm/packed_seq_helpers.py @@ -45,26 +45,15 @@ def build_packed_seq_idx(packed_seq_params: PackedSeqParams, total_tokens: int) the helper is agnostic to TP/SP/CP shapes upstream. """ cu_seqlens = get_cu_seqlens(packed_seq_params) - # Guard against a caller passing an upstream-sliced ``total_tokens`` - # (e.g. ``hidden_states.shape[0]`` from before the SP all-gather / CP - # all-to-all). Without this check, the trailing-chunk diff below goes - # negative and ``repeat_interleave`` fails with the unhelpful message - # ``repeats can not be negative``. - last_cu = int(cu_seqlens[-1].item()) - assert total_tokens >= last_cu, ( - f"build_packed_seq_idx: total_tokens={total_tokens} is smaller than " - f"cu_seqlens[-1]={last_cu}. This usually means the caller passed an " - f"upstream-sliced seq_len (SP-sharded or pre-CP-all-to-all). Pass the " - f"post-gather length instead — e.g. ``zVKQba.shape[0]`` taken after " - f"``self.cp.pre_conv_ssm(...)``." - ) total_tokens_tensor = torch.tensor( [total_tokens], dtype=cu_seqlens.dtype, device=cu_seqlens.device ) cu_seqlens_with_max = torch.cat([cu_seqlens, total_tokens_tensor]) seq_lengths = cu_seqlens_with_max[1:] - cu_seqlens_with_max[:-1] seq_idx = torch.repeat_interleave( - torch.arange(seq_lengths.numel(), device=cu_seqlens.device), seq_lengths + torch.arange(seq_lengths.numel(), device=cu_seqlens.device), + seq_lengths, + output_size=total_tokens, ) return seq_idx.to(torch.int32).unsqueeze(0) From cd159648b805efc95338070da219aa01577200ee Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Wed, 5 Aug 2026 17:49:10 -0700 Subject: [PATCH 18/22] refactor(ssm): name GDP mixer submodules explicitly Rename the GDP-specific submodule dataclass at its definition site and update its type annotations and call sites. This removes the alias required to distinguish it from the unrelated Mamba mixer submodule dataclass and reduces the risk of wrong imports. Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- megatron/core/models/hybrid/hybrid_layer_specs.py | 4 ++-- megatron/core/ssm/gated_delta_product.py | 4 ++-- tests/unit_tests/ssm/test_gdp_packed_seq.py | 7 +++++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 0ee36323c95..8e91f442e13 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -16,9 +16,9 @@ ) from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules -from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer from megatron.core.ssm.gated_delta_product import ( - MambaMixerSubmodules as GatedDeltaProductMixerSubmodules, + GatedDeltaProductMixer, + GatedDeltaProductMixerSubmodules, ) from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 366b8b43a2b..62f86e842d3 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -89,7 +89,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): @dataclass -class MambaMixerSubmodules: +class GatedDeltaProductMixerSubmodules: """ Contains the module specs for the input and output linear layers. """ @@ -143,7 +143,7 @@ class GatedDeltaProductMixer(MegatronModule): def __init__( self, config: TransformerConfig, - submodules: MambaMixerSubmodules, + submodules: GatedDeltaProductMixerSubmodules, d_model, d_conv=4, conv_init=None, diff --git a/tests/unit_tests/ssm/test_gdp_packed_seq.py b/tests/unit_tests/ssm/test_gdp_packed_seq.py index f38b7c5b67b..a39a2b8a23a 100644 --- a/tests/unit_tests/ssm/test_gdp_packed_seq.py +++ b/tests/unit_tests/ssm/test_gdp_packed_seq.py @@ -30,7 +30,10 @@ ) from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer, MambaMixerSubmodules +from megatron.core.ssm.gated_delta_product import ( + GatedDeltaProductMixer, + GatedDeltaProductMixerSubmodules, +) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -127,7 +130,7 @@ def _build_mixer(cp_group): """Construct a v4 GDP mixer wired to the given CP group.""" config = _make_config(cp_group.size()) pg = ProcessGroupCollection(tp=parallel_state.get_tensor_model_parallel_group(), cp=cp_group) - submodules = MambaMixerSubmodules( + submodules = GatedDeltaProductMixerSubmodules( in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear ) mixer = GatedDeltaProductMixer( From 25386da400083cf6576888ed2dd691b55e6c02fb Mon Sep 17 00:00:00 2001 From: Kezhi Kong Date: Thu, 6 Aug 2026 11:03:19 -0700 Subject: [PATCH 19/22] test(models): update hybrid MoE config golden Record the new gdp_num_householder TransformerConfig field and its backward-compatible default of three in the Mamba MoE golden configuration. This keeps the config drift test aligned with the intentional GDP configuration surface. Signed-off-by: Kezhi Kong Signed-off-by: Keshav Santhanam --- tests/unit_tests/models/test_hybrid_moe_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index eb871568046..45482ca5a20 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -134,6 +134,7 @@ "fused_residual_rmsnorm": False, "fused_single_qkv_rope": False, "gated_linear_unit": False, + "gdp_num_householder": 3, "gtp_weight_remat_size": 1, "glu_linear_offset": 0.0, "grad_scale_func": None, From 6878aeb60195aa80b22774189a56d8b7fe15fc97 Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Fri, 7 Aug 2026 15:04:06 -0700 Subject: [PATCH 20/22] Linting Signed-off-by: Keshav Santhanam --- tests/unit_tests/ssm/test_gdp_packed_seq.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/ssm/test_gdp_packed_seq.py b/tests/unit_tests/ssm/test_gdp_packed_seq.py index a39a2b8a23a..417f74b6183 100644 --- a/tests/unit_tests/ssm/test_gdp_packed_seq.py +++ b/tests/unit_tests/ssm/test_gdp_packed_seq.py @@ -15,6 +15,7 @@ torchrun --nproc_per_node=2 -m pytest \\ tests/unit_tests/ssm/test_gdp_packed_seq.py -m internal -v """ + from __future__ import annotations import os From 052282e3d258b40e0dff365de8ca47caecce39ad Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Mon, 10 Aug 2026 15:17:32 -0700 Subject: [PATCH 21/22] refactor(ssm): move GDP dynamic inference onto the shared SSM interface Gated Delta Product carried its own `_dynamic_inference` / `_ssm_decode` / `_ssm_prefill` trio, duplicating the request-level control flow that `SSMDynamicInferenceMixin` already owns for Mamba2: fetch the per-layer (conv_state, ssm_state) slabs, project, split the packed batch into decode and prefill partitions, run each through its kernels, merge back into packed token order, and project out. Subclass the mixin and implement only the two variant hooks. `ssm_decode` now takes the mixin's batch-first `[n, seq_len, proj_dim]` layout and rejects the speculative-decoding intermediate buffers explicitly rather than silently ignoring them; `ssm_prefill` reads its varlen metadata off the context instead of taking it as an argument list, and owns the chunked-prefill assertion, as the interface prescribes. Quarantine the static-batching path the same way MambaMixer does. Static decode moves out of the shared `forward` body into `_static_decode`, which delegates to `ssm_decode` with `batch_indices=None`; `forward` keeps only the training and static-prefill body. This removes the two `seqlen_offset > 0` branches that threaded static-batching bookkeeping through the training math. No functional change: `pre_conv_ssm` / `post_conv_ssm` are identity at cp_size == 1, which static decode already required. Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 314 +++++++++-------------- 1 file changed, 122 insertions(+), 192 deletions(-) diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 62f86e842d3..1baaea0ea3a 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -17,9 +17,7 @@ from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( - tensor_get_slice_after, tensor_masked_update, - tensor_merge, ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gdp_context_parallel import GDPContextParallel @@ -28,6 +26,7 @@ check_fla_sequence_packing_support, get_cu_seqlens, ) +from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.transformer import TransformerConfig from megatron.core.transformer.module import MegatronModule @@ -36,7 +35,7 @@ make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import deprecate_inference_params, is_using_quantization_scales +from megatron.core.utils import deprecate_inference_params try: from causal_conv1d import causal_conv1d_fn, causal_conv1d_update @@ -98,7 +97,7 @@ class GatedDeltaProductMixerSubmodules: out_proj: Union[ModuleSpec, type] = None -class GatedDeltaProductMixer(MegatronModule): +class GatedDeltaProductMixer(SSMDynamicInferenceMixin, MegatronModule): """Gated Delta Product (GDP) sequence mixer for hybrid models. The mixer accepts hidden states with shape ``[sequence, batch, hidden]`` and returns @@ -381,7 +380,12 @@ def forward( conv_state, ssm_state = None, None if inference_context is not None: if inference_context.is_dynamic_batching(): - return self._dynamic_inference(hidden_states, inference_context) + ok, reason = check_fla_sequence_packing_support() + assert ok, reason + assert ( + self.cp.cp_size == 1 + ), "Context parallel is not supported for GDP dynamic inference" + return self.ssm_dynamic_inference(hidden_states, inference_context) assert ( inference_context.is_static_batching() ), "GDP inference must be either static or dynamic batching." @@ -391,6 +395,9 @@ def forward( "Packing is only wired through the training/prefill (chunk) path." ) conv_state, ssm_state = self._get_states_from_cache(inference_context, batch_size) + if inference_context.seqlen_offset > 0: + # The states are updated in place. + return self._static_decode(hidden_states, conv_state, ssm_state) # Build cu_seqlens for the chunked recurrence (FLA) when running with # packed (THD) sequences on the training/prefill path. @@ -437,32 +444,19 @@ def forward( VKQ = VKQ.contiguous() VKQ = rearrange(VKQ, "b l d -> b d l") - # Decode - if inference_context is not None and inference_context.seqlen_offset > 0: - VKQ = causal_conv1d_update( - VKQ, - conv_state, - rearrange(self.conv1d.weight, "d 1 w -> d w"), - self.conv1d.bias, - self.activation, - ) - else: - # Prefill - if conv_state is not None: - # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv - # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. - conv_state.copy_( - F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0)) - ) # Update state (B D W) - # Train - # causal_conv1d uses seq_idx_packed to reset the convolution boundaries - VKQ = causal_conv1d_fn( - x=VKQ, - weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), - bias=self.cp.get_conv1d_bias(), - activation=self.activation, - seq_idx=seq_idx_packed, - ) + if conv_state is not None: + # Static-batching prefill: seed the conv state from the prompt's tail. + # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv + # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. + conv_state.copy_(F.pad(VKQ, (self.d_conv - VKQ.shape[-1], 0))) # Update state (B D W) + # causal_conv1d uses seq_idx_packed to reset the convolution boundaries + VKQ = causal_conv1d_fn( + x=VKQ, + weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + bias=self.cp.get_conv1d_bias(), + activation=self.activation, + seq_idx=seq_idx_packed, + ) VKQ = rearrange(VKQ, "b d l -> b l d").contiguous() @@ -506,46 +500,18 @@ def forward( self.cp.nheads_local_tpcp // self.cp.ngroups_local_tpcp, dim=2 ) - # Decode - if inference_context is not None and inference_context.seqlen_offset > 0: - - g_new = g.new_zeros(g.shape[0], g.shape[1], self.num_householder, g.shape[2]) - g_new[:, :, 0] = g - g = rearrange(g_new, '... t n h -> ... (t n) h') - - query_new = query.new_zeros( - query.shape[0], query.shape[1], self.num_householder, query.shape[2], query.shape[3] - ) - query_new[:, :, -1] = query - query = rearrange(query_new, '... t n h d-> ... (t n) h d') - - core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=ssm_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - ) - core_attn_out = rearrange( - core_attn_out, '... (t n) h d -> ... t n h d', n=self.num_householder - )[..., -1, :, :].contiguous() - # Train or Prefill - else: - core_attn_out, last_recurrent_state = chunk_gated_delta_product( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=(ssm_state is not None), - num_householder=self.num_householder, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens_packed, - ) + core_attn_out, last_recurrent_state = chunk_gated_delta_product( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=(ssm_state is not None), + num_householder=self.num_householder, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens_packed, + ) if ssm_state is not None: ssm_state.copy_(last_recurrent_state) @@ -561,112 +527,70 @@ def forward( return out, out_bias + # ================================================================== + # Static / eager inference + # + # ``_static_decode`` implements legacy static-batching decode. It is + # deliberately kept separate from the dynamic inference hooks below so that + # static-batching bookkeeping does not pollute the interface defined by + # ``SSMDynamicInferenceMixin``. Static-batching prefill shares the training + # body in ``forward``, which seeds the conv/SSM state when the caches are + # present. Mirrors ``MambaMixer._static_decode``. + # ================================================================== + def _static_decode( + self, hidden_states: torch.Tensor, conv_state: torch.Tensor, ssm_state: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Single-token static-batching decode step (updates state in place).""" + assert hidden_states.shape[0] == 1, "Only support decoding with 1 token at a time for now" + assert self.cp.cp_size == 1, "Context parallel not supported for GDP inference decode" + + # (1, b, d_model) -> (1, b, proj_dim) + zVKQba, _ = self.in_proj(hidden_states) + + # The decode kernels are batch-first: (1, b, proj_dim) -> (b, 1, proj_dim). + # Static batching has no slot remapping, so batch_indices is None. + y = self.ssm_decode( + zVKQba.transpose(0, 1), conv_state=conv_state, ssm_state=ssm_state, batch_indices=None + ) + + # (b, 1, d_inner) -> (1, b, d_inner), which is what out_proj expects. + return self.out_proj(y.transpose(0, 1)) + # ------------------------------------------------------------------ # Dynamic-batching inference. # - # Mirrors ``MambaMixer._dynamic_inference`` / ``_ssm_decode`` / ``_ssm_prefill`` - # (same ``_ssm_`` naming and the same request-level control flow), but runs - # the Gated Delta Product kernels instead of the Mamba2 scan. The per-request - # recurrent state (short-conv state + matrix-valued SSM state) is read/written - # through the slot-indexed caches owned by ``DynamicInferenceContext``. + # These are the two hooks required by ``SSMDynamicInferenceMixin``; the + # mixin owns the surrounding decode/prefill partitioning and merge. They run + # the Gated Delta Product kernels instead of the Mamba2 scan. The + # per-request recurrent state (short-conv state + matrix-valued SSM state) + # is read/written through the slot-indexed caches owned by + # ``DynamicInferenceContext``. # # MVP scope: this path does not yet support context parallelism (cp_size > 1), # speculative decoding, chunked prefill, Mamba prefix caching, or CUDA-graph # capture. The reshapes mirror the static ``forward`` math with batch/seq # repurposed for the packed dynamic layout. # ------------------------------------------------------------------ - def _dynamic_inference( - self, hidden_states: torch.Tensor, context: DynamicInferenceContext - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Execute one dynamic inference step by separating decode and prefill - requests, running each through the GDP kernels independently, and merging - the results back into packed token order.""" - ok, reason = check_fla_sequence_packing_support() - assert ok, reason - assert self.cp.cp_size == 1, "Context parallel is not supported for GDP dynamic inference" - assert ( - not context.is_chunked_prefill_enabled() - ), "GDP dynamic inference does not support chunked prefill yet." - - # GDP-style layers register as Mamba layers, so the same (conv_state, - # ssm_state) accessor and per-layer slab layout apply. - conv_state, ssm_state = context.mamba_states_cache(self.layer_number - self.pp_layer_offset) - - padded_dims = context.padded_batch_dimensions - token_count = padded_dims.token_count - decode_req_count = padded_dims.decode_req_count - prefill_req_count = padded_dims.prefill_req_count - metadata = context.mamba_metadata - - # Input projection over the full packed batch. - zVKQba, _ = self.in_proj(hidden_states) - - y_decode = None - y_prefill = None - - # --- Decode partition (placed first in the packed batch) --------- - if decode_req_count > 0: - # MVP: exactly one token per decode request (no speculative tokens). - zVKQba_decode = zVKQba[:decode_req_count] if prefill_req_count > 0 else zVKQba - y_decode = self._ssm_decode( - zVKQba_decode.transpose(0, 1), conv_state, ssm_state, metadata.batch_indices_decode - ).transpose(0, 1) - - # --- Prefill partition ------------------------------------------- - if prefill_req_count > 0: - if decode_req_count > 0: - # Mixed batch: gather the prefill tokens out of the packed tensor. - zVKQba_prefill = torch.empty_like(zVKQba) - tensor_get_slice_after( - zVKQba, zVKQba_prefill, metadata.device_decode_prefill, check_bounds=False - ) - else: - zVKQba_prefill = zVKQba - y_prefill = self._ssm_prefill( - zVKQba_prefill, - conv_state=conv_state, - ssm_state=ssm_state, - seq_idx=metadata.seq_idx, - cu_seqlens=metadata.cu_seqlens, - batch_indices=metadata.batch_indices_prefill, - ) - - # --- Merge back into packed token order -------------------------- - if y_decode is not None and y_prefill is not None: - y = torch.empty( - [token_count, 1, y_prefill.shape[-1]], - dtype=y_prefill.dtype, - device=y_prefill.device, - ) - tensor_merge(y_decode, y_prefill, metadata.device_decode_prefill, output_tensor=y) - elif y_decode is not None: - y = y_decode - elif y_prefill is not None: - y = y_prefill - else: - raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests") - - # Zero padding positions to avoid corrupting quantization amax calculations. - if is_using_quantization_scales(self.config): - y[context.padding_slice] = 0.0 - - out, out_bias = self.out_proj(y) - return out, out_bias - - def _ssm_decode( + def ssm_decode( self, zVKQba: torch.Tensor, conv_state: torch.Tensor, ssm_state: torch.Tensor, batch_indices: Optional[torch.Tensor] = None, + intermediate_conv_state: Optional[torch.Tensor] = None, + intermediate_ssm_state: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Single-token-per-request decode. ``zVKQba`` is ``[1, decode_req_count, - proj_dim]``; returns ``[1, decode_req_count, d_inner]``. The conv and SSM - states are read/written in place at the slots named by ``batch_indices`` - (``-1`` marks padding slots).""" - seq_len, _, _ = zVKQba.shape + """Single-token-per-request decode. ``zVKQba`` is ``[n, seq_len, + proj_dim]``; returns ``[n, seq_len, d_inner]``. The conv and SSM states + are read/written in place at the slots named by ``batch_indices`` + (``-1`` marks padding slots); ``batch_indices=None`` means static + batching, where the caches are already in request order.""" + _, seq_len, _ = zVKQba.shape assert seq_len == 1, "GDP decode supports one token per request" - zVKQba = zVKQba.squeeze(0) # [n, proj_dim] + assert ( + intermediate_conv_state is None and intermediate_ssm_state is None + ), "GDP decode does not support speculative decoding yet" + zVKQba = zVKQba.squeeze(1) # [n, proj_dim] M = self.num_householder z, VKQ, ba = torch.split( @@ -729,11 +653,14 @@ def _ssm_decode( query_new[:, :, -1] = query query = rearrange(query_new, "n t m h d -> n (t m) h d") - # Gather this step's per-request initial states. ``.clamp`` (NOT in-place) - # returns a new tensor, so ``batch_indices`` keeps its -1 padding sentinels - # for the scatter below; the padding rows' outputs are never scattered back. - gather_idx = batch_indices.clamp(min=0) - initial_state = ssm_state[gather_idx] + if batch_indices is None: + # Static batching: the cache rows are already in request order. + initial_state = ssm_state + else: + # Gather this step's per-request initial states. ``.clamp`` (NOT in-place) + # returns a new tensor, so ``batch_indices`` keeps its -1 padding sentinels + # for the scatter below; the padding rows' outputs are never scattered back. + initial_state = ssm_state[batch_indices.clamp(min=0)] core_attn_out, last_recurrent_state = fused_recurrent_gated_delta_rule( query, @@ -749,29 +676,37 @@ def _ssm_decode( ..., -1, :, : ].contiguous() # [n, 1, h, d] - # Scatter updated states back into the cache (skips -1 padding slots). - tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) + if batch_indices is None: + ssm_state.copy_(last_recurrent_state) + else: + # Scatter updated states back into the cache (skips -1 padding slots). + tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) - y = rearrange(core_attn_out, "n t h p -> t n (h p)").contiguous() # [1, n, d_inner] + y = rearrange(core_attn_out, "n t h p -> n t (h p)").contiguous() # [n, 1, d_inner] if self.rmsnorm: - z = rearrange(z, "n t h p -> t n (h p)").contiguous() + z = rearrange(z, "n t h p -> n t (h p)").contiguous() y = self.norm(y, z) return y - def _ssm_prefill( + def ssm_prefill( self, zVKQba: torch.Tensor, - conv_state: Optional[torch.Tensor] = None, - ssm_state: Optional[torch.Tensor] = None, - seq_idx: Optional[torch.Tensor] = None, - cu_seqlens: Optional[torch.Tensor] = None, - batch_indices: Optional[torch.Tensor] = None, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + context: DynamicInferenceContext, ) -> torch.Tensor: """Variable-length prefill over all prefill requests in one varlen call. ``zVKQba`` is ``[l, 1, proj_dim]``; returns ``[l, 1, d_inner]``. Fresh requests start from a zero recurrent state (no prefix caching in the MVP); the resulting final conv/SSM states are written back into the caches.""" - is_dynamic_batching = seq_idx is not None + assert ( + not context.is_chunked_prefill_enabled() + ), "GDP dynamic inference does not support chunked prefill yet." + + metadata = context.mamba_metadata + seq_idx = metadata.seq_idx + cu_seqlens = metadata.cu_seqlens + batch_indices = metadata.batch_indices_prefill M = self.num_householder # l b d -> b l d @@ -788,18 +723,14 @@ def _ssm_prefill( dim=-1, ) - if conv_state is not None and is_dynamic_batching: - assert batch_indices is not None - # Capture per-request final conv states (before the conv consumes the - # inputs) and write them into the prefill requests' cache rows. - conv_varlen_states = causal_conv1d_varlen_states( - VKQ.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] - ) - tensor_masked_update(conv_state, batch_indices, conv_varlen_states) - # Maintain channels-last memory layout so causal_conv1d_fn can use seq_idx. - VKQ = VKQ.transpose(1, 2) - else: - VKQ = rearrange(VKQ, "b l d -> b d l").contiguous() + # Capture per-request final conv states (before the conv consumes the + # inputs) and write them into the prefill requests' cache rows. + conv_varlen_states = causal_conv1d_varlen_states( + VKQ.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] + ) + tensor_masked_update(conv_state, batch_indices, conv_varlen_states) + # Maintain channels-last memory layout so causal_conv1d_fn can use seq_idx. + VKQ = VKQ.transpose(1, 2) seqlen = VKQ.size(2) if causal_conv1d_fn is None: @@ -846,15 +777,14 @@ def _ssm_prefill( g=g, beta=beta, initial_state=None, - output_final_state=ssm_state is not None, + output_final_state=True, num_householder=M, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, ) # Write per-request final SSM states into the cache for subsequent decode. - if ssm_state is not None and is_dynamic_batching: - tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) + tensor_masked_update(ssm_state, batch_indices, last_recurrent_state) y = rearrange(core_attn_out, "b l h p -> l b (h p)").contiguous() if self.rmsnorm: From 3b8a54f88b04f125d09fdf30caede6de61cabdbe Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Tue, 11 Aug 2026 16:29:55 -0700 Subject: [PATCH 22/22] Add GDP inference unit tests Signed-off-by: Keshav Santhanam --- megatron/core/ssm/gated_delta_product.py | 11 +- .../ssm/test_gdp_dynamic_inference.py | 629 ++++++++++++++++++ 2 files changed, 635 insertions(+), 5 deletions(-) create mode 100644 tests/unit_tests/ssm/test_gdp_dynamic_inference.py diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 1baaea0ea3a..3fa69813680 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -21,6 +21,9 @@ ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gdp_context_parallel import GDPContextParallel + +# Decode uses the in-repo Triton conv update, which accepts int64 slot indices. +from megatron.core.ssm.ops.causal_conv1d_triton import causal_conv1d_update from megatron.core.ssm.packed_seq_helpers import ( build_packed_seq_idx, check_fla_sequence_packing_support, @@ -38,11 +41,10 @@ from megatron.core.utils import deprecate_inference_params try: - from causal_conv1d import causal_conv1d_fn, causal_conv1d_update + from causal_conv1d import causal_conv1d_fn from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states except ImportError: causal_conv1d_fn = None - causal_conv1d_update = None causal_conv1d_varlen_states = None try: @@ -604,9 +606,8 @@ def ssm_decode( dim=-1, ) - # Indexed conv update: reads/writes the per-request conv state rows - # selected by ``batch_indices``, in place. ``self.activation`` must be the - # activation *string* so the kernel enables SiLU (a bool would disable it). + # Indexed conv update into the per-request state rows (``batch_indices`` + # is None for static batching, where the cache is already in order). VKQ = causal_conv1d_update( VKQ, conv_state, diff --git a/tests/unit_tests/ssm/test_gdp_dynamic_inference.py b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py new file mode 100644 index 00000000000..cbb4fbc3821 --- /dev/null +++ b/tests/unit_tests/ssm/test_gdp_dynamic_inference.py @@ -0,0 +1,629 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""GDP (Gated Delta Product) inference equivalence tests. + +`GatedDeltaProductMixer` supports two inference paths: static batching +(`StaticInferenceContext`) and dynamic batching (`DynamicInferenceContext`). +These tests assert that both are numerically equivalent to a plain +full-sequence forward, and therefore to each other. + +The reference is a single full-sequence `model.forward` (the +`chunk_gated_delta_product` path). Dynamic prefill runs the same chunk kernel +over a packed var-len layout, and static-batching prefill runs the same chunk +kernel with the recurrent cache seeded; both must reproduce the reference's +last-token logits. + +The single-forward equivalence tests run at TP=1 (they compare raw logits that +are sequence-sharded under sequence-parallel, and the static path does not +support SP). The end-to-end engine tests sweep TP (`_TP_SIZES`) with SP enabled +at TP>1, covering dynamic inference under tensor + sequence parallelism; TP>1 +variants skip when the world has too few GPUs. +""" + +from __future__ import annotations + +import random +import types +from typing import Dict, List, Optional, Sequence, Tuple + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig +from megatron.core.inference.contexts import StaticInferenceContext +from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext +from megatron.core.inference.engines import DynamicInferenceEngine +from megatron.core.inference.inference_request import DynamicInferenceRequest, Status +from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( + GPTInferenceWrapper, +) +from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) +from megatron.core.inference.utils import InferenceMode +from megatron.core.models.hybrid.hybrid_layer_specs import gated_delta_product_inference_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer +from megatron.core.ssm.packed_seq_helpers import check_fla_sequence_packing_support +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.cuda_graphs import delete_cuda_graphs +from megatron.core.utils import is_fa_min_version +from tests.unit_tests.test_utilities import Utils, clear_nvte_env_vars + +try: + import einops # noqa: F401 + import fla # noqa: F401 + import mamba_ssm # noqa: F401 + + HAVE_GDP_DEPS = True +except ImportError: + HAVE_GDP_DEPS = False + +# GDP dynamic inference relies on the same packed-sequence conv1d kernel as the +# training/prefill path (`causal_conv1d_fn(seq_idx=...)`, added in 1.4.0). +_PACKING_OK, _PACKING_REASON = check_fla_sequence_packing_support() + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif(not HAVE_GDP_DEPS, reason="GDP requires fla, mamba_ssm, and einops"), + pytest.mark.skipif(not _PACKING_OK, reason=_PACKING_REASON or "packed-seq support missing"), +] + + +# A short single-chunk prompt is enough to exercise the packed-varlen dynamic +# path; the sizes are kept small to keep the test fast. +_VOCAB_SIZE = 128 +_MAX_SEQ_LEN = 512 +_PROMPT_LEN = 64 + +# bf16 chunk-vs-chunk tolerance. Full-forward and prefill run the same +# `chunk_gated_delta_product` kernel, so they differ only by the packed var-len +# layout and floating-point accumulation order. +_ATOL = 5e-2 +_RTOL = 5e-2 + +# Looser tolerance for the decode-step check: it compares the recurrent decode +# kernel against a full-sequence chunk-kernel recompute (different kernels), so +# it drifts more than the chunk-vs-chunk prefill comparison. Still far tighter +# than the O(1)+ deviations a genuinely broken decode/state-handoff would show. +_DECODE_ATOL = 1e-1 +_DECODE_RTOL = 1e-1 + +# `DynamicInferenceContext` requires at least one attention layer, so the model +# pattern is GDP mixer + attention + MLP. Dynamic batching needs a recent +# flash-attention. +_LAYER_PATTERN = "M*-" +_NUM_LAYERS = len(_LAYER_PATTERN) +requires_dynamic_batching = pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need flash-attn >= 2.7.3 for dynamic batching" +) + +# Tensor-parallel sizes swept by the end-to-end engine tests. TP>1 requires +# sequence-parallel (the inference-optimized linears assert it), which GDP's +# dynamic path supports; static inference does not, so the single-forward +# equivalence tests above stay TP=1. +_TP_SIZES = [1, 2] + + +def _make_config(tp: int = 1) -> TransformerConfig: + """A small but shape-valid GDP config, sharded across `tp` tensor-parallel ranks. + + The in_proj output width (`zVKQba`) is column-parallel, so each rank sees + `proj_dim / tp` channels. The packed-prefill conv slices a channels-last view + out of that and `causal_conv1d_fn` requires its stride (the per-rank width) to + be a multiple of 8. With mamba_num_heads=16 the full width is + `(1+M)*d_inner + (M+1)*ngroups*d_state + (M+1)*nheads = 3*256 + 3*64 + 3*16 + = 1008`, so per-rank widths are 1008 (tp=1) and 504 (tp=2), both aligned. + Production configs satisfy this by having much larger, aligned dimensions. + """ + return TransformerConfig( + num_layers=_NUM_LAYERS, + hidden_size=64, + num_attention_heads=4, + num_query_groups=4, + ffn_hidden_size=128, + normalization="RMSNorm", + bf16=True, + params_dtype=torch.bfloat16, + mamba_num_heads=16, + mamba_head_dim=16, + mamba_num_groups=4, + mamba_state_dim=16, + gdp_num_householder=2, + is_hybrid_model=True, # needed for correct out_proj init + tensor_model_parallel_size=tp, + sequence_parallel=False, + context_parallel_size=1, + ) + + +def _build_model(tp: int = 1) -> HybridModel: + """Build a small GDP hybrid model (mixer + attention + MLP), eval on CUDA.""" + model_parallel_cuda_manual_seed(123) + model = HybridModel( + config=_make_config(tp), + hybrid_stack_spec=gated_delta_product_inference_stack_spec, + vocab_size=_VOCAB_SIZE, + max_sequence_length=_MAX_SEQ_LEN, + hybrid_layer_pattern=_LAYER_PATTERN, + ) + return model.cuda().eval() + + +@requires_dynamic_batching +class TestGDPDynamicInference: + """Static/dynamic GDP inference equivalence against a full-sequence forward. + + These compare raw `model.forward` logits, which are sequence-sharded under + sequence-parallel; combined with the static path not supporting SP, they run + at TP=1 only. TP>1 dynamic inference is covered end-to-end by the engine + tests below, which handle SP through the inference wrapper. + """ + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + self.model = _build_model(tp=1) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _input_ids(self) -> torch.Tensor: + """A deterministic single-request prompt: shape [1, _PROMPT_LEN].""" + return torch.arange(_PROMPT_LEN, device="cuda", dtype=torch.long).unsqueeze(0) + + @torch.inference_mode() + def _full_forward_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Reference: plain full-sequence forward -> last-token logits [1, V].""" + # No inference_context: GDP runs the training / chunk_gated_delta_product + # path. This is the ground truth both inference modes must reproduce. + InferenceMode.unset_active() + position_ids = torch.arange(input_ids.shape[1], device="cuda").unsqueeze(0) + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + runtime_gather_output=True, + ) + return logits[:, -1, :].float() + + def _build_dynamic_context(self) -> DynamicInferenceContext: + mamba_config = MambaInferenceStateConfig.from_model(self.model) + assert mamba_config is not None, "GDP hybrid model should expose Mamba inference state" + return DynamicInferenceContext( + model_config=self.model.config, + inference_config=InferenceConfig( + max_sequence_length=_MAX_SEQ_LEN, + buffer_size_gb=1.0, + block_size_tokens=256, + # Materialize all tokens so we can read the prompt's last-token + # logits directly (static batching always uses last-token only). + materialize_only_last_token_logits=False, + mamba_inference_state_config=mamba_config, + num_cuda_graphs=0, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=4, + max_tokens=128, + ), + ) + + @torch.inference_mode() + def _dynamic_prefill_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Dynamic-batching prefill -> last-token logits [1, V].""" + ctx = self._build_dynamic_context() + request = DynamicInferenceRequest( + request_id=0, + prompt_tokens=input_ids.cpu().squeeze(0), + sampling_params=SamplingParams(num_tokens_to_generate=1, termination_id=-1), + ) + ctx.add_request(request) + ctx.initialize_attention_state() + with InferenceMode.active(): + logits = self.model( + input_ids=input_ids, + position_ids=None, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + # materialize_only_last_token_logits=False -> [1, prompt_len, V]. + return logits[:, -1, :].float() + + @torch.inference_mode() + def _static_prefill_last_logits(self, input_ids: torch.Tensor) -> torch.Tensor: + """Static-batching prefill -> last-token logits [1, V].""" + ctx = StaticInferenceContext(max_batch_size=1, max_sequence_length=_MAX_SEQ_LEN) + ctx.sequence_len_offset = 0 + position_ids = torch.arange(input_ids.shape[1], device="cuda").unsqueeze(0) + with InferenceMode.active(): + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + # StaticInferenceContext forces materialize_only_last_token_logits=True, + # so the sequence dimension is already collapsed to the last token. + assert logits.shape[1] == 1 + return logits[:, 0, :].float() + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + def test_constructor(self): + """The GDP stack spec wires a GatedDeltaProductMixer into the mamba layer.""" + assert isinstance(self.model, HybridModel) + mixers = [ + layer.mixer + for layer in self.model.decoder.layers + if hasattr(layer, "mixer") and layer.mixer is not None + ] + assert len(mixers) == 1, f"pattern {_LAYER_PATTERN!r} should yield exactly one mixer layer" + assert isinstance(mixers[0], GatedDeltaProductMixer) + + def test_full_forward_shape(self): + """Sanity check: plain forward returns [batch, seq, vocab].""" + input_ids = self._input_ids() + InferenceMode.unset_active() + position_ids = torch.arange(_PROMPT_LEN, device="cuda").unsqueeze(0) + with torch.inference_mode(): + logits = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + runtime_gather_output=True, + ) + assert logits.shape == (1, _PROMPT_LEN, _VOCAB_SIZE) + + def test_dynamic_prefill_matches_full_forward(self): + """Dynamic-batching prefill reproduces the full-sequence forward.""" + input_ids = self._input_ids() + reference = self._full_forward_last_logits(input_ids) + dynamic = self._dynamic_prefill_last_logits(input_ids) + torch.testing.assert_close(dynamic, reference, atol=_ATOL, rtol=_RTOL) + + def test_static_prefill_matches_full_forward(self): + """Static-batching prefill reproduces the full-sequence forward.""" + input_ids = self._input_ids() + reference = self._full_forward_last_logits(input_ids) + static = self._static_prefill_last_logits(input_ids) + torch.testing.assert_close(static, reference, atol=_ATOL, rtol=_RTOL) + + def test_static_and_dynamic_prefill_agree(self): + """Static and dynamic inference produce equivalent logits. + + This is the central invariant: the two batching strategies must agree. + Anchoring each to the full-sequence forward (above) guarantees this + transitively, but assert it directly as well so a regression in either + path that happens to drift in the same direction is still caught. + """ + input_ids = self._input_ids() + static = self._static_prefill_last_logits(input_ids) + dynamic = self._dynamic_prefill_last_logits(input_ids) + torch.testing.assert_close(static, dynamic, atol=_ATOL, rtol=_RTOL) + + @torch.inference_mode() + def test_decode_step_matches_recompute(self): + """One decode step matches a full-sequence recompute (decode-path check). + + This validates the recurrent decode kernel and the prefill->decode + conv/SSM state handoff independently of any golden snapshot: after + prefilling the prompt, decoding one more token must produce the same + next-token logits as a plain forward over prompt+token. Compared at the + logit level with tolerance, so it is robust to the bf16 numerics that + make exact greedy token equality across the recurrent/chunk kernels + fragile. Uses `StaticInferenceContext`, whose decode calls the same + `ssm_decode` recurrent kernel as the dynamic engine. + """ + prompt = self._input_ids() # [1, P] + + # Ground truth: the next token from the prompt, and the full-recompute + # distribution for the token after it (chunk kernel over prompt+token). + next_token = int(self._full_forward_last_logits(prompt).argmax(dim=-1).item()) + extended = torch.cat( + [prompt, torch.tensor([[next_token]], dtype=torch.int64, device="cuda")], dim=1 + ) + recompute = self._full_forward_last_logits(extended) # [1, V] + + # Incremental path: prefill the prompt, then a single decode step. + ctx = StaticInferenceContext(max_batch_size=1, max_sequence_length=_MAX_SEQ_LEN) + prompt_length = prompt.shape[1] + with InferenceMode.active(): + ctx.sequence_len_offset = 0 + self.model( + input_ids=prompt, + position_ids=torch.arange(prompt_length, device="cuda").unsqueeze(0), + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + ctx.sequence_len_offset = prompt_length + decode_logits = self.model( + input_ids=torch.tensor([[next_token]], dtype=torch.int64, device="cuda"), + position_ids=torch.tensor([[prompt_length]], dtype=torch.int64, device="cuda"), + attention_mask=None, + inference_context=ctx, + runtime_gather_output=True, + ) + decode_last = decode_logits[:, -1, :].float() + + torch.testing.assert_close(decode_last, recompute, atol=_DECODE_ATOL, rtol=_DECODE_RTOL) + + +# ====================================================================== +# End-to-end engine tests. +# +# The tests above exercise a single forward pass. These drive the full +# `DynamicInferenceEngine` (add requests -> schedule -> prefill -> decode -> +# finish) so that GDP is validated through the same runtime path production +# inference uses: the text-generation controller, the inference-wrapped model, +# the KV/Mamba-state cache, and the request scheduler. +# +# Decoding is greedy (`top_k=1`) so outputs are deterministic. Decode +# correctness (the recurrent kernel `fused_recurrent_gated_delta_rule` plus the +# slot-indexed conv/SSM cache, distinct from the prefill chunk kernel) is +# validated by a byte-for-byte match against committed golden token ids, +# mirroring the Mamba2 `test_dynamic_engine.py::test_simple` style. The golden +# constant is captured from a reference GPU run (see `_GOLDEN_*` below); until it +# is populated the test self-captures and skips with the observed ids. +# ====================================================================== + +# Fixed prompts for the golden-token test. Deterministic (not random) so the +# committed golden ids below are reproducible across machines. Varying lengths +# exercise the scheduler's mixed-length prefill batching. +_GOLDEN_PROMPTS: List[List[int]] = [ + [3, 14, 15, 92, 65, 35, 89, 79], + [2, 71, 82, 81, 8], + [11, 22, 33, 44, 55, 66], + [7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7], +] +_GOLDEN_NUM_TOKENS_TO_GENERATE = 12 + +# Golden generated-token ids per TP size, one list per prompt in `_GOLDEN_PROMPTS`, +# captured from a reference GPU run. TP shards the weights differently, so each TP +# size has its own goldens. Environment-sensitive (FLA / causal_conv1d kernel +# build, GPU arch); re-capture if the kernels or config change. While an entry is +# None, the test self-captures: it prints the observed ids and skips instead of +# failing. Paste them in (see the skip message) to turn it into a hard assertion. +_GOLDEN_GENERATED_TOKENS: Dict[int, Optional[List[List[int]]]] = { + 1: [ + [32, 126, 35, 125, 52, 116, 55, 38, 39, 4, 53, 100], + [88, 10, 105, 95, 105, 44, 2, 100, 127, 59, 23, 18], + [29, 9, 61, 2, 100, 2, 69, 75, 36, 80, 103, 26], + [2, 55, 4, 108, 116, 120, 24, 113, 100, 48, 111, 22], + ], + 2: [ + [12, 51, 63, 22, 2, 40, 45, 30, 55, 10, 31, 29], + [38, 54, 29, 17, 33, 13, 10, 45, 1, 22, 21, 37], + [24, 55, 38, 55, 35, 61, 26, 25, 31, 20, 62, 56], + [53, 26, 31, 61, 16, 19, 42, 41, 49, 18, 53, 26], + ], +} + + +def _make_engine_config(tp: int = 1) -> TransformerConfig: + """GDP config for the engine tests: same shape as `_make_config`, plus the + deterministic inference sampling knobs greedy decoding needs. TP>1 turns on + sequence-parallel, which the inference-optimized linears require.""" + config = _make_config(tp) + config.sequence_parallel = tp > 1 + config.inference_rng_tracker = True + config.inference_sampling_seed = 123 + return config + + +@pytest.mark.internal +@requires_dynamic_batching +@pytest.mark.skipif(not HAVE_GDP_DEPS, reason="GDP requires fla, mamba_ssm, and einops") +@pytest.mark.skipif(not _PACKING_OK, reason=_PACKING_REASON or "packed-seq support missing") +class TestGDPDynamicInferenceEngine: + """End-to-end GDP decoding through `DynamicInferenceEngine`.""" + + SEED = 123 + VOCAB_SIZE = _VOCAB_SIZE + + def teardown_method(self, method): + delete_cuda_graphs() + Utils.destroy_model_parallel() + + # ------------------------------------------------------------------ + # Harness + # ------------------------------------------------------------------ + + def _build_engine( + self, + *, + num_tokens_to_generate: int, + tp: int = 1, + num_requests: Optional[int] = None, + prompt_length: Optional[int] = None, + prompts: Optional[Sequence[Sequence[int]]] = None, + ) -> Tuple[DynamicInferenceEngine, List[DynamicInferenceRequest]]: + """Build a greedy GDP engine plus its requests at TP=`tp`. + + Either pass explicit `prompts` (deterministic token lists) or + `num_requests` + `prompt_length` (random prompts of a fixed length). + Skips if the world is too small for the requested TP size. + """ + if Utils.world_size < tp: + pytest.skip(f"TP={tp} requires at least {tp} GPUs") + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp, pipeline_model_parallel_size=1 + ) + clear_nvte_env_vars() + random.seed(self.SEED) + torch.manual_seed(self.SEED) + model_parallel_cuda_manual_seed( + seed=self.SEED, inference_rng_tracker=True, force_reset_rng=True + ) + + if prompts is not None: + prompt_tensors = [torch.tensor(p, dtype=torch.int64, device="cuda") for p in prompts] + else: + assert num_requests is not None and prompt_length is not None + prompt_tensors = [ + torch.randint( + 0, self.VOCAB_SIZE - 1, (prompt_length,), dtype=torch.int64, device="cuda" + ) + for _ in range(num_requests) + ] + + max_prompt_length = max(int(p.numel()) for p in prompt_tensors) + max_sequence_length = max_prompt_length + num_tokens_to_generate + config = _make_engine_config(tp) + model = HybridModel( + config=config, + hybrid_stack_spec=gated_delta_product_inference_stack_spec, + vocab_size=self.VOCAB_SIZE, + max_sequence_length=max_sequence_length, + parallel_output=True, + hybrid_layer_pattern=_LAYER_PATTERN, + pre_process=parallel_state.is_pipeline_first_stage(), + post_process=parallel_state.is_pipeline_last_stage(), + ).cuda() + for param in model.parameters(): + param.data = param.data.to(config.params_dtype) + model.eval() + + context = DynamicInferenceContext( + model_config=config, + inference_config=InferenceConfig( + max_sequence_length=max_sequence_length, + buffer_size_gb=0.1, + block_size_tokens=256, + materialize_only_last_token_logits=True, + mamba_inference_state_config=MambaInferenceStateConfig.from_model(model), + num_cuda_graphs=None, + use_cuda_graphs_for_non_decode_steps=False, + max_requests=32, + max_tokens=1024, + ), + ) + + wrapped_model = GPTInferenceWrapper(model, context) + wrapped_model.model_is_pipeline_parallel = not ( + parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage() + ) + controller = TextGenerationController( + inference_wrapped_model=wrapped_model, + tokenizer=types.SimpleNamespace( + vocab_size=self.VOCAB_SIZE, detokenize=lambda tokens: "tokenized_prompt" + ), + ) + delete_cuda_graphs() + engine = DynamicInferenceEngine(controller, context) + + requests = [ + DynamicInferenceRequest( + request_id=request_id, + prompt_tokens=prompt_tokens, + sampling_params=SamplingParams( + num_tokens_to_generate=num_tokens_to_generate, + termination_id=-1, # never terminate early -> fixed output length + top_k=1, # greedy -> deterministic + ), + ) + for request_id, prompt_tokens in enumerate(prompt_tensors) + ] + return engine, requests + + @staticmethod + @torch.inference_mode() + def _run_to_completion( + engine: DynamicInferenceEngine, requests: List[DynamicInferenceRequest] + ) -> Dict[int, DynamicInferenceRequest]: + """Add every request, step until the engine drains, return finished requests by id.""" + for request in requests: + engine._add_request(request) + + finished: Dict[int, DynamicInferenceRequest] = {} + # Bound the loop so a scheduling regression fails loudly instead of hanging. + for _ in range(1000): + result = engine.step_modern() + for record in result["finished_request_records"]: + merged = record.merge() + finished[merged.request_id] = merged + if not engine.has_unfinished_requests(): + break + assert not engine.has_unfinished_requests(), "engine did not drain within step budget" + return finished + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_engine_runs_to_completion(self, tp): + """Every request completes and yields exactly the requested token count.""" + num_tokens_to_generate = 8 + engine, requests = self._build_engine( + tp=tp, num_requests=4, prompt_length=8, num_tokens_to_generate=num_tokens_to_generate + ) + finished = self._run_to_completion(engine, requests) + + assert len(finished) == len(requests) + for request in requests: + merged = finished[request.request_id] + assert merged.status == Status.COMPLETED + # termination_id=-1 disables early stop, so the length is exact. + assert len(merged.generated_tokens) == num_tokens_to_generate + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_engine_greedy_matches_golden(self, tp): + """Greedy decode reproduces committed golden token ids (Mamba2-style). + + Deterministic fixed prompts + greedy sampling make the output a stable + fingerprint of the GDP prefill+decode path. Until the TP entry in + `_GOLDEN_GENERATED_TOKENS` is captured from a reference GPU run, the test + prints the observed ids and skips instead of failing. + """ + engine, requests = self._build_engine( + tp=tp, prompts=_GOLDEN_PROMPTS, num_tokens_to_generate=_GOLDEN_NUM_TOKENS_TO_GENERATE + ) + finished = self._run_to_completion(engine, requests) + observed = [finished[r.request_id].generated_tokens for r in requests] + + golden = _GOLDEN_GENERATED_TOKENS.get(tp) + if golden is None: + pytest.skip( + f"golden tokens for TP={tp} not captured yet; paste the following into " + f"_GOLDEN_GENERATED_TOKENS[{tp}]:\n{observed!r}" + ) + + assert observed == golden, ( + f"generated tokens != golden (TP={tp}):\n golden = {golden}\n" + f" observed = {observed}" + ) + + @pytest.mark.parametrize("tp", _TP_SIZES) + def test_generate_over_multiple_prompts(self, tp): + """`engine.generate` drives several prompts through to completion at once.""" + engine, requests = self._build_engine( + tp=tp, num_requests=4, prompt_length=8, num_tokens_to_generate=4 + ) + + prompts = [f"prompt{i}" for i in range(len(requests))] + + def mock_tokenize_prompt(tokenizer, prompt, add_BOS=False): + prompt_num = int(prompt[-1]) + return [10 + i for i in range(prompt_num + 2)] + + engine.controller.tokenize_prompt = mock_tokenize_prompt + + finished_records = engine.generate(prompts, requests[0].sampling_params) + finished = [record.merge() for record in finished_records] + + assert len(finished) == len(prompts) + # generate() returns finished requests in request-id order. + assert [r.request_id for r in finished] == sorted(r.request_id for r in finished) + for request in finished: + assert request.status == Status.COMPLETED + assert len(request.generated_tokens) > 0