From ce49c47cbb241b59a32ef110079a5261dbd08e49 Mon Sep 17 00:00:00 2001 From: Akshat Anand Date: Sat, 17 Oct 2026 07:56:57 +0530 Subject: [PATCH] [Bugfix][Sampler] Carry the boundary tie count across blocks in the split top-p mask kernel (#62046) Signed-off-by: cipheraxat Co-authored-by: cipheraxat Co-authored-by: Bowen Wang --- tests/v1/sample/test_topk_topp_sampler.py | 30 +++++++++++++++++++++++ vllm/v1/sample/ops/topk_topp_triton.py | 1 + 2 files changed, 31 insertions(+) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 190a4e3fca..f78f48b3e8 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/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 diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 4aa837e9d7..d5849a8d67 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/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)