[DFLASH] Support grammar-constrained decoding in speculative verify

[hsthe29]

Motivation

Under --speculative-algorithm DFLASH, any grammar-constrained request is rejected with HTTP 400:

DFLASH speculative decoding does not support grammar-constrained decoding yet.

This makes DFLASH unusable for structured / forced tool calling and JSON output — a very common serving path (agents, function calling, document extraction). Verified trigger surface:

Request	before	after
tool_choice: "auto"	✅ OK	✅ OK
tool_choice: "required"	❌ 400	✅ OK
tool_choice: {named function}	❌ 400	✅ OK
response_format: {type: json_object}	❌ 400	✅ OK
response_format: {type: json_schema}	❌ 400	✅ OK
regex / ebnf / structural_tag	❌ 400	✅ OK
plain chat (no constraint)	✅ OK	✅ OK

validate_dflash_request() rejects any request whose sampling_params carries json_schema / regex / ebnf / structural_tag, because the DFLASH verify path never applied the grammar vocab mask (the DFLASH worker had no grammar handling at all). Simply removing the guard would be unsafe — grammar requests would then run through an unconstrained verify and emit malformed output silently.

Modifications

Two observations keep the change small and let it reuse the existing, proven EAGLE machinery:

1. DFLASH verify is a linear draft chain: draft_tokens[bs, block_size], where column 0 is the current (already-committed) token and columns 1: are the draft proposals. A chain is a degenerate case of EAGLE's draft tree — each node's only child is the next position, no siblings — so the existing helper sglang.srt.speculative.spec_utils.generate_token_bitmask() / traverse_tree() is reused as-is by constructing retrieve_next_token = [1, 2, …, -1] and retrieve_next_sibling = [-1, …].
2. Grammar FSM advancement over committed tokens is already generic for spec-v2 in SchedulerOutputProcessorMixin._resolve_spec_v2_tokens → _accept_grammar_tokens. DFLASH is spec-v2, so no change is needed there; the only missing piece was the mask application during verify.

Changes:

• python/sglang/srt/speculative/dflash_utils.py — remove the grammar-constraint rejection in validate_dflash_request (the return_logprob / return_hidden_states checks are kept).
• python/sglang/srt/speculative/dflash_worker_v2.py — in the verify block, right after apply_dflash_verify_logits_adjustments, when batch.has_grammar: build the linear-chain topology, call generate_token_bitmask(...), and grammar.apply_vocab_mask(next_token_logits, vocab_mask) before the accept/argmax step. This masks both the greedy-argmax and the sampling-softmax accept paths, since both read next_token_logits.

The masking mirrors EAGLE's grammar-aware verify (eagle_worker_v2.py + eagle_utils.py), so accepted tokens along the chain are grammar-valid by construction, and the downstream spec-v2 grammar FSM advancement stays correct.

Accuracy Tests

Model: Qwen3.6-35B-A3B, --speculative-algorithm DFLASH --speculative-dflash-block-size 8, --grammar-backend xgrammar, --tool-call-parser qwen3_coder (Tested on AMD ROCm / triton attention backend).

• All previously-400 request shapes now return 200 / well-formed output.
• Complex nested extraction (document-KIE shape, tool_choice: required + nested object/array schema), 10 trials @ temperature=0, streaming: 10/10 schema-valid, deterministic (1 distinct output), and 7/7 ground-truth field checks (invoice number, grand total, currency, line-item count, seller, SWIFT, Σ line_totals == grand_total).
• Diverse tool-call suite (all pass): multi-tool auto selection, tool_choice: required with mixed types (int / bool / enum / array), parallel tool calls (3 correct calls for a 3-city prompt), multi-turn with a returned tool result, Vietnamese unicode arguments preserved, streaming reassembly, and the no-tool-needed case (model does not emit a spurious call).
• Server log is clean: no accept_token failed, no traceback / abort.

Tests added:

• Unit (CPU CI) — test/registered/unit/spec/test_spec_utils_traverse_tree.py: two cases for the linear-chain topology that DFlashWorkerV2 feeds to generate_token_bitmask (retrieve_next_token = [1, 2, …, -1], no siblings) — one asserting the whole chain is walked in order, one asserting traversal stops when a chain token is disallowed by the grammar.
• Integration — test/registered/spec/dflash/test_dflash.py::test_grammar_constrained_decoding: launches the DFLASH server and sends a response_format: json_schema request (asserts schema-valid JSON) and a regex request (asserts the output matches the pattern). Runs across all DFLASH server configs in the file (page size, chunked prefill, no-cuda-graph, spec-v2, plan-stream).

Benchmarking and Profiling

DFLASH remains genuinely active on grammar-constrained requests — it is not a silent fallback to non-speculative decoding:

• spec_accept_rate ≈ 0.79, spec_accept_length ≈ 6.5 / 8, single-request decode throughput ~620–660 tok/s.

So grammar requests keep the speculative-decoding speedup while now producing correct, constrained output.

Reproduce

Launch the server with DFLASH + a grammar backend (--log-requests so per-request spec metrics are printed):

python3 -m sglang.launch_server \
  --model-path Qwen/Qwen3.6-35B-A3B \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path modal-labs/Qwen3.6-35B-A3B-DFlash \
  --speculative-dflash-block-size 8 \
  --grammar-backend xgrammar \
  --tool-call-parser qwen3_coder \
  --reasoning-parser qwen3 \
  --attention-backend triton \
  --log-requests --log-requests-level 1 \
  --host 0.0.0.0 --port 30000

Send a grammar-constrained request (before this PR: HTTP 400; after: schema-valid output):

curl -s http://localhost:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen3.6-35B-A3B",
    "messages": [{"role": "user", "content": "What is the capital of France? Reply as JSON."}],
    "max_tokens": 128, "temperature": 0,
    "chat_template_kwargs": {"enable_thinking": false},
    "response_format": {"type": "json_schema", "json_schema": {"name": "ans",
      "schema": {"type": "object", "properties": {"capital": {"type": "string"}},
                 "required": ["capital"], "additionalProperties": false}}}
  }'
# -> {"capital": "Paris"}

Read the speculative acceptance for that request from the server log (confirms DFLASH is still speculating under grammar, not falling back to 1 token/step):

# a `Finish:` line reports e.g. spec_accept_length ~6.5, spec_accept_rate ~0.79
grep "spec_accept_length" server.log | tail -1

Aggregate metrics are also available via --enable-metrics (sglang:spec_accept_length in Prometheus). Swapping response_format for tool_choice: "required", a named tool_choice, response_format: {type: json_object}, or a top-level "regex": "..." exercises the other previously-rejected grammar paths.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to the Documentation Guide.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to the Benchmark and Profiling Guide and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #30155825631
Latest PR Test (Extra): ❌ Run #30155825531

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[hsthe29]

Hi @Ying1123 @merrymercy @hnyls2002 @Qiaolin-Yu, could you please trigger CI for this PR? Thank you!

[shanemort1982]

Heads up from deploying this in production. We rolled DSpark onto DeepSeek-V4-Flash (TP2) and DeepSeek-V4-Pro (TP8) on B300 and this PR is the exact fix we needed, but as-is it is incomplete for DSPARK and silently corrupts output rather than fixing it. Two issues:

1. The guard removal covers DSPARK, the masking does not.
The grammar check is deleted from the shared validate_dflash_request(), which the scheduler gates on is_dflash_family() (is_dflash() or is_dspark()), so it stops rejecting DSPARK requests too. But the masking is only added to DFlashWorkerV2.forward_batch_generation. DSPARK runs DSparkWorkerV2, a sibling under BaseSpecWorker with its own _forward_decode, so no mask is ever applied. Net result on DSPARK: grammar requests are accepted and return unconstrained output. We measured 20/20 response_format: json_schema requests coming back as truncated invalid JSON. A silent correctness failure is worse than the original 400.

2. Likely off-by-one in the DFlash masking itself.
generate_token_bitmask is built over draft_tokens (block_size columns), but the TARGET_VERIFY forward returns logits over the full chain verify_ids_2d = cat([anchor, drafts]) = block_size + 1 columns. The bitmask rows and the logit rows do not line up, so the FSM advances over positions that were never committed. The symptom is exactly the mid-value truncation we saw. Suggest building the mask over verify_ids_2d.

Our DSPARK fix (validated in production, 0/20 -> 20/20 schema-valid, GSM8K unchanged, speculation still active): three changes to DSparkWorkerV2._forward_decode -

• build the bitmask over verify_ids_2d so mask rows align 1:1 with the target logit rows;
• force the eager path when batch.has_grammar (the folded DsparkVerifyEpilogue runs accept inside the CUDA graph and silently ignores a mask applied to next_token_logits in Python);
• apply the mask before accept_and_finalize, mirroring EAGLE.

Running live on real client traffic (response_format json_schema / json_object, tool_choice: required / named) with zero grammar rejections. Happy to open a PR with the DSPARK side on top of this, or fold it in here, whichever you prefer.

cc @hsthe29 @merrymercy @Ying1123 @Qiaolin-Yu

[hsthe29]

Heads up from deploying this in production. We rolled DSpark onto DeepSeek-V4-Flash (TP2) and DeepSeek-V4-Pro (TP8) on B300 and this PR is the exact fix we needed, but as-is it is incomplete for DSPARK and silently corrupts output rather than fixing it. Two issues:

1. The guard removal covers DSPARK, the masking does not.
The grammar check is deleted from the shared validate_dflash_request(), which the scheduler gates on is_dflash_family() (is_dflash() or is_dspark()), so it stops rejecting DSPARK requests too. But the masking is only added to DFlashWorkerV2.forward_batch_generation. DSPARK runs DSparkWorkerV2, a sibling under BaseSpecWorker with its own _forward_decode, so no mask is ever applied. Net result on DSPARK: grammar requests are accepted and return unconstrained output. We measured 20/20 response_format: json_schema requests coming back as truncated invalid JSON. A silent correctness failure is worse than the original 400.

2. Likely off-by-one in the DFlash masking itself.
generate_token_bitmask is built over draft_tokens (block_size columns), but the TARGET_VERIFY forward returns logits over the full chain verify_ids_2d = cat([anchor, drafts]) = block_size + 1 columns. The bitmask rows and the logit rows do not line up, so the FSM advances over positions that were never committed. The symptom is exactly the mid-value truncation we saw. Suggest building the mask over verify_ids_2d.

Our DSPARK fix (validated in production, 0/20 -> 20/20 schema-valid, GSM8K unchanged, speculation still active): three changes to DSparkWorkerV2._forward_decode -

• build the bitmask over verify_ids_2d so mask rows align 1:1 with the target logit rows;
• force the eager path when batch.has_grammar (the folded DsparkVerifyEpilogue runs accept inside the CUDA graph and silently ignores a mask applied to next_token_logits in Python);
• apply the mask before accept_and_finalize, mirroring EAGLE.

Running live on real client traffic (response_format json_schema / json_object, tool_choice: required / named) with zero grammar rejections. Happy to open a PR with the DSPARK side on top of this, or fold it in here, whichever you prefer.

cc @hsthe29 @merrymercy @Ying1123 @Qiaolin-Yu

Thanks for the detailed investigation and for validating this in production!

Feel free to cherry-pick my commit and build the DSPARK changes on top of it. That should preserve the original commit history/authorship while letting you iterate on the DSPARK-specific fixes. I'd be happy to review the follow-up PR once it's ready.

[hsthe29]

@shanemort1982

[shanemort1982]

Opened the DSPARK counterpart as #31753.

Following on from my note above: this PR lifts the grammar guard out of the
shared validate_dflash_request (gated on is_dflash_family(), so it covers
both DFLASH and DSPARK) but only masks in DFlashWorkerV2. A DSPARK deployment
routes to the sibling DSparkWorkerV2, so with the guard gone it would accept
grammar requests and then return unconstrained output. #31753 keeps the guard
for DFLASH and adds the masking in DSparkWorkerV2 (linear verify chain, mask
built over verify_ids_2d so rows line up, plus a has_grammar exclusion from
the folded epilogue). Validated on a DeepSeek-V4 DSPARK deployment: 0/20 to 20/20
schema-conforming.

Happy to reconcile whichever of the two lands first; the only overlap is the one
if in validate_dflash_request.

[hnyls2002]

/rerun-test test_dflash.py test_spec_utils_traverse_tree.py

[github-actions]

Results for /rerun-test test_dflash.py test_spec_utils_traverse_tree.py:

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/dflash/test_dflash.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/spec/test_spec_utils_traverse_tree.py

[hnyls2002]

/tag-and-rerun-ci