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

cortiq_engine/
lib.rs
