eredu_core/cache/
policy.rs

1//! Durable cache identity, geometry, and state-residency contracts.
2
3use std::{collections::BTreeSet, num::NonZeroU32};
4
5use serde::{Deserialize, Serialize};
6
7use crate::attention::AttentionPolicy;
8
9/// Representation stored atomically in one cache block.
10#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum CacheRepresentation {
13    /// Standard attention keys and values.
14    KeyValue,
15    /// Compressed latent state and rotary keys.
16    CompressedLatentRotary,
17}
18
19/// Optional rank identity included in a stable cache block identifier.
20#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
21pub struct CacheRankIdentity {
22    /// Ordered-stage rank, when stage partitioning is active.
23    stage_rank: Option<usize>,
24    /// State-shard rank, when cache state is partitioned.
25    shard_rank: Option<usize>,
26    /// Addressable-group rank for replicated cache state.
27    addressable_rank: Option<usize>,
28}
29
30impl CacheRankIdentity {
31    /// Creates a generic rank identity for persisted cache state.
32    pub const fn new(
33        stage_rank: Option<usize>,
34        shard_rank: Option<usize>,
35        addressable_rank: Option<usize>,
36    ) -> Self {
37        Self {
38            stage_rank,
39            shard_rank,
40            addressable_rank,
41        }
42    }
43
44    /// Returns the ordered-stage rank, when present.
45    pub const fn stage_rank(&self) -> Option<usize> {
46        self.stage_rank
47    }
48
49    /// Returns the state-shard rank, when present.
50    pub const fn shard_rank(&self) -> Option<usize> {
51        self.shard_rank
52    }
53
54    /// Returns the addressable-group rank, when present.
55    pub const fn addressable_rank(&self) -> Option<usize> {
56        self.addressable_rank
57    }
58}
59
60/// Stable identity for one immutable sealed cache block.
61#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
62pub struct CacheBlockId {
63    /// Identity shared by every block in one live cache.
64    pub session_id: u64,
65    /// Architecture-global decoder layer index.
66    pub global_layer: usize,
67    /// Stored attention representation.
68    pub representation: CacheRepresentation,
69    /// Inclusive absolute token position.
70    pub start: i64,
71    /// Exclusive absolute token position.
72    pub end: i64,
73    /// Rank-local ownership identity.
74    pub rank: Option<CacheRankIdentity>,
75}
76
77/// Logical location of a sealed cache block.
78#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
79#[serde(rename_all = "snake_case")]
80pub enum CacheTier {
81    /// Available to execution without a catalog load.
82    Device,
83    /// Evaluated CPU-resident state without an execution-device copy.
84    Host,
85    /// Stored in a backend-owned persistent shard.
86    Disk,
87}
88
89/// Exact state kind, attention policy, and tensor geometry for one decoder layer.
90#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
91#[serde(rename_all = "snake_case")]
92pub enum LayerCachePolicy {
93    /// This layer contributes no independently persisted state.
94    NoState,
95    /// Ordinary attention keys and values.
96    KeyValue {
97        /// Exact full or sliding attention range.
98        attention: AttentionPolicy,
99        /// Rank-local key/value head count.
100        num_key_value_heads: NonZeroU32,
101        /// Per-head key/value dimension.
102        head_dim: NonZeroU32,
103    },
104    /// Attention history whose value payload is intentionally empty.
105    KeyOnly {
106        /// Exact full or sliding attention range.
107        attention: AttentionPolicy,
108        /// Rank-local key head count.
109        num_key_heads: NonZeroU32,
110        /// Per-head key dimension.
111        head_dim: NonZeroU32,
112    },
113    /// Compressed latent state plus rotary keys.
114    CompressedLatentRotary {
115        /// Exact full or sliding attention range.
116        attention: AttentionPolicy,
117        /// Compressed latent width.
118        latent_dim: NonZeroU32,
119        /// Rotary-key width.
120        rotary_dim: NonZeroU32,
121    },
122    /// Fixed-size recurrent or convolution state without attention.
123    FixedState {
124        /// Ordered tensors required to resume this layer.
125        tensors: Vec<StateTensorPolicy>,
126    },
127    /// Ordinary attention plus fixed-size state.
128    KeyValueWithFixedState {
129        /// Exact full or sliding attention range.
130        attention: AttentionPolicy,
131        /// Rank-local key/value head count.
132        num_key_value_heads: NonZeroU32,
133        /// Per-head key/value dimension.
134        head_dim: NonZeroU32,
135        /// Ordered additional tensors.
136        tensors: Vec<StateTensorPolicy>,
137    },
138    /// Key-only attention plus fixed-size state.
139    KeyOnlyWithFixedState {
140        /// Exact full or sliding attention range.
141        attention: AttentionPolicy,
142        /// Rank-local key head count.
143        num_key_heads: NonZeroU32,
144        /// Per-head key dimension.
145        head_dim: NonZeroU32,
146        /// Ordered additional tensors.
147        tensors: Vec<StateTensorPolicy>,
148    },
149}
150
151/// Semantic role of one non-attention cache tensor.
152#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
153#[serde(rename_all = "snake_case")]
154pub enum StateTensorRole {
155    /// Bounded causal-convolution history.
156    Convolution {
157        /// Stable slot within the layer's convolution states.
158        slot: u32,
159    },
160    /// Recurrent transition or linear-attention state.
161    Recurrent,
162    /// Prepared multimodal prefix embeddings.
163    PrefixEmbedding,
164    /// Model-global multimodal position offset.
165    PositionDelta,
166    /// One tensor in an append-only token-pooling stream.
167    Pooling {
168        /// Stable stream slot within the owning layer.
169        stream: u32,
170        /// Exact component of the pooling state.
171        component: PoolingStateComponent,
172    },
173}
174
175/// Semantic component of one append-only pooling stream.
176#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
177#[serde(rename_all = "snake_case")]
178pub enum PoolingStateComponent {
179    /// Source values waiting for a complete pooling group.
180    PendingValues,
181    /// Source gate logits waiting for a complete pooling group.
182    PendingGates,
183    /// Complete pooled output history.
184    Pooled,
185    /// Source values retained for an overlapping group.
186    OverlapValues,
187    /// Source gate logits retained for an overlapping group.
188    OverlapGates,
189}
190
191/// Runtime ownership behavior for live model state.
192#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
193#[serde(rename_all = "snake_case")]
194pub enum StateResidencyClass {
195    /// Small mutable state that remains on the execution device.
196    AlwaysDeviceMutable,
197    /// Append-only state that becomes immutable blocks before paging.
198    SealablePaged,
199    /// Mutable state promoted only for its owning layer.
200    LayerScopedOffloadable,
201}
202
203/// Residency behaviors valid for mutable fixed-state tensors.
204#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
205#[serde(rename_all = "snake_case")]
206pub enum MutableStateResidency {
207    /// Small mutable state that remains on the execution device.
208    AlwaysDeviceMutable,
209    /// Mutable state promoted only for its owning layer.
210    LayerScopedOffloadable,
211}
212
213impl From<MutableStateResidency> for StateResidencyClass {
214    fn from(value: MutableStateResidency) -> Self {
215        match value {
216            MutableStateResidency::AlwaysDeviceMutable => Self::AlwaysDeviceMutable,
217            MutableStateResidency::LayerScopedOffloadable => Self::LayerScopedOffloadable,
218        }
219    }
220}
221
222/// One dimension in a persisted fixed-state tensor.
223#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
224#[serde(rename_all = "snake_case")]
225pub enum StateTensorDimension {
226    /// Manifest batch size.
227    Batch,
228    /// Exact prompt token count.
229    PrefixTokens,
230    /// Quotient of prompt tokens and a positive divisor.
231    PrefixTokensDiv(NonZeroU32),
232    /// Remainder of prompt tokens and a positive divisor.
233    PrefixTokensRem(NonZeroU32),
234    /// Positive architecture-defined dimension.
235    Fixed(NonZeroU32),
236    /// Scalar dimension list marker; valid only as the sole entry.
237    Scalar,
238}
239
240/// Condition under which a state tensor must be materialized.
241#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
242#[serde(rename_all = "snake_case")]
243pub enum StateTensorPresence {
244    /// Every persisted cache materializes the tensor.
245    Required,
246    /// The tensor may be present independently of prefix geometry.
247    Optional,
248    /// Present exactly when the prefix has a non-zero remainder.
249    PrefixRemainderNonZero(NonZeroU32),
250    /// Present exactly when the prefix contains one complete group.
251    PrefixAtLeast(NonZeroU32),
252}
253
254/// Accepted dtype family for one fixed-state tensor.
255#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
256#[serde(rename_all = "snake_case")]
257pub enum StateTensorDtype {
258    /// Any floating dtype.
259    Floating,
260    /// Exactly IEEE F32.
261    Float32,
262    /// Signed 32-bit integer.
263    Int32,
264    /// Unsigned 32-bit integer.
265    Uint32,
266}
267
268/// Exact semantic role, symbolic shape, and dtype contract for a state tensor.
269#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
270pub struct StateTensorPolicy {
271    /// Meaning of this tensor within its owner.
272    pub role: StateTensorRole,
273    /// Symbolic shape resolved from batch and prefix geometry.
274    pub shape: Vec<StateTensorDimension>,
275    /// Accepted dtype family.
276    pub dtype: StateTensorDtype,
277    /// Authoritative live-state residency behavior.
278    pub residency: StateResidencyClass,
279    /// Condition under which persisted caches materialize this tensor.
280    pub presence: StateTensorPresence,
281}
282
283/// Owner of one persisted non-attention state tensor.
284#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
285#[serde(rename_all = "snake_case")]
286pub enum StateTensorOwner {
287    /// Architecture-global decoder layer index.
288    Layer(usize),
289}
290
291/// Stable semantic identity for one independently managed runtime component.
292#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
293#[serde(rename_all = "snake_case")]
294pub enum StateComponentRole {
295    /// Ordinary attention keys.
296    AttentionKeys,
297    /// Ordinary attention values.
298    AttentionValues,
299    /// Head-independent compressed latent key/value state.
300    CompressedLatent,
301    /// Rotary keys paired with compressed latent state.
302    RotaryKeys,
303    /// One fixed-size or pooling tensor.
304    Fixed(StateTensorRole),
305}
306
307impl StateComponentRole {
308    /// Returns a stable checkpoint/runtime component name.
309    pub fn stable_name(self) -> String {
310        match self {
311            Self::AttentionKeys => "attention.keys".into(),
312            Self::AttentionValues => "attention.values".into(),
313            Self::CompressedLatent => "attention.compressed_latent".into(),
314            Self::RotaryKeys => "attention.rotary_keys".into(),
315            Self::Fixed(StateTensorRole::Convolution { slot }) => {
316                format!("state.convolution.{slot}")
317            }
318            Self::Fixed(StateTensorRole::Recurrent) => "state.recurrent".into(),
319            Self::Fixed(StateTensorRole::PrefixEmbedding) => "state.prefix_embedding".into(),
320            Self::Fixed(StateTensorRole::PositionDelta) => "state.position_delta".into(),
321            Self::Fixed(StateTensorRole::Pooling { stream, component }) => {
322                let component = match component {
323                    PoolingStateComponent::PendingValues => "pending_values",
324                    PoolingStateComponent::PendingGates => "pending_gates",
325                    PoolingStateComponent::Pooled => "pooled",
326                    PoolingStateComponent::OverlapValues => "overlap_values",
327                    PoolingStateComponent::OverlapGates => "overlap_gates",
328                };
329                format!("state.pooling.{stream}.{component}")
330            }
331        }
332    }
333}
334
335/// Symbolic geometry and persistence behavior of one named state component.
336#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
337pub struct StateComponentPolicy {
338    /// Stable semantic role.
339    role: StateComponentRole,
340    /// Symbolic shape resolved from batch and prefix geometry.
341    shape: Vec<StateTensorDimension>,
342    /// Accepted persisted dtype family.
343    dtype: StateTensorDtype,
344    /// Runtime residency behavior.
345    residency: StateResidencyClass,
346    /// Condition under which persisted state materializes this component.
347    presence: StateTensorPresence,
348}
349
350impl StateComponentPolicy {
351    /// Returns the stable semantic component role.
352    pub const fn role(&self) -> StateComponentRole {
353        self.role
354    }
355
356    /// Returns the symbolic component shape.
357    pub fn shape(&self) -> &[StateTensorDimension] {
358        &self.shape
359    }
360
361    /// Returns the accepted persisted dtype family.
362    pub const fn dtype(&self) -> StateTensorDtype {
363        self.dtype
364    }
365
366    /// Returns the runtime residency class.
367    pub const fn residency(&self) -> StateResidencyClass {
368        self.residency
369    }
370
371    /// Returns the conditional persistence rule.
372    pub const fn presence(&self) -> StateTensorPresence {
373        self.presence
374    }
375}
376
377impl LayerCachePolicy {
378    /// Returns the residency behavior of this layer's attention payload.
379    pub const fn attention_residency_class(&self) -> Option<StateResidencyClass> {
380        match self {
381            Self::NoState | Self::FixedState { .. } => None,
382            Self::KeyValue { .. }
383            | Self::KeyOnly { .. }
384            | Self::CompressedLatentRotary { .. }
385            | Self::KeyValueWithFixedState { .. }
386            | Self::KeyOnlyWithFixedState { .. } => Some(StateResidencyClass::SealablePaged),
387        }
388    }
389
390    /// Constructs validated ordinary key/value state geometry.
391    pub fn key_value(
392        attention: AttentionPolicy,
393        num_key_value_heads: i32,
394        head_dim: i32,
395    ) -> Result<Self, CachePolicyError> {
396        let policy = Self::KeyValue {
397            attention,
398            num_key_value_heads: positive_u32(num_key_value_heads, "key/value head count")?,
399            head_dim: positive_u32(head_dim, "key/value head dimension")?,
400        };
401        policy.validate()?;
402        Ok(policy)
403    }
404
405    /// Constructs validated key-only state geometry.
406    pub fn key_only(
407        attention: AttentionPolicy,
408        num_key_heads: i32,
409        head_dim: i32,
410    ) -> Result<Self, CachePolicyError> {
411        let policy = Self::KeyOnly {
412            attention,
413            num_key_heads: positive_u32(num_key_heads, "key head count")?,
414            head_dim: positive_u32(head_dim, "key head dimension")?,
415        };
416        policy.validate()?;
417        Ok(policy)
418    }
419
420    /// Constructs validated compressed-latent state geometry.
421    pub fn compressed_latent_rotary(
422        attention: AttentionPolicy,
423        latent_dim: i32,
424        rotary_dim: i32,
425    ) -> Result<Self, CachePolicyError> {
426        let policy = Self::CompressedLatentRotary {
427            attention,
428            latent_dim: positive_u32(latent_dim, "compressed latent dimension")?,
429            rotary_dim: positive_u32(rotary_dim, "rotary-key dimension")?,
430        };
431        policy.validate()?;
432        Ok(policy)
433    }
434
435    /// Constructs a validated fixed-state-only policy.
436    pub fn fixed_only(tensors: Vec<StateTensorPolicy>) -> Result<Self, CachePolicyError> {
437        let policy = Self::FixedState { tensors };
438        policy.validate()?;
439        Ok(policy)
440    }
441
442    /// Constructs validated key/value plus fixed-state geometry.
443    pub fn key_value_with_fixed_state(
444        attention: AttentionPolicy,
445        num_key_value_heads: i32,
446        head_dim: i32,
447        tensors: Vec<StateTensorPolicy>,
448    ) -> Result<Self, CachePolicyError> {
449        let policy = Self::KeyValueWithFixedState {
450            attention,
451            num_key_value_heads: positive_u32(num_key_value_heads, "key/value head count")?,
452            head_dim: positive_u32(head_dim, "key/value head dimension")?,
453            tensors,
454        };
455        policy.validate()?;
456        Ok(policy)
457    }
458
459    /// Constructs validated key-only plus fixed-state geometry.
460    pub fn key_only_with_fixed_state(
461        attention: AttentionPolicy,
462        num_key_heads: i32,
463        head_dim: i32,
464        tensors: Vec<StateTensorPolicy>,
465    ) -> Result<Self, CachePolicyError> {
466        let policy = Self::KeyOnlyWithFixedState {
467            attention,
468            num_key_heads: positive_u32(num_key_heads, "key head count")?,
469            head_dim: positive_u32(head_dim, "key head dimension")?,
470            tensors,
471        };
472        policy.validate()?;
473        Ok(policy)
474    }
475
476    /// Returns the exact attention policy when present.
477    pub const fn attention(&self) -> Option<AttentionPolicy> {
478        match self {
479            Self::NoState | Self::FixedState { .. } => None,
480            Self::KeyValue { attention, .. }
481            | Self::KeyOnly { attention, .. }
482            | Self::CompressedLatentRotary { attention, .. }
483            | Self::KeyValueWithFixedState { attention, .. }
484            | Self::KeyOnlyWithFixedState { attention, .. } => Some(*attention),
485        }
486    }
487
488    /// Returns ordered non-attention tensor policies.
489    pub fn fixed_state(&self) -> &[StateTensorPolicy] {
490        match self {
491            Self::FixedState { tensors }
492            | Self::KeyValueWithFixedState { tensors, .. }
493            | Self::KeyOnlyWithFixedState { tensors, .. } => tensors,
494            _ => &[],
495        }
496    }
497
498    /// Expands this layer policy into ordered, stably named semantic
499    /// components shared by runtime residency and prompt persistence.
500    pub fn components(&self) -> Vec<StateComponentPolicy> {
501        let mut components = Vec::new();
502        let floating = StateTensorDtype::Floating;
503        let required = StateTensorPresence::Required;
504        match self {
505            Self::NoState | Self::FixedState { .. } => {}
506            Self::KeyValue {
507                num_key_value_heads,
508                head_dim,
509                ..
510            }
511            | Self::KeyValueWithFixedState {
512                num_key_value_heads,
513                head_dim,
514                ..
515            } => {
516                let shape = vec![
517                    StateTensorDimension::Batch,
518                    StateTensorDimension::Fixed(*num_key_value_heads),
519                    StateTensorDimension::PrefixTokens,
520                    StateTensorDimension::Fixed(*head_dim),
521                ];
522                for role in [
523                    StateComponentRole::AttentionKeys,
524                    StateComponentRole::AttentionValues,
525                ] {
526                    components.push(StateComponentPolicy {
527                        role,
528                        shape: shape.clone(),
529                        dtype: floating,
530                        residency: StateResidencyClass::SealablePaged,
531                        presence: required,
532                    });
533                }
534            }
535            Self::KeyOnly {
536                num_key_heads,
537                head_dim,
538                ..
539            }
540            | Self::KeyOnlyWithFixedState {
541                num_key_heads,
542                head_dim,
543                ..
544            } => components.push(StateComponentPolicy {
545                role: StateComponentRole::AttentionKeys,
546                shape: vec![
547                    StateTensorDimension::Batch,
548                    StateTensorDimension::Fixed(*num_key_heads),
549                    StateTensorDimension::PrefixTokens,
550                    StateTensorDimension::Fixed(*head_dim),
551                ],
552                dtype: floating,
553                residency: StateResidencyClass::SealablePaged,
554                presence: required,
555            }),
556            Self::CompressedLatentRotary {
557                latent_dim,
558                rotary_dim,
559                ..
560            } => {
561                for (role, dimension) in [
562                    (StateComponentRole::CompressedLatent, *latent_dim),
563                    (StateComponentRole::RotaryKeys, *rotary_dim),
564                ] {
565                    components.push(StateComponentPolicy {
566                        role,
567                        shape: vec![
568                            StateTensorDimension::Batch,
569                            StateTensorDimension::PrefixTokens,
570                            StateTensorDimension::Fixed(dimension),
571                        ],
572                        dtype: floating,
573                        residency: StateResidencyClass::SealablePaged,
574                        presence: required,
575                    });
576                }
577            }
578        }
579        components.extend(
580            self.fixed_state()
581                .iter()
582                .map(|tensor| StateComponentPolicy {
583                    role: StateComponentRole::Fixed(tensor.role),
584                    shape: tensor.shape.clone(),
585                    dtype: tensor.dtype,
586                    residency: tensor.residency_class(),
587                    presence: tensor.presence,
588                }),
589        );
590        components
591    }
592
593    /// Validates dimensions and fixed-state invariants after deserialization.
594    pub fn validate(&self) -> Result<(), CachePolicyError> {
595        if let Some(attention) = self.attention() {
596            attention
597                .sliding_window_i32()
598                .map_err(|error| CachePolicyError::Invalid(error.to_string()))?;
599        }
600        let validate_dimension = |dimension: NonZeroU32| {
601            (dimension.get() <= i32::MAX as u32)
602                .then_some(())
603                .ok_or_else(|| {
604                    CachePolicyError::Invalid(format!(
605                        "prompt-cache layer dimension {dimension} exceeds the runtime i32 range"
606                    ))
607                })
608        };
609        match self {
610            Self::NoState | Self::FixedState { .. } => {}
611            Self::KeyValue {
612                num_key_value_heads,
613                head_dim,
614                ..
615            }
616            | Self::KeyValueWithFixedState {
617                num_key_value_heads,
618                head_dim,
619                ..
620            } => {
621                validate_dimension(*num_key_value_heads)?;
622                validate_dimension(*head_dim)?;
623            }
624            Self::KeyOnly {
625                num_key_heads,
626                head_dim,
627                ..
628            }
629            | Self::KeyOnlyWithFixedState {
630                num_key_heads,
631                head_dim,
632                ..
633            } => {
634                validate_dimension(*num_key_heads)?;
635                validate_dimension(*head_dim)?;
636            }
637            Self::CompressedLatentRotary {
638                latent_dim,
639                rotary_dim,
640                ..
641            } => {
642                validate_dimension(*latent_dim)?;
643                validate_dimension(*rotary_dim)?;
644            }
645        }
646        let tensors = self.fixed_state();
647        if tensors.is_empty()
648            && matches!(
649                self,
650                Self::FixedState { .. }
651                    | Self::KeyValueWithFixedState { .. }
652                    | Self::KeyOnlyWithFixedState { .. }
653            )
654        {
655            return Err(CachePolicyError::Invalid(
656                "fixed-state cache policy must contain at least one tensor".into(),
657            ));
658        }
659        validate_state_tensor_policies(tensors)
660    }
661}
662
663impl StateTensorDimension {
664    /// Constructs a positive fixed dimension.
665    pub fn fixed(value: i32) -> Result<Self, CachePolicyError> {
666        positive_u32(value, "fixed-state tensor dimension").map(Self::Fixed)
667    }
668}
669
670impl StateTensorPolicy {
671    /// Constructs and validates a state-tensor policy.
672    pub fn new(
673        role: StateTensorRole,
674        shape: Vec<StateTensorDimension>,
675        dtype: StateTensorDtype,
676        residency: MutableStateResidency,
677    ) -> Result<Self, CachePolicyError> {
678        Self::new_with_residency(role, shape, dtype, residency.into())
679    }
680
681    /// Constructs state with an explicit pageable or mutable residency class.
682    pub fn new_with_residency(
683        role: StateTensorRole,
684        shape: Vec<StateTensorDimension>,
685        dtype: StateTensorDtype,
686        residency: StateResidencyClass,
687    ) -> Result<Self, CachePolicyError> {
688        let policy = Self {
689            role,
690            shape,
691            dtype,
692            residency,
693            presence: StateTensorPresence::Required,
694        };
695        validate_state_tensor_policies(std::slice::from_ref(&policy))?;
696        Ok(policy)
697    }
698
699    /// Marks this tensor as optional.
700    pub const fn optional(mut self) -> Self {
701        self.presence = StateTensorPresence::Optional;
702        self
703    }
704
705    /// Requires the tensor when the prefix has a non-zero remainder.
706    pub const fn when_prefix_remainder_nonzero(mut self, divisor: NonZeroU32) -> Self {
707        self.presence = StateTensorPresence::PrefixRemainderNonZero(divisor);
708        self
709    }
710
711    /// Requires the tensor when the prefix contains a complete group.
712    pub const fn when_prefix_at_least(mut self, divisor: NonZeroU32) -> Self {
713        self.presence = StateTensorPresence::PrefixAtLeast(divisor);
714        self
715    }
716
717    /// Returns whether the tensor is required for an exact prefix length.
718    pub fn is_required_for(&self, prefix_tokens: usize) -> bool {
719        match self.presence {
720            StateTensorPresence::Required => true,
721            StateTensorPresence::Optional => false,
722            StateTensorPresence::PrefixRemainderNonZero(divisor) => {
723                !prefix_tokens.is_multiple_of(divisor.get() as usize)
724            }
725            StateTensorPresence::PrefixAtLeast(divisor) => prefix_tokens >= divisor.get() as usize,
726        }
727    }
728
729    /// Returns the unified residency classification.
730    pub fn residency_class(&self) -> StateResidencyClass {
731        self.residency
732    }
733
734    /// Resolves symbolic dimensions for an exact batch and prefix length.
735    pub fn resolved_shape(
736        &self,
737        batch_size: usize,
738        prefix_tokens: usize,
739    ) -> Result<Vec<i32>, CachePolicyError> {
740        self.shape
741            .iter()
742            .map(|dimension| match dimension {
743                StateTensorDimension::Batch => i32::try_from(batch_size),
744                StateTensorDimension::PrefixTokens => i32::try_from(prefix_tokens),
745                StateTensorDimension::PrefixTokensDiv(divisor) => {
746                    i32::try_from(prefix_tokens / divisor.get() as usize)
747                }
748                StateTensorDimension::PrefixTokensRem(divisor) => {
749                    i32::try_from(prefix_tokens % divisor.get() as usize)
750                }
751                StateTensorDimension::Fixed(value) => i32::try_from(value.get()),
752                StateTensorDimension::Scalar => Ok(1),
753            })
754            .collect::<Result<Vec<_>, _>>()
755            .map_err(|_| {
756                CachePolicyError::Invalid(
757                    "fixed-state tensor dimension exceeds runtime i32 range".into(),
758                )
759            })
760    }
761
762    /// Tests a stable serialized dtype name against this policy.
763    pub fn accepts_dtype_name(&self, dtype: &str) -> bool {
764        match self.dtype {
765            StateTensorDtype::Floating => {
766                matches!(dtype, "Float16" | "Bfloat16" | "Float32" | "Float64")
767            }
768            StateTensorDtype::Float32 => dtype == "Float32",
769            StateTensorDtype::Int32 => dtype == "Int32",
770            StateTensorDtype::Uint32 => dtype == "Uint32",
771        }
772    }
773}
774
775fn positive_u32(value: i32, field: &str) -> Result<NonZeroU32, CachePolicyError> {
776    u32::try_from(value)
777        .ok()
778        .and_then(NonZeroU32::new)
779        .ok_or_else(|| {
780            CachePolicyError::Invalid(format!(
781                "prompt-cache {field} must be positive and fit u32, got {value}"
782            ))
783        })
784}
785
786fn validate_state_tensor_policies(tensors: &[StateTensorPolicy]) -> Result<(), CachePolicyError> {
787    let mut roles = BTreeSet::new();
788    for tensor in tensors {
789        if !roles.insert(tensor.role) {
790            return Err(CachePolicyError::Invalid(format!(
791                "duplicate fixed-state tensor role {:?}",
792                tensor.role
793            )));
794        }
795        if tensor.shape.is_empty()
796            || (tensor.shape.contains(&StateTensorDimension::Scalar)
797                && tensor.shape.as_slice() != [StateTensorDimension::Scalar])
798        {
799            return Err(CachePolicyError::Invalid(format!(
800                "invalid fixed-state tensor shape for role {:?}",
801                tensor.role
802            )));
803        }
804        let expected = match tensor.role {
805            StateTensorRole::Recurrent => StateResidencyClass::LayerScopedOffloadable,
806            StateTensorRole::Convolution { .. }
807            | StateTensorRole::PrefixEmbedding
808            | StateTensorRole::PositionDelta => StateResidencyClass::AlwaysDeviceMutable,
809            StateTensorRole::Pooling {
810                component: PoolingStateComponent::Pooled,
811                ..
812            } => StateResidencyClass::SealablePaged,
813            StateTensorRole::Pooling { .. } => StateResidencyClass::AlwaysDeviceMutable,
814        };
815        if tensor.residency != expected {
816            return Err(CachePolicyError::Invalid(format!(
817                "fixed-state tensor role {:?} requires {:?} residency, got {:?}",
818                tensor.role, expected, tensor.residency
819            )));
820        }
821    }
822    Ok(())
823}
824
825/// Invalid cache geometry or state-residency contract.
826#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
827pub enum CachePolicyError {
828    /// A cache policy violates a structural or representational invariant.
829    #[error("{0}")]
830    Invalid(String),
831}
832
833#[cfg(test)]
834mod tests {
835    use super::*;
836
837    #[test]
838    fn validates_layer_and_fixed_state_contracts() {
839        let recurrent = StateTensorPolicy::new(
840            StateTensorRole::Recurrent,
841            vec![
842                StateTensorDimension::Batch,
843                StateTensorDimension::fixed(16).unwrap(),
844            ],
845            StateTensorDtype::Floating,
846            MutableStateResidency::LayerScopedOffloadable,
847        )
848        .unwrap();
849        let layer = LayerCachePolicy::key_value_with_fixed_state(
850            AttentionPolicy::sliding(128).unwrap(),
851            8,
852            64,
853            vec![recurrent.clone()],
854        )
855        .unwrap();
856        assert_eq!(
857            layer.attention_residency_class(),
858            Some(StateResidencyClass::SealablePaged)
859        );
860        assert_eq!(recurrent.resolved_shape(2, 9).unwrap(), vec![2, 16]);
861        assert!(recurrent.accepts_dtype_name("Float16"));
862        assert!(!recurrent.accepts_dtype_name("Int32"));
863    }
864
865    #[test]
866    fn rejects_invalid_policy_without_a_backend() {
867        assert!(LayerCachePolicy::key_value(AttentionPolicy::Full, 0, 64).is_err());
868        assert!(StateTensorPolicy::new(
869            StateTensorRole::Recurrent,
870            vec![StateTensorDimension::Scalar, StateTensorDimension::Batch],
871            StateTensorDtype::Floating,
872            MutableStateResidency::LayerScopedOffloadable,
873        )
874        .is_err());
875    }
876
877    #[test]
878    fn policy_schema_round_trips() {
879        let policy = LayerCachePolicy::key_only(AttentionPolicy::Full, 4, 32).unwrap();
880        let json = serde_json::to_string(&policy).unwrap();
881        assert_eq!(
882            serde_json::from_str::<LayerCachePolicy>(&json).unwrap(),
883            policy
884        );
885    }
886}
eredu_core/cache/policy.rs

eredu_core/cache/
policy.rs
