Module scale_layout

Module scale_layout 

Source
Expand description

SSOT for the cuBLASLt block-scaling factor layouts — what the library documents, mirrored as index math the CUDA adapter kernel and the CPU tests share.

WHY THIS FILE EXISTS. Measured on 1xH100 (2026-09-11 07:15Z, native_fp8_ffn_w8a8_microtest, tip 5f78270dc, cuBLASLt 13.1): the super::fp8_gemm_act_weight_t_blkscaled arm ran at 1140 TFLOP/s but disagreed with the in-tree fp8_gemm_t_blockscaled on the SAME quantized inputs — rel_rms 1.1e-2 / cosine 0.99994 at M=64, 7.7e-2-8.8e-2 / cosine 0.996 at M=1193, ~33000 BF16 ULP, 84-94% of elements unequal. Two correct W8A8 implementations over identical FP8 bytes with FP32 accumulation agree to ~1e-3, and the error grew with M, so the defect was a scale-tensor layout, not arithmetic. It was: the activation (VEC128) scale tensor was handed over in the quantizer’s [M, K/128] order, and cuBLASLt reads that operand’s scales MN-major.

THE DOCUMENTED RULES (cuBLAS 13.4 manual, “128-element 1D and 128x128 2D Block Scaling For FP8 Data Types” and its “Scaling factors layouts” subsection):

  • Supported mode pairs are VEC128/VEC128, VEC128/BLK128x128 and BLK128x128/VEC128; BLK128x128 on BOTH A and B is listed unsupported. The A=BLK128x128 (weight) + B=VEC128 (activation) pairing Atlas uses is therefore legal as written — the pairing was never the bug.
  • Scaling-factor start addresses must be 16 B aligned, and the matmul’s M and N “must be multiples of 4” — which is what the caller’s ceil16(M) pad satisfies for the token dimension.
  • VEC128_32F: the factors are “M-major for A with shape M x L” and “N-major for B with shape N x L”, where L = ceil(K/128) and major means that dimension is contiguous. So for B the token index is contiguous and the K-group index strides by the (padded) token count — the TRANSPOSE of the [M, K/128] the quantizer writes.
  • BLK128x128_32F: the factors are K-major, “the stride between the consecutive columns must be a multiple of 4”, shape L4 x ceil(M/128) for A (L4 x ceil(N/128) for B) with L4 = L rounded up to a multiple of 4. K-major with ceil(M/128) columns IS the checkpoint’s row-major [N/128, K/128] weight-scale grid whenever L is already a multiple of 4 — see blk128x128_stride_ok. The weight side needed no change.

Atlas maps out[M,N] = act[M,K] @ weight[N,K]ᵀ onto cuBLASLt as D[N,M] = opT(weightᶜ[K,N]) · opN(actᶜ[K,M]), so the library’s M is the weight’s N and the library’s N is the token count. Read the doc quotes above with that substitution: the VEC128 “N-major” operand is the activation, and its contiguous dimension is tokens.

Functions§

act_scale_rowmajor_to_kmajor
CPU reference for the fp8_act_scale_to_kmajor CUDA kernel: read the quantizer’s row-major [m, l] scales, write cuBLASLt’s [l, m_pad], with the m..m_pad pad slots zeroed (their FP8 activation bytes are zeroed too, so the phantom rows contribute a defined zero).
blk128x128_stride_ok
Whether a checkpoint’s row-major [N/128, K/128] weight-scale grid already satisfies the BLK128x128 column-stride rule (“must be a multiple of 4”), i.e. whether L = ceil(K/128) needs no padding to L4.
k_groups
Number of 128-wide K groups a K extent carries (L in the cuBLAS docs).
rowmajor_index
Offset of the same (token, k_group) scale in the layout per_token_group_quant_fp8 writes: row-major [M, K/128], K-group contiguous. The in-tree fp8_gemm_t_blockscaled indexes this one.
vec128_b_elems
FP32 element count of the VEC128 B-scale tensor cuBLASLt reads.
vec128_b_index
Offset of the VEC128 scale for (token, k_group) in the layout cuBLASLt documents for the B operand: shape N x L, N-major, N = m_pad tokens.