cortiq_engine/
gpu.rs

1//! Facade for GPU backends: a single call entry point for qtensor/pipeline/
2//! linear_core. Job types and the threshold are canonical HERE; behind the
3//! facade dispatch goes to a platform backend:
4//!   - `gpu_metal` (Apple Silicon, unified memory + no-copy buffers);
5//!   - `gpu_wgpu` (C1: Vulkan/DX12/Metal — NVIDIA/Radeon/Intel/Apple,
6//!     weights resident in VRAM), available under `--features gpu`.
7//!
8//! Runtime selection via `CMF_GPU`: `1` — native Metal (macOS) or wgpu
9//! (other OSes); `wgpu` — force wgpu (including for the local
10//! Metal-via-wgpu parity test). Any backend refusal — `false` and the honest
11//! CPU path, no partial results.
12
13use cortiq_core::CmfModel;
14use std::cell::Cell;
15use std::sync::{Arc, OnceLock};
16
17thread_local! {
18    /// Index of the current forward layer (−1 = outside a numbered layer:
19    /// lm_head/embed — always allowed). The pipeline sets it before
20    /// each layer so that the GPU/CPU layer-split works.
21    static CUR_LAYER: Cell<i64> = const { Cell::new(-1) };
22}
23
24/// Pipeline: mark the current layer (or −1 outside layers) for layer-split.
25pub fn set_layer(l: i64) {
26    CUR_LAYER.with(|c| c.set(l));
27}
28
29/// Parse `CMF_GPU_LAYERS` («0-19», «0,2,4», «0-9,30-39») once.
30/// None = no restriction (all layers on GPU). Garbage → also no restriction.
31fn layer_ranges() -> &'static Option<Vec<(i64, i64)>> {
32    static R: OnceLock<Option<Vec<(i64, i64)>>> = OnceLock::new();
33    R.get_or_init(|| {
34        let s = std::env::var("CMF_GPU_LAYERS").ok()?;
35        let mut v = Vec::new();
36        for part in s.split(',') {
37            let part = part.trim();
38            match part.split_once('-') {
39                Some((a, b)) => v.push((a.trim().parse().ok()?, b.trim().parse().ok()?)),
40                None => {
41                    let x: i64 = part.parse().ok()?;
42                    v.push((x, x));
43                }
44            }
45        }
46        Some(v)
47    })
48}
49
50fn layer_allowed() -> bool {
51    match layer_ranges() {
52        None => true,
53        Some(ranges) => {
54            let cur = CUR_LAYER.with(|c| c.get());
55            cur < 0 || ranges.iter().any(|(a, b)| cur >= *a && cur <= *b)
56        }
57    }
58}
59
60/// GPU allowed FOR THE CURRENT LAYER: backend is initialized AND the layer
61/// falls within `CMF_GPU_LAYERS` (GPU/CPU layer-split). Op gates call this.
62pub fn enabled_here() -> bool {
63    enabled() && layer_allowed()
64}
65
66/// Default row threshold: the GPU takes only larger matrices (lm_head
67/// class). Below it, the dispatch/readback cost does not pay off on unified memory.
68pub const GPU_MIN_ROWS: usize = 65_536;
69
70/// Effective threshold: `CMF_GPU_MIN_ROWS` overrides the default — on a
71/// discrete card it is worth lowering it (VRAM bandwidth pays off even for FFN),
72/// on unified memory — raising it. A «squeeze out the maximum» tuning on the server.
73pub fn min_rows() -> usize {
74    std::env::var("CMF_GPU_MIN_ROWS")
75        .ok()
76        .and_then(|v| v.parse().ok())
77        .unwrap_or(GPU_MIN_ROWS)
78}
79
80/// A single MoE-FFN job (an expert with its own weight), executed in one
81/// submission: (rows, cols, idx, row_scale) for gate/up/down + prescaled
82/// inputs + the down θ-field + the blending weight.
83pub struct MoeJob<'a> {
84    pub gate: (usize, usize, usize, &'a [f32]),
85    pub up: (usize, usize, usize, &'a [f32]),
86    pub down: (usize, usize, usize, &'a [f32]),
87    pub xs_gate: Vec<f32>,
88    pub xs_up: Vec<f32>,
89    pub down_col: &'a [f32],
90    pub w: f32,
91}
92
93/// A single independent batch matvec (GDN projections of one input).
94pub struct BatchJob<'a> {
95    pub idx: usize,
96    pub rows: usize,
97    pub cols: usize,
98    pub row_scale: &'a [f32],
99    pub xs: Vec<f32>,
100}
101
102#[derive(Clone, Copy, PartialEq, Eq)]
103enum Backend {
104    None,
105    #[cfg(target_os = "macos")]
106    Metal,
107    #[cfg(feature = "gpu")]
108    Wgpu,
109}
110
111fn backend() -> Backend {
112    #[cfg(feature = "gpu")]
113    if crate::gpu_wgpu::selected() {
114        return if crate::gpu_wgpu::enabled() { Backend::Wgpu } else { Backend::None };
115    }
116    #[cfg(target_os = "macos")]
117    if crate::gpu_metal::enabled() {
118        return Backend::Metal;
119    }
120    Backend::None
121}
122
123/// GPU enabled and initialized on the selected backend?
124pub fn enabled() -> bool {
125    backend() != Backend::None
126}
127
128/// q8_row/q8_2f matvec, rows [row0, row0+rows). `xs` — prescaled by the θ-field.
129#[allow(clippy::too_many_arguments, unused_variables)]
130pub fn q8_matvec_range(
131    model: &Arc<CmfModel>,
132    idx: usize,
133    row0: usize,
134    row_scale: &[f32],
135    xs: &[f32],
136    rows: usize,
137    cols: usize,
138    out: &mut [f32],
139) -> bool {
140    match backend() {
141        #[cfg(target_os = "macos")]
142        Backend::Metal => {
143            crate::gpu_metal::q8_matvec_range(model, idx, row0, row_scale, xs, rows, cols, out)
144        }
145        #[cfg(feature = "gpu")]
146        Backend::Wgpu => {
147            crate::gpu_wgpu::q8_matvec_range(model, idx, row0, row_scale, xs, rows, cols, out)
148        }
149        Backend::None => false,
150    }
151}
152
153/// GEMM of a prefill batch: `pre` — prescaled inputs row-major [b, cols],
154/// out — row-major [b, rows].
155#[allow(clippy::too_many_arguments, unused_variables)]
156pub fn q8_matmat(
157    model: &Arc<CmfModel>,
158    idx: usize,
159    row_scale: &[f32],
160    pre: &[f32],
161    b: usize,
162    rows: usize,
163    cols: usize,
164    out: &mut [f32],
165) -> bool {
166    match backend() {
167        #[cfg(target_os = "macos")]
168        Backend::Metal => {
169            crate::gpu_metal::q8_matmat(model, idx, row_scale, pre, b, rows, cols, out)
170        }
171        #[cfg(feature = "gpu")]
172        Backend::Wgpu => {
173            crate::gpu_wgpu::q8_matmat(model, idx, row_scale, pre, b, rows, cols, out)
174        }
175        Backend::None => false,
176    }
177}
178
179/// A layer's MoE-FFN in one submission (amortizing the dispatch cost).
180#[allow(unused_variables)]
181pub fn moe_block(model: &Arc<CmfModel>, jobs: &[MoeJob], out: &mut [f32]) -> bool {
182    match backend() {
183        #[cfg(target_os = "macos")]
184        Backend::Metal => crate::gpu_metal::moe_block(model, jobs, out),
185        #[cfg(feature = "gpu")]
186        Backend::Wgpu => crate::gpu_wgpu::moe_block(model, jobs, out),
187        Backend::None => false,
188    }
189}
190
191/// Independent matvecs of one input in a single submission (GDN projections).
192#[allow(unused_variables)]
193pub fn matvec_batch(model: &Arc<CmfModel>, jobs: &[BatchJob], out: &mut [&mut [f32]]) -> bool {
194    match backend() {
195        #[cfg(target_os = "macos")]
196        Backend::Metal => crate::gpu_metal::matvec_batch(model, jobs, out),
197        #[cfg(feature = "gpu")]
198        Backend::Wgpu => crate::gpu_wgpu::matvec_batch(model, jobs, out),
199        Backend::None => false,
200    }
201}
cortiq_engine/gpu.rs

cortiq_engine/
gpu.rs
