[None][fix] spec-decode NaN guard (~isnan, not isfinite) + draft_len_schedule CUDA-graph crash - #10
Merged
sopwg612 merged 2 commits intoJun 17, 2026
Conversation
…schedule reaching 0 crashes during CUDA graph capture, "bitmask must have the same batch size as logits" (NVIDIA#15023) Signed-off-by: chungen04 <b09901027@ntu.edu.tw>
…ided-decoding -inf is legal) Standardize the spec-decode target-logits guard on (~torch.isnan(logits)).all() so guided-decoding -inf (legitimate grammar masking) never trips it, while real NaN (argmax->PAD corruption) still does. Aligns deep-main with the priority build branch and the shipped rc16 images. Signed-off-by: Sihan Wang <sihan@deepinfra.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Restore the spec-decode target-logits NaN guard to a NaN-only check, and pull in the draft_len_schedule CUDA-graph crash fix.
Why
The guard in
_torch/speculative/interface.py(SpecWorkerBase) andmtp.pywas rewritten in 23939ea ("FP8 KV-cache + spec-decode robustness") totorch.isfinite(logits).all().isfiniteis False for-infas well as NaN. Guided decoding (xgrammar) writes-infinto masked-token logits as normal grammar masking, so the guard false-positivelyos._exits the engine on ANY structured-output (response_format) request under speculative decoding (Eagle3/MTP). This took down guided+spec models on rc16 (gpt-oss-120b-Turbo, MiniMax-M2.5, Kimi-K2.5, ...).Fix
(~torch.isnan(logits)).all(): fire only on real NaN (the argmax->PAD corruption the guard targets).-inffrom masking is legitimate and handled correctly by softmax/argmax. Restores the NaN-only intent of the original 6f5c96a.Validated in prod (Turbo + MiniMax): guided+spec requests no longer crash, structured output returns valid JSON, NaN protection retained.