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.

docs/developer/cuda_kernels.md

@@ -1,9 +1,9 @@
 # CUDA Kernels
 
 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.
+BF16 GEMV/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and
+CUDA is detected. BF16 GEMV and SwiGLU are directly callable and can be
+selected by guarded model dispatchers described below.
 
 ## Overview
 
@@ -15,35 +15,96 @@ model linear dispatcher described below.
 | `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..8 BF16 linear with FP32 accumulation (sm_80+) |
+| `bf16_swiglu` | `gemv/bf16_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 GEMV primitive
 
 `astrai.extension.bf16_gemv(x, weight, bias=None)` accepts a contiguous BF16
 input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and
-row-major weights `[N, K]`. One CTA reduces each output row and computes all M
-results together, reusing the weight row across tokens. The weight stream uses
-128-bit vectorized loads anchored 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; x loads are vectorized when every row base is
-16-byte aligned (always true for K % 8 == 0 with allocator-aligned tensors)
-and scalar otherwise. Accumulation is FP32; 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.
+row-major weights `[N, K]`. The general path assigns one 256-thread CTA to an
+output row and computes all M results together, reusing the weight row across
+tokens. For measured aligned M=4 medium projections, a 128-thread CTA instead
+assigns one output to each of four warps. That removes the CTA-wide reduction
+barrier and exposes four neighboring outputs without changing accumulation.
+
+The weight stream uses 128-bit vectorized loads anchored 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. The warp-tiled path is used
+only when both tensors and every row are 16-byte aligned; all other calls keep
+the general arbitrary-K path. Accumulation is FP32; 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 in [1, 8], 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:
+Measured SM89 small-M bands are enabled as follows:
 
-| M | Automatic `(N, K)` bands | Engine throughput |
+| M | Automatic `(N, K)` bands | Validated gain |
 |---:|---|---:|
-| 2 | `(256,1536)`, `(1536,1536)`, `(100000,1536)` | +14.0% |
-| 4 | `(256,1536)`, `(1536,1536)` | +11.8% |
+| 1 | OPT-1.3B Q/K/V/O and MLP | +4.54% OPT projection chain |
+| 2 | AstrAI `(256,1536)`, `(1536,1536)`, `(100000,1536)` plus all common shapes below | +14.0% on AstrAI 1B; +5.66% to +25.20% common chains |
+| 4 | AstrAI `(256,1536)`, `(1536,1536)` plus gated common shapes below | +11.8% on AstrAI 1B; +5.67% to +7.71% common chains |
+| 8 | none | at least one projection in every measured family missed the per-shape gate |
 
-These A→B→B→A results use the real `InferenceEngine`, including scheduler,
+The common set covers LLaMA 2 7B Q/O, gate/up, and down; LLaMA 3 8B K/V,
+gate/up, and down; LLaMA 2 13B Q/K/V/O, gate/up, and down; and GPT-NeoX MLP
+up/down. In `(N,K)` form it is `(1024,4096)`, `(4096,4096)`,
+`(11008,4096)`, `(4096,11008)`, `(14336,4096)`, `(4096,14336)`,
+`(5120,5120)`, `(13824,5120)`, `(5120,13824)`, `(16384,4096)`, and
+`(4096,16384)`. M=2 enables all eleven. M=4 excludes the three LLaMA 2 7B
+bands `(4096,4096)`, `(11008,4096)`, and `(4096,11008)` because their combined
+projection chain reached only +1.89%, below the 3% automatic-dispatch gate.
+
+The extended common set adds Qwen2-7B `(512,3584)`, `(3584,3584)`,
+`(18944,3584)`, and `(3584,18944)`; LLaMA 3 70B `(1024,8192)`,
+`(8192,8192)`, `(28672,8192)`, and `(8192,28672)`; and OPT-1.3B
+`(2048,2048)`, `(8192,2048)`, and `(2048,8192)`. Qwen2 and LLaMA 3 70B are
+enabled at M=2/4. OPT-1.3B is enabled at M=1/2. Other rows retain their
+previous policy or fall back to PyTorch.
+
+Inside the primitive, a templated cooperative kernel uses either 256 threads
+or a shape-gated 128-thread CTA. The smaller CTA is enabled only where an
+interleaved direct-module comparison against the original 256-thread kernel
+cleared 5%: OPT up at M=1; selected LLaMA 2 7B, Qwen2, and OPT projections at
+M=2; LLaMA 2 13B Q/O, Qwen2 Q/O, and selected OPT projections at M=4; and
+selected LLaMA 2, Qwen2, LLaMA 3 KV, and OPT projections at M=8. Confirmed
+direct-kernel gains range from +5.37% to +48.54%. Long-K and saturated shapes
+keep the 256-thread fallback. This internal selector is separate from model
+automatic dispatch, whose Python/wrapper overhead is included in the gates
+above.
+
+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 `scripts/tools/benchmark_gemv_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
@@ -55,6 +116,30 @@ 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 contains two output-row tilings. A CTA-reuse path reads each up/gate
+weight chunk once and applies it to all M rows. The native AstrAI 1B shape
+`(N,K)=(6912,1536)` uses one warp per decode row for M=2/4/8; on L20 this
+removes the shared reductions and barrier and reduces M=4 CUDA-Graph latency
+from 0.0324 ms to 0.0181 ms. Wider LLaMA/GPT-NeoX matrices keep CTA reuse,
+because duplicating their weight reads across row warps regressed 1.3-4.2%.
+
+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. See [the benchmark protocol](./swiglu_benchmark.md) for raw
+operator, engine, and checkpoint evidence.
+
 Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
 
 | Variant | File | Optimization |
@@ -245,10 +330,12 @@ astrai/extension/
 │   ├── attention.py        # Stateless attention kernel wrappers
 │   ├── rotary.py           # Stateless rotary kernel wrapper
 │   ├── gemv.py             # Stateless BF16 GEMV 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
 ```
 

docs/developer/decode_linear_benchmark.md

@@ -20,3 +20,22 @@ 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.
+
+For direct A/B coverage of the custom kernel and guarded dispatcher across
+traditional LLaMA and GPT-NeoX decode shapes, use:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python scripts/tools/benchmark_gemv_common.py \
+  --suite all --family traditional --m 2 4 \
+  --output results/gemv_common.json
+```
+
+The kernel suite compares the directly callable primitive with `F.linear`.
+Use repeatable `--shape-label` and `--chain-label` filters for a focused run.
+The synthetic-chain suite alternates `ASTRAI_GEMV=0` and `auto`, includes
+dependent MLP work and Python dispatch, and rotates through distinct weights.
+Pass `--candidate-mode 1` to characterize a family before adding it to the
+automatic shape table; the checked-in final evidence always uses `auto`.
+It is deliberately not labeled a whole-model throughput benchmark. Both
+suites report median/p90 CUDA-event latency plus maximum absolute error,
+relative L2 error, and row-wise argmax parity.

docs/developer/swiglu_benchmark.md

@@ -0,0 +1,58 @@
+# Fused SwiGLU benchmark
+
+`scripts/tools/benchmark_swiglu.py` compares the directly callable fused BF16
+SwiGLU primitive with both `F.linear` and the existing two-GEMV chain. It covers
+the native AstrAI 1B MLP plus LLaMA 2 7B/13B, LLaMA 3 8B, and GPT-NeoX 20B
+up/gate shapes at M=1/2/4/8 in eager and CUDA Graph modes.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python scripts/tools/benchmark_swiglu.py \
+  --output results/swiglu.json \
+  --markdown-output results/swiglu.md \
+  --m-values 1,2,4,8 --mode both \
+  --warmup 20 --iterations 100 --trials 10
+```
+
+Each trial uses A-B-C-C-B-A ordering to balance clock, cache, and temperature
+drift. The generated JSON records every timing sample, p50/p90/p99, CUDA launch
+count, maximum/mean absolute error, and cosine similarity.
+
+## L20 findings
+
+Hardware was one NVIDIA L20 (sm_89), PyTorch 2.11.0+cu128, CUDA 12.8. The
+existing GPU5 inference service remained resident (15.4 GiB) but idle at the
+sampling boundaries; no process or container was stopped.
+
+For AstrAI 1B `(N,K)=(6912,1536)`, CUDA Graph medians were:
+
+| M | torch (ms) | GEMV chain (ms) | fused (ms) | vs best unfused |
+|---:|---:|---:|---:|---:|
+| 1 | 0.02564 | 0.02298 | 0.01375 | +67.13% |
+| 2 | 0.02484 | 0.02628 | 0.01416 | +75.40% |
+| 4 | 0.02507 | 0.03839 | 0.01806 | +38.82% |
+| 8 | 0.02563 | 0.07007 | 0.03339 | -23.24% |
+
+The wide traditional shapes are weight-bandwidth dominated. CTA reuse keeps
+the fused primitive within roughly -1.2% to +0.9% of the best unfused chain,
+so none is eligible for automatic selection. This negative crossover is kept
+in the raw evidence rather than hidden by a favorable subset.
+
+The real 24-layer AstrAI checkpoint was then run through `InferenceEngine`,
+including scheduler, sampling, and CUDA Graph. A-B-B-A medians were:
+
+| Batch | unfused (ms/step) | forced fused (ms/step) | throughput gain |
+|---:|---:|---:|---:|
+| 1 | 4.125 | 3.925 | +5.10% |
+| 2 | 4.245 | 4.055 | +4.69% |
+| 4 | 4.475 | 4.305 | +3.95% |
+
+## Dispatch decision
+
+Direct correctness stayed close (`max_abs <= 2.4e-4`, cosine approximately
+1.0), but deterministic greedy generations changed at M=1, M=2, and M=4.
+
+For that reason no SM89 shape is enabled in `auto`. The default path stays on
+the existing unfused linear backend, including any independently qualified
+GEMV dispatch. `ASTRAI_SWIGLU=1` remains an explicit benchmark/experimentation
+switch for callers that accept normal BF16 reduction-order variation. A future
+automatic band must repeat both the performance and checkpoint-output gates.