metrale_gpu_runtime/cutlass/
gemm.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Dense CUTLASS GEMM wrappers: BF16, and NVFP4 with BF16 output.
4//!
5//! Owner: gpu-runtime.
6//! Invariants: none beyond the types.
7
8use anyhow::{Result, bail};
9
10#[cfg(metrale_cutlass)]
11use std::ffi::c_void;
12
13#[cfg(metrale_cutlass)]
14use super::*;
15
16/// 2026-09-25: Row-major `out[M,N] = act[M,K] @ weight[N,K]^T`, all BF16.
17#[allow(clippy::too_many_arguments)]
18pub fn bf16_gemm_act_weight_t(
19    act: u64,
20    weight: u64,
21    out: u64,
22    m: u32,
23    n: u32,
24    k: u32,
25    stream: u64,
26) -> Result<()> {
27    #[cfg(metrale_cutlass)]
28    {
29        let ctx = ctx()?;
30        let status = unsafe {
31            metrale_cutlass_bf16_gemm_act_weight_t(
32                act as *const c_void,
33                weight as *const c_void,
34                out as *mut c_void,
35                m as i32,
36                n as i32,
37                k as i32,
38                ctx.workspace as *mut c_void,
39                ctx.ws_size,
40                stream as *mut c_void,
41            )
42        };
43        if status != 0 {
44            bail!("CUTLASS bf16 GEMM failed: status {status} for {m}x{n}x{k}");
45        }
46        Ok(())
47    }
48    #[cfg(not(metrale_cutlass))]
49    {
50        let _ = (act, weight, out, m, n, k, stream);
51        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
52    }
53}
54
55/// 2026-09-25: NVFP4 dense projection with BF16 output:
56/// `out[M,N] = quant_nvfp4(act[M,K]) @ weight[N,K]^T`, scaled by `weight_scale_2`.
57///
58/// `weight_packed_t` is the CUTLASS `[N,K/2]` K-contiguous byte layout that
59/// `pack_bf16_weight_to_nvfp4_t` writes and `transpose_nvfp4_packed_kton`
60/// produces from `[K/2,N]`; `weight_scale_t` is E4M3 `[K/16,N]`. The wrapper
61/// quantizes the activation and swizzles the weight scales into the shared
62/// workspace before the GEMM.
63#[allow(clippy::too_many_arguments)]
64pub fn nvfp4_gemm_bf16_act_weight_t(
65    act: u64,
66    weight_packed_t: u64,
67    weight_scale_t: u64,
68    weight_scale_2: f32,
69    out: u64,
70    m: u32,
71    n: u32,
72    k: u32,
73    stream: u64,
74) -> Result<()> {
75    #[cfg(metrale_cutlass)]
76    {
77        let ctx = ctx()?;
78        let status = unsafe {
79            metrale_cutlass_nvfp4_gemm_bf16_act_weight_t(
80                act as *const c_void,
81                weight_packed_t as *const c_void,
82                weight_scale_t as *const c_void,
83                weight_scale_2,
84                out as *mut c_void,
85                m as i32,
86                n as i32,
87                k as i32,
88                ctx.workspace as *mut c_void,
89                ctx.ws_size,
90                stream as *mut c_void,
91            )
92        };
93        if status != 0 {
94            bail!("CUTLASS nvfp4 GEMM failed: status {status} for {m}x{n}x{k}");
95        }
96        Ok(())
97    }
98    #[cfg(not(metrale_cutlass))]
99    {
100        let _ = (
101            act,
102            weight_packed_t,
103            weight_scale_t,
104            weight_scale_2,
105            out,
106            m,
107            n,
108            k,
109            stream,
110        );
111        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
112    }
113}
metrale_gpu_runtime/cutlass/gemm.rs

metrale_gpu_runtime/cutlass/
gemm.rs
