From a4ae7d17fb1f26ea2d7447b6fe52bc9471641d39 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 11 Jul 2026 11:26:42 +0800 Subject: [PATCH] perf: increase decode split-K parallelism for short sequences - 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 --- csrc/build.py | 1 + csrc/kernels/attn_decode.cu | 17 +++++------------ csrc/tests/attn_decode_test.cu | 13 +++++-------- 3 files changed, 11 insertions(+), 20 deletions(-) diff --git a/csrc/build.py b/csrc/build.py index 8605830..6797898 100644 --- a/csrc/build.py +++ b/csrc/build.py @@ -27,6 +27,7 @@ NVCC_FLAGS = [ "--use_fast_math", "--ptxas-options=-O3,-v", "--extra-device-vectorization", + "--threads=8", ] diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu index f37cc81..d9c0947 100644 --- a/csrc/kernels/attn_decode.cu +++ b/csrc/kernels/attn_decode.cu @@ -5,25 +5,18 @@ #include "attn_decode_split_kv_mma.cuh" #endif -// Decode has only batch*kv_head independent tasks; without split-K the grid is -// tiny (e.g. 16 blocks) and leaves most SMs idle. Pick the smallest split count -// that fills the device (~2 blocks/SM), capped by the tile count, min work per -// split (at least 8 tiles), and 32. -static int decode_num_splits(const AttentionParams& 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 std::max(1, std::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))); } // Scalar fallback: one warp per query head, split-KV across grid.z. static void launch_scalar_decode(AttentionParams& p) { int group_size = 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); auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); @@ -46,7 +39,7 @@ static bool decode_use_mma(const AttentionParams& p) { template static void launch_mma_decode(AttentionParams& p) { 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); auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); diff --git a/csrc/tests/attn_decode_test.cu b/csrc/tests/attn_decode_test.cu index 38a3333..dfefda6 100644 --- a/csrc/tests/attn_decode_test.cu +++ b/csrc/tests/attn_decode_test.cu @@ -28,14 +28,11 @@ struct DecodeScratch { float* ml_part = nullptr; }; -static int decode_num_splits(const AttentionParams& 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& p) { template static void launch_mma_decode(AttentionParams& 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& p, DecodeScratch& sc) { static void launch_scalar_decode(AttentionParams& 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;