spark_runtime/
cublaslt.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Minimal cuBLASLt FFI for the high-efficiency GEMM path (`AVAROK_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 `avarok_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    gemm_act_weight_t_out(act, weight, out, m, n, k, CUDA_R_16BF, stream)
223}
224
225/// [`bf16_gemm_act_weight_t`] with BF16 inputs and an **FP32** output buffer.
226///
227/// Exists for consumers whose downstream kernel reads FP32 and which therefore could not
228/// use the batched BF16-out path at all. GLM's DSA indexer is the motivating case: its
229/// `wq_b` projection writes FP32, so `select_rows_batched` fell back to one M=1 GEMV PER
230/// ROW — 256 full sweeps of the same weight per DSA layer per prefill chunk — under a
231/// comment reading "no FP32-out batchm twin exists". This is that twin.
232///
233/// Identical to the BF16 form except the D layout: cuBLASLt accumulates in FP32 either
234/// way (`CUBLAS_COMPUTE_32F`), so this actually stores MORE of the accumulator than the
235/// BF16 output does, rather than less.
236pub fn bf16_gemm_act_weight_t_f32_out(
237    act: u64,
238    weight: u64,
239    out: u64,
240    m: u32,
241    n: u32,
242    k: u32,
243    stream: u64,
244) -> Result<()> {
245    gemm_act_weight_t_out(act, weight, out, m, n, k, CUDA_R_32F, stream)
246}
247
248/// Shared body of the two wrappers above; `out_dtype` selects the D layout.
249#[allow(clippy::too_many_arguments)]
250fn gemm_act_weight_t_out(
251    act: u64,
252    weight: u64,
253    out: u64,
254    m: u32,
255    n: u32,
256    k: u32,
257    out_dtype: i32,
258    stream: u64,
259) -> Result<()> {
260    let ctx = ctx()?;
261    unsafe {
262        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
263        chk(
264            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
265            "DescCreate",
266        )?;
267        let ta = CUBLAS_OP_T;
268        let tb = CUBLAS_OP_N;
269        chk(
270            cublasLtMatmulDescSetAttribute(
271                desc,
272                DESC_TRANSA,
273                &ta as *const i32 as *const c_void,
274                4,
275            ),
276            "TRANSA",
277        )?;
278        chk(
279            cublasLtMatmulDescSetAttribute(
280                desc,
281                DESC_TRANSB,
282                &tb as *const i32 as *const c_void,
283                4,
284            ),
285            "TRANSB",
286        )?;
287        // A = weight stored row-major [N,K] == col-major [K,N], ld=K, opT → [N,K]
288        // B = act    stored row-major [M,K] == col-major [K,M], ld=K, opN → [K,M]
289        // D = out    row-major [M,N]        == col-major [N,M], ld=N
290        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
291        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
292        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
293        chk(
294            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_16BF, k as u64, n as u64, k as i64),
295            "LayoutA",
296        )?;
297        chk(
298            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_16BF, k as u64, m as u64, k as i64),
299            "LayoutB",
300        )?;
301        chk(
302            cublasLtMatrixLayoutCreate(&mut ld_, out_dtype, n as u64, m as u64, n as i64),
303            "LayoutD",
304        )?;
305        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
306        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
307        let ws_size = ctx.ws_size;
308        chk(
309            cublasLtMatmulPreferenceSetAttribute(
310                pref,
311                PREF_MAX_WORKSPACE_BYTES,
312                &ws_size as *const usize as *const c_void,
313                std::mem::size_of::<usize>(),
314            ),
315            "PrefWorkspace",
316        )?;
317        // cublasLtMatmulHeuristicResult_t = { algo[64B], workspaceSize, state,
318        // wavesCount, reserved[4] } ≈ 96B; algo at offset 0. 128B for margin.
319        let mut result = [0u8; 128];
320        let mut returned: i32 = 0;
321        chk(
322            cublasLtMatmulAlgoGetHeuristic(
323                ctx.handle,
324                desc,
325                la,
326                lb,
327                ld_,
328                ld_,
329                pref,
330                1,
331                result.as_mut_ptr() as *mut c_void,
332                &mut returned,
333            ),
334            "AlgoGetHeuristic",
335        )?;
336        if returned < 1 {
337            bail!("cuBLASLt: no algorithm for {m}x{n}x{k}");
338        }
339        let alpha: f32 = 1.0;
340        let beta: f32 = 0.0;
341        let status = cublasLtMatmul(
342            ctx.handle,
343            desc,
344            &alpha as *const f32 as *const c_void,
345            weight as *const c_void,
346            la,
347            act as *const c_void,
348            lb,
349            &beta as *const f32 as *const c_void,
350            out as *const c_void,
351            ld_,
352            out as *mut c_void,
353            ld_,
354            result.as_ptr() as *const c_void,
355            ctx.workspace as *mut c_void,
356            ctx.ws_size,
357            stream as *mut c_void,
358        );
359        cublasLtMatmulPreferenceDestroy(pref);
360        cublasLtMatrixLayoutDestroy(la);
361        cublasLtMatrixLayoutDestroy(lb);
362        cublasLtMatrixLayoutDestroy(ld_);
363        cublasLtMatmulDescDestroy(desc);
364        chk(status, "Matmul")?;
365    }
366    Ok(())
367}
spark_runtime/cublaslt.rs

spark_runtime/
cublaslt.rs
