[Bugfix][Structured Output][Spec Decode] Fix async grammar bitmask alignment after draft trimming

[yanghui1-arch]

Purpose

Fixes #49694.

With async scheduling, ngram_gpu can cause the engine core and model runner to observe different speculative-token lengths for the same request. The model runner trims invalid draft slots based on the GPU-side n-gram result, while the engine core constructs the structured-output grammar bitmask using the original scheduled span.

Root cause

apply_grammar_bitmask previously used the model runner's trimmed draft length
for both:

1. locating the destination logits
2. advancing through the flattened bitmask generated by the engine core.

When a request's drafts were trimmed, the source offset became incorrect. Consequent requests in the same batch could receive bitmask rows belonging to the previous request, allowing grammar-invalid draft tokens to pass speculative verification. This explains why the failure rate increased with concurrency.

How does PR work

This PR records the engine core's per-request speculative-token count in GrammarOutput and uses it to advance through the source grammar bitmask. Worker-side lengths are still used to locate the actual logits, keeping draft and bonus-token rows correctly aligned when the two lengths differ.

Test Plan

Run the same structured-output serving workload before and after the fix using Qwen3.6-35B-A3B-FP8 with async scheduling, the xgrammar backend, and ngram_gpu speculative decoding with three speculative tokens.

The reproduction script is available here.

CUDA_VISIBLE_DEVICES=0 VLLM_LOGGING_LEVEL=INFO vllm serve Qwen/Qwen3.6-35B-A3B-FP8 \
  --served-model-name qwen3.6-35b-a3b-fp8 \
  --host 0.0.0.0 \
  --port 8000 \
  --tensor-parallel-size 1 \
  --max-model-len 4096 \
  --max-num-seqs 32 \
  --gpu-memory-utilization 0.90 \
  --async-scheduling \
  --structured-outputs-config.backend xgrammar \
  --speculative-config '{
    "method": "ngram_gpu",
    "num_speculative_tokens": 3,
    "prompt_lookup_min": 1,
    "prompt_lookup_max": 5
  }'

python repro_issue_49694.py \
  --base-url http://127.0.0.1:8000 \
  --model qwen3.6-35b-a3b-fp8 \
  --concurrency 1 4 8 16 32 \
  --requests 256

Test Result

The following table compares the overlapping concurrency levels from the reproduction before the fix and the same workload after the fix.

Concurrency	HTTP 5xx before	HTTP 5xx after
1	0 / 256	0 / 256
4	24 / 256	0 / 256
8	77 / 256	0 / 256
16	153 / 256	0 / 256
32	183 / 256	0 / 256

After the fix, all 1,280 requests completed without HTTP, transport, or internal errors. All outputs had the expected constrained length of 225 characters, with no wrong or truncated outputs.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[yanghui1-arch]

cc @WoosukKwon @njhill @benchislett @yewentao256 @orozery
PLTK. Thanks

[yanghui1-arch]

@mgoin Can you review my PR? Thanks

[yanghui1-arch]

@ivanium Do you have free time to review my PR? Thanks a lot

[yanghui1-arch]

@yewentao256 @njhill Can you take a look this PR? Thanks

[yanghui1-arch]

Please take a look. Thanks for your review. @ApostaC @robertgshaw2-redhat

[yanghui1-arch]

@chaunceyjiang Can you help review this pr? Thanks

[chaunceyjiang]

This might need a review from @benchislett.

[yanghui1-arch]

Hi, can you help review this PR? Thanks for your help. @benchislett

[yanghui1-arch]

This might need a review from @benchislett.

ok. Thanks

[yanghui1-arch]

@mgoin @aarnphm Please take a look

[yanghui1-arch]

Can you help review this pr? Thanks @yewentao256

[yewentao256]

Thanks for the work! I think @benchislett would be the best person to review this PR

[yanghui1-arch]

Can you help review this pr? Thanks.... @benchislett

[yanghui1-arch]

cc @benchislett

[yanghui1-arch]

cc @benchislett

[yanghui1-arch]

cc @benchislett @ivanium

[yanghui1-arch]

Please take a look @benchislett

[yanghui1-arch]

Hello, please take a look @benchislett

[sfeng33]

Superseded by #52477 as a more general fix.

[seanyourhighness]

Related MTP async validation and a complementary alignment path

I completed an async structured-output stress run on MTP-3 that may be useful for this PR's test matrix. This is not a test of the PR branch: my local stack used a different but complementary scheduler-side change.

The distinction as I understand it:

• [Bugfix][Structured Output][Spec Decode] Fix async grammar bitmask alignment after draft trimming #49738 preserves the engine-core grammar-bitmask source span when the worker has already shortened its draft list, preventing later requests from reading the wrong mask rows.
• My local MTP path kept the raw draft IDs in Scheduler.update_draft_token_ids_in_output and used grammar.validate_tokens only to count invalid positions, rather than replacing the invalid suffix with -1 before bitmask/verifier processing. This kept the grammar rows aligned with the exact raw positions the sampler saw.

With that local behavior plus the reasoning-boundary and XGrammar termination fixes:

• Two 900-request async structured hammers: 1,800/1,800 HTTP 200, zero FSM/grammar/server errors, approximately 84% MTP acceptance.
• Forced eight-active MTP-3 replay: 857.4 tok/s, zero errors.
• Real parent + seven-child soak: 58 requests, 1.46M prompt tokens, zero structured-output failures.

The approaches appear complementary because worker-side ngram/drafter trimming can still occur independently of grammar validation. Before proposing a separate MTP change, I plan to test four arms on current main: main, #49738, raw-draft preservation, and both together. If #49738 alone covers the MTP failure, I will not open another PR; otherwise I will document the materially different invariant and keep the follow-up narrowly scoped.

Disclosure: AI-assisted analysis and comment posting; the runs and measurements were produced and verified by me on the hardware described.