[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>

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.

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: