fix: propagate throws from LanguageModel.newCache and adapt to mlx-swift-lm PR #62

.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

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<DFlashEvent>`, 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,

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)

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)
                             }
                         }
                     }

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 ──
@@ -1888,7 +1895,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 ──
@@ -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()
@@ -1973,7 +1992,7 @@ func handleChatCompletion(
 
     // ── Cache-aware generation (standard path) ──
     let (stream, onPrefillDone) = try await container.perform { context -> (AsyncStream<Generation>, (() 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
@@ -2478,6 +2497,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
@@ -2626,6 +2653,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
         }
@@ -2926,7 +2956,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()
@@ -2991,7 +3022,7 @@ func handleTextNonStreaming(
             if completionTokenCount % 8 == 0 {
                 try? await Task.sleep(for: .microseconds(50))
             }
-        case .toolCall, .info:
+        case .toolCall, .rejectedToolCall, .info:
             break
         }
     }
@@ -3516,27 +3547,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)
             }
         }
     }
@@ -3814,7 +3855,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(
@@ -3834,7 +3875,15 @@ public struct ALMUserInputProcessor: UserInputProcessor, @unchecked Sendable {
     }
 }
 
-public final class ALMModelFactory: ModelFactory, @unchecked Sendable {
+// `class X: ModelFactory` (the constrained `GenericModelFactory<ModelContext,
+// ModelContainer>` 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
@@ -3889,7 +3938,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,
@@ -3903,7 +3952,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