spark_runtime/
cublaslt.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Minimal cuBLASLt FFI for the high-efficiency GEMM path (`ATLAS_CUBLAS_GEMM`).
3//!
4//! The hand-written mma.sync projection/MoE GEMMs reach only ~30% of the cuBLAS
5//! ceiling on GB10 (measured: 32 vs 85 TFLOPS bf16, 152 fp8, on the SSM-qkvz
6//! shape 3537×12288×2048). This routes those GEMMs through cuBLASLt instead.
7//! BF16 only for now — correctness-clean (no scale-format issues); native fp8
8//! block-scaled is the follow-up once the end-to-end win is proven.
9
10use anyhow::{Result, bail};
11use std::ffi::c_void;
12use std::sync::OnceLock;
13
14// Native FP8 (E4M3) GEMM paths live in the `fp8` sibling (≤500 LoC split);
15// re-exported so `spark_runtime::cublaslt::fp8_gemm_*` paths are unchanged.
16mod fp8;
17pub use fp8::{fp8_gemm_act_weight_t_blkscaled, fp8_gemm_act_weight_t_rowwise};
18
19#[allow(non_camel_case_types)]
20type cublasLtHandle_t = *mut c_void;
21#[allow(non_camel_case_types)]
22type cublasLtMatmulDesc_t = *mut c_void;
23#[allow(non_camel_case_types)]
24type cublasLtMatrixLayout_t = *mut c_void;
25#[allow(non_camel_case_types)]
26type cublasLtMatmulPreference_t = *mut c_void;
27
28const CUDA_R_16BF: i32 = 14;
29const CUDA_R_32F: i32 = 0;
30const CUDA_R_8F_E4M3: i32 = 28;
31const CUBLAS_COMPUTE_32F: i32 = 68;
32const CUBLAS_OP_N: i32 = 0;
33const CUBLAS_OP_T: i32 = 1;
34const DESC_TRANSA: u32 = 3;
35const DESC_TRANSB: u32 = 4;
36const DESC_A_SCALE_POINTER: u32 = 17;
37const DESC_B_SCALE_POINTER: u32 = 18;
38const DESC_A_SCALE_MODE: u32 = 31;
39const DESC_B_SCALE_MODE: u32 = 32;
40const SCALE_MODE_OUTER_VEC_32F: i32 = 3;
41const SCALE_MODE_VEC128_32F: i32 = 4;
42const SCALE_MODE_BLK128X128_32F: i32 = 5;
43const PREF_MAX_WORKSPACE_BYTES: u32 = 1;
44
45unsafe extern "C" {
46    fn cublasLtCreate(handle: *mut cublasLtHandle_t) -> i32;
47    fn cublasLtMatmulDescCreate(
48        desc: *mut cublasLtMatmulDesc_t,
49        compute_type: i32,
50        scale_type: i32,
51    ) -> i32;
52    fn cublasLtMatmulDescSetAttribute(
53        desc: cublasLtMatmulDesc_t,
54        attr: u32,
55        buf: *const c_void,
56        size: usize,
57    ) -> i32;
58    fn cublasLtMatmulDescDestroy(desc: cublasLtMatmulDesc_t) -> i32;
59    fn cublasLtMatrixLayoutCreate(
60        layout: *mut cublasLtMatrixLayout_t,
61        dtype: i32,
62        rows: u64,
63        cols: u64,
64        ld: i64,
65    ) -> i32;
66    fn cublasLtMatrixLayoutDestroy(layout: cublasLtMatrixLayout_t) -> i32;
67    fn cublasLtMatmulPreferenceCreate(pref: *mut cublasLtMatmulPreference_t) -> i32;
68    fn cublasLtMatmulPreferenceSetAttribute(
69        pref: cublasLtMatmulPreference_t,
70        attr: u32,
71        buf: *const c_void,
72        size: usize,
73    ) -> i32;
74    fn cublasLtMatmulPreferenceDestroy(pref: cublasLtMatmulPreference_t) -> i32;
75    #[allow(clippy::too_many_arguments)]
76    fn cublasLtMatmulAlgoGetHeuristic(
77        handle: cublasLtHandle_t,
78        desc: cublasLtMatmulDesc_t,
79        a: cublasLtMatrixLayout_t,
80        b: cublasLtMatrixLayout_t,
81        c: cublasLtMatrixLayout_t,
82        d: cublasLtMatrixLayout_t,
83        pref: cublasLtMatmulPreference_t,
84        requested: i32,
85        results: *mut c_void,
86        returned: *mut i32,
87    ) -> i32;
88    #[allow(clippy::too_many_arguments)]
89    fn cublasLtMatmul(
90        handle: cublasLtHandle_t,
91        desc: cublasLtMatmulDesc_t,
92        alpha: *const c_void,
93        a: *const c_void,
94        layout_a: cublasLtMatrixLayout_t,
95        b: *const c_void,
96        layout_b: cublasLtMatrixLayout_t,
97        beta: *const c_void,
98        c: *const c_void,
99        layout_c: cublasLtMatrixLayout_t,
100        d: *mut c_void,
101        layout_d: cublasLtMatrixLayout_t,
102        algo: *const c_void,
103        workspace: *mut c_void,
104        workspace_size: usize,
105        stream: *mut c_void,
106    ) -> i32;
107    fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
108    fn cuMemFree_v2(dptr: u64) -> i32;
109    fn cuStreamSynchronize(stream: u64) -> i32;
110}
111
112struct Ctx {
113    handle: cublasLtHandle_t,
114    workspace: u64,
115    ws_size: usize,
116}
117// cuBLASLt handle + device workspace are process-global; matmul is invoked
118// serially from the single-threaded scheduler forward.
119unsafe impl Send for Ctx {}
120unsafe impl Sync for Ctx {}
121
122/// STATIC, DELIBERATELY — CUDA host. This is a workspace allocated in THE
123/// process CUDA context (see `atlas_core::cuda_host`, which establishes one
124/// per process) and sized by a fixed budget, not by any model's shapes: the
125/// bounds below are generous upper limits chosen to fit any realistic serving
126/// configuration, so a swap needs no reallocation and re-allocating per model
127/// would churn hundreds of megabytes for no change in what is mapped.
128///
129/// It survives a model swap for the same reason the context does. Nothing in
130/// it is derived from a model — no token ids, no weight pointers, no shapes —
131/// only scratch the library plans within.
132static CTX: OnceLock<Ctx> = OnceLock::new();
133
134fn ctx() -> Result<&'static Ctx> {
135    if let Some(c) = CTX.get() {
136        return Ok(c);
137    }
138    let mut handle: cublasLtHandle_t = std::ptr::null_mut();
139    let st = unsafe { cublasLtCreate(&mut handle) };
140    if st != 0 {
141        bail!("cublasLtCreate failed: {st}");
142    }
143    let ws_size = 64 * 1024 * 1024;
144    let mut ws: u64 = 0;
145    let st = unsafe { cuMemAlloc_v2(&mut ws, ws_size) };
146    if st != 0 {
147        bail!("cuMemAlloc cuBLASLt workspace failed: {st}");
148    }
149    let _ = CTX.set(Ctx {
150        handle,
151        workspace: ws,
152        ws_size,
153    });
154    Ok(CTX.get().unwrap())
155}
156
157/// Force cuBLASLt's one-time costs at MODEL LOAD instead of on request 1.
158///
159/// The lazy `ctx()` means the first GEMM pays `cublasLtCreate`, the 64 MB
160/// workspace alloc, and — the expensive part — the library's kernel-image
161/// load and heuristic warm-up. Measured on the 35B flagship (2026-08-22,
162/// dgx1): the first in-serve request read ~0.9 s slower than warm requests
163/// once QKVZ routed through cuBLASLt, and cold TTFT is a headline metric.
164/// One 64x64x64 BF16 GEMM here is trivial GPU work and moves that cost to
165/// load time, where it overlaps the operator's mental model of "loading".
166///
167/// Never fails the serve: a pre-warm failure is logged and swallowed — the
168/// lazy path remains and request 1 simply pays the old cost.
169pub fn prewarm(stream: u64) {
170    let r = (|| -> Result<()> {
171        let bytes = 64usize * 64 * 2;
172        let mut a = 0u64;
173        let mut b = 0u64;
174        let mut d = 0u64;
175        unsafe {
176            chk(cuMemAlloc_v2(&mut a, bytes), "prewarm alloc a")?;
177            chk(cuMemAlloc_v2(&mut b, bytes), "prewarm alloc b")?;
178            chk(cuMemAlloc_v2(&mut d, bytes), "prewarm alloc d")?;
179        }
180        let res = bf16_gemm_act_weight_t(a, b, d, 64, 64, 64, stream);
181        unsafe {
182            chk(cuStreamSynchronize(stream), "prewarm sync")?;
183            let _ = cuMemFree_v2(a);
184            let _ = cuMemFree_v2(b);
185            let _ = cuMemFree_v2(d);
186        }
187        res
188    })();
189    match r {
190        Ok(()) => tracing::info!("cuBLASLt pre-warmed (handle + workspace + kernel images)"),
191        Err(e) => tracing::warn!("cuBLASLt pre-warm failed (request 1 pays lazy init): {e}"),
192    }
193}
194
195fn chk(status: i32, what: &str) -> Result<()> {
196    if status != 0 {
197        bail!("cuBLASLt {what} failed: status {status}");
198    }
199    Ok(())
200}
201
202/// Row-major `out[M,N] = act[M,K] @ weight[N,K]ᵀ`, all BF16 — the standard
203/// projection GEMM (activation × transposed weight). Maps to cuBLASLt's
204/// column-major convention as `D[N,M] = opT(weightᶜ[K,N]) · opN(actᶜ[K,M])`.
205pub fn bf16_gemm_act_weight_t(
206    act: u64,
207    weight: u64,
208    out: u64,
209    m: u32,
210    n: u32,
211    k: u32,
212    stream: u64,
213) -> Result<()> {
214    let ctx = ctx()?;
215    unsafe {
216        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
217        chk(
218            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
219            "DescCreate",
220        )?;
221        let ta = CUBLAS_OP_T;
222        let tb = CUBLAS_OP_N;
223        chk(
224            cublasLtMatmulDescSetAttribute(
225                desc,
226                DESC_TRANSA,
227                &ta as *const i32 as *const c_void,
228                4,
229            ),
230            "TRANSA",
231        )?;
232        chk(
233            cublasLtMatmulDescSetAttribute(
234                desc,
235                DESC_TRANSB,
236                &tb as *const i32 as *const c_void,
237                4,
238            ),
239            "TRANSB",
240        )?;
241        // A = weight stored row-major [N,K] == col-major [K,N], ld=K, opT → [N,K]
242        // B = act    stored row-major [M,K] == col-major [K,M], ld=K, opN → [K,M]
243        // D = out    row-major [M,N]        == col-major [N,M], ld=N
244        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
245        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
246        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
247        chk(
248            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_16BF, k as u64, n as u64, k as i64),
249            "LayoutA",
250        )?;
251        chk(
252            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_16BF, k as u64, m as u64, k as i64),
253            "LayoutB",
254        )?;
255        chk(
256            cublasLtMatrixLayoutCreate(&mut ld_, CUDA_R_16BF, n as u64, m as u64, n as i64),
257            "LayoutD",
258        )?;
259        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
260        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
261        let ws_size = ctx.ws_size;
262        chk(
263            cublasLtMatmulPreferenceSetAttribute(
264                pref,
265                PREF_MAX_WORKSPACE_BYTES,
266                &ws_size as *const usize as *const c_void,
267                std::mem::size_of::<usize>(),
268            ),
269            "PrefWorkspace",
270        )?;
271        // cublasLtMatmulHeuristicResult_t = { algo[64B], workspaceSize, state,
272        // wavesCount, reserved[4] } ≈ 96B; algo at offset 0. 128B for margin.
273        let mut result = [0u8; 128];
274        let mut returned: i32 = 0;
275        chk(
276            cublasLtMatmulAlgoGetHeuristic(
277                ctx.handle,
278                desc,
279                la,
280                lb,
281                ld_,
282                ld_,
283                pref,
284                1,
285                result.as_mut_ptr() as *mut c_void,
286                &mut returned,
287            ),
288            "AlgoGetHeuristic",
289        )?;
290        if returned < 1 {
291            bail!("cuBLASLt: no algorithm for {m}x{n}x{k}");
292        }
293        let alpha: f32 = 1.0;
294        let beta: f32 = 0.0;
295        let status = cublasLtMatmul(
296            ctx.handle,
297            desc,
298            &alpha as *const f32 as *const c_void,
299            weight as *const c_void,
300            la,
301            act as *const c_void,
302            lb,
303            &beta as *const f32 as *const c_void,
304            out as *const c_void,
305            ld_,
306            out as *mut c_void,
307            ld_,
308            result.as_ptr() as *const c_void,
309            ctx.workspace as *mut c_void,
310            ctx.ws_size,
311            stream as *mut c_void,
312        );
313        cublasLtMatmulPreferenceDestroy(pref);
314        cublasLtMatrixLayoutDestroy(la);
315        cublasLtMatrixLayoutDestroy(lb);
316        cublasLtMatrixLayoutDestroy(ld_);
317        cublasLtMatmulDescDestroy(desc);
318        chk(status, "Matmul")?;
319    }
320    Ok(())
321}
spark_runtime/cublaslt.rs

spark_runtime/
cublaslt.rs
