decode_flat/
decode_flat.rs

1//! Per-token decode-latency probe (scratch diagnostic, not a product).
2//!
3//! `bench` reports ONE aggregate decode tok/s, which cannot distinguish
4//! "the operator is algorithmically slower at depth" from "the machine
5//! throttled during the long prefill that precedes the decode". This
6//! prints the per-token wall time so the shape is visible: a thermal
7//! effect DECAYS across the run, an algorithmic cost is flat-but-higher
8//! from the first token.
9//!
10//! Usage:
11//!   cargo run --release --example decode_flat -- <model.cmf> <ctx> <tokens> [o1spec]
12
13use cortiq_core::CmfModel;
14use cortiq_engine::{Pipeline, SamplerConfig};
15use std::sync::Arc;
16use std::time::Instant;
17
18fn main() {
19    let args: Vec<String> = std::env::args().collect();
20    if args.len() < 4 {
21        eprintln!("usage: decode_flat <model.cmf> <ctx> <tokens> [o1spec|off]");
22        std::process::exit(2);
23    }
24    let path = &args[1];
25    let ctx: usize = args[2].parse().expect("ctx");
26    let tokens: usize = args[3].parse().expect("tokens");
27    let o1spec = args.get(4).cloned().unwrap_or_else(|| "off".to_string());
28
29    let model = Arc::new(CmfModel::open_sharded(path).expect("open model"));
30    let mut pipeline = Pipeline::from_model(&model, SamplerConfig::default()).expect("pipeline");
31    pipeline.set_o1(cortiq_engine::nystrom::O1Cfg::from_spec(
32        &o1spec, None, None, None, None,
33    ));
34    eprintln!("o1_active = {}", pipeline.o1_active());
35
36    // Same synthetic prompt shape as `cortiq bench --ctx`.
37    let prompt = "The quick brown fox jumps over the lazy dog. ".repeat(ctx / 8 + 2);
38    let mut ids = pipeline.tokenizer.encode(&prompt);
39    ids.truncate(ctx);
40    assert_eq!(ids.len(), ctx, "prompt too short for ctx");
41
42    // Warm the mmap so the first token is not billed for page faults.
43    let _ = pipeline.forward_ids(&ids[..2], None).expect("warm");
44
45    let stamps: Arc<std::sync::Mutex<Vec<Instant>>> = Arc::default();
46    let st = stamps.clone();
47    let cb: cortiq_engine::TokenCallback = Box::new(move |_t| {
48        st.lock().unwrap().push(Instant::now());
49        true
50    });
51    let t0 = Instant::now();
52    let _ = pipeline
53        .generate_from_ids(&ids, tokens, None, Some(cb))
54        .expect("generate");
55
56    let stamps = stamps.lock().unwrap();
57    // stamps[0] fires after generation's prefill + the one-off o1 seal:
58    // that gap is TTFT, not a decode step.
59    println!("ctx={ctx} o1={o1spec} ttft_s={:.3}", stamps[0].duration_since(t0).as_secs_f64());
60    println!("# idx  ms");
61    for i in 1..stamps.len() {
62        println!("{:4}  {:7.2}", i, (stamps[i] - stamps[i - 1]).as_secs_f64() * 1e3);
63    }
64    // Mean of the last half — past any startup transient.
65    let half = stamps.len() / 2;
66    if stamps.len() > 2 && half >= 1 {
67        let dt = (stamps[stamps.len() - 1] - stamps[half]).as_secs_f64();
68        let n = (stamps.len() - 1 - half) as f64;
69        println!("last-half: {:.2} ms/tok = {:.2} tok/s", dt / n * 1e3, n / dt);
70    }
71}
decode_flat/decode_flat.rs

decode_flat/
decode_flat.rs
