From 251cb523bded270899d75a268df59c2bc968de14 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Wed, 12 Aug 2026 09:22:30 +0000 Subject: [PATCH 1/2] [None][fix] GVR indexer top-K: repair the non-converged threshold search MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/heuristic_topk.cuh | 306 ++++++++++++++++-- .../_torch/thop/parallel/test_indexer_topk.py | 102 ++++++ 2 files changed, 376 insertions(+), 32 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh index f21a4d7cf2c6..973ed8a38d90 100644 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh @@ -164,6 +164,13 @@ constexpr int SAFETY_MARGIN = 2048; constexpr int MAX_CANDIDATES = TOP_K + SAFETY_MARGIN * 2; // 6144 constexpr int MAX_REFINE_ITERS = 15; +// Phase-3 repair bisection budget. The repair bisects on the order-preserving +// uint32 image of the float key space (see floatToOrderedKey), so the bracket +// provably collapses to adjacent representable values in <= 32 steps; 40 is +// that bound plus slack. Only rows whose Phase-2 secant did NOT converge +// (done != 1) ever enter the loop, and it exits as soon as the candidate +// count lands in [kK, kCC] — the converged fast path is untouched. +constexpr int MAX_REPAIR_ITERS = 40; constexpr int NUM_BINS = 2048; static_assert(TOP_K % BLOCK_SIZE == 0); @@ -419,6 +426,26 @@ __device__ __forceinline__ float warpReduceMax(float val) #endif +// ============================================================================ +// Order-preserving float <-> uint32 map (arch-independent) +// ============================================================================ +// Same bijection as floatToOrderedUint/orderedUintToFloat above, but defined +// for every __CUDA_ARCH__ (those are inside the >= 800 reduction block). Used +// by the Phase-3 repair to bisect on the key space itself: `a < b` for finite +// floats iff `gvrOrderKey(a) < gvrOrderKey(b)`, so a uint32 midpoint always +// makes progress and the bracket collapses to adjacent representable values +// in at most 32 steps — a float-average midpoint has no such bound. +__device__ __forceinline__ unsigned gvrOrderKey(float f) +{ + unsigned u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +__device__ __forceinline__ float gvrOrderKeyToFloat(unsigned u) +{ + return __uint_as_float((u & 0x80000000u) ? (u & ~0x80000000u) : ~u); +} + // ============================================================================ // Device: Block count ≥ threshold in GLOBAL memory (1-sync pattern) // ============================================================================ @@ -661,15 +688,25 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con } __syncthreads(); + // 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(); } // ================================================================ @@ -775,23 +812,61 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con // When done==1, Phase 2 already verified the candidate count is in // [kK, kCC]; skip the redundant full-N blockCountGE re-check. + // + // Otherwise the Phase-2 secant did not converge and `threshold` carries + // no guarantee at all. The repair below restores the invariant the + // collect depends on — cand_count >= kK — on BOTH sides: + // + // cand_count > kCC : candidates overflow smem->keys[]; the collect + // silently drops the excess (my_write_pos < kCC). + // cand_count < kK : the collect emits fewer than K entries and the + // Phase-4 tail pads the rest with index -1, i.e. a + // silently WRONG top-K. The previous loop guarded + // only the overflow side (`cand_count > kCC`), so + // an undershooting threshold — which the `done=2` + // fallback above can pick outright via val_hi — + // went straight through. Reproduced on production + // DSv4 decode captures: V4-Flash K=512 N=131075 + // layers 22/24 (283 / 87 slots left at -1) and + // V4-Pro K=1024 N=262127 layer 40 (550 slots), + // all on rows whose temporal hint was poor + // (hit-rate 0.02 - 0.12), which starts the secant + // from a bracket far off the true K-th value. if (smem->done != 1) { blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0 && smem->cand_count > kCC) - smem->val_lo = smem->threshold; + // Reset the bracket to endpoints whose counts are KNOWN. Phase 1 seeds + // val_lo/val_hi from the min/max of the *hinted* values with invented + // counts (M + M/4, 1); neither is measured, so a poor hint can leave + // both ends on the same side of the K-th value. The collapse handling + // below relies on count(val_lo) >= kK > count(val_hi), so anchor the + // untested end at a float extreme: count(-FLT_MAX) = #finite >= kK and + // count(FLT_MAX) = 0 < kK for any normal row. + if (tid == 0) + { + int c = smem->cand_count; + if (c > kCC) + { + smem->val_lo = smem->threshold; + smem->val_hi = FLT_MAX; + } + else if (c < kK) + { + smem->val_hi = smem->threshold; + smem->val_lo = -FLT_MAX; + } + } __syncthreads(); - for (int retry = 0; retry < 10 && smem->cand_count > kCC; retry++) + // Invariant maintained below: count(val_lo) >= kK. + for (int retry = 0; retry < MAX_REPAIR_ITERS && (smem->cand_count > kCC || smem->cand_count < kK); retry++) { + unsigned const klo = gvrOrderKey(smem->val_lo); + unsigned const khi = gvrOrderKey(smem->val_hi); + if (khi <= klo + 1u) + break; // bracket collapsed to adjacent representable values if (tid == 0) - { - float lo = smem->val_lo, hi = smem->val_hi; - float mid = (lo + hi) * 0.5f; - if (mid == lo) - mid = hi; - smem->threshold = mid; - } + smem->threshold = gvrOrderKeyToFloat(klo + ((khi - klo) >> 1)); __syncthreads(); blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); if (tid == 0) @@ -804,6 +879,80 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con } __syncthreads(); } + + // Still short of kK => the bisection collapsed. Fall back to val_lo, + // which by the invariant admits >= kK elements (or the row simply has + // fewer than kK finite entries, in which case the Phase-4 tail pad is + // the correct answer). blockCountGE also refreshes per_thread_counts, + // which the collect below consumes. + if (smem->cand_count < kK) + { + if (tid == 0) + smem->threshold = smem->val_lo; + __syncthreads(); + blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); + } + // blockCountGE publishes cand_count from tid 0 only; the branch below + // must be uniform across the block. + __syncthreads(); + + // Collapsed bracket with more than kCC elements at the threshold: + // every value in [val_lo, val_hi) equals val_lo, so the answer is + // "all elements strictly above val_lo" (fewer than kK of them, since + // count(val_hi) < kK) plus arbitrary ties at val_lo. The candidate + // buffer cannot hold them all, so emit directly instead — any tie + // subset is a valid top-K. + // The direct emit below is only valid once the bracket has collapsed: + // it assumes count(> thr) < kK, which is exactly "val_hi is the next + // representable value above val_lo and count(val_hi) < kK". If the + // loop ran out of iterations without collapsing (it cannot, given + // MAX_REPAIR_ITERS >= 32, but the guard keeps that an invariant rather + // than an assumption) fall through to the ordinary collect. + if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) + { + float const thr = smem->threshold; + if (tid == 0) + smem->out_count = 0; + __syncthreads(); + for (int i = tid; i < N; i += BLOCK_SIZE) + { + float const v = __ldg(&input[i]); + if (v > thr) + { + int const p = atomicAdd(&smem->out_count, 1); + if (p < kK) + { + outputValues[p] = v; + outputIndices[p] = i; + } + } + } + __syncthreads(); + int const n_gt = min(smem->out_count, kK); + if (tid == 0) + smem->out_count = n_gt; + __syncthreads(); + for (int i = tid; i < N && smem->out_count < kK; i += BLOCK_SIZE) + { + float const v = __ldg(&input[i]); + if (v == thr) + { + int const p = atomicAdd(&smem->out_count, 1); + if (p < kK) + { + outputValues[p] = v; + outputIndices[p] = i; + } + } + } + __syncthreads(); + for (int i = min(smem->out_count, kK) + tid; i < kK; i += BLOCK_SIZE) + { + outputValues[i] = -FLT_MAX; + outputIndices[i] = -1; + } + return; + } } // Reuse per-thread counts cached by the last blockCountGE call (saves @@ -1227,15 +1376,25 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i } __syncthreads(); + // 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] = __ldg(&input[i]); // both InputT, no convert - } - 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(); } // ================================================================ @@ -1339,23 +1498,40 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i // Phase 3 — Ballot-free candidate collect // ================================================================ + // Mirror of the fp32 Phase-3 repair in gvrTopKJob — see the comment block + // there for why the undershoot side (cand_count < kK) must be repaired: + // without it the collect emits < K entries and the tail is padded with + // index -1, i.e. a silently wrong top-K. if (smem->done != 1) { blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0 && smem->cand_count > kCC) - smem->val_lo = smem->threshold; + // See the fp32 path: anchor the untested bracket end at a float extreme + // so count(val_lo) >= kK > count(val_hi) holds by construction. + if (tid == 0) + { + int c = smem->cand_count; + if (c > kCC) + { + smem->val_lo = smem->threshold; + smem->val_hi = FLT_MAX; + } + else if (c < kK) + { + smem->val_hi = smem->threshold; + smem->val_lo = -FLT_MAX; + } + } __syncthreads(); - for (int retry = 0; retry < 10 && smem->cand_count > kCC; retry++) + // Invariant maintained below: count(val_lo) >= kK. + for (int retry = 0; retry < MAX_REPAIR_ITERS && (smem->cand_count > kCC || smem->cand_count < kK); retry++) { + unsigned const klo = gvrOrderKey(smem->val_lo); + unsigned const khi = gvrOrderKey(smem->val_hi); + if (khi <= klo + 1u) + break; // bracket collapsed to adjacent representable values if (tid == 0) - { - float lo = smem->val_lo, hi = smem->val_hi; - float mid = (lo + hi) * 0.5f; - if (mid == lo) - mid = hi; - smem->threshold = mid; - } + smem->threshold = gvrOrderKeyToFloat(klo + ((khi - klo) >> 1)); __syncthreads(); blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); if (tid == 0) @@ -1368,6 +1544,72 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i } __syncthreads(); } + + if (smem->cand_count < kK) + { + if (tid == 0) + smem->threshold = smem->val_lo; + __syncthreads(); + blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); + } + // blockCountGEDtype publishes cand_count from tid 0 only; the branch + // below must be uniform across the block. + __syncthreads(); + + // Collapsed bracket with > kCC elements at the threshold: emit the + // strictly-greater set plus arbitrary ties directly (see fp32 path). + // The direct emit below is only valid once the bracket has collapsed: + // it assumes count(> thr) < kK, which is exactly "val_hi is the next + // representable value above val_lo and count(val_hi) < kK". If the + // loop ran out of iterations without collapsing (it cannot, given + // MAX_REPAIR_ITERS >= 32, but the guard keeps that an invariant rather + // than an assumption) fall through to the ordinary collect. + if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) + { + float const thr = smem->threshold; + if (tid == 0) + smem->out_count = 0; + __syncthreads(); + for (int i = tid; i < N; i += BLOCK_SIZE) + { + float const v = Trait::to_fp32(__ldg(&input[i])); + if (v > thr) + { + int const p = atomicAdd(&smem->out_count, 1); + if (p < kK) + { + outputValues[p] = Trait::from_fp32(v); + outputIndices[p] = i; + } + } + } + __syncthreads(); + int const n_gt = min(smem->out_count, kK); + if (tid == 0) + smem->out_count = n_gt; + __syncthreads(); + for (int i = tid; i < N && smem->out_count < kK; i += BLOCK_SIZE) + { + float const v = Trait::to_fp32(__ldg(&input[i])); + if (v == thr) + { + int const p = atomicAdd(&smem->out_count, 1); + if (p < kK) + { + outputValues[p] = Trait::from_fp32(v); + outputIndices[p] = i; + } + } + } + __syncthreads(); + InputT const neg_max = Trait::from_fp32(-FLT_MAX); + for (int i = min(smem->out_count, kK) + tid; i < kK; i += BLOCK_SIZE) + { + outputValues[i] = neg_max; + outputIndices[i] = -1; + } + return; + } } int my_total_qual = smem->per_thread_counts[tid]; diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index aa9a327d2f89..cbd56d228f5d 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -2350,3 +2350,105 @@ def test_prefill_overflow_policy_overflow( dtype, row_start_offset=row_start_offset, ) + + +# ============================================================================ +# GVR Phase-3 threshold-repair regressions +# ============================================================================ +# The heuristic (GVR) decode path locates a value threshold whose candidate +# count lands in [K, kC] and then selects the top-K out of those candidates. +# Three inputs used to defeat that search and produce a silently WRONG top-K +# (no error, no -1-free output guarantee): +# +# 1. undershoot — the search ends with fewer than K candidates, the +# collect emits them all and pads the tail with -1. +# 2. degenerate hint— every hinted value identical, so Phase 1 builds an +# empty bracket; the kernel emitted the first K entries +# of the row verbatim (the head of the row, not a top-K). +# 3. tie plateau — more than kC elements share the K-th value, so NO +# threshold yields a count in [K, kC]; the candidate +# buffer overflowed and dropped strictly-greater entries. +# +# All three are hint-quality driven, i.e. they need no special logits — only a +# hint that points away from the true top-K, which production hits whenever a +# layer's temporal locality breaks down. + + +def _gvr_decode_exact_check(logits_row, pre_idx_row, index_topk, tag): + """Run indexer_topk_decode (cr=4, BS=1) and assert a tie-aware exact top-K.""" + n = logits_row.shape[-1] + dtype = logits_row.dtype + logits = logits_row.view(1, n).contiguous() + pre_idx = pre_idx_row.view(1, index_topk).to(torch.int32).contiguous() + seq_lens = torch.full((1,), n * 4, dtype=torch.int32, device="cuda") + 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" + + +@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}") From 303781b0a4896cb5b9fdd96840d8d02f7296056f Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:24:59 +0000 Subject: [PATCH 2/2] [None][fix] GVR indexer top-K: skip the secant after a degenerate-hint reset; harden the repair tests Review follow-up (three of the four automated-review findings; the fourth, a CI test-list registration note, needs no change since the file's suite is already in the l0 lists): * Degenerate-hint reset now sets done = 2 instead of re-entering the secant. The secant interpolates on the LINEAR float scale, so across the reset's (-FLT_MAX, FLT_MAX) trusted bracket it cannot converge within MAX_REFINE_ITERS and burned 15 full-N counting passes before the Phase-3 repair fixed the row anyway. Phase 2's single seed probe is kept (it can still promote to the done = 1 fast path); a non-converging row now goes straight to the ordered-key bisection, which collapses in <= 32 passes. Both drivers. No behavior change for rows with a usable hint bracket. * Tie-plateau test parameterized over fp32/bf16/fp16: the reduced-precision driver has its own collapsed-bracket direct-emit block, which previously had no test coverage. The plateau (1.0) and floor (-1.0) are exact in all three dtypes; casting may merge strictly-greater values with each other, which the top-K contract tolerates. * The exactness helper now allocates outputIndices with a -1 sentinel (unwritten slots must trip the assertions, not inherit torch.empty garbage) and asserts index distinctness: the repair paths emit through an atomic counter, and a duplicated index paired with an omitted one on a tie plateau would leave the sorted value multiset unchanged. Validation (B200, standalone build of the exact repo kernel): 135/135 adversarial cells and 353/353 real-capture cells (both models x 5 ISL rungs x all layers, plus a 98-cell hint-degradation grid on a real 1M row) exact; 27/27 tie-plateau cells across the three dtypes return K distinct indices. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/heuristic_topk.cuh | 107 +++++------------- .../_torch/thop/parallel/test_indexer_topk.py | 49 ++++---- 2 files changed, 53 insertions(+), 103 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh index 973ed8a38d90..d75933bd5d38 100644 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh @@ -164,12 +164,8 @@ constexpr int SAFETY_MARGIN = 2048; constexpr int MAX_CANDIDATES = TOP_K + SAFETY_MARGIN * 2; // 6144 constexpr int MAX_REFINE_ITERS = 15; -// Phase-3 repair bisection budget. The repair bisects on the order-preserving -// uint32 image of the float key space (see floatToOrderedKey), so the bracket -// provably collapses to adjacent representable values in <= 32 steps; 40 is -// that bound plus slack. Only rows whose Phase-2 secant did NOT converge -// (done != 1) ever enter the loop, and it exits as soon as the candidate -// count lands in [kK, kCC] — the converged fast path is untouched. +// Phase-3 repair budget: bisecting on the uint32 key image collapses any +// bracket to adjacent floats in <= 32 steps; 40 adds slack. constexpr int MAX_REPAIR_ITERS = 40; constexpr int NUM_BINS = 2048; @@ -429,12 +425,9 @@ __device__ __forceinline__ float warpReduceMax(float val) // ============================================================================ // Order-preserving float <-> uint32 map (arch-independent) // ============================================================================ -// Same bijection as floatToOrderedUint/orderedUintToFloat above, but defined -// for every __CUDA_ARCH__ (those are inside the >= 800 reduction block). Used -// by the Phase-3 repair to bisect on the key space itself: `a < b` for finite -// floats iff `gvrOrderKey(a) < gvrOrderKey(b)`, so a uint32 midpoint always -// makes progress and the bracket collapses to adjacent representable values -// in at most 32 steps — a float-average midpoint has no such bound. +// Same bijection as floatToOrderedUint above but defined for every +// __CUDA_ARCH__; the Phase-3 repair bisects on this key image so the +// bracket provably collapses (a float-average midpoint has no such bound). __device__ __forceinline__ unsigned gvrOrderKey(float f) { unsigned u = __float_as_uint(f); @@ -688,12 +681,10 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con } __syncthreads(); - // 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. + // Degenerate hint (all gathered values identical or out of range): + // reset to a trusted bracket instead of emitting row[0:K]; done = 2 + // skips the secant (it cannot converge on a full-range bracket) and + // hands the row to the Phase-3 repair. The hint only affects speed. if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) { if (tid == 0) @@ -704,7 +695,7 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con smem->cnt_lo = N; smem->cnt_hi = 0; smem->threshold = seed; - smem->done = 0; + smem->done = 2; } __syncthreads(); } @@ -810,38 +801,17 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con // Phase 3 (GVR Verify) — Ballot-free candidate collect // ================================================================ - // When done==1, Phase 2 already verified the candidate count is in - // [kK, kCC]; skip the redundant full-N blockCountGE re-check. - // - // Otherwise the Phase-2 secant did not converge and `threshold` carries - // no guarantee at all. The repair below restores the invariant the - // collect depends on — cand_count >= kK — on BOTH sides: - // - // cand_count > kCC : candidates overflow smem->keys[]; the collect - // silently drops the excess (my_write_pos < kCC). - // cand_count < kK : the collect emits fewer than K entries and the - // Phase-4 tail pads the rest with index -1, i.e. a - // silently WRONG top-K. The previous loop guarded - // only the overflow side (`cand_count > kCC`), so - // an undershooting threshold — which the `done=2` - // fallback above can pick outright via val_hi — - // went straight through. Reproduced on production - // DSv4 decode captures: V4-Flash K=512 N=131075 - // layers 22/24 (283 / 87 slots left at -1) and - // V4-Pro K=1024 N=262127 layer 40 (550 slots), - // all on rows whose temporal hint was poor - // (hit-rate 0.02 - 0.12), which starts the secant - // from a bracket far off the true K-th value. + // done==1: Phase 2 verified cand_count in [kK, kCC]; skip the re-check. + // Otherwise the secant did not converge and `threshold` carries no + // guarantee: repair BOTH sides (the old loop only handled overflow, so + // an undershooting threshold shipped a -1-padded, silently wrong top-K). if (smem->done != 1) { blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - // Reset the bracket to endpoints whose counts are KNOWN. Phase 1 seeds - // val_lo/val_hi from the min/max of the *hinted* values with invented - // counts (M + M/4, 1); neither is measured, so a poor hint can leave - // both ends on the same side of the K-th value. The collapse handling - // below relies on count(val_lo) >= kK > count(val_hi), so anchor the - // untested end at a float extreme: count(-FLT_MAX) = #finite >= kK and - // count(FLT_MAX) = 0 < kK for any normal row. + // Anchor the untested bracket end at a float extreme: Phase 1 seeds + // both ends from HINTED values with invented counts, so they can sit + // on the same side of the K-th value. count(-FLT_MAX) >= kK, + // count(FLT_MAX) = 0. if (tid == 0) { int c = smem->cand_count; @@ -880,11 +850,9 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con __syncthreads(); } - // Still short of kK => the bisection collapsed. Fall back to val_lo, - // which by the invariant admits >= kK elements (or the row simply has - // fewer than kK finite entries, in which case the Phase-4 tail pad is - // the correct answer). blockCountGE also refreshes per_thread_counts, - // which the collect below consumes. + // Still short of kK: the bracket collapsed; val_lo admits >= kK by + // the anchor invariant (or the row has < kK finite entries and the + // -1 tail pad is the correct answer). if (smem->cand_count < kK) { if (tid == 0) @@ -896,18 +864,10 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con // must be uniform across the block. __syncthreads(); - // Collapsed bracket with more than kCC elements at the threshold: - // every value in [val_lo, val_hi) equals val_lo, so the answer is - // "all elements strictly above val_lo" (fewer than kK of them, since - // count(val_hi) < kK) plus arbitrary ties at val_lo. The candidate - // buffer cannot hold them all, so emit directly instead — any tie - // subset is a valid top-K. - // The direct emit below is only valid once the bracket has collapsed: - // it assumes count(> thr) < kK, which is exactly "val_hi is the next - // representable value above val_lo and count(val_hi) < kK". If the - // loop ran out of iterations without collapsing (it cannot, given - // MAX_REPAIR_ITERS >= 32, but the guard keeps that an invariant rather - // than an assumption) fall through to the ordinary collect. + // Collapsed bracket still over kCC = a tie plateau wider than the + // candidate buffer: emit everything strictly above val_lo (< kK by + // construction) plus arbitrary ties — a valid tie-aware top-K. The + // adjacency guard keeps the emit sound if the loop ever ran dry. if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) { float const thr = smem->threshold; @@ -1376,12 +1336,10 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i } __syncthreads(); - // 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. + // Degenerate hint (all gathered values identical or out of range): + // reset to a trusted bracket instead of emitting row[0:K]; done = 2 + // skips the secant (it cannot converge on a full-range bracket) and + // hands the row to the Phase-3 repair. The hint only affects speed. if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) { if (tid == 0) @@ -1392,7 +1350,7 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i smem->cnt_lo = N; smem->cnt_hi = 0; smem->threshold = seed; - smem->done = 0; + smem->done = 2; } __syncthreads(); } @@ -1498,10 +1456,7 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i // Phase 3 — Ballot-free candidate collect // ================================================================ - // Mirror of the fp32 Phase-3 repair in gvrTopKJob — see the comment block - // there for why the undershoot side (cand_count < kK) must be repaired: - // without it the collect emits < K entries and the tail is padded with - // index -1, i.e. a silently wrong top-K. + // Mirror of the fp32 Phase-3 repair in gvrTopKJob (see comments there). if (smem->done != 1) { blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index cbd56d228f5d..33ebdb2c82ca 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -2353,25 +2353,10 @@ def test_prefill_overflow_policy_overflow( # ============================================================================ -# GVR Phase-3 threshold-repair regressions +# GVR Phase-3 threshold-repair regressions: hints that defeat the threshold +# search (undershoot / degenerate hint / tie plateau wider than kC) used to +# produce a silently wrong top-K (-1 pads or row[0:K]). # ============================================================================ -# The heuristic (GVR) decode path locates a value threshold whose candidate -# count lands in [K, kC] and then selects the top-K out of those candidates. -# Three inputs used to defeat that search and produce a silently WRONG top-K -# (no error, no -1-free output guarantee): -# -# 1. undershoot — the search ends with fewer than K candidates, the -# collect emits them all and pads the tail with -1. -# 2. degenerate hint— every hinted value identical, so Phase 1 builds an -# empty bracket; the kernel emitted the first K entries -# of the row verbatim (the head of the row, not a top-K). -# 3. tie plateau — more than kC elements share the K-th value, so NO -# threshold yields a count in [K, kC]; the candidate -# buffer overflowed and dropped strictly-greater entries. -# -# All three are hint-quality driven, i.e. they need no special logits — only a -# hint that points away from the true top-K, which production hits whenever a -# layer's temporal locality breaks down. def _gvr_decode_exact_check(logits_row, pre_idx_row, index_topk, tag): @@ -2381,7 +2366,8 @@ def _gvr_decode_exact_check(logits_row, pre_idx_row, index_topk, tag): logits = logits_row.view(1, n).contiguous() pre_idx = pre_idx_row.view(1, index_topk).to(torch.int32).contiguous() seq_lens = torch.full((1,), n * 4, dtype=torch.int32, device="cuda") - indices = torch.empty((1, index_topk), dtype=torch.int32, device="cuda") + # -1 sentinel so unwritten slots trip the assertions below. + 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( @@ -2401,6 +2387,12 @@ def _gvr_decode_exact_check(logits_row, pre_idx_row, index_topk, tag): assert int((indices < 0).sum()) == 0, ( f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1" ) + # Distinctness: a duplicate+omission pair on a tie plateau would leave + # the sorted value multiset below unchanged. + n_unique = int(torch.unique(indices[0]).numel()) + assert n_unique == index_topk, ( + f"{tag}: only {n_unique} of {index_topk} output indices are distinct" + ) flat = logits[0].float() got = flat[indices[0].long()].sort().values ref = flat.topk(index_topk).values.sort().values @@ -2436,19 +2428,22 @@ def test_indexer_topk_decode_gvr_hostile_hint(index_topk, num_tokens, dtype, hin @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. - """ +@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): + """More ties at the K-th value than the candidate buffer can hold: no + threshold lands in [K, kC], so the repair must emit the strictly-greater + set plus arbitrary ties. bf16/fp16 cover the reduced-precision driver's + separate direct-emit block.""" 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() + # Plateau (1.0) and floor (-1.0) are exact in every dtype; casting can + # only merge strictly-greater values with each other, which is tolerated. + logits = logits[torch.randperm(num_tokens, device="cuda")].contiguous().to(dtype) pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie}")