Skip to main content

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