cortiq_engine/
lib.rs

1//! Cortiq inference engine — sparse forward pass, attention, tokenization, sampling.
2
3pub mod attention;
4pub mod dit;
5pub mod dsv4;
6pub mod fcd;
7pub mod fcd_ops;
8pub mod g3n;
9pub mod gptq_capture;
10pub mod gpu;
11#[cfg(target_os = "macos")]
12pub mod gpu_metal;
13#[cfg(feature = "gpu")]
14pub mod gpu_wgpu;
15pub mod imagegen;
16pub mod inference;
17pub mod kv_cache;
18pub mod linear_core;
19pub mod loader;
20pub mod nystrom;
21pub mod pin;
22pub mod pipeline;
23pub mod pool;
24pub mod qtensor;
25pub mod router;
26pub mod runtime;
27pub mod sampler;
28pub mod skillbake;
29pub mod swarm;
30pub mod textenc;
31pub mod tokenizer;
32pub mod vae;
33
34pub use nystrom::NystromState;
35pub use pipeline::{GenerateResult, Pipeline, TokenCallback, TokenTrace};
36pub use runtime::CortiqRuntime;
37
38/// Test-only: N empty Metal command-buffer round trips, total seconds.
39#[doc(hidden)]
40#[cfg(target_os = "macos")]
41pub fn gpu_empty_submit_for_test(n: usize) -> f64 {
42    gpu_metal::empty_submit_bench(n)
43}
44
45/// Test-only: N pipelined empty submits, one final wait.
46#[doc(hidden)]
47#[cfg(target_os = "macos")]
48pub fn gpu_pipelined_submit_for_test(n: usize) -> f64 {
49    gpu_metal::pipelined_submit_bench(n)
50}
51
52/// Test-only: build a q1 MoeJob trio (weight 1.0).
53#[doc(hidden)]
54#[cfg(target_os = "macos")]
55pub fn gpu_moe_job_for_test(
56    gi: usize,
57    ui: usize,
58    di: usize,
59    inter: usize,
60    hidden: usize,
61    x: Vec<f32>,
62) -> gpu::MoeJob<'static> {
63    gpu::MoeJob {
64        gate: (gi, inter, hidden, &[]),
65        up: (ui, inter, hidden, &[]),
66        down: (di, hidden, inter, &[]),
67        xs_gate: x.clone(),
68        xs_up: x,
69        down_col: &[],
70        w: 1.0,
71        q1: true,
72        q4t: false,
73        q4tp: false,
74            swiglu_limit: 0.0,
75    }
76}
77
78/// Test-only: run the metal moe_block on one job.
79#[doc(hidden)]
80#[cfg(target_os = "macos")]
81pub fn gpu_moe_block_for_test(
82    model: &std::sync::Arc<cortiq_core::CmfModel>,
83    job: gpu::MoeJob<'_>,
84    out: &mut [f32],
85) -> bool {
86    gpu_metal::moe_block(model, &[job], out)
87}
88
89/// Test-only: q1 matvec_batch — jobs (idx, rows, cols); first two share
90/// x, the third takes xi.
91#[doc(hidden)]
92#[cfg(target_os = "macos")]
93pub fn gpu_batch_q1_for_test(
94    model: &std::sync::Arc<cortiq_core::CmfModel>,
95    shapes: &[(usize, usize, usize)],
96    x: &[f32],
97    xi: &[f32],
98    outs: &mut [&mut [f32]],
99) -> bool {
100    let jobs: Vec<gpu::BatchJob> = shapes
101        .iter()
102        .enumerate()
103        .map(|(k, &(idx, rows, cols))| gpu::BatchJob {
104            idx,
105            rows,
106            cols,
107            row_scale: &[],
108            xs: if k < 2 { x.to_vec() } else { xi.to_vec() },
109            layout: gpu::BatchLayout::Q1,
110        })
111        .collect();
112    gpu_metal::matvec_batch(model, &jobs, outs)
113}
114
115/// Test-only direct handle to the Metal q1 matvec (micro-benchmarks).
116#[doc(hidden)]
117#[cfg(target_os = "macos")]
118pub fn gpu_q1_matvec_for_test(
119    model: &std::sync::Arc<cortiq_core::CmfModel>,
120    idx: usize,
121    xs: &[f32],
122    rows: usize,
123    cols: usize,
124    out: &mut [f32],
125) -> bool {
126    gpu_metal::q1_matvec(model, idx, xs, rows, cols, out)
127}
128pub use sampler::SamplerConfig;
cortiq_engine/lib.rs

cortiq_engine/
lib.rs
