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/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)