From dece066498a7d36bb7aceb2135f285a1fe1b9ab5 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Mon, 14 Sep 2026 15:19:42 -0700 Subject: [PATCH 1/4] fix: propagate throws from LanguageModel.newCache and adapt to mlx-swift-lm PR #62 API changes Bump the mlx-swift-lm submodule to 348ff97 (upstream-sync PR #62) and update all call sites broken by its breaking API changes: - LanguageModel.newCache(parameters:) is now `throws`. Propagated `try` through every caller: DFlashRuntime's makeTargetCache/generateStreaming/ generateSync chain (with a documented do/catch in generate() since a plain AsyncStream cannot re-throw to its consumer), Gemma4MTPBench, InferenceEngine, and Server.swift. - GenerateCompletionInfo.totalDraftTokens was renamed/optional-ized to proposedDraftTokens/acceptedDraftTokens; updated InferenceEngine to unwrap safely instead of assuming non-optional Int. - ModelFactory is now a constrained typealias (GenericModelFactory< ModelContext, ModelContainer>) with primary associated types, which a class can no longer inherit from directly. ALMModelFactory and OmniModelFactory now conform to GenericModelFactory instead, matching upstream's own LLMModelFactory/VLMModelFactory pattern. - UserInput.audio was renamed to audios; updated the three genuine call sites (left LMInput.audio/vlmInput.audio alone, which are unrelated, unrenamed types). - Chat.Message's role cases were restructured (audios: label, ToolCall- based tool calls instead of raw dictionaries, tool(_, id:) instead of tool(_, toolCallId:)); rewrote toChatMessage() to build ToolCall values and decode JSON-string tool arguments into [String: JSONValue]. - Generation gained a new .rejectedToolCall(RejectedToolCall) case; added handling to all four previously-exhaustive switches, logging only reason/toolName/detail (never rawTextPreview, per its doc comment's privacy note). Verified both ways: `swift build -c release` and `swift build --build-tests` are clean (no errors) against the new pin (348ff97), and also clean when the submodule is temporarily rolled back to the previous pin (0e0cb47) with these same source changes in place (the added `try` keywords are harmless there). Submodule is left at 348ff97. TurboQuant C++ tests (9/9), SwiftLMTests (161/161), and SwiftBuddyTests (127 executed, 9 skipped) all pass against the final state. Co-Authored-By: Claude Sonnet 5 --- Sources/DFlash/DFlashRuntime.swift | 57 +++++++++----- Sources/Gemma4MTPBench/main.swift | 4 +- .../MLXInferenceCore/InferenceEngine.swift | 11 ++- Sources/SwiftLM/Server.swift | 78 +++++++++++++------ mlx-swift-lm | 2 +- 5 files changed, 102 insertions(+), 50 deletions(-) diff --git a/Sources/DFlash/DFlashRuntime.swift b/Sources/DFlash/DFlashRuntime.swift index 9c4439a..3a00c8f 100644 --- a/Sources/DFlash/DFlashRuntime.swift +++ b/Sources/DFlash/DFlashRuntime.swift @@ -149,8 +149,11 @@ public enum DFlashRuntime { /// - dflashUseTapeRollback=false → MambaSnapshotCache (snapshot-only, O(1) overhead) public static func makeTargetCache( targetModel: any DFlashTargetModel - ) -> [KVCache] { - var cache = targetModel.newCache(parameters: nil) + ) throws -> [KVCache] { + // `LanguageModel.newCache(parameters:)` is `throws` as of mlx-swift-lm + // commit 348ff97; propagate rather than swallow since a model's cache + // construction can legitimately fail (e.g. unsupported cache config). + var cache = try targetModel.newCache(parameters: nil) if targetModel.dflashIsHybridGDN { for i in 0 ..< cache.count { if cache[i] is MambaCache { @@ -252,21 +255,33 @@ public enum DFlashRuntime { // via a Continuation, avoiding the buffered-array bottleneck. AsyncStream(bufferingPolicy: .unbounded) { continuation in let task = Task { - generateStreaming( - targetModel: targetModel, - draftModel: draftModel, - promptTokens: promptTokens, - maxNewTokens: maxNewTokens, - blockTokens: blockTokens, - stopTokenIDs: stopTokenIDs, - suppressTokenIDs: suppressTokenIDs, - draftSinkSize: draftSinkSize, - draftWindowSize: draftWindowSize, - yield: { event in - guard !Task.isCancelled else { return } - continuation.yield(event) - } - ) + do { + try generateStreaming( + targetModel: targetModel, + draftModel: draftModel, + promptTokens: promptTokens, + maxNewTokens: maxNewTokens, + blockTokens: blockTokens, + stopTokenIDs: stopTokenIDs, + suppressTokenIDs: suppressTokenIDs, + draftSinkSize: draftSinkSize, + draftWindowSize: draftWindowSize, + yield: { event in + guard !Task.isCancelled else { return } + continuation.yield(event) + } + ) + } catch { + // `generate()` returns a plain `AsyncStream`, not + // an `AsyncThrowingStream`, so a failure here (e.g. cache + // construction failing inside `makeTargetCache`) cannot be + // re-thrown to the consumer. Log it and end the stream early; + // this mirrors the pre-existing behavior of any other early + // return from this loop (the consumer just sees no more + // events, exactly as if generation stopped normally). + FileHandle.standardError.write( + Data("[DFlashRuntime] generate() aborted: \(error)\n".utf8)) + } continuation.finish() } continuation.onTermination = { _ in task.cancel() } @@ -285,9 +300,9 @@ public enum DFlashRuntime { suppressTokenIDs: [Int]? = nil, draftSinkSize: Int = 64, draftWindowSize: Int = 1024 - ) -> [DFlashEvent] { + ) throws -> [DFlashEvent] { var events: [DFlashEvent] = [] - generateStreaming( + try generateStreaming( targetModel: targetModel, draftModel: draftModel, promptTokens: promptTokens, @@ -316,7 +331,7 @@ public enum DFlashRuntime { draftSinkSize: Int, draftWindowSize: Int, yield: (DFlashEvent) -> Void - ) { + ) throws { let promptLen = promptTokens.count guard promptLen > 0 && maxNewTokens > 0 else { return } @@ -329,7 +344,7 @@ public enum DFlashRuntime { let draftBackend = DFlashDraftBackend() - let targetCache = makeTargetCache(targetModel: targetModel) + let targetCache = try makeTargetCache(targetModel: targetModel) let draftCache = draftBackend.makeCache( draftModel: draftModel, diff --git a/Sources/Gemma4MTPBench/main.swift b/Sources/Gemma4MTPBench/main.swift index 714cfd4..84da3a4 100644 --- a/Sources/Gemma4MTPBench/main.swift +++ b/Sources/Gemma4MTPBench/main.swift @@ -147,7 +147,7 @@ struct Gemma4MTPBench: AsyncParsableCommand { let t0 = Date() var it = try TokenIterator( input: input, model: mainCtx.model, - cache: mainCtx.model.newCache(parameters: params), + cache: try mainCtx.model.newCache(parameters: params), parameters: params) while let tok = it.next() { baseOut.append(tok) @@ -179,7 +179,7 @@ struct Gemma4MTPBench: AsyncParsableCommand { let mtpT0 = Date() var mtpIt = try MTPTokenIterator( input: input, model: asstModel, - cache: mainCtx.model.newCache(parameters: params), + cache: try mainCtx.model.newCache(parameters: params), parameters: params, numMTPTokens: numDraft) while let tok = mtpIt.next() { mtpOut.append(tok) diff --git a/Sources/MLXInferenceCore/InferenceEngine.swift b/Sources/MLXInferenceCore/InferenceEngine.swift index fc137b1..8157a8f 100644 --- a/Sources/MLXInferenceCore/InferenceEngine.swift +++ b/Sources/MLXInferenceCore/InferenceEngine.swift @@ -749,7 +749,7 @@ extension InferenceEngine { // TurboKV: enable 3-bit PolarQuant+QJL on every KVCacheSimple cache layer. // KVCacheSimple is a cache object (not a neural-network Module), so we // iterate the cache array — mirroring the pattern in Server.swift. - let cache = await container.perform { ctx in ctx.model.newCache(parameters: params) } + let cache = try await container.perform { ctx in try ctx.model.newCache(parameters: params) } if config.turboKV { for layer in cache { if let simple = layer as? KVCacheSimple { @@ -821,8 +821,13 @@ extension InferenceEngine { continuation.yield(GenerationToken(text: text, isThinking: thinkingActive)) } else if case .info(let info) = generation { - if info.totalDraftTokens > 0 { - mtpAcceptanceRate = Double(info.acceptedDraftTokens) / Double(info.totalDraftTokens) + // `proposedDraftTokens`/`acceptedDraftTokens` are `Int?` as of + // mlx-swift-lm 348ff97 (nil for non-MTP iterators, renamed from + // the previously non-optional `totalDraftTokens`/`acceptedDraftTokens`). + if let proposed = info.proposedDraftTokens, proposed > 0, + let accepted = info.acceptedDraftTokens + { + mtpAcceptanceRate = Double(accepted) / Double(proposed) } } } diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 1f41f43..69f3b19 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -1821,7 +1821,7 @@ func handleChatCompletion( // true, it evaluates to false and still breaks. We MUST explicitly pass the boolean. let templateContext: [String: any Sendable] = ["enable_thinking": enableThinking] let userInput = UserInput(chat: chatMessages, tools: toolSpecs, additionalContext: templateContext) - print("[Server Debug] Created UserInput with \(userInput.images.count) images and \(userInput.audio.count) audio inputs.") + print("[Server Debug] Created UserInput with \(userInput.images.count) images and \(userInput.audios.count) audio inputs.") let lmInput = try await container.prepare(input: userInput) // ── Prompt caching: full token sequence for prefix matching ── @@ -1914,7 +1914,7 @@ func handleChatCompletion( // ── Cache-aware generation (standard path) ── let (stream, onPrefillDone) = try await container.perform { context -> (AsyncStream, (() async -> Void)?) in - let cache = context.model.newCache(parameters: params) + let cache = try context.model.newCache(parameters: params) // ── TurboQuant: enable 3-bit KV compression on every KVCacheSimple layer ── // This compresses cache history older than 8192 tokens into 3.5-bit Polar+QJL @@ -2395,6 +2395,14 @@ func handleChatStreaming( cont.yield(sseToolCallChunk(modelId: modelId, index: toolCallIndex, name: tc.function.name, arguments: argsJson)) toolCallIndex += 1 + case .rejectedToolCall(let rejection): + // `.rejectedToolCall` is new as of mlx-swift-lm 348ff97: a tool-call-shaped + // model output that failed parsing/authorization. There's no OpenAI wire + // shape for it, so it isn't forwarded to the client — just logged. Per + // `RejectedToolCall`'s doc comment, never log `rawTextPreview`; it may + // contain sensitive argument text. + print("[SwiftLM] Rejected tool call: reason=\(rejection.reason) tool=\(rejection.toolName ?? "?") detail=\(rejection.detail ?? "n/a")") + case .info(let info): heartbeatTask?.cancel() heartbeatTask = nil @@ -2540,6 +2548,9 @@ func handleChatNonStreaming( function: ToolCallFunction(name: tc.function.name, arguments: argsJson) )) tcIndex += 1 + case .rejectedToolCall(let rejection): + // See the matching comment in `handleChatStreaming`: log only, no wire shape. + print("[SwiftLM] Rejected tool call: reason=\(rejection.reason) tool=\(rejection.toolName ?? "?") detail=\(rejection.detail ?? "n/a")") case .info(let info): generationStopReason = info.stopReason } @@ -2823,7 +2834,8 @@ func handleTextStreaming( cont.yield(sseTextChunk(modelId: modelId, text: releasable, finishReason: nil)) } } - case .toolCall: + case .toolCall, .rejectedToolCall: + // Text-completion endpoint: tool calling has no wire representation here. break case .info(let info): heartbeatTask?.cancel() @@ -2888,7 +2900,7 @@ func handleTextNonStreaming( if completionTokenCount % 8 == 0 { try? await Task.sleep(for: .microseconds(50)) } - case .toolCall, .info: + case .toolCall, .rejectedToolCall, .info: break } } @@ -3353,27 +3365,37 @@ struct ChatCompletionRequest: Decodable { switch role { case "system", "developer": - return .system(text, images: imgs, audio: aud) + return .system(text, images: imgs, audios: aud) case "assistant": - var formattedToolCalls: [[String: any Sendable]]? = nil + // `Chat.Message.assistant(...)` takes `[ToolCall]?` as of + // mlx-swift-lm 348ff97 (previously a raw `[[String: any Sendable]]?` + // dictionary array), and no longer accepts `audios:` — assistant + // messages don't carry input audio. + var formattedToolCalls: [ToolCall]? = nil if let tc = tool_calls, !tc.isEmpty { - formattedToolCalls = tc.enumerated().map { (index, call) in - [ - "index": index, - "id": call.id, - "type": call.type, - "function": [ - "name": call.function.name, - "arguments": call.function.arguments - ] as [String: any Sendable] - ] as [String: any Sendable] + formattedToolCalls = tc.map { call in + // `call.function.arguments` is the raw JSON-string form (as sent + // by an OpenAI-style client); `ToolCall.Function` wants it decoded + // into `[String: JSONValue]`. Fall back to an empty dict if it + // isn't valid JSON rather than dropping the whole tool call. + let argsDict: [String: JSONValue] + if let data = call.function.arguments.data(using: .utf8), + let decoded = try? JSONDecoder().decode([String: JSONValue].self, from: data) + { + argsDict = decoded + } else { + argsDict = [:] + } + return ToolCall( + function: .init(name: call.function.name, arguments: argsDict), + id: call.id) } } - return .assistant(text, images: imgs, audio: aud, toolCalls: formattedToolCalls) + return .assistant(text, images: imgs, toolCalls: formattedToolCalls) case "tool": - return .tool(text, toolCallId: tool_call_id) + return .tool(text, id: tool_call_id) default: - return .user(text, images: imgs, audio: aud) + return .user(text, images: imgs, audios: aud) } } } @@ -3651,7 +3673,7 @@ public struct ALMUserInputProcessor: UserInputProcessor, @unchecked Sendable { messages: messages, tools: input.tools, additionalContext: input.additionalContext) // Check if there is audio to interleave - if !input.audio.isEmpty { + if !input.audios.isEmpty { print("[ALM] Interleaving Audio Tokens into prompt.") // Mock num audio embeddings for now - typically derived from the model or audio lengths let rawSequence = fusionProcessor.interleave( @@ -3671,7 +3693,15 @@ public struct ALMUserInputProcessor: UserInputProcessor, @unchecked Sendable { } } -public final class ALMModelFactory: ModelFactory, @unchecked Sendable { +// `class X: ModelFactory` (the constrained `GenericModelFactory` typealias) is no longer a legal inheritance clause as of +// mlx-swift-lm 348ff97 — a class can't inherit from a protocol type that +// supplies primary associated-type arguments. Upstream's own factories +// (`LLMModelFactory`, `VLMModelFactory`) switched to conforming to the +// unconstrained `GenericModelFactory` protocol directly, letting `ContextType`/ +// `ContainerType` be inferred as `ModelContext`/`ModelContainer` from the +// `_load`/`_wrap` implementations below; do the same here. +public final class ALMModelFactory: GenericModelFactory, @unchecked Sendable { public static let shared = ALMModelFactory() public let typeRegistry: ModelTypeRegistry = LLMTypeRegistry.shared public let modelRegistry: AbstractModelRegistry = LLMRegistry.shared @@ -3726,7 +3756,7 @@ public struct OmniUserInputProcessor: UserInputProcessor, @unchecked Sendable { return vlmInput } - if !input.audio.isEmpty && !tokens.isEmpty { + if !input.audios.isEmpty && !tokens.isEmpty { print("[Omni] Interleaving Audio Tokens into VLM prompt structure.") let rawSequence = fusionProcessor.interleave( textTokens: tokens, @@ -3740,7 +3770,9 @@ public struct OmniUserInputProcessor: UserInputProcessor, @unchecked Sendable { } } -public final class OmniModelFactory: ModelFactory, @unchecked Sendable { +// See the comment on `ALMModelFactory` above: conform to the unconstrained +// `GenericModelFactory` protocol, not the constrained `ModelFactory` typealias. +public final class OmniModelFactory: GenericModelFactory, @unchecked Sendable { public static let shared = OmniModelFactory() public let typeRegistry: ModelTypeRegistry = VLMTypeRegistry.shared public let modelRegistry: AbstractModelRegistry = VLMRegistry.shared diff --git a/mlx-swift-lm b/mlx-swift-lm index 50d35c1..348ff97 160000 --- a/mlx-swift-lm +++ b/mlx-swift-lm @@ -1 +1 @@ -Subproject commit 50d35c1b0b4232105ac15100d0713130407f71fe +Subproject commit 348ff97b23e7464f44f3244341d51d2b066e5013 From 402fd336afd2b433232f9f385b9e6ee10429dfda Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sat, 19 Sep 2026 08:37:37 -0700 Subject: [PATCH 2/4] chore: bump mlx-swift-lm to merged PR #62 (460ff81) Point the mlx-swift-lm submodule at the actual merged PR #62 commit (460ff81) instead of the interim pin (348ff97) this branch was originally adapted against. 460ff81 is a strict content superset of 348ff97 (confirmed via diff --stat: 81 files changed, 10594 insertions, 673 deletions, including the same Gemma4Unified fixes and extensive additional test coverage), so no regression from repointing. Package.resolved's swift-syntax pin moved to 603.0.2 as a byproduct of re-resolving against the new submodule state. Verified: `swift build -c release` and `swift test --skip-build` both pass (288 tests, 9 skipped, 0 failures). Co-Authored-By: Claude Sonnet 5 --- Package.resolved | 4 ++-- mlx-swift-lm | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Package.resolved b/Package.resolved index e35107a..af80570 100644 --- a/Package.resolved +++ b/Package.resolved @@ -248,8 +248,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/swiftlang/swift-syntax.git", "state" : { - "revision" : "0687f71944021d616d34d922343dcef086855920", - "version" : "600.0.1" + "revision" : "79e4b74a295b6eb74a8b585e3a39d29e70c1dbd1", + "version" : "603.0.2" } }, { diff --git a/mlx-swift-lm b/mlx-swift-lm index 348ff97..460ff81 160000 --- a/mlx-swift-lm +++ b/mlx-swift-lm @@ -1 +1 @@ -Subproject commit 348ff97b23e7464f44f3244341d51d2b066e5013 +Subproject commit 460ff8115f41d792cf6326880ca1d4a14defe34d From 2cfe9d7403406a67637046b3ca4be84f12dc189f Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sat, 19 Sep 2026 09:07:11 -0700 Subject: [PATCH 3/4] fix(ci): select Xcode 26.3 before building, matching mlx-swift-lm's own CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mlx-swift-lm's Package.swift declares swift-tools-version 6.2 (already true at the previous pin, 348ff97). SwiftLM's CI never selected a newer Xcode, so it silently depended on the default macos-15 runner image happening to ship a toolchain new enough — until this bump surfaced it: `swift package resolve` failed with "package 'mlx-swift-lm' is using Swift tools version 6.2.0 but the installed version is 6.1.0". Added the same `xcode-select -s /Applications/Xcode_26.3.app` step mlx-swift-lm's own ci.yml uses, to every job that resolves or builds the package (build_and_unit_test, speculative-decoding, dflash-speculative-decoding, speculative-decoding-eval, and ssd-draft-memory-guard's artifact-missing fallback build). Co-Authored-By: Claude Sonnet 5 --- .github/workflows/ci.yml | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4ee6a2f..c52799a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,6 +19,9 @@ jobs: with: submodules: recursive + - name: Select Xcode 26 (Swift 6.2) + run: sudo xcode-select -s /Applications/Xcode_26.3.app + - name: Install Metal Toolchain run: xcodebuild -downloadComponent MetalToolchain || true @@ -200,6 +203,9 @@ jobs: with: submodules: recursive + - name: Select Xcode 26 (Swift 6.2) + run: sudo xcode-select -s /Applications/Xcode_26.3.app + - name: Install Metal Toolchain run: xcodebuild -downloadComponent MetalToolchain || true @@ -300,6 +306,9 @@ jobs: with: submodules: recursive + - name: Select Xcode 26 (Swift 6.2) + run: sudo xcode-select -s /Applications/Xcode_26.3.app + - name: Install Metal Toolchain run: xcodebuild -downloadComponent MetalToolchain || true @@ -398,7 +407,10 @@ jobs: - uses: actions/checkout@v4 with: submodules: recursive - + + - name: Select Xcode 26 (Swift 6.2) + run: sudo xcode-select -s /Applications/Xcode_26.3.app + - name: Install Metal Toolchain run: xcodebuild -downloadComponent MetalToolchain || true @@ -557,6 +569,10 @@ jobs: name: swiftlm-architecture path: .build/release/ + - name: Select Xcode 26 (Swift 6.2) + if: hashFiles('.build/release/SwiftLM') == '' + run: sudo xcode-select -s /Applications/Xcode_26.3.app + - name: Build (Release) if artifact missing run: | if [ ! -f ".build/release/SwiftLM" ]; then From 572acecc1fe9b5054c1662c466cc9c4ca8b4fba9 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Thu, 24 Sep 2026 07:30:20 -0700 Subject: [PATCH 4/4] fix: emit DFlash .info and clear MLX cache after model load DFlash's .summary never yielded .info, so streams ending on EOS or max_tokens closed without finish_reason, usage or [DONE] (test-dflash Test 3). mlx-swift-lm's concurrent loader materializes every checkpoint tensor before sanitize; dropped tensors land in MLX's buffer cache, which under --stream-experts is capped only by the SSD budget. Clear it once loading finishes (ssd-draft-memory-guard was ~0.5 GB over on the new submodule). Co-Authored-By: Claude Opus 5.5 --- Sources/SwiftLM/Server.swift | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 4ba6685..da3e190 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -1189,6 +1189,13 @@ struct MLXServer: AsyncParsableCommand { // Same lenient pass as above: no template, or context the probe lacks. } + // mlx-swift-lm's concurrent loader materializes every tensor in the + // checkpoint before `sanitize` runs, so tensors the model drops (e.g. a + // vision tower when loading text-only) are freed into MLX's buffer cache. + // With --stream-experts the cache limit is the SSD budget, so those + // buffers would otherwise stay resident for the life of the server. + Memory.clearCache() + print("[SwiftLM] Model loaded. Starting HTTP server on \(host):\(port)") // ── Capture CLI defaults into a shared config ── @@ -1962,6 +1969,18 @@ func handleChatCompletion( 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()