metrale_gpu_runtime/cublaslt/
scale_layout.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Index math for the FP32 scale tensors of the block-scaled
4//! FP8 GEMM ([`super::fp8_gemm_act_weight_t_blkscaled`]): the activation
5//! scales in the layout it passes to cuBLASLt, in the quantizer's layout, and
6//! the weight-scale stride check.
7//!
8//! In cuBLASLt's terms the GEMM is `D[N,M] = opT(weightᶜ[K,N]) ·
9//! opN(actᶜ[K,M])`, so the library's M is the weight's N and its N is the
10//! token count. The activation (B, `VEC128_32F`) scales are laid out
11//! `[L, m_pad]`, token index contiguous, with `L = ceil(K/128)`; the
12//! quantizer `per_token_group_quant_fp8` writes `[M, L]`, K group
13//! contiguous. The weight (A, `BLK128x128_32F`) scales are the checkpoint's
14//! `[ceil(N/128), L]` grid, used as they are when `L` is a multiple of 4.
15//! The layouts follow the cuBLAS manual section "128-element 1D and 128x128
16//! 2D Block Scaling For FP8 Data Types".
17//!
18//! Measured 2026-09-11 on 1xH100 (`native_fp8_ffn_w8a8_microtest`, cuBLASLt
19//! 13.1): with the quantizer's `[M, L]` scales passed as B, the GEMM
20//! disagreed with the in-tree `fp8_gemm_t_blockscaled` on the same inputs,
21//! rel_rms 1.1e-2 at M=64 and 7.7e-2 to 8.8e-2 at M=1193.
22//!
23//! The metal build compiles this file too (`cublaslt_metal_stub.rs`).
24//!
25//! Owner: gpu-runtime (cuBLASLt wrapper).
26//! Invariants: none beyond the types.
27
28/// 2026-09-25: Number of 128-wide K groups in `k`: `L = ceil(k / 128)`.
29#[must_use]
30pub fn k_groups(k: usize) -> usize {
31    k.div_ceil(128)
32}
33
34/// 2026-09-25: Offset of the scale for `(token, k_group)` in the B-operand
35/// layout `[L, m_pad]`, token index contiguous. The CUDA adapter
36/// `fp8_act_scale_to_kmajor` (`kernels/gb10/common/fp8_scale_transpose.cu`)
37/// writes the same index.
38#[must_use]
39pub const fn vec128_b_index(m_pad: usize, token: usize, k_group: usize) -> usize {
40    k_group * m_pad + token
41}
42
43/// 2026-09-25: Offset of the same `(token, k_group)` scale in the layout
44/// `per_token_group_quant_fp8` writes: row-major `[M, L]`, K group
45/// contiguous. `fp8_gemm_t_blockscaled` reads this layout.
46#[must_use]
47pub const fn rowmajor_index(l: usize, token: usize, k_group: usize) -> usize {
48    token * l + k_group
49}
50
51/// 2026-09-25: FP32 element count of the B-operand scale tensor:
52/// `m_pad * L`.
53#[must_use]
54pub fn vec128_b_elems(m_pad: usize, k: usize) -> usize {
55    m_pad * k_groups(k)
56}
57
58/// 2026-09-25: Whether `L = ceil(k / 128)` is a multiple of 4, the condition
59/// under which the checkpoint's row-major weight-scale grid is passed to the
60/// `BLK128x128_32F` operand unchanged. The cuBLASLt dispatch gates check it
61/// (`dense_ffn_w8a8_prefill.rs`, `dispatch_proj_decode.rs`).
62#[must_use]
63pub fn blk128x128_stride_ok(k: usize) -> bool {
64    k_groups(k).is_multiple_of(4)
65}
66
67/// 2026-09-25: CPU reference for the `fp8_act_scale_to_kmajor` kernel: read
68/// row-major `[m, l]` scales and write `[l, m_pad]`, with the pad tokens
69/// `m..m_pad` set to 0.0, as the kernel does.
70#[must_use]
71pub fn act_scale_rowmajor_to_kmajor(src: &[f32], m: usize, m_pad: usize, l: usize) -> Vec<f32> {
72    debug_assert!(m_pad >= m, "pad cannot shrink the token extent");
73    debug_assert!(src.len() >= m * l, "source holds fewer than m x l scales");
74    let mut dst = vec![0.0f32; vec128_b_elems(m_pad, l * 128)];
75    for token in 0..m {
76        for kg in 0..l {
77            dst[vec128_b_index(m_pad, token, kg)] = src[rowmajor_index(l, token, kg)];
78        }
79    }
80    dst
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    /// 2026-09-25: Hand-placed values: `src` is row-major `[M=3, L=2]`, and
88    /// the output groups tokens by K group, with one zero pad token each.
89    #[test]
90    fn kmajor_adapter_transposes_and_zero_fills_the_pad() {
91        let src = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
92        let dst = act_scale_rowmajor_to_kmajor(&src, 3, 4, 2);
93        assert_eq!(dst, vec![1.0, 3.0, 5.0, 0.0, 2.0, 4.0, 6.0, 0.0]);
94    }
95
96    #[test]
97    fn adapter_is_identity_free_only_when_one_dimension_is_trivial() {
98        // 2026-09-25: With one K group, or one token, the two layouts
99        // coincide.
100        let src = [7.0f32, 8.0, 9.0];
101        assert_eq!(
102            act_scale_rowmajor_to_kmajor(&src, 3, 3, 1),
103            vec![7.0, 8.0, 9.0]
104        );
105        assert_eq!(
106            act_scale_rowmajor_to_kmajor(&src, 1, 1, 3),
107            vec![7.0, 8.0, 9.0]
108        );
109    }
110
111    #[test]
112    fn the_two_readings_disagree_at_the_measured_shape() {
113        // 2026-09-25: At M=64, K=5120 (L=40) the two offsets agree only where
114        // 63*kg == 39*token, i.e. at 4 of the 2560 scales.
115        let (m_pad, l) = (64usize, 40usize);
116        let same = (0..m_pad)
117            .flat_map(|token| (0..l).map(move |kg| (token, kg)))
118            .filter(|(token, kg)| {
119                vec128_b_index(m_pad, *token, *kg) == rowmajor_index(l, *token, *kg)
120            })
121            .count();
122        assert_eq!(same, 4, "only (0,0), (21,13), (42,26), (63,39) coincide");
123        assert_eq!(m_pad * l - same, 2556);
124    }
125
126    #[test]
127    fn documented_extents_match_the_shapes_we_serve() {
128        // 2026-09-25: K = 5120 and 17408, the `H` and `INTER` of
129        // `dense_ffn_w8a8_prefill_tests.rs`.
130        assert_eq!(k_groups(5120), 40);
131        assert_eq!(k_groups(17408), 136);
132        assert!(blk128x128_stride_ok(5120));
133        assert!(blk128x128_stride_ok(17408));
134        assert!(!blk128x128_stride_ok(128 * 5));
135        assert_eq!(vec128_b_elems(1200, 5120), 1200 * 40);
136    }
137}
metrale_gpu_runtime/cublaslt/scale_layout.rs

metrale_gpu_runtime/cublaslt/
scale_layout.rs
