For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

Mamba2SSDChunkScanVarlenFwdInplace

struct Mamba2SSDChunkScanVarlenFwdInplace[dt_softplus: Bool = True]

Varlen Mamba-2 SSD chunked-scan: in-place SSM-pool write-back.

Identical to Mamba2SSDChunkScanVarlenFwd except final states are written directly into the ssm_pool buffer at ssm_pool[cache_indices[b], ...] instead of producing a separate final_states output tensor. This eliminates the graph-side buffer_load → gather → scatter_nd → buffer_store whole-pool RMW that otherwise dominates decode GPU time (~30 % wall-clock on B200).

The ssm_pool is declared as a MutableInputTensor (slot-indexed in/out), matching the causal_conv1d_varlen_fwd / gated_delta_recurrence_fwd precedent. initial_states is also read from ssm_pool when has_initial_state[b] is true (no separate initial-states input needed).

Tensor shapes: - y: (total_len, nheads, head_dim), output (dtype) [OUT] - x: (total_len, nheads, head_dim), input (dtype) - dt: (total_len, nheads) - A: (nheads,) - B: (total_len, ngroups, dstate) - C: (total_len, ngroups, dstate) - D: (nheads,) optional/empty - dt_bias: (nheads,) optional/empty - ssm_pool: (max_slots, nheads, head_dim, dstate) fp32, or bf16 on Apple GPUs (storage dtype only; the scan accumulates in fp32) [MUT] - query_start_loc: (batch + 1,) int32 - has_initial_state: (batch,) bool optional/empty - cache_indices: (batch,) uint32, slot indices into ssm_pool

Parameters

• ​dt_softplus (Bool): If True (default), apply softplus to dt + dt_bias.

Implemented traits

AnyType, ImplicitlyDeletable

Methods

execute

static def execute[dtype: DType, state_dtype: DType, target: StringSlice[ImmStaticOrigin]](y: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=y.static_spec], x: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=x.static_spec], dt: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt.static_spec], A: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=A.static_spec], B: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=B.static_spec], C: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=C.static_spec], D: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=D.static_spec], dt_bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=dt_bias.static_spec], ssm_pool: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=ssm_pool.static_spec], query_start_loc: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_start_loc.static_spec], has_initial_state: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=has_initial_state.static_spec], cache_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_indices.static_spec], ctx: DeviceContext)