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 cast at unit scale (post-RMSNorm acts sit
148/// in fp8's range). ~1.8× the bf16 path (152 vs 85 TFLOPS on GB10).
149#[allow(clippy::too_many_arguments)]
150pub fn fp8_gemm_act_weight_t_blkscaled(
151    act_fp8: u64,
152    act_scale: u64,
153    weight_fp8: u64,
154    weight_block_scale: u64,
155    out: u64,
156    m: u32,
157    n: u32,
158    k: u32,
159    stream: u64,
160) -> Result<()> {
161    let ctx = ctx()?;
162    unsafe {
163        let mut desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
164        chk(
165            cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F),
166            "DescCreate",
167        )?;
168        let ta = CUBLAS_OP_T;
169        let tb = CUBLAS_OP_N;
170        let set = |attr: u32, val: *const c_void, sz: usize, what: &str| -> Result<()> {
171            chk(cublasLtMatmulDescSetAttribute(desc, attr, val, sz), what)
172        };
173        set(DESC_TRANSA, &ta as *const i32 as *const c_void, 4, "TRANSA")?;
174        set(DESC_TRANSB, &tb as *const i32 as *const c_void, 4, "TRANSB")?;
175        // FP8 block scaling requires BOTH operands use a 128-block mode (SCALAR
176        // is rejected → status 7). Weight = per-128×128 block, activation =
177        // per-[token,128-of-K] VEC128 (DeepSeek block-fp8 scheme).
178        let a_mode = SCALE_MODE_BLK128X128_32F;
179        let b_mode = SCALE_MODE_VEC128_32F;
180        set(
181            DESC_A_SCALE_MODE,
182            &a_mode as *const i32 as *const c_void,
183            4,
184            "A_SCALE_MODE",
185        )?;
186        set(
187            DESC_B_SCALE_MODE,
188            &b_mode as *const i32 as *const c_void,
189            4,
190            "B_SCALE_MODE",
191        )?;
192        set(
193            DESC_A_SCALE_POINTER,
194            &weight_block_scale as *const u64 as *const c_void,
195            8,
196            "A_SCALE_POINTER",
197        )?;
198        set(
199            DESC_B_SCALE_POINTER,
200            &act_scale as *const u64 as *const c_void,
201            8,
202            "B_SCALE_POINTER",
203        )?;
204
205        let mut la: cublasLtMatrixLayout_t = std::ptr::null_mut();
206        let mut lb: cublasLtMatrixLayout_t = std::ptr::null_mut();
207        let mut ld_: cublasLtMatrixLayout_t = std::ptr::null_mut();
208        chk(
209            cublasLtMatrixLayoutCreate(&mut la, CUDA_R_8F_E4M3, k as u64, n as u64, k as i64),
210            "LayoutA",
211        )?;
212        chk(
213            cublasLtMatrixLayoutCreate(&mut lb, CUDA_R_8F_E4M3, k as u64, m as u64, k as i64),
214            "LayoutB",
215        )?;
216        chk(
217            cublasLtMatrixLayoutCreate(&mut ld_, CUDA_R_16BF, n as u64, m as u64, n as i64),
218            "LayoutD",
219        )?;
220        let mut pref: cublasLtMatmulPreference_t = std::ptr::null_mut();
221        chk(cublasLtMatmulPreferenceCreate(&mut pref), "PrefCreate")?;
222        let ws_size = ctx.ws_size;
223        chk(
224            cublasLtMatmulPreferenceSetAttribute(
225                pref,
226                PREF_MAX_WORKSPACE_BYTES,
227                &ws_size as *const usize as *const c_void,
228                std::mem::size_of::<usize>(),
229            ),
230            "PrefWorkspace",
231        )?;
232        let mut result = [0u8; 128];
233        let mut returned: i32 = 0;
234        chk(
235            cublasLtMatmulAlgoGetHeuristic(
236                ctx.handle,
237                desc,
238                la,
239                lb,
240                ld_,
241                ld_,
242                pref,
243                1,
244                result.as_mut_ptr() as *mut c_void,
245                &mut returned,
246            ),
247            "AlgoGetHeuristic",
248        )?;
249        if returned < 1 {
250            bail!("cuBLASLt fp8: no algorithm for {m}x{n}x{k}");
251        }
252        let alpha: f32 = 1.0;
253        let beta: f32 = 0.0;
254        let status = cublasLtMatmul(
255            ctx.handle,
256            desc,
257            &alpha as *const f32 as *const c_void,
258            weight_fp8 as *const c_void,
259            la,
260            act_fp8 as *const c_void,
261            lb,
262            &beta as *const f32 as *const c_void,
263            out as *const c_void,
264            ld_,
265            out as *mut c_void,
266            ld_,
267            result.as_ptr() as *const c_void,
268            ctx.workspace as *mut c_void,
269            ctx.ws_size,
270            stream as *mut c_void,
271        );
272        cublasLtMatmulPreferenceDestroy(pref);
273        cublasLtMatrixLayoutDestroy(la);
274        cublasLtMatrixLayoutDestroy(lb);
275        cublasLtMatrixLayoutDestroy(ld_);
276        cublasLtMatmulDescDestroy(desc);
277        chk(status, "Matmul")?;
278    }
279    Ok(())
280}
spark_runtime/cublaslt/fp8.rs

spark_runtime/cublaslt/
fp8.rs
