superfluid_executor/
fake.rs

1//! `FakePrimitives`.
2
3use std::collections::{HashMap, HashSet};
4
5use superfluid_abi::encoding;
6
7use crate::primitives::{SampleSpec,
8    Feed, Input, MemCounters, PrimError, RuntimeDescriptor, RuntimePrimitives, SamplingDefaults, Seq, Vocabulary,
9};
10
11#[derive(Debug, Clone)]
12pub struct FakeConfig {
13    pub vocab: u32,
14    pub eos: u32,
15    pub page: u32,
16    pub max_batch: u32,
17    pub max_seq_len: u64,
18    pub kv_bytes_per_token: u64,
19    pub cells_total: u64,
20    pub truncate_partial: bool,
21    pub max_seqs: u32,
22    pub identity: u8,
23    pub engine_sampling: bool,
24    pub copy_shares_cells: bool,
25    pub takeover_preferred: bool,
26    pub seed_min_tokens: u64,
27    /// Cells the cache may keep in the pool, and tokens it may hold exported
28    /// (the descriptor's limits of the same names); 0 and 0, the default, is
29    /// a runtime whose resting sequences cost nothing.
30    pub cache_resident_cells: u64,
31    pub cache_exported_cells: u64,
32    pub vocabulary: Option<Vocabulary>,
33    pub sampling_defaults: SamplingDefaults,
34    /// Wall time each model step takes, for tests that need ticks to last.
35    pub step_delay: std::time::Duration,
36    /// Serve `step_rows` (a row after every token of a feed).
37    pub verify_rows: bool,
38    /// Predict the token that followed the last earlier occurrence of the
39    /// current one: text that repeats itself, as a prompt lookup drafts.
40    pub copy_model: bool,
41    /// With `copy_model`, every this-many positions predict a fresh token
42    /// instead (0: never): drafts that are taken in part.
43    pub copy_miss_every: usize,
44    /// The descriptor's `prefill_step_tokens`.
45    pub prefill_step_tokens: u32,
46}
47
48impl Default for FakeConfig {
49    fn default() -> Self {
50        FakeConfig {
51            vocab: 32000,
52            eos: 31999,
53            page: 16,
54            max_batch: 8,
55            max_seq_len: 4096,
56            kv_bytes_per_token: 64,
57            cells_total: 0,
58            truncate_partial: true,
59            max_seqs: 0,
60            identity: 0xFA,
61            engine_sampling: true,
62            copy_shares_cells: false,
63            takeover_preferred: false,
64            seed_min_tokens: 0,
65            cache_resident_cells: 0,
66            cache_exported_cells: 0,
67            vocabulary: None,
68            sampling_defaults: SamplingDefaults::default(),
69            step_delay: std::time::Duration::ZERO,
70            verify_rows: false,
71            copy_model: false,
72            copy_miss_every: 0,
73            prefill_step_tokens: 0,
74        }
75    }
76}
77
78pub struct FakePrimitives {
79    pub cfg: FakeConfig,
80    seqs: HashMap<Seq, Vec<u32>>,
81    next_seq: Seq,
82    scripted: HashMap<Seq, Vec<u32>>,
83    fail_next: Option<PrimError>,
84    poison_import: bool,
85    pub feeds: Vec<(usize, bool)>,
86    /// For every step, the cells the pool held for sequences the step did
87    /// not feed: what a runtime that attends over its whole pool (llama.cpp's
88    /// unified cache) pays for on each one.
89    pub beside: Vec<u64>,
90    cells: HashMap<Seq, Vec<u64>>,
91    holders: HashMap<u64, u32>,
92    next_cell: u64,
93    /// What the executor last asked of `sync_steps`.
94    pub synced: Option<bool>,
95}
96
97impl Default for FakePrimitives {
98    fn default() -> Self {
99        FakePrimitives::new(FakeConfig::default())
100    }
101}
102
103const SEED: u64 = 0xCBF2_9CE4_8422_2325;
104
105fn hash(content: &[u32], salt: u64) -> u64 {
106    let mut h = superfluid_fingerprint::fnv1a64(&salt.to_le_bytes(), SEED);
107    for &t in content {
108        h = superfluid_fingerprint::fnv1a64(&t.to_le_bytes(), h);
109    }
110    h
111}
112
113impl FakePrimitives {
114    pub fn new(cfg: FakeConfig) -> FakePrimitives {
115        FakePrimitives {
116            cfg,
117            seqs: HashMap::new(),
118            next_seq: 1,
119            scripted: HashMap::new(),
120            fail_next: None,
121            poison_import: false,
122            feeds: Vec::new(),
123            beside: Vec::new(),
124            cells: HashMap::new(),
125            holders: HashMap::new(),
126            next_cell: 1,
127            synced: None,
128        }
129    }
130
131    fn room_for(&self, more: u64, freed: u64) -> Result<(), PrimError> {
132        let used = (self.holders.len() as u64).saturating_sub(freed);
133        if self.cfg.cells_total > 0 && used + more > self.cfg.cells_total {
134            return Err(PrimError::Capacity);
135        }
136        Ok(())
137    }
138
139    fn take_cells(&mut self, seq: Seq, n: u64) {
140        let ids: Vec<u64> = (0..n).map(|k| self.next_cell + k).collect();
141        self.next_cell += n;
142        for &id in &ids {
143            self.holders.insert(id, 1);
144        }
145        self.cells.entry(seq).or_default().extend(ids);
146    }
147
148    fn drop_cells(&mut self, seq: Seq, from: usize) {
149        let Some(own) = self.cells.get_mut(&seq) else { return };
150        let gone: Vec<u64> = own.drain(from.min(own.len())..).collect();
151        for id in gone {
152            if let Some(n) = self.holders.get_mut(&id) {
153                *n -= 1;
154                if *n == 0 {
155                    self.holders.remove(&id);
156                }
157            }
158        }
159    }
160
161    pub fn token_for(&self, content: &[u32]) -> u32 {
162        let miss = self.cfg.copy_miss_every > 0 && content.len().is_multiple_of(self.cfg.copy_miss_every);
163        if self.cfg.copy_model && !miss {
164            if let Some((&last, before)) = content.split_last() {
165                if let Some(i) = before.iter().rposition(|&t| t == last) {
166                    return content[i + 1];
167                }
168            }
169        }
170        (hash(content, 1) % (self.cfg.vocab as u64 - 1)) as u32
171    }
172
173    pub fn row_for(&self, content: &[u32]) -> Vec<f32> {
174        let best = self.token_for(content);
175        let base = hash(content, 2);
176        let mut row: Vec<f32> = (0..self.cfg.vocab)
177            .map(|i| ((base.wrapping_mul(i as u64 + 1) >> 40) % 1000) as f32 / 1000.0 - 0.5)
178            .collect();
179        row[best as usize] = 10.0;
180        row
181    }
182
183    pub fn contents(&self, seq: Seq) -> Option<&[u32]> {
184        self.seqs.get(&seq).map(|v| v.as_slice())
185    }
186
187    pub fn script(&mut self, seq: Seq, tokens: Vec<u32>) {
188        self.scripted.insert(seq, tokens);
189    }
190
191    pub fn fault_next_step(&mut self, code: u32) {
192        self.fail_next = Some(PrimError::Fault(code));
193    }
194
195    pub fn fail_next_step(&mut self, e: PrimError) {
196        self.fail_next = Some(e);
197    }
198
199    fn count_beside(&mut self, feeds: &[Feed<'_>]) {
200        let fed: HashSet<u64> = feeds.iter().flat_map(|f| self.cells.get(&f.seq).into_iter().flatten().copied()).collect();
201        self.beside.push(self.holders.keys().filter(|id| !fed.contains(id)).count() as u64);
202    }
203
204    pub fn poison_next_import(&mut self) {
205        self.poison_import = true;
206    }
207
208    fn cut_ok(&self, len: u64, cur: u64) -> bool {
209        len == 0 || (len <= cur && (self.cfg.truncate_partial || len == cur))
210    }
211}
212
213impl RuntimePrimitives for FakePrimitives {
214    fn sync_steps(&mut self, on: bool) {
215        self.synced = Some(on);
216    }
217
218    fn describe(&self) -> RuntimeDescriptor {
219        RuntimeDescriptor {
220            runtime_id: "fake".into(),
221            runtime_version: "1".into(),
222            max_batch: self.cfg.max_batch,
223            max_seqs: self.cfg.max_seqs,
224            max_seq_len: self.cfg.max_seq_len,
225            vocab_size: self.cfg.vocab,
226            page_size_tokens: self.cfg.page,
227            kv_bytes_per_token: self.cfg.kv_bytes_per_token,
228            cells_total: self.cfg.cells_total,
229            truncate_partial: self.cfg.truncate_partial,
230            copy_shares_cells: self.cfg.copy_shares_cells,
231            takeover_preferred: self.cfg.takeover_preferred,
232            seed_min_tokens: self.cfg.seed_min_tokens,
233            cache_resident_cells: self.cfg.cache_resident_cells,
234            cache_exported_cells: self.cfg.cache_exported_cells,
235            recurrent: !self.cfg.truncate_partial,
236            export_encodings: vec![encoding::LOSSLESS],
237            engine_sampling: self.cfg.engine_sampling,
238            engine_draws: false,
239            verify_rows: self.cfg.verify_rows,
240            verify_batches: true,
241            prefill_step_tokens: self.cfg.prefill_step_tokens,
242            weights_identity: [self.cfg.identity; 32],
243            architecture: "fake".into(),
244            backend: "cpu".into(),
245            sampling_defaults: self.cfg.sampling_defaults,
246        }
247    }
248
249    fn is_eos(&self, token: u32) -> bool {
250        token == self.cfg.eos
251    }
252
253    fn mem_counters(&self) -> MemCounters {
254        let cells_used = self.holders.len() as u64;
255        MemCounters {
256            allocated_bytes: cells_used * self.cfg.kv_bytes_per_token,
257            cells_used,
258            cells_total: self.cfg.cells_total,
259        }
260    }
261
262    fn seq_create(&mut self) -> Result<Seq, PrimError> {
263        if self.cfg.max_seqs > 0 && self.seqs.len() as u32 >= self.cfg.max_seqs {
264            return Err(PrimError::Capacity);
265        }
266        let s = self.next_seq;
267        self.next_seq += 1;
268        self.seqs.insert(s, Vec::new());
269        self.cells.insert(s, Vec::new());
270        Ok(s)
271    }
272
273    fn seq_free(&mut self, seq: Seq) {
274        self.drop_cells(seq, 0);
275        self.cells.remove(&seq);
276        self.seqs.remove(&seq);
277        self.scripted.remove(&seq);
278    }
279
280    fn seq_len(&self, seq: Seq) -> u64 {
281        self.seqs.get(&seq).map(|v| v.len() as u64).unwrap_or(0)
282    }
283
284    fn cells_released(&self, seqs: &[Seq]) -> u64 {
285        let set: HashSet<Seq> = seqs.iter().copied().collect();
286        let mut held: HashMap<u64, u32> = HashMap::new();
287        for s in &set {
288            for &id in self.cells.get(s).into_iter().flatten() {
289                *held.entry(id).or_default() += 1;
290            }
291        }
292        held.iter().filter(|(id, n)| self.holders.get(id) == Some(n)).count() as u64
293    }
294
295    fn seq_copy(&mut self, src: Seq, dst: Seq, len: u64) -> Result<(), PrimError> {
296        let s = self.seqs.get(&src).ok_or(PrimError::UnknownSeq)?;
297        if !self.cut_ok(len, s.len() as u64) {
298            return Err(PrimError::OutOfBoundary);
299        }
300        if !self.seqs.contains_key(&dst) {
301            return Err(PrimError::UnknownSeq);
302        }
303        let prefix = s[..len as usize].to_vec();
304        if self.cfg.copy_shares_cells {
305            let shared: Vec<u64> = self.cells[&src][..len as usize].to_vec();
306            self.drop_cells(dst, 0);
307            for &id in &shared {
308                *self.holders.get_mut(&id).expect("held by src") += 1;
309            }
310            self.cells.insert(dst, shared);
311        } else {
312            self.room_for(len, self.cells_released(&[dst]))?;
313            self.drop_cells(dst, 0);
314            self.take_cells(dst, len);
315        }
316        *self.seqs.get_mut(&dst).expect("checked") = prefix;
317        Ok(())
318    }
319
320    fn seq_truncate(&mut self, seq: Seq, new_len: u64) -> Result<(), PrimError> {
321        let cur = self.seq_len(seq);
322        if !self.seqs.contains_key(&seq) {
323            return Err(PrimError::UnknownSeq);
324        }
325        if !self.cut_ok(new_len, cur) {
326            return Err(PrimError::OutOfBoundary);
327        }
328        self.seqs.get_mut(&seq).expect("checked").truncate(new_len as usize);
329        self.drop_cells(seq, new_len as usize);
330        Ok(())
331    }
332
333    fn seq_boundary(&self, seq: Seq, cap: u64) -> u64 {
334        let len = self.seq_len(seq);
335        if self.cfg.truncate_partial {
336            cap.min(len)
337        } else if len <= cap {
338            len
339        } else {
340            0
341        }
342    }
343
344    fn seq_export(&mut self, seq: Seq, len: u64, enc: u8) -> Result<Vec<u8>, PrimError> {
345        if enc != encoding::LOSSLESS {
346            return Err(PrimError::Unsupported);
347        }
348        let s = self.seqs.get(&seq).ok_or(PrimError::UnknownSeq)?;
349        if !self.cut_ok(len, s.len() as u64) {
350            return Err(PrimError::OutOfBoundary);
351        }
352        Ok(s[..len as usize].iter().flat_map(|t| t.to_le_bytes()).collect())
353    }
354
355    fn seq_import(&mut self, seq: Seq, payload: &[u8]) -> Result<u64, PrimError> {
356        let s = self.seqs.get(&seq).ok_or(PrimError::UnknownSeq)?;
357        if !s.is_empty() || !payload.len().is_multiple_of(4) {
358            return Err(PrimError::OutOfBoundary);
359        }
360        self.room_for(payload.len() as u64 / 4, 0)?;
361        self.take_cells(seq, payload.len() as u64 / 4);
362        let s = self.seqs.get_mut(&seq).expect("checked");
363        *s = payload.as_chunks::<4>().0.iter().map(|c| u32::from_le_bytes(*c)).collect();
364        if std::mem::take(&mut self.poison_import) {
365            return Err(PrimError::Fatal);
366        }
367        Ok(s.len() as u64)
368    }
369
370    fn step(&mut self, feeds: &[Feed<'_>]) -> Result<Vec<Vec<f32>>, PrimError> {
371        if let Some(e) = self.fail_next.take() {
372            return Err(e);
373        }
374        if !self.cfg.step_delay.is_zero() {
375            std::thread::sleep(self.cfg.step_delay);
376        }
377        let mut more = 0u64;
378        for f in feeds {
379            let Input::Tokens(t) = f.input else { return Err(PrimError::Unsupported) };
380            if !self.seqs.contains_key(&f.seq) {
381                return Err(PrimError::UnknownSeq);
382            }
383            more += t.len() as u64;
384        }
385        self.room_for(more, 0)?;
386        self.count_beside(feeds);
387        let mut rows = Vec::new();
388        for f in feeds {
389            let Input::Tokens(t) = f.input else { return Err(PrimError::Unsupported) };
390            self.take_cells(f.seq, t.len() as u64);
391            self.feeds.push((t.len(), f.wants_row));
392            let s = self.seqs.get_mut(&f.seq).expect("checked");
393            s.extend_from_slice(t);
394            if f.wants_row {
395                let content = s.clone();
396                let mut row = self.row_for(&content);
397                if let Some(script) = self.scripted.get_mut(&f.seq) {
398                    if !script.is_empty() {
399                        let next = script.remove(0);
400                        for x in row.iter_mut() {
401                            *x = -1.0;
402                        }
403                        row[next as usize] = 10.0;
404                    }
405                }
406                rows.push(row);
407            }
408        }
409        Ok(rows)
410    }
411
412    fn step_rows(&mut self, feeds: &[Feed<'_>]) -> Result<Option<Vec<Vec<Vec<f32>>>>, PrimError> {
413        if !self.cfg.verify_rows {
414            return Ok(None);
415        }
416        if let Some(e) = self.fail_next.take() {
417            return Err(e);
418        }
419        if !self.cfg.step_delay.is_zero() {
420            std::thread::sleep(self.cfg.step_delay);
421        }
422        let mut more = 0u64;
423        for f in feeds {
424            let Input::Tokens(t) = f.input else { return Err(PrimError::Unsupported) };
425            if !self.seqs.contains_key(&f.seq) {
426                return Err(PrimError::UnknownSeq);
427            }
428            more += t.len() as u64;
429        }
430        self.room_for(more, 0)?;
431        self.count_beside(feeds);
432        let mut out = Vec::new();
433        for f in feeds {
434            let Input::Tokens(t) = f.input else { return Err(PrimError::Unsupported) };
435            self.take_cells(f.seq, t.len() as u64);
436            self.feeds.push((t.len(), f.wants_row));
437            let s = self.seqs.get_mut(&f.seq).expect("checked");
438            let base = s.len();
439            s.extend_from_slice(t);
440            if f.wants_row {
441                let content = s.clone();
442                out.push((1..=t.len()).map(|i| self.row_for(&content[..base + i])).collect());
443            }
444        }
445        Ok(Some(out))
446    }
447
448    fn step_sampled(&mut self, feeds: &[Feed<'_>], specs: &[SampleSpec]) -> Result<Option<Vec<u32>>, PrimError> {
449        if !self.cfg.engine_sampling || specs.len() != feeds.len() || specs.iter().any(|s| s.temperature > 0.0) {
450            return Ok(None);
451        }
452        let rows = self.step(feeds)?;
453        Ok(Some(rows.iter().map(|r| crate::sampling::argmax(r)).collect()))
454    }
455
456    fn vocabulary(&self) -> Option<Vocabulary> {
457        self.cfg.vocabulary.clone()
458    }
459}
superfluid_executor/fake.rs

superfluid_executor/
fake.rs
