Skip to main content

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