core.ssm.packed_seq_helpers

Shared helpers for SSM mixers handling packed (THD-format) sequences.

Lifted from MambaMixer._create_packed_seq_idx so GDP, KDA, DPv2, GDN can share a single reference implementation (avoids drift across mixers).

Module Contents

Functions

get_cu_seqlens	Pick the right cu_seqlens tensor (padded if available).
build_packed_seq_idx	Build the per-token sequence index tensor used by varlen kernels.
check_fla_sequence_packing_support	Lighter sibling of _check_mamba_sequence_packing_support for FLA-backed mixers.

API

core.ssm.packed_seq_helpers.get_cu_seqlens(

packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams,

) → torch.Tensor

Pick the right cu_seqlens tensor (padded if available).

core.ssm.packed_seq_helpers.build_packed_seq_idx(

packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams,

total_tokens: int,

) → torch.Tensor

Build the per-token sequence index tensor used by varlen kernels.

For packed_seq_params.cu_seqlens_q[_padded] of the form [0, 5, 7, 11] and total_tokens=16 returns [0,0,0,0,0, 1,1, 2,2,2,2, 3,3,3,3,3] (shape [1, total_tokens], int32). The trailing chunk after cu_seqlens[-1] is treated as one extra sequence so the output covers every token in the pack. If cu_seqlens[-1] == total_tokens no extra index is added.

This is the per-token tensor consumed by causal_conv1d_fn(seq_idx=...) and by Mamba’s fused conv+SSM kernel as seq_idx.

total_tokens must equal the post-parallelism-gather sequence length that the kernel will actually consume — not the caller’s hidden_states.shape[0] which may be sequence-parallel-sharded and/or context-parallel-sliced. The robust pattern (mirrors mamba_mixer.py) is to call this after in_proj (SP all-gather) and pre_conv_ssm (CP all-to-all), passing the post-gather tensor’s seq dim — that way the helper is agnostic to TP/SP/CP shapes upstream.

core.ssm.packed_seq_helpers.check_fla_sequence_packing_support() → Tuple[bool, Optional[str]]

Lighter sibling of _check_mamba_sequence_packing_support for FLA-backed mixers.

GDP/KDA/DPv2/GDN reach into FLA’s chunk_kda / chunk_gated_delta_product / chunk_gated_delta_rule, all of which manage their own variable-length state internally. The only shared external dependency is the causal conv1d kernel — causal_conv1d_fn(seq_idx=...) was added in 1.4.0 and is required to reset the conv state at packed-document boundaries.

Mamba2’s stricter mamba_ssm minimums (used by mamba_split_conv1d_scan_combined) do not apply.