Add GDN dynamic inference

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

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

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(

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

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

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