AstrAI/csrc/kernels
ViperEkura fd65b9bc23 feat: support HEAD_DIM=32 and split extension into loader/ops
- 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
2026-07-08 14:14:11 +08:00
..
gqa_common.cuh fix: correct prefill mask index, unify GQA kernel interface 2026-07-06 17:21:23 +08:00
gqa_decode_attn.cu feat: support HEAD_DIM=32 and split extension into loader/ops 2026-07-08 14:14:11 +08:00
gqa_decode_attn.cuh fix: correct prefill mask index, unify GQA kernel interface 2026-07-06 17:21:23 +08:00
gqa_decode_attn_mma.cuh refactor: extract shared MMA utils into gqa_mma_utils.cuh 2026-07-07 23:01:15 +08:00
gqa_mma_utils.cuh feat: support HEAD_DIM=32 and split extension into loader/ops 2026-07-08 14:14:11 +08:00
gqa_prefill_attn.cu feat: support HEAD_DIM=32 and split extension into loader/ops 2026-07-08 14:14:11 +08:00
gqa_prefill_attn.cuh feat: tensor-core mma prefill with build-time dispatch 2026-07-06 20:33:24 +08:00
gqa_prefill_attn_mma.cuh feat: support HEAD_DIM=32 and split extension into loader/ops 2026-07-08 14:14:11 +08:00