[Bugfix][Test] Fix off-by-one error in sampled token rank causing flaky logprobs test (#53976)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Qiming Zhang committed
Aug 31, 2026 at 11:49 UTC
f5e441de10bd55149c4211473dc9fcffc84ad2ce
1 file changed
+6
-4
tests/v1/engine/utils.py
+6
-4
@@ -130,9 +130,11 @@ def _create_random_top_token_test_vector(
130
131
# Check if the sampled_token_id occurs in choice_tensor[1:]
132
if sampled_token_id in choice_tensor[1:]:
133
- sampled_token_rank = (
134
- (choice_tensor[1:] == sampled_token_id).nonzero(as_tuple=True)[0].item()
135
- )
133
+ # Use 1-based rank to match the top-k rank convention used by the
134
+ # consumer (logprobs.py assigns ranks starting from 1).
135
+ sampled_token_rank = (choice_tensor[1:] == sampled_token_id).nonzero(
136
+ as_tuple=True
137
+ )[0].item() + 1
138
else:
139
# If not found, assign a random int between num_logprobs and 50700
140
sampled_token_rank = random.randint(num_logprobs, 50700)
@@ -184,7 +186,7 @@ def _create_random_top_token_test_matrix(
186
row = matrix[rdx, 1:] # Skip the first column as it contains the token list
187
token_index = (row == tokens_list[rdx]).nonzero(as_tuple=True)[0]
188
if token_index.numel() > 0:
187
- prompt_token_ranks[rdx] = token_index.item()
189
+ prompt_token_ranks[rdx] = token_index.item() + 1
190
else:
191
prompt_token_ranks[rdx] = random.randint(shape[1], 50700)
192