logicaffeine_compile/spatial/decode_engine/
lower_graph.rs

1//! The engine's program read off a model's step graph: the decode step [`crate::model::author::decode_step`] writes
2//! for a [`ModelSpec`], lifted into the one IR and fused by the shared table under the engine's shim
3//! ([`super::shim::EngineShim`]), read block by block into the [`Command`]s of [`super::program`].
4//!
5//! [`lower_graph`] reads the step as the blocks the engine has units for: the embedding, then each layer's norm, its
6//! mixer (a Gated-DeltaNet block or an attention block), the residual add, the FFN's norm, the SwiGLU FFN and its add,
7//! then the step's advance, the final norm and the head. Every op of a block is checked three ways before the block
8//! is written: its kind, the tensors it reads and writes (each block reads the residual the block before left and the
9//! next one reads what it wrote), and the checkpoint tensor and view each constant it reads is bound to
10//! ([`crate::model::bind`]). The engine's merges are the shared table's fused symbols, each read as the unit that runs
11//! it: the projections of one checkpoint matrix (`merged_qkv`) — row ranges of one tensor that follow one another from
12//! its first row (a Gated-DeltaNet block's query, key and value of `attn_qkv`), or the interleaved halves of each head
13//! (an attention block's queries and output gate of `attn_q`) — as one matvec; each convolution with its SiLU
14//! (`conv_bias_silu`) as a channel of the conv unit; the sigmoid and the decay a delta rule reads (`gdn_gating`) as the
15//! gates; and an attention's output gate (`glu_act_mul`) as the sigmoid gate. A block is written in the order the
16//! engine's units take it, which is not the order the author writes its ops in. What a block does that the engine has
17//! no unit for — a norm after the mixer, an FFN that is not SwiGLU, experts, latent attention, attention sinks, a
18//! sliding window, a scaled embedding — is refused by name, so a program the engine runs computes what the step graph
19//! says.
20
21use crate::model::bind::{self, Binding, Source, View};
22use crate::model::spec::{Act, Attention, Ffn, GatedDelta, Mixer, ModelSpec};
23use crate::spatial::passes::recognize_stages::TensorOpSpec;
24use crate::tensor_graph::graph::{Fused, Op, Sym, TensorGraph};
25use crate::tensor_graph::symbols::ClassId;
26use crate::tensor_recognize::{ElementwiseOp, RowReduceKind};
27
28use super::model::{CacheFormat, EngineNumerics};
29use super::program::{Command, Program, Slot};
30
31/// Why a step graph is not a program the engine runs.
32#[derive(Clone, Debug, PartialEq, Eq)]
33pub enum LowerError {
34    /// An op other than the one the block expects next: where, what it is, and what was expected.
35    Unexpected { op: usize, kind: &'static str, block: String, expected: String },
36    /// The step ended inside a block.
37    Ended { block: String, expected: String },
38    /// An op reads a tensor other than the one the block's dataflow says it reads.
39    Dataflow { op: usize, block: String, expected: String, found: String },
40    /// A constant the checkpoint does not bind, or binds through a view the block does not read it by.
41    Binding { constant: String, why: String },
42    /// A part of the model the engine has no unit for, by name.
43    Unsupported { layer: Option<usize>, what: String },
44    /// Ops left after the head.
45    Trailing { op: usize, kind: &'static str },
46}
47
48impl std::fmt::Display for LowerError {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match self {
51            LowerError::Unexpected { op, kind, block, expected } => write!(f, "op {op} ({kind}) in {block}, where {expected} comes next"),
52            LowerError::Ended { block, expected } => write!(f, "the step ends in {block}, before {expected}"),
53            LowerError::Dataflow { op, block, expected, found } => write!(f, "op {op} in {block} reads {found}, where it reads {expected}"),
54            LowerError::Binding { constant, why } => write!(f, "{constant}: {why}"),
55            LowerError::Unsupported { layer: Some(l), what } => write!(f, "layer {l}: the engine has no unit for {what}"),
56            LowerError::Unsupported { layer: None, what } => write!(f, "the engine has no unit for {what}"),
57            LowerError::Trailing { op, kind } => write!(f, "op {op} ({kind}) after the head"),
58        }
59    }
60}
61
62impl std::error::Error for LowerError {}
63
64/// An op's kind as a refusal names it: a lifted op's sugar, or a fused symbol's name.
65fn kind(op: &Op) -> &'static str {
66    match &op.sym {
67        Sym::Prim(spec) => spec.kind(),
68        Sym::Fused(f) => f.class.id(),
69        Sym::View { .. } => "View",
70        Sym::Quantize(_) => "Quantize",
71    }
72}
73
74/// The lifted op `op` is, which a block reads as a part of a unit.
75fn prim(op: &Op) -> &TensorOpSpec {
76    op.prim().expect("a part a unit runs is a lifted op")
77}
78
79/// The program of `graph`, the step `spec` authors as the engine's shim leaves it, at `numerics`.
80pub fn lower_graph(graph: &TensorGraph, spec: &ModelSpec, numerics: EngineNumerics) -> Result<Program, LowerError> {
81    let refuse = |what: &str| Err(LowerError::Unsupported { layer: None, what: what.into() });
82    if spec.embedding_scale.is_some() {
83        return refuse("a scaled embedding");
84    }
85    if spec.final_logit_softcap.is_some() {
86        return refuse("capped logits");
87    }
88    let mut r = Reader { graph, spec, at: 0, block: String::new(), commands: Vec::new() };
89    let mut residual = r.embed()?;
90    for l in 0..spec.layers.len() {
91        residual = r.layer(l, &residual, numerics)?;
92    }
93    r.head(&residual)?;
94    if let Some(op) = graph.ops.get(r.at) {
95        return Err(LowerError::Trailing { op: r.at, kind: kind(op) });
96    }
97    Ok(Program { numerics, commands: r.commands })
98}
99
100/// A walk over the step's ops in order, writing commands block by block.
101struct Reader<'a> {
102    graph: &'a TensorGraph,
103    spec: &'a ModelSpec,
104    at: usize,
105    block: String,
106    commands: Vec<Command>,
107}
108
109impl<'a> Reader<'a> {
110    /// The next op, which must be of `kind`.
111    fn take(&mut self, kind: &'static str) -> Result<&'a Op, LowerError> {
112        let Some(op) = self.graph.ops.get(self.at) else {
113            return Err(LowerError::Ended { block: self.block.clone(), expected: kind.into() });
114        };
115        if self::kind(op) != kind {
116            return Err(LowerError::Unexpected { op: self.at, kind: self::kind(op), block: self.block.clone(), expected: kind.into() });
117        }
118        self.at += 1;
119        Ok(op)
120    }
121
122    /// The next op, which must be the fused symbol `class`, and its members.
123    fn take_fused(&mut self, class: ClassId) -> Result<(&'a Op, &'a [Op]), LowerError> {
124        let op = self.take(class.id())?;
125        let Sym::Fused(Fused { parts, .. }) = &op.sym else { unreachable!("taken as a fused symbol") };
126        Ok((op, parts))
127    }
128
129    /// Whether the next op is of `kind`.
130    fn next_is(&self, kind: &str) -> bool {
131        self.graph.ops.get(self.at).is_some_and(|op| self::kind(op) == kind)
132    }
133
134    /// That `op` (the one just taken) reads `expected` as its operand `k`.
135    fn reads(&self, op: &Op, k: usize, expected: &str) -> Result<(), LowerError> {
136        match op.ins.get(k) {
137            Some(found) if found == expected => Ok(()),
138            found => Err(LowerError::Dataflow {
139                op: self.at - 1,
140                block: self.block.clone(),
141                expected: expected.into(),
142                found: found.cloned().unwrap_or_else(|| "nothing".into()),
143            }),
144        }
145    }
146
147    /// The checkpoint tensor and view `constant` is bound to.
148    fn bound(&self, constant: &str) -> Result<Binding, LowerError> {
149        match bind::bind(self.spec, constant) {
150            Ok(Source::Weight(b)) => Ok(b),
151            Ok(Source::State) => Err(LowerError::Binding { constant: constant.into(), why: "carried state where a weight is read".into() }),
152            Err(bind::Unbound(_)) => Err(LowerError::Binding { constant: constant.into(), why: "the checkpoint binds nothing to it".into() }),
153        }
154    }
155
156    /// The tensor `constant` reads whole, as the matrix its projection multiplies by (its rows the outputs), in the
157    /// checkpoint's own row order or reordered within heads for the author's rotation.
158    fn matrix(&self, constant: &str) -> Result<String, LowerError> {
159        let b = self.bound(constant)?;
160        match b.view {
161            View::Transposed | View::NeoxPairs { transpose: true, .. } => Ok(b.tensor),
162            view => Err(LowerError::Binding { constant: constant.into(), why: format!("bound through {view:?}, not as a whole matrix") }),
163        }
164    }
165
166    /// The tensor `constant` reads whole, element for element.
167    fn vector(&self, constant: &str) -> Result<String, LowerError> {
168        let b = self.bound(constant)?;
169        match b.view {
170            View::Same | View::NeoxPairs { transpose: false, .. } => Ok(b.tensor),
171            view => Err(LowerError::Binding { constant: constant.into(), why: format!("bound through {view:?}, not as a whole vector") }),
172        }
173    }
174
175    /// The one tensor `constants` split into row ranges, each following the one before from the tensor's first row,
176    /// and the rows each takes.
177    fn row_ranges(&self, constants: &[&str]) -> Result<(String, Vec<usize>), LowerError> {
178        let (mut tensor, mut next, mut counts) = (None::<String>, 0usize, Vec::new());
179        for c in constants {
180            let b = self.bound(c)?;
181            let View::RowRange { first, count, transpose: true } = b.view else {
182                return Err(LowerError::Binding { constant: (*c).into(), why: format!("bound through {:?}, not as a row range", b.view) });
183            };
184            match &tensor {
185                Some(t) if *t != b.tensor => {
186                    return Err(LowerError::Binding { constant: (*c).into(), why: format!("a range of {}, not of {t}", b.tensor) });
187                }
188                _ => tensor = Some(b.tensor),
189            }
190            if first != next {
191                return Err(LowerError::Binding { constant: (*c).into(), why: format!("rows from {first}, where the range before ends at {next}") });
192            }
193            next = first + count;
194            counts.push(count);
195        }
196        Ok((tensor.expect("at least one constant"), counts))
197    }
198
199    /// `x` RMS-normed by the weight it reads: the norm op and its weight.
200    fn norm(&mut self, x: &str) -> Result<(&'a Op, String), LowerError> {
201        let op = self.take("FusedRow")?;
202        match prim(op) {
203            TensorOpSpec::FusedRow(r) if r.reduce == RowReduceKind::RmsNorm && r.bias.is_none() && r.residual.is_none() => {}
204            other => {
205                return Err(LowerError::Unexpected { op: self.at - 1, kind: other.kind(), block: self.block.clone(), expected: "an RMS norm with a weight".into() });
206            }
207        }
208        self.reads(op, 0, x)?;
209        let weight = op.ins.get(1).cloned().ok_or_else(|| LowerError::Dataflow {
210            op: self.at - 1,
211            block: self.block.clone(),
212            expected: "a norm weight".into(),
213            found: "nothing".into(),
214        })?;
215        Ok((op, weight))
216    }
217
218    /// A projection of `x`: the op and the constant it multiplies by.
219    fn projection(&mut self, x: &str) -> Result<(&'a Op, String), LowerError> {
220        let op = self.take("Matmul")?;
221        self.reads(op, 0, x)?;
222        Ok((op, op.ins.get(1).cloned().unwrap_or_default()))
223    }
224
225    /// `a op b` into the op's result.
226    fn elementwise(&mut self, op: ElementwiseOp, a: &str, b: &str) -> Result<&'a Op, LowerError> {
227        let e = self.take("Elementwise")?;
228        self.product(e, op, a, b)?;
229        Ok(e)
230    }
231
232    /// That `e`, the op just taken or a member of it, is `a op b`.
233    fn product(&self, e: &Op, op: ElementwiseOp, a: &str, b: &str) -> Result<(), LowerError> {
234        if !matches!(e.prim(), Some(TensorOpSpec::Elementwise(s)) if s.op == op && s.row.is_none()) {
235            return Err(LowerError::Unexpected { op: self.at - 1, kind: kind(e), block: self.block.clone(), expected: format!("an elementwise {op:?}") });
236        }
237        self.reads(e, 0, a)?;
238        self.reads(e, 1, b)
239    }
240
241    fn embed(&mut self) -> Result<String, LowerError> {
242        self.block = "the embedding".into();
243        let op = self.take("RowGather")?;
244        let table = self.vector(&op.ins[0])?;
245        self.commands.push(Command::Embed { table });
246        Ok(op.out.clone())
247    }
248
249    /// Layer `l`, reading the residual `x`; the residual it leaves.
250    fn layer(&mut self, l: usize, x: &str, numerics: EngineNumerics) -> Result<String, LowerError> {
251        let spec: &'a ModelSpec = self.spec;
252        let layer = &spec.layers[l];
253        let refuse = |what: &str| Err(LowerError::Unsupported { layer: Some(l), what: what.into() });
254        if layer.post_mixer_norm || layer.post_ffn_norm {
255            return refuse("a norm after the mixer or the FFN");
256        }
257        self.block = format!("layer {l}'s norm");
258        let (normed, weight) = self.norm(x)?;
259        let normed = normed.out.clone();
260        let attn_norm = self.vector(&weight)?;
261        self.commands.push(Command::Norm { weight: attn_norm, from: Slot::Residual, to: Slot::Normed });
262        let out = match &layer.mixer {
263            Mixer::GatedDelta(gd) => self.delta(l, &normed, gd)?,
264            Mixer::Attention(a) => self.attention(l, &normed, a, numerics)?,
265            Mixer::Mla(_) => return refuse("multi-head latent attention"),
266        };
267        self.block = format!("layer {l}'s mixer add");
268        let mixed = self.elementwise(ElementwiseOp::Add, x, &out)?.out.clone();
269        self.commands.push(Command::Add { from: Slot::Out, to: Slot::Residual });
270        match &layer.ffn {
271            Ffn::Dense { act: Act::Silu, .. } => {}
272            Ffn::Dense { .. } => return refuse("an FFN gated by other than SiLU"),
273            Ffn::Moe(_) => return refuse("a mixture of experts"),
274        }
275        self.block = format!("layer {l}'s FFN");
276        let (normed, weight) = self.norm(&mixed)?;
277        let normed = normed.out.clone();
278        let ffn_norm = self.vector(&weight)?;
279        self.commands.push(Command::Norm { weight: ffn_norm, from: Slot::Residual, to: Slot::Normed });
280        let (gate, wg) = self.projection(&normed)?;
281        let (up, wu) = self.projection(&normed)?;
282        let swiglu = self.take("SwiGlu")?;
283        self.reads(swiglu, 0, &gate.out)?;
284        self.reads(swiglu, 1, &up.out)?;
285        let (down, wd) = self.projection(&swiglu.out)?;
286        let (ffn_gate, ffn_up, ffn_down) = (self.matrix(&wg)?, self.matrix(&wu)?, self.matrix(&wd)?);
287        self.commands.push(Command::Matvec { tensor: ffn_gate, from: Slot::Normed, to: Slot::FfnGate });
288        self.commands.push(Command::Matvec { tensor: ffn_up, from: Slot::Normed, to: Slot::FfnUp });
289        self.commands.push(Command::SwiGlu { gate: Slot::FfnGate, up: Slot::FfnUp, to: Slot::Hidden });
290        self.commands.push(Command::Matvec { tensor: ffn_down, from: Slot::Hidden, to: Slot::Out });
291        self.block = format!("layer {l}'s FFN add");
292        let out = self.elementwise(ElementwiseOp::Add, &mixed, &down.out)?.out.clone();
293        self.commands.push(Command::Add { from: Slot::Out, to: Slot::Residual });
294        Ok(out)
295    }
296
297    /// A Gated-DeltaNet block over the normed `x`: the block's output.
298    fn delta(&mut self, l: usize, x: &str, gd: &GatedDelta) -> Result<String, LowerError> {
299        self.block = format!("layer {l}'s Gated-DeltaNet block");
300        let (hd, nk, nv) = (gd.head_dim, gd.k_heads, gd.v_heads);
301        let (_, projections) = self.take_fused(ClassId::MergedQkv)?;
302        let [q, k, v] = projections else {
303            return Err(LowerError::Unexpected { op: self.at - 1, kind: ClassId::MergedQkv.id(), block: self.block.clone(), expected: "the query, key and value as one matvec".into() });
304        };
305        for p in [q, k, v] {
306            self.reads(p, 0, x)?;
307        }
308        let (wq, wk, wv) = (q.ins[1].clone(), k.ins[1].clone(), v.ins[1].clone());
309        let (z, wz) = self.projection(x)?;
310        let (qkv, rows) = self.row_ranges(&[&wq, &wk, &wv])?;
311        if rows != [nk * hd, nk * hd, nv * hd] {
312            return Err(LowerError::Binding { constant: wq.clone(), why: format!("query, key and value of {rows:?} rows, where the spec says {}, {} and {}", nk * hd, nk * hd, nv * hd) });
313        }
314        let gate = self.matrix(&wz)?;
315        let (mut convs, mut silus) = (Vec::new(), Vec::new());
316        for from in [q, k, v] {
317            let (_, channel) = self.take_fused(ClassId::ConvBiasSilu)?;
318            let [c, s] = channel else { unreachable!("a convolution and its SiLU") };
319            self.reads(c, 0, &from.out)?;
320            self.reads(s, 0, &c.out)?;
321            convs.push(c);
322            silus.push(s);
323        }
324        let kernels: Vec<&str> = convs.iter().map(|c| c.ins[1].as_str()).collect();
325        let (kernel, parts) = self.row_ranges(&kernels)?;
326        if parts != rows {
327            return Err(LowerError::Binding { constant: kernels[0].into(), why: format!("the convolution's parts of {parts:?} rows, where its operands are {rows:?}") });
328        }
329        let taps = match prim(convs[0]) {
330            TensorOpSpec::CausalConv(c) => c.taps,
331            _ => unreachable!("a channel's convolution"),
332        };
333        if taps != gd.conv_taps {
334            return Err(LowerError::Binding { constant: kernels[0].into(), why: format!("{taps} taps, where the spec says {}", gd.conv_taps) });
335        }
336        let mut norms = Vec::new();
337        for s in &silus[..2] {
338            let n = self.take("L2Norm")?;
339            self.reads(n, 0, &s.out)?;
340            norms.push(n);
341        }
342        let (beta, wb) = self.projection(x)?;
343        let (alpha, wa) = self.projection(x)?;
344        let (_, gates) = self.take_fused(ClassId::GdnGating)?;
345        let [sigmoid, decay] = gates else { unreachable!("a write strength and a forget gate") };
346        self.reads(sigmoid, 0, &beta.out)?;
347        self.reads(decay, 0, &alpha.out)?;
348        let (dt, a) = (self.vector(&decay.ins[1])?, self.vector(&decay.ins[2])?);
349        let delta = self.take("GatedDelta")?;
350        for (k, from) in [&norms[0].out, &norms[1].out, &silus[2].out, &z.out, &sigmoid.out, &decay.out].into_iter().enumerate() {
351            self.reads(delta, k, from)?;
352        }
353        let ssm_norm = self.vector(&delta.ins[7])?;
354        let (out, wo) = self.projection(&delta.out)?;
355        let (ssm_beta, ssm_alpha, ssm_out) = (self.matrix(&wb)?, self.matrix(&wa)?, self.matrix(&wo)?);
356        self.commands.push(Command::Matvec { tensor: qkv, from: Slot::Normed, to: Slot::Qkv });
357        self.commands.push(Command::Matvec { tensor: gate, from: Slot::Normed, to: Slot::Z });
358        self.commands.push(Command::Matvec { tensor: ssm_beta, from: Slot::Normed, to: Slot::Beta });
359        self.commands.push(Command::Matvec { tensor: ssm_alpha, from: Slot::Normed, to: Slot::Alpha });
360        self.commands.push(Command::Conv { layer: l, kernel, taps, from: Slot::Qkv, to: Slot::Mixed });
361        self.commands.push(Command::L2Heads { from: Slot::Mixed, at: 0, heads: nk, width: hd, to: Slot::Query });
362        self.commands.push(Command::L2Heads { from: Slot::Mixed, at: nk * hd, heads: nk, width: hd, to: Slot::Key });
363        self.commands.push(Command::Gates {
364            alpha: Slot::Alpha,
365            beta: Slot::Beta,
366            dt,
367            a,
368            log_decay: Slot::LogDecay,
369            decay: Slot::Decay,
370            gate: Slot::Gate,
371        });
372        self.commands.push(Command::Delta {
373            layer: l,
374            query: Slot::Query,
375            key: Slot::Key,
376            value: Slot::Mixed,
377            value_at: 2 * nk * hd,
378            log_decay: Slot::LogDecay,
379            decay: Slot::Decay,
380            gate: Slot::Gate,
381            k_heads: nk,
382            v_heads: nv,
383            width: hd,
384            to: Slot::Read,
385        });
386        self.commands.push(Command::GatedNorm { weight: ssm_norm, from: Slot::Read, z: Slot::Z, width: hd, to: Slot::Read });
387        self.commands.push(Command::Matvec { tensor: ssm_out, from: Slot::Read, to: Slot::Out });
388        Ok(out.out.clone())
389    }
390
391    /// An attention block over the normed `x`: the block's output.
392    fn attention(&mut self, l: usize, x: &str, a: &Attention, numerics: EngineNumerics) -> Result<String, LowerError> {
393        self.block = format!("layer {l}'s attention block");
394        let refuse = |what: &str| Err(LowerError::Unsupported { layer: Some(l), what: what.into() });
395        if a.sinks {
396            return refuse("attention sinks");
397        }
398        if a.window.is_some() {
399            return refuse("a sliding window");
400        }
401        if a.bias || a.output_bias {
402            return refuse("biased projections");
403        }
404        if a.softcap.is_some() {
405            return refuse("capped scores");
406        }
407        if a.fused_qkv {
408            return refuse("one projection split into queries, keys and values");
409        }
410        let (hd, vd) = (a.head_dim, a.value_dim);
411        let rotated = numerics.cache == CacheFormat::Int8Hadamard;
412        let (query_out, gate_out, attn_q) = if a.output_gate {
413            let (_, projections) = self.take_fused(ClassId::MergedQkv)?;
414            let [q, g] = projections else {
415                return Err(LowerError::Unexpected { op: self.at - 1, kind: ClassId::MergedQkv.id(), block: self.block.clone(), expected: "the queries and their output gate as one matvec".into() });
416            };
417            self.reads(q, 0, x)?;
418            self.reads(g, 0, x)?;
419            let (wq, wqg) = (q.ins[1].clone(), g.ins[1].clone());
420            let (bq, bg) = (self.bound(&wq)?, self.bound(&wqg)?);
421            let halves = matches!(
422                (&bq.view, &bg.view),
423                (View::HeadPart { stride: s1, offset: 0, width: w1, transpose: true }, View::HeadPart { stride: s2, offset: o2, width: w2, transpose: true })
424                    if *s1 == 2 * hd && *s2 == 2 * hd && *w1 == hd && *o2 == hd && *w2 == hd
425            );
426            if !halves || bq.tensor != bg.tensor {
427                return Err(LowerError::Binding { constant: wq, why: format!("queries {:?} of {} and gates {:?} of {}, not one head's halves of one matrix", bq.view, bq.tensor, bg.view, bg.tensor) });
428            }
429            (q.out.clone(), Some(g.out.clone()), bq.tensor)
430        } else {
431            let (q, wq) = self.projection(x)?;
432            (q.out.clone(), None, self.matrix(&wq)?)
433        };
434        let (k, wk) = self.projection(x)?;
435        let (v, wv) = self.projection(x)?;
436        let (attn_k, attn_v) = (self.matrix(&wk)?, self.matrix(&wv)?);
437        let (mut query, mut key) = (query_out, k.out.clone());
438        let mut heads = Vec::new();
439        if a.qk_norm {
440            for (slot, from) in [(Slot::Query, &mut query), (Slot::Key, &mut key)] {
441                let (op, weight) = self.norm(from)?;
442                *from = op.out.clone();
443                heads.push((slot, Some(self.vector(&weight)?)));
444            }
445        } else {
446            heads = vec![(Slot::Query, None), (Slot::Key, None)];
447        }
448        let mut turned = [false, false];
449        if a.rope.is_some() {
450            for (i, from) in [&mut query, &mut key].into_iter().enumerate() {
451                let op = self.take("Rope")?;
452                self.reads(op, 0, from)?;
453                *from = op.out.clone();
454                turned[i] = true;
455            }
456        }
457        let set_k = self.take("SetRow")?;
458        self.reads(set_k, 1, &key)?;
459        let set_v = self.take("SetRow")?;
460        self.reads(set_v, 1, &v.out)?;
461        let attend = self.take("Attention")?;
462        self.reads(attend, 0, &query)?;
463        let mut read = attend.out.clone();
464        if let Some(gate) = &gate_out {
465            let (gated, parts) = self.take_fused(ClassId::GluActMul)?;
466            let [s, product] = parts else { unreachable!("an activation and the product it gates") };
467            self.reads(s, 0, gate)?;
468            self.product(product, ElementwiseOp::Mul, &read, &s.out)?;
469            read = gated.out.clone();
470        }
471        let (out, wo) = self.projection(&read)?;
472        let attn_output = self.matrix(&wo)?;
473        let projected = if a.output_gate { Slot::Projected } else { Slot::Query };
474        self.commands.push(Command::Matvec { tensor: attn_q, from: Slot::Normed, to: projected });
475        self.commands.push(Command::Matvec { tensor: attn_k, from: Slot::Normed, to: Slot::Key });
476        self.commands.push(Command::Matvec { tensor: attn_v, from: Slot::Normed, to: Slot::Value });
477        if a.output_gate {
478            self.commands.push(Command::Split { from: Slot::Projected, heads: a.heads, width: hd, query: Slot::Query, gate: Slot::OutGate });
479        }
480        for (i, (slot, weight)) in heads.into_iter().enumerate() {
481            if let Some(weight) = weight {
482                self.commands.push(Command::RmsHeads { weight, slot, width: hd });
483            }
484            if let (true, Some(rope)) = (turned[i], &a.rope) {
485                self.commands.push(Command::Rotate { slot, width: hd, rope: rope.clone() });
486            }
487        }
488        if rotated {
489            self.commands.push(Command::Hadamard { slot: Slot::Key, width: hd });
490            self.commands.push(Command::Hadamard { slot: Slot::Value, width: vd });
491        }
492        self.commands.push(Command::CacheAppend { layer: l, key: Slot::Key, value: Slot::Value, kv_heads: a.kv_heads, width: hd, value_width: vd });
493        if rotated {
494            self.commands.push(Command::Hadamard { slot: Slot::Query, width: hd });
495        }
496        self.commands.push(Command::Attend {
497            layer: l,
498            query: Slot::Query,
499            heads: a.heads,
500            kv_heads: a.kv_heads,
501            width: hd,
502            value_width: vd,
503            scale: a.scale as f32,
504            to: Slot::Read,
505        });
506        if a.output_gate {
507            self.commands.push(Command::SigmoidGate { slot: Slot::Read, gate: Slot::OutGate });
508        }
509        self.commands.push(Command::Matvec { tensor: attn_output, from: Slot::Read, to: Slot::Out });
510        Ok(out.out.clone())
511    }
512
513    /// The step's advance, the final norm of the residual `x` and the head; the next token's choice after it.
514    fn head(&mut self, x: &str) -> Result<(), LowerError> {
515        self.block = "the head".into();
516        self.commands.push(Command::Advance);
517        let (normed, weight) = self.norm(x)?;
518        let normed = normed.out.clone();
519        let output_norm = self.vector(&weight)?;
520        self.commands.push(Command::Norm { weight: output_norm, from: Slot::Residual, to: Slot::Normed });
521        let (logits, wout) = self.projection(&normed)?;
522        let head = self.matrix(&wout)?;
523        self.commands.push(Command::Matvec { tensor: head, from: Slot::Normed, to: Slot::Logits });
524        if self.next_is("ArgMax") {
525            let pick = self.take("ArgMax")?;
526            self.reads(pick, 0, &logits.out)?;
527        }
528        Ok(())
529    }
530}
531
532#[cfg(test)]
533#[path = "../../../../logicaffeine_tests/unit/logicaffeine_compile/spatial/decode_engine/lower_graph/tests.rs"]
534mod tests;
logicaffeine_compile/spatial/decode_engine/lower_graph.rs

logicaffeine_compile/spatial/decode_engine/
lower_graph.rs
