cortiq_engine/
lib.rs

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

cortiq_engine/
lib.rs
