spark_runtime/cutlass/
pack.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! CUTLASS NVFP4 weight pack / scale-swizzle / transpose host wrappers.
3
4use anyhow::{Result, bail};
5
6#[cfg(atlas_cutlass)]
7use std::ffi::c_void;
8
9#[cfg(atlas_cutlass)]
10use super::*;
11
12/// Repack an Atlas E4M3 weight scale into the CUTLASS SM120 blockscaled SFB
13/// swizzle atom (`tile_atom_to_shape_SFB`, ue4m3) that the grouped collective
14/// reads. M-independent (the SFB atom depends only on N,K) so this runs once
15/// per expert at load. `scale_out` must hold the swizzled SFB region the
16/// grouped kernel consumes.
17///
18/// `src_n_major` selects the SOURCE layout: `false` = Atlas-transposed
19/// `[K/16,N]`, `true` = checkpoint-native `[N,K/16]`. The N-major mode lets a
20/// checkpoint that already ships `[N,K/16]` scales (Laguna) build SFB without
21/// first materialising an Atlas-transposed copy. Output layout is identical
22/// either way.
23pub fn pack_weight_sfb(
24    scale_in: u64,
25    scale_out: u64,
26    n: u32,
27    k: u32,
28    src_n_major: bool,
29    stream: u64,
30) -> Result<()> {
31    #[cfg(atlas_cutlass)]
32    {
33        let status = unsafe {
34            atlas_cutlass_pack_weight_sfb(
35                scale_in as *const c_void,
36                scale_out as *mut c_void,
37                n as i32,
38                k as i32,
39                i32::from(src_n_major),
40                stream as *mut c_void,
41            )
42        };
43        if status != 0 {
44            bail!("CUTLASS weight SFB pack failed: status {status} for {n}x{k}");
45        }
46        Ok(())
47    }
48    #[cfg(not(atlas_cutlass))]
49    {
50        let _ = (scale_in, scale_out, n, k, src_n_major, stream);
51        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
52    }
53}
54
55/// Pack BF16 row-major weight `[N,K]` into the native CUTLASS NVFP4 layout:
56/// packed `[N,K/2]` (N-major, K-contiguous — NOT the Atlas transposed `[K/2,N]`)
57/// and E4M3 scales `[K/16,N]`. `weight_scale_2` is assumed to be 1.0 by the
58/// caller when feeding this into the native CUTLASS wrapper.
59pub fn pack_bf16_weight_to_nvfp4_t(
60    weight_bf16: u64,
61    packed_t: u64,
62    scale_t: u64,
63    n: u32,
64    k: u32,
65    stream: u64,
66) -> Result<()> {
67    #[cfg(atlas_cutlass)]
68    {
69        let status = unsafe {
70            atlas_cutlass_pack_bf16_weight_to_nvfp4_t(
71                weight_bf16 as *const c_void,
72                packed_t as *mut c_void,
73                scale_t as *mut c_void,
74                n as i32,
75                k as i32,
76                stream as *mut c_void,
77            )
78        };
79        if status != 0 {
80            bail!("CUTLASS BF16->NVFP4 weight pack failed: status {status} for {n}x{k}");
81        }
82        Ok(())
83    }
84    #[cfg(not(atlas_cutlass))]
85    {
86        let _ = (weight_bf16, packed_t, scale_t, n, k, stream);
87        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
88    }
89}
90
91/// Transpose an Atlas-packed NVFP4 weight from the checkpoint/hand-kernel
92/// `[K/2, N]` layout into CUTLASS's `[N, K/2]` layout (the byte order the
93/// native NVFP4 GEMM consumes for the ColumnMajor B operand). Pure byte
94/// transpose; nibble pairing within each byte is preserved. `dst_packed` must
95/// have `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(atlas_cutlass)]
104    {
105        let status = unsafe {
106            atlas_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(atlas_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}
spark_runtime/cutlass/pack.rs

spark_runtime/cutlass/
pack.rs
