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

cortiq_engine/
lib.rs
