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
@@ -5,7 +5,7 @@ logger = logging.getLogger(__name__)

 available: dict[str, bool] = {}

-for _name in ["gqa_decode_attn"]:
+for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
    try:
        importlib.import_module(f".{_name}", package=__package__)
        available[_name] = True