Mamba Operators¶

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Full forward¶

tileops.mamba.Mamba2FwdOp ¶

Mamba-2 State-Space Dual (SSD) full forward pass operator.

Combines SSDChunkCumsum → SSDChunkCoupling → SSDChunkState → SSDStatePassing → SSDChunkScan into a single callable whose interface matches mamba_chunk_scan_combined from the official mamba_ssm library, except that final_states is always returned.

init ¶

__init__(
    chunk_size=256,
    dt_softplus=True,
    *,
    target=None
)

Build the op and its sub-ops. Shapes and dtype are taken from each call.

Parameters:

chunk_size (int, default: 256 ) –

Tokens per chunk (default 256).
dt_softplus (bool, default: True ) –

Apply softplus to (dt + dt_bias) before use.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
    dt,
    A,
    B,
    C,
    dt_bias=None,
    initial_states=None,
)

Run the full Mamba-2 SSD forward pass.

Parameters:

x (Tensor) –

(batch, seqlen, n_heads, d_head) dtype
dt (Tensor) –

(batch, seqlen, n_heads) float32
A (Tensor) –

(n_heads,) float32 (log-space, ≤ 0)
B (Tensor) –

(batch, seqlen, n_groups, d_state) dtype
C (Tensor) –

(batch, seqlen, n_groups, d_state) dtype
dt_bias (Optional[Tensor], default: None ) –

(n_heads,) float32, optional
initial_states (Optional[Tensor], default: None ) –

(batch, n_heads, d_head, d_state) float32, optional

Returns:

y ( Tensor ) –

(batch, seqlen, n_heads, d_head) float32
final_states ( Tensor ) –

(batch, n_heads, d_head, d_state) float32

SSD stages¶

tileops.mamba.SSDChunkCumsumFwdOp ¶

Mamba-2 dA_cumsum forward operator.

Applies optional per-head bias, optional softplus activation, and clamping to raw dt values, then computes the chunk-local inclusive prefix sum of dA = dt * A.

Note: dt_out is cast to the target dtype for storage efficiency, but dA_cumsum is computed from the fp32 dt values before casting, ensuring numerical precision.

init ¶

__init__(
    chunk_len,
    out_dtype=torch.float32,
    dt_softplus=False,
    dt_min=0.0,
    dt_max=float(
        "inf"
    ),
    *,
    target=None
)

Build the op. Shapes are taken from each call.

Parameters:

chunk_len (int) –

Tokens per chunk.
out_dtype (dtype, default: float32 ) –

Storage dtype of dt_out.
dt_softplus (bool, default: False ) –

Whether to apply softplus (with bypass for dt > 20) to dt.
dt_min (float, default: 0.0 ) –

Lower clamp bound applied after bias and softplus.
dt_max (float, default: float('inf') ) –

Upper clamp bound applied after bias and softplus.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    dt,
    A,
    dt_bias=None,
)

Run the dA_cumsum forward pass.

Parameters:

dt (Tensor) –

(batch, seq_len, n_heads) float32 — raw dt values.
A (Tensor) –

(n_heads,) float32 — SSM decay parameters.
dt_bias (Optional[Tensor], default: None ) –

(n_heads,) float32, optional — per-head dt bias.

Returns:

dt_out ( Tensor ) –

(batch, n_heads, num_chunks, chunk_len) out_dtype — processed dt.
dA_cumsum ( Tensor ) –

(batch, n_heads, num_chunks, chunk_len) float32 — inclusive prefix sum of dA = dt_val * A, computed from fp32 dt_val before casting dt_out.

tileops.mamba.SSDChunkCouplingFwdOp ¶

CB (C@B) matrix producer operator.

Computes cb[b,c,g,l,s] = sum_n C[b,cQ+l,g,n] * B[b,cQ+s,g,n] with causal masking (cb[l,s] = 0 if s > l).

init ¶

__init__(
    chunk_len,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

chunk_len (int) –

Chunk length (Q).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    C_mat,
    B_mat,
)

Parameters:

C_mat (Tensor) –

[B, S, G, N] dtype
B_mat (Tensor) –

[B, S, G, N] dtype

Returns:

cb ( Tensor ) –

[B, C, G, Q, Q] dtype

tileops.mamba.SSDChunkStateFwdOp ¶

Mamba-2 State-Space Dual (SSD) chunk state forward operator.

Computes the chunk-end State Space Model (SSM) state for each chunk:

out[b, c, h, p, n] = sum_{l=0}^{Q-1} x[b, cQ+l, h, p] * B[b, cQ+l, g(h), n] * exp(dA_cumsum[b,h,c,Q-1] - dA_cumsum[b,h,c,l]) * dt[b, h, c, l] * (1 if seq_idx is None else (seq_idx[b,cQ+Q-1] >= 0 and seq_idx[b,cQ+l] == seq_idx[b,c*Q+Q-1]))

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
    Bmat,
    dt,
    dA_cumsum,
    seq_idx=None,
)

Run the SSD chunk state forward pass.

Parameters:

x (Tensor) –

(batch, seq_len, n_heads, d_head)
Bmat (Tensor) –

(batch, seq_len, n_groups, d_state)
dt (Tensor) –

(batch, n_heads, num_chunks, chunk_len), x's dtype or float32
dA_cumsum (Tensor) –

(batch, n_heads, num_chunks, chunk_len) float32
seq_idx (Optional[Tensor], default: None ) –

(batch, seq_len) int32, optional

Returns:

states ( Tensor ) –

(batch, num_chunks, n_heads, d_head, d_state) float32

tileops.mamba.SSDStatePassingFwdOp ¶

Mamba-2 State-Space Dual (SSD) state passing forward operator.

Performs the inter-chunk recurrent scan:

s_c[m] = exp(dA_chunk_cumsum[b, h, c]) * s_{c-1}[m] + states[b, c, h, m]

with s_{-1} = initial_states, or 0 when it is not passed.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    states,
    dA_chunk_cumsum,
    initial_states=None,
)

Run the SSD state passing forward pass.

Parameters:

states (Tensor) –

(batch, num_chunks, n_heads, d_state)
dA_chunk_cumsum (Tensor) –

(batch, n_heads, num_chunks) float32
initial_states (Optional[Tensor], default: None ) –

(batch, n_heads, d_state) float32, optional

Returns:

prev_states ( Tensor ) –

(batch, num_chunks, n_heads, d_state) float32
final_states ( Tensor ) –

(batch, n_heads, d_state) float32

tileops.mamba.SSDChunkScanFwdOp ¶

Mamba-2 State-Space Dual (SSD) fused chunk output operator.

Fuses the history (prev_states) contribution and intra-chunk causal decay into a single pass, computing:

out[l, p] = exp(dA_cumsum[l]) * (C[l] @ prev_states) + sum_{s <= l} cb[l, s] * exp(dA_cumsum[l] - dA_cumsum[s]) * dt[s] * x[s, p]

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
    cb,
    dA_cumsum,
    C,
    prev_states,
    dt,
)

Run the fused SSD chunk output pass.

Parameters:

x (Tensor) –

(batch, seqlen, n_heads, d_head) dtype
cb (Tensor) –

(batch, num_chunks, n_groups, chunk_len, chunk_len) dtype
dA_cumsum (Tensor) –

(batch, n_heads, num_chunks, chunk_len) float32
C (Tensor) –

(batch, seqlen, n_groups, d_state) dtype
prev_states (Tensor) –

(batch, num_chunks, n_heads, d_head, d_state) float32
dt (Tensor) –

(batch, n_heads, num_chunks, chunk_len) dtype

Returns:

y ( Tensor ) –

(batch, seqlen, n_heads, d_head) float32

Decode¶

tileops.mamba.SSDRecurrentFwdOp ¶

Mamba-2 State-Space Dual (SSD) recurrent decode (step) operator.

Performs a single decode step of the Mamba-2 State Space Model (SSM) core: updates the recurrent state in-place and returns the output y for the current token:

g = h // (n_heads // n_groups) dA[b, h, p, n] = exp(dt[b, h, p] * A[h, p, n]) state[b,h,p,n] <- dA[b,h,p,n] * state[b,h,p,n] + dt[b,h,p] * B_in[b,g,n] * x[b,h,p] y_out[b, h, p] = sum_n state[b, h, p, n] * C_in[b, g, n]

The skip connection (D * x) and output gate (z * silu) are not fused here and must be applied by the caller if needed.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    A,
    dt,
    x,
    B_in,
    C_in,
    state,
)

Run a single Mamba-2 decode step.

Parameters:

A (Tensor) –

(n_heads, d_head, d_state) float32 -- SSM decay parameter (A <= 0)
dt (Tensor) –

(batch, n_heads, d_head) float32 -- discretization step (post-softplus)
x (Tensor) –

(batch, n_heads, d_head) dtype -- input features per head
B_in (Tensor) –

(batch, n_groups, d_state) dtype -- SSM B matrix (per group)
C_in (Tensor) –

(batch, n_groups, d_state) dtype -- SSM C matrix (per group)
state (Tensor) –

(batch, n_heads, d_head, d_state) float32 -- recurrent state (mutated in-place)

Returns:

y_out ( Tensor ) –

(batch, n_heads, d_head) float32

Mamba Operators¶

Full forward¶

tileops.mamba.Mamba2FwdOp ¶

__init__ ¶

forward ¶

SSD stages¶

tileops.mamba.SSDChunkCumsumFwdOp ¶

__init__ ¶

forward ¶

tileops.mamba.SSDChunkCouplingFwdOp ¶

__init__ ¶

forward ¶

tileops.mamba.SSDChunkStateFwdOp ¶

__init__ ¶

forward ¶

tileops.mamba.SSDStatePassingFwdOp ¶

__init__ ¶

forward ¶

tileops.mamba.SSDChunkScanFwdOp ¶

__init__ ¶

forward ¶

Decode¶

tileops.mamba.SSDRecurrentFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶