spark_model/layers/ple/
aux_state.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! PLE Marconi aux-state: serialize / restore the per-sequence lexical
4//! carry (token history + conv state) that rides the SSM snapshots.
5//! Split from `layer.rs` for the ≤500 LoC cap.
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9
10use super::{PleLayer, PleSeqState};
11use crate::layers::ple::ids::ple_ngram_ids;
12
13impl PleLayer {
14    /// Marconi aux blob: `[hist_len u32][history u32s][conv f32 bytes]`.
15    /// The whole per-sequence carry — a prefix hit restoring KV+SSM without
16    /// this would run the n-gram hash on the PREVIOUS request's history.
17    pub fn snapshot_aux(
18        &self,
19        st: &PleSeqState,
20        gpu: &dyn GpuBackend,
21        stream: u64,
22    ) -> Result<Vec<u8>> {
23        let conv_bytes = self.state_len * self.hc_mult * self.hidden * 4;
24        let mut blob = Vec::with_capacity(4 + st.history.len() * 4 + conv_bytes);
25        blob.extend_from_slice(&(st.history.len() as u32).to_le_bytes());
26        for t in &st.history {
27            blob.extend_from_slice(&t.to_le_bytes());
28        }
29        let off = blob.len();
30        blob.resize(off + conv_bytes, 0);
31        gpu.copy_d2h_on_stream(st.conv, &mut blob[off..], stream)?;
32        Ok(blob)
33    }
34
35    /// Restore the blob from [`Self::snapshot_aux`] on a prefix-cache hit.
36    pub fn restore_aux(
37        &self,
38        st: &mut PleSeqState,
39        blob: &[u8],
40        gpu: &dyn GpuBackend,
41        stream: u64,
42    ) -> Result<()> {
43        anyhow::ensure!(blob.len() >= 4, "PLE aux blob truncated");
44        let n = u32::from_le_bytes(blob[..4].try_into().unwrap()) as usize;
45        let conv_bytes = self.state_len * self.hc_mult * self.hidden * 4;
46        anyhow::ensure!(
47            blob.len() == 4 + n * 4 + conv_bytes,
48            "PLE aux blob size mismatch"
49        );
50        st.history = blob[4..4 + n * 4]
51            .chunks_exact(4)
52            .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
53            .collect();
54        st.prestaged_va = None;
55        gpu.copy_h2d_async(&blob[4 + n * 4..], st.conv, stream)?;
56        Ok(())
57    }
58}
59
60impl PleLayer {
61    /// Fresh sequence: EOS-filled history and a zeroed conv state.
62    pub(super) fn reset(
63        &self,
64        st: &mut PleSeqState,
65        gpu: &dyn GpuBackend,
66        stream: u64,
67    ) -> Result<()> {
68        st.history = vec![self.dims.eos_token_id; self.dims.context_len()];
69        st.prestaged_va = None;
70        let zeros = vec![0u8; self.state_len * self.hc_mult * self.hidden * 4];
71        gpu.copy_h2d_async(&zeros, st.conv, stream)?;
72        Ok(())
73    }
74
75    /// Hoisted per-step HOST work for decode under CUDA graphs: the n-gram
76    /// hash, the NVMe fault-in and the slot upload into the stable
77    /// `slots_dev` buffer. All three are capture-illegal (the upload reads
78    /// pageable memory, which invalidates a recording graph with status
79    /// 901), so the scheduler calls this BEFORE graph replay/capture — the
80    /// same phasing decode_a already gives the `token_ids` upload. `forward`
81    /// then consumes `prestaged_va` and enqueues only stable-buffer kernels.
82    ///
83    /// History advances HERE; the prestaged `forward` must not advance it
84    /// 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    /// Release one sequence's PLE carry.
109    ///
110    /// Same shape of defect as the QSA indexer carry: `conv` is a bare
111    /// `DevicePtr`, so dropping `PleSeqState` frees nothing. Individually
112    /// small (~147 KB) and below the 32 MB allocation-trace threshold, which
113    /// is exactly why it stayed invisible — but it is one per SSM layer (36
114    /// on qwen4_exp) per sequence, and it never comes back.
115    ///
116    /// Idempotent: `conv` is nulled once freed.
117    pub fn release_seq_state(&self, st: &mut PleSeqState, gpu: &dyn GpuBackend) -> Result<()> {
118        if st.conv.is_null() {
119            return Ok(());
120        }
121        let r = gpu.free(st.conv);
122        st.conv = DevicePtr(0);
123        st.history.clear();
124        st.prestaged_va = None;
125        st.last_staged_va = 0;
126        r
127    }
128}
spark_model/layers/ple/aux_state.rs

spark_model/layers/ple/
aux_state.rs
