cortiq_engine/
lib.rs

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

cortiq_engine/
lib.rs
