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

cortiq_engine/
lib.rs
