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

cortiq_engine/
lib.rs
