memra-engine 0.84.1

From-scratch CUDA LLM inference engine for NVIDIA RTX 50-series (sm_120a) and Hopper (sm_90a) - custom kernels, no frameworks
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

//! DFlash draft-forward parity gate (bring-up step 2, DFLASH-BRINGUP-PLAN.md): the memra
//! forward on the real checkpoint must reproduce the torch reference (tools/
//! dflash_oracle.py flat dumps in /data/cache) on the fixed-seed synthetic inputs.
//! PASS bar: final hidden rel maxdiff < 1e-3 class (f32-vs-f32, same math different
//! kernel FP order) AND per-layer drift monotone (bisect handle if final fails).
use memra_engine::Engine;
use memra_engine::dflash::DflashDraft;

fn read_f32(p: &str) -> Vec<f32> {
    let b = std::fs::read(p).unwrap_or_else(|e| panic!("{p}: {e}"));
    b.chunks_exact(4)
        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect()
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let ckpt = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "/data/ai-ml/hf-models/dspark-gemma4-31b-draft/backbone-only".into());
    let cache = std::env::args()
        .nth(2)
        .unwrap_or_else(|| "/data/cache".into());
    let e = Engine::new(0)?;
    let m = DflashDraft::load(&e, std::path::Path::new(&ckpt))?;
    let c = &m.cfg;
    println!(
        "loaded dflash draft: {} layers, hidden {}, block {}, taps {:?}",
        c.n_layer, c.hidden, c.block_size, c.target_layer_ids
    );

    let th = read_f32(&format!("{cache}/dflash-target_hidden.f32"));
    let ne = read_f32(&format!("{cache}/dflash-noise_embedding.f32"));
    let ctx = th.len() / (c.target_layer_ids.len() * c.hidden);
    assert_eq!(ne.len(), c.block_size * c.hidden);
    let th_d = e.htod(&th)?;
    let ne_d = e.htod(&ne)?;
    let pos: Vec<i32> = (0..(ctx + c.block_size) as i32).collect();

    let out = m.forward(&e, &th_d, &ne_d, &pos, ctx)?;
    let got = e.dtoh(&out)?;
    let want = read_f32(&format!("{cache}/dflash-final.f32"));
    assert_eq!(got.len(), want.len());
    let (mut md, mut mi) = (0f32, 0usize);
    for (i, (a, b)) in got.iter().zip(&want).enumerate() {
        let d = (a - b).abs();
        if d > md {
            md = d;
            mi = i;
        }
    }
    // Bar calibration (bisect 2026-07-13): every stage isolated — xn 1e-6, ctx_features
    // 1e-3, per-layer drift FLAT at 3-7e-4 rel-to-max across all 5 layers (no structural
    // bug); the noise seed is the cuBLASLt f32 GEMM riding TF32-class compute (q0 maxdiff
    // 1.5e-2 on a 5376-K dot with bit-identical inputs/weights) amplified by the draft's
    // ~10k-scale activations. Bar = maxdiff vs max|ref| < 2e-3 (TF32 class); the ROUND
    // gates (acceptance + verify exactness) are the real oracle downstream.
    let mx = want.iter().fold(0f32, |a, v| a.max(v.abs()));
    let rel = md / mx;
    println!(
        "final: maxdiff {md:.3e} (idx {mi}: got {} want {}), max|ref| {mx:.2}, rel-to-max {rel:.3e}",
        got[mi], want[mi]
    );
    let pass = rel < 2e-3;
    println!("DFLASH-PARITY: {}", if pass { "PASS" } else { "FAIL" });
    std::process::exit(if pass { 0 } else { 1 });
}