Commit Graph
9 Commits
Author SHA1 Message Date
ViperEkura 4b25664c79 perf: precompute kv_indptr once per decode step
- bind_tasks builds kv_indptr (prefix sum of seq_lens) a single time
- fwd_decode/fwd_prefill reuse it instead of rebuilding per layer
- Removes 24 cumsum launches per decode step (was ~1ms/step at B=4)
- Decode B=4: 9.60 -> 7.82 ms/step (-18.5%), +22.8% tok/s
2026-08-01 16:09:26 +08:00
ViperEkura 91acaf4b0b refactor: unify attention mask to single attn_mask tensor
- CudaBackend.fwd_decode passes attn_mask directly instead of kv_cache.decode_mask
- TorchNativeBackend derives pos_mask from attn_mask[:,0,0] on decode
- Drop decode_mask and page_table fields from KVCache and bind_tasks
2026-08-01 15:49: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 530d280e33 perf: remove split partials memset and overlap decode tile loads
- alloc_split_partials now uses torch::empty: the split kernel writes every slot it owns, so the per-call zeros/full memset was pure overhead (2 kernels per layer per step)
- decode split-KV MMA kernels now run a true multi-stage cp.async pipeline (wait_group<STAGES-1> instead of wait_group<0>), keeping STAGES-1 tile loads in flight; the old wait_group<0> serialized load and compute so deeper STAGES made no difference
- add a fallback path when ntiles < STAGES to avoid a race on the last tile
2026-07-31 22:37:44 +08:00
ViperEkura 50cfd0d555 perf: reduce decode overhead in scheduler and executor
- Precompute page_table and decode_mask on KVCache once per step in PagePool.bind_tasks, instead of per-layer in CudaBackend/TorchNativeBackend
- Skip frequency penalty history tensor construction when all penalties are 0 in Executor.execute_decode
- Omit FrequencyPenaltyStrategy from sampling pipeline when penalty is 0
- Deduplicate get_active_tasks calls in scheduler loop (3 to 1), remove redundant sorted() on decode tasks
- Benchmark (L20, bf16, CUDA backend): B=1 9.48->9.40ms (+1%), B=4 10.73->9.89ms (+8.6%), B=8 10.77->10.13ms (+6.4%)
2026-07-31 14:50:16 +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
ViperEkura 32fd03a025 feat: add CudaBackend and rename to fwd_decode/fwd_prefill
- CudaBackend: paged decode via attn_paged_decode, prefill via attn_prefill
- Decode uses req_to_token as page_table with page_size=1
- Falls back to TorchNativeBackend when kernel unavailable
- Rename forward_decode/forward_extend to fwd_decode/fwd_prefill
- Register ATTN_BACKEND.CUDA in _BACKEND_REGISTRY
2026-07-30 18:45:33 +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
ViperEkura 5b67d5865a feat: add AttentionBackend ABC with context manager
- AttentionBackend ABC with forward_decode/forward_extend dispatch
- TorchNativeBackend: SDPA with indirect KV cache gather
- attn_backend() context manager + ATTN_BACKEND enum (mirrors sdpa_kernel)
- ContextVar-based thread-safe backend switching
- get_backend() falls back to default TorchNativeBackend singleton
2026-07-30 18:20:27 +08:00