From 00eb7cefa31e32585fb419db5bb945f6a42480fe Mon Sep 17 00:00:00 2001 From: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Date: Thu, 2 Jul 2026 00:26:03 +0800 Subject: [PATCH] [Bugfix] Prevent padding placeholders from reaching embeddings (#47029) Signed-off-by: qianlihuang <91178480+qianlihuang@users.noreply.github.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- tests/v1/worker/test_gpu_input_batch.py | 34 ++++++++++++++++++++++++ tests/v1/worker/test_gpu_model_runner.py | 22 +++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 4 +++ 3 files changed, 60 insertions(+) diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index bfd4016c9fe..4d0a1698a70 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -435,6 +435,40 @@ def test_pooling_prompt_lens_not_aliased(device: str): ) +def test_placeholder_spec_token_ids_written_verbatim(): + input_batch = InputBatch( + max_num_reqs=1, + max_model_len=8, + max_num_batched_tokens=8, + device=torch.device("cpu"), + vocab_size=VOCAB_SIZE, + block_sizes=[16], + kernel_block_sizes=[16], + ) + req = CachedRequestState( + req_id="req", + prompt_token_ids=[10, 11], + mm_features=[], + sampling_params=SamplingParams(), + block_ids=([],), + generator=None, + num_computed_tokens=3, + output_token_ids=[12], + ) + input_batch.add_request(req) + + input_batch.update_req_spec_token_ids( + req, + {"req": [13, -1, -1]}, + ) + + # Placeholders (-1) are kept verbatim in both the spec_token_ids list and + # the token buffer; they are clamped to 0 only at the embedding boundary + # (GPUModelRunner._preprocess). + assert input_batch.spec_token_ids[0] == [13, -1, -1] + assert input_batch.token_ids_cpu[0, 3:6].tolist() == [13, -1, -1] + + @pytest.mark.parametrize( ("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"), [ diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 6d538bc69d4..00722ccc244 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -862,6 +862,28 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler(): assert torch.equal(passed_draft_probs, expected_draft_probs) +def test_invalid_draft_suffixes_remain_rejected_in_metadata(): + runner = object.__new__(GPUModelRunner) + runner.device = torch.device("cpu") + runner.arange_np = np.arange(64, dtype=np.int64) + runner._arange_scratch = np.empty(64, dtype=np.int64) + # Placeholder (-1) drafts are kept in input_ids (clamped to 0 only at the + # embedding boundary). For num_draft_tokens=[2, 1, 2] the draft positions + # are [1, 2, 4, 6, 7], so the gather carries the -1s straight into the + # rejection-sampling metadata. + runner.input_ids = SimpleNamespace( + gpu=torch.tensor([99, 10, -1, 99, 12, 99, 13, -1], dtype=torch.int32), + ) + + metadata = GPUModelRunner._calc_spec_decode_metadata( + runner, + np.array([2, 1, 2], dtype=np.int32), + np.array([3, 5, 8], dtype=np.int32), + ) + + assert metadata.draft_token_ids.tolist() == [10, -1, 12, 13, -1] + + def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config): torch.set_default_dtype(torch.float16) layer_0 = "model.layers.0.self_attn.attn" diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 37e0c9f80a1..e8d472e9eb0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3440,6 +3440,10 @@ class GPUModelRunner( is_first_rank = get_pp_group().is_first_rank is_encoder_decoder = self.model_config.is_encoder_decoder + # Clamp speculative scheduler placeholders (-1) before embedding lookup. + if self.speculative_config is not None: + self.input_ids.gpu[:num_input_tokens].clamp_(min=0) + # _prepare_inputs may reorder the batch, so we must gather multi # modal outputs after that to ensure the correct order ec_connector_output = None