Files
AstrAI/docs/developer/decode_linear_benchmark.md
T
0z5a a144d7f306 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
2026-09-02 13:11:25 +08:00

1.1 KiB

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.

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.