Fix top-k ternary search timeout fallback producing inconsistent pivot/stats in Triton kernel

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

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