Struct SsmLayerState 

pub struct SsmLayerState {
    pub h_state: DevicePtr,
    pub conv_state: DevicePtr,
    pub h_state_checkpoint: Option<DevicePtr>,
    pub conv_state_checkpoint: Option<DevicePtr>,
    pub h_state_intermediates: Vec<DevicePtr>,
    pub conv_state_intermediates: Vec<DevicePtr>,
    pub h_is_f16: bool,
    pub h_prefill_stage: Option<DevicePtr>,
    pub ple: Option<PleSeqState>,
}

SSM layer state: recurrent hidden state + conv1d sliding window.

Used by Mamba, Gated Delta Net (GDN), and similar recurrent layers.

Fields

h_state: DevicePtr

Recurrent hidden state: [num_v_heads, v_dim, k_dim] in f32.

conv_state: DevicePtr

Conv1d sliding window state: [d_inner, d_conv] in f32.

h_state_checkpoint: Option<DevicePtr>

Checkpoint buffer for h_state (allocated lazily for speculative decode).

conv_state_checkpoint: Option<DevicePtr>

Checkpoint buffer for conv_state (allocated lazily for speculative decode).

h_state_intermediates: Vec<DevicePtr>

Intermediate h_state snapshots during batched verification. Element i holds h_state after processing verification token i. Used by rollback_ssm_states to restore to the correct position.

conv_state_intermediates: Vec<DevicePtr>

Intermediate conv_state snapshots during batched verification.

h_is_f16: bool

Storage dtype of h_state: false = FP32, true = FP16 (--ssm-h-dtype f16).

This is the single source of truth for the h-state format. Which edge sets it depends on the POOL width:

• FP32-sized pool (stage 1/2, h_prefill_stage == None): prefill is the only FP32 writer and writes the slot in place, so the flag starts false and the decode mixer (TransformerModel::ssm_h_to_f16_dispatch) flips it exactly once per sequence, on the first decode step. No caller has to know where the prefill->decode edge is.
• f16-SIZED pool (stage 3, h_prefill_stage == Some): the slot is physically 2 bytes/element and can NEVER hold FP32, so the flag is true from allocation onwards and the decode mixer is a no-op. Prefill’s FP32 kernels run over Self::h_prefill_stage instead.

It rides through swap-out/swap-in because state_io mutates these states in place rather than rebuilding them.

h_prefill_stage: Option<DevicePtr>

Stage-3 f16-SIZED pool ONLY (--ssm-h-dtype f16-pool): the FP32 staging blob for THIS sequence’s slot, which the GDN prefill widens h_state into before its FP32 kernels run and narrows back after.

None — every configuration before stage 3 — means “the slot IS FP32-wide”: prefill writes h_state in place exactly as it always has, and not one byte moves. The same blob is shared by every layer of the sequence (see SsmStatePool::h_prefill_stage).

ple: Option<PleSeqState>

PLE per-sequence carry (n-gram history + dilated-conv state), present only on the layer that hosts a PleLayer (Avarok #753 item B: one per in-flight sequence, lazily created on the sequence’s first pass).

Trait Implementations

impl LayerState for SsmLayerState

fn as_any(&self) -> &dyn Any

fn as_any_mut(&mut self) -> &mut dyn Any