From c18db7faaf75c607b4a2d4be441cd74444c7059b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 23 Jun 2026 06:42:44 +0000 Subject: [PATCH 1/2] Initial plan From eba3e1e73fea0acbd4d54c86cc8f45fc2645e91f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 23 Jun 2026 06:45:17 +0000 Subject: [PATCH 2/2] Fix top-k ternary search timeout fallback in Triton kernel --- tests/v1/sample/test_topk_topp_sampler.py | 46 +++++++++++++++++++++++ vllm/v1/sample/ops/topk_topp_triton.py | 39 +++++++++++++++---- 2 files changed, 78 insertions(+), 7 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 047e2b754ef0..07e769a19e3b 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -727,6 +727,52 @@ def test_mixed_neginf_and_normal_rows(self): if finite_in > 0: assert kept > 0, f"Row {i}: no tokens kept" + @pytest.mark.parametrize("k_val", [5, 10, 20]) + def test_topk_tied_logits_exact_count(self, k_val: int): + """Top-k with many tied logits forces ternary-search timeout fallback. + + When many tokens share the same logit value at the top-k boundary, + the ternary search interval collapses quickly and the fallback branch + is triggered. This is a regression test for the bug where k_pivot, + k_pivots_num, min_larger, and num_min_larger were inconsistent in + the fallback branch, causing the surviving-token count to differ + from k. + + The test asserts that exactly k tokens survive for each row where + there are at least k finite tokens, matching the PyTorch reference. + """ + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton + + batch_size, vocab_size = 16, 32000 + # Construct logits so that the top-(k+50) tokens all share the same + # value, forcing the ternary search to collapse and use the fallback. + logits = torch.full( + (batch_size, vocab_size), float("-inf"), dtype=torch.float32 + ) + tied_count = k_val + 50 # more tied tokens than k + for i in range(batch_size): + indices = torch.randperm(vocab_size, generator=self.generator)[ + :tied_count + ] + logits[i, indices] = 1.0 # all equal — forces range collapse + + k = torch.full((batch_size,), k_val, dtype=torch.int32) + + # PyTorch reference + logits_pytorch = logits.clone().float() + result_pytorch = apply_top_k_top_p_pytorch(logits_pytorch, k, p=None) + pytorch_kept = (result_pytorch > float("-inf")).sum(dim=-1) + + # Triton kernel + result_triton = apply_top_k_top_p_triton(logits.clone(), k, p=None) + triton_kept = (result_triton > float("-inf")).sum(dim=-1) + + assert torch.equal(pytorch_kept, triton_kept), ( + f"Top-k tied-logit count mismatch (k={k_val}): " + f"PyTorch kept {pytorch_kept.tolist()}, " + f"Triton kept {triton_kept.tolist()}" + ) + # ============================================================================= # FlashInfer top-k/top-p robustness tests diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index d20cac37fcd7..f026c00c6abd 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -278,9 +278,21 @@ def _topk_topp_kernel( num_iters += 1 if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9: - k_pivot = (max_range + min_range) / 2.0 - min_larger = min_larger_0 - num_min_larger = num_min_larger_0 + # Snap to an already-evaluated pivot so that + # k_pivot, k_pivots_num, min_larger and + # num_min_larger all describe the SAME threshold. + # Prefer the larger pivot that still keeps >= k + # tokens; otherwise fall back to the smaller one. + if k_pivots_num_1 >= k: + k_pivot = k_pivot_1 + k_pivots_num = k_pivots_num_1 + min_larger = min_larger_1 + num_min_larger = num_min_larger_1 + else: + k_pivot = k_pivot_0 + k_pivots_num = k_pivots_num_0 + min_larger = min_larger_0 + num_min_larger = num_min_larger_0 found_pivot = 1 else: # If top-k outlier gathering failed, search whole logit space @@ -360,14 +372,27 @@ def _topk_topp_kernel( num_iters += 1 if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9: - k_pivot = (max_range + min_range) / 2.0 - min_larger = min_larger_0 - num_min_larger = num_min_larger_0 + # Snap to an already-evaluated pivot so that + # k_pivot, k_pivots_num, min_larger and + # num_min_larger all describe the SAME threshold. + # Prefer the larger pivot that still keeps >= k + # tokens; otherwise fall back to the smaller one. + if k_pivots_num_1 >= k: + k_pivot = k_pivot_1 + k_pivots_num = k_pivots_num_1 + min_larger = min_larger_1 + num_min_larger = num_min_larger_1 + else: + k_pivot = k_pivot_0 + k_pivots_num = k_pivots_num_0 + min_larger = min_larger_0 + num_min_larger = num_min_larger_0 found_pivot = 1 duplicate_logit = min_larger num_duplicate_logit = num_min_larger - num_keep = num_duplicate_logit - (k_pivots_num - k) + excess = tl.where(k_pivots_num > k, k_pivots_num - k, 0) + num_keep = num_duplicate_logit - excess num_kept = tl.zeros((), dtype=tl.uint32) # Top-k only path. If there are fewer finite values