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!(
60        "ctx={ctx} o1={o1spec} ttft_s={:.3}",
61        stamps[0].duration_since(t0).as_secs_f64()
62    );
63    println!("# idx  ms");
64    for i in 1..stamps.len() {
65        println!(
66            "{:4}  {:7.2}",
67            i,
68            (stamps[i] - stamps[i - 1]).as_secs_f64() * 1e3
69        );
70    }
71    // Mean of the last half — past any startup transient.
72    let half = stamps.len() / 2;
73    if stamps.len() > 2 && half >= 1 {
74        let dt = (stamps[stamps.len() - 1] - stamps[half]).as_secs_f64();
75        let n = (stamps.len() - 1 - half) as f64;
76        println!(
77            "last-half: {:.2} ms/tok = {:.2} tok/s",
78            dt / n * 1e3,
79            n / dt
80        );
81    }
82}
decode_flat/decode_flat.rs

decode_flat/
decode_flat.rs
