Skip to main content

ClikaRT::ops::ssd_update

function

ssd_update()

Tensor ssd_update(
    Tensor x,
    Tensor dt,
    Tensor a_rate,
    Tensor b_mat,
    Tensor c_mat,
    OptionalTensor d_skip,
    OptionalTensor dt_bias,
    OptionalTensor z,
    Tensor state,
    OptionalTensor seq_lens = {},
    OptionalTensor slot_ids = {},
    bool dt_softplus = false
)

Mamba2 / SSD selective-state serving step over a per-sequence [H, dim, dstate] Float32 state (updated IN PLACE; the returned tensor is out [B, T, H, dim], dtype following x). Per token: the state decays by exp(dt'*A[h]) (dt' = softplus(dt + dt_bias[h]) when dt_softplus), accumulates dt'*(x (outer) B), and emits S*C + D[h]*x (optionally silu-gated by z). A/D/dt_bias are per-head Float32 parameters (A carries the NEGATIVE decay rate; fold -exp(A_log) at bind); B/C are [B, T, G, dstate] with H % G == 0. Decode is T == 1; a prefill runs the same sequential law. seq_lens ([B] Int32) bounds ragged rows. slot_ids ([B] Int32, device-resident) addresses state as a SLAB [num_slots, H, dim, dstate]: batch row b reads/updates slab row slot_ids[b] in place (ids in range and DISTINCT per call, the caller's contract); absent keeps state row b.

Declared in ClikaRT/compute/ops.h, line 3787