qualia_core_db/inference/
metal_lane.rs

1//! Metal mega-pass orchestrator: chain all transformer layers into one Metal
2//! command buffer with a single fence at the end. No per-layer readback.
3//!
4//! This mirrors the CUDA `cuda_lane/mega_pass.rs` architecture but targets
5//! Apple Silicon via `metal-rs` or wgpu's Metal backend. On non-Apple
6//! platforms, all functions return `None`/`false` — the wgpu WGSL path
7//! handles inference instead.
8//!
9//! ## Architecture
10//!
11//! 1. Upload hidden state to device buffer (once).
12//! 2. For each layer:
13//!    a. RMSNorm + QKV + RoPE (fused dispatch via `fused-qkv-rope` MSL kernel)
14//!    b. KV cache write
15//!    c. SDPA decode (via `sdpa-decode` MSL kernel)
16//!    d. O-proj GEMV + residual add
17//!    e. RMSNorm + SwiGLU (fused)
18//!    f. Down GEMV + residual add
19//! 3. Output norm + logits GEMV + argmax
20//! 4. Single readback of the final token
21//!
22//! All dispatches share one command buffer — the GPU never idles waiting
23//! for CPU between layers.
24
25#![allow(dead_code, unused_variables)]
26
27/// Per-layer weight references for the Metal mega-pass.
28pub struct MetalPassLayerWeights<'a> {
29    pub attn_norm: &'a [f32],
30    pub q_raw: &'a [u8],
31    pub k_raw: &'a [u8],
32    pub v_raw: &'a [u8],
33    pub o_raw: &'a [u8],
34    pub ffn_norm: &'a [f32],
35    pub gate_raw: &'a [u8],
36    pub up_raw: &'a [u8],
37    pub down_raw: &'a [u8],
38}
39
40/// Per-layer matmul dimensions (pre-computed by caller).
41pub struct MetalPassLayerDims {
42    pub q_in: usize,
43    pub q_out: usize,
44    pub kv_in: usize,
45    pub kv_out: usize,
46    pub o_in: usize,
47    pub o_out: usize,
48    pub gate_in: usize,
49    pub gate_out: usize,
50    pub up_in: usize,
51    pub up_out: usize,
52    pub down_in: usize,
53    pub down_out: usize,
54}
55
56/// Metal mega-pass: chain all transformer layers into one Metal command buffer
57/// with a **single fence** at the end. No per-layer readback. The hidden state
58/// stays on-device for the entire forward pass; only the final logits/token
59/// come back to host.
60///
61/// This is the Apple Silicon equivalent of `try_cuda_mega_pass`.
62///
63/// **Requirements:**
64/// - `metal-rs` crate (or wgpu Metal backend) must be available.
65/// - Device KV cache must be initialized.
66/// - All weights must be resident on-device.
67///
68/// Returns `Some(token)` on success, `None` if Metal is unavailable or
69/// any kernel fails.
70pub fn try_metal_mega_pass(
71    n_embd: usize,
72    n_head: usize,
73    n_kv: usize,
74    head_dim: usize,
75    n_layer: u32,
76    token_idx: u32,
77    max_context: u32,
78    layer_stride: u32,
79    slot_kv_elems: u32,
80    rope_base: f32,
81    rope_scale: f32,
82    rms_eps: f32,
83    hidden: &[f32],
84    layers: &[MetalPassLayerWeights<'_>],
85    layer_dims: &[MetalPassLayerDims],
86    output_norm: Option<&[f32]>,
87    lm_head_raw: Option<&[u8]>,
88    lm_head_in: usize,
89    lm_head_out: usize,
90) -> Option<u32> {
91    // Metal is only available on macOS. On other platforms, return None
92    // and let the wgpu WGSL path handle inference.
93    #[cfg(not(target_os = "macos"))]
94    {
95        None
96    }
97    #[cfg(target_os = "macos")]
98    {
99        // TODO: implement via metal-rs when available on the build target.
100        // The MSL kernels (rmsnorm, fused-qkv-rope, sdpa-decode, gemv-simd-matrix)
101        // are already emitted by the MSL emitter. The execution bridge needs
102        // `WgpuPipeline::compile_msl` to be implemented with metal-rs.
103        log::info!("metal_mega_pass: not yet implemented on this build");
104        None
105    }
106}
107
108/// Check if the Metal mega-pass is available on this platform.
109pub fn metal_mega_pass_available() -> bool {
110    cfg!(target_os = "macos")
111}
112
113/// Warm the Metal context (pre-compile kernels, pre-allocate arenas).
114/// On non-Apple platforms, this is a no-op.
115pub fn warm_metal_context() {
116    #[cfg(target_os = "macos")]
117    {
118        // TODO: pre-compile MSL kernels and cache them
119    }
120}
121
122/// MSL double-buffered weight streaming: overlap H2D weight copies with
123/// compute using `MTLBlitCommandEncoder` on a separate Metal command buffer.
124/// This is the Metal equivalent of CUDA's `write_view_prefetch` + `join_prefetch`
125/// on a secondary stream.
126///
127/// **Architecture:**
128/// 1. Issue async H2D copies for next layer's weights via `MTLBlitCommandEncoder`
129///    on a dedicated command buffer (prefetch buffer).
130/// 2. Encode compute kernels for current layer on the main command buffer.
131/// 3. Insert a `MTLEvent` signal so compute waits for prefetch to complete.
132/// 4. Repeat for each layer — compute and prefetch overlap.
133///
134/// **Requirements:** `metal-rs` crate (macOS only). Not implementable on
135/// non-Apple platforms.
136pub fn metal_double_buffered_prefetch(_weight_data: &[u8], _dst_offset: u64) -> bool {
137    #[cfg(target_os = "macos")]
138    {
139        // TODO: implement via metal-rs MTLBlitCommandEncoder
140        false
141    }
142    #[cfg(not(target_os = "macos"))]
143    {
144        false
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn metal_mega_pass_unavailable_on_non_macos() {
154        if cfg!(not(target_os = "macos")) {
155            assert!(!metal_mega_pass_available());
156            assert!(try_metal_mega_pass(
157                0,
158                0,
159                0,
160                0,
161                0,
162                0,
163                0,
164                0,
165                0,
166                0.0,
167                0.0,
168                0.0,
169                &[],
170                &[],
171                &[],
172                None,
173                None,
174                0,
175                0,
176            )
177            .is_none());
178        }
179    }
180}
qualia_core_db/inference/metal_lane.rs

qualia_core_db/inference/
metal_lane.rs
