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

cortiq_engine/
lib.rs
