ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
718 Commits 1 Branch 14 Tags
a5a3cc1fc2f068efa998872f45ed343d0f4deeb5
Commit Graph
12 Commits
Author SHA1 Message Date
ViperEkura d565d44c43 fix: harden attention kernel boundaries 
- fix scalar prefill head_dim=32 out-of-bounds via G=4 dispatch
- fix MMA decode 4D mask head indexing and invalid-row mask access
- add q_head/kv_head divisibility and head-dim contiguity checks
- validate split-KV scratch and decode out_buf layout in bindings
- set max dynamic shared memory for scalar decode D=256
- cover scalar prefill D=32 in pure C test
 2026-08-09 14:53:24 +08:00
ViperEkura d0c5debbab perf: preload V in decode split-kv shared mem and cache partial tensors 
- Preload V into shared memory alongside K to eliminate per-element KV address lookups in the inner softmax/accum loop (doubles smem)
- Cache split-KV partial tensors (o_part, ml_part) with static tensors instead of per-call allocation in both decode and paged-decode paths
- Force is_causal=True in CUDA decode backend (decode is always causal)
 2026-08-06 18:27:00 +08:00
ViperEkura 2667b8116d refactor: unify paged and contiguous attention kernels via KVSource policy 
- merge AttentionParams and PagedAttentionParams into one struct
- add attn_kv_source.cuh with ContigKV/PagedKV addressing policies
- template prefill/decode kernels (MMA + scalar) on the KV policy, deleting the four duplicated attn_paged_*.cuh variants
- template dispatcher launchers on KV; single combine kernel
- verify: all correctness tests pass and SASS matches baseline (no perf regression)
 2026-08-05 14:06:13 +08:00
ViperEkura b1b65a657e perf: target 512 grid blocks for decode split-K 
- compute_num_splits used 2*sm/base, undersplitting at large batch
- single-warp decode blocks host ~11/SM, not 1/2-SM, so B=16 got 3 splits when 8 was optimal
- Grid search on L20: bandwidth saturates near 256-512 total blocks; target 512
- Pass num_passes into base_blocks for the non-paged decode to match the paged path
- B=16 kv=2048: 0.0230->0.0157ms (-32%); paged B=16: 0.0527->0.0243ms (-54%); B=32: 0.0406->0.0241ms (-41%)
 2026-08-02 16:10:40 +08:00
ViperEkura 3439e3104e perf: launch CUDA kernels on torch's current stream 
- Thread a cudaStream_t through attn dispatchers onto torch's current stream
- Scope the device guard to the entry function so kernels run on tensor device
- DISPATCH_HEAD_DIM now forwards varargs so stream reaches each dispatch
- Parallelize CPU reference kernels with OpenMP (paged test 31s -> 7s)
- Merge decode/prefill standalone tests into attn_test.cu with correctness tables
- Drop bench error column (CPU ref too slow at large sizes)
- Update cuda_kernels.md for the merged test layout
 2026-08-02 13:20:14 +08:00
ViperEkura 925cbedc93 feat: scalar paged prefill fallback and decode causal fix 
- Add scalar paged prefill kernel mirroring split-Q MMA indexing for sm<80
- Wire scalar path into dispatch_paged_prefill under ASTRAI_NO_MMA
- Fix paged decode scalar causal mask dropping all kv>0 for decode
 2026-08-01 16:52:01 +08:00
ViperEkura 41dcf0feb9 feat: SGLang-style paged attention kernels replace page-table path 
- PagedAttentionParams uses flat KV pool + req_to_token + kv_indptr/qo_indptr instead of page_table
- MMA split-KV decode and split-Q prefill kernels with indirect ragged-batch addressing
- Prefill kernel accepts 4D mask (causal-aware); decode kernel supports 2D mask
- CudaBackend is inference-only: kv_cache=None raises, no torch fallback
- benchmark.py: required --ckpt, --backend/--compare options
- Parallel build isolates build-temp/build-lib per subprocess
- Standalone test covers decode/prefill with mask, 27 cases pass
 2026-08-01 15:41:25 +08:00
ViperEkura 7feeb0b93e refactor: replace magic layout ints with TensorLayout enum 
- Add TensorLayout enum (C++ + Python) to replace magic layout ints
- Add C10_CUDA_CHECK post-launch error checking to all kernel entries
- Add CUDAGuard + freqs_cis shape validation to rotary_emb.cu
- Cache SM count to eliminate per-call cudaDeviceGetAttribute
- Add DISPATCH_CAUSAL_MASK macro to deduplicate dispatcher if/else
- Convert mask type hints from X|None to Optional[X]
 2026-08-01 11:05:52 +08:00
ViperEkura 8055027df7 perf: enable paged MMA kernel for page_size=1 
- Replace per-tile page lookup with per-element lookup in load_tile
- Remove page_ok gate and scalar fallback in launch_paged_decode_mma
- Unified path works for any page_size (L1-cached when page_size >= BC)
- HBM BW: 12% → 73%, decode throughput: 2,250 → 2,606 tok/s (B=32)
- Scales to 5,232 tok/s at B=128 (2.54x vs torch native)
 2026-07-30 22:06:41 +08:00
ViperEkura 2e29ed45d3 perf: shrink decode tile to BC=16 for higher occupancy 
- BC=32→16 halves smem (32KB→16KB for D=128), doubling blocks/SM (3→6)
- D=256 now fits STAGES=2 double-buffer in 32KB, eliminating 176-byte spill
- min_tiles_per_split=2 avoids excessive split overhead on small kv
- paged decode: require page_size multiple of BC so tiles stay page-aligned

Benchmark (L20 sm_89, D=128):
- B=1 kv=4096: 0.0134→0.0122ms (+9% BW)
- B=16 kv=2048: 0.0434→0.0352ms (+23% BW)
- B=32 kv=1024: 0.0343→0.0282ms (+22% BW)
 2026-07-27 22:44:02 +08:00
ViperEkura 20041d7aa9 perf: extend MMA decode to arbitrary GQA ratio, add launch bounds, vectorize combine 
- Multi-pass MMA: encode pass in grid blockIdx.x, compute q_head0/G in-kernel
- Fixes crash for G>32 (previously block(32,G) exceeded 1024 threads)
- Fixes alloc_split_partials using uninitialized num_splits (MAX_SPLITS=32)
- __launch_bounds__ on all MMA and prefill kernels for better register allocation
- 4x vectorized combine kernel (4 head_dim per thread)
- uint4 vectorized K loads in scalar decode kernels
- cp.async .L2::128B cache hint for K/V tile streaming
- Extract warp_reduce_sum, bf16, MAX_SPLITS to attn_warp_utils.cuh
 2026-07-27 00:35:34 +08:00
ViperEkura f7a16efc9d refactor: extract shared dispatcher header, unify MMA/scalar dispatch format 
- Merge 3 duplicated dispatch blocks into single attn_dispatchers.cuh
- Merge compute_num_splits from attn_utils.cuh into dispatcher header
- All dim3 grid/block declarations and <<<>>> launches are single-line
- Production .cu files (35-42 loc) only handle torch wrapping + pybind11
- Test files include dispatcher header directly, removing all #ifndef ASTRAI_NO_MMA duplication
 2026-07-21 22:21:39 +08:00