[Bugfix][Model] Fix Nemotron diffusion scores and async sampling

[pst2154]

Overview

Fix Nemotron diffusion scores changing when requests with different logprob settings share a batch. Capture each position's distribution when it is revealed, preserving the native model's diffusion generation, AR, and linear speculation implementation.

Draft follow-up to vllm-project/vllm#58387, targeting its implementation branch rather than introducing a competing model PR. It retains native checkpoint loading, diffusion text generation, AR, and linear speculation. The delta is one four-file commit above 943c956: c2b4b58.

Claims

• Sampled-token and top-k scores come from the same reveal-time distribution and are isolated from neighbouring requests' logprob settings.
• Diffusion output uses fixed per-request offsets and GPU sample counts, eliminating the sampler's host reads of convergence flags.
• Unsupported diffusion/linear-spec pipeline configurations fail during configuration. Request initialization and mask filtering use GPU fills.
• Linear-spec logprob output keeps acceptance counts and carried-seed selection on GPU. H100/B200 eager/compiled and A100 compiled AR-parity checks passed.
• Stochastic generation uses a compiled graph boundary to avoid an SM100 Inductor reduction-scope error, validated on A100, H100, and B200.

Validation

The original implementation was reproduced on H100 with the public 8B checkpoint: a sampled token's logprob changed from -0.535999 alone to -11.367270 in a mixed batch while generated tokens remained identical.

The final regression compares equal-sized three-request batches, resetting prefix cache between calls. One batch requests logprobs=0 throughout; the other requests None, 0, and 2. Output budgets are 17, 41, and 9 tokens with EOS ignored: the top-k neighbour finishes before the 41-token request's second block. It checks generated-token equality, complete scores, and sampled-token probabilities for both requesting rows, preserving tolerance 0.02.

Public checkpoint: nvidia/Nemotron-Labs-Diffusion-8B, revision 16c67f0560b912e93e0cabb6e0c4f5c3086d95fc. BF16, TP=1, PP=1, Triton attention, maximum model length 2,048. Validation uses VLLM_GPU_SYNC_CHECK=error.

GPU	Validation	Result
A100 PCIe 80GB	Four configuration cases, greedy diffusion generation, mixed-temperature rollouts, eager and compiled mixed-logprob regression	8 passed; strengthened two-case regression also passed
H100 PCIe 80GB	Same subset	8 passed; strengthened two-case regression also passed
B200	Same eight diffusion/configuration cases, across runs	8 passed after the stochastic compiler workaround
H100 / B200	Eager and compiled linear speculation versus AR	2 passed on each GPU
A100	Eager and compiled linear speculation versus AR	Compiled passed; eager matched tokens but failed logprob tolerance, identically in the unchanged parent sampler
H100 FP32	Diffusion-logit parity and both AR selectors versus native HF	3 passed

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 attention_backend_supports_mixed_causality or stochastic_rollouts or greedy_generation' \
  -v

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 linear_spec_greedy_ar_parity -v

# Manual reference-forward diagnostics intentionally copy logits to CPU.
NEMOTRON_DLM_MODEL_PATH=/path/to/Nemotron-Labs-Diffusion-8B \
VLLM_USE_V2_MODEL_RUNNER=1 \
  .venv/bin/python -m pytest \
  --confcutdir=tests/models/language/generation \
  tests/models/language/generation/test_nemotron_dllm.py \
  -k 'test_logits_parity or test_ar_logits_and_generation' -v

.venv/bin/pre-commit run --files \
  vllm/model_executor/models/nemotron_dllm.py \
  vllm/model_executor/models/config.py \
  tests/models/language/generation/test_nemotron_dllm.py \
  docs/models/nemotron_labs_diffusion.md

Pre-commit passed. CPU/configuration tests run in the existing model suite. Weight-dependent tests require a supplied local checkpoint; default CI skips them unless NEMOTRON_DLM_MODEL_PATH points to one. The GPU results above provide manual coverage for this gap. These tests do not establish a new model accuracy score, speedup, quantization support, or multi-GPU correctness. Stochastic generation was rerun after the compiler workaround on all three GPUs and passed. A100 eager linear-spec logprob parity differs by 0.12671864 (tolerance 0.12), with the exact same scores in the unchanged sampler; no assertion was relaxed. All three FP32 native HF-reference checks passed on H100 (158.68 seconds). RTX validation is not available.

Details

Previously, the batch's maximum requested top-k count selected the output path for every request. A neighbour asking for top-k scores could therefore switch a logprobs=0 request from reveal-time scores to final-step scores. Scores could also be absent if that neighbour finished before another request committed its block.

The sampler now records top-k scores per requesting slot and per newly revealed position. Completed-token scores persist until the next reveal overwrites that position. Sampled-only requests use their reveal buffers regardless of the batch's maximum top-k count. Fixed canvas-stride output offsets let the output processor use GPU-produced sample counts without synchronizing on phase flags. Request-slot reuse clears stored scores through the existing lifecycle.

Accurate top-k capture requires computing requested scores at reveal steps rather than only at convergence; this may increase work for requests that explicitly ask for top-k logprobs. No performance claim is made. The default path without logprobs avoids this capture.

Linear-spec score rows are computed for the scheduled block, padded to canvas stride, and emitted according to GPU-produced acceptance counts. Carried seed logits are selected and preserved on GPU, replacing the blocking host read. Computing scores for positions that are not accepted can increase work when logprobs are requested. Stochastic selection has a graph boundary before downstream confidence/rank reductions to avoid an Inductor-generated temporary escaping its loop scope on SM100; both sides remain compiled, but the boundary can add kernel launches.

AI assistance: OpenAI Codex assisted with implementation, diagnosis, tests, and documentation. The detailed results above were recorded by Codex. The submitter has confirmed reviewing the changed lines and running the relevant tests. This follow-up is ready for maintainer review; the disclosed A100 eager numerical discrepancy and CI coverage limits remain unchanged.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing

Tracking issue for the full three-part stack: vllm-project#60283

[pst2154]

Codex assisting @pst2154: this incorrectly targeted fork PR is superseded by vllm-project#60287, targeting upstream main. The branch and reviewed commit are preserved. Tracking: vllm-project#60283

[pst2154]

Codex assisting @pst2154: reopened at the submitter’s request so this remains available for integration into the model author’s branch. The same fixes are also submitted upstream in vllm-project#60287; these are two review routes for the same commit, not separate changes. Tracking: vllm-project#60283