ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: rebuild decode gemm dispatch around shape-driven tile configs

Browse Source 
- split-K removed entirely: tiled kernel walks K in one pass, no partials/semas workspace, no memset, single launch per call
- skinny GEMM (M<=8) dispatch table replaces the hand-written switch
- shape-driven four-family table replaces plan_gemm: wide-N (n>=4096) default {16,64,64,3,128} with BM=32 at M>16; narrow-N deep-K rings {16,32,256,2,64} while the grid fits one wave, {16,32,128,2,64} past it
- narrow-N is K-serial: widening the grid measurably does nothing (BN 64->32 ties, doubled m_tiles tie, kv at 4 blocks ties q/o at 24); deeper K chunks win until 72KB smem forces one CTA per SM and past one wave the 2-wave quantization loses to BK=128
- launch-check macros in common/launch.cuh; smem opt-in for the 72KB/60KB rings
- rename kernels/bf16_*.cu to gemm.cu/swiglu.cu; module names unchanged
- Python gate: lm_head (N>32768) falls back to cuBLAS, band narrows to M<=32
- drop the stale per-op benchmark narratives; fold the live numbers into cuda_kernels.md

Benchmark: NVIDIA L20 (sm_89, 92 SMs), CUDA 12.8, bf16, L2-thrash weight rotation, per-call medians at M=16: q/o 9.5us, kv 8.6us, gate/up 33.3us, down 33.7us (down -29% vs prior default). End-to-end 1B decode (gen 128, 3 trials, tokens/s vs cuBLAS): B=1 260 vs 252, B=8 1660 vs 1446, B=16 2464 vs 2437, B=32 3620 vs 3690. Prior split-K dispatch measured B=16 2243 / B=32 3393.
This commit is contained in:
ViperEkura
2026-09-04 22:41:39 +08:00
parent 8e39d9d8c9
commit 1798474316
21 changed files with 1098 additions and 697 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/developer/cuda_kernels.md
+47 -42
View File
@@ -1,8 +1,8 @@
# CUDA Kernels
AstrAI includes optional custom CUDA kernels for attention, rotary embedding,
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
BF16 GEMM/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and
CUDA is detected. BF16 GEMM and SwiGLU are directly callable and can be
selected by guarded model dispatchers described below.
## Overview
@@ -14,49 +14,55 @@ selected by guarded model dispatchers described below.
| `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` | `bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) |
| `bf16_swiglu` | `bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
| `bf16_gemm` | `gemm.cu` | M=1..64 BF16 linear, skinny GEMM path for M<=8 and tensor-core tiled path above (sm_80+) |
| `bf16_swiglu` | `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
### BF16 GEMM 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 computes an output row for all M tokens
together, reusing the weight row across tokens. CTA size is 256 threads,
except for small weight matrices (`N*K <= 12 MiB`) at `M=8`, where a
128-thread CTA measured 5-9% faster on L20. Variant selection is otherwise
intentionally shape-free: under HBM-streaming conditions (weights rotated
through L2, as in real decode) the kernel is bandwidth-bound and block-size
choice measures within noise, so earlier per-shape variant tables were
removed along with the warp-tiled kernel.
`astrai.extension.bf16_gemm(x, weight, bias=None)` accepts a contiguous BF16
input shaped `[K]` or `[M, K]`, with `M` in `[1, 64]`, and row-major weights
`[N, K]`. One entry point selects the internal path by M:
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. 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.
- `M` in `[1, 8]` — register-resident GEMV kernel, any positive `K`: one
flat 128-thread CTA per weight row computes the output column for all M
tokens together. Under HBM-streaming conditions (weights rotated through
L2, as in real decode) the kernel is bandwidth-bound; on L20 the flat
128-thread CTA measured within ~1% of a per-shape tuned 128/256/512 mix at
M=1 and M=8 and at most ~2% below it at M=2-4, while dodging the register
cliff the larger CTAs hit on the 307MB lm_head at M>=7 (+21% DRAM
throughput vs 256t). Simplicity was chosen over the last ~2%.
- `M` in `(8, 64]` — CUTLASS-style tiled kernel (`tiled_gemm_kernel`,
fully parameterized by template parameters BM/BN/BK/stages/threads):
multistage cp.async staging with XOR-swizzled 16B chunks, no split-K —
every block walks the whole K range in one pass and the grid fills the
SMs along N. Dispatch is shape-driven: wide N (n >= 4096, grid already
full) uses the bandwidth-optimal default (BN=64, BK=64, 3 stages; BM
widens to 32 at M>16); narrow N is K-serial, where widening the grid
measurably does nothing and deeper K chunks win (BN=32, BK=256 while
the grid fits one wave, BK=128 past it). Requires `K % 8 == 0` and
16-byte-aligned tensors.
Both paths accumulate in FP32, fuse the optional BF16 bias before the BF16
store, use the current CUDA stream, are CUDA Graph capture-safe, and require
sm_80 or newer. The GEMV path anchors 128-bit weight loads 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.
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). Automatic
dispatch is keyed on the decode batch size alone: the kernel streams each
weight exactly once, so once a batch size is profitable it is profitable
across projection shapes. On compute capability 8.0+, `auto` selects the
kernel for `M` in [2, 4], where every measured model family beat the cuBLAS
small-M path at the HBM bandwidth floor (AstrAI 1B chain +11.8% to +14.0%,
common LLaMA/Qwen/OPT chains +5.66% to +25.20%). `M=1` keeps cuBLAS, whose
GEMV path is already at the floor, and `M >= 5` approaches the cuBLAS
tensor-core crossover (M=8 regressed at wrapper level in every measured
family). Out-of-band, training, prefill-sized, or unsupported calls fall
back to PyTorch.
`ASTRAI_GEMM=0` for an unconditional `F.linear` fallback, `1` to force the
kernel for any capable M (1-64, with K%8==0 and alignment for M>8), or
`auto` (the default). On compute capability 8.0+, `auto` routes all capable
decode batches M in [1, 64] through the kernel. The GEMV path (M≤8, any K)
measured universal wins at the HBM bandwidth floor (AstrAI 1B end-to-end
+8.8% at M=1 rising to +17% at M=8). The tiled path (M 9-64) measured on
L20 across five shape families wins or ties four (q/o, kv, gate/up, lm_head)
and regresses one (down-projection M>16: +25% latency; narrow-N long-K where
cuBLAS is strong). Out-of-band, training, prefill-sized, or unsupported
calls fall back to PyTorch.
Inside the primitive, a templated cooperative kernel uses 256 threads,
except at the largest decode batch where a 128-thread CTA wins 5-9% on
small weight matrices: `M=8` with 16-byte-aligned inputs, `K % 8 == 0`,
and `N*K <= 12 MiB` selects the smaller CTA. This internal selector is
separate from model automatic dispatch, whose Python/wrapper overhead is
included in the gates above.
The measurements below date from the original `[2, 4]` automatic band; the
band has since widened to `{1, ..., 8}` with the flat 128-thread kernel:
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
@@ -84,7 +90,7 @@ 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 `csrc/bench/benchmark_gemv_common.py`.
claims. Reproduce them with `csrc/bench/benchmark_gemm_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
@@ -121,8 +127,7 @@ the unfused linear backend, and `1` explicitly forces the fused primitive.
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.
existing path.
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
@@ -313,7 +318,7 @@ astrai/extension/
├── ops/
│ ├── attention.py # Stateless attention kernel wrappers
│ ├── rotary.py # Stateless rotary kernel wrapper
│ ├── gemv.py # Stateless BF16 GEMV primitive
│ ├── gemm.py # Stateless BF16 GEMM primitive
│ ├── swiglu.py # Stateless fused BF16 SwiGLU primitive
│ └── fp8.py # Stateless FP8 primitives (custom_op)
├── fp8.py # FP8 strategy layer (fp8_autocast, recipes)
docs/developer/decode_linear_benchmark.md
-42
View File
@@ -1,42 +0,0 @@
# Decode linear shape benchmark
`csrc/bench/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 csrc/bench/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.
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 csrc/bench/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.
Automatic dispatch is keyed on the decode batch size alone (`M` in `[2, 4]` on
compute capability 8.0+); use `--candidate-mode 1` to characterize a family
before widening that band. 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
-83
View File
@@ -1,83 +0,0 @@
# Fused SwiGLU benchmark
`csrc/bench/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 csrc/bench/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.
## HBM re-measurement and kernel simplification
The operator numbers above are L2-resident: the AstrAI pair is 40.5 MB,
smaller than the 96 MB L2, so a tight timing loop re-reads warm weights
(13.75 us implies ~3.1 TB/s, far above the 864 GB/s spec). Real decode rotates
~1 GB of per-layer weights through L2 every step, so every call is cold.
Re-measuring with rotated weight copies (>= 240 MB working set) on the same
L20 showed:
- The fused CTA-reuse kernel sits at the dual-stream cold-read floor
(702 vs 699 GB/s at (6912,1536); 369 vs 370 GB/s at (11008,4096)). Wide
LLaMA matrices cap at ~370-400 GB/s regardless of kernel, even for a
pure-read loop, so the old per-variant gaps there were noise.
- The `(6912,1536)` warp-per-row variant (formerly M=2/4/8) is 2-6% slower
than CTA reuse at M=2/4 under cold weights and no longer wins at M=8 once
the CTA drops to 128 threads. It and its dispatch table were deleted.
- New rule: 256 threads for M in [1, 7], 128 threads for M=8. End-to-end
through the built module at (6912,1536): 738-752 GB/s for M in [1, 4] and
702 GB/s at M=8 (+6% over the removed warp path).
The M=8 CUDA-Graph regression reported above (`-23.24%`) does not survive the
cold-weight regime: cuBLAS reaches L2 bandwidth in the warm loop while both
fused paths converge to the same HBM floor.