monkeypatch.models.mamba_utils

monkeypatch.models.mamba_utils

Shared utilities for Mamba2 SSM sample-packing and context-parallelism patches.

Used by: nemotron_h, falcon_h1, granite_moe_hybrid

Functions

Name	Description
assert_mamba2_scan_honours_seq_idx	Probe the resolved chunk-scan once, and raise if it drops `seq_idx`.
ensure_mamba_kernels_loaded	Eagerly resolve mamba-ssm and causal-conv1d globals on target_module.
get_seq_idx	Convert position_ids [B, T] → seq_idx [B, T] int32 for mamba-ssm kernels.
is_cp_active	Return True if context parallelism (ring attention) is active on this rank.
kernel_accepts	Whether a hub-wrapped kernel’s live implementation takes `name`.
mamba2_seq_idx_kernels_available	True when the installed Mamba2 kernels honour `seq_idx`.
patch_model_forward_seq_idx	Inject `seq_idx` (from `position_ids`) into `model_cls.forward` kwargs.
require_seq_idx_kernels	Raise unless every kernel in `names` honours `seq_idx` (or Hub kernels will).

assert_mamba2_scan_honours_seq_idx

monkeypatch.models.mamba_utils.assert_mamba2_scan_honours_seq_idx(target_module)

Probe the resolved chunk-scan once, and raise if it drops seq_idx.

Which implementation ends up behind mamba2_chunk_scan is only settled once the model is built (pip package, Hub kernel, or the torch fallback), and the fallback drops unknown kwargs silently — training would then mix SSM state across packed samples. The probe is skipped when it cannot run (no CUDA device, or a kernel that rejects the probe shapes).

ensure_mamba_kernels_loaded

monkeypatch.models.mamba_utils.ensure_mamba_kernels_loaded(target_module)

Eagerly resolve mamba-ssm and causal-conv1d globals on target_module.

Transformers >= 5.5 lazily loads these inside Mixer.__init__ via lazy_load_kernel. Our monkeypatches run before model instantiation, so the module globals are still None. This helper triggers the kernel resolution early so the patched cuda_kernels_forward can reference them.

get_seq_idx

monkeypatch.models.mamba_utils.get_seq_idx(position_ids)

Convert position_ids [B, T] → seq_idx [B, T] int32 for mamba-ssm kernels.

Example: position_ids [[0,1,2,3,0,1,2]] → seq_idx [[0,0,0,0,1,1,1]]

Under context parallelism a rank may receive a chunk that begins mid-sample (position_ids[0] != 0), so the raw cumsum starts at 0 and subtracting 1 would yield -1 — an invalid value for the Mamba kernels. Subtracting the first element of the cumsum instead normalises every chunk to start at 0 while still correctly incrementing at every intra-chunk sample boundary.

Example (CP rank 1, chunk starts mid-sample): position_ids [[3,4,5,0,1,2]] → seq_idx [[0,0,0,1,1,1]]

is_cp_active

monkeypatch.models.mamba_utils.is_cp_active()

Return True if context parallelism (ring attention) is active on this rank.

Zero-cost when CP is not configured: the import guard ensures we only touch the distributed group if ring_flash_attn is installed.

kernel_accepts

monkeypatch.models.mamba_utils.kernel_accepts(fn, name)

Whether a hub-wrapped kernel’s live implementation takes name.

transformers resolves the pip kernel at import time and filters kwargs to its signature, so a torch fallback (which takes **kwargs) drops seq_idx without a word. Once kernelize has swapped in a Hub kernel, its own signature is checked instead; the Hub layers declare **kwargs and pop the name inside, so their source is searched. Returns None when nothing can be told.

mamba2_seq_idx_kernels_available

monkeypatch.models.mamba_utils.mamba2_seq_idx_kernels_available()

True when the installed Mamba2 kernels honour seq_idx.

transformers resolves mamba-ssm / causal-conv1d ahead of its own torch fallbacks, and those fallbacks drop unknown kwargs — so without the real kernels seq_idx is silently ignored and SSM state leaks across packed samples.

patch_model_forward_seq_idx

monkeypatch.models.mamba_utils.patch_model_forward_seq_idx(model_cls)

Inject seq_idx (from position_ids) into model_cls.forward kwargs.

For architectures whose mixers already accept seq_idx through **kwargs but whose model forward never derives it, this is the only piece missing for sample packing.

require_seq_idx_kernels

monkeypatch.models.mamba_utils.require_seq_idx_kernels(
    mod,
    names,
    model_type,
    kernels_enabled,
)

Raise unless every kernel in names honours seq_idx (or Hub kernels will).