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

cortiq_engine/
lib.rs
