monkeypatch.models.mamba.modeling

monkeypatch.models.mamba.modeling

Sample-packing patches for the pure-SSM transformers models: Mamba, Mamba2, Falcon-Mamba.

None of these forwards accept position_ids, so the ForCausalLM wrapper turns them into seq_idx and stashes the boundaries on every block, which passes them to its mixer as a kwarg. The mixers forward kwargs into their kernels, so Mamba2 needs nothing more than a live kernel that takes seq_idx.

Mamba1’s selective scan has no such argument, so a packed row is scattered into right-padded per-document batches for the causal conv and the scan and gathered back: exact for real tokens, padding costs only those two ops. A conv kernel that takes seq_idx is used as is; the torch fallback goes per document too, so no kernel is required. Its fused training kernel bakes in the whole row, so it is disabled for packed batches and the forward continues on the unfused branch.

Classes

Name	Description
PackedSegments	Document boundaries of a packed batch, with a lazily built scatter plan.

PackedSegments

monkeypatch.models.mamba.modeling.PackedSegments(seq_idx, _plan=None)

Document boundaries of a packed batch, with a lazily built scatter plan.

Attributes

Name	Description
plan	[(index, mask)] groups, each [docs, max_len] into the flat B*T axis.

Functions

Name	Description
build_segment_plan	Group documents by length so no group pads beyond pad_factor of its tokens.
packed_causal_conv	Run a causal conv on each packed document separately; returns [B, C, T].
packed_selective_scan	Run scan_fn on each packed document separately; returns [B, D, T].

build_segment_plan

monkeypatch.models.mamba.modeling.build_segment_plan(
    seq_idx,
    pad_factor=PAD_FACTOR,
)

Group documents by length so no group pads beyond pad_factor of its tokens.

packed_causal_conv

monkeypatch.models.mamba.modeling.packed_causal_conv(
    conv_fn,
    segments,
    x,
    *args,
    **kwargs,
)

Run a causal conv on each packed document separately; returns [B, C, T].

Right padding never reaches earlier positions of a causal conv, so the per-document outputs are exact.

packed_selective_scan

monkeypatch.models.mamba.modeling.packed_selective_scan(
    scan_fn,
    segments,
    u,
    delta,
    A,
    B,
    C,
    D,
    z,
    delta_bias,
    **kwargs,
)

Run scan_fn on each packed document separately; returns [B, D, T].