diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index 7c6c9ba417..f37e824e6c 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -372,6 +372,62 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis _compare_objs(input_batch, ref_input_batch) +@pytest.mark.parametrize("device", DEVICES) +def test_condense_clears_stale_allowed_token_ids_mask(device: str): + """condense() must clear the mask row a constrained request is moved out + of. Otherwise a later request that reuses that row without + allowed_token_ids inherits the stale whitelist. + See https://github.com/vllm-project/vllm/issues/43894. + """ + batch_size = 4 + allowed_token_id = 13 + input_batch = InputBatch( + max_num_reqs=batch_size, + max_model_len=1024, + max_num_batched_tokens=1024, + device=torch.device(device), + vocab_size=VOCAB_SIZE, + block_sizes=[1], + kernel_block_sizes=[1], + max_num_blocks_per_req=[1024], + ) + + def _make_req(suffix: int, allowed_token_ids=None) -> CachedRequestState: + return CachedRequestState( + req_id=f"req_id_{suffix}", + prompt_token_ids=[1, 2, 3], + sampling_params=SamplingParams(allowed_token_ids=allowed_token_ids), + pooling_params=None, + mm_features=[], + block_ids=([],), + generator=None, + num_computed_tokens=0, + output_token_ids=[], + ) + + # Only req_id_2 is constrained, and it lands on the highest row. + assert input_batch.add_request(_make_req(0)) == 0 + assert input_batch.add_request(_make_req(1)) == 1 + assert input_batch.add_request(_make_req(2, [allowed_token_id])) == 2 + + mask = input_batch.allowed_token_ids_mask_cpu_tensor + assert mask is not None + # The constrained row masks every token except the single allowed id. + assert not mask[2][allowed_token_id].item() + assert int(mask[2].sum().item()) == VOCAB_SIZE - 1 + + # Free row 0, then condense: req_id_2 slides from row 2 down into row 0. + input_batch.remove_request("req_id_0") + input_batch.condense() + assert input_batch.req_id_to_index["req_id_2"] == 0 + assert int(mask[2].sum().item()) == 0 + + # A new unrestricted request reuses row 2 and must stay unconstrained. + assert input_batch.add_request(_make_req(3)) == 2 + assert "req_id_3" not in input_batch.has_allowed_token_ids + assert int(mask[2].sum().item()) == 0 + + @pytest.mark.parametrize("device", DEVICES) def test_swap_states_preserves_allowed_token_ids_mask(device: str): """swap_states() must exchange the two allowed_token_ids_mask rows. diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 2a1a3534c9..4707e4f2ca 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -817,6 +817,7 @@ class InputBatch: self.allowed_token_ids_mask_cpu_tensor[empty_index] = ( self.allowed_token_ids_mask_cpu_tensor[last_req_index] ) + self.allowed_token_ids_mask_cpu_tensor[last_req_index].fill_(False) bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None) if bad_words_token_ids is not None: