1use std::{
4 collections::{BTreeMap, BTreeSet},
5 ops::Range,
6};
7
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10
11use crate::attention::{AttentionPolicy, LayerSchedule};
12
13use super::{
14 CachePolicyError, CacheRankIdentity, CacheRepresentation, LayerCachePolicy, StateTensorOwner,
15 StateTensorRole,
16};
17
18pub const PROMPT_CACHE_SCHEMA_VERSION: u32 = 8;
20
21#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
23pub struct PromptCacheStateSegment {
24 id: String,
25 layers: Range<usize>,
26}
27
28impl PromptCacheStateSegment {
29 pub fn new(id: impl Into<String>, layers: Range<usize>) -> Result<Self, PromptCacheError> {
31 let id = id.into();
32 if id.trim().is_empty() {
33 return Err(PromptCacheError::Malformed(
34 "prompt-cache state segment identity must not be empty".into(),
35 ));
36 }
37 if layers.is_empty() {
38 return Err(PromptCacheError::Malformed(format!(
39 "prompt-cache state segment {id:?} has an empty range"
40 )));
41 }
42 Ok(Self { id, layers })
43 }
44
45 pub fn id(&self) -> &str {
47 &self.id
48 }
49
50 pub fn layers(&self) -> Range<usize> {
52 self.layers.clone()
53 }
54}
55
56#[derive(Debug, Clone, Eq, Hash, PartialEq)]
58pub struct PromptCacheDescriptor {
59 model_family: String,
61 effective_model_type: String,
63 checkpoint_fingerprint: String,
65 prefix_content_fingerprint: String,
67 architecture_fingerprint: String,
69 layer_count: usize,
71 global_layer_start: usize,
73 global_layer_end: usize,
75 batch_size: usize,
77 layer_layout: LayerSchedule<LayerCachePolicy>,
79 layer_prefix_offsets: Vec<i32>,
81 state_segments: Vec<PromptCacheStateSegment>,
83 sink_tokens: usize,
85 topology: PromptCacheTopology,
87}
88
89impl PromptCacheDescriptor {
90 #[allow(clippy::too_many_arguments)]
92 pub fn new(
93 model_family: impl Into<String>,
94 effective_model_type: impl Into<String>,
95 checkpoint_fingerprint: impl Into<String>,
96 prefix_content_fingerprint: impl Into<String>,
97 architecture_fingerprint: impl Into<String>,
98 layer_count: usize,
99 global_layer_start: usize,
100 global_layer_end: usize,
101 batch_size: usize,
102 layer_layout: LayerSchedule<LayerCachePolicy>,
103 layer_prefix_offsets: Vec<i32>,
104 state_segments: Vec<PromptCacheStateSegment>,
105 sink_tokens: usize,
106 topology: PromptCacheTopology,
107 ) -> Result<Self, PromptCacheError> {
108 let descriptor = Self {
109 model_family: model_family.into(),
110 effective_model_type: effective_model_type.into(),
111 checkpoint_fingerprint: checkpoint_fingerprint.into(),
112 prefix_content_fingerprint: prefix_content_fingerprint.into(),
113 architecture_fingerprint: architecture_fingerprint.into(),
114 layer_count,
115 global_layer_start,
116 global_layer_end,
117 batch_size,
118 layer_layout,
119 layer_prefix_offsets,
120 state_segments,
121 sink_tokens,
122 topology,
123 };
124 for value in [
125 &descriptor.model_family,
126 &descriptor.effective_model_type,
127 &descriptor.checkpoint_fingerprint,
128 &descriptor.prefix_content_fingerprint,
129 &descriptor.architecture_fingerprint,
130 ] {
131 if value.trim().is_empty() {
132 return Err(PromptCacheError::Malformed(
133 "prompt-cache identity strings must be non-empty".into(),
134 ));
135 }
136 }
137 descriptor.validate()?;
138 Ok(descriptor)
139 }
140
141 pub fn model_family(&self) -> &str {
143 &self.model_family
144 }
145 pub fn effective_model_type(&self) -> &str {
147 &self.effective_model_type
148 }
149 pub fn checkpoint_fingerprint(&self) -> &str {
151 &self.checkpoint_fingerprint
152 }
153 pub fn prefix_content_fingerprint(&self) -> &str {
155 &self.prefix_content_fingerprint
156 }
157 pub fn architecture_fingerprint(&self) -> &str {
159 &self.architecture_fingerprint
160 }
161 pub const fn layer_count(&self) -> usize {
163 self.layer_count
164 }
165 pub const fn global_layer_start(&self) -> usize {
167 self.global_layer_start
168 }
169 pub const fn global_layer_end(&self) -> usize {
171 self.global_layer_end
172 }
173 pub const fn batch_size(&self) -> usize {
175 self.batch_size
176 }
177 pub const fn layer_layout(&self) -> &LayerSchedule<LayerCachePolicy> {
179 &self.layer_layout
180 }
181 pub fn layer_prefix_offsets(&self) -> &[i32] {
183 &self.layer_prefix_offsets
184 }
185 pub fn state_segments(&self) -> &[PromptCacheStateSegment] {
187 &self.state_segments
188 }
189 pub const fn sink_tokens(&self) -> usize {
191 self.sink_tokens
192 }
193 pub const fn topology(&self) -> &PromptCacheTopology {
195 &self.topology
196 }
197 pub fn with_topology(
199 mut self,
200 topology: PromptCacheTopology,
201 ) -> Result<Self, PromptCacheError> {
202 self.topology = topology;
203 self.validate()?;
204 Ok(self)
205 }
206 pub fn with_architecture_fingerprint(
208 mut self,
209 architecture_fingerprint: impl Into<String>,
210 ) -> Result<Self, PromptCacheError> {
211 self.architecture_fingerprint = architecture_fingerprint.into();
212 if self.architecture_fingerprint.trim().is_empty() {
213 return Err(PromptCacheError::Malformed(
214 "prompt-cache architecture fingerprint must be non-empty".into(),
215 ));
216 }
217 self.validate()?;
218 Ok(self)
219 }
220 pub fn with_layer_count(mut self, layer_count: usize) -> Result<Self, PromptCacheError> {
222 self.layer_count = layer_count;
223 self.validate()?;
224 Ok(self)
225 }
226 pub fn from_model_identity(
232 model: PromptCacheModelIdentity,
233 checkpoint_fingerprint: impl Into<String>,
234 prefix_content_fingerprint: impl Into<String>,
235 batch_size: usize,
236 ) -> Result<Self, PromptCacheError> {
237 let descriptor = Self {
238 model_family: model.model_family,
239 effective_model_type: model.effective_model_type,
240 checkpoint_fingerprint: checkpoint_fingerprint.into(),
241 prefix_content_fingerprint: prefix_content_fingerprint.into(),
242 architecture_fingerprint: model.architecture_fingerprint,
243 layer_count: model.layer_count,
244 global_layer_start: model.global_layer_start,
245 global_layer_end: model.global_layer_end,
246 batch_size,
247 layer_layout: model.layer_layout,
248 layer_prefix_offsets: model.layer_prefix_offsets,
249 state_segments: model.state_segments,
250 sink_tokens: model.sink_tokens,
251 topology: model.topology,
252 };
253 descriptor.validate()?;
254 Ok(descriptor)
255 }
256
257 pub fn validate(&self) -> Result<(), PromptCacheError> {
259 IdentityLayout {
260 layer_count: self.layer_count,
261 global_layer_start: self.global_layer_start,
262 global_layer_end: self.global_layer_end,
263 batch_size: self.batch_size,
264 layer_layout: &self.layer_layout,
265 layer_prefix_offsets: &self.layer_prefix_offsets,
266 state_segments: &self.state_segments,
267 topology: &self.topology,
268 }
269 .validate("prompt-cache descriptor")
270 }
271}
272
273#[derive(Debug, Clone, Eq, Hash, PartialEq)]
275pub struct PromptCacheModelIdentity {
276 model_family: String,
278 effective_model_type: String,
280 architecture_fingerprint: String,
282 layer_count: usize,
284 global_layer_start: usize,
286 global_layer_end: usize,
288 sink_tokens: usize,
290 topology: PromptCacheTopology,
292 layer_layout: LayerSchedule<LayerCachePolicy>,
294 layer_prefix_offsets: Vec<i32>,
296 state_segments: Vec<PromptCacheStateSegment>,
298}
299
300impl PromptCacheModelIdentity {
301 #[allow(clippy::too_many_arguments)]
303 pub fn new(
304 model_family: impl Into<String>,
305 effective_model_type: impl Into<String>,
306 architecture_fingerprint: impl Into<String>,
307 layer_count: usize,
308 global_layer_start: usize,
309 global_layer_end: usize,
310 sink_tokens: usize,
311 topology: PromptCacheTopology,
312 layer_layout: LayerSchedule<LayerCachePolicy>,
313 layer_prefix_offsets: Vec<i32>,
314 state_segments: Vec<PromptCacheStateSegment>,
315 ) -> Result<Self, PromptCacheError> {
316 let identity = Self {
317 model_family: model_family.into(),
318 effective_model_type: effective_model_type.into(),
319 architecture_fingerprint: architecture_fingerprint.into(),
320 layer_count,
321 global_layer_start,
322 global_layer_end,
323 sink_tokens,
324 topology,
325 layer_layout,
326 layer_prefix_offsets,
327 state_segments,
328 };
329 for value in [
330 &identity.model_family,
331 &identity.effective_model_type,
332 &identity.architecture_fingerprint,
333 ] {
334 if value.trim().is_empty() {
335 return Err(PromptCacheError::Malformed(
336 "prompt-cache model identity strings must be non-empty".into(),
337 ));
338 }
339 }
340 identity.validate()?;
341 Ok(identity)
342 }
343
344 pub fn model_family(&self) -> &str {
346 &self.model_family
347 }
348 pub fn effective_model_type(&self) -> &str {
350 &self.effective_model_type
351 }
352 pub fn architecture_fingerprint(&self) -> &str {
354 &self.architecture_fingerprint
355 }
356 pub const fn layer_count(&self) -> usize {
358 self.layer_count
359 }
360 pub const fn global_layer_start(&self) -> usize {
362 self.global_layer_start
363 }
364 pub const fn global_layer_end(&self) -> usize {
366 self.global_layer_end
367 }
368 pub const fn sink_tokens(&self) -> usize {
370 self.sink_tokens
371 }
372 pub const fn topology(&self) -> &PromptCacheTopology {
374 &self.topology
375 }
376 pub const fn layer_layout(&self) -> &LayerSchedule<LayerCachePolicy> {
378 &self.layer_layout
379 }
380 pub fn layer_prefix_offsets(&self) -> &[i32] {
382 &self.layer_prefix_offsets
383 }
384 pub fn state_segments(&self) -> &[PromptCacheStateSegment] {
386 &self.state_segments
387 }
388 pub fn key_value_layouts(
390 sliding_windows: impl IntoIterator<Item = Option<i32>>,
391 num_key_value_heads: i32,
392 head_dim: i32,
393 ) -> Result<LayerSchedule<LayerCachePolicy>, PromptCacheError> {
394 let policies = sliding_windows
395 .into_iter()
396 .map(|window| {
397 let attention = AttentionPolicy::from_sliding_window(window)
398 .map_err(|error| PromptCacheError::Malformed(error.to_string()))?;
399 LayerCachePolicy::key_value(attention, num_key_value_heads, head_dim)
400 .map_err(PromptCacheError::from)
401 })
402 .collect::<Result<Vec<_>, _>>()?;
403 LayerSchedule::new(policies.len(), policies)
404 .map_err(|error| PromptCacheError::Malformed(error.to_string()))
405 }
406
407 pub fn compressed_layouts(
409 layer_count: usize,
410 latent_dim: i32,
411 rotary_dim: i32,
412 ) -> Result<LayerSchedule<LayerCachePolicy>, PromptCacheError> {
413 let policies = (0..layer_count)
414 .map(|_| {
415 LayerCachePolicy::compressed_latent_rotary(
416 AttentionPolicy::Full,
417 latent_dim,
418 rotary_dim,
419 )
420 .map_err(PromptCacheError::from)
421 })
422 .collect::<Result<Vec<_>, _>>()?;
423 LayerSchedule::new(layer_count, policies)
424 .map_err(|error| PromptCacheError::Malformed(error.to_string()))
425 }
426
427 pub fn validate(&self) -> Result<(), PromptCacheError> {
429 IdentityLayout {
430 layer_count: self.layer_count,
431 global_layer_start: self.global_layer_start,
432 global_layer_end: self.global_layer_end,
433 batch_size: 1,
434 layer_layout: &self.layer_layout,
435 layer_prefix_offsets: &self.layer_prefix_offsets,
436 state_segments: &self.state_segments,
437 topology: &self.topology,
438 }
439 .validate("loaded model")
440 }
441
442 pub fn state_segment(&self, id: &str) -> Result<&PromptCacheStateSegment, PromptCacheError> {
444 self.validate()?;
445 self.state_segments
446 .iter()
447 .find(|segment| segment.id() == id)
448 .ok_or_else(|| {
449 PromptCacheError::Incompatible(format!(
450 "loaded model has no prompt-cache state segment {id:?}"
451 ))
452 })
453 }
454
455 pub fn select_state_segment(&self, id: &str) -> Result<Self, PromptCacheError> {
457 let layers = self.state_segment(id)?.layers();
458 let length = layers.len();
459 let global_layer_start = self
460 .global_layer_start
461 .checked_add(layers.start)
462 .ok_or_else(|| PromptCacheError::Malformed("state segment range overflowed".into()))?;
463 let global_layer_end = global_layer_start
464 .checked_add(length)
465 .ok_or_else(|| PromptCacheError::Malformed("state segment range overflowed".into()))?;
466 let layer_layout = LayerSchedule::new(
467 length,
468 self.layer_layout
469 .iter()
470 .skip(layers.start)
471 .take(length)
472 .cloned()
473 .collect(),
474 )
475 .map_err(|error| PromptCacheError::Malformed(error.to_string()))?;
476 let layer_prefix_offsets = self
477 .layer_prefix_offsets
478 .get(layers.clone())
479 .ok_or_else(|| PromptCacheError::Malformed("state segment range is invalid".into()))?
480 .to_vec();
481 let selected = Self {
482 model_family: self.model_family.clone(),
483 effective_model_type: self.effective_model_type.clone(),
484 architecture_fingerprint: self.architecture_fingerprint.clone(),
485 layer_count: self.layer_count,
486 global_layer_start,
487 global_layer_end,
488 sink_tokens: self.sink_tokens,
489 topology: self.topology.clone(),
490 layer_layout,
491 layer_prefix_offsets,
492 state_segments: vec![PromptCacheStateSegment::new(id, 0..length)?],
493 };
494 selected.validate()?;
495 Ok(selected)
496 }
497}
498
499struct IdentityLayout<'a> {
500 layer_count: usize,
501 global_layer_start: usize,
502 global_layer_end: usize,
503 batch_size: usize,
504 layer_layout: &'a LayerSchedule<LayerCachePolicy>,
505 layer_prefix_offsets: &'a [i32],
506 state_segments: &'a [PromptCacheStateSegment],
507 topology: &'a PromptCacheTopology,
508}
509
510impl IdentityLayout<'_> {
511 fn validate(&self, subject: &str) -> Result<(), PromptCacheError> {
512 let owned = self
513 .global_layer_end
514 .checked_sub(self.global_layer_start)
515 .ok_or_else(|| {
516 PromptCacheError::Incompatible(format!("{subject} has an invalid layer range"))
517 })?;
518 if self.layer_count == 0
519 || self.global_layer_start >= self.global_layer_end
520 || self.global_layer_end > self.layer_count
521 || self.batch_size == 0
522 || self.batch_size > i32::MAX as usize
523 || self.layer_layout.len() != owned
524 || self.layer_prefix_offsets.len() != owned
525 || self.layer_prefix_offsets.iter().any(|offset| *offset > 0)
526 {
527 return Err(PromptCacheError::Incompatible(format!(
528 "{subject} supplied {} cache layouts and {} layer prefix offsets for {owned} owned layers",
529 self.layer_layout.len(),
530 self.layer_prefix_offsets.len()
531 )));
532 }
533 self.topology.validate()?;
534 validate_state_segments(self.state_segments, owned)
535 .map_err(|error| PromptCacheError::Incompatible(format!("{subject} {error}")))?;
536 for policy in self.layer_layout.iter() {
537 policy.validate()?;
538 }
539 Ok(())
540 }
541}
542
543pub fn validate_prompt_cache_model_identity(
545 expected: &PromptCacheDescriptor,
546 model: &PromptCacheModelIdentity,
547) -> Result<(), PromptCacheError> {
548 expected.validate()?;
549 model.validate()?;
550 macro_rules! require_equal {
551 ($field:ident) => {
552 if expected.$field != model.$field {
553 return Err(PromptCacheError::Incompatible(format!(
554 "caller descriptor {} does not match the loaded model",
555 stringify!($field)
556 )));
557 }
558 };
559 }
560 require_equal!(model_family);
561 require_equal!(effective_model_type);
562 require_equal!(architecture_fingerprint);
563 require_equal!(layer_count);
564 require_equal!(global_layer_start);
565 require_equal!(global_layer_end);
566 require_equal!(sink_tokens);
567 require_equal!(topology);
568 require_equal!(layer_layout);
569 require_equal!(layer_prefix_offsets);
570 require_equal!(state_segments);
571 Ok(())
572}
573
574fn validate_state_segments(
575 segments: &[PromptCacheStateSegment],
576 owned: usize,
577) -> Result<(), String> {
578 if segments.is_empty() {
579 return Err("has no named state segments".into());
580 }
581 let mut ids = BTreeSet::new();
582 let mut next = 0;
583 for segment in segments {
584 if segment.id.trim().is_empty() {
585 return Err("has an empty state segment identity".into());
586 }
587 if !ids.insert(segment.id.as_str()) {
588 return Err(format!(
589 "has duplicate state segment identity {:?}",
590 segment.id
591 ));
592 }
593 if segment.layers.start != next
594 || segment.layers.end <= segment.layers.start
595 || segment.layers.end > owned
596 {
597 return Err(format!(
598 "state segment {:?} range {}..{} does not continue an exact partition of {owned} owned layers",
599 segment.id, segment.layers.start, segment.layers.end
600 ));
601 }
602 next = segment.layers.end;
603 }
604 if next != owned {
605 return Err(format!(
606 "state segments cover {next} of {owned} owned layers"
607 ));
608 }
609 Ok(())
610}
611
612#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
614pub struct PromptCacheTopology {
615 stage: Option<(usize, usize)>,
617 shard: Option<(usize, usize)>,
619 addressable: Option<(usize, usize)>,
621 addressable_state_replicated: bool,
623}
624
625impl Default for PromptCacheTopology {
626 fn default() -> Self {
627 Self {
628 stage: None,
629 shard: None,
630 addressable: None,
631 addressable_state_replicated: true,
632 }
633 }
634}
635
636impl PromptCacheTopology {
637 pub fn new(
639 stage: Option<(usize, usize)>,
640 shard: Option<(usize, usize)>,
641 addressable: Option<(usize, usize)>,
642 addressable_state_replicated: bool,
643 ) -> Result<Self, PromptCacheError> {
644 let topology = Self {
645 stage,
646 shard,
647 addressable,
648 addressable_state_replicated,
649 };
650 topology.validate()?;
651 Ok(topology)
652 }
653
654 pub const fn stage(&self) -> Option<(usize, usize)> {
656 self.stage
657 }
658
659 pub const fn shard(&self) -> Option<(usize, usize)> {
661 self.shard
662 }
663
664 pub const fn addressable(&self) -> Option<(usize, usize)> {
666 self.addressable
667 }
668
669 pub const fn addressable_state_replicated(&self) -> bool {
671 self.addressable_state_replicated
672 }
673
674 pub fn validate(&self) -> Result<(), PromptCacheError> {
676 for (name, axis) in [
677 ("stage", self.stage),
678 ("state shard", self.shard),
679 ("addressable group", self.addressable),
680 ] {
681 if axis.is_some_and(|(size, rank)| size == 0 || rank >= size) {
682 return Err(PromptCacheError::Malformed(format!(
683 "invalid {name} topology"
684 )));
685 }
686 }
687 Ok(())
688 }
689
690 pub fn cache_rank_identity(&self) -> Option<CacheRankIdentity> {
692 (self.stage.is_some() || self.shard.is_some() || self.addressable.is_some()).then(|| {
693 CacheRankIdentity::new(
694 self.stage.map(|(_, rank)| rank),
695 self.shard.map(|(_, rank)| rank),
696 self.addressable.map(|(_, rank)| rank),
697 )
698 })
699 }
700}
701
702#[derive(Debug, Clone, Default)]
704pub struct PromptCacheOptions {
705 application_namespace: Option<String>,
707 replace_existing: bool,
709}
710
711impl PromptCacheOptions {
712 pub fn new(
714 application_namespace: Option<String>,
715 replace_existing: bool,
716 ) -> Result<Self, PromptCacheError> {
717 if application_namespace
718 .as_deref()
719 .is_some_and(|namespace| namespace.trim().is_empty())
720 {
721 return Err(PromptCacheError::Malformed(
722 "prompt-cache application namespace must not be empty".into(),
723 ));
724 }
725 Ok(Self {
726 application_namespace,
727 replace_existing,
728 })
729 }
730
731 pub fn application_namespace(&self) -> Option<&str> {
733 self.application_namespace.as_deref()
734 }
735
736 pub const fn replace_existing(&self) -> bool {
738 self.replace_existing
739 }
740}
741
742#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
744pub struct PromptCacheManifest {
745 pub schema_version: u32,
747 pub model_family: String,
749 pub effective_model_type: String,
751 pub checkpoint_fingerprint: String,
753 pub prefix_content_fingerprint: String,
755 pub architecture_fingerprint: String,
757 pub layer_count: usize,
759 pub global_layer_start: usize,
761 pub global_layer_end: usize,
763 pub block_size_tokens: i32,
765 pub batch_size: usize,
767 pub total_prefix_tokens: usize,
769 pub prefix_sha256: String,
771 pub layer_layout: LayerSchedule<LayerCachePolicy>,
773 pub layer_prefix_offsets: Vec<i32>,
775 pub state_segments: Vec<PromptCacheStateSegment>,
777 pub sink_tokens: usize,
779 pub topology: PromptCacheTopology,
781 pub application_namespace: Option<String>,
783 pub blocks: Vec<PromptCacheBlock>,
785 pub state_tensors: Vec<PromptCacheStateTensor>,
787}
788
789#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
791pub struct PromptCacheStateTensor {
792 pub owner: StateTensorOwner,
794 pub role: StateTensorRole,
796 pub shard: String,
798 pub array: String,
800 pub shape: Vec<i32>,
802 pub dtype: String,
804 pub logical_bytes: u64,
806 pub payload_sha256: String,
808}
809
810#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
812pub struct PromptCacheBlock {
813 pub global_layer: usize,
815 pub representation: CacheRepresentation,
817 pub start: i64,
819 pub end: i64,
821 pub rank: Option<CacheRankIdentity>,
823 pub shard: String,
825 pub first_array: String,
827 pub second_array: String,
829 pub first_shape: Vec<i32>,
831 pub second_shape: Vec<i32>,
833 pub first_dtype: String,
835 pub second_dtype: String,
837 pub logical_bytes: u64,
839 pub payload_sha256: String,
841}
842
843impl PromptCacheManifest {
844 pub fn validate(&self) -> Result<(), PromptCacheError> {
846 if self.schema_version != PROMPT_CACHE_SCHEMA_VERSION {
847 return Err(PromptCacheError::UnsupportedSchema(self.schema_version));
848 }
849 let owned = self.global_layer_end.checked_sub(self.global_layer_start);
850 if self.prefix_content_fingerprint.is_empty()
851 || self.block_size_tokens <= 0
852 || self.layer_count == 0
853 || self.global_layer_start >= self.global_layer_end
854 || self.global_layer_end > self.layer_count
855 || owned != Some(self.layer_layout.len())
856 || owned != Some(self.layer_prefix_offsets.len())
857 || self.batch_size == 0
858 || self.batch_size > i32::MAX as usize
859 || self.total_prefix_tokens == 0
860 || !is_sha256_hex(&self.prefix_sha256)
861 {
862 return Err(PromptCacheError::Malformed(
863 "invalid global cache dimensions".into(),
864 ));
865 }
866 self.topology.validate()?;
867 validate_state_segments(&self.state_segments, self.layer_layout.len())
868 .map_err(PromptCacheError::Malformed)?;
869 for (index, offset) in self.layer_prefix_offsets.iter().enumerate() {
870 layer_prefix_tokens(self.total_prefix_tokens, *offset).map_err(|error| {
871 PromptCacheError::Malformed(format!(
872 "invalid prefix frontier for global layer {}: {error}",
873 self.global_layer_start + index
874 ))
875 })?;
876 }
877 for (index, policy) in self.layer_layout.iter().enumerate() {
878 policy.validate().map_err(|error| {
879 PromptCacheError::Malformed(format!(
880 "invalid policy for global layer {}: {error}",
881 self.global_layer_start + index
882 ))
883 })?;
884 }
885 self.validate_blocks()?;
886 self.validate_state_tensors()?;
887 self.validate_coverage()
888 }
889
890 pub fn validate_compatibility(
892 &self,
893 expected: &PromptCacheDescriptor,
894 prefix_token_ids: &[u32],
895 ) -> Result<(), PromptCacheError> {
896 self.validate()?;
897 expected.validate()?;
898 macro_rules! require_equal {
899 ($field:ident) => {
900 if self.$field != expected.$field {
901 return Err(PromptCacheError::Incompatible(format!(
902 "{} mismatch",
903 stringify!($field)
904 )));
905 }
906 };
907 }
908 require_equal!(model_family);
909 require_equal!(effective_model_type);
910 require_equal!(checkpoint_fingerprint);
911 require_equal!(prefix_content_fingerprint);
912 require_equal!(architecture_fingerprint);
913 require_equal!(layer_count);
914 require_equal!(global_layer_start);
915 require_equal!(global_layer_end);
916 require_equal!(batch_size);
917 require_equal!(layer_layout);
918 require_equal!(layer_prefix_offsets);
919 require_equal!(state_segments);
920 require_equal!(sink_tokens);
921 require_equal!(topology);
922 if self.total_prefix_tokens != prefix_token_ids.len()
923 || self.prefix_sha256 != prompt_cache_token_fingerprint(prefix_token_ids)
924 {
925 return Err(PromptCacheError::PrefixIdentityMismatch);
926 }
927 Ok(())
928 }
929
930 fn validate_blocks(&self) -> Result<(), PromptCacheError> {
931 let mut previous = None;
932 for block in &self.blocks {
933 let layer_index = block
934 .global_layer
935 .checked_sub(self.global_layer_start)
936 .filter(|index| *index < self.layer_layout.len())
937 .ok_or_else(|| {
938 PromptCacheError::Malformed(format!(
939 "cache block layer {} is outside the owned range",
940 block.global_layer
941 ))
942 })?;
943 let layer_tokens = layer_prefix_tokens(
944 self.total_prefix_tokens,
945 self.layer_prefix_offsets[layer_index],
946 )?;
947 if block.start < 0
948 || block.end <= block.start
949 || block.end > layer_tokens as i64
950 || block.logical_bytes == 0
951 || block.first_shape.is_empty()
952 || block.second_shape.is_empty()
953 || !is_sha256_hex(&block.payload_sha256)
954 || !safe_relative_path(&block.shard)
955 {
956 return Err(PromptCacheError::Malformed(format!(
957 "invalid block at layer {} range {}..{}",
958 block.global_layer, block.start, block.end
959 )));
960 }
961 let order = (block.global_layer, block.start, block.end);
962 if previous.is_some_and(|value| value >= order) {
963 return Err(PromptCacheError::Malformed(format!(
964 "prompt-cache blocks are reordered or duplicated at layer {} range {}..{}",
965 block.global_layer, block.start, block.end
966 )));
967 }
968 previous = Some(order);
969 let policy = self.layer_layout.get(layer_index).expect("bounded");
970 let (representation, first_shape, second_shape) =
971 block_geometry(policy, self.batch_size, block.end - block.start)?;
972 if block.representation != representation
973 || block.first_shape != first_shape
974 || block.second_shape != second_shape
975 {
976 return Err(PromptCacheError::Malformed(format!(
977 "global layer {} payload geometry does not match its policy: actual {:?}/{:?}/{:?}, expected {:?}/{first_shape:?}/{second_shape:?}",
978 block.global_layer,
979 block.representation,
980 block.first_shape,
981 block.second_shape,
982 representation,
983 )));
984 }
985 if block.rank != self.topology.cache_rank_identity() {
986 return Err(PromptCacheError::Malformed(
987 "block rank identity does not match the recorded topology".into(),
988 ));
989 }
990 let names = array_names(block.representation);
991 if block.first_array != names.0
992 || block.second_array != names.1
993 || block.first_dtype != block.second_dtype
994 {
995 return Err(PromptCacheError::Malformed(
996 "block array names or dtypes do not match its representation".into(),
997 ));
998 }
999 }
1000 Ok(())
1001 }
1002
1003 fn validate_state_tensors(&self) -> Result<(), PromptCacheError> {
1004 let actual = self
1005 .state_tensors
1006 .iter()
1007 .map(|entry| (entry.owner, entry.role))
1008 .collect::<BTreeSet<_>>();
1009 if actual.len() != self.state_tensors.len() {
1010 return Err(PromptCacheError::Malformed(
1011 "fixed-state tensors contain duplicate owner/role entries".into(),
1012 ));
1013 }
1014 let mut expected = Vec::new();
1015 for (index, layer) in self.layer_layout.iter().enumerate() {
1016 let owner = StateTensorOwner::Layer(self.global_layer_start + index);
1017 let tokens =
1018 layer_prefix_tokens(self.total_prefix_tokens, self.layer_prefix_offsets[index])?;
1019 for policy in layer.fixed_state() {
1020 if (tokens != 0 && policy.is_required_for(tokens))
1023 || actual.contains(&(owner, policy.role))
1024 {
1025 expected.push((owner, policy, tokens));
1026 }
1027 }
1028 }
1029 if self.state_tensors.len() != expected.len() {
1030 return Err(PromptCacheError::Malformed(format!(
1031 "fixed-state tensor count {} does not match layout count {}",
1032 self.state_tensors.len(),
1033 expected.len()
1034 )));
1035 }
1036 for (entry, (owner, policy, tokens)) in self.state_tensors.iter().zip(expected) {
1037 if entry.owner != owner
1038 || entry.role != policy.role
1039 || entry.shape != policy.resolved_shape(self.batch_size, tokens)?
1040 || !policy.accepts_dtype_name(&entry.dtype)
1041 || entry.logical_bytes == 0
1042 || !is_sha256_hex(&entry.payload_sha256)
1043 || entry.array != "state"
1044 || !safe_relative_path(&entry.shard)
1045 {
1046 return Err(PromptCacheError::Malformed(format!(
1047 "fixed-state tensor {:?} for {:?} does not match its policy: shape {:?} and dtype {}, expected shape {:?}",
1048 entry.role,
1049 entry.owner,
1050 entry.shape,
1051 entry.dtype,
1052 policy.resolved_shape(self.batch_size, tokens)?,
1053 )));
1054 }
1055 }
1056 Ok(())
1057 }
1058
1059 fn validate_coverage(&self) -> Result<(), PromptCacheError> {
1060 let mut by_layer: BTreeMap<usize, Vec<&PromptCacheBlock>> = BTreeMap::new();
1061 for block in &self.blocks {
1062 by_layer.entry(block.global_layer).or_default().push(block);
1063 }
1064 for (index, policy) in self.layer_layout.iter().enumerate() {
1065 let layer = self.global_layer_start + index;
1066 let tokens =
1067 layer_prefix_tokens(self.total_prefix_tokens, self.layer_prefix_offsets[index])?;
1068 let mut blocks = by_layer.remove(&layer).unwrap_or_default();
1069 if policy.attention().is_none() {
1070 if !blocks.is_empty() {
1071 return Err(PromptCacheError::Malformed(format!(
1072 "stateless global layer {layer} has unexpected blocks"
1073 )));
1074 }
1075 continue;
1076 }
1077 if blocks.is_empty() {
1078 if tokens == 0 {
1079 continue;
1080 }
1081 return Err(PromptCacheError::Malformed(format!(
1082 "missing blocks for global layer {layer}"
1083 )));
1084 }
1085 blocks.sort_by_key(|block| block.start);
1086 let required = required_persisted_start(policy, tokens)?;
1087 let mut end = blocks[0].start;
1088 if end > required
1089 || (matches!(policy.attention(), Some(AttentionPolicy::Full)) && end != 0)
1090 {
1091 return Err(PromptCacheError::Malformed(format!(
1092 "global layer {layer} starts at {end}, but its policy requires history from {required}"
1093 )));
1094 }
1095 for block in blocks {
1096 if block.start != end {
1097 return Err(PromptCacheError::Malformed(format!(
1098 "gap or overlap at global layer {layer}: expected {end}, found {}",
1099 block.start
1100 )));
1101 }
1102 end = block.end;
1103 }
1104 if end != tokens as i64 {
1105 return Err(PromptCacheError::Malformed(format!(
1106 "global layer {layer} ends at {end}, expected {tokens}"
1107 )));
1108 }
1109 }
1110 Ok(())
1111 }
1112}
1113
1114fn block_geometry(
1115 policy: &LayerCachePolicy,
1116 batch_size: usize,
1117 token_count: i64,
1118) -> Result<(CacheRepresentation, Vec<i32>, Vec<i32>), PromptCacheError> {
1119 let batch = i32::try_from(batch_size)
1120 .map_err(|_| PromptCacheError::Malformed("prompt-cache batch exceeds i32".into()))?;
1121 let tokens = i32::try_from(token_count)
1122 .map_err(|_| PromptCacheError::Malformed("cache block token count exceeds i32".into()))?;
1123 match policy {
1124 LayerCachePolicy::NoState | LayerCachePolicy::FixedState { .. } => Err(
1125 PromptCacheError::Malformed("stateless layer has an attention payload".into()),
1126 ),
1127 LayerCachePolicy::KeyValue {
1128 num_key_value_heads,
1129 head_dim,
1130 ..
1131 }
1132 | LayerCachePolicy::KeyValueWithFixedState {
1133 num_key_value_heads,
1134 head_dim,
1135 ..
1136 } => {
1137 let shape = vec![
1138 batch,
1139 num_key_value_heads.get() as i32,
1140 tokens,
1141 head_dim.get() as i32,
1142 ];
1143 Ok((CacheRepresentation::KeyValue, shape.clone(), shape))
1144 }
1145 LayerCachePolicy::KeyOnly {
1146 num_key_heads,
1147 head_dim,
1148 ..
1149 }
1150 | LayerCachePolicy::KeyOnlyWithFixedState {
1151 num_key_heads,
1152 head_dim,
1153 ..
1154 } => Ok((
1155 CacheRepresentation::KeyValue,
1156 vec![
1157 batch,
1158 num_key_heads.get() as i32,
1159 tokens,
1160 head_dim.get() as i32,
1161 ],
1162 vec![batch, num_key_heads.get() as i32, tokens, 1],
1163 )),
1164 LayerCachePolicy::CompressedLatentRotary {
1165 latent_dim,
1166 rotary_dim,
1167 ..
1168 } => Ok((
1169 CacheRepresentation::CompressedLatentRotary,
1170 vec![batch, tokens, latent_dim.get() as i32],
1171 vec![batch, tokens, rotary_dim.get() as i32],
1172 )),
1173 }
1174}
1175
1176fn required_persisted_start(
1177 policy: &LayerCachePolicy,
1178 total_prefix_tokens: usize,
1179) -> Result<i64, PromptCacheError> {
1180 let total = i64::try_from(total_prefix_tokens).map_err(|_| {
1181 PromptCacheError::Malformed("prompt-cache prefix length exceeds i64".into())
1182 })?;
1183 match policy.attention() {
1184 None | Some(AttentionPolicy::Full) => Ok(0),
1185 Some(AttentionPolicy::Sliding { window }) => {
1186 Ok((total - i64::from(window.get() - 1)).max(0))
1187 }
1188 }
1189}
1190
1191fn layer_prefix_tokens(total: usize, offset: i32) -> Result<usize, PromptCacheError> {
1192 if offset > 0 {
1193 return Err(PromptCacheError::Malformed(
1194 "layer prefix offsets must not advance beyond the persisted prefix".into(),
1195 ));
1196 }
1197 total
1198 .checked_sub(offset.unsigned_abs() as usize)
1199 .ok_or_else(|| {
1200 PromptCacheError::Malformed(format!(
1201 "layer prefix offset {offset} precedes the start of a {total}-token prefix"
1202 ))
1203 })
1204}
1205
1206fn array_names(representation: CacheRepresentation) -> (&'static str, &'static str) {
1207 match representation {
1208 CacheRepresentation::KeyValue => ("keys", "values"),
1209 CacheRepresentation::CompressedLatentRotary => ("latent", "rotary_key"),
1210 }
1211}
1212
1213fn safe_relative_path(value: &str) -> bool {
1214 !value.is_empty()
1215 && !value.starts_with('/')
1216 && value
1217 .split('/')
1218 .all(|part| !part.is_empty() && part != "." && part != "..")
1219 && !value.contains('\\')
1220}
1221
1222fn is_sha256_hex(value: &str) -> bool {
1223 value.len() == 64
1224 && value
1225 .bytes()
1226 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
1227}
1228
1229pub fn derive_prompt_cache_architecture_fingerprint<I, K, V>(
1231 model_family: &str,
1232 fields: I,
1233) -> String
1234where
1235 I: IntoIterator<Item = (K, V)>,
1236 K: Into<String>,
1237 V: Into<String>,
1238{
1239 let mut fields = fields
1240 .into_iter()
1241 .map(|(key, value)| (key.into(), value.into()))
1242 .collect::<Vec<_>>();
1243 fields.sort_unstable();
1244 let mut hasher = Sha256::new();
1245 hash_component(&mut hasher, b"eredu-prompt-cache-architecture-v1");
1246 hash_component(&mut hasher, model_family.as_bytes());
1247 for (key, value) in fields {
1248 hash_component(&mut hasher, key.as_bytes());
1249 hash_component(&mut hasher, value.as_bytes());
1250 }
1251 format!("sha256:{}", hex(hasher.finalize()))
1252}
1253
1254pub fn prompt_cache_token_fingerprint(tokens: &[u32]) -> String {
1256 let mut hasher = Sha256::new();
1257 for token in tokens {
1258 hasher.update(token.to_le_bytes());
1259 }
1260 hex(hasher.finalize())
1261}
1262
1263fn hash_component(hasher: &mut Sha256, value: &[u8]) {
1264 hasher.update((value.len() as u64).to_le_bytes());
1265 hasher.update(value);
1266}
1267
1268fn hex(digest: impl AsRef<[u8]>) -> String {
1269 const HEX: &[u8; 16] = b"0123456789abcdef";
1270 let mut encoded = String::with_capacity(digest.as_ref().len() * 2);
1271 for &byte in digest.as_ref() {
1272 encoded.push(HEX[usize::from(byte >> 4)] as char);
1273 encoded.push(HEX[usize::from(byte & 0x0f)] as char);
1274 }
1275 encoded
1276}
1277
1278#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1280pub enum PromptCacheError {
1281 #[error(transparent)]
1283 Policy(#[from] CachePolicyError),
1284 #[error("unsupported prompt cache schema version {0}")]
1286 UnsupportedSchema(u32),
1287 #[error("malformed prompt cache manifest: {0}")]
1289 Malformed(String),
1290 #[error("incompatible prompt cache: {0}")]
1292 Incompatible(String),
1293 #[error("prompt cache prefix token identity does not match")]
1295 PrefixIdentityMismatch,
1296}
1297
1298#[cfg(test)]
1299mod tests {
1300 use super::*;
1301
1302 fn manifest() -> PromptCacheManifest {
1303 let layout = LayerSchedule::new(
1304 1,
1305 vec![LayerCachePolicy::key_value(AttentionPolicy::Full, 2, 4).unwrap()],
1306 )
1307 .unwrap();
1308 PromptCacheManifest {
1309 schema_version: PROMPT_CACHE_SCHEMA_VERSION,
1310 model_family: "llama".into(),
1311 effective_model_type: "llama".into(),
1312 checkpoint_fingerprint: "checkpoint".into(),
1313 prefix_content_fingerprint: "content".into(),
1314 architecture_fingerprint: "architecture".into(),
1315 layer_count: 1,
1316 global_layer_start: 0,
1317 global_layer_end: 1,
1318 block_size_tokens: 2,
1319 batch_size: 1,
1320 total_prefix_tokens: 2,
1321 prefix_sha256: prompt_cache_token_fingerprint(&[7, 8]),
1322 layer_layout: layout,
1323 layer_prefix_offsets: vec![0],
1324 state_segments: vec![PromptCacheStateSegment::new("state", 0..1).unwrap()],
1325 sink_tokens: 0,
1326 topology: PromptCacheTopology::default(),
1327 application_namespace: None,
1328 blocks: vec![PromptCacheBlock {
1329 global_layer: 0,
1330 representation: CacheRepresentation::KeyValue,
1331 start: 0,
1332 end: 2,
1333 rank: None,
1334 shard: "blocks/layer-0.safetensors".into(),
1335 first_array: "keys".into(),
1336 second_array: "values".into(),
1337 first_shape: vec![1, 2, 2, 4],
1338 second_shape: vec![1, 2, 2, 4],
1339 first_dtype: "Float16".into(),
1340 second_dtype: "Float16".into(),
1341 logical_bytes: 64,
1342 payload_sha256: "0".repeat(64),
1343 }],
1344 state_tensors: vec![],
1345 }
1346 }
1347
1348 #[test]
1349 fn manifest_round_trips_and_validates_without_a_backend() {
1350 let manifest = manifest();
1351 manifest.validate().unwrap();
1352 let json = serde_json::to_string(&manifest).unwrap();
1353 let restored: PromptCacheManifest = serde_json::from_str(&json).unwrap();
1354 restored.validate().unwrap();
1355 assert_eq!(restored, manifest);
1356 }
1357
1358 #[test]
1359 fn descriptor_derives_every_model_owned_field_from_identity() {
1360 let manifest = manifest();
1361 let identity = PromptCacheModelIdentity {
1362 model_family: manifest.model_family.clone(),
1363 effective_model_type: manifest.effective_model_type.clone(),
1364 architecture_fingerprint: manifest.architecture_fingerprint.clone(),
1365 layer_count: manifest.layer_count,
1366 global_layer_start: manifest.global_layer_start,
1367 global_layer_end: manifest.global_layer_end,
1368 sink_tokens: manifest.sink_tokens,
1369 topology: manifest.topology.clone(),
1370 layer_layout: manifest.layer_layout.clone(),
1371 layer_prefix_offsets: manifest.layer_prefix_offsets.clone(),
1372 state_segments: manifest.state_segments.clone(),
1373 };
1374
1375 let descriptor = PromptCacheDescriptor::from_model_identity(
1376 identity.clone(),
1377 "caller-checkpoint",
1378 "caller-prefix-content",
1379 3,
1380 )
1381 .unwrap();
1382
1383 validate_prompt_cache_model_identity(&descriptor, &identity).unwrap();
1384 assert_eq!(descriptor.checkpoint_fingerprint, "caller-checkpoint");
1385 assert_eq!(
1386 descriptor.prefix_content_fingerprint,
1387 "caller-prefix-content"
1388 );
1389 assert_eq!(descriptor.batch_size, 3);
1390 assert!(
1391 PromptCacheDescriptor::from_model_identity(identity, "checkpoint", "prefix", 0)
1392 .is_err()
1393 );
1394 }
1395
1396 #[test]
1397 fn architecture_fingerprint_uses_the_eredu_domain() {
1398 let fingerprint = derive_prompt_cache_architecture_fingerprint(
1399 "llama",
1400 [("layers", "32"), ("hidden_size", "4096")],
1401 );
1402 assert_eq!(
1403 fingerprint,
1404 "sha256:9ee0b30ea8687d04eb4b65db3a58ccfff0a72bdd502805e9fdd6edb223ca5949"
1405 );
1406 }
1407
1408 #[test]
1409 fn zero_frontier_prediction_state_needs_no_materialized_tensor() {
1410 let recurrent = crate::cache::StateTensorPolicy::new(
1411 StateTensorRole::Recurrent,
1412 vec![crate::cache::StateTensorDimension::Batch],
1413 crate::cache::StateTensorDtype::Floating,
1414 crate::cache::MutableStateResidency::LayerScopedOffloadable,
1415 )
1416 .unwrap();
1417 let mut value = manifest();
1418 value.total_prefix_tokens = 1;
1419 value.prefix_sha256 = prompt_cache_token_fingerprint(&[7]);
1420 value.layer_prefix_offsets = vec![-1];
1421 value.layer_layout = LayerSchedule::new(
1422 1,
1423 vec![LayerCachePolicy::fixed_only(vec![recurrent]).unwrap()],
1424 )
1425 .unwrap();
1426 value.blocks.clear();
1427 value.state_tensors.clear();
1428 value.validate().unwrap();
1429 }
1430
1431 #[test]
1432 fn rejects_bad_topology_geometry_coverage_and_paths() {
1433 let mut value = manifest();
1434 value.topology.shard = Some((1, 1));
1435 assert!(value.validate().is_err());
1436 let mut value = manifest();
1437 value.blocks[0].first_shape[2] = 1;
1438 assert!(value.validate().is_err());
1439 let mut value = manifest();
1440 value.blocks[0].shard = "../escape".into();
1441 assert!(value.validate().is_err());
1442 }
1443
1444 #[test]
1445 fn identity_and_prefix_compatibility_fail_closed() {
1446 let manifest = manifest();
1447 let descriptor = PromptCacheDescriptor {
1448 model_family: manifest.model_family.clone(),
1449 effective_model_type: manifest.effective_model_type.clone(),
1450 checkpoint_fingerprint: manifest.checkpoint_fingerprint.clone(),
1451 prefix_content_fingerprint: manifest.prefix_content_fingerprint.clone(),
1452 architecture_fingerprint: manifest.architecture_fingerprint.clone(),
1453 layer_count: 1,
1454 global_layer_start: 0,
1455 global_layer_end: 1,
1456 batch_size: 1,
1457 layer_layout: manifest.layer_layout.clone(),
1458 layer_prefix_offsets: vec![0],
1459 state_segments: manifest.state_segments.clone(),
1460 sink_tokens: 0,
1461 topology: PromptCacheTopology::default(),
1462 };
1463 manifest
1464 .validate_compatibility(&descriptor, &[7, 8])
1465 .unwrap();
1466 assert!(manifest
1467 .validate_compatibility(&descriptor, &[8, 7])
1468 .is_err());
1469 let mut renamed = descriptor.clone();
1470 renamed.state_segments = vec![PromptCacheStateSegment::new("renamed", 0..1).unwrap()];
1471 assert!(matches!(
1472 manifest.validate_compatibility(&renamed, &[7, 8]),
1473 Err(PromptCacheError::Incompatible(_))
1474 ));
1475 let mut invalid = descriptor;
1476 invalid.layer_prefix_offsets[0] = 1;
1477 assert!(matches!(
1478 invalid.validate(),
1479 Err(PromptCacheError::Incompatible(_))
1480 ));
1481
1482 let mut malformed = manifest.clone();
1483 malformed.state_segments = vec![PromptCacheStateSegment::new("state", 0..2).unwrap()];
1484 assert!(matches!(
1485 malformed.validate(),
1486 Err(PromptCacheError::Malformed(_))
1487 ));
1488 }
1489}