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

@@ -27,6 +27,7 @@ NVCC_FLAGS = [
     "--use_fast_math",
     "--ptxas-options=-O3,-v",
     "--extra-device-vectorization",
+    "--threads=8",
 ]
 
 

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
+static int decode_num_splits(int base_blocks, int tiles_total) {
-// 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<bf16>& p, int tiles_total) {
     int sm_count = 0;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
-    int base_blocks = p.kv_head * p.batch;
+    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
-    int desired = 2 * (sm_count > 0 ? sm_count : 64);
+    return std::max(1, std::min(n, std::min(tiles_total, 32)));
-    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}));
 }
 
 // Scalar fallback: one warp per query head, split-KV across grid.z.
 static void launch_scalar_decode(AttentionParams<bf16>& 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<bf16>& p) {
 template <int HEAD_DIM, int BC>
 static void launch_mma_decode(AttentionParams<bf16>& 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);

csrc/tests/attn_decode_test.cu

@@ -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 n = (2 * sm_count + base_blocks - 1) / base_blocks;
-    int desired = 2 * (sm_count > 0 ? sm_count : 64);
+    return std::max(1, std::min(n, std::min(tiles_total, 32)));
-    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));
 }
 
 // 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;