Skip to content

fix: stop DFlash at end-of-sequence, and cancel it when the stream ends early - #225

Merged
solderzzc merged 1 commit into
SharpAI:mainfrom
HamedShams:fix/dflash-stop-and-cancel
Oct 10, 2026
Merged

solderzzc merged 1 commit into
SharpAI:mainfrom
HamedShams:fix/dflash-stop-and-cancel

Conversation

@HamedShams

Copy link
Copy Markdown
Contributor

Closes #224.

What

Two things kept a DFlash generation alive after its response had ended, and a third made that fatal:

  1. DFlash now stops at end-of-sequence. DFlashRuntime.generate was called without stopTokenIDs, so the loop ran to max_tokens past <|im_end|> (<|endoftext|><|im_start|>user… in the log). It now gets the same set the standard path stops on: the configuration's eosTokenIds, the tokenizer's EOS and extraEOSTokens, built by a small stopTokenIDs(…) helper.
  2. Ending the stream cancels the generation. The DFlash→Generation bridge was a bare Task nobody cancelled, so when handleChatStreaming ended a response on a text-level stop (or the client went away) and released the slot, the bridge kept draining and DFlashRuntime's own onTermination → task.cancel() never fired. The bridge is now dflashGenerationStream(_:tokenLimit:nextText:) with an onTermination that cancels it; tearing down the consumer side propagates to the DFlash stream.
  3. The DFlash loop honours cancellation. generateStreaming only checked Task.isCancelled in its yield closure, 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-decoding job 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 (its onTermination fires), and that tokens and the summary still map to .chunk / .info as before. 4 tests, and the full SwiftLMTests suite is green (284) on an M1 Max, Xcode 27 / Swift 6.4.
  • tests/test-dflash.sh gained Test 4: a streamed one-word answer with max_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 the dflash-speculative-decoding job on this PR is the run for it.
  • For the model CI uses, the stop set is non-empty: the tokenizer's EOS is <|im_end|> and text_config.eos_token_id is set.

Notes

  • DFlash commits a block per cycle, so up to one block of tokens after EOS can still be emitted; the text-level stop trims them as before. Stopping inside a block is a runtime change I kept out of this PR.
  • A client stop string 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.

…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.
@solderzzc
solderzzc merged commit 8811253 into SharpAI:main Oct 10, 2026
30 checks passed
@solderzzc

Copy link
Copy Markdown
Member

Thank you for the fix, and for the tests that pin each half without needing a model. Merged as 8811253. The dflash-speculative-decoding job passed on its first attempt in both runs we made on this branch, where it previously needed all three attempts or failed. One cycle between slot release and cancellation remains, as you noted, so we will keep an eye on that job.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

DFlash: generation ignores EOS and outlives the stream; the next request overlaps it and crashes the server (flaky dflash-speculative-decoding job)

2 participants