Mamba / SSU¶

MATE implements Mamba2 selective state update (SSU) on MUSA. Mamba2 is the multi-head successor to Mamba: it keeps one decay per head (A[h]) instead of a per-channel decay, which enables the chunked SSD formulation and makes the recurrence matmul-friendly. Two TileLang kernels share the selective_state_update entry point: a one-token kernel for the plain decode step (one token per sequence, no slot table) and a packed kernel for everything else – packed variable-length rows, multi-token (MTP) decoding with num_accepted_tokens, and the dense 4-D multi-token form. selective_state_update applies decode steps of the SSM recurrence, and is the MUSA-native implementation behind the FlashInfer-shaped flashinfer.mamba.selective_state_update compatibility surface. The packed variable-length SSD prefill path is covered in SSD packed prefill.

Minimal SSU example:

import torch
from mate.mamba import selective_state_update

batch, heads, dim, dstate, groups, slots = 2, 8, 8, 128, 2, 8
state = torch.zeros(slots, heads, dim, dstate, device="musa", dtype=torch.float16)
x = torch.randn(batch, heads, dim, device="musa", dtype=torch.bfloat16)
dt = torch.rand(batch, heads, dim, device="musa", dtype=torch.float32) * 0.1
A = -torch.rand(heads, device="musa", dtype=torch.float32)
B = torch.randn(batch, groups, dstate, device="musa", dtype=torch.bfloat16)
C = torch.randn(batch, groups, dstate, device="musa", dtype=torch.bfloat16)
slots_index = torch.tensor([0, 1], device="musa", dtype=torch.int32)

y = selective_state_update(state, x, dt, A, B, C, None)
y = selective_state_update(
    state,
    x,
    dt,
    A,
    B,
    C,
    None,
    state_batch_indices=slots_index,
    dst_state_batch_indices=slots_index,
)

Speculative decoding (MTP) example, batch sequences of steps tokens:

steps = 7
rows = batch * steps
x = torch.randn(rows, heads, dim, device="musa", dtype=torch.bfloat16)
# One state slot per (sequence, speculative position) on the destination side,
# and the position the accepted token stopped at on the read side.
read_slots = torch.arange(rows, device="musa", dtype=torch.int32).view(batch, steps)
write_slots = (read_slots + rows).contiguous()
cu_seqlens = torch.arange(0, rows + 1, steps, device="musa", dtype=torch.int32)
accepted = torch.full((batch,), steps, device="musa", dtype=torch.int32)

y = selective_state_update(
    state,
    x,
    dt,
    A,
    B,
    C,
    None,
    state_batch_indices=read_slots,
    dst_state_batch_indices=write_slots,
    num_accepted_tokens=accepted,
    cu_seqlens=cu_seqlens,
    pad_slot_id=0,
)

vLLM’s MTP6 decode passes the slot tables flat instead: one entry per packed row, which the Triton kernel addresses by unsqueezing to [rows, 1]. Both forms are accepted, and with one sequence per call they name the same entries:

rows = 7
slots = torch.arange(rows, device="musa", dtype=torch.int32)
y = selective_state_update(
    state,
    x,
    dt,
    A,
    B,
    C,
    None,
    state_batch_indices=slots,
    dst_state_batch_indices=slots,
    num_accepted_tokens=torch.tensor([1], device="musa", dtype=torch.int32),
    cu_seqlens=torch.tensor([0, rows], device="musa", dtype=torch.int32),
    pad_slot_id=0,
)

SSU at a glance¶

Capability	Supported scope
Device	MUSA through the TileLang backend
Sequence layout	One token per sequence (`x`, `dt` and `z` are `[batch, heads, dim]`) on the one-token path; the packed `[rows, heads, dim]` layout with `cu_seqlens`, or the dense `[batch, steps, heads, dim]` form, on the packed path
State	`[slots, heads, dim, dstate]`, contiguous, fp16/bf16/fp32, updated in place through the destination slots
`A` / `D` / `dt_bias`	per-head fp32, broadcasting over `dim`/`dstate`
`B` / `C`	`[batch, groups, dstate]` with the `x` dtype; `heads` must be divisible by `groups`
Output	`[batch, heads, dim]` in the `x` dtype, optionally written into a caller-provided `out`
`dt_softplus`	supported
Slot semantics	`state_batch_indices` and `dst_state_batch_indices` may be the flat `[rows]` form or the 2-D `[sequences, steps]` table; both are addressed by strides, so a flat table puts sequence `b`’s position `t` at entry `b + t` and a contiguous 2-D table at `b * steps + t`. `pad_slot_id` reads as a zero state and is never written; `disable_state_update` leaves the pool unchanged
Variable-length and MTP decode	supported: `cu_seqlens` splits the packed rows between sequences and `num_accepted_tokens` seeds each sequence’s read slot at `count - 1` (floored at 0). Every token then publishes its state to its own destination entry — or to its own entry of the read table when no destination table is given, the default the Triton kernel applies. `pad_slot_id` destinations are never written; empty sequences are skipped
Stochastic rounding	not implemented; `rand_seed` raises `NotImplementedError`

SSU toolchain requirements¶

TileLang for MUSA (tilelang-musa) is required; the kernel is compiled at run time for the shape and dtype combination it is first called with.
The recurrence runs in fp32 with the state updated in the state dtype.

Graph capture¶

Call prewarm_selective_state_update during warmup, before any graph capture: the TileLang kernel is compiled on first use, and a first-call compile inside a captured graph stalls the worker. The packed path issues a single kernel launch and performs no host/device scalar copies; the one-token path makes one comparison against cu_seqlens (the check that decides which kernel a call takes, and the only synchronization in the entry point).

SSD packed prefill¶

ssd_combined_fwd_varlen is the prefill-side counterpart of the decode entrypoint: one packed call over a variable-length batch, running the five SSD stages (chunk-local dt/decay cumsum, intra-chunk states, state passing, the C·Bᵀ products, and the chunk scan) as five TileLang launches. It is the symbol the FlashInfer-shaped compatibility wrapper dispatches to for --mamba-backend flashinfer on MUSA, and it drives the same implementation for every Mamba2/Nemotron-H prefill shape.

The per-shape intermediates come from a cached workspace, so a steady-state call allocates nothing beyond the returned state tensor. cu_chunk_seqlens is the only authority for the chunk-to-token mapping, and seq_idx – not the last_chunk_indices arithmetic – decides which chunks open a sequence and therefore where their entering state comes from.

dt_limit is a (min, max) tuple that clamps the processed timestep after softplus(dt + dt_bias). The default is (0.0, float("inf")), which leaves the result unclamped above and floors it at zero.

Everything outside the supported contract raises NotImplementedError or ValueError naming the missing capability rather than silently narrowing semantics: the checkpoint_* replay arguments are not implemented.

API reference¶

mate.mamba.selective_state_update(state: torch.Tensor, x: torch.Tensor, dt: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, D: torch.Tensor | None, z: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, dt_softplus: bool = False, state_batch_indices: torch.Tensor | None = None, pad_slot_id: int = -1, state_scale: torch.Tensor | None = None, out: torch.Tensor | None = None, disable_state_update: bool = False, intermediate_states_buffer: torch.Tensor | None = None, intermediate_state_indices: torch.Tensor | None = None, intermediate_state_scales: torch.Tensor | None = None, rand_seed: torch.Tensor | None = None, philox_rounds: int = 10, cache_steps: int = 0, algorithm: str = 'auto', dst_state_batch_indices: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, backend: str = 'auto') → torch.Tensor[source]¶

Apply one decode-step selective state update and return the output y.

The argument order, defaults and semantics follow the FlashInfer-shaped contract the MUSA consumers already call, so this function can serve as the implementation behind that surface without an adapter layer.

x is the packed [rows, heads, dim] form together with cu_seqlens, or the dense [batch, steps, heads, dim] form of the same recurrence. Both are supported, and the call shape – not a flag – decides which kernel runs.

mate.mamba.prewarm_selective_state_update(*, state_dtype: torch.dtype, io_dtype: torch.dtype, batch: int, heads: int, dim: int, dstate: int, groups: int, slot_dtype: torch.dtype, rows_per_cta: int = 4, lanes_per_row: int = 32, packed: bool = False, meta_dtype: torch.dtype = torch.int32, dt_dim: int | None = None) → None[source]¶

Compile the native kernel for one configuration, outside capture.

Call this during warmup (before any graph capture) for each shape/dtype combination the serving configuration will use. Compiling inside a captured graph stalls the worker, so prewarming is required rather than optional.

packed selects the kernel a multi-token serving configuration will actually run – the packed/MTP one – because that is a different compilation from the one-token fast path. dt_dim is what the configuration passes as dt: 1 for the per-head step the FlashInfer-shaped wrapper hands over (one value per head, broadcast over dim), None for this module’s default of a full [rows, heads, dim] dt. meta_dtype is the dtype of cu_seqlens/num_accepted_tokens, which the packed kernel specializes on.

mate.mamba.ssd_combined_fwd_varlen(x: torch.Tensor, dt: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, chunk_size: int, cu_seqlens: torch.Tensor, cu_chunk_seqlens: torch.Tensor, last_chunk_indices: torch.Tensor, seq_idx: torch.Tensor, out: torch.Tensor | None = None, D: torch.Tensor | None = None, z: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, dt_softplus: bool = False, dt_limit: tuple[float, float] = (0.0, inf), initial_states: torch.Tensor | None = None, return_intermediate_states: bool = False, state_dtype: torch.dtype | None = None, checkpoint_token_indices: torch.Tensor | None = None, checkpoint_state_slots: torch.Tensor | None = None, checkpoint_states: torch.Tensor | None = None) → torch.Tensor[source]¶

Packed variable-length SSD prefill, native on MUSA.

The five stages run as five tilelang launches with per-shape scratch taken from a cached workspace, so a steady-state call allocates nothing but the returned state tensor.

Contract, matching vLLM’s Mamba2 prefill caller:

x is [tokens, heads, dim] and dt is [tokens, heads], both packed over the flattened (variable-length) batch;
cu_chunk_seqlens is the only authority for the chunk-to-token mapping: chunk c covers [cu_chunk_seqlens[c], cu_chunk_seqlens[c+1]). A chunk never spans two sequences, and every chunk but the last of a sequence spans chunk_size tokens;
last_chunk_indices[b] is the ordinal of sequence b’s last logical chunk and seq_idx[c] is the request owning chunk c. A chunk opens a sequence when seq_idx differs from the previous chunk’s, and that transition – not the last_chunk_indices arithmetic – decides where the entering state comes from. The two can be made to disagree only by an inconsistent caller, and the two independent implementations of this contract resolve that disagreement differently, so seq_idx is the authority;
out is written in place when supplied;
the return value is a single tensor: the full [nchunks, heads, dim, dstate] state buffer when return_intermediate_states, otherwise states[last_chunk_indices] with shape [batch, heads, dim, dstate], both in state_dtype (which defaults to initial_states.dtype, else C.dtype).

checkpoint_* is not supported and raises rather than silently ignoring a capability the caller asked for.

mate.mamba.prewarm_ssd_combined_fwd_varlen(heads: int, groups: int, dim: int, dstate: int, chunk_size: int, io_dtype: torch.dtype = torch.bfloat16, state_dtype: torch.dtype = torch.bfloat16, dt_dtype: torch.dtype = torch.float32, block_M: int = 32) → None[source]¶

Compile every SSD stage for one shape, outside any captured graph.

Call during warmup for each configuration the server will serve: a first-call compile inside a captured graph stalls the worker. The compile-time degrees of freedom are the four shape parameters, the three dtypes and the scan’s block_M; the token and chunk counts are runtime dimensions.