metrale_model_layers/layers/ple/
aux_state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The PLE per-sequence carry (token history and conv state):
4//! snapshot and restore for the layer's aux state, reset, the decode
5//! prestage, and release.
6//!
7//! Owner: model-layers (PLE).
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12
13use super::{PleLayer, PleSeqState};
14use crate::layers::ple::ids::ple_ngram_ids;
15
16impl PleLayer {
17    /// 2026-09-25: The aux blob: `hist_len` (u32), the history (u32 each),
18    /// then the conv state's f32 bytes, all little-endian.
19    pub fn snapshot_aux(
20        &self,
21        st: &PleSeqState,
22        gpu: &dyn GpuBackend,
23        stream: u64,
24    ) -> Result<Vec<u8>> {
25        let conv_bytes = self.state_len * self.hc_mult * self.hidden * 4;
26        let mut blob = Vec::with_capacity(4 + st.history.len() * 4 + conv_bytes);
27        blob.extend_from_slice(&(st.history.len() as u32).to_le_bytes());
28        for t in &st.history {
29            blob.extend_from_slice(&t.to_le_bytes());
30        }
31        let off = blob.len();
32        blob.resize(off + conv_bytes, 0);
33        gpu.copy_d2h_on_stream(st.conv, &mut blob[off..], stream)?;
34        Ok(blob)
35    }
36
37    /// 2026-09-25: Restore a blob from [`Self::snapshot_aux`] and clear
38    /// `prestaged_va`. Errors when the blob's size does not match.
39    pub fn restore_aux(
40        &self,
41        st: &mut PleSeqState,
42        blob: &[u8],
43        gpu: &dyn GpuBackend,
44        stream: u64,
45    ) -> Result<()> {
46        anyhow::ensure!(blob.len() >= 4, "PLE aux blob truncated");
47        let n = u32::from_le_bytes(blob[..4].try_into().unwrap()) as usize;
48        let conv_bytes = self.state_len * self.hc_mult * self.hidden * 4;
49        anyhow::ensure!(
50            blob.len() == 4 + n * 4 + conv_bytes,
51            "PLE aux blob size mismatch"
52        );
53        st.history = blob[4..4 + n * 4]
54            .chunks_exact(4)
55            .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
56            .collect();
57        st.prestaged_va = None;
58        gpu.copy_h2d_async(&blob[4 + n * 4..], st.conv, stream)?;
59        Ok(())
60    }
61}
62
63impl PleLayer {
64    /// 2026-09-25: Fresh sequence: EOS-filled history and a zeroed conv
65    /// state.
66    pub(super) fn reset(
67        &self,
68        st: &mut PleSeqState,
69        gpu: &dyn GpuBackend,
70        stream: u64,
71    ) -> Result<()> {
72        st.history = vec![self.dims.eos_token_id; self.dims.context_len()];
73        st.prestaged_va = None;
74        let zeros = vec![0u8; self.state_len * self.hc_mult * self.hidden * 4];
75        gpu.copy_h2d_async(&zeros, st.conv, stream)?;
76        Ok(())
77    }
78
79    /// 2026-09-25: The per-step host work of a decode step: the n-gram hash,
80    /// the row gather (`gather_host`, which uploads the slots into
81    /// `slots_dev` from host memory) and the history advance. It runs
82    /// through `decode_prestage` before any graph replay or capture; the
83    /// next `forward` consumes `prestaged_va` and does not advance the
84    /// history again.
85    pub fn prestage(
86        &self,
87        st: &mut PleSeqState,
88        tokens: &[u32],
89        gpu: &dyn GpuBackend,
90        stream: u64,
91    ) -> Result<()> {
92        if st.history.len() != self.dims.context_len() {
93            self.reset(st, gpu, stream)?;
94        }
95        let mut window = st.history.clone();
96        window.extend_from_slice(tokens);
97        let all = ple_ngram_ids(&self.dims, &window);
98        let rows = &all[all.len() - tokens.len()..];
99        let flat: Vec<u64> = rows.iter().flat_map(|r| r.iter().copied()).collect();
100        let va = self.gather_host(&flat, gpu, stream)?;
101        let keep = self.dims.context_len();
102        st.history = window[window.len() - keep..].to_vec();
103        st.prestaged_va = Some(va);
104        st.last_staged_va = va;
105        Ok(())
106    }
107
108    /// 2026-09-25: Free one sequence's PLE carry; `conv` is a bare
109    /// `DevicePtr`, so dropping `PleSeqState` frees nothing. Idempotent:
110    /// `conv` is nulled, and the state cleared, even when the free fails,
111    /// whose error is returned.
112    pub fn release_seq_state(&self, st: &mut PleSeqState, gpu: &dyn GpuBackend) -> Result<()> {
113        if st.conv.is_null() {
114            return Ok(());
115        }
116        let r = gpu.free(st.conv);
117        st.conv = DevicePtr(0);
118        st.history.clear();
119        st.prestaged_va = None;
120        st.last_staged_va = 0;
121        r
122    }
123}
metrale_model_layers/layers/ple/aux_state.rs

metrale_model_layers/layers/ple/
aux_state.rs
