Merge branch 'main' of github.com:evilsocket/cake

Author: evilsocket

Cargo.lock

@@ -42,10 +42,13 @@ candle-transformers = { version = "0.9" }
 candle-flash-attn = { version = "0.9", optional = true }
 candle-metal-kernels = { version = "0.9", optional = true }
 objc2-metal = { version = "0.3", optional = true }
+ash = { version = "0.38", optional = true, default-features = false, features = ["linked", "debug", "std"] }
+gpu-allocator = { version = "0.27", optional = true, default-features = false, features = ["vulkan"] }
 wgpu = { version = "24", optional = true }
 libloading = { version = "0.8", optional = true }
 pollster = { version = "0.4", optional = true }
 bytemuck = { version = "1", optional = true, features = ["derive"] }
+naga = { version = "24", optional = true, features = ["wgsl-in", "spv-out"] }
 half = "2"
 image = "0.25.2"
 hf-hub = "0.5"
@@ -69,7 +72,7 @@ default = ["master", "llama", "qwen2", "qwen3_5", "qwen3", "qwen3_moe", "qwen3_5
 metal = ["candle-core/metal", "candle-nn/metal", "candle-transformers/metal", "dep:candle-metal-kernels", "dep:objc2-metal"]
 cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda", "dep:bindgen_cuda"]
 flash-attn = ["cuda", "dep:candle-flash-attn"]
-vulkan = ["dep:wgpu", "dep:pollster", "dep:bytemuck"]
+vulkan = ["dep:ash", "dep:gpu-allocator", "dep:bytemuck"]
 rocm = ["dep:libloading"]
 
 master = ["dep:actix-web", "dep:async-stream", "dep:uuid"]
@@ -92,6 +95,7 @@ luxtts = ["dep:rustfft"]
 
 [build-dependencies]
 bindgen_cuda = { version = "0.1.6", optional = true }
+naga = { version = "24", features = ["wgsl-in", "spv-out"] }
 
 [dev-dependencies]
 tempfile = "3"

cake-core/Cargo.toml

@@ -1,5 +1,7 @@
-/// Vulkan backend benchmarks — GPU GEMV vs CPU matmul at model-realistic sizes,
-/// dispatch overhead, upload/download costs, and elementwise ops.
+/// Vulkan backend benchmarks — GPU matmul, elementwise ops, and full MLP.
+///
+/// All "vulkan_*" benchmarks exercise the GPU compute path.
+/// All "cpu_*" benchmarks use CPU-only candle ops for comparison.
 ///
 /// Run on Steam Deck: `cargo bench -p cake-core --features vulkan -- vulkan`
 
@@ -25,27 +27,47 @@ fn vulkan_dispatch_overhead(bencher: divan::Bencher) {
     bencher.bench_local(|| backend.silu_mul(&a, &b).unwrap());
 }
 
-// ── GPU GEMV vs CPU matmul at model sizes ────────────────────────────
+// ── GPU matmul at model sizes ────────────────────────────────────────
 // Qwen3-0.6B: hidden=1024, intermediate=3072, head_dim=128
-// QKV: (1,1024) × (1024,4096), O: (1,1024) × (1024,1024)
-// gate_up: (1,1024) × (1024,6144), down: (1,3072) × (3072,1024)
+// M=2 is the smallest prefill batch that goes through GPU (M>1).
 
-#[divan::bench(args = [1024, 4096, 6144])]
-fn vulkan_gemv_1024xN(bencher: divan::Bencher, n: usize) {
+#[divan::bench(args = [2, 8, 32, 64])]
+fn vulkan_gemm_Mx1024x4096(bencher: divan::Bencher, m: usize) {
     let backend = vk();
-    let a = cpu_tensor(&[1, 1024], 1100);
-    let b = cpu_tensor(&[1024, n], 1101);
+    let a = cpu_tensor(&[m, 1024], 1300);
+    let b = cpu_tensor(&[1024, 4096], 1301);
     bencher.bench_local(|| backend.matmul(&a, &b).unwrap());
 }
 
-#[divan::bench]
-fn vulkan_gemv_3072x1024(bencher: divan::Bencher) {
+#[divan::bench(args = [2, 8, 32, 64])]
+fn cpu_gemm_Mx1024x4096(bencher: divan::Bencher, m: usize) {
+    let a = cpu_tensor(&[m, 1024], 1300);
+    let b = cpu_tensor(&[1024, 4096], 1301);
+    bencher.bench_local(|| a.matmul(&b).unwrap());
+}
+
+// ── GPU matmul at other model shapes ─────────────────────────────────
+// gate_up: Mx1024x6144, down: Mx3072x1024
+
+#[divan::bench(args = [2, 8, 32])]
+fn vulkan_gemm_Mx1024x6144(bencher: divan::Bencher, m: usize) {
     let backend = vk();
-    let a = cpu_tensor(&[1, 3072], 1200);
+    let a = cpu_tensor(&[m, 1024], 1100);
+    let b = cpu_tensor(&[1024, 6144], 1101);
+    bencher.bench_local(|| backend.matmul(&a, &b).unwrap());
+}
+
+#[divan::bench(args = [2, 8, 32])]
+fn vulkan_gemm_Mx3072x1024(bencher: divan::Bencher, m: usize) {
+    let backend = vk();
+    let a = cpu_tensor(&[m, 3072], 1200);
     let b = cpu_tensor(&[3072, 1024], 1201);
     bencher.bench_local(|| backend.matmul(&a, &b).unwrap());
 }
 
+// ── CPU generation (M=1) baseline ────────────────────────────────────
+// M=1 uses CPU fallback (dispatch overhead > compute gain).
+
 #[divan::bench(args = [1024, 4096, 6144])]
 fn cpu_gemv_1024xN(bencher: divan::Bencher, n: usize) {
     let a = cpu_tensor(&[1, 1024], 1100);
@@ -60,34 +82,18 @@ fn cpu_gemv_3072x1024(bencher: divan::Bencher) {
     bencher.bench_local(|| a.matmul(&b).unwrap());
 }
 
-// ── GPU GEMM (prefill) at model sizes ────────────────────────────────
-
-#[divan::bench(args = [8, 32, 64])]
-fn vulkan_gemm_Mx1024x4096(bencher: divan::Bencher, m: usize) {
-    let backend = vk();
-    let a = cpu_tensor(&[m, 1024], 1300);
-    let b = cpu_tensor(&[1024, 4096], 1301);
-    bencher.bench_local(|| backend.matmul(&a, &b).unwrap());
-}
+// ── Elementwise ops — GPU path (large tensors) ──────────────────────
+// Above 8192 element threshold to ensure GPU dispatch.
 
-#[divan::bench(args = [8, 32, 64])]
-fn cpu_gemm_Mx1024x4096(bencher: divan::Bencher, m: usize) {
-    let a = cpu_tensor(&[m, 1024], 1300);
-    let b = cpu_tensor(&[1024, 4096], 1301);
-    bencher.bench_local(|| a.matmul(&b).unwrap());
-}
-
-// ── Elementwise ops at model sizes ───────────────────────────────────
-
-#[divan::bench(args = [1024, 3072])]
+#[divan::bench(args = [16384, 32768])]
 fn vulkan_silu_mul(bencher: divan::Bencher, size: usize) {
     let backend = vk();
     let gate = cpu_tensor(&[1, 1, size], 1400);
     let up = cpu_tensor(&[1, 1, size], 1401);
     bencher.bench_local(|| backend.silu_mul(&gate, &up).unwrap());
 }
 
-#[divan::bench(args = [1024, 3072])]
+#[divan::bench(args = [16384, 32768])]
 fn cpu_silu_mul(bencher: divan::Bencher, size: usize) {
     let gate = cpu_tensor(&[1, 1, size], 1400);
     let up = cpu_tensor(&[1, 1, size], 1401);
@@ -96,7 +102,7 @@ fn cpu_silu_mul(bencher: divan::Bencher, size: usize) {
     });
 }
 
-#[divan::bench(args = [1024, 3072])]
+#[divan::bench(args = [16384, 32768])]
 fn vulkan_add3(bencher: divan::Bencher, size: usize) {
     let backend = vk();
     let a = cpu_tensor(&[1, 1, size], 1500);
@@ -105,7 +111,7 @@ fn vulkan_add3(bencher: divan::Bencher, size: usize) {
     bencher.bench_local(|| backend.add3(&a, &b, &c).unwrap());
 }
 
-// ── RMS norm (CPU-only in current backend) ───────────────────────────
+// ── RMS norm (CPU fallback in current backend) ──────────────────────
 
 #[divan::bench(args = [1024, 3072])]
 fn vulkan_rms_norm_gated(bencher: divan::Bencher, size: usize) {
@@ -126,32 +132,34 @@ fn vulkan_add_rms_norm(bencher: divan::Bencher, size: usize) {
 }
 
 // ── Full MLP pass (gate_up + silu_mul + down) ────────────────────────
+// Prefill MLP at M=8: all ops go through GPU.
 
 #[divan::bench]
 fn vulkan_mlp_full(bencher: divan::Bencher) {
     let backend = vk();
-    let x = cpu_tensor(&[1, 1024], 1800);
-    let gate_up_w = cpu_tensor(&[6144, 1024], 1801);
-    let down_w = cpu_tensor(&[1024, 3072], 1802);
+    let x = cpu_tensor(&[8, 1024], 1800);
+    // Pre-transpose weights so they get cached on GPU across iterations
+    let gate_up_wt = cpu_tensor(&[6144, 1024], 1801).t().unwrap().contiguous().unwrap();
+    let down_wt = cpu_tensor(&[1024, 3072], 1802).t().unwrap().contiguous().unwrap();
     bencher.bench_local(|| {
-        let fused = backend.matmul(&x, &gate_up_w.t().unwrap()).unwrap();
+        let fused = backend.matmul(&x, &gate_up_wt).unwrap();
         let gate = fused.narrow(1, 0, 3072).unwrap().contiguous().unwrap();
         let up = fused.narrow(1, 3072, 3072).unwrap().contiguous().unwrap();
         let act = backend.silu_mul(&gate, &up).unwrap();
-        backend.matmul(&act, &down_w.t().unwrap()).unwrap()
+        backend.matmul(&act, &down_wt).unwrap()
     });
 }
 
 #[divan::bench]
 fn cpu_mlp_full(bencher: divan::Bencher) {
-    let x = cpu_tensor(&[1, 1024], 1800);
-    let gate_up_w = cpu_tensor(&[6144, 1024], 1801);
-    let down_w = cpu_tensor(&[1024, 3072], 1802);
+    let x = cpu_tensor(&[8, 1024], 1800);
+    let gate_up_wt = cpu_tensor(&[6144, 1024], 1801).t().unwrap().contiguous().unwrap();
+    let down_wt = cpu_tensor(&[1024, 3072], 1802).t().unwrap().contiguous().unwrap();
     bencher.bench_local(|| {
-        let fused = x.matmul(&gate_up_w.t().unwrap()).unwrap();
+        let fused = x.matmul(&gate_up_wt).unwrap();
         let gate = fused.narrow(1, 0, 3072).unwrap();
         let up = fused.narrow(1, 3072, 3072).unwrap();
         let act = (candle_nn::ops::silu(&gate).unwrap() * &up).unwrap();
-        act.matmul(&down_w.t().unwrap()).unwrap()
+        act.matmul(&down_wt).unwrap()
     });
 }

cake-core/benches/bench_vulkan.rs

@@ -1,4 +1,64 @@
 fn main() {
+    #[cfg(feature = "vulkan")]
+    {
+        println!("cargo::rerun-if-changed=src/backends/vulkan/ops.wgsl");
+
+        let wgsl_src = std::fs::read_to_string("src/backends/vulkan/ops.wgsl")
+            .expect("failed to read ops.wgsl");
+        let module = naga::front::wgsl::parse_str(&wgsl_src)
+            .expect("failed to parse WGSL");
+        let info = naga::valid::Validator::new(
+            naga::valid::ValidationFlags::all(),
+            naga::valid::Capabilities::empty(),
+        )
+        .validate(&module)
+        .expect("WGSL validation failed");
+
+        let options = naga::back::spv::Options {
+            lang_version: (1, 3),
+            ..Default::default()
+        };
+        // Generate one SPIR-V module per entry point
+        let entry_points: Vec<String> = module
+            .entry_points
+            .iter()
+            .map(|ep| ep.name.clone())
+            .collect();
+
+        let out_dir = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
+        let mut includes = String::new();
+
+        for ep_name in &entry_points {
+            let pipeline_options = naga::back::spv::PipelineOptions {
+                shader_stage: naga::ShaderStage::Compute,
+                entry_point: ep_name.clone(),
+            };
+            let spv_words = naga::back::spv::write_vec(
+                &module,
+                &info,
+                &options,
+                Some(&pipeline_options),
+            )
+            .unwrap_or_else(|e| panic!("SPIR-V generation failed for {ep_name}: {e}"));
+
+            // Write SPIR-V binary
+            let spv_path = out_dir.join(format!("ops_{ep_name}.spv"));
+            let bytes: Vec<u8> = spv_words.iter().flat_map(|w| w.to_le_bytes()).collect();
+            std::fs::write(&spv_path, &bytes).unwrap();
+
+            includes.push_str(&format!(
+                "(\"{ep_name}\", include_bytes!(concat!(env!(\"OUT_DIR\"), \"/ops_{ep_name}.spv\"))),\n"
+            ));
+        }
+
+        // Write a Rust file with all SPIR-V modules
+        let rs_path = out_dir.join("spirv_ops.rs");
+        let code = format!(
+            "static SPIRV_MODULES: &[(&str, &[u8])] = &[\n{includes}];\n"
+        );
+        std::fs::write(&rs_path, &code).expect("failed to write SPIR-V module list");
+    }
+
     #[cfg(feature = "cuda")]
     {
         println!("cargo::rerun-if-changed=src/backends/cuda/ops.cu");