[None][fix] Make indexer top-K launch policy device-aware by mingyangHao · Pull Request #13886 · NVIDIA/TensorRT-LLM

mingyangHao · 2026-05-08T06:44:01Z

Description

Magic number audit

Constant	Before	After
`kDecodeTargetTotalBlocks`	`132` (H100 SM count)	deleted — replaced by `getSchemeXBounds().smCount`
`kNumBins`	hardcoded `2048` × 2 function-local copies	hoisted to namespace scope; single source of truth
`kSortingAlgorithmThreshold`	`12288` (file-scope + 2 function-local copies)	`6 * kNumBins` (namespace scope only; duplicates removed)
`kDecodeMinColsPerSubBlock`	`2048`	`kNumBins`
`kDefaultSplitWorkThreshold`	`200 * 1000` (file-scope + 2 function-local copies)	unchanged value; duplicates removed
`kMaxBlocksPerRowDecode`	`10`	unchanged (empirical cap on aux-buffer / merge SMEM)

Perf data (B200, fp32 logits, cols=196608, next_n=1, no preIdx)

Production batch sizes (BS=1/32/128 — covering the original PR's perf table):

topK	BS=1	BS=32	BS=128
512 baseline (us)	51.1	65.2	95.0
512 new (us)	51.1	67.3	95.2
512 speedup	1.00x	0.97x	1.00x
1024 baseline (us)	56.0	72.6	105.7
1024 new (us)	55.9	72.9	109.9
1024 speedup	1.00x	1.00x	0.96x
2048 baseline (us)	63.1	79.8	122.8
2048 new (us)	63.1	79.3	121.9
2048 speedup	1.00x	1.01x	1.01x

Cliff band (BS=100-300):

topK	BS=100	BS=132	BS=148	BS=200	BS=300
512 baseline (us)	93.5	288.4	384.0	409.8	460.9
512 new (us)	89.3	107.6	112.2	144.0	228.4
512 speedup	1.05x	2.68x	3.42x	2.85x	2.02x
1024 baseline (us)	93.6	289.0	386.3	413.3	465.3
1024 new (us)	90.2	109.0	113.3	145.1	230.7
1024 speedup	1.04x	2.65x	3.41x	2.85x	2.02x
2048 baseline (us)	107.7	296.6	392.1	419.4	473.3
2048 new (us)	105.6	124.7	130.3	172.9	274.1
2048 speedup	1.02x	2.38x	3.01x	2.43x	1.73x

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.

lfr-0531 · 2026-05-08T09:30:36Z

/bot run

tensorrt-cicd · 2026-05-08T09:38:12Z

PR_Github #47374 [ run ] triggered by Bot. Commit: e72ae10 Link to invocation

tensorrt-cicd · 2026-05-08T09:56:55Z

PR_Github #47374 [ run ] completed with state FAILURE. Commit: e72ae10
/LLM/main/L0_MergeRequest_PR pipeline #37305 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

mingyangHao · 2026-05-11T01:36:08Z

/bot run

mingyangHao · 2026-05-11T06:13:41Z

/bot run

lfr-0531 · 2026-05-11T06:15:44Z

/bot run

tensorrt-cicd · 2026-05-11T06:22:43Z

PR_Github #47675 [ run ] triggered by Bot. Commit: a948f52 Link to invocation

tensorrt-cicd · 2026-05-11T07:29:39Z

PR_Github #47675 [ run ] completed with state SUCCESS. Commit: a948f52
/LLM/main/L0_MergeRequest_PR pipeline #37574 completed with status: 'SUCCESS'

CI Report

Link to invocation

lfr-0531 · 2026-05-11T07:30:39Z

/bot run --add-multi-gpu-test --disable-fail-fast

tensorrt-cicd · 2026-05-11T07:37:37Z

PR_Github #47682 [ run ] triggered by Bot. Commit: a373c1a Link to invocation

tensorrt-cicd · 2026-05-11T09:34:58Z

PR_Github #47682 [ run ] completed with state SUCCESS. Commit: a373c1a
/LLM/main/L0_MergeRequest_PR pipeline #37578 completed with status: 'SUCCESS'

CI Report

Link to invocation

longcheng-nv

Reviewed end-to-end with focus on substantive correctness / perf / coverage concerns. Posted 4 inline comments + this top-level point on perf table coverage. Detailed analysis is in the inline threads.

Perf table coverage gaps (top-level / PR description)

Table is K ∈ {512, 1024, 2048} × BS ∈ {1, 32} × cols ∈ {131K, 196K} = 9 cells. The policy change affects considerably more state space:

BS ∈ {64, 96, 128, 148} — exactly the wave-fit edge and the bump-to-2 band (see inline thread on indexerTopK.cu). Nothing currently shows whether the win generalizes here or whether it regresses.
BS=1, N≈8K — production GVR heuristic path on DeepSeek decode; this PR refactors the single-block dispatcher and the aux-skip logic that this path traverses. One row would prove no regression.
Prefill — see inline thread on invokeIndexerTopKPrefill. Algorithm-selection change deserves at least 1-2 prefill perf rows.

Adding ~3-4 rows would close the loop.

Overall: solid framing on device-aware launch + wave-fit. The headline BS=32 × cols=131K = 1.27× wins are clean — the inline comments are about widening confidence the win generalizes (1-3) and maintainability hardening (4).

computeIndexerTopKDecodeBlocksPerRow used a fixed kDecodeTargetTotalBlocks = 132 for smTarget. On B200 (148 SMs) the fixed constant leaves 16 SMs idle in the ceil-rounded band (numRows in [133, 148]): smTarget rounds to 1 (132/numRows < 1) so each row launches a single block, even though two sub-blocks per row would fit in one wave. Querying the device via getSchemeXBounds().smCount lifts that band to smTarget = 2 and saturates both waves' worth of SMs. Adds test_indexer_topk_decode_launch_policy_transitions covering batch_size in {1, 64, 100, 133, 256} so the smTarget transitions stay regression-protected; complements the existing decode_sm_saturation case. The new helper _run_indexer_topk_prefill_check extracts the prefill check body so future prefill regression tests can call it directly. Measured on B200, fp32 logits, topK = 2048, next_n = 1: numRows | numColumns | baseline | after | speedup --------|------------|---------:|-------:|--------: 133 | 196608 | 261.4us | 124.6us| 2.10x 133 | 131072 | 175.3us | 92.2us| 1.90x 64/100 | 131072+ | - | - | unchanged 148/200 | 196608 | - | - | unchanged Supersedes PR NVIDIA#13886. The competing computeIndexerTopKDecodeLaunchPolicy struct and canUseHeuristicTopKDecode helper in that PR are made redundant by the SchemeX v1.2 dispatcher that landed independently; the bump-to-2 path is similarly subsumed by smTarget = ceil(smCount / numRows). Only the device-aware change is retained here. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

lfr-0531 · 2026-05-14T15:43:59Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-14T15:50:23Z

PR_Github #48396 [ run ] triggered by Bot. Commit: 2ec03c9 Link to invocation

tensorrt-cicd · 2026-05-14T19:48:15Z

PR_Github #48396 [ run ] completed with state SUCCESS. Commit: 2ec03c9
/LLM/main/L0_MergeRequest_PR pipeline #38199 completed with status: 'SUCCESS'

CI Report

Link to invocation

Remove the hardcoded constants in computeIndexerTopKDecodeBlocksPerRow and route the cliff band through the multi-block path: * kDecodeTargetTotalBlocks = 132 (H100 SM count) is deleted. The smTarget heuristic now queries the actual device SM count via getSchemeXBounds().smCount (cudaDeviceGetAttribute, cached) so the dispatch tracks the hardware on B200 (148 SMs) and beyond. * kNumBins is hoisted to namespace scope (was duplicated as two function-local copies inside topKPerRowDecode / topKPerRowPrefill). kDecodeMinColsPerSubBlock = kNumBins and kSortingAlgorithmThreshold = 6 * kNumBins are now derived from it instead of being independent magic numbers. * Function-local copies of kSortingAlgorithmThreshold = 12288 and kDefaultSplitWorkThreshold = 200000 in invokeIndexerTopKDecodeDtype, invokeIndexerTopKPrefill, and canIndexerTopKDecodeUseGvr are removed in favor of the file-scope constants. The new dispatch heuristic uses two device-derived branches: numRows < smCount / 2 → wave-quantization-aware sweep over bp ∈ [2, maxBp] minimizing waves(bp) / bp; this avoids the +1-block-per-row wave-boundary spill that a naive ceil(smCount / numRows) target would produce on B200 (e.g. BS=15 stays at 1 wave instead of 2). numRows >= smCount / 2 → bp = 2. bp=1 in this band lands on the single-block radix kernel, which pays a sharp wave-scheduling cliff once gridDim.x approaches smCount (B200, cols=196608, topK=2048: BS=131=125us, BS=132=312us, BS=148=390us). The multi-block split+merge path is a different kernel instantiation and does not have this cliff; bp=2 is the cheapest split (smallest merge input). Measured speedup on B200 vs baseline (cols=196608, fp32 logits, next_n=1, no preIdx), across the three production topK values: topK BS BASELINE NEW speedup 512 100 93.5 89.3 1.05x 512 132 288.4 107.6 2.68x 512 148 384.0 112.2 3.42x 512 200 409.8 144.0 2.85x 512 300 460.9 228.4 2.02x 1024 100 93.6 90.2 1.04x 1024 132 289.0 109.0 2.65x 1024 148 386.3 113.3 3.41x 1024 200 413.3 145.1 2.85x 1024 300 465.3 230.7 2.02x 2048 100 107.7 105.6 1.02x 2048 132 296.6 124.7 2.38x 2048 148 392.1 130.3 3.01x 2048 200 419.4 172.9 2.43x 2048 300 473.3 274.1 1.73x BS < smCount / 2 is unchanged within measurement noise. Adds test_indexer_topk_decode_launch_policy_transitions covering batch_size in {1, 64, 100, 132, 148, 256} so the path-switch boundary stays regression-protected; the existing test_indexer_topk_decode_sm_saturation parametrization is kept verbatim. The new helper _run_indexer_topk_prefill_check extracts the prefill check body so future prefill regression tests can call it directly. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

longcheng-nv · 2026-05-15T02:49:59Z

Looking at the force-pushed fb2be5e2 (vs the prior 2ec03c9 that I'd previously cleared on substance) — the new commit makes the dispatch materially different, and I have four asks before I can re-ACK.

1. The numRows ∈ [smCount/2, smCount) force-bp=2 branch reintroduces the multi-wave bump I'd originally flagged.

The new computeIndexerTopKDecodeBlocksPerRow splits into if (numRows < smCount/2) sweep else bp = min(2, maxByCols). On B200 (smCount=148) the else branch fires for BS ∈ [74, 256] and produces:

BS=100 → 200 blocks → 1.35 waves
BS=128 → 256 blocks → 1.73 waves
BS=148 → 296 blocks → 2.0 waves

The new in-code comment justifies the bump with BS=131=125us, BS=132=312us, BS=148=390us — the cliff disappears entirely with bp=2, which establishes that bp=1 hits a wave-scheduling cliff, but does not establish that bp=2 in this band is faster than bp=1 post-cliff. The post-cliff bp=1 number is 390us at BS=148 — what is bp=2 at BS=100/128/148? Please add the bp=2 absolute timings to the PR description / comment so the bump is justified end-to-end, not just on cliff-avoidance.

2. Test comment contradicts the code.

The new comment above test_indexer_topk_decode_launch_policy_transitions says:

Covers the single-block / multi-block path transition at batch_size == kDecodeTargetTotalBlocks (= 132). batch_size < 132 keeps the single-block-per-row path; at and above 132 the launch policy routes to the multi-block split-and-merge path with bp = 2 …

But the actual threshold in the code is numRows < smCount / 2, which is 74 on B200 and 66 on H100 — not 132. Also kDecodeTargetTotalBlocks is deleted in this PR, so the comment references a non-existent constant. Please update the comment to match the real threshold (smCount / 2).

3. Undocumented behavior change: SM-count query failure now hard-aborts.

- int const smCount = bounds.smCount > 0 ? bounds.smCount : 132;
+ TLLM_CHECK_WITH_INFO(bounds.smCount > 0, "indexerTopK: failed to query device SM count");
+ int const smCount = bounds.smCount;

Previous revs had a soft 132 fallback; this rev throws. Two questions: (a) is there any test/CI environment (CPU-only sandbox, mocked CUDA) where cudaDeviceGetAttribute(MultiProcessorCount) can fail or return ≤0 and that previously kept working via the fallback? (b) If yes, please restore the fallback. If no, please call this out explicitly in the commit message — it is currently silent.

4. CI has not run on fb2be5e2.

The last PR_Github / L0_MergeRequest_PR pipeline ran on 2ec03c9 (#48396, completed 2026-05-14T19:48Z SUCCESS). The current head is fb2be5e2 (force-pushed 2026-05-15T02:42Z) with a substantively different dispatch function — please request a fresh /bot run so we have CI signal on the actual code under review.

The PR-type bump ([chore] → [perf]) is appropriate given the dispatch change, but it also means this isn't a pure refactor anymore and the perf table in the PR description should be expanded accordingly. cc @mingyangHao.

mingyangHao · 2026-05-15T03:04:56Z

/bot run --disable-fail-fast

lfr-0531 · 2026-05-15T03:11:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-15T03:17:50Z

PR_Github #48497 [ run ] triggered by Bot. Commit: fb2be5e Link to invocation

tensorrt-cicd · 2026-05-15T07:02:18Z

PR_Github #48497 [ run ] completed with state SUCCESS. Commit: fb2be5e
/LLM/main/L0_MergeRequest_PR pipeline #38292 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> (cherry picked from commit 4aad70b) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com> (cherry picked from commit 4aad70b) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

github-actions Bot assigned mingyangHao May 8, 2026

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch from e35aea3 to e72ae10 Compare May 8, 2026 06:50

lfr-0531 added the deepseek-v4 label May 9, 2026

mingyangHao marked this pull request as ready for review May 11, 2026 01:35

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch 2 times, most recently from 959db42 to a948f52 Compare May 11, 2026 05:48

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch from a948f52 to a373c1a Compare May 11, 2026 06:31

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch 3 times, most recently from 72c3df9 to 1f8e33c Compare May 13, 2026 05:25

mingyangHao requested a review from longcheng-nv May 13, 2026 06:06

longcheng-nv reviewed May 13, 2026

View reviewed changes

Comment thread cpp/tensorrt_llm/kernels/indexerTopK.cu Outdated

Comment thread tests/unittest/_torch/thop/parallel/test_indexer_topk.py Outdated

Comment thread cpp/tensorrt_llm/thop/IndexerTopKOp.cpp Outdated

Comment thread cpp/tensorrt_llm/kernels/indexerTopK.cu

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch from 352fa49 to c5d85fb Compare May 14, 2026 06:13

lfr-0531 force-pushed the feat/deepseek_v4 branch from 0a93d10 to 118e7a5 Compare May 14, 2026 07:44

lfr-0531 requested a review from a team as a code owner May 14, 2026 07:44

lfr-0531 requested a review from a team May 14, 2026 07:44

lfr-0531 requested a review from a team as a code owner May 14, 2026 07:44

lfr-0531 requested review from tburt-nv and removed request for a team May 14, 2026 07:44

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch 2 times, most recently from 93c46ff to af25d9d Compare May 14, 2026 13:35

mingyangHao removed request for a team, mzweilz and tburt-nv May 14, 2026 13:37

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch from af25d9d to 2ec03c9 Compare May 14, 2026 13:44

mingyangHao force-pushed the mingyangh/indexer-topk-launch-policy branch from 2ec03c9 to fb2be5e Compare May 15, 2026 02:42

lfr-0531 merged commit 4aad70b into NVIDIA:feat/deepseek_v4 May 15, 2026
8 checks passed

longcheng-nv mentioned this pull request May 24, 2026

[None][perf] DSv4 indexer Top-K: align kFTarget=kK, uniform kSeqSmall=4096, diagnostic env knobs #14413

Merged

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 10, 2026

[None][fix] Make indexer top-K launch policy device-aware (NVIDIA#13886)

6af944f

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 11, 2026

[None][fix] Make indexer top-K launch policy device-aware (NVIDIA#13886)

a6250d7

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 12, 2026

[None][fix] Make indexer top-K launch policy device-aware (NVIDIA#13886)

142a99a

Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

Uh oh!

Conversation

mingyangHao commented May 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Magic number audit

Perf data (B200, fp32 logits, cols=196608, next_n=1, no preIdx)

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

lfr-0531 commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

mingyangHao commented May 11, 2026

Uh oh!

mingyangHao commented May 11, 2026

Uh oh!

lfr-0531 commented May 11, 2026

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

lfr-0531 commented May 11, 2026

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

longcheng-nv left a comment

Choose a reason for hiding this comment

Perf table coverage gaps (top-level / PR description)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lfr-0531 commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

tensorrt-cicd commented May 14, 2026

Uh oh!

longcheng-nv commented May 15, 2026

Uh oh!

mingyangHao commented May 15, 2026

Uh oh!

lfr-0531 commented May 15, 2026

Uh oh!

tensorrt-cicd commented May 15, 2026

Uh oh!

tensorrt-cicd commented May 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mingyangHao commented May 8, 2026 •

edited

Loading