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
This commit is contained in:
0z5a
2026-09-02 13:11:25 +08:00
parent 9c3ef0c2a1
commit a144d7f306
14 changed files with 1242 additions and 3 deletions
+40 -1
View File
@@ -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/
+22
View File
@@ -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.