rustytorch/crates/models/rtx-csm/examples/qmm_layer_diff.rs

//! Diagnose where qmatmul diverges in our actual model: load layer-0 q_proj
//! from a real CSM GGUF, run through both Self::QTensor and Self::TensorF16
//! variants on the SAME input, compare.

use anyhow::Result;
use candle_core::quantized::{QMatMul, QTensor};
use candle_core::{DType, Device, Module, Tensor};
use candle_transformers::quantized_var_builder::VarBuilder as QVarBuilder;
use std::sync::Arc;

fn rms(a: &[f32], b: &[f32]) -> f32 {
    let n = a.len() as f32;
    let mut s = 0.0f32;
    for (x, y) in a.iter().zip(b) {
        s += (x - y).powi(2);
    }
    (s / n).sqrt()
}

fn main() -> Result<()> {
    let backend = std::env::args().nth(1).unwrap_or_else(|| "metal".into());
    let gguf = std::env::args()
        .nth(2)
        .unwrap_or_else(|| "/tmp/csm-q8-v2.gguf".into());
    let device = match backend.as_str() {
        "cpu" => Device::Cpu,
        _ => Device::new_metal(0)?,
    };
    println!("device: {device:?}");
    println!("gguf:   {gguf}");

    // Build TWO QVarBuilders: one with DEQUANTIZE_ALL_F16, one without.
    // Each call to vb.get returns an Arc<QTensor> regardless; what differs is
    // how QMatMul::from_arc wraps it later.
    let vb = QVarBuilder::from_gguf(&gguf, &device)?;
    let key = "backbone.layers.0.attn.q_proj.weight";
    let qt_arc = vb
        .pp("backbone.layers.0.attn.q_proj")
        .get_no_shape("weight")?;
    println!(
        "{key} shape={:?} dtype={:?}",
        qt_arc.shape(),
        qt_arc.dtype()
    );

    // Single QMatMul: candle's thread-locals initialize from env vars on first
    // access in this process. Caller controls via shell env. The variant is
    // private but we can introspect indirectly: TensorF16 path internally
    // casts xs.to_dtype(F16); we can detect by feeding F32 xs and checking if
    // the output is bit-identical to the manual F16-then-matmul path.
    let qmm_qtensor = QMatMul::from_arc(qt_arc.clone())?;
    let qmm_f16 = QMatMul::from_arc(qt_arc.clone())?;
    println!(
        "DEQUANTIZE_ALL_F16 env: {:?}",
        std::env::var("CANDLE_DEQUANTIZE_ALL_F16")
    );

    // Run with a fixed input so both paths see the exact same xs.
    let (n, k) = qt_arc.shape().dims2()?;
    let m = 8usize;
    let xs = Tensor::randn(0.0f32, 1.0, (m, k), &device)?;
    println!("xs shape=({m}, {k})  weight shape=({n}, {k})");

    // F32 input, output dtype tells us which variant we got.
    let out_q = qmm_qtensor.forward(&xs)?;
    let out_f = qmm_f16.forward(&xs)?;
    // Manual F16 reference: dequantize to F16, transpose, matmul.
    let deq_f16 = qt_arc.dequantize_f16(&device)?;
    let manual_f16 = xs
        .to_dtype(DType::F16)?
        .matmul(&deq_f16.t()?)?
        .to_dtype(DType::F32)?;

    let v_q = out_q
        .flatten_all()?
        .to_dtype(DType::F32)?
        .to_vec1::<f32>()?;
    let v_f = out_f
        .flatten_all()?
        .to_dtype(DType::F32)?
        .to_vec1::<f32>()?;
    let v_m = manual_f16.flatten_all()?.to_vec1::<f32>()?;
    let r = rms(&v_q, &v_f);
    let r_qm = rms(&v_q, &v_m);
    let r_fm = rms(&v_f, &v_m);
    println!("manual F16 ref vs out_q: rms={r_qm:.6}");
    println!("manual F16 ref vs out_f: rms={r_fm:.6}");
    println!(
        "QTensor variant out dtype={:?}, F16-deq variant out dtype={:?}",
        out_q.dtype(),
        out_f.dtype()
    );
    println!("L2 distance (rms) between QTensor and F16-dequant outputs: {r:.6}");
    println!("first 8 values:");
    println!("  QTensor:    {:?}", &v_q[..8.min(v_q.len())]);
    println!("  TensorF16:  {:?}", &v_f[..8.min(v_f.len())]);

    // Also dequantize directly and compare to QTensor path.
    let deq_arc: Arc<QTensor> = qt_arc.clone();
    let deq = deq_arc.dequantize(&device)?.to_dtype(DType::F32)?;
    let out_deq = xs.matmul(&deq.t()?)?;
    let v_d = out_deq.flatten_all()?.to_vec1::<f32>()?;
    let rd = rms(&v_q, &v_d);
    println!("L2 between QTensor and manual dequant matmul: {rd:.6}");
    println!("  ManualDeq: {:?}", &v_d[..8.min(v_d.len())]);
    Ok(())
}