metrale_gpu_runtime/cublaslt/
fp8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: FP8 E4M3 cuBLASLt GEMMs with BF16 output: row-wise scaled, and
4//! 128-block scaled with a contiguous or pitched output.
5//!
6//! Owner: gpu-runtime (cuBLASLt wrapper).
7//! Invariants: none beyond the types.
8
9use anyhow::{Result, bail};
10use std::ffi::c_void;
11
12use super::*;
13
14/// 2026-09-25: FP8 E4M3 `out[M,N] = act[M,K] @ weight[N,K]ᵀ`, BF16 output,
15/// with FP32 `OUTER_VEC` scales on both operands: `weight_scale[N]` per
16/// weight row (cuBLASLt's A) and `act_scale[M]` per token (B).
17#[allow(clippy::too_many_arguments)]
18pub fn fp8_gemm_act_weight_t_rowwise(
19    act_fp8: u64,
20    act_scale: u64,
21    weight_fp8: u64,
22    weight_scale: u64,
23    out: u64,
24    m: u32,
25    n: u32,
26    k: u32,
27    stream: u64,
28) -> Result<()> {
29    let ctx = ctx()?;
30    unsafe {
31        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
32        chk(
33            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
34            "DescCreate",
35        )?;
36        let ta = CUBLAS_OP_T;
37        let tb = CUBLAS_OP_N;
38        let set = |attr: u32, val: *const c_void, sz: usize, what: &str| -> Result<()> {
39            chk(cublasLtMatmulDescSetAttribute(desc, attr, val, sz), what)
40        };
41        set(DESC_TRANSA, &ta as *const i32 as *const c_void, 4, "TRANSA")?;
42        set(DESC_TRANSB, &tb as *const i32 as *const c_void, 4, "TRANSB")?;
43        let mode = SCALE_MODE_OUTER_VEC_32F;
44        set(
45            DESC_A_SCALE_MODE,
46            &mode as *const i32 as *const c_void,
47            4,
48            "A_SCALE_MODE",
49        )?;
50        set(
51            DESC_B_SCALE_MODE,
52            &mode as *const i32 as *const c_void,
53            4,
54            "B_SCALE_MODE",
55        )?;
56        set(
57            DESC_A_SCALE_POINTER,
58            &weight_scale as *const u64 as *const c_void,
59            8,
60            "A_SCALE_POINTER",
61        )?;
62        set(
63            DESC_B_SCALE_POINTER,
64            &act_scale as *const u64 as *const c_void,
65            8,
66            "B_SCALE_POINTER",
67        )?;
68
69        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
70        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
71        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
72        chk(
73            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_8F_E4M3, k as u64, n as u64, k as i64),
74            "LayoutA",
75        )?;
76        chk(
77            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_8F_E4M3, k as u64, m as u64, k as i64),
78            "LayoutB",
79        )?;
80        chk(
81            cublasLtMatrixLayoutCreate(&mut ld_, CUDA_R_16BF, n as u64, m as u64, n as i64),
82            "LayoutD",
83        )?;
84        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
85        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
86        let ws_size = ctx.ws_size;
87        chk(
88            cublasLtMatmulPreferenceSetAttribute(
89                pref,
90                PREF_MAX_WORKSPACE_BYTES,
91                &ws_size as *const usize as *const c_void,
92                std::mem::size_of::<usize>(),
93            ),
94            "PrefWorkspace",
95        )?;
96        let mut result = [0u8; 128];
97        let mut returned: i32 = 0;
98        chk(
99            cublasLtMatmulAlgoGetHeuristic(
100                ctx.handle,
101                desc,
102                la,
103                lb,
104                ld_,
105                ld_,
106                pref,
107                1,
108                result.as_mut_ptr() as *mut c_void,
109                &mut returned,
110            ),
111            "AlgoGetHeuristic",
112        )?;
113        if returned < 1 {
114            bail!("cuBLASLt fp8 rowwise: no algorithm for {m}x{n}x{k}");
115        }
116        let alpha: f32 = 1.0;
117        let beta: f32 = 0.0;
118        let status = cublasLtMatmul(
119            ctx.handle,
120            desc,
121            &alpha as *const f32 as *const c_void,
122            weight_fp8 as *const c_void,
123            la,
124            act_fp8 as *const c_void,
125            lb,
126            &beta as *const f32 as *const c_void,
127            out as *const c_void,
128            ld_,
129            out as *mut c_void,
130            ld_,
131            result.as_ptr() as *const c_void,
132            ctx.workspace as *mut c_void,
133            ctx.ws_size,
134            stream as *mut c_void,
135        );
136        cublasLtMatmulPreferenceDestroy(pref);
137        cublasLtMatrixLayoutDestroy(la);
138        cublasLtMatrixLayoutDestroy(lb);
139        cublasLtMatrixLayoutDestroy(ld_);
140        cublasLtMatmulDescDestroy(desc);
141        chk(status, "Matmul")?;
142    }
143    Ok(())
144}
145
146/// 2026-09-25: FP8 E4M3 `out[M,N] = act[M,K] @ weight[N,K]ᵀ`, BF16 output,
147/// with FP32 block scales: `BLK128x128_32F` on the weight (cuBLASLt's A) and
148/// `VEC128_32F` on the activation (B). The scale layouts are in
149/// [`super::scale_layout`]:
150///
151/// * `weight_block_scale` is the checkpoint's row-major `[N/128, K/128]`
152///   grid, passed unchanged, which requires `blk128x128_stride_ok(k)`; this
153///   function does not check it.
154/// * `act_scale` is expected as `[K/128, m]`, token index contiguous
155///   (`vec128_b_index`): the transpose of the `[M, K/128]` that
156///   `per_token_group_quant_fp8` writes. The `fp8_act_scale_to_kmajor`
157///   kernel converts one to the other.
158///
159/// `m` goes to cuBLASLt unchanged; the callers pass a padded count
160/// (`m_pad`), and `act_fp8` and `act_scale` must cover it. The output is
161/// contiguous `[m, N]`;
162/// [`fp8_gemm_act_weight_t_blkscaled_ldc`] takes an output row pitch.
163#[allow(clippy::too_many_arguments)]
164pub fn fp8_gemm_act_weight_t_blkscaled(
165    act_fp8: u64,
166    act_scale: u64,
167    weight_fp8: u64,
168    weight_block_scale: u64,
169    out: u64,
170    m: u32,
171    n: u32,
172    k: u32,
173    stream: u64,
174) -> Result<()> {
175    fp8_gemm_act_weight_t_blkscaled_ldc(
176        act_fp8,
177        act_scale,
178        weight_fp8,
179        weight_block_scale,
180        out,
181        m,
182        n,
183        k,
184        n,
185        stream,
186    )
187}
188
189/// 2026-09-25: [`fp8_gemm_act_weight_t_blkscaled`] with output rows `ldc`
190/// BF16 elements apart: D is column-major `[N, m]` with leading dimension
191/// `ldc`, i.e. row-major `[m, N]` with row pitch `ldc`. The multi-seq decode
192/// Q/K/V projections use it to write each row into its sequence's slot of
193/// the QKV buffer (`w8a8_decode.rs` `qkv_decode_w8a8_plans`, `ldc` =
194/// `per_seq_qkv` in BF16 elements).
195///
196/// All `m` rows are written, padding rows included, so the output extent is
197/// `(m - 1) * ldc + n` elements; callers bound it against their buffer with
198/// `metrale_model_layers::layers::ops::strided_out_extent_elems`.
199///
200/// Fails before calling cuBLASLt when `ldc < n`.
201#[allow(clippy::too_many_arguments)]
202pub fn fp8_gemm_act_weight_t_blkscaled_ldc(
203    act_fp8: u64,
204    act_scale: u64,
205    weight_fp8: u64,
206    weight_block_scale: u64,
207    out: u64,
208    m: u32,
209    n: u32,
210    k: u32,
211    ldc: u32,
212    stream: u64,
213) -> Result<()> {
214    if ldc < n {
215        bail!("cuBLASLt fp8: output row pitch ldc={ldc} is shorter than N={n}");
216    }
217    let ctx = ctx()?;
218    unsafe {
219        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
220        chk(
221            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
222            "DescCreate",
223        )?;
224        let ta = CUBLAS_OP_T;
225        let tb = CUBLAS_OP_N;
226        let set = |attr: u32, val: *const c_void, sz: usize, what: &str| -> Result<()> {
227            chk(cublasLtMatmulDescSetAttribute(desc, attr, val, sz), what)
228        };
229        set(DESC_TRANSA, &ta as *const i32 as *const c_void, 4, "TRANSA")?;
230        set(DESC_TRANSB, &tb as *const i32 as *const c_void, 4, "TRANSB")?;
231        // 2026-09-25: Weight: one scale per 128x128 block. Activation: one
232        // scale per token per 128 elements of K.
233        let a_mode = SCALE_MODE_BLK128X128_32F;
234        let b_mode = SCALE_MODE_VEC128_32F;
235        set(
236            DESC_A_SCALE_MODE,
237            &a_mode as *const i32 as *const c_void,
238            4,
239            "A_SCALE_MODE",
240        )?;
241        set(
242            DESC_B_SCALE_MODE,
243            &b_mode as *const i32 as *const c_void,
244            4,
245            "B_SCALE_MODE",
246        )?;
247        set(
248            DESC_A_SCALE_POINTER,
249            &weight_block_scale as *const u64 as *const c_void,
250            8,
251            "A_SCALE_POINTER",
252        )?;
253        set(
254            DESC_B_SCALE_POINTER,
255            &act_scale as *const u64 as *const c_void,
256            8,
257            "B_SCALE_POINTER",
258        )?;
259
260        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
261        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
262        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
263        chk(
264            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_8F_E4M3, k as u64, n as u64, k as i64),
265            "LayoutA",
266        )?;
267        chk(
268            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_8F_E4M3, k as u64, m as u64, k as i64),
269            "LayoutB",
270        )?;
271        chk(
272            cublasLtMatrixLayoutCreate(&mut ld_, CUDA_R_16BF, n as u64, m as u64, ldc as i64),
273            "LayoutD",
274        )?;
275        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
276        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
277        let ws_size = ctx.ws_size;
278        chk(
279            cublasLtMatmulPreferenceSetAttribute(
280                pref,
281                PREF_MAX_WORKSPACE_BYTES,
282                &ws_size as *const usize as *const c_void,
283                std::mem::size_of::<usize>(),
284            ),
285            "PrefWorkspace",
286        )?;
287        let mut result = [0u8; 128];
288        let mut returned: i32 = 0;
289        chk(
290            cublasLtMatmulAlgoGetHeuristic(
291                ctx.handle,
292                desc,
293                la,
294                lb,
295                ld_,
296                ld_,
297                pref,
298                1,
299                result.as_mut_ptr() as *mut c_void,
300                &mut returned,
301            ),
302            "AlgoGetHeuristic",
303        )?;
304        if returned < 1 {
305            bail!("cuBLASLt fp8: no algorithm for {m}x{n}x{k}");
306        }
307        let alpha: f32 = 1.0;
308        let beta: f32 = 0.0;
309        let status = cublasLtMatmul(
310            ctx.handle,
311            desc,
312            &alpha as *const f32 as *const c_void,
313            weight_fp8 as *const c_void,
314            la,
315            act_fp8 as *const c_void,
316            lb,
317            &beta as *const f32 as *const c_void,
318            out as *const c_void,
319            ld_,
320            out as *mut c_void,
321            ld_,
322            result.as_ptr() as *const c_void,
323            ctx.workspace as *mut c_void,
324            ctx.ws_size,
325            stream as *mut c_void,
326        );
327        cublasLtMatmulPreferenceDestroy(pref);
328        cublasLtMatrixLayoutDestroy(la);
329        cublasLtMatrixLayoutDestroy(lb);
330        cublasLtMatrixLayoutDestroy(ld_);
331        cublasLtMatmulDescDestroy(desc);
332        chk(status, "Matmul")?;
333    }
334    Ok(())
335}
metrale_gpu_runtime/cublaslt/fp8.rs

metrale_gpu_runtime/cublaslt/
fp8.rs
