Filter post-boundary speculative tokens before commit

[simon-mo]

Purpose

Mitigate DFlash performance collapse on structured-output/tool-guided requests by disabling DFlash draft proposal for scheduled batches with active structured-output grammar. DFlash drafts are currently unconstrained while the target logits are grammar-constrained, so draft work can be repeatedly invalidated or rejected for tool-call/JSON-style outputs. Both model runner paths now return empty draft-token lists for this unsupported combination and log a one-time warning.

Why this is not duplicating an existing PR: searched open PRs for DFlash structured output, dflash tool call, and speculative structured output. Related PRs exist for structured-output speculative decoding correctness and a prototype grammar-aware spec decode path, but none specifically mitigate DFlash by suppressing unconstrained DFlash drafts on structured-output batches.

AI assistance was used to investigate the bug report and prepare this change.

Test Plan

• .venv/bin/pre-commit run ruff-check --files vllm/v1/spec_decode/utils.py vllm/v1/worker/gpu/spec_decode/utils.py vllm/v1/worker/gpu_model_runner.py vllm/v1/worker/gpu/model_runner.py tests/v1/spec_decode/test_mtp_structured_output.py
• .venv/bin/python -m pytest tests/v1/spec_decode/test_mtp_structured_output.py -q

Test Result

• Ruff check: passed.
• Structured-output spec-decode tests: 17 passed, 25 warnings in 30.57s.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

  

[simon-mo]

The agent also investigated some ongoing PRs

Here’s the handoff summary.

Current finding

The original-like issue is reproducible with public models:

• Target: Qwen/Qwen3-8B
• DFlash head: z-lab/Qwen3-8B-DFlash-b16
• Config: V1 runner, --reasoning-parser qwen3, complex forced JSON schema
• Result:
  • DFlash acceptance collapses to ~8-9%
  • DFlash becomes ~1.9x slower than no-spec under concurrent load

This suggests the root is not simply “DFlash + structured output is incompatible.” Simple JSON/tool-call cases can still have high acceptance. The bad case is complex grammar-constrained generation with reasoning parser active, where DFlash drafts are unconstrained but target verification is grammar/reasoning constrained.

Existing PRs

vLLM PR vllm-project#36138

“Grammar was ignored when reasoning ended within speculated tokens.”

• Broad fix attempt for reasoning + structured output + speculative decoding.
• Tries to make grammar handling correct when speculative tokens cross the reasoning-to-answer boundary.
• More invasive: touches multiple grammar/spec-decode interaction points.
• Addresses correctness around reasoning boundary, not the full DFlash low-acceptance/perf issue.

vLLM PR vllm-project#43424

“structured_output × spec-decode: pre-commit grammar filter for boundary-step bonus tokens.”

• Narrower fix than [Bugfix] Grammar was ignored when reasoning ended within speculated tokens vllm-project/vllm#36138.
• Filters tokens before committing them to request output.
• Detects reasoning-end inside the generated token block.
• Validates only post-boundary tokens against grammar.
• Truncates invalid trailing speculative/bonus tokens before they leak into output.
• This is the approach I ported locally.

vLLM PR vllm-project#40898

“Add Sliding Window Attention support to DFlash drafter.”

• Relevant because exact Qwen3.6 NVFP4+DFlash reproduction was blocked locally by:
  • NotImplementedError: DFlash does not yet support mixed sliding/full attention via layer_types
• Needed to reproduce the exact original Qwen3.6 DFlash model pair cleanly.
• Separate from grammar/tool-call acceptance collapse, but required for exact model parity.

Local branch changes

Branch: cursor/skip-dflash-structured-cb76

Net diff against base is now only:

• vllm/v1/structured_output/__init__.py
• vllm/v1/core/sched/scheduler.py
• tests/v1/structured_output/test_reasoning_structured_output.py

Earlier incorrect skip-DFlash mitigation was reverted.

Proposed/local change

Add a pre-commit grammar filter:

1. Before _update_request_with_output, call:
   • StructuredOutputManager.precommit_filter_tokens(request, new_token_ids)
2. If reasoning-end completes inside new_token_ids:
   • Split into pre-boundary and post-boundary tokens.
   • Dry-run grammar validation on post-boundary suffix.
   • Keep only valid prefix.
   • Truncate invalid trailing tokens before output commit.
   • Rewind num_computed_tokens / num_output_placeholders for rejected tokens.
3. Avoid double-advancing structural-tag grammars when spec decode is active.

This preserves DFlash instead of skipping it.

What this fixes

Correctness issue:

• Prevents unconstrained post-reasoning speculative/bonus tokens from leaking into structured output.
• Aligns with PR [V1][Bugfix] structured_output × spec-decode: pre-commit grammar filter for boundary-step bonus tokens vllm-project/vllm#43424.
• Keeps DFlash enabled for structured-output/tool-call requests.

Validated locally:

• Unit test passed: tests/v1/structured_output/test_reasoning_structured_output.py
• Simple JSON/tool-call workloads still get high DFlash acceptance and speedups.

What this does not fix

It does not solve the full low-acceptance slowdown.

The clean repro still shows:

• DFlash + complex JSON + reasoning parser: ~8-9% acceptance
• DFlash slower than no-spec under concurrency

So the remaining issue likely needs one of:

1. Grammar-aware DFlash proposal

   • Feed grammar masks/constraints into the drafter sampling path.
   • Avoid proposing tokens that grammar will reject.

2. Adaptive DFlash disable/throttle

   • Per-request or per-batch dynamic fallback when observed useful acceptance is too low.
   • Should not be the old static “skip DFlash for any structured output” approach.
   • Needs to preserve high-acceptance tool-call/simple-JSON cases.

3. Exact-model enablement

   • Land/coordinate with [Spec Decode] Add Sliding Window Attention support to DFlash drafter vllm-project/vllm#40898 so Qwen3.6 NVFP4+DFlash mixed sliding/full attention works.
   • Then rerun the original reported model pair.

Recommended next steps

1. Do not submit the local pre-commit filter as an independent upstream PR without coordinating with [V1][Bugfix] structured_output × spec-decode: pre-commit grammar filter for boundary-step bonus tokens vllm-project/vllm#43424.
2. Use the clean repro script/report as the benchmark for future fixes.
3. Prototype grammar-aware DFlash sampling or adaptive per-request fallback.
4. Re-test:
   • simple JSON
   • complex forced JSON
   • real tool-call parser path
   • thinking enabled/disabled
   • concurrent load
5. After [Spec Decode] Add Sliding Window Attention support to DFlash drafter vllm-project/vllm#40898 or equivalent lands, rerun exact Qwen3.6 NVFP4+DFlash reproduction.