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

cortiq_engine/
lib.rs
