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

cortiq_engine/
lib.rs
