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

Mamba2SSDChunkScanVarlenFwd

struct Mamba2SSDChunkScanVarlenFwd[dt_softplus: Bool = True]

Varlen Mamba-2 SSD chunked-scan prefill forward.

Matches mamba_chunk_scan_combined semantics for the Nemotron-H NemotronHMamba2Mixer. Per-head scalar A, grouped B/C, per-head dt

• dt_bias softplus. State resets at each query_start_loc boundary (no cross-sequence bleed). Gating z + MambaRMSNormGated are applied OUTSIDE this op (norm_before_gate=False).

The registration lives here in the built-in kernel library (mirroring the gated_delta_conv1d_fwd / gated_delta_recurrence_fwd precedent) so the graph compiler / serve path can resolve the op with no out-of-tree custom_extensions. The kernel math lives in state_space.mamba2_ssd_scan (B200 / sm_100, bf16 in/out, fp32 states).

Tensor shapes (varlen / ragged; time dim is the packed total_len): - y: (total_len, nheads, head_dim) - output (dtype) - final_states: (batch, nheads, head_dim, dstate) - out, fp32 - x: (total_len, nheads, head_dim) - input (dtype) - dt: (total_len, nheads) - per-head time deltas (dtype) - A: (nheads,) - per-head scalar (dtype) - B: (total_len, ngroups, dstate) - grouped input proj (dtype) - C: (total_len, ngroups, dstate) - grouped output proj (dtype) - D: (nheads,) - skip connection (dtype, optional / empty) - dt_bias: (nheads,) - dt bias (dtype, optional / empty) - initial_states: (batch, nheads, head_dim, dstate) - in, fp32 (optional / empty) - query_start_loc: (batch + 1,) - cumulative sequence lengths (int32) - has_initial_state: (batch,) - whether to load initial_states (bool, optional / empty)

Parameters

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

Implemented traits

AnyType, ImplicitlyDeletable

Methods

execute

static def execute[dtype: DType, target: StringSlice[ImmStaticOrigin]](y: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=y.static_spec], final_states: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=final_states.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], initial_states: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=initial_states.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], ctx: DeviceContext)