qualia_core_db/inference/runtime/scheduler/
batch.rs

1//! Caller-buffered ragged decode batches.
2//!
3//! The scheduler lowers active sequences and their page tables into flat ABI records. A backend
4//! receives one call per scheduling round and must return one output for every item in slot order.
5
6#[repr(C)]
7#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)]
8pub struct RaggedBatchItem {
9    pub request_id: u64,
10    pub slot: u32,
11    pub token_id: u32,
12    pub position: u32,
13    pub block_table_offset: u32,
14    pub logical_pages: u32,
15    pub _reserved: u32,
16}
17
18#[repr(C)]
19#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)]
20pub struct RaggedBatchOutput {
21    pub request_id: u64,
22    pub slot: u32,
23    pub next_token_id: u32,
24}
25
26#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
27pub struct RaggedBatchReceipt {
28    pub batch_size: u32,
29    pub backend_launches: u32,
30    pub device_to_host_bytes: u64,
31}
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq)]
34pub enum RaggedBackendError {
35    OutputTooSmall,
36    Rejected,
37}
38
39/// Prepared backend boundary for one continuous-batching decode round.
40///
41/// `items` and `block_tables` are immutable caller storage. `out` has at least `items.len()`
42/// entries. Implementations may not retain any slice after returning.
43pub trait RaggedDecodeBackend {
44    fn execute_ragged(
45        &mut self,
46        items: &[RaggedBatchItem],
47        block_tables: &[u32],
48        out: &mut [RaggedBatchOutput],
49    ) -> Result<RaggedBatchReceipt, RaggedBackendError>;
50}
qualia_core_db/inference/runtime/scheduler/batch.rs

qualia_core_db/inference/runtime/scheduler/
batch.rs
