fp8_gemm_t_blockscaled

metrale_model_layers::layers::ops

Function fp8_gemm_t_blockscaled 

Source 
pub fn fp8_gemm_t_blockscaled(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    a_fp8: DevicePtr,
    a_scale: DevicePtr,
    b_fp8: DevicePtr,
    b_scale: DevicePtr,
    output: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: W8A8 GEMM with per-token activation scales and per-block weight scales, applied in an FP32 epilogue per 128-element K-group:

C[M, N] = bf16( Σ_g (FP8 MMA over K-group g) × a_scale[M, g] × b_scale[N/128, g] )

Inputs:

  • a_fp8 [M, K] FP8 E4M3
  • a_scale [M, K/128] FP32 (from per_token_group_quant_fp8)
  • b_fp8 [N, K] FP8 E4M3
  • b_scale [N/128, K/128] FP32
  • output [M, N] BF16