From f5e441de10bd55149c4211473dc9fcffc84ad2ce Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Mon, 31 Aug 2026 11:49:43 -0700 Subject: [PATCH] [Bugfix][Test] Fix off-by-one error in sampled token rank causing flaky logprobs test (#53976) Signed-off-by: mayuyuace --- tests/v1/engine/utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/v1/engine/utils.py b/tests/v1/engine/utils.py index 324c9c9ad5..b1e8c612ae 100644 --- a/tests/v1/engine/utils.py +++ b/tests/v1/engine/utils.py @@ -130,9 +130,11 @@ def _create_random_top_token_test_vector( # Check if the sampled_token_id occurs in choice_tensor[1:] if sampled_token_id in choice_tensor[1:]: - sampled_token_rank = ( - (choice_tensor[1:] == sampled_token_id).nonzero(as_tuple=True)[0].item() - ) + # Use 1-based rank to match the top-k rank convention used by the + # consumer (logprobs.py assigns ranks starting from 1). + sampled_token_rank = (choice_tensor[1:] == sampled_token_id).nonzero( + as_tuple=True + )[0].item() + 1 else: # If not found, assign a random int between num_logprobs and 50700 sampled_token_rank = random.randint(num_logprobs, 50700) @@ -184,7 +186,7 @@ def _create_random_top_token_test_matrix( row = matrix[rdx, 1:] # Skip the first column as it contains the token list token_index = (row == tokens_list[rdx]).nonzero(as_tuple=True)[0] if token_index.numel() > 0: - prompt_token_ranks[rdx] = token_index.item() + prompt_token_ranks[rdx] = token_index.item() + 1 else: prompt_token_ranks[rdx] = random.randint(shape[1], 50700)