Support thinking budgets during MTP speculative decoding - #2230
YuhuaBillChen wants to merge 3 commits into
Conversation
|
One idea that would remove the zero-accept rounds. Right now a forced token is only injected at the start of the next round with accepted_list[row] = 0, so each token of the close sequence costs one full verify, and with uniform acceptance that sets every row to 0 for that round. If the budget exposed a pure lookup, forced_token(prefix_ids) -> Optional[int], _apply_mtp_token_controls could resolve it inside the current round: speculative/mtp.py, per row (non-uniform case)hist = history[original_row] # committed ids for this row A draft that already equals the forced token is simply accepted. A mismatch looks exactly like an ordinary rejection with a replacement token, so the cache commit path doesn't change. Because the lookup has no state, rollback and EOS truncation can't desync it. For uniform acceptance you'd still take the min across rows, as now. We ship this pattern downstream as a logits processor (ThinkingBudgetProcessor, tests), including multi-token close markers. To be upfront: that coverage is CPU/unit-level, and we haven't GPU-qualified it under MTP. |
Summary
The server currently rejects
thinking_budgetwhen an MTP drafter is enabled. This change allows MTP to use the existingThinkingBudgetCriteria: it shortens verified blocks when the criterion requests a forced token, then feeds the closing sequence through the model so generation and cache state remain consistent.For targets requiring uniform batch acceptance, a boundary in one row limits the commit depth for every row. Token controls observe only the retained output. Requests without a thinking budget keep the existing fast path. DFlash and Eagle3 retain their existing rejection.
The existing criterion's counting and closing-token semantics are unchanged; this is not a new definition of an exact reasoning-token limit.
Validation
test_generate.py,test_speculative.py, andtest_server.pyon M4 Pro: 1,269 passed, with no skipped or deselected tests. The previously reported NVFP4 failures did not reproduce in this run.5fa214f6by SHA-256.thinking_start_token="<|think|>",thinking_end_token="<channel|>",enable_thinking=true. Counting starts from the prompt marker and includes the generated channel header.The Gemma runs are short functional smoke tests and ended at their output caps. They do not establish long-generation stability, answer accuracy, or a throughput improvement. Exact forced-token boundaries are covered by focused tests; the HTTP smoke responses alone do not provide a token-level proof.