cortiq_engine/
lib.rs

1//! Cortiq inference engine — sparse forward pass, attention, tokenization, sampling.
2
3pub mod attention;
4pub mod chat_template;
5pub mod cpuprof;
6pub mod audiovae;
7pub mod blr;
8pub mod bounded;
9pub mod dit;
10pub mod dsv4;
11pub mod dsv41;
12pub mod dsv41_encoding;
13pub mod dsv41_vision;
14pub mod egemma2;
15pub mod egemma2_audio;
16pub mod egemma2_mm;
17pub mod egemma2_vision;
18pub(crate) mod expert_store;
19pub mod fcd;
20pub mod fcd_ops;
21#[cfg(feature = "gpu")]
22pub mod nvg_train;
23pub mod f32_backend;
24pub mod g3n;
25pub mod g4e;
26pub mod gptq_capture;
27pub mod gpu;
28#[cfg(target_os = "macos")]
29pub mod gpu_metal;
30#[cfg(feature = "gpu")]
31pub mod gpu_wgpu;
32/// The wgpu crate the device types of `gpu_wgpu` come from (offline tools
33/// hold buffers across calls).
34#[cfg(feature = "gpu")]
35pub use wgpu;
36pub mod imagegen;
37pub mod inference;
38pub mod kv_cache;
39pub mod linear_core;
40pub mod loader;
41pub mod lookup;
42pub mod mimo_moe;
43pub(crate) mod minicpm_tail;
44pub mod ltxaudio;
45pub mod ltxdit;
46pub mod ltxdur;
47pub mod ltxenc;
48pub mod ltxlora;
49pub mod ltxpipe;
50pub mod ltxte;
51pub mod ltxups;
52pub mod ltxvae;
53pub mod media;
54pub mod mimo_audio;
55pub mod mimo_mm;
56pub mod mimo_ingress;
57pub mod mimo_vision;
58pub mod mm_ab;
59pub mod mmh3;
60pub mod mmh3ups;
61pub mod music3;
62pub mod moss_image;
63pub mod moss_imagegen;
64pub mod nystrom;
65pub mod pin;
66pub mod pipeline;
67pub mod pool;
68pub mod prism;
69pub mod qtensor;
70pub mod qwen3te;
71pub mod qwen3vis;
72pub mod qwen4_exp;
73pub mod qwen_image;
74pub mod qwen_image21;
75pub mod qwen_image21_vae;
76pub mod qwen_image21gen;
77pub mod qwen_image_encoder;
78pub mod qwen_image_ops;
79pub mod qwen_image_vae;
80pub mod qwen_image_vision;
81pub mod qwen_imagegen;
82pub mod router;
83pub mod runtime;
84pub mod sam3;
85pub mod sam3vid;
86pub(crate) mod sam3vid_net;
87pub mod sampler;
88pub mod skillbake;
89pub mod swarm;
90pub mod textenc;
91pub mod tokenizer;
92pub mod vae;
93pub mod vae3d;
94pub mod videogen;
95/// Native Whisper encoder/decoder inference over CMF checkpoints.
96pub mod whisper;
97pub mod zimage;
98pub mod zimagegen;
99/// The native Vulkan lane — an accelerator behind a capability probe,
100/// present only where Vulkan is.
101#[cfg(all(
102    feature = "gpu",
103    any(target_os = "linux", target_os = "windows", target_os = "android")
104))]
105pub use nystrom::NystromState;
106pub use pipeline::{GenerateResult, Pipeline, TokenCallback, TokenTrace};
107pub use runtime::CortiqRuntime;
108
109/// Test-only: N empty Metal command-buffer round trips, total seconds.
110#[doc(hidden)]
111#[cfg(target_os = "macos")]
112pub fn gpu_empty_submit_for_test(n: usize) -> f64 {
113    gpu_metal::empty_submit_bench(n)
114}
115
116/// Test-only: N pipelined empty submits, one final wait.
117#[doc(hidden)]
118#[cfg(target_os = "macos")]
119pub fn gpu_pipelined_submit_for_test(n: usize) -> f64 {
120    gpu_metal::pipelined_submit_bench(n)
121}
122
123/// Test-only: build a q1 MoeJob trio (weight 1.0).
124#[doc(hidden)]
125#[cfg(target_os = "macos")]
126pub fn gpu_moe_job_for_test(
127    gi: usize,
128    ui: usize,
129    di: usize,
130    inter: usize,
131    hidden: usize,
132    x: Vec<f32>,
133) -> gpu::MoeJob<'static> {
134    gpu::MoeJob {
135        gate: (gi, inter, hidden, &[]),
136        up: (ui, inter, hidden, &[]),
137        down: (di, hidden, inter, &[]),
138        xs_gate: x.clone(),
139        xs_up: x,
140        down_col: &[],
141        w: 1.0,
142        q1: true,
143        q4t: false,
144        q4tp: false,
145        gu_q2: false,
146        swiglu_limit: 0.0,
147        act: 0,
148    }
149}
150
151/// Test-only: run the metal moe_block on one job.
152#[doc(hidden)]
153#[cfg(target_os = "macos")]
154pub fn gpu_moe_block_for_test(
155    model: &std::sync::Arc<cortiq_core::CmfModel>,
156    job: gpu::MoeJob<'_>,
157    out: &mut [f32],
158) -> bool {
159    gpu_metal::moe_block(model, &[job], out)
160}
161
162/// Test-only: q1 matvec_batch — jobs (idx, rows, cols); first two share
163/// x, the third takes xi.
164#[doc(hidden)]
165#[cfg(target_os = "macos")]
166pub fn gpu_batch_q1_for_test(
167    model: &std::sync::Arc<cortiq_core::CmfModel>,
168    shapes: &[(usize, usize, usize)],
169    x: &[f32],
170    xi: &[f32],
171    outs: &mut [&mut [f32]],
172) -> bool {
173    let jobs: Vec<gpu::BatchJob> = shapes
174        .iter()
175        .enumerate()
176        .map(|(k, &(idx, rows, cols))| gpu::BatchJob {
177            idx,
178            rows,
179            cols,
180            row_scale: &[],
181            xs: if k < 2 { x.to_vec() } else { xi.to_vec() },
182            layout: gpu::BatchLayout::Q1,
183        })
184        .collect();
185    gpu_metal::matvec_batch(model, &jobs, outs)
186}
187
188/// Test-only direct handle to the Metal q1 matvec (micro-benchmarks).
189#[doc(hidden)]
190#[cfg(target_os = "macos")]
191pub fn gpu_q1_matvec_for_test(
192    model: &std::sync::Arc<cortiq_core::CmfModel>,
193    idx: usize,
194    xs: &[f32],
195    rows: usize,
196    cols: usize,
197    out: &mut [f32],
198) -> bool {
199    gpu_metal::q1_matvec(model, idx, xs, rows, cols, out)
200}
201pub use sampler::SamplerConfig;
202/// Re-exported for offline tools (the NVG encoder's f64 GEMMs in the CLI).
203pub use matrixmultiply;
cortiq_engine/lib.rs

cortiq_engine/
lib.rs
