pub fn fp8_gemm_n128(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
input: DevicePtr,
b_fp8: DevicePtr,
output: DevicePtr,
m: u32,
n: u32,
k: u32,
stream: u64,
) -> Result<()>Expand description
2026-09-25: FP8-weight prefill GEMM, C = A @ B_fp8^T: A [M, K] BF16,
B_fp8 [N, K] FP8 E4M3 (from predequant_nvfp4_to_fp8), C [M, N]
BF16.
When k is a multiple of 32 and METRALE_FP8_LDMAB is not 0, it casts A
to FP8 into a scratch buffer (bf16_to_fp8) and launches
fp8_fp8_gemm_ldmab (module w4a16_fp8_ldmab, 128x256 tiles, two 128-byte K
stages) instead of kernel.