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

cortiq_engine/
lib.rs
