From e71cdf31155b3c82a91e8dd17cc3a24c97f598d5 Mon Sep 17 00:00:00 2001 From: Hamed Shams Date: Sat, 10 Oct 2026 13:41:43 +0300 Subject: [PATCH] fix: stop DFlash at end-of-sequence, and cancel it when the stream ends 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 #224. --- Sources/DFlash/DFlashRuntime.swift | 4 + Sources/SwiftLM/Server.swift | 116 ++++++++++++------ .../DFlashStreamBridgeTests.swift | 106 ++++++++++++++++ tests/test-dflash.sh | 20 +++ 4 files changed, 209 insertions(+), 37 deletions(-) create mode 100644 tests/SwiftLMTests/DFlashStreamBridgeTests.swift diff --git a/Sources/DFlash/DFlashRuntime.swift b/Sources/DFlash/DFlashRuntime.swift index 3a00c8f..08fd94f 100644 --- a/Sources/DFlash/DFlashRuntime.swift +++ b/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)) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index a4bb8da..9aeef3c 100644 --- a/Sources/SwiftLM/Server.swift +++ b/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, tokenizerEOS: Int?, extraEOSTokens: Set, + 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, + tokenLimit: Int, + nextText: @escaping @Sendable (Int) -> String? +) -> AsyncStream { + AsyncStream { 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 { 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) diff --git a/tests/SwiftLMTests/DFlashStreamBridgeTests.swift b/tests/SwiftLMTests/DFlashStreamBridgeTests.swift new file mode 100644 index 0000000..2f6daf0 --- /dev/null +++ b/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 + + init() { + let terminated = self.terminated + var counter = 0 + stream = AsyncStream { 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.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") + } +} diff --git a/tests/test-dflash.sh b/tests/test-dflash.sh index 92a4a6d..7d674e5 100755 --- a/tests/test-dflash.sh +++ b/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)"