ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
693 Commits 1 Branch 14 Tags
056c1382ff4682f3b5bbac78cad928d434f61caa
Commit Graph
7 Commits
Author SHA1 Message Date
ViperEkura 55ee258e95 style: fix ruff lint warnings 
- Remove unused local variable b in attention_backend.py
- Remove unused variable rank0_sd in test_broadcast_state_dict.py
- Remove unused imports across test files
 2026-08-07 14:17:48 +08:00
ViperEkura 6f49738991 feat: auto-select best available attention backend 
- Default backend resolves to highest-priority available: flash -> cuda -> torch
- attention() falls back per-call for training/fp32/unsupported head_dim
- Re-apply index_copy_ for CUDA KV cache writes (index_put_ race mitigation)
 2026-08-07 13:48:59 +08:00
ViperEkura 6f67ba8942 perf: move decode split partials to InferenceWorkspace 
- Replace per-.cu-file static cached tensors with workspace-managed pre-allocated buffers

- InferenceWorkspace now owns decode_o_part / decode_ml_part (mirrors FlashInfer's workspace pattern)

- KVCache carries the buffers through the backend -> C++ kernel chain

- C++ kernels accept optional pre-allocated buffers; fallback to alloc_split_partials for backward compat

- Pre-allocates once at Executor init, zero allocation in the decode hot loop

- Prerequisite for CUDA-graph capture (all kernel addresses are stable)
 2026-08-06 19:12:09 +08:00
ViperEkura a03504a280 perf: preallocate inference decode buffers 
- add InferenceWorkspace with fixed-shape per-step buffers (input_ids, decode mask, KV bind metadata) for CUDA-graph capture
- bind_tasks derives seq_lens from the pool's own _task_len tracking, dropping the seq_lens parameter
- update decode metadata in-place (position_ids, seq_lens, kv_indptr) instead of re-allocating per step
- task_extend advances _task_len in contiguous mode so the pool tracks current length
- skip log_softmax when logprobs are not requested
 2026-08-03 00:55:26 +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 738cb8f128 fix: broadcast ref/old model state_dict for FSDP 
- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks
- Fix create_ref_model returning None on non-rank-0 under FSDP
- Fix sync_old_model only updating old_model on rank-0 under FSDP
- Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest
- Add distributed tests for broadcast_state_dict and create_ref_model
 2026-07-31 08:32:22 +08:00
ViperEkura 3067a8e1a6 feat: unify attention backend with multi-dim mask support 
- Add attention() functional entry delegating to active backend
- GQA/MLA forward calls attention() instead of inline cache/SDPA
- CUDA kernels support 2D/3D/4D mask via mask_h_stride field
- CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens
- KVCache.max_len precomputed in bind_tasks to avoid GPU sync
- batch==1 decode short-circuits mask=None
- Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask
- 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native
 2026-07-30 20:38:34 +08:00