Skip to main content

eredu_core/cache/
prompt.rs

1//! Portable reusable prompt-cache identity, catalog, and validation.
2
3use 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
18/// Current reusable prompt-cache schema version.
19pub const PROMPT_CACHE_SCHEMA_VERSION: u32 = 8;
20
21/// One named contiguous state range in a portable prompt-cache identity.
22#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
23pub struct PromptCacheStateSegment {
24    id: String,
25    layers: Range<usize>,
26}
27
28impl PromptCacheStateSegment {
29    /// Creates a named non-empty local state range.
30    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    /// Returns the architecture-declared stable segment identity.
46    pub fn id(&self) -> &str {
47        &self.id
48    }
49
50    /// Returns the segment's local range in the identity's ordered layout.
51    pub fn layers(&self) -> Range<usize> {
52        self.layers.clone()
53    }
54}
55
56/// Caller-supplied identity and geometry for a reusable prefix cache.
57#[derive(Debug, Clone, Eq, Hash, PartialEq)]
58pub struct PromptCacheDescriptor {
59    /// Stable architecture family.
60    model_family: String,
61    /// Effective normalized model type.
62    effective_model_type: String,
63    /// Caller-verified checkpoint identity.
64    checkpoint_fingerprint: String,
65    /// Identity of all content that produced the cached activations.
66    prefix_content_fingerprint: String,
67    /// Cache-relevant architecture identity.
68    architecture_fingerprint: String,
69    /// Total model layer count.
70    layer_count: usize,
71    /// Inclusive first global layer stored by this rank.
72    global_layer_start: usize,
73    /// Exclusive global layer boundary stored by this rank.
74    global_layer_end: usize,
75    /// Prefix batch size.
76    batch_size: usize,
77    /// Ordered cache layout for the owned layer range.
78    layer_layout: LayerSchedule<LayerCachePolicy>,
79    /// Per-layer processed-token delta relative to the persisted prefix.
80    layer_prefix_offsets: Vec<i32>,
81    /// Architecture-declared named ranges in the ordered state layout.
82    state_segments: Vec<PromptCacheStateSegment>,
83    /// Attention sink or pinned-prefix token count.
84    sink_tokens: usize,
85    /// Distributed rank-local layout.
86    topology: PromptCacheTopology,
87}
88
89impl PromptCacheDescriptor {
90    /// Creates and validates a complete reusable prefix-cache descriptor.
91    #[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    /// Stable architecture family.
142    pub fn model_family(&self) -> &str {
143        &self.model_family
144    }
145    /// Effective normalized model type.
146    pub fn effective_model_type(&self) -> &str {
147        &self.effective_model_type
148    }
149    /// Caller-verified checkpoint identity.
150    pub fn checkpoint_fingerprint(&self) -> &str {
151        &self.checkpoint_fingerprint
152    }
153    /// Prefix-content identity.
154    pub fn prefix_content_fingerprint(&self) -> &str {
155        &self.prefix_content_fingerprint
156    }
157    /// Cache-relevant architecture identity.
158    pub fn architecture_fingerprint(&self) -> &str {
159        &self.architecture_fingerprint
160    }
161    /// Total model layer count.
162    pub const fn layer_count(&self) -> usize {
163        self.layer_count
164    }
165    /// Inclusive first global layer stored by this rank.
166    pub const fn global_layer_start(&self) -> usize {
167        self.global_layer_start
168    }
169    /// Exclusive global layer boundary stored by this rank.
170    pub const fn global_layer_end(&self) -> usize {
171        self.global_layer_end
172    }
173    /// Prefix batch size.
174    pub const fn batch_size(&self) -> usize {
175        self.batch_size
176    }
177    /// Ordered cache layout.
178    pub const fn layer_layout(&self) -> &LayerSchedule<LayerCachePolicy> {
179        &self.layer_layout
180    }
181    /// Per-layer processed-token deltas.
182    pub fn layer_prefix_offsets(&self) -> &[i32] {
183        &self.layer_prefix_offsets
184    }
185    /// Named state-layout ranges.
186    pub fn state_segments(&self) -> &[PromptCacheStateSegment] {
187        &self.state_segments
188    }
189    /// Attention sink token count.
190    pub const fn sink_tokens(&self) -> usize {
191        self.sink_tokens
192    }
193    /// Distributed rank-local layout.
194    pub const fn topology(&self) -> &PromptCacheTopology {
195        &self.topology
196    }
197    /// Replaces the distributed topology and revalidates the descriptor.
198    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    /// Replaces the cache-relevant architecture fingerprint.
207    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    /// Replaces the total model layer count while preserving the owned range.
221    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    /// Derives every model-owned field from a prepared model identity.
227    ///
228    /// The checkpoint and prefix-content fingerprints remain caller-owned
229    /// because they identify the concrete weights and processed input rather
230    /// than model structure.
231    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    /// Validates the complete portable identity and cache geometry.
258    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/// Cache-relevant structure derived from a prepared model.
274#[derive(Debug, Clone, Eq, Hash, PartialEq)]
275pub struct PromptCacheModelIdentity {
276    /// Stable architecture family.
277    model_family: String,
278    /// Effective normalized model type.
279    effective_model_type: String,
280    /// Cache-relevant architecture identity.
281    architecture_fingerprint: String,
282    /// Total model layer count.
283    layer_count: usize,
284    /// Inclusive first global layer owned by this model instance.
285    global_layer_start: usize,
286    /// Exclusive global layer boundary owned by this model instance.
287    global_layer_end: usize,
288    /// Attention sink or pinned-prefix token count.
289    sink_tokens: usize,
290    /// Distributed rank-local layout.
291    topology: PromptCacheTopology,
292    /// Ordered cache layout for the owned layer range.
293    layer_layout: LayerSchedule<LayerCachePolicy>,
294    /// Per-layer processed-token delta relative to the persisted prefix.
295    layer_prefix_offsets: Vec<i32>,
296    /// Architecture-declared named ranges in the ordered state layout.
297    state_segments: Vec<PromptCacheStateSegment>,
298}
299
300impl PromptCacheModelIdentity {
301    /// Creates and validates cache-relevant prepared-model identity.
302    #[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    /// Stable architecture family.
345    pub fn model_family(&self) -> &str {
346        &self.model_family
347    }
348    /// Effective normalized model type.
349    pub fn effective_model_type(&self) -> &str {
350        &self.effective_model_type
351    }
352    /// Cache-relevant architecture identity.
353    pub fn architecture_fingerprint(&self) -> &str {
354        &self.architecture_fingerprint
355    }
356    /// Total model layer count.
357    pub const fn layer_count(&self) -> usize {
358        self.layer_count
359    }
360    /// Inclusive first global layer owned locally.
361    pub const fn global_layer_start(&self) -> usize {
362        self.global_layer_start
363    }
364    /// Exclusive global layer boundary owned locally.
365    pub const fn global_layer_end(&self) -> usize {
366        self.global_layer_end
367    }
368    /// Attention sink token count.
369    pub const fn sink_tokens(&self) -> usize {
370        self.sink_tokens
371    }
372    /// Distributed rank-local layout.
373    pub const fn topology(&self) -> &PromptCacheTopology {
374        &self.topology
375    }
376    /// Ordered cache layout.
377    pub const fn layer_layout(&self) -> &LayerSchedule<LayerCachePolicy> {
378        &self.layer_layout
379    }
380    /// Per-layer processed-token deltas.
381    pub fn layer_prefix_offsets(&self) -> &[i32] {
382        &self.layer_prefix_offsets
383    }
384    /// Named state-layout ranges.
385    pub fn state_segments(&self) -> &[PromptCacheStateSegment] {
386        &self.state_segments
387    }
388    /// Builds an ordered ordinary key/value layout from runtime window values.
389    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    /// Builds a uniform compressed-latent layout.
408    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    /// Validates the owned layer range and every policy.
428    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    /// Returns one architecture-declared state segment by stable identity.
443    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    /// Selects one named state segment as a validated standalone identity.
456    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
543/// Verifies that a caller descriptor was derived from the prepared model.
544pub 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/// Rank-local topology recorded in a prompt-cache manifest.
613#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
614pub struct PromptCacheTopology {
615    /// Ordered-stage partition size and rank.
616    stage: Option<(usize, usize)>,
617    /// State-shard partition size and rank.
618    shard: Option<(usize, usize)>,
619    /// Addressable-group size and rank.
620    addressable: Option<(usize, usize)>,
621    /// Whether cache state is replicated on the addressable axis.
622    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    /// Creates and validates an exact cache-placement topology.
638    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    /// Returns ordered-stage size and rank when partitioned.
655    pub const fn stage(&self) -> Option<(usize, usize)> {
656        self.stage
657    }
658
659    /// Returns state-shard size and rank when partitioned.
660    pub const fn shard(&self) -> Option<(usize, usize)> {
661        self.shard
662    }
663
664    /// Returns addressable-group size and rank when partitioned.
665    pub const fn addressable(&self) -> Option<(usize, usize)> {
666        self.addressable
667    }
668
669    /// Returns whether cache state is replicated on the addressable axis.
670    pub const fn addressable_state_replicated(&self) -> bool {
671        self.addressable_state_replicated
672    }
673
674    /// Validates every optional world-size/rank pair.
675    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    /// Returns the rank identity stored on cache blocks, if distributed.
691    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/// Explicit publication behavior for a reusable prefix cache.
703#[derive(Debug, Clone, Default)]
704pub struct PromptCacheOptions {
705    /// Optional application grouping label; never used for compatibility.
706    application_namespace: Option<String>,
707    /// Allows atomically replacing an existing destination.
708    replace_existing: bool,
709}
710
711impl PromptCacheOptions {
712    /// Creates validated prompt-cache publication options.
713    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    /// Returns the optional application grouping label.
732    pub fn application_namespace(&self) -> Option<&str> {
733        self.application_namespace.as_deref()
734    }
735
736    /// Returns whether an existing destination may be replaced atomically.
737    pub const fn replace_existing(&self) -> bool {
738        self.replace_existing
739    }
740}
741
742/// Versioned metadata inspectable without loading backend arrays.
743#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
744pub struct PromptCacheManifest {
745    /// Persistence schema version.
746    pub schema_version: u32,
747    /// Model architecture family.
748    pub model_family: String,
749    /// Effective normalized model type.
750    pub effective_model_type: String,
751    /// Caller-selected checkpoint identity.
752    pub checkpoint_fingerprint: String,
753    /// Identity of all content that produced this prefix.
754    pub prefix_content_fingerprint: String,
755    /// Cache-relevant architecture identity.
756    pub architecture_fingerprint: String,
757    /// Total model layer count.
758    pub layer_count: usize,
759    /// Inclusive first global layer represented locally.
760    pub global_layer_start: usize,
761    /// Exclusive global layer boundary represented locally.
762    pub global_layer_end: usize,
763    /// Block size used by the producer.
764    pub block_size_tokens: i32,
765    /// Prefix batch size.
766    pub batch_size: usize,
767    /// Exact prefix token count.
768    pub total_prefix_tokens: usize,
769    /// SHA-256 over little-endian prefix token IDs.
770    pub prefix_sha256: String,
771    /// Ordered cache layout for the owned layer range.
772    pub layer_layout: LayerSchedule<LayerCachePolicy>,
773    /// Per-layer processed-token delta relative to the prefix.
774    pub layer_prefix_offsets: Vec<i32>,
775    /// Architecture-declared named ranges in the ordered state layout.
776    pub state_segments: Vec<PromptCacheStateSegment>,
777    /// Pinned prefix or sink token count.
778    pub sink_tokens: usize,
779    /// Distributed rank-local representation.
780    pub topology: PromptCacheTopology,
781    /// Optional non-authoritative application grouping label.
782    pub application_namespace: Option<String>,
783    /// Ordered immutable cache blocks.
784    pub blocks: Vec<PromptCacheBlock>,
785    /// Ordered fixed-size state tensors.
786    pub state_tensors: Vec<PromptCacheStateTensor>,
787}
788
789/// One independently validated fixed-size state tensor catalog entry.
790#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
791pub struct PromptCacheStateTensor {
792    /// Layer owner.
793    pub owner: StateTensorOwner,
794    /// Semantic role declared by the canonical layout.
795    pub role: StateTensorRole,
796    /// Safe relative backend shard path.
797    pub shard: String,
798    /// Array name within the shard.
799    pub array: String,
800    /// Exact stored shape.
801    pub shape: Vec<i32>,
802    /// Exact stored dtype.
803    pub dtype: String,
804    /// Logical bytes in the array.
805    pub logical_bytes: u64,
806    /// SHA-256 of the exact payload bytes.
807    pub payload_sha256: String,
808}
809
810/// One cache block catalog entry in a prompt-cache manifest.
811#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
812pub struct PromptCacheBlock {
813    /// Architecture-global layer identity.
814    pub global_layer: usize,
815    /// Stored attention representation.
816    pub representation: CacheRepresentation,
817    /// Inclusive absolute token position.
818    pub start: i64,
819    /// Exclusive absolute token position.
820    pub end: i64,
821    /// Optional rank identity.
822    pub rank: Option<CacheRankIdentity>,
823    /// Safe relative backend shard path.
824    pub shard: String,
825    /// First array name.
826    pub first_array: String,
827    /// Second array name.
828    pub second_array: String,
829    /// First array shape.
830    pub first_shape: Vec<i32>,
831    /// Second array shape.
832    pub second_shape: Vec<i32>,
833    /// First array dtype.
834    pub first_dtype: String,
835    /// Second array dtype.
836    pub second_dtype: String,
837    /// Logical bytes in both arrays.
838    pub logical_bytes: u64,
839    /// SHA-256 of the exact payload bytes.
840    pub payload_sha256: String,
841}
842
843impl PromptCacheManifest {
844    /// Validates all backend-independent schema, geometry, and coverage rules.
845    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    /// Validates compatibility with a caller descriptor and exact prefix IDs.
891    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                // A zero-token frontier has no materialized recurrent value,
1021                // even when that value is required once execution begins.
1022                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
1229/// Derives a stable cache architecture fingerprint from ordered semantic fields.
1230pub 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
1254/// Hashes exact prefix token IDs as little-endian `u32` values.
1255pub 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/// Invalid reusable prompt-cache identity, schema, or catalog.
1279#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1280pub enum PromptCacheError {
1281    /// A layer or state policy is invalid.
1282    #[error(transparent)]
1283    Policy(#[from] CachePolicyError),
1284    /// The persistence schema version is unsupported.
1285    #[error("unsupported prompt cache schema version {0}")]
1286    UnsupportedSchema(u32),
1287    /// The portable manifest structure is malformed.
1288    #[error("malformed prompt cache manifest: {0}")]
1289    Malformed(String),
1290    /// The prepared model or caller identity differs from the producer.
1291    #[error("incompatible prompt cache: {0}")]
1292    Incompatible(String),
1293    /// Exact prefix IDs differ from the persisted identity.
1294    #[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}