ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
736 Commits 1 Branch 14 Tags
3fb4b8ab13fbf7fc900631d3fa6730ddf43c9577
Commit Graph
8 Commits
Author SHA1 Message Date
ViperEkura 3fb4b8ab13 perf: use int32 paged KV indices 
- store page-table, request-row, and cache-location indices as int32
- preserve CUDA graph replay with bit-exact logits and KV cache coverage
- improve B=1 decode latency by 1-6% across 1K-32K contexts on L20
 2026-08-15 13:17:06 +08:00
ViperEkura c5fba9c238 perf: flatten paged prefill tile dispatch 
- remove the host-provided max_q_len argument
- dispatch only the ragged prefill tile upper bound
- validate the rebuilt CUDA backend end to end
 2026-08-09 23:12:53 +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 af25833fab fix: add out_buf to attn_paged_decode for CUDA graph capture compatibility 
- Pre-allocate decode_out in InferenceWorkspace so attn_paged_decode does not call torch::empty inside graph capture
- Wire decode_out through KVCache, PagePool.bind_tasks, and CudaBackend.fwd_decode
- Run live forward before graph capture to get valid output (graph pool memory is zeroed after capture block exits)
- Greedy generation with graph replay is bit-exact across all batch sizes
- Decode speedups vs no-graph: B=1 2.09x, B=4 1.80x, B=8 1.94x, B=16 1.76x
 2026-08-07 19:45: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 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 21bf37dd83 refactor: unify extension API to blhd layout and is_causal 
- Rename ops.py to attention_ops.py
- Remove layout/scale params: fixed blhd, auto scale
- Replace causal_offset with is_causal bool
- Move SDPA fallback to backend, ops only calls CUDA kernels
- Update __init__.py exports
 2026-07-30 18:39:20 +08:00