[Bugfix] V1: clear stale allowed_token_ids mask in InputBatch.condense (#43931)
Signed-off-by: Varshith <kvarshithgowda@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Varshith committed
Sep 26, 2026 at 22:44 UTC
fba47397f304993741d2b76d73fb42e5be3e8eac
2 files changed
+57
tests/v1/worker/test_gpu_input_batch.py
+56
@@ -372,6 +372,62 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
372
_compare_objs(input_batch, ref_input_batch)
373
374
375
+@pytest.mark.parametrize("device", DEVICES)
376
+def test_condense_clears_stale_allowed_token_ids_mask(device: str):
377
+ """condense() must clear the mask row a constrained request is moved out
378
+ of. Otherwise a later request that reuses that row without
379
+ allowed_token_ids inherits the stale whitelist.
380
+ See https://github.com/vllm-project/vllm/issues/43894.
381
+ """
382
+ batch_size = 4
383
+ allowed_token_id = 13
384
+ input_batch = InputBatch(
385
+ max_num_reqs=batch_size,
386
+ max_model_len=1024,
387
+ max_num_batched_tokens=1024,
388
+ device=torch.device(device),
389
+ vocab_size=VOCAB_SIZE,
390
+ block_sizes=[1],
391
+ kernel_block_sizes=[1],
392
+ max_num_blocks_per_req=[1024],
393
+ )
394
+
395
+ def _make_req(suffix: int, allowed_token_ids=None) -> CachedRequestState:
396
+ return CachedRequestState(
397
+ req_id=f"req_id_{suffix}",
398
+ prompt_token_ids=[1, 2, 3],
399
+ sampling_params=SamplingParams(allowed_token_ids=allowed_token_ids),
400
+ pooling_params=None,
401
+ mm_features=[],
402
+ block_ids=([],),
403
+ generator=None,
404
+ num_computed_tokens=0,
405
+ output_token_ids=[],
406
+ )
407
+
408
+ # Only req_id_2 is constrained, and it lands on the highest row.
409
+ assert input_batch.add_request(_make_req(0)) == 0
410
+ assert input_batch.add_request(_make_req(1)) == 1
411
+ assert input_batch.add_request(_make_req(2, [allowed_token_id])) == 2
412
+
413
+ mask = input_batch.allowed_token_ids_mask_cpu_tensor
414
+ assert mask is not None
415
+ # The constrained row masks every token except the single allowed id.
416
+ assert not mask[2][allowed_token_id].item()
417
+ assert int(mask[2].sum().item()) == VOCAB_SIZE - 1
418
+
419
+ # Free row 0, then condense: req_id_2 slides from row 2 down into row 0.
420
+ input_batch.remove_request("req_id_0")
421
+ input_batch.condense()
422
+ assert input_batch.req_id_to_index["req_id_2"] == 0
423
+ assert int(mask[2].sum().item()) == 0
424
+
425
+ # A new unrestricted request reuses row 2 and must stay unconstrained.
426
+ assert input_batch.add_request(_make_req(3)) == 2
427
+ assert "req_id_3" not in input_batch.has_allowed_token_ids
428
+ assert int(mask[2].sum().item()) == 0
429
+
430
+
431
@pytest.mark.parametrize("device", DEVICES)
432
def test_swap_states_preserves_allowed_token_ids_mask(device: str):
433
"""swap_states() must exchange the two allowed_token_ids_mask rows.
vllm/v1/worker/gpu_input_batch.py
+1
@@ -817,6 +817,7 @@ class InputBatch:
817
self.allowed_token_ids_mask_cpu_tensor[empty_index] = (
818
self.allowed_token_ids_mask_cpu_tensor[last_req_index]
819
)
820
+ self.allowed_token_ids_mask_cpu_tensor[last_req_index].fill_(False)
821
822
bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None)
823
if bad_words_token_ids is not None: