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 diff --git a/Package.resolved b/Package.resolved index 39c9452..c636ff4 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/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 7072f2d..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 ── @@ -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, (() 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` 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 diff --git a/mlx-swift-lm b/mlx-swift-lm index 50d35c1..460ff81 160000 --- a/mlx-swift-lm +++ b/mlx-swift-lm @@ -1 +1 @@ -Subproject commit 50d35c1b0b4232105ac15100d0713130407f71fe +Subproject commit 460ff8115f41d792cf6326880ca1d4a14defe34d