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

cortiq_engine/
lib.rs
