perf: tune BF16 GEMV for common transformer shapes

[0z5a]

Summary

• retain the barrier-free four-warps-per-CTA M=4 kernel for its four measured LLaMA bands, and add a templated 128/256-thread cooperative GEMV with an exact (M,N,K) selector
• extend the traditional-shape portfolio from LLaMA/GPT-NeoX to Qwen2-7B, LLaMA 3 70B, and OPT-1.3B; add focused shape/chain filters and a forced-candidate mode for pre-dispatch characterization
• expand SM89 auto only where every projection clears the 5% per-shape gate and the dependent chain clears 3%; all other M/family combinations fail closed to the existing kernel or PyTorch
• check in raw paired/interleaved L20 evidence, CUDA correctness/dispatch tests, documentation, and an updated explicit AstrAI InfraSWE D3 Draft

L20 benchmark

Environment: NVIDIA L20 (SM89), PyTorch 2.11.0+cu128, CUDA 12.8. GPU5 retained the co-resident AstrAI service allocation throughout; it was not stopped or modified. The large LLaMA 3 70B chain started only after verifying 0% instantaneous utilization and 2,477 MiB free. CUDA compilation completed before timed cases.

CTA selection

The production kernel now instantiates 128- and 256-thread cooperative CTAs for M=1..8. All specializations compile with zero spills (26–52 registers, one barrier); the existing M=4 warp-tiled kernel remains at 42 registers, zero barriers, zero spills.

An identical-input, separately named, interleaved module comparison against the original 256-thread kernel used 21 samples × 200 inner iterations over 22 common shapes. The final selector retains only cells above the 5% gate:

M	selected direct-kernel gain vs 256-thread baseline
1	+16.68%
2	+5.37% to +30.10%
4	+11.04% to +36.84%
8	+8.67% to +48.54%

Long-K, saturated, and sub-threshold cells retain the 256-thread fallback. All 88 comparison rows preserve row argmax.

Final primitive vs F.linear

For the bands considered for new automatic dispatch (21 paired/interleaved samples):

family	M	per-projection speedup range	row argmax
Qwen2-7B	2	+6.11% to +117.80%	exact
Qwen2-7B	4	+6.17% to +120.43%	exact
LLaMA 3 70B	2	+5.56% to +83.44%	exact
LLaMA 3 70B	4	+6.15% to +48.22%	exact
OPT-1.3B	1	+24.00% to +62.20%	exact
OPT-1.3B	2	+38.47% to +194.59%	exact

Guarded dependent chains

Final ASTRAI_GEMV=auto versus 0, using distinct weights and including Python dispatch plus dependent MLP work:

chain	M=1	M=2	M=4	M=8	enabled row argmax
Qwen2-7B	fallback	+7.48%	+7.48%	fallback	exact
LLaMA 3 70B	fallback	+7.77%	+7.69%	fallback	exact
OPT-1.3B	+4.54%	+25.20%	fallback	fallback	exact

M=8 remains disabled even where the aggregate chain is positive because at least one projection misses the per-shape gate (Qwen Q/O, LLaMA 3 70B K/V, and OPT up). Existing LLaMA/GPT-NeoX policy remains unchanged, including the deliberate LLaMA 2 7B M=4 fallback.

These are synthetic projection-chain measurements, not whole-model throughput claims. Raw environment, parameters, median/p90 timings, numerical error, and argmax diagnostics are under docs/benchmarks/.

Validation

• local full suite: 625 passed, 185 skipped
• L20 built-extension GEMV/dispatch suite: 140 passed
• final SM89 bf16_gemv build: successful; all cooperative specializations zero-spill
• ruff format --check .: passed
• ruff check . --select I: passed
• InfraSWE Draft v2 validate/resolve: passed; Draft engine 53 passed
• frozen-formula diagnostic ProjectFit 92.64, BenchmarkTrust 95.87

The InfraSWE numbers are explicitly diagnostic and non-official. Official ProjectFit remains unresolved because the Draft is unsealed and fresh-process/system-trace/hidden-probe/evidence-manifest gates are incomplete.

[ViperEkura]

Landed on main as d4a292b (squashed with #54)