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

cortiq_engine/
lib.rs
