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
parent d7da51569f
commit 579b8c3129
4 changed files with 181 additions and 31 deletions
@@ -27,3 +27,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs):


 register("gqa_decode_attn")
+register("gqa_prefill_attn")