metrale_gpu_runtime/
cutlass.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host-side FFI to the CUTLASS GEMM wrappers in `cuda/cutlass_*.cu`:
4//! the shared `extern` block and workspace here; dense BF16 and NVFP4 GEMMs in
5//! `gemm`, grouped per-expert NVFP4 MoE GEMMs in `grouped`, and NVFP4 weight
6//! pack, SFB swizzle and transpose in `pack`.
7//!
8//! Owner: gpu-runtime.
9//! Invariants:
10//! - The FFI and the workspace exist only under `cfg(metrale_cutlass)`, which
11//!   `build.rs` sets only when `CUTLASS_HOME` is set at build time. Without it
12//!   every public wrapper returns an error.
13//! - A wrapper that returns `Ok` got status 0 from its C entry point.
14
15#[cfg(metrale_cutlass)]
16use anyhow::{Result, bail};
17
18#[cfg(metrale_cutlass)]
19use std::ffi::c_void;
20#[cfg(metrale_cutlass)]
21use std::sync::OnceLock;
22
23mod gemm;
24mod grouped;
25mod pack;
26
27pub use gemm::{bf16_gemm_act_weight_t, nvfp4_gemm_bf16_act_weight_t};
28pub use grouped::{nvfp4_grouped_down, nvfp4_grouped_gate_up, nvfp4_grouped_gate_up_fused};
29pub use pack::{pack_bf16_weight_to_nvfp4_t, pack_weight_sfb, transpose_nvfp4_packed_kton};
30
31#[cfg(all(test, metrale_cutlass))]
32mod tests;
33
34#[cfg(metrale_cutlass)]
35unsafe extern "C" {
36    pub(crate) fn metrale_cutlass_bf16_gemm_act_weight_t(
37        act: *const c_void,
38        weight: *const c_void,
39        out: *mut c_void,
40        m: i32,
41        n: i32,
42        k: i32,
43        workspace: *mut c_void,
44        workspace_size: usize,
45        stream: *mut c_void,
46    ) -> i32;
47    pub(crate) fn metrale_cutlass_nvfp4_gemm_bf16_act_weight_t(
48        act: *const c_void,
49        weight_packed_t: *const c_void,
50        weight_scale_t: *const c_void,
51        weight_scale_2: f32,
52        out: *mut c_void,
53        m: i32,
54        n: i32,
55        k: i32,
56        workspace: *mut c_void,
57        workspace_size: usize,
58        stream: *mut c_void,
59    ) -> i32;
60    pub(crate) fn metrale_cutlass_pack_bf16_weight_to_nvfp4_t(
61        weight_bf16: *const c_void,
62        packed_t: *mut c_void,
63        scale_t: *mut c_void,
64        n: i32,
65        k: i32,
66        stream: *mut c_void,
67    ) -> i32;
68    pub(crate) fn metrale_cutlass_nvfp4_grouped_gate_up(
69        a_bf16: *const c_void,
70        gate_packed_ptrs: *const u64,
71        gate_scale_ptrs: *const u64,
72        gate_scale2_vals: *const f32,
73        up_packed_ptrs: *const u64,
74        up_scale_ptrs: *const u64,
75        up_scale2_vals: *const f32,
76        c_gate_bf16: *mut c_void,
77        c_up_bf16: *mut c_void,
78        expert_offsets_host: *const i32,
79        num_experts: i32,
80        n: i32,
81        k: i32,
82        workspace: *mut c_void,
83        workspace_size: usize,
84        stream: *mut c_void,
85    ) -> i32;
86    pub(crate) fn metrale_cutlass_nvfp4_grouped_gate_up_fused(
87        a_bf16: *const c_void,
88        sorted_token_ids: *const i32,
89        gate_packed_ptrs: *const u64,
90        gate_sfb_ptrs: *const u64,
91        gate_scale2_vals: *const f32,
92        up_packed_ptrs: *const u64,
93        up_sfb_ptrs: *const u64,
94        up_scale2_vals: *const f32,
95        c_gate_bf16: *mut c_void,
96        c_up_bf16: *mut c_void,
97        expert_offsets_host: *const i32,
98        num_experts: i32,
99        n: i32,
100        k: i32,
101        workspace: *mut c_void,
102        workspace_size: usize,
103        stream: *mut c_void,
104    ) -> i32;
105    pub(crate) fn metrale_cutlass_nvfp4_grouped_down(
106        a_bf16: *const c_void,
107        packed_ptrs: *const u64,
108        sfb_ptrs: *const u64,
109        scale2_vals: *const f32,
110        c_bf16: *mut c_void,
111        expert_offsets_host: *const i32,
112        num_experts: i32,
113        n: i32,
114        k: i32,
115        workspace: *mut c_void,
116        workspace_size: usize,
117        stream: *mut c_void,
118    ) -> i32;
119    pub(crate) fn metrale_cutlass_pack_weight_sfb(
120        scale_in: *const c_void,
121        scale_out: *mut c_void,
122        n: i32,
123        k: i32,
124        src_n_major: i32,
125        stream: *mut c_void,
126    ) -> i32;
127    pub(crate) fn metrale_cutlass_transpose_nvfp4_packed_kton(
128        src_packed_t: *const c_void,
129        dst_packed: *mut c_void,
130        n: i32,
131        k: i32,
132        stream: *mut c_void,
133    ) -> i32;
134    #[cfg(test)]
135    pub(crate) fn metrale_cutlass_bf16_gemm_act_weight_t_128x256(
136        act: *const c_void,
137        weight: *const c_void,
138        out: *mut c_void,
139        m: i32,
140        n: i32,
141        k: i32,
142        workspace: *mut c_void,
143        workspace_size: usize,
144        stream: *mut c_void,
145    ) -> i32;
146    #[cfg(test)]
147    pub(crate) fn metrale_cutlass_bf16_gemm_act_weight_t_256x128(
148        act: *const c_void,
149        weight: *const c_void,
150        out: *mut c_void,
151        m: i32,
152        n: i32,
153        k: i32,
154        workspace: *mut c_void,
155        workspace_size: usize,
156        stream: *mut c_void,
157    ) -> i32;
158    #[cfg(test)]
159    pub(crate) fn metrale_cutlass_bf16_gemm_act_weight_t_64x128(
160        act: *const c_void,
161        weight: *const c_void,
162        out: *mut c_void,
163        m: i32,
164        n: i32,
165        k: i32,
166        workspace: *mut c_void,
167        workspace_size: usize,
168        stream: *mut c_void,
169    ) -> i32;
170    #[cfg(test)]
171    pub(crate) fn metrale_cutlass_bf16_gemm_act_weight_t_128x64(
172        act: *const c_void,
173        weight: *const c_void,
174        out: *mut c_void,
175        m: i32,
176        n: i32,
177        k: i32,
178        workspace: *mut c_void,
179        workspace_size: usize,
180        stream: *mut c_void,
181    ) -> i32;
182    #[cfg(test)]
183    pub(crate) fn metrale_cutlass_bf16_gemm_act_weight_t_64x64(
184        act: *const c_void,
185        weight: *const c_void,
186        out: *mut c_void,
187        m: i32,
188        n: i32,
189        k: i32,
190        workspace: *mut c_void,
191        workspace_size: usize,
192        stream: *mut c_void,
193    ) -> i32;
194    #[cfg(test)]
195    pub(crate) fn metrale_cublaslt_bf16_gemm_act_weight_t_algo(
196        act: *const c_void,
197        weight: *const c_void,
198        out: *mut c_void,
199        m: i32,
200        n: i32,
201        k: i32,
202        workspace: *mut c_void,
203        workspace_size: usize,
204        stream: *mut c_void,
205        algo_index: i32,
206        returned_count: *mut i32,
207    ) -> i32;
208    pub(crate) fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
209}
210
211#[cfg(metrale_cutlass)]
212pub(crate) struct Ctx {
213    pub(crate) workspace: u64,
214    pub(crate) ws_size: usize,
215}
216
217#[cfg(metrale_cutlass)]
218unsafe impl Send for Ctx {}
219#[cfg(metrale_cutlass)]
220unsafe impl Sync for Ctx {}
221
222#[cfg(metrale_cutlass)]
223/// 2026-09-25: The shared CUTLASS workspace, allocated on first use and never
224/// freed. Static like the process CUDA context (`crate::cuda_host`): its size
225/// comes from a fixed budget, not from any model, so it is kept across model
226/// loads.
227static CTX: OnceLock<Ctx> = OnceLock::new();
228
229#[cfg(metrale_cutlass)]
230pub(crate) fn ctx() -> Result<&'static Ctx> {
231    if let Some(c) = CTX.get() {
232        return Ok(c);
233    }
234    // 2026-09-25: Scratch for every CUTLASS wrapper; the grouped NVFP4 path
235    // carves packed A, SFA, per-group arrays and the GEMM workspace from it.
236    // 512 MiB unless `METRALE_CUTLASS_WORKSPACE_MB` parses as a number.
237    let ws_size = std::env::var("METRALE_CUTLASS_WORKSPACE_MB")
238        .ok()
239        .and_then(|v| v.parse::<usize>().ok())
240        .unwrap_or(512)
241        * 1024
242        * 1024;
243    let mut workspace = 0u64;
244    let status = unsafe { cuMemAlloc_v2(&mut workspace, ws_size) };
245    if status != 0 {
246        bail!("cuMemAlloc CUTLASS workspace failed: {status}");
247    }
248    let _ = CTX.set(Ctx { workspace, ws_size });
249    Ok(CTX.get().unwrap())
250}
metrale_gpu_runtime/cutlass.rs

metrale_gpu_runtime/
cutlass.rs
