spark_model/layers/
qsa_snapshot.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! What the indexer must put back after a rejected draft.
4//!
5//! Split from `qsa.rs` on the 500-line cap, along the seam the pair already
6//! forms: everything else in that file computes a selection, and these two
7//! preserve and restore the state that computing it consumed. That is a
8//! different question, and the one speculative decoding gets wrong — a draft
9//! that is rejected must leave the indexer exactly as it found it, or the
10//! next step selects against a prefix that never happened.
11
12use anyhow::Result;
13use spark_runtime::gpu::GpuBackend;
14
15use super::{QsaIndexer, QsaSeqState};
16use crate::layers::ops;
17
18impl QsaIndexer {
19    /// Marconi aux blob: `[ingested u64][pooled u64][raw_keys bf16 bytes]`.
20    /// Raw keys are a deterministic function of the token prefix, so the
21    /// snapshot IS the indexer state; block keys are re-pooled on restore
22    /// (one kernel) rather than serialized.
23    pub fn snapshot_aux(
24        &self,
25        st: &QsaSeqState,
26        gpu: &dyn GpuBackend,
27        stream: u64,
28    ) -> Result<Vec<u8>> {
29        let hd = self.hd as usize;
30        let key_bytes = st.ingested * hd * 2;
31        let mut blob = Vec::with_capacity(16 + key_bytes);
32        blob.extend_from_slice(&(st.ingested as u64).to_le_bytes());
33        blob.extend_from_slice(&(st.pooled as u64).to_le_bytes());
34        let off = blob.len();
35        blob.resize(off + key_bytes, 0);
36        if key_bytes > 0 {
37            gpu.copy_d2h_on_stream(st.raw_keys, &mut blob[off..], stream)?;
38        }
39        Ok(blob)
40    }
41
42    /// Restore the blob from [`Self::snapshot_aux`] on a prefix-cache hit:
43    /// upload the raw keys, reset the counters, re-pool the block keys.
44    pub fn restore_aux(
45        &self,
46        st: &mut QsaSeqState,
47        blob: &[u8],
48        gpu: &dyn GpuBackend,
49        stream: u64,
50    ) -> Result<()> {
51        anyhow::ensure!(blob.len() >= 16, "QSA aux blob truncated");
52        let ingested = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize;
53        let pooled = u64::from_le_bytes(blob[8..16].try_into().unwrap()) as usize;
54        let hd = self.hd as usize;
55        anyhow::ensure!(
56            blob.len() == 16 + ingested * hd * 2,
57            "QSA aux blob size mismatch"
58        );
59        anyhow::ensure!(ingested <= self.max_tokens, "QSA aux exceeds key cache");
60        if ingested > 0 {
61            gpu.copy_h2d_async(&blob[16..], st.raw_keys, stream)?;
62        }
63        st.ingested = ingested;
64        st.pooled = 0;
65        if pooled > 0 {
66            ops::qsa_block_pool(
67                gpu,
68                self.k_pool_k,
69                st.raw_keys,
70                self.k_norm_w,
71                st.block_keys,
72                0,
73                pooled as u32,
74                self.ratio,
75                self.hd,
76                self.rot,
77                self.theta,
78                self.eps,
79                stream,
80            )?;
81            st.pooled = pooled;
82        }
83        Ok(())
84    }
85}
spark_model/layers/qsa_snapshot.rs

spark_model/layers/
qsa_snapshot.rs
