ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: extract shared dispatcher header, unify MMA/scalar dispatch format

Browse Source 
- Merge 3 duplicated dispatch blocks into single attn_dispatchers.cuh
- Merge compute_num_splits from attn_utils.cuh into dispatcher header
- All dim3 grid/block declarations and <<<>>> launches are single-line
- Production .cu files (35-42 loc) only handle torch wrapping + pybind11
- Test files include dispatcher header directly, removing all #ifndef ASTRAI_NO_MMA duplication
This commit is contained in:
ViperEkura
2026-07-21 22:21:39 +08:00
parent a01e8bbe98
commit f7a16efc9d
9 changed files with 245 additions and 406 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_prefill.cu
+1 -48
View File
@@ -1,53 +1,6 @@
#include "attn_prefill_split_q.cuh"
#include "attn_dispatchers.cuh"
#include "attn_entry_utils.cuh"
#ifndef ASTRAI_NO_MMA
#include "attn_prefill_split_q_mma.cuh"
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static void launch_mma_prefill(AttentionParams<bf16>& p) {
constexpr int WARPS = 4;
constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
dim3 grid((p.q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS),
p.q_head, p.batch);
dim3 block(Traits::NUM_THREADS, 1, 1);
attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
}
#endif
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static void launch_scalar_prefill(AttentionParams<bf16>& p) {
constexpr int G = 8, ROWS = 32, P_BC = 32;
dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
dim3 block(G, ROWS, 1);
attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask><<<grid, block>>>(p);
}
template <int HEAD_DIM>
static void dispatch_prefill(AttentionParams<bf16>& p) {
bool is_causal = (p.causal_offset >= 0);
bool has_mask = (p.use_mask && p.mask);
#ifndef ASTRAI_NO_MMA
if (is_causal) {
if (has_mask) launch_mma_prefill<HEAD_DIM, true, true>(p);
else launch_mma_prefill<HEAD_DIM, true, false>(p);
} else {
if (has_mask) launch_mma_prefill<HEAD_DIM, false, true>(p);
else launch_mma_prefill<HEAD_DIM, false, false>(p);
}
#else
if (is_causal) {
if (has_mask) launch_scalar_prefill<HEAD_DIM, true, true>(p);
else launch_scalar_prefill<HEAD_DIM, true, false>(p);
} else {
if (has_mask) launch_scalar_prefill<HEAD_DIM, false, true>(p);
else launch_scalar_prefill<HEAD_DIM, false, false>(p);
}
#endif
}
torch::Tensor attn_prefill(
torch::Tensor q,
torch::Tensor k,