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

cortiq_engine/
lib.rs
