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

The dflash-speculative-decoding CI job fails intermittently at Test 5 ("Sequential request 1: empty response") with the server dying on Trace/BPT trap: 5. On main's own last run (748c4bf) it needed all three attempts to pass; on #223's merge commit (run 38034726754) all three attempts failed. The dflash-test-logs artifact shows the cause, and the code confirms it.

What the server log shows

Test 3 (streaming, max_tokens: 30) finishes its answer and emits <|im_end|>. The server's text-level default stop ends the SSE stream there, the client sees [DONE], and the generation slot is released. But the DFlash loop keeps going:

 juicy,[DFlash] Cycle 6: blockLen=16, verifyLen=16, acceptanceLen=1, commitCount=2
<|im_end|>[Server Debug] Created UserInput with 0 images and 0 audio inputs.
srv  slot_launch: id 0 | prompt=18t | thinking=false | prefilling...
[DFlash] Cycle 7: blockLen=16, verifyLen=16, acceptanceLen=0, commitCount=1
…
[DFlash] Cycle 9: blockLen=16, verifyLen=16, acceptanceLen=1, commitCount=2
[DFlash] Cycle 1: blockLen=9, verifyLen=9, acceptanceLen=2, commitCount=3
<|im_end|>
[DFlash] Cycle 10: blockLen=14, verifyLen=14, acceptanceLen=0, commitCount=1
[DFlash] Cycle 2: blockLen=7, verifyLen=7, acceptanceLen=0, commitCount=1
<|endoftext|>[DFlash] Cycle 3: …
<|im_start|>[DFlash] Cycle 4: …
user<|im_end|>
MLX/ErrorHandler.swift:345: Fatal error: [eval] Attempting to eval an array without a primitive.

Two DFlash generations are interleaved on the same model (cycles 7–13 of the finished request, cycles 1–5 of the new one), and MLX aborts. Whether the overlap happens depends on how many of the 30 tokens were still unused when <|im_end|> appeared, which is why the job is flaky rather than red.

Why, in the code

1. DFlash never sees EOS. The server calls DFlashRuntime.generate without stopTokenIDs, so the loop runs to max_tokens and keeps emitting <|endoftext|><|im_start|>user… after the answer:
   

   SwiftLM/Sources/SwiftLM/Server.swift

   Lines 2525 to 2530 in 748c4bf

   	let dflashStream = DFlashRuntime.generate(
   	targetModel: targetModel,
   	draftModel: dflashDraft,
   	promptTokens: promptTokens,
   	maxNewTokens: tokenLimit,
   	blockTokens: dflashBlockSize

   
   The stop check in the runtime is there but gets an empty set:

   SwiftLM/Sources/DFlash/DFlashRuntime.swift

   Line 622 in 748c4bf

   let hit = committedIDs.contains { stopTokenSet.contains($0) }

   . The standard path stops on eosTokenIds + tokenizer.eosTokenId + extraEOSTokens (mlx-swift-lm Evaluate.swift L2338–L2378).

2. Nothing cancels it when the stream ends early. The DFlash→Generation wrapper is an inner Task with no onTermination:
   

   SwiftLM/Sources/SwiftLM/Server.swift

   Lines 2540 to 2571 in 748c4bf

   	let genStream = AsyncStream<Generation> { continuation in
   	Task {
   	for await event in dflashStream {
   	switch event {
   	case .token(let tokenID, _, _, _):
   	box.detokenizer.append(token: tokenID)
   	if let chunk = box.detokenizer.next() {
   	continuation.yield(.chunk(chunk, tokenId: tokenID))
   	}
   	case .prefill, .prefillProgress:
   	break
   	case .summary(let summary):
   	print("[SwiftLM] DFlash summary: \(summary.generationTokens) tokens, \(String(format: "%.1f", summary.tokensPerSecond)) tok/s, acceptance=\(String(format: "%.1f%%", summary.acceptanceRatio * 100)), \(summary.cyclesCompleted) cycles")
   	// The SSE/non-streaming handlers emit finish_reason, usage and
   	// `[DONE]` from `.info`. Without it, a DFlash run that ends on EOS
   	// or max_tokens (rather than a textual stop sequence) closes the
   	// stream with no `[DONE]` sentinel.
   	let prefillSec = summary.phaseTimingsUs.prefill / 1_000_000.0
   	continuation.yield(.info(GenerateCompletionInfo(
   	promptTokenCount: summary.promptTokenCount,
   	generationTokenCount: summary.generationTokens,
   	promptTime: prefillSec,
   	generationTime: summary.elapsedUs / 1_000_000.0 - prefillSec,
   	stopReason: summary.generationTokens >= tokenLimit ? .length : .stop
   	)))
   	}
   	}
   	continuation.finish()
   	}
   	}
   	return (genStream, nil)

   
   so when handleChatStreaming hits the text-level stop (L3216) and its defer releases the slot (L3067), DFlashRuntime's own onTermination → task.cancel() (L287) never fires. And even if it did, generateStreaming never checks Task.isCancelled inside its cycle loop (L467); only the yield closure does (L270), so cancellation would mute the output but not stop the compute.

3. With the slot free, the next request starts a second DFlash generation on the same model, and the two race.

The standard path is not affected: its stream cancels the producer on termination and the token loop checks Task.isCancelled.

Suggested fix

• Pass the same stop IDs the standard path uses into DFlashRuntime.generate(stopTokenIDs:), so DFlash ends at EOS like plain decoding.
• Give the wrapper stream an onTermination that cancels its inner task, and check Task.isCancelled at the top of the DFlash cycle loop, so an early stop or a client disconnect actually stops the work.

Either half alone would make the CI flake far rarer; both together close it. Happy to send a PR.