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

cortiq_engine/
lib.rs
