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

cortiq_engine/
lib.rs
