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

refactor: accept arbitrary K in bf16 gemv with aligned head-tail sweeps

Browse Source 
- Drop the K % 2 entry rejection and the per-K if/else load-width branch: the weight stream now anchors uint4 loads at each row's first 16-byte-aligned address, with scalar head/tail sweeps covering at most 14 remainder elements, so any positive K and any storage offset is correct
- Keep one pure-uint4 loop (no branching inside the loop) for the production case where every x row base is 16-byte aligned (K % 8 == 0 with allocator-aligned tensors) and a scalar-x pairing loop only for unaligned K, where per-row uint4 loads are not addressable; measured cost of scalar x everywhere was up to 2.5x on multi-row shapes (down M=4 28.4us vs 11.3us)
- Remove the now-obsolete k_aligned axis and K divisibility gate from the linear dispatch spec since the primitive no longer rejects any K
- Add test coverage for unaligned K (7, 12, 100, 1534) at M=1 and M=3

Benchmark: 8x L20 (sm_89, CUDA 12.8), L2-resident microbench, 300 iters; hot path unchanged within noise vs the pure-uint4 kernel (q M=2 5.8us, down M=4 11.3us, lm M=1 391us); full gate green
This commit is contained in:
ViperEkura
2026-09-02 14:48:01 +08:00
committed by 0z5a
parent 1c3515714f
commit 800981d85a
4 changed files with 90 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/developer/cuda_kernels.md
+10 -6
View File
@@ -14,22 +14,26 @@ model linear dispatcher 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` | `gemv/bf16_gemv.cu` | M=1/2/4/8 BF16 linear with FP32 accumulation (sm_80+) |
| `bf16_gemv` | `gemv/bf16_gemv.cu` | M=1..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
input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and
row-major weights `[N, K]`. One CTA reduces each output row and computes all M
results together, reusing the weight row across tokens. 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; x loads are vectorized when every row base is
16-byte aligned (always true for K % 8 == 0 with allocator-aligned tensors)
and scalar otherwise. 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.
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
kernel for any supported M in [1, 8], 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: