ViperEkura
53ed52b4b8
refactor: extension dispatch layer with CUDA/torch fallback
...
- Add gqa_decode_attn/gqa_prefill_attn dispatch functions
- Internal _available/__modules with underscore prefix
- CUDA kernel path with F.scaled_dot_product_attention fallback
- GQA head expansion in fallback path
2026-07-06 21:07:16 +08:00
ViperEkura
579b8c3129
fix: correct gqa_decode_attn reduction + add gqa_prefill_attn
...
- gqa_decode_attn: rewrite to per-KV-head, K in smem
- gqa_prefill_attn: new kernel for Q_len > 1 with GQA
2026-07-06 13:45:18 +08:00
ViperEkura
e8e228d035
feat: add optional CUDA kernel system (csrc/) + fused GQA decode attention
...
Structure:
csrc/ -- .cu sources + build.py registry
astrai/extension/ -- compiled .so + __init__.py (import dispatcher)
setup.py -- CUDAExtension from csrc/build.py REGISTRY
Control: CSRC_KERNELS=true|false env var at install time.
Fallback: astrai.extension.available dict for runtime detection.
2026-07-06 12:09:58 +08:00
ViperEkura