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 scan alone (the conv resets through seq_idx) and gathered back: exact for real tokens, padding costs only the scan. 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_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_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].