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/__init__.py

@@ -1,11 +1,10 @@
 """CUDA kernel wrappers, operator dispatch, and backend selection.
 
 Public API:
-    - ``attention``, ``linear``, ``swiglu``, ``apply_rotary_emb`` — op
+    - ``attention``, ``apply_rotary_emb`` — op families with safe torch
-      families with safe torch fallbacks (see ``astrai.extension.backend``)
+      fallbacks (see ``astrai.extension.backend``)
     - ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` /
       ``attn_paged_prefill`` — direct attention kernel wrappers
-    - ``bf16_gemm`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
     - ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` /
       ``FlashAttnBackend`` — attention backend strategies
     - ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared
@@ -14,6 +13,9 @@ Public API:
 Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 (blhd). Scale is always ``1/sqrt(head_dim)``. Wrapper functions call their
 compiled CUDA kernels directly; fallback is the backend's responsibility.
+Linear projections and dense-MLP SwiGLU run plain torch (``F.linear`` /
+``Linear`` / ``MLP``); the former bf16_gemm / bf16_swiglu kernels and their
+backends were removed.
 """
 
 from astrai.extension.backend import (
@@ -27,8 +29,6 @@ from astrai.extension.backend import (
     attention,
     attn_backend,
     get_backend,
-    linear,
-    swiglu,
 )
 from astrai.extension.dispatch import (
     Axes,
@@ -52,9 +52,8 @@ from astrai.extension.ops import (
     TensorLayout,
     attn_decode,
     attn_paged_decode,
+    attn_paged_prefill,
     attn_prefill,
-    bf16_gemm,
-    bf16_swiglu,
 )
 
 __all__ = [
@@ -68,13 +67,10 @@ __all__ = [
     "attention",
     "attn_backend",
     "get_backend",
-    "linear",
-    "swiglu",
     "attn_decode",
     "attn_paged_decode",
     "attn_prefill",
-    "bf16_gemm",
+    "attn_paged_prefill",
-    "bf16_swiglu",
     "is_available",
     "KERNEL_NAMES",
     "apply_rotary_emb",

astrai/extension/backend/__init__.py

@@ -11,9 +11,7 @@ from astrai.extension.backend.attention import (
     attn_backend,
     get_backend,
 )
-from astrai.extension.backend.linear import linear
 from astrai.extension.backend.rotary import apply_rotary_emb
-from astrai.extension.backend.swiglu import swiglu
 
 __all__ = [
     "ATTN_BACKEND",
@@ -26,6 +24,4 @@ __all__ = [
     "attention",
     "attn_backend",
     "get_backend",
-    "linear",
-    "swiglu",
 ]

astrai/extension/backend/linear.py

@@ -1,215 +0,0 @@
-"""Inference-only dispatch for AstrAI linear layers.
-
-The CUDA GEMM path is sized by the decode batch M. M in [1, 8] uses the
-register-resident GEMV kernel (any K); M in (8, 64] uses the tiled
-kernel (K % 8 == 0, 16-byte-aligned tensors). Automatic mode
-selects GEMM for M in [1, 64] where the primitive is capable; every
-training, prefill-sized, or unsupported call falls back to PyTorch.
-
-The family stays registered with the shared operator dispatcher, so
-``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
-``explain`` keep working. The per-layer hot path only consults the
-dispatcher when one of those selections is active.
-"""
-
-from typing import Any, Dict, List, Optional
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-
-from astrai.extension.dispatch import (
-    ImplRecord,
-    Spec,
-    axis,
-    env_mode,
-    env_selection,
-    get_override,
-    register_family,
-    resolve,
-    tensor_axes,
-)
-from astrai.extension.loader import is_available
-from astrai.extension.ops.gemm import bf16_gemm
-
-
-def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
-    return F.linear(x, weight, bias)
-
-
-def _inference_bf16_gemm(
-    x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
-) -> Tensor:
-    return bf16_gemm(
-        x.detach(),
-        weight.detach(),
-        bias.detach() if bias is not None else None,
-    )
-
-
-def _gemm_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
-    """Check whether bf16_gemm can safely handle the call."""
-    if (
-        torch.is_grad_enabled()
-        or not x.is_cuda
-        or x.dtype != torch.bfloat16
-        or weight.dtype != torch.bfloat16
-        or weight.ndim != 2
-        or x.ndim not in (1, 2)
-        or x.shape[-1] != weight.shape[1]
-        or x.device != weight.device
-        or not x.is_contiguous()
-        or not weight.is_contiguous()
-        or torch.cuda.get_device_capability(x.get_device()) < (8, 0)
-        or not is_available("bf16_gemm")
-    ):
-        return False
-    m = 1 if x.ndim == 1 else x.shape[0]
-    # M <= 32 is where the kernel wins: L2-rotation measurements on L20
-    # show every production shape at M=24-32 winning or tying, while
-    # M=48-64 loses the long-K down_proj by 8-10% (cuBLAS switches to a
-    # wider tile there). The kernel itself still accepts M <= 64 when
-    # called directly through astrai.extension.ops.gemm.
-    if not (1 <= m <= 32):
-        return False
-    # Vocabulary-sized lm_head weights (N in the tens of thousands+) stream
-    # better through cuBLAS: our skinny path ties it at M<=8 and the tiled
-    # path loses ~4% at M=9-16 (L2-rotation measurements on L20). Gate the
-    # whole shape family out instead of splitting hairs per M band.
-    if weight.shape[0] > 32768:
-        return False
-    # M > 8 tiled path requires K % 8 == 0 and 16-byte alignment.
-    k = x.shape[-1]
-    if m > 8 and (
-        k % 8 != 0 or (x.data_ptr() & 15) != 0 or (weight.data_ptr() & 15) != 0
-    ):
-        return False
-    return bias is None or (
-        bias.device == x.device
-        and bias.dtype == torch.bfloat16
-        and bias.ndim == 1
-        and bias.shape[0] == weight.shape[0]
-        and bias.is_contiguous()
-    )
-
-
-def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
-    weight_shape = tuple(weight.shape)
-    m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
-    supported_m = m is not None and 1 <= m <= 64
-    shape_matches = (
-        weight.ndim == 2
-        and x.ndim in (1, 2)
-        and bool(x.shape)
-        and x.shape[-1] == weight_shape[-1]
-    )
-    same_device = x.device == weight.device and (
-        bias is None or bias.device == x.device
-    )
-    bias_supported = bias is None or (
-        bias.ndim == 1
-        and weight.ndim == 2
-        and bias.shape[0] == weight_shape[0]
-        and bias.dtype == torch.bfloat16
-        and bias.is_contiguous()
-    )
-    capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None
-    return tensor_axes(
-        x,
-        mode=env_mode("ASTRAI_GEMM"),
-        m=m,
-        supported_m=supported_m,
-        auto_m=supported_m,
-        shape_matches=shape_matches,
-        same_device=same_device,
-        weight_dtype=weight.dtype,
-        x_contiguous=x.is_contiguous(),
-        weight_contiguous=weight.is_contiguous(),
-        bias_supported=bias_supported,
-        capability=capability,
-    )
-
-
-_SPEC_CAPABLE = (
-    axis("device_cuda").truthy()
-    & axis("dtype").in_(torch.bfloat16)
-    & axis("weight_dtype").in_(torch.bfloat16)
-    & axis("grad_enabled").eq(False)
-    & axis("supported_m").truthy()
-    & axis("shape_matches").truthy()
-    & axis("same_device").truthy()
-    & axis("x_contiguous").truthy()
-    & axis("weight_contiguous").truthy()
-    & axis("bias_supported").truthy()
-    & Spec.of(
-        lambda ax: ax.get("capability") is not None and ax.get("capability") >= (8, 0),
-        "capability>=sm_80",
-    )
-)
-
-_SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()
-
-
-def _linear_records() -> List[ImplRecord]:
-    mode = env_mode("ASTRAI_GEMM")
-    gemm_priority = 0 if mode == "1" else 100
-    auto_priority = 0 if mode == "auto" else 90
-    torch_priority = 0 if mode == "0" else 50
-    return [
-        ImplRecord(
-            family="linear",
-            name="gemm",
-            obj=_inference_bf16_gemm,
-            spec=_SPEC_CAPABLE,
-            available=lambda: is_available("bf16_gemm"),
-            priority=gemm_priority,
-        ),
-        ImplRecord(
-            family="linear",
-            name="auto_gemm",
-            obj=_inference_bf16_gemm,
-            spec=_SPEC_AUTO,
-            available=lambda: is_available("bf16_gemm"),
-            priority=auto_priority,
-        ),
-        ImplRecord(
-            family="linear",
-            name="torch",
-            obj=_torch_linear,
-            spec=Spec.always(),
-            priority=torch_priority,
-        ),
-    ]
-
-
-def _fallback_record() -> ImplRecord:
-    return ImplRecord(
-        family="linear",
-        name="torch",
-        obj=_torch_linear,
-        spec=Spec.always(),
-        priority=999,
-    )
-
-
-register_family("linear", _axes, _linear_records, _fallback_record)
-
-
-def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
-    """Apply a linear projection with safe inference-only GEMM dispatch.
-
-    ``ASTRAI_GEMM=0`` always uses PyTorch, ``1`` forces GEMM whenever the
-    primitive can safely handle the call (M in [1, 64], K % 8 == 0 and
-    16-byte-aligned for M > 8), and ``auto`` (the default) selects GEMM
-    for all capable decode batches.
-    """
-    if get_override("linear") is not None or env_selection("linear") is not None:
-        return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
-
-    mode = env_mode("ASTRAI_GEMM")
-    if mode != "0" and _gemm_capable(x, weight, bias):
-        return _inference_bf16_gemm(x, weight, bias)
-    return _torch_linear(x, weight, bias)
-
-
-__all__ = ["linear"]

astrai/extension/backend/swiglu.py

@@ -1,67 +0,0 @@
-"""Inference-only fused SwiGLU selection for dense MLP layers."""
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-
-from astrai.extension.backend.linear import linear
-from astrai.extension.dispatch import env_mode
-from astrai.extension.loader import is_available
-from astrai.extension.ops.swiglu import bf16_swiglu
-
-
-def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
-    # Keep the existing linear backend in the fallback chain. This preserves
-    # any independently qualified GEMV batches instead of making the fusion
-    # decision suppress linear-level optimizations.
-    return linear(x, up_weight) * F.silu(linear(x, gate_weight))
-
-
-def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
-    return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())
-
-
-def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
-    return not (
-        torch.is_grad_enabled()
-        or not x.is_cuda
-        or x.dtype != torch.bfloat16
-        or up_weight.dtype != torch.bfloat16
-        or gate_weight.dtype != torch.bfloat16
-        or x.ndim not in (1, 2)
-        or up_weight.ndim != 2
-        or gate_weight.ndim != 2
-        or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
-        or up_weight.shape != gate_weight.shape
-        or x.shape[-1] != up_weight.shape[1]
-        or x.shape[-1] % 8 != 0
-        or x.device != up_weight.device
-        or x.device != gate_weight.device
-        or not x.is_contiguous()
-        or not up_weight.is_contiguous()
-        or not gate_weight.is_contiguous()
-        # The fused kernel reads all streams as uint4; contiguous-but-offset
-        # views are routed to the unfused chain instead of failing.
-        or (x.data_ptr() & 15) != 0
-        or (up_weight.data_ptr() & 15) != 0
-        or (gate_weight.data_ptr() & 15) != 0
-        or not is_available("bf16_swiglu")
-    )
-
-
-def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
-    """Apply the dense-MLP SwiGLU projection with a safe torch fallback.
-
-    ``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain; ``1`` forces
-    the fused primitive for supported inputs; ``auto`` (the default) uses
-    the fused primitive for decode batches with M in ``{1, ..., 8}``. The
-    fused kernel reads x once and covers both projections plus the SiLU
-    gate-multiply in a single launch, measured 9-15% faster than the
-    unfused chain per MLP call on L20 with L2-thrashing weight rotation.
-    """
-    if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
-        return _fused_swiglu(x, up_weight, gate_weight)
-    return _unfused_swiglu(x, up_weight, gate_weight)
-
-
-__all__ = ["swiglu"]

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"]

astrai/model/components/linear.py

@@ -1,9 +1,8 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.extension.backend.linear import linear
-
 
 class Linear(nn.Module):
     def __init__(
@@ -22,4 +21,4 @@ class Linear(nn.Module):
             nn.init.uniform_(self.bias, -bound, bound)
 
     def forward(self, x: Tensor) -> Tensor:
-        return linear(x, self.weight, self.bias)
+        return F.linear(x, self.weight, self.bias)

astrai/model/components/mlp.py

@@ -5,7 +5,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.extension.backend.swiglu import swiglu
 from astrai.factory import BaseFactory
 from astrai.model.components.linear import Linear
 
@@ -39,7 +38,7 @@ class MLP(nn.Module):
         self.down = Linear(dim_ffn, dim, init_std=down_init_std)
 
     def forward(self, x: Tensor) -> FFNOutput:
-        gated = swiglu(x, self.up.weight, self.gate.weight)
+        gated = self.up(x) * F.silu(self.gate(x))
         out = self.down(gated)
         return {"hidden_states": out, "aux_loss": None, "router_stats": None}
 

csrc/CMakeLists.txt

@@ -61,8 +61,6 @@ set(KERNEL_NAMES
     attn_prefill
     attn_paged_decode
     attn_paged_prefill
-    bf16_gemm
-    bf16_swiglu
     rotary_emb
 )
 set(KERNEL_SRCS
@@ -70,8 +68,6 @@ set(KERNEL_SRCS
     attention/prefill.cu
     attention/paged_decode.cu
     attention/paged_prefill.cu
-    gemm.cu
-    swiglu.cu
     rotary_emb.cu
 )
 

csrc/bench/benchmark_attention.py

@@ -0,0 +1,649 @@
+"""Benchmark the four attention kernels against single-launch torch SDPA.
+
+Suites (--suite): decode, prefill, paged_decode, paged_prefill, all. The
+torch side times one SDPA call per step over dense tensors: GQA expansion,
+page-table gathers, padding, and masks are built once outside the timed
+region, and masked calls prefer the cuDNN backend (the default masked path
+is the slow math backend). The kernel's timed work still includes its fused
+paged reads and current-token K/V append. Agreement is checked against the
+same reference.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import statistics
+from dataclasses import asdict, dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Callable, Optional
+
+import click
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from astrai.extension import is_available
+from astrai.extension.ops import (
+    attn_decode,
+    attn_paged_decode,
+    attn_paged_prefill,
+    attn_prefill,
+)
+from astrai.inference.workspace import MAX_SPLITS, Q_TILE_ROWS
+
+
+@dataclass(frozen=True)
+class GqaConfig:
+    """One model family's attention geometry (llama-style GQA)."""
+
+    name: str
+    hq: int
+    hkv: int
+    head_dim: int
+
+
+DEFAULT_CONFIGS = (
+    GqaConfig("llama2_7b", 32, 8, 128),
+    GqaConfig("llama3_70b", 64, 8, 128),
+    GqaConfig("qwen2_7b", 28, 4, 128),
+    GqaConfig("llama3_8b_d64", 32, 8, 64),
+)
+
+# (batch, per-request context length); context includes the token being
+# decoded (kv_len = context, the last slot written in-kernel).
+DECODE_CASES = ((1, 4096), (8, 4096), (32, 2048), (64, 1024))
+# (batch, q_len); prefill from scratch so kv_len == q_len.
+PREFILL_CASES = ((1, 2048), (1, 4096), (4, 1024), (8, 512))
+
+
+def parse_config(value: str) -> GqaConfig:
+    parts = value.split(":")
+    if len(parts) != 4 or not parts[0]:
+        raise click.BadParameter("config must use NAME:HQ:HKV:HEAD_DIM")
+    try:
+        hq, hkv, head_dim = (int(item) for item in parts[1:])
+    except ValueError as exc:
+        raise click.BadParameter("HQ/HKV/HEAD_DIM must be integers") from exc
+    if hq <= 0 or hkv <= 0 or head_dim <= 0 or hq % hkv or head_dim % 32:
+        raise click.BadParameter(
+            "HQ/HKV positive with HQ % HKV == 0; HEAD_DIM % 32 == 0"
+        )
+    return GqaConfig(parts[0], hq, hkv, head_dim)
+
+
+def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        operation()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) / iterations
+
+
+def summarize(values: list[float]) -> dict[str, float]:
+    ordered = sorted(values)
+    return {
+        "median_ms": statistics.median(ordered),
+        "p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
+    }
+
+
+def measure_operations(
+    operations: dict[str, Callable[[], torch.Tensor]],
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, list[float]]:
+    for operation in operations.values():
+        for _ in range(warmup):
+            operation()
+    torch.cuda.synchronize()
+
+    samples: dict[str, list[float]] = {name: [] for name in operations}
+    order = tuple(operations)
+    # A-B-B-A order balances cache, clock, and temperature drift.
+    for _ in range(trials):
+        for name in (*order, *reversed(order)):
+            samples[name].append(time_operation(operations[name], iterations))
+    return samples
+
+
+def repeat_kv_heads(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Expand [*, n_kv_heads, head_dim] to [*, n_kv_heads * n_rep, head_dim]
+    with the backend's grouping (kv head = q head // n_rep)."""
+    if n_rep == 1:
+        return x
+    n_heads, head_dim = x.shape[-2:]
+    return (
+        x.unsqueeze(-2)
+        .expand(*x.shape[:-2], n_heads, n_rep, head_dim)
+        .reshape(*x.shape[:-2], n_heads * n_rep, head_dim)
+    )
+
+
+def sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs) -> torch.Tensor:
+    """SDPA over blhd tensors: [batch, seq, heads, head_dim] -> blhd."""
+    out = F.scaled_dot_product_attention(
+        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), **kwargs
+    )
+    return out.transpose(1, 2)
+
+
+def prefer_cudnn_sdpa(
+    call: Callable[[], torch.Tensor],
+) -> Callable[[], torch.Tensor]:
+    """Return a closure running the masked SDPA ``call`` on cuDNN attention
+    when the backend accepts the bool mask, else torch's default (the
+    default masked path falls back to the much slower math backend)."""
+
+    def with_cudnn() -> torch.Tensor:
+        with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
+            return call()
+
+    try:
+        with_cudnn()
+    except Exception:
+        return call
+    return with_cudnn
+
+
+def ragged_lens(batch: int, span: int) -> list[int]:
+    """Deterministic mixed lengths spanning [span // 2, span]."""
+    if batch == 1:
+        return [span]
+    step = max(span // 2 // (batch - 1), 1)
+    return [span - (batch - 1 - i) * step for i in range(batch)]
+
+
+def cumsum_indptr(lens: list[int]) -> torch.Tensor:
+    return torch.tensor(
+        [0, *torch.tensor(lens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
+    )
+
+
+@dataclass(frozen=True)
+class PagedInputs:
+    """Standalone replicas of the PagePool / InferenceWorkspace tensors."""
+
+    q: torch.Tensor
+    k_cache: torch.Tensor
+    v_cache: torch.Tensor
+    req_to_token: torch.Tensor
+    req_pool_indices: torch.Tensor
+    kv_indptr: torch.Tensor
+
+
+def build_paged_inputs(
+    config: GqaConfig, kv_lens: list[int], q_lens: Optional[list[int]]
+) -> PagedInputs:
+    """Flat pool + page table: request ``i`` owns the contiguous slot range
+    ``[offset_i, offset_i + kv_len_i)``. Q is packed across requests when
+    ``q_lens`` is given (ragged prefill), else [B, Hq, D] (decode)."""
+    batch = len(kv_lens)
+    pool = torch.randn(
+        sum(kv_lens),
+        config.hkv,
+        config.head_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    req_to_token = torch.zeros(batch, max(kv_lens), dtype=torch.int32, device="cuda")
+    offset = 0
+    for i, length in enumerate(kv_lens):
+        req_to_token[i, :length] = torch.arange(
+            offset, offset + length, dtype=torch.int32, device="cuda"
+        )
+        offset += length
+    return PagedInputs(
+        q=torch.randn(
+            sum(q_lens) if q_lens is not None else batch,
+            config.hq,
+            config.head_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        ),
+        k_cache=pool,
+        v_cache=torch.randn_like(pool),
+        req_to_token=req_to_token,
+        req_pool_indices=torch.arange(batch, dtype=torch.int32, device="cuda"),
+        kv_indptr=cumsum_indptr(kv_lens),
+    )
+
+
+def report_result(
+    suite: str,
+    config: GqaConfig,
+    case: dict[str, int],
+    operations: dict[str, Callable[[], torch.Tensor]],
+    samples: dict[str, list[float]],
+    io_bytes: int,
+    reference: torch.Tensor,
+    actual: torch.Tensor,
+) -> dict[str, object]:
+    difference = actual.float() - reference.float()
+    result: dict[str, object] = {
+        "suite": suite,
+        "config": config.name,
+        "agreement": {
+            "max_abs_error": float(difference.abs().max()),
+            "cosine_similarity": float(
+                F.cosine_similarity(
+                    actual.float().flatten(), reference.float().flatten(), dim=0
+                )
+            ),
+        },
+        "estimated_io_bytes": io_bytes,
+        **case,
+    }
+    for name, operation in operations.items():
+        latency = summarize(samples[name])
+        result[name] = {
+            "effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
+            **latency,
+        }
+    speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
+    label = f"B={case.get('batch')}" + (
+        f" ctx={case['context']}" if "context" in case else f" q={case['q_len']}"
+    )
+    print(
+        f"{suite},{config.name},{label},{result['torch']['median_ms']:.4f},"
+        f"{result['cuda']['median_ms']:.4f},{speedup:+.1f}%,"
+        f"{result['agreement']['max_abs_error']:.4f}"
+    )
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Suites
+# ---------------------------------------------------------------------------
+
+
+def benchmark_decode(
+    config: GqaConfig,
+    batch: int,
+    context: int,
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, object]:
+    q = torch.randn(
+        batch, 1, config.hq, config.head_dim, device="cuda", dtype=torch.bfloat16
+    )
+    k = torch.randn(
+        batch,
+        context,
+        config.hkv,
+        config.head_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    v = torch.randn_like(k)
+    # GQA expansion is data preparation, not attention compute — build it
+    # once so the timed torch side is a single SDPA launch.
+    k_expanded = repeat_kv_heads(k, config.hq // config.hkv)
+    v_expanded = repeat_kv_heads(v, config.hq // config.hkv)
+
+    def torch_op() -> torch.Tensor:
+        return sdpa(q, k_expanded, v_expanded)
+
+    def cuda_op() -> torch.Tensor:
+        return attn_decode(q, k, v, is_causal=True)
+
+    operations = {"torch": torch_op, "cuda": cuda_op}
+    samples = measure_operations(
+        operations, warmup=warmup, iterations=iterations, trials=trials
+    )
+    io_bytes = (2 * q.numel() + 2 * k.numel()) * q.element_size()
+    return report_result(
+        "decode",
+        config,
+        {"batch": batch, "context": context},
+        operations,
+        samples,
+        io_bytes,
+        torch_op(),
+        cuda_op(),
+    )
+
+
+def benchmark_prefill(
+    config: GqaConfig,
+    batch: int,
+    q_len: int,
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, object]:
+    q = torch.randn(
+        batch, q_len, config.hq, config.head_dim, device="cuda", dtype=torch.bfloat16
+    )
+    k = torch.randn(
+        batch, q_len, config.hkv, config.head_dim, device="cuda", dtype=torch.bfloat16
+    )
+    v = torch.randn_like(k)
+    k_expanded = repeat_kv_heads(k, config.hq // config.hkv)
+    v_expanded = repeat_kv_heads(v, config.hq // config.hkv)
+
+    def torch_op() -> torch.Tensor:
+        return sdpa(q, k_expanded, v_expanded, is_causal=True)
+
+    def cuda_op() -> torch.Tensor:
+        return attn_prefill(q, k, v, is_causal=True)
+
+    operations = {"torch": torch_op, "cuda": cuda_op}
+    samples = measure_operations(
+        operations, warmup=warmup, iterations=iterations, trials=trials
+    )
+    io_bytes = (2 * q.numel() + 2 * k.numel()) * q.element_size()
+    return report_result(
+        "prefill",
+        config,
+        {"batch": batch, "q_len": q_len},
+        operations,
+        samples,
+        io_bytes,
+        torch_op(),
+        cuda_op(),
+    )
+
+
+def benchmark_paged_decode(
+    config: GqaConfig,
+    batch: int,
+    context: int,
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, object]:
+    n_rep = config.hq // config.hkv
+    kv_lens = [length + 1 for length in ragged_lens(batch, context)]
+    inputs = build_paged_inputs(config, kv_lens, None)
+    new_k = torch.randn(
+        batch, config.hkv, config.head_dim, device="cuda", dtype=torch.bfloat16
+    )
+    new_v = torch.randn_like(new_k)
+    o_part = torch.empty(
+        batch,
+        config.hq,
+        MAX_SPLITS,
+        config.head_dim,
+        dtype=torch.float32,
+        device="cuda",
+    )
+    ml_part = torch.empty(
+        batch, config.hq, MAX_SPLITS, 2, dtype=torch.float32, device="cuda"
+    )
+    out_buf = torch.empty(
+        batch, config.hq, config.head_dim, dtype=torch.bfloat16, device="cuda"
+    )
+
+    def cuda_op() -> torch.Tensor:
+        return attn_paged_decode(
+            inputs.q,
+            inputs.k_cache,
+            inputs.v_cache,
+            inputs.req_to_token,
+            inputs.req_pool_indices,
+            inputs.kv_indptr,
+            new_k=new_k,
+            new_v=new_v,
+            is_causal=True,
+            o_part_buf=o_part,
+            ml_part_buf=ml_part,
+            out_buf=out_buf,
+        )
+
+    # Reference-side data preparation happens once, outside the timed
+    # closure: append the current-token K/V into the pool (the kernel does
+    # this fused inside its launch), gather padded K/V, expand GQA heads.
+    max_len = max(kv_lens)
+    slots = inputs.req_to_token[:, :max_len].long()
+    last_slots = inputs.req_to_token[
+        torch.arange(batch, device="cuda"), torch.tensor(kv_lens) - 1
+    ].long()
+    inputs.k_cache[last_slots] = new_k
+    inputs.v_cache[last_slots] = new_v
+    k_expanded = repeat_kv_heads(inputs.k_cache[slots], n_rep)
+    v_expanded = repeat_kv_heads(inputs.v_cache[slots], n_rep)
+    position = torch.arange(max_len, device="cuda")
+    lengths = torch.tensor(kv_lens, device="cuda", dtype=torch.long)
+    keep_mask = (position[None, :] < lengths[:, None])[:, None, None, :]
+    q_batched = inputs.q.unsqueeze(1)
+
+    sdpa_call = prefer_cudnn_sdpa(
+        lambda: sdpa(q_batched, k_expanded, v_expanded, attn_mask=keep_mask)
+    )
+
+    def torch_op() -> torch.Tensor:
+        return sdpa_call().squeeze(1)  # [B, Hq, D]
+
+    operations = {"torch": torch_op, "cuda": cuda_op}
+    samples = measure_operations(
+        operations, warmup=warmup, iterations=iterations, trials=trials
+    )
+    io_bytes = (
+        2 * inputs.q.numel()  # q read + out write
+        + 2 * sum(kv_lens) * config.hkv * config.head_dim  # k/v reads
+        + 2 * new_k.numel()  # new k/v writes
+    ) * inputs.q.element_size()
+    return report_result(
+        "paged_decode",
+        config,
+        {"batch": batch, "context": context},
+        operations,
+        samples,
+        io_bytes,
+        torch_op(),
+        cuda_op(),
+    )
+
+
+def benchmark_paged_prefill(
+    config: GqaConfig,
+    batch: int,
+    q_len: int,
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, object]:
+    n_rep = config.hq // config.hkv
+    q_lens = ragged_lens(batch, q_len)
+    inputs = build_paged_inputs(config, q_lens, q_lens)
+    qo_indptr = cumsum_indptr(q_lens)
+
+    tile_batches, tile_indices = [], []
+    for request, length in enumerate(q_lens):
+        n_tiles = (length + Q_TILE_ROWS - 1) // Q_TILE_ROWS
+        tile_batches.extend([request] * n_tiles)
+        tile_indices.extend(range(n_tiles))
+    q_tile_to_batch = torch.tensor(tile_batches, dtype=torch.int32, device="cuda")
+    q_tile_to_index = torch.tensor(tile_indices, dtype=torch.int32, device="cuda")
+
+    def cuda_op() -> torch.Tensor:
+        return attn_paged_prefill(
+            inputs.q,
+            inputs.k_cache,
+            inputs.v_cache,
+            inputs.req_to_token,
+            inputs.req_pool_indices,
+            inputs.kv_indptr,
+            qo_indptr,
+            q_tile_to_batch,
+            q_tile_to_index,
+            is_causal=True,
+        )
+
+    # Same once-only preparation: gather padded K/V, expand GQA heads, pad Q,
+    # build the causal + validity mask. The timed reference is one SDPA call
+    # plus the packed-row unpack.
+    max_len = max(q_lens)
+    slots = inputs.req_to_token[:, :max_len].long()
+    k_expanded = repeat_kv_heads(inputs.k_cache[slots], n_rep)
+    v_expanded = repeat_kv_heads(inputs.v_cache[slots], n_rep)
+    position = torch.arange(max_len, device="cuda")
+    lengths = torch.tensor(q_lens, device="cuda", dtype=torch.long)
+    causal = position[None, :, None] >= position[None, None, :]
+    keep = position[None, None, :] < lengths[:, None, None]
+    attn_mask = (causal & keep).unsqueeze(1)
+    q_padded = torch.zeros(
+        batch,
+        max_len,
+        config.hq,
+        config.head_dim,
+        device="cuda",
+        dtype=inputs.q.dtype,
+    )
+    for i, length in enumerate(q_lens):
+        q_padded[i, :length] = inputs.q[int(qo_indptr[i]) : int(qo_indptr[i + 1])]
+
+    sdpa_call = prefer_cudnn_sdpa(
+        lambda: sdpa(q_padded, k_expanded, v_expanded, attn_mask=attn_mask)
+    )
+
+    def torch_op() -> torch.Tensor:
+        out = sdpa_call()
+        return torch.cat([out[i, :length] for i, length in enumerate(q_lens)])
+
+    operations = {"torch": torch_op, "cuda": cuda_op}
+    samples = measure_operations(
+        operations, warmup=warmup, iterations=iterations, trials=trials
+    )
+    io_bytes = (
+        2 * inputs.q.numel() + 2 * sum(q_lens) * config.hkv * config.head_dim
+    ) * inputs.q.element_size()
+    return report_result(
+        "paged_prefill",
+        config,
+        {"batch": batch, "q_len": q_len},
+        operations,
+        samples,
+        io_bytes,
+        torch_op(),
+        cuda_op(),
+    )
+
+
+@click.command(help=__doc__)
+@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
+@click.option(
+    "--suite",
+    "suites",
+    type=click.Choice(("decode", "prefill", "paged_decode", "paged_prefill", "all")),
+    multiple=True,
+    default=("all",),
+    show_default=True,
+)
+@click.option(
+    "--config",
+    "config_values",
+    multiple=True,
+    help="Filter defaults by bare name, or add/override with NAME:HQ:HKV:HEAD_DIM.",
+)
+@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
+@click.option("--iterations", type=click.IntRange(min=1), default=50, show_default=True)
+@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
+@click.option("--seed", type=int, default=0, show_default=True)
+def benchmark_command(
+    output: Path | None,
+    suites: tuple[str, ...],
+    config_values: tuple[str, ...],
+    warmup: int,
+    iterations: int,
+    trials: int,
+    seed: int,
+) -> None:
+    if not torch.cuda.is_available():
+        raise click.ClickException("CUDA is required")
+    kernel_for_suite = {
+        "decode": "attn_decode",
+        "prefill": "attn_prefill",
+        "paged_decode": "attn_paged_decode",
+        "paged_prefill": "attn_paged_prefill",
+    }
+    selected = (
+        tuple(kernel_for_suite) if "all" in suites else tuple(dict.fromkeys(suites))
+    )
+    missing = [
+        kernel_for_suite[suite]
+        for suite in selected
+        if not is_available(kernel_for_suite[suite])
+    ]
+    if missing:
+        raise click.ClickException(f"built kernels required: {', '.join(missing)}")
+
+    # A bare name filters the matching default; a full spec overrides or appends.
+    chosen: dict[str, GqaConfig] = {}
+    for value in config_values:
+        config = (
+            next((c for c in DEFAULT_CONFIGS if c.name == value), None)
+            if ":" not in value
+            else parse_config(value)
+        )
+        if config is None:
+            raise click.BadParameter(f"unknown default config {value!r}")
+        chosen[config.name] = config
+    configs = tuple(chosen.values()) or DEFAULT_CONFIGS
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    runners = {
+        "decode": (benchmark_decode, DECODE_CASES),
+        "prefill": (benchmark_prefill, PREFILL_CASES),
+        "paged_decode": (benchmark_paged_decode, DECODE_CASES),
+        "paged_prefill": (benchmark_paged_prefill, PREFILL_CASES),
+    }
+    print("suite,config,case,torch_ms,cuda_ms,speedup,max_abs")
+    results = []
+    with torch.inference_mode():
+        for suite in selected:
+            runner, cases = runners[suite]
+            for config in configs:
+                for case in cases:
+                    results.append(
+                        runner(
+                            config,
+                            *case,
+                            warmup=warmup,
+                            iterations=iterations,
+                            trials=trials,
+                        )
+                    )
+                torch.cuda.empty_cache()
+
+    if output is not None:
+        props = torch.cuda.get_device_properties(0)
+        payload = {
+            "metadata": {
+                "gpu_name": props.name,
+                "compute_capability": f"{props.major}.{props.minor}",
+                "torch_version": torch.__version__,
+                "cuda_version": torch.version.cuda,
+                "timestamp_utc": datetime.now(timezone.utc).isoformat(),
+            },
+            "settings": {
+                "warmup": warmup,
+                "iterations": iterations,
+                "trials": trials,
+                "seed": seed,
+                "order": "A-B-B-A",
+                "suites": list(selected),
+                "configs": [asdict(config) for config in configs],
+            },
+            "results": results,
+        }
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+
+
+if __name__ == "__main__":
+    benchmark_command()

csrc/bench/benchmark_fp8.py

@@ -0,0 +1,394 @@
+"""Benchmark the FP8 quantize and GEMM kernels against torch baselines.
+
+Suites (--suite): quantize (plain / delayed-scaling ring / dual-orientation
+entries vs the aten float8 cast) and gemm (pre-quantized ``mm_fp8`` in the
+NT orientation the fp8 linear path uses, vs bf16 ``F.linear``). GEMM
+agreement reports both kernel error (vs the fp32 dequantized fp8 product)
+and format error (that product vs the bf16 matmul). FP8 MMA requires
+compute capability 89+.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import statistics
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Callable
+
+import click
+import torch
+import torch.nn.functional as F
+
+from astrai.extension import is_available
+from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
+
+FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
+
+
+@dataclass(frozen=True)
+class MatrixShape:
+    name: str
+    rows: int
+    cols: int
+
+
+QUANTIZE_SHAPES = (
+    MatrixShape("astrai_1b_act", 2048, 1536),
+    MatrixShape("llama2_7b_act", 2048, 4096),
+    MatrixShape("llama2_7b_down_w", 4096, 11008),
+    MatrixShape("llama3_70b_act", 2048, 8192),
+)
+
+# GEMM shapes as (N, K) weight mats; M comes from --m-values.
+GEMM_SHAPES = (
+    MatrixShape("llama2_7b_qkv", 4096, 4096),
+    MatrixShape("llama2_7b_up_gate", 11008, 4096),
+    MatrixShape("llama2_7b_down", 4096, 11008),
+    MatrixShape("llama3_70b_up_gate", 28672, 8192),
+)
+
+
+def parse_positive_ints(value: str) -> tuple[int, ...]:
+    try:
+        values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
+    except ValueError as exc:
+        raise click.BadParameter("expected comma-separated integers") from exc
+    if not values or any(item <= 0 for item in values):
+        raise click.BadParameter("values must be positive integers")
+    return values
+
+
+def parse_shape(value: str) -> MatrixShape:
+    parts = value.split(":")
+    if len(parts) != 3 or not parts[0]:
+        raise click.BadParameter("shape must use NAME:ROWS:COLS")
+    try:
+        rows, cols = (int(item) for item in parts[1:])
+    except ValueError as exc:
+        raise click.BadParameter("ROWS and COLS must be integers") from exc
+    if rows <= 0 or cols <= 0:
+        raise click.BadParameter("ROWS and COLS must be positive")
+    return MatrixShape(parts[0], rows, cols)
+
+
+def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        operation()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) / iterations
+
+
+def summarize(values: list[float]) -> dict[str, float]:
+    ordered = sorted(values)
+    return {
+        "median_ms": statistics.median(ordered),
+        "p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
+    }
+
+
+def measure_operations(
+    operations: dict[str, Callable[[], torch.Tensor]],
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, list[float]]:
+    for operation in operations.values():
+        for _ in range(warmup):
+            operation()
+    torch.cuda.synchronize()
+
+    samples: dict[str, list[float]] = {name: [] for name in operations}
+    order = tuple(operations)
+    # A-B-C-C-B-A order balances cache, clock, and temperature drift.
+    for _ in range(trials):
+        for name in (*order, *reversed(order)):
+            samples[name].append(time_operation(operations[name], iterations))
+    return samples
+
+
+def quant_step(x: torch.Tensor, fmt: str) -> torch.Tensor:
+    """Quantization step (dequant scale) from the current amax; ``quantize``
+    takes its reciprocal as the multiplier."""
+    amax = x.abs().amax().to(torch.float32).clamp_min(1e-12)
+    return amax / FP8_MAX[fmt]
+
+
+def benchmark_quantize(
+    shape: MatrixShape,
+    fmt: str,
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, object]:
+    x = torch.randn(shape.rows, shape.cols, device="cuda", dtype=torch.bfloat16) * 0.1
+    multiplier = quant_step(x, fmt).reciprocal()
+    fp8_dtype = torch.float8_e4m3fn if fmt == "e4m3" else torch.float8_e5m2
+    ring_state = torch.zeros(16 + 4, dtype=torch.float32, device="cuda")
+
+    operations: dict[str, Callable[[], torch.Tensor]] = {
+        "torch_cast": lambda: x.to(fp8_dtype),
+        "plain": lambda: quantize(x, multiplier, fmt)[0],
+        "ring": lambda: quantize(
+            x,
+            multiplier,
+            fmt,
+            ring_state=ring_state,
+            hist_idx=0,
+            fp8_max=FP8_MAX[fmt],
+        )[0],
+        "dual": lambda: quantize_dual(x, multiplier, fmt)[0],
+    }
+    samples = measure_operations(
+        operations, warmup=warmup, iterations=iterations, trials=trials
+    )
+
+    with torch.no_grad():
+        step = quant_step(x, fmt)
+        x8, _ = quantize(x, multiplier, fmt)
+        dequant_error = float((x8.to(torch.float32) * step - x.float()).abs().max())
+        dual_t = quantize_dual(x, multiplier, fmt)[1]
+        dual_matches = bool(torch.equal(dual_t.t().contiguous(), x8.contiguous()))
+
+    io_bytes = 3 * x.numel() + 4  # bf16 read + fp8 write + f32 amax
+    result: dict[str, object] = {
+        "suite": "quantize",
+        "shape": shape.name,
+        "rows": shape.rows,
+        "cols": shape.cols,
+        "fmt": fmt,
+        "estimated_io_bytes": io_bytes,
+        "dequant_max_abs_error": dequant_error,
+        "dual_transpose_matches": dual_matches,
+    }
+    for name, samples_ms in samples.items():
+        latency = summarize(samples_ms)
+        bytes_per_call = io_bytes + x.numel() if name == "dual" else io_bytes
+        result[name] = {
+            "effective_bandwidth_gbps": bytes_per_call
+            / (latency["median_ms"] / 1000)
+            / 1e9,
+            **latency,
+        }
+    speedup = (
+        result["torch_cast"]["median_ms"] / result["plain"]["median_ms"] - 1.0
+    ) * 100.0
+    print(
+        f"quantize,{shape.name},{shape.rows}x{shape.cols},"
+        f"{result['torch_cast']['median_ms']:.4f},{result['plain']['median_ms']:.4f},"
+        f"{result['ring']['median_ms']:.4f},{result['dual']['median_ms']:.4f},"
+        f"{speedup:+.1f}%,{dequant_error:.5f},{dual_matches}"
+    )
+    return result
+
+
+def benchmark_gemm(
+    shape: MatrixShape,
+    m: int,
+    fmt: str,
+    *,
+    warmup: int,
+    iterations: int,
+    trials: int,
+) -> dict[str, object]:
+    x = torch.randn(m, shape.cols, device="cuda", dtype=torch.bfloat16) * 0.1
+    w = (
+        torch.randn(shape.rows, shape.cols, device="cuda", dtype=torch.bfloat16)
+        * shape.cols**-0.5
+    )
+    sx, sw = quant_step(x, fmt), quant_step(w, fmt)
+    with torch.no_grad():
+        x8, _ = quantize(x, sx.reciprocal(), fmt)
+        w8, _ = quantize(w, sw.reciprocal(), fmt)
+    dequant_scale = sx * sw
+
+    def torch_op() -> torch.Tensor:
+        return F.linear(x, w)
+
+    def fp8_op() -> torch.Tensor:
+        return mm_fp8(x8, w8, dequant_scale, trans_b=True)
+
+    operations = {"torch_bf16": torch_op, "fp8": fp8_op}
+    samples = measure_operations(
+        operations, warmup=warmup, iterations=iterations, trials=trials
+    )
+
+    with torch.no_grad():
+        actual = fp8_op().float()
+        dequant_reference = (
+            x8.to(torch.float32) @ w8.to(torch.float32).t() * dequant_scale
+        )
+        bf16_reference = torch_op().float()
+        kernel_difference = actual - dequant_reference
+        format_difference = dequant_reference - bf16_reference
+    io_bytes = m * shape.cols + shape.rows * shape.cols + 2 * m * shape.rows
+
+    result: dict[str, object] = {
+        "suite": "gemm",
+        "shape": shape.name,
+        "m": m,
+        "n": shape.rows,
+        "k": shape.cols,
+        "fmt": fmt,
+        "estimated_io_bytes": io_bytes,
+        "kernel_max_abs": float(kernel_difference.abs().max()),
+        "kernel_rel_l2": float(
+            kernel_difference.norm() / dequant_reference.norm().clamp_min(1e-12)
+        ),
+        "format_rel_l2": float(
+            format_difference.norm() / bf16_reference.norm().clamp_min(1e-12)
+        ),
+    }
+    for name, samples_ms in samples.items():
+        latency = summarize(samples_ms)
+        result[name] = {
+            "effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
+            **latency,
+        }
+    speedup = (
+        result["torch_bf16"]["median_ms"] / result["fp8"]["median_ms"] - 1.0
+    ) * 100.0
+    print(
+        f"gemm,{shape.name},{m}x{shape.rows}x{shape.cols},"
+        f"{result['torch_bf16']['median_ms']:.4f},{result['fp8']['median_ms']:.4f},"
+        f"{speedup:+.1f}%,{result['kernel_max_abs']:.4f},"
+        f"{result['kernel_rel_l2']:.6f},{result['format_rel_l2']:.6f}"
+    )
+    return result
+
+
+@click.command(help=__doc__)
+@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
+@click.option(
+    "--suite",
+    "suites",
+    type=click.Choice(("quantize", "gemm", "all")),
+    multiple=True,
+    default=("all",),
+    show_default=True,
+)
+@click.option("--fmt", type=click.Choice(("e4m3", "e5m2")), default="e4m3")
+@click.option("--m-values", default="512,2048,4096", show_default=True)
+@click.option(
+    "--shape",
+    "shape_values",
+    multiple=True,
+    help="Filter defaults by bare name (either suite), or add/override with "
+    "NAME:ROWS:COLS.",
+)
+@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
+@click.option(
+    "--iterations", type=click.IntRange(min=1), default=100, show_default=True
+)
+@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
+@click.option("--seed", type=int, default=0, show_default=True)
+def benchmark_command(
+    output: Path | None,
+    suites: tuple[str, ...],
+    fmt: str,
+    m_values: str,
+    shape_values: tuple[str, ...],
+    warmup: int,
+    iterations: int,
+    trials: int,
+    seed: int,
+) -> None:
+    if not torch.cuda.is_available():
+        raise click.ClickException("CUDA is required")
+    if not is_available("fp8_ops"):
+        raise click.ClickException(
+            "the built fp8_ops kernel is required (compute capability 89+)"
+        )
+    selected = ("quantize", "gemm") if "all" in suites else tuple(dict.fromkeys(suites))
+    m_values_parsed = parse_positive_ints(m_values)
+
+    # Any --shape selection replaces the defaults for both suites: a bare
+    # name keeps that suite's matching default, a NAME:ROWS:COLS spec
+    # overrides the same-name default or adds a new one.
+    bare_names = {value for value in shape_values if ":" not in value}
+    known = {shape.name for shape in QUANTIZE_SHAPES + GEMM_SHAPES}
+    unknown = sorted(bare_names - known)
+    if unknown:
+        raise click.BadParameter(f"unknown default shape names: {', '.join(unknown)}")
+    specs = [parse_shape(value) for value in shape_values if ":" in value]
+
+    def resolve_shapes(defaults: tuple[MatrixShape, ...]) -> list[MatrixShape]:
+        if not shape_values:
+            return list(defaults)
+        by_name = {shape.name: shape for shape in defaults if shape.name in bare_names}
+        for spec in specs:
+            by_name[spec.name] = spec
+        return list(by_name.values())
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    results = []
+    with torch.inference_mode():
+        if "quantize" in selected:
+            print(
+                "suite,shape,size,cast_ms,plain_ms,ring_ms,dual_ms,vs_cast,"
+                "dequant_max,dual_ok"
+            )
+            for shape in resolve_shapes(QUANTIZE_SHAPES):
+                results.append(
+                    benchmark_quantize(
+                        shape, fmt, warmup=warmup, iterations=iterations, trials=trials
+                    )
+                )
+                torch.cuda.empty_cache()
+        if "gemm" in selected:
+            print(
+                "suite,shape,mxn_xk,bf16_ms,fp8_ms,speedup,kernel_max,"
+                "kernel_rel_l2,format_rel_l2"
+            )
+            for shape in resolve_shapes(GEMM_SHAPES):
+                for m in m_values_parsed:
+                    results.append(
+                        benchmark_gemm(
+                            shape,
+                            m,
+                            fmt,
+                            warmup=warmup,
+                            iterations=iterations,
+                            trials=trials,
+                        )
+                    )
+                torch.cuda.empty_cache()
+
+    if output is not None:
+        props = torch.cuda.get_device_properties(0)
+        payload = {
+            "metadata": {
+                "gpu_name": props.name,
+                "compute_capability": f"{props.major}.{props.minor}",
+                "torch_version": torch.__version__,
+                "cuda_version": torch.version.cuda,
+                "timestamp_utc": datetime.now(timezone.utc).isoformat(),
+                "fmt": fmt,
+            },
+            "settings": {
+                "warmup": warmup,
+                "iterations": iterations,
+                "trials": trials,
+                "seed": seed,
+                "order": "A-B-C-C-B-A",
+                "suites": list(selected),
+                "m_values": list(m_values_parsed),
+            },
+            "results": results,
+        }
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+
+
+if __name__ == "__main__":
+    benchmark_command()

csrc/bench/benchmark_gemm.py

@@ -1,333 +0,0 @@
-"""Benchmark decode-time linear shapes before enabling custom GEMM dispatch.
-
-The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
-establishes the per-architecture cuBLAS baseline that later GEMM primitives and
-dispatch decisions must beat.
-"""
-
-from __future__ import annotations
-
-import json
-import math
-import statistics
-from dataclasses import asdict, dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Callable, Iterable
-
-import click
-import torch
-import torch.nn.functional as F
-
-
-@dataclass(frozen=True)
-class LinearShape:
-    name: str
-    n: int
-    k: int
-
-
-DEFAULT_SHAPES = (
-    LinearShape("q_proj", 1536, 1536),
-    LinearShape("k_proj", 256, 1536),
-    LinearShape("v_proj", 256, 1536),
-    LinearShape("attn_out", 1536, 1536),
-    LinearShape("mlp_up", 6912, 1536),
-    LinearShape("mlp_gate", 6912, 1536),
-    LinearShape("mlp_down", 1536, 6912),
-    LinearShape("lm_head", 100000, 1536),
-)
-DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16}
-
-
-def parse_positive_ints(value: str) -> tuple[int, ...]:
-    """Parse a comma-separated, duplicate-free list of positive integers."""
-    try:
-        values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
-    except ValueError as exc:
-        raise click.BadParameter("expected comma-separated integers") from exc
-    if not values or any(item <= 0 for item in values):
-        raise click.BadParameter("values must be positive integers")
-    return values
-
-
-def parse_shape(value: str) -> LinearShape:
-    """Parse NAME:N:K into a benchmark shape."""
-    parts = value.split(":")
-    if len(parts) != 3 or not parts[0]:
-        raise click.BadParameter("shape must use NAME:N:K")
-    try:
-        n, k = (int(item) for item in parts[1:])
-    except ValueError as exc:
-        raise click.BadParameter("N and K must be integers") from exc
-    if n <= 0 or k <= 0:
-        raise click.BadParameter("N and K must be positive")
-    return LinearShape(parts[0], n, k)
-
-
-def estimate_io_bytes(
-    m: int, n: int, k: int, element_size: int, *, has_bias: bool
-) -> int:
-    """Estimate bytes touched once by Y[M,N] = X[M,K] @ W[N,K].T."""
-    elements = m * k + n * k + m * n
-    if has_bias:
-        elements += n
-    return elements * element_size
-
-
-def percentile(values: Iterable[float], quantile: float) -> float:
-    ordered = sorted(values)
-    if not ordered:
-        raise ValueError("percentile requires at least one sample")
-    rank = (len(ordered) - 1) * quantile
-    lower = math.floor(rank)
-    upper = math.ceil(rank)
-    if lower == upper:
-        return ordered[lower]
-    fraction = rank - lower
-    return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
-
-
-def summarize_latency(samples_ms: list[float]) -> dict[str, float]:
-    return {
-        "median_ms": statistics.median(samples_ms),
-        "p90_ms": percentile(samples_ms, 0.90),
-        "p99_ms": percentile(samples_ms, 0.99),
-        "min_ms": min(samples_ms),
-        "max_ms": max(samples_ms),
-    }
-
-
-def measure_cuda_ms(
-    operation: Callable[[], torch.Tensor], *, warmup: int, iterations: int, trials: int
-) -> list[float]:
-    for _ in range(warmup):
-        operation()
-    torch.cuda.synchronize()
-
-    samples = []
-    for _ in range(trials):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-        start.record()
-        for _ in range(iterations):
-            operation()
-        end.record()
-        end.synchronize()
-        samples.append(start.elapsed_time(end) / iterations)
-    return samples
-
-
-def count_cuda_kernels(
-    operation: Callable[[], torch.Tensor], repeats: int = 5
-) -> float:
-    """Profile a few calls and return the average device events per call."""
-    with torch.profiler.profile(
-        activities=[
-            torch.profiler.ProfilerActivity.CPU,
-            torch.profiler.ProfilerActivity.CUDA,
-        ],
-        acc_events=True,
-    ) as profile:
-        for _ in range(repeats):
-            operation()
-        torch.cuda.synchronize()
-
-    device_type = torch.autograd.DeviceType.CUDA
-    events = [event for event in profile.events() if event.device_type == device_type]
-    return len(events) / repeats
-
-
-def capture_linear(
-    x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None
-) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
-    for _ in range(3):
-        F.linear(x, weight, bias)
-    torch.cuda.synchronize()
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        output = F.linear(x, weight, bias)
-    return graph, output
-
-
-def benchmark_case(
-    shape: LinearShape,
-    m: int,
-    *,
-    dtype: torch.dtype,
-    mode: str,
-    bias_enabled: bool,
-    warmup: int,
-    iterations: int,
-    trials: int,
-) -> dict[str, object]:
-    x = torch.randn((m, shape.k), device="cuda", dtype=dtype)
-    weight = torch.randn((shape.n, shape.k), device="cuda", dtype=dtype)
-    bias = torch.randn(shape.n, device="cuda", dtype=dtype) if bias_enabled else None
-
-    graph = None
-    graph_output = None
-    if mode == "graph":
-        graph, graph_output = capture_linear(x, weight, bias)
-
-        def operation() -> torch.Tensor:
-            graph.replay()
-            return graph_output
-
-    else:
-
-        def operation() -> torch.Tensor:
-            return F.linear(x, weight, bias)
-
-    samples_ms = measure_cuda_ms(
-        operation, warmup=warmup, iterations=iterations, trials=trials
-    )
-    latency = summarize_latency(samples_ms)
-    io_bytes = estimate_io_bytes(
-        m, shape.n, shape.k, x.element_size(), has_bias=bias is not None
-    )
-    median_seconds = latency["median_ms"] / 1000
-
-    result: dict[str, object] = {
-        "name": shape.name,
-        "m": m,
-        "n": shape.n,
-        "k": shape.k,
-        "mode": mode,
-        "bias": bias is not None,
-        "estimated_io_bytes": io_bytes,
-        "effective_bandwidth_gbps": io_bytes / median_seconds / 1e9,
-        "cuda_kernel_launches_per_call": count_cuda_kernels(operation),
-        **latency,
-        "samples_ms": samples_ms,
-    }
-    return result
-
-
-def render_markdown(payload: dict[str, object]) -> str:
-    metadata = payload["metadata"]
-    assert isinstance(metadata, dict)
-    results = payload["results"]
-    assert isinstance(results, list)
-
-    lines = [
-        "# Decode linear baseline",
-        "",
-        f"- GPU: {metadata['gpu_name']}",
-        f"- Compute capability: {metadata['compute_capability']}",
-        f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
-        f"- Dtype: {metadata['dtype']}",
-        "",
-        "| Layer | M | N | K | Mode | Median (ms) | p99 (ms) | GB/s | CUDA kernels/call |",
-        "|---|---:|---:|---:|---|---:|---:|---:|---:|",
-    ]
-    for item in results:
-        assert isinstance(item, dict)
-        lines.append(
-            "| {name} | {m} | {n} | {k} | {mode} | {median_ms:.4f} | "
-            "{p99_ms:.4f} | {effective_bandwidth_gbps:.1f} | "
-            "{cuda_kernel_launches_per_call:.2f} |".format(**item)
-        )
-    lines.append("")
-    return "\n".join(lines)
-
-
-def device_metadata(dtype_name: str) -> dict[str, object]:
-    props = torch.cuda.get_device_properties(0)
-    return {
-        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
-        "gpu_name": props.name,
-        "compute_capability": f"{props.major}.{props.minor}",
-        "total_memory_bytes": props.total_memory,
-        "torch_version": torch.__version__,
-        "cuda_version": torch.version.cuda,
-        "dtype": dtype_name,
-    }
-
-
-@click.command(help=__doc__)
-@click.option("--output", type=click.Path(path_type=Path), required=True)
-@click.option("--markdown-output", type=click.Path(path_type=Path))
-@click.option("--m-values", default="1,2,4,8,16,32", show_default=True)
-@click.option(
-    "--shape",
-    "shape_values",
-    multiple=True,
-    help="Override defaults with repeatable NAME:N:K shapes.",
-)
-@click.option("--dtype", type=click.Choice(tuple(DTYPES)), default="bfloat16")
-@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
-@click.option("--bias/--no-bias", default=False)
-@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
-@click.option(
-    "--iterations", type=click.IntRange(min=1), default=100, show_default=True
-)
-@click.option("--trials", type=click.IntRange(min=1), default=20, show_default=True)
-@click.option("--seed", type=int, default=0, show_default=True)
-def benchmark_command(
-    output: Path,
-    markdown_output: Path | None,
-    m_values: str,
-    shape_values: tuple[str, ...],
-    dtype: str,
-    mode: str,
-    bias: bool,
-    warmup: int,
-    iterations: int,
-    trials: int,
-    seed: int,
-) -> None:
-    if not torch.cuda.is_available():
-        raise click.ClickException("CUDA is required")
-
-    parsed_m = parse_positive_ints(m_values)
-    shapes = tuple(parse_shape(item) for item in shape_values) or DEFAULT_SHAPES
-    modes = ("eager", "graph") if mode == "both" else (mode,)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-    results = []
-    for shape in shapes:
-        for m in parsed_m:
-            for current_mode in modes:
-                click.echo(
-                    f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
-                )
-                results.append(
-                    benchmark_case(
-                        shape,
-                        m,
-                        dtype=DTYPES[dtype],
-                        mode=current_mode,
-                        bias_enabled=bias,
-                        warmup=warmup,
-                        iterations=iterations,
-                        trials=trials,
-                    )
-                )
-
-    payload: dict[str, object] = {
-        "schema_version": 1,
-        "metadata": device_metadata(dtype),
-        "parameters": {
-            "m_values": list(parsed_m),
-            "shapes": [asdict(shape) for shape in shapes],
-            "modes": list(modes),
-            "bias": bias,
-            "warmup": warmup,
-            "iterations": iterations,
-            "trials": trials,
-            "seed": seed,
-        },
-        "results": results,
-    }
-    output.parent.mkdir(parents=True, exist_ok=True)
-    output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
-    if markdown_output is not None:
-        markdown_output.parent.mkdir(parents=True, exist_ok=True)
-        markdown_output.write_text(render_markdown(payload), encoding="utf-8")
-
-
-if __name__ == "__main__":
-    benchmark_command()

csrc/bench/benchmark_gemm_common.py

@@ -1,450 +0,0 @@
-"""Benchmark the BF16 GEMM primitive and guarded linear dispatcher.
-
-The kernel suite covers AstrAI's native projections plus common LLaMA and
-GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain;
-it measures dispatcher overhead and dependent MLP work, but is deliberately
-not presented as a whole-model throughput benchmark.
-"""
-
-import argparse
-import gc
-import json
-import math
-import os
-import statistics
-from collections.abc import Callable
-from dataclasses import dataclass
-from pathlib import Path
-
-import torch
-import torch.nn.functional as F
-
-from astrai.extension import bf16_gemm, is_available, linear
-
-
-@dataclass(frozen=True)
-class Shape:
-    label: str
-    n: int
-    k: int
-
-
-@dataclass(frozen=True)
-class Chain:
-    label: str
-    hidden: int
-    kv: int
-    intermediate: int
-    fused_qkv: bool = False
-    gated_mlp: bool = True
-
-
-@dataclass(frozen=True)
-class Timing:
-    median_ms: float
-    p90_ms: float
-
-
-ASTRAI_SHAPES = (
-    Shape("astrai_qkv", 256, 1536),
-    Shape("astrai_square", 1536, 1536),
-    Shape("astrai_up_gate", 6912, 1536),
-    Shape("astrai_down", 1536, 6912),
-    Shape("astrai_lm_head", 100000, 1536),
-)
-
-TRADITIONAL_SHAPES = (
-    Shape("llama2_7b_qo", 4096, 4096),
-    Shape("llama2_7b_up_gate", 11008, 4096),
-    Shape("llama2_7b_down", 4096, 11008),
-    Shape("llama3_8b_kv", 1024, 4096),
-    Shape("llama3_8b_up_gate", 14336, 4096),
-    Shape("llama3_8b_down", 4096, 14336),
-    Shape("llama2_13b_qo", 5120, 5120),
-    Shape("llama2_13b_up_gate", 13824, 5120),
-    Shape("llama2_13b_down", 5120, 13824),
-    Shape("gpt_neox_up", 16384, 4096),
-    Shape("gpt_neox_down", 4096, 16384),
-    Shape("qwen2_7b_kv", 512, 3584),
-    Shape("qwen2_7b_qo", 3584, 3584),
-    Shape("qwen2_7b_up_gate", 18944, 3584),
-    Shape("qwen2_7b_down", 3584, 18944),
-    Shape("llama3_70b_kv", 1024, 8192),
-    Shape("llama3_70b_qo", 8192, 8192),
-    Shape("llama3_70b_up_gate", 28672, 8192),
-    Shape("llama3_70b_down", 8192, 28672),
-    Shape("opt_1_3b_qkvo", 2048, 2048),
-    Shape("opt_1_3b_up", 8192, 2048),
-    Shape("opt_1_3b_down", 2048, 8192),
-)
-
-CHAINS = (
-    Chain("llama2_7b", 4096, 4096, 11008),
-    Chain("llama3_8b", 4096, 1024, 14336),
-    Chain("llama2_13b", 5120, 5120, 13824),
-    Chain("gpt_neox_20b", 4096, 4096, 16384, fused_qkv=True),
-    Chain("qwen2_7b", 3584, 512, 18944),
-    Chain("llama3_70b", 8192, 1024, 28672),
-    Chain("opt_1_3b", 2048, 2048, 8192, gated_mlp=False),
-)
-
-
-def _elapsed_ms(fn: Callable[[], torch.Tensor], inner: int) -> float:
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(inner):
-        fn()
-    end.record()
-    end.synchronize()
-    return start.elapsed_time(end) / inner
-
-
-def _timing(values: list[float]) -> Timing:
-    ordered = sorted(values)
-    p90_index = max(0, math.ceil(0.9 * len(ordered)) - 1)
-    return Timing(statistics.median(ordered), ordered[p90_index])
-
-
-def _measure_pair(
-    baseline: Callable[[], torch.Tensor],
-    candidate: Callable[[], torch.Tensor],
-    *,
-    warmup: int,
-    samples: int,
-    inner: int,
-    prepare_baseline: Callable[[], None] = lambda: None,
-    prepare_candidate: Callable[[], None] = lambda: None,
-) -> tuple[Timing, Timing]:
-    cases = (
-        ("baseline", prepare_baseline, baseline),
-        ("candidate", prepare_candidate, candidate),
-    )
-    for iteration in range(warmup):
-        _, prepare, fn = cases[iteration % 2]
-        prepare()
-        fn()
-    torch.cuda.synchronize()
-
-    values: dict[str, list[float]] = {"baseline": [], "candidate": []}
-    for sample in range(samples):
-        order = cases if sample % 2 == 0 else tuple(reversed(cases))
-        for label, prepare, fn in order:
-            prepare()
-            values[label].append(_elapsed_ms(fn, inner))
-    return _timing(values["baseline"]), _timing(values["candidate"])
-
-
-def _print_header() -> None:
-    print(
-        "suite,label,m,n,k,torch_median_ms,torch_p90_ms,"
-        "candidate_median_ms,candidate_p90_ms,speedup_pct,"
-        "max_abs,relative_l2,argmax_equal"
-    )
-
-
-def _print_result(
-    suite: str,
-    label: str,
-    m: int,
-    n: int,
-    k: int,
-    baseline: Timing,
-    candidate: Timing,
-    reference: torch.Tensor,
-    actual: torch.Tensor,
-) -> dict[str, object]:
-    difference = actual.float() - reference.float()
-    max_abs = difference.abs().max().item()
-    relative_l2 = difference.norm().item() / max(reference.float().norm().item(), 1e-12)
-    argmax_equal = torch.equal(actual.argmax(dim=-1), reference.argmax(dim=-1))
-    speedup = (baseline.median_ms / candidate.median_ms - 1.0) * 100.0
-    result: dict[str, object] = {
-        "suite": suite,
-        "label": label,
-        "m": m,
-        "n": n,
-        "k": k,
-        "torch_median_ms": baseline.median_ms,
-        "torch_p90_ms": baseline.p90_ms,
-        "candidate_median_ms": candidate.median_ms,
-        "candidate_p90_ms": candidate.p90_ms,
-        "speedup_pct": speedup,
-        "max_abs": max_abs,
-        "relative_l2": relative_l2,
-        "argmax_equal": argmax_equal,
-    }
-    print(
-        f"{suite},{label},{m},{n},{k},"
-        f"{baseline.median_ms:.6f},{baseline.p90_ms:.6f},"
-        f"{candidate.median_ms:.6f},{candidate.p90_ms:.6f},"
-        f"{speedup:+.2f},{max_abs:.6f},{relative_l2:.8f},"
-        f"{str(argmax_equal).lower()}",
-        flush=True,
-    )
-    return result
-
-
-def _weight(n: int, k: int, device: torch.device, std: float) -> torch.Tensor:
-    weight = torch.empty((n, k), device=device, dtype=torch.bfloat16)
-    weight.normal_(mean=0.0, std=std)
-    return weight.requires_grad_(True)
-
-
-def _kernel_functions(
-    x: torch.Tensor, weight: torch.Tensor
-) -> tuple[Callable[[], torch.Tensor], Callable[[], torch.Tensor]]:
-    def baseline() -> torch.Tensor:
-        return F.linear(x, weight)
-
-    def candidate() -> torch.Tensor:
-        return bf16_gemm(x, weight.detach())
-
-    return baseline, candidate
-
-
-def benchmark_kernels(
-    args: argparse.Namespace, device: torch.device
-) -> list[dict[str, object]]:
-    if args.family == "astrai":
-        shapes = ASTRAI_SHAPES
-    elif args.family == "traditional":
-        shapes = TRADITIONAL_SHAPES
-    else:
-        shapes = ASTRAI_SHAPES + TRADITIONAL_SHAPES
-    if args.shape_label:
-        requested = set(args.shape_label)
-        shapes = tuple(shape for shape in shapes if shape.label in requested)
-        missing = requested - {shape.label for shape in shapes}
-        if missing:
-            raise ValueError(f"unknown shape labels: {', '.join(sorted(missing))}")
-
-    results: list[dict[str, object]] = []
-    for shape in shapes:
-        weight = _weight(shape.n, shape.k, device, args.weight_std)
-        for m in args.m:
-            x = torch.randn((m, shape.k), device=device, dtype=torch.bfloat16)
-            baseline_fn, candidate_fn = _kernel_functions(x, weight)
-            with torch.inference_mode():
-                reference = baseline_fn()
-                actual = candidate_fn()
-                baseline, candidate = _measure_pair(
-                    baseline_fn,
-                    candidate_fn,
-                    warmup=args.warmup,
-                    samples=args.samples,
-                    inner=args.inner,
-                )
-            results.append(
-                _print_result(
-                    "kernel",
-                    shape.label,
-                    m,
-                    shape.n,
-                    shape.k,
-                    baseline,
-                    candidate,
-                    reference,
-                    actual,
-                )
-            )
-            del baseline_fn, candidate_fn, x, reference, actual
-        del weight
-        gc.collect()
-        torch.cuda.empty_cache()
-    return results
-
-
-def _set_mode(mode: str) -> None:
-    os.environ["ASTRAI_GEMM"] = mode
-
-
-def _chain_weights(
-    spec: Chain, device: torch.device, std: float
-) -> dict[str, torch.Tensor]:
-    weights = {
-        "o": _weight(spec.hidden, spec.hidden, device, std),
-        "up": _weight(spec.intermediate, spec.hidden, device, std),
-        "down": _weight(spec.hidden, spec.intermediate, device, std),
-    }
-    if spec.fused_qkv:
-        weights["qkv"] = _weight(3 * spec.hidden, spec.hidden, device, std)
-    else:
-        weights.update(
-            {
-                "q": _weight(spec.hidden, spec.hidden, device, std),
-                "k": _weight(spec.kv, spec.hidden, device, std),
-                "v": _weight(spec.kv, spec.hidden, device, std),
-            }
-        )
-        if spec.gated_mlp:
-            weights["gate"] = _weight(spec.intermediate, spec.hidden, device, std)
-    return weights
-
-
-def _chain_fn(
-    x: torch.Tensor, weights: dict[str, torch.Tensor], spec: Chain
-) -> Callable[[], torch.Tensor]:
-    def run() -> torch.Tensor:
-        output_projection = linear(x, weights["o"])
-        up = linear(x, weights["up"])
-        if spec.fused_qkv:
-            attention_projection = linear(x, weights["qkv"])[..., : x.shape[-1]]
-            hidden = F.gelu(up)
-        else:
-            attention_projection = linear(x, weights["q"])
-            linear(x, weights["k"])
-            linear(x, weights["v"])
-            if spec.gated_mlp:
-                gate = linear(x, weights["gate"])
-                hidden = F.silu(gate) * up
-            else:
-                hidden = F.gelu(up)
-        down = linear(hidden, weights["down"])
-        return attention_projection + output_projection + down
-
-    return run
-
-
-def benchmark_chains(
-    args: argparse.Namespace, device: torch.device
-) -> list[dict[str, object]]:
-    results: list[dict[str, object]] = []
-    chains = CHAINS
-    if args.chain_label:
-        requested = set(args.chain_label)
-        chains = tuple(chain for chain in chains if chain.label in requested)
-        missing = requested - {chain.label for chain in chains}
-        if missing:
-            raise ValueError(f"unknown chain labels: {', '.join(sorted(missing))}")
-    for spec in chains:
-        weights = _chain_weights(spec, device, args.weight_std)
-        for m in args.m:
-            x = torch.randn((m, spec.hidden), device=device, dtype=torch.bfloat16)
-            run = _chain_fn(x, weights, spec)
-            with torch.inference_mode():
-                _set_mode("0")
-                reference = run()
-                _set_mode(args.candidate_mode)
-                actual = run()
-                baseline, candidate = _measure_pair(
-                    run,
-                    run,
-                    warmup=args.warmup,
-                    samples=args.samples,
-                    inner=args.chain_inner,
-                    prepare_baseline=lambda: _set_mode("0"),
-                    prepare_candidate=lambda: _set_mode(args.candidate_mode),
-                )
-            results.append(
-                _print_result(
-                    "synthetic_chain",
-                    spec.label,
-                    m,
-                    spec.hidden,
-                    spec.intermediate,
-                    baseline,
-                    candidate,
-                    reference,
-                    actual,
-                )
-            )
-            del x, reference, actual
-        del weights
-        gc.collect()
-        torch.cuda.empty_cache()
-    return results
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--suite", choices=("kernel", "chain", "all"), default="all")
-    parser.add_argument(
-        "--family", choices=("astrai", "traditional", "all"), default="all"
-    )
-    parser.add_argument(
-        "--m", type=int, nargs="+", choices=(1, 2, 4, 8), default=(1, 2, 4, 8)
-    )
-    parser.add_argument(
-        "--shape-label",
-        action="append",
-        help="limit the kernel suite to one or more named shape labels",
-    )
-    parser.add_argument(
-        "--chain-label",
-        action="append",
-        help="limit the chain suite to one or more named model families",
-    )
-    parser.add_argument("--device", type=int, default=0)
-    parser.add_argument("--warmup", type=int, default=20)
-    parser.add_argument("--samples", type=int, default=9)
-    parser.add_argument("--inner", type=int, default=100)
-    parser.add_argument("--chain-inner", type=int, default=20)
-    parser.add_argument(
-        "--candidate-mode",
-        choices=("auto", "1"),
-        default="auto",
-        help="dispatcher mode for the candidate side of the chain suite",
-    )
-    parser.add_argument("--weight-std", type=float, default=0.02)
-    parser.add_argument("--seed", type=int, default=20260902)
-    parser.add_argument(
-        "--output",
-        type=Path,
-        help="optional JSON output; stdout always retains the compact CSV table",
-    )
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    if not torch.cuda.is_available() or not is_available("bf16_gemm"):
-        raise RuntimeError("benchmark requires CUDA and the built bf16_gemm extension")
-    if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1:
-        raise ValueError("warmup must be non-negative and sample/inner counts positive")
-
-    torch.cuda.set_device(args.device)
-    device = torch.device("cuda", args.device)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    properties = torch.cuda.get_device_properties(device)
-    print(
-        f"# device={properties.name}, capability={properties.major}.{properties.minor}, "
-        f"seed={args.seed}, weight_std={args.weight_std}"
-    )
-    _print_header()
-    results: list[dict[str, object]] = []
-    if args.suite in ("kernel", "all"):
-        results.extend(benchmark_kernels(args, device))
-    if args.suite in ("chain", "all"):
-        results.extend(benchmark_chains(args, device))
-    if args.output is not None:
-        payload = {
-            "environment": {
-                "device": properties.name,
-                "capability": f"{properties.major}.{properties.minor}",
-                "torch": torch.__version__,
-                "cuda": torch.version.cuda,
-            },
-            "parameters": {
-                "suite": args.suite,
-                "family": args.family,
-                "m": args.m,
-                "shape_labels": args.shape_label,
-                "chain_labels": args.chain_label,
-                "candidate_mode": args.candidate_mode,
-                "seed": args.seed,
-                "weight_std": args.weight_std,
-                "warmup": args.warmup,
-                "samples": args.samples,
-                "inner": args.inner,
-                "chain_inner": args.chain_inner,
-            },
-            "results": results,
-        }
-        args.output.parent.mkdir(parents=True, exist_ok=True)
-        args.output.write_text(json.dumps(payload, indent=2) + "\n")
-
-
-if __name__ == "__main__":
-    main()

csrc/bench/benchmark_rotary.py

@@ -0,0 +1,252 @@
+"""Benchmark the fused rotary-embedding kernel against the torch fallback.
+
+The baseline is the complex-multiply fallback from
+``astrai.extension.backend.rotary``. Layouts mirror the production call
+shapes: packed 3D [tokens, n_heads, head_dim] and dense 4D
+[batch, seq_len, n_heads, head_dim]; positions are random integers so every
+row exercises a distinct cos/sin gather.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import statistics
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Callable
+
+import click
+import torch
+import torch.nn.functional as F
+
+from astrai.extension import is_available
+from astrai.extension.ops import rotary_emb
+
+
+@dataclass(frozen=True)
+class RotaryCase:
+    name: str
+    layout: str
+    batch: int
+    seq_len: int
+    heads: int
+    head_dim: int
+
+
+DEFAULT_CASES = (
+    RotaryCase("decode_bs1", "packed", 1, 1, 32, 128),
+    RotaryCase("decode_bs32", "packed", 32, 1, 32, 128),
+    RotaryCase("prefill_4k_llama7b", "packed", 1, 4096, 32, 128),
+    RotaryCase("prefill_4k_llama70b", "packed", 1, 4096, 64, 128),
+    RotaryCase("train_8x2k_llama7b", "dense", 8, 2048, 32, 128),
+    RotaryCase("train_4x1k_d64", "dense", 4, 1024, 32, 64),
+)
+
+
+def parse_case(value: str) -> RotaryCase:
+    parts = value.split(":")
+    if len(parts) != 6 or not parts[0]:
+        raise click.BadParameter("case must use NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM")
+    name, layout, batch, seq_len, heads, head_dim = parts
+    try:
+        fields = (int(batch), int(seq_len), int(heads), int(head_dim))
+    except ValueError as exc:
+        raise click.BadParameter("fields must be integers") from exc
+    if layout not in ("packed", "dense"):
+        raise click.BadParameter("layout must be 'packed' or 'dense'")
+    if any(field <= 0 for field in fields) or head_dim % 2:
+        raise click.BadParameter("fields must be positive; HEAD_DIM even")
+    return RotaryCase(name, layout, fields[0], fields[1], fields[2], fields[3])
+
+
+def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        operation()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) / iterations
+
+
+def summarize(values: list[float]) -> dict[str, float]:
+    ordered = sorted(values)
+    return {
+        "median_ms": statistics.median(ordered),
+        "p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
+    }
+
+
+def torch_apply(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """The complex-multiply fallback (mirrors backend.rotary._torch_apply)."""
+    cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
+    dtype = x.dtype
+    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
+    x_complex = torch.view_as_complex(x_)
+    freqs_cis_complex = torch.complex(cos, sin).unsqueeze(-2)
+    x_rotated = x_complex * freqs_cis_complex
+    return torch.view_as_real(x_rotated).flatten(-2).to(dtype)
+
+
+def build_freqs(head_dim: int, positions: torch.Tensor) -> torch.Tensor:
+    """[cos, sin] pairs for the given positions, laid out [..., head_dim/2, 2]."""
+    theta = 10000.0 ** (
+        -torch.arange(0, head_dim, 2, dtype=torch.float64, device=positions.device)
+        / head_dim
+    )
+    freqs = positions.double().unsqueeze(-1) * theta
+    return torch.stack([freqs.cos(), freqs.sin()], dim=-1).float()
+
+
+def benchmark_case(
+    case: RotaryCase, *, warmup: int, iterations: int, trials: int
+) -> dict[str, object]:
+    if case.layout == "packed":
+        tokens = case.batch * case.seq_len
+        x = torch.randn(
+            tokens, case.heads, case.head_dim, device="cuda", dtype=torch.bfloat16
+        )
+        positions = torch.randint(0, 65536, (tokens,), device="cuda")
+    else:
+        x = torch.randn(
+            case.batch,
+            case.seq_len,
+            case.heads,
+            case.head_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        positions = torch.randint(0, 65536, (case.batch, case.seq_len), device="cuda")
+    freqs_cis = build_freqs(case.head_dim, positions)
+    operations: dict[str, Callable[[], torch.Tensor]] = {
+        "torch": lambda: torch_apply(x, freqs_cis),
+        "cuda": lambda: rotary_emb(x, freqs_cis),
+    }
+    for operation in operations.values():
+        for _ in range(warmup):
+            operation()
+    torch.cuda.synchronize()
+
+    samples: dict[str, list[float]] = {name: [] for name in operations}
+    order = tuple(operations)
+    # A-B-B-A order balances cache, clock, and temperature drift.
+    for _ in range(trials):
+        for name in (*order, *reversed(order)):
+            samples[name].append(time_operation(operations[name], iterations))
+
+    with torch.no_grad():
+        expected = operations["torch"]().float()
+        actual = operations["cuda"]().float()
+    difference = actual - expected
+    io_bytes = (
+        2 * x.numel() * x.element_size() + freqs_cis.numel() * freqs_cis.element_size()
+    )
+
+    result: dict[str, object] = {
+        "case": case.name,
+        "layout": case.layout,
+        "heads": case.heads,
+        "head_dim": case.head_dim,
+        "estimated_io_bytes": io_bytes,
+        "max_abs_error": float(difference.abs().max()),
+        "cosine_similarity": float(
+            F.cosine_similarity(actual.flatten(), expected.flatten(), dim=0)
+        ),
+    }
+    for name, samples_ms in samples.items():
+        latency = summarize(samples_ms)
+        result[name] = {
+            "effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
+            **latency,
+        }
+    speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
+    print(
+        f"{case.name},{case.layout},{result['torch']['median_ms']:.5f},"
+        f"{result['cuda']['median_ms']:.5f},{speedup:+.1f}%,"
+        f"{result['max_abs_error']:.5f}"
+    )
+    return result
+
+
+@click.command(help=__doc__)
+@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
+@click.option(
+    "--case",
+    "case_values",
+    multiple=True,
+    help="Filter defaults by bare name, or add/override with "
+    "NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM.",
+)
+@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
+@click.option(
+    "--iterations", type=click.IntRange(min=1), default=100, show_default=True
+)
+@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
+@click.option("--seed", type=int, default=0, show_default=True)
+def benchmark_command(
+    output: Path | None,
+    case_values: tuple[str, ...],
+    warmup: int,
+    iterations: int,
+    trials: int,
+    seed: int,
+) -> None:
+    if not torch.cuda.is_available():
+        raise click.ClickException("CUDA is required")
+    if not is_available("rotary_emb"):
+        raise click.ClickException("the built rotary_emb kernel is required")
+
+    # A bare name filters the matching default; a full spec overrides or appends.
+    chosen: dict[str, RotaryCase] = {}
+    for value in case_values:
+        case = (
+            next((c for c in DEFAULT_CASES if c.name == value), None)
+            if ":" not in value
+            else parse_case(value)
+        )
+        if case is None:
+            raise click.BadParameter(f"unknown default case {value!r}")
+        chosen[case.name] = case
+    cases = tuple(chosen.values()) or DEFAULT_CASES
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    print("case,layout,torch_ms,cuda_ms,speedup,max_abs")
+    results = []
+    with torch.inference_mode():
+        for case in cases:
+            results.append(
+                benchmark_case(
+                    case, warmup=warmup, iterations=iterations, trials=trials
+                )
+            )
+            torch.cuda.empty_cache()
+
+    if output is not None:
+        props = torch.cuda.get_device_properties(0)
+        payload = {
+            "metadata": {
+                "gpu_name": props.name,
+                "compute_capability": f"{props.major}.{props.minor}",
+                "torch_version": torch.__version__,
+                "cuda_version": torch.version.cuda,
+                "timestamp_utc": datetime.now(timezone.utc).isoformat(),
+            },
+            "settings": {
+                "warmup": warmup,
+                "iterations": iterations,
+                "trials": trials,
+                "seed": seed,
+                "order": "A-B-B-A",
+            },
+            "results": results,
+        }
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+
+
+if __name__ == "__main__":
+    benchmark_command()

csrc/bench/benchmark_swiglu.py

@@ -1,325 +0,0 @@
-"""Benchmark fused BF16 SwiGLU against torch and unfused GEMM chains."""
-
-from __future__ import annotations
-
-import json
-import math
-import statistics
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Callable, Iterable
-
-import click
-import torch
-import torch.nn.functional as F
-
-from astrai.extension import bf16_gemm, bf16_swiglu, is_available
-
-
-@dataclass(frozen=True)
-class SwiGLUShape:
-    name: str
-    n: int
-    k: int
-
-
-DEFAULT_SHAPES = (
-    SwiGLUShape("astrai_1b", 6912, 1536),
-    SwiGLUShape("llama2_7b", 11008, 4096),
-    SwiGLUShape("llama3_8b", 14336, 4096),
-    SwiGLUShape("llama2_13b", 13824, 5120),
-    SwiGLUShape("gpt_neox_20b", 16384, 6144),
-)
-
-
-def parse_positive_ints(value: str) -> tuple[int, ...]:
-    try:
-        values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
-    except ValueError as exc:
-        raise click.BadParameter("expected comma-separated integers") from exc
-    if not values or any(item <= 0 for item in values):
-        raise click.BadParameter("values must be positive integers")
-    return values
-
-
-def parse_shape(value: str) -> SwiGLUShape:
-    parts = value.split(":")
-    if len(parts) != 3 or not parts[0]:
-        raise click.BadParameter("shape must use NAME:N:K")
-    try:
-        n, k = (int(item) for item in parts[1:])
-    except ValueError as exc:
-        raise click.BadParameter("N and K must be integers") from exc
-    if n <= 0 or k <= 0 or k % 8:
-        raise click.BadParameter("N must be positive and K positive/divisible by 8")
-    return SwiGLUShape(parts[0], n, k)
-
-
-def percentile(values: Iterable[float], quantile: float) -> float:
-    ordered = sorted(values)
-    rank = (len(ordered) - 1) * quantile
-    lower = math.floor(rank)
-    upper = math.ceil(rank)
-    if lower == upper:
-        return ordered[lower]
-    fraction = rank - lower
-    return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
-
-
-def summarize(values: list[float]) -> dict[str, float]:
-    return {
-        "median_ms": statistics.median(values),
-        "p90_ms": percentile(values, 0.90),
-        "p99_ms": percentile(values, 0.99),
-        "min_ms": min(values),
-        "max_ms": max(values),
-    }
-
-
-def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(iterations):
-        operation()
-    end.record()
-    end.synchronize()
-    return start.elapsed_time(end) / iterations
-
-
-def count_cuda_kernels(
-    operation: Callable[[], torch.Tensor], repeats: int = 5
-) -> float:
-    with torch.profiler.profile(
-        activities=[
-            torch.profiler.ProfilerActivity.CPU,
-            torch.profiler.ProfilerActivity.CUDA,
-        ],
-        acc_events=True,
-    ) as profile:
-        for _ in range(repeats):
-            operation()
-        torch.cuda.synchronize()
-    device_type = torch.autograd.DeviceType.CUDA
-    events = [event for event in profile.events() if event.device_type == device_type]
-    return len(events) / repeats
-
-
-def capture(operation: Callable[[], torch.Tensor]):
-    for _ in range(3):
-        operation()
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        output = operation()
-
-    def replay() -> torch.Tensor:
-        graph.replay()
-        return output
-
-    return replay
-
-
-def make_operations(x, up_weight, gate_weight, mode: str):
-    operations: dict[str, Callable[[], torch.Tensor]] = {
-        "torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)),
-        "gemm_chain": lambda: (
-            bf16_gemm(x, up_weight) * F.silu(bf16_gemm(x, gate_weight))
-        ),
-        "fused": lambda: bf16_swiglu(x, up_weight, gate_weight),
-    }
-    if mode == "graph":
-        operations = {name: capture(op) for name, op in operations.items()}
-    return operations
-
-
-def benchmark_case(
-    shape: SwiGLUShape,
-    m: int,
-    mode: str,
-    *,
-    warmup: int,
-    iterations: int,
-    trials: int,
-) -> list[dict[str, object]]:
-    x = torch.randn((m, shape.k), device="cuda", dtype=torch.bfloat16) * 0.1
-    scale = shape.k**-0.5
-    up_weight = (
-        torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
-    )
-    gate_weight = (
-        torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
-    )
-    operations = make_operations(x, up_weight, gate_weight, mode)
-    for operation in operations.values():
-        for _ in range(warmup):
-            operation()
-    torch.cuda.synchronize()
-
-    samples = {name: [] for name in operations}
-    forward_order = tuple(operations)
-    # A-B-C-C-B-A order balances cache, clock, and temperature drift.
-    for _ in range(trials):
-        for name in (*forward_order, *reversed(forward_order)):
-            samples[name].append(time_operation(operations[name], iterations))
-
-    with torch.no_grad():
-        expected = operations["torch"]().clone()
-        actual = operations["fused"]().clone()
-    difference = (actual.float() - expected.float()).abs()
-    max_abs_error = float(difference.max())
-    mean_abs_error = float(difference.mean())
-    cosine_similarity = float(
-        F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0)
-    )
-
-    results = []
-    for name, operation in operations.items():
-        result: dict[str, object] = {
-            "shape": shape.name,
-            "m": m,
-            "n": shape.n,
-            "k": shape.k,
-            "mode": mode,
-            "implementation": name,
-            "cuda_kernel_launches_per_call": count_cuda_kernels(operation),
-            **summarize(samples[name]),
-        }
-        if name == "fused":
-            result.update(
-                max_abs_error=max_abs_error,
-                mean_abs_error=mean_abs_error,
-                cosine_similarity=cosine_similarity,
-            )
-        results.append(result)
-    return results
-
-
-def device_metadata() -> dict[str, object]:
-    props = torch.cuda.get_device_properties(0)
-    return {
-        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
-        "gpu_name": props.name,
-        "compute_capability": f"{props.major}.{props.minor}",
-        "total_memory_bytes": props.total_memory,
-        "torch_version": torch.__version__,
-        "cuda_version": torch.version.cuda,
-        "dtype": "bfloat16",
-    }
-
-
-def render_markdown(payload: dict[str, object]) -> str:
-    metadata = payload["metadata"]
-    results = payload["results"]
-    assert isinstance(metadata, dict)
-    assert isinstance(results, list)
-    by_case = {
-        (item["shape"], item["m"], item["mode"], item["implementation"]): item
-        for item in results
-    }
-    cases = sorted({(item["shape"], item["m"], item["mode"]) for item in results})
-    lines = [
-        "# Fused SwiGLU benchmark",
-        "",
-        f"- GPU: {metadata['gpu_name']}",
-        f"- Compute capability: {metadata['compute_capability']}",
-        f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
-        "",
-        "| Shape | M | Mode | torch ms | GEMM chain ms | fused ms | "
-        "vs best unfused | fused kernels | max abs | cosine |",
-        "|---|---:|---|---:|---:|---:|---:|---:|---:|---:|",
-    ]
-    for shape, m, mode in cases:
-        torch_item = by_case[(shape, m, mode, "torch")]
-        gemm_item = by_case[(shape, m, mode, "gemm_chain")]
-        fused_item = by_case[(shape, m, mode, "fused")]
-        best = min(torch_item["median_ms"], gemm_item["median_ms"])
-        improvement = (best / fused_item["median_ms"] - 1) * 100
-        lines.append(
-            f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | "
-            f"{gemm_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
-            f"{improvement:+.2f}% | "
-            f"{fused_item['cuda_kernel_launches_per_call']:.1f} | "
-            f"{fused_item['max_abs_error']:.5f} | "
-            f"{fused_item['cosine_similarity']:.8f} |"
-        )
-    lines.append("")
-    return "\n".join(lines)
-
-
-@click.command(help=__doc__)
-@click.option("--output", type=click.Path(path_type=Path), required=True)
-@click.option("--markdown-output", type=click.Path(path_type=Path))
-@click.option("--m-values", default="1,2,4,8", show_default=True)
-@click.option("--shape", "shape_values", multiple=True, help="Repeat NAME:N:K.")
-@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
-@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
-@click.option(
-    "--iterations", type=click.IntRange(min=1), default=100, show_default=True
-)
-@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
-@click.option("--seed", type=int, default=0, show_default=True)
-def benchmark_command(
-    output: Path,
-    markdown_output: Path | None,
-    m_values: str,
-    shape_values: tuple[str, ...],
-    mode: str,
-    warmup: int,
-    iterations: int,
-    trials: int,
-    seed: int,
-) -> None:
-    if not torch.cuda.is_available():
-        raise click.ClickException("CUDA is required")
-    if not is_available("bf16_gemm") or not is_available("bf16_swiglu"):
-        raise click.ClickException("built bf16_gemm and bf16_swiglu are required")
-    shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES
-    m_values_parsed = parse_positive_ints(m_values)
-    if any(m > 8 for m in m_values_parsed):
-        raise click.BadParameter("fused primitive supports M up to 8")
-    modes = ("eager", "graph") if mode == "both" else (mode,)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-    results = []
-    with torch.inference_mode():
-        for shape in shapes:
-            for m in m_values_parsed:
-                for current_mode in modes:
-                    click.echo(
-                        f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
-                    )
-                    results.extend(
-                        benchmark_case(
-                            shape,
-                            m,
-                            current_mode,
-                            warmup=warmup,
-                            iterations=iterations,
-                            trials=trials,
-                        )
-                    )
-            torch.cuda.empty_cache()
-
-    payload: dict[str, object] = {
-        "metadata": device_metadata(),
-        "settings": {
-            "warmup": warmup,
-            "iterations": iterations,
-            "trials": trials,
-            "seed": seed,
-            "order": "A-B-C-C-B-A",
-        },
-        "results": results,
-    }
-    output.parent.mkdir(parents=True, exist_ok=True)
-    output.write_text(json.dumps(payload, indent=2) + "\n")
-    if markdown_output is not None:
-        markdown_output.parent.mkdir(parents=True, exist_ok=True)
-        markdown_output.write_text(render_markdown(payload))
-
-
-if __name__ == "__main__":
-    benchmark_command()

csrc/kernels/gemm.cu

@@ -1,650 +0,0 @@
-// BF16 decode-time GEMM primitive for linear layers: one entry point whose
-// internal path is sized by the decode batch M.
-//
-//   M in [1, 8]  — one CTA per weight row, the M rows held in registers
-//                  (any K).
-//   M in (8, 32] — BM=16 tensor-core tiles; shape-driven configs (see
-//                  the dispatch table at the entry point) cover every
-//                  production shape (K % 8 == 0 and 16-byte-aligned
-//                  tensors).
-//
-// Both paths take row-major [N, K] weights, accumulate in FP32, fuse an
-// optional bias, and are inference-only.
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAException.h>
-#include <cuda_bf16.h>
-#include <torch/extension.h>
-
-#include <cstdint>
-#include <limits>
-
-// Route kernel-launch failures through the torch error check instead of
-// common/launch.cuh's print+abort default. Must precede the kernel code.
-#define ASTRAI_LAUNCH_FAIL(err, what) C10_CUDA_CHECK(err)
-
-#include "common/cp_async.cuh"
-#include "common/launch.cuh"
-#include "common/mma.cuh"
-
-namespace {
-
-constexpr int kHalfCtaThreads = 128;
-constexpr int kWarpSize = 32;
-
-__device__ __forceinline__ float warp_sum(float value) {
-#pragma unroll
-    for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
-        value += __shfl_down_sync(0xffffffff, value, offset);
-    }
-    return value;
-}
-
-template <int Rows, int Threads>
-__global__ void skinny_gemm_kernel(
-    const __nv_bfloat16* __restrict__ x,
-    const __nv_bfloat16* __restrict__ weight,
-    const __nv_bfloat16* __restrict__ bias,
-    __nv_bfloat16* __restrict__ output,
-    int n,
-    int k
-) {
-    const int output_index = blockIdx.x;
-    const int lane = threadIdx.x & (kWarpSize - 1);
-    const int warp = threadIdx.x / kWarpSize;
-
-    float sums[Rows] = {};
-    __shared__ float warp_sums[Rows][Threads / kWarpSize];
-    // Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so
-    // any K is accepted while keeping 128-bit weight loads, which dominate
-    // bandwidth on decode shapes. x pairs with scalar loads: it is a tiny
-    // L1/L2-resident matrix, consecutive threads still touch contiguous
-    // addresses, and no per-row alignment case analysis is needed.
-    const __nv_bfloat16* __restrict__ wrow =
-        weight + static_cast<int64_t>(output_index) * k;
-    const unsigned whead_raw =
-        ((16u - (reinterpret_cast<uintptr_t>(wrow) & 15u)) & 15u) >> 1;
-    const int whead = static_cast<int>(min(whead_raw, static_cast<unsigned>(k)));
-    const int wvecs = (k - whead) / 8;
-    const int wtail_start = whead + wvecs * 8;
-    const uint4* __restrict__ w4 = reinterpret_cast<const uint4*>(wrow + whead);
-
-    // x chunks pair element-for-element with the aligned weight middle:
-    // the uint4 view is rooted at ``x + whead`` (16-byte aligned by the
-    // branch guard), and each row strides by ``k / 8`` vectors because its
-    // first middle element sits ``whead`` scalars past ``row * k``. When
-    // K % 8 == 0 and the weight row is already aligned (whead == 0, the
-    // production case) this reduces to one pure uint4 loop with an empty
-    // head/tail. Otherwise per-row uint4 loads are not 16-byte addressable,
-    // and scalar x pairing keeps the kernel correct for any K while the
-    // weight stream stays vectorized.
-    if (k % 8 == 0 &&
-        ((reinterpret_cast<uintptr_t>(x) + 2u * static_cast<unsigned>(whead)) & 15u) == 0u) {
-        const auto* x4 = reinterpret_cast<const uint4*>(x + whead);
-        for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
-            const uint4 wv_raw = w4[v];
-            const auto* wv =
-                reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
-#pragma unroll
-            for (int row = 0; row < Rows; ++row) {
-                const uint4 xv_raw =
-                    x4[(static_cast<int64_t>(row) * (k / 8)) + v];
-                const auto* xv =
-                    reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
-#pragma unroll
-                for (int p = 0; p < 4; ++p) {
-                    sums[row] = fmaf(
-                        __bfloat162float(__low2bfloat16(xv[p])),
-                        __bfloat162float(__low2bfloat16(wv[p])),
-                        sums[row]
-                    );
-                    sums[row] = fmaf(
-                        __bfloat162float(__high2bfloat16(xv[p])),
-                        __bfloat162float(__high2bfloat16(wv[p])),
-                        sums[row]
-                    );
-                }
-            }
-        }
-    } else {
-        for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
-            const uint4 wv_raw = w4[v];
-            const __nv_bfloat16* wv_s =
-                reinterpret_cast<const __nv_bfloat16*>(&wv_raw);
-#pragma unroll
-            for (int row = 0; row < Rows; ++row) {
-                const __nv_bfloat16* xv =
-                    x + static_cast<int64_t>(row) * k + whead + 8 * v;
-#pragma unroll
-                for (int s = 0; s < 8; ++s) {
-                    sums[row] = fmaf(
-                        __bfloat162float(xv[s]),
-                        __bfloat162float(wv_s[s]),
-                        sums[row]
-                    );
-                }
-            }
-        }
-    }
-    // Head and tail remainders: plain scalar pairing, at most 14 elements.
-    for (int i = threadIdx.x; i < whead; i += blockDim.x) {
-        const float wv = __bfloat162float(wrow[i]);
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            sums[row] = fmaf(
-                __bfloat162float(x[static_cast<int64_t>(row) * k + i]),
-                wv,
-                sums[row]
-            );
-        }
-    }
-    for (int i = wtail_start + threadIdx.x; i < k; i += blockDim.x) {
-        const float wv = __bfloat162float(wrow[i]);
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            sums[row] = fmaf(
-                __bfloat162float(x[static_cast<int64_t>(row) * k + i]),
-                wv,
-                sums[row]
-            );
-        }
-    }
-
-#pragma unroll
-    for (int row = 0; row < Rows; ++row) {
-        sums[row] = warp_sum(sums[row]);
-    }
-    if (lane == 0) {
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            warp_sums[row][warp] = sums[row];
-        }
-    }
-    __syncthreads();
-
-    if (warp == 0) {
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            float sum =
-                lane < (Threads / kWarpSize) ? warp_sums[row][lane] : 0.0f;
-            sum = warp_sum(sum);
-            if (lane == 0) {
-                if (bias != nullptr) {
-                    sum += __bfloat162float(bias[output_index]);
-                }
-                output[row * n + output_index] = __float2bfloat16_rn(sum);
-            }
-        }
-    }
-}
-
-template <int Rows, int Threads>
-void launch_skinny_gemm(
-    const __nv_bfloat16* x,
-    const __nv_bfloat16* weight,
-    const __nv_bfloat16* bias,
-    __nv_bfloat16* output,
-    int n,
-    int k,
-    cudaStream_t stream
-) {
-    // Split-K was tried and rejected for the small-N shapes (GQA kv
-    // projections): their ~48K uint4 loads already saturate thread-level
-    // parallelism one load deep, so they sit on the launch+HBM latency
-    // floor, and the fence+atomic+last-CTA partial round trip adds ~0.6us
-    // of fixed sync cost (measured -27% at M=1, -113% at M=8 on L20).
-    // The real fix for those shapes is fusing the QKV projections so the
-    // tiny kv rows stop launching as standalone kernels at all.
-    //
-    // Decode is HBM weight-streaming bound. Every shape launches with a
-    // single 128-thread CTA size: measured on L20 (sm_89) with L2-thrashing
-    // weight rotation, flat 128t is within ~1% of a per-shape tuned
-    // 128/256/512 mix at M=1 and M=8 and gives up at most ~2% at M=2-4
-    // (512t down_proj, 256t gate/up/lm_head cells), while dodging the
-    // Rows>=7 register cliff on the 307MB lm_head (+21% DRAM throughput
-    // vs 256t at M=8). Re-measured
-    // after the table refactor: 256t on wide-N gate/up wins only ~2.8% at
-    // M=2-4 and ties at M=1/6, inside run-to-run drift. Simplicity keeps
-    // winning over the last ~2%.
-    skinny_gemm_kernel<Rows, Threads>
-        <<<n, Threads, 0, stream>>>(
-            x, weight, bias, output, n, k
-        );
-    ASTRAI_LAUNCH_CHECK();
-}
-
-// Compile-time dispatch tables: replace a hand-written switch over M with
-// function-pointer tables indexed by M-1. Adding a Rows variant means adding
-// one table entry, not a new case block at the call site.
-using SkinnyGemmFn = void (*)(
-    const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*,
-    __nv_bfloat16*, int, int, cudaStream_t
-);
-
-constexpr SkinnyGemmFn kSkinnyGemm[8] = {
-    &launch_skinny_gemm<1, kHalfCtaThreads>,
-    &launch_skinny_gemm<2, kHalfCtaThreads>,
-    &launch_skinny_gemm<3, kHalfCtaThreads>,
-    &launch_skinny_gemm<4, kHalfCtaThreads>,
-    &launch_skinny_gemm<5, kHalfCtaThreads>,
-    &launch_skinny_gemm<6, kHalfCtaThreads>,
-    &launch_skinny_gemm<7, kHalfCtaThreads>,
-    &launch_skinny_gemm<8, kHalfCtaThreads>,
-};
-
-// ---------------------------------------------------------------------------
-// M > 8 path: small-M (M <= 64) tiled GEMM.
-//
-// F.linear for decode batches in (8, 64]: cuBLAS tiles the small M as a
-// single tile row, which starves the grid (measured on L20, sm_89: 24-54
-// CTAs on 92 SMs, tensor pipe 24-42%, DRAM <= 72%). This kernel keeps the
-// M rows in one CTA tile and fills the SMs along N and K instead.
-//
-// CUTLASS-style configuration: the kernel is parameterized by template
-// parameters — CTA tile (BM x BN x BK), pipeline depth, thread count —
-// with the warp layout derived inside (kMt M fragments, kNt n16 tiles per
-// warp). Four families are instantiated (see the dispatch table at the
-// entry point): the default (BN=64, BK=64, 3 stages) and its BM=32 wide-N
-// variant, plus narrow-N deep-K rings (BK=256/128, BN=32, 64 threads); a
-// new shape is an instantiation, not a rewrite.
-//
-// Operand staging reuses the FP8 GEMM's scheme (see fp8/gemm/*.cuh and
-// docs/developer/cuda_kernels.md): 16B chunks XOR-swizzled with row & 7,
-// one barrier per k-tile, and a kStages+1 ring whose prefetch for tile
-// i+kStages lands in the slot tile i-1 released — no post-compute barrier.
-// ---------------------------------------------------------------------------
-
-using bf16 = __nv_bfloat16;
-
-// Logical (row, byte column) -> byte offset in one flat [rows * BK*2B]
-// staging tile. The 16B chunk index is XORed with row & (chunks - 1) so an
-// ldmatrix fragment load (8 consecutive rows x 16B) hits all 32 banks once.
-template <int BK>
-__device__ __forceinline__ int tile_off(int row, int byte_col) {
-    constexpr int kRowBytes = BK * 2;
-    constexpr int kChunks = kRowBytes / 16;
-    static_assert(
-        (kChunks & (kChunks - 1)) == 0, "swizzle needs a power-of-two chunk count"
-    );
-    return row * kRowBytes +
-           (((byte_col >> 4) ^ (row & (kChunks - 1))) << 4) + (byte_col & 15);
-}
-
-// Predicated staging of one [RowsTile x BK] operand slice from a
-// row-major [total_rows, K] tensor into its swizzled ring slot. Chunks past
-// the row count or past K zero-fill (the wrapper guarantees K % 8 == 0 and
-// 16B-aligned rows, so a misaligned *valid* chunk cannot occur).
-template <int RowsTile, int BK, int kThreads>
-__device__ __forceinline__ void stage_tile(
-    char* slot, const bf16* __restrict__ src, int total_rows, int64_t k,
-    int row0, int tid, int kt
-) {
-    constexpr int kRowBytes = BK * 2;
-    constexpr int kChunks = kRowBytes / 16;
-#pragma unroll
-    for (int c = tid; c < RowsTile * kChunks; c += kThreads) {
-        const int r = c / kChunks;
-        const int c8 = c % kChunks;
-        const int64_t kbase = (int64_t)kt * BK;
-        const bool ok = row0 + r < total_rows && kbase + c8 * 8 + 8 <= k;
-        char* dst = slot + tile_off<BK>(r, c8 * 16);
-        if (ok) {
-            astrai::cp_async_16(
-                reinterpret_cast<bf16*>(dst),
-                src + (int64_t)(row0 + r) * k + kbase + c8 * 8
-            );
-        } else {
-            unsigned* w = reinterpret_cast<unsigned*>(dst);
-#pragma unroll
-            for (int i = 0; i < 4; ++i)
-                w[i] = 0u;
-        }
-    }
-}
-
-template <int BM, int BN, int BK, int kStages, int kThreads>
-__global__ __launch_bounds__(kThreads, 1) void tiled_gemm_kernel(
-    const bf16* __restrict__ x,
-    const bf16* __restrict__ w,
-    const bf16* __restrict__ bias,
-    bf16* __restrict__ out,
-    int m,
-    int n,
-    int k
-) {
-    constexpr int kMt = BM / 16;              // m16 fragments
-    constexpr int kWarps = kThreads / 32;
-    constexpr int kNt = BN / (kWarps * 16);   // n16 tiles per warp
-    constexpr int kRowBytes = BK * 2;
-    constexpr int kRing = kStages + 1;        // ring slots per operand
-    constexpr int kSegs = BK / 16;            // m16k16 segments per tile
-    constexpr int kSegXor = 32;               // bytes: +2 chunks per segment
-    static_assert(BM % 16 == 0, "BM must be a multiple of 16");
-    static_assert(BN % (kWarps * 16) == 0, "warps must tile BN in n16 units");
-
-    extern __shared__ __align__(16) char smem[];
-    char* const a_ring = smem;
-    char* const b_ring = smem + kRing * BM * kRowBytes;
-
-    const int tid = threadIdx.x;
-    const int warp = tid >> 5;
-    const int lane = tid & 31;
-    // Standard 2D grid mapping: grid.x covers M tiles, grid.y covers N
-    // tiles. The grid fills the SMs along N; every block walks the whole
-    // K range in a single pass.
-    const int m0 = blockIdx.x * BM;
-    const int n0 = blockIdx.y * BN;
-    const int total_tiles = (k + BK - 1) / BK;
-
-    auto a_slot = [&](int t) {
-        return a_ring + (t % kRing) * BM * kRowBytes;
-    };
-    auto b_slot = [&](int t) {
-        return b_ring + (t % kRing) * BN * kRowBytes;
-    };
-
-#pragma unroll
-    for (int s = 0; s < kStages; ++s) {
-        if (s < total_tiles) {
-            stage_tile<BM, BK, kThreads>(
-                a_slot(s), x, m, k, m0, tid, s
-            );
-            stage_tile<BN, BK, kThreads>(b_slot(s), w, n, k, n0, tid, s);
-        }
-        astrai::cp_async_commit_group();
-    }
-
-    // Per-lane ldmatrix fragment addresses (relative to each ring slot):
-    //   A x4: lanes 0-7 mat0 (rows 0-7, chunk k-lo), 8-15 mat1 (rows 8-15,
-    //   k-lo), 16-23 mat2 (rows 0-7, k-hi), 24-31 mat3 (rows 8-15, k-hi).
-    //   B x4 pair: lanes 0-7/8-15 the first n8 tile's k-lo/k-hi chunks,
-    //   16-23/24-31 the second n8 tile's. Warp w owns the n16 tiles at
-    //   (w + j * kWarps) * 16 for j in [0, kNt).
-    const int a_row = ((lane >> 3) & 1) * 8 + (lane & 7);
-    const unsigned a_off = tile_off<BK>(a_row, (lane >> 4) * 16);
-    unsigned b_off[kNt];
-#pragma unroll
-    for (int j = 0; j < kNt; ++j) {
-        const int b_row =
-            (warp + j * kWarps) * 16 + (lane & 7) + (lane >> 4) * 8;
-        b_off[j] = tile_off<BK>(b_row, ((lane >> 3) & 1) * 16);
-    }
-    const unsigned a_base0 = __cvta_generic_to_shared(a_ring) + a_off;
-    unsigned b_base0[kNt];
-#pragma unroll
-    for (int j = 0; j < kNt; ++j)
-        b_base0[j] = __cvta_generic_to_shared(b_ring) + b_off[j];
-
-    float acc[kMt][kNt][2][4] = {};
-    for (int i = 0; i < total_tiles; ++i) {
-        astrai::cp_async_wait_group<kStages - 1>();
-        __syncthreads();
-        const unsigned a_base =
-            a_base0 + (unsigned)((i % kRing) * BM * kRowBytes);
-        unsigned b_base[kNt];
-#pragma unroll
-        for (int j = 0; j < kNt; ++j)
-            b_base[j] = b_base0[j] + (unsigned)((i % kRing) * BN * kRowBytes);
-#pragma unroll
-        for (int seg = 0; seg < kSegs; ++seg) {
-            const unsigned a_seg = a_base ^ (unsigned)(seg * kSegXor);
-            unsigned a4[kMt][4], b4[kNt][4];
-#pragma unroll
-            for (int mt = 0; mt < kMt; ++mt)
-                astrai::ldmatrix_x4_lane(
-                    a4[mt], a_seg + (unsigned)(mt * 16 * kRowBytes)
-                );
-#pragma unroll
-            for (int j = 0; j < kNt; ++j)
-                astrai::ldmatrix_x4_lane(
-                    b4[j], (b_base[j] ^ (unsigned)(seg * kSegXor))
-                );
-#pragma unroll
-            for (int mt = 0; mt < kMt; ++mt)
-#pragma unroll
-                for (int j = 0; j < kNt; ++j)
-#pragma unroll
-                    for (int nt = 0; nt < 2; ++nt)
-                        astrai::mma_sync<bf16>(
-                            acc[mt][j][nt], a4[mt], b4[j] + nt * 2,
-                            acc[mt][j][nt]
-                        );
-        }
-        // Prefetch tile i+kStages into the slot tile i-1 released. The
-        // barrier at the top of the next iteration separates every
-        // thread's reads of that slot (iteration i-1) from these writes.
-        const int pf = i + kStages;
-        if (pf < total_tiles) {
-            stage_tile<BM, BK, kThreads>(
-                a_slot(pf), x, m, k, m0, tid, pf
-            );
-            stage_tile<BN, BK, kThreads>(b_slot(pf), w, n, k, n0, tid, pf);
-        }
-        astrai::cp_async_commit_group();
-    }
-
-    const int row0 = lane >> 2;
-    const int col0 = (lane & 3) * 2;
-#pragma unroll
-    for (int mt = 0; mt < kMt; ++mt) {
-        if (m0 + mt * 16 + row0 >= m)
-            continue;
-        const int64_t orow = (int64_t)(m0 + mt * 16 + row0) * n;
-#pragma unroll
-        for (int j = 0; j < kNt; ++j) {
-#pragma unroll
-            for (int nt = 0; nt < 2; ++nt) {
-                const int col =
-                    n0 + (warp + j * kWarps) * 16 + col0 + nt * 8;
-                if (col >= n)
-                    continue;
-                float2 v, v8;
-                v.x = acc[mt][j][nt][0];
-                v.y = acc[mt][j][nt][1];
-                v8.x = acc[mt][j][nt][2];
-                v8.y = acc[mt][j][nt][3];
-                if (bias != nullptr) {
-                    v.x += __bfloat162float(bias[col]);
-                    v.y += __bfloat162float(bias[col + 1]);
-                    v8.x += __bfloat162float(bias[col]);
-                    v8.y += __bfloat162float(bias[col + 1]);
-                }
-                if (col + 1 < n) {
-                    *reinterpret_cast<__nv_bfloat162*>(out + orow + col) =
-                        __floats2bfloat162_rn(v.x, v.y);
-                    if (m0 + mt * 16 + row0 + 8 < m)
-                        *reinterpret_cast<__nv_bfloat162*>(
-                            out + orow + 8 * n + col) =
-                            __floats2bfloat162_rn(v8.x, v8.y);
-                } else {
-                    out[orow + col] = __float2bfloat16_rn(v.x);
-                    if (m0 + mt * 16 + row0 + 8 < m)
-                        out[orow + 8 * n + col] =
-                            __float2bfloat16_rn(v8.x);
-                }
-            }
-        }
-    }
-}
-
-template <int BM, int BN, int BK, int kStages, int kThreads>
-void launch_tiled_gemm(
-    const bf16* x,
-    const bf16* w,
-    const bf16* bias,
-    bf16* out,
-    int m,
-    int n,
-    int k,
-    cudaStream_t stream
-) {
-    constexpr int kRing = kStages + 1;
-    constexpr int smem = kRing * (BM + BN) * BK * 2;
-    // 99KB is the sm_86/89 per-block opt-in ceiling; configs above 48KB
-    // (long-K BK=128) run one CTA per SM and pay a one-time attribute opt-in.
-    static_assert(smem <= 99 * 1024, "family must fit the sm_86/89 opt-in ceiling");
-    if constexpr (smem > 48 * 1024) {
-        static const bool opted_in = [] {
-            ASTRAI_CUDA_CHECK(cudaFuncSetAttribute(
-                tiled_gemm_kernel<BM, BN, BK, kStages, kThreads>,
-                cudaFuncAttributeMaxDynamicSharedMemorySize, smem
-            ));
-            return true;
-        }();
-        (void)opted_in;
-    }
-    dim3 grid((m + BM - 1) / BM, (n + BN - 1) / BN);
-    tiled_gemm_kernel<BM, BN, BK, kStages, kThreads>
-        <<<grid, kThreads, smem, stream>>>(
-            x, w, bias, out, m, n, k
-        );
-    ASTRAI_LAUNCH_CHECK();
-}
-
-using TiledGemmFn = void (*)(
-    const bf16*, const bf16*, const bf16*, bf16*, int, int, int,
-    cudaStream_t
-);
-
-// CTA tile configuration pairing the shape parameters with their matched
-// instantiation. Field order is canonical everywhere it appears — template
-// arguments, this struct, the dispatch table — as BM, BN, BK, then
-// pipeline depth (stages) and CTA size (threads). A new family is one
-// row here plus one select branch.
-struct TileConfig {
-    int bm;
-    int bn;
-    int bk;
-    int stages;
-    int threads;
-    TiledGemmFn launch;
-};
-
-// Shape -> tile config, measured on L20 with L2-thrashing weight rotation.
-// The selector trades grid fill against K-loop serial latency:
-//  - Wide N (n >= 4096): ceil(n/64) tiles already cover the SMs, the GEMM
-//    is HBM-bound and the default config wins; BM widens to 32 at M>16 to
-//    halve the re-staged weight stream.
-//  - Narrow N: few N tiles leave the grid K-serial — widening the grid
-//    does not help (measured: BN 64->32 ties, doubled m_tiles tie, kv at
-//    4 blocks ties q/o at 24); fewer, deeper K chunks do. BK=256 with a
-//    72KB two-stage ring is the winner while the grid fits one wave
-//    (72KB smem means one CTA per SM); past one wave its 2-wave
-//    quantization loses to BK=128's 36.9KB two-CTA ring.
-inline TileConfig select_tile_config(int m, int n, int k) {
-    if (n >= 4096) {
-        if (m > 16) {
-            return {32, 64, 64, 3, 128,
-                    &launch_tiled_gemm<32, 64, 64, 3, 128>};
-        }
-        return {16, 64, 64, 3, 128, &launch_tiled_gemm<16, 64, 64, 3, 128>};
-    }
-    const int n_tiles = (n + 31) / 32;
-    const int m_tiles = (m + 15) / 16;
-    if (n_tiles * m_tiles <= 92) {  // one wave on the 92-SM L20
-        return {16, 32, 256, 2, 64,  &launch_tiled_gemm<16, 32, 256, 2, 64>};
-    }
-    return {16, 32, 128, 2, 64, &launch_tiled_gemm<16, 32, 128, 2, 64>};
-}
-
-}  // namespace
-
-// Single entry point: M in [1, 8] routes to the register-resident skinny
-// GEMM kernel (any K), M in (8, 64] to the tiled kernel (K % 8 == 0 and
-// 16-byte-aligned tensors, checked at the branch).
-torch::Tensor bf16_gemm(
-    torch::Tensor x,
-    torch::Tensor weight,
-    py::object bias_object
-) {
-    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "x and weight must be CUDA tensors");
-    TORCH_CHECK(x.device() == weight.device(), "x and weight must share device");
-    TORCH_CHECK(
-        x.scalar_type() == torch::kBFloat16 &&
-            weight.scalar_type() == torch::kBFloat16,
-        "x and weight must be bf16"
-    );
-    TORCH_CHECK(
-        x.dim() == 1 || x.dim() == 2,
-        "x must have shape [K] or [M, K]"
-    );
-    TORCH_CHECK(weight.dim() == 2, "weight must have shape [N, K]");
-    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(), "x and weight must be contiguous");
-    TORCH_CHECK(
-        !x.requires_grad() && !weight.requires_grad(),
-        "bf16_gemm is inference-only and does not support autograd"
-    );
-
-    const int64_t m = x.dim() == 1 ? 1 : x.size(0);
-    const int64_t k = x.size(-1);
-    const int64_t n = weight.size(0);
-    TORCH_CHECK(weight.size(1) == k, "weight K must match x K");
-    TORCH_CHECK(m >= 1 && m <= 64, "M must be in [1, 64]");
-    TORCH_CHECK(k > 0 && n > 0, "N and K must be positive");
-    TORCH_CHECK(
-        k <= std::numeric_limits<int>::max() &&
-            n <= std::numeric_limits<int>::max(),
-        "N or K exceeds the CUDA launcher limit"
-    );
-
-    torch::Tensor bias;
-    const __nv_bfloat16* bias_ptr = nullptr;
-    if (!bias_object.is_none()) {
-        bias = bias_object.cast<torch::Tensor>();
-        TORCH_CHECK(bias.is_cuda() && bias.device() == x.device(), "bias must share the CUDA device");
-        TORCH_CHECK(bias.scalar_type() == torch::kBFloat16, "bias must be bf16");
-        TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n, "bias must have shape [N]");
-        TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
-        TORCH_CHECK(!bias.requires_grad(), "bf16_gemm bias does not support autograd");
-        bias_ptr = reinterpret_cast<const __nv_bfloat16*>(bias.data_ptr());
-    }
-
-    const at::cuda::OptionalCUDAGuard guard(x.device());
-    const auto* properties = at::cuda::getDeviceProperties(x.device().index());
-    TORCH_CHECK(properties->major >= 8, "bf16_gemm requires compute capability 8.0+");
-    auto stream = at::cuda::getCurrentCUDAStream();
-    auto output = x.dim() == 1 ? torch::empty({n}, x.options())
-                               : torch::empty({m, n}, x.options());
-
-    const auto* x_ptr = reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
-    const auto* weight_ptr =
-        reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr());
-    auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
-    const int m_int = static_cast<int>(m);
-    const int n_int = static_cast<int>(n);
-    const int k_int = static_cast<int>(k);
-
-    if (m <= 8) {
-        kSkinnyGemm[m_int - 1](
-            x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int,
-            stream.stream()
-        );
-    } else {
-        TORCH_CHECK(
-            k % 8 == 0 &&
-                (reinterpret_cast<uintptr_t>(x.data_ptr()) & 15) == 0u &&
-                (reinterpret_cast<uintptr_t>(weight.data_ptr()) & 15) == 0u,
-            "M > 8 requires K to be a multiple of 8 and x/weight 16-byte aligned"
-        );
-        const TileConfig cfg = select_tile_config(m_int, n_int, k_int);
-        cfg.launch(
-            x_ptr, weight_ptr, bias_ptr, output_ptr, m_int, n_int, k_int,
-            stream.stream()
-        );
-    }
-    C10_CUDA_CHECK(cudaGetLastError());
-    return output;
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
-    module.def(
-        "bf16_gemm",
-        &bf16_gemm,
-        py::arg("x"),
-        py::arg("weight"),
-        py::arg("bias") = py::none(),
-        "M in [1, 64] BF16 GEMM with optional fused bias "
-        "(register-resident skinny GEMM path for M <= 8, tensor-core tiles above)"
-    );
-}

csrc/kernels/swiglu.cu

@@ -1,297 +0,0 @@
-// Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers.
-// One CTA per output column; each weight pair is read once and reused across
-// all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant
-// selection beyond the M=8 block-size rule is noise.
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAException.h>
-#include <cuda_bf16.h>
-#include <torch/extension.h>
-
-#include <cstdint>
-#include <limits>
-
-namespace {
-
-constexpr int kWarpSize = 32;
-
-__device__ __forceinline__ float warp_sum(float value) {
-#pragma unroll
-    for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
-        value += __shfl_down_sync(0xffffffff, value, offset);
-    }
-    return value;
-}
-
-__device__ __forceinline__ float round_bf16(float value) {
-    return __bfloat162float(__float2bfloat16_rn(value));
-}
-
-template <int Threads, int Rows>
-__global__ void bf16_swiglu_kernel(
-    const __nv_bfloat16* __restrict__ x,
-    const __nv_bfloat16* __restrict__ up_weight,
-    const __nv_bfloat16* __restrict__ gate_weight,
-    __nv_bfloat16* __restrict__ output,
-    int n,
-    int k
-) {
-    constexpr int kWarps = Threads / kWarpSize;
-    const int output_index = blockIdx.x;
-    const int lane = threadIdx.x & (kWarpSize - 1);
-    const int warp = threadIdx.x / kWarpSize;
-    const int vector_count = k / 8;
-
-    float up_sums[Rows] = {};
-    float gate_sums[Rows] = {};
-    __shared__ float up_warp_sums[Rows][kWarps];
-    __shared__ float gate_warp_sums[Rows][kWarps];
-
-    const auto* x4 = reinterpret_cast<const uint4*>(x);
-    const auto* up4 = reinterpret_cast<const uint4*>(
-        up_weight + static_cast<int64_t>(output_index) * k
-    );
-    const auto* gate4 = reinterpret_cast<const uint4*>(
-        gate_weight + static_cast<int64_t>(output_index) * k
-    );
-
-    // Read each pair of up/gate weight chunks once per CTA, then reuse it for
-    // every active decode row. The fused epilogue removes two [M, N]
-    // intermediates and the standalone SiLU and multiply launches.
-    for (int vector_index = threadIdx.x;
-         vector_index < vector_count;
-         vector_index += blockDim.x) {
-        const uint4 up_raw = up4[vector_index];
-        const uint4 gate_raw = gate4[vector_index];
-        const auto* up_values =
-            reinterpret_cast<const __nv_bfloat162*>(&up_raw);
-        const auto* gate_values =
-            reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
-
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            const uint4 x_raw =
-                x4[static_cast<int64_t>(row) * vector_count + vector_index];
-            const auto* x_values =
-                reinterpret_cast<const __nv_bfloat162*>(&x_raw);
-#pragma unroll
-            for (int pair = 0; pair < 4; ++pair) {
-                const float2 xv = __bfloat1622float2(x_values[pair]);
-                const float2 uv = __bfloat1622float2(up_values[pair]);
-                const float2 gv = __bfloat1622float2(gate_values[pair]);
-                up_sums[row] = fmaf(xv.x, uv.x, up_sums[row]);
-                up_sums[row] = fmaf(xv.y, uv.y, up_sums[row]);
-                gate_sums[row] = fmaf(xv.x, gv.x, gate_sums[row]);
-                gate_sums[row] = fmaf(xv.y, gv.y, gate_sums[row]);
-            }
-        }
-    }
-
-#pragma unroll
-    for (int row = 0; row < Rows; ++row) {
-        up_sums[row] = warp_sum(up_sums[row]);
-        gate_sums[row] = warp_sum(gate_sums[row]);
-    }
-    if (lane == 0) {
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            up_warp_sums[row][warp] = up_sums[row];
-            gate_warp_sums[row][warp] = gate_sums[row];
-        }
-    }
-    __syncthreads();
-
-    if (warp == 0) {
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            float up = lane < kWarps ? up_warp_sums[row][lane] : 0.0f;
-            float gate = lane < kWarps ? gate_warp_sums[row][lane] : 0.0f;
-            up = warp_sum(up);
-            gate = warp_sum(gate);
-            if (lane == 0) {
-                // Match the public composition's BF16 rounding boundaries:
-                // BF16 linear outputs, BF16 SiLU output, then BF16 multiply.
-                up = round_bf16(up);
-                gate = round_bf16(gate);
-                const float silu = round_bf16(gate / (1.0f + expf(-gate)));
-                output[static_cast<int64_t>(row) * n + output_index] =
-                    __float2bfloat16_rn(up * silu);
-            }
-        }
-    }
-}
-
-template <int Threads, int Rows>
-void launch_bf16_swiglu(
-    const __nv_bfloat16* x,
-    const __nv_bfloat16* up_weight,
-    const __nv_bfloat16* gate_weight,
-    __nv_bfloat16* output,
-    int n,
-    int k,
-    cudaStream_t stream
-) {
-    bf16_swiglu_kernel<Threads, Rows><<<n, Threads, 0, stream>>>(
-        x, up_weight, gate_weight, output, n, k
-    );
-}
-
-torch::Tensor bf16_swiglu(
-    torch::Tensor x,
-    torch::Tensor up_weight,
-    torch::Tensor gate_weight
-) {
-    TORCH_CHECK(
-        x.is_cuda() && up_weight.is_cuda() && gate_weight.is_cuda(),
-        "x, up_weight, and gate_weight must be CUDA tensors"
-    );
-    TORCH_CHECK(
-        x.device() == up_weight.device() && x.device() == gate_weight.device(),
-        "x and weights must share a device"
-    );
-    TORCH_CHECK(
-        x.scalar_type() == torch::kBFloat16 &&
-            up_weight.scalar_type() == torch::kBFloat16 &&
-            gate_weight.scalar_type() == torch::kBFloat16,
-        "x and weights must be bf16"
-    );
-    TORCH_CHECK(
-        x.dim() == 1 || x.dim() == 2,
-        "x must have shape [K] or [M, K]"
-    );
-    TORCH_CHECK(
-        up_weight.dim() == 2 && gate_weight.dim() == 2,
-        "weights must have shape [N, K]"
-    );
-    TORCH_CHECK(
-        x.is_contiguous() && up_weight.is_contiguous() &&
-            gate_weight.is_contiguous(),
-        "x and weights must be contiguous"
-    );
-    // The kernel loads all three streams as uint4; contiguous-but-offset
-    // views would fault with an opaque "misaligned address" CUDA error, so
-    // reject them here with an actionable message.
-    TORCH_CHECK(
-        (reinterpret_cast<uintptr_t>(x.data_ptr()) & 15u) == 0u,
-        "bf16_swiglu requires 16-byte aligned x (storage_offset must keep "
-        "data_ptr divisible by 16); clone the tensor or use the torch path"
-    );
-    TORCH_CHECK(
-        (reinterpret_cast<uintptr_t>(up_weight.data_ptr()) & 15u) == 0u,
-        "bf16_swiglu requires 16-byte aligned up_weight (storage_offset "
-        "must keep data_ptr divisible by 16); clone the tensor or use the "
-        "torch path"
-    );
-    TORCH_CHECK(
-        (reinterpret_cast<uintptr_t>(gate_weight.data_ptr()) & 15u) == 0u,
-        "bf16_swiglu requires 16-byte aligned gate_weight (storage_offset "
-        "must keep data_ptr divisible by 16); clone the tensor or use the "
-        "torch path"
-    );
-    TORCH_CHECK(
-        !x.requires_grad() && !up_weight.requires_grad() &&
-            !gate_weight.requires_grad(),
-        "bf16_swiglu is inference-only and does not support autograd"
-    );
-
-    const int64_t m = x.dim() == 1 ? 1 : x.size(0);
-    const int64_t k = x.size(-1);
-    const int64_t n = up_weight.size(0);
-    TORCH_CHECK(m >= 1 && m <= 8, "M must be in [1, 8]");
-    TORCH_CHECK(
-        gate_weight.sizes() == up_weight.sizes(),
-        "up_weight and gate_weight must have identical shapes"
-    );
-    TORCH_CHECK(up_weight.size(1) == k, "weight K must match x K");
-    TORCH_CHECK(k > 0 && n > 0, "N and K must be positive");
-    TORCH_CHECK(k % 8 == 0, "K must be divisible by 8");
-    TORCH_CHECK(
-        k <= std::numeric_limits<int>::max() &&
-            n <= std::numeric_limits<int>::max(),
-        "N or K exceeds the CUDA launcher limit"
-    );
-
-    const at::cuda::OptionalCUDAGuard guard(x.device());
-    const auto* properties = at::cuda::getDeviceProperties(x.device().index());
-    TORCH_CHECK(
-        properties->major >= 8,
-        "bf16_swiglu requires compute capability 8.0+"
-    );
-    auto stream = at::cuda::getCurrentCUDAStream();
-    auto output = x.dim() == 1 ? torch::empty({n}, x.options())
-                               : torch::empty({m, n}, x.options());
-
-    const auto* x_ptr =
-        reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
-    const auto* up_ptr =
-        reinterpret_cast<const __nv_bfloat16*>(up_weight.data_ptr());
-    const auto* gate_ptr =
-        reinterpret_cast<const __nv_bfloat16*>(gate_weight.data_ptr());
-    auto* output_ptr =
-        reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
-    const int n_int = static_cast<int>(n);
-    const int k_int = static_cast<int>(k);
-
-    // Block size 256 keeps the weight streams at the HBM bandwidth floor for
-    // M in [1, 7]; M=8 halves the CTA so each thread owns more of the row
-    // and the shared-memory reduction tree shrinks (measured on L20 with
-    // rotated cold weights; larger CTAs only add idle warps).
-    switch (m) {
-        case 1:
-            launch_bf16_swiglu<256, 1>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 2:
-            launch_bf16_swiglu<256, 2>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 3:
-            launch_bf16_swiglu<256, 3>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 4:
-            launch_bf16_swiglu<256, 4>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 5:
-            launch_bf16_swiglu<256, 5>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 6:
-            launch_bf16_swiglu<256, 6>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 7:
-            launch_bf16_swiglu<256, 7>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-        case 8:
-            launch_bf16_swiglu<128, 8>(
-                x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-            );
-            break;
-    }
-    C10_CUDA_CHECK(cudaGetLastError());
-    return output;
-}
-
-}  // namespace
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
-    module.def(
-        "bf16_swiglu",
-        &bf16_swiglu,
-        py::arg("x"),
-        py::arg("up_weight"),
-        py::arg("gate_weight"),
-        "M in [1, 8] fused BF16 up/gate projection and SwiGLU"
-    );
-}

docs/developer/cuda_kernels.md

@@ -1,9 +1,7 @@
 # CUDA Kernels
 
 AstrAI includes optional custom CUDA kernels for attention, rotary embedding,
-BF16 GEMM/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and
+and FP8 GEMM. These are built when `nvcc` is available and CUDA is detected.
-CUDA is detected. BF16 GEMM and SwiGLU are directly callable and can be
-selected by guarded model dispatchers described below.
 
 ## Overview
 
@@ -14,121 +12,8 @@ selected by guarded model dispatchers described below.
 | `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention |
 | `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
 | `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
-| `bf16_gemm` | `gemm.cu` | M=1..64 BF16 linear, skinny GEMM path for M<=8 and tensor-core tiled path above (sm_80+) |
-| `bf16_swiglu` | `swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
 | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
 
-### BF16 GEMM primitive
-
-`astrai.extension.bf16_gemm(x, weight, bias=None)` accepts a contiguous BF16
-input shaped `[K]` or `[M, K]`, with `M` in `[1, 64]`, and row-major weights
-`[N, K]`. One entry point selects the internal path by M:
-
-- `M` in `[1, 8]` — register-resident GEMV kernel, any positive `K`: one
-  flat 128-thread CTA per weight row computes the output column for all M
-  tokens together. Under HBM-streaming conditions (weights rotated through
-  L2, as in real decode) the kernel is bandwidth-bound; on L20 the flat
-  128-thread CTA measured within ~1% of a per-shape tuned 128/256/512 mix at
-  M=1 and M=8 and at most ~2% below it at M=2-4, while dodging the register
-  cliff the larger CTAs hit on the 307MB lm_head at M>=7 (+21% DRAM
-  throughput vs 256t). Simplicity was chosen over the last ~2%.
-- `M` in `(8, 64]` — CUTLASS-style tiled kernel (`tiled_gemm_kernel`,
-  fully parameterized by template parameters BM/BN/BK/stages/threads):
-  multistage cp.async staging with XOR-swizzled 16B chunks, no split-K —
-  every block walks the whole K range in one pass and the grid fills the
-  SMs along N. Dispatch is shape-driven: wide N (n >= 4096, grid already
-  full) uses the bandwidth-optimal default (BN=64, BK=64, 3 stages; BM
-  widens to 32 at M>16); narrow N is K-serial, where widening the grid
-  measurably does nothing and deeper K chunks win (BN=32, BK=256 while
-  the grid fits one wave, BK=128 past it). Requires `K % 8 == 0` and
-  16-byte-aligned tensors.
-
-Both paths accumulate in FP32, fuse the optional BF16 bias before the BF16
-store, use the current CUDA stream, are CUDA Graph capture-safe, and require
-sm_80 or newer. The GEMV path anchors 128-bit weight loads at each row's
-first 16-byte-aligned address with scalar head/tail sweeps for unaligned
-remainders, so arbitrary `K` and storage offsets stay correct.
-
-Model `Linear` calls route through the lightweight linear backend. Set
-`ASTRAI_GEMM=0` for an unconditional `F.linear` fallback, `1` to force the
-kernel for any capable M (1-64, with K%8==0 and alignment for M>8), or
-`auto` (the default). On compute capability 8.0+, `auto` routes all capable
-decode batches M in [1, 64] through the kernel. The GEMV path (M≤8, any K)
-measured universal wins at the HBM bandwidth floor (AstrAI 1B end-to-end
-+8.8% at M=1 rising to +17% at M=8). The tiled path (M 9-64) measured on
-L20 across five shape families wins or ties four (q/o, kv, gate/up, lm_head)
-and regresses one (down-projection M>16: +25% latency; narrow-N long-K where
-cuBLAS is strong). Out-of-band, training, prefill-sized, or unsupported
-calls fall back to PyTorch.
-
-The measurements below date from the original `[2, 4]` automatic band; the
-band has since widened to `{1, ..., 8}` with the flat 128-thread kernel:
-
-On NVIDIA L20 (SM89), the common-shape microbenchmark reports +5.37% to
-+114.39% for M=2 and +5.38% to +115.26% for M=4 versus `F.linear`. The paired
-main-versus-warp-tiling run used identical interleaved settings; for the four
-M=4 selected shapes, candidate latency changed from 0.016292 to 0.016108 ms
-for `(4096,4096)`, 0.037939 to 0.028539 ms for `(11008,4096)`, 0.043407 to
-0.039803 ms for `(4096,11008)`, and 0.006697 to 0.006390 ms for
-`(1024,4096)`.
-
-The dependent projection-chain gate, which includes Python dispatch and
-rotates through distinct weights instead of repeatedly warming one matrix,
-measured:
-
-| Synthetic chain | M=2 | M=4 | Row argmax parity |
-|---|---:|---:|---|
-| LLaMA 2 7B | +8.49% | fallback (M=4 bands excluded) | exact |
-| LLaMA 3 8B | +8.50% | +6.44% | exact |
-| LLaMA 2 13B | +5.66% | +5.67% | exact |
-| GPT-NeoX 20B | +6.95% | +5.93% | exact |
-| Qwen2 7B | +7.48% | +7.48% | exact |
-| LLaMA 3 70B | +7.77% | +7.69% | exact |
-| OPT 1.3B | +25.20% | fallback (M=4 up projection regresses) | exact |
-
-OPT 1.3B M=1 is +4.54%. Qwen2 and LLaMA 3 70B M=1, and all three new
-families at M=8, remain exact PyTorch fallbacks.
-
-These are synthetic projection-chain measurements, not whole-model throughput
-claims. Reproduce them with `csrc/bench/benchmark_gemm_common.py`.
-
-The AstrAI 1B A→B→B→A results use the real `InferenceEngine`, including scheduler,
-sampling, and CUDA Graph. M=8 stays on PyTorch because its remaining
-greedy-stable winners missed the 3% end-to-end gate. Long-K MLP-down bands are
-also excluded because their valid BF16 error changed a checkpoint greedy
-argmax; the enabled M=2/4 bands matched the baseline greedy output exactly.
-
-Training, prefill, unmeasured architectures, and losing shape bands always
-remain on PyTorch. Use mode `1` only for explicit A/B runs outside this table.
-
-The primitive remains directly callable and deliberately has no internal
-`F.linear` fallback. The model-level backend owns fallback and dispatch policy.
-
-### BF16 SwiGLU primitive
-
-`astrai.extension.bf16_swiglu(x, up_weight, gate_weight)` fuses the two dense
-MLP projections with `up * silu(gate)` into one CUDA launch for contiguous
-BF16 inputs with `M` in `[1, 8]` and K divisible by 8. It preserves the BF16
-rounding boundaries of the two projection outputs, SiLU output, and final
-product while accumulating dot products in FP32.
-
-The kernel is a single CTA-reuse tiling: one CTA per output column reads each
-up/gate weight chunk once and applies it to all M rows. Block size is 256
-threads for M in [1, 7] and 128 for M=8, where the shorter shared-memory
-reduction tree wins under cold-HBM decode traffic. An earlier per-shape
-`(6912,1536)` warp-per-row variant and its dispatch table were removed: HBM
-measurements with rotated weights showed the table was tuned against L2-cache
-regime timing and was up to 6% slower than CTA reuse at M=2/4; the kernel is
-bandwidth-bound, so finer variant selection is noise.
-
-Dense `MLP` modules route through the SwiGLU backend. `ASTRAI_SWIGLU=0` keeps
-the unfused linear backend, and `1` explicitly forces the fused primitive.
-`auto` is the default but currently has no enabled bands: although direct
-errors are small (maximum absolute error at most 2.4e-4 in the L20 matrix),
-the different FP32 reduction order changed greedy checkpoint output for
-M=1/2/4. Automatic dispatch therefore remains numerically identical to the
-existing path.
-
 Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
 
 | Variant | File | Optimization |
@@ -318,13 +203,10 @@ astrai/extension/
 ├── ops/
 │   ├── attention.py        # Stateless attention kernel wrappers
 │   ├── rotary.py           # Stateless rotary kernel wrapper
-│   ├── gemm.py             # Stateless BF16 GEMM primitive
-│   ├── swiglu.py           # Stateless fused BF16 SwiGLU primitive
 │   └── fp8.py              # Stateless FP8 primitives (custom_op)
 ├── fp8.py                  # FP8 strategy layer (fp8_autocast, recipes)
 └── backend/
     ├── attention.py        # Backend selection, KV cache I/O, and fallback
-    ├── swiglu.py           # Inference-only fused/unfused SwiGLU policy
     └── rotary.py           # Per-call CUDA/torch rotary dispatch
 ```
 

setup.py

@@ -121,8 +121,6 @@ class _CMakeBuildExt(_build_ext):
             "attn_prefill",
             "attn_paged_decode",
             "attn_paged_prefill",
-            "bf16_gemm",
-            "bf16_swiglu",
             "rotary_emb",
         )
         missing = [name for name in required if not any(lib_dir.glob(f"{name}.*.so"))]

tests/extension/test_gemm.py

@@ -1,396 +0,0 @@
-import pytest
-import torch
-import torch.nn.functional as F
-
-from astrai.extension import bf16_gemm, is_available
-
-GEMM_AVAILABLE = (
-    torch.cuda.is_available()
-    and is_available("bf16_gemm")
-    and torch.cuda.get_device_capability() >= (8, 0)
-)
-skip_no_gemm = pytest.mark.skipif(
-    not GEMM_AVAILABLE,
-    reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
-)
-
-
-def _assert_close_fp64(actual, x, weight, bias=None):
-    """Compare bf16 kernel output vs fp64-exact with ulp-scaled tolerance.
-
-    Avoids false failures from cuBLAS default bf16 split-K partial reduction
-    (which can introduce ~2 ulp diffs on near-tie rounding at long K). Used
-    for tiled-path tests (M > 8, K >= 4096) where the bf16 accumulation tie
-    pattern may differ from cuBLAS's."""
-    exact = x.double() @ weight.double().T
-    if bias is not None:
-        exact = exact + bias.double()
-    ulp = (exact.abs() * 2**-9).clamp(min=2**-9)
-    max_ulp = ((actual.double() - exact).abs() / ulp).max().item()
-    assert max_ulp < 8, (
-        f"max_ulp={max_ulp:.1f} exceeds 8 (exact fp32 accumulation should stay within ~2 ulps)"
-    )
-
-
-@skip_no_gemm
-@pytest.mark.parametrize(
-    "n,k",
-    [(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)],
-)
-def test_bf16_gemm_matches_linear_shape_families(n, k):
-    torch.manual_seed(17)
-    x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight)
-    expected = F.linear(x, weight)
-    assert actual.shape == (n,)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [2, 3, 4, 5, 6, 7, 8])
-@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)])
-def test_bf16_gemm_matches_small_decode_batches(m, n, k):
-    torch.manual_seed(19 + m)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight)
-    expected = F.linear(x, weight)
-    assert actual.shape == (m, n)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [2, 4])
-@pytest.mark.parametrize(
-    "n,k",
-    [
-        (1024, 4096),
-        (4096, 4096),
-        (11008, 4096),
-        (4096, 11008),
-        (14336, 4096),
-        (4096, 14336),
-        (5120, 5120),
-        (13824, 5120),
-        (5120, 13824),
-        (16384, 4096),
-        (4096, 16384),
-        (512, 3584),
-        (3584, 3584),
-        (18944, 3584),
-        (3584, 18944),
-        (1024, 8192),
-        (8192, 8192),
-        (28672, 8192),
-        (8192, 28672),
-        (2048, 2048),
-        (8192, 2048),
-        (2048, 8192),
-    ],
-)
-def test_bf16_gemm_matches_common_transformer_shapes(m, n, k):
-    torch.manual_seed(2026 + m + n + k)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
-    weight.normal_(mean=0.0, std=0.02)
-    actual = bf16_gemm(x, weight)
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize(
-    "m,n,k",
-    [
-        (1, 8192, 2048),
-        (8, 4096, 11008),
-        (8, 512, 3584),
-        (8, 1024, 8192),
-        (8, 2048, 8192),
-    ],
-)
-def test_bf16_gemm_matches_m8_edge_bands(m, n, k):
-    torch.manual_seed(2026 + m + n + k)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
-    weight.normal_(mean=0.0, std=0.02)
-    actual = bf16_gemm(x, weight)
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-def test_bf16_gemm_preserves_singleton_batch_and_fuses_bias():
-    torch.manual_seed(23)
-    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    bias = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight, bias)
-    expected = F.linear(x, weight, bias)
-    assert actual.shape == (1, 1536)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-def test_bf16_gemm_small_batch_fuses_bias():
-    torch.manual_seed(25)
-    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight, bias)
-    expected = F.linear(x, weight, bias)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-def test_bf16_gemm_uses_current_stream():
-    x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    torch.cuda.synchronize()
-    stream = torch.cuda.Stream()
-    with torch.cuda.stream(stream):
-        actual = bf16_gemm(x, weight)
-        expected = F.linear(x, weight)
-    stream.synchronize()
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-def test_bf16_gemm_cuda_graph_replay():
-    torch.manual_seed(29)
-    x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    for _ in range(3):
-        bf16_gemm(x, weight)
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        actual = bf16_gemm(x, weight)
-
-    x.copy_(torch.randn_like(x))
-    graph.replay()
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("n,k", [(64, 7), (64, 12), (33, 100), (256, 1534)])
-def test_bf16_gemm_handles_unaligned_k(n, k):
-    torch.manual_seed(29)
-    x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight)
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-    x3 = torch.randn(3, k, device="cuda", dtype=torch.bfloat16)
-    actual3 = bf16_gemm(x3, weight)
-    torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [1, 2, 3, 4])
-def test_bf16_gemm_accepts_complementary_misalignment(m):
-    """Misaligned weight rows plus an x base chosen so the vectorized branch
-    is entered with a non-16B-aligned ``x`` pointer (regression: the branch
-    guard checked ``x + whead`` alignment but the uint4 view was rooted at
-    ``x`` itself, faulting with a misaligned-address CUDA error)."""
-    torch.manual_seed(37)
-    n, k = 256, 1536
-    # offset 5 elements = +10 bytes: weight rows land at 10 % 16 (whead=3)
-    # and x at 10 % 16, so (x + 2*whead) % 16 == 0 selects the fast path.
-    big_w = torch.randn(n * k + 8, device="cuda", dtype=torch.bfloat16)
-    weight = big_w[5 : 5 + n * k].view(n, k)
-    big_x = torch.randn(m * k + 8, device="cuda", dtype=torch.bfloat16)
-    x = big_x[5 : 5 + m * k].view(m, k) if m > 1 else big_x[5 : 5 + k]
-    assert (x.data_ptr() & 15) == 10 and (weight.data_ptr() & 15) == 10
-
-    actual = bf16_gemm(x, weight)
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5 if m > 1 else 0.25)
-
-
-@skip_no_gemm
-def test_bf16_gemm_scalar_path_handles_misaligned_weight_only():
-    """Weight rows misaligned while x stays 16B-aligned take the scalar-x
-    middle and must stay exact."""
-    torch.manual_seed(41)
-    n, k = 256, 1536
-    big_w = torch.randn(n * k + 8, device="cuda", dtype=torch.bfloat16)
-    weight = big_w[5 : 5 + n * k].view(n, k)
-    x = torch.randn(2, k, device="cuda", dtype=torch.bfloat16)
-    assert (weight.data_ptr() & 15) == 10 and (x.data_ptr() & 15) == 0
-
-    actual = bf16_gemm(x, weight)
-    torch.testing.assert_close(actual, F.linear(x, weight), rtol=0.02, atol=0.5)
-
-
-@skip_no_gemm
-def test_bf16_gemm_small_batch_cuda_graph_replay():
-    torch.manual_seed(31)
-    x = torch.randn(8, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    for _ in range(3):
-        bf16_gemm(x, weight)
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        actual = bf16_gemm(x, weight)
-
-    x.copy_(torch.randn_like(x))
-    graph.replay()
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize(
-    "make_args,error",
-    [
-        (
-            lambda: (
-                torch.randn(65, 16, device="cuda", dtype=torch.bfloat16),
-                torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
-            ),
-            "M must",
-        ),
-        (
-            lambda: (
-                torch.randn(16, device="cuda", dtype=torch.float16),
-                torch.randn(8, 16, device="cuda", dtype=torch.float16),
-            ),
-            "bf16",
-        ),
-        (
-            lambda: (
-                torch.randn(
-                    16, device="cuda", dtype=torch.bfloat16, requires_grad=True
-                ),
-                torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
-            ),
-            "autograd",
-        ),
-    ],
-)
-def test_bf16_gemm_rejects_unsupported_inputs(make_args, error):
-    with pytest.raises(RuntimeError, match=error):
-        bf16_gemm(*make_args())
-
-
-# ---------------------------------------------------------------------------
-# Tiled path: M in (8, 64]
-# ---------------------------------------------------------------------------
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [9, 12, 16, 17, 24, 32, 33, 48, 64])
-@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912)])
-def test_bf16_gemm_tiled_matches_decode_batches(m, n, k):
-    torch.manual_seed(31 + m)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
-    weight.normal_(mean=0.0, std=0.02)
-    actual = bf16_gemm(x, weight)
-    assert actual.shape == (m, n)
-    _assert_close_fp64(actual, x, weight)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [12, 64])
-def test_bf16_gemm_tiled_matches_lm_head(m):
-    # N=100000 fills the SMs with N tiles alone: the splits=1 epilogue.
-    torch.manual_seed(37 + m)
-    x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.empty(100000, 1536, device="cuda", dtype=torch.bfloat16)
-    weight.normal_(mean=0.0, std=0.02)
-    actual = bf16_gemm(x, weight)
-    _assert_close_fp64(actual, x, weight)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize(
-    "m,n,k",
-    [
-        (12, 100, 72),
-        (12, 96, 1536),
-        (12, 1632, 1536),
-        (64, 100, 72),
-        (17, 200, 8),
-        (33, 160, 152),
-    ],
-)
-def test_bf16_gemm_tiled_handles_remainder_tiles(m, n, k):
-    # N not a multiple of 64 (predicated epilogue columns) and K not a
-    # multiple of 64 (zero-filled staging chunks).
-    torch.manual_seed(41 + m + n + k)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight)
-    _assert_close_fp64(actual, x, weight)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m,n,k", [(12, 1536, 1536), (64, 6912, 1536)])
-def test_bf16_gemm_tiled_fuses_bias(m, n, k):
-    # (12, 1536) exercises the narrow-N deep-K config; (64, 6912) the
-    # wide-N default.
-    torch.manual_seed(43 + m)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
-    actual = bf16_gemm(x, weight, bias)
-    _assert_close_fp64(actual, x, weight, bias)
-
-
-@skip_no_gemm
-def test_bf16_gemm_tiled_deterministic_across_runs():
-    # Single-pass K accumulation with no atomics: reruns are bitwise
-    # identical — CUDA Graph replay relies on this.
-    torch.manual_seed(47)
-    x = torch.randn(16, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    first = bf16_gemm(x, weight)
-    second = bf16_gemm(x, weight)
-    assert torch.equal(first, second)
-
-
-@skip_no_gemm
-def test_bf16_gemm_tiled_cuda_graph_replay():
-    torch.manual_seed(53)
-    x = torch.randn(16, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    for _ in range(3):
-        bf16_gemm(x, weight)
-
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        actual = bf16_gemm(x, weight)
-
-    x.copy_(torch.randn_like(x))
-    graph.replay()
-    _assert_close_fp64(actual, x, weight)
-
-
-@skip_no_gemm
-def test_bf16_gemm_tiled_rejects_k_not_multiple_of_8():
-    x = torch.randn(12, 12, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(64, 12, device="cuda", dtype=torch.bfloat16)
-    with pytest.raises(RuntimeError, match="multiple of 8"):
-        bf16_gemm(x, weight)
-
-
-@skip_no_gemm
-def test_bf16_gemm_tiled_rejects_misaligned_x():
-    # A 2-byte storage offset breaks the 16B alignment the tiled path
-    # stages chunks on; the M <= 8 GEMV path still accepts it.
-    k = 1536
-    storage = torch.randn(12 * k + 1, device="cuda", dtype=torch.bfloat16)
-    x8 = storage[1 : 1 + 8 * k].view(8, k)
-    x12 = storage[1 : 1 + 12 * k].view(12, k)
-    weight = torch.randn(1536, k, device="cuda", dtype=torch.bfloat16)
-    torch.testing.assert_close(
-        bf16_gemm(x8, weight), F.linear(x8, weight), rtol=0.02, atol=0.5
-    )
-    with pytest.raises(RuntimeError, match="16-byte"):
-        bf16_gemm(x12, weight)

tests/extension/test_linear_dispatch.py

@@ -1,233 +0,0 @@
-import importlib
-import logging
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from astrai.extension import is_available, linear
-from astrai.extension.dispatch import explain, op_backend, resolve
-
-# The package attribute ``linear`` is the dispatched function; reach the
-# module object explicitly for monkeypatching its private helpers.
-linear_module = importlib.import_module("astrai.extension.backend.linear")
-
-GEMM_AVAILABLE = (
-    torch.cuda.is_available()
-    and is_available("bf16_gemm")
-    and torch.cuda.get_device_capability() >= (8, 0)
-)
-skip_no_gemm = pytest.mark.skipif(
-    not GEMM_AVAILABLE,
-    reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
-)
-
-
-def _routes_to_gemm(monkeypatch, x, weight, bias=None) -> bool:
-    """Patch the GEMM entry point to a sentinel and report whether
-    ``linear`` selected it (torch fallback would compute a real tensor)."""
-    sentinel = object()
-
-    def fake_gemm(x, weight, bias):
-        return sentinel
-
-    monkeypatch.setattr(linear_module, "_inference_bf16_gemm", fake_gemm)
-    return linear(x, weight, bias) is sentinel
-
-
-def test_model_linear_routes_through_backend(monkeypatch):
-    sentinel = torch.randn(2, 4)
-
-    def fake_linear(x, weight, bias):
-        assert x.shape == (2, 3)
-        assert weight.shape == (4, 3)
-        assert bias is None
-        return sentinel
-
-    monkeypatch.setattr("astrai.model.components.linear.linear", fake_linear)
-    from astrai.model.components.linear import Linear
-
-    layer = Linear(3, 4)
-    assert layer(torch.randn(2, 3)) is sentinel
-
-
-def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
-    monkeypatch.setenv("ASTRAI_GEMM", "invalid-test-mode")
-    x = torch.randn(2, 8)
-    weight = torch.randn(4, 8)
-    with caplog.at_level(logging.WARNING):
-        actual = linear(x, weight)
-    assert "using auto" in caplog.text
-    torch.testing.assert_close(actual, F.linear(x, weight))
-
-
-def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
-    monkeypatch.setenv("ASTRAI_GEMM", "1")
-    x = torch.randn(2, 8, requires_grad=True)
-    weight = torch.randn(4, 8, requires_grad=True)
-    actual = linear(x, weight)
-    expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected)
-    actual.sum().backward()
-    assert x.grad is not None
-    assert weight.grad is not None
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [1, 2, 3, 4, 5, 6, 7, 8])
-def test_auto_selects_small_decode_batches(monkeypatch, m):
-    monkeypatch.setenv("ASTRAI_GEMM", "auto")
-    x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert _routes_to_gemm(monkeypatch, x, weight)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [12, 16, 24, 32])
-def test_auto_selects_larger_decode_batches(monkeypatch, m):
-    """Auto covers M up to 32; 48+ loses to cuBLAS on long-K shapes."""
-    monkeypatch.setenv("ASTRAI_GEMM", "auto")
-    x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert _routes_to_gemm(monkeypatch, x, weight)
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [48, 64, 65])
-def test_auto_falls_back_outside_band(monkeypatch, m):
-    """M beyond 32 falls back to cuBLAS (measured regression at M=48+)."""
-    monkeypatch.setenv("ASTRAI_GEMM", "auto")
-    x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert not _routes_to_gemm(monkeypatch, x, weight)
-        torch.testing.assert_close(
-            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
-        )
-
-
-@skip_no_gemm
-def test_mode_zero_disables_gemm(monkeypatch):
-    monkeypatch.setenv("ASTRAI_GEMM", "0")
-    x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert not _routes_to_gemm(monkeypatch, x, weight)
-        torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
-
-
-@skip_no_gemm
-@pytest.mark.parametrize("m", [1, 2, 8, 16, 32])
-def test_mode_one_forces_every_capable_batch(monkeypatch, m):
-    monkeypatch.setenv("ASTRAI_GEMM", "1")
-    x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert _routes_to_gemm(monkeypatch, x, weight)
-
-
-@skip_no_gemm
-def test_mode_one_rejects_oversized_batch_and_grad(monkeypatch):
-    monkeypatch.setenv("ASTRAI_GEMM", "1")
-    weight = torch.randn(
-        256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
-    )
-    with torch.no_grad():
-        oversized = torch.randn(65, 1536, device="cuda", dtype=torch.bfloat16)
-        assert not _routes_to_gemm(monkeypatch, oversized, weight)
-    assert not _routes_to_gemm(
-        monkeypatch, torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16), weight
-    )
-
-
-@skip_no_gemm
-def test_mode_one_supports_bias_and_vector_input(monkeypatch):
-    monkeypatch.setenv("ASTRAI_GEMM", "1")
-    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert _routes_to_gemm(monkeypatch, x, weight, bias)
-        monkeypatch.undo()
-        torch.testing.assert_close(
-            linear(x, weight, bias),
-            F.linear(x, weight, bias),
-            rtol=0.02,
-            atol=0.25,
-        )
-
-
-@skip_no_gemm
-def test_dispatched_linear_cuda_graph_replay(monkeypatch):
-    monkeypatch.setenv("ASTRAI_GEMM", "1")
-    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        for _ in range(3):
-            linear(x, weight)
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            actual = linear(x, weight)
-        x.copy_(torch.randn_like(x))
-        graph.replay()
-        expected = F.linear(x, weight)
-    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
-
-
-def test_linear_family_is_registered_with_shared_dispatcher():
-    from astrai.extension.dispatch import _FAMILIES
-
-    assert "linear" in _FAMILIES
-    x = torch.randn(2, 8)
-    weight = torch.randn(4, 8)
-    resolution = resolve("linear", x, weight)
-    assert resolution.record.family == "linear"
-    assert resolution.origin in ("chain", "fallback")
-    assert "linear" in explain("linear", x, weight)
-
-
-@skip_no_gemm
-def test_ops_env_override_forces_torch_for_capable_call(monkeypatch):
-    """ASTR_OPS=linear=torch must keep working after the M-band rewrite
-    (regression: the family was silently dropped from the dispatcher, so
-    the override warned, fell through, and the gemm kernel still ran)."""
-    monkeypatch.setenv("ASTR_OPS", "linear=torch")
-    x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        assert not _routes_to_gemm(monkeypatch, x, weight)
-        torch.testing.assert_close(
-            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
-        )
-
-
-@skip_no_gemm
-def test_ops_env_override_forces_gemm(monkeypatch):
-    monkeypatch.setenv("ASTR_OPS", "linear=gemm")
-    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        # M=1 is outside the auto band but inside the forced gemm record.
-        assert _routes_to_gemm(monkeypatch, x, weight)
-
-
-@skip_no_gemm
-def test_op_backend_context_selects_torch(monkeypatch):
-    monkeypatch.setenv("ASTRAI_GEMM", "1")
-    x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad(), op_backend(linear="torch"):
-        assert not _routes_to_gemm(monkeypatch, x, weight)
-        torch.testing.assert_close(
-            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
-        )
-    # The override is scoped: the forced mode applies again afterwards.
-    with torch.no_grad():
-        assert _routes_to_gemm(monkeypatch, x, weight)
-
-
-def test_op_backend_rejects_unknown_linear_handle():
-    with pytest.raises(ValueError, match="Unknown linear implementation"):
-        op_backend(linear="nonexistent").__enter__()

tests/extension/test_swiglu.py

@@ -1,112 +0,0 @@
-import pytest
-import torch
-import torch.nn.functional as F
-
-from astrai.extension import bf16_swiglu, is_available
-
-SWIGLU_AVAILABLE = (
-    torch.cuda.is_available()
-    and is_available("bf16_swiglu")
-    and torch.cuda.get_device_capability() >= (8, 0)
-)
-skip_no_swiglu = pytest.mark.skipif(
-    not SWIGLU_AVAILABLE,
-    reason="BF16 SwiGLU requires a built kernel and compute capability 8.0+",
-)
-
-
-def reference_swiglu(x, up_weight, gate_weight):
-    return F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight))
-
-
-@skip_no_swiglu
-@pytest.mark.parametrize("m", [1, 2, 4, 8])
-@pytest.mark.parametrize("n,k", [(6912, 1536), (4096, 4096), (11008, 4096)])
-def test_bf16_swiglu_matches_common_dense_mlp_shapes(m, n, k):
-    torch.manual_seed(37 + m)
-    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.1
-    up_weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
-    gate_weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
-    actual = bf16_swiglu(x, up_weight, gate_weight)
-    expected = reference_swiglu(x, up_weight, gate_weight)
-    assert actual.shape == (m, n)
-    torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
-
-
-@skip_no_swiglu
-def test_bf16_swiglu_preserves_vector_shape():
-    x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) * 0.1
-    up_weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
-    gate_weight = torch.randn_like(up_weight) * 0.02
-    actual = bf16_swiglu(x, up_weight, gate_weight)
-    expected = reference_swiglu(x, up_weight, gate_weight)
-    assert actual.shape == (256,)
-    torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
-
-
-@skip_no_swiglu
-def test_bf16_swiglu_uses_current_stream_and_cuda_graph():
-    torch.manual_seed(43)
-    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1
-    up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
-    gate_weight = torch.randn_like(up_weight) * 0.02
-    stream = torch.cuda.Stream()
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            bf16_swiglu(x, up_weight, gate_weight)
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            actual = bf16_swiglu(x, up_weight, gate_weight)
-        x.copy_(torch.randn_like(x) * 0.1)
-        graph.replay()
-    stream.synchronize()
-    expected = reference_swiglu(x, up_weight, gate_weight)
-    torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
-
-
-@skip_no_swiglu
-@pytest.mark.parametrize(
-    "make_args,error",
-    [
-        (
-            lambda: (
-                torch.randn(9, 16, device="cuda", dtype=torch.bfloat16),
-                torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
-                torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
-            ),
-            "M must",
-        ),
-        (
-            lambda: (
-                torch.randn(2, 15, device="cuda", dtype=torch.bfloat16),
-                torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
-                torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
-            ),
-            "divisible by 8",
-        ),
-        (
-            lambda: (
-                torch.randn(2, 16, device="cuda", dtype=torch.bfloat16),
-                torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
-                torch.randn(7, 16, device="cuda", dtype=torch.bfloat16),
-            ),
-            "identical shapes",
-        ),
-    ],
-)
-def test_bf16_swiglu_rejects_unsupported_inputs(make_args, error):
-    with pytest.raises(RuntimeError, match=error):
-        bf16_swiglu(*make_args())
-
-
-@skip_no_swiglu
-def test_bf16_swiglu_rejects_misaligned_storage_with_clear_error():
-    """Contiguous-but-offset views must fail the wrapper's TORCH_CHECK with
-    an actionable message instead of a sticky CUDA misaligned-address error
-    (regression: the kernel casts x directly to uint4 without checking)."""
-    k = 1536
-    base = torch.randn(k + 1, device="cuda", dtype=torch.bfloat16)
-    x = base[1:]  # +2 bytes: contiguous but not 16B-aligned
-    weights = torch.randn(8, k, device="cuda", dtype=torch.bfloat16)
-    with pytest.raises(RuntimeError, match="16-byte aligned"):
-        bf16_swiglu(x, weights, weights)

tests/extension/test_swiglu_dispatch.py

@@ -1,118 +0,0 @@
-import logging
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from astrai.extension import is_available, swiglu
-from astrai.model.components.mlp import MLP
-
-SWIGLU_AVAILABLE = (
-    torch.cuda.is_available()
-    and is_available("bf16_swiglu")
-    and torch.cuda.get_device_capability() >= (8, 0)
-)
-skip_no_swiglu = pytest.mark.skipif(
-    not SWIGLU_AVAILABLE,
-    reason="BF16 SwiGLU requires a built kernel and compute capability 8.0+",
-)
-
-
-def reference_swiglu(x, up_weight, gate_weight):
-    return F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight))
-
-
-def test_cpu_and_training_calls_fall_back_with_gradients(monkeypatch):
-    monkeypatch.setenv("ASTRAI_SWIGLU", "1")
-    x = torch.randn(2, 8, requires_grad=True)
-    up_weight = torch.randn(4, 8, requires_grad=True)
-    gate_weight = torch.randn(4, 8, requires_grad=True)
-    actual = swiglu(x, up_weight, gate_weight)
-    expected = reference_swiglu(x, up_weight, gate_weight)
-    torch.testing.assert_close(actual, expected)
-    actual.sum().backward()
-    assert x.grad is not None
-    assert up_weight.grad is not None
-    assert gate_weight.grad is not None
-
-
-def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
-    monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode")
-    x = torch.randn(2, 8)
-    up_weight = torch.randn(4, 8)
-    gate_weight = torch.randn(4, 8)
-    with caplog.at_level(logging.WARNING):
-        actual = swiglu(x, up_weight, gate_weight)
-    assert "using auto" in caplog.text
-    torch.testing.assert_close(actual, reference_swiglu(x, up_weight, gate_weight))
-
-
-def test_mlp_routes_through_swiglu_backend(monkeypatch):
-    sentinel = torch.randn(2, 4)
-
-    def fake_swiglu(x, up_weight, gate_weight):
-        assert x.shape == (2, 3)
-        assert up_weight.shape == gate_weight.shape == (4, 3)
-        return sentinel
-
-    monkeypatch.setattr("astrai.model.components.mlp.swiglu", fake_swiglu)
-    layer = MLP(3, 4)
-    output = layer(torch.randn(2, 3))
-    assert output["hidden_states"].shape == (2, 3)
-
-
-@skip_no_swiglu
-def test_mode_zero_disables_fused_kernel(monkeypatch):
-    monkeypatch.setenv("ASTRAI_SWIGLU", "0")
-    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1
-    up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
-    gate_weight = torch.randn_like(up_weight) * 0.02
-    with torch.no_grad():
-        actual = swiglu(x, up_weight, gate_weight)
-        expected = reference_swiglu(x, up_weight, gate_weight)
-    torch.testing.assert_close(actual, expected)
-
-
-@skip_no_swiglu
-def test_mode_one_forces_supported_shape(monkeypatch):
-    monkeypatch.setenv("ASTRAI_SWIGLU", "1")
-    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1
-    up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
-    gate_weight = torch.randn_like(up_weight) * 0.02
-    with torch.no_grad():
-        actual = swiglu(x, up_weight, gate_weight)
-        expected = reference_swiglu(x, up_weight, gate_weight)
-    torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
-
-
-@skip_no_swiglu
-def test_auto_uses_fused_chain_for_decode_batches(monkeypatch):
-    # Auto adopts the fused primitive for the decode band; numerics match
-    # the unfused linear-backend chain within BF16 accumulation-order noise.
-    monkeypatch.setenv("ASTRAI_SWIGLU", "auto")
-    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
-    up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16)
-    gate_weight = torch.randn_like(up_weight)
-    with torch.no_grad():
-        actual = swiglu(x, up_weight, gate_weight)
-        expected = reference_swiglu(x, up_weight, gate_weight)
-    torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.1)
-
-
-@skip_no_swiglu
-def test_mode_one_falls_back_for_misaligned_storage(monkeypatch):
-    """Contiguous-but-offset views must route to the unfused torch chain
-    even with ASTRAI_SWIGLU=1 instead of reaching the uint4-only kernel
-    (regression: the fused primitive faulted with a misaligned-address
-    CUDA error for such inputs)."""
-    monkeypatch.setenv("ASTRAI_SWIGLU", "1")
-    k = 1536
-    x_base = torch.randn(2 * k + 8, device="cuda", dtype=torch.bfloat16) * 0.1
-    x = x_base[1 : 1 + 2 * k].view(2, k)
-    assert x.is_contiguous() and (x.data_ptr() & 15) != 0
-    up_weight = torch.randn(64, k, device="cuda", dtype=torch.bfloat16) * 0.02
-    gate_weight = torch.randn_like(up_weight)
-    with torch.no_grad():
-        actual = swiglu(x, up_weight, gate_weight)
-        expected = reference_swiglu(x, up_weight, gate_weight)
-    torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)