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

perf: vectorize rotary kernel loads and halve index math

Browse Source 
- walk exact 2-pair chunks (8B x access, 16B cos/sin float4) and decompose the flat index per chunk instead of per pair, halving integer div/mod work
- enforce head_dim % 4 == 0 at the binding instead of carrying a scalar fallback path
- raise the grid-stride block cap from 1024 to 2048 for full SM coverage on streaming shapes
- hoist kernels/rotary/rotary_emb.cu to kernels/rotary_emb.cu (single-file directory)

Benchmark: NVIDIA L20 (sm_89, shared GPU), interleaved A/B of old and new module, 500-iter means
- (32768 tokens, 8 heads, D=64): 72.7 -> 37.5 us (1.94x)
- (32768 tokens, 32 heads, D=256): 4420 -> 3630 us (1.22x)
- (32768 tokens, 32 heads, D=128): 2120 -> 1824 us (1.16x)
- (32 tokens, 32 heads, D=128) decode size: unchanged at ~1.9 us
This commit is contained in:
ViperEkura
2026-09-01 15:30:07 +08:00
parent 1c17e80882
commit aabf366633
3 changed files with 36 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/developer/cuda_kernels.md
+3 -4
View File
@@ -10,7 +10,7 @@ AstrAI includes optional custom CUDA kernels for attention, rotary embedding, an
| `attn_prefill` | `attention/prefill.cu` | GQA prefill attention (split-Q) |
| `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/rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
| `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
@@ -27,7 +27,7 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac
### Rotary Embedding Kernel
The `rotary_emb` kernel (`csrc/kernels/rotary/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
- One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store
- f32 cos/sin input, bf16 compute and output
@@ -460,8 +460,7 @@ csrc/
│ │ ├── prefill.cu # → module attn_prefill
│ │ ├── paged_decode.cu # → module attn_paged_decode
│ │ └── paged_prefill.cu # → module attn_paged_prefill
│ ├── rotary/
│ │ └── rotary_emb.cu # rotary embedding (kernel + binding in one file) → module rotary_emb
│ ├── rotary_emb.cu # rotary embedding (kernel + binding in one file) → module rotary_emb
│ └── fp8/ # FP8 family (module name fp8_ops)
│ ├── common.h # FP8Format enum, Fp8GemmTraits, FP8Params / FP8QuantizeParams PODs, layout tags (no torch)
│ ├── quantize.cuh # quantize kernels: vectorized + 32×32-tile transpose (out_layout 0/1/2) (no torch)