spark_runtime/
cutlass.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Optional CUTLASS host-wrapper FFI for de-risking GB10 GEMM replacements.
3//!
4//! ★ REFERENCE IMPLEMENTATION — A BENCHMARK TARGET, NOT A DEPENDENCY.
5//!
6//! Atlas ships its OWN kernels. CUTLASS is wrapped here for exactly one
7//! purpose: to be the opponent we measure against and beat. **Nothing in a
8//! default build or a default serve calls a single line of it.**
9//!
10//! Two independent gates keep that true, and BOTH must survive any edit:
11//!
12//!   1. COMPILE TIME — every item in this module is `#[cfg(atlas_cutlass)]`, and
13//!      `build.rs` sets that cfg only when `CUTLASS_HOME` is exported. A build
14//!      without it links no CUTLASS object at all.
15//!   2. RUNTIME — the dispatch arms are opt-in behind `ATLAS_CUTLASS_GEMM=1`.
16//!      The default is OFF.
17//!
18//! So the honest reading of an Atlas performance number is that Atlas kernels
19//! produced it, because a default binary cannot reach this code. Export the
20//! env var and you are measuring CUTLASS — label the number that way.
21//!
22//! ★ WHY KEEP IT COMPILED-BUT-DARK. Agentic benchmarking. An optimisation
23//! claim needs a credible opponent: "faster than our own previous commit" is
24//! a far weaker statement than "faster than CUTLASS on this shape". Keeping the
25//! wrapper one env var away lets any agent A/B a shape against the industry
26//! reference on the same box, same checkpoint, same stream — which is the
27//! only comparison worth quoting.
28//!
29//! Do NOT promote any of this to a default path. If a CUTLASS shape beats ours,
30//! the correct response is to make OUR kernel faster and re-measure.
31//!
32//! Split for the ≤500 LoC cap: this root holds the shared FFI `extern` block,
33//! the workspace `Ctx`, and module wiring; the public wrappers live in the
34//! `gemm` (dense BF16 + NVFP4), `grouped` (per-expert MoE), and `pack`
35//! (weight pack / SFB swizzle / transpose) siblings. The public API
36//! (`spark_runtime::cutlass::<fn>`) is preserved via the re-exports below.
37
38#[cfg(atlas_cutlass)]
39use anyhow::{Result, bail};
40
41#[cfg(atlas_cutlass)]
42use std::ffi::c_void;
43#[cfg(atlas_cutlass)]
44use std::sync::OnceLock;
45
46mod gemm;
47mod grouped;
48mod pack;
49
50pub use gemm::{bf16_gemm_act_weight_t, nvfp4_gemm_bf16_act_weight_t};
51pub use grouped::{nvfp4_grouped_down, nvfp4_grouped_gate_up, nvfp4_grouped_gate_up_fused};
52pub use pack::{pack_bf16_weight_to_nvfp4_t, pack_weight_sfb, transpose_nvfp4_packed_kton};
53
54#[cfg(all(test, atlas_cutlass))]
55mod tests;
56
57#[cfg(atlas_cutlass)]
58unsafe extern "C" {
59    pub(crate) fn atlas_cutlass_bf16_gemm_act_weight_t(
60        act: *const c_void,
61        weight: *const c_void,
62        out: *mut c_void,
63        m: i32,
64        n: i32,
65        k: i32,
66        workspace: *mut c_void,
67        workspace_size: usize,
68        stream: *mut c_void,
69    ) -> i32;
70    pub(crate) fn atlas_cutlass_nvfp4_gemm_bf16_act_weight_t(
71        act: *const c_void,
72        weight_packed_t: *const c_void,
73        weight_scale_t: *const c_void,
74        weight_scale_2: f32,
75        out: *mut c_void,
76        m: i32,
77        n: i32,
78        k: i32,
79        workspace: *mut c_void,
80        workspace_size: usize,
81        stream: *mut c_void,
82    ) -> i32;
83    pub(crate) fn atlas_cutlass_pack_bf16_weight_to_nvfp4_t(
84        weight_bf16: *const c_void,
85        packed_t: *mut c_void,
86        scale_t: *mut c_void,
87        n: i32,
88        k: i32,
89        stream: *mut c_void,
90    ) -> i32;
91    pub(crate) fn atlas_cutlass_nvfp4_grouped_gate_up(
92        a_bf16: *const c_void,
93        gate_packed_ptrs: *const u64,
94        gate_scale_ptrs: *const u64,
95        gate_scale2_vals: *const f32,
96        up_packed_ptrs: *const u64,
97        up_scale_ptrs: *const u64,
98        up_scale2_vals: *const f32,
99        c_gate_bf16: *mut c_void,
100        c_up_bf16: *mut c_void,
101        expert_offsets_host: *const i32,
102        num_experts: i32,
103        n: i32,
104        k: i32,
105        workspace: *mut c_void,
106        workspace_size: usize,
107        stream: *mut c_void,
108    ) -> i32;
109    pub(crate) fn atlas_cutlass_nvfp4_grouped_gate_up_fused(
110        a_bf16: *const c_void,
111        sorted_token_ids: *const i32,
112        gate_packed_ptrs: *const u64,
113        gate_sfb_ptrs: *const u64,
114        gate_scale2_vals: *const f32,
115        up_packed_ptrs: *const u64,
116        up_sfb_ptrs: *const u64,
117        up_scale2_vals: *const f32,
118        c_gate_bf16: *mut c_void,
119        c_up_bf16: *mut c_void,
120        expert_offsets_host: *const i32,
121        num_experts: i32,
122        n: i32,
123        k: i32,
124        workspace: *mut c_void,
125        workspace_size: usize,
126        stream: *mut c_void,
127    ) -> i32;
128    pub(crate) fn atlas_cutlass_nvfp4_grouped_down(
129        a_bf16: *const c_void,
130        packed_ptrs: *const u64,
131        sfb_ptrs: *const u64,
132        scale2_vals: *const f32,
133        c_bf16: *mut c_void,
134        expert_offsets_host: *const i32,
135        num_experts: i32,
136        n: i32,
137        k: i32,
138        workspace: *mut c_void,
139        workspace_size: usize,
140        stream: *mut c_void,
141    ) -> i32;
142    pub(crate) fn atlas_cutlass_pack_weight_sfb(
143        scale_in: *const c_void,
144        scale_out: *mut c_void,
145        n: i32,
146        k: i32,
147        src_n_major: i32,
148        stream: *mut c_void,
149    ) -> i32;
150    pub(crate) fn atlas_cutlass_transpose_nvfp4_packed_kton(
151        src_packed_t: *const c_void,
152        dst_packed: *mut c_void,
153        n: i32,
154        k: i32,
155        stream: *mut c_void,
156    ) -> i32;
157    #[cfg(test)]
158    pub(crate) fn atlas_cutlass_bf16_gemm_act_weight_t_128x256(
159        act: *const c_void,
160        weight: *const c_void,
161        out: *mut c_void,
162        m: i32,
163        n: i32,
164        k: i32,
165        workspace: *mut c_void,
166        workspace_size: usize,
167        stream: *mut c_void,
168    ) -> i32;
169    #[cfg(test)]
170    pub(crate) fn atlas_cutlass_bf16_gemm_act_weight_t_256x128(
171        act: *const c_void,
172        weight: *const c_void,
173        out: *mut c_void,
174        m: i32,
175        n: i32,
176        k: i32,
177        workspace: *mut c_void,
178        workspace_size: usize,
179        stream: *mut c_void,
180    ) -> i32;
181    #[cfg(test)]
182    pub(crate) fn atlas_cutlass_bf16_gemm_act_weight_t_64x128(
183        act: *const c_void,
184        weight: *const c_void,
185        out: *mut c_void,
186        m: i32,
187        n: i32,
188        k: i32,
189        workspace: *mut c_void,
190        workspace_size: usize,
191        stream: *mut c_void,
192    ) -> i32;
193    #[cfg(test)]
194    pub(crate) fn atlas_cutlass_bf16_gemm_act_weight_t_128x64(
195        act: *const c_void,
196        weight: *const c_void,
197        out: *mut c_void,
198        m: i32,
199        n: i32,
200        k: i32,
201        workspace: *mut c_void,
202        workspace_size: usize,
203        stream: *mut c_void,
204    ) -> i32;
205    #[cfg(test)]
206    pub(crate) fn atlas_cutlass_bf16_gemm_act_weight_t_64x64(
207        act: *const c_void,
208        weight: *const c_void,
209        out: *mut c_void,
210        m: i32,
211        n: i32,
212        k: i32,
213        workspace: *mut c_void,
214        workspace_size: usize,
215        stream: *mut c_void,
216    ) -> i32;
217    #[cfg(test)]
218    pub(crate) fn atlas_cublaslt_bf16_gemm_act_weight_t_algo(
219        act: *const c_void,
220        weight: *const c_void,
221        out: *mut c_void,
222        m: i32,
223        n: i32,
224        k: i32,
225        workspace: *mut c_void,
226        workspace_size: usize,
227        stream: *mut c_void,
228        algo_index: i32,
229        returned_count: *mut i32,
230    ) -> i32;
231    pub(crate) fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
232}
233
234#[cfg(atlas_cutlass)]
235pub(crate) struct Ctx {
236    pub(crate) workspace: u64,
237    pub(crate) ws_size: usize,
238}
239
240#[cfg(atlas_cutlass)]
241unsafe impl Send for Ctx {}
242#[cfg(atlas_cutlass)]
243unsafe impl Sync for Ctx {}
244
245#[cfg(atlas_cutlass)]
246/// STATIC, DELIBERATELY — CUDA host. This is a workspace allocated in THE
247/// process CUDA context (see `atlas_core::cuda_host`, which establishes one
248/// per process) and sized by a fixed budget, not by any model's shapes: the
249/// bounds below are generous upper limits chosen to fit any realistic serving
250/// configuration, so a swap needs no reallocation and re-allocating per model
251/// would churn hundreds of megabytes for no change in what is mapped.
252///
253/// It survives a model swap for the same reason the context does. Nothing in
254/// it is derived from a model — no token ids, no weight pointers, no shapes —
255/// only scratch the library plans within.
256static CTX: OnceLock<Ctx> = OnceLock::new();
257
258#[cfg(atlas_cutlass)]
259pub(crate) fn ctx() -> Result<&'static Ctx> {
260    if let Some(c) = CTX.get() {
261        return Ok(c);
262    }
263    // Shared scratch for all CUTLASS host wrappers. The grouped NVFP4 MoE path
264    // (single-launch kGrouped over up to 256 experts) stages packed-A + SFA +
265    // per-group arrays + the gemm workspace here; at large prefill M the 256-group
266    // gemm workspace alone exceeds the old 64 MB (-> status -2 + an OOB context
267    // corruption). 512 MB by default; override with ATLAS_CUTLASS_WORKSPACE_MB.
268    let ws_size = std::env::var("ATLAS_CUTLASS_WORKSPACE_MB")
269        .ok()
270        .and_then(|v| v.parse::<usize>().ok())
271        .unwrap_or(512)
272        * 1024
273        * 1024;
274    let mut workspace = 0u64;
275    let status = unsafe { cuMemAlloc_v2(&mut workspace, ws_size) };
276    if status != 0 {
277        bail!("cuMemAlloc CUTLASS workspace failed: {status}");
278    }
279    let _ = CTX.set(Ctx { workspace, ws_size });
280    Ok(CTX.get().unwrap())
281}
spark_runtime/cutlass.rs

spark_runtime/
cutlass.rs
