Skip to main content

logicaffeine_compile/codegen_cuda/lane/
plan.rs

1//! A step graph lowered to launches of the kernel library. The lane decides no fusion of its own: the step, lifted into
2//! the one IR, is optimized by the shared table under the lane's shim ([`LaneShim`]), which answers whether a fused
3//! class has a launch; each op the optimizer leaves takes the launches its link names ([`link`]), a class the lane
4//! runs whole one launch, a class it runs as its parts each part's. Then the lane's schedule sets how the launches run
5//! — four rows a warp for a wide product (`rows_a_warp`), the shared-row attention, the split counts, a batch's
6//! sequences as one launch (`sequences`) — and every tensor is given its buffers.
7//!
8//! A row read at a literal index of a computed tensor (`row 4 of bxo`) is no launch at all: a view into the tensor —
9//! unless the tensor is a table written in place before the row's last reader, where the row is copied when taken.
10//! A view a projection reads quantized is a view into its tensor's quantized form too (when it starts on a whole
11//! block), so the tensor is quantized once, in its producer, for every row a projection reads of it.
12
13use std::collections::{BTreeMap, BTreeSet};
14
15use crate::codegen_cuda::knobs::{CachePrecision, CudaKnobs, StatePrecision};
16use crate::codegen_cuda::qwen_emit::Q8MatKind;
17use crate::model::quant::GgmlType;
18use crate::spatial::passes::recognize_stages::TensorOpSpec;
19use crate::tensor_recognize::qwen::{Activation, MoeAct};
20use crate::tensor_recognize::{ElementwiseOp, Epsilon};
21
22use crate::codegen_cuda::lane_link::link;
23use crate::codegen_cuda::lane_shim::{chunks_apart, LaneShim};
24use crate::tensor_graph::fuse::{gates_rows, optimize, steps};
25use crate::tensor_graph::graph::{Storage, Sym, TensorGraph};
26use crate::tensor_graph::symbols::ClassId;
27
28pub use crate::tensor_graph::graph::{SeqIo, StepIo};
29
30use super::graph::{Role, StepGraph};
31use super::kernels;
32use super::weights::Placement;
33use super::LaneError;
34
35/// The most rows the GEMV kernel multiplies in one launch (its warp per output row holds an accumulator per row); a
36/// projection of more rows runs on the tensor cores, or (under the precise numerics) as slabs of this many rows.
37pub const GEMV_ROWS: usize = 16;
38
39/// The largest prefill chunk a lane compiles.
40pub const MAX_CHUNK: usize = 512;
41
42/// The most positions one query of an attention over a cache of `len` rows attends: the cache, or its window.
43pub fn attended_span(len: usize, window: Option<usize>) -> usize {
44    window.map_or(len, |w| w.min(len))
45}
46
47/// The values of an activation's quantized form that share one scale.
48pub const QBLOCK: usize = 32;
49
50/// A view `offset` elements into a tensor reads the tensor's quantized form in place: it starts on a whole block.
51pub fn quant_in_place(offset: usize) -> bool {
52    offset % QBLOCK == 0
53}
54
55/// Which row of a table a write starts at: the row a parameter names (1-based), or a fixed row.
56#[derive(Clone, Debug, PartialEq, Eq)]
57pub enum RowSel {
58    Param(String),
59    Fixed(usize),
60}
61
62/// A part of every unit of a wider tensor: unit `u`'s values start `u·pitch + offset` in.
63#[derive(Clone, Copy, Debug, PartialEq, Eq)]
64pub struct Part {
65    pub pitch: usize,
66    pub offset: usize,
67}
68
69/// Where a kernel writes its result: a buffer, or the rows of a table of rows `pitch` wide from a starting row — the
70/// whole of it, or (`unit`) a part of every unit of it.
71#[derive(Clone, Debug, PartialEq, Eq)]
72pub struct Dst {
73    pub buf: String,
74    pub row: Option<(RowSel, usize)>,
75    pub unit: Option<Part>,
76}
77
78impl Dst {
79    pub fn to(buf: &str) -> Dst {
80        Dst { buf: buf.to_string(), row: None, unit: None }
81    }
82
83    /// The tensor the result is the whole of, in its own buffer.
84    pub fn own(&self) -> Option<&str> {
85        (self.row.is_none() && self.unit.is_none()).then_some(self.buf.as_str())
86    }
87}
88
89/// What the arg-max that closes a step writes besides its result: the token the next decode step reads (every form's
90/// arg-max writes the decode step's; a batch's sequence `item` writes item `item` of its tokens) and the position the
91/// sequence's next step reads; and the run generated so far, which only the host reads.
92#[derive(Clone, Debug, PartialEq)]
93pub struct Advance {
94    pub token: String,
95    pub position: String,
96    pub item: usize,
97}
98
99/// The norm a norm launch takes of its added operand in the same block before adding it (`sandwich_norm`): the weight
100/// and epsilon of the norm that operand's own launch took.
101#[derive(Clone, Debug, PartialEq)]
102pub struct PreNorm {
103    pub w: Option<String>,
104    pub eps: Epsilon,
105}
106
107/// A part of a GEMM in parts: its matrix, bias, destination and outputs a row.
108#[derive(Clone, Debug, PartialEq)]
109pub struct GemmPart {
110    pub w: String,
111    pub bias: Option<String>,
112    pub y: Dst,
113    pub n: usize,
114}
115
116/// One launch of a library kernel. A producer names the tensor it computes (`out`); `f32` when it writes the tensor's
117/// `f32` buffer, `quant` when it writes its quantized form. Row counts are the block's (1 for a decode step). A `part`
118/// is where a launch reads its input `x` unit by unit: a part of every unit of a wider tensor (none, the whole of `x`).
119#[derive(Clone, Debug, PartialEq)]
120pub enum Launch {
121    Gather { table: String, index: String, out: String, cols: usize, rows: usize, format: Option<GgmlType>, scale: Option<f64> },
122    View { src: String, offset: usize, out: String },
123    /// `into`: where its `f32` result lands when not in `out`'s own buffer.
124    Norm { x: String, part: Option<Part>, add: Option<String>, pre: Option<PreNorm>, sum: Option<String>, w: Option<String>, out: String, into: Option<Dst>, f32: bool, quant: bool, rows: usize, cols: usize, eps: Epsilon },
125    Quant { x: String, n: usize },
126    /// A projection of `m` rows, `per_warp` output rows a warp; one of four rows a warp quantizes its output in its
127    /// epilogue (`quant`, keeping the `f32` result when `f32`). On the tensor cores, `tail` is how its last wave is
128    /// filled (`kernels::gemm_blk_tail`): the tiles a block each, and the parts every later tile is cut into.
129    Gemv { x: String, w: String, format: GgmlType, bias: Option<String>, y: Dst, k: usize, n: usize, m: usize, per_warp: usize, f32: bool, quant: bool, tail: Option<(usize, usize)>, batch: bool },
130    /// An `f32` projection of `m` rows; `split` is the parts it cuts `k` into: each tile's in tiles
131    /// (`Launch::multiplies_in_tiles`, `kernels::gemm_f32_splits`), each output row's over a block's warps otherwise
132    /// (`kernels::gemv_f32_parts`).
133    GemvF32 { x: String, w: String, bias: Option<String>, y: Dst, k: usize, n: usize, m: usize, split: usize, batch: bool },
134    /// Projections of one activation's `m` rows by matrices in one format, each into a buffer of its own or rows of a
135    /// table, as one launch over their outputs side by side: a tensor-core GEMM above the GEMV's rows
136    /// (`merged_projections`; `tail` as a projection's, over the parts' tiles together), the GEMV in parts at or below
137    /// them (`gemv_parts`; `per_warp` output rows a warp, as each part's own GEMV took them).
138    GemmParts { x: String, format: GgmlType, k: usize, m: usize, parts: Vec<GemmPart>, tail: Option<(usize, usize)>, per_warp: usize, batch: bool },
139    Binary { op: ElementwiseOp, a: String, b: String, gate: Option<Activation>, row: Option<usize>, out: String, f32: bool, quant: bool, n: usize },
140    Scale { x: String, y: String, n: usize, factor: f64 },
141    SwiGlu { g: String, u: String, out: String, f32: bool, quant: bool, n: usize },
142    Unary { act: Activation, x: String, y: String, n: usize },
143    Rope { x: String, part: Option<Part>, norm: Option<(String, Epsilon)>, cos: String, sin: String, pos: Option<String>, y: Dst, heads: usize, per_row: usize, hd: usize, half: usize },
144    /// Rotations of a step's one row by one table at one position — a layer's query heads and key heads, each normed
145    /// first where the model norms them — as one launch over their heads (`rope_parts`), each head computed as its own
146    /// rotation's launch computes it.
147    /// `seqs`: a batch's rotations over its sequences (`sequences`), a side's row `r` sequence `r`, its position item
148    /// `r` of `pos`, its destination row in its own table `RopeSide::yz` values on from the last row's.
149    RopeParts { sides: Vec<RopeSide>, cos: String, sin: String, pos: Option<String>, hd: usize, half: usize, seqs: bool },
150    /// `seq`: a batch's rows set over its sequences (`sequences`): row `r` of `src` into its own table, `seq` values on
151    /// from the last row's, at its own index, item `r` of the index.
152    /// `f32` and `quant`: the rows' values are set, and their quantized form too, into the table's — a batch's
153    /// sequence's row, quantized where its producer stored it (`row_write_quant`, PARITY-BATCH-ROW-QUANT), its rows whole
154    /// 32-value blocks, its values set only where something reads the table `f32`.
155    SetRow { table: String, index: RowSel, src: String, cols: usize, count: usize, seq: Option<usize>, f32: bool, quant: bool },
156    /// `vd` the value head width; the value rows `vstride` apart, the heads `vd` wide from `voff`; `scale` the scores'.
157    /// A block takes `per_block` query heads of one group together; a row's span is cut into at most `splits` chunks,
158    /// a block each when `parallel`, all in its row's block otherwise. With a `split`, each query and key is read in
159    /// two parts. `shared`: the most chunks each row's span is cut into when every head reads its one key/value group's
160    /// rows together, each cached row read once for all of them (`lk_attention_shared`). `seq`: a batch's attention over
161    /// its sequences (`sequences`), row `r` sequence `r`, its caches `seq` values on from the last row's, its position
162    /// item `r` of `pos`.
163    /// `quant`: the attention quantizes its output in its own store for the projection that reads it (keeping it `f32`
164    /// too when `f32`), as `lk_quant` would — the decode step's and a short chunk's kernels (`lk_attention`,
165    /// `lk_attention_shared` and its fold), never those in tiles.
166    /// `gate`: the attention multiplies each output by its activation of the gate's value at the same place as it stores
167    /// it, before it quantizes it (`attn_output_gate`) — the gate's values read from `goff` on: a batch's sequence's row
168    /// of the batch's gate, 0 otherwise.
169    Attention { q: String, kc: String, vc: String, pos: Option<String>, sinks: Option<String>, out: String, heads: usize, group: usize, hd: usize, vd: usize, len: usize, rows: usize, window: Option<usize>, vstride: usize, voff: usize, scale: f64, per_block: usize, splits: usize, parallel: bool, split: Option<Split>, shared: Option<usize>, seq: Option<usize>, f32: bool, quant: bool, gate: Option<(String, Activation)>, goff: usize },
170    /// `advance`: what the arg-max that closes the step moves on besides its result (`argmax_advance`).
171    /// `rows`: a batch's arg-maxes over its `rows` sequences (`sequences`), each row's logits, result and advance its
172    /// own; 1 otherwise.
173    ArgMax { x: String, out: String, n: usize, advance: Option<Advance>, rows: usize },
174    Conv { x: String, w: String, ring: String, y: String, chan: usize, taps: usize, rows: usize, silu: bool, l2: Option<(usize, Epsilon)> },
175    /// Causal convolutions of a step's one row through the SiLU, each over its own input, weights and ring, as one launch
176    /// (`conv_parts`): a block `block` channels of one part, the parts' blocks one after another; a normed part's heads
177    /// `block` wide, each normed to unit length with `eps`.
178    /// `seqs`: a batch's convolutions over its `seqs` sequences (`sequences`), each part's input row, ring and output
179    /// the sequence's part of one array of them all; 1 otherwise.
180    ConvParts { parts: Vec<ConvPart>, taps: usize, block: usize, eps: Option<Epsilon>, seqs: usize },
181    L2Norm { x: String, y: String, rows: usize, cols: usize, eps: Epsilon },
182    Decay { ar: String, dt: String, al: String, y: String, n: usize, heads: usize },
183    Gates { br: String, ar: String, dt: String, al: String, beta: String, decay: String, n: usize, heads: usize },
184    /// A step's two Gated DeltaNet gates from its one activation row (`gate_projections`): each head's write strength
185    /// `sigmoid(x·wb)` and forget gate `(1 + e^(x·wa + dt))^al`, the head's two `f32` dots summed as its own projections
186    /// sum them (`parts` warps a head, `kernels::gemv_f32_parts`) and its gates as `Gates` computes them, in one launch.
187    GateProj { x: String, wb: String, wa: String, dt: String, al: String, beta: String, decay: String, k: usize, heads: usize, parts: usize },
188    /// `seqs`: a batch's delta rules over its `seqs` sequences, a row and a memory each (`sequences`), each input,
189    /// gate, memory and output the sequence's part of one array of them all; 1 otherwise.
190    Delta { q: String, k: String, v: String, z: String, b: String, a: String, st: String, gn: String, out: String, f32: bool, quant: bool, hv: usize, hk: usize, d: usize, rows: usize, eps: Epsilon, seqs: usize },
191    /// A mixture of `experts` experts of `dff` keeping `k`, and the `shared` experts every token takes, summed after —
192    /// its gate and up rows `per_warp` a warp — over `rows` rows of `x` routed each on its own, every expert reading its
193    /// weights once for all the rows that chose it. `router_split`: the parts an `f32` router's logits GEMM cuts each
194    /// tile's `k` into (`kernels::gemm_f32_splits`) on the tensor cores.
195    Moe { x: String, wr: String, wg: String, wu: String, wd: String, formats: [GgmlType; 3], y: Dst, d: usize, dff: usize, experts: usize, k: usize, router: Router, rb: Option<String>, biases: [Option<String>; 3], renormalize: bool, act: MoeAct, shared: Option<SharedExperts>, per_warp: usize, rows: usize, router_split: usize, batch: bool },
196    /// `y[g·width + i] = x[g·every + from + i]` over `groups` groups.
197    Slice { x: String, y: String, groups: usize, every: usize, from: usize, width: usize },
198    /// Each of the `rows · heads` units' `ha` values of `a`, then its `hb` values of `b`. `yz`: a batch's join into its
199    /// sequences' caches (`sequences`), row `r` of `rows` into its own table, `yz` values on from the last's, at its own
200    /// position; 0 otherwise.
201    Join { a: String, b: String, y: Dst, rows: usize, heads: usize, ha: usize, hb: usize, yz: usize },
202    /// The per-head matmul over `m` rows and `heads` heads (`a` inputs, `b` outputs a head): `rows_out` when the matrix's
203    /// rows are the heads' outputs, its rows its inputs otherwise. `shared`: every head reads its row's one unit of `x`
204    /// (the `part` of each row). `upto`: the rows multiplied are those before the position the parameter names plus
205    /// the count given, at most `m`.
206    HeadGemv { x: String, part: Option<Part>, w: String, format: GgmlType, y: Dst, m: usize, heads: usize, a: usize, b: usize, rows_out: bool, shared: bool, upto: Option<(String, usize)>, batch: bool },
207    /// Two projections of `x` through a gate: `out = act(x·wg + bg)·(x·wu + bu)`, SwiGLU when `act` is none — each
208    /// matrix `n` rows of `k`, `m` activation rows, `per_warp` output rows a warp. On the tensor cores, one GEMM of the
209    /// gate's and up's tiles side by side (`lk_gemm_blk`'s `G` 5), `tail` the gate's own GEMM's over twice its tiles.
210    GemvGlu { x: String, wg: String, wu: String, formats: [GgmlType; 2], biases: [Option<String>; 2], act: Option<Activation>, out: String, f32: bool, quant: bool, k: usize, n: usize, m: usize, per_warp: usize, tail: Option<(usize, usize)>, batch: bool },
211    /// One launch running `first`, then `then` over what `first` wrote, in the blocks that finish what `then` reads: a
212    /// projection with the residual add and the next norm in its epilogue (`gemv_norm_epilogue`), run by its last block;
213    /// projections of one activation in parts with their heads' norms and rotations (`qkv_head_epilogue`), each head
214    /// turned by the block that writes its last rows; the embedding's gather with the first norm (`embed_norm`), each
215    /// row normed by the block that gathers it; or a mixture — alone, or folded with its sum with another array — with
216    /// the residual add and the next norm (`moe_deferred_finalize`), run by the decode step's last down block or by a
217    /// chunk's combine, a block a row. A projection's blocks count in as they finish, and the finishing block reads their
218    /// results back through L2; every block runs `then`'s work as `then`'s own launch would, so each computes what it
219    /// computes apart.
220    Fold { first: Box<Launch>, then: Box<Launch> },
221}
222
223/// One rotation of a merged launch (`Launch::RopeParts`): its input (read in place as a part of every unit of a wider
224/// tensor where it names one), its norm, its destination, and its heads, `per_row` of them a row.
225#[derive(Clone, Debug, PartialEq)]
226pub struct RopeSide {
227    pub x: String,
228    pub part: Option<Part>,
229    pub norm: Option<(String, Epsilon)>,
230    pub y: Dst,
231    pub heads: usize,
232    pub per_row: usize,
233    /// A batch's side (`RopeParts::seqs`): the values between one sequence's table and the next's; 0 otherwise.
234    pub yz: usize,
235}
236
237/// One causal convolution of a merged launch (`Launch::ConvParts`): its input, weights, ring and output, its channels,
238/// and whether each head of its output is normed to unit length.
239#[derive(Clone, Debug, PartialEq)]
240pub struct ConvPart {
241    pub x: String,
242    pub w: String,
243    pub ring: String,
244    pub y: String,
245    pub chan: usize,
246    pub l2: bool,
247}
248
249/// An attention's queries and keys in two parts, split at `own`: a query's first `own` values its unit's in `q`
250/// (`q_pitch` apart a (row, head)), the rest its unit's in `q_tail`; a key's first `own` its group's in the keys, the
251/// rest its row's in `k_tail`, every group's alike — an expanded latent attention's queries, each head's unrotated part
252/// read in place in the query projection and its rotated part, and its keys, each head's expansion and the rotated key
253/// the heads share, read in place in the cache.
254#[derive(Clone, Debug, PartialEq, Eq)]
255pub struct Split {
256    pub own: usize,
257    pub q_pitch: usize,
258    pub q_tail: (String, Part),
259    pub k_tail: (String, Part),
260}
261
262/// How an expert mixture's router weights are held: quantized checkpoint rows of a format, or `f32` values (an
263/// expert's row of the input's width at a time) — the router reads the input `f32` then, under every numerics.
264#[derive(Clone, Copy, Debug, PartialEq, Eq)]
265pub enum Router {
266    Rows(GgmlType),
267    F32,
268}
269
270/// The shared experts a mixture takes into its own launches (`shared_expert`): one SwiGLU FFN of `dff`, its gate, up and
271/// down projections read in place in their `formats`.
272#[derive(Clone, Debug, PartialEq, Eq)]
273pub struct SharedExperts {
274    pub wg: String,
275    pub wu: String,
276    pub wd: String,
277    pub formats: [GgmlType; 3],
278    pub dff: usize,
279}
280
281/// How a launch touches a tensor: through its `f32` (or `i32`) buffer, or its quantized activation.
282#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
283pub enum Touch {
284    Plain,
285    Quant,
286}
287
288fn strs<'a>(v: impl IntoIterator<Item = &'a String>) -> Vec<&'a str> {
289    v.into_iter().map(String::as_str).collect()
290}
291
292impl Launch {
293    /// The tensors the launch reads (a parameter naming a row included).
294    pub fn reads(&self) -> Vec<&str> {
295        fn row_param(d: &Dst) -> Vec<&str> {
296            match &d.row {
297                Some((RowSel::Param(p), _)) => vec![p.as_str()],
298                _ => vec![],
299            }
300        }
301        let mut v: Vec<&str> = match self {
302            Launch::Gather { table, index, .. } => vec![table, index],
303            Launch::View { src, .. } => vec![src],
304            Launch::Norm { x, add, pre, w, into, .. } => {
305                let mut r = strs([Some(x), add.as_ref(), w.as_ref(), pre.as_ref().and_then(|p| p.w.as_ref())].into_iter().flatten());
306                r.extend(into.iter().flat_map(row_param));
307                r
308            }
309            Launch::Quant { x, .. } => vec![x],
310            Launch::Gemv { x, w, bias, y, .. } | Launch::GemvF32 { x, w, bias, y, .. } => {
311                let mut r = vec![x.as_str(), w.as_str()];
312                r.extend(bias.as_deref());
313                r.extend(row_param(y));
314                r
315            }
316            Launch::GemmParts { x, parts, .. } => {
317                let mut r = vec![x.as_str()];
318                r.extend(parts.iter().flat_map(|p| std::iter::once(p.w.as_str()).chain(p.bias.as_deref()).chain(row_param(&p.y))));
319                r
320            }
321            Launch::Binary { a, b, .. } => vec![a, b],
322            Launch::Scale { x, .. } => vec![x],
323            Launch::SwiGlu { g, u, .. } => vec![g, u],
324            Launch::Unary { x, .. } => vec![x],
325            Launch::Rope { x, norm, cos, sin, pos, y, .. } => {
326                let mut r = strs([Some(x), norm.as_ref().map(|(w, _)| w), Some(cos), Some(sin), pos.as_ref()].into_iter().flatten());
327                r.extend(row_param(y));
328                r
329            }
330            Launch::RopeParts { sides, cos, sin, pos, .. } => {
331                let mut r = Vec::new();
332                for side in sides {
333                    r.extend(strs([Some(&side.x), side.norm.as_ref().map(|(w, _)| w)].into_iter().flatten()));
334                    r.extend(row_param(&side.y));
335                }
336                r.extend(strs([Some(cos), Some(sin), pos.as_ref()].into_iter().flatten()));
337                r
338            }
339            Launch::SetRow { table, index, src, .. } => match index {
340                RowSel::Param(p) => vec![table, p, src],
341                RowSel::Fixed(_) => vec![table, src],
342            },
343            Launch::Attention { q, kc, vc, pos, sinks, split, gate, .. } => {
344                let mut r = strs([Some(q), Some(kc), Some(vc), pos.as_ref(), sinks.as_ref(), gate.as_ref().map(|(g, _)| g)].into_iter().flatten());
345                r.extend(split.iter().flat_map(|s| [s.q_tail.0.as_str(), s.k_tail.0.as_str()]));
346                r
347            }
348            Launch::ArgMax { x, .. } => vec![x],
349            Launch::Conv { x, w, ring, .. } => vec![x, w, ring],
350            Launch::ConvParts { parts, .. } => parts.iter().flat_map(|p| [p.x.as_str(), p.w.as_str(), p.ring.as_str()]).collect(),
351            Launch::L2Norm { x, .. } => vec![x],
352            Launch::Decay { ar, dt, al, .. } => vec![ar, dt, al],
353            Launch::Gates { br, ar, dt, al, .. } => vec![br, ar, dt, al],
354            Launch::GateProj { x, wb, wa, dt, al, .. } => vec![x, wb, wa, dt, al],
355            Launch::Delta { q, k, v, z, b, a, st, gn, .. } => vec![q, k, v, z, b, a, st, gn],
356            Launch::Moe { x, wr, wg, wu, wd, y, rb, biases, shared, .. } => {
357                let mut r = vec![x.as_str(), wr.as_str(), wg.as_str(), wu.as_str(), wd.as_str()];
358                r.extend(shared.iter().flat_map(|e| [e.wg.as_str(), e.wu.as_str(), e.wd.as_str()]));
359                r.extend(rb.as_deref());
360                r.extend(biases.iter().flatten().map(String::as_str));
361                r.extend(row_param(y));
362                r
363            }
364            Launch::Slice { x, .. } => vec![x],
365            Launch::Join { a, b, y, .. } => {
366                let mut r = vec![a.as_str(), b.as_str()];
367                r.extend(row_param(y));
368                r
369            }
370            Launch::HeadGemv { x, w, y, upto, .. } => {
371                let mut r = vec![x.as_str(), w.as_str()];
372                r.extend(row_param(y));
373                r.extend(upto.as_ref().map(|(p, _)| p.as_str()));
374                r
375            }
376            Launch::GemvGlu { x, wg, wu, biases, .. } => {
377                let mut r = vec![x.as_str(), wg.as_str(), wu.as_str()];
378                r.extend(biases.iter().flatten().map(String::as_str));
379                r
380            }
381            // What `then` reads of what `first` wrote is handed over within the launch, read by none of its own.
382            Launch::Fold { first, then } => {
383                let made = first.writes();
384                first.reads().into_iter().chain(then.reads().into_iter().filter(|t| !made.contains(t))).collect()
385            }
386        };
387        v.dedup();
388        v
389    }
390
391    /// The arrays the launch's kernels take in any storage type, reading each value widened exactly and writing it
392    /// rounded to the nearest: an attention's keys, values and split keys' rest, a delta rule's memory, a slice's source,
393    /// a tiled per-head matmul's activation, and every destination a producer writes into a table's rows (`lk_out`).
394    /// A view touches nothing itself.
395    pub fn typed(&self, numerics: Q8MatKind) -> Vec<&str> {
396        match self {
397            Launch::View { .. } => self.reads().into_iter().chain(self.writes()).collect(),
398            Launch::Attention { kc, vc, split, .. } => {
399                let mut t = vec![kc.as_str(), vc.as_str()];
400                t.extend(split.iter().map(|s| s.k_tail.0.as_str()));
401                t
402            }
403            Launch::Delta { st, .. } => vec![st],
404            Launch::Slice { x, .. } => vec![x],
405            Launch::HeadGemv { x, y, .. } if self.heads_in_tiles(numerics) => vec![x, y.buf.as_str()],
406            Launch::HeadGemv { y, .. } | Launch::Gemv { y, .. } | Launch::GemvF32 { y, .. } | Launch::Rope { y, .. } | Launch::Join { y, .. } => {
407                vec![y.buf.as_str()]
408            }
409            Launch::GemmParts { parts, .. } => parts.iter().map(|p| p.y.buf.as_str()).collect(),
410            Launch::RopeParts { sides, .. } => sides.iter().map(|s| s.y.buf.as_str()).collect(),
411            Launch::Norm { into: Some(d), .. } => vec![d.buf.as_str()],
412            Launch::SetRow { table, .. } => vec![table],
413            Launch::Fold { first, then } => first.typed(numerics).into_iter().chain(then.typed(numerics)).collect(),
414            _ => Vec::new(),
415        }
416    }
417
418    /// The tensors the launch writes (a cache it appends to, a memory it updates, included).
419    pub fn writes(&self) -> Vec<&str> {
420        match self {
421            Launch::Gather { out, .. }
422            | Launch::View { out, .. }
423            | Launch::Binary { out, .. }
424            | Launch::SwiGlu { out, .. }
425            | Launch::Attention { out, .. }
426            | Launch::GemvGlu { out, .. } => vec![out],
427            Launch::ArgMax { out, advance, .. } => {
428                let mut w = vec![out.as_str()];
429                w.extend(advance.iter().flat_map(|a| [a.token.as_str(), a.position.as_str()]));
430                w
431            }
432            Launch::Slice { y, .. } => vec![y],
433            Launch::Norm { out, into, sum, .. } => {
434                let mut w = vec![into.as_ref().map_or(out.as_str(), |d| d.buf.as_str())];
435                w.extend(sum.as_deref());
436                w
437            }
438            Launch::Delta { out, st, .. } => vec![out, st],
439            Launch::Quant { x, .. } => vec![x],
440            Launch::GemmParts { parts, .. } => parts.iter().map(|p| p.y.buf.as_str()).collect(),
441            Launch::RopeParts { sides, .. } => sides.iter().map(|s| s.y.buf.as_str()).collect(),
442            Launch::Gemv { y, .. }
443            | Launch::GemvF32 { y, .. }
444            | Launch::Rope { y, .. }
445            | Launch::Moe { y, .. }
446            | Launch::Join { y, .. }
447            | Launch::HeadGemv { y, .. } => vec![y.buf.as_str()],
448            Launch::Unary { y, .. } | Launch::Scale { y, .. } | Launch::L2Norm { y, .. } | Launch::Decay { y, .. } => vec![y],
449            Launch::SetRow { table, .. } => vec![table],
450            Launch::Conv { y, ring, .. } => vec![y, ring],
451            Launch::ConvParts { parts, .. } => parts.iter().flat_map(|p| [p.y.as_str(), p.ring.as_str()]).collect(),
452            Launch::Gates { beta, decay, .. } | Launch::GateProj { beta, decay, .. } => vec![beta, decay],
453            Launch::Fold { first, then } => first.writes().into_iter().chain(then.writes()).collect(),
454        }
455    }
456
457    /// The tensors the launch reads and writes, each in the form it touches: a projection reads its activation
458    /// quantized unless the numerics keep it `f32`; a producer writes its result `f32`, quantized, or both. A view
459    /// touches nothing: what its readers touch of it, they touch of its tensor.
460    pub fn touches(&self, numerics: Q8MatKind) -> (Vec<(&str, Touch)>, Vec<(&str, Touch)>) {
461        fn plain(v: Vec<&str>) -> Vec<(&str, Touch)> {
462            v.into_iter().map(|t| (t, Touch::Plain)).collect()
463        }
464        fn produced(out: &str, f32: bool, quant: bool) -> Vec<(&str, Touch)> {
465            [(f32, Touch::Plain), (quant, Touch::Quant)].into_iter().filter(|(on, _)| *on).map(|(_, f)| (out, f)).collect()
466        }
467        let reads = match (self, self.projected_input()) {
468            (Launch::View { .. }, _) => return (Vec::new(), Vec::new()),
469            (Launch::Fold { first, then }, _) => {
470                let made = first.touches(numerics).1;
471                first.touches(numerics).0.into_iter().chain(then.touches(numerics).0.into_iter().filter(|t| !made.contains(t))).collect()
472            }
473            (_, Some(x)) if numerics != Q8MatKind::Dequant => {
474                let mut r = plain(self.reads());
475                if !self.reads_f32_input() {
476                    r.retain(|(t, _)| *t != x);
477                }
478                r.push((x, Touch::Quant));
479                r
480            }
481            (Launch::SetRow { table, src, f32, quant, .. }, _) => {
482                let mut r: Vec<(&str, Touch)> = plain(self.reads()).into_iter().filter(|(t, _)| *f32 || (*t != table.as_str() && *t != src.as_str())).collect();
483                if *quant {
484                    r.push((src.as_str(), Touch::Quant));
485                }
486                r
487            }
488            _ => plain(self.reads()),
489        };
490        let writes = match self {
491            Launch::Norm { out, into, f32, quant, sum, .. } => {
492                let mut w = match into {
493                    Some(d) => [(*f32, d.buf.as_str(), Touch::Plain), (*quant, out.as_str(), Touch::Quant)]
494                        .into_iter()
495                        .filter(|(on, ..)| *on)
496                        .map(|(_, t, f)| (t, f))
497                        .collect(),
498                    None => produced(out, *f32, *quant),
499                };
500                w.extend(sum.iter().map(|s| (s.as_str(), Touch::Plain)));
501                w
502            }
503            Launch::Binary { out, f32, quant, .. }
504            | Launch::SwiGlu { out, f32, quant, .. }
505            | Launch::GemvGlu { out, f32, quant, .. }
506            | Launch::Attention { out, f32, quant, .. } => produced(out, *f32, *quant),
507            Launch::Delta { out, f32, quant, st, .. } => {
508                let mut w = produced(out, *f32, *quant);
509                w.push((st.as_str(), Touch::Plain));
510                w
511            }
512            Launch::Quant { x, .. } => vec![(x.as_str(), Touch::Quant)],
513            Launch::SetRow { table, f32, quant, .. } => produced(table, *f32, *quant),
514            Launch::Fold { first, then } => first.touches(numerics).1.into_iter().chain(then.touches(numerics).1).collect(),
515            _ => plain(self.writes()),
516        };
517        (reads, writes)
518    }
519
520    /// Whether the launch multiplies on the int8 tensor cores under `numerics`: a quantized projection of more rows than
521    /// the GEMV takes, its activation quantized, by `lk_gemm_blk` in any format, Q8_0's split arrays included — several
522    /// side by side in parts, a gated pair as one paired GEMM — and an expert mixture over as many rows, each expert's
523    /// matrices by the grouped `lk_gemm_blk`.
524    pub fn on_tensor_cores(&self, numerics: Q8MatKind) -> bool {
525        matches!(self, Launch::Gemv { m, batch: false, .. } | Launch::GemmParts { m, batch: false, .. } | Launch::GemvGlu { m, batch: false, .. }
526            | Launch::Moe { rows: m, batch: false, .. } if *m > GEMV_ROWS)
527            && numerics != Q8MatKind::Dequant
528    }
529
530    /// Whether a batch's projection runs on the tensor cores in its rows' own order (`lk_gemv_mma`): a projection, its
531    /// parts or its gated pair over Q8_0 matrices of a batch's `GEMV_MMA_MIN_ROWS` or more rows, under the quantized
532    /// numerics, every part a whole number of 8-output tiles and `k` a whole number of 32-value blocks.
533    pub fn batch_on_tensor_cores(&self, numerics: Q8MatKind) -> bool {
534        let rows = |m: usize| m >= kernels::GEMV_MMA_MIN_ROWS;
535        numerics != Q8MatKind::Dequant
536            && match self {
537                Launch::Gemv { format: GgmlType::Q8_0, k, n, m, f32: true, quant: false, batch: true, .. } => rows(*m) && n % 8 == 0 && k % 32 == 0,
538                Launch::GemmParts { format: GgmlType::Q8_0, k, m, parts, batch: true, .. } => rows(*m) && k % 32 == 0 && parts.iter().all(|p| p.n % 8 == 0),
539                Launch::GemvGlu { formats: [GgmlType::Q8_0, GgmlType::Q8_0], k, n, m, batch: true, .. } => rows(*m) && n % 32 == 0 && k % 32 == 0,
540                _ => false,
541            }
542    }
543
544    /// Whether the launch is a batch's projection (`batch`): its rows are the batch's sequences, each summed as the
545    /// decode step sums its one row — row by row whatever their count, never in tiles.
546    pub fn is_batch(&self) -> bool {
547        matches!(self, Launch::Gemv { batch: true, .. } | Launch::GemvF32 { batch: true, .. } | Launch::GemmParts { batch: true, .. }
548            | Launch::GemvGlu { batch: true, .. } | Launch::Moe { batch: true, .. } | Launch::HeadGemv { batch: true, .. })
549    }
550
551    /// A projection's outputs a row: a GEMM in parts' over its parts together; none for any other launch.
552    pub fn gemm_cols(&self) -> usize {
553        match self {
554            Launch::Gemv { n, .. } => *n,
555            Launch::GemmParts { parts, .. } => parts.iter().map(|p| p.n).sum(),
556            _ => 0,
557        }
558    }
559
560    /// The launches this launch runs the work of: a fold's first's and then's, any other launch itself — what a question
561    /// about one kernel's reads or kind asks of each.
562    pub fn parts(&self) -> Vec<&Launch> {
563        match self {
564            Launch::Fold { first, then } => first.parts().into_iter().chain(then.parts()).collect(),
565            l => vec![l],
566        }
567    }
568
569    /// The mixture a fold finishes with the next norm, for the schedule to set how it runs: the launch itself where it is
570    /// a mixture.
571    pub fn mixture_mut(&mut self) -> Option<&mut Launch> {
572        if matches!(self, Launch::Moe { .. }) {
573            return Some(self);
574        }
575        let Launch::Fold { first, then } = self else { return None };
576        if !matches!(then.as_ref(), Launch::Norm { .. }) {
577            return None;
578        }
579        match first.as_mut() {
580            m @ Launch::Moe { .. } => Some(m),
581            Launch::Fold { first: m, then: sum } if matches!((m.as_ref(), sum.as_ref()), (Launch::Moe { .. }, Launch::Binary { .. })) => Some(m.as_mut()),
582            _ => None,
583        }
584    }
585
586    /// The mixture a fold finishes with the next norm, its norm, and — where the norm adds the mixture's sum with another
587    /// array — that array and the sum.
588    pub fn finished_mixture(&self) -> Option<(&Launch, &Launch, Option<(&str, &str)>)> {
589        let Launch::Fold { first, then } = self else { return None };
590        if !matches!(then.as_ref(), Launch::Norm { .. }) {
591            return None;
592        }
593        match first.as_ref() {
594            m @ Launch::Moe { .. } => Some((m, then, None)),
595            Launch::Fold { first: m, then: sum } => match (m.as_ref(), sum.as_ref()) {
596                (Launch::Moe { y, .. }, Launch::Binary { a, b, out, .. }) => {
597                    Some((m.as_ref(), then, Some((if *a == y.buf { b.as_str() } else { a.as_str() }, out.as_str()))))
598                }
599                _ => None,
600            },
601            _ => None,
602        }
603    }
604
605    /// The activation rows a block of a dense tensor-core GEMM takes under `numerics` (`kernels::gemm_blk_rows`).
606    pub fn gemm_rows(&self, numerics: Q8MatKind) -> usize {
607        let format = match self {
608            Launch::Gemv { format, .. } | Launch::GemmParts { format, .. } | Launch::GemvGlu { formats: [format, _], .. } => *format,
609            _ => return kernels::GEMM_BLK_TILE_M,
610        };
611        kernels::gemm_blk_rows(kernels::format_arg(format), numerics == Q8MatKind::TwoPassInt8)
612    }
613
614    /// The tiles a tensor-core GEMM's outputs a row take: of `GEMM_TILE_N` outputs, a gated pair's of half as many
615    /// hidden values (its gate's and up's side by side).
616    pub fn gemm_tiles_wide(&self) -> usize {
617        match self {
618            Launch::GemvGlu { n, .. } => n.div_ceil(kernels::GEMM_TILE_N / 2),
619            _ => self.gemm_cols().div_ceil(kernels::GEMM_TILE_N),
620        }
621    }
622
623    /// Whether an `f32` projection multiplies its rows as one `f32` GEMM of 64 × 64 tiles under `numerics`
624    /// (`lk_gemm_f32`): a chunk of more rows than the GEMV takes under the quantized numerics — the precise numerics sum
625    /// each row's dots in `f64` — its rows whole `float4`s, its result its own buffer.
626    pub fn multiplies_in_tiles(&self, numerics: Q8MatKind) -> bool {
627        numerics != Q8MatKind::Dequant
628            && matches!(self, Launch::GemvF32 { k, y, m, batch: false, .. } if *m > GEMV_ROWS && k % 4 == 0 && y.own().is_some())
629    }
630
631    /// Whether an expert mixture on the tensor cores multiplies each expert's gate and up together (`lk_gemm_blk`'s
632    /// paired mode, the gated activation and its quantizer in the GEMM's epilogue): its gate and up in one format.
633    pub fn pairs_gate_and_up(&self, numerics: Q8MatKind) -> bool {
634        matches!(self, Launch::Moe { formats, .. } if formats[0] == formats[1]) && self.on_tensor_cores(numerics)
635    }
636
637    /// Whether the launch walks its delta rule a chunk of rows at a time under `numerics` (`lk_delta_prep`, then
638    /// `lk_delta_state`): a delta rule over a chunk of more rows than the GEMV takes, under the quantized numerics, its
639    /// heads a width the chunked rule takes (`kernels::delta_chunks_take`).
640    pub fn delta_in_chunks(&self, numerics: Q8MatKind) -> bool {
641        numerics != Q8MatKind::Dequant && matches!(self, Launch::Delta { rows, d, .. } if *rows > GEMV_ROWS && kernels::delta_chunks_take(*d))
642    }
643
644    /// Whether the launch attends on the f16 tensor cores under `numerics` (`lk_attention_mma`): it attends in tiles
645    /// under the one-word numerics — llama.cpp's precision, whose attention rounds its operands to f16 — its heads and
646    /// values whole multiples of 16 and at most 256 wide.
647    pub fn attends_on_tensor_cores(&self, numerics: Q8MatKind) -> bool {
648        numerics == Q8MatKind::Int8Act
649            && self.attends_in_tiles(numerics)
650            && matches!(self, Launch::Attention { hd, vd, .. } if hd % 16 == 0 && vd % 16 == 0 && *hd <= 256 && *vd <= 256)
651    }
652
653    /// Whether the launch attends on the int8 tensor cores in two words under `numerics` (`lk_attention_i8`, its keys and
654    /// values taken into two words once by `lk_kv_two_words`): it attends in tiles under the two-word numerics, its heads
655    /// and values whole multiples of 32 and at most 256 wide.
656    pub fn attends_in_two_words(&self, numerics: Q8MatKind) -> bool {
657        numerics == Q8MatKind::TwoPassInt8
658            && self.attends_in_tiles(numerics)
659            && matches!(self, Launch::Attention { hd, vd, .. } if hd % 32 == 0 && vd % 32 == 0 && *hd <= 256 && *vd <= 256)
660    }
661
662    /// Whether the launch attends in tiles in flash style under `numerics` (`lk_attention_flash`): it attends in tiles
663    /// and on neither the f16 nor the int8 tensor cores, its heads and values at most 256 wide; wider ones take
664    /// `lk_attention_tiled`.
665    pub fn attends_flash(&self, numerics: Q8MatKind) -> bool {
666        self.attends_in_tiles(numerics)
667            && !self.attends_on_tensor_cores(numerics)
668            && !self.attends_in_two_words(numerics)
669            && matches!(self, Launch::Attention { hd, vd, .. } if *hd <= 256 && *vd <= 256)
670    }
671
672    /// Whether the launch attends in tiles under `numerics` (`lk_attention_flash`, or `lk_attention_tiled` for wider
673    /// heads): an attention over a chunk of more rows than the GEMV takes, under the quantized numerics, its keys,
674    /// values and outputs read and written four at a time and its values at most 512 wide.
675    pub fn attends_in_tiles(&self, numerics: Q8MatKind) -> bool {
676        let fits = |rows: usize, hd: usize, vd: usize, vstride: usize, voff: usize, kvw: usize| {
677            rows > GEMV_ROWS && hd % 4 == 0 && vd % 4 == 0 && vd <= 512 && vstride % 4 == 0 && voff % 4 == 0 && kvw % 4 == 0
678        };
679        // A split query and key: each part four-value words.
680        let split_fits = |s: &Option<Split>| {
681            s.as_ref().map_or(true, |s| {
682                [s.own, s.q_pitch, s.q_tail.1.pitch, s.q_tail.1.offset, s.k_tail.1.pitch, s.k_tail.1.offset].iter().all(|v| v % 4 == 0)
683            })
684        };
685        matches!(self, Launch::Attention { rows, hd, vd, vstride, voff, heads, group, split, seq: None, .. }
686            if fits(*rows, *hd, *vd, *vstride, *voff, heads / group * split.as_ref().map_or(*hd, |s| s.own)) && split_fits(split))
687            && numerics != Q8MatKind::Dequant
688    }
689
690    /// Whether the launch is a per-head matmul run in tiles under `numerics` (`lk_head_tiled`): over a chunk of more rows
691    /// than the GEMV takes, under the quantized numerics, its inputs read four words at a time.
692    pub fn heads_in_tiles(&self, numerics: Q8MatKind) -> bool {
693        let fits = |m: usize, a: usize, pitch: usize, offset: usize| m > GEMV_ROWS && a % 4 == 0 && pitch % 4 == 0 && offset % 4 == 0;
694        matches!(self, Launch::HeadGemv { m, a, part, batch: false, .. } if fits(*m, *a, part.map_or(*a, |p| p.pitch), part.map_or(0, |p| p.offset)))
695            && numerics != Q8MatKind::Dequant
696    }
697
698    /// Whether the launch writes all `len` elements of tensor `t`'s array `touch` — a result in its own buffer, every
699    /// element of it — and not rows of a table, a part of every unit of a wider tensor, or carried state it updates.
700    pub fn writes_whole(&self, t: &str, touch: Touch, len: usize) -> bool {
701        let own = |d: &Dst| d.own() == Some(t);
702        let count = match self {
703            Launch::View { .. } | Launch::SetRow { .. } => return false,
704            Launch::Gather { rows, cols, .. } | Launch::L2Norm { rows, cols, .. } => rows * cols,
705            Launch::Norm { out, into, sum, rows, cols, .. } => {
706                let whole = match touch {
707                    Touch::Quant => out == t,
708                    Touch::Plain => sum.as_deref() == Some(t) || into.as_ref().map_or(out == t, own),
709                };
710                if !whole {
711                    return false;
712                }
713                rows * cols
714            }
715            Launch::Quant { n, .. }
716            | Launch::Binary { n, .. }
717            | Launch::Scale { n, .. }
718            | Launch::SwiGlu { n, .. }
719            | Launch::Unary { n, .. }
720            | Launch::Decay { n, .. }
721            | Launch::Gates { n, .. } => *n,
722            Launch::GateProj { heads, .. } => *heads,
723            Launch::ArgMax { .. } => 1,
724            Launch::Gemv { y, n, m, .. } | Launch::GemvF32 { y, n, m, .. } if own(y) => n * m,
725            Launch::GemmParts { parts, m, .. } => match parts.iter().find(|p| own(&p.y)) {
726                Some(p) => p.n * m,
727                None => return false,
728            },
729            Launch::GemvGlu { n, m, .. } => n * m,
730            Launch::Rope { y, heads, hd, .. } if own(y) => heads * hd,
731            Launch::RopeParts { sides, hd, .. } => match sides.iter().find(|s| own(&s.y)) {
732                Some(s) => s.heads * hd,
733                None => return false,
734            },
735            Launch::Attention { rows, heads, vd, .. } => rows * heads * vd,
736            Launch::Conv { ring, rows, chan, .. } if ring != t => rows * chan,
737            Launch::ConvParts { parts, seqs, .. } => match parts.iter().find(|p| p.y == t) {
738                Some(p) => p.chan * seqs,
739                None => return false,
740            },
741            Launch::Delta { st, rows, hv, d, seqs, .. } if st != t => rows * hv * d * seqs,
742            Launch::Moe { y, rows, d, .. } if own(y) => rows * d,
743            Launch::Slice { groups, width, .. } => groups * width,
744            Launch::Join { y, rows, heads, ha, hb, .. } if own(y) => rows * heads * (ha + hb),
745            Launch::HeadGemv { y, m, heads, b, .. } if own(y) => m * heads * b,
746            Launch::Fold { first, then } => return first.writes_whole(t, touch, len) || then.writes_whole(t, touch, len),
747            _ => return false,
748        };
749        count == len
750    }
751
752    /// The part of every unit of tensor `t` (`len` elements) the launch writes, `(pitch, from, width)`: a result written
753    /// in place as a part of every unit of a wider tensor (`join_in_place`), all of its units.
754    pub fn writes_part(&self, t: &str, len: usize) -> Option<(usize, usize, usize)> {
755        let d = match self {
756            Launch::Norm { into: Some(d), .. } => d,
757            Launch::Rope { y, .. } | Launch::HeadGemv { y, .. } => y,
758            _ => return None,
759        };
760        let (units, width) = self.unit_output()?;
761        match (&d.row, &d.unit) {
762            (None, Some(p)) if d.buf == t && units * p.pitch == len && p.offset + width <= p.pitch => Some((p.pitch, p.offset, width)),
763            _ => None,
764        }
765    }
766
767    /// The activation a product by checkpoint rows reads quantized, under every numerics but the precise one. A
768    /// per-head matmul — a latent's expansion half — reads its activation `f32` under every numerics
769    /// ([`kernels::head_mode`]).
770    pub fn projected_input(&self) -> Option<&str> {
771        match self {
772            Launch::Gemv { x, .. } | Launch::GemmParts { x, .. } | Launch::Moe { x, .. } | Launch::GemvGlu { x, .. } => Some(x),
773            Launch::Fold { first, .. } => first.projected_input(),
774            _ => None,
775        }
776    }
777
778    /// Whether a projection that reads its input quantized reads it `f32` as well: an expert mixture's `f32` router
779    /// does.
780    fn reads_f32_input(&self) -> bool {
781        match self {
782            Launch::Fold { first, .. } => first.reads_f32_input(),
783            l => matches!(l, Launch::Moe { router: Router::F32, .. }),
784        }
785    }
786
787    /// The tensor a producer computes into a buffer of its own — `None` for a launch whose result lands in rows of a
788    /// table, in a part of another tensor, or in carried state alone.
789    fn result(&self) -> Option<&str> {
790        match self {
791            Launch::Gemv { y, .. }
792            | Launch::GemvF32 { y, .. }
793            | Launch::Rope { y, .. }
794            | Launch::Moe { y, .. }
795            | Launch::Join { y, .. }
796            | Launch::HeadGemv { y, .. } => y.own(),
797            Launch::Norm { out, into, .. } => into.is_none().then_some(out.as_str()),
798            Launch::SetRow { .. } | Launch::Quant { .. } | Launch::Gates { .. } | Launch::GateProj { .. } => None,
799            Launch::Fold { then, .. } => then.result(),
800            _ => self.writes().first().copied(),
801        }
802    }
803
804    /// The units a launch writes its result in, `(units, width)`, when it can write them as a part of every unit of a
805    /// wider tensor (`join_in_place`).
806    fn unit_output(&self) -> Option<(usize, usize)> {
807        match self {
808            Launch::Norm { rows, cols, f32: true, quant: false, .. } => Some((*rows, *cols)),
809            Launch::Rope { heads, hd, .. } => Some((*heads, *hd)),
810            Launch::HeadGemv { m, heads, b, .. } => Some((m * heads, *b)),
811            _ => None,
812        }
813    }
814}
815
816/// A fusion rule that fired: its name and the graph ops it joined.
817#[derive(Clone, Debug, PartialEq, Eq)]
818pub struct Fusion {
819    pub rule: &'static str,
820    pub ops: Vec<usize>,
821}
822
823/// The element type of a buffer.
824#[derive(Clone, Copy, Debug, PartialEq, Eq)]
825pub enum DType {
826    F32,
827    I32,
828    /// A recurrent memory stored in bfloat16 between steps (`CudaKnobs::recurrent_state`).
829    Bf16,
830    /// An attention's cache stored in binary16 (`CudaKnobs::cache_precision`).
831    F16,
832}
833
834impl DType {
835    /// The CUDA element type a buffer of it is declared with.
836    pub fn c_type(self) -> &'static str {
837        match self {
838            DType::F32 => "float",
839            DType::I32 => "int",
840            DType::Bf16 => "__nv_bfloat16",
841            DType::F16 => "__half",
842        }
843    }
844
845    /// The bytes of one element.
846    pub fn bytes(self) -> usize {
847        match self {
848            DType::F32 | DType::I32 => 4,
849            DType::Bf16 | DType::F16 => 2,
850        }
851    }
852}
853
854/// What a buffer holds and how it is filled before the first step.
855#[derive(Clone, Debug, PartialEq)]
856pub enum Fill {
857    /// A run-time input the generation loop sets (the tokens, the position).
858    Param,
859    /// Carried state, or a table a step fills row by row: zero before the first step.
860    Zero,
861    /// A constant, from the checkpoint or the compiler.
862    Constant(Placement),
863    /// Computed by the step.
864    Computed,
865    /// A row of another tensor, read in place: `offset` elements into `of` — its quantized form too, `offset` into
866    /// `of`'s, when the offset is whole blocks ([`quant_in_place`]).
867    Alias { of: String, offset: usize },
868}
869
870/// One tensor's buffers.
871#[derive(Clone, Debug, PartialEq)]
872pub struct Buffer {
873    pub name: String,
874    pub len: usize,
875    pub dtype: DType,
876    pub fill: Fill,
877    /// The tensor's `f32` (or `i32`) buffer is allocated.
878    pub plain: bool,
879    /// Its quantized form is allocated: int8 values (two words in the two-pass numerics) and a scale per 32.
880    pub quantized: bool,
881}
882
883/// A step lowered: its launches in order (with the graph ops each computes — none for a quantize), its buffers, and
884/// the fusions that made the launches.
885#[derive(Clone, Debug, PartialEq)]
886pub struct Schedule {
887    pub launches: Vec<Launch>,
888    pub ops: Vec<Vec<usize>>,
889    pub buffers: BTreeMap<String, Buffer>,
890    pub fusions: Vec<Fusion>,
891    pub numerics: Q8MatKind,
892    pub io: StepIo,
893}
894
895impl Schedule {
896    /// How many launches of each kind the step makes (a view is none).
897    pub fn census(&self) -> BTreeMap<&'static str, usize> {
898        let mut c = BTreeMap::new();
899        for l in self.launches.iter().filter(|l| !matches!(l, Launch::View { .. })) {
900            *c.entry(kind(l)).or_insert(0) += 1;
901        }
902        c
903    }
904
905    /// How many times `rule` fired.
906    pub fn fired(&self, rule: &str) -> usize {
907        self.fusions.iter().filter(|f| f.rule == rule).count()
908    }
909}
910
911/// A launch's kind, as a census counts it.
912pub fn kind(l: &Launch) -> &'static str {
913    match l {
914        Launch::Gather { .. } => "gather",
915        Launch::View { .. } => "view",
916        Launch::Norm { .. } => "rmsnorm",
917        Launch::Quant { .. } => "quant",
918        Launch::Gemv { .. } => "gemv_q8",
919        Launch::GemvF32 { .. } => "gemv_f32",
920        Launch::GemmParts { m, batch, .. } if *m <= GEMV_ROWS || *batch => "gemv_parts",
921        Launch::GemmParts { .. } => "gemm_parts",
922        Launch::Binary { .. } => "binary",
923        Launch::Scale { .. } => "scale",
924        Launch::SwiGlu { .. } => "swiglu",
925        Launch::Unary { .. } => "unary",
926        Launch::Rope { .. } | Launch::RopeParts { .. } => "rope",
927        Launch::SetRow { .. } => "set_row",
928        Launch::Attention { .. } => "attention",
929        Launch::ArgMax { .. } => "argmax",
930        Launch::Conv { .. } | Launch::ConvParts { .. } => "conv",
931        Launch::L2Norm { .. } => "l2norm",
932        Launch::Decay { .. } => "decay",
933        Launch::Gates { .. } | Launch::GateProj { .. } => "gates",
934        Launch::Delta { .. } => "delta",
935        Launch::Moe { .. } => "moe",
936        Launch::Slice { .. } => "slice",
937        Launch::Join { .. } => "join",
938        Launch::HeadGemv { .. } => "head_gemv",
939        Launch::GemvGlu { .. } => "gemv_glu",
940        Launch::Fold { first, .. } if matches!(first.as_ref(), Launch::Gather { .. }) => "gather_norm",
941        Launch::Fold { then, .. } if matches!(then.as_ref(), Launch::RopeParts { .. }) => "gemv_rope",
942        l @ Launch::Fold { .. } if l.finished_mixture().is_some() => "moe_norm",
943        Launch::Fold { .. } => "gemv_norm",
944    }
945}
946
947use slots::Slots;
948
949mod slots {
950    use super::Launch;
951
952    /// The launches being scheduled, each with the graph ops it computes: the schedule sets how each runs, never which
953    /// launches there are.
954    pub(super) struct Slots {
955        slots: Vec<(Launch, Vec<usize>)>,
956    }
957
958    impl Slots {
959        pub(super) fn new(slots: Vec<(Launch, Vec<usize>)>) -> Self {
960            Slots { slots }
961        }
962
963        pub(super) fn len(&self) -> usize {
964            self.slots.len()
965        }
966
967        pub(super) fn slot(&self, i: usize) -> Option<&(Launch, Vec<usize>)> {
968            self.slots.get(i)
969        }
970
971        pub(super) fn slot_mut(&mut self, i: usize) -> Option<&mut (Launch, Vec<usize>)> {
972            self.slots.get_mut(i)
973        }
974
975        pub(super) fn get(&self, i: usize) -> &Launch {
976            &self.slots[i].0
977        }
978
979        pub(super) fn get_mut(&mut self, i: usize) -> &mut Launch {
980            &mut self.slots[i].0
981        }
982
983        pub(super) fn live(&self) -> impl Iterator<Item = (usize, &Launch)> {
984            self.slots.iter().enumerate().map(|(i, (l, _))| (i, l))
985        }
986
987        /// The slot that computes `t` into its own buffer.
988        pub(super) fn producer(&self, t: &str) -> Option<usize> {
989            self.slots.iter().position(|(l, _)| l.result() == Some(t))
990        }
991
992        pub(super) fn into_inner(self) -> Vec<(Launch, Vec<usize>)> {
993            self.slots
994        }
995    }
996}
997
998/// The step graph `g` lowered under `numerics` and `knobs`: `constants` holds the placement of every constant the
999/// graph declares, `io` the step's carried names, `held` the steps of [`RULES`] held back — a shared rule, or a step of
1000/// the lane's schedule.
1001pub fn plan(
1002    g: &StepGraph,
1003    _names: &dyn Fn(crate::intern::Symbol) -> String,
1004    constants: BTreeMap<String, Placement>,
1005    numerics: Q8MatKind,
1006    knobs: &CudaKnobs,
1007    io: StepIo,
1008    held: &[&str],
1009) -> Result<Schedule, LaneError> {
1010    assert!(held.iter().all(|r| RULES.contains(r)), "every rule held back is one of {RULES:?}: {held:?}");
1011    let knobs = CudaKnobs { activations: numerics, ..*knobs };
1012    let shim = LaneShim { knobs, batch: !io.more.is_empty(), held: RULES.iter().copied().filter(|r| held.contains(r)).collect() };
1013    let steps = steps(&shim);
1014    let shared: Vec<&str> = held.iter().copied().filter(|r| steps.contains(r)).collect();
1015    let lifted = g.lifted.placed(
1016        &|t| match constants.get(t) {
1017            Some(Placement::Rows(r)) => Storage::Rows { format: r.ty, transposed: r.transposed },
1018            _ => Storage::Values,
1019        },
1020        Some(io.clone()),
1021    );
1022    let o = optimize(&lifted, &shim, &shared);
1023    let linked = link(&o.graph, &constants, &shim, &io).map_err(|(k, why)| {
1024        let op = o.graph.ops[k].covers.iter().copied().min().unwrap_or(0);
1025        LaneError::NoKernel { op, kind: g.ops[op].spec.kind(), out: g.ops[op].out.clone(), why }
1026    })?;
1027    if numerics != Q8MatKind::Dequant {
1028        whole_blocks(g, &linked)?;
1029    }
1030    // The firings in the order the steps take them — the merge of products on the tensor cores but a schedule of the
1031    // link's, which records none — each with the lifted ops of the launch computing its result, as that launch lists
1032    // them: a latent attention taken whole runs as its halves and the attention, and records the attention's; a batch's
1033    // gated attentions run as each sequence's attention and its row's set, and record them all.
1034    let recorded = |covers: &[usize]| -> Vec<usize> {
1035        let op = o.graph.ops.iter().find(|op| covered(op) == covers);
1036        if op.is_some_and(gates_batch_rows) {
1037            return covers.to_vec();
1038        }
1039        let result = op.map(|op| op.out.as_str());
1040        linked
1041            .iter()
1042            .find(|(l, ops)| {
1043                let mut set = ops.clone();
1044                set.sort_unstable();
1045                set == covers || (ops.iter().all(|k| covers.contains(k)) && result.is_some_and(|out| l.writes().contains(&out)))
1046            })
1047            .map_or_else(|| covers.to_vec(), |(_, ops)| ops.clone())
1048    };
1049    let mut fusions: Vec<Fusion> =
1050        o.fired.iter().filter(|f| f.rule != "merged_projections").map(|f| Fusion { rule: f.rule, ops: recorded(&f.covers) }).collect();
1051    let mut slots = Slots::new(linked);
1052    if !held.contains(&"rows_a_warp") {
1053        rows_a_warp(&mut slots, &mut fusions, &o.graph);
1054    }
1055    shared_rows(&mut slots, numerics, &knobs);
1056    for i in 0..slots.len() {
1057        if let Some((l, _)) = slots.slot_mut(i) {
1058            let l = if l.finished_mixture().is_some() { l.mixture_mut().expect("a finished mixture") } else { l };
1059            split_counts(l, numerics, &knobs);
1060        }
1061    }
1062    let (mut launches, mut ops): (Vec<Launch>, Vec<Vec<usize>>) = slots.into_inner().into_iter().unzip();
1063    // The tensors the link's own launches compute, which the step graph does not hold: a latent's expanded keys and
1064    // values, every cached row through each half.
1065    let planned: BTreeMap<String, usize> = launches
1066        .iter()
1067        .filter_map(|l| match l {
1068            Launch::HeadGemv { y, m, heads, b, shared: true, .. } => Some((y.buf.clone(), m * heads * b)),
1069            _ => None,
1070        })
1071        .collect();
1072    // The matrices read as values — a product's, a mixture's router — which the lane stores transposed.
1073    let transposed_values: BTreeSet<String> = g
1074        .ops
1075        .iter()
1076        .filter(|op| matches!(op.spec, TensorOpSpec::Matmul(_) | TensorOpSpec::Moe(_)))
1077        .map(|op| &op.ins[1])
1078        .filter(|w| matches!(constants.get(*w), Some(Placement::Values(_))))
1079        .cloned()
1080        .collect();
1081    let mut buffers = buffers(g, &planned, &launches, constants, &transposed_values, numerics, &io)?;
1082    // A gated-delta memory in the precision the `recurrent_state` knob stores it in, before a batch's memories are laid
1083    // out as one array of them.
1084    if knobs.recurrent_state == StatePrecision::Bf16 {
1085        for l in &launches {
1086            if let Launch::Delta { st, .. } = l {
1087                if let Some(b) = buffers.get_mut(st) {
1088                    b.dtype = DType::Bf16;
1089                }
1090            }
1091        }
1092    }
1093    // An attention's cache in the precision the `cache_precision` knob stores it in, and every view into it alike.
1094    if knobs.cache_precision == CachePrecision::F16 {
1095        for c in g.caches() {
1096            if let Some(b) = buffers.get_mut(&c) {
1097                b.dtype = DType::F16;
1098            }
1099        }
1100    }
1101    let views: Vec<(String, DType)> = buffers
1102        .values()
1103        .filter(|b| matches!(b.fill, Fill::Alias { .. }))
1104        .map(|b| (b.name.clone(), buffers[&span(&buffers, &b.name).0].dtype))
1105        .collect();
1106    for (name, dtype) in views {
1107        buffers.get_mut(&name).expect("a view").dtype = dtype;
1108    }
1109    if !io.more.is_empty() && !held.contains(&"sequences") {
1110        sequences(&mut launches, &mut ops, &mut buffers, &mut fusions, &io, knobs.chip.sms);
1111    }
1112    // An array stored in 16 bits is read and written only by kernels that take it so.
1113    for (i, l) in launches.iter().enumerate() {
1114        let typed = l.typed(numerics);
1115        if l.reads().into_iter().chain(l.writes()).any(|t| buffers.get(t).is_some_and(|b| b.dtype.bytes() == 2) && !typed.contains(&t)) {
1116            let op = ops[i].first().copied().unwrap_or(0);
1117            return Err(LaneError::NoKernel {
1118                op,
1119                kind: g.ops[op].spec.kind(),
1120                out: g.ops[op].out.clone(),
1121                why: "its kernel takes an array the step stores in 16 bits only as f32".to_string(),
1122            });
1123        }
1124    }
1125    Ok(Schedule { launches, ops, buffers, fusions, numerics, io })
1126}
1127
1128/// The array a tensor's values live in and where, `(array, from, to)`: its own, or through its views those it is part of.
1129fn span(buffers: &BTreeMap<String, Buffer>, name: &str) -> (String, usize, usize) {
1130    let len = buffers.get(name).map_or(0, |b| b.len);
1131    match buffers.get(name).map(|b| &b.fill) {
1132        Some(Fill::Alias { of, offset }) => {
1133            let (root, from, _) = span(buffers, of);
1134            (root, from + offset, from + offset + len)
1135        }
1136        _ => (name.to_string(), 0, len),
1137    }
1138}
1139
1140/// Where a batch's sequence `b`'s array lives in the batch's one array of them: `(array, sequence, values a sequence)`,
1141/// for a tensor that is a part of another at `b` times one width — a sequence's row of a stacked tensor, or its part of
1142/// one array of every sequence's (`sequences`).
1143fn sequence_part(buffers: &BTreeMap<String, Buffer>, name: &str, width: usize) -> Option<(String, usize)> {
1144    match &buffers.get(name)?.fill {
1145        Fill::Alias { of, offset } if width > 0 && offset % width == 0 => Some((of.clone(), offset / width)),
1146        _ => None,
1147    }
1148}
1149
1150/// A batch's step laid out and launched over its sequences. First every array a sequence of its own computes or carries
1151/// (`X_s<b>`, one for each of the `B` sequences, alike in length) becomes sequence `b`'s part of one array of them all:
1152/// its row of the table each sequence's array is set into whole, when there is one — the attended rows, set into the
1153/// output projection's input — or of a new array `X_s`. That changes no value the step computes. Then each set of `B`
1154/// launches alike but for their sequence — a layer's rotations, its value rows' writes into the caches, its attentions —
1155/// runs as one launch whose rows are the sequences, at the last one's place, each row reading its own position (item
1156/// `b` of the positions, `pos_s0` on) and its own caches; a set moves only where nothing between its first and its last
1157/// reads what it writes or writes what it reads, but the other sets.
1158fn sequences(launches: &mut Vec<Launch>, ops: &mut Vec<Vec<usize>>, buffers: &mut BTreeMap<String, Buffer>, fusions: &mut Vec<Fusion>, io: &StepIo, sms: usize) {
1159    let bsz = 1 + io.more.len();
1160    let positions: Vec<String> = io.positions().into_iter().map(String::from).collect();
1161    // The positions are the program's own array already, and the logits rows of the head's: the rest stack.
1162    let io_names: BTreeSet<String> = io.positions().into_iter().chain(io.more.iter().map(|q| q.logits.as_str())).chain([io.logits.as_str()]).map(String::from).collect();
1163    // The families: every sequence's array alike, carried or computed — plain, quantized where each member's producer
1164    // quantizes it, or both, its rows whole blocks where quantized, so each member's forms are its parts of the array's.
1165    let families: Vec<(String, Vec<String>)> = buffers
1166        .keys()
1167        .filter_map(|name| name.strip_suffix("_s0"))
1168        .map(|base| (base.to_string(), (0..bsz).map(|b| format!("{base}_s{b}")).collect::<Vec<String>>()))
1169        .filter(|(_, members)| {
1170            let first = &buffers[&members[0]];
1171            !io_names.contains(&members[0])
1172                && matches!(first.fill, Fill::Zero | Fill::Computed)
1173                && (first.plain || first.quantized)
1174                && (!first.quantized || quant_in_place(first.len))
1175                && members.iter().all(|m| {
1176                    buffers.get(m).is_some_and(|b| {
1177                        b.len == first.len && b.dtype == first.dtype && b.fill == first.fill && b.plain == first.plain && b.quantized == first.quantized
1178                    })
1179                })
1180        })
1181        .collect();
1182    // Each name's readers and writers, in order. A set a family's member is set by leaves the step once every family
1183    // is laid out; until then it is passed over, as gone.
1184    let (mut readers, mut writers): (std::collections::HashMap<String, Vec<usize>>, std::collections::HashMap<String, Vec<usize>>) = Default::default();
1185    for (i, l) in launches.iter().enumerate() {
1186        for t in l.reads().into_iter().collect::<BTreeSet<_>>() {
1187            readers.entry(t.to_string()).or_default().push(i);
1188        }
1189        for t in l.writes().into_iter().collect::<BTreeSet<_>>() {
1190            writers.entry(t.to_string()).or_default().push(i);
1191        }
1192    }
1193    let mut gone: BTreeSet<usize> = BTreeSet::new();
1194    for (base, members) in families {
1195        let len = buffers[&members[0]].len;
1196        // Each member set whole into its row of one table, by its only reader: the table's rows are the members.
1197        let into: Option<(String, Vec<usize>)> = {
1198            let sets: Vec<Option<(usize, String)>> = members
1199                .iter()
1200                .enumerate()
1201                .map(|(b, m)| {
1202                    let readers: Vec<usize> = readers.get(m).map_or_else(Vec::new, |at| at.iter().copied().filter(|i| !gone.contains(i)).collect());
1203                    match readers[..] {
1204                        [j] => match &launches[j] {
1205                            Launch::SetRow { table, index: RowSel::Fixed(r), src, cols, count: 1, seq: None, f32, quant }
1206                                if src == m && *cols == len && *r == b + 1 && *f32 == buffers[m].plain && *quant == buffers[m].quantized =>
1207                            {
1208                                Some((j, table.clone()))
1209                            }
1210                            _ => None,
1211                        },
1212                        _ => None,
1213                    }
1214                })
1215                .collect();
1216            match sets.into_iter().collect::<Option<Vec<_>>>() {
1217                Some(sets)
1218                    if sets.iter().all(|(_, t)| *t == sets[0].1)
1219                        && buffers.get(&sets[0].1).is_some_and(|t| {
1220                            let first = &buffers[&members[0]];
1221                            t.len == bsz * len && (t.plain || !first.plain) && (t.quantized || !first.quantized)
1222                        }) =>
1223                {
1224                    Some((sets[0].1.clone(), sets.iter().map(|(j, _)| *j).collect()))
1225                }
1226                _ => None,
1227            }
1228        };
1229        let array = match &into {
1230            Some((table, _)) => table.clone(),
1231            None => {
1232                let array = format!("{base}_s");
1233                if buffers.contains_key(&array) {
1234                    continue;
1235                }
1236                let first = buffers[&members[0]].clone();
1237                buffers.insert(array.clone(), Buffer { name: array.clone(), len: bsz * len, ..first });
1238                array
1239            }
1240        };
1241        for (b, m) in members.iter().enumerate() {
1242            buffers.get_mut(m).expect("a member").fill = Fill::Alias { of: array.clone(), offset: b * len };
1243        }
1244        if let Some((_, sets)) = into {
1245            // Each member's set made by its producer, which now writes the table's row: the set's graph op is its.
1246            for (b, &j) in sets.iter().enumerate() {
1247                if let Some(p) = writers.get(&members[b]).and_then(|at| at.iter().copied().find(|i| !gone.contains(i))) {
1248                    let set = ops[j].clone();
1249                    ops[p].extend(set);
1250                    ops[p].sort_unstable();
1251                }
1252            }
1253            gone.extend(sets);
1254        }
1255    }
1256    for j in gone.into_iter().rev() {
1257        launches.remove(j);
1258        ops.remove(j);
1259    }
1260    // Every name's span, as `span` finds it: the buffers are laid out by here.
1261    let mut spans = SpanCache::default();
1262    // A sequence's position: item `b` of the positions.
1263    let seq_of_pos = |p: &Option<String>| p.as_ref().and_then(|p| positions.iter().position(|q| q == p));
1264    // Each kind of per-sequence launch: its sequence and the launch over all of them, built from sequence 0's — the
1265    // attentions first, then the cache rows' writes, then the rotations, so each set moves past only the later
1266    // sequences' launches of the kinds before it, never one its own sequence's launch reads.
1267    type Merge = fn(&[&Launch], &BTreeMap<String, Buffer>, &str) -> Option<Launch>;
1268    let kinds: [(&str, Merge); 8] = [
1269        ("delta", merge_deltas),
1270        ("gates", merge_gates),
1271        ("conv", merge_convs),
1272        ("attention", merge_attentions),
1273        ("join", merge_joins),
1274        ("set_row", merge_row_sets),
1275        ("rope", merge_rotations),
1276        ("argmax", merge_argmaxes),
1277    ];
1278    // The launches as slots: a merged set's members but its last leave theirs empty, closed up once every kind is merged.
1279    let mut slots: Vec<Option<(Launch, Vec<usize>)>> = std::mem::take(launches).into_iter().zip(std::mem::take(ops)).map(Some).collect();
1280    let live = |slots: &[Option<(Launch, Vec<usize>)>], from: usize, to: usize| -> Vec<usize> {
1281        (from..to).filter(|&j| slots[j].is_some()).collect()
1282    };
1283    for (what, merge) in kinds {
1284        // The per-sequence launches of this kind, by sequence, in order: a run of `bsz` from sequence 0 to the last.
1285        for i in 0..slots.len() {
1286            let Some((first, _)) = &slots[i] else { continue };
1287            if kind(first) != what || launch_sequence(first, &seq_of_pos) != Some(0) {
1288                continue;
1289            }
1290            let mut at = vec![i];
1291            for j in i + 1..slots.len() {
1292                if at.len() == bsz {
1293                    break;
1294                }
1295                if let Some((l, _)) = &slots[j] {
1296                    if kind(l) == what && launch_sequence(l, &seq_of_pos) == Some(at.len()) {
1297                        at.push(j);
1298                    }
1299                }
1300            }
1301            let members: Vec<&Launch> = at.iter().map(|&j| &slots[j].as_ref().expect("a member").0).collect();
1302            let merged = if at.len() == bsz { merge(&members, buffers, &positions[0]) } else { None };
1303            let Some(mut merged) = merged else { continue };
1304            // A merged attention's rows are every sequence's: whether its chunks take blocks of their own is decided
1305            // for them all, the arithmetic the same either way.
1306            if let Launch::Attention { rows, heads, per_block, splits, parallel, .. } = &mut merged {
1307                *parallel = chunks_apart(*splits, *rows, *heads / *per_block, sms);
1308            }
1309            let last = *at.last().expect("a run");
1310            // Each member moves to the last's place past the launches between: none may write what it reads, or read
1311            // or write what it writes, an array's part overlapping the whole of it. One sweep from the first member to
1312            // the last: each launch between met against every member before it, what they touch and what they write.
1313            let (mut touched, mut written) = (Spans::default(), Spans::default());
1314            let mut crossed = false;
1315            for j in live(&slots, at[0], last) {
1316                let l = &slots[j].as_ref().expect("live").0;
1317                let (reads, writes) = (l.reads(), l.writes());
1318                if at.binary_search(&j).is_ok() {
1319                    for t in reads.iter().chain(&writes) {
1320                        touched.add(spans.of(buffers, t));
1321                    }
1322                    for t in &writes {
1323                        written.add(spans.of(buffers, t));
1324                    }
1325                } else if writes.iter().any(|t| touched.meets(spans.of(buffers, t))) || reads.iter().any(|t| written.meets(spans.of(buffers, t))) {
1326                    crossed = true;
1327                    break;
1328                }
1329            }
1330            if crossed {
1331                continue;
1332            }
1333            let mut joined: Vec<usize> = at.iter().flat_map(|&j| slots[j].as_ref().expect("a member").1.clone()).collect();
1334            joined.sort_unstable();
1335            fusions.push(Fusion { rule: "sequences", ops: joined.clone() });
1336            slots[last] = Some((merged, joined));
1337            for &j in &at[..at.len() - 1] {
1338                slots[j] = None;
1339            }
1340        }
1341    }
1342    (*launches, *ops) = slots.into_iter().flatten().unzip();
1343}
1344
1345/// Whether `op` is a batch's gated attentions (`attn_output_gate` over its sequences' rows), or their fold with the
1346/// quantize of the product: the op the lane runs as each sequence's attention and its row's set.
1347pub fn gates_batch_rows(op: &crate::tensor_graph::graph::Op) -> bool {
1348    let Sym::Fused(f) = &op.sym else { return false };
1349    match f.class {
1350        ClassId::AttnOutputGate => gates_rows(&f.parts),
1351        ClassId::GluActMulQuant => gates_batch_rows(&f.parts[0]),
1352        _ => false,
1353    }
1354}
1355
1356/// Each name's span (`span`), found once.
1357#[derive(Default)]
1358struct SpanCache(std::collections::HashMap<String, (String, usize, usize)>);
1359
1360impl SpanCache {
1361    fn of(&mut self, buffers: &BTreeMap<String, Buffer>, name: &str) -> (String, usize, usize) {
1362        if let Some(s) = self.0.get(name) {
1363            return s.clone();
1364        }
1365        let s = span(buffers, name);
1366        self.0.insert(name.to_string(), s.clone());
1367        s
1368    }
1369}
1370
1371/// Spans gathered by array, each array's parts as a set of positions — the overlapping ones joined, so a span meets
1372/// the set where it meets one of them — and its empty ones apart: an empty span at `p` meets a part around `p` but
1373/// never another empty one, as `span`'s parts overlap.
1374#[derive(Default)]
1375struct Spans(std::collections::HashMap<String, (BTreeMap<usize, usize>, BTreeSet<usize>)>);
1376
1377impl Spans {
1378    fn add(&mut self, (root, from, to): (String, usize, usize)) {
1379        let (parts, empty) = self.0.entry(root).or_default();
1380        if from == to {
1381            empty.insert(from);
1382            return;
1383        }
1384        let (mut from, mut to) = (from, to);
1385        // Every part overlapping the new one joins it.
1386        let joined: Vec<usize> = parts.range(..to).rev().take_while(|(_, &e)| e > from).map(|(&s, _)| s).collect();
1387        for s in joined {
1388            let e = parts.remove(&s).expect("a part");
1389            (from, to) = (from.min(s), to.max(e));
1390        }
1391        parts.insert(from, to);
1392    }
1393
1394    /// Whether `[from, to)` overlaps a span gathered: some part `[s, e)` with `from < e` and `s < to`.
1395    fn meets(&self, (root, from, to): (String, usize, usize)) -> bool {
1396        let Some((parts, empty)) = self.0.get(&root) else { return false };
1397        let part = parts.range(..to.max(from + 1)).next_back().is_some_and(|(&s, &e)| s < to && from < e);
1398        part || (from < to && empty.range(from + 1..to).next().is_some())
1399    }
1400}
1401
1402/// The batch sequence a per-sequence launch is for: the item of the positions it reads — or, for one that reads no
1403/// position, the sequence whose array it writes, `X_s<b>`.
1404fn launch_sequence(l: &Launch, seq_of_pos: &dyn Fn(&Option<String>) -> Option<usize>) -> Option<usize> {
1405    let of_name = |t: &str| t.rsplit_once("_s").and_then(|(_, b)| (!b.is_empty() && b.bytes().all(|c| c.is_ascii_digit())).then(|| b.parse().ok()).flatten());
1406    match l {
1407        Launch::RopeParts { pos, seqs: false, .. } | Launch::Attention { pos, seq: None, rows: 1, .. } => seq_of_pos(pos),
1408        Launch::SetRow { index: RowSel::Param(p), seq: None, count: 1, .. } => seq_of_pos(&Some(p.clone())),
1409        Launch::Delta { out, rows: 1, seqs: 1, .. } => of_name(out),
1410        Launch::ArgMax { advance: Some(a), rows: 1, .. } => seq_of_pos(&Some(a.position.clone())),
1411        Launch::Join { y, rows: 1, yz: 0, .. } => match &y.row {
1412            Some((RowSel::Param(p), _)) => seq_of_pos(&Some(p.clone())),
1413            Some(_) => None,
1414            None => of_name(&y.buf),
1415        },
1416        Launch::Gates { beta, .. } => of_name(beta),
1417        Launch::ConvParts { parts, seqs: 1, .. } => of_name(&parts[0].y),
1418        _ => None,
1419    }
1420}
1421
1422/// A layer's joins, one launch a sequence joining its row's units into its own tensor or into its own cache's row at its
1423/// position, as one launch whose rows are the sequences.
1424fn merge_joins(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, positions: &str) -> Option<Launch> {
1425    let Launch::Join { y, rows: 1, heads, ha, hb, yz: 0, .. } = members[0] else { return None };
1426    let width = heads * (ha + hb);
1427    let table = match &y.row {
1428        None => width,
1429        Some((RowSel::Param(_), _)) => buffers.get(&y.buf)?.len,
1430        Some(_) => return None,
1431    };
1432    let parts = |l: &Launch| match l {
1433        Launch::Join { a, b, y: Dst { buf, row, unit: None }, rows: 1, heads: h, ha: x, hb: z, yz: 0 }
1434            if (h, x, z) == (heads, ha, hb) && row.as_ref().map(|(_, p)| p) == y.row.as_ref().map(|(_, p)| p) =>
1435        {
1436            Some((sequence_part(buffers, a, heads * ha)?, sequence_part(buffers, b, heads * hb)?, sequence_part(buffers, buf, table)?))
1437        }
1438        _ => None,
1439    };
1440    let ((ar, _), (br, _), (yr, _)) = parts(members[0])?;
1441    for (s, l) in members.iter().enumerate() {
1442        if parts(l)? != ((ar.clone(), s), (br.clone(), s), (yr.clone(), s)) {
1443            return None;
1444        }
1445    }
1446    let (row, yz) = match &y.row {
1447        Some((_, pitch)) => (Some((RowSel::Param(positions.to_string()), *pitch)), table),
1448        None => (None, 0),
1449    };
1450    Some(Launch::Join { a: ar, b: br, y: Dst { buf: yr, row, unit: None }, rows: members.len(), heads: *heads, ha: *ha, hb: *hb, yz })
1451}
1452
1453/// A batch's arg-maxes, one launch a sequence over its row of the head's logits advancing its own token, position and
1454/// run, as one launch over the sequences.
1455fn merge_argmaxes(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, positions: &str) -> Option<Launch> {
1456    let Launch::ArgMax { x, out, n, advance: Some(a), rows: 1 } = members[0] else { return None };
1457    let (xa, _) = sequence_part(buffers, x, *n)?;
1458    let (oa, _) = sequence_part(buffers, out, 1)?;
1459    for (s, l) in members.iter().enumerate() {
1460        let Launch::ArgMax { x, out, n: m, advance: Some(b), rows: 1 } = l else { return None };
1461        if m != n || b.token != a.token || b.item != s || sequence_part(buffers, x, *n) != Some((xa.clone(), s)) || sequence_part(buffers, out, 1) != Some((oa.clone(), s)) {
1462            return None;
1463        }
1464    }
1465    let advance = Some(Advance { token: a.token.clone(), position: positions.to_string(), item: 0 });
1466    Some(Launch::ArgMax { x: xa, out: oa, n: *n, advance, rows: members.len() })
1467}
1468
1469/// A recurrent layer's delta rules, one launch a sequence over its one row and its own memory, as one launch over the
1470/// sequences: each input, gate, memory and output the sequence's part of one array of them all.
1471fn merge_deltas(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, _: &str) -> Option<Launch> {
1472    let Launch::Delta { q, k, v, z, b, a, st, gn, out, f32, quant, hv, hk, d, rows: 1, eps, .. } = members[0] else { return None };
1473    let widths = [(q, hk * d), (k, hk * d), (v, hv * d), (z, hv * d), (b, *hv), (a, *hv), (st, hv * d * d), (out, hv * d)];
1474    let arrays: Vec<String> = widths.iter().map(|(t, w)| sequence_part(buffers, t, *w).map(|(arr, _)| arr)).collect::<Option<_>>()?;
1475    for (s, l) in members.iter().enumerate() {
1476        let Launch::Delta { q, k, v, z, b, a, st, gn: g, out, f32: f, quant: qu, hv: h, hk: kk, d: dd, rows: 1, eps: e, .. } = l else { return None };
1477        let names = [q, k, v, z, b, a, st, out];
1478        if (g, f, qu, h, kk, dd, e) != (gn, f32, quant, hv, hk, d, eps)
1479            || names.iter().zip(&widths).zip(&arrays).any(|((t, (_, w)), arr)| sequence_part(buffers, t, *w) != Some((arr.clone(), s)))
1480        {
1481            return None;
1482        }
1483    }
1484    let [q, k, v, z, b, a, st, out]: [String; 8] = arrays.try_into().ok()?;
1485    Some(Launch::Delta { q, k, v, z, b, a, st, gn: gn.clone(), out, f32: *f32, quant: *quant, hv: *hv, hk: *hk, d: *d, rows: 1, eps: *eps, seqs: members.len() })
1486}
1487
1488/// A recurrent layer's gates, one launch a sequence over its row of the gate projections, as one launch over every
1489/// sequence's row: each gate's value is its element's alone.
1490fn merge_gates(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, _: &str) -> Option<Launch> {
1491    let Launch::Gates { br, ar, dt, al, beta, decay, n, heads } = members[0] else { return None };
1492    let names = [br, ar, beta, decay];
1493    let arrays: Vec<String> = names.iter().map(|t| sequence_part(buffers, t, *n).map(|(arr, _)| arr)).collect::<Option<_>>()?;
1494    for (s, l) in members.iter().enumerate() {
1495        let Launch::Gates { br, ar, dt: d, al: a, beta, decay, n: m, heads: h } = l else { return None };
1496        if (d, a, m, h) != (dt, al, n, heads) || [br, ar, beta, decay].iter().zip(&arrays).any(|(t, arr)| sequence_part(buffers, t, *n) != Some((arr.clone(), s))) {
1497            return None;
1498        }
1499    }
1500    let [br, ar, beta, decay]: [String; 4] = arrays.try_into().ok()?;
1501    Some(Launch::Gates { br, ar, dt: dt.clone(), al: al.clone(), beta, decay, n: n * members.len(), heads: *heads })
1502}
1503
1504/// A recurrent layer's convolutions, one launch a sequence of its parts, as one launch over the sequences: each part's
1505/// input row, ring and output the sequence's part of one array of them all.
1506fn merge_convs(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, _: &str) -> Option<Launch> {
1507    let Launch::ConvParts { parts, taps, block, eps, seqs: 1 } = members[0] else { return None };
1508    let mut merged = Vec::new();
1509    for (k, part) in parts.iter().enumerate() {
1510        let ring = taps * part.chan;
1511        if buffers.get(&part.ring)?.len != ring {
1512            return None;
1513        }
1514        let (x, _) = sequence_part(buffers, &part.x, part.chan)?;
1515        let (r, _) = sequence_part(buffers, &part.ring, ring)?;
1516        let (y, _) = sequence_part(buffers, &part.y, part.chan)?;
1517        for (s, l) in members.iter().enumerate() {
1518            let Launch::ConvParts { parts: ps, taps: t, block: bl, eps: e, seqs: 1 } = l else { return None };
1519            let o = ps.get(k)?;
1520            if ps.len() != parts.len() || (t, bl, e) != (taps, block, eps) || (o.chan, o.l2, &o.w) != (part.chan, part.l2, &part.w)
1521                || sequence_part(buffers, &o.x, part.chan) != Some((x.clone(), s))
1522                || sequence_part(buffers, &o.ring, ring) != Some((r.clone(), s))
1523                || sequence_part(buffers, &o.y, part.chan) != Some((y.clone(), s))
1524            {
1525                return None;
1526            }
1527        }
1528        merged.push(ConvPart { x, w: part.w.clone(), ring: r, y, chan: part.chan, l2: part.l2 });
1529    }
1530    Some(Launch::ConvParts { parts: merged, taps: *taps, block: *block, eps: *eps, seqs: members.len() })
1531}
1532
1533/// A layer's rotations, one launch a sequence of its query and its key — the query's row of the stacked projection into
1534/// its row of the queries, the key's into its cache's row at its position — as one launch over the sequences.
1535fn merge_rotations(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, positions: &str) -> Option<Launch> {
1536    let Launch::RopeParts { sides, cos, sin, hd, half, .. } = members[0] else { return None };
1537    let mut merged: Vec<RopeSide> = Vec::new();
1538    for (k, side) in sides.iter().enumerate() {
1539        if side.part.is_some() || side.y.unit.is_some() || side.heads != side.per_row {
1540            return None;
1541        }
1542        let width = side.heads * hd;
1543        // Each sequence's input row, at its place in one stacked tensor.
1544        let (x, _) = sequence_part(buffers, &side.x, width)?;
1545        let cache = side.y.row.as_ref();
1546        let (y, table) = match cache {
1547            None => (sequence_part(buffers, &side.y.buf, width)?.0, 0),
1548            Some(_) => {
1549                let len = buffers.get(&side.y.buf)?.len;
1550                (sequence_part(buffers, &side.y.buf, len)?.0, len)
1551            }
1552        };
1553        for (b, l) in members.iter().enumerate() {
1554            let Launch::RopeParts { sides: s, cos: c, sin: sn, hd: h, half: hf, .. } = l else { return None };
1555            let o = s.get(k)?;
1556            let width_ok = sequence_part(buffers, &o.x, width) == Some((x.clone(), b));
1557            let y_ok = match (cache, &o.y.row) {
1558                (None, None) => sequence_part(buffers, &o.y.buf, width) == Some((y.clone(), b)),
1559                (Some((_, pitch)), Some((RowSel::Param(_), p))) => p == pitch && sequence_part(buffers, &o.y.buf, table) == Some((y.clone(), b)),
1560                _ => false,
1561            };
1562            if s.len() != sides.len() || c != cos || sn != sin || h != hd || hf != half || o.norm != side.norm || o.heads != side.heads || !width_ok || !y_ok {
1563                return None;
1564            }
1565        }
1566        let row = cache.map(|(_, pitch)| (RowSel::Param(positions.to_string()), *pitch));
1567        merged.push(RopeSide {
1568            x,
1569            part: None,
1570            norm: side.norm.clone(),
1571            y: Dst { buf: y, row, unit: None },
1572            heads: side.heads * members.len(),
1573            per_row: side.per_row,
1574            yz: table,
1575        });
1576    }
1577    Some(Launch::RopeParts { sides: merged, cos: cos.clone(), sin: sin.clone(), pos: Some(positions.to_string()), hd: *hd, half: *half, seqs: true })
1578}
1579
1580/// A layer's value rows, one launch a sequence setting its row of the stacked projection into its cache's row at its
1581/// position, as one launch over the sequences.
1582fn merge_row_sets(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, positions: &str) -> Option<Launch> {
1583    let Launch::SetRow { table, src, cols, f32: true, quant: false, .. } = members[0] else { return None };
1584    let len = buffers.get(table)?.len;
1585    let (t, _) = sequence_part(buffers, table, len)?;
1586    let (x, _) = sequence_part(buffers, src, *cols)?;
1587    for (b, l) in members.iter().enumerate() {
1588        let Launch::SetRow { table: tb, src: sb, cols: cb, count: 1, f32: true, quant: false, .. } = l else { return None };
1589        if cb != cols || sequence_part(buffers, tb, len) != Some((t.clone(), b)) || sequence_part(buffers, sb, *cols) != Some((x.clone(), b)) {
1590            return None;
1591        }
1592    }
1593    Some(Launch::SetRow { table: t, index: RowSel::Param(positions.to_string()), src: x, cols: *cols, count: members.len(), seq: Some(len), f32: true, quant: false })
1594}
1595
1596/// A layer's attentions, one launch a sequence over its own caches up to its position, as one launch whose rows are the
1597/// sequences: each row's query and result its row of one array, its caches its part of one array of them.
1598fn merge_attentions(members: &[&Launch], buffers: &BTreeMap<String, Buffer>, positions: &str) -> Option<Launch> {
1599    let Launch::Attention { q, kc, vc, out, heads, hd, vd, split: None, gate, goff, .. } = members[0] else { return None };
1600    let len = buffers.get(kc)?.len;
1601    if buffers.get(vc)?.len != len {
1602        return None;
1603    }
1604    let (qa, _) = sequence_part(buffers, q, heads * hd)?;
1605    let (oa, _) = sequence_part(buffers, out, heads * vd)?;
1606    let (ka, _) = sequence_part(buffers, kc, len)?;
1607    let (va, _) = sequence_part(buffers, vc, len)?;
1608    for (b, l) in members.iter().enumerate() {
1609        let Launch::Attention { q: qb, kc: kb, vc: vb, out: ob, pos: _, gate: gb, goff: gob, .. } = l else { return None };
1610        let same = |x: &Launch| match (x, members[0]) {
1611            (
1612                Launch::Attention { sinks, heads, group, hd, vd, len, window, vstride, voff, scale, per_block, splits, parallel, split, shared, f32, quant, .. },
1613                Launch::Attention { sinks: s0, heads: h0, group: g0, hd: d0, vd: v0, len: l0, window: w0, vstride: vs0, voff: vo0, scale: sc0, per_block: p0, splits: sp0, parallel: pa0, split: st0, shared: sh0, f32: f0, quant: qu0, .. },
1614            ) => {
1615                (sinks, heads, group, hd, vd, len, window) == (s0, h0, g0, d0, v0, l0, w0)
1616                    && (vstride, voff, per_block, splits, parallel, split, shared) == (vs0, vo0, p0, sp0, pa0, st0, sh0)
1617                    && (f32, quant) == (f0, qu0)
1618                    && scale.to_bits() == sc0.to_bits()
1619            }
1620            _ => false,
1621        };
1622        // A gated sequence's gate is its row of one gate, as its output is its row of one table.
1623        let gated_alike = gb == gate && (gate.is_none() || *gob == goff + b * heads * vd);
1624        if !same(l)
1625            || !gated_alike
1626            || sequence_part(buffers, qb, heads * hd) != Some((qa.clone(), b))
1627            || sequence_part(buffers, ob, heads * vd) != Some((oa.clone(), b))
1628            || sequence_part(buffers, kb, len) != Some((ka.clone(), b))
1629            || sequence_part(buffers, vb, len) != Some((va.clone(), b))
1630        {
1631            return None;
1632        }
1633    }
1634    let mut merged = members[0].clone();
1635    if let Launch::Attention { q, kc, vc, out, pos, rows, seq, .. } = &mut merged {
1636        *q = qa;
1637        *kc = ka;
1638        *vc = va;
1639        *out = oa;
1640        *pos = Some(positions.to_string());
1641        *rows = members.len();
1642        *seq = Some(len);
1643    }
1644    Some(merged)
1645}
1646
1647/// An attention below the tiles whose values are its key rows' own part, under the quantized numerics, every query
1648/// head reading the one key/value group (a latent attention's absorbed form): one block takes every head and each
1649/// cached row is read once for all of them (`lk_attention_shared`), where head blocks of their own would each read the
1650/// whole span. Its span is cut into the decode step's chunks in every form, a block each, folded by a launch of their
1651/// own (`lk_attention_fold`). Taken where two steps' rows fit one block's shared memory. The values are the key rows'
1652/// part whether the attention reads them in place (`strided_values`) or a launch slices them out first, so the choice
1653/// and the results are the same with that rule held back.
1654fn shared_rows(s: &mut Slots, numerics: Q8MatKind, knobs: &CudaKnobs) {
1655    if numerics == Q8MatKind::Dequant {
1656        return;
1657    }
1658    for i in 0..s.len() {
1659        let Some((l, _)) = s.slot(i) else { continue };
1660        let fits = !l.attends_in_tiles(numerics)
1661            && matches!(l, Launch::Attention { heads, group, hd, vd, split: None, .. }
1662                if group == heads && *heads <= kernels::ATTN_SHARED_HEADS && hd % 4 == 0 && vd % 2 == 0)
1663            && values_in_key_rows(s, i);
1664        if !fits {
1665            continue;
1666        }
1667        if let Launch::Attention { heads, hd, len, window, shared, .. } = s.get_mut(i) {
1668            if kernels::attn_shared_smem_bytes(*heads, *hd) <= knobs.chip.smem_kib_per_block * 1024 {
1669                *shared = Some(kernels::attn_shared_splits(attended_span(*len, *window), knobs.chip.sms));
1670            }
1671        }
1672    }
1673}
1674
1675/// Whether attention `j`'s values are its key rows' own part, an even number of values from an even place: read in
1676/// place in the key rows, or sliced out of them by a launch of their own with the keys untouched between.
1677fn values_in_key_rows(s: &Slots, j: usize) -> bool {
1678    let Launch::Attention { kc, vc, heads, group, hd, vd, vstride, voff, .. } = s.get(j) else { return false };
1679    let (kvw, vw) = (heads / group * hd, heads / group * vd);
1680    if kc == vc {
1681        return *vstride == kvw && voff % 2 == 0 && voff + vd <= *hd;
1682    }
1683    let Some(i) = s.producer(vc) else { return false };
1684    matches!(s.get(i), Launch::Slice { x, every, from, width, .. } if x == kc && *every == kvw && from % 2 == 0 && from + vd <= *hd && *width == vw)
1685        && *voff == 0
1686        && *vstride == vw
1687        && s.live().filter(|(k, _)| *k > i && *k < j).all(|(_, l)| !l.writes().contains(&kc.as_str()))
1688}
1689
1690/// The steps of a lane a plan can hold back ([`plan`]'s `held`), in the order the lane takes them: the shared rules
1691/// its kernels link, by their names in the shared table ([`crate::tensor_graph::fuse`]); `rows_a_warp`, the schedule's
1692/// four rows a warp for a wide product; the decode step's products of one activation as one GEMV in parts
1693/// (`gemv_parts`, the lane's own rule in the table); and `sequences`, a batch's launches over its sequences as one. The
1694/// arg-max's advance closes the step and is always taken; the activation quantizer's fusion is the `fuse_quant` knob.
1695pub const RULES: [&str; 27] = [
1696    "add_norm",
1697    "sandwich_norm",
1698    "gemv_bias",
1699    "gemv_norm_epilogue",
1700    "head_norm_rope",
1701    "slice_read",
1702    "row_write",
1703    "qkv_head_epilogue",
1704    "join_in_place",
1705    "conv_silu",
1706    "conv_silu_l2",
1707    "conv_parts",
1708    "gates",
1709    "gate_projections",
1710    "act_gate",
1711    "attn_output_gate",
1712    "glu_pair",
1713    "shared_expert",
1714    "moe_rows",
1715    "moe_deferred_finalize",
1716    "gather_scale",
1717    "embed_norm",
1718    "strided_values",
1719    "rows_a_warp",
1720    "rope_parts",
1721    "gemv_parts",
1722    "sequences",
1723];
1724
1725/// A projection, gated pair or mixture whose shape outruns one output row a warp takes several
1726/// (`kernels::rows_a_warp`), each row's dot unchanged. A GEMV in parts takes its parts' rows a warp as the link set
1727/// them, each part recorded as taking them: the parts of the op of `g` that computes the launch's ops.
1728fn rows_a_warp(s: &mut Slots, fusions: &mut Vec<Fusion>, g: &TensorGraph) {
1729    for i in 0..s.len() {
1730        let Some((l, ops)) = s.slot_mut(i) else { continue };
1731        // A mixture with the next norm in its last step takes its rows as the mixture alone takes them.
1732        let l = if l.finished_mixture().is_some() { l.mixture_mut().expect("a finished mixture") } else { l };
1733        let (wide, per_warp) = match l {
1734            Launch::Gemv { format, k, n, m, per_warp, .. } => (kernels::rows_a_warp(*format, *n, *k, *m), per_warp),
1735            Launch::GemvGlu { formats: [fg, fu], k, n, m, per_warp, .. } => {
1736                (kernels::rows_a_warp(*fg, *n, *k, *m).min(kernels::rows_a_warp(*fu, *n, *k, *m)), per_warp)
1737            }
1738            Launch::Moe { formats, d, dff, k, shared, per_warp, rows: 1, .. } => {
1739                let hidden = *k * *dff + shared.as_ref().map_or(0, |e| e.dff);
1740                let formats = [formats[0], formats[1]].into_iter().chain(shared.iter().flat_map(|e| [e.formats[0], e.formats[1]]));
1741                (formats.map(|f| kernels::rows_a_warp(f, hidden, *d, 1)).min().expect("the gate's format"), per_warp)
1742            }
1743            Launch::GemmParts { per_warp, .. } if *per_warp > 1 => {
1744                for part in products_of(g, ops) {
1745                    fusions.push(Fusion { rule: "rows_a_warp", ops: part });
1746                }
1747                continue;
1748            }
1749            Launch::Fold { first, .. } if matches!(first.as_ref(), Launch::GemmParts { per_warp, .. } | Launch::Gemv { per_warp, .. } if *per_warp > 1) => {
1750                for part in products_of(g, ops) {
1751                    fusions.push(Fusion { rule: "rows_a_warp", ops: part });
1752                }
1753                continue;
1754            }
1755            _ => continue,
1756        };
1757        if wide > 1 {
1758            *per_warp = wide;
1759            fusions.push(Fusion { rule: "rows_a_warp", ops: ops.clone() });
1760        }
1761    }
1762}
1763
1764/// The lifted ops of each product a GEMV in parts computes — a part, or the projection of a part that turns its heads
1765/// in its epilogue — from the op of `g` that computes the launch's `ops`.
1766fn products_of(g: &TensorGraph, ops: &[usize]) -> Vec<Vec<usize>> {
1767    products_in(g.ops.iter().find(|op| covered(op) == ops).expect("a GEMV in parts is the op computing its ops"))
1768}
1769
1770/// The products a fused op computes: a fold's projection, a GEMV in parts' every part, the producer's where a quantize
1771/// rides its epilogue.
1772fn products_in(op: &crate::tensor_graph::graph::Op) -> Vec<Vec<usize>> {
1773    use ClassId::*;
1774    let Sym::Fused(f) = &op.sym else { panic!("a GEMV in parts is a fold of its parts") };
1775    let product = |p: &crate::tensor_graph::graph::Op| match &p.sym {
1776        Sym::Fused(e) if e.class == QkvHeadEpilogue => covered(&e.parts[0]),
1777        _ => covered(p),
1778    };
1779    match f.class {
1780        QkvHeadEpilogue | GemvNormEpilogue => vec![covered(&f.parts[0])],
1781        NormQuant | AddNormQuant | GluActMulQuant | ElementwiseQuant | GdnOutputQuant | GemvQuant | AttnOutputQuant
1782            if matches!(f.parts[0].sym, Sym::Fused(_)) =>
1783        {
1784            products_in(&f.parts[0])
1785        }
1786        _ => f.parts.iter().map(product).collect(),
1787    }
1788}
1789
1790/// The lifted ops an op of the IR computes, in order.
1791fn covered(op: &crate::tensor_graph::graph::Op) -> Vec<usize> {
1792    let mut c = op.covers.clone();
1793    c.sort_unstable();
1794    c
1795}
1796
1797/// A launch's split counts: a wide chunk's `f32` router and `f32` projections multiply their rows as one GEMM, its `k`
1798/// cut to give the chip's blocks work; any other `f32` projection cuts each output row's `k` over a block's warps; the
1799/// gates from their activation cut each head's dots as the projections they replace cut them.
1800fn split_counts(l: &mut Launch, numerics: Q8MatKind, knobs: &CudaKnobs) {
1801    let (on_cores, in_tiles) = (l.on_tensor_cores(numerics), l.multiplies_in_tiles(numerics));
1802    match l {
1803        Launch::Moe { router: Router::F32, rows, experts, d, router_split, .. } if on_cores => {
1804            let tiles = experts.div_ceil(64) * rows.div_ceil(64);
1805            *router_split = kernels::gemm_f32_splits(tiles, knobs.chip.sms, d.div_ceil(32));
1806        }
1807        Launch::GemvF32 { k, n, m, split, .. } if in_tiles => {
1808            *split = kernels::gemm_f32_splits(n.div_ceil(64) * m.div_ceil(64), knobs.chip.sms, k.div_ceil(32));
1809        }
1810        Launch::GemvF32 { k, n, split, .. } => *split = kernels::gemv_f32_parts(*n, *k, knobs.chip.sms),
1811        Launch::GateProj { k, heads, parts, .. } => *parts = kernels::gemv_f32_parts(*heads, *k, knobs.chip.sms),
1812        _ => {}
1813    }
1814}
1815
1816/// Every activation a projection reads quantized is whole 32-value blocks, its quantized form's, and so is the tensor
1817/// a view of it lies in where a view's rows start on a whole block: the tensor is quantized once, for all its views.
1818fn whole_blocks(g: &StepGraph, launches: &[(Launch, Vec<usize>)]) -> Result<(), LaneError> {
1819    let whole = |x: &str| {
1820        let n = g.tensors[x].len();
1821        match n % QBLOCK {
1822            0 => Ok(()),
1823            _ => Err(LaneError::Shape { var: x.to_string(), declared: n, bound: n.next_multiple_of(QBLOCK) }),
1824        }
1825    };
1826    let views: BTreeMap<&str, (&str, usize)> = launches
1827        .iter()
1828        .filter_map(|(l, _)| if let Launch::View { src, offset, out } = l { Some((out.as_str(), (src.as_str(), *offset))) } else { None })
1829        .collect();
1830    let read: BTreeSet<&str> = launches.iter().filter_map(|(l, _)| l.projected_input()).collect();
1831    let mut roots = BTreeSet::new();
1832    for x in read {
1833        whole(x)?;
1834        let mut root = x;
1835        while let Some(&(src, _)) = views.get(root).filter(|(_, offset)| quant_in_place(*offset)) {
1836            root = src;
1837        }
1838        roots.insert(root);
1839    }
1840    roots.into_iter().try_for_each(whole)
1841}
1842
1843/// Every tensor the launches touch, with the buffers it needs; `planned` holds the lengths of the tensors the plan's
1844/// own launches compute, which the step graph does not hold.
1845fn buffers(
1846    g: &StepGraph,
1847    planned: &BTreeMap<String, usize>,
1848    launches: &[Launch],
1849    mut constants: BTreeMap<String, Placement>,
1850    transposed: &BTreeSet<String>,
1851    numerics: Q8MatKind,
1852    io: &StepIo,
1853) -> Result<BTreeMap<String, Buffer>, LaneError> {
1854    let ints: BTreeSet<&str> = launches
1855        .iter()
1856        .filter_map(|l| if let Launch::ArgMax { out, .. } = l { Some(out.as_str()) } else { None })
1857        .collect();
1858    let views: BTreeMap<&str, (&str, usize)> = launches
1859        .iter()
1860        .filter_map(|l| if let Launch::View { src, offset, out } = l { Some((out.as_str(), (src.as_str(), *offset))) } else { None })
1861        .collect();
1862    let mut out: BTreeMap<String, Buffer> = BTreeMap::new();
1863    let mut mark = |out: &mut BTreeMap<String, Buffer>, name: &str, form: Touch| -> Result<(), LaneError> {
1864        if !out.contains_key(name) {
1865            if let Some(&len) = planned.get(name) {
1866                out.insert(name.to_string(), Buffer { name: name.to_string(), len, dtype: DType::F32, fill: Fill::Computed, plain: false, quantized: false });
1867                let b = out.get_mut(name).expect("inserted");
1868                match form {
1869                    Touch::Plain => b.plain = true,
1870                    Touch::Quant => b.quantized = true,
1871                }
1872                return Ok(());
1873            }
1874            let t = g.tensors.get(name).ok_or_else(|| LaneError::Unbound(name.to_string()))?;
1875            let (dtype, fill) = match t.role {
1876                Role::Param => (DType::I32, Fill::Param),
1877                Role::State => (DType::F32, Fill::Zero),
1878                Role::Constant => {
1879                    let mut p = constants.remove(name).ok_or_else(|| LaneError::Unbound(name.to_string()))?;
1880                    if let (true, Placement::Values(v)) = (transposed.contains(name), &mut p) {
1881                        *v = (0..t.cols).flat_map(|c| (0..t.rows).map(move |r| (r, c))).map(|(r, c)| v[r * t.cols + c]).collect();
1882                    }
1883                    (DType::F32, Fill::Constant(p))
1884                }
1885                Role::Activation => match views.get(name) {
1886                    Some(&(of, offset)) => (DType::F32, Fill::Alias { of: of.to_string(), offset }),
1887                    None if ints.contains(name) => (DType::I32, Fill::Computed),
1888                    None => (DType::F32, Fill::Computed),
1889                },
1890            };
1891            out.insert(name.to_string(), Buffer { name: name.to_string(), len: t.len(), dtype, fill, plain: false, quantized: false });
1892        }
1893        let b = out.get_mut(name).expect("inserted");
1894        match form {
1895            Touch::Plain => b.plain = true,
1896            Touch::Quant => b.quantized = true,
1897        }
1898        Ok(())
1899    };
1900    // A view's `f32` buffer is its tensor's, so reading one reads the other. (Its quantized form is its tensor's in place
1901    // too, and that tensor is quantized by the launch `quantize` gave it, which marks it.)
1902    let mut touch = |name: &str, form: Touch| -> Result<(), LaneError> {
1903        let mut at = Some(name);
1904        while let Some(name) = at {
1905            mark(&mut out, name, form)?;
1906            at = views.get(name).filter(|_| form == Touch::Plain).map(|&(of, _)| of);
1907        }
1908        Ok(())
1909    };
1910    for l in launches {
1911        let (reads, writes) = l.touches(numerics);
1912        for (t, form) in reads.into_iter().chain(writes) {
1913            touch(t, form)?;
1914        }
1915    }
1916    for name in io.names() {
1917        touch(name, Touch::Plain)?;
1918    }
1919    Ok(out)
1920}