[jperezdealgaba]

Overview

DiffusionGemma can finish prefill with prompt scores and no sampled tokens. The base scheduler treats scores on a no-output step as an invariant violation, so a normal prompt-scoring request terminates EngineCore. Keep the scores in the Diffusion scheduler until the request emits output, and discard them if that request is aborted or preempted.

Claims

• Prompt logprobs and prompt-token-ID scores survive a zero-token prefill and appear on the first emitted output in synchronous and asynchronous scheduling.
• Aborted, preempted, and dropped-stale requests do not retain scores for a later output.
• The ordinary no-score path continues through the base scheduler unchanged.

Validation

• VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest tests/v1/core/test_scheduler.py -q --tb=short: 216 passed on macOS (using the existing workspace virtual environment).
• pre-commit run --files vllm/v1/core/sched/diffusion_scheduler.py tests/v1/core/test_scheduler.py: passed, including ruff and mypy.
• pre-commit run --all-files: failed on two unrelated macOS checks: mypy references Linux-only MADV_HUGEPAGE / posix_fadvise constants, and the non-root entrypoint test compares /var with its /private/var path.
• Served DiffusionGemma validation is pending on a host with Triton and the model available. The scheduler test reproduces and fixes the failing step, but this workstation cannot run the model end to end.

Details

The model runner returns prompt scores at final prefill; DiffusionGemma samples zero tokens in that step. Scheduler.update_from_output then asserts that no prompt scores exist because it emits no output. The fix is confined to DiffusionScheduler: it holds both score forms by request ID, attaches them to the first emitted output, and clears them when a request is aborted or preempted. Scores from an output the base scheduler would drop as stale are ignored. No API request fields or model computation change.

Duplicate-work check: #60527 changes DiffusionGemma prompt-logit computation but leaves the scheduler assertion in place. The diffs of #58668, #52054, #50795, and #57597 address other scheduler or prompt-logprob paths. No open or merged PR in the sweep contains this missing scheduler control.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill.
• AI assistance was used during the creation of this PR.
• Design Fit: Keeps model-specific state in the Diffusion scheduler and preserves the base scheduler's output invariant.
• Testing and Validation: Scheduler tests pass; served-model validation remains pending.
• Code Quality and Style: The changed files pass lint and type checks.
• Pull Request Contents: Includes root cause, codepath coverage, duplicate-work check, test results, and the remaining validation requirement.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.