ViperEkura
1798474316
- split-K removed entirely: tiled kernel walks K in one pass, no partials/semas workspace, no memset, single launch per call
- skinny GEMM (M<=8) dispatch table replaces the hand-written switch
- shape-driven four-family table replaces plan_gemm: wide-N (n>=4096) default {16,64,64,3,128} with BM=32 at M>16; narrow-N deep-K rings {16,32,256,2,64} while the grid fits one wave, {16,32,128,2,64} past it
- narrow-N is K-serial: widening the grid measurably does nothing (BN 64->32 ties, doubled m_tiles tie, kv at 4 blocks ties q/o at 24); deeper K chunks win until 72KB smem forces one CTA per SM and past one wave the 2-wave quantization loses to BK=128
- launch-check macros in common/launch.cuh; smem opt-in for the 72KB/60KB rings
- rename kernels/bf16_*.cu to gemm.cu/swiglu.cu; module names unchanged
- Python gate: lm_head (N>32768) falls back to cuBLAS, band narrows to M<=32
- drop the stale per-op benchmark narratives; fold the live numbers into cuda_kernels.md
Benchmark: NVIDIA L20 (sm_89, 92 SMs), CUDA 12.8, bf16, L2-thrash weight rotation, per-call medians at M=16: q/o 9.5us, kv 8.6us, gate/up 33.3us, down 33.7us (down -29% vs prior default). End-to-end 1B decode (gen 128, 3 trials, tokens/s vs cuBLAS): B=1 260 vs 252, B=8 1660 vs 1446, B=16 2464 vs 2437, B=32 3620 vs 3690. Prior split-K dispatch measured B=16 2243 / B=32 3393.
28 lines
865 B
Python
28 lines
865 B
Python
"""Stateless wrapper for the directly callable BF16 GEMM primitive."""
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from astrai.extension.loader import get_module
|
|
|
|
|
|
def bf16_gemm(
|
|
x: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
"""Compute ``F.linear(x, weight, bias)`` for up to 64 BF16 rows.
|
|
|
|
``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 64]``, and
|
|
``weight`` must be a contiguous row-major ``[N, K]`` tensor. M in
|
|
``[1, 8]`` uses the register-resident skinny GEMM kernel (any K);
|
|
larger M uses the tiled kernel (K must be a multiple of 8 with
|
|
16-byte-aligned tensors). This primitive is inference-only and
|
|
intentionally performs no fallback or model-level dispatch.
|
|
"""
|
|
return get_module("bf16_gemm").bf16_gemm(x, weight, bias)
|
|
|
|
|
|
__all__ = ["bf16_gemm"]
|