diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 2f13694bccad..e7813d4a01e3 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1341,6 +1341,62 @@ def test_spec_decode_logprobs( assert ref_logprob.decoded_token == spec_logprob.decoded_token +def test_prompt_logprobs_stop_before_speculative_padding(monkeypatch): + """A one-token prompt has no scored tokens, even in a padded decode batch.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + # Step the engine explicitly so the warm request cannot finish before admission. + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + with VllmRunner( + "Qwen/Qwen3-0.6B", + revision="c1899de289a04d12100db370d81485cdf75e47ca", + dtype="bfloat16", + max_model_len=128, + max_num_seqs=2, + max_num_batched_tokens=64, + enable_chunked_prefill=True, + enable_prefix_caching=False, + async_scheduling=False, + enforce_eager=True, + gpu_memory_utilization=0.25, + speculative_config={ + "method": "ngram", + "prompt_lookup_min": 3, + "prompt_lookup_max": 5, + "num_speculative_tokens": 4, + }, + ) as vllm_model: + engine = vllm_model.llm.llm_engine + engine.add_request( + "warm", + TokensPrompt(prompt_token_ids=[100, 101, 102, 103] * 2), + SamplingParams(temperature=0, max_tokens=32, ignore_eos=True), + ) + try: + (warm,) = engine.step() + assert warm.request_id == "warm" + assert not warm.finished + assert len(warm.outputs[0].token_ids) == 1 + + engine.add_request( + "one-token", + TokensPrompt(prompt_token_ids=[100]), + SamplingParams(temperature=0, max_tokens=1, prompt_logprobs=0), + ) + # With an existing decode, admission pads this prompt to 1 + 4 rows. + outputs = {output.request_id: output for output in engine.step()} + assert set(outputs) == {"warm", "one-token"} + result = outputs["one-token"] + assert result.finished + assert len(result.outputs[0].token_ids) == 1 + assert result.prompt_token_ids == [100] + assert result.prompt_logprobs is not None + assert len(result.prompt_logprobs) == 1 + assert result.prompt_logprobs == [None] + finally: + engine.abort_request(["warm", "one-token"]) + + def test_prompt_logprobs_with_chunking_and_preemption(): """Test that prompt logprobs are correctly returned when using both chunked prefill and preemption. diff --git a/tests/v1/worker/test_gpu_prompt_logprobs.py b/tests/v1/worker/test_gpu_prompt_logprobs.py index 549a93502ff3..6e1de61100dd 100644 --- a/tests/v1/worker/test_gpu_prompt_logprobs.py +++ b/tests/v1/worker/test_gpu_prompt_logprobs.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace +from dataclasses import replace import numpy as np import pytest @@ -9,24 +9,24 @@ from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams from vllm.triton_utils import HAS_TRITON +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker NUM_DRAFTS = 4 def _batch(computed, prompt_len, query_len, device): - query_start_loc = np.array([0, query_len], dtype=np.int32) - return SimpleNamespace( - num_reqs=1, - num_tokens=query_len, + buffers = InputBuffers(1, query_len, device) + batch = InputBatch.make_dummy(1, query_len, buffers, is_padding=False) + computed_tokens = np.array([computed], dtype=np.int32) + return replace( + batch, req_ids=["0"], - idx_mapping_np=np.arange(1), - idx_mapping=torch.zeros(1, dtype=torch.int32, device=device), - num_computed_prefill_tokens_np=np.array([computed]), - prefill_len_np=np.array([prompt_len]), - num_scheduled_tokens=np.array([query_len]), - query_start_loc_np=query_start_loc, - query_start_loc=torch.from_numpy(query_start_loc).to(device), + num_computed_tokens_np=computed_tokens, + num_computed_prefill_tokens_np=computed_tokens, + prefill_len_np=np.array([prompt_len], dtype=np.int32), + is_prefilling_np=np.array([computed < prompt_len]), + has_prefill=computed < prompt_len, ) @@ -64,3 +64,51 @@ def test_scores_stop_at_the_end_of_the_prompt( ) num_rows = result["0"].logprobs.shape[0] if result else 0 assert num_rows == max(prompt_len - max(computed, prompt_logprob_start) - 1, 0) + + @pytest.mark.parametrize("first_chunk_len", [3, 7]) + def test_chunked_scores_preserved_at_prompt_end(self, first_chunk_len): + """Keep earlier scores when the final chunk adds some or no prompt rows.""" + device = torch.device(current_platform.device_type) + prompt_len = 8 + vocab_size = 32 + all_token_ids = torch.arange( + vocab_size, dtype=torch.int32, device=device + ).unsqueeze(0) + hidden_states = torch.arange( + prompt_len + NUM_DRAFTS, dtype=torch.float32, device=device + ).unsqueeze(1) + vocabulary = torch.arange(vocab_size, device=device).unsqueeze(0) + + def logits_fn(hidden): + # Distinct scores per position expose lost, repeated or reordered rows. + return -(vocabulary - hidden[:, :1]).square() + + def compute(worker, computed, query_len): + return worker.compute_prompt_logprobs( + logits_fn, + hidden_states[computed : computed + query_len], + _batch(computed, prompt_len, query_len, device), + all_token_ids, + torch.tensor([computed], dtype=torch.int32, device=device), + np.array([prompt_len]), + ) + + reference_worker = PromptLogprobsWorker(1, device) + reference_worker.add_request("0", 0, SamplingParams(prompt_logprobs=0)) + # A whole prompt without placeholders is correct even before this fix. + expected = compute(reference_worker, 0, prompt_len)["0"] + assert expected.logprobs.shape == (prompt_len - 1, 1) + + worker = PromptLogprobsWorker(1, device) + worker.add_request("0", 0, SamplingParams(prompt_logprobs=0)) + assert compute(worker, 0, first_chunk_len) == {} + result = compute( + worker, first_chunk_len, prompt_len - first_chunk_len + NUM_DRAFTS + ) + + assert set(result) == {"0"} + actual = result["0"] + for field in ("logprob_token_ids", "logprobs", "selected_token_ranks"): + torch.testing.assert_close( + getattr(actual, field), getattr(expected, field), rtol=0, atol=0 + )