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

Sources/DFlash/DFlashRuntime.swift

@@ -465,6 +465,10 @@ public enum DFlashRuntime {
         var prefetchedBlockLen: Int?
 
         while generatedTokenIDs.count < maxNewTokens {
+            // The consumer may be gone (stop sequence hit, client disconnected): the
+            // wrapper task is cancelled, and this loop must stop computing, not just
+            // stop yielding (#224).
+            if Task.isCancelled { break }
             let remaining = maxNewTokens - generatedTokenIDs.count
             let blockLen = max(1, min(effectiveBlockTokens, remaining))
 

Sources/SwiftLM/Server.swift

@@ -2321,6 +2321,73 @@ func collectBody(_ request: Request) async throws -> Data {
     return Data(bodyBytes)
 }
 
+/// The end-of-sequence set the standard path stops on (`buildStopTokenIds` in
+/// mlx-swift-lm): the configuration's ids, the tokenizer's EOS and the extra EOS tokens
+/// resolved through the tokenizer. DFlash takes it as a list.
+func stopTokenIDs(
+    eosTokenIDs: Set<Int>, tokenizerEOS: Int?, extraEOSTokens: Set<String>,
+    tokenID: (String) -> Int?
+) -> [Int] {
+    var ids = eosTokenIDs
+    if let tokenizerEOS { ids.insert(tokenizerEOS) }
+    for token in extraEOSTokens {
+        if let id = tokenID(token) { ids.insert(id) }
+    }
+    return ids.sorted()
+}
+
+/// Mutable detokenizer shared with the bridge task in `dflashGenerationStream`.
+private final class DetokenizerBox: @unchecked Sendable {
+    var detokenizer: NaiveStreamingDetokenizer
+    init(_ d: NaiveStreamingDetokenizer) { self.detokenizer = d }
+}
+
+/// Bridges a DFlash event stream to the `Generation` stream the handlers consume.
+///
+/// When the consumer stops early — a text-level stop sequence, a client disconnect — the
+/// returned stream is torn down, which cancels the bridge task; that ends the DFlash
+/// stream and, through its own `onTermination`, the DFlash generation task. Before this
+/// the bridge was a bare `Task` nobody cancelled, so the DFlash loop ran on to
+/// `max_tokens` after the response had ended and the next request's generation
+/// overlapped it (#224).
+func dflashGenerationStream(
+    _ events: AsyncStream<DFlashEvent>,
+    tokenLimit: Int,
+    nextText: @escaping @Sendable (Int) -> String?
+) -> AsyncStream<Generation> {
+    AsyncStream<Generation> { continuation in
+        let bridge = Task {
+            for await event in events {
+                if Task.isCancelled { break }
+                switch event {
+                case .token(let tokenID, _, _, _):
+                    if let chunk = nextText(tokenID) {
+                        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()
+        }
+        continuation.onTermination = { _ in bridge.cancel() }
+    }
+}
+
 // ── Chat Completions Handler ─────────────────────────────────────────────────
 
 func handleChatCompletion(
@@ -2522,50 +2589,25 @@ func handleChatCompletion(
         fflush(stdout)
         // Convert DFlashEvent stream to Generation stream with proper streaming detokenizer
         let dflashTokenizer = await container.tokenizer
+        let dflashConfiguration = await container.configuration
         let dflashStream = DFlashRuntime.generate(
             targetModel: targetModel,
             draftModel: dflashDraft,
             promptTokens: promptTokens,
             maxNewTokens: tokenLimit,
-            blockTokens: dflashBlockSize
+            blockTokens: dflashBlockSize,
+            // The end-of-sequence set the standard path stops on. Without it DFlash ran
+            // to max_tokens past `<|im_end|>` (#224).
+            stopTokenIDs: stopTokenIDs(
+                eosTokenIDs: dflashConfiguration.eosTokenIds,
+                tokenizerEOS: dflashTokenizer.eosTokenId,
+                extraEOSTokens: dflashConfiguration.extraEOSTokens,
+                tokenID: dflashTokenizer.convertTokenToId)
         )
-
-        // Use a class wrapper so the detokenizer can be mutated inside the closure
-        final class DetokenizerBox: @unchecked Sendable {
-            var detokenizer: NaiveStreamingDetokenizer
-            init(_ d: NaiveStreamingDetokenizer) { self.detokenizer = d }
-        }
         let box = DetokenizerBox(NaiveStreamingDetokenizer(tokenizer: dflashTokenizer))
-
-        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()
-            }
+        let genStream = dflashGenerationStream(dflashStream, tokenLimit: tokenLimit) { tokenID in
+            box.detokenizer.append(token: tokenID)
+            return box.detokenizer.next()
         }
 
         return (genStream, nil)

tests/SwiftLMTests/DFlashStreamBridgeTests.swift

@@ -0,0 +1,106 @@
+import XCTest
+import Foundation
+@testable import DFlash
+import MLXLMCommon
+@testable import SwiftLM
+
+// MARK: - The DFlash → Generation bridge (issue #224)
+//
+// The bridge used to be a bare `Task` that nobody cancelled. When the SSE handler ended
+// a response early (a text-level stop sequence, a client disconnect) the handler released
+// its slot, but the bridge kept draining the DFlash stream and the DFlash loop kept
+// generating up to `max_tokens`; the next request then ran a second generation on the
+// same model and MLX aborted. These pin the two halves of the fix that need no model:
+// the stop-token set handed to DFlash, and cancellation reaching the DFlash stream.
+final class DFlashStreamBridgeTests: XCTestCase {
+
+    // MARK: stopTokenIDs(…) — same union the standard path stops on
+
+    func testStopTokenIDsUnionsConfigurationTokenizerAndExtraTokens() {
+        let ids = stopTokenIDs(
+            eosTokenIDs: [151645, 151643],
+            tokenizerEOS: 151645,
+            extraEOSTokens: ["<|endoftext|>", "<|im_end|>", "<|not_in_vocab|>"],
+            tokenID: { ["<|endoftext|>": 151643, "<|im_end|>": 151645][$0] })
+        XCTAssertEqual(ids, [151643, 151645], "deduplicated; unknown extra tokens are dropped")
+    }
+
+    func testStopTokenIDsWithoutTokenizerEOS() {
+        XCTAssertEqual(stopTokenIDs(eosTokenIDs: [2], tokenizerEOS: nil, extraEOSTokens: [], tokenID: { _ in nil }), [2])
+        XCTAssertEqual(stopTokenIDs(eosTokenIDs: [], tokenizerEOS: nil, extraEOSTokens: [], tokenID: { _ in nil }), [])
+    }
+
+    // MARK: dflashGenerationStream(…) — tearing down the consumer cancels the producer
+
+    /// A DFlash-shaped stream that never ends on its own, and reports whether the
+    /// consumer side tore it down.
+    private final class EndlessEvents: @unchecked Sendable {
+        let terminated = XCTestExpectation(description: "DFlash stream terminated")
+        private(set) var yielded = 0
+        let stream: AsyncStream<DFlashEvent>
+
+        init() {
+            let terminated = self.terminated
+            var counter = 0
+            stream = AsyncStream<DFlashEvent> { continuation in
+                let producer = Task {
+                    while !Task.isCancelled {
+                        counter += 1
+                        continuation.yield(.token(tokenID: counter, generatedTokens: counter, acceptanceRatio: 0, cyclesCompleted: counter))
+                        try? await Task.sleep(for: .milliseconds(5))
+                    }
+                    continuation.finish()
+                }
+                continuation.onTermination = { _ in
+                    producer.cancel()
+                    terminated.fulfill()
+                }
+            }
+        }
+    }
+
+    func testStreamTornDownAfterConsumerStopsEarly() async throws {
+        let events = EndlessEvents()
+        var received: [String] = []
+        do {
+            // Scoped like the handler's `for await` over the generation stream: once the
+            // consumer breaks out and the stream value goes away, the bridge must be
+            // torn down and the DFlash stream with it.
+            let generation = dflashGenerationStream(events.stream, tokenLimit: 1_000) { "t\($0) " }
+            for await item in generation {
+                if case .chunk(let text, _) = item { received.append(text) }
+                if received.count == 3 { break }
+            }
+        }
+        XCTAssertEqual(received, ["t1 ", "t2 ", "t3 "])
+        await fulfillment(of: [events.terminated], timeout: 2)
+    }
+
+    func testBridgeMapsTokensAndSummary() async throws {
+        let summary = DFlashSummary(
+            elapsedUs: 2_000_000, promptTokenCount: 7, generatedTokenIDs: [1, 2, 3],
+            acceptedFromDraft: 1, acceptanceRatio: 0.33, blockTokens: 16, cyclesCompleted: 2,
+            phaseTimingsUs: .init(prefill: 500_000, draft: 0, verify: 0, replay: 0))
+        let (events, continuation) = AsyncStream<DFlashEvent>.makeStream()
+        continuation.yield(.prefill(promptTokenCount: 7, prefillUs: 500_000))
+        continuation.yield(.token(tokenID: 1, generatedTokens: 1, acceptanceRatio: 0, cyclesCompleted: 1))
+        continuation.yield(.token(tokenID: 2, generatedTokens: 2, acceptanceRatio: 0, cyclesCompleted: 1))
+        continuation.yield(.summary(summary))
+        continuation.finish()
+
+        var chunks: [String] = []
+        var info: GenerateCompletionInfo?
+        // The detokenizer seam may hold text back (nil) and release it on a later token.
+        for await item in dflashGenerationStream(events, tokenLimit: 3, nextText: { $0 == 1 ? nil : "one two" }) {
+            switch item {
+            case .chunk(let text, _): chunks.append(text)
+            case .info(let i): info = i
+            default: break
+            }
+        }
+        XCTAssertEqual(chunks, ["one two"])
+        XCTAssertEqual(info?.generationTokenCount, 3)
+        XCTAssertEqual(info?.promptTokenCount, 7)
+        XCTAssertEqual(info?.stopReason, .length, "3 generated against a limit of 3 is a length stop")
+    }
+}

tests/test-dflash.sh

@@ -189,6 +189,26 @@ else
     fail "Speculative decoding path not activated (missing log line)"
 fi
 
+# ── Test 4: A stream that ends early must not leave a generation running ──
+# Issue #224: the text-level stop ended the SSE stream and released the slot, but the
+# DFlash loop ran on to max_tokens; the next request then overlapped it and MLX aborted.
+# A one-word answer against a large budget leaves plenty of orphaned work for the next
+# request to collide with.
+log "Test 4: early-ended stream does not leave a generation running"
+EARLY=$(curl -sf -N --max-time 120 -X POST "$URL/v1/chat/completions" \
+    -H "Content-Type: application/json" \
+    -d "{\"model\":\"$MAIN_MODEL\",\"stream\":true,\"max_tokens\":200,\"messages\":[{\"role\":\"user\",\"content\":\"Reply with the single word: ok\"}]}" \
+    2>/dev/null || true)
+AFTER=$(curl -sf --max-time 120 -X POST "$URL/v1/chat/completions" \
+    -H "Content-Type: application/json" \
+    -d "{\"model\":\"$MAIN_MODEL\",\"max_tokens\":10,\"messages\":[{\"role\":\"user\",\"content\":\"Say the word: done\"}]}" 2>/dev/null || echo "")
+AFTER_CONTENT=$(echo "$AFTER" | jq -r '.choices[0].message.content // empty' 2>/dev/null || echo "")
+if echo "$EARLY" | grep -q "data: \[DONE\]" && [ -n "$AFTER_CONTENT" ] && curl -sf "$URL/health" >/dev/null 2>&1; then
+    pass "request after an early-ended stream completed and the server is still up"
+else
+    fail "request after an early-ended stream failed (server alive: $(curl -sf "$URL/health" >/dev/null 2>&1 && echo yes || echo no))"
+fi
+
 # ── Test 5: Multiple sequential requests (stability) ────────────────
 log "Test 5: Sequential request stability (3 requests)"