int8_gemm_faith2_prefill

Function int8_gemm_faith2_prefill 

Source
pub fn int8_gemm_faith2_prefill(
    gpu: &dyn GpuBackend,
    faith2_kernel: KernelHandle,
    requant_a_kernel: KernelHandle,
    a_bf16: DevicePtr,
    w_i8: DevicePtr,
    w_scale: DevicePtr,
    a_i8_scratch: DevicePtr,
    a_scale_scratch: DevicePtr,
    out: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    stream: u64,
) -> Result<()>
Expand description

int8 W4A8 prefill GEMM: requant BF16 activations to int8 (per-32 F32 scale) then C = (A_i8 * W_i8) folded with per-32 A/W block scales via int8_gemm_faith2. The weight is already int8 (see requant_w_nvfp4_int8).

A_bf16: [M, K] BF16 activations. W_i8: [N, K] int8 weights. W_scale: [N, K/32] F32. a_i8_scratch / a_scale_scratch are caller-owned scratch buffers of at least M*K bytes and M*(K/32)*4 bytes respectively. Out: [M, N] BF16.

Two launches on stream (stream-ordered): requant_a → faith2. requant_a grid: (ceil(M*(K/32) / 128), 1, 1) block: (128, 1, 1) faith2 grid: (ceil(N/128), ceil(M/128), 1) block: (256, 1, 1)