perf: tune bf16 gemv and add opt-in fused swiglu

- deepen common-shape BF16 GEMV tuning with warp-row tiling for LLaMA/Qwen2/GPT-NeoX/OPT decode projections
- add fused BF16 up/gate SwiGLU CUDA primitive with ASTRAI_SWIGLU=0/1/auto dispatch
- keep the unfused linear backend as the default path; auto enables no shape until per-architecture checkpoint gates pass
- fall back to the linear/torch chain when kernels are absent, on CPU, in training, or outside supported M/K/dtype shapes
- add gemv/swiglu benchmark scripts, dispatch and parity tests, and kernel documentation

Benchmark: NVIDIA L20 (sm_89), CUDA 12.8, PyTorch 2.11.0+cu128, idle GPU. AstrAI 1B config (24 layers, hidden 1536, vocab 100000), BF16, prompt 128, 32 greedy decode tokens, CUDA graphs enabled, A/B in separate interleaved processes (3 rounds, 8 trials each, medians). Default vs ASTRAI_SWIGLU=1 per generate call: batch 1 134.8->129.1 ms (+4.44%), batch 2 136.2->130.9 ms (+4.06%), batch 4 145.5->140.3 ms (+3.66%). Greedy output identical at batch 1, differs at batch 2/4, so auto stays unfused by default; kernelless fallback verified bit-identical greedy.

astrai/extension/ops/__init__.py

@@ -9,6 +9,7 @@ from astrai.extension.ops.attention import (
 )
 from astrai.extension.ops.gemv import bf16_gemv
 from astrai.extension.ops.rotary import rotary_emb
+from astrai.extension.ops.swiglu import bf16_swiglu
 
 __all__ = [
     "TensorLayout",
@@ -17,5 +18,6 @@ __all__ = [
     "attn_paged_prefill",
     "attn_prefill",
     "bf16_gemv",
+    "bf16_swiglu",
     "rotary_emb",
 ]

astrai/extension/ops/swiglu.py

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