[Bugfix][Sampler] Carry the boundary tie count across blocks in the split top-p mask kernel (#62046)

Signed-off-by: cipheraxat <cipheraxat@users.noreply.github.com>
Co-authored-by: cipheraxat <cipheraxat@users.noreply.github.com>
Co-authored-by: Bowen Wang <bowen@inferact.ai>

tests/v1/sample/test_topk_topp_sampler.py

@@ -821,6 +821,36 @@ class TestTritonTopkTopp:
         )
         assert len(kept_idx) == pytorch_kept
 
+    def test_topp_ties_across_blocks(self):
+        """The split top-p tie budget must carry across a slice's blocks.
+
+        Regression: the split top-p mask kernel restarted its count of
+        boundary-value copies at every 8192-entry block, so each block of a
+        slice could keep up to `numkeep` copies of its own.
+        """
+        from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
+
+        batch_size, vocab_size, num_ties = 64, 151936, 4096
+        # 4096 tied tokens spread over many blocks, between 64 tokens above
+        # and a low tail; the p=0.75 boundary lands mid-copy in the tie group.
+        tie_idx = torch.arange(num_ties) * 37
+        logits = torch.full((batch_size, vocab_size), -8.0)
+        logits[:, tie_idx] = 0.0
+        logits[:, -64:] = 4.0
+        p = torch.full((batch_size,), 0.75, dtype=torch.float32)
+
+        kept = apply_top_k_top_p_triton(logits.clone(), None, p) != float("-inf")
+        expected = apply_top_k_top_p_pytorch(logits.clone(), None, p) != float("-inf")
+        num_kept, num_expected = kept.sum(dim=-1), expected.sum(dim=-1)
+        assert torch.equal(num_kept, num_expected), (
+            f"kept {num_kept.unique().tolist()}, "
+            f"expected {num_expected.unique().tolist()}"
+        )
+        # Kept copies of the boundary value are the first ones in index order.
+        kept_ties = kept[:, tie_idx]
+        num_kept_ties = kept_ties.sum(dim=-1, keepdim=True)
+        assert torch.equal(kept_ties, torch.arange(num_ties) < num_kept_ties)
+
 
 # =============================================================================
 # FlashInfer top-k/top-p robustness tests

vllm/v1/sample/ops/topk_topp_triton.py

@@ -1513,6 +1513,7 @@ def _topp_sb_mask_kernel(
             cum = tl.cumsum(dmask.to(tl.int32)).to(tl.float32) + num_preceding
             keep_dup = dmask & (cum <= numkeep)
             keep = keep & (~dmask | keep_dup)
+            num_preceding += tl.sum(dmask.to(tl.float32))
         tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n)