metrale_gpu_runtime/cutlass/
pack.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: CUTLASS NVFP4 weight pack, scale swizzle and transpose wrappers.
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: Swizzle an E4M3 weight scale into the CUTLASS SFB layout
17/// (`tile_atom_to_shape_SFB`, ue4m3) that the grouped GEMM reads. The layout
18/// depends only on `n` and `k`, so it is built once per expert at load
19/// (`MoeLayer::build_cutlass_grouped_sfb`). `scale_out` must hold the whole
20/// swizzled region.
21///
22/// `src_n_major` selects the source layout: `false` reads `[K/16,N]`, `true`
23/// reads `[N,K/16]`. The output is the same either way.
24pub fn pack_weight_sfb(
25    scale_in: u64,
26    scale_out: u64,
27    n: u32,
28    k: u32,
29    src_n_major: bool,
30    stream: u64,
31) -> Result<()> {
32    #[cfg(metrale_cutlass)]
33    {
34        let status = unsafe {
35            metrale_cutlass_pack_weight_sfb(
36                scale_in as *const c_void,
37                scale_out as *mut c_void,
38                n as i32,
39                k as i32,
40                i32::from(src_n_major),
41                stream as *mut c_void,
42            )
43        };
44        if status != 0 {
45            bail!("CUTLASS weight SFB pack failed: status {status} for {n}x{k}");
46        }
47        Ok(())
48    }
49    #[cfg(not(metrale_cutlass))]
50    {
51        let _ = (scale_in, scale_out, n, k, src_n_major, stream);
52        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
53    }
54}
55
56/// 2026-09-25: Pack a row-major BF16 weight `[N,K]` into CUTLASS NVFP4: packed
57/// `[N,K/2]` (K-contiguous) and E4M3 scales `[K/16,N]`, each scale the group's
58/// max magnitude / 6. The scales carry no second-level factor, so pass
59/// `weight_scale_2 = 1.0` to `nvfp4_gemm_bf16_act_weight_t`.
60pub fn pack_bf16_weight_to_nvfp4_t(
61    weight_bf16: u64,
62    packed_t: u64,
63    scale_t: u64,
64    n: u32,
65    k: u32,
66    stream: u64,
67) -> Result<()> {
68    #[cfg(metrale_cutlass)]
69    {
70        let status = unsafe {
71            metrale_cutlass_pack_bf16_weight_to_nvfp4_t(
72                weight_bf16 as *const c_void,
73                packed_t as *mut c_void,
74                scale_t as *mut c_void,
75                n as i32,
76                k as i32,
77                stream as *mut c_void,
78            )
79        };
80        if status != 0 {
81            bail!("CUTLASS BF16->NVFP4 weight pack failed: status {status} for {n}x{k}");
82        }
83        Ok(())
84    }
85    #[cfg(not(metrale_cutlass))]
86    {
87        let _ = (weight_bf16, packed_t, scale_t, n, k, stream);
88        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
89    }
90}
91
92/// 2026-09-25: Transpose a packed NVFP4 weight from `[K/2, N]` into the CUTLASS
93/// `[N, K/2]` layout that `nvfp4_gemm_bf16_act_weight_t` reads. A byte
94/// transpose: the two nibbles of each byte stay together. `dst_packed` must hold
95/// `N * K/2` bytes.
96pub fn transpose_nvfp4_packed_kton(
97    src_packed_t: u64,
98    dst_packed: u64,
99    n: u32,
100    k: u32,
101    stream: u64,
102) -> Result<()> {
103    #[cfg(metrale_cutlass)]
104    {
105        let status = unsafe {
106            metrale_cutlass_transpose_nvfp4_packed_kton(
107                src_packed_t as *const c_void,
108                dst_packed as *mut c_void,
109                n as i32,
110                k as i32,
111                stream as *mut c_void,
112            )
113        };
114        if status != 0 {
115            bail!("CUTLASS NVFP4 weight transpose failed: status {status} for {n}x{k}");
116        }
117        Ok(())
118    }
119    #[cfg(not(metrale_cutlass))]
120    {
121        let _ = (src_packed_t, dst_packed, n, k, stream);
122        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
123    }
124}
metrale_gpu_runtime/cutlass/pack.rs

metrale_gpu_runtime/cutlass/
pack.rs
