ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
811 Commits 1 Branch 14 Tags
main
Commit Graph
10 Commits
Author SHA1 Message Date
ViperEkura f6db546578 fix: zero-init AttentionParams in pure C tests 
- paged decode test left new_k_ptr/new_v_ptr as stack garbage; PagedKV::decode_addr then took the new-KV path on wild pointers (illegal access or wrong last-token K/V)
- value-init the POD (= {}) at every construction site
 2026-08-24 14:56:19 +08:00
ViperEkura 31ca357c61 refactor: namespace csrc kernels and extract common helpers 
- attention family -> astrai::attention; fp8 family -> astrai::fp8
- new common/reduce.cuh (warp/group reductions, atomic_max_float)
- new common/cp_async.cuh (predicated cp_async_16, commit/wait group)
- move MAX_SPLITS into attention/common.h; delete warp_utils.cuh
- .cu bindings and pure C tests open family namespaces via using
 2026-08-24 14:49:55 +08:00
ViperEkura 16a55bb474 refactor: reorganize CUDA kernels into per-family directories 
- move attention kernels to csrc/kernels/attention/ and rotary to rotary/
- add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks)
- split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu
- fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads
- update extension ops, loader, and kernel tests
 2026-08-22 20:40:31 +08:00
ViperEkura b5afe3d7a4 perf: optimize small-head causal prefill 
- map D=32 and D=64 causal prefill to BC=64 tiles

- add small-head correctness and benchmark coverage
 2026-08-14 23:25:49 +08:00
ViperEkura cd31f1f62f refactor: tidy attention params and launcher interfaces 
- rename output pointer field o to o_ptr for consistency with q_ptr/k_ptr/v_ptr
- regroup AttentionParams fields by responsibility and fix misleading comments
- drop unused max_seq_len/total_q fields and paged decode max_seq_len arg
- drop redundant group_size param from decode launchers (computed from p)
 2026-08-09 20:52:06 +08:00
ViperEkura a5a3cc1fc2 refactor: unify attention param field names 
- rename q_stride_* to q_*_stride to match mask stride convention
- rename mask_q_stride to mask_l_stride for consistent l-dim naming
- merge k/v and k_cache/v_cache into k_ptr/v_ptr; rename q to q_ptr
- KVSource policy selects contiguous vs paged mode at compile time
 2026-08-09 20:23:58 +08:00
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 d28552f878 refactor: use C++17 struct dispatch in csrc tests, tighten paged tolerances to 0.01 
- Replace C++20 explicit lambda template parameters with file-scope structs (DecodeDispatch/PrefillDispatch etc.)
- Remove unused gs variable in run_decode_test
- Tighten paged test atol/rtol from 0.02 to 0.01 to match contiguous tests
 2026-08-09 10:23:12 +08:00
ViperEkura 49a9c6b3d2 build: migrate CUDA kernel build to CMake 
Replace torch CUDAExtension/ParallelBuildExtension with a CMake-based build. Each kernel compiles as an independent pybind11 module in parallel via cmake --build -j, outputting to astrai/extension/lib.

- Add csrc/CMakeLists.txt (5 kernel targets, torch/pybind11 linking)
- setup.py: _CMakeBuildExt invokes cmake; auto-detect CUDA arch via torch
- Remove csrc/build.py (REGISTRY/build flags now in CMakeLists)
- Fix rel-err eps in attn_test.cu (1e-8 -> 1e-4, bf16 scale)
- Update docs/developer/cuda_kernels.md build section
- .gitignore: allow csrc/CMakeLists.txt
 2026-08-04 21:27:22 +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