ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
779 Commits 1 Branch 14 Tags
057c0d33df2a2bbbf01e0142804fa05047a879a9
Commit Graph
9 Commits
Author SHA1 Message Date
ViperEkura 2bc4d2b8a8 refactor: unify kernel module loading and packaging 
- loader.py: lazy/cached import; is_available defers the actual load; get_module raises on unavailable
- ops/{attention,rotary,fp8}: use get_module instead of touching private _modules or their own _mod() cache
- package-data: ship astrai.extension.lib *.so in built wheels (non-editable installs previously lost every kernel)
 2026-08-23 15:57:04 +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 a5b238dd86 feat: add fp8 training via cublasLt dispatch 
- fp8_mm kernel (csrc): cublasLt fp8 e4m3 gemm, TN layout mapped zero-copy
- custom::fp8_mm custom op: meta/cuda/cpu kernels + scale-corrected bf16 autograd
- aten::linear and linear_backward dispatch on CUDA key, zero model changes
- per-tensor scale or raw cast; single-GPU smoke loss matches bf16
 2026-08-14 00:39:49 +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 9f83d982ec refactor: move compiled kernel .so files into extension/lib 
- CUDAExtension module names changed from astrai.extension.<name> to astrai.extension.lib.<name>
- Compiled .so files now land in astrai/extension/lib/ instead of alongside Python source
- loader.py imports from .lib.<name> subpackage
- Add astrai/extension/lib/__init__.py to make lib a proper package
- Separates compiled artifacts from Python source for cleaner directory structure
 2026-07-31 15:36:32 +08:00
ViperEkura 3e67b4f88d perf: add fused CUDA rotary embedding kernel 
- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call)
- RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor
- apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit)
- Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out
- Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16)
- Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)
 2026-07-31 15:27:31 +08:00
ViperEkura 2c3cef1c87 feat: wire up paged decode CUDA kernel to Python extension 
- Add attn_paged_decode wrapper in ops.py with gather fallback
- Register kernel in loader.py and export from __init__.py
- Extract test_utils.cuh shared by all attention unit tests
- Rename attn_paged_vs_contiguous.cu to attn_paged_decode_test.cu
- Refactor decode/prefill tests to use common bf16 helpers and cpu ref
- Fix k_cache dim check in attn_paged_decode.cu
 2026-07-11 18:40:49 +08:00
ViperEkura d923ebe38d refactor: rename gqa_* to attn_*, split-KV for all decode paths 
- Rename all csrc/kernels/gqa_*.cuh/cu to attn_*, with _split_q / _split_kv
  strategy suffix and optional _mma compute suffix
- Remove non-split MMA decode kernel, keep only split-KV path
- Convert scalar decode fallback to split-KV (o_part/ml_part + combine)
- Move combine kernel to attn_decode_split_kv.cuh (shared by both paths)
- Rename GQAParams to AttentionParams
- Update all C++ #include, PYBIND11, and Python extension references
 2026-07-10 23:35:14 +08:00
ViperEkura fd65b9bc23 feat: support HEAD_DIM=32 and split extension into loader/ops 
- add case 32 to decode/prefill dispatch switch
- fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory
- restructure decode dispatch to #ifndef/#else/#endif matching prefill
- split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API
 2026-07-08 14:14:11 +08:00