fp8_gemm_t_blockscaled

spark_model::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

W8A8 + FP32 epilogue GEMM with per-token activation scales and per-block weight scales — vLLM-equivalent FP8 numerics.

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] BF16 (existing checkpoint layout)
  • output [M, N] BF16

Grid: (ceil(N/128), ceil(M/64), 1) Block: (128, 1, 1)