cherenkov/runner/
diagnostics.rs

1//! Optional CPU-reference comparisons, research dumps and per-step traces.
2
3use crate::qwen4_exp;
4use anyhow::Result;
5
6fn argmax(logits: &[f32]) -> u32 {
7    logits
8        .iter()
9        .enumerate()
10        .max_by(|a, b| a.1.total_cmp(b.1))
11        .map(|(i, _)| i as u32)
12        .unwrap()
13}
14
15fn logits_diff(a: &[f32], b: &[f32]) -> (f32, f32) {
16    let max_abs = a
17        .iter()
18        .zip(b)
19        .map(|(x, y)| (x - y).abs())
20        .fold(0.0f32, f32::max);
21    let scale = b.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
22
23    (max_abs, max_abs / scale)
24}
25
26pub(super) fn dump_run(gpu: &qwen4_exp::gpu::Gpu<'_>) -> Result<()> {
27    if let Ok(path) = std::env::var("CHERENKOV_DUMP_EXPERTS") {
28        std::fs::write(&path, serde_json::to_vec(&gpu.expert_history)?)?;
29        eprintln!(
30            "expert history ({} steps) written to {path}",
31            gpu.expert_history.len()
32        );
33    }
34
35    if let Ok(path) = std::env::var("CHERENKOV_DUMP_STATES") {
36        // [step][block] = (router input, top-k), row 0 only; f32 little
37        // endian after a small JSON header line.
38        use std::io::Write as _;
39
40        let mut f = std::io::BufWriter::new(std::fs::File::create(&path)?);
41        let steps = gpu.state_history.len();
42        let blocks = gpu.state_history.first().map_or(0, |s| s.len());
43        let hidden = gpu
44            .state_history
45            .first()
46            .and_then(|s| s.first())
47            .map_or(0, |b| b.0.len());
48        let k = gpu
49            .state_history
50            .first()
51            .and_then(|s| s.first())
52            .map_or(0, |b| b.1.len());
53
54        // Steps carry 48 or 49 blocks (the folded MTP block only runs on
55        // drafting steps), so each step is prefixed with its count.
56        writeln!(
57            f,
58            "{{\"format\":2,\"steps\":{steps},\"blocks\":{blocks},\"hidden\":{hidden},\"k\":{k}}}"
59        )?;
60
61        for step in &gpu.state_history {
62            f.write_all(&(step.len() as u32).to_le_bytes())?;
63
64            for (x, ids) in step {
65                for v in x {
66                    f.write_all(&v.to_le_bytes())?;
67                }
68
69                for id in ids {
70                    f.write_all(&id.to_le_bytes())?;
71                }
72            }
73        }
74
75        eprintln!("router states ({steps} steps x {blocks} blocks x {hidden}) written to {path}");
76    }
77
78    if let Ok(path) = std::env::var("CHERENKOV_DUMP_ROUTES") {
79        std::fs::write(&path, serde_json::to_vec(&gpu.route_history)?)?;
80        eprintln!(
81            "route history ({} steps) written to {path}",
82            gpu.route_history.len()
83        );
84    }
85
86    if let Ok(path) = std::env::var("CHERENKOV_DUMP_LA") {
87        std::fs::write(&path, serde_json::to_vec(&gpu.la_log)?)?;
88        eprintln!(
89            "lookahead log ({} predictions) written to {path}",
90            gpu.la_log.len()
91        );
92    }
93
94    Ok(())
95}
96
97/// Compare the GPU's rows against the CPU reference, which forwards the
98/// same tokens one at a time. With `next` (the token after each row), the
99/// CPU also runs the MTP head per row and the last row's draft logits are
100/// compared with the GPU's.
101#[allow(clippy::too_many_arguments)]
102pub(super) fn qwen4_exp_check_rows(
103    m: &qwen4_exp::cpu::CpuModel<'_>,
104    st: &mut qwen4_exp::cpu::State,
105    gpu: &qwen4_exp::gpu::Gpu<'_>,
106    rows: &[u32],
107    next: Option<&[u32]>,
108    pos0: usize,
109    label: &str,
110    prefill: bool,
111) -> Result<()> {
112    for (r, &t) in rows.iter().enumerate() {
113        let ref_logits = m.forward_token(t, st)?;
114        let g = if prefill {
115            gpu.pf_logits_row(r)
116        } else {
117            gpu.logits_row(r)
118        };
119        let (abs, rel) = logits_diff(g, &ref_logits);
120        let (ga, ra) = (argmax(g), argmax(&ref_logits));
121
122        eprintln!(
123            "  {label} {}: gpu argmax {ga} cpu argmax {ra} | max abs diff {abs:.4} (rel {rel:.2e}){}",
124            pos0 + r,
125            if ga == ra { "" } else { "  <-- MISMATCH" }
126        );
127
128        if let Some(next) = next {
129            let hyper = st.last_hyper.clone();
130            let (ml, _) = m.mtp_forward(next[r], &hyper, pos0 + r, st)?;
131
132            if r + 1 == rows.len() {
133                let g = gpu.mtp_logits_row(if prefill { 0 } else { r });
134                let (abs, rel) = logits_diff(g, &ml);
135                let (ga, ra) = (argmax(g), argmax(&ml));
136
137                eprintln!(
138                    "  {label} {} mtp: gpu draft {ga} cpu draft {ra} | max abs diff {abs:.4} (rel {rel:.2e}){}",
139                    pos0 + r,
140                    if ga == ra { "" } else { "  <-- MISMATCH" }
141                );
142            }
143        }
144    }
145
146    Ok(())
147}
148
149/// Record prompt argmax rows and the final long-context attention selection.
150pub(super) fn dump_prefill_chunk(
151    gpu: &qwen4_exp::gpu::Gpu<'_>,
152    engine: bool,
153    pos: usize,
154    rows: usize,
155    prompt_tokens: usize,
156    lines: &mut Vec<String>,
157) {
158    for row in 0..rows {
159        let logits = if engine {
160            gpu.pf_logits_row(row)
161        } else {
162            gpu.logits_row(row)
163        };
164        let token = argmax(logits);
165
166        lines.push(format!(
167            "{} {token} {:.4}",
168            pos + row,
169            logits[token as usize]
170        ));
171    }
172
173    let cfg = &gpu.p.cfg;
174
175    if pos + rows != prompt_tokens
176        || prompt_tokens <= cfg.indexer_budget + cfg.indexer_compress_ratio
177    {
178        return;
179    }
180
181    if engine {
182        let blocks = gpu.debug_engine_blocks(rows - 1);
183
184        eprintln!(
185            "engine last row: {} blocks selected, last 8 {:?}",
186            blocks.len(),
187            &blocks[blocks.len().saturating_sub(8)..]
188        );
189
190        return;
191    }
192
193    let vis = gpu.debug_row_vis(rows - 1);
194    let blocks: Vec<u32> = vis
195        .iter()
196        .filter(|&&t| t % 4 == 0)
197        .map(|&t| t / 4)
198        .collect();
199
200    eprintln!(
201        "row path last row: {} visible tokens, last 8 {:?}; {} block starts, last 8 {:?}",
202        vis.len(),
203        &vis[vis.len().saturating_sub(8)..],
204        blocks.len(),
205        &blocks[blocks.len().saturating_sub(8)..]
206    );
207}
208
209pub(super) fn dump_decode(
210    gpu: &qwen4_exp::gpu::Gpu<'_>,
211    ids: &[u32],
212    out: &[u32],
213    prefill_steps: usize,
214) -> Result<()> {
215    if let Ok(path) = std::env::var("CHERENKOV_DUMP_TOKENS") {
216        std::fs::write(
217            &path,
218            serde_json::to_vec(&serde_json::json!({ "prompt": ids, "out": out }))?,
219        )?;
220        eprintln!("tokens written to {path}");
221    }
222
223    if std::env::var_os("CHERENKOV_TRACE").is_some() {
224        for i in prefill_steps..gpu.step_ms.len() {
225            eprintln!(
226                "  step {i}: {} rows, {:.1} ms wall, {:.1} ms gpu, {:.1} ms io wait, {} sync fetches ({:.0} MB), lookahead hit {:.2} fetched {}",
227                gpu.rows[i],
228                gpu.step_ms[i],
229                gpu.gpu_ms[i],
230                gpu.io_ms[i],
231                gpu.misses[i],
232                gpu.miss_bytes[i] as f64 / 1e6,
233                gpu.lookahead_hit[i],
234                gpu.lookahead_issued[i]
235            );
236        }
237    }
238
239    Ok(())
240}
cherenkov/runner/diagnostics.rs

cherenkov/runner/
diagnostics.rs
