dense_mm_bf16

Function dense_mm_bf16 

Source
pub fn dense_mm_bf16(
    gpu: &dyn GpuBackend,
    k: &DenseMmKernels,
    a: DevicePtr,
    b: DevicePtr,
    c: DevicePtr,
    m: usize,
    n: usize,
    kk: usize,
    stream: u64,
) -> Result<()>
Expand description

C[M, N] = A[M, K] @ B[N, K]^T, BF16 in and out, output row stride N.

🔴 The M dispatch is the whole point. At M == 1 the tile GEMM’s grid collapses (73 GB/s against a 254 GB/s part). At 2 ..= 8 it is ~94 % padding and measured 3.6× SLOWER than the batched GEMV on this exact workload (multi_seq/qkv.rs::wide_verify_gemm). batchm reads the weight matrix ONCE for all M rows — which is what makes a K-token speculative verify cost one weight sweep instead of K.

🪤 batchm is bit-identical to M separate dense_gemv calls (same K-iteration order and reduction tree per row, --fmad=false), so batching K rows that were previously K serial single-row decodes does not move a single bit. The tile-GEMM arm is NOT bit-identical to either — it reassociates. Widening a site past 8 rows changes numerics.