rlx-gemma 0.2.11

// RLX — versatile ML compiler + runtime. GPLv3.
//! Diagnose MLX/wgpu E2B drift by comparing per-position logits (all 16
//! prefill rows) against CPU on a single forward pass. CPU is bit-exact
//! to HF transformers, so it's the ground truth here.
//!
//! Run:
//! ```bash
//! cargo run --release -p rlx-gemma --features "apple-silicon" --example e2b_logits_diff
//! ```

use std::collections::HashMap;
use std::path::PathBuf;

use anyhow::{Context, Result, bail};
use rlx_core::flow_util::compile_graph_gemma_prefill_with_params;
use rlx_gemma::config::GemmaConfig;
use rlx_gemma::qat_loader::GemmaQatLoader;
use rlx_runtime::Device;

const PROMPT_IDS: &[u32] = &[818, 5279, 529, 7001, 563];
const BUCKET: usize = 16;

fn ckpt_dir() -> Result<PathBuf> {
    if let Some(d) = std::env::var_os("RLX_GEMMA4_E2B_DIR") {
        let p = PathBuf::from(d);
        return Ok(p);
    }
    let home = std::env::var_os("HOME").context("HOME")?;
    let base = std::path::Path::new(&home).join(
        ".cache/huggingface/hub/\
         models--google--gemma-4-E2B-it-qat-mobile-transformers/snapshots",
    );
    Ok(std::fs::read_dir(&base)?
        .flatten()
        .next()
        .context("snapshot")?
        .path())
}

fn run(device: Device, ple: &[f32], ids_f32: &[f32]) -> Result<Vec<f32>> {
    let dir = ckpt_dir()?;
    let cfg = GemmaConfig::from_file(&dir.join("config.json"))?;
    let mut bld = GemmaQatLoader::open(&dir)?;
    let mut packed = HashMap::new();
    let (graph, params) = rlx_gemma::builder::build_gemma_graph_sized_packed_ext(
        &cfg,
        &mut bld,
        1,
        BUCKET,
        true,
        false,
        false,
        &mut packed,
        None,
        None,
    )?;
    let mut compiled = compile_graph_gemma_prefill_with_params(device, graph, params)?;
    let outs = compiled.run(&[("input_ids", ids_f32), ("per_layer_inputs", ple)]);
    if outs.is_empty() {
        bail!("no outputs");
    }
    Ok(outs.into_iter().next().unwrap())
}

fn argmax(v: &[f32]) -> (u32, f32) {
    v.iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .map(|(i, &x)| (i as u32, x))
        .unwrap_or((0, 0.0))
}

fn main() -> Result<()> {
    let dir = ckpt_dir()?;
    let cfg = GemmaConfig::from_file(&dir.join("config.json"))?;
    let loader = GemmaQatLoader::open(&dir)?;
    let vocab = cfg.vocab_size;
    // Use the post-tok#1 ids: this is exactly what the GENERATE loop hands
    // to MLX on its second call. CPU gets tok#2=563 from this. MLX does not.
    let mut ids = vec![0u32; BUCKET];
    for (i, &t) in PROMPT_IDS.iter().enumerate() {
        ids[i] = t;
    }
    ids[5] = 7001; // tok#1, generated by CPU/Metal/MLX on the first pass
    let ple = loader.compute_per_layer_inputs(&cfg, &ids)?;
    let ids_f32: Vec<f32> = ids.iter().map(|&i| i as f32).collect();

    println!(
        "→ Prompt+tok#1: {:?} (positions 0..6 valid; 6..16 padded 0)",
        &ids[..6]
    );

    let cpu = run(Device::Cpu, &ple, &ids_f32)?;
    let mlx = run(Device::Mlx, &ple, &ids_f32).ok();
    let metal = run(Device::Metal, &ple, &ids_f32).ok();

    println!("\n── Per-position argmax (positions 0..6 — position p predicts tok p+1) ──");
    println!("    pos │   CPU argmax (max)   │   Metal              │   MLX              ");
    println!("    ────┼──────────────────────┼──────────────────────┼─────────────────────");
    for pos in 0..6 {
        let c = argmax(&cpu[pos * vocab..(pos + 1) * vocab]);
        let m = metal
            .as_ref()
            .map(|v| argmax(&v[pos * vocab..(pos + 1) * vocab]));
        let x = mlx
            .as_ref()
            .map(|v| argmax(&v[pos * vocab..(pos + 1) * vocab]));
        let m_s = m
            .map(|(t, v)| format!("{t:>6} ({v:>+8.3})"))
            .unwrap_or_else(|| "—".into());
        let x_s = x
            .map(|(t, v)| format!("{t:>6} ({v:>+8.3})"))
            .unwrap_or_else(|| "—".into());
        let same_m = m.map(|(t, _)| t == c.0).unwrap_or(false);
        let same_x = x.map(|(t, _)| t == c.0).unwrap_or(false);
        let m_tag = if same_m { "✓" } else { "✗" };
        let x_tag = if same_x { "✓" } else { "✗" };
        println!(
            "    {pos:>3} │ {:>6} ({:>+8.3}) │ {m_s} {m_tag} │ {x_s} {x_tag}",
            c.0, c.1
        );
    }

    // L2 distance per row vs CPU
    println!("\n── L2(MLX-CPU) / L2(Metal-CPU) per position ──");
    for pos in 0..6 {
        let cpu_row = &cpu[pos * vocab..(pos + 1) * vocab];
        let metal_l2 = metal.as_ref().map(|v| {
            let r = &v[pos * vocab..(pos + 1) * vocab];
            let d: f64 = r
                .iter()
                .zip(cpu_row.iter())
                .map(|(a, b)| (*a as f64 - *b as f64).powi(2))
                .sum();
            d.sqrt()
        });
        let mlx_l2 = mlx.as_ref().map(|v| {
            let r = &v[pos * vocab..(pos + 1) * vocab];
            let d: f64 = r
                .iter()
                .zip(cpu_row.iter())
                .map(|(a, b)| (*a as f64 - *b as f64).powi(2))
                .sum();
            d.sqrt()
        });
        println!(
            "    pos {pos}: metal={:.3e}  mlx={:.3e}",
            metal_l2.unwrap_or(f64::NAN),
            mlx_l2.unwrap_or(f64::NAN)
        );
    }

    Ok(())
}