spark_runtime/cutlass/
gemm.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Dense CUTLASS GEMM host wrappers (BF16 + native NVFP4 projection).
3
4use anyhow::{Result, bail};
5
6#[cfg(atlas_cutlass)]
7use std::ffi::c_void;
8
9#[cfg(atlas_cutlass)]
10use super::*;
11
12/// Row-major `out[M,N] = act[M,K] @ weight[N,K]^T`, all BF16.
13#[allow(clippy::too_many_arguments)]
14pub fn bf16_gemm_act_weight_t(
15    act: u64,
16    weight: u64,
17    out: u64,
18    m: u32,
19    n: u32,
20    k: u32,
21    stream: u64,
22) -> Result<()> {
23    #[cfg(atlas_cutlass)]
24    {
25        let ctx = ctx()?;
26        let status = unsafe {
27            atlas_cutlass_bf16_gemm_act_weight_t(
28                act as *const c_void,
29                weight as *const c_void,
30                out as *mut c_void,
31                m as i32,
32                n as i32,
33                k as i32,
34                ctx.workspace as *mut c_void,
35                ctx.ws_size,
36                stream as *mut c_void,
37            )
38        };
39        if status != 0 {
40            bail!("CUTLASS bf16 GEMM failed: status {status} for {m}x{n}x{k}");
41        }
42        Ok(())
43    }
44    #[cfg(not(atlas_cutlass))]
45    {
46        let _ = (act, weight, out, m, n, k, stream);
47        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
48    }
49}
50
51/// Native CUTLASS NVFP4 dense projection:
52/// `out[M,N] = quant_nvfp4(act[M,K]) @ weight_t[N,K]^T -> BF16`.
53///
54/// `weight_packed_t` and `weight_scale_t` are Atlas's transposed NVFP4
55/// prefill layout: packed data `[K/2,N]`, scales `[K/16,N]`. The wrapper
56/// repacks activation and scale tensors into CUTLASS's SM120 blockscaled
57/// layouts in the shared CUTLASS workspace before dispatch.
58#[allow(clippy::too_many_arguments)]
59pub fn nvfp4_gemm_bf16_act_weight_t(
60    act: u64,
61    weight_packed_t: u64,
62    weight_scale_t: u64,
63    weight_scale_2: f32,
64    out: u64,
65    m: u32,
66    n: u32,
67    k: u32,
68    stream: u64,
69) -> Result<()> {
70    #[cfg(atlas_cutlass)]
71    {
72        let ctx = ctx()?;
73        let status = unsafe {
74            atlas_cutlass_nvfp4_gemm_bf16_act_weight_t(
75                act as *const c_void,
76                weight_packed_t as *const c_void,
77                weight_scale_t as *const c_void,
78                weight_scale_2,
79                out as *mut c_void,
80                m as i32,
81                n as i32,
82                k as i32,
83                ctx.workspace as *mut c_void,
84                ctx.ws_size,
85                stream as *mut c_void,
86            )
87        };
88        if status != 0 {
89            bail!("CUTLASS nvfp4 GEMM failed: status {status} for {m}x{n}x{k}");
90        }
91        Ok(())
92    }
93    #[cfg(not(atlas_cutlass))]
94    {
95        let _ = (
96            act,
97            weight_packed_t,
98            weight_scale_t,
99            weight_scale_2,
100            out,
101            m,
102            n,
103            k,
104            stream,
105        );
106        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
107    }
108}
spark_runtime/cutlass/gemm.rs

spark_runtime/cutlass/
gemm.rs
