selective_scan¶
SelectiveScan: a Mamba-1 mixer's values, at the selective-scan kernel call.
Bases: RecurrentMixer
A Mamba-1 mixer (Mamba, Falcon-Mamba, Jamba's Mamba blocks): a selective state-space scan.
Each channel of the mixer's inner width carries a state of state_dim
numbers. For each token t the mixer computes a step size per channel,
dt = softplus(dt_proj(x) + dt_bias), and two input-dependent vectors
B and C of state_dim numbers shared by every channel, and
runs, per channel::
h_t = exp(dt_t * A) * h_{t-1} + dt_t * B_t * x_t
y_t = C_t . h_t + D * x_t
with A a learned [channels, state_dim] matrix of negative rates.
The output is y gated by silu(z) and projected back by
out_proj. In the shared vocabulary: C reads the state as a
query does, B writes the input into it as a key does, x is what
gets written (the values), dt is how strongly each token writes
(betas), and dt * A is the log of how much of the state each
token keeps (decays, per channel and state dimension, where a
DeltaNet's is one number per head).
The values are read at the kernel call the forward makes:
mamba_selective_scan(x, dt, A, B, C, D=, z=, delta_bias=, ...) on a
prompt, mamba_selective_state_update(state, x, dt, A, B, C, D, z=,
dt_bias=, ...) on a decode step (one token). The kernel's tensors are
channel-first (x and dt [batch, channels, seq], B and
C [batch, state_dim, seq]; a decode step drops the sequence
axis); each value is a view of the argument laid out tokens first, and
an assignment is laid back out, so reads, writes and in-place edits
reach the kernel.
attention_queries/attention_keys:C/B,[batch, seq, 1, state_dim](one group).attention_values:x,[batch, seq, channels].betas:softplus(dt + dt_bias), the step size, anddecays:A * betas,[batch, seq, channels, state_dim]; derived, read-only.state_input: a decode step's cached state,[batch, channels, state_dim](a copy; assign to replace it),Noneon a prompt, whose scan starts from zeros.attention_head_outputs:ybefore the gate andout_proj,[batch, seq, channels]: a binding inside the kernel.state_output: the state after the call's last token, as the cache receives it: the scan's returned state on a prompt, the updated state the decode kernel copies into the cache.state/states/state_after/set_state_after: the state after each token, from the scan's token loop (STATE_OP) on a prompt and the single update (STEP_STATE_OP) on a decode step.
The pure-torch selective scan is the token loop, so
route_kernels(model.family, "torch") binds each kernel name to its
own pure-torch function. With mamba_ssm installed the forward
dispatches both to its compiled kernels, which have no source to read
inside (and need CUDA): route before the first trace.
class-attribute
instance-attribute
¶
class-attribute
instance-attribute
¶
ARGUMENTS = {'x': (0, 1), 'dt': (1, 2), 'A': (2, 3), 'B': (3, 4), 'C': (4, 5), 'dt_bias': ('delta_bias', 'dt_bias')}
class-attribute
instance-attribute
¶
betas = DerivedEProperty(_betas, description='The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only', unavailable=needs_torch_kernels)
class-attribute
instance-attribute
¶
decays = DerivedEProperty(_decays, description='The per-token log decay dt * A of the state; derived, read-only', unavailable=needs_torch_kernels)
class-attribute
instance-attribute
¶
states = DerivedEProperty(_scan_states, description="The state after every token of this call; needs route_kernels(family, 'torch')", unavailable=needs_token_loop)
state_output(value: Tensor) -> ScanState
The state after this call's last token, [batch, channels, state_dim]: what the next decode step starts from.
The scan's returned state on a prompt; on a decode step the updated state as the kernel copies it into the cache. Either way it is what the cache receives: assigning it changes what the next step starts from, not this call's output.
Why a value read inside the kernels' own bodies is unavailable: they are not bound to the pure-torch functions.
Why the state after each token is unavailable: the kernels are not routed, or the scan runs mambapy's parallel scan.