From c5dd8bfde3fb1ec8c8537e218e201dd0199189f1 Mon Sep 17 00:00:00 2001 From: David Orman Date: Sat, 25 Jul 2026 04:51:09 -0500 Subject: [PATCH] sm120: skip split-K for single-token decode shapes The SM120 split-K heuristic fills idle SMs by partitioning K whenever the MN block count is below half the SM count. At decode-sized M this trades a cheap GEMM for an expensive reduction: the reduce pass reads split_k * M * N floats to produce M * N, so its cost is dominated by launch and write-back rather than by the arithmetic it parallelises. Measured on SM120 (RTX PRO 6000, DeepSeek-V4-Flash single-user decode, M = 4 after padding, N = 8192, K = 1024, BLOCK_M = 32): the split GEMM plus sm120_split_k_reduce_impl costs about 22.4 us per projection, against about 13.2 us for the same projection with split_k = 1. Across 43 layers this is roughly 0.4 ms per decode graph. Skip split-K when M is at least 4x smaller than BLOCK_M, which leaves prefill and larger-batch decode shapes unchanged. Signed-off-by: David Orman --- csrc/jit_kernels/heuristics/sm120.hpp | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp index 0dcdcc8675..50ebf6e27f 100644 --- a/csrc/jit_kernels/heuristics/sm120.hpp +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -326,6 +326,18 @@ struct SM120ArchSpec { if (num_mn_blocks >= desc.num_sms / 2) return 1; + // Single-token decode (M well below block_m) leaves the split-K + // reduction dominated by launch and write-back cost rather than by the + // partial-sum arithmetic it parallelises: the reduction reads + // split_k * M * N floats to produce M * N, and at M = 4 that is a + // separate kernel whose runtime exceeds the GEMM time it saves. + // Measured on SM120 (RTX PRO 6000, DeepSeek-V4-Flash decode, M = 4, + // N = 8192, K = 1024): the split GEMM plus reduction costs about + // 22.4 us against about 13.2 us for the same projection without a + // split. Keep split-K for shapes with a real M tile. + if (desc.get_expected_m() * 4 <= layout.block_m) + return 1; + const int target_blocks = desc.num_sms * 3 / 4; int split_k = ceil_div(target_blocks, num_mn_blocks);