Conversation
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
ParBproject
left a comment
There was a problem hiding this comment.
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>
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_maskfield.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 > 0validation 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 ...andpre-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.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 importingtests/conftest.pybecausetblibis 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.