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

spark_runtime/
cublaslt.rs
