metrale_model_layers/layers/
qsa_snapshot.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Serialise and restore one sequence's QSA indexer state, the aux blob
4//! stored with an SSM prefix-cache snapshot.
5//!
6//! Owner: model-layers (QSA).
7//! Invariants:
8//! - `restore_aux` sets the counters only after the size checks pass and the
9//!   raw-key upload is issued; `pooled` stays 0 unless re-pooling was issued.
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::GpuBackend;
13
14use super::{QsaIndexer, QsaSeqState};
15use crate::layers::ops;
16
17impl QsaIndexer {
18    /// 2026-09-25: Aux blob: `[ingested u64][pooled u64][raw_keys bf16 bytes]`.
19    /// Block keys are not serialised: `restore_aux` re-pools them from the raw
20    /// keys with one `qsa_block_pool` launch.
21    pub fn snapshot_aux(
22        &self,
23        st: &QsaSeqState,
24        gpu: &dyn GpuBackend,
25        stream: u64,
26    ) -> Result<Vec<u8>> {
27        let hd = self.hd as usize;
28        let key_bytes = st.ingested * hd * 2;
29        let mut blob = Vec::with_capacity(16 + key_bytes);
30        blob.extend_from_slice(&(st.ingested as u64).to_le_bytes());
31        blob.extend_from_slice(&(st.pooled as u64).to_le_bytes());
32        let off = blob.len();
33        blob.resize(off + key_bytes, 0);
34        if key_bytes > 0 {
35            gpu.copy_d2h_on_stream(st.raw_keys, &mut blob[off..], stream)?;
36        }
37        Ok(blob)
38    }
39
40    /// 2026-09-25: Restore the blob from [`Self::snapshot_aux`] on a prefix-cache hit:
41    /// upload the raw keys, reset the counters, re-pool the block keys.
42    pub fn restore_aux(
43        &self,
44        st: &mut QsaSeqState,
45        blob: &[u8],
46        gpu: &dyn GpuBackend,
47        stream: u64,
48    ) -> Result<()> {
49        anyhow::ensure!(blob.len() >= 16, "QSA aux blob truncated");
50        let ingested = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize;
51        let pooled = u64::from_le_bytes(blob[8..16].try_into().unwrap()) as usize;
52        let hd = self.hd as usize;
53        anyhow::ensure!(
54            blob.len() == 16 + ingested * hd * 2,
55            "QSA aux blob size mismatch"
56        );
57        anyhow::ensure!(ingested <= self.max_tokens, "QSA aux exceeds key cache");
58        if ingested > 0 {
59            gpu.copy_h2d_async(&blob[16..], st.raw_keys, stream)?;
60        }
61        st.ingested = ingested;
62        st.pooled = 0;
63        if pooled > 0 {
64            ops::qsa_block_pool(
65                gpu,
66                self.k_pool_k,
67                st.raw_keys,
68                self.k_norm_w,
69                st.block_keys,
70                0,
71                pooled as u32,
72                self.ratio,
73                self.hd,
74                self.rot,
75                self.theta,
76                self.eps,
77                stream,
78            )?;
79            st.pooled = pooled;
80        }
81        Ok(())
82    }
83}
metrale_model_layers/layers/qsa_snapshot.rs

metrale_model_layers/layers/
qsa_snapshot.rs
