pub fn w4a16_gemm_n128_m128_bf16_ldb(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
input: DevicePtr,
weight: &QuantizedWeight,
output: DevicePtr,
m: u32,
n: u32,
k: u32,
ldb: u32,
stream: u64,
) -> Result<()>Expand description
W4A16 GEMM — LOSSLESS BF16 prefill variant of w4a16_gemm_n128_m128.
Identical launch config (grid/block/SMEM, M_TILE2=128) and weight layout
(transposed NVFP4) to w4a16_gemm_n128_m128, but launches the
w4a16_gemm_t_m128_bf16 kernel: FP4→BF16 dequant + BF16 m16n8k16 MMA
(FP32 accum), i.e. the base w4a16_gemm math at the fast 128x128 tiling.
Unlike the default t_m128 (which crushes weights+acts to FP8 E4M3 on
NVIDIA), this preserves prefill outputs bit-for-bit vs the base kernel.
Grid: (ceil(N/128), ceil(M/128), 1) Block: (128, 1, 1)
w4a16_gemm_n128_m128_bf16 with an explicit transposed-B row stride, for the
LOSSLESS BF16-MMA path. Needed for the same reason as w4a16_gemm_n128_ldb:
the B loads are 16-byte cp.async and lm_head’s N is the vocab size (248077,
odd), so an unpadded stride misaligns 15 of every 16 k-rows.