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}")