pub fn fp8_gemm_act_weight_t_blkscaled(
act_fp8: u64,
act_scale: u64,
weight_fp8: u64,
weight_block_scale: u64,
out: u64,
m: u32,
n: u32,
k: u32,
stream: u64,
) -> Result<()>Expand description
Native FP8 (E4M3) out[M,N] = act[M,K] @ weight[N,K]ᵀ → BF16, with the
weight per-128×128-block FP32-scaled (matches Atlas’s Fp8Weight.row_scale
layout exactly) and the activation per-[token,128-of-K] FP32-scaled.
~1.8× the bf16 path (152 vs 85 TFLOPS on GB10).
⚠ SCALE-TENSOR LAYOUTS — the two operands do NOT agree, and getting this
wrong is silent (see super::scale_layout for the doc quotes, the H100
measurement that caught it, and the index math):
weight_block_scale(A, BLK128x128_32F) is K-major,L4 × ⌈N/128⌉— the checkpoint’s row-major[N/128, K/128]grid as-is, valid while⌈K/128⌉is a multiple of 4 (scale_layout::blk128x128_stride_ok).act_scale(B, VEC128_32F) is N-major,M × ⌈K/128⌉with the TOKEN index contiguous — i.e.[K/128, M], the TRANSPOSE of whatper_token_group_quant_fp8writes. Callers adapt it with thefp8_act_scale_to_kmajorkernel; passing the quantizer’s buffer straight through permutes the scales and costs ~8% relative RMS at M≈1200.
m must already include the caller’s pad (the docs require the matmul’s M
and N to be multiples of 4), and act_fp8/act_scale must cover it.
The output is CONTIGUOUS [M, N]; fp8_gemm_act_weight_t_blkscaled_ldc
is the same GEMM with a caller-chosen output row pitch.