Skip to main content

cherenkov/
kernels.rs

1//! Metal translation units shared by the engine and GPU tests.
2//!
3//! Files are fragments, assembled at Rust compile time. Keep shared types
4//! and helpers before their users; no filesystem includes are needed at
5//! runtime. Each fragment gets its own filename/line numbers in diagnostics.
6
7macro_rules! metal_library {
8    ($($path:literal),+ $(,)?) => {
9        concat!(
10            "#include <metal_stdlib>\nusing namespace metal;\n",
11            $(
12                "\n#line 1 \"", $path, "\"\n",
13                include_str!(concat!("../kernels/", $path)),
14                "\n",
15            )+
16        )
17    };
18}
19
20pub(crate) const FORWARD_MSL: &str = metal_library!(
21    "common/quantized.metal",
22    "common/gemm.metal",
23    "common/attention.metal",
24    "common/deltanet.metal",
25    "common/elementwise.metal",
26    "common/sampling.metal",
27);
28
29pub(crate) const BATCH_MSL: &str = metal_library!(
30    "qwen4_exp/types.metal",
31    // Q4 helpers serve both hyper-connections and experts. Q2/Q3 are
32    // expert-only and stay with the expert kernels.
33    "qwen4_exp/quantized_rows.metal",
34    "qwen4_exp/hyperconnection.metal",
35    "qwen4_exp/experts.metal",
36    "qwen4_exp/expert_gemm.metal",
37    "qwen4_exp/ple.metal",
38    "qwen4_exp/mtp.metal",
39    "qwen4_exp/deltanet.metal",
40    "qwen4_exp/qsa.metal",
41    "qwen4_exp/rows.metal",
42);
43
44/// Device diagnostic, compiled separately from model kernels.
45pub(crate) const CLOCK_PROBE_MSL: &str = metal_library!("device/clock.metal");