spark_runtime/cublaslt/
fp8.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Native FP8 (E4M3) cuBLASLt GEMM paths (row-wise + 128-block scaled).
4
5use anyhow::{Result, bail};
6use std::ffi::c_void;
7
8use super::*;
9
10/// Native FP8 (E4M3) `out[M,N] = act[M,K] @ weight[N,K]ᵀ` → BF16 with ROW-WISE
11/// scaling (OUTER_VEC): per-output-row weight scale `weight_scale[N]` and
12/// per-token activation scale `act_scale[M]`. This is the fp8 path GB10/sm_121
13/// actually supports (128-block fp8 is B200-only). ~1.8× the bf16 path.
14/// cuBLAS folds `A_scale[i]·B_scale[j]` into the FP32 epilogue; with D=`[N,M]`,
15/// i indexes weight rows (N) and j indexes tokens (M) — exactly row-wise.
16#[allow(clippy::too_many_arguments)]
17pub fn fp8_gemm_act_weight_t_rowwise(
18    act_fp8: u64,
19    act_scale: u64,
20    weight_fp8: u64,
21    weight_scale: u64,
22    out: u64,
23    m: u32,
24    n: u32,
25    k: u32,
26    stream: u64,
27) -> Result<()> {
28    let ctx = ctx()?;
29    unsafe {
30        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
31        chk(
32            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
33            "DescCreate",
34        )?;
35        let ta = CUBLAS_OP_T;
36        let tb = CUBLAS_OP_N;
37        let set = |attr: u32, val: *const c_void, sz: usize, what: &str| -> Result<()> {
38            chk(cublasLtMatmulDescSetAttribute(desc, attr, val, sz), what)
39        };
40        set(DESC_TRANSA, &ta as *const i32 as *const c_void, 4, "TRANSA")?;
41        set(DESC_TRANSB, &tb as *const i32 as *const c_void, 4, "TRANSB")?;
42        let mode = SCALE_MODE_OUTER_VEC_32F;
43        set(
44            DESC_A_SCALE_MODE,
45            &mode as *const i32 as *const c_void,
46            4,
47            "A_SCALE_MODE",
48        )?;
49        set(
50            DESC_B_SCALE_MODE,
51            &mode as *const i32 as *const c_void,
52            4,
53            "B_SCALE_MODE",
54        )?;
55        set(
56            DESC_A_SCALE_POINTER,
57            &weight_scale as *const u64 as *const c_void,
58            8,
59            "A_SCALE_POINTER",
60        )?;
61        set(
62            DESC_B_SCALE_POINTER,
63            &act_scale as *const u64 as *const c_void,
64            8,
65            "B_SCALE_POINTER",
66        )?;
67
68        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
69        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
70        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
71        chk(
72            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_8F_E4M3, k as u64, n as u64, k as i64),
73            "LayoutA",
74        )?;
75        chk(
76            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_8F_E4M3, k as u64, m as u64, k as i64),
77            "LayoutB",
78        )?;
79        chk(
80            cublasLtMatrixLayoutCreate(&mut ld_, CUDA_R_16BF, n as u64, m as u64, n as i64),
81            "LayoutD",
82        )?;
83        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
84        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
85        let ws_size = ctx.ws_size;
86        chk(
87            cublasLtMatmulPreferenceSetAttribute(
88                pref,
89                PREF_MAX_WORKSPACE_BYTES,
90                &ws_size as *const usize as *const c_void,
91                std::mem::size_of::<usize>(),
92            ),
93            "PrefWorkspace",
94        )?;
95        let mut result = [0u8; 128];
96        let mut returned: i32 = 0;
97        chk(
98            cublasLtMatmulAlgoGetHeuristic(
99                ctx.handle,
100                desc,
101                la,
102                lb,
103                ld_,
104                ld_,
105                pref,
106                1,
107                result.as_mut_ptr() as *mut c_void,
108                &mut returned,
109            ),
110            "AlgoGetHeuristic",
111        )?;
112        if returned < 1 {
113            bail!("cuBLASLt fp8 rowwise: no algorithm for {m}x{n}x{k}");
114        }
115        let alpha: f32 = 1.0;
116        let beta: f32 = 0.0;
117        let status = cublasLtMatmul(
118            ctx.handle,
119            desc,
120            &alpha as *const f32 as *const c_void,
121            weight_fp8 as *const c_void,
122            la,
123            act_fp8 as *const c_void,
124            lb,
125            &beta as *const f32 as *const c_void,
126            out as *const c_void,
127            ld_,
128            out as *mut c_void,
129            ld_,
130            result.as_ptr() as *const c_void,
131            ctx.workspace as *mut c_void,
132            ctx.ws_size,
133            stream as *mut c_void,
134        );
135        cublasLtMatmulPreferenceDestroy(pref);
136        cublasLtMatrixLayoutDestroy(la);
137        cublasLtMatrixLayoutDestroy(lb);
138        cublasLtMatrixLayoutDestroy(ld_);
139        cublasLtMatmulDescDestroy(desc);
140        chk(status, "Matmul")?;
141    }
142    Ok(())
143}
144
145/// Native FP8 (E4M3) `out[M,N] = act[M,K] @ weight[N,K]ᵀ` → BF16, with the
146/// weight per-128×128-block FP32-scaled (matches Atlas's `Fp8Weight.row_scale`
147/// layout exactly) and the activation per-[token,128-of-K] FP32-scaled.
148/// ~1.8× the bf16 path (152 vs 85 TFLOPS on GB10).
149///
150/// ⚠ SCALE-TENSOR LAYOUTS — the two operands do NOT agree, and getting this
151/// wrong is silent (see [`super::scale_layout`] for the doc quotes, the H100
152/// measurement that caught it, and the index math):
153///
154/// * `weight_block_scale` (A, BLK128x128_32F) is K-major, `L4 × ⌈N/128⌉` —
155///   the checkpoint's row-major `[N/128, K/128]` grid as-is, valid while
156///   `⌈K/128⌉` is a multiple of 4 (`scale_layout::blk128x128_stride_ok`).
157/// * `act_scale` (B, VEC128_32F) is N-major, `M × ⌈K/128⌉` with the TOKEN
158///   index contiguous — i.e. `[K/128, M]`, the TRANSPOSE of what
159///   `per_token_group_quant_fp8` writes. Callers adapt it with the
160///   `fp8_act_scale_to_kmajor` kernel; passing the quantizer's buffer straight
161///   through permutes the scales and costs ~8% relative RMS at M≈1200.
162///
163/// `m` must already include the caller's pad (the docs require the matmul's M
164/// and N to be multiples of 4), and `act_fp8`/`act_scale` must cover it.
165///
166/// The output is CONTIGUOUS `[M, N]`; [`fp8_gemm_act_weight_t_blkscaled_ldc`]
167/// is the same GEMM with a caller-chosen output row pitch.
168#[allow(clippy::too_many_arguments)]
169pub fn fp8_gemm_act_weight_t_blkscaled(
170    act_fp8: u64,
171    act_scale: u64,
172    weight_fp8: u64,
173    weight_block_scale: u64,
174    out: u64,
175    m: u32,
176    n: u32,
177    k: u32,
178    stream: u64,
179) -> Result<()> {
180    fp8_gemm_act_weight_t_blkscaled_ldc(
181        act_fp8,
182        act_scale,
183        weight_fp8,
184        weight_block_scale,
185        out,
186        m,
187        n,
188        k,
189        n,
190        stream,
191    )
192}
193
194/// [`fp8_gemm_act_weight_t_blkscaled`] with an explicit output ROW PITCH.
195///
196/// WHY (#927, the 5..16-row decode projections). cuBLASLt's D operand is a
197/// column-major `[N, M]` layout with leading dimension `ldc`, which is exactly
198/// a row-major `[M, N]` whose rows are `ldc` BF16 elements apart — so one
199/// parameter is the whole difference between a contiguous `[M, N]` output and
200/// writing straight into a strided slot. The multi-seq decode QKV buffer is
201/// `[n, per_seq_qkv]` with Q at 0, K at `q_proj_bytes` and V after it, so
202/// `ldc = per_seq_qkv / 2` puts each row's `n` outputs in its own sequence's
203/// slot with the gaps left alone — the same thing the `_strided` GEMV entry
204/// points do, without a staging buffer or a scatter kernel.
205///
206/// ⚠ WRITE EXTENT. The library writes `n` elements of EACH of the `m` columns,
207/// i.e. the last byte touched is at element `(m - 1) * ldc + n`. `m` here is
208/// the caller's PADDED row count, so the phantom rows are written too; callers
209/// must bound that extent against their buffer — see
210/// `spark_model::layers::ops::strided_out_extent_elems`, which is the SSOT for
211/// the arithmetic and is unit-tested on the CPU.
212///
213/// `ldc >= n` is required by the library (a leading dimension shorter than the
214/// column is rejected); it is checked here so the failure names itself instead
215/// of arriving as a cuBLAS status code.
216#[allow(clippy::too_many_arguments)]
217pub fn fp8_gemm_act_weight_t_blkscaled_ldc(
218    act_fp8: u64,
219    act_scale: u64,
220    weight_fp8: u64,
221    weight_block_scale: u64,
222    out: u64,
223    m: u32,
224    n: u32,
225    k: u32,
226    ldc: u32,
227    stream: u64,
228) -> Result<()> {
229    if ldc < n {
230        bail!("cuBLASLt fp8: output row pitch ldc={ldc} is shorter than N={n}");
231    }
232    let ctx = ctx()?;
233    unsafe {
234        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
235        chk(
236            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
237            "DescCreate",
238        )?;
239        let ta = CUBLAS_OP_T;
240        let tb = CUBLAS_OP_N;
241        let set = |attr: u32, val: *const c_void, sz: usize, what: &str| -> Result<()> {
242            chk(cublasLtMatmulDescSetAttribute(desc, attr, val, sz), what)
243        };
244        set(DESC_TRANSA, &ta as *const i32 as *const c_void, 4, "TRANSA")?;
245        set(DESC_TRANSB, &tb as *const i32 as *const c_void, 4, "TRANSB")?;
246        // FP8 block scaling requires BOTH operands use a 128-block mode (SCALAR
247        // is rejected → status 7). Weight = per-128×128 block, activation =
248        // per-[token,128-of-K] VEC128 (DeepSeek block-fp8 scheme).
249        let a_mode = SCALE_MODE_BLK128X128_32F;
250        let b_mode = SCALE_MODE_VEC128_32F;
251        set(
252            DESC_A_SCALE_MODE,
253            &a_mode as *const i32 as *const c_void,
254            4,
255            "A_SCALE_MODE",
256        )?;
257        set(
258            DESC_B_SCALE_MODE,
259            &b_mode as *const i32 as *const c_void,
260            4,
261            "B_SCALE_MODE",
262        )?;
263        set(
264            DESC_A_SCALE_POINTER,
265            &weight_block_scale as *const u64 as *const c_void,
266            8,
267            "A_SCALE_POINTER",
268        )?;
269        set(
270            DESC_B_SCALE_POINTER,
271            &act_scale as *const u64 as *const c_void,
272            8,
273            "B_SCALE_POINTER",
274        )?;
275
276        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
277        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
278        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
279        chk(
280            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_8F_E4M3, k as u64, n as u64, k as i64),
281            "LayoutA",
282        )?;
283        chk(
284            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_8F_E4M3, k as u64, m as u64, k as i64),
285            "LayoutB",
286        )?;
287        chk(
288            cublasLtMatrixLayoutCreate(&mut ld_, CUDA_R_16BF, n as u64, m as u64, ldc as i64),
289            "LayoutD",
290        )?;
291        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
292        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
293        let ws_size = ctx.ws_size;
294        chk(
295            cublasLtMatmulPreferenceSetAttribute(
296                pref,
297                PREF_MAX_WORKSPACE_BYTES,
298                &ws_size as *const usize as *const c_void,
299                std::mem::size_of::<usize>(),
300            ),
301            "PrefWorkspace",
302        )?;
303        let mut result = [0u8; 128];
304        let mut returned: i32 = 0;
305        chk(
306            cublasLtMatmulAlgoGetHeuristic(
307                ctx.handle,
308                desc,
309                la,
310                lb,
311                ld_,
312                ld_,
313                pref,
314                1,
315                result.as_mut_ptr() as *mut c_void,
316                &mut returned,
317            ),
318            "AlgoGetHeuristic",
319        )?;
320        if returned < 1 {
321            bail!("cuBLASLt fp8: no algorithm for {m}x{n}x{k}");
322        }
323        let alpha: f32 = 1.0;
324        let beta: f32 = 0.0;
325        let status = cublasLtMatmul(
326            ctx.handle,
327            desc,
328            &alpha as *const f32 as *const c_void,
329            weight_fp8 as *const c_void,
330            la,
331            act_fp8 as *const c_void,
332            lb,
333            &beta as *const f32 as *const c_void,
334            out as *const c_void,
335            ld_,
336            out as *mut c_void,
337            ld_,
338            result.as_ptr() as *const c_void,
339            ctx.workspace as *mut c_void,
340            ctx.ws_size,
341            stream as *mut c_void,
342        );
343        cublasLtMatmulPreferenceDestroy(pref);
344        cublasLtMatrixLayoutDestroy(la);
345        cublasLtMatrixLayoutDestroy(lb);
346        cublasLtMatrixLayoutDestroy(ld_);
347        cublasLtMatmulDescDestroy(desc);
348        chk(status, "Matmul")?;
349    }
350    Ok(())
351}
spark_runtime/cublaslt/fp8.rs

spark_runtime/cublaslt/
fp8.rs
