Skip to content

[Bugfix] Preserve sampling masks in DELTA and TITO streaming outputs - #55935

Draft
aoshen02 wants to merge 5 commits into
vllm-project:mainfrom
aoshen02:codex/stream-sampling-masks
Draft

aoshen02 wants to merge 5 commits into
vllm-project:mainfrom
aoshen02:codex/stream-sampling-masks

Conversation

@aoshen02

@aoshen02 aoshen02 commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

With sampling masks enabled, DELTA consumers and TITO SSE clients can lose the per-token support sets that are available in final output. Emit masks with each DELTA chunk, merge them when the collector merges token chunks, and serialize them in the existing TITO sampling_mask field.

The production diff is limited to three files: 15 added lines and one removed line. FINAL_ONLY and final CUMULATIVE behavior retain the existing aggregation. The existing top_k > 0 validation remains in place. Empty-token SSE handling is unchanged. No sampler, kernel, sampling distribution, or speculative-method support is changed.

The regression tests cover stream intervals, collector coalescing, all three output kinds, and TITO SSE with masks enabled and disabled.

Review scope

Address the unbounded-mask concern in review 5149564830 by restoring the original top-k validation and its original test. Remove the unrelated empty-terminal-output change. Remove the mask-length slice: mask chunks and emitted token IDs accumulate over the same interval, so serialization should not silently truncate the mask history.

Duplicate check

This updates the existing PR rather than opening another one. The previous duplicate check identified #50520 (non-streaming HTTP masks), #53826 (batch-sharded sampling), and #54166 (MTP mask generation); those changes address different boundaries from DELTA aggregation and TITO streaming.

Validation

Current revision, on Windows with an isolated Python environment:

  • python -m ruff check <the seven original PR Python files> using the repository-pinned Ruff 0.14.0: passed.
  • python -m ruff format --check <the same files>: passed, all seven files already formatted.
  • git diff --check: passed.
  • pre-commit run ruff-check --files ... and pre-commit run ruff-format --files ... on all five files in the final PR diff: passed. The full pre-commit run was interrupted during actionlint environment setup; a complete hook-suite pass is not claimed.
  • Supplemental isolated execution of the existing test_sampling_masks_follow_output_token_boundaries, loading the actual output classes/methods without the full vLLM import graph: base revision 8 passed / 4 failed; revised code 12 passed / 0 failed. This is not a full pytest or GPU run.
  • python -m pytest tests/v1/engine/test_output_processor.py -k sampling_masks_follow_output_token_boundaries -q: blocked while importing tests/conftest.py because tblib is missing. This local environment does not contain the complete vLLM/test/GPU dependencies; the normal test suite and TITO SSE tests have not run on this revision.

The earlier PR description reported a backport run on AMD64/H200-host and ARM64/GB200-host containers (15 passed each), and a real GB200 Qwen3-0.6B comparison with top-k=16/top-p=0.9 in which streaming/non-streaming token IDs and masks matched. Those are historical backport results, not new validation of this revision. Model/GPU validation has not been rerun here.

AI assistance: implemented and checked with OpenAI Codex. This remains a draft for human review; no completed human sign-off or full upstream CI pass is claimed.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

@ParBproject ParBproject left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The removal of the top_k > 0 guard re-opens the unbounded-mask case that the existing validation was explicitly preventing. The device representation is bounded/packed, but SamplingMaskTensors.tolists() expands any row wider than MAX_COMPACT_SUPPORT back into every supported token ID, and RequestState.sampling_mask_chunks keeps those expanded arrays until the final output for FINAL_ONLY/CUMULATIVE requests. With top_k=-1, top_p=1, that is effectively vocab_size int32 IDs per generated token (e.g. ~512 KiB/token for a 128k vocab, before Python/output serialization overhead), so a long request can accumulate gigabytes and OOM the API process. The new top_k=-1 test only generates a tiny response and does not exercise this scaling behavior. If unbounded top_k is now intended to be supported, could we keep the mask packed through final aggregation/transport or enforce another explicit support-size/output bound before dropping this guard?

Restore the existing top-k size guard, remove unrelated empty-terminal SSE changes, and avoid silently truncating accumulated sampling masks.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working frontend

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants