inferencelayer 0.2.13

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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

//! Kernel spike: tiled shader-ALU GEMM (`enc_gemm3`, f32) vs cooperative-matrix GEMM
//! (`enc_gemm4_f16w`, tensor cores) at LLM prefill shapes. Answers the one question behind
//! BEATING_VLLM.md P0: how much does routing prefill matmuls through the tensor cores we ALREADY
//! ship (encoder/vision) buy us — the ceiling for wiring coop-matrix into forward.rs.
//!
//! Reports GFLOP/s for each and the ratio. Synthetic data, no model load. Run on the target GPU.
use anyhow::Result;
use inferencelayer::encoder::{enc_gemm3_src, enc_gemm4_f16w_src};
use inferencelayer::GpuCtx;
use wgpu::util::DeviceExt;

fn pipeline(ctx: &GpuCtx, label: &str, src: &str) -> wgpu::ComputePipeline {
    let m = ctx.device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some(label),
        source: wgpu::ShaderSource::Wgsl(src.into()),
    });
    ctx.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: Some(label),
        layout: None,
        module: &m,
        entry_point: Some("main"),
        compilation_options: Default::default(),
        cache: None,
    })
}

/// Times a GEMM pipeline at (m,n,k) with tile (bm,bn,bk). `bias_len` = 1 for gemm3, 8*n for
/// gemm4_f16w (its bias broadcasts over 8 rows). Returns GFLOP/s (best of 5).
#[allow(clippy::too_many_arguments)]
fn bench(ctx: &GpuCtx, label: &str, src: &str, m: usize, n: usize, k: usize,
         bm: usize, bn: usize, bias_len: usize, reps: usize) -> f64 {
    let pl = pipeline(ctx, label, src);
    let xb = ctx.storage(&vec![0.5f32; m * k]);
    let wb = ctx.storage(&vec![0.25f32; k * n]); // [k,n] transposed
    let bb = ctx.storage(&vec![0.0f32; bias_len]);
    let yb = ctx.storage(&vec![0f32; m * n]);
    let meta = ctx.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("meta"),
        contents: bytemuck::cast_slice(&[m as u32, n as u32, k as u32, 0u32]),
        usage: wgpu::BufferUsages::UNIFORM,
    });
    let bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &pl.get_bind_group_layout(0),
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: xb.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: wb.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: bb.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 3, resource: yb.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 4, resource: meta.as_entire_binding() },
        ],
    });
    let (gx, gy) = ((n as u32).div_ceil(bn as u32), (m as u32).div_ceil(bm as u32));
    let run = |reps: usize| -> f64 {
        let mut enc = ctx.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            pass.set_pipeline(&pl);
            pass.set_bind_group(0, &bg, &[]);
            for _ in 0..reps { pass.dispatch_workgroups(gx, gy, 1); }
        }
        let t0 = std::time::Instant::now();
        ctx.queue.submit([enc.finish()]);
        let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
        t0.elapsed().as_secs_f64()
    };
    run(2); // warm
    let mut best = f64::MAX;
    for _ in 0..5 { best = best.min(run(reps)); }
    2.0 * m as f64 * n as f64 * k as f64 * reps as f64 / best / 1e9
}

fn main() -> Result<()> {
    let ctx = GpuCtx::new()?;
    println!("caps: coop_matrix={} f16={} subgroups={}", ctx.coop_matrix, ctx.f16, ctx.subgroups);
    if !ctx.coop_matrix {
        println!("!! coop_matrix NOT available on this adapter — coop numbers will be absent");
    }
    // Qwen3-30B-A3B prefill-ish GEMM shapes (hidden 2048). Square shapes read cleanly as TFLOPS.
    let shapes = [
        (2048usize, 2048usize, 2048usize),
        (2624, 2048, 2048),   // prefill tokens × hidden
        (4096, 4096, 4096),   // larger — better tensor-core utilization
    ];
    let reps = 20;
    println!("{:>18} | {:>12} | {:>12} | {:>7}", "shape (m,n,k)", "ALU GFLOP/s", "coop GFLOP/s", "ratio");
    for (m, n, k) in shapes {
        let alu = bench(&ctx, "gemm3_alu", &enc_gemm3_src(false, 64, 64, 16), m, n, k, 64, 64, 1, reps);
        let coop = if ctx.coop_matrix {
            bench(&ctx, "gemm4_coop", &enc_gemm4_f16w_src(128, 128, 8), m, n, k, 128, 128, 8 * n, reps)
        } else { f64::NAN };
        println!("{:>6},{:>5},{:>5} | {:>12.0} | {:>12.0} | {:>6.1}x", m, n, k, alu, coop, coop / alu);
    }
    Ok(())
}