 ViperEkura
|
ddc4bd1cf6
|
feat: tensor-core mma prefill with build-time dispatch
- add register-resident flash-attention kernel using mma.sync.m16n8k16
- dispatch mma vs scalar at build time: pre-sm_80 defines
-DASTRAI_NO_MMA, else defaults to mma
- scalar path vectorized with float4 smem loads (ld8)
|
2026-07-06 20:33:24 +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 |