From 7900d2eabc4f9a7facb7a196a04c2c32803a79c8 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Mon, 12 Oct 2026 08:18:44 -0500 Subject: [PATCH] [Test] Accept tied tokens at the k-th value in the Triton top-k test (#60726) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/v1/sample/test_topk_topp_sampler.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index c0b9792456..190a4e3fca 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -251,8 +251,17 @@ class TestTritonTopkTopp: triton_kept = (result_triton != float("-inf")).sum(dim=-1) if p is None: - # Top-k only: expect exact match - assert torch.equal(pytorch_kept, triton_kept), ( + # Top-k only: Triton keeps exactly k (index-ordered tie-break at + # the k-th value) while the sort-based reference keeps all ties, + # so PyTorch may only keep extra tokens tied with the k-th value. + pytorch_mask = result_pytorch != float("-inf") + triton_mask = result_triton != float("-inf") + kth = torch.where(triton_mask, logits.float(), float("inf")) + kth = kth.amin(dim=-1, keepdim=True).expand_as(logits) + extra = pytorch_mask & ~triton_mask + assert not (triton_mask & ~pytorch_mask).any() and torch.equal( + logits.float()[extra], kth[extra] + ), ( f"Top-k mask mismatch: PyTorch kept {pytorch_kept.tolist()}, " f"Triton kept {triton_kept.tolist()}" )