1use std::collections::{BTreeMap, BTreeSet};
15use std::fmt::Write as _;
16
17use crate::codegen_cuda::knobs::CudaKnobs;
18use crate::codegen_cuda::provenance::{header, Provenance};
19use crate::codegen_cuda::qwen_emit::Q8MatKind;
20use crate::model::quant::GgmlType;
21use crate::tensor_recognize::qwen::{Activation, MoeAct};
22use crate::tensor_recognize::ElementwiseOp;
23
24use super::arena;
25use super::branches;
26use super::kernels::{self, Kernel};
27use super::plan::{kind, quant_in_place, Buffer, Dst, Fill, Launch, Part, RopeSide, Router, RowSel, Touch, GEMV_ROWS, QBLOCK};
28use super::residence::{shared_names, Plan, POOL_ALIGN};
29use super::weights::Placement;
30use super::graph::Role;
31use super::{Form, FormKind, LaneError};
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq)]
35pub struct Checkpoint {
36 pub size: u64,
37 pub data_start: u64,
38 pub digest: u64,
39}
40
41#[derive(Clone, Debug, PartialEq, Eq)]
43pub struct Program {
44 pub cu: String,
45 pub sidecar: Vec<u8>,
46}
47
48fn c(name: &str) -> String {
50 name.chars().map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '_' }).collect()
51}
52
53fn cdiv(a: usize, b: usize) -> usize {
54 a.div_ceil(b)
55}
56
57fn flit(v: f32) -> String {
59 format!("{v:e}f")
60}
61
62struct Names<'a> {
65 tag: String,
66 rows: usize,
67 batch: bool,
68 shared: &'a BTreeSet<String>,
69 positions: Vec<&'a str>,
70 two: bool,
71 sms: usize,
72 buffers: &'a BTreeMap<String, Buffer>,
73}
74
75impl<'a> Names<'a> {
76 fn of(f: &'a Form, shared: &'a BTreeSet<String>, two: bool, sms: usize) -> Names<'a> {
77 Names {
78 tag: f.tag(),
79 rows: f.rows,
80 batch: f.kind == FormKind::Batch,
81 shared,
82 positions: f.schedule.io.positions(),
83 two,
84 sms,
85 buffers: &f.schedule.buffers,
86 }
87 }
88
89 fn elems(&self, arrays: &[Option<&str>]) -> String {
91 elems(self.buffers, arrays)
92 }
93
94 fn elem(&self, t: &str) -> &'static str {
96 self.buffers.get(t).map_or("float", |b| b.dtype.c_type())
97 }
98
99 fn plain(&self, t: &str) -> String {
100 match self.positions.iter().position(|p| *p == t) {
101 Some(b) if self.batch => format!("(d_bpos + {b})"),
102 Some(_) => "d_pos".into(),
103 None if self.shared.contains(t) => format!("t_{}", c(t)),
104 None => format!("t{}_{}", self.tag, c(t)),
105 }
106 }
107
108 fn conv_pos(&self, ring: &str, width: usize, seqs: usize) -> String {
112 if !self.batch {
113 return "d_pos".into();
114 }
115 if seqs > 1 {
116 return "d_bpos".into();
117 }
118 let named = ring.rsplit_once("_s").and_then(|(_, b)| b.parse::<usize>().ok());
119 let b = named.or_else(|| match self.buffers.get(ring).map(|b| &b.fill) {
120 Some(Fill::Alias { offset, .. }) if width > 0 && offset % width == 0 => Some(offset / width),
121 _ => None,
122 });
123 format!("(d_bpos + {})", b.unwrap_or(0))
124 }
125
126 fn opt(&self, t: &Option<String>) -> String {
127 t.as_deref().map_or_else(|| "nullptr".to_string(), |t| self.plain(t))
128 }
129
130 fn qname(&self, t: &str) -> String {
133 if self.shared.contains(t) {
134 format!("_{}", c(t))
135 } else {
136 format!("{}_{}", self.tag, c(t))
137 }
138 }
139
140 fn quant(&self, t: &str, on: bool) -> String {
141 if !on {
142 return "nullptr, nullptr, nullptr".into();
143 }
144 let n = self.qname(t);
145 let lo = if self.two { format!("l{n}") } else { "nullptr".into() };
146 format!("q{n}, {lo}, s{n}")
147 }
148
149 fn quant_at(&self, t: &str, at: usize) -> String {
151 let n = self.qname(t);
152 let lo = if self.two { format!("l{n} + {at}") } else { "nullptr".into() };
153 format!("q{n} + {at}, {lo}, s{n} + {}", at / QBLOCK)
154 }
155
156 fn scratch(&self, kind: &str, t: &str) -> String {
157 format!("{kind}{}_{}", self.tag, c(t))
158 }
159
160 fn dst(&self, d: &Dst) -> String {
162 self.dst_at(d, 0)
163 }
164
165 fn dst_at(&self, d: &Dst, off: usize) -> String {
167 let (buf, row, pitch) = self.dst_split(d, off);
168 format!("{buf}, {row}, {pitch}")
169 }
170
171 fn dst_split(&self, d: &Dst, off: usize) -> (String, String, usize) {
173 let at = |o: usize| if o == 0 { self.plain(&d.buf) } else { format!("{} + {o}", self.plain(&d.buf)) };
174 match &d.row {
175 Some((RowSel::Param(p), pitch)) => (at(off), self.plain(p), *pitch),
176 Some((RowSel::Fixed(r), pitch)) => (at((r - 1) * pitch + off), "nullptr".into(), 0),
177 None => (at(off), "nullptr".into(), 0),
178 }
179 }
180
181 fn dst_ptr(&self, d: &Dst) -> String {
183 match &d.row {
184 Some((RowSel::Fixed(r), pitch)) => format!("{} + {}", self.plain(&d.buf), (r - 1) * pitch),
185 _ => self.plain(&d.buf),
186 }
187 }
188}
189
190const QWORD_ROWS: &str = "w_";
191const QWORD_SCALES: &str = "ws_";
192
193fn rows_q(t: &str) -> String {
194 format!("{QWORD_ROWS}{}", c(t))
195}
196
197fn rows_s(t: &str) -> String {
198 format!("{QWORD_SCALES}{}", c(t))
199}
200
201fn rows_b(t: &str) -> String {
203 format!("wb_{}", c(t))
204}
205
206fn handle(t: &str, ty: GgmlType, cols: usize) -> String {
209 if ty == GgmlType::Q8_0 {
210 format!("lk_w{{{}, {}, nullptr, 0}}", rows_q(t), rows_s(t))
211 } else {
212 format!("lk_w{{nullptr, nullptr, {}, {}L}}", rows_b(t), ty.bytes_for(cols))
213 }
214}
215
216fn gemm_blk_name(f: u32, m: u32, group: u32, bm: usize, act: i32) -> String {
220 let default = if matches!(group, 1..=3) { 64 } else { kernels::GEMM_BLK_TILE_M };
221 match (group, bm == default, act) {
222 (0, true, _) => format!("lk_gemm_blk<{f}, {m}>"),
223 (_, true, 0) => format!("lk_gemm_blk<{f}, {m}, {group}>"),
224 (_, _, 0) => format!("lk_gemm_blk<{f}, {m}, {group}, {bm}>"),
225 _ => format!("lk_gemm_blk<{f}, {m}, {group}, {bm}, {act}>"),
226 }
227}
228
229fn kernels_used(forms: &[Form]) -> Vec<Kernel> {
231 let mut ks: Vec<Kernel> = forms.iter().flat_map(|f| f.schedule.launches.iter().flat_map(move |l| launch_kernels(l, f.schedule.numerics))).collect();
232 ks.sort();
233 ks.dedup();
234 ks
235}
236
237fn launch_kernels(l: &Launch, numerics: Q8MatKind) -> Vec<Kernel> {
239 match l {
240 Launch::Gather { format: Some(GgmlType::Q8_0), .. } => vec![Kernel::GatherQ8],
241 Launch::Gather { format: Some(_), .. } => vec![Kernel::Blocks, Kernel::GatherBlk],
242 Launch::Gather { .. } => vec![Kernel::GatherF32],
243 Launch::View { .. } => vec![],
244 Launch::Norm { .. } => vec![Kernel::RmsNorm],
245 Launch::Quant { .. } => vec![Kernel::Quant],
246 Launch::Moe { router, .. } if l.on_tensor_cores(numerics) => {
247 let mut k = vec![Kernel::Blocks, Kernel::Mma, Kernel::GemmBlk, Kernel::Moe, Kernel::Quant];
248 k.extend((*router == Router::F32).then_some(Kernel::GemmF32));
249 k
250 }
251 l if l.on_tensor_cores(numerics) => vec![Kernel::Blocks, Kernel::Mma, Kernel::GemmBlk],
252 Launch::GemvGlu { .. } if l.batch_on_tensor_cores(numerics) => {
253 vec![Kernel::Blocks, Kernel::Mma, Kernel::GemvMma, Kernel::SwiGlu, Kernel::Binary]
254 }
255 l if l.batch_on_tensor_cores(numerics) => vec![Kernel::Blocks, Kernel::Mma, Kernel::GemvMma],
256 Launch::GemmParts { .. } => vec![Kernel::Blocks, Kernel::GemvParts],
257 Launch::Gemv { format: GgmlType::Q8_0, .. } => vec![Kernel::GemvQ8],
258 Launch::Gemv { .. } => vec![Kernel::Blocks, Kernel::GemvBlk],
259 l if l.multiplies_in_tiles(numerics) => vec![Kernel::GemmF32],
260 Launch::GemvF32 { .. } => vec![Kernel::GemvF32],
261 Launch::Binary { .. } => vec![Kernel::Binary],
262 Launch::Scale { .. } => vec![Kernel::Scale],
263 Launch::SwiGlu { .. } => vec![Kernel::SwiGlu],
264 Launch::Unary { .. } => vec![Kernel::Unary],
265 Launch::Rope { .. } | Launch::RopeParts { .. } => vec![Kernel::Rope],
266 Launch::SetRow { .. } => vec![Kernel::SetRow],
267 Launch::Attention { .. } if l.attends_on_tensor_cores(numerics) => vec![Kernel::AttentionMma],
268 Launch::Attention { .. } if l.attends_in_two_words(numerics) => vec![Kernel::Mma, Kernel::AttentionI8],
269 Launch::Attention { .. } if l.attends_flash(numerics) => vec![Kernel::Mma, Kernel::AttentionFlash],
270 Launch::Attention { .. } if l.attends_in_tiles(numerics) => vec![Kernel::AttentionTiled],
271 Launch::Attention { shared: Some(_), .. } => vec![Kernel::Mma, Kernel::AttentionShared],
272 Launch::Attention { .. } => vec![Kernel::Attention],
273 Launch::ArgMax { .. } => vec![Kernel::ArgMax],
274 Launch::Conv { .. } | Launch::ConvParts { .. } => vec![Kernel::Conv],
275 Launch::L2Norm { .. } => vec![Kernel::L2Norm],
276 Launch::Decay { .. } => vec![Kernel::Decay],
277 Launch::Gates { .. } | Launch::GateProj { .. } => vec![Kernel::Gates],
278 Launch::Delta { .. } if l.delta_in_chunks(numerics) => vec![Kernel::DeltaChunks, Kernel::DeltaNorm],
279 Launch::Delta { .. } => vec![Kernel::DeltaStep, Kernel::DeltaNorm, Kernel::DeltaStepNorm],
280 Launch::Moe { .. } if numerics != Q8MatKind::Dequant => vec![Kernel::Blocks, Kernel::Moe, Kernel::Quant],
281 Launch::Moe { .. } => vec![Kernel::Blocks, Kernel::Moe],
282 Launch::Slice { .. } => vec![Kernel::Slice],
283 Launch::Join { .. } => vec![Kernel::Join],
284 Launch::HeadGemv { .. } if l.heads_in_tiles(numerics) => vec![Kernel::Blocks, Kernel::HeadTiled],
285 Launch::HeadGemv { .. } => vec![Kernel::Blocks, Kernel::Heads],
286 Launch::GemvGlu { .. } => vec![Kernel::Blocks, Kernel::Glu],
287 Launch::Fold { first, .. } if matches!(first.as_ref(), Launch::Gather { .. }) => {
288 vec![Kernel::Blocks, Kernel::GatherQ8, Kernel::GatherBlk, Kernel::GatherF32, Kernel::RmsNorm, Kernel::GatherNorm]
289 }
290 Launch::Fold { then, .. } if matches!(then.as_ref(), Launch::RopeParts { .. }) => vec![Kernel::Blocks, Kernel::Rope, Kernel::GemvPartsRope],
291 Launch::Fold { first, .. } if l.finished_mixture().is_some() => {
293 let mut k = launch_kernels(first.parts()[0], numerics);
294 k.extend([Kernel::RmsNorm, Kernel::MoeNorm]);
295 k
296 }
297 Launch::Fold { first, .. } if matches!(first.as_ref(), Launch::Gemv { format: GgmlType::Q8_0, .. }) => vec![Kernel::RmsNorm, Kernel::GemvQ8Norm],
298 Launch::Fold { .. } => vec![Kernel::Blocks, Kernel::GemvBlk, Kernel::RmsNorm, Kernel::GemvBlkNorm],
299 }
300}
301
302fn reading(part: &Option<Part>, width: usize) -> (usize, usize) {
305 part.map_or((width, 0), |q| (q.pitch, q.offset))
306}
307
308fn unit(d: &Dst, width: usize) -> (usize, usize) {
310 reading(&d.unit, width)
311}
312
313fn per_warp_arg(rows: usize) -> String {
315 if rows == 1 {
316 String::new()
317 } else {
318 format!(", {rows}")
319 }
320}
321
322fn rope_threads(normed: bool, hd: usize) -> usize {
325 if normed {
326 kernels::norm_threads(hd)
327 } else {
328 128
329 }
330}
331
332fn rope_side(r: &RopeSide, hd: usize, n: &Names) -> String {
334 let ((xp, xo), (yp, yo)) = (reading(&r.part, hd), unit(&r.y, hd));
335 format!(
336 "lk_rope_side{{{}, {}, {}, {}, {xp}, {xo}, {yp}, {yo}, {}, {}, {}}}",
337 n.plain(&r.x),
338 r.norm.as_ref().map_or_else(|| "nullptr".to_string(), |(w, _)| n.plain(w)),
339 n.dst(&r.y),
340 flit(r.norm.as_ref().map_or(0.0, |(_, e)| e.get() as f32)),
341 r.heads,
342 r.per_row,
343 r.yz
344 )
345}
346
347fn elems(buffers: &BTreeMap<String, Buffer>, arrays: &[Option<&str>]) -> String {
352 let types: Vec<&str> = arrays.iter().map(|a| a.and_then(|a| buffers.get(a)).map_or("float", |b| b.dtype.c_type())).collect();
353 if types.iter().all(|t| *t == "float") {
354 String::new()
355 } else {
356 types.iter().map(|t| format!(", {t}")).collect()
357 }
358}
359
360#[allow(clippy::too_many_arguments)]
365fn attention_kernel(per_block: usize, vd: usize, ty: &str, mb: usize, kv: &str, qo: &str, path: &str, hd: usize) -> String {
366 let kv = if kv.is_empty() { ", float, float" } else { kv };
367 let qo = if qo.is_empty() { ", 0" } else { qo };
368 format!("lk_attention<{per_block}, {}, {ty}, {mb}{kv}{qo}, {path}, {}>", kernels::attn_values_per_thread(vd), kernels::attn_run_width(hd))
369}
370
371fn attention_shared_kernel(heads: usize, vd: usize, hd: usize, vk: bool, kv: &str, qo: &str) -> String {
375 let kv = if !qo.is_empty() && kv.is_empty() { ", float, float" } else { kv };
376 format!("lk_attention_shared<{heads}, {}, {}, {vk}{kv}{qo}>", kernels::attn_values_per_thread(vd), hd / 4)
377}
378
379fn attention_qo(quant: bool, gate: &Option<(String, Activation)>, m: u32) -> String {
382 if quant || gate.is_some() {
383 format!(", {}", if quant { m } else { 0 })
384 } else {
385 String::new()
386 }
387}
388
389fn attention_quant(n: &Names, out: &str, quant: bool, gate: &Option<(String, Activation)>, goff: usize, m: u32) -> (String, String) {
392 let words = if quant { n.quant(out, true) } else { "nullptr, nullptr, nullptr".into() };
393 let (g, ga) = gate.as_ref().map_or(("nullptr".to_string(), -1), |(g, act)| match goff {
394 0 => (n.plain(g), act_arg(*act)),
395 off => (format!("({} + {off})", n.plain(g)), act_arg(*act)),
396 });
397 let qo = attention_qo(quant, gate, m);
398 let args = if qo.is_empty() { String::new() } else { format!(", lk_qout{{{words}, {g}, {ga}}}") };
399 (qo, args)
400}
401
402fn attention_smem(per_block: usize, hd: usize, vd: usize, splits: usize, parallel: bool, acc: usize) -> (usize, usize) {
406 let side = if !parallel && kernels::attn_holds(per_block, vd, splits, acc) { kernels::attn_side_chunks(vd, splits) } else { 0 };
407 let held = if side > 0 { kernels::attn_held_bytes(per_block, vd, splits, acc) + kernels::attn_side_bytes(per_block, side, acc) } else { 0 };
408 (kernels::attn_smem_bytes(per_block, hd, splits, acc) + held, side)
409}
410
411struct PartOut {
414 w: String,
415 dst: (String, String, usize),
416 bias: String,
417 n: usize,
418}
419
420fn lk_parts(parts: &[PartOut]) -> String {
422 let pad = |v: Vec<String>, with: &str| (0..kernels::GEMM_PARTS).map(|i| v.get(i).cloned().unwrap_or_else(|| with.to_string())).collect::<Vec<_>>().join(", ");
423 let cols: usize = parts.iter().map(|p| p.n).sum();
424 let ends: Vec<String> = parts.iter().scan(0, |end, p| {
425 *end += p.n;
426 Some(end.to_string())
427 }).collect();
428 let groups = [
429 pad(parts.iter().map(|p| p.w.clone()).collect(), "lk_w{}"),
430 pad(parts.iter().map(|p| p.dst.0.clone()).collect(), "nullptr"),
431 pad(parts.iter().map(|p| p.bias.clone()).collect(), "nullptr"),
432 pad(ends, &cols.to_string()),
433 pad(parts.iter().map(|p| p.dst.1.clone()).collect(), "nullptr"),
434 pad(parts.iter().map(|p| p.dst.2.to_string()).collect(), "0"),
435 ];
436 format!("lk_parts{{{}}}", groups.iter().map(|g| format!("{{{g}}}")).collect::<Vec<_>>().join(", "))
437}
438
439fn gemv_mma(n: &Names, x: &str, k: usize, rows: usize, parts: &[PartOut], m: u32) -> String {
443 let cols: usize = parts.iter().map(|p| p.n).sum();
444 (0..rows.div_ceil(kernels::GEMV_MMA_ROWS))
445 .map(|slab| {
446 let r = kernels::GEMV_MMA_ROWS.min(rows - slab * kernels::GEMV_MMA_ROWS);
447 let at = slab * kernels::GEMV_MMA_ROWS;
448 let (mt, nt, mw, nw) = kernels::gemv_mma_launch(r, cols, k, n.sms);
449 let slab_parts: Vec<PartOut> = parts
450 .iter()
451 .map(|p| {
452 let (buf, row, pitch) = &p.dst;
453 let base = if at == 0 { buf.clone() } else { format!("{buf} + {}", at * p.n) };
454 PartOut { w: p.w.clone(), dst: (base, row.clone(), *pitch), bias: p.bias.clone(), n: p.n }
455 })
456 .collect();
457 format!(
458 "lk_gemv_mma<{m}, {mt}, {nt}, {mw}, {nw}><<<{}, {}, {}, st>>>({}, {k}, {cols}, {r}, {});",
459 cdiv(cols, 8 * nt),
460 kernels::gemv_mma_threads(mt, nt, mw, nw),
461 kernels::gemv_mma_smem_bytes(n.two, mt, nt),
462 n.quant_at(x, at * k),
463 lk_parts(&slab_parts)
464 )
465 })
466 .collect::<Vec<_>>()
467 .join("\n CKL();\n ")
468}
469
470fn act_arg(act: Activation) -> i32 {
471 match act {
472 Activation::Silu => 0,
473 Activation::Sigmoid => 1,
474 Activation::Gelu => 2,
475 }
476}
477
478fn launch(l: &Launch, n: &Names, numerics: Q8MatKind, token: &str, out: &mut String) {
481 let m = kernels::mode(numerics);
482 let qm = |on: bool| if on { m } else { 0 };
483 let sum_ty = kernels::accumulator(m).1;
485 let b = |v: bool| if v { "true" } else { "false" };
486 let p = |t: &str| n.plain(t);
487 let (l, finish) = match l.finished_mixture() {
490 Some((mixture, norm, sum)) => (mixture, Some((norm, sum, n.scratch("gc", l.writes()[0])))),
491 None => (l, None),
492 };
493 let line = match l {
494 Launch::View { .. } => return,
495 Launch::Gather { table, index, out, cols, rows, format: Some(GgmlType::Q8_0), scale } => format!(
496 "lk_gather_q8<<<dim3({}, {rows}), 256, 0, st>>>({}, {}, {}, {}, {cols}, {});",
497 cdiv(*cols, 256),
498 rows_q(table),
499 rows_s(table),
500 p(index),
501 p(out),
502 flit(scale.unwrap_or(1.0) as f32)
503 ),
504 Launch::Gather { table, index, out, cols, rows, format: Some(ty), scale } => format!(
505 "lk_gather_blk<{}><<<dim3({}, {rows}), 256, 0, st>>>({}, {}, {}, {cols}, {});",
506 kernels::format_arg(*ty),
507 cdiv(cols / 32, 256),
508 handle(table, *ty, *cols),
509 p(index),
510 p(out),
511 flit(scale.unwrap_or(1.0) as f32)
512 ),
513 Launch::Gather { table, index, out, cols, rows, format: None, scale } => format!(
514 "lk_gather_f32<<<dim3({}, {rows}), 256, 0, st>>>({}, {}, {}, {cols}, {});",
515 cdiv(*cols, 256),
516 p(table),
517 p(index),
518 p(out),
519 flit(scale.unwrap_or(1.0) as f32)
520 ),
521 Launch::Scale { x, y, n: len, factor } => format!(
522 "lk_scale<<<{}, 256, 0, st>>>({}, {}, {len}, {});",
523 cdiv(*len, 256),
524 p(x),
525 p(y),
526 flit(*factor as f32)
527 ),
528 Launch::Norm { x, part, add, pre, sum, w, out, into, f32, quant, rows, cols, eps } => {
529 let own = Dst::to(out);
530 let d = into.as_ref().unwrap_or(&own);
531 let (y, yrow, ypitch) = if *f32 { n.dst_split(d, 0) } else { ("nullptr".into(), "nullptr".into(), 0) };
532 let ((xp, xo), (yp, yo)) = (reading(part, *cols), unit(d, *cols));
533 let threads = kernels::norm_threads(*cols);
534 let (pre_arg, pre_tail) = match pre {
536 Some(pn) => (", true", format!(", {}, {}", n.opt(&pn.w), flit(pn.eps.get() as f32))),
537 None => ("", String::new()),
538 };
539 format!(
540 "lk_rmsnorm<{}, {}, {}, {}, {}, {sum_ty}{pre_arg}><<<{rows}, {threads}, 0, st>>>({}, {}, {}, {}, {y}, {}, {cols}, {}, {xp}, {xo}, {yrow}, {ypitch}, {yp}, {yo}{pre_tail});",
541 b(add.is_some()),
542 b(sum.is_some()),
543 b(*f32),
544 qm(*quant),
545 cols.div_ceil(threads),
546 p(x),
547 n.opt(add),
548 n.opt(sum),
549 n.opt(w),
550 n.quant(out, *quant),
551 flit(eps.get() as f32)
552 )
553 }
554 Launch::Quant { x, n: len } => {
555 format!("lk_quant<{m}><<<{}, 256, 0, st>>>({}, {}, {len});", cdiv(*len, 256), p(x), n.quant(x, true))
556 }
557 Launch::Fold { first, then } if matches!(first.as_ref(), Launch::Gather { .. }) => {
559 let (Launch::Gather { table, index, out: raw, cols, rows, format, scale }, Launch::Norm { w, out, f32, quant, eps, .. }) = (first.as_ref(), then.as_ref()) else {
560 unreachable!("the lane norms a gathered row in its gather alone")
561 };
562 let (y, yrow, ypitch) = if *f32 { n.dst_split(&Dst::to(out), 0) } else { ("nullptr".into(), "nullptr".into(), 0) };
563 let threads = kernels::norm_threads(*cols);
564 let (w8, t) = match format {
565 Some(ty) => (handle(table, *ty, *cols), "nullptr".to_string()),
566 None => ("lk_w{}".to_string(), p(table)),
567 };
568 format!(
569 "lk_gather_norm<{}, {}, {}, {}, {sum_ty}><<<{rows}, {threads}, 0, st>>>({w8}, {t}, {}, {}, {cols}, {}, {}, {y}, {}, {}, {yrow}, {ypitch});",
570 format.map_or(0, kernels::format_arg),
571 b(*f32),
572 qm(*quant),
573 cols.div_ceil(threads),
574 p(index),
575 p(raw),
576 flit(scale.unwrap_or(1.0) as f32),
577 n.opt(w),
578 n.quant(out, *quant),
579 flit(eps.get() as f32)
580 )
581 }
582 Launch::Fold { first, then } if matches!((first.as_ref(), then.as_ref()), (Launch::Gemv { format, .. }, Launch::Norm { .. }) if *format != GgmlType::Q8_0) => {
585 let Launch::Gemv { x: ax, w, format, bias, y: gy, k, n: cols, per_warp, .. } = first.as_ref() else { unreachable!() };
586 let warps = kernels::gemv_blk_norm_threads(*per_warp) / 32;
587 let (norm_t, nm) = norm_of(then, n, m, *cols, sum_ty);
588 let xin = if m == 0 { format!("{}, nullptr, nullptr, nullptr", p(ax)) } else { format!("nullptr, {}", n.quant_at(ax, 0)) };
589 format!(
590 "lk_gemv_blk_norm<{}, {m}, {per_warp}, {norm_t}, {}><<<{}, {}, 0, st>>>({}, {xin}, {}, {}, {k}, {cols}, {nm}, {}, {});",
591 kernels::format_arg(*format),
592 (kernels::norm_threads(*cols) / 32).div_ceil(warps),
593 cdiv(*cols, warps * per_warp),
594 32 * warps,
595 handle(w, *format, *k),
596 n.opt(bias),
597 p(&gy.buf),
598 kernels::norm_threads(*cols),
599 n.scratch("gc", l.writes()[0])
600 )
601 }
602 Launch::Fold { first, then } if matches!(then.as_ref(), Launch::Norm { .. }) => {
603 let (Launch::Gemv { x: ax, w, bias, y: gy, k, n: cols, m: rows, .. }, Launch::Norm { x, add, pre, sum, w: nw, out, into, f32, quant, eps, .. }) =
604 (first.as_ref(), then.as_ref())
605 else {
606 unreachable!("the lane folds a norm into a one-row projection only")
607 };
608 let own = Dst::to(out);
609 let (y, yrow, ypitch) = if *f32 { n.dst_split(into.as_ref().unwrap_or(&own), 0) } else { ("nullptr".into(), "nullptr".into(), 0) };
610 let threads = kernels::norm_threads(*cols);
611 let (pre_arg, pre_tail) = match pre {
612 Some(pn) => ("true", format!(", {}, {}", n.opt(&pn.w), flit(pn.eps.get() as f32))),
613 None => ("false", String::new()),
614 };
615 format!(
616 "lk_gemv_q8_norm<{m}, {rows}, {}, {}, {}, {}, {}, {sum_ty}, {pre_arg}><<<{}, {threads}, 0, st>>>({}, {}, {}, {}, {}, {}, {k}, {cols}, {}, {}, {}, {}, {y}, {yrow}, {ypitch}, {}, {}, {}{pre_tail});",
617 b(add.is_some()),
618 b(sum.is_some()),
619 b(*f32),
620 qm(*quant),
621 cols.div_ceil(threads),
622 cdiv(*cols, threads / 32),
623 rows_q(w),
624 rows_s(w),
625 if m == 0 { p(ax) } else { "nullptr".into() },
626 n.quant(ax, m > 0),
627 n.opt(bias),
628 p(&gy.buf),
629 p(x),
630 n.opt(add),
631 n.opt(sum),
632 n.opt(nw),
633 n.quant(out, *quant),
634 flit(eps.get() as f32),
635 n.scratch("gc", l.writes()[0])
636 )
637 }
638 Launch::Gemv { x, w, bias, y, k, n: cols, m: rows, .. } if l.batch_on_tensor_cores(numerics) => {
641 let part = [PartOut { w: handle(w, GgmlType::Q8_0, *k), dst: n.dst_split(y, 0), bias: n.opt(bias), n: *cols }];
642 gemv_mma(n, x, *k, *rows, &part, m)
643 }
644 Launch::GemmParts { x, k, m: rows, parts, .. } if l.batch_on_tensor_cores(numerics) => {
645 let parts: Vec<PartOut> = parts
646 .iter()
647 .map(|p| PartOut { w: handle(&p.w, GgmlType::Q8_0, *k), dst: n.dst_split(&p.y, 0), bias: n.opt(&p.bias), n: p.n })
648 .collect();
649 gemv_mma(n, x, *k, *rows, &parts, m)
650 }
651 Launch::GemvGlu { x, wg, wu, biases: [bg, bu], act, out, f32, quant, k, n: cols, m: rows, .. } if l.batch_on_tensor_cores(numerics) => {
652 let (hg, hu) = (n.scratch("mg", out), n.scratch("mu", out));
653 let scratch = |w: &str, s: &str, bias: &Option<String>| PartOut {
654 w: handle(w, GgmlType::Q8_0, *k),
655 dst: (s.to_string(), "nullptr".to_string(), 0),
656 bias: n.opt(bias),
657 n: *cols,
658 };
659 let parts = [scratch(wg, &hg, bg), scratch(wu, &hu, bu)];
660 let len = rows * cols;
661 let gate = match act {
662 None => format!(
663 "lk_swiglu<{}, {}><<<{}, 256, 0, st>>>({hg}, {hu}, {}, {}, {len});",
664 b(*f32),
665 qm(*quant),
666 cdiv(len, 256),
667 if *f32 { p(out) } else { "nullptr".into() },
668 n.quant(out, *quant)
669 ),
670 Some(a) => format!(
671 "lk_binary<1, {}, {}, {}><<<{}, 256, 0, st>>>({hu}, {hg}, {}, {}, {len}, {len});",
672 act_arg(*a),
673 b(*f32),
674 qm(*quant),
675 cdiv(len, 256),
676 if *f32 { p(out) } else { "nullptr".into() },
677 n.quant(out, *quant)
678 ),
679 };
680 format!("{}\n CKL();\n {gate}", gemv_mma(n, x, *k, *rows, &parts, m))
681 }
682 Launch::Gemv { x, w, format, bias, y, k, n: cols, m: rows, tail, .. } if l.on_tensor_cores(numerics) => {
683 let bm = l.gemm_rows(numerics);
684 let tiles = cdiv(*cols, kernels::GEMM_TILE_N) * cdiv(*rows, bm);
685 let (grid, line) = match tail {
687 Some((whole, split)) => (
688 format!("{}", whole + (tiles - whole) * split),
689 format!(
690 ", nullptr, nullptr, nullptr, nullptr, {{}}, nullptr, nullptr, nullptr, nullptr, 0.f, 0.f, {whole}, {split}, {}, {}",
691 n.scratch("tp", &y.buf),
692 n.scratch("tc", &y.buf)
693 ),
694 ),
695 None => (format!("dim3({}, {})", cdiv(*cols, kernels::GEMM_TILE_N), cdiv(*rows, bm)), String::new()),
696 };
697 format!(
698 "{}<<<{grid}, {}, {}, st>>>({}, {}, {}, {}, {k}, {cols}, {rows}{line});",
699 gemm_blk_name(kernels::format_arg(*format), m, 0, bm, 0),
700 kernels::gemm_blk_launch_threads(kernels::format_arg(*format), bm, *k),
701 kernels::gemm_blk_launch_smem(kernels::format_arg(*format), m == 2, bm, *k),
702 handle(w, *format, *k),
703 n.quant(x, true),
704 n.opt(bias),
705 n.dst(y)
706 )
707 }
708 Launch::GemmParts { x, format, k, m: rows, parts, tail, .. } if l.on_tensor_cores(numerics) => {
709 let bm = l.gemm_rows(numerics);
710 let cols = l.gemm_cols();
711 let tiles = cdiv(cols, kernels::GEMM_TILE_N) * cdiv(*rows, bm);
712 let key = &parts[0].y.buf;
713 let (grid, line) = match tail {
714 Some((whole, split)) => (format!("{}", whole + (tiles - whole) * split), format!("{whole}, {split}, {}, {}", n.scratch("tp", key), n.scratch("tc", key))),
715 None => (format!("dim3({}, {})", cdiv(cols, kernels::GEMM_TILE_N), cdiv(*rows, bm)), "0, 1, nullptr, nullptr".to_string()),
716 };
717 let outs: Vec<PartOut> = parts
718 .iter()
719 .map(|p| PartOut { w: handle(&p.w, *format, *k), dst: n.dst_split(&p.y, 0), bias: n.opt(&p.bias), n: p.n })
720 .collect();
721 format!(
722 "{}<<<{grid}, {}, {}, st>>>({{}}, {}, nullptr, nullptr, nullptr, 0, {k}, {cols}, {rows}, nullptr, nullptr, nullptr, nullptr, {{}}, nullptr, nullptr, nullptr, nullptr, 0.f, 0.f, {line}, {});",
723 gemm_blk_name(kernels::format_arg(*format), m, 4, bm, 0),
724 kernels::gemm_blk_launch_threads(kernels::format_arg(*format), bm, *k),
725 kernels::gemm_blk_launch_smem(kernels::format_arg(*format), m == 2, bm, *k),
726 n.quant(x, true),
727 lk_parts(&outs)
728 )
729 }
730 Launch::GemmParts { x, format, k, m: rows, parts, per_warp, .. } => {
731 let cols = l.gemm_cols();
732 (0..rows.div_ceil(GEMV_ROWS))
734 .map(|slab| {
735 let r = GEMV_ROWS.min(rows - slab * GEMV_ROWS);
736 let at = slab * GEMV_ROWS * k;
737 let outs: Vec<PartOut> = parts
738 .iter()
739 .map(|p| PartOut {
740 w: handle(&p.w, *format, *k),
741 dst: n.dst_split(&p.y, slab * GEMV_ROWS * p.n),
742 bias: n.opt(&p.bias),
743 n: p.n,
744 })
745 .collect();
746 let xin = if m == 0 { format!("{} + {at}, nullptr, nullptr, nullptr", p(x)) } else { format!("nullptr, {}", n.quant_at(x, at)) };
747 format!(
749 "lk_gemv_parts<{}, {m}, {r}{}><<<{}, 256, 0, st>>>({xin}, {k}, {cols}, {});",
750 kernels::format_arg(*format),
751 per_warp_arg(*per_warp),
752 cdiv(cols, 8 * *per_warp),
753 lk_parts(&outs)
754 )
755 })
756 .collect::<Vec<_>>()
757 .join("\n CKL();\n ")
758 }
759 Launch::Gemv { x, w, format, bias, y, k, n: cols, m: rows, per_warp, f32, quant, .. } if *format != GgmlType::Q8_0 => (0..rows.div_ceil(GEMV_ROWS))
760 .map(|slab| {
761 let r = GEMV_ROWS.min(rows - slab * GEMV_ROWS);
762 let (at, oat) = (slab * GEMV_ROWS * k, slab * GEMV_ROWS * cols);
763 let xin = if m == 0 { format!("{} + {at}, nullptr, nullptr, nullptr", p(x)) } else { format!("nullptr, {}", n.quant_at(x, at)) };
764 let rows_a_warp = *per_warp;
765 let warps = match (*quant, rows_a_warp) {
767 (true, w) => 32 / w,
768 (false, 1) => 8,
769 (false, _) => 4,
770 };
771 let dst = if *f32 { n.dst_at(y, oat) } else { "nullptr, nullptr, 0".into() };
772 let yq = if *quant { n.quant_at(&y.buf, oat) } else { "nullptr, nullptr, nullptr".into() };
773 format!(
774 "lk_gemv_blk<{}, {m}, {r}{}><<<{}, {}, 0, st>>>({}, {xin}, {}, {dst}, {k}, {cols}, {yq});",
775 kernels::format_arg(*format),
776 match (rows_a_warp, *quant) {
777 (1, false) => String::new(),
778 (w, q) => format!(", {w}, {}", qm(q)),
779 },
780 cdiv(*cols, warps * rows_a_warp),
781 32 * warps,
782 handle(w, *format, *k),
783 n.opt(bias)
784 )
785 })
786 .collect::<Vec<_>>()
787 .join("\n CKL();\n "),
788 Launch::Gemv { x, w, bias, y, k, n: cols, m: rows, .. } if *rows > GEMV_ROWS => (0..rows.div_ceil(GEMV_ROWS))
789 .map(|slab| {
790 let r = GEMV_ROWS.min(rows - slab * GEMV_ROWS);
791 format!(
792 "lk_gemv_q8<0, {r}><<<{}, 256, 0, st>>>({}, {}, {} + {}, nullptr, nullptr, nullptr, {}, {}, {k}, {cols});",
793 cdiv(*cols, 8),
794 rows_q(w),
795 rows_s(w),
796 p(x),
797 slab * GEMV_ROWS * k,
798 n.opt(bias),
799 n.dst_at(y, slab * GEMV_ROWS * cols)
800 )
801 })
802 .collect::<Vec<_>>()
803 .join("\n CKL();\n "),
804 Launch::Gemv { x, w, bias, y, k, n: cols, m: rows, .. } => format!(
805 "lk_gemv_q8<{m}, {rows}><<<{}, 256, 0, st>>>({}, {}, {}, {}, {}, {}, {k}, {cols});",
806 cdiv(*cols, 8),
807 rows_q(w),
808 rows_s(w),
809 if m == 0 { p(x) } else { "nullptr".into() },
810 n.quant(x, m > 0),
811 n.opt(bias),
812 n.dst(y)
813 ),
814 Launch::GemvF32 { x, w, bias, y, k, n: cols, m: rows, split, .. } if l.multiplies_in_tiles(numerics) => format!(
815 "lk_gemm_f32<<<dim3({}, {}, {split}), 256, 0, st>>>({}, {}, {}, {}, {rows}, {cols}, {k}{});",
816 cdiv(*cols, 64),
817 cdiv(*rows, 64),
818 p(x),
819 p(w),
820 n.opt(bias),
821 n.dst_ptr(y),
822 if *split > 1 { format!(", {split}, {}, {}", n.scratch("fp", &y.buf), n.scratch("fq", &y.buf)) } else { String::new() }
823 ),
824 Launch::GemvF32 { x, w, bias, y, k, n: cols, m: rows, split, .. } => format!(
825 "lk_gemv_f32<<<dim3({}, {}), 256, 0, st>>>({}, {}, {}, {}, {k}, {cols}, {rows}{});",
826 cdiv(*cols, 8 / split),
827 cdiv(*rows, kernels::GEMV_F32_ROWS),
828 p(w),
829 p(x),
830 n.opt(bias),
831 n.dst(y),
832 if *split > 1 { format!(", {split}") } else { String::new() }
833 ),
834 Launch::Binary { op, a, b: bb, gate, row, out, f32, quant, n: len } => format!(
835 "lk_binary<{}, {}, {}, {}><<<{}, 256, 0, st>>>({}, {}, {}, {}, {len}, {});",
836 match op {
837 ElementwiseOp::Add => 0,
838 ElementwiseOp::Mul => 1,
839 },
840 gate.map_or(-1, act_arg),
841 b(*f32),
842 qm(*quant),
843 cdiv(*len, 256),
844 p(a),
845 p(bb),
846 if *f32 { p(out) } else { "nullptr".into() },
847 n.quant(out, *quant),
848 row.unwrap_or(*len)
849 ),
850 Launch::SwiGlu { g, u, out, f32, quant, n: len } => format!(
851 "lk_swiglu<{}, {}><<<{}, 256, 0, st>>>({}, {}, {}, {}, {len});",
852 b(*f32),
853 qm(*quant),
854 cdiv(*len, 256),
855 p(g),
856 p(u),
857 if *f32 { p(out) } else { "nullptr".into() },
858 n.quant(out, *quant)
859 ),
860 Launch::GemvGlu { x, wg, wu, formats: [fg, _], biases: [bg, bu], act, out, f32, quant, k, n: cols, m: rows, tail, .. } if l.on_tensor_cores(numerics) => {
862 let bm = l.gemm_rows(numerics);
863 let tw = kernels::GEMM_TILE_N / 2;
864 let tiles = cdiv(*cols, tw) * cdiv(*rows, bm);
865 let (grid, line) = match tail {
866 Some((whole, split)) => (format!("{}", whole + (tiles - whole) * split), format!("{whole}, {split}, {}, {}", n.scratch("tp", out), n.scratch("tc", out))),
867 None => (format!("dim3({}, {})", cdiv(*cols, tw), cdiv(*rows, bm)), "0, 1, nullptr, nullptr".to_string()),
868 };
869 format!(
870 "{}<<<{grid}, {}, {}, st>>>({}, {}, {}, {}, nullptr, 0, {k}, {cols}, {rows}, nullptr, nullptr, nullptr, nullptr, {}, {}, {}, 0.f, 0.f, {line});",
871 gemm_blk_name(kernels::format_arg(*fg), m, 5, bm, act.map_or(-1, act_arg)),
872 kernels::gemm_blk_launch_threads(kernels::format_arg(*fg), bm, *k),
873 kernels::gemm_blk_launch_smem(kernels::format_arg(*fg), m == 2, bm, *k),
874 handle(wg, *fg, *k),
875 n.quant(x, true),
876 n.opt(bg),
877 if *f32 { p(out) } else { "nullptr".into() },
878 handle(wu, *fg, *k),
879 n.opt(bu),
880 if *quant { n.quant(out, true) } else { "nullptr, nullptr, nullptr".into() }
881 )
882 }
883 Launch::GemvGlu { x, wg, wu, formats: [fg, fu], biases: [bg, bu], act, out, f32, quant, k, n: cols, m: rows, per_warp, .. } => (0..rows.div_ceil(GEMV_ROWS))
884 .map(|slab| {
885 let r = GEMV_ROWS.min(rows - slab * GEMV_ROWS);
886 let (at, oat) = (slab * GEMV_ROWS * k, slab * GEMV_ROWS * cols);
887 let xin = if m == 0 { format!("{} + {at}, nullptr, nullptr, nullptr", p(x)) } else { format!("nullptr, {}", n.quant_at(x, at)) };
888 let y = if *f32 { format!("{} + {oat}", p(out)) } else { "nullptr".into() };
889 let yq = if *quant { n.quant_at(out, oat) } else { "nullptr, nullptr, nullptr".into() };
890 let rows_a_warp = *per_warp;
891 let warps = if *quant { 32 / rows_a_warp } else { 8 };
893 format!(
894 "lk_gemv_glu<{}, {}, {m}, {r}, {}, {}, {}{}><<<{}, {}, 0, st>>>({}, {}, {xin}, {}, {}, {y}, {yq}, {k}, {cols});",
895 kernels::format_arg(*fg),
896 kernels::format_arg(*fu),
897 act.map_or(-1, act_arg),
898 b(*f32),
899 qm(*quant),
900 per_warp_arg(rows_a_warp),
901 cdiv(*cols, warps * rows_a_warp),
902 32 * warps,
903 handle(wg, *fg, *k),
904 handle(wu, *fu, *k),
905 n.opt(bg),
906 n.opt(bu)
907 )
908 })
909 .collect::<Vec<_>>()
910 .join("\n CKL();\n "),
911 Launch::Unary { act, x, y, n: len } => format!(
912 "lk_unary<{}><<<{}, 256, 0, st>>>({}, {}, {len});",
913 act_arg(*act),
914 cdiv(*len, 256),
915 p(x),
916 p(y)
917 ),
918 Launch::Rope { x, part, norm, cos, sin, pos, y, heads, per_row, hd, half } => {
919 let ((xp, xo), (yp, yo)) = (reading(part, *hd), unit(y, *hd));
920 let threads = rope_threads(norm.is_some(), *hd);
921 format!(
922 "lk_rope<{}, {sum_ty}><<<{heads}, {threads}, 0, st>>>({}, {}, {}, {}, {}, {}, {hd}, {half}, {per_row}, {}, {xp}, {xo}, {yp}, {yo});",
923 b(norm.is_some()),
924 p(x),
925 norm.as_ref().map_or_else(|| "nullptr".to_string(), |(w, _)| p(w)),
926 p(cos),
927 p(sin),
928 n.opt(pos),
929 n.dst(y),
930 flit(norm.as_ref().map_or(0.0, |(_, e)| e.get() as f32))
931 )
932 }
933 Launch::RopeParts { sides, cos, sin, pos, hd, half, seqs } => {
934 let normed = sides[0].norm.is_some();
935 format!(
936 "lk_rope_parts<{}, {sum_ty}><<<{}, {}, 0, st>>>({}, {}, {}, {}, {}, {hd}, {half}, {});",
937 b(normed),
938 sides.iter().map(|r| r.heads).sum::<usize>(),
939 rope_threads(normed, *hd),
940 rope_side(&sides[0], *hd, n),
941 rope_side(&sides[1], *hd, n),
942 p(cos),
943 p(sin),
944 n.opt(pos),
945 u8::from(*seqs)
946 )
947 }
948 Launch::Fold { first, then } => {
952 let (Launch::GemmParts { x, format, k, parts, per_warp, .. }, Launch::RopeParts { sides, cos, sin, pos, hd, half, .. }) = (first.as_ref(), then.as_ref()) else {
953 unreachable!("the lane turns heads only in a GEMV in parts")
954 };
955 let normed = sides[0].norm.is_some();
956 let threads = rope_threads(normed, *hd).max(256);
957 let outs: Vec<PartOut> = parts.iter().map(|q| PartOut { w: handle(&q.w, *format, *k), dst: n.dst_split(&q.y, 0), bias: n.opt(&q.bias), n: q.n }).collect();
958 let read = |side: Option<&RopeSide>| side.map_or(-1, |r| parts.iter().position(|q| q.y.own() == Some(r.x.as_str())).expect("a side reads a part") as i64);
959 let none = RopeSide { x: sides[0].x.clone(), part: None, norm: sides[0].norm.clone(), y: sides[0].y.clone(), heads: 0, per_row: 1, yz: 0 };
960 let xin = if m == 0 { format!("{}, nullptr, nullptr, nullptr", p(x)) } else { format!("nullptr, {}", n.quant_at(x, 0)) };
961 format!(
962 "lk_gemv_parts_rope<{}, {m}, 1, {per_warp}, {}, {sum_ty}><<<{}, {threads}, 0, st>>>({xin}, {k}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {hd}, {half}, {});",
963 kernels::format_arg(*format),
964 b(normed),
965 cdiv(first.gemm_cols(), threads / 32 * per_warp),
966 first.gemm_cols(),
967 lk_parts(&outs),
968 rope_side(&sides[0], *hd, n),
969 rope_side(sides.get(1).unwrap_or(&none), *hd, n),
970 read(sides.first()),
971 read(sides.get(1)),
972 p(cos),
973 p(sin),
974 n.opt(pos),
975 n.scratch("gc", l.writes()[0])
976 )
977 }
978 Launch::SetRow { table, index, src, cols, count, seq, f32, quant } => {
979 let (idx, fixed) = match index {
980 RowSel::Param(i) => (p(i), 0),
981 RowSel::Fixed(r) => ("nullptr".to_string(), *r),
982 };
983 let (blocks, n_set, tz) = (cdiv(count * cols, 256), count * cols, seq.unwrap_or(0));
984 let mut sets = Vec::new();
985 if *f32 {
986 sets.push(format!("lk_set_row<<<{blocks}, 256, 0, st>>>({}, {idx}, {fixed}, {}, {n_set}, {cols}, {tz});", p(table), p(src)));
987 }
988 if *quant {
989 sets.push(format!(
990 "lk_set_row_q<<<{blocks}, 256, 0, st>>>({}, {idx}, {fixed}, {}, {n_set}, {cols}, {tz});",
991 n.quant(table, true),
992 n.quant(src, true)
993 ));
994 }
995 sets.join("\n CKL();\n ")
996 }
997 Launch::Attention { q, kc, vc, pos, sinks, out, heads, group, hd, vd, len, rows, window, vstride, voff, scale, split, .. } if l.attends_in_tiles(numerics) => {
998 let hg = kernels::attn_tile_heads(*group);
999 let (mma, flash) = (l.attends_on_tensor_cores(numerics), l.attends_flash(numerics));
1000 let kv = n.elems(&[Some(kc), Some(vc), split.as_ref().map(|s| s.k_tail.0.as_str())]);
1001 let parts = split.as_ref().map_or(String::new(), |s| {
1003 format!(
1004 ", {}, {}, {}, {}, {}, {}, {}",
1005 p(&s.q_tail.0),
1006 s.q_pitch,
1007 s.q_tail.1.pitch,
1008 p(&s.k_tail.0),
1009 s.k_tail.1.pitch,
1010 s.k_tail.1.offset,
1011 s.own
1012 )
1013 });
1014 if l.attends_in_two_words(numerics) {
1015 let (lpad, kvh) = (kernels::attn_i8_rows(*len), heads / group);
1017 let (k8, w8, s8) = (n.scratch("ak", out), n.scratch("aw", out), n.scratch("as", out));
1018 let (kparts, qparts) = split.as_ref().map_or((String::new(), String::new()), |s| {
1019 (
1020 format!(", {}, {}, {}, {}", p(&s.k_tail.0), s.k_tail.1.pitch, s.k_tail.1.offset, s.own),
1021 format!(", {}, {}, {}, {}", p(&s.q_tail.0), s.q_pitch, s.q_tail.1.pitch, s.own),
1022 )
1023 });
1024 format!(
1025 "lk_kv_two_words<{hd}, {vd}{kv}><<<dim3({}, {kvh}), 256, 0, st>>>({}, {}, {}, {rows}, {len}, {}, {vstride}, {voff}, {lpad}, \
1026 {k8}, {k8} + {kb}, {s8}, {w8}, {w8} + {vb}, {s8} + {sb}{kparts});\n CKL();\n \
1027 lk_attention_i8<{hg}, {hd}, {vd}><<<dim3({}, {}), 128, {}, st>>>({}, {k8}, {k8} + {kb}, {s8}, {w8}, {w8} + {vb}, {s8} + {sb}, {lpad}, \
1028 {}, {}, {}, {len}, {}, {group}, {heads}, {}, {rows}{qparts});",
1029 lpad / 32,
1030 p(kc),
1031 p(vc),
1032 n.opt(pos),
1033 heads / group * split.as_ref().map_or(*hd, |s| s.own),
1034 rows.div_ceil(kernels::ATTN_I8_QUERIES / hg),
1035 heads / hg,
1036 kernels::attn_i8_smem_bytes(*hd, *vd),
1037 p(q),
1038 n.opt(pos),
1039 n.opt(sinks),
1040 p(out),
1041 window.unwrap_or(0),
1042 flit(*scale as f32),
1043 kb = kvh * lpad * hd,
1044 vb = kvh * lpad * vd,
1045 sb = kvh * lpad,
1046 )
1047 } else if mma {
1048 format!(
1049 "lk_attention_mma<{hg}, {hd}, {vd}{kv}><<<dim3({}, {}), 128, {}, st>>>({}, {}, {}, {}, {}, {}, {len}, {}, {group}, {heads}, {}, {vstride}, {voff}, {}, {rows}{parts});",
1050 rows.div_ceil(kernels::ATTN_MMA_QUERIES / hg),
1051 heads / hg,
1052 kernels::attn_mma_smem_bytes(*hd, *vd),
1053 p(q),
1054 p(kc),
1055 p(vc),
1056 n.opt(pos),
1057 n.opt(sinks),
1058 p(out),
1059 window.unwrap_or(0),
1060 heads / group * split.as_ref().map_or(*hd, |s| s.own),
1061 flit(*scale as f32)
1062 )
1063 } else if flash {
1064 let (bq, nt) = (kernels::attn_flash_queries(*hd, *vd), kernels::attn_flash_threads(*hd, *vd));
1065 format!(
1066 "lk_attention_flash<{hg}, {hd}, {vd}, {bq}, {nt}{kv}><<<dim3({}, {}), {nt}, {}, st>>>({}, {}, {}, {}, {}, {}, {len}, {}, {group}, {heads}, {}, {vstride}, {voff}, {}, {rows}{parts});",
1067 rows.div_ceil(bq / hg),
1068 heads / hg,
1069 kernels::attn_flash_smem_bytes(*hd, *vd),
1070 p(q),
1071 p(kc),
1072 p(vc),
1073 n.opt(pos),
1074 n.opt(sinks),
1075 p(out),
1076 window.unwrap_or(0),
1077 heads / group * split.as_ref().map_or(*hd, |s| s.own),
1078 flit(*scale as f32)
1079 )
1080 } else {
1081 format!(
1082 "lk_attention_tiled<{hg}, {}{kv}><<<dim3({}, {}), 256, {}, st>>>({}, {}, {}, {}, {}, {}, {len}, {}, {hd}, {vd}, {group}, {heads}, {}, {vstride}, {voff}, {}, {rows}{parts});",
1083 kernels::attn_tile_values(*vd),
1084 rows.div_ceil(kernels::ATTN_TILE_Q / hg),
1085 heads / hg,
1086 kernels::attn_tile_smem_bytes(*vd),
1087 p(q),
1088 p(kc),
1089 p(vc),
1090 n.opt(pos),
1091 n.opt(sinks),
1092 p(out),
1093 window.unwrap_or(0),
1094 heads / group * split.as_ref().map_or(*hd, |s| s.own),
1095 flit(*scale as f32)
1096 )
1097 }
1098 }
1099 Launch::Attention { q, kc, vc, pos, sinks, out, heads, group, hd, vd, len, rows, window, vstride, voff, scale, shared: Some(splits), seq, f32, quant, gate, goff, .. } => {
1101 let part = if *splits > 1 { n.scratch("ap", out) } else { "nullptr".into() };
1102 let (win, minc) = (window.unwrap_or(0), kernels::ATTN_MIN_SPLIT);
1103 let (qo, qargs) = attention_quant(n, out, *quant, gate, *goff, m);
1104 let kernel = attention_shared_kernel(*heads, *vd, *hd, kc == vc, &n.elems(&[Some(kc), Some(vc)]), &qo);
1105 let dst = if *f32 { p(out) } else { "nullptr".into() };
1106 let attend = format!(
1107 "{kernel}<<<dim3({}, {rows}, {splits}), 256, {}, st>>>({}, {}, {}, {}, {}, {dst}, {part}, {len}, {win}, {vd}, {heads}, {}, {vstride}, {voff}, {}, {minc}, {splits}, {}{qargs});",
1108 heads / group,
1109 kernels::attn_shared_smem_bytes(*heads, *hd),
1110 p(q),
1111 p(kc),
1112 p(vc),
1113 n.opt(pos),
1114 n.opt(sinks),
1115 heads / group * hd,
1116 flit(*scale as f32),
1117 seq.unwrap_or(0)
1118 );
1119 if *splits == 1 {
1120 attend
1121 } else {
1122 let fold = if qo.is_empty() { "lk_attention_fold".to_string() } else { format!("lk_attention_fold<{}>", if *quant { m } else { 0 }) };
1123 format!(
1124 "{attend}\n CKL();\n {fold}<<<dim3({heads}, {rows}, {}), {}, {}, st>>>({part}, {}, {}, {dst}, {len}, {win}, {vd}, {heads}, {minc}, {splits}, {}{qargs});",
1125 cdiv(*vd, 2 * kernels::ATTN_FOLD_PAIRS),
1126 kernels::ATTN_FOLD_PAIRS,
1127 splits * 4,
1128 n.opt(pos),
1129 n.opt(sinks),
1130 u8::from(seq.is_some())
1131 )
1132 }
1133 }
1134 Launch::Attention { q, kc, vc, pos, sinks, out, heads, group, hd, vd, len, rows, window, vstride, voff, scale, per_block, splits, parallel, split, shared: None, seq, f32, quant, gate, goff } => {
1135 assert!(split.is_none(), "a split attention attends in tiles");
1136 let (part, cnt) = if *parallel { (n.scratch("ap", out), n.scratch("ac", out)) } else { ("nullptr".into(), "nullptr".into()) };
1137 let (acc, ty) = kernels::accumulator(m);
1138 let (smem, side) = attention_smem(*per_block, *hd, *vd, *splits, *parallel, acc);
1139 let z = if *parallel { *splits } else { 1 };
1140 let grid = heads / per_block * rows * z;
1141 let (qo, qargs) = attention_quant(n, out, *quant, gate, *goff, m);
1142 let mb = kernels::attn_blocks_a_multiprocessor(grid, n.sms, acc);
1143 let kernel = attention_kernel(*per_block, *vd, ty, mb, &n.elems(&[Some(kc), Some(vc)]), &qo, kernels::attn_path(z, side, *splits), *hd);
1144 let dst = if *f32 { p(out) } else { "nullptr".into() };
1145 format!(
1146 "{kernel}<<<dim3({}, {rows}, {z}), 256, {smem}, st>>>({}, {}, {}, {}, {}, {dst}, {part}, {cnt}, {len}, {}, {hd}, {vd}, {group}, {heads}, {}, {vstride}, {voff}, {}, {}, {splits}, {}, {}{qargs});",
1147 heads / per_block,
1148 p(q),
1149 p(kc),
1150 p(vc),
1151 n.opt(pos),
1152 n.opt(sinks),
1153 window.unwrap_or(0),
1154 heads / group * hd,
1155 flit(*scale as f32),
1156 kernels::ATTN_MIN_SPLIT,
1157 seq.unwrap_or(0),
1158 side
1159 )
1160 }
1161 Launch::Slice { x, y, groups, every, from, width } => format!(
1162 "lk_slice<<<{}, 256, 0, st>>>({}, {}, {}, {every}, {from}, {width});",
1163 cdiv(groups * width, 256),
1164 p(x),
1165 p(y),
1166 groups * width
1167 ),
1168 Launch::Join { a, b: bb, y, rows, heads, ha, hb, yz } => format!(
1169 "lk_join<<<{}, 256, 0, st>>>({}, {}, {}, {}, {ha}, {hb}{});",
1170 cdiv(rows * heads * (ha + hb), 256),
1171 p(a),
1172 p(bb),
1173 n.dst(y),
1174 rows * heads * (ha + hb),
1175 if *yz > 0 { format!(", {heads}, {yz}") } else { String::new() }
1176 ),
1177 Launch::HeadGemv { x, part, w, format, y, m: rows, heads, a, b: bw, rows_out, shared, upto, .. } if l.heads_in_tiles(numerics) => {
1179 let ((xp, xo), (yp, yo)) = (reading(part, *a), unit(y, *bw));
1180 let (xp, xr) = if *shared { (0, xp) } else { (xp, 0) };
1182 let bound = upto.as_ref().map_or(String::new(), |(at, after)| format!(", {}, {after}", p(at)));
1183 let (yat, yrow, ypitch) = n.dst_split(y, 0);
1184 format!(
1185 "lk_head_tiled<{}, {}, {rows_out}><<<dim3({}, {}, {heads}), 256, 0, st>>>({}, nullptr, nullptr, nullptr, {yat}, {heads}, {a}, {bw}, {xp}, {xo}, {yrow}, {ypitch}, {yp}, {yo}, {rows}, {}, {xr}{bound});",
1186 kernels::format_arg(*format),
1187 kernels::head_mode(numerics),
1188 cdiv(*bw, 64),
1189 cdiv(*rows, kernels::HEAD_TILE_M),
1190 handle(w, *format, if *rows_out { *a } else { *bw }),
1191 p(x)
1192 )
1193 }
1194 Launch::HeadGemv { shared: true, .. } | Launch::HeadGemv { upto: Some(_), .. } => {
1196 unreachable!("a per-head matmul over every head's one unit or bounded rows runs in tiles")
1197 }
1198 Launch::HeadGemv { x, part, w, format, y, m: rows, heads, a, b: bw, rows_out, .. } => (0..rows.div_ceil(GEMV_ROWS))
1199 .map(|slab| {
1200 let (r, m) = (GEMV_ROWS.min(rows - slab * GEMV_ROWS), kernels::head_mode(numerics));
1201 let ((xp, xo), (yp, yo)) = (reading(part, *a), unit(y, *bw));
1202 let at = slab * GEMV_ROWS * heads * xp;
1203 let xin = format!("{} + {at}, nullptr, nullptr, nullptr", p(x));
1204 let (yat, yrow, ypitch) = n.dst_split(y, slab * GEMV_ROWS * heads * yp);
1205 let f = kernels::format_arg(*format);
1206 let shape = format!("{heads}, {a}, {bw}, {xp}, {xo}, {yrow}, {ypitch}, {yp}, {yo}");
1207 if *rows_out {
1208 format!(
1209 "lk_head_gemv<{f}, {m}, {r}><<<{}, 256, 0, st>>>({}, {xin}, {yat}, {shape});",
1210 cdiv(heads * bw, 8),
1211 handle(w, *format, *a)
1212 )
1213 } else if kernels::axpy_staged(*format, m) {
1214 let pitch = kernels::axpy_pitch(*format);
1217 let tile = (*a).min((48 * 1024 - kernels::AXPY_SUMS) / pitch);
1218 format!(
1219 "lk_head_axpy_staged<{f}, {m}, {r}, {}, {}><<<dim3({}, {heads}), {}, {}, st>>>({}, {xin}, {yat}, {shape}, {tile});",
1220 kernels::AXPY_W,
1221 kernels::AXPY_R,
1222 cdiv(bw / 32, kernels::AXPY_W),
1223 32 * kernels::AXPY_W * kernels::AXPY_R,
1224 tile * pitch,
1225 handle(w, *format, *bw)
1226 )
1227 } else {
1228 format!(
1229 "lk_head_axpy<{f}, {m}, {r}><<<dim3({}, {heads}), 256, 0, st>>>({}, {xin}, {yat}, {shape});",
1230 bw / 32,
1231 handle(w, *format, *bw)
1232 )
1233 }
1234 })
1235 .collect::<Vec<_>>()
1236 .join("\n CKL();\n "),
1237 Launch::ArgMax { x, out, n: len, advance, rows } => format!(
1238 "lk_argmax<{}><<<dim3({}, {rows}), 256, 0, st>>>({}, {len}, {}, {}, {}, {}, {}, {}, {});",
1239 b(advance.is_some()),
1240 kernels::argmax_blocks(*len),
1241 p(x),
1242 p(out),
1243 match advance {
1244 Some(a) if n.batch => format!("{} + {1}, d_bpos + {1}, 1, d_bgen + {1} * CTX, d_bcount + {1}", n.plain(&a.token), a.item),
1246 Some(_) => format!("{token}, d_pos, {}, d_gen, d_count", n.rows),
1247 None => "nullptr, nullptr, 0, nullptr, nullptr".into(),
1248 },
1249 n.scratch("av", out),
1250 n.scratch("ai", out),
1251 n.scratch("ac", out),
1252 u8::from(*rows > 1),
1253 if *rows > 1 { "CTX" } else { "0" }
1254 ),
1255 Launch::Conv { x, w, ring, y, chan, taps, rows, silu, l2 } if *rows > 1 => {
1257 let (grid, block, nl2, eps) = match l2 {
1258 Some((d, e)) => (chan / d, *d, chan / d, e.get() as f32),
1259 None => (cdiv(*chan, 256), 256, 0, 0.0),
1260 };
1261 format!(
1262 "lk_conv_rows<{}, {}, {sum_ty}><<<dim3({grid}, {rows}), {block}, 0, st>>>({}, {}, {}, {}, {chan}, {taps}, {rows}, {nl2}, {}, {pos});\n CKL();\n \
1263 lk_conv_carry<<<{}, 256, 0, st>>>({}, {}, {chan}, {taps}, {rows}, {pos});",
1264 b(*silu),
1265 b(l2.is_some()),
1266 p(x),
1267 p(w),
1268 p(ring),
1269 p(y),
1270 flit(eps),
1271 cdiv(*chan, 256),
1272 p(x),
1273 p(ring),
1274 pos = n.conv_pos(ring, taps * chan, 1)
1275 )
1276 }
1277 Launch::Conv { x, w, ring, y, chan, taps, rows, silu, l2 } => {
1278 let (grid, block, nl2, eps) = match l2 {
1279 Some((d, e)) => (chan / d, *d, chan / d, e.get() as f32),
1280 None => (cdiv(*chan, 256), 256, 0, 0.0),
1281 };
1282 format!(
1283 "lk_conv<{}, {}, {sum_ty}><<<{grid}, {block}, 0, st>>>({}, {}, {}, {}, {chan}, {taps}, {rows}, {nl2}, {}, {});",
1284 b(*silu),
1285 b(l2.is_some()),
1286 p(x),
1287 p(w),
1288 p(ring),
1289 p(y),
1290 flit(eps),
1291 n.conv_pos(ring, taps * chan, 1)
1292 )
1293 }
1294 Launch::ConvParts { parts, taps, block, eps, seqs } => {
1295 let entries: Vec<String> = parts
1296 .iter()
1297 .map(|q| format!("{{{}, {}, {}, {}, {}, {}, {}}}", p(&q.x), p(&q.w), p(&q.ring), p(&q.y), q.chan, cdiv(q.chan, *block), u8::from(q.l2)))
1298 .collect();
1299 let blocks: usize = parts.iter().map(|q| cdiv(q.chan, *block)).sum();
1300 format!(
1301 "lk_conv_parts<{sum_ty}><<<dim3({blocks}, {seqs}), {block}, 0, st>>>(lk_convs{{{{{}}}}}, {}, {taps}, {}, {}, {});",
1302 entries.join(", "),
1303 parts.len(),
1304 flit(eps.map_or(0.0, |e| e.get() as f32)),
1305 n.conv_pos(&parts[0].ring, taps * parts[0].chan, *seqs),
1306 u8::from(*seqs > 1)
1307 )
1308 }
1309 Launch::L2Norm { x, y, rows, cols, eps } => {
1310 format!("lk_l2norm<{sum_ty}><<<{rows}, 256, 0, st>>>({}, {}, {cols}, {});", p(x), p(y), flit(eps.get() as f32))
1311 }
1312 Launch::Decay { ar, dt, al, y, n: len, heads } => format!(
1313 "lk_decay<<<{}, 256, 0, st>>>({}, {}, {}, {}, {len}, {heads});",
1314 cdiv(*len, 256),
1315 p(ar),
1316 p(dt),
1317 p(al),
1318 p(y)
1319 ),
1320 Launch::GateProj { x, wb, wa, dt, al, beta, decay, k, heads, parts } => format!(
1321 "lk_gate_proj<<<{}, 256, 0, st>>>({}, {}, {}, {}, {}, {}, {}, {k}, {heads}, {parts});",
1322 cdiv(*heads, 8 / parts),
1323 p(wb),
1324 p(wa),
1325 p(x),
1326 p(dt),
1327 p(al),
1328 p(beta),
1329 p(decay)
1330 ),
1331 Launch::Gates { br, ar, dt, al, beta, decay, n: len, heads } => format!(
1332 "lk_gates<<<{}, 256, 0, st>>>({}, {}, {}, {}, {}, {}, {len}, {heads});",
1333 cdiv(*len, 256),
1334 p(br),
1335 p(ar),
1336 p(dt),
1337 p(al),
1338 p(beta),
1339 p(decay)
1340 ),
1341 Launch::Delta { q, k, v, z, b: bb, a, st, gn, out, f32, quant, hv, hk, d, rows, eps, .. } if l.delta_in_chunks(numerics) => {
1343 let chunks = rows.div_ceil(kernels::DELTA_CHUNK);
1344 let (dk, dq, dt, dp, dg) = (n.scratch("dk", out), n.scratch("dq", out), n.scratch("dt", out), n.scratch("dp", out), n.scratch("dg", out));
1345 format!(
1346 "lk_delta_prep<{d}><<<dim3({hv}, {chunks}), 256, {}, st>>>({}, {}, {}, {}, {dk}, {dq}, {dt}, {dp}, {dg}, {hk}, {hv}, {rows});\n CKL();\n \
1347 lk_delta_state<{d}, {}{}><<<dim3({hv}, {}), {}, 0, st>>>({}, {}, {dk}, {dq}, {dt}, {dp}, {dg}, {}, {}, {hk}, {hv}, {rows});\n CKL();\n \
1348 lk_delta_norm<{}, {}, {sum_ty}><<<{}, {d}, 0, st>>>({}, {}, {}, {}, {}, {d}, {});",
1349 kernels::DELTA_PREP_SMEM,
1350 p(q),
1351 p(k),
1352 p(bb),
1353 p(a),
1354 kernels::DELTA_STATE_ROWS,
1355 n.elems(&[Some(st)]),
1356 d / kernels::DELTA_STATE_ROWS,
1357 kernels::DELTA_STATE_THREADS,
1358 p(k),
1359 p(v),
1360 p(st),
1361 n.scratch("o", out),
1362 b(*f32),
1363 qm(*quant),
1364 rows * hv,
1365 p(z),
1366 p(gn),
1367 n.scratch("o", out),
1368 if *f32 { p(out) } else { "nullptr".into() },
1369 n.quant(out, *quant),
1370 flit(eps.get() as f32)
1371 )
1372 }
1373 Launch::Delta { q, k, v, z, b: bb, a, st, gn, out, f32, quant, hv, hk, d, rows, eps, seqs } => format!(
1376 "lk_delta_step<{sum_ty}, {}, {}, {}, {}, {}><<<dim3({hv}, {}, {seqs}), 128, 0, st>>>({}, {}, {}, {}, {}, {}, {}, {d}, {hk}, {hv}, {rows}, {}, {}, {}, {}, {}, {}, {});",
1377 d.div_ceil(32),
1378 b(*f32),
1379 qm(*quant),
1380 d.div_ceil(32).div_ceil(4),
1381 n.elem(st),
1382 cdiv(*d, 4),
1383 p(q),
1384 p(k),
1385 p(v),
1386 p(bb),
1387 p(a),
1388 p(st),
1389 n.scratch("o", out),
1390 u8::from(*seqs > 1),
1391 p(z),
1392 p(gn),
1393 if *f32 { p(out) } else { "nullptr".into() },
1394 n.quant(out, *quant),
1395 flit(eps.get() as f32),
1396 n.scratch("dc", out)
1397 ),
1398 Launch::Moe { x, wr, wg, wu, wd, formats, y, d, dff, experts, k, router, rb, biases, renormalize, act, shared, per_warp, rows, router_split, .. } => {
1399 let key = &y.buf;
1400 let xin = format!("{}, {}", if m == 0 { p(x) } else { "nullptr".into() }, n.quant(x, m > 0));
1401 let (rf, fr, route_in, router_rows) = match router {
1403 Router::F32 => {
1404 (1, 8, format!("{}, {}", p(x), n.quant(x, m > 0)), format!("lk_w{{nullptr, nullptr, nullptr, 0}}, {}", p(wr)))
1405 }
1406 Router::Rows(ty) => (0, kernels::format_arg(*ty), xin.clone(), format!("{}, nullptr", handle(wr, *ty, *d))),
1407 };
1408 let [fg, fu, fd] = formats.map(kernels::format_arg);
1409 let none = "lk_w{nullptr, nullptr, nullptr, 0}".to_string();
1411 let (sdff, [fsg, fsu, fsd], [sg, su, sd]) = match shared {
1412 Some(e) => (
1413 e.dff,
1414 e.formats.map(kernels::format_arg),
1415 [handle(&e.wg, e.formats[0], *d), handle(&e.wu, e.formats[1], *d), handle(&e.wd, e.formats[2], e.dff)],
1416 ),
1417 None => (0, [fg, fu, fd], [none.clone(), none.clone(), none]),
1418 };
1419 let (hidden, rows_a_warp) = (k * dff + sdff, *per_warp);
1420 let wpe = kernels::route_warps(*router == Router::F32, *d);
1421 let (act_arg, limit, slope) = match act {
1422 MoeAct::Swiglu => (0, 0.0, 0.0),
1423 MoeAct::ClampedSwiglu { limit, slope } => (1, f64::from_bits(*limit), f64::from_bits(*slope)),
1424 };
1425 let (hid, i, g) = (n.scratch("h", key), n.scratch("i", key), n.scratch("g", key));
1426 let (lg, cnt) = (n.scratch("rl", key), n.scratch("rc", key));
1427 let hq = if m > 0 {
1428 let lo = if n.two { n.scratch("hl", key) } else { "nullptr".into() };
1429 format!("{}, {lo}, {}", n.scratch("hq", key), n.scratch("hs", key))
1430 } else {
1431 "nullptr, nullptr, nullptr".into()
1432 };
1433 if *rows > 1 {
1434 let (rows, choices) = (*rows, rows * k);
1436 let names = ["ec", "eo", "el", "ep", "gx", "gq", "gl", "gs", "pa"].map(|t| n.scratch(t, key));
1437 let [ec, eo, el, ep, gx, gq, gl, gs, pa] = &names;
1438 let gathered = if m == 0 {
1439 format!("{gx}, nullptr, nullptr, nullptr")
1440 } else {
1441 format!("nullptr, {gq}, {}, {gs}", if n.two { gl.as_str() } else { "nullptr" })
1442 };
1443 let hidden_rows = choices * dff;
1444 let quant = if m > 0 {
1445 format!("lk_quant<{m}><<<{}, 256, 0, st>>>({hid}, {hq}, {hidden_rows});\n CKL();\n ", cdiv(hidden_rows, 256))
1446 } else {
1447 String::new()
1448 };
1449 let experts_pass = if l.pairs_gate_and_up(numerics) {
1450 let words = format!("{gq}, {}, {gs}", if n.two { gl.as_str() } else { "nullptr" });
1454 let bm = kernels::gemm_blk_expert_rows(choices, *experts);
1455 let shape = |f: u32, k: usize| format!("{}, {}", kernels::gemm_blk_launch_threads(f, bm, k), kernels::gemm_blk_launch_smem(f, n.two, bm, k));
1456 format!(
1457 "{}<<<dim3({}, {}, {experts}), {}, st>>>({}, {words}, {}, nullptr, nullptr, 0, {d}, {dff}, 0, {ec}, {eo}, nullptr, nullptr, \
1458 {}, {}, {hq}, {}, {});\n CKL();\n \
1459 {}<<<dim3({}, {}, {experts}), {}, st>>>({}, {hq}, {}, {pa}, nullptr, 0, {dff}, {d}, 0, {ec}, {eo}, {el}, {g});\n CKL();\n ",
1460 gemm_blk_name(fg, m, 3, bm, i32::from(act_arg)),
1461 cdiv(*dff, kernels::GEMM_TILE_N / 2),
1462 cdiv(rows, bm),
1463 shape(fg, *d),
1464 handle(wg, formats[0], *d),
1465 n.opt(&biases[0]),
1466 handle(wu, formats[1], *d),
1467 n.opt(&biases[1]),
1468 flit(limit as f32),
1469 flit(slope as f32),
1470 gemm_blk_name(fd, m, 2, bm, 0),
1471 cdiv(*d, kernels::GEMM_TILE_N),
1472 cdiv(rows, bm),
1473 shape(fd, *dff),
1474 handle(wd, formats[2], *dff),
1475 n.opt(&biases[2])
1476 )
1477 } else if l.on_tensor_cores(numerics) {
1478 let (hg, hu) = (n.scratch("hg", key), n.scratch("hu", key));
1481 let words = format!("{gq}, {}, {gs}", if n.two { gl.as_str() } else { "nullptr" });
1482 let bm = kernels::gemm_blk_expert_rows(choices, *experts);
1483 let grid = |cols: usize| format!("dim3({}, {}, {experts})", cdiv(cols, kernels::GEMM_TILE_N), cdiv(rows, bm));
1484 let shape = |f: u32, k: usize| format!("{}, {}", kernels::gemm_blk_launch_threads(f, bm, k), kernels::gemm_blk_launch_smem(f, n.two, bm, k));
1485 format!(
1486 "{}<<<{}, {}, st>>>({}, {words}, nullptr, {hg}, nullptr, 0, {d}, {dff}, 0, {ec}, {eo});\n CKL();\n \
1487 {}<<<{}, {}, st>>>({}, {words}, nullptr, {hu}, nullptr, 0, {d}, {dff}, 0, {ec}, {eo});\n CKL();\n \
1488 lk_moe_glu<{act_arg}><<<{}, 256, 0, st>>>({hg}, {hu}, {}, {}, {i}, {el}, {hid}, {hidden_rows}L, {dff}, {}, {});\n CKL();\n \
1489 {quant}{}<<<{}, {}, st>>>({}, {hq}, {}, {pa}, nullptr, 0, {dff}, {d}, 0, {ec}, {eo}, {el}, {g});\n CKL();\n ",
1490 gemm_blk_name(fg, m, 1, bm, 0),
1491 grid(*dff),
1492 shape(fg, *d),
1493 handle(wg, formats[0], *d),
1494 gemm_blk_name(fu, m, 1, bm, 0),
1495 grid(*dff),
1496 shape(fu, *d),
1497 handle(wu, formats[1], *d),
1498 cdiv(hidden_rows, 256),
1499 n.opt(&biases[0]),
1500 n.opt(&biases[1]),
1501 flit(limit as f32),
1502 flit(slope as f32),
1503 gemm_blk_name(fd, m, 2, bm, 0),
1504 grid(*d),
1505 shape(fd, *dff),
1506 handle(wd, formats[2], *dff),
1507 n.opt(&biases[2])
1508 )
1509 } else {
1510 format!(
1512 "lk_moe_up_rows<{m}, {act_arg}, {fg}, {fu}><<<dim3({}, {experts}), 256, 0, st>>>({gathered}, {}, {}, {}, {}, {ec}, {eo}, {hid}, {d}, {dff}, {}, {});\n CKL();\n \
1513 {quant}lk_moe_down_rows<{m}, {fd}><<<dim3({}, {experts}), 256, 0, st>>>({hid}, {hq}, {}, {}, {ec}, {eo}, {el}, {g}, {pa}, {d}, {dff});\n CKL();\n ",
1514 cdiv(*dff, 8),
1515 handle(wg, formats[0], *d),
1516 handle(wu, formats[1], *d),
1517 n.opt(&biases[0]),
1518 n.opt(&biases[1]),
1519 flit(limit as f32),
1520 flit(slope as f32),
1521 cdiv(*d, 8),
1522 handle(wd, formats[2], *dff),
1523 n.opt(&biases[2])
1524 )
1525 };
1526 let route = if l.on_tensor_cores(numerics) {
1527 let logits = match router {
1529 Router::F32 => format!(
1530 "lk_gemm_f32<<<dim3({}, {}, {router_split}), 256, 0, st>>>({}, {}, {}, {lg}, {rows}, {experts}, {d}{});",
1531 cdiv(*experts, 64),
1532 cdiv(rows, 64),
1533 p(x),
1534 p(wr),
1535 n.opt(rb),
1536 if *router_split > 1 {
1537 format!(", {router_split}, {}, {}", n.scratch("rp", key), n.scratch("rq", key))
1538 } else {
1539 String::new()
1540 }
1541 ),
1542 Router::Rows(ty) => format!(
1543 "{}<<<dim3({}, {}), {}, {}, st>>>({}, {}, {}, {lg}, nullptr, 0, {d}, {experts}, {rows});",
1544 gemm_blk_name(kernels::format_arg(*ty), m, 0, kernels::GEMM_BLK_TILE_M, 0),
1545 cdiv(*experts, kernels::GEMM_TILE_N),
1546 cdiv(rows, kernels::GEMM_BLK_TILE_M),
1547 kernels::gemm_blk_launch_threads(kernels::format_arg(*ty), kernels::GEMM_BLK_TILE_M, *d),
1548 kernels::gemm_blk_launch_smem(kernels::format_arg(*ty), n.two, kernels::GEMM_BLK_TILE_M, *d),
1549 handle(wr, *ty, *d),
1550 n.quant(x, true),
1551 n.opt(rb)
1552 ),
1553 };
1554 format!(
1555 "{logits}\n CKL();\n lk_moe_pick<{m}><<<{rows}, 256, {}, st>>>({lg}, {experts}, {k}, {}, {i}, {g});",
1556 experts * 4,
1557 u8::from(*renormalize)
1558 )
1559 } else {
1560 format!(
1561 "lk_moe_route<{m}, {rf}, {fr}{}><<<dim3({}, {rows}), 256, {}, st>>>({route_in}, {router_rows}, {}, {d}, {experts}, {k}, {}, {i}, {g}, {lg}, {cnt});",
1562 per_warp_arg(wpe),
1563 cdiv(*experts * wpe, 8),
1564 experts * 4,
1565 n.opt(rb),
1566 u8::from(*renormalize)
1567 )
1568 };
1569 let combine = match &finish {
1571 None => format!("lk_moe_combine<{m}><<<dim3({}, {rows}), 256, 0, st>>>({pa}, {ep}, {k}, {d}, {});", cdiv(*d, 256), n.dst_ptr(y)),
1572 Some((norm, sum, _)) => {
1573 let (norm_t, nm) = norm_of(norm, n, m, *d, sum_ty);
1574 let (plus, dst) = sum.map_or(("nullptr".into(), n.dst_ptr(y)), |(plus, to)| (p(plus), p(to)));
1575 format!("lk_moe_combine_norm<{m}, {norm_t}><<<{rows}, {}, 0, st>>>({pa}, {ep}, {k}, {d}, {plus}, {dst}, {nm});", kernels::norm_threads(*d))
1576 }
1577 };
1578 let _ = write!(
1579 out,
1580 " {route}\n CKL();\n \
1581 lk_moe_group<<<1, {}, {}, st>>>({i}, {choices}, {experts}, {ec}, {eo}, {el}, {ep});\n CKL();\n \
1582 lk_moe_gather<{m}><<<{choices}, 256, 0, st>>>({xin}, {el}, {k}, {d}, {gathered});\n CKL();\n \
1583 {experts_pass}{combine}\n CKL();\n",
1584 kernels::moe_group_threads(*experts),
1585 kernels::moe_group_smem_bytes(*experts, kernels::moe_group_threads(*experts)),
1586 );
1587 return;
1588 }
1589 let quant = if m > 0 {
1590 format!("lk_quant<{m}><<<{}, 256, 0, st>>>({hid}, {hq}, {hidden});\n CKL();\n ", cdiv(hidden, 256))
1591 } else {
1592 String::new()
1593 };
1594 let (down_rows, shared_warps) = kernels::moe_down_rows(*d, *dff, sdff);
1595 let (down_blocks, down_threads) = kernels::moe_down_launch(*d, *dff, *k, sdff);
1596 let down = match &finish {
1598 None => format!("lk_moe_down<{m}, {fd}, {fsd}, {down_rows}, {shared_warps}><<<{down_blocks}, {down_threads}, 0, st>>>"),
1599 Some((norm, ..)) => {
1600 let vt = kernels::norm_threads(*d);
1601 let norm_t = norm_of(norm, n, m, *d, sum_ty).0;
1602 format!(
1603 "lk_moe_down_norm<{m}, {fd}, {fsd}, {down_rows}, {shared_warps}, {norm_t}, {}><<<{down_blocks}, {down_threads}, 0, st>>>",
1604 (vt / 32).div_ceil(down_threads / 32)
1605 )
1606 }
1607 };
1608 let tail = match &finish {
1609 None => String::new(),
1610 Some((norm, sum, gc)) => format!(", {}, {}, {}, {gc}", sum.map_or("nullptr".into(), |(plus, _)| p(plus)), norm_of(norm, n, m, *d, sum_ty).1, kernels::norm_threads(*d)),
1611 };
1612 let dst = finish.as_ref().and_then(|(_, sum, _)| sum.map(|(_, to)| p(to))).unwrap_or_else(|| n.dst_ptr(y));
1613 format!(
1614 "lk_moe_route<{m}, {rf}, {fr}{}><<<{}, 256, {}, st>>>({route_in}, {router_rows}, {}, {d}, {experts}, {k}, {}, {i}, {g}, {lg}, {cnt});\n CKL();\n \
1615 lk_moe_up<{m}, {act_arg}, {fg}, {fu}, {fsg}, {fsu}{}><<<{}, 256, 0, st>>>({xin}, {}, {}, {sg}, {su}, {}, {}, {i}, {hid}, {d}, {dff}, {k}, {sdff}, {}, {});\n CKL();\n \
1616 {quant}{down}({hid}, {hq}, {}, {sd}, {}, {i}, {g}, {dst}, {d}, {dff}, {k}, {sdff}{tail});",
1617 per_warp_arg(wpe),
1618 cdiv(*experts * wpe, 8),
1619 experts * 4,
1620 n.opt(rb),
1621 u8::from(*renormalize),
1622 per_warp_arg(rows_a_warp),
1623 cdiv(hidden, 8 * rows_a_warp),
1624 handle(wg, formats[0], *d),
1625 handle(wu, formats[1], *d),
1626 n.opt(&biases[0]),
1627 n.opt(&biases[1]),
1628 flit(limit as f32),
1629 flit(slope as f32),
1630 handle(wd, formats[2], *dff),
1631 n.opt(&biases[2])
1632 )
1633 }
1634 };
1635 let _ = write!(out, " {line}\n CKL();\n");
1636}
1637
1638fn norm_of(norm: &Launch, n: &Names, m: u32, cols: usize, sum_ty: &str) -> (String, String) {
1642 let Launch::Norm { x, add, pre, sum, w, out, into, f32, quant, eps, .. } = norm else { unreachable!("a fold's norm is a norm") };
1643 let own = Dst::to(out);
1644 let (y, yrow, ypitch) = if *f32 { n.dst_split(into.as_ref().unwrap_or(&own), 0) } else { ("nullptr".into(), "nullptr".into(), 0) };
1645 let b = |v: bool| if v { "true" } else { "false" };
1646 let (pw, peps) = match pre {
1647 Some(pn) => (n.opt(&pn.w), flit(pn.eps.get() as f32)),
1648 None => ("nullptr".into(), flit(0.0)),
1649 };
1650 (
1651 format!(
1652 "{}, {}, {}, {}, {}, {sum_ty}, {}",
1653 b(add.is_some()),
1654 b(sum.is_some()),
1655 b(*f32),
1656 if *quant { m } else { 0 },
1657 cols.div_ceil(kernels::norm_threads(cols)),
1658 b(pre.is_some())
1659 ),
1660 format!(
1661 "lk_norm_of{{{}, {}, {}, {}, {y}, {yrow}, {ypitch}, {}, {}, {pw}, {peps}}}",
1662 n.plain(x),
1663 n.opt(add),
1664 n.opt(sum),
1665 n.opt(w),
1666 n.quant(out, *quant),
1667 flit(eps.get() as f32)
1668 ),
1669 )
1670}
1671
1672
1673pub fn emit(forms: &[Form], knobs: &CudaKnobs, ctx: usize, ckpt: Checkpoint, plan: &Plan) -> Result<Program, LaneError> {
1675 let numerics = forms[0].schedule.numerics;
1676 let m = kernels::mode(numerics);
1677 let two = numerics == Q8MatKind::TwoPassInt8;
1678 let shared = shared_names(forms);
1679 let decode = forms.iter().find(|f| f.rows == 1).expect("a lane has a decode step");
1680 let token = format!("t1_{}", c(&decode.schedule.io.token));
1681
1682 let mut cu = header(Provenance::Compiled);
1683 for f in forms {
1684 let _ = writeln!(cu, "// {}", f.identity.line());
1685 for line in f.source.lines() {
1686 let _ = writeln!(cu, "// {line}");
1687 }
1688 }
1689 cu.push_str(
1690 "#include <cstdio>\n#include <cstdlib>\n#include <cstdint>\n#include <cstring>\n#include <cmath>\n\
1691 #include <vector>\n#include <string>\n#include <fcntl.h>\n#include <sys/mman.h>\n#include <sys/stat.h>\n\
1692 #include <unistd.h>\n#include <cuda_runtime.h>\n#include <cuda_fp16.h>\n#include <cuda_bf16.h>\n\
1693 #define CK(x) do { cudaError_t e_ = (x); if(e_ != cudaSuccess){ fprintf(stderr, \"cuda %s:%d %s\\n\", __FILE__, __LINE__, cudaGetErrorString(e_)); exit(1); } } while(0)\n\
1694 #define CKL() CK(cudaGetLastError())\n",
1695 );
1696 cu.push_str(kernels::PRELUDE);
1697 cu.push_str(kernels::STALL);
1698 for k in kernels_used(forms) {
1699 cu.push_str(k.source());
1700 }
1701 let _ = write!(
1702 cu,
1703 "\nstatic const uint64_t CKPT_SIZE = {}ULL, CKPT_DATA = {}ULL, CKPT_DIGEST = {}ULL;\nstatic const int CTX = {ctx};\n",
1704 ckpt.size, ckpt.data_start, ckpt.digest
1705 );
1706 let _ = write!(
1707 cu,
1708 "static const size_t LK_ALIGN = {POOL_ALIGN}UL, LK_ARENA = {}UL, LK_FIXED = {}UL, LK_SPARE = {}UL;\n\
1709 static const int LK_MOVABLE = {};\nstatic const size_t LK_ORDER[] = {{{}}};\n",
1710 plan.arena,
1711 plan.fixed,
1712 plan.spare,
1713 plan.order.len(),
1714 if plan.order.is_empty() { "0".to_string() } else { plan.order.iter().map(|(_, b)| format!("{b}UL")).collect::<Vec<_>>().join(", ") }
1715 );
1716 let rank: BTreeMap<&str, usize> = plan.order.iter().enumerate().map(|(i, (m, _))| (m.as_str(), i)).collect();
1717 let moved = |name: &str| format!("lk_moved > {}", rank[name]);
1718
1719 let mut sidecar: Vec<u8> = Vec::new();
1720 let mut decls = String::new();
1721 let mut loads = String::new();
1722 let mut aliases = String::new();
1723 let mut resets = String::new();
1724 let arena_bytes = forms.iter().map(|f| f.arena.bytes).max().unwrap_or(0);
1727 let mut arena = format!(
1728 " unsigned char* lk_arena = nullptr; CK(cudaMalloc(&lk_arena, {arena_bytes}UL)); CK(cudaMemset(lk_arena, 0xff, {arena_bytes}UL)); lk_device_bytes += {arena_bytes}UL;\n"
1729 );
1730 let mut seen: BTreeMap<String, &Buffer> = BTreeMap::new();
1731 for f in forms {
1732 let n = Names::of(f, &shared, two, knobs.chip.sms);
1733 for buf in f.schedule.buffers.values() {
1734 if f.schedule.io.positions().contains(&buf.name.as_str()) {
1735 continue;
1736 }
1737 if shared.contains(&buf.name) {
1738 if let Some(prev) = seen.get(&buf.name) {
1739 if prev.len != buf.len || prev.fill != buf.fill {
1740 return Err(LaneError::Shape { var: buf.name.clone(), declared: prev.len, bound: buf.len });
1741 }
1742 continue;
1743 }
1744 seen.insert(buf.name.clone(), buf);
1745 }
1746 let len = buf.len;
1747 match &buf.fill {
1748 Fill::Constant(Placement::Rows(r)) if plan.shares.contains_key(&buf.name) => {
1749 let of = &plan.shares[&buf.name];
1750 if r.ty == GgmlType::Q8_0 {
1751 let _ = writeln!(decls, " signed char* {} = nullptr; __half* {} = nullptr;", rows_q(&buf.name), rows_s(&buf.name));
1752 let _ = writeln!(aliases, " {} = {}; {} = {};", rows_q(&buf.name), rows_q(of), rows_s(&buf.name), rows_s(of));
1753 } else {
1754 let _ = writeln!(decls, " unsigned char* {} = nullptr;", rows_b(&buf.name));
1755 let _ = writeln!(aliases, " {} = {};", rows_b(&buf.name), rows_b(of));
1756 }
1757 }
1758 Fill::Constant(Placement::Rows(r)) if r.ty != GgmlType::Q8_0 => {
1759 let runs: Vec<String> = r.runs.iter().map(|run| format!("{}ULL, {}ULL", run.offset, run.rows)).collect();
1760 let _ = writeln!(cu, "static const uint64_t R_{}[] = {{{}}};", c(&buf.name), runs.join(", "));
1761 let _ = writeln!(decls, " unsigned char* {} = nullptr;", rows_b(&buf.name));
1762 let _ = writeln!(
1763 loads,
1764 " upload_raw(file, R_{}, {}, {}L, {}L, &{}, {});",
1765 c(&buf.name),
1766 r.runs.len(),
1767 r.rows,
1768 r.ty.bytes_for(r.cols),
1769 rows_b(&buf.name),
1770 moved(&buf.name)
1771 );
1772 }
1773 Fill::Constant(Placement::Rows(r)) => {
1774 let runs: Vec<String> = r.runs.iter().map(|run| format!("{}ULL, {}ULL", run.offset, run.rows)).collect();
1775 let _ = writeln!(cu, "static const uint64_t R_{}[] = {{{}}};", c(&buf.name), runs.join(", "));
1776 let _ = writeln!(decls, " signed char* {} = nullptr; __half* {} = nullptr;", rows_q(&buf.name), rows_s(&buf.name));
1777 let _ = writeln!(
1778 loads,
1779 " upload_rows(file, R_{}, {}, {}L, {}, &{}, &{}, {});",
1780 c(&buf.name),
1781 r.runs.len(),
1782 r.rows,
1783 r.cols,
1784 rows_q(&buf.name),
1785 rows_s(&buf.name),
1786 moved(&buf.name)
1787 );
1788 }
1789 Fill::Constant(Placement::Values(v)) => {
1790 let off = sidecar.len();
1791 sidecar.extend(v.iter().flat_map(|x| x.to_le_bytes()));
1792 let _ = writeln!(decls, " float* {} = nullptr;", n.plain(&buf.name));
1793 let _ = writeln!(
1794 loads,
1795 " lk_malloc(&{0}, {1}UL * 4); CK(cudaMemcpy({0}, side.data() + {off}, {1}UL * 4, cudaMemcpyHostToDevice));",
1796 n.plain(&buf.name),
1797 v.len()
1798 );
1799 }
1800 Fill::Alias { of, offset } => {
1801 if buf.plain {
1802 let _ = writeln!(aliases, " {}* {} = {} + {offset};", buf.dtype.c_type(), n.plain(&buf.name), n.plain(of));
1803 }
1804 if buf.quantized && quant_in_place(*offset) {
1805 let (v, t) = (n.qname(&buf.name), n.qname(of));
1806 let lo = if two { format!("l{t} + {offset}") } else { "nullptr".into() };
1807 let _ = writeln!(
1808 aliases,
1809 " signed char* q{v} = q{t} + {offset}; signed char* l{v} = {lo}; float* s{v} = s{t} + {};",
1810 offset / QBLOCK
1811 );
1812 continue;
1813 }
1814 }
1815 Fill::Param | Fill::Zero | Fill::Computed => {
1816 if buf.plain {
1817 let ty = buf.dtype.c_type();
1818 let _ = writeln!(decls, " {ty}* {} = nullptr;", n.plain(&buf.name));
1819 match f.arena.at(&buf.name, Touch::Plain) {
1820 Some(p) => {
1821 let _ = writeln!(arena, " {} = ({ty}*)(lk_arena + {});", n.plain(&buf.name), p.offset);
1822 }
1823 None => {
1824 let _ = writeln!(loads, " lk_malloc(&{0}, {len}UL * {1}); CK(cudaMemset({0}, 0, {len}UL * {1}));", n.plain(&buf.name), buf.dtype.bytes());
1825 }
1826 }
1827 if buf.fill == Fill::Zero {
1828 let _ = writeln!(resets, " CK(cudaMemsetAsync({}, 0, {len}UL * {}, st));", n.plain(&buf.name), buf.dtype.bytes());
1829 }
1830 }
1831 }
1832 }
1833 if buf.quantized {
1834 let q = n.qname(&buf.name);
1835 let _ = writeln!(decls, " signed char* q{q} = nullptr; signed char* l{q} = nullptr; float* s{q} = nullptr;");
1836 match f.arena.at(&buf.name, Touch::Quant) {
1837 Some(p) => {
1838 let (words, low, scales) = arena::quant_parts(p.offset, len, two);
1839 let _ = write!(arena, " q{q} = (signed char*)(lk_arena + {words}); s{q} = (float*)(lk_arena + {scales});");
1840 if let Some(low) = low {
1841 let _ = write!(arena, " l{q} = (signed char*)(lk_arena + {low});");
1842 }
1843 arena.push('\n');
1844 }
1845 None => {
1846 let _ = writeln!(loads, " lk_malloc(&q{q}, {len}); lk_malloc(&s{q}, {}UL * 4);", len / QBLOCK);
1847 if two {
1848 let _ = writeln!(loads, " lk_malloc(&l{q}, {len});");
1849 }
1850 }
1851 }
1852 }
1853 }
1854 }
1855
1856 let mut scratch = String::new();
1857 let mut attrs = String::new();
1858 let mut steps = String::new();
1859 let mut graphs = String::new();
1860 let mut runs = String::new();
1861 let mut batch = String::new();
1863 let mut batch_rows = 0usize;
1864 let mut attn_smem: BTreeMap<String, usize> = BTreeMap::new();
1865 let mut shared_smem: BTreeMap<String, usize> = BTreeMap::new();
1866 let mut flash: BTreeSet<(usize, usize, usize, String)> = BTreeSet::new();
1867 let mut mma: BTreeSet<(usize, usize, usize, String)> = BTreeSet::new();
1868 let mut gemm_blk: BTreeSet<(u32, u32, usize, i32)> = BTreeSet::new();
1869 let mut gemv_mma: BTreeSet<(usize, usize, usize, usize)> = BTreeSet::new();
1871 let mut prof: Vec<(i64, &str, String, usize)> = Vec::new();
1874 let mut prof_events = 0usize;
1875 let mut route_smem: BTreeMap<(u32, u32, usize), usize> = BTreeMap::new();
1876 for f in forms {
1877 let n = Names::of(f, &shared, two, knobs.chip.sms);
1878 for (at, l) in f.schedule.launches.iter().enumerate() {
1879 for sc in arena::scratch(l, f.schedule.numerics) {
1881 let name = n.scratch(sc.tag, l.writes()[0]);
1882 let _ = write!(scratch, " {}* {name} = nullptr;", sc.ty);
1883 match f.arena.scratch_at(at, sc.tag) {
1884 Some(p) if !sc.counter => {
1885 let _ = writeln!(scratch, " {name} = ({}*)(lk_arena + {});", sc.ty, p.offset);
1886 }
1887 _ => {
1888 let _ = writeln!(scratch, " lk_malloc(&{name}, {0}UL); CK(cudaMemset({name}, 0, {0}UL));", sc.bytes);
1889 }
1890 }
1891 }
1892 for l in l.parts() {
1893 match l {
1894 Launch::Moe { d, experts, router, formats, rows, k, act, .. } => {
1895 let (rf, fr) = match router {
1896 Router::F32 => (1, 8),
1897 Router::Rows(ty) => (0, kernels::format_arg(*ty)),
1898 };
1899 let wpe = kernels::route_warps(*router == Router::F32, *d);
1900 let bytes = route_smem.entry((rf, fr, wpe)).or_insert(0);
1901 *bytes = (*bytes).max(experts * 4);
1902 if l.on_tensor_cores(f.schedule.numerics) {
1903 let [fg, fu, fd] = formats.map(kernels::format_arg);
1904 let bm = kernels::gemm_blk_expert_rows(rows * k, *experts);
1905 if l.pairs_gate_and_up(f.schedule.numerics) {
1906 gemm_blk.insert((fg, 3, bm, i32::from(*act != MoeAct::Swiglu)));
1907 } else {
1908 gemm_blk.extend([(fg, 1, bm, 0), (fu, 1, bm, 0)]);
1909 }
1910 gemm_blk.insert((fd, 2, bm, 0));
1911 if let Router::Rows(ty) = router {
1912 gemm_blk.insert((kernels::format_arg(*ty), 0, kernels::GEMM_BLK_TILE_M, 0));
1913 }
1914 }
1915 }
1916 Launch::Attention { kc, vc, group, hd, vd, split, .. } if l.attends_on_tensor_cores(f.schedule.numerics) => {
1917 let kv = elems(&f.schedule.buffers, &[Some(kc), Some(vc), split.as_ref().map(|s| s.k_tail.0.as_str())]);
1918 mma.insert((kernels::attn_tile_heads(*group), *hd, *vd, kv));
1919 }
1920 Launch::Attention { kc, vc, group, hd, vd, split, .. } if l.attends_flash(f.schedule.numerics) => {
1921 let kv = elems(&f.schedule.buffers, &[Some(kc), Some(vc), split.as_ref().map(|s| s.k_tail.0.as_str())]);
1922 flash.insert((kernels::attn_tile_heads(*group), *hd, *vd, kv));
1923 }
1924 Launch::Attention { kc, vc, heads, hd, vd, shared: Some(_), quant, gate, .. } => {
1925 let kv = elems(&f.schedule.buffers, &[Some(kc), Some(vc)]);
1926 let kernel = attention_shared_kernel(*heads, *vd, *hd, kc == vc, &kv, &attention_qo(*quant, gate, m));
1927 shared_smem.insert(kernel, kernels::attn_shared_smem_bytes(*heads, *hd));
1928 }
1929 Launch::Attention { kc, vc, hd, vd, heads, rows, per_block, splits, parallel, quant, gate, .. } if !l.attends_in_tiles(f.schedule.numerics) => {
1930 let (acc, ty) = kernels::accumulator(m);
1931 let z = if *parallel { *splits } else { 1 };
1932 let mb = kernels::attn_blocks_a_multiprocessor(heads / per_block * rows * z, knobs.chip.sms, acc);
1933 let (smem, side) = attention_smem(*per_block, *hd, *vd, *splits, *parallel, acc);
1934 let kv = elems(&f.schedule.buffers, &[Some(kc), Some(vc)]);
1935 let kernel = attention_kernel(*per_block, *vd, ty, mb, &kv, &attention_qo(*quant, gate, m), kernels::attn_path(z, side, *splits), *hd);
1936 let bytes = attn_smem.entry(kernel).or_insert(0);
1937 *bytes = (*bytes).max(smem);
1938 }
1939 Launch::Gemv { format, .. } if l.on_tensor_cores(f.schedule.numerics) => {
1940 gemm_blk.insert((kernels::format_arg(*format), 0, l.gemm_rows(f.schedule.numerics), 0));
1941 }
1942 Launch::Gemv { m: rows, .. } | Launch::GemmParts { m: rows, .. } | Launch::GemvGlu { m: rows, .. }
1943 if l.batch_on_tensor_cores(f.schedule.numerics) =>
1944 {
1945 let (cols, k) = match l {
1946 Launch::Gemv { n, k, .. } | Launch::GemvGlu { n, k, .. } => (*n, *k),
1947 Launch::GemmParts { parts, k, .. } => (parts.iter().map(|p| p.n).sum(), *k),
1948 _ => unreachable!(),
1949 };
1950 gemv_mma.extend((0..rows.div_ceil(kernels::GEMV_MMA_ROWS)).map(|s| {
1951 kernels::gemv_mma_launch(kernels::GEMV_MMA_ROWS.min(rows - s * kernels::GEMV_MMA_ROWS), cols, k, knobs.chip.sms)
1952 }));
1953 }
1954 Launch::GemmParts { format, .. } if l.on_tensor_cores(f.schedule.numerics) => {
1955 gemm_blk.insert((kernels::format_arg(*format), 4, l.gemm_rows(f.schedule.numerics), 0));
1956 }
1957 Launch::GemvGlu { formats: [fg, _], act, .. } if l.on_tensor_cores(f.schedule.numerics) => {
1958 gemm_blk.insert((kernels::format_arg(*fg), 5, l.gemm_rows(f.schedule.numerics), act.map_or(-1, act_arg)));
1959 }
1960 _ => {}
1961 }
1962 }
1963 }
1964 let b = branches::branch(&f.schedule, &f.arena);
1969 let r = f.tag();
1972 let key = if f.kind == FormKind::Batch { -(f.rows as i64) } else { f.rows as i64 };
1973 let _ = writeln!(
1974 steps,
1975 " cudaEvent_t ev{r}[{}], fork{r}; for(auto& e : ev{r}) CK(cudaEventCreateWithFlags(&e, cudaEventDisableTiming)); CK(cudaEventCreateWithFlags(&fork{r}, cudaEventDisableTiming));",
1976 b.events.max(1)
1977 );
1978 let mut body = format!(
1979 " cudaStream_t S[{}] = {{st{}}};\n CK(cudaEventRecord(fork{r}, st));\n",
1980 branches::STREAMS,
1981 (1..branches::STREAMS).map(|t| format!(", prof ? st : lk_side[{}]", t - 1)).collect::<String>()
1982 );
1983 for t in &b.sides {
1984 let _ = writeln!(body, " CK(cudaStreamWaitEvent(S[{t}], fork{r}, 0));");
1985 }
1986 for (i, l) in f.schedule.launches.iter().enumerate() {
1987 if !matches!(l, Launch::View { .. }) {
1989 let _ = writeln!(body, " PROF({prof_events});");
1990 prof.push((key, kind(l), l.writes()[0].to_string(), prof_events));
1991 prof_events += 1;
1992 }
1993 let mut stmts = String::new();
1994 launch(l, &n, f.schedule.numerics, &token, &mut stmts);
1995 if stmts.is_empty() {
1996 continue;
1997 }
1998 let t = b.stream[i];
1999 for e in &b.waits[i] {
2000 let _ = writeln!(body, " CK(cudaStreamWaitEvent(S[{t}], ev{r}[{e}], 0));");
2001 }
2002 let _ = writeln!(body, " if(stagger){{ lk_stall<<<1, 1, 0, S[{t}]>>>(lk_stall_ns(stagger_seed, {key}, {i})); CKL(); }}");
2003 let _ = write!(body, " {{ cudaStream_t st = S[{t}];\n{stmts} }}\n");
2004 if let Some(e) = b.records[i] {
2005 let _ = writeln!(body, " CK(cudaEventRecord(ev{r}[{e}], S[{t}]));");
2006 }
2007 }
2008 for e in &b.joins {
2009 let _ = writeln!(body, " CK(cudaStreamWaitEvent(st, ev{r}[{e}], 0));");
2010 }
2011 let _ = writeln!(body, " PROF({prof_events});");
2012 prof_events += 1;
2013 let _ = write!(steps, " auto step{r} = [&](cudaStream_t st){{\n{body} }};\n");
2014 let _ = writeln!(
2015 graphs,
2016 " cudaGraphExec_t exec{r} = nullptr;\n if(captured){{ cudaGraph_t g; CK(cudaStreamBeginCapture(st, cudaStreamCaptureModeGlobal)); step{r}(st); \
2017 CK(cudaStreamEndCapture(st, &g)); CK(cudaGraphInstantiate(&exec{r}, g, 0)); }}"
2018 );
2019 if f.kind == FormKind::Batch {
2020 let logits: Vec<String> = std::iter::once(&f.schedule.io.logits).chain(f.schedule.io.more.iter().map(|q| &q.logits)).map(|t| n.plain(t)).collect();
2022 let _ = writeln!(batch, " auto run_batch = [&]{{ if(captured) CK(cudaGraphLaunch(exec{r}, st)); else step{r}(st); }};");
2023 let _ = writeln!(batch, " float* blogits[BATCH] = {{{}}};", logits.join(", "));
2024 let _ = writeln!(batch, " int* d_btoks = {};", n.plain(&f.schedule.io.token));
2025 let state: Vec<String> = f.graph.of_role(Role::State).map(|t| t.name.clone()).collect();
2027 let mut cases = String::new();
2028 for seq in 0..f.rows {
2029 let suffix = format!("_s{seq}");
2030 let mut body = String::new();
2031 for name in state.iter().filter(|t| t.ends_with(&suffix)) {
2032 let single = &name[..name.len() - suffix.len()];
2033 if shared.contains(single) {
2034 let _ = write!(body, " t_{} = {};", c(single), n.plain(name));
2035 let buf = &f.schedule.buffers[name];
2036 if buf.quantized {
2037 let (to, from) = (format!("_{}", c(single)), n.qname(name));
2038 let _ = write!(body, " q{to} = q{from}; l{to} = l{from}; s{to} = s{from};");
2039 }
2040 }
2041 }
2042 let _ = writeln!(cases, " case {seq}:{body} break;");
2043 }
2044 let _ = writeln!(batch, " auto slot = [&](int b){{\n switch(b){{\n{cases} }}\n }};");
2045 batch_rows = f.rows;
2046 } else {
2047 let _ = writeln!(
2048 runs,
2049 " if(rows == {0}){{ if(captured && !lk_slot) CK(cudaGraphLaunch(exec{r}, st)); else step{r}(st); logits = {1}; return; }}",
2050 f.rows,
2051 n.plain(&f.schedule.io.logits)
2052 );
2053 }
2054 }
2055 if batch_rows == 0 {
2056 batch.push_str(" auto run_batch = [&]{};\n float* blogits[1] = {nullptr};\n int* d_btoks = nullptr;\n auto slot = [&](int){};\n");
2057 }
2058 for (kernel, bytes) in attn_smem.into_iter().chain(shared_smem).filter(|(_, b)| *b > 48 * 1024) {
2059 let _ = writeln!(attrs, " CK(cudaFuncSetAttribute({kernel}, cudaFuncAttributeMaxDynamicSharedMemorySize, {bytes}));");
2060 }
2061 for (hg, hd, vd, kv) in mma.into_iter().filter(|(_, hd, vd, _)| kernels::attn_mma_smem_bytes(*hd, *vd) > 48 * 1024) {
2062 let _ = writeln!(
2063 attrs,
2064 " CK(cudaFuncSetAttribute(lk_attention_mma<{hg}, {hd}, {vd}{kv}>, cudaFuncAttributeMaxDynamicSharedMemorySize, {}));",
2065 kernels::attn_mma_smem_bytes(hd, vd)
2066 );
2067 }
2068 for (hg, hd, vd, kv) in flash {
2069 let (bq, nt) = (kernels::attn_flash_queries(hd, vd), kernels::attn_flash_threads(hd, vd));
2070 let _ = writeln!(
2071 attrs,
2072 " CK(cudaFuncSetAttribute(lk_attention_flash<{hg}, {hd}, {vd}, {bq}, {nt}{kv}>, cudaFuncAttributeMaxDynamicSharedMemorySize, {}));",
2073 kernels::attn_flash_smem_bytes(hd, vd)
2074 );
2075 }
2076 for (f, group, bm, act) in gemm_blk {
2077 let _ = writeln!(
2078 attrs,
2079 " CK(cudaFuncSetAttribute({}, cudaFuncAttributeMaxDynamicSharedMemorySize, {}));",
2080 gemm_blk_name(f, m, group, bm, act),
2081 kernels::gemm_blk_smem_bytes(f, two, bm)
2082 );
2083 }
2084 for (mt, nt, mw, nw) in gemv_mma {
2085 if kernels::gemv_mma_smem_bytes(two, mt, nt) > 48 * 1024 {
2086 let _ = writeln!(
2087 attrs,
2088 " CK(cudaFuncSetAttribute(lk_gemv_mma<{m}, {mt}, {nt}, {mw}, {nw}>, cudaFuncAttributeMaxDynamicSharedMemorySize, {}));",
2089 kernels::gemv_mma_smem_bytes(two, mt, nt)
2090 );
2091 }
2092 }
2093 for ((rf, fr, wpe), bytes) in route_smem.into_iter().filter(|(_, b)| *b > 48 * 1024) {
2094 let route = format!("lk_moe_route<{m}, {rf}, {fr}{}>", per_warp_arg(wpe));
2095 let _ = writeln!(attrs, " CK(cudaFuncSetAttribute({route}, cudaFuncAttributeMaxDynamicSharedMemorySize, {bytes}));");
2096 }
2097 let list = |f: &dyn Fn(&(i64, &str, String, usize)) -> String| prof.iter().map(f).collect::<Vec<_>>().join(", ");
2098 let profile = format!(
2099 " static const int PROF_ROWS[] = {{{}}};\n static const char* PROF_KIND[] = {{{}}};\n static const char* PROF_OUT[] = {{{}}};\n\
2100 \x20 static const int PROF_EV[] = {{{}}};\n const int PROF_N = {};\n\
2101 \x20 std::vector<cudaEvent_t> pev({prof_events});\n if(prof) for(auto& e : pev) CK(cudaEventCreate(&e));\n\
2102 \x20 std::vector<double> pms(PROF_N, 0.0);\n std::vector<int> pcount(PROF_N, 0);\n\
2103 #define PROF(k) do{{ if(prof) CK(cudaEventRecord(pev[k], st)); }}while(0)\n",
2104 list(&|p| p.0.to_string()),
2105 list(&|p| format!("\"{}\"", p.1)),
2106 list(&|p| format!("\"{}\"", p.2)),
2107 list(&|p| p.3.to_string()),
2108 prof.len()
2109 );
2110 let mut chunks: Vec<usize> = forms.iter().filter(|f| f.kind == FormKind::Prefill).map(|f| f.rows).collect();
2111 chunks.sort_unstable_by(|a, b| b.cmp(a));
2112 let chunk_tokens: String = forms
2113 .iter()
2114 .filter(|f| f.kind == FormKind::Prefill)
2115 .map(|f| format!(" if(T == {0}) CK(cudaMemcpyAsync(t{0}_{1}, d_prompt + at, {0} * 4, cudaMemcpyDeviceToDevice, st));\n", f.rows, c(&f.schedule.io.token)))
2116 .collect();
2117 let vocab = decode.schedule.buffers[&decode.schedule.io.logits].len;
2118 let _ = write!(
2119 cu,
2120 "static const size_t SIDECAR_BYTES = {}UL;\nstatic const int CHUNKS[] = {{{}}};\nstatic const int BATCH = {batch_rows};\n{HOST}\nint main(int argc, char** argv){{\n{ARGS}{PLACE}{decls}{loads}{arena}{aliases}{scratch}{attrs}\
2121 \x20 int* d_pos = nullptr; lk_malloc(&d_pos, 4);\n\
2122 \x20 int* d_gen = nullptr; int* d_count = nullptr; int* d_prompt = nullptr;\n\
2123 \x20 lk_malloc(&d_gen, (size_t)CTX * 4); lk_malloc(&d_count, 4); CK(cudaMemset(d_count, 0, 4));\n\
2124 \x20 lk_malloc(&d_prompt, (size_t)CTX * 4);\n\
2125 \x20 int *d_bpos = nullptr, *d_bgen = nullptr, *d_bcount = nullptr;\n\
2126 \x20 if(BATCH > 0){{ lk_malloc(&d_bpos, (size_t)BATCH * 4); lk_malloc(&d_bgen, (size_t)BATCH * CTX * 4); lk_malloc(&d_bcount, (size_t)BATCH * 4); }}\n\
2127 \x20 if(lk_pool_at != lk_pool_bytes || lk_hpool_at != lk_hpool_bytes){{\n\
2128 \x20 fprintf(stderr, \"the program allocates %zu of its %zu pooled device bytes and %zu of its %zu host bytes\\n\", lk_pool_at, lk_pool_bytes, lk_hpool_at, lk_hpool_bytes);\n\
2129 \x20 return 5;\n\
2130 \x20 }}\n\
2131 \x20 printf(\"DEVICE_BYTES: %zu\\nHOST_BYTES: %zu\\nMOVED: %d\\n\", lk_device_bytes, lk_host_bytes, lk_moved);\n\
2132 \x20 cudaStream_t st; CK(cudaStreamCreate(&st));\n\
2133 \x20 cudaStream_t lk_side[{}]; for(auto& s : lk_side) CK(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking));\n{profile}{steps}\
2134 \x20 const bool graphs = {};\n const bool captured = graphs && !prof;\n{graphs}\
2135 \x20 float* logits = nullptr;\n\
2136 \x20 bool lk_slot = false;\n\
2137 \x20 auto run = [&](int rows){{\n{runs} }};\n{batch}\
2138 \x20 auto reset = [&]{{\n{resets} CK(cudaMemsetAsync(d_count, 0, 4, st));\n CK(cudaStreamSynchronize(st));\n\
2139 \x20 int one = 1; CK(cudaMemcpy(d_pos, &one, 4, cudaMemcpyHostToDevice));\n }};\n\
2140 \x20 const int vocab = {vocab};\n{}}}\n",
2141 sidecar.len(),
2142 if chunks.is_empty() { "0".to_string() } else { chunks.iter().map(usize::to_string).collect::<Vec<_>>().join(", ") },
2143 branches::STREAMS - 1,
2144 if knobs.cuda_graphs { "true" } else { "false" },
2145 DRIVE.replace("t_tok_", &token).replace("/*CHUNK_TOKENS*/\n", &chunk_tokens),
2146 HOST = HOST,
2147 ARGS = ARGS,
2148 PLACE = PLACE,
2149 );
2150 Ok(Program { cu, sidecar })
2151}
2152
2153const HOST: &str = r#"
2155static uint64_t fnv1a64(const uint8_t* p, size_t n){
2156 uint64_t h = 0xcbf29ce484222325ULL;
2157 for(size_t i = 0; i < n; ++i){ h ^= p[i]; h *= 0x100000001b3ULL; }
2158 return h;
2159}
2160// Every standing allocation the program makes, in bytes: on the device (the arena and the pool), and in pinned host
2161// memory the device maps.
2162static size_t lk_device_bytes = 0, lk_host_bytes = 0;
2163// The device pool every array but the arena is carved from, and the mapped host pool every moved matrix is, each at
2164// `LK_ALIGN`: their bytes, and how far each is carved.
2165static unsigned char *lk_pool = nullptr, *lk_hpool = nullptr;
2166static size_t lk_pool_bytes = 0, lk_pool_at = 0, lk_hpool_bytes = 0, lk_hpool_at = 0;
2167static void* lk_carve(unsigned char* pool, size_t bytes, size_t& at, size_t& counted, size_t n, const char* what){
2168 n = (n + LK_ALIGN - 1) / LK_ALIGN * LK_ALIGN;
2169 if(at + n > bytes){ fprintf(stderr, "the program allocates past its %zu %s bytes\n", bytes, what); exit(5); }
2170 void* p = pool + at;
2171 at += n;
2172 counted += n;
2173 return p;
2174}
2175template<class T> static void lk_malloc(T** p, size_t n){
2176 *p = (T*)lk_carve(lk_pool, lk_pool_bytes, lk_pool_at, lk_device_bytes, n, "pooled device");
2177}
2178// A matrix past the device's memory, in the mapped host pool, by the pointer a kernel reads it through.
2179template<class T> static void lk_mapped(T** p, size_t n){
2180 *p = (T*)lk_carve(lk_hpool, lk_hpool_bytes, lk_hpool_at, lk_host_bytes, n, "host");
2181}
2182template<class T> static void lk_alloc(T** p, size_t n, bool host){ if(host) lk_mapped(p, n); else lk_malloc(p, n); }
2183// Copy a Q8_0 matrix's rows out of the checkpoint (34-byte blocks: an f16 scale, 32 int8) into int8 rows and their
2184// scales, the layout the GEMV reads, and upload both — to the device, or to mapped host memory when `host`.
2185static void upload_rows(const uint8_t* file, const uint64_t* runs, int nruns, long rows, int cols,
2186 signed char** q, __half** s, bool host){
2187 int nb = cols / 32;
2188 std::vector<signed char> hq((size_t)rows * cols);
2189 std::vector<__half> hs((size_t)rows * nb);
2190 long r = 0;
2191 for(int k = 0; k < nruns; ++k){
2192 const uint8_t* src = file + runs[2 * k];
2193 for(uint64_t i = 0; i < runs[2 * k + 1]; ++i, ++r)
2194 for(int b = 0; b < nb; ++b){
2195 const uint8_t* blk = src + (i * nb + b) * 34;
2196 memcpy(&hs[(size_t)r * nb + b], blk, 2);
2197 memcpy(&hq[(size_t)r * cols + b * 32], blk + 2, 32);
2198 }
2199 }
2200 lk_alloc(q, (size_t)rows * cols, host); lk_alloc(s, (size_t)rows * nb * 2, host);
2201 CK(cudaMemcpy(*q, hq.data(), (size_t)rows * cols, cudaMemcpyDefault));
2202 CK(cudaMemcpy(*s, hs.data(), (size_t)rows * nb * 2, cudaMemcpyDefault));
2203}
2204// The rows of a matrix kept in its checkpoint format, `row_bytes` a row, as the runs of the file lay them out, then 32
2205// zero bytes: a staged read of a row's units runs up to 30 bytes past them. On the device, or in mapped host memory
2206// when `host`.
2207static void upload_raw(const uint8_t* file, const uint64_t* runs, int nruns, long rows, long row_bytes, unsigned char** b,
2208 bool host){
2209 lk_alloc(b, (size_t)rows * row_bytes + 32, host);
2210 static const unsigned char zeros[32] = {0};
2211 CK(cudaMemcpy(*b + (size_t)rows * row_bytes, zeros, 32, cudaMemcpyDefault));
2212 long r = 0;
2213 for(int k = 0; k < nruns; ++k){
2214 CK(cudaMemcpy(*b + (size_t)r * row_bytes, file + runs[2 * k], (size_t)runs[2 * k + 1] * row_bytes, cudaMemcpyDefault));
2215 r += runs[2 * k + 1];
2216 }
2217}
2218"#;
2219
2220const ARGS: &str = r#" setenv("CUDA_MODULE_LOADING", "EAGER", 1);
2222 bool prof = false, stagger = false;
2223 unsigned long long stagger_seed = 0;
2224 size_t device_bytes = SIZE_MAX;
2225 std::vector<char*> av;
2226 for(int i = 0; i < argc; ++i){
2227 std::string s = argv[i];
2228 if(s == "--profile") prof = true;
2229 else if(s == "--stagger" && i + 1 < argc){ stagger = true; stagger_seed = strtoull(argv[++i], nullptr, 10); }
2230 else if(s == "--device-bytes" && i + 1 < argc) device_bytes = strtoull(argv[++i], nullptr, 10);
2231 else av.push_back(argv[i]);
2232 }
2233 argc = (int)av.size(); argv = av.data();
2234 if(argc < 5){
2235 fprintf(stderr, "usage: %s <checkpoint.gguf> <sidecar> <n_gen> [--logits <file>] <token> ...\n"
2236 " %s <checkpoint.gguf> <sidecar> --prompts <file> [--logits <file>] (a line a prompt: <n_gen> <token> ...)\n"
2237 " %s <checkpoint.gguf> <sidecar> --batch <file> [--logits <file>] (the batch step's prompts, decoded together)\n"
2238 " --profile anywhere times every step, the decode step and each prefill chunk, launch by launch\n"
2239 " --stagger <seed> anywhere holds each launch of a step back by a stall drawn from the seed\n"
2240 " --device-bytes <n> anywhere keeps the arrays to n bytes of the device, the matrices past them in host memory\n",
2241 argv[0], argv[0], argv[0]);
2242 return 2;
2243 }
2244 std::vector<std::pair<int, std::vector<int>>> jobs;
2245 FILE* logits_out = nullptr;
2246 int a = 4;
2247 const bool batch_mode = std::string(argv[3]) == "--batch";
2248 if(std::string(argv[3]) == "--prompts" || batch_mode){
2249 FILE* pf = fopen(argv[4], "r");
2250 if(!pf){ perror("prompts"); return 2; }
2251 std::vector<char> line(1 << 20);
2252 while(fgets(line.data(), (int)line.size(), pf)){
2253 std::vector<int> v;
2254 char* p = line.data(); char* e = nullptr;
2255 for(long x = strtol(p, &e, 10); e != p; x = strtol(p, &e, 10)){ v.push_back((int)x); p = e; }
2256 if(v.empty()) continue;
2257 std::vector<int> prompt;
2258 for(size_t i = 1; i < v.size(); ++i) prompt.push_back(v[i] + 1);
2259 jobs.push_back({v[0], prompt});
2260 }
2261 fclose(pf);
2262 a = 5;
2263 if(a + 1 < argc && std::string(argv[a]) == "--logits"){ logits_out = fopen(argv[a + 1], "wb"); a += 2; }
2264 } else {
2265 const int n_gen = atoi(argv[3]);
2266 if(a + 1 < argc && std::string(argv[a]) == "--logits"){ logits_out = fopen(argv[a + 1], "wb"); a += 2; }
2267 std::vector<int> prompt;
2268 for(; a < argc; ++a) prompt.push_back(atoi(argv[a]) + 1);
2269 jobs.push_back({n_gen, prompt});
2270 }
2271 if(batch_mode){
2272 if(BATCH == 0){ fprintf(stderr, "this program has no batch step\n"); return 2; }
2273 if((int)jobs.size() != BATCH){ fprintf(stderr, "the batch step decodes %d sequences: %zu prompts given\n", BATCH, jobs.size()); return 2; }
2274 for(auto& job : jobs) if(job.first != jobs[0].first){ fprintf(stderr, "a batch's sequences generate as many tokens each\n"); return 2; }
2275 }
2276 for(auto& job : jobs){
2277 const std::vector<int>& prompt = job.second;
2278 int n_gen = job.first;
2279 if(prompt.empty() || n_gen < 1 || (int)prompt.size() + n_gen - 1 > CTX){
2280 fprintf(stderr, "a prompt of %zu tokens and %d generated need %zu positions; the caches hold %d\n",
2281 prompt.size(), n_gen, prompt.size() + n_gen - 1, CTX);
2282 return 2;
2283 }
2284 }
2285 int fd = open(argv[1], O_RDONLY); if(fd < 0){ perror("checkpoint"); return 3; }
2286 struct stat sb; fstat(fd, &sb);
2287 if((uint64_t)sb.st_size != CKPT_SIZE){ fprintf(stderr, "%s is not the checkpoint this program was compiled against\n", argv[1]); return 3; }
2288 const uint8_t* file = (const uint8_t*)mmap(nullptr, sb.st_size, PROT_READ, MAP_SHARED, fd, 0);
2289 if(file == MAP_FAILED){ perror("mmap"); return 3; }
2290 if(fnv1a64(file, CKPT_DATA) != CKPT_DIGEST){ fprintf(stderr, "%s is not the checkpoint this program was compiled against\n", argv[1]); return 3; }
2291 std::vector<uint8_t> side(SIDECAR_BYTES);
2292 FILE* sf = fopen(argv[2], "rb");
2293 if(!sf || fread(side.data(), 1, SIDECAR_BYTES, sf) != SIDECAR_BYTES || fgetc(sf) != EOF){ fprintf(stderr, "%s is not this program's sidecar\n", argv[2]); return 3; }
2294 fclose(sf);
2295"#;
2296
2297const PLACE: &str = r#" CK(cudaFree(0));
2302 size_t lk_free = 0, lk_total = 0;
2303 CK(cudaMemGetInfo(&lk_free, &lk_total));
2304 const size_t lk_budget = device_bytes != SIZE_MAX ? device_bytes : (lk_free > LK_SPARE ? lk_free - LK_SPARE : 0);
2305 if(LK_ARENA + LK_FIXED > lk_budget){
2306 fprintf(stderr, "the caches, arena and buffers take %zu bytes of the device's %zu, every matrix moved to the host\n", LK_ARENA + LK_FIXED, lk_budget);
2307 return 4;
2308 }
2309 size_t lk_device = LK_ARENA + LK_FIXED, lk_host = 0;
2310 for(int i = 0; i < LK_MOVABLE; ++i) lk_device += LK_ORDER[i];
2311 int lk_moved = 0;
2312 while(lk_device > lk_budget){ lk_device -= LK_ORDER[lk_moved]; lk_host += LK_ORDER[lk_moved]; ++lk_moved; }
2313 lk_pool_bytes = lk_device - LK_ARENA;
2314 CK(cudaMalloc(&lk_pool, lk_pool_bytes));
2315 if(lk_host){
2316 void* h = nullptr;
2317 CK(cudaHostAlloc(&h, lk_host, cudaHostAllocMapped));
2318 CK(cudaHostGetDevicePointer((void**)&lk_hpool, h, 0));
2319 lk_hpool_bytes = lk_host;
2320 }
2321"#;
2322
2323const DRIVE: &str = r#" cudaEvent_t e0, e1, e2; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1)); CK(cudaEventCreate(&e2));
2326 if(batch_mode){
2327 // Each sequence's prompt prefilled into its own state, as a prompt alone runs; then every sequence a token a step
2328 // in the batch step. Each sequence's logits are written in its turn, as `--prompts` writes a prompt's.
2329 const int n_gen = jobs[0].first, one = 1;
2330 std::vector<std::vector<float>> bl(BATCH);
2331 auto bdump = [&](int b, const float* p){
2332 if(!logits_out) return;
2333 size_t at = bl[b].size();
2334 bl[b].resize(at + vocab);
2335 CK(cudaMemcpyAsync(bl[b].data() + at, p, (size_t)vocab * 4, cudaMemcpyDeviceToHost, st));
2336 CK(cudaStreamSynchronize(st));
2337 };
2338 reset();
2339 lk_slot = true;
2340 CK(cudaEventRecord(e0, st));
2341 for(int b = 0; b < BATCH; ++b){
2342 const std::vector<int>& prompt = jobs[b].second;
2343 slot(b);
2344 CK(cudaMemsetAsync(d_count, 0, 4, st));
2345 CK(cudaMemcpyAsync(d_pos, &one, 4, cudaMemcpyHostToDevice, st));
2346 CK(cudaMemcpyAsync(d_prompt, prompt.data(), prompt.size() * 4, cudaMemcpyHostToDevice, st));
2347 for(int at = 0; at < (int)prompt.size(); ){
2348 int T = 1;
2349 for(int c : CHUNKS) if(c > 1 && c <= (int)prompt.size() - at){ T = c; break; }
2350 if(T == 1) CK(cudaMemcpyAsync(t_tok_, d_prompt + at, 4, cudaMemcpyDeviceToDevice, st));
2351/*CHUNK_TOKENS*/
2352 run(T);
2353 at += T;
2354 }
2355 // The sequence's first generated token, its next position and its run into the batch's slot.
2356 CK(cudaMemcpyAsync(d_btoks + b, t_tok_, 4, cudaMemcpyDeviceToDevice, st));
2357 CK(cudaMemcpyAsync(d_bpos + b, d_pos, 4, cudaMemcpyDeviceToDevice, st));
2358 CK(cudaMemcpyAsync(d_bgen + (size_t)b * CTX, t_tok_, 4, cudaMemcpyDeviceToDevice, st));
2359 CK(cudaMemcpyAsync(d_bcount + b, &one, 4, cudaMemcpyHostToDevice, st));
2360 bdump(b, logits);
2361 }
2362 lk_slot = false;
2363 CK(cudaEventRecord(e1, st));
2364 for(int g = 1; g < n_gen; ++g){
2365 run_batch();
2366 for(int b = 0; b < BATCH; ++b) bdump(b, blogits[b]);
2367 }
2368 CK(cudaEventRecord(e2, st));
2369 CK(cudaStreamSynchronize(st));
2370 CK(cudaGetLastError());
2371 float prefill_ms = 0.f, decode_ms = 0.f;
2372 CK(cudaEventElapsedTime(&prefill_ms, e0, e1)); CK(cudaEventElapsedTime(&decode_ms, e1, e2));
2373 std::vector<int> gen((size_t)BATCH * CTX);
2374 CK(cudaMemcpy(gen.data(), d_bgen, gen.size() * 4, cudaMemcpyDeviceToHost));
2375 for(int b = 0; b < BATCH; ++b){
2376 printf("TOKENS:");
2377 for(int g = 0; g < n_gen; ++g) printf(" %d", gen[(size_t)b * CTX + g] - 1);
2378 printf("\n");
2379 if(logits_out) fwrite(bl[b].data(), 4, bl[b].size(), logits_out);
2380 }
2381 printf("PREFILL_MS: %.3f\n", prefill_ms);
2382 if(n_gen > 1) printf("BATCH: %d\nDECODE_MS_PER_STEP: %.4f\nDECODE_TOK_PER_S: %.2f\n", BATCH, decode_ms / (n_gen - 1), 1000.0 * BATCH * (n_gen - 1) / decode_ms);
2383 if(logits_out) fclose(logits_out);
2384 return 0;
2385 }
2386 std::vector<float> host_logits(logits_out ? vocab : 0);
2387 auto dump = [&]{
2388 if(!logits_out) return;
2389 CK(cudaMemcpyAsync(host_logits.data(), logits, (size_t)vocab * 4, cudaMemcpyDeviceToHost, st));
2390 CK(cudaStreamSynchronize(st));
2391 fwrite(host_logits.data(), 4, vocab, logits_out);
2392 };
2393 // A step of `rows`, its launches timed one by one under `--profile`.
2394 auto timed = [&](int rows){
2395 run(rows);
2396 if(!prof) return;
2397 CK(cudaStreamSynchronize(st));
2398 for(int k = 0; k < PROF_N; ++k) if(PROF_ROWS[k] == rows){
2399 float ms = 0.f; CK(cudaEventElapsedTime(&ms, pev[PROF_EV[k]], pev[PROF_EV[k] + 1])); pms[k] += ms; ++pcount[k];
2400 }
2401 };
2402 for(auto& job : jobs){
2403 const std::vector<int>& prompt = job.second;
2404 const int n_gen = job.first;
2405 reset();
2406 CK(cudaMemcpy(d_prompt, prompt.data(), prompt.size() * 4, cudaMemcpyHostToDevice));
2407 CK(cudaEventRecord(e0, st));
2408 int steps = 0;
2409 for(int at = 0; at < (int)prompt.size(); ){
2410 int T = 1;
2411 for(int c : CHUNKS) if(c > 1 && c <= (int)prompt.size() - at){ T = c; break; }
2412 if(T == 1) CK(cudaMemcpyAsync(t_tok_, d_prompt + at, 4, cudaMemcpyDeviceToDevice, st));
2413/*CHUNK_TOKENS*/
2414 timed(T);
2415 at += T; ++steps;
2416 }
2417 CK(cudaEventRecord(e1, st));
2418 dump();
2419 for(int g = 1; g < n_gen; ++g){
2420 timed(1);
2421 dump(); ++steps;
2422 }
2423 CK(cudaEventRecord(e2, st));
2424 CK(cudaStreamSynchronize(st));
2425 CK(cudaGetLastError());
2426 float prefill_ms = 0.f, decode_ms = 0.f;
2427 CK(cudaEventElapsedTime(&prefill_ms, e0, e1)); CK(cudaEventElapsedTime(&decode_ms, e1, e2));
2428 std::vector<int> gen(steps);
2429 CK(cudaMemcpy(gen.data(), d_gen, (size_t)steps * 4, cudaMemcpyDeviceToHost));
2430 printf("TOKENS:");
2431 for(int g = steps - n_gen; g < steps; ++g) printf(" %d", gen[g] - 1);
2432 printf("\nPREFILL_MS: %.3f\n", prefill_ms);
2433 if(n_gen > 1) printf("DECODE_MS_PER_TOKEN: %.4f\nDECODE_TOK_PER_S: %.2f\n", decode_ms / (n_gen - 1), 1000.0 * (n_gen - 1) / decode_ms);
2434 fflush(stdout);
2435 }
2436 if(prof){
2437 // The decode step's launches, then each chunk's that ran: a launch's mean time over its step's runs.
2438 double total = 0.0;
2439 for(int k = 0; k < PROF_N; ++k) if(PROF_ROWS[k] == 1 && pcount[k]){
2440 double us = 1000.0 * pms[k] / pcount[k];
2441 printf("PROFILE: %s %s %.3f\n", PROF_KIND[k], PROF_OUT[k], us); total += us;
2442 }
2443 if(total > 0.0) printf("PROFILE_TOTAL_US: %.3f\n", total);
2444 for(int c : CHUNKS){
2445 if(c <= 1) continue;
2446 double chunk = 0.0;
2447 for(int k = 0; k < PROF_N; ++k) if(PROF_ROWS[k] == c && pcount[k]){
2448 double us = 1000.0 * pms[k] / pcount[k];
2449 printf("PROFILE_CHUNK: %d %s %s %.3f\n", c, PROF_KIND[k], PROF_OUT[k], us); chunk += us;
2450 }
2451 if(chunk > 0.0) printf("PROFILE_CHUNK_TOTAL_US: %d %.3f\n", c, chunk);
2452 }
2453 }
2454 if(logits_out) fclose(logits_out);
2455 return 0;
2456"#;