From b6ff5f4d764b0056451d27b69659057c64588d27 Mon Sep 17 00:00:00 2001 From: Anil Thomas Date: Fri, 14 Aug 2026 07:38:38 -0700 Subject: [PATCH] Add GDN dynamic inference Integrate GDN with the shared recurrent-state cache and FLA kernels for packed prefill and decode. Add configuration validation plus state-continuity coverage. Co-authored-by: OpenAI Codex Signed-off-by: Anil Thomas --- megatron/core/inference/config.py | 36 +++- .../inference/contexts/dynamic_context.py | 27 +-- megatron/core/models/hybrid/hybrid_block.py | 7 +- .../core/models/hybrid/hybrid_layer_specs.py | 14 ++ megatron/core/ssm/gated_delta_net/common.py | 3 + megatron/core/ssm/gated_delta_net/gdn.py | 170 ++++++++++++++++- pyproject.toml | 2 +- .../inference/engines/test_dynamic_engine.py | 133 ++++++++++++- .../inference/test_inference_config.py | 33 +++- tests/unit_tests/ssm/conftest.py | 14 ++ tests/unit_tests/ssm/test_gated_delta_net.py | 8 + .../ssm/test_gated_delta_net_inference.py | 180 ++++++++++++++++++ tests/unit_tests/ssm/test_hybrid_block.py | 11 +- uv.lock | 4 +- 14 files changed, 601 insertions(+), 41 deletions(-) create mode 100644 tests/unit_tests/ssm/conftest.py create mode 100644 tests/unit_tests/ssm/test_gated_delta_net_inference.py diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index c0e3ed4a1d1..01ecdc1c896 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -15,24 +15,24 @@ @dataclass class MambaInferenceStateConfig: """ - Config for initializing Mamba model inference state tensors. + Config for initializing recurrent mixer inference state tensors. Note that we maintain separate metadata for decode, regular prefill, and - chunked prefill requests because the Mamba kernels do not yet support mixing - these. Once the kernels have been updated we can simplify this code. + chunked prefill requests because the recurrent kernels do not yet support + mixing these. Once the kernels have been updated we can simplify this code. """ layer_type_list: List[str] """ - A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer. + A list of strings that indicates the layer type (Mamba / GDN / Attention / MLP) for each layer. See `megatron/core/models/hybrid/hybrid_layer_allocation.py` for the list of symbols. """ conv_states_shape: Tuple[int] - """Mamba conv states shape per request.""" + """Recurrent mixer's conv state shape per request.""" ssm_states_shape: Tuple[int] - """Mamba SSM states shape per request.""" + """Recurrent mixer state shape per request.""" conv_states_dtype: torch.dtype """The dtype to use for the Mamba conv state tensor. Defaults to the model dtype.""" @@ -55,12 +55,29 @@ def from_model( conv_states_dtype: Optional[torch.dtype] = None, ssm_states_dtype: Optional[torch.dtype] = None, ) -> Optional["MambaInferenceStateConfig"]: - """Returns Mamba inference state config from the model if it is a hybrid model.""" + """Return recurrent inference state config for a Mamba or GDN hybrid model.""" from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols decoder = get_attr_wrapped_model(model, "decoder") layer_type_list = getattr(decoder, "layer_type_list", None) - if layer_type_list is not None and Symbols.MAMBA in layer_type_list: + recurrent_symbols = (Symbols.MAMBA, Symbols.GDN) + if layer_type_list is not None and any( + symbol in layer_type_list for symbol in recurrent_symbols + ): + present_recurrent_symbols = { + symbol for symbol in recurrent_symbols if symbol in layer_type_list + } + if len(present_recurrent_symbols) > 1: + raise ValueError( + "Dynamic inference does not support mixing Mamba and GDN layers; " + "the recurrent-state cache and prefill metadata use one shared shape " + "and chunk size." + ) + if ( + Symbols.GDN in present_recurrent_symbols + and model.config.experimental_attention_variant == "gdn2" + ): + raise NotImplementedError("GDN2 does not support dynamic inference.") mamba_conv_states_shape, mamba_ssm_states_shape = ( decoder.mamba_state_shapes_per_request() ) @@ -82,6 +99,9 @@ def from_model( if layer_type == Symbols.MAMBA and hasattr(layer, 'mixer'): mamba_chunk_size = layer.mixer.chunk_size break + if layer_type == Symbols.GDN and hasattr(layer, 'self_attention'): + mamba_chunk_size = layer.self_attention.chunk_size + break # Gated Delta Product layers register as Mamba layers but carry a # Householder count, which sizes their (separate) chunk descriptors. gdp_num_householder = 0 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4de2c64f06f..4c74c803a3d 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -441,21 +441,22 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC "boundaries are not rounded between decode chunks." ) - # For hybrid models, the layer map converts the global layer index to the - # corresponding attention layer index or Mamba layer index depending on the - # layer type. - attention_layer_map, dsa_layer_map, gdn_layer_map, mamba_layer_map = ( - operator.itemgetter( - Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.MAMBA - )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) - ) - - if len(gdn_layer_map) > 0: - raise NotImplementedError("GDN layers are not supported for inference.") + # Mamba and GDN use the same slot-indexed recurrent-state cache contract. Build + # one map in global layer order; independently generated per-symbol maps both + # start at zero and would alias if they were simply unioned. + attention_layer_map, dsa_layer_map = operator.itemgetter( + Symbols.ATTENTION, Symbols.DS_ATTENTION + )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) + recurrent_layer_map = {} + for global_layer_idx, layer_type in enumerate( + mamba_inference_state_config.layer_type_list + ): + if layer_type in (Symbols.MAMBA, Symbols.GDN): + recurrent_layer_map[global_layer_idx] = len(recurrent_layer_map) self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map) - self.num_mamba_layers = len(mamba_layer_map) - self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map + self.num_mamba_layers = len(recurrent_layer_map) + self.layer_map = attention_layer_map | dsa_layer_map | recurrent_layer_map else: # The layer map is the identity function for pure Transformer models. # Use the same per-PP-rank layer count as TransformerBlock (handles diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 49cc938eb02..8729a887d57 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -206,6 +206,7 @@ def __init__( pg_collection=pg_collection, # Set to False as we do not want to change offset. add_layer_offset=False, + pp_layer_offset=pp_layer_offset, name=(name + f".layers.{i}") if name is not None else None, ) else: @@ -258,12 +259,14 @@ def set_input_tensor(self, input_tensor: Tensor): def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: """ - Returns the Mamba conv and ssm states shapes per input sequence - if this block contains Mamba layers (this may not be the case with PP > 1). + Returns the recurrent mixer's conv and SSM state shapes per input sequence + if this block contains Mamba or GDN layers (this may not be the case with PP > 1). """ for layer_type, layer in zip(self.layer_type_list, self.layers): if layer_type == LayerSymbols.MAMBA: return layer.mamba_state_shapes_per_request() + if layer_type == LayerSymbols.GDN: + return layer.self_attention.mamba_state_shapes_per_request() return None def forward( diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 8e91f442e13..e8b1ebc9a56 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -266,6 +266,20 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj): mamba_bda=get_bias_dropout_add, ), ), + gdn_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=GatedDeltaNet, + submodules=GatedDeltaNetSubmodules( + in_proj=InferenceLayerNormColumnParallelLinear, + out_norm=TENorm, + out_proj=InferenceRowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), # Started with spec from gpt_layer_specs.py (with MLP removed) # Using the TE spec because we had problems getting the non-TE spec # working diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 6802cfa3ded..63f2d244840 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -124,6 +124,7 @@ def __init__( *, name: str | None = None, cp_comm_type: str | None = None, + pp_layer_offset: int = 0, ): """ Args: @@ -141,6 +142,7 @@ def __init__( cp_comm_type (Optional[str]): Accepted for TransformerLayer compatibility and ignored; GDN implements context parallelism with its own all-to-alls rather than the attention CP communication schemes. + pp_layer_offset: Offset of this pipeline stage's first global layer. """ if not HAVE_FLA: raise ImportError( @@ -152,6 +154,7 @@ def __init__( # Attributes from arguments self.layer_number = layer_number + self.pp_layer_offset = pp_layer_offset self.bias = bias self.conv_bias = conv_bias self.conv_init = conv_init diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index a28478218de..c7a983841cb 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -12,7 +12,10 @@ import torch.nn.functional as F from megatron.core import tensor_parallel -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_masked_update, +) from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.ssm.gated_delta_net.common import ( @@ -23,10 +26,18 @@ get_parameter_local_cp, l2norm, ) +from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push +try: + from fla.modules.convolution import causal_conv1d_update + from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule +except ImportError: + causal_conv1d_update = None + fused_recurrent_gated_delta_rule = None + -class GatedDeltaNet(_GDNBase): +class GatedDeltaNet(SSMDynamicInferenceMixin, _GDNBase): # pylint: disable=missing-class-docstring def _setup_variant_attrs(self): """Set the GDN in_proj sizing, split tables, gate parameter dims, and kernel.""" @@ -59,6 +70,7 @@ def _setup_variant_attrs(self): self.gated_delta_rule = torch_chunk_gated_delta_rule else: self.gated_delta_rule = chunk_gated_delta_rule + self.chunk_size = 64 @jit_fuser def _compute_gates( @@ -100,12 +112,26 @@ def forward( seq_len = seq_len * self.sp_size * self.cp_size if inference_context is not None: - assert ( - inference_context.is_static_batching() - ), "GDN does not currently support dynamic inference batching." + if inference_context.is_dynamic_batching(): + assert ( + not self.config.deterministic_mode + ), "GDN dynamic inference requires the FLA recurrent kernels." + assert ( + not self.config.batch_invariant_mode + ), "GDN dynamic inference does not support batch-invariant mode." + assert ( + self.cp_size == 1 + ), "Context parallelism is not supported for GDN dynamic inference." + assert ( + inference_context.num_speculative_tokens == 0 + ), "GDN dynamic inference does not support speculative decoding." + assert ( + not inference_context.enable_prefix_caching + ), "GDN dynamic inference does not support prefix caching." + return self.ssm_dynamic_inference(hidden_states, inference_context) + assert inference_context.is_static_batching() assert not self.config.sequence_parallel - # TODO: support inference - raise NotImplementedError("GDN does not support inference for now.") + raise NotImplementedError("GDN static-batching inference is not supported.") if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': assert batch == 1, "Packed sequence expects batch dimension to be 1" @@ -162,8 +188,7 @@ def forward( # Split the tensor into q, k, v, gate (z), and the variant-specific gate features # (beta, alpha for GDN; f, b, w for GDN2) - qkv, gate, beta, alpha = torch.split(qkvzba, self.feat_dim_split, dim=-1) - gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + qkv, gate, beta, alpha = self._split_projection(qkvzba, batch, seq_len) # Convolution on qkv nvtx_range_push(suffix="conv1d") @@ -263,6 +288,133 @@ def forward( return out, out_bias + def _split_projection( + self, projected: torch.Tensor, batch: int, seq_len: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Split the fused projection into qkv, output gate, beta, and alpha.""" + qkv, gate, beta, alpha = torch.split(projected, self.feat_dim_split, dim=-1) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + return qkv, gate, beta, alpha + + def _prepare_inference_inputs( + self, qkv: torch.Tensor, beta: torch.Tensor, alpha: torch.Tensor, batch: int, seq_len: int + ) -> dict[str, torch.Tensor]: + """Prepare raw FLA inputs while leaving normalization and gates fused in-kernel.""" + query_key, value = torch.split(qkv, [2 * self.qk_dim_local_tp, self.v_dim_local_tp], dim=-1) + query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) + query, key = torch.chunk(query_key, 2, dim=2) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + return { + "q": query.contiguous(), + "k": key.contiguous(), + "v": value.contiguous(), + "g": alpha.contiguous(), + "beta": beta.contiguous(), + } + + def mamba_state_shapes_per_request(self) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Return the TP-local convolution and delta-rule cache shapes.""" + return ( + (self.conv_dim_local_tp, self.conv_kernel_dim), + (self.num_v_heads_local_tp, self.key_head_dim, self.value_head_dim), + ) + + def ssm_decode( + self, + projected: torch.Tensor, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + batch_indices: torch.Tensor, + intermediate_conv_state: torch.Tensor | None = None, + intermediate_ssm_state: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run one CUDA-graph-compatible GDN decode token per request.""" + batch, seq_len, _ = projected.shape + assert seq_len == 1, "GDN speculative decoding is not supported." + assert ( + intermediate_conv_state is None and intermediate_ssm_state is None + ), "GDN speculative decoding state capture is not supported." + assert causal_conv1d_update is not None and fused_recurrent_gated_delta_rule is not None + + qkv, gate, beta, alpha = self._split_projection(projected, batch, seq_len) + read_indices = batch_indices.clamp(min=0) + + active_conv_state = conv_state[read_indices].contiguous() + qkv_dtype = qkv.dtype + qkv, active_conv_state = causal_conv1d_update( + x=qkv.to(conv_state.dtype), + cache=active_conv_state, + weight=self.conv1d.weight.squeeze(1).to(conv_state.dtype), + bias=self.conv1d.bias.to(conv_state.dtype) if self.conv1d.bias is not None else None, + activation=self.activation, + ) + qkv = qkv.to(qkv_dtype) + tensor_masked_update(conv_state, batch_indices, active_conv_state) + + kernel_inputs = self._prepare_inference_inputs(qkv, beta, alpha, batch, seq_len) + active_ssm_state = ssm_state[read_indices].contiguous() + core_attn_out, final_ssm_state = fused_recurrent_gated_delta_rule( + **kernel_inputs, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=active_ssm_state, + output_final_state=True, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + ) + tensor_masked_update(ssm_state, batch_indices, final_ssm_state) + return self._apply_gated_norm(core_attn_out, gate).reshape(batch, seq_len, -1) + + def ssm_prefill( + self, + projected: torch.Tensor, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + context: DynamicInferenceContext, + ) -> torch.Tensor: + """Run packed variable-length GDN prefill and populate request states.""" + assert ( + not context.is_chunked_prefill_enabled() + ), "GDN dynamic inference does not support chunked prefill." + metadata = context.mamba_metadata + cu_seqlens = metadata.cu_seqlens + batch_indices = metadata.batch_indices_prefill + token_count = projected.shape[0] + + projected = projected.transpose(0, 1).contiguous() + qkv, gate, beta, alpha = self._split_projection(projected, 1, token_count) + read_indices = batch_indices.clamp(min=0) + + qkv_dtype = qkv.dtype + qkv, final_conv_state = causal_conv1d( + x=qkv.to(conv_state.dtype), + weight=self.conv1d.weight.squeeze(1).to(conv_state.dtype), + bias=self.conv1d.bias.to(conv_state.dtype) if self.conv1d.bias is not None else None, + activation=self.activation, + initial_state=conv_state[read_indices].contiguous(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + qkv = qkv.to(qkv_dtype) + tensor_masked_update(conv_state, batch_indices, final_conv_state) + + kernel_inputs = self._prepare_inference_inputs(qkv, beta, alpha, 1, token_count) + core_attn_out, final_ssm_state = chunk_gated_delta_rule( + **kernel_inputs, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=ssm_state[read_indices].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + tensor_masked_update(ssm_state, batch_indices, final_ssm_state) + y = self._apply_gated_norm(core_attn_out, gate) + return y.reshape(1, token_count, -1).transpose(0, 1).contiguous() + #################### # Torch native gated delta rule diff --git a/pyproject.toml b/pyproject.toml index 4b4ad8c1c86..ec01cd75c3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,7 +98,6 @@ dev = [ "multi-storage-client~=0.50", "opentelemetry-api>=1.33.1,<2", "nemo-lens", - "flash-linear-attention==0.5.1", "megatron-energon[av_decode]~=7.0", "av", "flashinfer-python>=0.5.0,<0.7.0", @@ -134,6 +133,7 @@ te = [ ssm = [ "mamba-ssm", "causal-conv1d~=1.5", + "flash-linear-attention==0.5.1", ] [dependency-groups] diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 093e4062cc8..6d5d5a56ad6 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -16,6 +16,7 @@ import msgpack import pytest import torch +import torch.nn.functional as F from tqdm import tqdm from transformer_engine.pytorch.fp8 import check_fp8_support @@ -58,6 +59,7 @@ from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.ssm.gated_delta_net import HAVE_FLA from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import delete_cuda_graphs @@ -433,7 +435,8 @@ def _build_test_env(cls, test_config): mtp_block_spec=mtp_block_spec, position_embedding_type=test_config.position_embedding_type, ).cuda() - elif test_config.model_provider == "hybrid": + elif test_config.model_provider in ("hybrid", "gdn"): + is_gdn = test_config.model_provider == "gdn" pp_size = test_config.pipeline_model_parallel_size # Transformer config. transformer_config = TransformerConfig( @@ -445,6 +448,11 @@ def _build_test_env(cls, test_config): hidden_size=256, # The Mamba layer places several constraints on this mamba_num_heads=16, num_attention_heads=16, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=8, use_cpu_initialization=True, cuda_graph_impl=effective_cuda_graph_impl, inference_rng_tracker=True, @@ -476,9 +484,11 @@ def _build_test_env(cls, test_config): ), normalization=( "RMSNorm" - if test_config.transformer_impl == "inference_optimized" + if is_gdn or test_config.transformer_impl == "inference_optimized" else "LayerNorm" ), + layernorm_zero_centered_gamma=is_gdn, + activation_func=F.silu if is_gdn else F.gelu, is_hybrid_model=True, # Needs to be set for correct out_proj init ) @@ -486,10 +496,11 @@ def _build_test_env(cls, test_config): # When speculative tokens are configured, append MTP depth sections # to the hybrid layer pattern so the model creates MTP blocks. mtp_suffix = "/M" * test_config.num_speculative_tokens + recurrent_symbol = "G" if is_gdn else "M" if pp_size == 1: - mamba_pattern = "M*-" + mtp_suffix + mamba_pattern = recurrent_symbol + "*-" + mtp_suffix else: - mamba_pattern = "M*-|M*-" + mtp_suffix + mamba_pattern = recurrent_symbol + "*-|" + recurrent_symbol + "*-" + mtp_suffix model = HybridModel( config=transformer_config, hybrid_stack_spec=hybrid_stack_spec, @@ -5739,6 +5750,75 @@ def mtp_with_rejection( assert env.engine.context.total_request_count == 0 +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" +) +class TestGDNDynamicInferenceEngine(DynamicInferenceEngineTestBase): + """Exercise GDN through the production scheduler and local CUDA graphs.""" + + @classmethod + def setup_class(cls): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=1, + ) + + @classmethod + def teardown_class(cls): + delete_cuda_graphs() + set_rounder(64) + Utils.destroy_model_parallel() + + @staticmethod + def _generated_tokens(env): + return [list(request.generated_tokens) for request in env.requests] + + def test_cuda_graph_parity(self): + common = dict( + model_provider="gdn", + num_requests=4, + min_prompt_length=8, + max_prompt_length=8, + num_tokens_to_generate=8, + num_gap_steps=0, + top_k=1, + context_max_requests=32, + use_cuda_graphs_for_non_decode_steps=False, + ) + eager = self._run_test(**common, num_cuda_graphs=None) + graphed = self._run_test( + **common, + num_cuda_graphs=3, + force_build_cuda_graphs=True, + inference_cuda_graph_scope=InferenceCudaGraphScope.block, + ) + + model = graphed.engine.controller.inference_wrapped_model.model + assert model.cudagraph_manager.cudagraph_runners + assert self._generated_tokens(graphed) == self._generated_tokens(eager) + + def test_scheduling_invariance(self): + """Staggered admission must not change per-request greedy output.""" + common = dict( + model_provider="gdn", + num_requests=4, + min_prompt_length=8, + max_prompt_length=8, + num_tokens_to_generate=8, + top_k=1, + context_max_requests=32, + num_cuda_graphs=None, + ) + dense = self._run_test(**common, num_gap_steps=0) + staggered = self._run_test(**common, num_gap_steps=3) + + assert self._generated_tokens(staggered) == self._generated_tokens(dense) + + class TestDynamicInferenceEngineParallel(DynamicInferenceEngineTestBase): """Tests that require non-default parallel configs (tp>1, pp>1, or ep>1). @@ -5762,6 +5842,51 @@ def _build_test_env(cls, test_config): ) return super()._build_test_env(test_config) + @pytest.mark.internal + @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") + @pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" + ) + @torch.inference_mode() + def test_gdn_tensor_parallel(self): + if int(os.environ.get("WORLD_SIZE", "1")) < 2: + pytest.skip("GDN TP=2 inference requires at least two GPUs.") + env = self._run_test( + model_provider="gdn", + tensor_model_parallel_size=2, + num_requests=4, + min_prompt_length=8, + max_prompt_length=8, + num_tokens_to_generate=4, + num_gap_steps=0, + top_k=1, + context_max_requests=16, + ) + assert all(request.status == Status.COMPLETED for request in env.requests) + + @pytest.mark.internal + @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") + @pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" + ) + @torch.inference_mode() + def test_gdn_pipeline_parallel(self): + if int(os.environ.get("WORLD_SIZE", "1")) < 2: + pytest.skip("GDN PP=2 inference requires at least two GPUs.") + env = self._run_test( + model_provider="gdn", + pipeline_model_parallel_size=2, + num_requests=4, + min_prompt_length=8, + max_prompt_length=8, + num_tokens_to_generate=4, + num_gap_steps=0, + top_k=1, + context_max_requests=16, + ) + assert all(request.status == Status.COMPLETED for request in env.requests) + assert all(len(request.generated_tokens) == 4 for request in env.requests) + @pytest.mark.internal @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" diff --git a/tests/unit_tests/inference/test_inference_config.py b/tests/unit_tests/inference/test_inference_config.py index adb05095c59..aa7f32f2250 100644 --- a/tests/unit_tests/inference/test_inference_config.py +++ b/tests/unit_tests/inference/test_inference_config.py @@ -5,10 +5,16 @@ from types import SimpleNamespace import pytest +import torch -from megatron.core.inference.config import AsyncScheduleMode, InferenceConfig +from megatron.core.inference.config import ( + AsyncScheduleMode, + InferenceConfig, + MambaInferenceStateConfig, +) from megatron.core.inference.moe import InferenceGroupedGemmBackend from megatron.core.inference.quantization.utils import resolve_mxfp8_backend +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.arguments import _add_inference_args from megatron.training.config.inference_config import InferenceSetupConfig @@ -32,6 +38,31 @@ def test_resolve_mxfp8_backend_rejects_unsupported_backend(self, grouped_gemm_ba with pytest.raises(ValueError, match="does not support inference_grouped_gemm_backend"): resolve_mxfp8_backend(grouped_gemm_backend) + @staticmethod + def _hybrid_model(layer_type_list, experimental_attention_variant="gdn"): + return SimpleNamespace( + config=SimpleNamespace( + params_dtype=torch.bfloat16, + batch_invariant_mode=False, + experimental_attention_variant=experimental_attention_variant, + ), + decoder=SimpleNamespace(layer_type_list=layer_type_list, layers=[]), + ) + + def test_mamba_inference_state_config_rejects_mixed_recurrent_layers(self): + """Mamba and GDN cannot share one state shape and prefill chunk size.""" + model = self._hybrid_model([Symbols.MAMBA, Symbols.GDN]) + + with pytest.raises(ValueError, match="mixing Mamba and GDN"): + MambaInferenceStateConfig.from_model(model) + + def test_mamba_inference_state_config_rejects_gdn2(self): + """GDN2 should fail explicitly instead of missing the GDN inference hooks.""" + model = self._hybrid_model([Symbols.GDN], experimental_attention_variant="gdn2") + + with pytest.raises(NotImplementedError, match="GDN2"): + MambaInferenceStateConfig.from_model(model) + def test_mutual_exclusivity_with_transformer_config(self): """ Ensure mutual exclusivity between fields in `InferenceConfig` and diff --git a/tests/unit_tests/ssm/conftest.py b/tests/unit_tests/ssm/conftest.py new file mode 100644 index 00000000000..678c197e334 --- /dev/null +++ b/tests/unit_tests/ssm/conftest.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import os + +import pytest +import torch + + +@pytest.fixture(scope="session", autouse=True) +def select_local_cuda_device(): + """Bind each torchrun process before an SSM test initializes CUDA or Triton.""" + if torch.cuda.is_available(): + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank % torch.cuda.device_count()) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 1e762ba5bbb..b63420b2fa9 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -317,6 +317,14 @@ def test_module_construction(self): assert gdn.A_log.shape == (gdn.num_value_heads // self.tp_size,) assert gdn.dt_bias.shape == (gdn.num_value_heads // self.tp_size,) + def test_inference_state_shapes(self): + if self.use_gdn2: + pytest.skip("GDN2 inference is not supported.") + assert self.gdn.mamba_state_shapes_per_request() == ( + (self.gdn.conv_dim_local_tp, self.gdn.conv_kernel_dim), + (self.gdn.num_v_heads_local_tp, self.gdn.key_head_dim, self.gdn.value_head_dim), + ) + def test_jit_compiled_helpers(self): import torch._dynamo diff --git a/tests/unit_tests/ssm/test_gated_delta_net_inference.py b/tests/unit_tests/ssm/test_gated_delta_net_inference.py new file mode 100644 index 00000000000..b5e5725cde1 --- /dev/null +++ b/tests/unit_tests/ssm/test_gated_delta_net_inference.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Focused kernel and CUDA-graph tests for Gated DeltaNet dynamic inference.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net import HAVE_FLA, GatedDeltaNet +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 + + +@pytest.mark.internal +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +class TestGatedDeltaNetInference: + """Validate state continuity and capture of the real FLA inference kernels.""" + + @pytest.fixture(scope="function", autouse=True) + def setup_method(self): + # Other tests in this multi-rank bucket may change the process's current + # device. Raw CUDA-graph capture and Triton launches must use LOCAL_RANK's + # device, matching Utils.initialize_distributed's initial assignment. + torch.cuda.set_device(Utils.rank % torch.cuda.device_count()) + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + model_parallel_cuda_manual_seed(123) + config = TransformerConfig( + hidden_size=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=8, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=8, + num_query_groups=2, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=1, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + transformer_impl="transformer_engine", + ) + spec = get_experimental_attention_variant_module_spec(config=config) + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + self.gdn = ( + spec.module( + config, + submodules=spec.submodules, + layer_number=1, + bias=False, + conv_bias=False, + use_qk_l2norm=True, + pg_collection=pg_collection, + ) + .cuda() + .bfloat16() + .eval() + ) + yield + Utils.destroy_model_parallel() + + @staticmethod + def _prefill_context(cu_seqlens, batch_indices): + return SimpleNamespace( + is_chunked_prefill_enabled=lambda: False, + mamba_metadata=SimpleNamespace( + cu_seqlens=cu_seqlens, batch_indices_prefill=batch_indices + ), + ) + + def _empty_states(self, slots=4): + conv_shape, ssm_shape = self.gdn.mamba_state_shapes_per_request() + conv_state = torch.zeros( + slots, *conv_shape, device="cuda", dtype=self.gdn.conv1d.weight.dtype + ) + ssm_state = torch.zeros( + slots, *ssm_shape, device="cuda", dtype=self.gdn.in_proj.weight.dtype + ) + return conv_state, ssm_state + + @torch.inference_mode() + def test_prefill_decode_matches_full_forward(self): + prompt_len, total_len = 9, 14 + hidden = torch.randn( + total_len, 1, self.gdn.hidden_size, device="cuda", dtype=torch.bfloat16 + ) + expected, _ = self.gdn(hidden, attention_mask=None) + projected, _ = self.gdn.in_proj(hidden) + conv_state, ssm_state = self._empty_states() + slot = torch.tensor([2], device="cuda", dtype=torch.int32) + cu_seqlens = torch.tensor([0, prompt_len], device="cuda", dtype=torch.long) + + outputs = [ + self.gdn.ssm_prefill( + projected[:prompt_len], + conv_state, + ssm_state, + self._prefill_context(cu_seqlens, slot), + ) + ] + for token in projected[prompt_len:]: + outputs.append( + self.gdn.ssm_decode( + token.view(1, 1, -1), conv_state, ssm_state, batch_indices=slot + ).transpose(0, 1) + ) + actual, _ = self.gdn.out_proj(torch.cat(outputs, dim=0)) + torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) + + @torch.inference_mode() + def test_padding_index_does_not_modify_state(self): + conv_state, ssm_state = self._empty_states() + projected = torch.randn( + 2, 1, self.gdn.in_proj_dim // self.gdn.tp_size, device="cuda", dtype=torch.bfloat16 + ) + indices = torch.tensor([1, -1], device="cuda", dtype=torch.int32) + conv_before = conv_state.clone() + ssm_before = ssm_state.clone() + + self.gdn.ssm_decode(projected, conv_state, ssm_state, batch_indices=indices) + + assert not torch.equal(conv_state[1], conv_before[1]) + assert not torch.equal(ssm_state[1], ssm_before[1]) + torch.testing.assert_close(conv_state[[0, 2, 3]], conv_before[[0, 2, 3]], atol=0, rtol=0) + torch.testing.assert_close(ssm_state[[0, 2, 3]], ssm_before[[0, 2, 3]], atol=0, rtol=0) + + @torch.inference_mode() + def test_decode_cuda_graph_capture_and_replay(self): + batch = 4 + projected = torch.randn( + batch, 1, self.gdn.in_proj_dim // self.gdn.tp_size, device="cuda", dtype=torch.bfloat16 + ) + indices = torch.arange(batch, device="cuda", dtype=torch.int32) + + eager_conv, eager_ssm = self._empty_states(slots=batch) + expected = self.gdn.ssm_decode(projected, eager_conv, eager_ssm, batch_indices=indices) + + graph_conv, graph_ssm = self._empty_states(slots=batch) + static_projected = projected.clone() + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + graph_conv.zero_() + graph_ssm.zero_() + self.gdn.ssm_decode(static_projected, graph_conv, graph_ssm, batch_indices=indices) + torch.cuda.current_stream().wait_stream(warmup_stream) + graph_conv.zero_() + graph_ssm.zero_() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = self.gdn.ssm_decode( + static_projected, graph_conv, graph_ssm, batch_indices=indices + ) + graph_conv.zero_() + graph_ssm.zero_() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(graph_output, expected, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(graph_conv, eager_conv, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(graph_ssm, eager_ssm, atol=3e-2, rtol=3e-2) diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 5d3c33264f4..322f1d78911 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -6,7 +6,10 @@ from megatron.core.extensions.transformer_engine import TEDotProductAttention from megatron.core.models.hybrid.hybrid_block import HybridStack from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols, validate_segment_layers -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_inference_stack_spec, + hybrid_stack_spec, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net import GatedDeltaNet from megatron.core.ssm.mamba_layer import MambaLayer @@ -265,6 +268,12 @@ def test_gdn_layer_types(self): assert isinstance(layers[1].self_attention, SelfAttention) assert isinstance(layers[2], MambaLayer) + def test_gdn_inference_spec(self): + """The inference stack must materialize GDN rather than its IdentityOp default.""" + gdn_spec = hybrid_inference_stack_spec.submodules.gdn_layer + assert gdn_spec.module is TransformerLayer + assert gdn_spec.submodules.self_attention.module is GatedDeltaNet + def test_gdn_gpu_forward(self): """Test GPU forward pass with GDN, attention, and Mamba layers.""" layer_pattern = Symbols.GDN + Symbols.ATTENTION + Symbols.MAMBA diff --git a/uv.lock b/uv.lock index 8a0dfaa9a3a..a1abafea131 100644 --- a/uv.lock +++ b/uv.lock @@ -2224,7 +2224,6 @@ dev = [ { name = "emerging-optimizers" }, { name = "fast-hadamard-transform" }, { name = "fastapi" }, - { name = "flash-linear-attention" }, { name = "flashinfer-python" }, { name = "hypercorn" }, { name = "megatron-energon", extra = ["av-decode"] }, @@ -2257,6 +2256,7 @@ otel = [ ] ssm = [ { name = "causal-conv1d" }, + { name = "flash-linear-attention" }, { name = "mamba-ssm" }, ] te = [ @@ -2339,7 +2339,7 @@ requires-dist = [ { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, - { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "==0.5.1" }, + { name = "flash-linear-attention", marker = "extra == 'ssm'", specifier = "==0.5.1" }, { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" },