[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 <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Yiliu Dong committed
Jul 2, 2026 at 00:26 UTC
00eb7cefa31e32585fb419db5bb945f6a42480fe
3 files changed
+60
tests/v1/worker/test_gpu_input_batch.py
+34
@@ -435,6 +435,40 @@ def test_pooling_prompt_lens_not_aliased(device: str):
435
)
436
437
438
+def test_placeholder_spec_token_ids_written_verbatim():
439
+ input_batch = InputBatch(
440
+ max_num_reqs=1,
441
+ max_model_len=8,
442
+ max_num_batched_tokens=8,
443
+ device=torch.device("cpu"),
444
+ vocab_size=VOCAB_SIZE,
445
+ block_sizes=[16],
446
+ kernel_block_sizes=[16],
447
+ )
448
+ req = CachedRequestState(
449
+ req_id="req",
450
+ prompt_token_ids=[10, 11],
451
+ mm_features=[],
452
+ sampling_params=SamplingParams(),
453
+ block_ids=([],),
454
+ generator=None,
455
+ num_computed_tokens=3,
456
+ output_token_ids=[12],
457
+ )
458
+ input_batch.add_request(req)
459
+
460
+ input_batch.update_req_spec_token_ids(
461
+ req,
462
+ {"req": [13, -1, -1]},
463
+ )
464
+
465
+ # Placeholders (-1) are kept verbatim in both the spec_token_ids list and
466
+ # the token buffer; they are clamped to 0 only at the embedding boundary
467
+ # (GPUModelRunner._preprocess).
468
+ assert input_batch.spec_token_ids[0] == [13, -1, -1]
469
+ assert input_batch.token_ids_cpu[0, 3:6].tolist() == [13, -1, -1]
470
+
471
+
472
@pytest.mark.parametrize(
473
("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"),
474
[
tests/v1/worker/test_gpu_model_runner.py
+22
@@ -862,6 +862,28 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler():
862
assert torch.equal(passed_draft_probs, expected_draft_probs)
863
864
865
+def test_invalid_draft_suffixes_remain_rejected_in_metadata():
866
+ runner = object.__new__(GPUModelRunner)
867
+ runner.device = torch.device("cpu")
868
+ runner.arange_np = np.arange(64, dtype=np.int64)
869
+ runner._arange_scratch = np.empty(64, dtype=np.int64)
870
+ # Placeholder (-1) drafts are kept in input_ids (clamped to 0 only at the
871
+ # embedding boundary). For num_draft_tokens=[2, 1, 2] the draft positions
872
+ # are [1, 2, 4, 6, 7], so the gather carries the -1s straight into the
873
+ # rejection-sampling metadata.
874
+ runner.input_ids = SimpleNamespace(
875
+ gpu=torch.tensor([99, 10, -1, 99, 12, 99, 13, -1], dtype=torch.int32),
876
+ )
877
+
878
+ metadata = GPUModelRunner._calc_spec_decode_metadata(
879
+ runner,
880
+ np.array([2, 1, 2], dtype=np.int32),
881
+ np.array([3, 5, 8], dtype=np.int32),
882
+ )
883
+
884
+ assert metadata.draft_token_ids.tolist() == [10, -1, 12, 13, -1]
885
+
886
+
887
def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):
888
torch.set_default_dtype(torch.float16)
889
layer_0 = "model.layers.0.self_attn.attn"
vllm/v1/worker/gpu_model_runner.py
+4
@@ -3440,6 +3440,10 @@ class GPUModelRunner(
3440
is_first_rank = get_pp_group().is_first_rank
3441
is_encoder_decoder = self.model_config.is_encoder_decoder
3442
3443
+ # Clamp speculative scheduler placeholders (-1) before embedding lookup.
3444
+ if self.speculative_config is not None:
3445
+ self.input_ids.gpu[:num_input_tokens].clamp_(min=0)
3446
+
3447
# _prepare_inputs may reorder the batch, so we must gather multi
3448
# modal outputs after that to ensure the correct order
3449
ec_connector_output = None