qualia_core_db/inference/runtime/kv/paged/
table.rs

1use super::config::{PagedKvConfig, INVALID_BLOCK};
2use super::pool::{BlockPool, CopyOnWrite, PoolError};
3
4#[derive(Clone, Copy, Debug, PartialEq, Eq)]
5pub enum TableError {
6    InvalidConfig,
7    OutputTooSmall,
8    OutOfRange,
9    TargetNotEmpty,
10    Pool(PoolError),
11}
12
13impl From<PoolError> for TableError {
14    fn from(value: PoolError) -> Self {
15        Self::Pool(value)
16    }
17}
18
19/// Fill the production single-sequence GPU table with deterministic identity pages.
20///
21/// Layout is `[layer][logical_page] -> physical_page`. With identity pages the byte layout is
22/// compatible with the former dense `[layer][token][K|V]` arena, which enables staged rollout and
23/// differential testing while kernels consume the real block table.
24pub fn fill_identity_block_table(
25    config: &PagedKvConfig,
26    out: &mut [u32],
27) -> Result<usize, TableError> {
28    if !config.is_valid() {
29        return Err(TableError::InvalidConfig);
30    }
31    let required = config.required_single_sequence_blocks() as usize;
32    if out.len() < required {
33        return Err(TableError::OutputTooSmall);
34    }
35    for (index, entry) in out[..required].iter_mut().enumerate() {
36        *entry = index as u32;
37    }
38    Ok(required)
39}
40
41/// Cold-owned immutable block table uploaded by a prepared GPU plan.
42#[derive(Debug)]
43pub struct GpuBlockTablePlan {
44    config: PagedKvConfig,
45    entries: Vec<u32>,
46}
47
48impl GpuBlockTablePlan {
49    pub fn identity(config: PagedKvConfig) -> Result<Self, TableError> {
50        let mut entries = vec![INVALID_BLOCK; config.required_single_sequence_blocks() as usize];
51        fill_identity_block_table(&config, &mut entries)?;
52        Ok(Self { config, entries })
53    }
54
55    pub fn config(&self) -> PagedKvConfig {
56        self.config
57    }
58
59    pub fn entries(&self) -> &[u32] {
60        &self.entries
61    }
62}
63
64/// One sequence's logical-page ownership. Capacity is fixed at construction.
65#[derive(Debug)]
66pub struct SequenceBlockTable {
67    entries: Vec<u32>,
68}
69
70impl SequenceBlockTable {
71    pub fn new(logical_pages: u32) -> Self {
72        Self {
73            entries: vec![INVALID_BLOCK; logical_pages as usize],
74        }
75    }
76
77    pub fn entries(&self) -> &[u32] {
78        &self.entries
79    }
80
81    pub fn get(&self, logical_page: u32) -> Option<u32> {
82        self.entries
83            .get(logical_page as usize)
84            .copied()
85            .filter(|block| *block != INVALID_BLOCK)
86    }
87
88    pub fn ensure_writable(
89        &mut self,
90        logical_page: u32,
91        pool: &mut BlockPool,
92    ) -> Result<CopyOnWrite, TableError> {
93        let entry = self
94            .entries
95            .get_mut(logical_page as usize)
96            .ok_or(TableError::OutOfRange)?;
97        let action = pool.writable(*entry)?;
98        *entry = match action {
99            CopyOnWrite::Existing(block) | CopyOnWrite::Allocated(block) => block,
100            CopyOnWrite::Copy { destination, .. } => destination,
101        };
102        Ok(action)
103    }
104
105    pub fn fork_into(&self, target: &mut Self, pool: &mut BlockPool) -> Result<(), TableError> {
106        if self.entries.len() != target.entries.len() {
107            return Err(TableError::OutOfRange);
108        }
109        if target.entries.iter().any(|block| *block != INVALID_BLOCK) {
110            return Err(TableError::TargetNotEmpty);
111        }
112        let mut retained = 0usize;
113        for &block in &self.entries {
114            if block != INVALID_BLOCK {
115                if let Err(error) = pool.retain(block) {
116                    for &rollback in &self.entries[..retained] {
117                        if rollback != INVALID_BLOCK {
118                            let _ = pool.release(rollback);
119                        }
120                    }
121                    return Err(error.into());
122                }
123            }
124            retained += 1;
125        }
126        target.entries.copy_from_slice(&self.entries);
127        Ok(())
128    }
129
130    /// Attach graph/prompt-cache prefix pages to an empty sequence table.
131    ///
132    /// Reference counts are retained before publishing entries; failures roll back, so a denied
133    /// or stale graph prefix cannot partially mutate request state.
134    pub fn install_shared_prefix(
135        &mut self,
136        pages: &[u32],
137        pool: &mut BlockPool,
138    ) -> Result<(), TableError> {
139        if pages.len() > self.entries.len() {
140            return Err(TableError::OutOfRange);
141        }
142        if self.entries.iter().any(|block| *block != INVALID_BLOCK) {
143            return Err(TableError::TargetNotEmpty);
144        }
145        for (index, &page) in pages.iter().enumerate() {
146            if let Err(error) = pool.retain(page) {
147                for &rollback in &pages[..index] {
148                    let _ = pool.release(rollback);
149                }
150                return Err(error.into());
151            }
152        }
153        self.entries[..pages.len()].copy_from_slice(pages);
154        Ok(())
155    }
156
157    pub fn release_all(&mut self, pool: &mut BlockPool) -> Result<(), TableError> {
158        for entry in &mut self.entries {
159            if *entry != INVALID_BLOCK {
160                pool.release(*entry)?;
161                *entry = INVALID_BLOCK;
162            }
163        }
164        Ok(())
165    }
166}
qualia_core_db/inference/runtime/kv/paged/table.rs

qualia_core_db/inference/runtime/kv/paged/
table.rs
