inferencelayer 0.2.4

Kortexya's engine-native inference layer — LLM generation + embedding/encoder family on wgpu (WGSL kernels, any adapter) with a pure-Rust CPU fallback
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
67
68
69
70
71
72
73

//! DeepEncoder GPU attention benchmark — for running on the V100 (Vulkan) to get the definitive,
//! unthrottled absolute numbers. Prints the wgpu backend (to confirm the V100, not llvmpipe) and the
//! naive-dense vs register-flash SAM-attention head-to-head at real scale.
use inferencelayer::deepencoder_gpu::bench_sam_attention;
use inferencelayer::GpuCtx;

fn main() -> anyhow::Result<()> {
    let ctx = GpuCtx::new()?;
    println!("backend: {}", ctx.backend);
    println!(
        "caps: subgroups={} coop_matrix={} f16={}",
        ctx.subgroups, ctx.coop_matrix, ctx.f16
    );

    // GLM vision tower bench: GLM_DIR=<checkpoint dir> times the GPU tower forward at real page
    // grids — isolates the tower's share of OCR latency from the decoder's.
    if let Ok(dir) = std::env::var("GLM_DIR") {
        use inferencelayer::vision::ImagePatches;
        let cpu = inferencelayer::vision_glm::GlmVisionTower::load(std::path::Path::new(&dir))?;
        let pd = cpu.cfg.patch_dim();
        let fill = |len: usize| -> Vec<f32> {
            let mut s = 7u32;
            (0..len)
                .map(|_| {
                    s = s.wrapping_mul(1664525).wrapping_add(1013904223);
                    ((s >> 8) as f32 / 16_777_216.0 - 0.5) * 0.2
                })
                .collect()
        };
        // gemm3 tile sweep at the swept-best rb=16, f16 weights (deploy config)
        for (bm, bn, bk) in [(128usize,128usize,16usize),(128,128,32),(64,128,16),(128,64,16),(64,64,16)] {
            let gpu = inferencelayer::vision_glm_gpu::GlmVisionGpu::new_full(&ctx, &cpu, 16, true, bm, bn, bk)?;
            for (label, gh, gw) in [("~760x980", 54u32, 70u32), ("~1240x1600", 88u32, 114u32)] {
                let n = (gh * gw) as usize;
                let img = ImagePatches { patches: fill(n * pd), grid: [1, gh, gw] };
                let _ = gpu.forward(&ctx, &[img.clone()])?;
                let t = std::time::Instant::now();
                for _ in 0..3 { let _ = gpu.forward(&ctx, &[img.clone()])?; }
                let ms = t.elapsed().as_secs_f64() * 1000.0 / 3.0;
                println!("GLM tower tile={bm}x{bn}x{bk} {label}: {ms:.0} ms");
            }
        }
        if std::env::var("GLM_RB_SWEEP").is_err() { return Ok(()); }
        for rb in [4usize, 8, 16, 32] {
            let gpu = inferencelayer::vision_glm_gpu::GlmVisionGpu::new_with_rb(&ctx, &cpu, rb)?;
            for (label, gh, gw) in [("~760x980", 54u32, 70u32), ("~1240x1600", 88u32, 114u32)] {
                let n = (gh * gw) as usize;
                let img = ImagePatches {
                    patches: fill(n * pd),
                    grid: [1, gh, gw],
                };
                let _ = gpu.forward(&ctx, &[img.clone()])?; // warm
                let t = std::time::Instant::now();
                let iters = 3;
                for _ in 0..iters {
                    let _ = gpu.forward(&ctx, &[img.clone()])?;
                }
                let ms = t.elapsed().as_secs_f64() * 1000.0 / iters as f64;
                println!("GLM tower rb={rb:<2} {label} ({n} patches): {ms:.0} ms");
            }
        }
        return Ok(());
    }
    // real SAM global-attention scale: 64×64 grid, 768d, 12 heads (hd=64)
    let (gh, gw, width, heads, iters) = (64usize, 64usize, 768usize, 12usize, 5usize);
    let (naive, flash) = bench_sam_attention(&ctx, gh, gw, width, heads, iters);
    let n = gh * gw;
    println!("=== SAM global attention ({}²,{}d,{}h) on this adapter ===", gh, width, heads);
    println!("  naive-dense (805MB n² buffer, 3 passes):  {:8.2} ms", naive);
    println!("  register-flash (no n² buffer, portable):  {:8.2} ms", flash);
    println!("  SPEEDUP: {:.2}×   (naive n² buffer = {:.0} MB)", naive / flash, (heads * n * n * 4) as f64 / 1.0e6);
    Ok(())
}