metrale_model_arch/glm5next_dsa/
aux_state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Snapshot and restore of the DSA indexer cache, the part of a GLM-5.3
4//! sequence that a KV-only prefix-cache hit cannot rebuild.
5//!
6//! Owner: model-arch (GLM-5.3 DSA).
7//! Invariants:
8//! - `restore_blob` checks the header, `index_head_dim`, total size and
9//!   capacity before any copy, and returns an error for any mismatch.
10//! - The cursor moves only after every copy is enqueued, so a failed restore
11//!   leaves `len` unchanged.
12//! - A blob is `HEADER_BYTES + len * (4 * index_head_dim + 1)` bytes: it
13//!   carries `len` rows, not `capacity`.
14//!
15//! # Why a blob and not a rewind
16//!
17//! `rewind_to` restores within a live sequence: the rows stay and only the
18//! cursor moves. A prefix-cache hit hands the prefix to a sequence whose
19//! indexer buffer was never written, and the rows cannot be rebuilt from the
20//! MLA latent: `k_normed` and `gate` are projections of the hidden state
21//! (`indexer.wk`, `index_kpool_compress_gate`). So the rows travel with the
22//! snapshot.
23//!
24//! `layers/qsa_snapshot.rs` in model-layers does the same for the QSA indexer,
25//! and re-pools its block keys on restore. The DSA state stores no pooled keys
26//! (pooling runs inside the selector kernel each step), so the blob is the
27//! whole reachable cache.
28//!
29//! Applying a mismatched blob would make the selector read another sequence's
30//! keys, which is why every mismatch is an error rather than a repair.
31//!
32//! # Cost
33//!
34//! `len * (4 * index_head_dim + 1)` bytes per DSA layer, 513 B per token and
35//! layer at `index_head_dim = 128`.
36
37use anyhow::{Result, ensure};
38use metrale_gpu_runtime::gpu::GpuBackend;
39
40use super::state::Glm5NextDsaState;
41
42/// 2026-09-25: `[len u64][index_head_dim u64]`, little-endian.
43const HEADER_BYTES: usize = 16;
44
45/// 2026-09-25: Bytes a blob occupies for `len` rows at `index_head_dim`: the
46/// per-row size of [`super::state::indexer_state_bytes`], over `len` rows, plus
47/// the header.
48fn blob_bytes(len: usize, index_head_dim: usize) -> usize {
49    HEADER_BYTES + len * (index_head_dim * 4 + 1)
50}
51
52impl Glm5NextDsaState {
53    /// 2026-09-25: Serialize the reachable indexer rows `[0, len)` for a
54    /// prefix-cache snapshot. Rows past `len` were never written or are
55    /// unreachable (`rewind_to` leaves them in place).
56    pub fn snapshot_blob(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<Vec<u8>> {
57        let len = self.len();
58        let d = self.index_head_dim();
59        let key_bytes = len * d * 2;
60
61        let mut blob = vec![0u8; blob_bytes(len, d)];
62        blob[..8].copy_from_slice(&(len as u64).to_le_bytes());
63        blob[8..16].copy_from_slice(&(d as u64).to_le_bytes());
64
65        if len > 0 {
66            let (k_off, g_off, v_off) = self.blob_offsets(len, d);
67            gpu.copy_d2h_on_stream(self.k_normed, &mut blob[k_off..k_off + key_bytes], stream)?;
68            gpu.copy_d2h_on_stream(self.gate, &mut blob[g_off..g_off + key_bytes], stream)?;
69            gpu.copy_d2h_on_stream(self.valid, &mut blob[v_off..v_off + len], stream)?;
70        }
71        Ok(blob)
72    }
73
74    /// 2026-09-25: Restore a snapshot's indexer rows into this state.
75    ///
76    /// A blob that is truncated, of another `index_head_dim`, inconsistent with
77    /// its header, or longer than this sequence's reservation is an `Err`;
78    /// model-engine's `apply_aux_states` propagates it with `?`, so the
79    /// prefix-cache hit fails.
80    pub fn restore_blob(&mut self, blob: &[u8], gpu: &dyn GpuBackend, stream: u64) -> Result<()> {
81        ensure!(
82            blob.len() >= HEADER_BYTES,
83            "DSA aux blob truncated: {} bytes, need at least {HEADER_BYTES} for the header",
84            blob.len()
85        );
86        let len = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize;
87        let d = u64::from_le_bytes(blob[8..16].try_into().unwrap()) as usize;
88
89        let want_d = self.index_head_dim();
90        ensure!(
91            d == want_d,
92            "DSA aux blob index_head_dim {d} != this layer's {want_d} — the snapshot was \
93             taken under a different model geometry"
94        );
95        ensure!(
96            blob.len() == blob_bytes(len, d),
97            "DSA aux blob size mismatch: {} bytes for {len} rows at head_dim {d}, expected {}",
98            blob.len(),
99            blob_bytes(len, d)
100        );
101        // 2026-09-25: Capacity comes from the configured context, so it can
102        // differ from the state that took the snapshot. A longer blob is refused,
103        // not truncated: a clamped `len` would select over a prefix while the MLA
104        // cache holds the full context.
105        self.ensure_room_through(len)?;
106
107        let key_bytes = len * d * 2;
108        if len > 0 {
109            let (k_off, g_off, v_off) = self.blob_offsets(len, d);
110            gpu.copy_h2d_async(&blob[k_off..k_off + key_bytes], self.k_normed, stream)?;
111            gpu.copy_h2d_async(&blob[g_off..g_off + key_bytes], self.gate, stream)?;
112            gpu.copy_h2d_async(&blob[v_off..v_off + len], self.valid, stream)?;
113        }
114        // 2026-09-25: The cursor moves last, through `rewind_to` and `advance`,
115        // so an early `?` above leaves `len` unchanged and `decode_k`'s lockstep
116        // check never sees a partial restore.
117        self.rewind_to(0)?;
118        self.advance(len)?;
119        Ok(())
120    }
121
122    /// 2026-09-25: `(k_normed, gate, valid)` byte offsets within a blob of `len`
123    /// rows.
124    fn blob_offsets(&self, len: usize, d: usize) -> (usize, usize, usize) {
125        let key_bytes = len * d * 2;
126        (
127            HEADER_BYTES,
128            HEADER_BYTES + key_bytes,
129            HEADER_BYTES + 2 * key_bytes,
130        )
131    }
132}
133
134#[cfg(test)]
135mod tests;
metrale_model_arch/glm5next_dsa/aux_state.rs

metrale_model_arch/glm5next_dsa/
aux_state.rs
