kime_tensor/
plan.rs

1//! The op graph a model builder emits and a backend lowers, see spec/07-engine.md.
2//!
3//! A graph names its activations as [`Val`]s whose size is a row count times a width, and the row
4//! count is one of the three a batch has: tokens, sequences or markers. Nothing in a graph knows the
5//! batch size. A backend lowers the graph once per [`Bucket`](crate::Bucket), which fixes the row
6//! counts, and [`layout`] gives every value an offset in one arena, reusing space once a value's
7//! last reader has run.
8
9/// An activation in a graph, an index into [`Graph::vals`].
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub struct Val(pub u32);
12
13/// A weight, the index of a tensor in the checkpoint the graph was built from.
14pub type W = usize;
15
16/// What a value has one row per.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum Rows {
19    /// One row per token of the batch.
20    Tokens,
21    /// One row per sequence.
22    Seqs,
23    /// One row per marker.
24    Markers,
25}
26
27/// The shape of a value: `rows` rows of `width` f32.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct Shape {
30    /// What there is one row per.
31    pub rows: Rows,
32    /// Elements per row.
33    pub width: usize,
34}
35
36/// What a GEMM does to each output element after the bias.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum Epilogue {
39    /// Store it.
40    None,
41    /// Store the exact GELU of it.
42    Gelu,
43    /// Store max(0, x).
44    Relu,
45    /// Add it to what `out` already holds, the residual connection.
46    Accumulate,
47}
48
49/// One step of a graph. Every op runs on the rows the batch has, not the bucket's padded count.
50#[derive(Debug, Clone, PartialEq)]
51pub enum Op {
52    /// `out[i] = table[ids[i]]`, token rows.
53    Embed {
54        /// `[vocab, width]`.
55        table: W,
56        /// Token rows.
57        out: Val,
58    },
59    /// LayerNorm over each row, statistics in f64.
60    LayerNorm {
61        /// Input.
62        x: Val,
63        /// Scale.
64        w: W,
65        /// Shift, if the norm has one.
66        b: Option<W>,
67        /// Epsilon.
68        eps: f64,
69        /// Output, the same shape as `x`.
70        out: Val,
71    },
72    /// `out = epilogue(a wᵀ + b)` with `w` as `[n, k]`.
73    Gemm {
74        /// Input, `k` wide.
75        a: Val,
76        /// `[n, k]`.
77        w: W,
78        /// `[n]`, if there is one.
79        b: Option<W>,
80        /// What happens to each result.
81        epilogue: Epilogue,
82        /// Output, `n` wide, with as many rows as `a`.
83        out: Val,
84    },
85    /// Rotary embedding applied in place to the q and k parts of fused `[q | k | v]` rows, with
86    /// positions counted from the start of each sequence.
87    Rope {
88        /// Token rows, `3 heads 64` wide.
89        qkv: Val,
90        /// The base.
91        theta: f64,
92    },
93    /// Self attention within each sequence over fused `[q | k | v]` rows.
94    Attention {
95        /// Token rows, `3 heads 64` wide.
96        qkv: Val,
97        /// Keys at most this far away on either side, or all of them.
98        window: Option<usize>,
99        /// Token rows, `heads 64` wide.
100        out: Val,
101    },
102    /// `out = gelu(x[:, ..n]) * x[:, n..]` with `x` twice as wide as `out`.
103    GeGlu {
104        /// Input.
105        x: Val,
106        /// Output.
107        out: Val,
108    },
109    /// Adds `table[qtype[s]]` to every row of sequence `s`, in place.
110    AddType {
111        /// Token rows.
112        h: Val,
113        /// `[types, width]`.
114        table: W,
115    },
116    /// The row of each marker.
117    GatherMarkers {
118        /// Token rows.
119        h: Val,
120        /// Marker rows.
121        out: Val,
122    },
123    /// Per sequence, its first row followed by `[top1, top1 - top2, entropy / ln k, k / 255]` over
124    /// the softmax of its markers' logits, Laya's act head input.
125    ActFeatures {
126        /// Token rows.
127        h: Val,
128        /// Marker rows, one wide.
129        logits: Val,
130        /// Sequence rows, `h` width plus four.
131        out: Val,
132    },
133    /// Per sequence, the mean of its token rows, Laya's `embed_fn_from_agent` pooling.
134    MeanPool {
135        /// Token rows.
136        h: Val,
137        /// Sequence rows, as wide as `h`.
138        out: Val,
139    },
140}
141
142impl Op {
143    /// Values the op reads, then values it writes. An in place op lists the value in both.
144    #[must_use]
145    pub fn uses(&self) -> (Vec<Val>, Vec<Val>) {
146        match *self {
147            Op::Embed { out, .. } => (vec![], vec![out]),
148            Op::LayerNorm { x, out, .. } | Op::GeGlu { x, out } => (vec![x], vec![out]),
149            Op::Gemm { a, epilogue, out, .. } => {
150                if epilogue == Epilogue::Accumulate {
151                    (vec![a, out], vec![out])
152                } else {
153                    (vec![a], vec![out])
154                }
155            }
156            Op::Rope { qkv, .. } => (vec![qkv], vec![qkv]),
157            Op::Attention { qkv, out, .. } => (vec![qkv], vec![out]),
158            Op::AddType { h, .. } => (vec![h], vec![h]),
159            Op::GatherMarkers { h, out } | Op::MeanPool { h, out } => (vec![h], vec![out]),
160            Op::ActFeatures { h, logits, out } => (vec![h, logits], vec![out]),
161        }
162    }
163}
164
165/// A model's forward pass as ops over values.
166#[derive(Debug, Clone, Default, PartialEq)]
167pub struct Graph {
168    /// The shape of every value.
169    pub vals: Vec<Shape>,
170    /// The ops, in the order they run.
171    pub ops: Vec<Op>,
172    /// Marker rows, one wide, one logit per option.
173    pub logits: Option<Val>,
174    /// Sequence rows, two wide, the act head.
175    pub act: Option<Val>,
176    /// Sequence rows, the pooled embedding of each sequence.
177    pub pooled: Option<Val>,
178}
179
180impl Graph {
181    /// A new value.
182    ///
183    /// # Panics
184    ///
185    /// Past 2^32 values.
186    pub fn val(&mut self, rows: Rows, width: usize) -> Val {
187        self.vals.push(Shape { rows, width });
188        Val(u32::try_from(self.vals.len() - 1).expect("fewer than 2^32 values"))
189    }
190
191    /// The shape of `v`.
192    #[must_use]
193    pub fn shape(&self, v: Val) -> Shape {
194        self.vals[v.0 as usize]
195    }
196
197    /// Appends an op.
198    pub fn push(&mut self, op: Op) {
199        self.ops.push(op);
200    }
201}
202
203/// Values packed into one arena, in f32 elements.
204#[derive(Debug, Clone, PartialEq, Eq)]
205pub struct Layout {
206    /// Where each value starts, indexed like [`Graph::vals`].
207    pub offsets: Vec<usize>,
208    /// The arena's length.
209    pub len: usize,
210}
211
212/// Offsets are multiples of this, 64 bytes, a cache line.
213pub const ALIGN: usize = 16;
214
215/// Lays the values of `graph` out in one arena for the row counts `rows`, by a greedy interval
216/// colouring over their live ranges. A value lives from the first op that touches it to the last,
217/// and the outputs live to the end. Two values whose ranges overlap never share space, so an op's
218/// inputs and outputs never alias unless the op is in place on one value. Values are placed
219/// largest first at the lowest offset that fits.
220#[must_use]
221pub fn layout(graph: &Graph, rows: impl Fn(Rows) -> usize) -> Layout {
222    let n = graph.vals.len();
223    let mut live = vec![(usize::MAX, 0usize); n];
224    for (i, op) in graph.ops.iter().enumerate() {
225        let (r, w) = op.uses();
226        for v in r.into_iter().chain(w) {
227            let l = &mut live[v.0 as usize];
228            l.0 = l.0.min(i);
229            l.1 = l.1.max(i);
230        }
231    }
232    for v in [graph.logits, graph.act, graph.pooled].into_iter().flatten() {
233        let l = &mut live[v.0 as usize];
234        l.0 = l.0.min(graph.ops.len());
235        l.1 = usize::MAX;
236    }
237    let size: Vec<usize> =
238        graph.vals.iter().map(|s| (rows(s.rows) * s.width).next_multiple_of(ALIGN)).collect();
239    let mut order: Vec<usize> = (0..n).filter(|&v| live[v].0 != usize::MAX).collect();
240    order.sort_by_key(|&v| (std::cmp::Reverse(size[v]), v));
241    let mut offsets = vec![0usize; n];
242    let mut placed: Vec<usize> = Vec::new();
243    let mut len = 0;
244    for v in order {
245        let (a, b) = live[v];
246        let mut taken: Vec<(usize, usize)> = placed
247            .iter()
248            .filter(|&&u| live[u].0 <= b && a <= live[u].1)
249            .map(|&u| (offsets[u], offsets[u] + size[u]))
250            .collect();
251        taken.sort_unstable();
252        let mut at = 0;
253        for (lo, hi) in taken {
254            if at + size[v] <= lo {
255                break;
256            }
257            at = at.max(hi);
258        }
259        offsets[v] = at;
260        len = len.max(at + size[v]);
261        placed.push(v);
262    }
263    Layout { offsets, len }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    fn chain(width: usize, steps: usize) -> Graph {
271        let mut g = Graph::default();
272        let mut x = g.val(Rows::Tokens, width);
273        g.push(Op::Embed { table: 0, out: x });
274        for _ in 0..steps {
275            let y = g.val(Rows::Tokens, width);
276            g.push(Op::LayerNorm { x, w: 1, b: None, eps: 1e-5, out: y });
277            x = y;
278        }
279        g.logits = Some(x);
280        g
281    }
282
283    #[test]
284    fn a_chain_needs_two_buffers() {
285        let g = chain(10, 20);
286        let l = layout(&g, |_| 7);
287        let one = (7 * 10usize).next_multiple_of(ALIGN);
288        assert_eq!(l.len, 2 * one);
289        for (i, op) in g.ops.iter().enumerate().skip(1) {
290            let Op::LayerNorm { x, out, .. } = *op else { unreachable!() };
291            assert_ne!(l.offsets[x.0 as usize], l.offsets[out.0 as usize], "op {i} aliases");
292        }
293    }
294
295    #[test]
296    fn live_values_never_overlap() {
297        // A residual stream that lives throughout, with temporaries of several sizes around it.
298        let mut g = Graph::default();
299        let h = g.val(Rows::Tokens, 8);
300        g.push(Op::Embed { table: 0, out: h });
301        for i in 0..6 {
302            let a = g.val(Rows::Tokens, 8 * (i % 3 + 1));
303            let b = g.val(Rows::Seqs, 3);
304            g.push(Op::Gemm { a: h, w: 0, b: None, epilogue: Epilogue::None, out: a });
305            g.push(Op::GatherMarkers { h: a, out: b });
306            g.push(Op::Gemm { a, w: 0, b: None, epilogue: Epilogue::Accumulate, out: h });
307        }
308        g.act = Some(h);
309        let rows = |r| match r {
310            Rows::Tokens => 33,
311            Rows::Seqs => 5,
312            Rows::Markers => 9,
313        };
314        let l = layout(&g, rows);
315        let mut live = vec![(usize::MAX, 0); g.vals.len()];
316        for (i, op) in g.ops.iter().enumerate() {
317            let (r, w) = op.uses();
318            for v in r.into_iter().chain(w) {
319                let x = &mut live[v.0 as usize];
320                *x = (x.0.min(i), x.1.max(i));
321            }
322        }
323        live[h.0 as usize].1 = usize::MAX;
324        for u in 0..g.vals.len() {
325            for v in 0..u {
326                let time = live[u].0 <= live[v].1 && live[v].0 <= live[u].1;
327                let su = rows(g.vals[u].rows) * g.vals[u].width;
328                let sv = rows(g.vals[v].rows) * g.vals[v].width;
329                let space = l.offsets[u] < l.offsets[v] + sv && l.offsets[v] < l.offsets[u] + su;
330                assert!(!(time && space), "values {u} and {v} overlap");
331                assert!(l.offsets[u] + su <= l.len);
332            }
333        }
334        assert!(l.len < g.vals.iter().map(|s| rows(s.rows) * s.width).sum::<usize>());
335    }
336}
kime_tensor/plan.rs

kime_tensor/
plan.rs
