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
This commit is contained in:
2026-07-11 18:40:49 +08:00
parent 89ece26c25
commit 2c3cef1c87
9 changed files with 510 additions and 389 deletions
+1 -1
View File
@@ -11,7 +11,7 @@ import logging
logger = logging.getLogger(__name__)
KERNEL_NAMES = ["attn_decode", "attn_prefill"]
KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode"]
_available: dict[str, bool] = {}
_modules: dict[str, object] = {}