perf(sm120): skip split-K when the reduction is slower than the single-token GEMM - #77
Open
ormandj wants to merge 1 commit into
Open
perf(sm120): skip split-K when the reduction is slower than the single-token GEMM#77ormandj wants to merge 1 commit into
ormandj wants to merge 1 commit into
Conversation
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 <ormandj@corenode.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
The SM120 split-K heuristic partitions K when the number of MN blocks is below half the SM count. For single-token decode shapes where M is much smaller than
BLOCK_M, the separate reduction kernel costs more than the work saved by splitting the GEMM.This change returns
split_k = 1when:expected_m * 4 <= block_mLarger decode and prefill shapes continue through the existing split-K heuristic.
Profiled result
Measured on an RTX PRO 6000 Blackwell SM120 using a DeepSeek-V4-Flash single-user decode projection:
sm120_split_k_reduce_implsplit_k = 1This projection appears in 43 model layers, corresponding to approximately 0.4 ms removed from the decode graph.
The profile comparison uses the same projection shape and changes only the split-K selection.
Numerical and serving validation
The source-equivalent serving stack containing the change:
The complete candidate produced these forward-pass medians:
The complete release stack also contained the separately reviewed SM120 compiled-dimension correction. These end-to-end results therefore establish integration and regression coverage; the isolated split-K performance attribution is the matched 22.4 µs versus 13.2 µs projection profile above.
Scope
The change is limited to the SM120 heuristic. It does not modify the kernels, reduction implementation, other architectures, or larger-M split-K behavior.
Prepared with AI assistance.