perf: rebuild decode gemm dispatch around shape-driven tile configs

- 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.

astrai/extension/ops/__init__.py

@@ -7,7 +7,7 @@ from astrai.extension.ops.attention import (
     attn_paged_prefill,
     attn_prefill,
 )
-from astrai.extension.ops.gemv import bf16_gemv
+from astrai.extension.ops.gemm import bf16_gemm
 from astrai.extension.ops.rotary import rotary_emb
 from astrai.extension.ops.swiglu import bf16_swiglu
 
@@ -17,7 +17,7 @@ __all__ = [
     "attn_paged_decode",
     "attn_paged_prefill",
     "attn_prefill",
-    "bf16_gemv",
+    "bf16_gemm",
     "bf16_swiglu",
     "rotary_emb",
 ]

astrai/extension/ops/gemm.py

@@ -0,0 +1,27 @@
+"""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"]

astrai/extension/ops/gemv.py

@@ -1,26 +0,0 @@
-"""Stateless wrapper for the directly callable BF16 GEMV primitive."""
-
-from typing import Optional
-
-import torch
-
-from astrai.extension.loader import get_module
-
-
-def bf16_gemv(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """Compute ``F.linear(x, weight, bias)`` for up to eight BF16 rows.
-
-    ``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 8]``,
-    and ``weight`` must be a contiguous row-major ``[N, K]`` tensor. The CUDA
-    kernel reuses each weight row across M, accumulates in FP32, and returns
-    BF16. This primitive is inference-only and intentionally performs no
-    fallback or model-level dispatch.
-    """
-    return get_module("bf16_gemv").bf16_gemv(x, weight, bias)
-
-
-__all__ = ["bf16_gemv"]