diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh index f21a4d7cf2c6..d75933bd5d38 100644 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh @@ -164,6 +164,9 @@ constexpr int SAFETY_MARGIN = 2048; constexpr int MAX_CANDIDATES = TOP_K + SAFETY_MARGIN * 2; // 6144 constexpr int MAX_REFINE_ITERS = 15; +// 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; static_assert(TOP_K % BLOCK_SIZE == 0); @@ -419,6 +422,23 @@ __device__ __forceinline__ float warpReduceMax(float val) #endif +// ============================================================================ +// Order-preserving float <-> uint32 map (arch-independent) +// ============================================================================ +// 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); + 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 +681,23 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con } __syncthreads(); + // 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) - 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 = 2; + } + __syncthreads(); } // ================================================================ @@ -773,25 +801,42 @@ __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. + // 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); - if (tid == 0 && smem->cand_count > kCC) - smem->val_lo = smem->threshold; + // 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; + 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 +849,70 @@ __device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int con } __syncthreads(); } + + // 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) + 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 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; + 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 +1336,23 @@ __device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, i } __syncthreads(); + // 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) - 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 = 2; + } + __syncthreads(); } // ================================================================ @@ -1339,23 +1456,37 @@ __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 comments there). 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 +1499,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..33ebdb2c82ca 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -2350,3 +2350,100 @@ def test_prefill_overflow_policy_overflow( dtype, row_start_offset=row_start_offset, ) + + +# ============================================================================ +# 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]). +# ============================================================================ + + +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") + # -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( + 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" + ) + # 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 + 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]) +@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 + # 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}")