Skip to main content

ferrox_core/
kv_signature.rs

1//! Compatibility marking for stored KV blocks: whether a block found
2//! in a cache may be *used*, as opposed to merely found.
3//!
4//! A [`BlockHash`](crate::kv_block::BlockHash) says which tokens a block
5//! covers. It says nothing about the shape of the tensors stored under
6//! it -- how many layers, how many KV heads, what head dimension, what
7//! dtype, how many token positions are really there. A cache that
8//! survives a restart will eventually be read by a process configured
9//! differently from the one that wrote it, and reading a block whose
10//! layout does not match is not a miss: it is silently wrong attention
11//! state, which produces confident wrong tokens.
12//!
13//! The one rule this module exists to enforce:
14//!
15//! > **A signature is derived from the block's own payload, never from
16//! > the manager's expectation, and an unmarked block is rejected
17//! > rather than trusted.**
18//!
19//! So [`CacheSignature::from_payload`] measures the tensors; it takes no
20//! "expected" shape to fill gaps with. [`UnverifiedBlock::verify`] then
21//! makes three separate checks, in order:
22//!
23//! 1. the block carries a signature at all ([`SignatureError::Unmarked`]
24//!    otherwise -- absence is never treated as agreement);
25//! 2. the recorded signature matches what the payload actually is
26//!    ([`SignatureError::PayloadMismatch`]) -- a stamp may not vouch for
27//!    a width the payload does not have;
28//! 3. only then, that this shape is the shape the reader wants
29//!    ([`SignatureError::Incompatible`]).
30//!
31//! Step 2 is the one that is easy to skip and expensive to skip: without
32//! it, "the stamp says 32 heads" and "there are 32 heads" are different
33//! claims that a cache would be treating as one.
34//!
35//! # The block layout is part of the signature
36//!
37//! A signature also carries the [`BlockLayout`] the block was cut
38//! under: its block size, and the sliding window that block size had to
39//! divide (see [`kv_swa`](crate::kv_swa) for why it must). Both are
40//! here rather than checked once at startup because a *durable* cache
41//! outlives the configuration that filled it. A block written by a
42//! build running gpt-oss at window 128 must not be handed to a build
43//! running it at window 256, and the failure mode if it were is not a
44//! crash -- it is an attention mask silently wider or narrower than the
45//! model's.
46//!
47//! `block_size` is payload-checkable and is checked: a stored block is
48//! exactly one whole block (`kv_block::chain` refuses to name a partial
49//! tail), so a stamp claiming `block_size = 64` over a 48-token payload
50//! is refused the same way a stamp claiming 32 heads over 16 is. The
51//! window is not provable from tensors -- like `model` -- so it is
52//! carried and compared against the reader's expectation.
53
54use crate::cache::KvCache;
55use crate::kv_swa::{BlockLayout, BlockLayoutError};
56
57/// Layout version of a stored block payload. A reader accepts only the
58/// versions in [`READABLE_FORMAT_VERSIONS`]; anything else is rejected
59/// rather than guessed at.
60pub const BLOCK_FORMAT_VERSION: u32 = 2;
61
62/// Versions this build can read. Kept explicit (rather than `<=
63/// BLOCK_FORMAT_VERSION`) so dropping support for an old layout is a
64/// deliberate edit and not an accident of arithmetic.
65///
66/// Version 1 is deliberately **not** readable. Its header had no block
67/// size and no sliding window, so a v1 block cannot say what layout it
68/// was cut under -- and this module's rule is that absence is never
69/// agreement. Reading one would mean assuming it happened to be
70/// aligned, which is the exact assumption `kv_swa` exists to refuse.
71pub const READABLE_FORMAT_VERSIONS: &[u32] = &[2];
72
73/// Element type of the stored K/V tensors.
74///
75/// Only `F32` exists today, because [`KvCache`] stores `Vec<f32>`. The
76/// enum is here so a future f16/quantized KV tier changes the signature
77/// -- and therefore invalidates blocks written by an f32 build -- rather
78/// than reinterpreting their bytes.
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
80pub enum KvDtype {
81    F32,
82}
83
84impl KvDtype {
85    pub fn as_str(self) -> &'static str {
86        match self {
87            KvDtype::F32 => "f32",
88        }
89    }
90}
91
92/// What a block's payload *is*: the shape any reader must match to use
93/// it. Construct it from a payload with [`Self::from_payload`], or as a
94/// reader's requirement with [`Self::expected`].
95#[derive(Clone, Debug, PartialEq, Eq)]
96pub struct CacheSignature {
97    pub format_version: u32,
98    /// Identifies the weights the KV state was computed under. The one
99    /// field no payload can prove about itself -- which is exactly why
100    /// it is checked against the reader's expectation in step 3.
101    pub model: String,
102    pub n_layers: usize,
103    pub n_kv_heads: usize,
104    pub head_dim: usize,
105    pub dtype: KvDtype,
106    /// Token positions actually stored, per layer. Equal to
107    /// `layout.block_size()` for any block this module will stamp or
108    /// verify.
109    pub tokens: usize,
110    /// How the sequence was cut into blocks, and the sliding window
111    /// that cut had to line up with. See the module note.
112    pub layout: BlockLayout,
113}
114
115impl CacheSignature {
116    /// Derives a signature by *measuring* `layers`. There is
117    /// deliberately no parameter to fill a gap from: every shape field
118    /// except `model` and the sliding window comes from the tensors
119    /// themselves, and even the declared `layout`'s block size is
120    /// checked against the depth the tensors actually have.
121    ///
122    /// Fails if the payload cannot describe itself coherently: no
123    /// layers at all, layers that disagree with each other, a layer
124    /// whose buffers do not match its own declared shape, or a token
125    /// depth that is not the block size the layout claims.
126    pub fn from_payload(
127        model: &str,
128        layout: BlockLayout,
129        layers: &[KvCache],
130    ) -> Result<Self, SignatureError> {
131        let first = layers.first().ok_or(SignatureError::EmptyPayload)?;
132        let n_kv_heads = first.n_kv_heads;
133        let head_dim = first.head_dim;
134        if n_kv_heads == 0 || head_dim == 0 {
135            return Err(SignatureError::DegenerateLayer {
136                layer: 0,
137                n_kv_heads,
138                head_dim,
139            });
140        }
141        let per_token = n_kv_heads * head_dim;
142        let tokens = measure_layer(0, first, per_token)?;
143
144        for (index, layer) in layers.iter().enumerate().skip(1) {
145            if layer.n_kv_heads != n_kv_heads || layer.head_dim != head_dim {
146                return Err(SignatureError::RaggedPayload {
147                    layer: index,
148                    field: "layer shape",
149                    expected: format!("{n_kv_heads}x{head_dim}"),
150                    found: format!("{}x{}", layer.n_kv_heads, layer.head_dim),
151                });
152            }
153            let layer_tokens = measure_layer(index, layer, per_token)?;
154            if layer_tokens != tokens {
155                return Err(SignatureError::RaggedPayload {
156                    layer: index,
157                    field: "token count",
158                    expected: tokens.to_string(),
159                    found: layer_tokens.to_string(),
160                });
161            }
162        }
163
164        // The block size is a claim like any other, and this one the
165        // payload can settle: a stored block is one whole block.
166        if tokens != layout.block_size() {
167            return Err(SignatureError::BlockSizeMismatch {
168                block_size: layout.block_size(),
169                tokens,
170            });
171        }
172
173        Ok(CacheSignature {
174            format_version: BLOCK_FORMAT_VERSION,
175            model: model.to_string(),
176            n_layers: layers.len(),
177            n_kv_heads,
178            head_dim,
179            dtype: KvDtype::F32,
180            tokens,
181            layout,
182        })
183    }
184
185    /// A reader's requirement: the shape this process would compute
186    /// itself. Never stamped onto a block -- only compared against one.
187    pub fn expected(
188        model: &str,
189        layout: BlockLayout,
190        n_layers: usize,
191        n_kv_heads: usize,
192        head_dim: usize,
193        tokens: usize,
194    ) -> Self {
195        CacheSignature {
196            format_version: BLOCK_FORMAT_VERSION,
197            model: model.to_string(),
198            n_layers,
199            n_kv_heads,
200            head_dim,
201            dtype: KvDtype::F32,
202            tokens,
203            layout,
204        }
205    }
206
207    /// Field-by-field comparison, naming the first field that differs
208    /// so an operator learns *what* changed rather than "cache miss".
209    fn compare(
210        &self,
211        other: &CacheSignature,
212        mismatch: fn(&'static str, String, String) -> SignatureError,
213    ) -> Result<(), SignatureError> {
214        if self.format_version != other.format_version {
215            return Err(mismatch(
216                "format_version",
217                self.format_version.to_string(),
218                other.format_version.to_string(),
219            ));
220        }
221        if self.model != other.model {
222            return Err(mismatch("model", self.model.clone(), other.model.clone()));
223        }
224        if self.n_layers != other.n_layers {
225            return Err(mismatch(
226                "n_layers",
227                self.n_layers.to_string(),
228                other.n_layers.to_string(),
229            ));
230        }
231        if self.n_kv_heads != other.n_kv_heads {
232            return Err(mismatch(
233                "n_kv_heads",
234                self.n_kv_heads.to_string(),
235                other.n_kv_heads.to_string(),
236            ));
237        }
238        if self.head_dim != other.head_dim {
239            return Err(mismatch(
240                "head_dim",
241                self.head_dim.to_string(),
242                other.head_dim.to_string(),
243            ));
244        }
245        if self.dtype != other.dtype {
246            return Err(mismatch(
247                "dtype",
248                self.dtype.as_str().to_string(),
249                other.dtype.as_str().to_string(),
250            ));
251        }
252        if self.tokens != other.tokens {
253            return Err(mismatch(
254                "tokens",
255                self.tokens.to_string(),
256                other.tokens.to_string(),
257            ));
258        }
259        if self.layout.block_size() != other.layout.block_size() {
260            return Err(mismatch(
261                "block_size",
262                self.layout.block_size().to_string(),
263                other.layout.block_size().to_string(),
264            ));
265        }
266        if self.layout.sliding_window() != other.layout.sliding_window() {
267            return Err(mismatch(
268                "sliding_window",
269                describe_window(self.layout.sliding_window()),
270                describe_window(other.layout.sliding_window()),
271            ));
272        }
273        Ok(())
274    }
275}
276
277/// Renders a window for an error message. `None` is spelled out rather
278/// than printed as an empty string: "this build uses no sliding window"
279/// and "this build did not say" must not look the same in a log.
280fn describe_window(window: Option<usize>) -> String {
281    match window {
282        Some(w) => w.to_string(),
283        None => "none (full causal)".to_string(),
284    }
285}
286
287/// Measures one layer, rejecting a layer whose buffers disagree with
288/// its own declared `seq_len` -- `seq_len` is a claim, `k.len()` is the
289/// evidence.
290fn measure_layer(index: usize, layer: &KvCache, per_token: usize) -> Result<usize, SignatureError> {
291    if layer.v_head_dim != layer.head_dim {
292        return Err(SignatureError::SplitKvHeadWidth {
293            layer: index,
294            head_dim: layer.head_dim,
295            v_head_dim: layer.v_head_dim,
296        });
297    }
298    if !layer.k.len().is_multiple_of(per_token) {
299        return Err(SignatureError::RaggedPayload {
300            layer: index,
301            field: "k length",
302            expected: format!("a multiple of {per_token}"),
303            found: layer.k.len().to_string(),
304        });
305    }
306    if layer.v.len() != layer.k.len() {
307        return Err(SignatureError::RaggedPayload {
308            layer: index,
309            field: "v length",
310            expected: layer.k.len().to_string(),
311            found: layer.v.len().to_string(),
312        });
313    }
314    let tokens = layer.k.len() / per_token;
315    // Positions against rows. They are equal for anything this codec
316    // can currently produce, and a payload where they disagree is
317    // ragged rather than windowed. When a store learns to evict (#61)
318    // a restored windowed layer will legitimately carry more positions
319    // than rows, and THIS CHECK IS WHERE THAT HAS TO BE TAUGHT: the
320    // codec will need to serialize the position count separately
321    // instead of deriving it from the byte length.
322    if layer.positions() != tokens {
323        return Err(SignatureError::RaggedPayload {
324            layer: index,
325            field: "positions",
326            expected: tokens.to_string(),
327            found: layer.positions().to_string(),
328        });
329    }
330    Ok(tokens)
331}
332
333/// Why a stored block was refused. Every variant is a refusal to guess.
334#[derive(Clone, Debug, PartialEq, Eq)]
335pub enum SignatureError {
336    /// The block carries no signature. Not treated as "probably fine":
337    /// an unmarked block was written by something whose layout is
338    /// unknown, which is precisely the case that must not be trusted.
339    Unmarked,
340    /// A block with no layers describes nothing and can vouch for
341    /// nothing.
342    EmptyPayload,
343    /// A layer with no heads or zero head dimension.
344    DegenerateLayer {
345        layer: usize,
346        n_kv_heads: usize,
347        head_dim: usize,
348    },
349    /// A layer whose V head width differs from its K head width
350    /// (MiMo-V2). The block format carries ONE `head_dim`, so such a
351    /// payload has no honest encoding; refused rather than written with
352    /// the K width and read back into V rows of the wrong shape.
353    SplitKvHeadWidth {
354        layer: usize,
355        head_dim: usize,
356        v_head_dim: usize,
357    },
358    /// The payload does not agree with itself: layers of different
359    /// shapes or lengths, or a layer whose buffers contradict its own
360    /// `seq_len`.
361    RaggedPayload {
362        layer: usize,
363        field: &'static str,
364        expected: String,
365        found: String,
366    },
367    /// The recorded signature claims something the payload is not. The
368    /// stamp is wrong (or was written by a build with a different
369    /// layout); the payload is the truth.
370    PayloadMismatch {
371        field: &'static str,
372        recorded: String,
373        actual: String,
374    },
375    /// The payload is coherent and honestly stamped, but it is not what
376    /// this reader needs -- a different model, a config change, a
377    /// different block size.
378    Incompatible {
379        field: &'static str,
380        expected: String,
381        found: String,
382    },
383    /// The stamp claims a block size the payload's token depth is not.
384    /// A stored block is exactly one whole block, so these are the same
385    /// number or the stamp is wrong.
386    BlockSizeMismatch { block_size: usize, tokens: usize },
387    /// The block layout itself is not usable -- most importantly, a
388    /// block size that does not divide the sliding window. See
389    /// [`kv_swa`](crate::kv_swa).
390    BadLayout(BlockLayoutError),
391    /// Written by a build whose payload layout this one cannot read.
392    UnsupportedFormat {
393        found: u32,
394        readable: &'static [u32],
395    },
396}
397
398impl From<BlockLayoutError> for SignatureError {
399    fn from(err: BlockLayoutError) -> Self {
400        SignatureError::BadLayout(err)
401    }
402}
403
404impl std::fmt::Display for SignatureError {
405    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
406        match self {
407            SignatureError::Unmarked => write!(
408                f,
409                "KV block carries no cache signature; refusing to trust an unmarked block"
410            ),
411            SignatureError::EmptyPayload => {
412                write!(f, "KV block has no layers; nothing to verify")
413            }
414            SignatureError::DegenerateLayer {
415                layer,
416                n_kv_heads,
417                head_dim,
418            } => write!(
419                f,
420                "KV block layer {layer} is degenerate: {n_kv_heads} kv heads x {head_dim} head dim"
421            ),
422            SignatureError::SplitKvHeadWidth {
423                layer,
424                head_dim,
425                v_head_dim,
426            } => write!(
427                f,
428                "KV block layer {layer} has a V head width ({v_head_dim}) that differs from its K \
429                 head width ({head_dim}); the block format carries one head_dim and cannot \
430                 encode it"
431            ),
432            SignatureError::RaggedPayload {
433                layer,
434                field,
435                expected,
436                found,
437            } => write!(
438                f,
439                "KV block payload is inconsistent at layer {layer}: {field} is {found}, expected {expected}"
440            ),
441            SignatureError::PayloadMismatch {
442                field,
443                recorded,
444                actual,
445            } => write!(
446                f,
447                "KV block signature vouches for {field}={recorded} but its payload has {field}={actual}"
448            ),
449            SignatureError::Incompatible {
450                field,
451                expected,
452                found,
453            } => write!(
454                f,
455                "KV block is incompatible: {field} is {found}, this server needs {expected}"
456            ),
457            SignatureError::BlockSizeMismatch { block_size, tokens } => write!(
458                f,
459                "KV block signature declares a block size of {block_size} but its payload holds \
460                 {tokens} token positions; a stored block is exactly one whole block"
461            ),
462            SignatureError::BadLayout(err) => write!(f, "KV block layout is unusable: {err}"),
463            SignatureError::UnsupportedFormat { found, readable } => write!(
464                f,
465                "KV block format version {found} is not readable by this build (readable: {readable:?})"
466            ),
467        }
468    }
469}
470
471impl std::error::Error for SignatureError {}
472
473/// A block whose signature has been verified against its own payload
474/// and against the reader's expectation. Only way to get one is
475/// [`KvBlock::stamp`] (writing) or [`UnverifiedBlock::verify`]
476/// (reading), so holding one is itself the proof.
477pub struct KvBlock {
478    signature: CacheSignature,
479    layers: Vec<KvCache>,
480}
481
482/// Summarizes rather than dumping tensors: a block's `Debug` is for a
483/// log line or a failing assertion, and printing every f32 in a KV
484/// block helps nobody.
485impl std::fmt::Debug for KvBlock {
486    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
487        f.debug_struct("KvBlock")
488            .field("signature", &self.signature)
489            .field("layers", &self.layers.len())
490            .finish()
491    }
492}
493
494impl std::fmt::Debug for UnverifiedBlock {
495    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
496        f.debug_struct("UnverifiedBlock")
497            .field("signature", &self.signature)
498            .field("layers", &self.layers.len())
499            .finish()
500    }
501}
502
503impl KvBlock {
504    /// Stamps a block from its own payload. The caller supplies the
505    /// model identity and the layout the block was cut under; every
506    /// shape field is measured, and the layout's block size is checked
507    /// against the depth measured.
508    pub fn stamp(
509        model: &str,
510        layout: BlockLayout,
511        layers: Vec<KvCache>,
512    ) -> Result<Self, SignatureError> {
513        let signature = CacheSignature::from_payload(model, layout, &layers)?;
514        Ok(KvBlock { signature, layers })
515    }
516
517    /// The layout this block was cut under.
518    pub fn layout(&self) -> BlockLayout {
519        self.signature.layout
520    }
521
522    pub fn signature(&self) -> &CacheSignature {
523        &self.signature
524    }
525
526    pub fn tokens(&self) -> usize {
527        self.signature.tokens
528    }
529
530    pub fn layers(&self) -> &[KvCache] {
531        &self.layers
532    }
533
534    pub fn into_layers(self) -> Vec<KvCache> {
535        self.layers
536    }
537}
538
539/// A block as it comes back from an untrusted source: a disk tier, a
540/// shared cache, or a build that is not this one. `signature` is an
541/// `Option` on purpose -- "no signature recorded" is a state a real
542/// stored block can be in, and it must be representable so it can be
543/// refused.
544pub struct UnverifiedBlock {
545    pub signature: Option<CacheSignature>,
546    pub layers: Vec<KvCache>,
547}
548
549impl UnverifiedBlock {
550    pub fn new(signature: Option<CacheSignature>, layers: Vec<KvCache>) -> Self {
551        UnverifiedBlock { signature, layers }
552    }
553
554    /// Verifies in the order the module doc describes: marked, honest,
555    /// then compatible. `expected` is used only in the last step -- it
556    /// never contributes a field to the signature being checked.
557    pub fn verify(self, expected: &CacheSignature) -> Result<KvBlock, SignatureError> {
558        let recorded = self.signature.ok_or(SignatureError::Unmarked)?;
559        if !READABLE_FORMAT_VERSIONS.contains(&recorded.format_version) {
560            return Err(SignatureError::UnsupportedFormat {
561                found: recorded.format_version,
562                readable: READABLE_FORMAT_VERSIONS,
563            });
564        }
565        // Measured from the payload. `recorded.model` is carried over
566        // because no payload can prove which weights produced it; the
567        // model check happens against `expected` below, where a wrong
568        // model is caught as an incompatibility.
569        // `recorded.layout` is carried over for the same reason
570        // `recorded.model` is: the reader's expectation must not get a
571        // vote before the payload has been checked against the stamp.
572        // Carrying it is not trusting it -- `from_payload` rejects a
573        // block size the token depth contradicts, and the window is
574        // settled against `expected` below.
575        let actual = CacheSignature::from_payload(&recorded.model, recorded.layout, &self.layers)?;
576        recorded.compare(&actual, |field, recorded, actual| {
577            SignatureError::PayloadMismatch {
578                field,
579                recorded,
580                actual,
581            }
582        })?;
583        expected.compare(&actual, |field, expected, found| {
584            SignatureError::Incompatible {
585                field,
586                expected,
587                found,
588            }
589        })?;
590        Ok(KvBlock {
591            signature: actual,
592            layers: self.layers,
593        })
594    }
595}
596
597#[cfg(test)]
598mod tests {
599    use super::*;
600
601    fn layer(n_kv_heads: usize, head_dim: usize, tokens: usize) -> KvCache {
602        let mut cache = KvCache::new(n_kv_heads, head_dim);
603        let step = vec![0.5f32; n_kv_heads * head_dim];
604        for _ in 0..tokens {
605            cache.push(&step, &step).expect("unpooled push cannot fail");
606        }
607        cache
608    }
609
610    fn payload(n_layers: usize, n_kv_heads: usize, head_dim: usize, tokens: usize) -> Vec<KvCache> {
611        (0..n_layers)
612            .map(|_| layer(n_kv_heads, head_dim, tokens))
613            .collect()
614    }
615
616    /// A full-causal layout whose block size is the payload depth --
617    /// what every test that is not about SWA wants.
618    fn flat(block_size: usize) -> BlockLayout {
619        BlockLayout::full_attention(block_size).expect("positive block size")
620    }
621
622    /// The core rule, in its positive form: every shape field comes
623    /// from the tensors. `stamp` is given a model name and nothing else.
624    #[test]
625    fn signature_is_measured_from_the_payload() {
626        let block = KvBlock::stamp("model-a", flat(4), payload(3, 2, 8, 4)).expect("stamp");
627        let sig = block.signature();
628        assert_eq!(sig.n_layers, 3);
629        assert_eq!(sig.n_kv_heads, 2);
630        assert_eq!(sig.head_dim, 8);
631        assert_eq!(sig.tokens, 4);
632        assert_eq!(sig.dtype, KvDtype::F32);
633        assert_eq!(sig.format_version, BLOCK_FORMAT_VERSION);
634        assert_eq!(block.tokens(), 4);
635        assert_eq!(block.layers().len(), 3);
636    }
637
638    #[test]
639    fn a_stamped_block_round_trips_through_verification() {
640        let layers = payload(3, 2, 8, 4);
641        let signature =
642            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
643        let expected = CacheSignature::expected("model-a", flat(4), 3, 2, 8, 4);
644        let block = UnverifiedBlock::new(Some(signature), layers)
645            .verify(&expected)
646            .expect("a block that is what it says it is must verify");
647        assert_eq!(block.layers().len(), 3);
648        assert_eq!(block.into_layers().len(), 3);
649    }
650
651    /// Absence of a signature is not agreement. This is the difference
652    /// between a persistent cache and silent corruption after a config
653    /// change: an unmarked block came from something whose layout is
654    /// unknown by definition.
655    #[test]
656    fn an_unmarked_block_is_rejected_not_trusted() {
657        let expected = CacheSignature::expected("model-a", flat(4), 3, 2, 8, 4);
658        let err = UnverifiedBlock::new(None, payload(3, 2, 8, 4))
659            .verify(&expected)
660            .expect_err("an unmarked block must be refused");
661        assert_eq!(err, SignatureError::Unmarked);
662    }
663
664    /// The rule the plan states as "a signature must never vouch for a
665    /// width the payload does not have". Here the stamp is exactly what
666    /// the reader expects -- and the payload is not. A cache that
667    /// trusted the stamp (or, equivalently, stamped from the manager's
668    /// expectation) would hand back tensors of the wrong width and
669    /// produce confident wrong tokens.
670    #[test]
671    fn a_signature_that_overstates_its_payload_is_rejected() {
672        let expected = CacheSignature::expected("model-a", flat(4), 3, 2, 16, 4);
673        let mut lying = expected.clone();
674        assert_eq!(lying.head_dim, 16);
675        let err = UnverifiedBlock::new(Some(lying.clone()), payload(3, 2, 8, 4))
676            .verify(&expected)
677            .expect_err("stamp claims head_dim 16 over an 8-wide payload");
678        assert_eq!(
679            err,
680            SignatureError::PayloadMismatch {
681                field: "head_dim",
682                recorded: "16".into(),
683                actual: "8".into(),
684            }
685        );
686
687        // Same shape, overstated depth: 8 token positions claimed over
688        // a 4-position payload.
689        lying.head_dim = 8;
690        lying.tokens = 8;
691        let expected = CacheSignature::expected("model-a", flat(8), 3, 2, 8, 8);
692        let err = UnverifiedBlock::new(Some(lying.clone()), payload(3, 2, 8, 4))
693            .verify(&expected)
694            .expect_err("stamp claims 8 tokens over a 4-token payload");
695        assert_eq!(
696            err,
697            SignatureError::PayloadMismatch {
698                field: "tokens",
699                recorded: "8".into(),
700                actual: "4".into(),
701            }
702        );
703
704        // And overstated layer count.
705        lying.tokens = 4;
706        lying.n_layers = 4;
707        let expected = CacheSignature::expected("model-a", flat(4), 4, 2, 8, 4);
708        let err = UnverifiedBlock::new(Some(lying), payload(3, 2, 8, 4))
709            .verify(&expected)
710            .expect_err("stamp claims 4 layers over a 3-layer payload");
711        assert_eq!(
712            err,
713            SignatureError::PayloadMismatch {
714                field: "n_layers",
715                recorded: "4".into(),
716                actual: "3".into(),
717            }
718        );
719    }
720
721    /// A payload that lies to itself: `seq_len` says 4, the buffers
722    /// hold 3. `seq_len` is a claim; `k.len()` is the evidence.
723    #[test]
724    fn a_layer_whose_seq_len_contradicts_its_buffers_is_rejected() {
725        let mut layers = payload(2, 2, 8, 4);
726        // The contradiction is the thing under test.
727        layers[1].force_positions_for_test(7);
728        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
729            .expect_err("seq_len must be verified, not believed");
730        assert_eq!(
731            err,
732            SignatureError::RaggedPayload {
733                layer: 1,
734                field: "positions",
735                expected: "4".into(),
736                found: "7".into(),
737            }
738        );
739    }
740
741    #[test]
742    fn a_ragged_payload_is_rejected() {
743        let mut layers = payload(3, 2, 8, 4);
744        layers[2] = layer(2, 4, 4);
745        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
746            .expect_err("shape disagreement");
747        assert!(matches!(
748            err,
749            SignatureError::RaggedPayload {
750                layer: 2,
751                field: "layer shape",
752                ..
753            }
754        ));
755
756        let mut layers = payload(3, 2, 8, 4);
757        layers[1] = layer(2, 8, 3);
758        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
759            .expect_err("depth disagreement");
760        assert!(matches!(
761            err,
762            SignatureError::RaggedPayload {
763                layer: 1,
764                field: "token count",
765                ..
766            }
767        ));
768
769        let mut layers = payload(2, 2, 8, 4);
770        layers[0].v.truncate(8);
771        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
772            .expect_err("k/v disagreement");
773        assert!(matches!(
774            err,
775            SignatureError::RaggedPayload {
776                layer: 0,
777                field: "v length",
778                ..
779            }
780        ));
781    }
782
783    #[test]
784    fn an_empty_payload_is_rejected() {
785        assert_eq!(
786            CacheSignature::from_payload("model-a", flat(4), &[])
787                .expect_err("nothing to vouch for"),
788            SignatureError::EmptyPayload
789        );
790    }
791
792    /// An honest block of the wrong shape is a *miss*, reported as an
793    /// incompatibility naming the field that changed -- not a
794    /// corruption, and not a silent fallback.
795    #[test]
796    fn an_honest_block_from_a_different_config_is_incompatible() {
797        let layers = payload(3, 2, 8, 4);
798        let signature =
799            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
800        let err = UnverifiedBlock::new(Some(signature.clone()), layers)
801            .verify(&CacheSignature::expected("model-b", flat(4), 3, 2, 8, 4))
802            .expect_err("a different model must not share KV state");
803        assert_eq!(
804            err,
805            SignatureError::Incompatible {
806                field: "model",
807                expected: "model-b".into(),
808                found: "model-a".into(),
809            }
810        );
811
812        let layers = payload(3, 2, 8, 4);
813        let err = UnverifiedBlock::new(Some(signature), layers)
814            .verify(&CacheSignature::expected("model-a", flat(4), 3, 4, 8, 4))
815            .expect_err("a different KV head count must not be reused");
816        assert_eq!(
817            err,
818            SignatureError::Incompatible {
819                field: "n_kv_heads",
820                expected: "4".into(),
821                found: "2".into(),
822            }
823        );
824    }
825
826    #[test]
827    fn an_unreadable_format_version_is_rejected() {
828        let layers = payload(2, 2, 8, 4);
829        let mut signature =
830            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
831        signature.format_version = 99;
832        let err = UnverifiedBlock::new(Some(signature), layers)
833            .verify(&CacheSignature::expected("model-a", flat(4), 2, 2, 8, 4))
834            .expect_err("an unknown layout must not be guessed at");
835        assert_eq!(
836            err,
837            SignatureError::UnsupportedFormat {
838                found: 99,
839                readable: READABLE_FORMAT_VERSIONS,
840            }
841        );
842    }
843
844    /// The `kv-swa-block-alignment` invariant at the signature layer.
845    ///
846    /// Two builds of the same model, same tensors, same block size --
847    /// one configured with a 128-token sliding window and one with 256.
848    /// The payload cannot tell them apart, which is precisely why the
849    /// window is stamped: without this check the second build reads the
850    /// first build's blocks back and runs a mask the model never had.
851    #[test]
852    fn a_block_written_under_a_different_window_is_refused_not_reused() {
853        let layout_128 = BlockLayout::new(4, Some(128)).expect("4 divides 128");
854        let layout_256 = BlockLayout::new(4, Some(256)).expect("4 divides 256");
855        let layers = payload(3, 2, 8, 4);
856        let signature =
857            CacheSignature::from_payload("model-a", layout_128, &layers).expect("signature");
858
859        let err = UnverifiedBlock::new(Some(signature.clone()), layers)
860            .verify(&CacheSignature::expected("model-a", layout_256, 3, 2, 8, 4))
861            .expect_err("a window change must invalidate the block, not be ignored");
862        assert_eq!(
863            err,
864            SignatureError::Incompatible {
865                field: "sliding_window",
866                expected: "256".into(),
867                found: "128".into(),
868            }
869        );
870
871        // And the same block under the same window still verifies --
872        // the check must invalidate on change, not on principle.
873        let layers = payload(3, 2, 8, 4);
874        UnverifiedBlock::new(Some(signature), layers)
875            .verify(&CacheSignature::expected("model-a", layout_128, 3, 2, 8, 4))
876            .expect("unchanged config must still hit");
877    }
878
879    /// Turning SWA off (or on) is a config change of exactly the same
880    /// kind, and `None` must not read as "matches anything".
881    #[test]
882    fn a_full_causal_reader_will_not_take_a_sliding_window_block() {
883        let sliding = BlockLayout::new(4, Some(128)).expect("aligned");
884        let layers = payload(2, 2, 8, 4);
885        let signature =
886            CacheSignature::from_payload("model-a", sliding, &layers).expect("signature");
887        let err = UnverifiedBlock::new(Some(signature), layers)
888            .verify(&CacheSignature::expected("model-a", flat(4), 2, 2, 8, 4))
889            .expect_err("no window and a 128 window are different configurations");
890        assert_eq!(
891            err,
892            SignatureError::Incompatible {
893                field: "sliding_window",
894                expected: "none (full causal)".into(),
895                found: "128".into(),
896            }
897        );
898    }
899
900    /// A block cut at a different block size cannot be spliced into a
901    /// sequence cut at this one: the hash chain would not line up and
902    /// the eviction unit would not either.
903    #[test]
904    fn a_block_cut_at_a_different_block_size_is_incompatible() {
905        let layers = payload(2, 2, 8, 4);
906        let signature =
907            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
908        let err = UnverifiedBlock::new(Some(signature), layers)
909            .verify(&CacheSignature::expected("model-a", flat(2), 2, 2, 8, 4))
910            .expect_err("a 4-token block is not a 2-token block");
911        assert_eq!(
912            err,
913            SignatureError::Incompatible {
914                field: "block_size",
915                expected: "2".into(),
916                found: "4".into(),
917            }
918        );
919    }
920
921    /// `block_size` is the one new field the payload can settle, so it
922    /// is settled: a stamp claiming 8-token blocks over a 4-token
923    /// payload is a lying stamp, exactly like an overstated head_dim.
924    #[test]
925    fn a_stamp_may_not_claim_a_block_size_the_payload_lacks() {
926        let err = KvBlock::stamp("model-a", flat(8), payload(2, 2, 8, 4))
927            .expect_err("8-token blocks over a 4-token payload");
928        assert_eq!(
929            err,
930            SignatureError::BlockSizeMismatch {
931                block_size: 8,
932                tokens: 4,
933            }
934        );
935
936        // Same on the read path, where the stamp comes from a file
937        // rather than from this process.
938        let honest =
939            CacheSignature::from_payload("model-a", flat(4), &payload(2, 2, 8, 4)).expect("sig");
940        let mut lying = honest.clone();
941        lying.layout = flat(8);
942        lying.tokens = 8;
943        let err = UnverifiedBlock::new(Some(lying), payload(2, 2, 8, 4))
944            .verify(&CacheSignature::expected("model-a", flat(8), 2, 2, 8, 8))
945            .expect_err("the payload settles the block size, not the stamp");
946        assert_eq!(
947            err,
948            SignatureError::BlockSizeMismatch {
949                block_size: 8,
950                tokens: 4,
951            }
952        );
953    }
954
955    /// v1 blocks recorded no window at all. Reading one would mean
956    /// assuming it was aligned -- the assumption this whole item
957    /// exists to refuse -- so the readable-set drops it.
958    #[test]
959    fn blocks_from_the_pre_layout_format_are_not_readable() {
960        assert!(!READABLE_FORMAT_VERSIONS.contains(&1));
961        let layers = payload(2, 2, 8, 4);
962        let mut signature =
963            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
964        signature.format_version = 1;
965        let err = UnverifiedBlock::new(Some(signature), layers)
966            .verify(&CacheSignature::expected("model-a", flat(4), 2, 2, 8, 4))
967            .expect_err("a v1 block cannot say what layout it was cut under");
968        assert_eq!(
969            err,
970            SignatureError::UnsupportedFormat {
971                found: 1,
972                readable: READABLE_FORMAT_VERSIONS,
973            }
974        );
975    }
976
977    #[test]
978    fn errors_name_the_field_that_changed() {
979        let text = SignatureError::Incompatible {
980            field: "head_dim",
981            expected: "128".into(),
982            found: "64".into(),
983        }
984        .to_string();
985        assert!(text.contains("head_dim"), "{text}");
986        assert!(text.contains("64"), "{text}");
987        assert!(text.contains("128"), "{text}");
988        assert!(SignatureError::Unmarked.to_string().contains("unmarked"));
989    }
990}