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}
134
135impl Op {
136    /// Values the op reads, then values it writes. An in place op lists the value in both.
137    #[must_use]
138    pub fn uses(&self) -> (Vec<Val>, Vec<Val>) {
139        match *self {
140            Op::Embed { out, .. } => (vec![], vec![out]),
141            Op::LayerNorm { x, out, .. } | Op::GeGlu { x, out } => (vec![x], vec![out]),
142            Op::Gemm { a, epilogue, out, .. } => {
143                if epilogue == Epilogue::Accumulate {
144                    (vec![a, out], vec![out])
145                } else {
146                    (vec![a], vec![out])
147                }
148            }
149            Op::Rope { qkv, .. } => (vec![qkv], vec![qkv]),
150            Op::Attention { qkv, out, .. } => (vec![qkv], vec![out]),
151            Op::AddType { h, .. } => (vec![h], vec![h]),
152            Op::GatherMarkers { h, out } => (vec![h], vec![out]),
153            Op::ActFeatures { h, logits, out } => (vec![h, logits], vec![out]),
154        }
155    }
156}
157
158/// A model's forward pass as ops over values.
159#[derive(Debug, Clone, Default, PartialEq)]
160pub struct Graph {
161    /// The shape of every value.
162    pub vals: Vec<Shape>,
163    /// The ops, in the order they run.
164    pub ops: Vec<Op>,
165    /// Marker rows, one wide, one logit per option.
166    pub logits: Option<Val>,
167    /// Sequence rows, two wide, the act head.
168    pub act: Option<Val>,
169}
170
171impl Graph {
172    /// A new value.
173    ///
174    /// # Panics
175    ///
176    /// Past 2^32 values.
177    pub fn val(&mut self, rows: Rows, width: usize) -> Val {
178        self.vals.push(Shape { rows, width });
179        Val(u32::try_from(self.vals.len() - 1).expect("fewer than 2^32 values"))
180    }
181
182    /// The shape of `v`.
183    #[must_use]
184    pub fn shape(&self, v: Val) -> Shape {
185        self.vals[v.0 as usize]
186    }
187
188    /// Appends an op.
189    pub fn push(&mut self, op: Op) {
190        self.ops.push(op);
191    }
192}
193
194/// Values packed into one arena, in f32 elements.
195#[derive(Debug, Clone, PartialEq, Eq)]
196pub struct Layout {
197    /// Where each value starts, indexed like [`Graph::vals`].
198    pub offsets: Vec<usize>,
199    /// The arena's length.
200    pub len: usize,
201}
202
203/// Offsets are multiples of this, 64 bytes, a cache line.
204pub const ALIGN: usize = 16;
205
206/// Lays the values of `graph` out in one arena for the row counts `rows`, by a greedy interval
207/// colouring over their live ranges. A value lives from the first op that touches it to the last,
208/// and the outputs live to the end. Two values whose ranges overlap never share space, so an op's
209/// inputs and outputs never alias unless the op is in place on one value. Values are placed
210/// largest first at the lowest offset that fits.
211#[must_use]
212pub fn layout(graph: &Graph, rows: impl Fn(Rows) -> usize) -> Layout {
213    let n = graph.vals.len();
214    let mut live = vec![(usize::MAX, 0usize); n];
215    for (i, op) in graph.ops.iter().enumerate() {
216        let (r, w) = op.uses();
217        for v in r.into_iter().chain(w) {
218            let l = &mut live[v.0 as usize];
219            l.0 = l.0.min(i);
220            l.1 = l.1.max(i);
221        }
222    }
223    for v in [graph.logits, graph.act].into_iter().flatten() {
224        let l = &mut live[v.0 as usize];
225        l.0 = l.0.min(graph.ops.len());
226        l.1 = usize::MAX;
227    }
228    let size: Vec<usize> =
229        graph.vals.iter().map(|s| (rows(s.rows) * s.width).next_multiple_of(ALIGN)).collect();
230    let mut order: Vec<usize> = (0..n).filter(|&v| live[v].0 != usize::MAX).collect();
231    order.sort_by_key(|&v| (std::cmp::Reverse(size[v]), v));
232    let mut offsets = vec![0usize; n];
233    let mut placed: Vec<usize> = Vec::new();
234    let mut len = 0;
235    for v in order {
236        let (a, b) = live[v];
237        let mut taken: Vec<(usize, usize)> = placed
238            .iter()
239            .filter(|&&u| live[u].0 <= b && a <= live[u].1)
240            .map(|&u| (offsets[u], offsets[u] + size[u]))
241            .collect();
242        taken.sort_unstable();
243        let mut at = 0;
244        for (lo, hi) in taken {
245            if at + size[v] <= lo {
246                break;
247            }
248            at = at.max(hi);
249        }
250        offsets[v] = at;
251        len = len.max(at + size[v]);
252        placed.push(v);
253    }
254    Layout { offsets, len }
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260
261    fn chain(width: usize, steps: usize) -> Graph {
262        let mut g = Graph::default();
263        let mut x = g.val(Rows::Tokens, width);
264        g.push(Op::Embed { table: 0, out: x });
265        for _ in 0..steps {
266            let y = g.val(Rows::Tokens, width);
267            g.push(Op::LayerNorm { x, w: 1, b: None, eps: 1e-5, out: y });
268            x = y;
269        }
270        g.logits = Some(x);
271        g
272    }
273
274    #[test]
275    fn a_chain_needs_two_buffers() {
276        let g = chain(10, 20);
277        let l = layout(&g, |_| 7);
278        let one = (7 * 10usize).next_multiple_of(ALIGN);
279        assert_eq!(l.len, 2 * one);
280        for (i, op) in g.ops.iter().enumerate().skip(1) {
281            let Op::LayerNorm { x, out, .. } = *op else { unreachable!() };
282            assert_ne!(l.offsets[x.0 as usize], l.offsets[out.0 as usize], "op {i} aliases");
283        }
284    }
285
286    #[test]
287    fn live_values_never_overlap() {
288        // A residual stream that lives throughout, with temporaries of several sizes around it.
289        let mut g = Graph::default();
290        let h = g.val(Rows::Tokens, 8);
291        g.push(Op::Embed { table: 0, out: h });
292        for i in 0..6 {
293            let a = g.val(Rows::Tokens, 8 * (i % 3 + 1));
294            let b = g.val(Rows::Seqs, 3);
295            g.push(Op::Gemm { a: h, w: 0, b: None, epilogue: Epilogue::None, out: a });
296            g.push(Op::GatherMarkers { h: a, out: b });
297            g.push(Op::Gemm { a, w: 0, b: None, epilogue: Epilogue::Accumulate, out: h });
298        }
299        g.act = Some(h);
300        let rows = |r| match r {
301            Rows::Tokens => 33,
302            Rows::Seqs => 5,
303            Rows::Markers => 9,
304        };
305        let l = layout(&g, rows);
306        let mut live = vec![(usize::MAX, 0); g.vals.len()];
307        for (i, op) in g.ops.iter().enumerate() {
308            let (r, w) = op.uses();
309            for v in r.into_iter().chain(w) {
310                let x = &mut live[v.0 as usize];
311                *x = (x.0.min(i), x.1.max(i));
312            }
313        }
314        live[h.0 as usize].1 = usize::MAX;
315        for u in 0..g.vals.len() {
316            for v in 0..u {
317                let time = live[u].0 <= live[v].1 && live[v].0 <= live[u].1;
318                let su = rows(g.vals[u].rows) * g.vals[u].width;
319                let sv = rows(g.vals[v].rows) * g.vals[v].width;
320                let space = l.offsets[u] < l.offsets[v] + sv && l.offsets[v] < l.offsets[u] + su;
321                assert!(!(time && space), "values {u} and {v} overlap");
322                assert!(l.offsets[u] + su <= l.len);
323            }
324        }
325        assert!(l.len < g.vals.iter().map(|s| rows(s.rows) * s.width).sum::<usize>());
326    }
327}
kime_tensor/plan.rs

kime_tensor/
plan.rs
