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:
@@ -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/
|
||||
|
||||
@@ -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.
|
||||
Reference in New Issue
Block a user