fix(spec decode): suppress EOS at draft positions in rejection sampler

[ToastyTheBot]

Summary

When using MTP speculative decoding, the rejection sampler's target model can produce EOS as the argmax at a draft position. The scheduler iterates the MTP burst tokens one-by-one via check_stop(), which immediately sets FINISHED_STOPPED when it encounters EOS — discarding all remaining tokens in the burst, including the bonus token that would have continued generation.

This manifests as premature stopping at reasoning-to-tool-call boundaries: the client receives finish_reason: "stop" with only reasoning_content and no tool_calls or content.

Observed hit rate: ~0.25% under concurrent load with num_speculative_tokens=3, 0% without MTP.

Context

We run Qwen3.6-27B-FP8 with MTP=3 in production and observed a suite of tool-calling issues with speculative decoding. After applying several open PRs together, tool call reliability improved dramatically:

• [Bugfix] Fix Qwen3 reasoning parser: raw text tags, transition loss, end detection, token counting, withhold recovery #40783 — fragmented <think/> tag handling in reasoning parser
• [Bugfix][ToolParser] Fix Qwen3 XML and Coder streaming tool call parser regressions #40861 — structural parsing fixes in Qwen3Coder tool parser
• Fix Qwen3 reasoning tool calls embedded inside think #39055 — embedded tool call recovery from reasoning
• [Bugfix] Grammar was ignored when reasoning ended within speculated tokens #36138 — grammar ignored when reasoning ends within speculated tokens
• [BugFix] Fix streaming tool call empty fields with MTP: Pydantic null serialization + qwen3coder early return #39598 — Pydantic null serialization in streaming tool call deltas

We also opened two other sibling draft PRs in an attempt to perfectly fix the issue:

• [Bugfix] Fix Qwen3Coder prev_tool_call_arr double-emission on parse failure #41466 — prev_tool_call_arr double-emission fallback
• [Bugfix] Detect MTP truncation at reasoning-to-tool-call boundary #41467 — Detect MTP truncation at reasoning-to-tool-call boundary

Root Cause Analysis

MTP speculative decoding produces two separate logit tensors:

• target_logits — covers draft positions (K tokens)
• bonus_logits — covers the position after the last accepted draft token

The rejection sampler compares draft tokens against target_logits's argmax. When the argmax at any draft position is EOS, the scheduler's _update_request_with_output calls check_stop() which immediately sets FINISHED_STOPPED, discarding all remaining tokens including the bonus token.

This is particularly problematic at the reasoning-to-tool-call boundary where the model's output transitions from reasoning content to tool-call XML. The target model correctly predicts the continuation at the bonus position, but EOS at a draft position causes the scheduler to stop before reaching it.

Changes

• RejectionSampler.__init__ accepts eos_token_ids (plural) parameter — collects all EOS variants
• RejectionSampler.forward suppresses all EOS tokens in target_logits after apply_sampling_constraints() and before rejection_sample(), using scalar column indexing ([:, eid].fill_()) to avoid the indexSelectSmallIndex CUDA kernel that asserts with large vocab sizes (observed with Qwen3.6-27B: vocab=248320, eos=248044)
• _collect_eos_token_ids helper gathers EOS IDs from multiple config sources:
  • model_config.hf_config.eos_token_id
  • model_config.hf_text_config.eos_token_id (multimodal models like Qwen3.6-27B nest EOS here)
  • generation_config.eos_token_id
  • Handles both int and list[int] — for Qwen3.6-27B, collects both 248044 (text_config) and 248046 (tokenizer)
• Only the large model runner (gpu_model_runner.py) is patched — the small runner uses a different RejectionSampler from vllm/v1/worker/gpu/model_runner.py with a different API
• Bounds check filters EOS IDs exceeding actual vocab dimension

Why This Is Safe

Generation is still bounded by multiple mechanisms:

• Bonus token: Sampled from a separate bonus_logits tensor — CAN still produce EOS for legitimate stops
• max_tokens: check_stop() enforces FINISHED_LENGTH_CAPPED
• min_tokens: check_stop() enforces minimum generation before any stop
• Worst case: At most K extra tokens (MTP=3 → 3 tokens) per burst before the bonus token handles the stop. At greedy temperature, the bonus token deterministically produces EOS when the model wants to stop.

Reproduction

Using Qwen3.6-27B-FP8 with MTP=3, send tool-calling requests with reasoning enabled under concurrent load. The truncation occurs at ~0.25% hit rate. Stress testing with --concurrent 4 or higher increases the hit rate.

Test Plan

• Verify non-MTP generation is unaffected
• Stress test with MTP=3 tool-calling workload
• Run vLLM CI test suite
• Verify normal (non-tool-calling) generation with MTP stops correctly

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

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

[mergify]

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

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

[gemini-code-assist]

Code Review

This pull request adds SiluAndMulWithClamp activation kernels for DeepSeek-V4 and refactors the pooling API to deprecate the score task and multitask support. It also improves V1 KV cache admission gating for sliding window and chunked-local attention to resolve potential deadlocks and updates the RejectionSampler to suppress EOS tokens at draft positions. Feedback notes that in-place logit modification for EOS suppression might impact logprobs observability, suggesting cloning the logits or modifying the scheduler logic.

[gemini-code-assist]

In-place modification of target_logits to suppress EOS tokens will affect logprobs if they are computed from this tensor later in the pipeline. While this correctly prevents premature stopping in speculative decoding, it results in -inf logprobs for EOS at draft positions, which may be unexpected for observability. Consider cloning the logits before masking if logprobs are required, or handling the EOS suppression in the scheduler's stop logic instead.

[mergify]

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

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