Is this page helpful?

core.ssm.mamba_mixer#

Module Contents#

Classes#

ExtendedRMSNorm

RMSNormGated with sharded state dict.

MambaMixerSubmodules

Contains the module specs for the input and output linear layers.

MambaMixer
param config:

The config of the model.

param submodules:

Contains the module specs for the input and output linear layers.

param d_model:

The hidden size of the model.

param d_state:

The state size of the SSM.

param d_conv:

The number of channels in the causal convolution.

param conv_init:

The initialization range for the causal convolution weights.

param expand:

The expansion factor for the SSM.

param headdim:

The hidden size of each attention head.

param ngroups:

The number of attention heads.

param A_init_range:

The initialization range for the attention weights.

param D_has_hdim:

Whether the D parameter has the same number of dimensions as the hidden state.

param rmsnorm:

Whether to use root mean square normalization.

param norm_before_gate:

Whether to apply normalization before the gating mechanism.

param dt_min:

The minimum value of the dt parameter.

param dt_max:

The maximum value of the dt parameter.

param dt_init:

The initialization value of the dt parameter.

param dt_scale:

The scaling factor for the dt parameter.

param dt_init_floor:

The minimum value of the dt parameter after initialization.

param bias:

Whether to use bias in the linear layers.

param conv_bias:

Whether to use bias in the causal convolution.

param chunk_size:

The chunk size for the Mamba SSM fused kernel.

param use_mem_eff_path:

Whether to use the memory-efficient path for the Mamba model.

param layer_number:

The layer number of this Mamba layer.

param pg_collection:

The required process groups to use for tensor model parallel and context parallel.

Functions#

_check_mamba_sequence_packing_support

Checks whether causal_conv1d and mamba_ssm support sequence packing.

Data#

MAMBA_HAS_STATE_DTYPE

logger

API#

core.ssm.mamba_mixer.MAMBA_HAS_STATE_DTYPE#

None

core.ssm.mamba_mixer.logger#

‘getLogger(…)’

class core.ssm.mamba_mixer.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.mamba_mixer.MambaMixerSubmodules#

Contains the module specs for the input and output linear layers.

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.mamba_mixer.MambaMixer(
config: megatron.core.transformer.TransformerConfig,
submodules: core.ssm.mamba_mixer.MambaMixerSubmodules,
d_model,
d_conv=4,
conv_init=None,
expand=2,
A_init_range=(1, 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=0.0001,
bias=False,
conv_bias=True,
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, megatron.core.transformer.module.TwoStageAttentionLayer

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_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 Mamba SSM 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.

Initialization

Parameters:

name (str | None) – module instance name passed top-down from its paranet module

_mamba_chunk(
hidden_states,
conv_state=None,
ssm_state=None,
packed_seq_params=None,
inference_mode=False,
)#

Run Mamba through its normalized SSM output, before output projection.

forward_pre_attn_and_core_attn(
hidden_states,
*,
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
packed_sequence_cp_metadata=None,
)#

Run the training pre-attention and core-attention stage.

forward_post_core_attn(y)#

Apply the Mamba output projection to an SSM output tensor.

forward(
hidden_states,
inference_context=None,
*,
inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None,
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
)#

Run the input projection/SSM phase followed by the output projection.

_static_decode(
hidden_states,
conv_state,
ssm_state,
) → Tuple[torch.Tensor, torch.Tensor]#

Single-token static-batching decode step (updates state in place).

_static_prefill(
zxBCdt: torch.Tensor,
conv_state: Optional[torch.Tensor],
ssm_state: Optional[torch.Tensor],
) → torch.Tensor#

Performs single-sequence SSM prefill for static-batching inference and the non-memory-efficient (use_mem_eff_path=False) training fallback.

conv_state / ssm_state are None for the training fallback and non-None for static-batching inference (updated in place).

Parameters:

  • zxBCdt – The input tensor of shape (l, b, d), a concatenation of z, x, B, C, and dt projections.

  • conv_state – The convolution state tensor, or None for training.

  • ssm_state – The selective scan state tensor, or None for training.

Returns:

Output tensor of shape (l, b, d).

_ssm_training(
zxBCdt: torch.Tensor,
packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,
) → torch.Tensor#

Performs SSM computation for training step.

Uses the memory-efficient kernel mamba_split_conv1d_scan_combined which reduces the size of forward activations stored for backprop and therefore reduces memory pressure during training.

ssm_prefill(
zxBCdt: torch.Tensor,
conv_state: torch.Tensor,
ssm_state: torch.Tensor,
context: megatron.core.inference.contexts.DynamicInferenceContext,
) → torch.Tensor#

Performs the variable-length SSM prefill for all dynamic-batching prefill requests.

All varlen metadata (cu_seqlens, seq_idx, batch_indices, chunk boundaries, intermediate extraction buffers, etc.) is read directly from context.mamba_metadata / context.mamba_slot_allocator – there is no intermediate layer that unpacks the metadata into a long argument list. All prefill requests (including chunked prefill) are processed together through a single varlen kernel call; the precomputed metadata avoids .item() calls and data-dependent control flow, enabling CUDA graph compatibility. Intermediate state extraction (for Mamba prefix caching) is performed via pre-allocated output buffers, also CUDA graph compatible.

When padded_prefill_count > 0 but real_prefill_count == 0 (e.g. a decode-only rank in expert parallelism that must match a mixed CUDA graph), the full kernel path still runs; the metadata reflects zero-length sequences (cu_seqlens all equal, batch_indices all -1) so the kernels produce a correctly-shaped zero output tensor, which is required by the merge logic in ssm_dynamic_inference.

Parameters:

  • zxBCdt – The input tensor of shape (l, b, d), which is a concatenation of z, x, B, C, and dt projections.

  • conv_state – The convolution state tensor for inference.

  • ssm_state – The selective scan state tensor for inference.

  • context – The dynamic inference context supplying all varlen metadata.

Returns:

Output tensor of shape (l, b, d). Intermediate states (if any) are written directly into the slot-allocator buffers held by context.

refresh_cache() → None#

Refresh the existing decode-cache storage from the current A_log.

_get_decode_A_neg_exp() → torch.Tensor#

Cached -exp(A_log.float()) pre-expanded to (nheads, headdim, dstate).

A_log is frozen during inference; recomputing it per token otherwise launches three small elementwise kernels (float cast, exp, neg) that rival selective_state_update itself in the decode profile. The stride-0 expand view also triggers the kernel’s TIE_HDIM fast path.

_get_batch_invariant_decoder() → megatron.core.ssm.ops.mamba2.batch_invariant_decode.MambaBatchInvariantDecode#

Batch-invariant decode adapter, created on first use.

train(mode: bool = True)#

Mark the decode cache stale in training and refresh it for evaluation.

ssm_decode(
zxBCdt: 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#

Performs SSM computation for inference decode step.

Parameters:

  • zxBCdt – The input tensor of shape (b, s, d), which is a concatenation of z, x, B, C, and dt projections. s is the sequence length (1 + num_speculative_tokens).

  • conv_state – The convolution state tensor for inference.

  • ssm_state – The selective scan state tensor for inference.

  • batch_indices – A map from batch id to position in the Mamba state tensors.

  • intermediate_conv_state – Optional buffer for storing conv state at each sequence step (for speculative decoding rollback).

  • intermediate_ssm_state – Optional buffer for storing SSM state at each sequence step (for speculative decoding rollback).

Returns:

The output tensor of shape (b, s, d).

mamba_state_shapes_per_request() → Tuple[Tuple[int], Tuple[int]]#

Returns the Mamba conv and ssm states shapes per request.

property ssm_inference_chunk_size: int#

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

The Mamba2 SSD kernels run at the same chunk size for training and inference, so this is just chunk_size. It exists so callers that need a chunk-aligned boundary can ask every SSM mixer the same question – Gated Delta Product answers 64 regardless of its 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.mamba_mixer._check_mamba_sequence_packing_support(
for_inference_not_training: bool = True,
) → Tuple[bool, Optional[str]]#

Checks whether causal_conv1d and mamba_ssm support sequence packing.