[None][feat] Indexer topk opt - #13811
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #47097 [ run ] triggered by Bot. Commit: |
|
PR_Github #47097 [ run ] completed with state |
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
cc @pcastonguay
Before zooming into the diff, I'd like to align on the broader trajectory of this kernel.
Strategic context
The TRT-LLM top-K kernel is migrating from single-CTA → multi-CTA. The strategic implementation of that step already landed on main as #12354 — single-pass multi-CTA cluster top-K in CuTe DSL, which uses cluster barriers + DSMEM to do the cross-CTA histogram merge in one pass and reuses the exact same radix micro-kernel that topKPerRowDecode uses internally.
This PR introduces a parallel-but-different multi-CTA strategy on the CUDA C++ side: a "split + global aux + merge" pattern with two kernel launches (the pre-cluster idiom). Because both implementations share the same inner micro-kernel, the gap between them is purely in the launch / scheduling / merge layer, which makes a head-to-head benchmark cleanly attributable.
Asks before merging
-
Head-to-head against
CuteDSLTopKDecodeSinglePassMultiCTAClusterRunner(#12354) on B200.
For the (BS × seq_len × cr) grid in (2). Since the micro-kernel is identical, this isolates "double-launch + global-aux merge" vs "cluster + DSMEM in-kernel". If #12354 wins on Blackwell, please reposition this PR's hardware scope (e.g., sm_90 / pre-cluster only) and state it explicitly in the description. -
Drive the dispatch constants from a systematic sweep, not from a single workload.
kDecodeMinColsPerSubBlock=2048,kDecodeTargetTotalBlocks=132, andkMaxBlocksPerRowDecode=10are currently picked ahead of time. Multi-CTA scheduling is tightly coupled with occupancy and L2 footprint; constants tuned from one (concurrency=256, ISL=8192, OSL=1024) trajectory will not generalize. Suggested grid:BS ∈ {1, 2, 4, 8, 16, 32, 64, 128, 132, 256}— covers thesmTarget=1transition at 132seq_len ∈ {4K, 8K, 16K, 32K, 64K, 128K, 256K}×cr ∈ {1, 4}— crosses both 12288 and 200000 boundaries- H100 and B200
Report
µs/call(before/after) as a heatmap + per-BS curves; pick the constants from the resulting envelope. Also:kDecodeTargetTotalBlocks=132is hard-coded for H100; B200 is 148 SM. Consider runtimecudaDeviceGetAttribute(cudaDevAttrMultiProcessorCount, ...). -
Isolate baseline / bug-fix / perf in a 3-way table.
The current "after" column bundles the bug fix in commitbf7f808111together with the perf change. Please report:Variant GSM8K kernel µs/call serving throughput baseline baseline + commit bf7f808111(bug fix only)baseline + all 4 commits Without this, the +5.13% throughput claim is partially attributable to the correctness fix.
-
End-to-end coverage at lower concurrency.
Commit messages 1–2 motivate the change with "BS=32", but the only serving run is atconcurrency=256, where decoder steady-state BS is much larger than 32 and the new heuristic does not engage on the hot decode path. Please add a run atconcurrency ∈ {32, 64}so the regime the PR targets is actually exercised end-to-end. -
Split commit
bf7f808111into a standalone[None][fix]PR tomain.
The two fixes (processHistogramStep<step==1>reset,finalIndices = -1sentinel pairing) are pre-existing bugs in the multi-block merge path. They affectmaintoo — large BS /numColumns ≥ 200000already routes to the same multi-block radix and inherits both. Landing the fix onmainshould not depend on the DSv4 → main merge timeline.
Follow-ups for the GVR top-K side (not gating this PR)
- GVR top-K is still a single-CTA implementation today. Once it evolves to a multi-CTA variant, the dispatch boundary between GVR and the multi-CTA radix path introduced here will need to be re-derived: rerun the same (BS × seq_len × cr) sweep with both kernels in their multi-CTA form, and pick the GVR ↔ radix crossover from the head-to-head envelope. Until that point, current single-CTA-GVR vs multi-CTA-radix numbers should be treated as an interim baseline only — the dispatch position cannot be considered final.
IndexerTopKOp.cppnow sizes aux buffers dynamically as{N, blocksPerRow, K}(range 1–10) instead of fixed{N, 10, K}. Please confirm CUDA Graph capture/replay on the GVR warm-up path is unaffected (no warmup-vs-replay shape diff).
|
PR_Github #47109 [ run ] triggered by Bot. Commit: |
|
PR_Github #47109 [ run ] completed with state |
The decode top-K kernel previously launched gridDim=(numRows,1,1) for
numColumns < 200K, leaving most SMs idle when numRows is small (e.g.
DeepSeek-V4 decode with batch=32 hits ~24% SM utilization on H100/B200).
The split-and-merge multipleBlocksPerRow path was only triggered for
numColumns >= 200K.
Introduce computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns) and
use it in both invokeIndexerTopKDecode and the IndexerTopKOp aux-buffer
allocation. The heuristic picks the smallest split that still saturates
SMs and keeps each sub-block large enough for the radix passes to be
efficient. It falls back to the original single-block path when numRows
is large enough to fill SMs and to the original 10-block path at >= 200K
columns, so existing behavior is unchanged at both ends.
Coverage extends down into the small-numColumns range: tiny rows
(numCols < 2048) still run as a single block, but moderate column
counts that previously short-circuited to blocksPerRow=1 now split
into 2-5 blocks per row to fill the GPU.
Two pre-existing bugs in the topKPerRowDecode<..., mergeBlocks=true>
merge step, only observable once the multi-block path is exercised at
smaller numColumns where rowLen <= topK seeds inputs with -FLT_MAX
padding, are also fixed:
1. processHistogramStep: when step 0 (fp16 fast path) cannot resolve
top-K because the threshold bin exceeds kNumFinalItems, the fp32
radix in steps 1-3 restarts from a full histogram. Reset
smemFoundTopKValues and smemFinalDstIdx at the entry of step 1 so
we don't double-count the candidates step 0 already wrote into
smemOutput. Without this, valid entries that fall below the
threshold in both fp16 and fp32 binnings were emitted twice,
producing 2x duplicated indices.
2. topKPerRowJob radix-sort fall-through: pair finalIndices[ii] = -1
with finalLogits[ii] = -FLT_MAX in the per-thread initialization
so unused sort slots survive SortDescendingBlockedToStriped as -1
rather than uninitialized register values. The bug surfaced when
the threshold bin was dominated by -FLT_MAX padding (all sort items
tied), causing the trailing top-K slots to receive garbage indices.
Adds test_indexer_topk_decode_sm_saturation that pins batch_size in
{16, 32, 128} and num_tokens up to 32K to exercise the multi-block
split path across the maxByCols guard, the moderate numColumns range,
and higher column counts.
Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
bf7f808 to
3f896a6
Compare
4 participants
@coderabbitai summary
Description
GSM8K accuracy (1319 samples, 5-shot, greedy)
exact_match,strict-matchexact_match,flexible-extractWithin stderr — accuracy unchanged.
Serving performance (concurrency 256, ISL 8192 / OSL 1024, gen_only)
Dataset:
random-v2/deepseek-v4-flash-8192-1024-20000-ratio-08_for_serve.json. 1× ctx server + 1× gen server, max-concurrency 256, 1 round, non-streaming.→ +5.1 % throughput, −7.4 % TTFT, −5.0 % E2EL. TPOT regresses ~1.3 %.
Before / After comparison — indexer top-K decode optimization
Target kernel: topKPerRowDecode
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.