superfluid_executor/
primitives.rs

1//! `RuntimePrimitives`.
2
3pub type Seq = u64;
4
5/// The most rows `step_rows` answers for one feed: a verified draft of
6/// `VERIFY_ROWS_MAX - 1` tokens and the token before it.
7pub const VERIFY_ROWS_MAX: u32 = 8;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum PrimError {
11    Unsupported,
12    UnknownSeq,
13    OutOfBoundary,
14    Capacity,
15    Fault(u32),
16    Fatal,
17}
18
19#[derive(Debug, Clone)]
20pub struct RuntimeDescriptor {
21    pub runtime_id: String,
22    pub runtime_version: String,
23    pub max_batch: u32,
24    pub max_seqs: u32,
25    pub max_seq_len: u64,
26    pub vocab_size: u32,
27    pub page_size_tokens: u32,
28    pub kv_bytes_per_token: u64,
29    pub cells_total: u64,
30    pub truncate_partial: bool,
31    pub copy_shares_cells: bool,
32    pub takeover_preferred: bool,
33    pub seed_min_tokens: u64,
34    /// Cells the prefix cache may keep in the pool: the state its entries
35    /// hold that no lane or bare sequence shares. `0` = no bound: an entry at
36    /// rest costs the other sequences nothing. A runtime whose steps cost by
37    /// what its pool holds names what it affords (llama.cpp's unified cache
38    /// runs every step's attention over the cells up to the highest one in
39    /// use, whoever holds them); past it the least recently used entries
40    /// leave the pool, exported or evicted.
41    pub cache_resident_cells: u64,
42    /// Tokens of state the prefix cache may hold out of the pool as lossless exports.
43    /// `0` = none: an entry that leaves the pool is evicted.
44    pub cache_exported_cells: u64,
45    pub recurrent: bool,
46    pub export_encodings: Vec<u8>,
47    pub engine_sampling: bool,
48    pub engine_draws: bool,
49    /// `step_rows` answers a row for every position of a feed: what
50    /// verifying a draft in one pass needs.
51    pub verify_rows: bool,
52    /// `step_rows` runs every lane's draft in one pass; a runtime that runs
53    /// them one lane at a time drafts only while one lane decodes.
54    pub verify_batches: bool,
55    /// The most prompt tokens one prefill step runs (the runtime's own
56    /// micro-batch); a tick asked to yield stops between such steps. 0: a
57    /// tick's prefill is one step.
58    pub prefill_step_tokens: u32,
59    pub weights_identity: [u8; 32],
60    pub architecture: String,
61    pub backend: String,
62    pub sampling_defaults: SamplingDefaults,
63}
64
65#[derive(Debug, Clone, Copy, Default, PartialEq)]
66pub struct SamplingDefaults {
67    pub temperature: Option<f64>,
68    pub top_p: Option<f64>,
69    pub top_k: Option<u32>,
70    pub min_p: Option<f64>,
71    pub repetition_penalty: Option<f64>,
72    pub do_sample: Option<bool>,
73}
74
75impl SamplingDefaults {
76    pub fn from_generation_config(v: &serde_json::Value) -> SamplingDefaults {
77        let f = |k: &str| v.get(k).and_then(serde_json::Value::as_f64);
78        SamplingDefaults {
79            temperature: f("temperature"),
80            top_p: f("top_p"),
81            top_k: v.get("top_k").and_then(serde_json::Value::as_u64).map(|x| x as u32),
82            min_p: f("min_p"),
83            repetition_penalty: f("repetition_penalty"),
84            do_sample: v.get("do_sample").and_then(serde_json::Value::as_bool),
85        }
86    }
87}
88
89#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
90pub struct MemCounters {
91    pub allocated_bytes: u64,
92    pub cells_used: u64,
93    pub cells_total: u64,
94}
95
96#[derive(Debug, Clone, Default)]
97pub struct Vocabulary {
98    pub tokens: Vec<Vec<u8>>,
99    pub specials: Vec<(String, u32)>,
100    pub eos: Vec<u32>,
101}
102
103#[derive(Debug, Clone, Copy)]
104pub struct Feed<'a> {
105    pub seq: Seq,
106    pub input: Input<'a>,
107    pub wants_row: bool,
108}
109
110#[derive(Debug, Clone, Copy)]
111pub enum Input<'a> {
112    Tokens(&'a [u32]),
113    Media { handle: u64, offset: u32 },
114}
115
116#[derive(Debug, Clone, Copy, PartialEq)]
117pub struct SampleSpec {
118    pub temperature: f32,
119    pub top_k: u32,
120    pub top_p: f32,
121    pub min_p: f32,
122    pub rng_position: u64,
123}
124
125pub trait RuntimePrimitives {
126    fn describe(&self) -> RuntimeDescriptor;
127    fn is_eos(&self, token: u32) -> bool;
128    fn mem_counters(&self) -> MemCounters;
129
130    fn seq_create(&mut self) -> Result<Seq, PrimError>;
131    fn seq_free(&mut self, seq: Seq);
132    fn seq_len(&self, seq: Seq) -> u64;
133    fn cells_released(&self, seqs: &[Seq]) -> u64 {
134        let mut seen = std::collections::HashSet::new();
135        seqs.iter().filter(|s| seen.insert(**s)).map(|&s| self.seq_len(s)).sum()
136    }
137    fn seq_copy(&mut self, src: Seq, dst: Seq, len: u64) -> Result<(), PrimError>;
138    fn seq_truncate(&mut self, seq: Seq, new_len: u64) -> Result<(), PrimError>;
139    fn seq_boundary(&self, seq: Seq, cap: u64) -> u64;
140    fn seq_export(&mut self, seq: Seq, len: u64, encoding: u8) -> Result<Vec<u8>, PrimError>;
141    fn seq_import(&mut self, seq: Seq, payload: &[u8]) -> Result<u64, PrimError>;
142
143    fn step(&mut self, feeds: &[Feed<'_>]) -> Result<Vec<Vec<f32>>, PrimError>;
144    fn step_sampled(&mut self, _feeds: &[Feed<'_>], _specs: &[SampleSpec]) -> Result<Option<Vec<u32>>, PrimError> {
145        Ok(None)
146    }
147
148    /// `step`, answering for each feed that wants rows a row after every
149    /// one of its tokens (in order), not only its last. `Ok(None)`: not
150    /// served, and nothing moved.
151    fn step_rows(&mut self, _feeds: &[Feed<'_>]) -> Result<Option<Vec<Vec<Vec<f32>>>>, PrimError> {
152        Ok(None)
153    }
154
155    /// Whether a step that reads no row is done when it returns, rather than queued.
156    fn sync_steps(&mut self, _on: bool) {}
157
158    fn shed(&mut self, _bytes_target: u64) -> u64 {
159        0
160    }
161
162    fn vocabulary(&self) -> Option<Vocabulary> {
163        None
164    }
165}
superfluid_executor/primitives.rs

superfluid_executor/
primitives.rs
