[Fix][Stacked] Preserve Nemotron diffusion scores and async sampling

[pst2154]

Overview

Fix Nemotron masked-diffusion logprobs changing or disappearing when neighbouring requests use different logprob settings. This stacked follow-up targets upstream main and depends on #58387. Tracking: #60283.

Review the isolated four-file fixes delta. The comparison against main includes the unmerged model-support dependency; this is not a competing model implementation.

Claims

• Preserve sampled-token and top-k scores from the distribution at each position's reveal step, independently of neighbouring requests.
• Use fixed per-request output offsets and GPU completion counts, avoiding host reads of convergence flags.
• Reject unsupported masked-diffusion pipeline parallelism during configuration.
• Retain the stochastic compiler workaround previously validated on SM100.

Linear-speculation bookkeeping is deferred with the removed linear-speculation feature rather than included in this fix.

Validation

Rebased onto main 1a3a7a8f01a8c1a3184151b8185072efda11daab and the reduced model-support dependency.

CPU-only tests of the complete stacked branch: 93 passed, 9 skipped. This runs the sampling-parameter suite and the Nemotron model unit/configuration suite with CUDA disabled. The nine skips are CUDA/checkpoint-gated cases, not GPU validation.

The complete applicable pre-commit suite passed. CUDA serving and model-quality checks have not been rerun on this base. Upstream CI remains subject to its author/label pre-run gate; local passing checks are not upstream CI success.

Historical manual validation used public nvidia/Nemotron-Labs-Diffusion-8B, revision 16c67f0560b912e93e0cabb6e0c4f5c3086d95fc, BF16, TP=1, PP=1, Triton attention and context 2,048. Diffusion/configuration subsets passed on A100, H100 and B200, including eager/compiled mixed-logprob tests and stochastic generation. Three FP32 HF-reference checks passed on H100. Those results predate this rebase; RTX was not validated.

The original H100 reproducer changed a sampled token's logprob from -0.535999 alone to -11.367270 in a mixed batch without changing its token. The regression compares equal-sized three-request batches, resets the prefix cache, requests None, 0 and 2 logprobs, and uses token budgets 17, 41 and 9. It checks token equality, complete scores and sampled-token probabilities within 0.02 absolute tolerance, including when the top-k neighbour finishes early.

.venv/bin/python -m pytest --noconftest tests/test_sampling_params.py \
  tests/models/language/generation/test_nemotron_dllm.py -q

NEMOTRON_DLM_MODEL_PATH=/path/to/Nemotron-Labs-Diffusion-8B \
VLLM_USE_V2_MODEL_RUNNER=1 VLLM_GPU_SYNC_CHECK=error \
  .venv/bin/python -m pytest --confcutdir=tests/models/language/generation \
  tests/models/language/generation/test_nemotron_dllm.py \
  -k 'logprobs_independent_of_batch_requests or stochastic_rollouts or greedy_generation' -v

Details

Previously, the batch-wide maximum top-k request selected the score-output path for every request, which could replace reveal-time scores with final-step scores or lose scores when that neighbour completed. Per-request reveal buffers and retained top-k entries remove that dependency. Requested top-k capture can add computation; no end-to-end speedup or new model-quality score is claimed.

The PP rejection test asserts the PP=1 restriction without depending on the old pluralized error message. No numeric tolerance was relaxed. After the dependency merges, rebase and rerun the affected CUDA checks before merge.

Codex assisted with implementation, review fixes and checks. Human review of the revised diff and confirmation of required validation remain necessary.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill.
• AI assistance was used during the creation of this PR.
• Design Fit: Changes are isolated to model-specific sampler/configuration behavior.
• Testing and Validation: Fresh CUDA serving verification remains outstanding; historical reproducer results and current checks are separated above.
• Code Quality and Style: Applicable checks pass; regression assertions retain their behavioral checks.
• Pull Request Contents: Dependencies, root cause, evidence and tradeoffs are documented.

[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.

[mergify]

Documentation preview: https://vllm--60287.org.readthedocs.build/en/60287/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @pst2154.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[pst2154]

Codex update (AI-assisted): rebased this follow-up and its dependency commits onto upstream main at 89ac8d9. The PR no longer has merge conflicts. Added a compatibility fix replacing legacy Transformers PretrainedConfig imports with PreTrainedConfig, as required by current main. The original model-support PR #58387 is unchanged and remains a dependency.

The complete local pre-commit suite passed, including type checks and forbidden-import checks. The System One OpenAPI schema also validated on the dependent decision branch. Historical GPU results in the description are not fresh validation of this rebase; GPU serving tests have not been rerun. Upstream CI remains behind its existing author/label pre-run gate, not a reported model-test failure.

Updated isolated fixes comparison: pst2154/vllm@88a6778...37e0476

[pst2154]

Codex update: restacked the fixes on the reduced #58387 implementation and current main. Linear-speculation bookkeeping is deferred with that feature; masked-diffusion score isolation, asynchronous completion handling and PP rejection remain. Updated the PP rejection test to match the PP=1 restriction rather than the former pluralized error message. The PR description and isolated comparison link are updated.

Current full-stack CPU/unit checks: 93 passed, 9 CUDA/checkpoint-gated skips. Full local pre-commit passed. Fresh CUDA serving checks remain outstanding; historical GPU results are explicitly labelled. No GPU deployments were changed.