metrale_gpu_runtime/
cublaslt.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: cuBLASLt FFI for `out[M,N] = act[M,K] @ weight[N,K]ᵀ`: BF16 GEMMs
4//! with BF16 or FP32 output here, FP8 E4M3 GEMMs in `fp8`, and the FP8
5//! scale-tensor layout math in [`scale_layout`].
6//!
7//! Owner: gpu-runtime (cuBLASLt wrapper).
8//! Invariants:
9//! - `CTX` holds a single cuBLASLt handle and 64 MiB workspace for the
10//!   process, allocated with `cuMemAlloc_v2`, outside the backend's
11//!   allocation ledger.
12//! - Every GEMM here passes the whole workspace to `cublasLtMatmul`, so two
13//!   GEMMs must not run concurrently.
14
15use anyhow::{Result, bail};
16use std::ffi::c_void;
17use std::sync::OnceLock;
18
19mod fp8;
20pub use fp8::{
21    fp8_gemm_act_weight_t_blkscaled, fp8_gemm_act_weight_t_blkscaled_ldc,
22    fp8_gemm_act_weight_t_rowwise,
23};
24
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// 2026-09-25: The handle and workspace are process-global and shared by
126// every caller; see the module invariants.
127unsafe impl Send for Ctx {}
128unsafe impl Sync for Ctx {}
129
130/// 2026-09-25: The process-wide cuBLASLt handle and workspace, created on
131/// first use by `ctx`. Its size is fixed, not derived from a model, so it
132/// stays valid across a model swap, like the process CUDA context
133/// (`crate::cuda_host`).
134static CTX: OnceLock<Ctx> = OnceLock::new();
135
136fn ctx() -> Result<&'static Ctx> {
137    if let Some(c) = CTX.get() {
138        return Ok(c);
139    }
140    let mut handle: cublasLtHandle_t = std::ptr::null_mut();
141    let st = unsafe { cublasLtCreate(&mut handle) };
142    if st != 0 {
143        bail!("cublasLtCreate failed: {st}");
144    }
145    let ws_size = 64 * 1024 * 1024;
146    let mut ws: u64 = 0;
147    let st = unsafe { cuMemAlloc_v2(&mut ws, ws_size) };
148    if st != 0 {
149        bail!("cuMemAlloc cuBLASLt workspace failed: {st}");
150    }
151    let _ = CTX.set(Ctx {
152        handle,
153        workspace: ws,
154        ws_size,
155    });
156    Ok(CTX.get().unwrap())
157}
158
159/// 2026-09-25: Run one 64x64x64 BF16 GEMM on `stream` at model load
160/// (`serve_load.rs`), so the first cuBLASLt call, including `ctx`'s
161/// `cublasLtCreate` and workspace allocation, happens before the first
162/// request.
163///
164/// Returns nothing: a failure is logged as a warning, and the first GEMM of
165/// a request then pays the same costs.
166pub fn prewarm(stream: u64) {
167    let r = (|| -> Result<()> {
168        let bytes = 64usize * 64 * 2;
169        let mut a = 0u64;
170        let mut b = 0u64;
171        let mut d = 0u64;
172        unsafe {
173            chk(cuMemAlloc_v2(&mut a, bytes), "prewarm alloc a")?;
174            chk(cuMemAlloc_v2(&mut b, bytes), "prewarm alloc b")?;
175            chk(cuMemAlloc_v2(&mut d, bytes), "prewarm alloc d")?;
176        }
177        let res = bf16_gemm_act_weight_t(a, b, d, 64, 64, 64, stream);
178        unsafe {
179            chk(cuStreamSynchronize(stream), "prewarm sync")?;
180            let _ = cuMemFree_v2(a);
181            let _ = cuMemFree_v2(b);
182            let _ = cuMemFree_v2(d);
183        }
184        res
185    })();
186    match r {
187        Ok(()) => tracing::info!("cuBLASLt pre-warmed (handle + workspace + kernel images)"),
188        Err(e) => tracing::warn!("cuBLASLt pre-warm failed (request 1 pays lazy init): {e}"),
189    }
190}
191
192fn chk(status: i32, what: &str) -> Result<()> {
193    if status != 0 {
194        bail!("cuBLASLt {what} failed: status {status}");
195    }
196    Ok(())
197}
198
199/// 2026-09-25: Row-major `out[M,N] = act[M,K] @ weight[N,K]ᵀ`, all BF16, FP32
200/// accumulation. In cuBLASLt's column-major terms:
201/// `D[N,M] = opT(weightᶜ[K,N]) · opN(actᶜ[K,M])`.
202pub fn bf16_gemm_act_weight_t(
203    act: u64,
204    weight: u64,
205    out: u64,
206    m: u32,
207    n: u32,
208    k: u32,
209    stream: u64,
210) -> Result<()> {
211    gemm_act_weight_t_out(act, weight, out, m, n, k, CUDA_R_16BF, stream)
212}
213
214/// 2026-09-25: [`bf16_gemm_act_weight_t`] with an FP32 `out`. Only the D
215/// layout's type differs; both use `CUBLAS_COMPUTE_32F`.
216pub fn bf16_gemm_act_weight_t_f32_out(
217    act: u64,
218    weight: u64,
219    out: u64,
220    m: u32,
221    n: u32,
222    k: u32,
223    stream: u64,
224) -> Result<()> {
225    gemm_act_weight_t_out(act, weight, out, m, n, k, CUDA_R_32F, stream)
226}
227
228/// 2026-09-25: Shared body of the two wrappers above; `out_dtype` is the D
229/// layout's type.
230#[allow(clippy::too_many_arguments)]
231fn gemm_act_weight_t_out(
232    act: u64,
233    weight: u64,
234    out: u64,
235    m: u32,
236    n: u32,
237    k: u32,
238    out_dtype: i32,
239    stream: u64,
240) -> Result<()> {
241    let ctx = ctx()?;
242    unsafe {
243        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
244        chk(
245            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
246            "DescCreate",
247        )?;
248        let ta = CUBLAS_OP_T;
249        let tb = CUBLAS_OP_N;
250        chk(
251            cublasLtMatmulDescSetAttribute(
252                desc,
253                DESC_TRANSA,
254                &ta as *const i32 as *const c_void,
255                4,
256            ),
257            "TRANSA",
258        )?;
259        chk(
260            cublasLtMatmulDescSetAttribute(
261                desc,
262                DESC_TRANSB,
263                &tb as *const i32 as *const c_void,
264                4,
265            ),
266            "TRANSB",
267        )?;
268        // 2026-09-25: A = weight, row-major [N,K] = col-major [K,N], ld K, opT.
269        // B = act, row-major [M,K] = col-major [K,M], ld K, opN.
270        // D = out, row-major [M,N] = col-major [N,M], ld N.
271        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
272        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
273        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
274        chk(
275            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_16BF, k as u64, n as u64, k as i64),
276            "LayoutA",
277        )?;
278        chk(
279            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_16BF, k as u64, m as u64, k as i64),
280            "LayoutB",
281        )?;
282        chk(
283            cublasLtMatrixLayoutCreate(&mut ld_, out_dtype, n as u64, m as u64, n as i64),
284            "LayoutD",
285        )?;
286        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
287        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
288        let ws_size = ctx.ws_size;
289        chk(
290            cublasLtMatmulPreferenceSetAttribute(
291                pref,
292                PREF_MAX_WORKSPACE_BYTES,
293                &ws_size as *const usize as *const c_void,
294                std::mem::size_of::<usize>(),
295            ),
296            "PrefWorkspace",
297        )?;
298        // 2026-09-25: Holds one cublasLtMatmulHeuristicResult_t (96 bytes in
299        // the CUDA 13 cublasLt.h: 64-byte algo first, then workspaceSize,
300        // state, wavesCount, reserved[4]); `algo` is read from offset 0.
301        let mut result = [0u8; 128];
302        let mut returned: i32 = 0;
303        chk(
304            cublasLtMatmulAlgoGetHeuristic(
305                ctx.handle,
306                desc,
307                la,
308                lb,
309                ld_,
310                ld_,
311                pref,
312                1,
313                result.as_mut_ptr() as *mut c_void,
314                &mut returned,
315            ),
316            "AlgoGetHeuristic",
317        )?;
318        if returned < 1 {
319            bail!("cuBLASLt: no algorithm for {m}x{n}x{k}");
320        }
321        let alpha: f32 = 1.0;
322        let beta: f32 = 0.0;
323        let status = cublasLtMatmul(
324            ctx.handle,
325            desc,
326            &alpha as *const f32 as *const c_void,
327            weight as *const c_void,
328            la,
329            act as *const c_void,
330            lb,
331            &beta as *const f32 as *const c_void,
332            out as *const c_void,
333            ld_,
334            out as *mut c_void,
335            ld_,
336            result.as_ptr() as *const c_void,
337            ctx.workspace as *mut c_void,
338            ctx.ws_size,
339            stream as *mut c_void,
340        );
341        cublasLtMatmulPreferenceDestroy(pref);
342        cublasLtMatrixLayoutDestroy(la);
343        cublasLtMatrixLayoutDestroy(lb);
344        cublasLtMatrixLayoutDestroy(ld_);
345        cublasLtMatmulDescDestroy(desc);
346        chk(status, "Matmul")?;
347    }
348    Ok(())
349}
metrale_gpu_runtime/cublaslt.rs

metrale_gpu_runtime/
cublaslt.rs
