- add case 32 to decode/prefill dispatch switch - fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory - restructure decode dispatch to #ifndef/#else/#endif matching prefill - split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API |
||
|---|---|---|
| .. | ||
| gqa_common.cuh | ||
| gqa_decode_attn.cu | ||
| gqa_decode_attn.cuh | ||
| gqa_decode_attn_mma.cuh | ||
| gqa_mma_utils.cuh | ||
| gqa_prefill_attn.cu | ||
| gqa_prefill_attn.cuh | ||
| gqa_prefill_attn_mma.cuh | ||