cortiq_engine/
lib.rs

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

cortiq_engine/
lib.rs
