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

- 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:
2026-07-21 22:21:39 +08:00
parent a01e8bbe98
commit f7a16efc9d
9 changed files with 245 additions and 406 deletions
-10
View File
@@ -18,16 +18,6 @@ inline double now_ms() {
return duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
}
inline int compute_num_splits(int base_blocks, int tiles_total) {
int sm_count = 0;
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
if (n > tiles_total) n = tiles_total;
if (n > 32) n = 32;
if (n < 1) n = 1;
return n;
}
#define CUDA_CHECK(call) \
do { \
cudaError_t _e = (call); \