[Bugfix] V1: fix allowed_token_ids_mask aliasing in InputBatch.swap_states (#48419)
Signed-off-by: Rui Zhu <rui.zhu.rz399@yale.edu> Co-authored-by: Claude <noreply@anthropic.com>
Rui Zhu committed
Sep 25, 2026 at 19:23 UTC
974cb65155303139cec4f5d8d8b68bbd8ef49a2a
2 files changed
+56
-6
tests/v1/worker/test_gpu_input_batch.py
+54
@@ -372,6 +372,60 @@ 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_swap_states_preserves_allowed_token_ids_mask(device: str):
377
+ """swap_states() must exchange the two allowed_token_ids_mask rows.
378
+
379
+ Swapping row views via a tuple assignment
380
+ (t[i1], t[i2] = t[i2], t[i1]) aliases them, so both rows collapse to the
381
+ second row's contents and one request silently loses its token
382
+ restriction. This is the sibling of the condense() leak fixed for
383
+ https://github.com/vllm-project/vllm/issues/43894.
384
+ """
385
+ allowed_token_id = 13
386
+ input_batch = InputBatch(
387
+ max_num_reqs=2,
388
+ max_model_len=1024,
389
+ max_num_batched_tokens=1024,
390
+ device=torch.device(device),
391
+ vocab_size=VOCAB_SIZE,
392
+ block_sizes=[1],
393
+ kernel_block_sizes=[1],
394
+ max_num_blocks_per_req=[1024],
395
+ )
396
+
397
+ def _make_req(suffix: int, allowed_token_ids=None) -> CachedRequestState:
398
+ return CachedRequestState(
399
+ req_id=f"req_id_{suffix}",
400
+ prompt_token_ids=[1, 2, 3],
401
+ sampling_params=SamplingParams(allowed_token_ids=allowed_token_ids),
402
+ pooling_params=None,
403
+ mm_features=[],
404
+ block_ids=([],),
405
+ generator=None,
406
+ num_computed_tokens=0,
407
+ output_token_ids=[],
408
+ )
409
+
410
+ # Row 0 is constrained to a single token; row 1 is unconstrained.
411
+ assert input_batch.add_request(_make_req(0, [allowed_token_id])) == 0
412
+ assert input_batch.add_request(_make_req(1)) == 1
413
+
414
+ mask = input_batch.allowed_token_ids_mask_cpu_tensor
415
+ assert mask is not None
416
+ assert int(mask[0].sum().item()) == VOCAB_SIZE - 1
417
+ assert not mask[0][allowed_token_id].item()
418
+ assert int(mask[1].sum().item()) == 0
419
+
420
+ input_batch.swap_states(0, 1)
421
+
422
+ # The constrained request now sits at row 1 and must keep its mask; the
423
+ # unconstrained request now sits at row 0 and must stay open.
424
+ assert int(mask[1].sum().item()) == VOCAB_SIZE - 1
425
+ assert not mask[1][allowed_token_id].item()
426
+ assert int(mask[0].sum().item()) == 0
427
+
428
+
429
def _construct_pooling_request(req_id_suffix: int, pooling_params=None):
430
from vllm.pooling_params import PoolingParams
431
vllm/v1/worker/gpu_input_batch.py
+2
-6
@@ -688,12 +688,8 @@ class InputBatch:
688
swap_dict_values(self.bad_words_token_ids, i1, i2)
689
690
if self.allowed_token_ids_mask_cpu_tensor is not None:
691
- (
692
- self.allowed_token_ids_mask_cpu_tensor[i1],
693
- self.allowed_token_ids_mask_cpu_tensor[i2],
694
- ) = (
695
- self.allowed_token_ids_mask_cpu_tensor[i2],
696
- self.allowed_token_ids_mask_cpu_tensor[i1],
691
+ self.allowed_token_ids_mask_cpu_tensor[[i1, i2]] = (
692
+ self.allowed_token_ids_mask_cpu_tensor[[i2, i1]]
693
)
694
695
def _get_active_token_count(self, req_index: int) -> int: