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

//! Generate from a GGUF and meter the decode on this process's CPU and GPU:
//! joules per token, CPU and GPU apart, with the GPU figure's grade and this
//! process's share of the GPU's time.
//!
//! ```text
//! Q8_PACKED=1 cargo run --release -p jouleclaw-runtime --example gguf_energy -- \
//!     /path/model.gguf "prompt" [n_tokens]
//! ```
//!
//! `Q8_PACKED=1` (Q8_0 weights), `Q4K_PACKED=1` / `Q6K_PACKED=1` (Q4_K / Q6_K)
//! keep the weights packed and run the decode GEMVs on the GPU
//! (`jouleclaw-backend-apple`); without them the decode runs on the CPU.
//!
//! The meter is `jouleclaw_energy::span::ProcessSpanMeter`: this process's CPU
//! energy (`process-cpu`, a hardware reading) plus its GPU energy — the IOReport
//! GPU rail when this process had at least 99% of every GPU client's time over
//! the span, `HwShunt` both. Where the GPU was shared (or a GPU client of
//! another process exited during the span) the GPU has no figure, and the span
//! is the CPU's alone, `Partial`, with why. Prefill (prompt through the first
//! token) and decode (every later token) are metered as two spans.

use jouleclaw_energy::span::{ProcessSpanMeter, SpanEnergy, SpanSource, SpanUses};
use jouleclaw_energy::EnergyError;
use jouleclaw_loader_gguf::read_gguf_file;
use jouleclaw_loader_gguf::sample::SamplingConfig;
use jouleclaw_loader_gguf::tokenizer::Vocab;
use jouleclaw_runtime::generate::{GenerateConfig, KvCacheKind, TokenizerKind};
use jouleclaw_runtime::streaming::Conversation;
use jouleclaw_runtime::Runtime;
use std::time::Instant;

fn report(what: &str, tokens: usize, secs: f64, e: &Result<SpanEnergy, EnergyError>) {
    let n = tokens.max(1) as f64;
    println!("{what}: {tokens} token(s) in {:.2} s ({:.1} ms/token)", secs, secs * 1e3 / n);
    match e {
        Ok(e) => {
            println!(
                "  total {:.4} J/token ({:.3} J), {:?}, coverage {:?}",
                e.uj as f64 / 1e6 / n,
                e.uj as f64 / 1e6,
                e.provenance,
                e.coverage
            );
            match &e.source {
                SpanSource::Process { cpu_uj, gpu } => match (gpu.all_ns, gpu.rail_uj) {
                    (Some(all_ns), Some(rail_uj)) => println!(
                        "  CPU {:.4} J/token (HwShunt) | GPU {:.4} J/token ({:?}): this process {:.1} ms of {:.1} ms \
                         GPU time ({:.2}%) of a {:.3} J GPU rail",
                        *cpu_uj as f64 / 1e6 / n,
                        gpu.uj as f64 / 1e6 / n,
                        gpu.provenance(),
                        gpu.self_ns as f64 / 1e6,
                        all_ns as f64 / 1e6,
                        gpu.share_ppm() as f64 / 1e4,
                        rail_uj as f64 / 1e6
                    ),
                    _ => println!(
                        "  CPU {:.4} J/token (HwShunt) | GPU 0 J/token, exact: this process spent no GPU time",
                        *cpu_uj as f64 / 1e6 / n
                    ),
                },
                SpanSource::CpuOnly { cpu_uj, why } => println!(
                    "  CPU alone {:.4} J/token (HwShunt), Partial: no GPU figure, because {why}",
                    *cpu_uj as f64 / 1e6 / n
                ),
                SpanSource::Fallback { meter, why } => println!("  fallback meter `{meter}`, because {why}"),
            }
        }
        Err(e) => println!("  no figure: {e}"),
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: gguf_energy <model.gguf> [prompt] [n_tokens]");
        std::process::exit(1);
    }
    let path = &args[1];
    let prompt = args.get(2).map(|s| s.as_str()).unwrap_or("The capital of France is");
    let n: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(64);

    let t0 = Instant::now();
    let model = read_gguf_file(path).expect("read_gguf_file");
    let vocab = Vocab::from_gguf(&model).expect("Vocab::from_gguf");
    let arch = model.metadata.get("general.architecture").and_then(|v| v.as_string()).unwrap_or("?").to_string();
    eprintln!("arch={arch} tensors={} | loaded in {:?}", model.tensors.len(), t0.elapsed());

    let meter = ProcessSpanMeter::new(jouleclaw_backend_apple::gpu_liveness_load);
    println!(
        "meter: process counters {} (GPU liveness {:?}); fallback `{}`",
        if meter.is_metered() { "live" } else { "NOT live" },
        meter.gpu().liveness(),
        meter.fallback_name()
    );

    let mut conv = Conversation::with_runtime(&model, &vocab, (n + 64).max(128), Runtime::boot())
        .expect("Conversation::with_runtime");
    let cfg = GenerateConfig {
        max_new_tokens: n,
        add_bos: true,
        tokenizer_kind: TokenizerKind::Auto,
        cache_kind: KvCacheKind::InPlace,
        sampling: SamplingConfig::greedy(),
        max_seq: Some((n + 64).max(128)),
        stop_strings: vec![],
    };

    let mut prefill_start = Some(meter.start());
    let t = Instant::now();
    let mut tokens: Vec<u32> = Vec::new();
    let mut prefill = None;
    let mut decode = None;
    for st in conv.extend(prompt, &cfg).expect("extend") {
        tokens.push(st.expect("decode step").id);
        if let Some(start) = prefill_start.take() {
            let secs = t.elapsed().as_secs_f64();
            prefill = Some((secs, meter.finish(start, SpanUses::CpuGpu)));
            decode = Some((Instant::now(), meter.start()));
        }
    }
    let decode_figure = decode.map(|(t, s)| (t.elapsed().as_secs_f64(), meter.finish(s, SpanUses::CpuGpu)));

    println!("\n>>> {}{}\n", prompt, vocab.decode_bpe(&tokens));
    if let Some((secs, e)) = &prefill {
        report("prefill (prompt + first token)", 1, *secs, e);
    }
    if let Some((secs, e)) = &decode_figure {
        report("decode", tokens.len().saturating_sub(1), *secs, e);
    }
}