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.k.len().is_multiple_of(per_token) {
292        return Err(SignatureError::RaggedPayload {
293            layer: index,
294            field: "k length",
295            expected: format!("a multiple of {per_token}"),
296            found: layer.k.len().to_string(),
297        });
298    }
299    if layer.v.len() != layer.k.len() {
300        return Err(SignatureError::RaggedPayload {
301            layer: index,
302            field: "v length",
303            expected: layer.k.len().to_string(),
304            found: layer.v.len().to_string(),
305        });
306    }
307    let tokens = layer.k.len() / per_token;
308    if layer.seq_len != tokens {
309        return Err(SignatureError::RaggedPayload {
310            layer: index,
311            field: "seq_len",
312            expected: tokens.to_string(),
313            found: layer.seq_len.to_string(),
314        });
315    }
316    Ok(tokens)
317}
318
319/// Why a stored block was refused. Every variant is a refusal to guess.
320#[derive(Clone, Debug, PartialEq, Eq)]
321pub enum SignatureError {
322    /// The block carries no signature. Not treated as "probably fine":
323    /// an unmarked block was written by something whose layout is
324    /// unknown, which is precisely the case that must not be trusted.
325    Unmarked,
326    /// A block with no layers describes nothing and can vouch for
327    /// nothing.
328    EmptyPayload,
329    /// A layer with no heads or zero head dimension.
330    DegenerateLayer {
331        layer: usize,
332        n_kv_heads: usize,
333        head_dim: usize,
334    },
335    /// The payload does not agree with itself: layers of different
336    /// shapes or lengths, or a layer whose buffers contradict its own
337    /// `seq_len`.
338    RaggedPayload {
339        layer: usize,
340        field: &'static str,
341        expected: String,
342        found: String,
343    },
344    /// The recorded signature claims something the payload is not. The
345    /// stamp is wrong (or was written by a build with a different
346    /// layout); the payload is the truth.
347    PayloadMismatch {
348        field: &'static str,
349        recorded: String,
350        actual: String,
351    },
352    /// The payload is coherent and honestly stamped, but it is not what
353    /// this reader needs -- a different model, a config change, a
354    /// different block size.
355    Incompatible {
356        field: &'static str,
357        expected: String,
358        found: String,
359    },
360    /// The stamp claims a block size the payload's token depth is not.
361    /// A stored block is exactly one whole block, so these are the same
362    /// number or the stamp is wrong.
363    BlockSizeMismatch { block_size: usize, tokens: usize },
364    /// The block layout itself is not usable -- most importantly, a
365    /// block size that does not divide the sliding window. See
366    /// [`kv_swa`](crate::kv_swa).
367    BadLayout(BlockLayoutError),
368    /// Written by a build whose payload layout this one cannot read.
369    UnsupportedFormat {
370        found: u32,
371        readable: &'static [u32],
372    },
373}
374
375impl From<BlockLayoutError> for SignatureError {
376    fn from(err: BlockLayoutError) -> Self {
377        SignatureError::BadLayout(err)
378    }
379}
380
381impl std::fmt::Display for SignatureError {
382    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
383        match self {
384            SignatureError::Unmarked => write!(
385                f,
386                "KV block carries no cache signature; refusing to trust an unmarked block"
387            ),
388            SignatureError::EmptyPayload => {
389                write!(f, "KV block has no layers; nothing to verify")
390            }
391            SignatureError::DegenerateLayer {
392                layer,
393                n_kv_heads,
394                head_dim,
395            } => write!(
396                f,
397                "KV block layer {layer} is degenerate: {n_kv_heads} kv heads x {head_dim} head dim"
398            ),
399            SignatureError::RaggedPayload {
400                layer,
401                field,
402                expected,
403                found,
404            } => write!(
405                f,
406                "KV block payload is inconsistent at layer {layer}: {field} is {found}, expected {expected}"
407            ),
408            SignatureError::PayloadMismatch {
409                field,
410                recorded,
411                actual,
412            } => write!(
413                f,
414                "KV block signature vouches for {field}={recorded} but its payload has {field}={actual}"
415            ),
416            SignatureError::Incompatible {
417                field,
418                expected,
419                found,
420            } => write!(
421                f,
422                "KV block is incompatible: {field} is {found}, this server needs {expected}"
423            ),
424            SignatureError::BlockSizeMismatch { block_size, tokens } => write!(
425                f,
426                "KV block signature declares a block size of {block_size} but its payload holds \
427                 {tokens} token positions; a stored block is exactly one whole block"
428            ),
429            SignatureError::BadLayout(err) => write!(f, "KV block layout is unusable: {err}"),
430            SignatureError::UnsupportedFormat { found, readable } => write!(
431                f,
432                "KV block format version {found} is not readable by this build (readable: {readable:?})"
433            ),
434        }
435    }
436}
437
438impl std::error::Error for SignatureError {}
439
440/// A block whose signature has been verified against its own payload
441/// and against the reader's expectation. Only way to get one is
442/// [`KvBlock::stamp`] (writing) or [`UnverifiedBlock::verify`]
443/// (reading), so holding one is itself the proof.
444pub struct KvBlock {
445    signature: CacheSignature,
446    layers: Vec<KvCache>,
447}
448
449/// Summarizes rather than dumping tensors: a block's `Debug` is for a
450/// log line or a failing assertion, and printing every f32 in a KV
451/// block helps nobody.
452impl std::fmt::Debug for KvBlock {
453    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
454        f.debug_struct("KvBlock")
455            .field("signature", &self.signature)
456            .field("layers", &self.layers.len())
457            .finish()
458    }
459}
460
461impl std::fmt::Debug for UnverifiedBlock {
462    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
463        f.debug_struct("UnverifiedBlock")
464            .field("signature", &self.signature)
465            .field("layers", &self.layers.len())
466            .finish()
467    }
468}
469
470impl KvBlock {
471    /// Stamps a block from its own payload. The caller supplies the
472    /// model identity and the layout the block was cut under; every
473    /// shape field is measured, and the layout's block size is checked
474    /// against the depth measured.
475    pub fn stamp(
476        model: &str,
477        layout: BlockLayout,
478        layers: Vec<KvCache>,
479    ) -> Result<Self, SignatureError> {
480        let signature = CacheSignature::from_payload(model, layout, &layers)?;
481        Ok(KvBlock { signature, layers })
482    }
483
484    /// The layout this block was cut under.
485    pub fn layout(&self) -> BlockLayout {
486        self.signature.layout
487    }
488
489    pub fn signature(&self) -> &CacheSignature {
490        &self.signature
491    }
492
493    pub fn tokens(&self) -> usize {
494        self.signature.tokens
495    }
496
497    pub fn layers(&self) -> &[KvCache] {
498        &self.layers
499    }
500
501    pub fn into_layers(self) -> Vec<KvCache> {
502        self.layers
503    }
504}
505
506/// A block as it comes back from an untrusted source: a disk tier, a
507/// shared cache, or a build that is not this one. `signature` is an
508/// `Option` on purpose -- "no signature recorded" is a state a real
509/// stored block can be in, and it must be representable so it can be
510/// refused.
511pub struct UnverifiedBlock {
512    pub signature: Option<CacheSignature>,
513    pub layers: Vec<KvCache>,
514}
515
516impl UnverifiedBlock {
517    pub fn new(signature: Option<CacheSignature>, layers: Vec<KvCache>) -> Self {
518        UnverifiedBlock { signature, layers }
519    }
520
521    /// Verifies in the order the module doc describes: marked, honest,
522    /// then compatible. `expected` is used only in the last step -- it
523    /// never contributes a field to the signature being checked.
524    pub fn verify(self, expected: &CacheSignature) -> Result<KvBlock, SignatureError> {
525        let recorded = self.signature.ok_or(SignatureError::Unmarked)?;
526        if !READABLE_FORMAT_VERSIONS.contains(&recorded.format_version) {
527            return Err(SignatureError::UnsupportedFormat {
528                found: recorded.format_version,
529                readable: READABLE_FORMAT_VERSIONS,
530            });
531        }
532        // Measured from the payload. `recorded.model` is carried over
533        // because no payload can prove which weights produced it; the
534        // model check happens against `expected` below, where a wrong
535        // model is caught as an incompatibility.
536        // `recorded.layout` is carried over for the same reason
537        // `recorded.model` is: the reader's expectation must not get a
538        // vote before the payload has been checked against the stamp.
539        // Carrying it is not trusting it -- `from_payload` rejects a
540        // block size the token depth contradicts, and the window is
541        // settled against `expected` below.
542        let actual = CacheSignature::from_payload(&recorded.model, recorded.layout, &self.layers)?;
543        recorded.compare(&actual, |field, recorded, actual| {
544            SignatureError::PayloadMismatch {
545                field,
546                recorded,
547                actual,
548            }
549        })?;
550        expected.compare(&actual, |field, expected, found| {
551            SignatureError::Incompatible {
552                field,
553                expected,
554                found,
555            }
556        })?;
557        Ok(KvBlock {
558            signature: actual,
559            layers: self.layers,
560        })
561    }
562}
563
564#[cfg(test)]
565mod tests {
566    use super::*;
567
568    fn layer(n_kv_heads: usize, head_dim: usize, tokens: usize) -> KvCache {
569        let mut cache = KvCache::new(n_kv_heads, head_dim);
570        let step = vec![0.5f32; n_kv_heads * head_dim];
571        for _ in 0..tokens {
572            cache.push(&step, &step).expect("unpooled push cannot fail");
573        }
574        cache
575    }
576
577    fn payload(n_layers: usize, n_kv_heads: usize, head_dim: usize, tokens: usize) -> Vec<KvCache> {
578        (0..n_layers)
579            .map(|_| layer(n_kv_heads, head_dim, tokens))
580            .collect()
581    }
582
583    /// A full-causal layout whose block size is the payload depth --
584    /// what every test that is not about SWA wants.
585    fn flat(block_size: usize) -> BlockLayout {
586        BlockLayout::full_attention(block_size).expect("positive block size")
587    }
588
589    /// The core rule, in its positive form: every shape field comes
590    /// from the tensors. `stamp` is given a model name and nothing else.
591    #[test]
592    fn signature_is_measured_from_the_payload() {
593        let block = KvBlock::stamp("model-a", flat(4), payload(3, 2, 8, 4)).expect("stamp");
594        let sig = block.signature();
595        assert_eq!(sig.n_layers, 3);
596        assert_eq!(sig.n_kv_heads, 2);
597        assert_eq!(sig.head_dim, 8);
598        assert_eq!(sig.tokens, 4);
599        assert_eq!(sig.dtype, KvDtype::F32);
600        assert_eq!(sig.format_version, BLOCK_FORMAT_VERSION);
601        assert_eq!(block.tokens(), 4);
602        assert_eq!(block.layers().len(), 3);
603    }
604
605    #[test]
606    fn a_stamped_block_round_trips_through_verification() {
607        let layers = payload(3, 2, 8, 4);
608        let signature =
609            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
610        let expected = CacheSignature::expected("model-a", flat(4), 3, 2, 8, 4);
611        let block = UnverifiedBlock::new(Some(signature), layers)
612            .verify(&expected)
613            .expect("a block that is what it says it is must verify");
614        assert_eq!(block.layers().len(), 3);
615        assert_eq!(block.into_layers().len(), 3);
616    }
617
618    /// Absence of a signature is not agreement. This is the difference
619    /// between a persistent cache and silent corruption after a config
620    /// change: an unmarked block came from something whose layout is
621    /// unknown by definition.
622    #[test]
623    fn an_unmarked_block_is_rejected_not_trusted() {
624        let expected = CacheSignature::expected("model-a", flat(4), 3, 2, 8, 4);
625        let err = UnverifiedBlock::new(None, payload(3, 2, 8, 4))
626            .verify(&expected)
627            .expect_err("an unmarked block must be refused");
628        assert_eq!(err, SignatureError::Unmarked);
629    }
630
631    /// The rule the plan states as "a signature must never vouch for a
632    /// width the payload does not have". Here the stamp is exactly what
633    /// the reader expects -- and the payload is not. A cache that
634    /// trusted the stamp (or, equivalently, stamped from the manager's
635    /// expectation) would hand back tensors of the wrong width and
636    /// produce confident wrong tokens.
637    #[test]
638    fn a_signature_that_overstates_its_payload_is_rejected() {
639        let expected = CacheSignature::expected("model-a", flat(4), 3, 2, 16, 4);
640        let mut lying = expected.clone();
641        assert_eq!(lying.head_dim, 16);
642        let err = UnverifiedBlock::new(Some(lying.clone()), payload(3, 2, 8, 4))
643            .verify(&expected)
644            .expect_err("stamp claims head_dim 16 over an 8-wide payload");
645        assert_eq!(
646            err,
647            SignatureError::PayloadMismatch {
648                field: "head_dim",
649                recorded: "16".into(),
650                actual: "8".into(),
651            }
652        );
653
654        // Same shape, overstated depth: 8 token positions claimed over
655        // a 4-position payload.
656        lying.head_dim = 8;
657        lying.tokens = 8;
658        let expected = CacheSignature::expected("model-a", flat(8), 3, 2, 8, 8);
659        let err = UnverifiedBlock::new(Some(lying.clone()), payload(3, 2, 8, 4))
660            .verify(&expected)
661            .expect_err("stamp claims 8 tokens over a 4-token payload");
662        assert_eq!(
663            err,
664            SignatureError::PayloadMismatch {
665                field: "tokens",
666                recorded: "8".into(),
667                actual: "4".into(),
668            }
669        );
670
671        // And overstated layer count.
672        lying.tokens = 4;
673        lying.n_layers = 4;
674        let expected = CacheSignature::expected("model-a", flat(4), 4, 2, 8, 4);
675        let err = UnverifiedBlock::new(Some(lying), payload(3, 2, 8, 4))
676            .verify(&expected)
677            .expect_err("stamp claims 4 layers over a 3-layer payload");
678        assert_eq!(
679            err,
680            SignatureError::PayloadMismatch {
681                field: "n_layers",
682                recorded: "4".into(),
683                actual: "3".into(),
684            }
685        );
686    }
687
688    /// A payload that lies to itself: `seq_len` says 4, the buffers
689    /// hold 3. `seq_len` is a claim; `k.len()` is the evidence.
690    #[test]
691    fn a_layer_whose_seq_len_contradicts_its_buffers_is_rejected() {
692        let mut layers = payload(2, 2, 8, 4);
693        layers[1].seq_len = 7;
694        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
695            .expect_err("seq_len must be verified, not believed");
696        assert_eq!(
697            err,
698            SignatureError::RaggedPayload {
699                layer: 1,
700                field: "seq_len",
701                expected: "4".into(),
702                found: "7".into(),
703            }
704        );
705    }
706
707    #[test]
708    fn a_ragged_payload_is_rejected() {
709        let mut layers = payload(3, 2, 8, 4);
710        layers[2] = layer(2, 4, 4);
711        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
712            .expect_err("shape disagreement");
713        assert!(matches!(
714            err,
715            SignatureError::RaggedPayload {
716                layer: 2,
717                field: "layer shape",
718                ..
719            }
720        ));
721
722        let mut layers = payload(3, 2, 8, 4);
723        layers[1] = layer(2, 8, 3);
724        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
725            .expect_err("depth disagreement");
726        assert!(matches!(
727            err,
728            SignatureError::RaggedPayload {
729                layer: 1,
730                field: "token count",
731                ..
732            }
733        ));
734
735        let mut layers = payload(2, 2, 8, 4);
736        layers[0].v.truncate(8);
737        let err = CacheSignature::from_payload("model-a", flat(4), &layers)
738            .expect_err("k/v disagreement");
739        assert!(matches!(
740            err,
741            SignatureError::RaggedPayload {
742                layer: 0,
743                field: "v length",
744                ..
745            }
746        ));
747    }
748
749    #[test]
750    fn an_empty_payload_is_rejected() {
751        assert_eq!(
752            CacheSignature::from_payload("model-a", flat(4), &[])
753                .expect_err("nothing to vouch for"),
754            SignatureError::EmptyPayload
755        );
756    }
757
758    /// An honest block of the wrong shape is a *miss*, reported as an
759    /// incompatibility naming the field that changed -- not a
760    /// corruption, and not a silent fallback.
761    #[test]
762    fn an_honest_block_from_a_different_config_is_incompatible() {
763        let layers = payload(3, 2, 8, 4);
764        let signature =
765            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
766        let err = UnverifiedBlock::new(Some(signature.clone()), layers)
767            .verify(&CacheSignature::expected("model-b", flat(4), 3, 2, 8, 4))
768            .expect_err("a different model must not share KV state");
769        assert_eq!(
770            err,
771            SignatureError::Incompatible {
772                field: "model",
773                expected: "model-b".into(),
774                found: "model-a".into(),
775            }
776        );
777
778        let layers = payload(3, 2, 8, 4);
779        let err = UnverifiedBlock::new(Some(signature), layers)
780            .verify(&CacheSignature::expected("model-a", flat(4), 3, 4, 8, 4))
781            .expect_err("a different KV head count must not be reused");
782        assert_eq!(
783            err,
784            SignatureError::Incompatible {
785                field: "n_kv_heads",
786                expected: "4".into(),
787                found: "2".into(),
788            }
789        );
790    }
791
792    #[test]
793    fn an_unreadable_format_version_is_rejected() {
794        let layers = payload(2, 2, 8, 4);
795        let mut signature =
796            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
797        signature.format_version = 99;
798        let err = UnverifiedBlock::new(Some(signature), layers)
799            .verify(&CacheSignature::expected("model-a", flat(4), 2, 2, 8, 4))
800            .expect_err("an unknown layout must not be guessed at");
801        assert_eq!(
802            err,
803            SignatureError::UnsupportedFormat {
804                found: 99,
805                readable: READABLE_FORMAT_VERSIONS,
806            }
807        );
808    }
809
810    /// The `kv-swa-block-alignment` invariant at the signature layer.
811    ///
812    /// Two builds of the same model, same tensors, same block size --
813    /// one configured with a 128-token sliding window and one with 256.
814    /// The payload cannot tell them apart, which is precisely why the
815    /// window is stamped: without this check the second build reads the
816    /// first build's blocks back and runs a mask the model never had.
817    #[test]
818    fn a_block_written_under_a_different_window_is_refused_not_reused() {
819        let layout_128 = BlockLayout::new(4, Some(128)).expect("4 divides 128");
820        let layout_256 = BlockLayout::new(4, Some(256)).expect("4 divides 256");
821        let layers = payload(3, 2, 8, 4);
822        let signature =
823            CacheSignature::from_payload("model-a", layout_128, &layers).expect("signature");
824
825        let err = UnverifiedBlock::new(Some(signature.clone()), layers)
826            .verify(&CacheSignature::expected("model-a", layout_256, 3, 2, 8, 4))
827            .expect_err("a window change must invalidate the block, not be ignored");
828        assert_eq!(
829            err,
830            SignatureError::Incompatible {
831                field: "sliding_window",
832                expected: "256".into(),
833                found: "128".into(),
834            }
835        );
836
837        // And the same block under the same window still verifies --
838        // the check must invalidate on change, not on principle.
839        let layers = payload(3, 2, 8, 4);
840        UnverifiedBlock::new(Some(signature), layers)
841            .verify(&CacheSignature::expected("model-a", layout_128, 3, 2, 8, 4))
842            .expect("unchanged config must still hit");
843    }
844
845    /// Turning SWA off (or on) is a config change of exactly the same
846    /// kind, and `None` must not read as "matches anything".
847    #[test]
848    fn a_full_causal_reader_will_not_take_a_sliding_window_block() {
849        let sliding = BlockLayout::new(4, Some(128)).expect("aligned");
850        let layers = payload(2, 2, 8, 4);
851        let signature =
852            CacheSignature::from_payload("model-a", sliding, &layers).expect("signature");
853        let err = UnverifiedBlock::new(Some(signature), layers)
854            .verify(&CacheSignature::expected("model-a", flat(4), 2, 2, 8, 4))
855            .expect_err("no window and a 128 window are different configurations");
856        assert_eq!(
857            err,
858            SignatureError::Incompatible {
859                field: "sliding_window",
860                expected: "none (full causal)".into(),
861                found: "128".into(),
862            }
863        );
864    }
865
866    /// A block cut at a different block size cannot be spliced into a
867    /// sequence cut at this one: the hash chain would not line up and
868    /// the eviction unit would not either.
869    #[test]
870    fn a_block_cut_at_a_different_block_size_is_incompatible() {
871        let layers = payload(2, 2, 8, 4);
872        let signature =
873            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
874        let err = UnverifiedBlock::new(Some(signature), layers)
875            .verify(&CacheSignature::expected("model-a", flat(2), 2, 2, 8, 4))
876            .expect_err("a 4-token block is not a 2-token block");
877        assert_eq!(
878            err,
879            SignatureError::Incompatible {
880                field: "block_size",
881                expected: "2".into(),
882                found: "4".into(),
883            }
884        );
885    }
886
887    /// `block_size` is the one new field the payload can settle, so it
888    /// is settled: a stamp claiming 8-token blocks over a 4-token
889    /// payload is a lying stamp, exactly like an overstated head_dim.
890    #[test]
891    fn a_stamp_may_not_claim_a_block_size_the_payload_lacks() {
892        let err = KvBlock::stamp("model-a", flat(8), payload(2, 2, 8, 4))
893            .expect_err("8-token blocks over a 4-token payload");
894        assert_eq!(
895            err,
896            SignatureError::BlockSizeMismatch {
897                block_size: 8,
898                tokens: 4,
899            }
900        );
901
902        // Same on the read path, where the stamp comes from a file
903        // rather than from this process.
904        let honest =
905            CacheSignature::from_payload("model-a", flat(4), &payload(2, 2, 8, 4)).expect("sig");
906        let mut lying = honest.clone();
907        lying.layout = flat(8);
908        lying.tokens = 8;
909        let err = UnverifiedBlock::new(Some(lying), payload(2, 2, 8, 4))
910            .verify(&CacheSignature::expected("model-a", flat(8), 2, 2, 8, 8))
911            .expect_err("the payload settles the block size, not the stamp");
912        assert_eq!(
913            err,
914            SignatureError::BlockSizeMismatch {
915                block_size: 8,
916                tokens: 4,
917            }
918        );
919    }
920
921    /// v1 blocks recorded no window at all. Reading one would mean
922    /// assuming it was aligned -- the assumption this whole item
923    /// exists to refuse -- so the readable-set drops it.
924    #[test]
925    fn blocks_from_the_pre_layout_format_are_not_readable() {
926        assert!(!READABLE_FORMAT_VERSIONS.contains(&1));
927        let layers = payload(2, 2, 8, 4);
928        let mut signature =
929            CacheSignature::from_payload("model-a", flat(4), &layers).expect("signature");
930        signature.format_version = 1;
931        let err = UnverifiedBlock::new(Some(signature), layers)
932            .verify(&CacheSignature::expected("model-a", flat(4), 2, 2, 8, 4))
933            .expect_err("a v1 block cannot say what layout it was cut under");
934        assert_eq!(
935            err,
936            SignatureError::UnsupportedFormat {
937                found: 1,
938                readable: READABLE_FORMAT_VERSIONS,
939            }
940        );
941    }
942
943    #[test]
944    fn errors_name_the_field_that_changed() {
945        let text = SignatureError::Incompatible {
946            field: "head_dim",
947            expected: "128".into(),
948            found: "64".into(),
949        }
950        .to_string();
951        assert!(text.contains("head_dim"), "{text}");
952        assert!(text.contains("64"), "{text}");
953        assert!(text.contains("128"), "{text}");
954        assert!(SignatureError::Unmarked.to_string().contains("unmarked"));
955    }
956}