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

cortiq_engine/
lib.rs
