refactor: template combine kernel, fix mask bug, unify dispatch

- Template combine kernel, share macros, extract entry_utils helpers
- Fix mask indexing (pass stride not pre-multiplied base)
- Remove !p.use_mask — MMA handles mask
This commit is contained in:
2026-07-15 21:44:17 +08:00
parent 1f0be382ad
commit 0654b4b916
7 changed files with 70 additions and 110 deletions
+1 -1
View File
@@ -47,7 +47,7 @@ torch::Tensor attn_prefill(
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_prefill<D>(p); });
DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
return O;
}