From b1b65a657ec473583eccbbc18bb9864010a340d4 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 2 Aug 2026 16:05:56 +0800 Subject: [PATCH] perf: target 512 grid blocks for decode split-K - compute_num_splits used 2*sm/base, undersplitting at large batch - single-warp decode blocks host ~11/SM, not 1/2-SM, so B=16 got 3 splits when 8 was optimal - Grid search on L20: bandwidth saturates near 256-512 total blocks; target 512 - Pass num_passes into base_blocks for the non-paged decode to match the paged path - B=16 kv=2048: 0.0230->0.0157ms (-32%); paged B=16: 0.0527->0.0243ms (-54%); B=32: 0.0406->0.0241ms (-41%) --- csrc/kernels/attn_dispatchers.cuh | 29 ++++++++++------------------- 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attn_dispatchers.cuh index c954cc9..acc5744 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attn_dispatchers.cuh @@ -16,29 +16,20 @@ #include "attn_paged_prefill_split_q_mma.cuh" #endif -// Cached SM count — cudaDeviceGetAttribute is a host-side call that was -// invoked on every decode/paged-decode launch. Cache per-device so multi-GPU -// setups with heterogeneous GPUs still get the right count, while the common -// single-GPU path hits the cache after the first call. -inline int get_sm_count() { - int dev = 0; - cudaGetDevice(&dev); - static int cached_dev = -1; - static int cached_count = 0; - if (dev != cached_dev) { - cudaDeviceGetAttribute(&cached_count, cudaDevAttrMultiProcessorCount, dev); - cached_dev = dev; - } - return cached_count; -} - // Split-KV: compute number of splits to fill all SMs for small-batch decode. // Caps splits so each split processes at least `min_tiles_per_split` tiles, // avoiding excessive loop/prologue overhead when tiles are small. +// +// Target total grid blocks (`TARGET_BLOCKS`) rather than scaling splits by SM +// count. Decode blocks are single-warp (32 threads) and a SM hosts ~11 of +// them, so the old `2*sm/base` cap badly undersplit at large batch (B=16 got +// 3 splits, optimal ~8). Measured (L20, grid search): bandwidth saturates +// near 256-512 total blocks; 512 minimizes worst-case latency across the +// B x kv grid; more is pure oversplit overhead. +constexpr int DECODE_TARGET_BLOCKS = 512; inline int compute_num_splits(int base_blocks, int tiles_total, int min_tiles_per_split = 1) { - int sm_count = get_sm_count(); - int n = (2 * sm_count + base_blocks - 1) / base_blocks; + int n = (DECODE_TARGET_BLOCKS + base_blocks - 1) / base_blocks; int max_by_work = tiles_total / min_tiles_per_split; return std::max(1, std::min(n, std::min(max_by_work, MAX_SPLITS))); } @@ -112,7 +103,7 @@ static inline void launch_decode_mma(AttentionParams& p, int group_size, c int num_passes = (G + MAX_G - 1) / MAX_G; constexpr int BC = 16; int tiles_total = (p.kv_len + BC - 1) / BC; - p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total, 2); + p.num_splits = compute_num_splits(p.batch * p.kv_head * num_passes, tiles_total, 2); constexpr int STAGES = 2; using Traits = KernelTraits; dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);