Repository navigation
fix: stop DFlash at end-of-sequence, and cancel it when the stream ends early - #225
Merged
Merged
Conversation
…ds early DFlashRuntime.generate was called without stopTokenIDs, so the loop ran to max_tokens past <|im_end|>. It now gets the same end-of-sequence set the standard path stops on. The DFlash -> Generation bridge was a bare Task nobody cancelled: when the SSE handler ended a response on a text-level stop and released its slot, the bridge kept draining and the DFlash task never saw the termination. The bridge is now dflashGenerationStream(_:tokenLimit: nextText:) with an onTermination that cancels it, and the DFlash cycle loop breaks on Task.isCancelled instead of only muting its output. With the slot free and the loop still running, the next request started a second DFlash generation on the same model and MLX aborted; that is the flaky dflash-speculative-decoding job. Closes SharpAI#224.
Member
|
Thank you for the fix, and for the tests that pin each half without needing a model. Merged as 8811253. The |
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
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.
Closes #224.
What
Two things kept a DFlash generation alive after its response had ended, and a third made that fatal:
DFlashRuntime.generatewas called withoutstopTokenIDs, so the loop ran tomax_tokenspast<|im_end|>(<|endoftext|><|im_start|>user…in the log). It now gets the same set the standard path stops on: the configuration'seosTokenIds, the tokenizer's EOS andextraEOSTokens, built by a smallstopTokenIDs(…)helper.Generationbridge was a bareTasknobody cancelled, so whenhandleChatStreamingended a response on a text-level stop (or the client went away) and released the slot, the bridge kept draining andDFlashRuntime's ownonTermination → task.cancel()never fired. The bridge is nowdflashGenerationStream(_:tokenLimit:nextText:)with anonTerminationthat cancels it; tearing down the consumer side propagates to the DFlash stream.generateStreamingonly checkedTask.isCancelledin itsyieldclosure, so cancellation muted output but not compute. The cycle loop now breaks on it.The standard path is untouched; it already does all three.
Why
The
dflash-speculative-decodingjob has been flaky: on main's last run it needed all three attempts, on #223's merge commit it failed all three. The uploaded server log showed two DFlash generations interleaving (Cycle 7… Cycle 1… Cycle 8… Cycle 2…) and MLX aborting with[eval] Attempting to eval an array without a primitive: the orphaned loop from a finished request racing the next request's generation on the same model. Details and line references in #224.Verified
DFlashStreamBridgeTests(new, no model needed): the stop-token union (deduplicated, unknown extra tokens dropped), that a consumer which stops early tears the DFlash stream down (itsonTerminationfires), and that tokens and the summary still map to.chunk/.infoas before. 4 tests, and the fullSwiftLMTestssuite is green (284) on an M1 Max, Xcode 27 / Swift 6.4.tests/test-dflash.shgained Test 4: a streamed one-word answer withmax_tokens: 200, immediately followed by a second request, which must complete with the server still answering/health. That is the failing scenario from the CI log in a deterministic shape. I don't have the 4B pair on this machine, so thedflash-speculative-decodingjob on this PR is the run for it.<|im_end|>andtext_config.eos_token_idis set.Notes
stopstring still releases the slot one DFlash cycle before the loop observes cancellation, the same shape as the standard path's one-token window. Holding the slot until the producer has wound down would be a separate change.