dense_gemm_tc

spark_model::layers::ops

Function dense_gemm_tc 

Source 
pub fn dense_gemm_tc(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    weight: &DenseWeight,
    output: DevicePtr,
    m: u32,
    n: u32,
    k: u32,
    stream: u64,
) -> Result<()>
Expand description

Dense BF16 GEMM: C = A @ B^T.

A: [M, K] row-major (activations) B: [N, K] row-major (weights, HuggingFace layout) C: [M, N] row-major (output)

Kernel: dense_gemm_bf16(A, B, C, M, N, K) Grid: (ceil(N/16), ceil(M/16), 1) Block: (16, 16, 1) Tensor-core BF16 GEMM: m16n8k16 MMA for 3-5x speedup over scalar. Grid: (ceil(N/64), ceil(M/16), 1), Block: (128, 1, 1)