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

perf: increase decode split-K parallelism for short sequences

Browse Source 
- Remove tiles_total/8 min-work cap that limited splits for small workloads
- Simplify decode_num_splits to only use base_blocks and tiles_total
- Short sequences now generate more blocks, improving SM utilization
This commit is contained in:
ViperEkura
2026-07-11 11:26:42 +08:00
parent 8a8550184f
commit a4ae7d17fb
3 changed files with 11 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/tests/attn_decode_test.cu
+5 -8
View File
@@ -28,14 +28,11 @@ struct DecodeScratch {
float* ml_part = nullptr;
};
static int decode_num_splits(const AttentionParams<bf16>& p, int tiles_total) {
static int decode_num_splits(int base_blocks, int tiles_total) {
int sm_count = 0;
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
int base_blocks = p.kv_head * p.batch;
int desired = 2 * (sm_count > 0 ? sm_count : 64);
int n = (desired + base_blocks - 1) / base_blocks;
int max_by_work = tiles_total / 8;
return max(1, min(min(min(n, tiles_total), 32), max_by_work));
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
return std::max(1, std::min(n, std::min(tiles_total, 32)));
}
// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
@@ -49,7 +46,7 @@ static bool decode_use_mma(const AttentionParams<bf16>& p) {
template <int HEAD_DIM, int BC>
static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
int tiles_total = (p.kv_len + BC - 1) / BC;
p.num_splits = decode_num_splits(p, tiles_total);
p.num_splits = decode_num_splits(p.batch * p.kv_head, tiles_total);
p.o_part = sc.o_part;
p.ml_part = sc.ml_part;
@@ -62,7 +59,7 @@ static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
static void launch_scalar_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
int gs = p.q_head / p.kv_head;
int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
p.num_splits = decode_num_splits(p, chunks_total);
p.num_splits = decode_num_splits(p.batch * p.kv_head, chunks_total);
p.o_part = sc.o_part;
p.ml_part = sc.ml_part;