[Test] Cover prompt-logprob boundaries through worker and real model

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.

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
+            )