[Security] Defer diffusion prompt scores until output

tests/v1/core/test_scheduler.py

@@ -66,6 +66,7 @@
     DraftTokenIds,
     ECConnectorOutput,
     KVConnectorOutput,
+    LogprobsTensors,
     ModelRunnerOutput,
     SamplingMaskLists,
     make_empty_encoder_model_runner_output,
@@ -7572,6 +7573,105 @@ def test_sync_diffusion_step_commits_only_emitted_tokens(sampled, expected_statu
     assert request.status == expected_status
 
 
+@pytest.mark.parametrize("async_scheduling", [False, True])
+@pytest.mark.parametrize("score_kind", ["prompt_logprobs", "token_ids"])
+@pytest.mark.skip_global_cleanup
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_prompt_scores_wait_for_emitted_tokens(async_scheduling, score_kind):
+    scheduler = _diffusion_scheduler(async_scheduling=async_scheduling)
+    request = _diffusion_request("scored", {})
+    scheduler.add_request(request)
+    scores = (
+        LogprobsTensors(
+            torch.tensor([[1, 2]], dtype=torch.int32),
+            torch.tensor([[-0.1, -0.2]]),
+            torch.tensor([0], dtype=torch.int32),
+        )
+        if score_kind == "prompt_logprobs"
+        else torch.tensor([[-0.1, -0.2]])
+    )
+    model_output = ModelRunnerOutput(
+        req_ids=[request.request_id],
+        req_id_to_index={request.request_id: 0},
+        sampled_token_ids=[[]],
+        prompt_logprobs_dict=(
+            {request.request_id: scores} if score_kind == "prompt_logprobs" else {}
+        ),
+        prompt_token_id_logprobs_dict=(
+            {request.request_id: scores} if score_kind == "token_ids" else {}
+        ),
+    )
+
+    prefill_outputs = scheduler.update_from_output(scheduler.schedule(), model_output)
+    assert all(not batch.outputs for batch in prefill_outputs.values())
+
+    if not async_scheduling:
+        scheduler.update_draft_token_ids(
+            DraftTokenIds([request.request_id], [list(range(8))])
+        )
+    decode_output = ModelRunnerOutput(
+        req_ids=[request.request_id],
+        req_id_to_index={request.request_id: 0},
+        sampled_token_ids=[[1] * 8],
+    )
+    outputs = scheduler.update_from_output(scheduler.schedule(), decode_output)
+    (emitted,) = next(iter(outputs.values())).outputs
+    if score_kind == "prompt_logprobs":
+        assert emitted.new_prompt_logprobs_tensors is scores
+        assert emitted.prompt_token_id_logprobs is None
+    else:
+        assert emitted.prompt_token_id_logprobs is scores
+        assert emitted.new_prompt_logprobs_tensors is None
+
+
+@pytest.mark.parametrize("exit_kind", ["abort", "preempt"])
+@pytest.mark.skip_global_cleanup
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_prompt_scores_cleared_before_request_restarts(exit_kind):
+    scheduler = _diffusion_scheduler(async_scheduling=False)
+    request = _diffusion_request("scored", {})
+    scheduler.add_request(request)
+    scores = torch.tensor([[-0.1]])
+    scheduler.update_from_output(
+        scheduler.schedule(),
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[]],
+            prompt_token_id_logprobs_dict={request.request_id: scores},
+        ),
+    )
+
+    if exit_kind == "abort":
+        scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
+    else:
+        scheduler._preempt_request(request, timestamp=0.0)
+    assert request.request_id not in scheduler._pending_prompt_token_id_logprobs
+
+
+@pytest.mark.skip_global_cleanup
+@pytest.mark.usefixtures("diffusion_model_runner")
+def test_diffusion_dropped_stale_prompt_scores_are_not_retained():
+    scheduler = _diffusion_scheduler(async_scheduling=False)
+    request = _diffusion_request("scored", {})
+    scheduler.add_request(request)
+    prefill = scheduler.schedule()
+    request.drop_stale_output = True
+    request.num_stale_output_tokens = prefill.num_scheduled_tokens[request.request_id]
+
+    scheduler.update_from_output(
+        prefill,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[]],
+            prompt_token_id_logprobs_dict={request.request_id: torch.tensor([[-0.1]])},
+        ),
+    )
+
+    assert request.request_id not in scheduler._pending_prompt_token_id_logprobs
+
+
 @pytest.mark.usefixtures("diffusion_model_runner")
 def test_diffusion_scheduler_defers_a_read_with_every_step_in_flight():
     scheduler = _diffusion_scheduler()

vllm/v1/core/sched/diffusion_scheduler.py

@@ -2,11 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Canvas-width handling and async read deferral for diffusion requests."""
 
+from collections.abc import Iterable
+
+import torch
+
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
-from vllm.v1.outputs import DraftTokenIds
-from vllm.v1.request import Request
+from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.outputs import DraftTokenIds, LogprobsTensors, ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
 
 
 def diffusion_canvas_width(request: Request, canvas_length: int) -> int:
@@ -35,6 +40,78 @@ def _read_in_flight(request: Request, width: int) -> bool:
 
 
 class DiffusionScheduler(Scheduler):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._pending_prompt_logprobs: dict[str, LogprobsTensors] = {}
+        self._pending_prompt_token_id_logprobs: dict[str, torch.Tensor] = {}
+
+    def update_from_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        model_runner_output: ModelRunnerOutput,
+    ) -> dict[int, EngineCoreOutputs]:
+        if not (
+            self._pending_prompt_logprobs
+            or self._pending_prompt_token_id_logprobs
+            or model_runner_output.prompt_logprobs_dict
+            or model_runner_output.prompt_token_id_logprobs_dict
+        ):
+            return super().update_from_output(scheduler_output, model_runner_output)
+
+        # Diffusion prefills score the prompt without emitting tokens.
+        for req_id, scores in model_runner_output.prompt_logprobs_dict.items():
+            request = self.requests.get(req_id)
+            if (
+                scores is not None
+                and request is not None
+                and not request.is_finished()
+                and not (
+                    request.drop_stale_output and request.num_stale_output_tokens > 0
+                )
+            ):
+                self._pending_prompt_logprobs[req_id] = scores
+        for req_id, scores in model_runner_output.prompt_token_id_logprobs_dict.items():
+            request = self.requests.get(req_id)
+            if (
+                request is not None
+                and not request.is_finished()
+                and not (
+                    request.drop_stale_output and request.num_stale_output_tokens > 0
+                )
+            ):
+                self._pending_prompt_token_id_logprobs[req_id] = scores
+
+        model_runner_output.prompt_logprobs_dict = {}
+        model_runner_output.prompt_token_id_logprobs_dict = {}
+        outputs = super().update_from_output(scheduler_output, model_runner_output)
+        for client_outputs in outputs.values():
+            for output in client_outputs.outputs:
+                output.new_prompt_logprobs_tensors = self._pending_prompt_logprobs.pop(
+                    output.request_id, None
+                )
+                output.prompt_token_id_logprobs = (
+                    self._pending_prompt_token_id_logprobs.pop(output.request_id, None)
+                )
+        return outputs
+
+    def finish_requests(
+        self,
+        request_ids: str | Iterable[str] | None,
+        finished_status: RequestStatus,
+    ) -> list[Request]:
+        finished = super().finish_requests(request_ids, finished_status)
+        for request in finished:
+            self._pending_prompt_logprobs.pop(request.request_id, None)
+            self._pending_prompt_token_id_logprobs.pop(request.request_id, None)
+        return finished
+
+    def _preempt_request(
+        self, request: Request, timestamp: float, drop_stale_output: bool = False
+    ) -> None:
+        self._pending_prompt_logprobs.pop(request.request_id, None)
+        self._pending_prompt_token_id_logprobs.pop(request.request_id, None)
+        super()._preempt_request(request, timestamp, drop_stale_output)
+
     def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
         for i, (req_id, token_ids) in enumerate(
             zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids)