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:
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user