ferrum_interfaces/
recurrent_state.rs

1//! Recurrent-state cache contracts for state-space / hybrid models.
2//!
3//! Recurrent state is intentionally separate from KV cache. KV grows with
4//! sequence length and is addressed through attention blocks; recurrent state is
5//! compact per-layer state that can coexist with KV in hybrid models.
6
7use ferrum_types::{DataType, Device, RequestId, Result};
8use serde::{Deserialize, Serialize};
9use std::{any::Any, sync::Arc, time::Instant};
10
11/// A single recurrent-state tensor owned by a model layer.
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct RecurrentStateTensorSpec {
14    /// Layer index that owns this state tensor.
15    pub layer_index: usize,
16    /// Backend/model-local state name, for example `delta_state`.
17    pub name: String,
18    /// Tensor shape excluding any request/batch slot dimension.
19    pub shape: Vec<usize>,
20    /// Storage dtype for this state tensor.
21    pub dtype: DataType,
22}
23
24impl RecurrentStateTensorSpec {
25    pub fn new(
26        layer_index: usize,
27        name: impl Into<String>,
28        shape: Vec<usize>,
29        dtype: DataType,
30    ) -> Self {
31        Self {
32            layer_index,
33            name: name.into(),
34            shape,
35            dtype,
36        }
37    }
38
39    pub fn checked_num_elements(&self) -> Option<usize> {
40        self.shape
41            .iter()
42            .copied()
43            .try_fold(1usize, usize::checked_mul)
44    }
45
46    pub fn num_elements(&self) -> usize {
47        self.checked_num_elements().unwrap_or(usize::MAX)
48    }
49}
50
51/// Allocation request for recurrent state.
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
53pub struct RecurrentStateSpec {
54    /// Request ID this allocation is for.
55    pub request_id: RequestId,
56    /// Number of model layers that may carry recurrent state.
57    pub num_layers: usize,
58    /// Concrete state tensors to allocate.
59    pub tensors: Vec<RecurrentStateTensorSpec>,
60    /// Target device.
61    pub device: Device,
62    /// Number of request/batch slots reserved by this handle.
63    pub max_batch_slots: usize,
64}
65
66impl RecurrentStateSpec {
67    pub fn estimated_memory_bytes(&self) -> usize {
68        let state_bytes_per_slot = self.tensors.iter().fold(0usize, |total, tensor| {
69            total.saturating_add(
70                tensor
71                    .num_elements()
72                    .saturating_mul(tensor.dtype.size_bytes()),
73            )
74        });
75        state_bytes_per_slot.saturating_mul(self.max_batch_slots)
76    }
77}
78
79/// Resume policy for state that has been preempted or evicted.
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
81pub enum RecurrentStateResumePolicy {
82    /// Rebuild state from request token history before the next decode.
83    RecomputeOnResume,
84    /// Save and restore backend-specific state bytes during preemption.
85    SnapshotOnPreempt,
86}
87
88/// Statistics for one recurrent-state handle.
89#[derive(Debug, Clone)]
90pub struct RecurrentStateHandleStats {
91    /// Total memory used by this handle.
92    pub memory_bytes: usize,
93    /// Number of state tensors represented by this handle.
94    pub state_tensors: usize,
95    /// Number of request/batch slots reserved by this handle.
96    pub batch_slots: usize,
97    /// Last access timestamp for eviction policies.
98    pub last_access: Instant,
99}
100
101/// Recurrent-state cache handle.
102pub trait RecurrentStateHandle: Send + Sync + std::fmt::Debug {
103    /// Request ID this state belongs to.
104    fn request_id(&self) -> RequestId;
105
106    /// Device where the state resides.
107    fn device(&self) -> Device;
108
109    /// Number of model layers represented by this handle.
110    fn num_layers(&self) -> usize;
111
112    /// Approximate memory used by this state.
113    fn state_bytes(&self) -> usize;
114
115    /// Clone handle reference. Implementations should not deep-copy state.
116    fn clone_handle(&self) -> Result<Arc<dyn RecurrentStateHandle>>;
117
118    /// Downcast support for backend-specific handles.
119    fn as_any(&self) -> &dyn Any;
120
121    /// Handle statistics.
122    fn stats(&self) -> RecurrentStateHandleStats;
123
124    /// Check whether state is still valid and accessible.
125    fn is_valid(&self) -> bool;
126
127    /// Unique identifier for this cache instance.
128    fn cache_id(&self) -> String;
129}
130
131/// Aggregate recurrent-state manager statistics.
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct RecurrentStateManagerStats {
134    /// Total memory budget visible to this manager.
135    pub total_memory_bytes: usize,
136    /// Memory currently used by active recurrent states.
137    pub used_memory_bytes: usize,
138    /// Number of active state handles.
139    pub active_states: usize,
140    /// Number of active state tensors.
141    pub active_state_tensors: usize,
142    /// Total slots visible to this manager.
143    pub total_batch_slots: usize,
144    /// Slots currently allocated to active state handles.
145    pub used_batch_slots: usize,
146    /// Number of successful allocations.
147    pub allocation_count: u64,
148    /// Number of failed allocations.
149    pub allocation_failures: u64,
150    /// Number of evictions performed.
151    pub eviction_count: u64,
152}
153
154/// Recurrent-state cache manager for allocation and lifecycle management.
155#[async_trait::async_trait]
156pub trait RecurrentStateManager: Send + Sync {
157    /// Allocate recurrent state for one request.
158    async fn allocate(&self, spec: &RecurrentStateSpec) -> Result<Arc<dyn RecurrentStateHandle>>;
159
160    /// Deallocate recurrent state for one request.
161    async fn deallocate(&self, request_id: RequestId) -> Result<()>;
162
163    /// Check whether the manager can satisfy the allocation.
164    fn can_allocate(&self, spec: &RecurrentStateSpec) -> bool;
165
166    /// Get handle for an existing request.
167    fn get_handle(&self, request_id: RequestId) -> Option<Arc<dyn RecurrentStateHandle>>;
168
169    /// List all active recurrent-state handles.
170    fn list_handles(&self) -> Vec<(RequestId, Arc<dyn RecurrentStateHandle>)>;
171
172    /// Get aggregate manager statistics.
173    fn stats(&self) -> RecurrentStateManagerStats;
174
175    /// Drop all state owned by this manager.
176    async fn reset(&self) -> Result<()>;
177}
ferrum_interfaces/recurrent_state.rs

ferrum_interfaces/
recurrent_state.rs
