qualia_core_db/inference/runtime/scheduler/
request_table.rs

1use crate::inference::runtime::graph_assist::PrefixIdentity;
2use crate::inference::runtime::kv::paged::{
3    BlockPool, CopyOnWrite, SequenceBlockTable, TableError,
4};
5use crate::inference::runtime::kv::prefix::{PrefixKvError, PrefixKvStore};
6
7use super::batch::{
8    RaggedBackendError, RaggedBatchItem, RaggedBatchOutput, RaggedBatchReceipt, RaggedDecodeBackend,
9};
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum RequestState {
13    Empty,
14    Prefill,
15    Decode,
16}
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
19pub struct Admission {
20    pub slot: u16,
21    pub prefix_hit: bool,
22    pub prefix_tokens: u32,
23}
24
25#[repr(C)]
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub struct RequestView {
28    pub request_id: u64,
29    pub slot: u16,
30    pub state: RequestState,
31    pub token_count: u32,
32    pub prefix_tokens: u32,
33}
34
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36pub enum SchedulerError {
37    DuplicateRequest,
38    Full,
39    UnknownRequest,
40    OutputTooSmall,
41    Prefix(PrefixKvError),
42    Table(TableError),
43}
44
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub enum DecodeRoundError {
47    Scheduler(SchedulerError),
48    ItemBufferTooSmall,
49    BlockTableBufferTooSmall,
50    OutputBufferTooSmall,
51    MissingInputToken,
52    Backend(RaggedBackendError),
53    BackendLaunchCount,
54    OutputIdentityMismatch,
55}
56
57impl From<SchedulerError> for DecodeRoundError {
58    fn from(value: SchedulerError) -> Self {
59        Self::Scheduler(value)
60    }
61}
62
63impl From<RaggedBackendError> for DecodeRoundError {
64    fn from(value: RaggedBackendError) -> Self {
65        Self::Backend(value)
66    }
67}
68
69impl From<PrefixKvError> for SchedulerError {
70    fn from(value: PrefixKvError) -> Self {
71        Self::Prefix(value)
72    }
73}
74
75impl From<TableError> for SchedulerError {
76    fn from(value: TableError) -> Self {
77        Self::Table(value)
78    }
79}
80
81struct RequestSlot {
82    request_id: u64,
83    state: RequestState,
84    token_count: u32,
85    prefix_tokens: u32,
86    next_token_id: u32,
87    token_ready: bool,
88    blocks: SequenceBlockTable,
89}
90
91impl RequestSlot {
92    fn new(logical_pages: u32) -> Self {
93        Self {
94            request_id: 0,
95            state: RequestState::Empty,
96            token_count: 0,
97            prefix_tokens: 0,
98            next_token_id: 0,
99            token_ready: false,
100            blocks: SequenceBlockTable::new(logical_pages),
101        }
102    }
103
104    fn clear(&mut self) {
105        self.request_id = 0;
106        self.state = RequestState::Empty;
107        self.token_count = 0;
108        self.prefix_tokens = 0;
109        self.next_token_id = 0;
110        self.token_ready = false;
111    }
112}
113
114/// Fixed-request-capacity scheduler. `REQUESTS` is the hard concurrent request bound.
115pub struct RequestScheduler<const REQUESTS: usize> {
116    slots: [RequestSlot; REQUESTS],
117}
118
119impl<const REQUESTS: usize> RequestScheduler<REQUESTS> {
120    /// Cold construction; allocates each fixed-capacity logical page table exactly once.
121    pub fn new(logical_pages: u32) -> Self {
122        Self {
123            slots: std::array::from_fn(|_| RequestSlot::new(logical_pages)),
124        }
125    }
126
127    pub fn active_count(&self) -> usize {
128        self.slots
129            .iter()
130            .filter(|slot| slot.state != RequestState::Empty)
131            .count()
132    }
133
134    /// Admit a request and atomically attach graph-derived prefix pages when available.
135    pub fn admit_with_prefix<const ENTRIES: usize, const PAGES: usize>(
136        &mut self,
137        request_id: u64,
138        identity: Option<PrefixIdentity>,
139        prefixes: &PrefixKvStore<ENTRIES, PAGES>,
140        pool: &mut BlockPool,
141    ) -> Result<Admission, SchedulerError> {
142        if self
143            .slots
144            .iter()
145            .any(|slot| slot.state != RequestState::Empty && slot.request_id == request_id)
146        {
147            return Err(SchedulerError::DuplicateRequest);
148        }
149        let index = self
150            .slots
151            .iter()
152            .position(|slot| slot.state == RequestState::Empty)
153            .ok_or(SchedulerError::Full)?;
154        let prefix_tokens = match identity {
155            Some(identity) => prefixes
156                .attach(identity, &mut self.slots[index].blocks, pool)?
157                .unwrap_or(0),
158            None => 0,
159        };
160        let slot = &mut self.slots[index];
161        slot.request_id = request_id;
162        slot.state = if prefix_tokens == 0 {
163            RequestState::Prefill
164        } else {
165            RequestState::Decode
166        };
167        slot.token_count = prefix_tokens;
168        slot.prefix_tokens = prefix_tokens;
169        Ok(Admission {
170            slot: index as u16,
171            prefix_hit: prefix_tokens != 0,
172            prefix_tokens,
173        })
174    }
175
176    /// Make a logical page writable. A `Copy` result instructs the backend to copy the physical
177    /// KV page before writing the next token.
178    pub fn ensure_writable(
179        &mut self,
180        request_id: u64,
181        logical_page: u32,
182        pool: &mut BlockPool,
183    ) -> Result<CopyOnWrite, SchedulerError> {
184        let slot = self.find_mut(request_id)?;
185        slot.blocks
186            .ensure_writable(logical_page, pool)
187            .map_err(Into::into)
188    }
189
190    pub fn mark_prefill_complete(&mut self, request_id: u64) -> Result<(), SchedulerError> {
191        self.find_mut(request_id)?.state = RequestState::Decode;
192        Ok(())
193    }
194
195    /// Publish the token consumed by the request's next decode step.
196    pub fn seed_decode_token(
197        &mut self,
198        request_id: u64,
199        token_id: u32,
200    ) -> Result<(), SchedulerError> {
201        let slot = self.find_mut(request_id)?;
202        slot.next_token_id = token_id;
203        slot.token_ready = true;
204        Ok(())
205    }
206
207    pub fn record_token(&mut self, request_id: u64) -> Result<u32, SchedulerError> {
208        let slot = self.find_mut(request_id)?;
209        slot.token_count = slot.token_count.saturating_add(1);
210        Ok(slot.token_count)
211    }
212
213    /// Write all runnable requests into caller storage in stable slot order.
214    pub fn runnable_into(&self, out: &mut [RequestView]) -> Result<usize, SchedulerError> {
215        let required = self.active_count();
216        if out.len() < required {
217            return Err(SchedulerError::OutputTooSmall);
218        }
219        let mut written = 0usize;
220        for (index, slot) in self.slots.iter().enumerate() {
221            if slot.state == RequestState::Empty {
222                continue;
223            }
224            out[written] = RequestView {
225                request_id: slot.request_id,
226                slot: index as u16,
227                state: slot.state,
228                token_count: slot.token_count,
229                prefix_tokens: slot.prefix_tokens,
230            };
231            written += 1;
232        }
233        Ok(written)
234    }
235
236    /// Execute one ragged decode round with exactly one backend call.
237    ///
238    /// Active decode requests are lowered in stable scheduler-slot order. Their logical page
239    /// tables are concatenated into `block_table_scratch`; each item carries its table offset.
240    /// Outputs are identity-checked as a complete batch before scheduler state is mutated.
241    pub fn execute_decode_round<B: RaggedDecodeBackend>(
242        &mut self,
243        backend: &mut B,
244        item_scratch: &mut [RaggedBatchItem],
245        block_table_scratch: &mut [u32],
246        output_scratch: &mut [RaggedBatchOutput],
247    ) -> Result<RaggedBatchReceipt, DecodeRoundError> {
248        let decode_count = self
249            .slots
250            .iter()
251            .filter(|slot| slot.state == RequestState::Decode)
252            .count();
253        if item_scratch.len() < decode_count {
254            return Err(DecodeRoundError::ItemBufferTooSmall);
255        }
256        if output_scratch.len() < decode_count {
257            return Err(DecodeRoundError::OutputBufferTooSmall);
258        }
259        let required_table_entries = self
260            .slots
261            .iter()
262            .filter(|slot| slot.state == RequestState::Decode)
263            .map(|slot| slot.blocks.entries().len())
264            .sum::<usize>();
265        if block_table_scratch.len() < required_table_entries {
266            return Err(DecodeRoundError::BlockTableBufferTooSmall);
267        }
268
269        let mut item_count = 0usize;
270        let mut table_offset = 0usize;
271        for (index, slot) in self.slots.iter().enumerate() {
272            if slot.state != RequestState::Decode {
273                continue;
274            }
275            if !slot.token_ready {
276                return Err(DecodeRoundError::MissingInputToken);
277            }
278            let table = slot.blocks.entries();
279            block_table_scratch[table_offset..table_offset + table.len()].copy_from_slice(table);
280            item_scratch[item_count] = RaggedBatchItem {
281                request_id: slot.request_id,
282                slot: index as u32,
283                token_id: slot.next_token_id,
284                position: slot.token_count,
285                block_table_offset: table_offset as u32,
286                logical_pages: table.len() as u32,
287                _reserved: 0,
288            };
289            item_count += 1;
290            table_offset += table.len();
291        }
292        if item_count == 0 {
293            return Ok(RaggedBatchReceipt::default());
294        }
295
296        let receipt = backend.execute_ragged(
297            &item_scratch[..item_count],
298            &block_table_scratch[..table_offset],
299            &mut output_scratch[..item_count],
300        )?;
301        if receipt.batch_size != item_count as u32 {
302            return Err(DecodeRoundError::OutputIdentityMismatch);
303        }
304        if receipt.backend_launches != 1 {
305            return Err(DecodeRoundError::BackendLaunchCount);
306        }
307        for (item, output) in item_scratch[..item_count]
308            .iter()
309            .zip(&output_scratch[..item_count])
310        {
311            if output.request_id != item.request_id || output.slot != item.slot {
312                return Err(DecodeRoundError::OutputIdentityMismatch);
313            }
314        }
315        for output in &output_scratch[..item_count] {
316            let slot = &mut self.slots[output.slot as usize];
317            slot.next_token_id = output.next_token_id;
318            slot.token_count = slot.token_count.saturating_add(1);
319            slot.token_ready = true;
320        }
321        Ok(receipt)
322    }
323
324    /// Cancellation and normal completion have identical deterministic resource release.
325    pub fn finish(&mut self, request_id: u64, pool: &mut BlockPool) -> Result<(), SchedulerError> {
326        let slot = self.find_mut(request_id)?;
327        slot.blocks.release_all(pool)?;
328        slot.clear();
329        Ok(())
330    }
331
332    fn find_mut(&mut self, request_id: u64) -> Result<&mut RequestSlot, SchedulerError> {
333        self.slots
334            .iter_mut()
335            .find(|slot| slot.state != RequestState::Empty && slot.request_id == request_id)
336            .ok_or(SchedulerError::UnknownRequest)
337    }
338}
qualia_core_db/inference/runtime/scheduler/request_table.rs

qualia_core_db/inference/runtime/scheduler/
request_table.rs
