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

cortiq_engine/
lib.rs
