open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/granite_decode.rs

//! Wave 52 bring-up: Granite-4.0-H hybrid end-to-end on the
//! real checkpoint — second hybrid family through the wave
//! 47 SSD primitives.
//!
//! ```bash
//! cargo run --release --features rayon --example granite_decode -- \
//!     /path/to/snapshot-dir [n_tokens] [prompt text]
//! ```

use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::granite_h::{GraniteHConfig, GraniteHModel};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: granite_decode <snapshot-dir> [n_tokens] [prompt]");
        std::process::exit(1);
    }
    let dir = PathBuf::from(&args[1]);
    let n_tokens: usize = args.get(2).map(|s| s.parse()).transpose()?.unwrap_or(8);
    let prompt_text = args.get(3).cloned()
        .unwrap_or_else(|| "The capital of France is".to_string());

    let cfg_bytes = std::fs::read(dir.join("config.json"))?;
    let config: GraniteHConfig = serde_json::from_slice(&cfg_bytes)?;
    let n_mamba = config.layer_types.iter().filter(|t| *t == "mamba").count();
    println!("config: {} layers ({} mamba + {} attention), hidden {}, vocab {}",
        config.num_hidden_layers, n_mamba,
        config.num_hidden_layers - n_mamba,
        config.hidden_size, config.vocab_size);

    // Single-file or sharded checkpoint: collect every
    // model*.safetensors in the snapshot dir.
    let t0 = Instant::now();
    let mut shard_paths: Vec<PathBuf> = std::fs::read_dir(&dir)?
        .filter_map(|e| e.ok().map(|e| e.path()))
        .filter(|p| {
            p.extension().is_some_and(|e| e == "safetensors")
                && p.file_name().is_some_and(
                    |n| n.to_string_lossy().starts_with("model"))
        })
        .collect();
    shard_paths.sort();
    let shards = shard_paths.iter()
        .map(|p| read_safetensors_file(p).map_err(|e| format!("{p:?}: {e:?}")))
        .collect::<Result<Vec<_>, _>>()?;
    let n_tensors: usize = shards.iter().map(|s| s.tensors.len()).sum();
    println!("checkpoint mapped in {:.2}s ({} shards, {} tensors)",
        t0.elapsed().as_secs_f64(), shards.len(), n_tensors);

    let model = GraniteHModel::bind_sharded(config, shards)?;
    let tokenizer = BpeTokenizer::from_file(dir.join("tokenizer.json"))?;
    println!("tokenizer: {} vocab", tokenizer.vocab_size());

    let prompt = tokenizer.encode(&prompt_text);
    println!("prompt {prompt_text:?} → {prompt:?}");

    let mut state = model.new_state();
    let t_pre = Instant::now();
    let mut logits_next = model.forward_prompt(&prompt, &mut state);
    let dt_pre = t_pre.elapsed();
    println!(
        "prefill: {} tokens in {:.0} ms ({:.1} tok/s)",
        prompt.len(),
        dt_pre.as_secs_f64() * 1e3,
        prompt.len() as f64 / dt_pre.as_secs_f64(),
    );
    let mut generated: Vec<u32> = Vec::new();
    println!("greedy decode:");
    for step in 0..n_tokens {
        let t = Instant::now();
        let logits = if step == 0 {
            std::mem::take(&mut logits_next)
        } else {
            model.forward_token(*generated.last().unwrap(), &mut state)
        };
        let dt = t.elapsed();
        let finite = logits.iter().all(|v| v.is_finite());
        let (argmax, max) = logits.iter().enumerate()
            .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
                if v > bv { (i, v) } else { (bi, bv) }
            });
        println!("  step {step:>2}: out={argmax:>6}  max={max:>8.3}  finite={finite}  {:.0} ms",
            dt.as_secs_f64() * 1e3);
        if !finite {
            eprintln!("FAILED: non-finite logits");
            std::process::exit(1);
        }
        generated.push(argmax as u32);
    }
    println!("\ngenerated text: {:?}", tokenizer.decode(&generated));
    #[cfg(feature = "fabric")]
    {
        let entries = jouleclaw_loader_gguf::fabric_dispatch::ledger_snapshot();
        if !entries.is_empty() {
            let total_gf: f64 = entries.iter().map(|e| e.gflops).sum();
            println!("\nfabric attribution (whole run):");
            for e in &entries {
                let uj = e
                    .projected_uj
                    .map(|u| format!("{:.1} J projected", u / 1e6))
                    .unwrap_or_else(|| "no calibration".into());
                println!(
                    "  {:>18}: {:>6} ops, {:>9.1} GFLOP ({:>4.1}%), {}",
                    e.member_id, e.ops, e.gflops,
                    100.0 * e.gflops / total_gf, uj,
                );
            }
        }
    }
    println!("PASS");
    Ok(())
}