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
ViperEkura