Merge pull request #96 from SharpAI/fix/swiftbuddy-model-loading-recovery-main

Harden SwiftBuddy model loading and align local server settings

Author: solderzzc

Sources/MLXInferenceCore/GenerationConfig.swift

@@ -12,6 +12,24 @@ public struct GenerationConfig: Sendable {
     public var seed: UInt64?
     public var enableThinking: Bool
 
+    // ── SwiftLM Engine Parameters ──────────────────────────────────────
+    /// Enable TurboQuant KV-cache compression (3-bit PolarQuant+QJL).
+    /// Compresses KV history > 8192 tokens to ~3.5 bits/token.
+    public var turboKV: Bool
+
+    /// Enable SSD expert streaming for MoE models.
+    public var streamExperts: Bool
+
+    /// Chunk size for prefill evaluation.
+    /// Lower values prevent GPU timeout on large models.
+    public var prefillSize: Int
+
+    /// KV-cache quantization bits (nil = no quantization, 4 or 8 typical).
+    public var kvBits: Int?
+
+    /// KV-cache quantization group size (default 64).
+    public var kvGroupSize: Int
+
     public init(
         maxTokens: Int = 2048,
         temperature: Float = 0.6,
@@ -20,7 +38,12 @@ public struct GenerationConfig: Sendable {
         minP: Float = 0.0,
         repetitionPenalty: Float = 1.05,
         seed: UInt64? = nil,
-        enableThinking: Bool = false
+        enableThinking: Bool = false,
+        turboKV: Bool = false,
+        streamExperts: Bool = false,
+        prefillSize: Int = 512,
+        kvBits: Int? = nil,
+        kvGroupSize: Int = 64
     ) {
         self.maxTokens = maxTokens
         self.temperature = temperature
@@ -30,6 +53,11 @@ public struct GenerationConfig: Sendable {
         self.repetitionPenalty = repetitionPenalty
         self.seed = seed
         self.enableThinking = enableThinking
+        self.turboKV = turboKV
+        self.streamExperts = streamExperts
+        self.prefillSize = prefillSize
+        self.kvBits = kvBits
+        self.kvGroupSize = kvGroupSize
     }
 
     public static let `default` = GenerationConfig()

Sources/MLXInferenceCore/InferenceEngine.swift

@@ -114,6 +114,10 @@ public final class InferenceEngine: ObservableObject {
     @Published public private(set) var activeContextTokens: Int = 0
     @Published public private(set) var maxContextWindow: Int = 0
 
+    /// Set when a corrupted/truncated model is detected during inference.
+    /// The UI should observe this and offer to delete & re-download.
+    @Published public var corruptedModelId: String? = nil
+
     /// Whether to automatically unload the model when the app backgrounds
     /// and reload it when returning to foreground.
     /// Defaults to true on iOS (prevents jetsam), false on macOS.
@@ -277,7 +281,44 @@ public final class InferenceEngine: ObservableObject {
             state = .error("Device is too hot. Let it cool before loading a model.")
             return
         }
+        corruptedModelId = nil
+
+        guard ModelStorage.verifyModelIntegrity(for: modelId) else {
+            await downloadThenLoad(modelId: modelId)
+            return
+        }
+
+        await loadVerifiedModel(modelId: modelId)
+    }
+
+    private func downloadThenLoad(modelId: String) async {
+        print("[InferenceEngine] Model \(modelId) is missing or incomplete. Starting download before load.")
+        releaseLoadedModelResources()
+        state = .downloading(progress: 0.0, speed: "Preparing...")
+
+        let task = downloadManager.startDownload(modelId: modelId)
+
+        do {
+            try await task.value
+            state = .downloading(progress: 1.0, speed: "Verifying...")
+
+            guard ModelStorage.verifyModelIntegrity(for: modelId) else {
+                markModelCorrupted(
+                    modelId: modelId,
+                    message: "Model files are incomplete after download. Choose a recovery option."
+                )
+                return
+            }
+
+            await loadVerifiedModel(modelId: modelId)
+        } catch is CancellationError {
+            state = .idle
+        } catch {
+            state = .error("Failed to download \(modelId): \(error.localizedDescription)")
+        }
+    }
 
+    private func loadVerifiedModel(modelId: String) async {
         state = .loading
         currentModelId = modelId
 
@@ -312,25 +353,29 @@ public final class InferenceEngine: ObservableObject {
                 downloader: downloader
             )
 
+            let speedTracker = DownloadSpeedTracker()
+
             if architecture.supportsVision {
                 container = try await VLMModelFactory.shared.loadContainer(
                     from: downloader,
                     using: TransformersTokenizerLoader(),
                     configuration: config
                 ) { [weak self] progress in
+                    speedTracker.record(totalBytes: progress.completedUnitCount)
+                    let smoothedSpeed = speedTracker.speedBytesPerSec
+
                     Task { @MainActor in
                         guard let self else { return }
                         let pct = progress.fractionCompleted
-                        let speedBytesPerSec = progress.userInfo[ProgressUserInfoKey("throughputKey")] as? Double
-                        let speedStr = speedBytesPerSec
+                        let speedStr = smoothedSpeed
                             .map { String(format: "%.1f MB/s", $0 / 1_000_000) } ?? ""
                         self.state = .downloading(progress: pct, speed: speedStr)
 
                         self.downloadManager.updateProgress(ModelDownloadProgress(
                             modelId: modelId,
                             fractionCompleted: pct,
                             currentFile: "",
-                            speedMBps: speedBytesPerSec.map { $0 / 1_000_000 }
+                            speedMBps: smoothedSpeed.map { $0 / 1_000_000 }
                         ))
                     }
                 }
@@ -340,19 +385,21 @@ public final class InferenceEngine: ObservableObject {
                     using: TransformersTokenizerLoader(),
                     configuration: config
                 ) { [weak self] progress in
+                    speedTracker.record(totalBytes: progress.completedUnitCount)
+                    let smoothedSpeed = speedTracker.speedBytesPerSec
+
                     Task { @MainActor in
                         guard let self else { return }
                         let pct = progress.fractionCompleted
-                        let speedBytesPerSec = progress.userInfo[ProgressUserInfoKey("throughputKey")] as? Double
-                        let speedStr = speedBytesPerSec
+                        let speedStr = smoothedSpeed
                             .map { String(format: "%.1f MB/s", $0 / 1_000_000) } ?? ""
                         self.state = .downloading(progress: pct, speed: speedStr)
 
                         self.downloadManager.updateProgress(ModelDownloadProgress(
                             modelId: modelId,
                             fractionCompleted: pct,
                             currentFile: "",
-                            speedMBps: speedBytesPerSec.map { $0 / 1_000_000 }
+                            speedMBps: smoothedSpeed.map { $0 / 1_000_000 }
                         ))
                     }
                 }
@@ -361,26 +408,85 @@ public final class InferenceEngine: ObservableObject {
             downloadManager.clearProgress(modelId: modelId)
             downloadManager.lastLoadedModelId = modelId
             downloadManager.refresh()
+
+            // Verify integrity to catch incomplete downloads before marking as ready
+            guard ModelStorage.verifyModelIntegrity(for: modelId) else {
+                throw NSError(domain: "InferenceEngine", code: 1, userInfo: [NSLocalizedDescriptionKey: "Model safetensors files are incomplete. Please delete and re-download."])
+            }
+
+            // Read the model's actual max context length from config.json
+            if let ctxLen = ModelStorage.readMaxContextLength(for: modelId) {
+                self.maxContextWindow = ctxLen
+                print("[InferenceEngine] Model context window: \(ctxLen) tokens")
+            } else {
+                self.maxContextWindow = 8192  // conservative fallback for models without explicit limits
+                print("[InferenceEngine] No explicit context limit found in config.json, defaulting to 8192")
+            }
+
             state = .ready(modelId: modelId)
 
         } catch {
             ExpertStreamingConfig.shared.deactivate()
             downloadManager.clearProgress(modelId: modelId)
             state = .error("Failed to load \(modelId): \(error.localizedDescription)")
+
+            // If the model is incomplete/corrupted, flag it so the UI shows the "Delete & Re-download" button
+            let nsError = error as NSError
+            if nsError.domain == "InferenceEngine" && nsError.code == 1 || Self.isModelCorruptionError(error) {
+                markModelCorrupted(
+                    modelId: modelId,
+                    message: "Model weights are corrupted or incomplete. Choose a recovery option."
+                )
+                return
+            }
+
             container = nil
+            self.maxContextWindow = 0
+            self.activeContextTokens = 0
         }
     }
 
     /// Unload the current model and free all GPU memory.
     public func unload() {
+        releaseLoadedModelResources()
+        corruptedModelId = nil
+        state = .idle
+    }
+
+    private func releaseLoadedModelResources() {
         generationTask?.cancel()
+        generationTask = nil
         container = nil
         currentModelId = nil
-        state = .idle
+        maxContextWindow = 0
+        activeContextTokens = 0
         ExpertStreamingConfig.shared.deactivate()
         MLX.Memory.cacheLimit = 0
     }
 
+    private func markModelCorrupted(modelId: String?, message: String) {
+        let failedModelId = modelId ?? currentModelId
+        releaseLoadedModelResources()
+        state = .error(message)
+        corruptedModelId = failedModelId
+    }
+
+    private static func isModelCorruptionError(_ error: Error) -> Bool {
+        let description = error.localizedDescription.lowercased()
+        return description.contains("ssd streaming")
+            || description.contains("pread")
+            || description.contains("safetensors")
+            || description.contains("corrupt")
+            || description.contains("incomplete")
+    }
+
+    public func clearCorruptionRecovery() {
+        corruptedModelId = nil
+        if case .error = state {
+            state = .idle
+        }
+    }
+
     // MARK: — Generation
 
     public nonisolated func generate(
@@ -422,11 +528,17 @@ public final class InferenceEngine: ObservableObject {
                     }
 
                     let mlxMessages = finalMessages
-                    var params = GenerateParameters(temperature: config.temperature)
-                    params.topP = config.topP
-                    params.topK = config.topK
-                    params.minP = config.minP
-                    params.repetitionPenalty = config.repetitionPenalty
+                    var params = GenerateParameters(
+                        maxTokens: config.maxTokens,
+                        kvBits: config.kvBits,
+                        kvGroupSize: config.kvGroupSize,
+                        temperature: config.temperature,
+                        topP: config.topP,
+                        topK: config.topK,
+                        minP: config.minP,
+                        repetitionPenalty: config.repetitionPenalty,
+                        prefillStepSize: config.prefillSize
+                    )
                     params.repetitionContextSize = 20
 
                     var thinkingActive = false
@@ -442,9 +554,7 @@ public final class InferenceEngine: ObservableObject {
                     let baseTokens = Int(Double(stringLength) / 3.5)
                     self.activeContextTokens = baseTokens
 
-                    // If we have a max length config, expose it
-                    // TODO: Safely extract from ModelConfiguration when MLX exposes it dynamically
-                    self.maxContextWindow = 8192
+                    // maxContextWindow is already set during loadModel() from config.json
 
                     let stream: AsyncStream<Generation> = try await container.generate(
                         input: lmInput,
@@ -485,11 +595,30 @@ public final class InferenceEngine: ObservableObject {
                             continuation.yield(GenerationToken(text: text, isThinking: thinkingActive))
                         }
                     }
+                } catch let ssdError as SSDStreamingError {
+                    // Corrupted/truncated safetensors — surface a clear, actionable error
+                    let msg = "Model weights are corrupted or incomplete. Please re-download the model."
+                    print("[InferenceEngine] SSD Streaming Error: \(ssdError.localizedDescription)")
+                    continuation.yield(GenerationToken(text: "\n\n[Error: \(msg)]"))
+                    self.markModelCorrupted(modelId: self.currentModelId, message: msg)
                 } catch {
+                    // Check if the generic error is also an SSD streaming issue
+                    if Self.isModelCorruptionError(error) {
+                        let msg = "Model weights are corrupted or incomplete. Please re-download the model."
+                        self.markModelCorrupted(modelId: self.currentModelId, message: msg)
+                    }
                     continuation.yield(GenerationToken(text: "\n\n[Error: \(error.localizedDescription)]"))
                 }
 
-                self.state = self.currentModelId.map { .ready(modelId: $0) } ?? .idle
+                if let latchedError = SSDStreamingErrorLatch.shared.consume() {
+                    let msg = "Model weights are corrupted or incomplete. Please re-download the model."
+                    print("[InferenceEngine] Latched SSD error after generation: \(latchedError.localizedDescription)")
+                    self.markModelCorrupted(modelId: self.currentModelId, message: msg)
+                } else if case .error = self.state {
+                    // Already in error state from catch block above
+                } else {
+                    self.state = self.currentModelId.map { .ready(modelId: $0) } ?? .idle
+                }
                 continuation.finish()
             }
         }
@@ -500,4 +629,29 @@ public final class InferenceEngine: ObservableObject {
         generationTask = nil
         if let id = currentModelId { state = .ready(modelId: id) }
     }
+
+    /// Delete corrupted model files and start a fresh download.
+    /// Called from the UI when the user confirms re-download after corruption is detected.
+    public func deleteCorruptedAndRedownload() {
+        guard let modelId = corruptedModelId else { return }
+
+        releaseLoadedModelResources()
+        state = .downloading(progress: 0.0, speed: "Deleting corrupted files...")
+
+        do {
+            try ModelStorage.delete(modelId)
+            print("[InferenceEngine] Successfully deleted corrupted cache directory for \(modelId).")
+        } catch {
+            print("[InferenceEngine] FAILED to delete corrupted cache: \(error.localizedDescription)")
+            state = .error("Failed to delete corrupted model: \(error.localizedDescription)")
+            return
+        }
+        downloadManager.refresh()
+        corruptedModelId = nil
+
+        print("[InferenceEngine] Deleted corrupted files for \(modelId), starting fresh download")
+        Task { @MainActor in
+            await downloadThenLoad(modelId: modelId)
+        }
+    }
 }