fp8_gemm_n128

Function fp8_gemm_n128 

Source
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.