1use crate::cache::KvCache;
55use crate::kv_swa::{BlockLayout, BlockLayoutError};
56
57pub const BLOCK_FORMAT_VERSION: u32 = 2;
61
62pub const READABLE_FORMAT_VERSIONS: &[u32] = &[2];
72
73#[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#[derive(Clone, Debug, PartialEq, Eq)]
96pub struct CacheSignature {
97 pub format_version: u32,
98 pub model: String,
102 pub n_layers: usize,
103 pub n_kv_heads: usize,
104 pub head_dim: usize,
105 pub dtype: KvDtype,
106 pub tokens: usize,
110 pub layout: BlockLayout,
113}
114
115impl CacheSignature {
116 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 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 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 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
277fn 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
287fn 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#[derive(Clone, Debug, PartialEq, Eq)]
321pub enum SignatureError {
322 Unmarked,
326 EmptyPayload,
329 DegenerateLayer {
331 layer: usize,
332 n_kv_heads: usize,
333 head_dim: usize,
334 },
335 RaggedPayload {
339 layer: usize,
340 field: &'static str,
341 expected: String,
342 found: String,
343 },
344 PayloadMismatch {
348 field: &'static str,
349 recorded: String,
350 actual: String,
351 },
352 Incompatible {
356 field: &'static str,
357 expected: String,
358 found: String,
359 },
360 BlockSizeMismatch { block_size: usize, tokens: usize },
364 BadLayout(BlockLayoutError),
368 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
440pub struct KvBlock {
445 signature: CacheSignature,
446 layers: Vec<KvCache>,
447}
448
449impl 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 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 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
506pub 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 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 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 fn flat(block_size: usize) -> BlockLayout {
586 BlockLayout::full_attention(block_size).expect("positive block size")
587 }
588
589 #[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 #[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 #[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 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 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 #[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 #[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 #[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 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 #[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 #[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 #[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 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 #[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}