kime_tensor/
backend.rs

1//! The `Backend` trait and the executor that picks a bucket and replays its plan.
2
3use std::fmt;
4
5use crate::bucket::{Bucket, Buckets};
6use crate::dtype::DType;
7use crate::plan::Graph;
8
9/// A weight as it sits in the checkpoint, before a backend converts it.
10#[derive(Debug, Clone, Copy)]
11pub struct HostTensor<'a> {
12    /// Element type.
13    pub dtype: DType,
14    /// Shape.
15    pub shape: &'a [usize],
16    /// Little endian, row major.
17    pub bytes: &'a [u8],
18}
19
20/// A batch, flattened. Sequence `s` is tokens `cu[s]..cu[s + 1]` and markers
21/// `mcu[s]..mcu[s + 1]`, and each marker is a position within its own sequence.
22#[derive(Debug, Clone, Copy, Default)]
23pub struct Batch<'a> {
24    /// Token ids of every sequence, end to end.
25    pub ids: &'a [u32],
26    /// Sequence starts in `ids`, with the total at the end.
27    pub cu: &'a [u32],
28    /// Marker positions of every sequence, end to end.
29    pub markers: &'a [u32],
30    /// Sequence starts in `markers`, with the total at the end.
31    pub mcu: &'a [u32],
32    /// The question type of each sequence.
33    pub qtype: &'a [u8],
34}
35
36impl Batch<'_> {
37    /// Sequences.
38    #[must_use]
39    pub fn seqs(&self) -> usize {
40        self.cu.len().saturating_sub(1)
41    }
42
43    /// Checks that the pieces agree with each other, so a backend can index without checks.
44    ///
45    /// # Errors
46    ///
47    /// [`Error::Batch`] naming the first problem.
48    pub fn check(&self, vocab: usize, types: usize) -> Result<()> {
49        let bad = |m: String| Err(Error::Batch(m));
50        let s = self.seqs();
51        if self.cu.first() != Some(&0) || self.mcu.first() != Some(&0) {
52            return bad("cu and mcu must start at 0".into());
53        }
54        if self.mcu.len() != s + 1 || self.qtype.len() != s {
55            return bad(format!(
56                "{s} sequences but {} mcu and {} qtype",
57                self.mcu.len(),
58                self.qtype.len()
59            ));
60        }
61        if self.cu[s] as usize != self.ids.len() || self.mcu[s] as usize != self.markers.len() {
62            return bad("cu and mcu must end at the totals".into());
63        }
64        for i in 0..s {
65            let (a, b) = (self.cu[i], self.cu[i + 1]);
66            if b < a || self.mcu[i + 1] < self.mcu[i] {
67                return bad(format!("sequence {i} ends before it starts"));
68            }
69            if usize::from(self.qtype[i]) >= types {
70                return bad(format!("sequence {i} has type {} of {types}", self.qtype[i]));
71            }
72            for &m in &self.markers[self.mcu[i] as usize..self.mcu[i + 1] as usize] {
73                if m >= b - a {
74                    return bad(format!("marker {m} is past sequence {i} of {} tokens", b - a));
75                }
76            }
77        }
78        if let Some(&id) = self.ids.iter().find(|&&id| id as usize >= vocab) {
79            return bad(format!("token id {id} is past the vocabulary of {vocab}"));
80        }
81        Ok(())
82    }
83}
84
85/// What a batch produces. Reused across batches, so after the first few it does not allocate.
86#[derive(Debug, Clone, Default, PartialEq)]
87pub struct Outputs {
88    /// One logit per marker, in batch order.
89    pub logits: Vec<f32>,
90    /// The act head, one pair per sequence.
91    pub act: Vec<[f32; 2]>,
92}
93
94/// What a backend can do.
95#[derive(Debug, Clone, PartialEq, Eq)]
96pub struct Caps {
97    /// `cpu`, `cuda`, `metal` or `ane`.
98    pub name: &'static str,
99    /// Worker threads, for the CPU.
100    pub threads: usize,
101    /// Whether plans are captured as device graphs, which makes padding to the bucket matter.
102    pub graphs: bool,
103    /// Whether host and device share memory.
104    pub unified_memory: bool,
105}
106
107/// Engine errors.
108#[derive(Debug, Clone, PartialEq, Eq)]
109pub enum Error {
110    /// The batch is inconsistent.
111    Batch(String),
112    /// No bucket of the stage holds the batch.
113    NoBucket {
114        /// The stage.
115        stage: String,
116        /// Tokens in the batch.
117        tokens: usize,
118        /// Sequences.
119        seqs: usize,
120        /// Markers.
121        markers: usize,
122    },
123    /// The graph uses something the backend does not have, or weights of the wrong shape.
124    Unsupported(String),
125    /// The device or its driver failed.
126    Device(String),
127}
128
129impl fmt::Display for Error {
130    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
131        match self {
132            Error::Batch(m) => write!(f, "bad batch: {m}"),
133            Error::NoBucket { stage, tokens, seqs, markers } => write!(
134                f,
135                "no {stage} bucket holds {tokens} tokens, {seqs} sequences and {markers} markers"
136            ),
137            Error::Unsupported(m) => write!(f, "unsupported: {m}"),
138            Error::Device(m) => write!(f, "device: {m}"),
139        }
140    }
141}
142
143impl std::error::Error for Error {}
144
145/// Engine results.
146pub type Result<T> = std::result::Result<T, Error>;
147
148/// A device that runs graphs. The engine is generic over it, so the path through it is
149/// monomorphized and there is no dynamic dispatch per op.
150pub trait Backend: Send + Sync {
151    /// Weights in the backend's own layout.
152    type Weights: Send + Sync;
153    /// A graph lowered for one bucket, with its arena.
154    type Plan: Send;
155
156    /// What the backend can do.
157    fn caps(&self) -> Caps;
158
159    /// Converts a checkpoint's tensors, indexed as the graph's weights are.
160    ///
161    /// # Errors
162    ///
163    /// When a tensor cannot be converted.
164    fn upload(&self, tensors: &[HostTensor<'_>], graph: &Graph) -> Result<Self::Weights>;
165
166    /// Lowers a graph for one bucket. This is where memory is allocated.
167    ///
168    /// # Errors
169    ///
170    /// [`Error::Unsupported`] for an op or a shape the backend cannot run.
171    fn lower(&self, w: &Self::Weights, graph: &Graph, bucket: Bucket) -> Result<Self::Plan>;
172
173    /// Runs a batch that fits the plan's bucket and has been checked.
174    ///
175    /// # Errors
176    ///
177    /// When the device fails.
178    fn run(&self, plan: &mut Self::Plan, batch: &Batch<'_>, out: &mut Outputs) -> Result<()>;
179}
180
181/// A graph, its weights on one backend, and a plan per bucket built on first use.
182#[derive(Debug)]
183pub struct Executor<B: Backend> {
184    backend: B,
185    weights: B::Weights,
186    graph: Graph,
187    stage: String,
188    buckets: Vec<Bucket>,
189    plans: Vec<(Bucket, B::Plan)>,
190    vocab: usize,
191    types: usize,
192}
193
194impl<B: Backend> Executor<B> {
195    /// Uploads the weights. `vocab` and `types` bound the token ids and question types a batch may
196    /// hold, and `stage` names the buckets to use from `buckets`.
197    ///
198    /// # Errors
199    ///
200    /// From [`Backend::upload`].
201    pub fn new(
202        backend: B,
203        tensors: &[HostTensor<'_>],
204        graph: Graph,
205        buckets: &Buckets,
206        stage: &str,
207        vocab: usize,
208        types: usize,
209    ) -> Result<Self> {
210        let weights = backend.upload(tensors, &graph)?;
211        Ok(Self {
212            backend,
213            weights,
214            graph,
215            stage: stage.to_string(),
216            buckets: buckets.stage(stage).to_vec(),
217            plans: Vec::new(),
218            vocab,
219            types,
220        })
221    }
222
223    /// The backend.
224    pub fn backend(&self) -> &B {
225        &self.backend
226    }
227
228    /// The buckets with a plan built.
229    pub fn warm(&self) -> impl Iterator<Item = Bucket> + '_ {
230        self.plans.iter().map(|p| p.0)
231    }
232
233    /// The plans built so far, to inspect or profile.
234    pub fn plans_mut(&mut self) -> impl Iterator<Item = &mut B::Plan> + '_ {
235        self.plans.iter_mut().map(|p| &mut p.1)
236    }
237
238    /// Builds the plan for `bucket` now rather than on first use.
239    ///
240    /// # Errors
241    ///
242    /// From [`Backend::lower`].
243    pub fn prepare(&mut self, bucket: Bucket) -> Result<usize> {
244        if let Some(i) = self.plans.iter().position(|p| p.0 == bucket) {
245            return Ok(i);
246        }
247        let plan = self.backend.lower(&self.weights, &self.graph, bucket)?;
248        self.plans.push((bucket, plan));
249        Ok(self.plans.len() - 1)
250    }
251
252    /// Runs a batch in the smallest bucket that holds it. Once that bucket's plan exists and `out`
253    /// has grown to the batch size, this allocates nothing.
254    ///
255    /// # Errors
256    ///
257    /// [`Error::Batch`] for an inconsistent batch, [`Error::NoBucket`] for one too big, and
258    /// anything the backend reports.
259    pub fn run(&mut self, batch: &Batch<'_>, out: &mut Outputs) -> Result<Bucket> {
260        batch.check(self.vocab, self.types)?;
261        let (t, s, m) = (batch.ids.len(), batch.seqs(), batch.markers.len());
262        let bucket = self.buckets.iter().copied().find(|b| b.holds(t, s, m)).ok_or_else(|| {
263            Error::NoBucket { stage: self.stage.clone(), tokens: t, seqs: s, markers: m }
264        })?;
265        let i = self.prepare(bucket)?;
266        self.backend.run(&mut self.plans[i].1, batch, out)?;
267        Ok(bucket)
268    }
269}
270
271/// Owned storage for a [`Batch`], reused from one batch to the next so that building one does
272/// not allocate once it has grown.
273#[derive(Debug, Clone, Default)]
274pub struct BatchBuf {
275    ids: Vec<u32>,
276    cu: Vec<u32>,
277    markers: Vec<u32>,
278    mcu: Vec<u32>,
279    qtype: Vec<u8>,
280}
281
282impl BatchBuf {
283    /// Empties it, keeping the memory.
284    pub fn clear(&mut self) {
285        for v in [&mut self.ids, &mut self.cu, &mut self.markers, &mut self.mcu] {
286            v.clear();
287        }
288        self.qtype.clear();
289    }
290
291    /// Appends a sequence.
292    ///
293    /// # Panics
294    ///
295    /// If the batch passes 2^32 tokens or markers.
296    pub fn push(&mut self, ids: &[u32], markers: &[u32], qtype: u8) {
297        if self.cu.is_empty() {
298            self.cu.push(0);
299            self.mcu.push(0);
300        }
301        self.ids.extend_from_slice(ids);
302        self.markers.extend_from_slice(markers);
303        self.cu.push(u32::try_from(self.ids.len()).expect("under 2^32 tokens"));
304        self.mcu.push(u32::try_from(self.markers.len()).expect("under 2^32 markers"));
305        self.qtype.push(qtype);
306    }
307
308    /// The batch.
309    #[must_use]
310    pub fn batch(&self) -> Batch<'_> {
311        const EMPTY: &[u32] = &[0];
312        let (cu, mcu) =
313            if self.cu.is_empty() { (EMPTY, EMPTY) } else { (&self.cu[..], &self.mcu[..]) };
314        Batch { ids: &self.ids, cu, markers: &self.markers, mcu, qtype: &self.qtype }
315    }
316}
kime_tensor/backend.rs

kime_tensor/
backend.rs
