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:
@@ -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] = {}
|
||||
|
||||
Reference in New Issue
Block a user