matvec_bw/
matvec_bw.rs

1//! Weight-path bandwidth probe (scratch diagnostic, not a product).
2//!
3//! Two modes, both isolating the weight path from GDN/attention/sampling:
4//!
5//! - `one <tensor>`: repeat ONE tensor. A single dispatch amortized over a
6//!   large matrix — measures the kernel's ceiling. NOTE: tensors small
7//!   enough to sit in the SLC report inflated GB/s; only the ~636 MB head
8//!   is a trustworthy DRAM number here.
9//! - `sweep` (default): walk EVERY 2-D q8_2f tensor once, in directory
10//!   order — the real decode access pattern (cold weights, one dispatch
11//!   per tensor, ~200 dispatches). This is the honest weight-path number.
12//!
13//! Usage: cargo run --release --example matvec_bw -- <model.cmf> [sweep|one <tensor>]
14
15use cortiq_core::TensorDtype;
16use cortiq_engine::pool::Pool;
17use cortiq_engine::qtensor::QTensor;
18use std::sync::Arc;
19use std::time::Instant;
20
21fn main() {
22    let mut args = std::env::args().skip(1);
23    let path = args
24        .next()
25        .expect("usage: matvec_bw <model.cmf> [sweep|one <tensor>]");
26    let mode = args.next().unwrap_or_else(|| "sweep".to_string());
27    let model = Arc::new(cortiq_core::CmfModel::open(&path).expect("open model"));
28
29    // Real LM activations carry a few heavy channels (>8·rms); measured
30    // mean on this model is ~3.7. NOUT models that distribution.
31    let nout: usize = std::env::var("NOUT")
32        .ok()
33        .and_then(|v| v.parse().ok())
34        .unwrap_or(4);
35    let mk_x = |cols: usize| -> Vec<f32> {
36        let mut x: Vec<f32> = (0..cols).map(|i| ((i % 17) as f32 - 8.0) / 8.0).collect();
37        for k in 0..nout {
38            x[k * 37 % cols] = 40.0;
39        }
40        x
41    };
42
43    if mode == "one" {
44        let name = args
45            .next()
46            .unwrap_or_else(|| "model.embed_tokens.weight".to_string());
47        let entry = model.tensor(&name).expect("tensor not found");
48        let (rows, cols) = (entry.shape[0], entry.shape[1]);
49        let nbytes = entry.nbytes as f64;
50        println!(
51            "tensor {name}: {rows}x{cols} {:?} = {:.1} MB, NOUT={nout}",
52            entry.dtype,
53            nbytes / 1e6
54        );
55        let t = QTensor::from_model(&model, &name).expect("wrap");
56        let x = mk_x(cols);
57        let mut out = vec![0f32; rows];
58        t.matvec(&x, &mut out, None);
59        for nt in [1usize, 2, 4, 6, 8, 10] {
60            let pool = if nt == 1 { None } else { Some(Pool::new(nt)) };
61            let iters = 8;
62            t.matvec(&x, &mut out, pool.as_ref());
63            let t0 = Instant::now();
64            for _ in 0..iters {
65                t.matvec(&x, &mut out, pool.as_ref());
66            }
67            let el = t0.elapsed().as_secs_f64();
68            println!(
69                "threads={nt:2}  {:6.2} ms/matvec  {:6.1} GB/s (sink {:.3})",
70                el / iters as f64 * 1e3,
71                nbytes * iters as f64 / el / 1e9,
72                out[0]
73            );
74        }
75        return;
76    }
77
78    // sweep: every 2-D q8_2f tensor once = one decode's worth of weights.
79    let names: Vec<String> = model
80        .tensors
81        .iter()
82        .filter(|t| t.dtype == TensorDtype::Q8_2f && t.shape.len() == 2)
83        .map(|t| t.name.clone())
84        .collect();
85    let total_bytes: f64 = model
86        .tensors
87        .iter()
88        .filter(|t| t.dtype == TensorDtype::Q8_2f && t.shape.len() == 2)
89        .map(|t| t.nbytes as f64)
90        .sum();
91    println!(
92        "sweep: {} q8_2f tensors, {:.2} GB total (= weights streamed per decode token), NOUT={nout}",
93        names.len(),
94        total_bytes / 1e9
95    );
96
97    let tensors: Vec<(QTensor, Vec<f32>, Vec<f32>)> = names
98        .iter()
99        .map(|n| {
100            let e = model.tensor(n).unwrap();
101            let (rows, cols) = (e.shape[0], e.shape[1]);
102            (
103                QTensor::from_model(&model, n).expect("wrap"),
104                mk_x(cols),
105                vec![0f32; rows],
106            )
107        })
108        .collect();
109    let mut tensors = tensors;
110
111    // Whole-model residency pass BEFORE any timing: the first touch of a
112    // 4.2 GB mmap faults ~260k pages, which would otherwise be charged to
113    // whichever thread count happens to run first. REVERSE=1 flips the
114    // order as a check that no first-touch cost is left in the table.
115    for _ in 0..2 {
116        for (t, x, out) in tensors.iter_mut() {
117            t.matvec(x, out, None);
118        }
119    }
120    let mut counts = vec![1usize, 2, 4, 6, 8, 10];
121    if std::env::var("REVERSE").is_ok() {
122        counts.reverse();
123    }
124    for nt in counts {
125        let pool = if nt == 1 { None } else { Some(Pool::new(nt)) };
126        // one warm pass, then two measured
127        for (t, x, out) in tensors.iter_mut() {
128            t.matvec(x, out, pool.as_ref());
129        }
130        let iters = 2;
131        let t0 = Instant::now();
132        for _ in 0..iters {
133            for (t, x, out) in tensors.iter_mut() {
134                t.matvec(x, out, pool.as_ref());
135            }
136        }
137        let el = t0.elapsed().as_secs_f64();
138        let per_tok = el / iters as f64;
139        println!(
140            "threads={nt:2}  {:7.1} ms/sweep  {:6.1} GB/s  -> weight-path-only ceiling {:5.1} tok/s",
141            per_tok * 1e3,
142            total_bytes * iters as f64 / el / 1e9,
143            1.0 / per_tok
144        );
145    }
146}
matvec_bw/matvec_bw.rs

matvec_bw/
matvec_bw.rs
