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    /// Bytes the weights take on the device.
181    fn weight_bytes(&self, _w: &Self::Weights) -> usize {
182        0
183    }
184
185    /// Bytes a plan takes on the device, its arena and staging buffers.
186    fn plan_bytes(&self, _p: &Self::Plan) -> usize {
187        0
188    }
189}
190
191/// A graph, its weights on one backend, and a plan per bucket built on first use.
192#[derive(Debug)]
193pub struct Executor<B: Backend> {
194    backend: B,
195    weights: B::Weights,
196    graph: Graph,
197    stage: String,
198    buckets: Vec<Bucket>,
199    plans: Vec<(Bucket, B::Plan)>,
200    vocab: usize,
201    types: usize,
202}
203
204impl<B: Backend> Executor<B> {
205    /// Uploads the weights. `vocab` and `types` bound the token ids and question types a batch may
206    /// hold, and `stage` names the buckets to use from `buckets`.
207    ///
208    /// # Errors
209    ///
210    /// From [`Backend::upload`].
211    pub fn new(
212        backend: B,
213        tensors: &[HostTensor<'_>],
214        graph: Graph,
215        buckets: &Buckets,
216        stage: &str,
217        vocab: usize,
218        types: usize,
219    ) -> Result<Self> {
220        let weights = backend.upload(tensors, &graph)?;
221        Ok(Self {
222            backend,
223            weights,
224            graph,
225            stage: stage.to_string(),
226            buckets: buckets.stage(stage).to_vec(),
227            plans: Vec::new(),
228            vocab,
229            types,
230        })
231    }
232
233    /// The backend.
234    pub fn backend(&self) -> &B {
235        &self.backend
236    }
237
238    /// The buckets with a plan built.
239    pub fn warm(&self) -> impl Iterator<Item = Bucket> + '_ {
240        self.plans.iter().map(|p| p.0)
241    }
242
243    /// Bytes on the device: the weights, and the plans built so far.
244    pub fn memory(&self) -> (usize, usize) {
245        let plans = self.plans.iter().map(|p| self.backend.plan_bytes(&p.1)).sum();
246        (self.backend.weight_bytes(&self.weights), plans)
247    }
248
249    /// The plans built so far, to inspect or profile.
250    pub fn plans_mut(&mut self) -> impl Iterator<Item = &mut B::Plan> + '_ {
251        self.plans.iter_mut().map(|p| &mut p.1)
252    }
253
254    /// Builds the plan for `bucket` now rather than on first use.
255    ///
256    /// # Errors
257    ///
258    /// From [`Backend::lower`].
259    pub fn prepare(&mut self, bucket: Bucket) -> Result<usize> {
260        if let Some(i) = self.plans.iter().position(|p| p.0 == bucket) {
261            return Ok(i);
262        }
263        let plan = self.backend.lower(&self.weights, &self.graph, bucket)?;
264        self.plans.push((bucket, plan));
265        Ok(self.plans.len() - 1)
266    }
267
268    /// Runs a batch in the smallest bucket that holds it. Once that bucket's plan exists and `out`
269    /// has grown to the batch size, this allocates nothing.
270    ///
271    /// # Errors
272    ///
273    /// [`Error::Batch`] for an inconsistent batch, [`Error::NoBucket`] for one too big, and
274    /// anything the backend reports.
275    pub fn run(&mut self, batch: &Batch<'_>, out: &mut Outputs) -> Result<Bucket> {
276        batch.check(self.vocab, self.types)?;
277        let (t, s, m) = (batch.ids.len(), batch.seqs(), batch.markers.len());
278        let bucket = self.buckets.iter().copied().find(|b| b.holds(t, s, m)).ok_or_else(|| {
279            Error::NoBucket { stage: self.stage.clone(), tokens: t, seqs: s, markers: m }
280        })?;
281        let i = self.prepare(bucket)?;
282        self.backend.run(&mut self.plans[i].1, batch, out)?;
283        Ok(bucket)
284    }
285}
286
287/// Owned storage for a [`Batch`], reused from one batch to the next so that building one does
288/// not allocate once it has grown.
289#[derive(Debug, Clone, Default)]
290pub struct BatchBuf {
291    ids: Vec<u32>,
292    cu: Vec<u32>,
293    markers: Vec<u32>,
294    mcu: Vec<u32>,
295    qtype: Vec<u8>,
296}
297
298impl BatchBuf {
299    /// Empties it, keeping the memory.
300    pub fn clear(&mut self) {
301        for v in [&mut self.ids, &mut self.cu, &mut self.markers, &mut self.mcu] {
302            v.clear();
303        }
304        self.qtype.clear();
305    }
306
307    /// Appends a sequence.
308    ///
309    /// # Panics
310    ///
311    /// If the batch passes 2^32 tokens or markers.
312    pub fn push(&mut self, ids: &[u32], markers: &[u32], qtype: u8) {
313        if self.cu.is_empty() {
314            self.cu.push(0);
315            self.mcu.push(0);
316        }
317        self.ids.extend_from_slice(ids);
318        self.markers.extend_from_slice(markers);
319        self.cu.push(u32::try_from(self.ids.len()).expect("under 2^32 tokens"));
320        self.mcu.push(u32::try_from(self.markers.len()).expect("under 2^32 markers"));
321        self.qtype.push(qtype);
322    }
323
324    /// The batch.
325    #[must_use]
326    pub fn batch(&self) -> Batch<'_> {
327        const EMPTY: &[u32] = &[0];
328        let (cu, mcu) =
329            if self.cu.is_empty() { (EMPTY, EMPTY) } else { (&self.cu[..], &self.mcu[..]) };
330        Batch { ids: &self.ids, cu, markers: &self.markers, mcu, qtype: &self.qtype }
331    }
332}
kime_tensor/backend.rs

kime_tensor/
backend.rs
