superfluid_engine/
rings.rs

1//! Shm-ring abstraction.
2
3use superfluid_abi::{RingRef, TokenRef};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
6pub enum RingError {
7    #[error("unknown ring id {0}")]
8    UnknownRing(u32),
9    #[error("stale generation (ref {reference}, current {current})")]
10    StaleGeneration { reference: u64, current: u64 },
11    #[error("span out of ring bounds")]
12    OutOfBounds,
13}
14
15#[derive(Debug, Clone, Copy)]
16pub struct RingAttachment {
17    pub ring_id: u32,
18    pub kind: u32,
19    pub role: u32,
20    pub slots: u32,
21    pub slot_bytes: u32,
22    pub base: *mut u8,
23    pub len: usize,
24}
25
26pub trait Rings {
27    fn read_tokens(&self, r: &TokenRef) -> Result<Vec<u32>, RingError>;
28
29    fn write_tokens(&mut self, tokens: &[u32]) -> Result<TokenRef, RingError>;
30
31    fn write_logits_row(&mut self, row: &[f32]) -> RingRef;
32
33    fn read_logits_row(&self, r: &RingRef) -> Result<Vec<f32>, RingError>;
34
35    /// The daemon wants the running tick to end at its next step boundary
36    /// (a request is waiting to be admitted, or a lane was cancelled).
37    fn yield_requested(&self) -> bool {
38        false
39    }
40
41    /// Latency-sensitive requests have arrived lately: run prefill in short
42    /// steps, so a yield comes soon after it is asked for.
43    fn latency_wanted(&self) -> bool {
44        false
45    }
46}
47
48pub const TOKEN_RING_IN: u32 = 1;
49pub const TOKEN_RING_OUT: u32 = 2;
50pub const LOGITS_RING: u32 = 3;
51
52#[derive(Debug)]
53pub struct InMemoryRings {
54    token_in: RingBuf<Vec<u32>>,
55    token_out: RingBuf<Vec<u32>>,
56    logits: RingBuf<Vec<f32>>,
57    yield_requested: bool,
58    latency_wanted: bool,
59}
60
61#[derive(Debug)]
62struct RingBuf<T> {
63    slots: Vec<Option<(u64, T)>>,
64    next_index: usize,
65    next_generation: u64,
66}
67
68impl<T: Clone> RingBuf<T> {
69    fn new(slots: usize) -> Self {
70        RingBuf {
71            slots: vec![None; slots],
72            next_index: 0,
73            next_generation: 1,
74        }
75    }
76
77    fn push(&mut self, payload: T) -> (u32, u64) {
78        let index = self.next_index;
79        let generation = self.next_generation;
80        self.slots[index] = Some((generation, payload));
81        self.next_index = (index + 1) % self.slots.len();
82        self.next_generation += 1;
83        (index as u32, generation)
84    }
85
86    fn get(&self, index: u32, generation: u64) -> Result<&T, RingError> {
87        let slot = self
88            .slots
89            .get(index as usize)
90            .ok_or(RingError::OutOfBounds)?;
91        match slot {
92            Some((gen, payload)) if *gen == generation => Ok(payload),
93            Some((gen, _)) => Err(RingError::StaleGeneration {
94                reference: generation,
95                current: *gen,
96            }),
97            None => Err(RingError::StaleGeneration {
98                reference: generation,
99                current: 0,
100            }),
101        }
102    }
103}
104
105impl InMemoryRings {
106    pub fn new() -> Self {
107        InMemoryRings {
108            token_in: RingBuf::new(64),
109            token_out: RingBuf::new(64),
110            logits: RingBuf::new(16),
111            yield_requested: false,
112            latency_wanted: false,
113        }
114    }
115
116    /// What `yield_requested` answers, for tests of an engine's early stop.
117    pub fn set_yield(&mut self, on: bool) {
118        self.yield_requested = on;
119    }
120
121    /// What `latency_wanted` answers, for tests.
122    pub fn set_latency(&mut self, on: bool) {
123        self.latency_wanted = on;
124    }
125
126    pub fn stage_prompt(&mut self, tokens: &[u32]) -> TokenRef {
127        let (index, generation) = self.token_in.push(tokens.to_vec());
128        TokenRef {
129            ring_id: TOKEN_RING_IN,
130            index,
131            count: tokens.len() as u32,
132            _pad0: 0,
133            generation,
134        }
135    }
136}
137
138impl Default for InMemoryRings {
139    fn default() -> Self {
140        Self::new()
141    }
142}
143
144impl Rings for InMemoryRings {
145    fn read_tokens(&self, r: &TokenRef) -> Result<Vec<u32>, RingError> {
146        let buf = match r.ring_id {
147            TOKEN_RING_IN => &self.token_in,
148            TOKEN_RING_OUT => &self.token_out,
149            other => return Err(RingError::UnknownRing(other)),
150        };
151        let tokens = buf.get(r.index, r.generation)?;
152        if tokens.len() != r.count as usize {
153            return Err(RingError::OutOfBounds);
154        }
155        Ok(tokens.clone())
156    }
157
158    fn write_tokens(&mut self, tokens: &[u32]) -> Result<TokenRef, RingError> {
159        let (index, generation) = self.token_out.push(tokens.to_vec());
160        Ok(TokenRef {
161            ring_id: TOKEN_RING_OUT,
162            index,
163            count: tokens.len() as u32,
164            _pad0: 0,
165            generation,
166        })
167    }
168
169    fn write_logits_row(&mut self, row: &[f32]) -> RingRef {
170        let (index, generation) = self.logits.push(row.to_vec());
171        RingRef {
172            ring_id: LOGITS_RING,
173            index,
174            generation,
175        }
176    }
177
178    fn read_logits_row(&self, r: &RingRef) -> Result<Vec<f32>, RingError> {
179        if r.ring_id != LOGITS_RING {
180            return Err(RingError::UnknownRing(r.ring_id));
181        }
182        self.logits.get(r.index, r.generation).cloned()
183    }
184
185    fn yield_requested(&self) -> bool {
186        self.yield_requested
187    }
188
189    fn latency_wanted(&self) -> bool {
190        self.latency_wanted
191    }
192}
superfluid_engine/rings.rs

superfluid_engine/
rings.rs
