refactor: remove bf16 gemm and swiglu kernels and rebuild csrc benchmarks

- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration
- remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly
- drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections
- add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines
- attention, rotary_emb, and fp8_ops kernels are unchanged

astrai/extension/ops/__init__.py

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

astrai/extension/ops/gemm.py

@@ -1,27 +0,0 @@
-"""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/swiglu.py

@@ -1,22 +0,0 @@
-"""Stateless wrapper for the directly callable fused BF16 SwiGLU primitive."""
-
-import torch
-
-from astrai.extension.loader import get_module
-
-
-def bf16_swiglu(
-    x: torch.Tensor,
-    up_weight: torch.Tensor,
-    gate_weight: torch.Tensor,
-) -> torch.Tensor:
-    """Compute ``linear(x, up) * silu(linear(x, gate))`` for M in [1, 8].
-
-    Inputs must be contiguous BF16 CUDA tensors. Both weights use row-major
-    ``[N, K]`` storage with identical shapes, and K must be divisible by 8.
-    The primitive is inference-only and performs no fallback.
-    """
-    return get_module("bf16_swiglu").bf16_swiglu(x, up_weight, gate_weight)
-
-
-__all__ = ["bf16_swiglu"]