metrale_model_layers/layer/transformer_layer/
aux_state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `LayerAuxState`, the host-serialized per-sequence state that travels with an SSM
4//! snapshot.
5//!
6//! Owner: model-layers.
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::GpuBackend;
11
12use crate::layer::LayerState;
13
14/// 2026-09-26: A supertrait of `TransformerLayer`; see the module header.
15pub trait LayerAuxState {
16    /// 2026-09-25: Host-serialized per-sequence state that must travel with an SSM
17    /// snapshot for a prefix-cache hit to be complete, such as the PLE and QSA carries.
18    /// Device-to-host copies inside must be ordered on `stream`. Default `None`.
19    fn snapshot_aux(
20        &self,
21        _state: &dyn LayerState,
22        _gpu: &dyn GpuBackend,
23        _stream: u64,
24    ) -> Result<Option<Vec<u8>>> {
25        Ok(None)
26    }
27
28    /// 2026-09-25: True when this layer produces aux state. When any layer does, restore
29    /// sites decline a snapshot that has no aux blobs (`requires_aux_state`).
30    fn has_aux_state(&self) -> bool {
31        false
32    }
33
34    /// 2026-09-25: Apply a blob from [`Self::snapshot_aux`] to this layer's state
35    /// (`apply_aux_states`). The default returns an error.
36    fn restore_aux(
37        &self,
38        _state: &mut dyn LayerState,
39        _blob: &[u8],
40        _gpu: &dyn GpuBackend,
41        _stream: u64,
42    ) -> Result<()> {
43        anyhow::bail!("restore_aux on a layer with no aux state")
44    }
45}

metrale_model_layers/layer/transformer_layer/aux_state.rs

metrale_model_layers/layer/transformer_layer/
aux_state.rs