`core.ssm.gated_delta_product`#

Module Contents#

Classes#

`ExtendedRMSNorm`	RMSNormGated with sharded state dict.
`GatedDeltaProductMixerSubmodules`	Contains the module specs for the input and output projections.
`GatedDeltaProductMixer`	Gated Delta Product (GDP) sequence mixer for hybrid models.

Functions#

`_kernel_accepts_kwarg`	Return True if `kernel` explicitly declares a keyword argument called `name`.
`_get_in_proj_checkpoint_split_layout`	Return TP-reshardable splits for the packed `[z,V,K,Q,b,a]` projection.
`_get_conv_checkpoint_split_layout`	Return TP-reshardable splits for the packed `[V,K,Q]` convolution.
`_split_tensor_factory`	Builds a factory that splits a given ShardedTensor into several independent chunks.

Data#

logger

API#

core.ssm.gated_delta_product.logger#: ‘getLogger(…)’

core.ssm.gated_delta_product._kernel_accepts_kwarg(kernel, name: str) → bool#: Return True if kernel explicitly declares a keyword argument called name.

class core.ssm.gated_delta_product.ExtendedRMSNorm(/, *args, **kw)#

Bases: mamba_ssm.ops.triton.layernorm_gated.RMSNorm

RMSNormGated with sharded state dict.

Initialization

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#: Sharding along axis 0, bias not sharded

class core.ssm.gated_delta_product.GatedDeltaProductMixerSubmodules#

Contains the module specs for the input and output projections.

in_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

out_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

class core.ssm.gated_delta_product.GatedDeltaProductMixer( config: megatron.core.transformer.TransformerConfig, submodules: core.ssm.gated_delta_product.GatedDeltaProductMixerSubmodules, d_model, d_conv=4, conv_init=None, 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_floor=0.0001, bias=False, conv_bias=False, chunk_size=128, layer_number=None, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection = None, pp_layer_offset: int = 0, name: str | None = None, )#

Bases: megatron.core.ssm.ssm_inference.SSMDynamicInferenceMixin, megatron.core.transformer.module.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 selected GDP kernel 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.

Parameters:

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_conv – The number of channels in the causal convolution.
conv_init – The initialization range for the causal convolution weights.
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_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.
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.

Initialization

forward( hidden_states, inference_context=None, *, inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, packed_seq_params=None, packed_sequence_cp_metadata: megatron.core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata | None = None, )#: Run the gated delta product mixer on hidden states.

_packed_metadata( packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, ) → tuple[torch.Tensor | None, torch.Tensor | None]#: Return sequence indices and cumulative lengths for packed input.

_make_uniform_cutedsl_cu_seqlens(VKQ: torch.Tensor) → torch.Tensor#

Build uniform sequence boundaries for an unpacked CuTeDSL input.

Under chunkwise CP, sequence_length is the rank-local shard length, so these boundaries describe only this rank. The CP metadata determines which boundary sequence continues on adjacent ranks.

_gdp_chunk_forward( hidden_states, conv_state=None, ssm_state=None, packed_seq_params=None, packed_sequence_cp_metadata: megatron.core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata | None = None, )#

Chunked-kernel forward, shared by training and static-batching prefill: input projection, causal conv, QKV preparation, and the chunked gated delta product kernel, with optional in_proj/QKV recompute and fine-grained activation offloading of the causal conv input.

Prefill passes conv_state and ssm_state so the trailing conv window and the final recurrent state are cached for the decode steps.

_run_gdp_kernel( query, key, value, g, beta, VKQ, output_final_state, cu_seqlens=None, preceding_rank_start=0, following_rank_stop=None, )#

Run the selected chunked gated delta product kernel and return (core_attn_out, final_state).

The CuTeDSL kernel uses a flattened token axis described by cu_seqlens and applies the query/key L2 norm itself. Its output is restored to (b, l, h, p) here. The FLA kernel uses the batched layout and receives query/key tensors already normalized by _prepare_qkv. Unpacked FLA leaves cu_seqlens unset.

_chunkwise_packed_metadata( packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None, local_sequence_length: int, metadata: megatron.core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata | None, ) → megatron.core.ssm.context_parallel.chunkwise.PackedSequenceCPMetadata | None#: Validate and return cached rank-local packed-sequence metadata.

_in_proj_preprocess(hidden_states, packed_seq_params=None)#

Run the input projection, gather its output across CP ranks, switch to (b, l, d) layout, and split it into the z, VKQ, and ba groups.

Kept as a single function so that it can be checkpointed as one unit: the in_proj output and the intermediate CP-gathered/transposed copies of it are then freed inside the checkpoint, leaving only hidden_states saved for the backward pass.

_preprocess(zVKQba)#: Switch the (l, b, proj_dim) input projection to the batch-first layout the causal conv and the kernels expect, and split it into the z, VKQ, and ba groups.

_prepare_qkv( x, conv_state=None, seq_idx=None, l2_norm_in_kernel=False, cu_seqlens=None, precomputed_seq_idx=None, precomputed_seq_start=None, conv_initial_states=None, )#

Run the causal conv on the VKQ slice and split/reshape it into query, key, and value.

x: (b, l, d). Keep the transposes to the conv layout inside this function so that the checkpointed subgraph takes the (b, l, d) slice as input and hands back a gradient that is already contiguous in (b, l, d) for the split backward’s concat.

Decode does not come through here: it runs _decode_conv followed by the fused gdp_decode_prepare kernel, which produces the same query, key, and value in one launch.

l2_norm_in_kernel leaves the query/key L2 norm to the caller’s kernel, which every caller whose kernel normalizes internally must set: the CuTeDSL path (use_qk_l2norm_in_kernel) and ssm_prefill’s varlen kernel.

Dynamic-batching prefill passes cu_seqlens plus the precomputed per-token conv metadata, which routes to the forked varlen conv instead of causal_conv1d_fn: it reads the request ids and start offsets rather than deriving them, which is what makes the step replayable from a CUDA graph. conv_initial_states is the left conv boundary for each request, which under chunked prefill is the tail of the previous chunk rather than zeros.

_decode_conv(x, conv_state, conv_state_indices)#

Step the cached short conv one token per request, in place.

Indexed conv update: reads/writes the per-request conv state rows selected by conv_state_indices. Unlike causal_conv1d_fn, the Triton causal_conv1d_update takes x as (batch, seq_len, dim) – a 2-D (batch, dim) input is unsqueezed at dim 1 – so decode keeps the (b, l, d) layout rather than transposing to [B, D, L]. Here l == 1.

_compute_gating(ba)#: Compute the beta and g gating tensors from the ba slice.

_postprocess(core_attn_out, z, packed_seq_params=None)#: Switch back to (l, b, d) layout, scatter across CP ranks, and apply the gated output norm.

_static_decode( 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).

ssm_decode( 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 [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, whose outputs are zeroed); batch_indices=None means static batching, where the caches are already in request order.

Every op here is CUDA-graph safe: no host synchronization, no data-dependent shapes, and the state caches are addressed by device-side indices rather than gathered and scattered.

ssm_prefill( zVKQba: torch.Tensor, conv_state: torch.Tensor, ssm_state: torch.Tensor, context: megatron.core.inference.contexts.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].

Each request enters carrying whatever conv/SSM state its cache slot holds and leaves having written the slot back, which is the whole of the chunked-prefill contract: cu_seqlens describes only the slice of the prompt scheduled this step, and the state cache is the only thing that carries across steps. A fresh request’s slot was zeroed at allocation, so the same code path serves first and continuation chunks.

The whole step is CUDA-graph capturable: it runs on the in-tree kernels, which take precomputed chunk descriptors and per-token conv metadata and so need neither a host synchronization nor a data-dependent shape. Padding requests are zero-length sequences with a -1 state slot; they produce zero output and touch no state.

allocate_inference_cache(batch_size, max_seqlen, dtype=None)#: allocate inference cache

mamba_state_shapes_per_request() → Tuple[Tuple[int], Tuple[int]]#: Returns the Mamba conv and SSM state shapes per request.

property ssm_inference_chunk_size: int#

Chunk length the dynamic-inference prefill kernels actually run at.

The forked GDP kernels chunk at a fixed 64, independent of the training-path chunk_size. Scheduling decisions that must land on a chunk boundary (batch-invariant chunked prefill, recurrent-state extraction for prefix caching) align to this, not to chunk_size.

_get_states_from_cache( 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.

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#: Provide a sharded state dictionary for distributed checkpointing.

core.ssm.gated_delta_product._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.

core.ssm.gated_delta_product._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.

core.ssm.gated_delta_product._split_tensor_factory( orig_sh_ten: megatron.core.dist_checkpointing.ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int, ) → megatron.core.dist_checkpointing.mapping.ShardedTensorFactory#: Builds a factory that splits a given ShardedTensor into several independent chunks.

core.ssm.gated_delta_product#