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

cortiq_engine/
lib.rs
