[None][fix] GVR indexer top-K: repair the non-converged threshold search - #17550
Open
longcheng-nv wants to merge 1 commit into
Open
[None][fix] GVR indexer top-K: repair the non-converged threshold search#17550longcheng-nv wants to merge 1 commit into
longcheng-nv wants to merge 1 commit into
Conversation
The heuristic (GVR) decode path picks a value threshold whose candidate count lands in [K, kC] and then selects the top-K from those candidates. Three input shapes defeated that search and produced a silently WRONG top-K — no error, no diagnostic, just wrong indices: 1. Undershoot. The Phase-2 secant is capped at 15 iterations; on non- convergence the `done = 2` fallback can pick `val_hi` outright, and the Phase-3 retry loop only ever triggered on the overflow side (`cand_count > kCC`). A threshold admitting fewer than K candidates went straight through: the collect emitted what it had and the Phase-4 tail padded the rest with index -1. 2. Degenerate hint. When every hinted value is identical, Phase 1 builds an empty bracket and the kernel returned `outputIndices[i] = i` — the head of the row, not a top-K at all. 3. Tie plateau. When more than kC elements share the K-th value, NO threshold yields a count in [K, kC]. The collect clamps at kC and dropped strictly-greater entries. All three are hint-quality driven and need no unusual logits, so they are reachable in production whenever a layer's temporal locality breaks down. On the shipped op (1.3.0rc21) an anti-correlated hint at N=65536, K=512 returns 512 of 512 wrong indices. Production DSv4 decode captures hit it unaided: V4-Flash K=512 N=131075 layers 22/24 (283 / 87 slots left at -1, hit-rate 0.023 / 0.057) and V4-Pro K=1024 N=262127 layer 40 (550 slots, hit-rate 0.122). N=131075 is inside the shipped GVR routing window, so this is a live defect, not a latent one. Fix, in both the fp32 and the bf16/fp16 job: * Phase 1's degenerate-bracket branch no longer emits row[0:K]; it resets to the widest trusted bracket and falls through. The hint may only affect speed, never the answer. * Phase 3 repairs BOTH sides. Entry anchors the untested bracket end at a float extreme (count(-FLT_MAX) = #finite >= K, count(FLT_MAX) = 0 < K), because Phase 1 seeds val_lo/val_hi from hinted min/max with invented counts (M + M/4, 1) that can leave both ends on the same side of the K-th value. The loop then bisects on the order-preserving uint32 image of the key space, so the bracket provably collapses to adjacent representable values in <= 32 steps instead of relying on a float average. * On collapse with more than kC elements at the threshold, emit directly: everything strictly above (fewer than K by construction) plus arbitrary ties, which is a valid top-K. Guarded on the collapse test so it can never run on a non-collapsed bracket. The converged fast path (`done == 1`) is untouched — only rows whose secant failed enter any of this. Measured on B200, cold-L2, real DSv4 decode captures (both models x 5 ISL rungs x 90 rows, fp32): geomean 0.972 vs the unfixed kernel, i.e. neutral to slightly faster, with one outlier at 1.30x — V4-Pro ISL=1M, the bucket that contains the non-converging layer 40, which now pays the bisection instead of returning a wrong answer. Exactness: 353/353 real-capture cells and 0/135 adversarial cells inexact, against 4 and 54 before. Tests: hostile-hint (bottom-K / uniform-argmax / random) x K x N and tie-plateau regressions; both classes fail on the pre-fix kernel. Made-with: Claude Code (Opus 5, 1M context) Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Contributor
WalkthroughThe GVR heuristic top-K kernel now repairs degenerate hints, candidate overflow, and undershoot across fp32, bf16, and fp16 paths. New tests cover hostile hints and large threshold-tie plateaus against exact ChangesGVR top-K repair
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant IndexerTopKTest
participant GVRDecode
participant HeuristicTopK
participant TorchTopK
IndexerTopKTest->>GVRDecode: submit compressed-ratio-4 inputs and heuristic hints
GVRDecode->>HeuristicTopK: decode candidate bracket
HeuristicTopK->>HeuristicTopK: repair threshold and emit selected values
GVRDecode->>TorchTopK: compute reference top-K
IndexerTopKTest->>IndexerTopKTest: compare decoded values with reference
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
cpp/tensorrt_llm/kernels/heuristic_topk.cuh (2)
429-448: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider hoisting a single ordered-key bijection.
gvrOrderKeyandgvrOrderKeyToFloatduplicatefloatToOrderedUintandorderedUintToFloatat lines 376-385. The only difference is the__CUDA_ARCH__ >= 800guard around the originals. You can define the pair once above the guard and letwarpReduceMinandwarpReduceMaxcall it. The emitted bit operations stay the same, so the SASS-sensitive fp32 K=2048 path is unaffected. Verify the byte-identity claim in CI if you apply this.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/heuristic_topk.cuh` around lines 429 - 448, Hoist the ordered float/uint32 conversion pair into a single definition before the __CUDA_ARCH__ >= 800 guard, then update warpReduceMin and warpReduceMax to reuse those helpers instead of the duplicate gvrOrderKey and gvrOrderKeyToFloat definitions. Preserve the existing bit operations and verify the fp32 K=2048 generated bytes remain identical in CI.
1527-1546: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winThe bisection uses fp32 keys on bf16/fp16 data, so the collapse test needs the full iteration budget.
gvrOrderKeymaps to the fp32 key space. The input elements are bf16 or fp16. Many fp32 keys lie between two adjacent representable input values, soblockCountGEDtypereturns the same count for all of them. The loop keeps halving the fp32 gap and issues a full-N counting pass per iteration, even though the count stopped changing much earlier. On the tie-plateau rows that this PR targets, the loop runs close toMAX_REPAIR_ITERSbeforekhi <= klo + 1ubecomes true.Consider adding an early exit when the bracket is already tight in the input dtype, for example by comparing
Trait::from_fp32(val_lo)withTrait::from_fp32(val_hi). That reduces the worst-case repair from about 32 full-N passes to about 16 for fp16 and about 8 for bf16.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/heuristic_topk.cuh` around lines 1527 - 1546, Update the repair loop’s collapse check in the heuristic top-k bisection to compare val_lo and val_hi after conversion to the input dtype via Trait::from_fp32, while retaining the existing fp32-key adjacency check. Exit when the converted bounds are equal or otherwise represent an already-tight input-dtype bracket, so bf16/fp16 inputs avoid unnecessary full-N counting passes while preserving the existing threshold and count-update flow.tests/unittest/_torch/thop/parallel/test_indexer_topk.py (1)
2377-2383: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse keyword arguments for optional tensors. The current order is correct, but keywords prevent breakage if the schema changes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py` around lines 2377 - 2383, Update the indexer_topk_decode invocation in _gvr_decode_exact_check to pass optional tensor arguments using their parameter names rather than positional ordering, while preserving the current argument mapping and behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tensorrt_llm/kernels/heuristic_topk.cuh`:
- Around line 691-710: Prevent the degenerate-hint reset from feeding an
infinite range into Phase 2: in the float path around lines 691-710 and the
bf16/fp16 path around lines 1379-1397 of
cpp/tensorrt_llm/kernels/heuristic_topk.cuh, update the corresponding secant
logic to fall back to a gvrOrderKey midpoint whenever vhi - vlo is non-finite,
or bypass Phase 2 and enter ordered-key bisection directly. Apply the same guard
to both secant implementations so threshold refinement progresses without
relying solely on Phase 3 repair.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 2384-2407: Initialize the `indices` tensor in the indexer_topk
test with a negative sentinel instead of `torch.empty`, then add an assertion
that `indices[0]` contains `index_topk` unique values before comparing selected
logits. Keep the existing negative-sentinel check and value comparison
unchanged.
- Around line 2436-2454: Parameterize test_indexer_topk_decode_gvr_tie_plateau
over fp32, bf16, and fp16, converting logits to the selected dtype before
execution. Replace torch.linspace with exactly representable power-of-two-based
values so strictly-greater entries remain distinct in every dtype, and keep the
tie plateau exactly representable; preserve the existing hostile tie sizes and
exact top-K validation so the bf16/fp16 collapsed-bracket direct-emit path is
exercised.
- Around line 2410-2454: The helper _gvr_decode_exact_check currently validates
selected values without ensuring indices are distinct. Update this helper to
explicitly assert that the returned top-K indices are unique, while preserving
its existing value comparison and coverage for tied logits and all dtypes.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/heuristic_topk.cuh`:
- Around line 429-448: Hoist the ordered float/uint32 conversion pair into a
single definition before the __CUDA_ARCH__ >= 800 guard, then update
warpReduceMin and warpReduceMax to reuse those helpers instead of the duplicate
gvrOrderKey and gvrOrderKeyToFloat definitions. Preserve the existing bit
operations and verify the fp32 K=2048 generated bytes remain identical in CI.
- Around line 1527-1546: Update the repair loop’s collapse check in the
heuristic top-k bisection to compare val_lo and val_hi after conversion to the
input dtype via Trait::from_fp32, while retaining the existing fp32-key
adjacency check. Exit when the converted bounds are equal or otherwise represent
an already-tight input-dtype bracket, so bf16/fp16 inputs avoid unnecessary
full-N counting passes while preserving the existing threshold and count-update
flow.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py`:
- Around line 2377-2383: Update the indexer_topk_decode invocation in
_gvr_decode_exact_check to pass optional tensor arguments using their parameter
names rather than positional ordering, while preserving the current argument
mapping and behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 72bcf753-589f-4eba-98b0-a170d6258442
📒 Files selected for processing (2)
cpp/tensorrt_llm/kernels/heuristic_topk.cuhtests/unittest/_torch/thop/parallel/test_indexer_topk.py
Comment on lines
+691
to
710
| // Degenerate hint (every hinted value identical, or none in range): | ||
| // Phase 1 produced no usable bracket. This used to emit the first K | ||
| // elements of the row verbatim, which is not a top-K at all — it is | ||
| // simply the head of the row. Fall through instead with the widest | ||
| // trusted bracket and let Phase 2 / the Phase-3 repair locate the | ||
| // threshold; the hint only ever affects speed, never the answer. | ||
| if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) | ||
| { | ||
| if (tid == 0) | ||
| for (int i = 0; i < topK && i < N; i++) | ||
| { | ||
| outputIndices[i] = i; | ||
| outputValues[i] = input[i]; | ||
| } | ||
| return; | ||
| { | ||
| float const seed = (smem->val_hi <= -FLT_MAX) ? 0.0f : smem->pmax_saved; | ||
| smem->val_lo = -FLT_MAX; | ||
| smem->val_hi = FLT_MAX; | ||
| smem->cnt_lo = N; | ||
| smem->cnt_hi = 0; | ||
| smem->threshold = seed; | ||
| smem->done = 0; | ||
| } | ||
| __syncthreads(); | ||
| } |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
The degenerate-hint reset writes a bracket whose width overflows to +inf, so Phase 2 degenerates before the repair runs. Both dtype paths set val_lo = -FLT_MAX and val_hi = FLT_MAX. The Phase-2 secant computes vhi - vlo, which overflows to +inf, then produces -inf and finally NaN for the threshold. The row consumes every MAX_REFINE_ITERS full-N counting pass without progress and depends entirely on the Phase-3 repair for a correct answer.
cpp/tensorrt_llm/kernels/heuristic_topk.cuh#L691-L710: after the reset, either skip Phase 2 and enter the ordered-key bisection directly, or make the secant at Line 745 fall back to agvrOrderKeymidpoint whenrangeis not finite.cpp/tensorrt_llm/kernels/heuristic_topk.cuh#L1379-L1397: apply the identical guard to the bf16/fp16 secant at Line 1433, which computes the samevhi - vlo.
📍 Affects 1 file
cpp/tensorrt_llm/kernels/heuristic_topk.cuh#L691-L710(this comment)cpp/tensorrt_llm/kernels/heuristic_topk.cuh#L1379-L1397
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/tensorrt_llm/kernels/heuristic_topk.cuh` around lines 691 - 710, Prevent
the degenerate-hint reset from feeding an infinite range into Phase 2: in the
float path around lines 691-710 and the bf16/fp16 path around lines 1379-1397 of
cpp/tensorrt_llm/kernels/heuristic_topk.cuh, update the corresponding secant
logic to fall back to a gvrOrderKey midpoint whenever vhi - vlo is non-finite,
or bypass Phase 2 and enter ordered-key bisection directly. Apply the same guard
to both secant implementations so threshold refinement progresses without
relying solely on Phase 3 repair.
Comment on lines
+2384
to
+2407
| indices = torch.empty((1, index_topk), dtype=torch.int32, device="cuda") | ||
| scratch = torch.empty(index_topk, dtype=dtype, device="cuda") | ||
| aux_indices, aux_logits = _build_radix_aux_buffers(1, index_topk) | ||
| torch.ops.trtllm.indexer_topk_decode( | ||
| logits, | ||
| seq_lens, | ||
| indices, | ||
| 1, | ||
| index_topk, | ||
| pre_idx, | ||
| scratch, | ||
| compress_ratio=4, | ||
| radix_aux_indices=aux_indices, | ||
| radix_aux_logits=aux_logits, | ||
| ) | ||
| torch.cuda.synchronize() | ||
|
|
||
| assert int((indices < 0).sum()) == 0, ( | ||
| f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1" | ||
| ) | ||
| flat = logits[0].float() | ||
| got = flat[indices[0].long()].sort().values | ||
| ref = flat.topk(index_topk).values.sort().values | ||
| assert torch.equal(got, ref), f"{tag}: selected values differ from torch.topk" |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Add an index-uniqueness assertion and initialize indices with a sentinel.
Two gaps weaken this check:
- The value comparison at Line 2407 sorts the selected logits and compares them against the reference. It does not verify that the returned indices are distinct. The repair paths added in this PR emit through
atomicAddon a shared counter (heuristic_topk.cuhLines 911-955 and Lines 1567-1612). A duplicated index combined with an omitted index on a tie plateau produces the same sorted value multiset, so this assertion passes. The tie-plateau test at Line 2439 builds exactly that input shape. - Line 2384 allocates
indiceswithtorch.empty. If the kernel leaves a slot unwritten, the assertion at Line 2401 reads uninitialized device memory. The result is nondeterministic instead of a clear failure.
💚 Proposed fix
- indices = torch.empty((1, index_topk), dtype=torch.int32, device="cuda")
+ indices = torch.full((1, index_topk), -1, dtype=torch.int32, device="cuda")
scratch = torch.empty(index_topk, dtype=dtype, device="cuda") assert int((indices < 0).sum()) == 0, (
f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1"
)
+ sel = indices[0].long()
+ assert int(sel.unique().numel()) == index_topk, (
+ f"{tag}: {index_topk - int(sel.unique().numel())} duplicate indices in the output"
+ )
flat = logits[0].float()
- got = flat[indices[0].long()].sort().values
+ got = flat[sel].sort().values
ref = flat.topk(index_topk).values.sort().values
assert torch.equal(got, ref), f"{tag}: selected values differ from torch.topk"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| indices = torch.empty((1, index_topk), dtype=torch.int32, device="cuda") | |
| scratch = torch.empty(index_topk, dtype=dtype, device="cuda") | |
| aux_indices, aux_logits = _build_radix_aux_buffers(1, index_topk) | |
| torch.ops.trtllm.indexer_topk_decode( | |
| logits, | |
| seq_lens, | |
| indices, | |
| 1, | |
| index_topk, | |
| pre_idx, | |
| scratch, | |
| compress_ratio=4, | |
| radix_aux_indices=aux_indices, | |
| radix_aux_logits=aux_logits, | |
| ) | |
| torch.cuda.synchronize() | |
| assert int((indices < 0).sum()) == 0, ( | |
| f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1" | |
| ) | |
| flat = logits[0].float() | |
| got = flat[indices[0].long()].sort().values | |
| ref = flat.topk(index_topk).values.sort().values | |
| assert torch.equal(got, ref), f"{tag}: selected values differ from torch.topk" | |
| indices = torch.full((1, index_topk), -1, dtype=torch.int32, device="cuda") | |
| scratch = torch.empty(index_topk, dtype=dtype, device="cuda") | |
| aux_indices, aux_logits = _build_radix_aux_buffers(1, index_topk) | |
| torch.ops.trtllm.indexer_topk_decode( | |
| logits, | |
| seq_lens, | |
| indices, | |
| 1, | |
| index_topk, | |
| pre_idx, | |
| scratch, | |
| compress_ratio=4, | |
| radix_aux_indices=aux_indices, | |
| radix_aux_logits=aux_logits, | |
| ) | |
| torch.cuda.synchronize() | |
| assert int((indices < 0).sum()) == 0, ( | |
| f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1" | |
| ) | |
| sel = indices[0].long() | |
| assert int(sel.unique().numel()) == index_topk, ( | |
| f"{tag}: {index_topk - int(sel.unique().numel())} duplicate indices in the output" | |
| ) | |
| flat = logits[0].float() | |
| got = flat[sel].sort().values | |
| ref = flat.topk(index_topk).values.sort().values | |
| assert torch.equal(got, ref), f"{tag}: selected values differ from torch.topk" |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py` around lines 2384 -
2407, Initialize the `indices` tensor in the indexer_topk test with a negative
sentinel instead of `torch.empty`, then add an assertion that `indices[0]`
contains `index_topk` unique values before comparing selected logits. Keep the
existing negative-sentinel check and value comparison unchanged.
Comment on lines
+2410
to
+2454
| @skip_pre_blackwell | ||
| @pytest.mark.parametrize("index_topk", [512, 1024, 2048]) | ||
| @pytest.mark.parametrize("num_tokens", [65536, 131072]) | ||
| @pytest.mark.parametrize( | ||
| "dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"] | ||
| ) | ||
| @pytest.mark.parametrize("hint", ["bottom_k", "uniform_max", "random"]) | ||
| def test_indexer_topk_decode_gvr_hostile_hint(index_topk, num_tokens, dtype, hint): | ||
| """A hint that points away from the top-K must not change the result. | ||
|
|
||
| ``uniform_max`` (every slot = argmax) additionally collapses Phase 1's | ||
| min/max bracket to a point, which used to short-circuit the kernel into | ||
| emitting row[0:K]. | ||
| """ | ||
| torch.manual_seed(1234) | ||
| logits = torch.randn(num_tokens, dtype=torch.float32, device="cuda").to(dtype) | ||
| flat = logits.float() | ||
| if hint == "bottom_k": | ||
| pre = flat.topk(index_topk, largest=False).indices | ||
| elif hint == "uniform_max": | ||
| pre = flat.argmax().repeat(index_topk) | ||
| else: | ||
| pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") | ||
| _gvr_decode_exact_check(logits, pre, index_topk, f"hint={hint}") | ||
|
|
||
|
|
||
| @skip_pre_blackwell | ||
| @pytest.mark.parametrize("index_topk", [512, 1024, 2048]) | ||
| @pytest.mark.parametrize("n_tie", [6000, 20000, 100000]) | ||
| def test_indexer_topk_decode_gvr_tie_plateau(index_topk, n_tie): | ||
| """More ties at the K-th value than the candidate buffer can hold. | ||
|
|
||
| No threshold yields a candidate count in [K, kC], so the search must | ||
| collapse the bracket and emit "everything strictly greater + arbitrary | ||
| ties" — dropping strictly-greater entries instead is a wrong top-K. | ||
| """ | ||
| torch.manual_seed(1234) | ||
| num_tokens = 131072 | ||
| n_above = index_topk // 2 | ||
| logits = torch.full((num_tokens,), -1.0, dtype=torch.float32, device="cuda") | ||
| logits[:n_above] = torch.linspace(2.0, 3.0, n_above, device="cuda") | ||
| logits[n_above : n_above + n_tie] = 1.0 | ||
| logits = logits[torch.randperm(num_tokens, device="cuda")].contiguous() | ||
| pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") | ||
| _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie}") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find how test_indexer_topk.py is registered in the CI and QA test lists.
fd -t f -e yml -e txt . tests/integration/test_lists | xargs rg -n 'test_indexer_topk|thop/parallel|_torch/thop'Repository: NVIDIA/TensorRT-LLM
Length of output: 3382
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- test file structure ---'
ast-grep outline tests/unittest/_torch/thop/parallel/test_indexer_topk.py
printf '%s\n' '--- relevant test code ---'
sed -n '2350,2460p' tests/unittest/_torch/thop/parallel/test_indexer_topk.py
printf '%s\n' '--- helper usages ---'
rg -n -C 4 '_gvr_decode_exact_check|test_indexer_topk_decode_gvr' tests/unittest/_torch/thop/parallel/test_indexer_topk.py
printf '%s\n' '--- kernel code ---'
sed -n '1540,1630p' cpp/tensorrt_llm/kernels/heuristic_topk.cuh
printf '%s\n' '--- CI list context ---'
for f in tests/integration/test_lists/test-db/l0_h100.yml \
tests/integration/test_lists/test-db/l0_b200.yml \
tests/integration/test_lists/test-db/l0_gb300.yml \
tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml; do
echo "### $f"
sed -n '1,75p' "$f" | nl -ba | sed -n '/thop\/parallel/,+3p'
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 14870
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- exact test-list registrations ---'
rg -n -C 2 'unittest/_torch/thop/parallel($|/)' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa
printf '%s\n' '--- uniqueness assertions in this test module ---'
rg -n -C 3 'unique|duplicate|indices.*sort|sort.*indices|torch\.unique' \
tests/unittest/_torch/thop/parallel/test_indexer_topk.py
printf '%s\n' '--- kernel template and dispatch context ---'
rg -n -C 8 'gvrOrderKey|Collapsed bracket|indexer_topk_decode|heuristicTopK|topk.*decode' \
cpp/tensorrt_llm/kernels/heuristic_topk.cuh cpp tests/unittest/_torch/thop/parallel/test_indexer_topk.py
printf '%s\n' '--- relevant list entries with line numbers ---'
for f in tests/integration/test_lists/test-db/l0_h100.yml \
tests/integration/test_lists/test-db/l0_b200.yml \
tests/integration/test_lists/test-db/l0_gb300.yml \
tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml; do
echo "### $f"
grep -n -A 3 -B 2 'unittest/_torch/thop/parallel' "$f" || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
python3 - <<'PY'
import torch
logits = torch.tensor([3.0, 2.0, 2.0, 1.0])
index_topk = 3
duplicate_indices = torch.tensor([0, 1, 1])
got = logits[duplicate_indices].sort().values
ref = logits.topk(index_topk).values.sort().values
print("value-only assertion passes:", torch.equal(got, ref))
print("indices are unique:", duplicate_indices.unique().numel() == index_topk)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 275
🏁 Script executed:
python3 - <<'PY'
logits = [3.0, 2.0, 2.0, 1.0]
index_topk = 3
duplicate_indices = [0, 1, 1]
got = sorted((logits[i] for i in duplicate_indices))[-index_topk:]
ref = sorted(logits)[-index_topk:]
print("value-only assertion passes:", got == ref)
print("indices are unique:", len(set(duplicate_indices)) == index_topk)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 217
Assert unique output indices in _gvr_decode_exact_check.
- Added tests:
test_indexer_topk_decode_gvr_hostile_hint(54 cases) andtest_indexer_topk_decode_gvr_tie_plateau(9 cases). - CI coverage:
tests/integration/test_lists/test-db/l0_h100.yml,l0_b200.yml,l0_gb300.yml, andl0_gb300_multi_gpus.ymlcollect this file through directory entries. - Coverage verdict: insufficient. The helper compares only selected values. Duplicate indices can pass when tied logits are present. Add an explicit uniqueness assertion.
- The hostile-hint test already covers fp32, bf16, and fp16.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py` around lines 2410 -
2454, The helper _gvr_decode_exact_check currently validates selected values
without ensuring indices are distinct. Update this helper to explicitly assert
that the returned top-K indices are unique, while preserving its existing value
comparison and coverage for tied logits and all dtypes.
Source: Path instructions
Comment on lines
+2436
to
+2454
| @skip_pre_blackwell | ||
| @pytest.mark.parametrize("index_topk", [512, 1024, 2048]) | ||
| @pytest.mark.parametrize("n_tie", [6000, 20000, 100000]) | ||
| def test_indexer_topk_decode_gvr_tie_plateau(index_topk, n_tie): | ||
| """More ties at the K-th value than the candidate buffer can hold. | ||
|
|
||
| No threshold yields a candidate count in [K, kC], so the search must | ||
| collapse the bracket and emit "everything strictly greater + arbitrary | ||
| ties" — dropping strictly-greater entries instead is a wrong top-K. | ||
| """ | ||
| torch.manual_seed(1234) | ||
| num_tokens = 131072 | ||
| n_above = index_topk // 2 | ||
| logits = torch.full((num_tokens,), -1.0, dtype=torch.float32, device="cuda") | ||
| logits[:n_above] = torch.linspace(2.0, 3.0, n_above, device="cuda") | ||
| logits[n_above : n_above + n_tie] = 1.0 | ||
| logits = logits[torch.randperm(num_tokens, device="cuda")].contiguous() | ||
| pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") | ||
| _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie}") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Parameterize the tie-plateau test over dtype to cover the bf16/fp16 repair path.
This test runs fp32 only. The PR adds a separate collapsed-bracket direct-emit block to the bf16/fp16 path in cpp/tensorrt_llm/kernels/heuristic_topk.cuh at Lines 1567-1612. The hostile-hint test at Line 2417 covers bf16 and fp16, but it uses torch.randn logits, which do not produce a tie plateau larger than kC. No test executes the bf16/fp16 direct-emit block.
If you add the dtype parameter, choose plateau and above-plateau values that are exactly representable in bf16 and fp16. torch.linspace(2.0, 3.0, n_above) rounds to about 256 distinct bf16 values, so the strictly-greater set gains its own ties. Powers of two above the plateau keep the values distinct in every dtype.
💚 Proposed change
`@skip_pre_blackwell`
`@pytest.mark.parametrize`("index_topk", [512, 1024, 2048])
`@pytest.mark.parametrize`("n_tie", [6000, 20000, 100000])
-def test_indexer_topk_decode_gvr_tie_plateau(index_topk, n_tie):
+@pytest.mark.parametrize(
+ "dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"]
+)
+def test_indexer_topk_decode_gvr_tie_plateau(index_topk, n_tie, dtype):
torch.manual_seed(1234)
num_tokens = 131072
n_above = index_topk // 2
logits = torch.full((num_tokens,), -1.0, dtype=torch.float32, device="cuda")
- logits[:n_above] = torch.linspace(2.0, 3.0, n_above, device="cuda")
+ # Exactly representable in fp32, bf16, and fp16; all values distinct.
+ logits[:n_above] = 2.0 + torch.arange(n_above, device="cuda", dtype=torch.float32) / 512.0
logits[n_above : n_above + n_tie] = 1.0
- logits = logits[torch.randperm(num_tokens, device="cuda")].contiguous()
+ logits = logits[torch.randperm(num_tokens, device="cuda")].to(dtype).contiguous()
pre = torch.randint(0, num_tokens, (index_topk,), device="cuda")
- _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie}")
+ _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie} dtype={dtype}")For fp16, confirm that 2.0 + n_above/512.0 stays below the fp16 spacing limit. With index_topk=2048, n_above=1024, the maximum is 4.0, and the fp16 spacing at 4.0 is 1/512. Reduce the divisor or n_above if a collision appears.
As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/thop/parallel/test_indexer_topk.py` around lines 2436 -
2454, Parameterize test_indexer_topk_decode_gvr_tie_plateau over fp32, bf16, and
fp16, converting logits to the selected dtype before execution. Replace
torch.linspace with exactly representable power-of-two-based values so
strictly-greater entries remain distinct in every dtype, and keep the tie
plateau exactly representable; preserve the existing hostile tie sizes and exact
top-K validation so the bf16/fp16 collapsed-bracket direct-emit path is
exercised.
Source: Path instructions
Collaborator
|
PR_Github #65551 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #65551 [ run ] completed with state
|
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
2 participants
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 heuristic (GVR) indexer top-K decode path picks a value threshold whose candidate count lands in
[K, kC]and then selects the top-K out of those candidates. Three input shapes defeat that search and make the kernel return a silently wrong top-K — no error, no diagnostic, just wrong indices:MAX_REFINE_ITERS = 15; on non-convergence thedone = 2fallback can pickval_hioutright, and the Phase-3 retry loop only ever triggered on the overflow side (cand_count > kCC). A threshold admitting fewer thanKcandidates went straight through: the collect emitted what it had and the Phase-4 tail padded the rest with index-1.outputIndices[i] = i— the head of the row, not a top-K at all.kCelements share the K-th value, no threshold yields a count in[K, kC]. The collect clamps atkCand dropped strictly-greater entries.All three are driven by hint quality, not by unusual logits, so they are reachable in production whenever a layer's temporal locality breaks down. On the shipped op an anti-correlated hint at
N = 65536, K = 512returns 512 of 512 wrong indices.Production DSv4 decode captures hit this unaided (fp32, BS=1, captured
preidx.inas the hint):-1N = 131075is inside the shipped GVR routing window (numColumns < 200000), so this is a live defect rather than a latent one.Fix
Applied to both
gvrTopKJob(fp32) andgvrTopKJobDtype(bf16/fp16):row[0:K]; it resets to the widest trusted bracket and falls through. The hint may only affect speed, never the answer.count(-FLT_MAX) = #finite >= K,count(FLT_MAX) = 0 < K), because Phase 1 seedsval_lo/val_hifrom the hinted min/max with invented counts (M + M/4,1) that can leave both ends on the same side of the K-th value. The loop then bisects on the order-preservinguint32image of the key space, so the bracket provably collapses to adjacent representable values in<= 32steps instead of relying on a float average that has no such bound.kCelements at the threshold, emit directly: everything strictly above (fewer thanKby construction) plus arbitrary ties, which is a valid top-K. Guarded on the collapse test so it can never run on a non-collapsed bracket.The converged fast path (
done == 1) is untouched — only rows whose secant failed enter any of this.Performance
B200, cold-L2, real DSv4 decode captures (both models x 5 ISL rungs x 90 rows, fp32), fixed vs unfixed kernel:
Neutral to slightly faster on nine of ten buckets. The single outlier is V4-Pro ISL=1M — the bucket containing the non-converging layer 40, which now pays the bisection instead of returning a wrong answer.
Exactness across the same measurement set: 353/353 real-capture cells and 135/135 adversarial cells now exact, against 4 and 54 inexact before.
Test Coverage
tests/unittest/_torch/thop/parallel/test_indexer_topk.py:test_indexer_topk_decode_gvr_hostile_hint— hints pointing away from the true top-K (bottom_k,uniform_max,random) xindex_topkin {512, 1024, 2048} xnum_tokensin {65536, 131072} x {fp32, bf16, fp16}.uniform_maxalso covers the degenerate-bracket path.test_indexer_topk_decode_gvr_tie_plateau— more ties at the K-th value than the candidate buffer holds (n_tiein {6000, 20000, 100000}).Both classes fail on the pre-fix kernel and pass after. Existing
test_indexer_topk_decode_dist*coverage is unchanged.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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.
🤖 Generated with Claude Code
Dev Engineer Review
heuristic_topk.cuhto repair incorrect GVR threshold refinement.uint32key conversion and reverse conversion.MAX_REPAIR_ITERSwith a bounded 40-iteration repair budget.QA Engineer Review
test_indexer_topk_decode_gvr_hostile_hint.test_indexer_topk_decode_gvr_tie_plateau.torch.topk.tests/integration/test_lists/were reported.