open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-runtime/examples/devstral_hybrid.rs

//! HYBRID decode: the graph prefill seeds the device-resident
//! decoder (via `seed_prefill`), which then runs the fast decode loop.
//!
//! MEASURED NEGATIVE RESULT (Devstral-24B-Q4_K_M, 43-tok prompt): the
//! graph prefill is SLOWER than the resident driver's own sequential
//! per-token prefill — ~334 ms/tok (graph) vs ~68 ms/tok (resident
//! sequential), because the graph executor's per-op command-buffer
//! overhead (the very thing the resident driver bypasses) dominates
//! even when batched. So this hybrid is a net LOSS for prefill; it
//! stays as a comparison harness + proof the `seed_prefill` mechanism
//! is correct (output is coherent). A genuinely faster prefill needs
//! BATCHED RESIDENT (m>1) device kernels, not the graph.
//!
//!   cargo run --release --features fabric -p jouleclaw-runtime \
//!     --example devstral_hybrid -- <model.gguf> [n_gen] [prompt]

#[cfg(all(target_os = "macos", feature = "fabric"))]
fn main() {
    use jouleclaw_loader_gguf::devstral_resident::DevstralResident;
    use jouleclaw_loader_gguf::read_gguf_file;
    use jouleclaw_loader_gguf::tokenizer::Vocab;
    use jouleclaw_runtime::{prefill_inplace, Runtime};
    use std::time::Instant;

    let path = std::env::args().nth(1).expect("usage: <model.gguf> [n_gen] [prompt]");
    let n_gen: usize = std::env::args().nth(2).and_then(|s| s.parse().ok()).unwrap_or(24);
    let prompt = std::env::args()
        .nth(3)
        .unwrap_or_else(|| "The capital of France is".to_string());

    eprintln!("loading {path} ...");
    let model = read_gguf_file(&path).expect("read gguf");
    let vocab = Vocab::from_gguf(&model).expect("vocab");
    let rt = Runtime::boot();
    let driver = DevstralResident::bind(&model).expect("bind resident driver");

    let ids = vocab.encode_bpe_regex(&prompt, true);
    eprintln!("prompt {} tokens", ids.len());
    let max_seq = ids.len() + n_gen + 8;

    // 1) Batched graph prefill → populated cache + last-pos logits.
    let t0 = Instant::now();
    let (cache, prefill_logits) =
        prefill_inplace(&model, &rt, &ids, max_seq).expect("prefill");
    let prefill_ms = t0.elapsed().as_secs_f64() * 1e3;

    // 2) Seed the device caches from the graph prefill.
    let ts = Instant::now();
    let mut state = driver.new_state(max_seq).expect("new_state");
    driver.seed_prefill(&mut state, &cache).expect("seed");
    let seed_ms = ts.elapsed().as_secs_f64() * 1e3;

    // 3) First token from the prefill logits, then resident decode.
    // The graph prefill returns logits for ALL prompt positions
    // ([new_seq, vocab]); only the LAST row drives generation.
    let vocab_size = driver.config().vocab_size;
    let mut out_ids = ids.clone();
    let mut logits = prefill_logits[prefill_logits.len() - vocab_size..].to_vec();
    let mut step_ms: Vec<f64> = Vec::new();
    for _ in 0..n_gen {
        let next = argmax(&logits);
        out_ids.push(next);
        if Some(next) == vocab.eos_id {
            break;
        }
        let s = Instant::now();
        logits = driver.forward_token(next, &mut state).expect("decode");
        step_ms.push(s.elapsed().as_secs_f64() * 1e3);
    }

    println!("\nOUT: {}", vocab.decode_bpe(&out_ids));
    let mut sorted = step_ms.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let median = sorted.get(sorted.len() / 2).copied().unwrap_or(0.0);
    eprintln!(
        "\nHYBRID prefill {prefill_ms:.0} ms ({} tok, {:.1} ms/tok) + seed {seed_ms:.1} ms",
        ids.len(),
        prefill_ms / ids.len().max(1) as f64,
    );
    eprintln!(
        "decode WARM median {median:.1} ms/tok = {:.1} tok/s over {} tok",
        1000.0 / median.max(0.001),
        step_ms.len(),
    );
}

#[cfg(all(target_os = "macos", feature = "fabric"))]
fn argmax(v: &[f32]) -> u32 {
    let mut bi = 0usize;
    let mut bv = f32::NEG_INFINITY;
    for (i, &x) in v.iter().enumerate() {
        if x > bv {
            bv = x;
            bi = i;
        }
    }
    bi as u32
}

#[cfg(not(all(target_os = "macos", feature = "fabric")))]
fn main() {
    eprintln!("requires macOS + --features fabric");
}