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