cherenkov/
runner.rs

1//! CLI generation, prompt processing and shared decode entry points.
2
3use crate::units::BYTES_PER_GB;
4use crate::{options::Options, prompt::Prompt, qwen4_exp, tok};
5use anyhow::{Context, Result};
6use std::path::Path;
7
8mod decode;
9mod diagnostics;
10
11pub(crate) use decode::Decode;
12use diagnostics::{dump_decode, dump_prefill_chunk, dump_run, qwen4_exp_check_rows};
13
14pub fn run(model_dir: &Path, prompt: &str, options: &Options) -> Result<()> {
15    options.validate()?;
16
17    let (max_tokens, max_ctx, raw, check, repeat) = (
18        options.max_tokens,
19        options.max_ctx,
20        options.raw,
21        options.check,
22        options.repeat,
23    );
24    let t0 = std::time::Instant::now();
25    let mut model_options = options.clone();
26    let packed = qwen4_exp::packed::Packed::open(model_dir)?;
27
28    if packed.cfg.mtp_num_hidden_layers == 0 {
29        model_options.drafts = 0;
30    }
31
32    let options = &model_options;
33    let tok = tok::ChatTokenizer::load(model_dir)?;
34    let prompt = if raw {
35        Prompt::raw(prompt.to_owned())
36    } else {
37        tok.template
38            .as_ref()
39            .context("checkpoint has no chat template")?
40            .user(prompt)?
41    };
42    let ids = tok.encode(&prompt.text)?;
43
44    check_budget(ids.len(), max_tokens, options.effective_drafts(), max_ctx)?;
45
46    if options.cut_weak > 0.0 {
47        eprintln!(
48            "WARNING: --cut-weak skips late weak experts; output depends on disk timing and is not reproducible."
49        );
50    }
51
52    let mut gpu = qwen4_exp::gpu::Gpu::load(&packed, max_ctx, options)?;
53    let cpu = if check {
54        Some(qwen4_exp::cpu::CpuModel::load(&packed)?)
55    } else {
56        None
57    };
58
59    eprintln!(
60        "cherenkov gpu: max_ctx {max_ctx}, {:.2} GB Metal (expert pool {:.1} GB = {} records, working-set limit {:.2} GB), load {:.2}s, clock probe {:.2} ms",
61        gpu.allocated_gb(),
62        gpu.pool_bytes() as f64 / BYTES_PER_GB as f64,
63        gpu.pool_slots(),
64        gpu.working_set_limit_gb(),
65        t0.elapsed().as_secs_f64(),
66        gpu.throttle_ms()?
67    );
68
69    for run in 0..repeat {
70        if run > 0 {
71            gpu.reset();
72            eprintln!("--- repeat {run}: clock probe {:.2} ms", gpu.throttle_ms()?);
73        }
74
75        let mut cpu_state = cpu.as_ref().map(|m| m.new_state());
76
77        qwen4_exp_gen_once(
78            &mut gpu,
79            cpu.as_ref(),
80            cpu_state.as_mut(),
81            &tok,
82            &ids,
83            max_tokens,
84            options,
85            None,
86            &mut |token| {
87                use std::io::Write as _;
88
89                print!("{}", tok.decode(&[token])?);
90                std::io::stdout().flush()?;
91
92                Ok(())
93            },
94        )?;
95        println!();
96    }
97
98    eprintln!("clock probe at end {:.2} ms", gpu.throttle_ms()?);
99
100    dump_run(&gpu)
101}
102
103/// Stream equal-sized prompt chunks so a short tail does not reread the whole store.
104fn prefill_engine(
105    gpu: &mut qwen4_exp::gpu::Gpu<'_>,
106    cpu: Option<&qwen4_exp::cpu::CpuModel<'_>>,
107    mut cpu_state: Option<&mut qwen4_exp::cpu::State>,
108    ids: &[u32],
109    n_draft: usize,
110    pf_chunk: usize,
111    mut argmax_lines: Option<&mut Vec<String>>,
112) -> Result<PrefillResume> {
113    let (mut p, mut cur, mut drafts) = (0, ids[0], Vec::new());
114    let check = cpu.is_some();
115    // Equal chunks: every chunk streams (nearly) the whole expert
116    // store, so a small trailing chunk would cost as much as a full one.
117    let n_chunks = ids.len().div_ceil(pf_chunk);
118    let chunk_len = ids.len().div_ceil(n_chunks);
119    let mut d1 = 0u32;
120
121    while p < ids.len() {
122        let n = (ids.len() - p).min(chunk_len);
123        let chunk = &ids[p..p + n];
124        let next_after = ids.get(p + n).copied();
125        let (c, d) = gpu.prefill_chunk(chunk, next_after, check || argmax_lines.is_some())?;
126        cur = c;
127        d1 = d;
128
129        if let Some(lines) = argmax_lines.as_deref_mut() {
130            dump_prefill_chunk(gpu, true, p, n, ids.len(), lines);
131        }
132
133        if let (Some(m), Some(st)) = (cpu, cpu_state.as_deref_mut()) {
134            let mut next: Vec<u32> = chunk[1..].to_vec();
135
136            next.push(next_after.unwrap_or(cur));
137            qwen4_exp_check_rows(
138                m,
139                st,
140                gpu,
141                chunk,
142                gpu.has_mtp().then_some(&next[..]),
143                p,
144                "prefill",
145                true,
146            )?;
147        }
148
149        p += n;
150    }
151
152    if n_draft > 0 {
153        drafts = vec![d1];
154
155        if n_draft >= 2 {
156            drafts.push(gpu.mtp_chain(d1)?);
157        }
158    }
159
160    gpu.prefill_release();
161
162    let st = &gpu.prefill_stats;
163    let recs: usize = st.iter().map(|s| s.fetched).sum();
164    let sum = |f: fn(&qwen4_exp::gpu::prefill::ChunkStats) -> f64| st.iter().map(f).sum::<f64>();
165    let bytes: usize = st.iter().map(|s| s.fetched_bytes).sum();
166
167    eprintln!(
168        "prefill engine: {} chunk(s) of up to {pf_chunk}, {recs} expert records streamed ({:.1} GB, {:.0} MB/token), {:.1}s waiting for ring reuse | GPU s: DeltaNet blocks {:.1}, attention blocks {:.1}, expert streams {:.1}, MTP {:.1} | n-gram gather {:.1}s CPU",
169        st.len(),
170        bytes as f64 / BYTES_PER_GB as f64,
171        bytes as f64 / 1e6 / ids.len() as f64,
172        sum(|s| s.wait_s),
173        sum(|s| s.gpu_delta_s),
174        sum(|s| s.gpu_attn_s),
175        sum(|s| s.gpu_experts_s),
176        sum(|s| s.gpu_mtp_s),
177        sum(|s| s.ngram_s),
178    );
179
180    Ok(PrefillResume {
181        next: cur,
182        drafts,
183        logits: None,
184    })
185}
186
187/// Use the decode row kernels for short prompts and explicit row-path checks.
188fn prefill_row_batches(
189    gpu: &mut qwen4_exp::gpu::Gpu<'_>,
190    cpu: Option<&qwen4_exp::cpu::CpuModel<'_>>,
191    mut cpu_state: Option<&mut qwen4_exp::cpu::State>,
192    ids: &[u32],
193    n_draft: usize,
194    mut argmax_lines: Option<&mut Vec<String>>,
195) -> Result<PrefillResume> {
196    use qwen4_exp::gpu::MAX_NB;
197
198    let (mut p, mut cur, mut drafts) = (0, ids[0], Vec::new());
199    // Debug: CHERENKOV_ROWS_MAX caps rows per step in this path.
200    let rows_max: usize = std::env::var("CHERENKOV_ROWS_MAX")
201        .ok()
202        .and_then(|v| v.parse().ok())
203        .unwrap_or(MAX_NB)
204        .clamp(1, MAX_NB);
205
206    while p < ids.len() {
207        let n = (ids.len() - p).min(rows_max);
208        let rows = &ids[p..p + n];
209        let res = gpu.step_rows(rows, false, false)?;
210
211        if let Some(lines) = argmax_lines.as_deref_mut() {
212            dump_prefill_chunk(gpu, false, p, n, ids.len(), lines);
213        }
214
215        gpu.commit(n)?;
216
217        let last = p + n == ids.len();
218        let mut next: Vec<u32> = ids[p + 1..p + n].to_vec();
219
220        next.push(if last { res[n - 1] } else { ids[p + n] });
221
222        if n_draft > 0 {
223            drafts = gpu.mtp_draft(&next, if last { n_draft } else { 1 })?;
224        }
225
226        if let (Some(m), Some(st)) = (cpu, cpu_state.as_deref_mut()) {
227            qwen4_exp_check_rows(
228                m,
229                st,
230                gpu,
231                rows,
232                (n_draft > 0).then_some(&next[..]),
233                p,
234                "prefill",
235                false,
236            )?;
237        }
238
239        p += n;
240        cur = res[n - 1];
241    }
242
243    Ok(PrefillResume {
244        next: cur,
245        drafts,
246        logits: None,
247    })
248}
249
250/// Fill the trunk and draft caches, then return the next token and draft chain.
251fn prefill_prompt(
252    gpu: &mut qwen4_exp::gpu::Gpu<'_>,
253    cpu: Option<&qwen4_exp::cpu::CpuModel<'_>>,
254    cpu_state: Option<&mut qwen4_exp::cpu::State>,
255    ids: &[u32],
256    n_draft: usize,
257    resume: Option<PrefillResume>,
258) -> Result<PrefillResume> {
259    use qwen4_exp::gpu::MAX_NB;
260
261    // Prefill in chunks of MAX_NB known tokens; the MTP head follows each
262    // chunk to fill its cache and drafts after the last one.
263    let t1 = std::time::Instant::now();
264    let prepared = resume.is_some();
265    let mut seed = resume.unwrap_or_else(|| PrefillResume {
266        next: ids[0],
267        drafts: Vec::new(),
268        logits: None,
269    });
270    // Prompts from CHERENKOV_PREFILL_MIN tokens (default 64) go through
271    // the prefill engine in adaptive chunks, overridden by
272    // CHERENKOV_PREFILL_CHUNK (at most 1024 when checking).
273    let check = cpu.is_some();
274    let pf_min: usize = std::env::var("CHERENKOV_PREFILL_MIN")
275        .ok()
276        .and_then(|v| v.parse().ok())
277        .unwrap_or(64);
278    let engine = !prepared && ids.len() >= pf_min;
279    let prefill_fit = if engine {
280        gpu.prefill_rows_fit(check || std::env::var_os("CHERENKOV_DUMP_ARGMAX").is_some())?
281    } else {
282        1
283    };
284    let mut pf_chunk: usize = std::env::var("CHERENKOV_PREFILL_CHUNK")
285        .ok()
286        .and_then(|v| v.parse().ok())
287        .unwrap_or(prefill_fit)
288        .min(prefill_fit)
289        .max(1);
290
291    if check {
292        pf_chunk = pf_chunk.min(1024);
293    }
294
295    // Debug: CHERENKOV_DUMP_ARGMAX=path writes "pos argmax maxlogit" for
296    // every prompt row, to diff the two prefill paths position by position.
297    let dump_argmax = std::env::var("CHERENKOV_DUMP_ARGMAX").ok();
298    let mut argmax_lines: Vec<String> = Vec::new();
299    let lines = dump_argmax.as_ref().map(|_| &mut argmax_lines);
300
301    if engine {
302        seed = prefill_engine(gpu, cpu, cpu_state, ids, n_draft, pf_chunk, lines)?;
303    } else if !prepared {
304        seed = prefill_row_batches(gpu, cpu, cpu_state, ids, n_draft, lines)?;
305    }
306
307    if let Some(path) = &dump_argmax {
308        std::fs::write(path, argmax_lines.join("\n") + "\n")?;
309    }
310
311    // Debug: CHERENKOV_DUMP_LOGITS=path writes the last prompt row's
312    // logits (bisecting the prefill engine against the row-batched path).
313    // Both paths leave the last row's logits in row 0 of the head buffer
314    // only when the last step had one row; the engine always puts them in
315    // row 0, the row path in row (rows in last step - 1).
316    if let Ok(path) = std::env::var("CHERENKOV_DUMP_LOGITS") {
317        let last_rows = if engine {
318            1
319        } else {
320            ((ids.len() - 1) % MAX_NB) + 1
321        };
322        let l = gpu.logits_row(last_rows - 1);
323        let bytes: Vec<u8> = l.iter().flat_map(|v| v.to_le_bytes()).collect();
324
325        std::fs::write(&path, bytes)?;
326    }
327
328    let prefill = t1.elapsed().as_secs_f64();
329
330    if !prepared {
331        eprintln!(
332            "prefill {} tokens in {:.2}s ({:.1} tok/s){}",
333            ids.len(),
334            prefill,
335            ids.len() as f64 / prefill,
336            if engine {
337                " [engine]"
338            } else {
339                " [row batches]"
340            }
341        );
342    }
343
344    Ok(seed)
345}
346
347#[allow(clippy::too_many_arguments)]
348pub(crate) fn qwen4_exp_gen_once(
349    gpu: &mut qwen4_exp::gpu::Gpu<'_>,
350    cpu: Option<&qwen4_exp::cpu::CpuModel<'_>>,
351    mut cpu_state: Option<&mut qwen4_exp::cpu::State>,
352    tok: &tok::ChatTokenizer,
353    ids: &[u32],
354    max_tokens: usize,
355    options: &Options,
356    resume: Option<PrefillResume>,
357    emit: &mut dyn FnMut(u32) -> Result<()>,
358) -> Result<()> {
359    // Up to two drafts per step by default, the second only after a step
360    // that accepted its whole batch (adaptive, below): a flat second
361    // draft costs more than its extra tokens are worth under the
362    // SSD-driven throttle, so the second draft is adaptive.
363    let n_draft = options.effective_drafts();
364
365    anyhow::ensure!(
366        n_draft == 0 || gpu.has_mtp(),
367        "drafting needs the checkpoint's MTP head"
368    );
369    // Always fold the first draft into the trunk command buffer: neutral
370    // in measured speed, but saves one submit/wait. Chain adaptively after
371    // fully accepted batches; this measured about 4% faster than always chaining.
372
373    let mut seed = prefill_prompt(gpu, cpu, cpu_state.as_deref_mut(), ids, n_draft, resume)?;
374
375    if !options.sampling.greedy() && seed.logits.is_none() {
376        seed.logits = Some(gpu.logits_row(gpu.last_logits_row()).to_vec());
377    }
378
379    let prefill_steps = gpu.step_ms.len();
380    let mut decoder = Decode::new(seed, ids, tok, options, max_tokens, None)?;
381    let t2 = std::time::Instant::now();
382
383    while decoder.finish_reason.is_none() {
384        decoder.step(gpu, cpu, cpu_state.as_deref_mut(), emit)?;
385    }
386
387    let steps = decoder.steps;
388    let accepted = decoder.accepted;
389    let out = decoder.tokens;
390    let decode = t2.elapsed().as_secs_f64();
391    let n = steps.max(1);
392    let recent = |v: &[f64]| -> Vec<f64> { v[v.len().saturating_sub(n)..].to_vec() };
393    let mean = |v: &[f64]| v.iter().sum::<f64>() / v.len().max(1) as f64;
394    let ms = recent(&gpu.step_ms);
395    let gpu_recent = recent(&gpu.gpu_ms);
396    let io_recent = recent(&gpu.io_ms);
397    let set_recent = recent(&gpu.set_ms);
398    let read_recent = recent(&gpu.read_ms);
399    let idle_recent = recent(&gpu.gpu_idle_ms);
400    let mtp_recent = recent(&gpu.mtp_ms);
401    let ngram_recent = recent(&gpu.ngram_ms);
402    let miss_recent = &gpu.misses[gpu.misses.len().saturating_sub(n)..];
403    let la_recent = &gpu.lookahead_issued[gpu.lookahead_issued.len().saturating_sub(n)..];
404    let warm_recent = &gpu.warm[gpu.warm.len().saturating_sub(n)..];
405    let cut_recent = &gpu.cut[gpu.cut.len().saturating_sub(n)..];
406
407    eprintln!(
408        "  dispatches/step {} | cpu turnaround {:.1} ms/step | drafts {n_draft}, accepted {:.2}/step ({:.2} tokens/step) | mtp {:.1} ms/step | n-gram gather {:.1} ms/step",
409        gpu.dispatches.last().copied().unwrap_or(0),
410        mean(&idle_recent),
411        accepted as f64 / n as f64,
412        (out.len()) as f64 / n as f64,
413        mean(&mtp_recent),
414        mean(&ngram_recent),
415    );
416    eprintln!(
417        "decode {} tokens in {:.2}s ({:.2} tok/s) | {} steps, mean step {:.1} ms (min {:.1}, max {:.1}) | gpu-span {:.1} ms | io wait {:.1} ms (set {:.1}, read {:.1}) | sync fetches/step {:.1} + lookahead {:.1} (warm {:.1}, cut {:.1}) | pool {}/{} | {:.2} GB Metal",
418        out.len(),
419        decode,
420        out.len() as f64 / decode.max(1e-9),
421        steps,
422        mean(&ms),
423        ms.iter().cloned().fold(f64::INFINITY, f64::min),
424        ms.iter().cloned().fold(0.0, f64::max),
425        mean(&gpu_recent),
426        mean(&io_recent),
427        mean(&set_recent),
428        mean(&read_recent),
429        miss_recent.iter().sum::<usize>() as f64 / miss_recent.len().max(1) as f64,
430        la_recent.iter().sum::<usize>() as f64 / la_recent.len().max(1) as f64,
431        warm_recent.iter().sum::<usize>() as f64 / warm_recent.len().max(1) as f64,
432        cut_recent.iter().sum::<usize>() as f64 / cut_recent.len().max(1) as f64,
433        gpu.pool_resident(),
434        gpu.pool_slots(),
435        gpu.allocated_gb()
436    );
437    eprintln!(
438        "memory_stats {}",
439        serde_json::to_string(&gpu.memory_stats())?
440    );
441    dump_decode(gpu, ids, &out, prefill_steps)?;
442
443    Ok(())
444}
445
446#[derive(Default)]
447pub(crate) struct PrefillResume {
448    pub next: u32,
449    pub drafts: Vec<u32>,
450    pub logits: Option<Vec<f32>>,
451}
452
453pub(crate) fn check_budget(
454    prompt: usize,
455    output: usize,
456    drafts: usize,
457    context: usize,
458) -> Result<()> {
459    anyhow::ensure!(prompt > 0, "prompt must contain at least one token");
460
461    let need = prompt
462        .checked_add(output)
463        .and_then(|n| n.checked_add(drafts))
464        .ok_or_else(|| anyhow::anyhow!("token budget overflows"))?;
465
466    anyhow::ensure!(
467        need <= context,
468        "prompt ({prompt}) + output ({output}) + draft lookahead ({drafts}) needs {need} tokens, exceeding --max-ctx {context}"
469    );
470
471    Ok(())
472}
473
474#[cfg(test)]
475#[path = "../tests/unit/runner.rs"]
476mod tests;
cherenkov/runner.rs

cherenkov/
runner.rs
