Is this page helpful?

core.ssm.ssm_inference#

Shared dynamic-batching inference scaffolding for linear-attention mixers.

A growing family of mixers in Megatron behave like “linear attention” / SSM recurrences for inference purposes: they carry a small per-request recurrent state (a short-convolution state plus a matrix-valued SSM state) instead of a growing KV cache. Mamba was the first; Gated Delta Net / Gated Delta Product (GDP) and friends are the same shape of computation with different kernels.

All of these variants share an identical request-level control flow for the dynamic inference engine:

1. Fetch this layer's (conv_state, ssm_state) slabs from the context.
2. Project the packed input (`in_proj`).
3. Split the packed batch into a decode partition (1 token per request,
   placed first) and a prefill partition (variable length, placed after).
   The kernels cannot mix the two, so they run independently.
4. Run the decode and prefill kernels on their respective partitions.
5. Merge the two partitions back into packed token order.
6. Apply the output projection (`out_proj`).

Only the kernels in step 4 differ between variants. This mixin owns the shared control flow (steps 1-3, 5, 6 and the orchestration) and delegates the variant-specific work to two hooks, ssm_decode and ssm_prefill. New linear-attention variants should subclass this mixin and implement those two hooks rather than re-deriving the decode/prefill bookkeeping.

Both hooks are given the DynamicInferenceContext directly and read whatever per-step metadata they need from context.mamba_metadata / context.mamba_slot_allocator themselves; there is deliberately no intermediate “unpack the metadata into a long argument list” layer.

Speculative decoding is supported by the shared orchestration: the decode path reshapes tokens into [batch, seq_len, d], fetches intermediate state buffers from the context, and passes them to ssm_decode. Variants that do not yet support speculative decoding should assert seq_len == 1 inside their ssm_decode implementation.

Chunked prefill and prefix caching are handled entirely inside ssm_prefill via context.mamba_metadata and context.mamba_slot_allocator; the mixin orchestration is unaware of them.

Note: static-batching (“legacy”) inference is intentionally not part of this interface. Concrete mixers keep any static/eager inference path separate so it does not pollute the dynamic decode/prefill hooks defined here.

Module Contents#

Classes#

SSMChunking

Chunk-related facts shared by every SSM layer in a stack.

SSMDynamicInferenceMixin

Mixin providing the shared decode/prefill orchestration for the dynamic inference engine. Concrete mixers implement the two ssm_* hooks below.

Functions#

ssm_chunking

Returns the chunking every SSM layer in a stack shares, or None.

API#

class core.ssm.ssm_inference.SSMChunking#

Bases: typing.NamedTuple

Chunk-related facts shared by every SSM layer in a stack.

chunk_size: int#

None

The mixer’s configured chunk size.

inference_chunk_size: int#

None

The chunk length the dynamic-inference prefill kernels actually run at.

num_householder: int#

None

Householder copies for Gated Delta Product layers; 0 for other mixers.

core.ssm.ssm_inference.ssm_chunking(
layer_type_list: List[str],
layers: Sequence,
) → Optional[core.ssm.ssm_inference.SSMChunking]#

Returns the chunking every SSM layer in a stack shares, or None.

None means the stack holds no recurrent layer, which happens on a pipeline stage made up entirely of attention and MLP layers.

The stack is assumed homogeneous: one mixer type, one chunking. A mixed stack would need a per-mixer alignment quantum and per-mixer chunk descriptors, and nothing downstream models that, so it is rejected here rather than silently taking the first layer’s answer for every layer.

Parameters:

  • layer_type_list – Per-layer symbols, positionally matching layers. See megatron/core/models/hybrid/hybrid_layer_allocation.py.

  • layers – The stack’s layers.

Returns:

The shared SSMChunking, or None if no layer is recurrent.

class core.ssm.ssm_inference.SSMDynamicInferenceMixin#

Mixin providing the shared decode/prefill orchestration for the dynamic inference engine. Concrete mixers implement the two ssm_* hooks below.

abstractmethod ssm_decode(
zxBCdt: torch.Tensor,
conv_state: torch.Tensor,
ssm_state: torch.Tensor,
batch_indices: torch.Tensor,
intermediate_conv_state: torch.Tensor = None,
intermediate_ssm_state: torch.Tensor = None,
) → torch.Tensor#

Run the single-token-per-request decode kernels.

Parameters:

  • zxBCdt – [decode_req_count, seq_len, proj_dim] projected decode tokens, where seq_len = 1 + num_speculative_tokens.

  • conv_state – [num_slots, conv_channels, d_conv] conv state cache.

  • ssm_state – [num_slots, *ssm_shape] SSM state cache.

  • batch_indices – [decode_req_count] slot index per decode request (-1 marks padding slots).

  • intermediate_conv_state – Optional buffer for storing conv states at intermediate sequence steps (speculative decoding).

  • intermediate_ssm_state – Optional buffer for storing SSM states at intermediate sequence steps (speculative decoding).

Returns [decode_req_count, seq_len, d_inner]; updates state in place. Variants that do not yet support speculative decoding should assert seq_len == 1 inside their implementation.

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

Run the variable-length prefill kernels for all prefill requests.

The implementation reads its varlen metadata (cu_seqlens, batch_indices_prefill, seq_idx, chunk boundaries, intermediate extraction buffers, etc.) directly from context.mamba_metadata and context.mamba_slot_allocator and processes every prefill request in one varlen call, writing the resulting final states back into the caches.

Returns [prefill_token_count, 1, d_inner]; updates state in place.

ssm_dynamic_inference(
hidden_states: torch.Tensor,
context: megatron.core.inference.contexts.DynamicInferenceContext,
) → Tuple[torch.Tensor, torch.Tensor]#

Execute one dynamic inference step for a linear-attention mixer.

Separates decode and prefill requests, runs them through the variant-specific kernels independently, and merges the results back into packed token order.