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

Mojo struct

DStateVecLoader

struct DStateVecLoader[state_dtype: DType, kernel_dtype: DType, ssm_pool_LT: TensorLayout, B_LT: TensorLayout, C_LT: TensorLayout, Storage: TensorStorage, VEC: Int, NCHUNK: Int]

Owner of the Apple SSD kernel's vectorized dstate state / B / C I/O.

Target hardware family: Apple silicon GPU (Metal 4). The B200 ..._dstate_split sibling performs the identical widen-on-load / round-on-store / vectorized-vs-scalar-fallback logic inline; it COULD reuse this owner later (it is parameterized on the storage dtypes / layouts / Storage), but is not wired to it here (Apple-scoped change).

Every DRAM <-> register transition on the innermost (dstate) axis has an owner here instead of a raw_load / raw_store scattered across the kernel body -- the owner-per-transition pattern (see VarlenConvIO in varlen_causal_conv1d.mojo, Fp4WeightLoader in matmul2d_fp4.mojo, and new-primitives/amd-tile-io-expert-objects). Three verbs:

load_state: widen ssm_pool storage -> fp32 register chunks (the NCHUNK VEC-wide initial-state read).
load_bc: widen one VEC-lane B and C chunk -> fp32.
store_state: round the fp32 state chunks -> ssm_pool storage.

Each verb owns the VEC-wide aligned raw_load / raw_store and the scalar fallback, so the recurrence in the kernel body indexes only fp32 registers and issues no raw load/store. load_state / store_state each own their contig branch, which wraps the whole NCHUNK loop (a verbatim relocation of the prior inline loops). load_bc runs inside the fused recurrence loop, so it takes the contig decision as a COMPTIME contig parameter: the caller hoists the runtime bc_contig branch ONCE outside the chunk loop, keeping the fast path a straight-line unrolled vec-load loop (a per-chunk runtime branch measured ~18% slower on M5). Every method is @always_inline. The vectorized fast path (unit innermost stride) serves every row-major caller; a non-unit stride selects the exact v1 scalar gather/scatter, so the result is bit-identical for any layout.

Parameters

state_dtype (DType): ssm_pool storage dtype (fp32 or bf16).
kernel_dtype (DType): B / C element dtype.
ssm_pool_LT (TensorLayout): Layout type of the ssm_pool view.
B_LT (TensorLayout): Layout type of the B view.
C_LT (TensorLayout): Layout type of the C view.
Storage (TensorStorage): Shared tensor-storage policy of the views.
VEC (Int): SIMD load/store width over the contiguous dstate axis.
NCHUNK (Int): Number of VEC-wide chunks spanning DSTATE.

Fields

ssm_pool (TileTensor[state_dtype, ssm_pool_LT, MutUntrackedOrigin, Storage=Storage]):
B (TileTensor[kernel_dtype, B_LT, MutUntrackedOrigin, Storage=Storage]):
C (TileTensor[kernel_dtype, C_LT, MutUntrackedOrigin, Storage=Storage]):
pool_dstate_stride (Int):
b_dstate_stride (Int):
c_dstate_stride (Int):
pool_contig (Bool):
bc_contig (Bool):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

Methods

`init`

def __init__(out self, ssm_pool: TileTensor[state_dtype, ssm_pool_LT, MutUntrackedOrigin, Storage=Storage], B: TileTensor[kernel_dtype, B_LT, MutUntrackedOrigin, Storage=Storage], C: TileTensor[kernel_dtype, C_LT, MutUntrackedOrigin, Storage=Storage], pool_dstate_stride: Int, b_dstate_stride: Int, c_dstate_stride: Int)

`load_state`

def load_state(self, mut state: Array[SIMD[DType.float32, VEC], NCHUNK], pool_base: UInt32)

Fill fp32 state from ssm_pool[.., pool_base + n] (widening).

The load widens to fp32 (a no-op when state_dtype is fp32); the recurrence downstream runs on the fp32 register copy.

`load_bc`

def load_bc[c: Int, contig: Bool](self, b_base: UInt32, c_base: UInt32, mut b_c: SIMD[DType.float32, VEC], mut c_c: SIMD[DType.float32, VEC])

Load B and C chunk c (VEC dstate lanes each), widening to fp32.

contig is a COMPTIME parameter: the caller reads self.bc_contig and hoists the runtime branch ONCE outside the NCHUNK chunk loop, then instantiates the pure vectorized (contig=True) or pure scalar-gather (contig=False) path here. A per-chunk runtime branch interleaves the cold scalar-gather code into the hot vectorized loop and measured ~18% slower on M5 (33.6 vs 28.4 us/launch at the decode shape), so the fast path must stay a straight-line unrolled vec-load loop.

`store_state`

def store_state(self, state: Array[SIMD[DType.float32, VEC], NCHUNK], pool_wb: UInt32)

Round fp32 state to state_dtype and write back to ssm_pool.

The round happens once here at the final write-back (a no-op when state_dtype is fp32); the recurrent accumulator never leaves fp32.

Parameters​

Fields​

Implemented traits​

Methods​

__init__​

load_state​

load_bc​

store_state​