Module scale_layout

Module scale_layout 

Source
Expand description

2026-09-25: Index math for the FP32 scale tensors of the block-scaled FP8 GEMM (super::fp8_gemm_act_weight_t_blkscaled): the activation scales in the layout it passes to cuBLASLt, in the quantizer’s layout, and the weight-scale stride check.

In cuBLASLt’s terms the GEMM is D[N,M] = opT(weightᶜ[K,N]) · opN(actᶜ[K,M]), so the library’s M is the weight’s N and its N is the token count. The activation (B, VEC128_32F) scales are laid out [L, m_pad], token index contiguous, with L = ceil(K/128); the quantizer per_token_group_quant_fp8 writes [M, L], K group contiguous. The weight (A, BLK128x128_32F) scales are the checkpoint’s [ceil(N/128), L] grid, used as they are when L is a multiple of 4. The layouts follow the cuBLAS manual section “128-element 1D and 128x128 2D Block Scaling For FP8 Data Types”.

Measured 2026-09-11 on 1xH100 (native_fp8_ffn_w8a8_microtest, cuBLASLt 13.1): with the quantizer’s [M, L] scales passed as B, the GEMM disagreed with the in-tree fp8_gemm_t_blockscaled on the same inputs, rel_rms 1.1e-2 at M=64 and 7.7e-2 to 8.8e-2 at M=1193.

The metal build compiles this file too (cublaslt_metal_stub.rs).

Owner: gpu-runtime (cuBLASLt wrapper). Invariants: none beyond the types.

Functions§

act_scale_rowmajor_to_kmajor
2026-09-25: CPU reference for the fp8_act_scale_to_kmajor kernel: read row-major [m, l] scales and write [l, m_pad], with the pad tokens m..m_pad set to 0.0, as the kernel does.
blk128x128_stride_ok
2026-09-25: Whether L = ceil(k / 128) is a multiple of 4, the condition under which the checkpoint’s row-major weight-scale grid is passed to the BLK128x128_32F operand unchanged. The cuBLASLt dispatch gates check it (dense_ffn_w8a8_prefill.rs, dispatch_proj_decode.rs).
k_groups
2026-09-25: Number of 128-wide K groups in k: L = ceil(k / 128).
rowmajor_index
2026-09-25: Offset of the same (token, k_group) scale in the layout per_token_group_quant_fp8 writes: row-major [M, L], K group contiguous. fp8_gemm_t_blockscaled reads this layout.
vec128_b_elems
2026-09-25: FP32 element count of the B-operand scale tensor: m_pad * L.
vec128_b_index
2026-09-25: Offset of the scale for (token, k_group) in the B-operand layout [L, m_pad], token index contiguous. The CUDA adapter fp8_act_scale_to_kmajor (kernels/gb10/common/fp8_scale_transpose.cu) writes the same index.