perf: accelerate decode linear with bf16 gemv

- add decode-shape benchmark harness
- add bf16 GEMV CUDA primitive with head-dim generic kernel
- dispatch decode-time linear layers to gemv for M=1
- extend gemv coverage to small decode batches

astrai/extension/__init__.py

@@ -26,6 +26,7 @@ from astrai.extension.backend import (
     attention,
     attn_backend,
     get_backend,
+    linear,
 )
 from astrai.extension.dispatch import (
     Axes,
@@ -49,6 +50,7 @@ from astrai.extension.ops import (
     attn_decode,
     attn_paged_decode,
     attn_prefill,
+    bf16_gemv,
 )
 
 __all__ = [
@@ -62,9 +64,11 @@ __all__ = [
     "attention",
     "attn_backend",
     "get_backend",
+    "linear",
     "attn_decode",
     "attn_paged_decode",
     "attn_prefill",
+    "bf16_gemv",
     "is_available",
     "KERNEL_NAMES",
     "apply_rotary_emb",

astrai/extension/backend/__init__.py

@@ -11,6 +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
 
 __all__ = [
@@ -24,4 +25,5 @@ __all__ = [
     "attention",
     "attn_backend",
     "get_backend",
+    "linear",
 ]

astrai/extension/backend/linear.py

@@ -0,0 +1,264 @@
+"""Inference-only dispatch for AstrAI linear layers.
+
+The CUDA GEMV path is deliberately narrow: automatic selection is enabled
+only for single-row BF16 shapes measured to beat ``F.linear`` on a supported
+architecture. Every training, prefill, unsupported-layout, and unmeasured
+call falls back to PyTorch.
+"""
+
+import logging
+import os
+from functools import lru_cache
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from astrai.extension.dispatch import (
+    ImplRecord,
+    Spec,
+    axis,
+    get_override,
+    register_family,
+    resolve,
+    tensor_axes,
+)
+from astrai.extension.loader import is_available
+from astrai.extension.ops.gemv import bf16_gemv
+
+logger = logging.getLogger(__name__)
+
+# Shape keys are (N, K) for Y[M, N] = X[M, K] @ W[N, K].T. A band is
+# automatic only after both the per-shape >=5% and end-to-end decode >=3%
+# gates pass and checkpoint greedy output remains stable. M=1 and M=8 remain
+# empty on SM89; the safe M=2/4 bands improve real-engine throughput by
+# 11.8-14.0%.
+_AUTO_GEMV_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = {
+    (8, 9): {
+        2: frozenset(
+            {
+                (256, 1536),
+                (1536, 1536),
+                (100000, 1536),
+            }
+        ),
+        4: frozenset({(256, 1536), (1536, 1536)}),
+    }
+}
+_AUTO_GEMV_M = frozenset(
+    m for architecture in _AUTO_GEMV_SHAPES.values() for m in architecture
+)
+
+_VALID_MODES = {"0", "1", "auto"}
+_WARNED_MODES: set[str] = set()
+
+
+def _gemv_mode() -> str:
+    mode = os.environ.get("ASTRAI_GEMV", "auto").strip().lower()
+    if mode in _VALID_MODES:
+        return mode
+    if mode not in _WARNED_MODES:
+        _WARNED_MODES.add(mode)
+        logger.warning(
+            "ASTRAI_GEMV=%r is invalid; expected 0, 1, or auto; using auto",
+            mode,
+        )
+    return "auto"
+
+
+def _axes(
+    x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
+) -> dict[str, object]:
+    x_shape = tuple(x.shape)
+    weight_shape = tuple(weight.shape)
+    m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
+    supported_m = m in (1, 2, 4, 8)
+    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
+    n = weight_shape[0] if weight.ndim == 2 else None
+    k = weight_shape[1] if weight.ndim == 2 else None
+    return tensor_axes(
+        x,
+        mode=_gemv_mode(),
+        capability=capability,
+        n=n,
+        k=k,
+        m=m,
+        supported_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,
+        k_even=k is not None and k % 2 == 0,
+    )
+
+
+_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()
+    & axis("k_even").truthy()
+)
+
+_SPEC_AUTO = _SPEC_CAPABLE & Spec.of(
+    lambda ax: (
+        (ax.get("n"), ax.get("k"))
+        in _AUTO_GEMV_SHAPES.get(ax.get("capability"), {}).get(ax.get("m"), ())
+    ),
+    "shape is a measured winner for this architecture",
+)
+
+
+def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
+    return F.linear(x, weight, bias)
+
+
+def _inference_bf16_gemv(
+    x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
+) -> Tensor:
+    # Model parameters retain requires_grad=True after eval().  Dispatch is
+    # already restricted to no-grad, so detached views preserve storage and
+    # layout while satisfying the primitive's explicit autograd guard.
+    return bf16_gemv(
+        x.detach(),
+        weight.detach(),
+        bias.detach() if bias is not None else None,
+    )
+
+
+@lru_cache(maxsize=None)
+def _device_capability(device_index: int) -> tuple[int, int]:
+    return torch.cuda.get_device_capability(device_index)
+
+
+def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
+    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.ndim == 2 and x.shape[0] not in (1, 2, 4, 8))
+        or x.shape[-1] != weight.shape[1]
+        or weight.shape[1] % 2 != 0
+        or x.device != weight.device
+        or not x.is_contiguous()
+        or not weight.is_contiguous()
+        or not is_available("bf16_gemv")
+    ):
+        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 _auto_gemv_shape(x: Tensor, weight: Tensor) -> bool:
+    capability = _device_capability(x.get_device())
+    m = 1 if x.ndim == 1 else x.shape[0]
+    return (weight.shape[0], weight.shape[1]) in _AUTO_GEMV_SHAPES.get(
+        capability, {}
+    ).get(m, ())
+
+
+def _linear_records() -> list[ImplRecord]:
+    mode = _gemv_mode()
+    gemv_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="gemv",
+            obj=_inference_bf16_gemv,
+            spec=_SPEC_CAPABLE,
+            available=lambda: is_available("bf16_gemv"),
+            priority=gemv_priority,
+        ),
+        ImplRecord(
+            family="linear",
+            name="auto_gemv",
+            obj=_inference_bf16_gemv,
+            spec=_SPEC_AUTO,
+            available=lambda: is_available("bf16_gemv"),
+            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 GEMV dispatch.
+
+    ``ASTRAI_GEMV=0`` always uses PyTorch, ``1`` forces GEMV whenever the
+    primitive can safely handle an M in ``{1, 2, 4, 8}``, and ``auto`` (the
+    default) uses only architecture/shape bands backed by benchmark evidence.
+    """
+    # Preserve the shared dispatcher for explicit/context selection and
+    # ASTR_OPS diagnostics, while keeping the default per-layer hot path free
+    # of axes dictionaries, record sorting, and repeated capability queries.
+    if get_override("linear") is not None or "linear" in os.environ.get("ASTR_OPS", ""):
+        return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
+
+    mode = _gemv_mode()
+    if mode == "0" or (mode == "auto" and not _AUTO_GEMV_SHAPES):
+        return _torch_linear(x, weight, bias)
+    if mode == "auto":
+        m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
+        if m not in _AUTO_GEMV_M:
+            return _torch_linear(x, weight, bias)
+    if mode != "0" and _gemv_capable(x, weight, bias):
+        if mode == "1" or _auto_gemv_shape(x, weight):
+            return _inference_bf16_gemv(x, weight, bias)
+    return _torch_linear(x, weight, bias)
+
+
+__all__ = ["linear"]

astrai/extension/ops/__init__.py

@@ -7,6 +7,7 @@ from astrai.extension.ops.attention import (
     attn_paged_prefill,
     attn_prefill,
 )
+from astrai.extension.ops.gemv import bf16_gemv
 from astrai.extension.ops.rotary import rotary_emb
 
 __all__ = [
@@ -15,5 +16,6 @@ __all__ = [
     "attn_paged_decode",
     "attn_paged_prefill",
     "attn_prefill",
+    "bf16_gemv",
     "rotary_emb",
 ]

astrai/extension/ops/gemv.py

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

astrai/model/components/linear.py

@@ -1,8 +1,9 @@
 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__(
@@ -21,4 +22,4 @@ class Linear(nn.Module):
             nn.init.uniform_(self.bias, -bound, bound)
 
     def forward(self, x: Tensor) -> Tensor:
-        return F.linear(x, self.weight, self.bias)
+        return linear(x, self.weight, self.bias)

csrc/CMakeLists.txt

@@ -61,6 +61,7 @@ set(KERNEL_NAMES
     attn_prefill
     attn_paged_decode
     attn_paged_prefill
+    bf16_gemv
     rotary_emb
 )
 set(KERNEL_SRCS
@@ -68,6 +69,7 @@ set(KERNEL_SRCS
     attention/prefill.cu
     attention/paged_decode.cu
     attention/paged_prefill.cu
+    gemv/bf16_gemv.cu
     rotary_emb.cu
 )
 

csrc/kernels/gemv/bf16_gemv.cu

@@ -0,0 +1,206 @@
+// Directly callable small-M BF16 GEMV primitive for decode-time linear layers.
+
+#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 kThreads = 256;
+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>
+__global__ void bf16_gemv_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;
+    const int pairs = k / 2;
+    const auto* x2 = reinterpret_cast<const __nv_bfloat162*>(x);
+    const auto* w2 =
+        reinterpret_cast<const __nv_bfloat162*>(weight) + output_index * pairs;
+
+    float sums[Rows] = {};
+    for (int pair = threadIdx.x; pair < pairs; pair += blockDim.x) {
+        const __nv_bfloat162 wv = w2[pair];
+#pragma unroll
+        for (int row = 0; row < Rows; ++row) {
+            const __nv_bfloat162 xv = x2[row * pairs + pair];
+            sums[row] = fmaf(
+                __bfloat162float(__low2bfloat16(xv)),
+                __bfloat162float(__low2bfloat16(wv)),
+                sums[row]
+            );
+            sums[row] = fmaf(
+                __bfloat162float(__high2bfloat16(xv)),
+                __bfloat162float(__high2bfloat16(wv)),
+                sums[row]
+            );
+        }
+    }
+
+    __shared__ float warp_sums[Rows][kThreads / kWarpSize];
+#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 < (kThreads / 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>
+void launch_bf16_gemv(
+    const __nv_bfloat16* x,
+    const __nv_bfloat16* weight,
+    const __nv_bfloat16* bias,
+    __nv_bfloat16* output,
+    int n,
+    int k,
+    cudaStream_t stream
+) {
+    bf16_gemv_kernel<Rows><<<n, kThreads, 0, stream>>>(
+        x, weight, bias, output, n, k
+    );
+}
+
+torch::Tensor bf16_gemv(
+    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_gemv 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(
+        m == 1 || m == 2 || m == 4 || m == 8,
+        "M must be one of 1, 2, 4, or 8"
+    );
+    TORCH_CHECK(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 % 2 == 0, "K must be even for vectorized bf16 loads");
+    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_gemv 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_gemv 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 n_int = static_cast<int>(n);
+    const int k_int = static_cast<int>(k);
+    switch (m) {
+        case 1:
+            launch_bf16_gemv<1>(
+                x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream()
+            );
+            break;
+        case 2:
+            launch_bf16_gemv<2>(
+                x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream()
+            );
+            break;
+        case 4:
+            launch_bf16_gemv<4>(
+                x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream()
+            );
+            break;
+        case 8:
+            launch_bf16_gemv<8>(
+                x_ptr, weight_ptr, bias_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_gemv",
+        &bf16_gemv,
+        py::arg("x"),
+        py::arg("weight"),
+        py::arg("bias") = py::none(),
+        "M in {1,2,4,8} BF16 GEMV with FP32 accumulation and optional fused bias"
+    );
+}

docs/developer/cuda_kernels.md

@@ -1,6 +1,9 @@
 # CUDA Kernels
 
-AstrAI includes optional custom CUDA kernels for attention, rotary embedding, and FP8 GEMM. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend, auto-dispatched for rotary, or invoked through the FP8 linear primitives.
+AstrAI includes optional custom CUDA kernels for attention, rotary embedding,
+BF16 GEMV, and FP8 GEMM. These are built when `nvcc` is available and CUDA is
+detected. BF16 GEMV is directly callable and can be selected by the guarded
+model linear dispatcher described below.
 
 ## Overview
 
@@ -11,8 +14,43 @@ AstrAI includes optional custom CUDA kernels for attention, rotary embedding, an
 | `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_gemv` | `gemv/bf16_gemv.cu` | M=1/2/4/8 BF16 linear with FP32 accumulation (sm_80+) |
 | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
 
+### BF16 GEMV primitive
+
+`astrai.extension.bf16_gemv(x, weight, bias=None)` accepts a contiguous BF16
+input shaped `[K]` or `[M, K]`, with `M` in `{1, 2, 4, 8}`, and row-major
+weights `[N, K]`. One CTA reduces each output row and computes all M results
+together, reusing the weight row across tokens. It uses vectorized
+`__nv_bfloat162` loads and FP32 accumulation; the optional BF16 bias is fused
+before the BF16 store. The launcher uses the current CUDA stream, is CUDA
+Graph capture-safe, and requires sm_80 or newer.
+
+Model `Linear` calls route through the lightweight linear backend. Set
+`ASTRAI_GEMV=0` for an unconditional `F.linear` fallback, `1` to force the
+kernel for any supported M=1/2/4/8 call, or `auto` (the default) to select only
+architecture/shape bands that pass both the per-shape and end-to-end gates.
+M=1 has no automatic SM89 band because isolated winners did not reach the 3%
+whole-graph gate. Measured SM89 small-M bands are enabled as follows:
+
+| M | Automatic `(N, K)` bands | Engine throughput |
+|---:|---|---:|
+| 2 | `(256,1536)`, `(1536,1536)`, `(100000,1536)` | +14.0% |
+| 4 | `(256,1536)`, `(1536,1536)` | +11.8% |
+
+These 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.
+
 Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
 
 | Variant | File | Optimization |
@@ -202,6 +240,7 @@ astrai/extension/
 ├── ops/
 │   ├── attention.py        # Stateless attention kernel wrappers
 │   ├── rotary.py           # Stateless rotary kernel wrapper
+│   ├── gemv.py             # Stateless BF16 GEMV primitive
 │   └── fp8.py              # Stateless FP8 primitives (custom_op)
 ├── fp8.py                  # FP8 strategy layer (fp8_autocast, recipes)
 └── backend/

docs/developer/decode_linear_benchmark.md

@@ -0,0 +1,22 @@
+# Decode linear shape benchmark
+
+`scripts/tools/benchmark_gemv.py` records the `F.linear` baseline used to decide
+whether a BF16 GEMV or small-M kernel should enter automatic inference dispatch.
+It does not change model execution or select a custom kernel.
+
+The default matrix covers the AstrAI 1B q/k/v/out projections, MLP up/gate/down,
+and LM head for `M=1,2,4,8,16,32`. Each shape runs in eager and CUDA Graph replay
+modes. Results include device-event latency samples, p50/p90/p99, estimated
+effective IO bandwidth, and CUDA kernel launches per call.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python scripts/tools/benchmark_gemv.py \
+  --output results/decode_linear.json \
+  --markdown-output results/decode_linear.md
+```
+
+Use `--shape NAME:N:K` repeatedly to override the preset and `--m-values` to
+change the decode batch sizes. Compare each GPU architecture only with its own
+baseline; do not use absolute A100-versus-L20 numbers as a dispatch criterion.
+Keep the raw JSON as the source of truth and generate tables with
+`--markdown-output` rather than transcribing measurements by hand.

scripts/tools/benchmark_gemv.py

@@ -0,0 +1,333 @@
+"""Benchmark decode-time linear shapes before enabling custom GEMV dispatch.
+
+The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
+establishes the per-architecture cuBLAS baseline that later GEMV 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()

setup.py

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

tests/extension/test_gemv.py

@@ -0,0 +1,155 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from astrai.extension import bf16_gemv, is_available
+
+GEMV_AVAILABLE = (
+    torch.cuda.is_available()
+    and is_available("bf16_gemv")
+    and torch.cuda.get_device_capability() >= (8, 0)
+)
+skip_no_gemv = pytest.mark.skipif(
+    not GEMV_AVAILABLE,
+    reason="BF16 GEMV requires a built kernel and compute capability 8.0+",
+)
+
+
+@skip_no_gemv
+@pytest.mark.parametrize(
+    "n,k",
+    [(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)],
+)
+def test_bf16_gemv_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_gemv(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_gemv
+@pytest.mark.parametrize("m", [2, 4, 8])
+@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)])
+def test_bf16_gemv_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_gemv(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_gemv
+def test_bf16_gemv_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_gemv(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_gemv
+def test_bf16_gemv_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_gemv(x, weight, bias)
+    expected = F.linear(x, weight, bias)
+    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
+
+
+@skip_no_gemv
+def test_bf16_gemv_uses_current_stream():
+    x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        actual = bf16_gemv(x, weight)
+        expected = F.linear(x, weight)
+    stream.synchronize()
+    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
+
+
+@skip_no_gemv
+def test_bf16_gemv_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_gemv(x, weight)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = bf16_gemv(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_gemv
+def test_bf16_gemv_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_gemv(x, weight)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = bf16_gemv(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_gemv
+@pytest.mark.parametrize(
+    "make_args,error",
+    [
+        (
+            lambda: (
+                torch.randn(3, 16, device="cuda", dtype=torch.bfloat16),
+                torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
+            ),
+            "M must",
+        ),
+        (
+            lambda: (
+                torch.randn(15, device="cuda", dtype=torch.bfloat16),
+                torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
+            ),
+            "even",
+        ),
+        (
+            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_gemv_rejects_unsupported_inputs(make_args, error):
+    with pytest.raises(RuntimeError, match=error):
+        bf16_gemv(*make_args())

tests/extension/test_linear_dispatch.py

@@ -0,0 +1,185 @@
+import logging
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from astrai.extension import explain, is_available, linear, op_backend
+from astrai.extension.backend import linear as public_linear
+from astrai.model.components.linear import Linear
+
+GEMV_AVAILABLE = (
+    torch.cuda.is_available()
+    and is_available("bf16_gemv")
+    and torch.cuda.get_device_capability() >= (8, 0)
+)
+skip_no_gemv = pytest.mark.skipif(
+    not GEMV_AVAILABLE,
+    reason="BF16 GEMV requires a built kernel and compute capability 8.0+",
+)
+
+
+def test_linear_backend_is_public():
+    assert linear is public_linear
+
+
+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)
+    layer = Linear(3, 4)
+    assert layer(torch.randn(2, 3)) is sentinel
+
+
+def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "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
+
+
+def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
+    monkeypatch.setenv("ASTRAI_GEMV", "invalid-test-mode")
+    with caplog.at_level(logging.WARNING):
+        trace = explain("linear", torch.randn(1, 8), torch.randn(4, 8))
+    assert "using auto" in caplog.text
+    assert "=> torch" in trace
+
+
+@skip_no_gemv
+def test_mode_zero_disables_gemv(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "0")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert "=> torch" in explain("linear", x, weight)
+        torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
+
+
+@skip_no_gemv
+def test_mode_one_forces_capable_unmeasured_shape(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "1")
+    x = torch.randn(1, 64, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert "=> gemv" in explain("linear", x, weight)
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
+        )
+
+
+@skip_no_gemv
+@pytest.mark.parametrize("m", [2, 4, 8])
+def test_mode_one_dispatches_supported_small_batches(monkeypatch, m):
+    monkeypatch.setenv("ASTRAI_GEMV", "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 "=> gemv" in explain("linear", x, weight)
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
+        )
+
+
+@skip_no_gemv
+def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "auto")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    winning = torch.randn(
+        1536,
+        1536,
+        device="cuda",
+        dtype=torch.bfloat16,
+        requires_grad=True,
+    )
+    with torch.no_grad():
+        assert "=> torch" in explain("linear", x, winning)
+        torch.testing.assert_close(linear(x, winning), F.linear(x, winning))
+
+
+@skip_no_gemv
+def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "auto")
+    x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        trace = explain("linear", x, weight)
+        if torch.cuda.get_device_capability() == (8, 9):
+            assert "=> auto_gemv" in trace
+        else:
+            assert "=> torch" in trace
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.5
+        )
+
+
+@skip_no_gemv
+@pytest.mark.parametrize(
+    "m,n,k",
+    [
+        (2, 6912, 1536),  # up/gate loses at every measured M
+        (2, 1536, 6912),  # long-K accumulation changed checkpoint greedy output
+        (4, 100000, 1536),  # LM head misses the 5% M=4 gate
+        (4, 1536, 6912),  # long-K accumulation changed checkpoint greedy output
+        (8, 256, 1536),  # remaining M=8 winners miss the 3% end-to-end gate
+    ],
+)
+def test_auto_rejects_measured_small_batch_losers(monkeypatch, m, n, k):
+    monkeypatch.setenv("ASTRAI_GEMV", "auto")
+    x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        assert "=> torch" in explain("linear", x, weight)
+
+
+@skip_no_gemv
+def test_grad_enabled_and_unsupported_multirow_always_fall_back(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "1")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(
+        256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
+    )
+    assert "=> torch" in explain("linear", x, weight)
+    with torch.no_grad():
+        multirow = x.expand(3, -1).contiguous()
+        assert "=> torch" in explain("linear", multirow, weight)
+
+
+@skip_no_gemv
+def test_explicit_gemv_selection_respects_capability(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "0")
+    x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad(), op_backend(linear="gemv"):
+        assert "=> gemv" in explain("linear", x, weight)
+        torch.testing.assert_close(
+            linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
+        )
+
+
+@skip_no_gemv
+def test_dispatched_linear_cuda_graph_replay(monkeypatch):
+    monkeypatch.setenv("ASTRAI_GEMV", "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)