Temporary reasoning proxy just to test

.mlx_typings/mlx_lm/models/deepseek_v4.pyi

@@ -108,10 +108,12 @@ class Compressor(nn.Module):
     def __call__(
         self,
         x: mx.array,
-        cache: "DeepseekV4Cache",
+        state: ArraysCache,
         offset: Any,
-        key: str = ...,
-    ) -> mx.array: ...
+        slot_compressed: int,
+        slot_kv_state: int,
+        slot_score_state: int,
+    ) -> Optional[mx.array]: ...
 
 class Indexer(nn.Module):
     def __init__(
@@ -121,29 +123,14 @@ class Indexer(nn.Module):
         rope: DeepseekV4RoPE,
     ) -> None: ...
 
-class _CompressorBranch:
-    buffer_kv: Optional[mx.array]
-    buffer_gate: Optional[mx.array]
-    prev_kv: Optional[mx.array]
-    prev_gate: Optional[mx.array]
-    pool: Optional[mx.array]
-    buffer_lengths: Optional[List[int]]
-    pool_lengths: Optional[List[int]]
-    buffer_count: int
-    _new_pool_lengths: Optional[List[int]]
-
-    def __init__(self) -> None: ...
-
 class DeepseekV4Cache:
-    local: RotatingKVCache
+    local: Any
     offset: int
     keys: Optional[mx.array]
     values: Optional[mx.array]
     state: Any
     meta_state: Any
     nbytes: int
-    _branches: Dict[str, _CompressorBranch]
-    _pending_lengths: Optional[List[int]]
 
     def __init__(self, sliding_window: int) -> None: ...
     def update_and_fetch(

app/EXO/EXO/ContentView.swift

@@ -16,13 +16,22 @@ struct ContentView: View {
     @EnvironmentObject private var updater: SparkleUpdater
     @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
     @EnvironmentObject private var settingsWindowController: SettingsWindowController
-    @EnvironmentObject private var bugReportWindowController: BugReportWindowController
     @State private var focusedNode: NodeViewModel?
     @State private var deletingInstanceIDs: Set<String> = []
     @State private var showAllNodes = false
     @State private var showAllInstances = false
     @State private var baseURLCopied = false
     @State private var showAdvanced = false
+    @State private var showDebugInfo = false
+    private enum BugReportPhase: Equatable {
+        case idle
+        case prompting
+        case sending(String)
+        case success(String)
+        case failure(String)
+    }
+    @State private var bugReportPhase: BugReportPhase = .idle
+    @State private var bugReportUserDescription: String = ""
     @State private var uninstallInProgress = false
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
@@ -285,13 +294,6 @@ struct ContentView: View {
             ) {
                 updater.checkForUpdates()
             }
-            HoverButton(
-                title: "Share Bug Report…",
-                tint: .primary,
-                trailingSystemImage: "ladybug"
-            ) {
-                bugReportWindowController.open()
-            }
             .padding(.bottom, 8)
             HoverButton(title: "Quit", tint: .secondary) {
                 controller.stop()
@@ -475,6 +477,40 @@ struct ContentView: View {
         }
     }
 
+    private var debugSection: some View {
+        VStack(alignment: .leading, spacing: 4) {
+            HoverButton(
+                title: "Debug Info",
+                tint: .primary,
+                trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
+                small: true
+            ) {
+                showDebugInfo.toggle()
+            }
+            if showDebugInfo {
+                VStack(alignment: .leading, spacing: 4) {
+                    Text("Version: \(buildTag)")
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                    Text("Commit: \(buildCommit)")
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                    Text(thunderboltStatusText)
+                        .font(.caption2)
+                        .foregroundColor(thunderboltStatusColor)
+                    clusterThunderboltBridgeView
+                    interfaceIpList
+                    rdmaStatusView
+                    sendBugReportButton
+                        .padding(.top, 6)
+                }
+                .padding(.leading, 8)
+                .transition(.opacity)
+            }
+        }
+        .animation(.easeInOut(duration: 0.25), value: showDebugInfo)
+    }
+
     private var rdmaStatusView: some View {
         let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
         let localNodeId = stateService.localNodeId
@@ -523,6 +559,127 @@ struct ContentView: View {
         }
     }
 
+    private var sendBugReportButton: some View {
+        VStack(alignment: .leading, spacing: 6) {
+            switch bugReportPhase {
+            case .idle:
+                Button {
+                    bugReportPhase = .prompting
+                    bugReportUserDescription = ""
+                } label: {
+                    HStack {
+                        Text("Send Bug Report")
+                            .font(.caption)
+                            .fontWeight(.semibold)
+                        Spacer()
+                    }
+                    .padding(.vertical, 6)
+                    .padding(.horizontal, 8)
+                    .background(
+                        RoundedRectangle(cornerRadius: 6)
+                            .fill(Color.accentColor.opacity(0.12))
+                    )
+                }
+                .buttonStyle(.plain)
+
+            case .prompting:
+                VStack(alignment: .leading, spacing: 6) {
+                    VStack(alignment: .leading, spacing: 2) {
+                        Text("Tell us what went wrong (optional)")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        Text(
+                            "A quick description of what you were doing and what happened helps us track down the bug for you."
+                        )
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                        .opacity(0.8)
+                        .fixedSize(horizontal: false, vertical: true)
+                    }
+                    TextEditor(text: $bugReportUserDescription)
+                        .font(.caption2)
+                        .frame(height: 60)
+                        .overlay(
+                            RoundedRectangle(cornerRadius: 4)
+                                .stroke(Color.secondary.opacity(0.3), lineWidth: 1)
+                        )
+                    HStack(spacing: 8) {
+                        Button("Send") {
+                            Task {
+                                await sendBugReport()
+                            }
+                        }
+                        .font(.caption2)
+                        .buttonStyle(.borderedProminent)
+                        .controlSize(.small)
+                        Button("Cancel") {
+                            bugReportPhase = .idle
+                        }
+                        .font(.caption2)
+                        .buttonStyle(.bordered)
+                        .controlSize(.small)
+                    }
+                }
+                .padding(8)
+                .background(
+                    RoundedRectangle(cornerRadius: 6)
+                        .fill(Color.accentColor.opacity(0.06))
+                )
+
+            case .sending(let message):
+                HStack(spacing: 6) {
+                    ProgressView()
+                        .scaleEffect(0.6)
+                    Text(message)
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                }
+
+            case .success(let message):
+                VStack(alignment: .leading, spacing: 6) {
+                    Text(message)
+                        .font(.caption2)
+                        .foregroundColor(.secondary)
+                        .fixedSize(horizontal: false, vertical: true)
+                    Button {
+                        openGitHubIssue()
+                    } label: {
+                        HStack(spacing: 4) {
+                            Image(systemName: "arrow.up.right.square")
+                                .imageScale(.small)
+                            Text("Create GitHub Issue")
+                                .font(.caption2)
+                        }
+                    }
+                    .buttonStyle(.bordered)
+                    .controlSize(.small)
+                    Button("Done") {
+                        bugReportPhase = .idle
+                        bugReportUserDescription = ""
+                    }
+                    .font(.caption2)
+                    .buttonStyle(.plain)
+                    .foregroundColor(.secondary)
+                }
+
+            case .failure(let message):
+                VStack(alignment: .leading, spacing: 4) {
+                    Text(message)
+                        .font(.caption2)
+                        .foregroundColor(.red)
+                        .fixedSize(horizontal: false, vertical: true)
+                    Button("Dismiss") {
+                        bugReportPhase = .idle
+                    }
+                    .font(.caption2)
+                    .buttonStyle(.plain)
+                    .foregroundColor(.secondary)
+                }
+            }
+        }
+        .animation(.easeInOut(duration: 0.2), value: bugReportPhase)
+    }
+
     private var processToggleBinding: Binding<Bool> {
         Binding(
             get: {
@@ -563,6 +720,61 @@ struct ContentView: View {
         )
     }
 
+    private func sendBugReport() async {
+        bugReportPhase = .sending("Collecting logs...")
+        let service = BugReportService()
+        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
+        do {
+            let outcome = try await service.sendReport(
+                isManual: true,
+                userDescription: description.isEmpty ? nil : description
+            )
+            if outcome.success {
+                bugReportPhase = .success(outcome.message)
+            } else {
+                bugReportPhase = .failure(outcome.message)
+            }
+        } catch {
+            bugReportPhase = .failure(error.localizedDescription)
+        }
+    }
+
+    private func openGitHubIssue() {
+        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
+
+        var bodyParts: [String] = []
+        bodyParts.append("## Describe the bug")
+        bodyParts.append("")
+        if !description.isEmpty {
+            bodyParts.append(description)
+        } else {
+            bodyParts.append("A clear and concise description of what the bug is.")
+        }
+        bodyParts.append("")
+        bodyParts.append("## Environment")
+        bodyParts.append("")
+        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
+        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
+        bodyParts.append("")
+        bodyParts.append("## Additional context")
+        bodyParts.append("")
+        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
+
+        let body = bodyParts.joined(separator: "\n")
+
+        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
+        components.queryItems = [
+            URLQueryItem(name: "template", value: "bug_report.md"),
+            URLQueryItem(name: "title", value: "[BUG] "),
+            URLQueryItem(name: "body", value: body),
+            URLQueryItem(name: "labels", value: "bug"),
+        ]
+
+        if let url = components.url {
+            NSWorkspace.shared.open(url)
+        }
+    }
+
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"
@@ -645,6 +857,13 @@ struct ContentView: View {
         }
     }
 
+    private var buildTag: String {
+        Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
+    }
+
+    private var buildCommit: String {
+        Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
+    }
 }
 
 private struct HoverButton: View {

app/EXO/EXO/EXOApp.swift

@@ -22,7 +22,6 @@ struct EXOApp: App {
     @StateObject private var updater: SparkleUpdater
     @StateObject private var thunderboltBridgeService: ThunderboltBridgeService
     @StateObject private var settingsWindowController: SettingsWindowController
-    @StateObject private var bugReportWindowController: BugReportWindowController
     private let terminationObserver: TerminationObserver
     private let firstLaunchPopout = FirstLaunchPopout()
     private let ciContext = CIContext(options: nil)
@@ -47,7 +46,6 @@ struct EXOApp: App {
         let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
         _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         _settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
-        _bugReportWindowController = StateObject(wrappedValue: BugReportWindowController())
         enableLaunchAtLoginIfNeeded()
         // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
         NetworkSetupHelper.promptAndInstallIfNeeded()
@@ -68,7 +66,6 @@ struct EXOApp: App {
                 .environmentObject(updater)
                 .environmentObject(thunderboltBridgeService)
                 .environmentObject(settingsWindowController)
-                .environmentObject(bugReportWindowController)
         } label: {
             menuBarIcon
                 .onReceive(controller.$isFirstLaunchReady) { ready in

app/EXO/EXO/Services/ClusterStateService.swift

@@ -17,7 +17,7 @@ final class ClusterStateService: ObservableObject {
 
     init(
         baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
-        session: URLSession = ClusterStateService.makeNonCachingSession()
+        session: URLSession = .shared
     ) {
         self.baseURL = baseURL
         self.endpoint = baseURL.appendingPathComponent("state")
@@ -27,23 +27,6 @@ final class ClusterStateService: ObservableObject {
         self.decoder = decoder
     }
 
-    /// `URLSession.shared` carries an on-disk `URLCache` that persists every
-    /// response body under `~/Library/Caches/exolabs.EXO/`. We poll `/state`
-    /// at 2 Hz from `startPolling`, so leaving the shared cache attached
-    /// dirties ~500–620 KB/sec of file-backed memory and trips macOS's
-    /// per-process `disk writes` resource limit (microstackshot reports
-    /// observed on M3 Ultra producing GBs of cached responses per hour).
-    /// Cluster-state polling responses are time-sensitive and small; they
-    /// gain nothing from being cached on disk. Use an ephemeral session
-    /// with `urlCache = nil` so neither response bodies nor metadata
-    /// touch disk.
-    private static func makeNonCachingSession() -> URLSession {
-        let config = URLSessionConfiguration.ephemeral
-        config.urlCache = nil
-        config.requestCachePolicy = .reloadIgnoringLocalCacheData
-        return URLSession(configuration: config)
-    }
-
     func startPolling(interval: TimeInterval = 0.5) {
         stopPolling()
         Task {

app/EXO/EXO/Views/BugReportWindowController.swift

@@ -1,242 +0,0 @@
-import AppKit
-import SwiftUI
-
-/// Manages a standalone window for the bug-report flow.
-/// Ensures only one instance exists and brings it to front on repeated opens.
-@MainActor
-final class BugReportWindowController: ObservableObject {
-    private var window: NSWindow?
-
-    func open() {
-        if let existing = window, existing.isVisible {
-            existing.makeKeyAndOrderFront(nil)
-            NSApp.activate()
-            return
-        }
-
-        let view = BugReportView(onDismiss: { [weak self] in
-            self?.window?.close()
-        })
-
-        let hostingController = NSHostingController(rootView: view)
-        hostingController.sizingOptions = [.preferredContentSize, .minSize]
-
-        let newWindow = NSWindow(contentViewController: hostingController)
-        newWindow.styleMask = [.titled, .closable, .resizable]
-        newWindow.title = "Send a Bug Report"
-        newWindow.center()
-        newWindow.setFrameAutosaveName("ExoBugReportWindow")
-        newWindow.isReleasedWhenClosed = false
-        newWindow.makeKeyAndOrderFront(nil)
-        NSApp.activate()
-
-        window = newWindow
-    }
-}
-
-private struct BugReportView: View {
-    fileprivate enum Phase: Equatable {
-        case prompting
-        case sending(String)
-        case success(String)
-        case failure(String)
-    }
-
-    let onDismiss: () -> Void
-
-    @State private var phase: Phase = .prompting
-    @State private var userDescription: String = ""
-    @FocusState private var descriptionFocused: Bool
-
-    var body: some View {
-        VStack(alignment: .leading, spacing: 12) {
-            switch phase {
-            case .prompting:
-                promptingView
-            case .sending(let message):
-                sendingView(message: message)
-            case .success(let message):
-                successView(message: message)
-            case .failure(let message):
-                failureView(message: message)
-            }
-        }
-        .padding(16)
-        .frame(minWidth: 380)
-        .animation(.easeInOut(duration: 0.2), value: phase)
-        .onAppear { descriptionFocused = true }
-    }
-
-    private var promptingView: some View {
-        VStack(alignment: .leading, spacing: 8) {
-            Text("Description (optional)")
-                .font(.subheadline)
-                .foregroundColor(.secondary)
-            ZStack(alignment: .topLeading) {
-                if userDescription.isEmpty {
-                    Text("What were you doing when it broke?")
-                        .font(.body)
-                        .foregroundColor(Color(nsColor: .placeholderTextColor))
-                        .padding(.horizontal, 10)
-                        .padding(.vertical, 8)
-                        .allowsHitTesting(false)
-                }
-                TextEditor(text: $userDescription)
-                    .font(.body)
-                    .scrollContentBackground(.hidden)
-                    .padding(4)
-                    .frame(height: 72)
-                    .focused($descriptionFocused)
-            }
-            .background(
-                RoundedRectangle(cornerRadius: 6)
-                    .fill(Color(nsColor: .textBackgroundColor))
-            )
-            .overlay(
-                RoundedRectangle(cornerRadius: 6)
-                    .strokeBorder(Color(nsColor: .separatorColor), lineWidth: 1)
-            )
-
-            Text("Diagnostic logs will be uploaded with your report.")
-                .font(.caption)
-                .foregroundColor(.secondary)
-
-            HStack {
-                Spacer()
-                Button("Cancel") { onDismiss() }
-                    .keyboardShortcut(.cancelAction)
-                Button("Send") {
-                    Task { await send() }
-                }
-                .keyboardShortcut(.defaultAction)
-            }
-            .padding(.top, 4)
-        }
-    }
-
-    private func sendingView(message: String) -> some View {
-        VStack(alignment: .leading, spacing: 12) {
-            HStack(spacing: 10) {
-                ProgressView().controlSize(.small)
-                Text(message)
-                    .foregroundColor(.secondary)
-            }
-            HStack {
-                Spacer()
-                Button("Cancel") { onDismiss() }
-                    .keyboardShortcut(.cancelAction)
-                    .disabled(true)
-                Button("Send") {}
-                    .disabled(true)
-            }
-        }
-    }
-
-    private func successView(message: String) -> some View {
-        VStack(alignment: .leading, spacing: 12) {
-            HStack(alignment: .top, spacing: 10) {
-                Image(systemName: "checkmark.circle.fill")
-                    .foregroundColor(.green)
-                    .font(.title2)
-                Text(message)
-                    .fixedSize(horizontal: false, vertical: true)
-            }
-            HStack {
-                Button {
-                    openGitHubIssue()
-                } label: {
-                    HStack(spacing: 4) {
-                        Image(systemName: "arrow.up.right.square")
-                        Text("Open GitHub Issue")
-                    }
-                }
-                Spacer()
-                Button("Done") { onDismiss() }
-                    .keyboardShortcut(.defaultAction)
-            }
-        }
-    }
-
-    private func failureView(message: String) -> some View {
-        VStack(alignment: .leading, spacing: 12) {
-            HStack(alignment: .top, spacing: 10) {
-                Image(systemName: "exclamationmark.triangle.fill")
-                    .foregroundColor(.orange)
-                    .font(.title2)
-                Text(message)
-                    .fixedSize(horizontal: false, vertical: true)
-            }
-            HStack {
-                Spacer()
-                Button("Try Again") {
-                    phase = .prompting
-                }
-                Button("Close") { onDismiss() }
-                    .keyboardShortcut(.defaultAction)
-            }
-        }
-    }
-
-    private func send() async {
-        phase = .sending("Collecting logs and uploading…")
-        let service = BugReportService()
-        let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-        do {
-            let outcome = try await service.sendReport(
-                isManual: true,
-                userDescription: description.isEmpty ? nil : description
-            )
-            if outcome.success {
-                phase = .success(outcome.message)
-            } else {
-                phase = .failure(outcome.message)
-            }
-        } catch {
-            phase = .failure(error.localizedDescription)
-        }
-    }
-
-    private func openGitHubIssue() {
-        let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-
-        var bodyParts: [String] = []
-        bodyParts.append("## Describe the bug")
-        bodyParts.append("")
-        if !description.isEmpty {
-            bodyParts.append(description)
-        } else {
-            bodyParts.append("A clear and concise description of what the bug is.")
-        }
-        bodyParts.append("")
-        bodyParts.append("## Environment")
-        bodyParts.append("")
-        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
-        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
-        bodyParts.append("")
-        bodyParts.append("## Additional context")
-        bodyParts.append("")
-        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
-
-        let body = bodyParts.joined(separator: "\n")
-
-        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
-        components.queryItems = [
-            URLQueryItem(name: "template", value: "bug_report.md"),
-            URLQueryItem(name: "title", value: "[BUG] "),
-            URLQueryItem(name: "body", value: body),
-            URLQueryItem(name: "labels", value: "bug"),
-        ]
-
-        if let url = components.url {
-            NSWorkspace.shared.open(url)
-        }
-    }
-
-    private var buildTag: String {
-        Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
-    }
-
-    private var buildCommit: String {
-        Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
-    }
-}

app/EXO/EXO/Views/SettingsView.swift

@@ -21,6 +21,8 @@ struct SettingsView: View {
     @State private var pendingReadOnlyModelsDirs: String = ""
     @State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
     @State private var needsRestart = false
+    @State private var bugReportInFlight = false
+    @State private var bugReportMessage: String?
     @State private var uninstallInProgress = false
 
     var body: some View {
@@ -200,6 +202,8 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
+
+                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -500,30 +504,63 @@ struct SettingsView: View {
         }
     }
 
+    private var sendBugReportButton: some View {
+        VStack(alignment: .leading, spacing: 4) {
+            Button {
+                Task {
+                    await sendBugReport()
+                }
+            } label: {
+                HStack {
+                    if bugReportInFlight {
+                        ProgressView()
+                            .scaleEffect(0.6)
+                    }
+                    Text("Send Bug Report")
+                        .font(.caption)
+                        .fontWeight(.semibold)
+                    Spacer()
+                }
+            }
+            .disabled(bugReportInFlight)
+
+            if let message = bugReportMessage {
+                Text(message)
+                    .font(.caption2)
+                    .foregroundColor(.secondary)
+                    .fixedSize(horizontal: false, vertical: true)
+            }
+        }
+    }
+
     // MARK: - Actions
 
+    private func sendBugReport() async {
+        bugReportInFlight = true
+        bugReportMessage = "Collecting logs..."
+        let service = BugReportService()
+        do {
+            let outcome = try await service.sendReport(isManual: true)
+            bugReportMessage = outcome.message
+        } catch {
+            bugReportMessage = error.localizedDescription
+        }
+        bugReportInFlight = false
+    }
+
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"
         alert.informativeText = """
-            This will remove EXO and all its components:
+            This will remove EXO and all its system components:
 
             • Network configuration daemon
             • Launch at login registration
             • EXO network location
-            • EXO data directory (~/.exo)
 
             The app will be moved to Trash.
             """
         alert.alertStyle = .warning
-
-        let checkbox = NSButton(
-            checkboxWithTitle: "Keep downloaded models (~/.exo/models)",
-            target: nil, action: nil)
-        checkbox.state = .off
-        checkbox.sizeToFit()
-        alert.accessoryView = checkbox
-
         alert.addButton(withTitle: "Uninstall")
         alert.addButton(withTitle: "Cancel")
 
@@ -533,11 +570,11 @@ struct SettingsView: View {
 
         let response = alert.runModal()
         if response == .alertFirstButtonReturn {
-            performUninstall(keepModels: checkbox.state == .on)
+            performUninstall()
         }
     }
 
-    private func performUninstall(keepModels: Bool) {
+    private func performUninstall() {
         uninstallInProgress = true
 
         controller.cancelPendingLaunch()
@@ -547,7 +584,6 @@ struct SettingsView: View {
         DispatchQueue.global(qos: .utility).async {
             do {
                 try NetworkSetupHelper.uninstall()
-                try Self.removeExoDirectory(keepModels: keepModels)
 
                 DispatchQueue.main.async {
                     LaunchAtLoginHelper.disable()
@@ -571,23 +607,6 @@ struct SettingsView: View {
         }
     }
 
-    private static func removeExoDirectory(keepModels: Bool) throws {
-        let fm = FileManager.default
-        let exoDir = ExoProcessController.exoDirectoryURL
-        guard fm.fileExists(atPath: exoDir.path) else { return }
-
-        if !keepModels {
-            try fm.removeItem(at: exoDir)
-            return
-        }
-
-        let contents = try fm.contentsOfDirectory(
-            at: exoDir, includingPropertiesForKeys: nil, options: [])
-        for entry in contents where entry.lastPathComponent != "models" {
-            try? fm.removeItem(at: entry)
-        }
-    }
-
     private func moveAppToTrash() {
         guard let appURL = Bundle.main.bundleURL as URL? else { return }
         do {

app/EXO/uninstall-exo.sh

@@ -3,55 +3,25 @@
 # EXO Uninstaller Script
 #
 # This script removes all EXO system components that persist after deleting the app.
-# Run with: sudo ./uninstall-exo.sh [--keep-models]
-#
-# Options:
-#   --keep-models   Preserve ~/.exo/models when removing the EXO data directory.
+# Run with: sudo ./uninstall-exo.sh
 #
 # Components removed:
 # - LaunchDaemon: /Library/LaunchDaemons/io.exo.networksetup.plist
 # - Network script: /Library/Application Support/EXO/
 # - Log files: /var/log/io.exo.networksetup.*
 # - Network location: "exo"
-# - EXO data directory: ~/.exo (or all of ~/.exo except models/ when --keep-models is set)
 # - Launch at login registration
 #
 
 set -euo pipefail
 
-KEEP_MODELS=0
-for arg in "$@"; do
-  case "$arg" in
-  --keep-models)
-    KEEP_MODELS=1
-    ;;
-  -h | --help)
-    echo "Usage: sudo ./uninstall-exo.sh [--keep-models]"
-    echo "  --keep-models   Preserve ~/.exo/models when removing the EXO data directory."
-    exit 0
-    ;;
-  *)
-    echo "Unknown argument: $arg" >&2
-    echo "Usage: sudo ./uninstall-exo.sh [--keep-models]" >&2
-    exit 2
-    ;;
-  esac
-done
-
 LABEL="io.exo.networksetup"
-# Current script path. Older installs used a different filename; keep the
-# legacy path here so a fresh uninstall still cleans up upgraded machines.
-CURRENT_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge.sh"
-LEGACY_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
+SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
 PLIST_DEST="/Library/LaunchDaemons/io.exo.networksetup.plist"
 LOG_OUT="/var/log/${LABEL}.log"
 LOG_ERR="/var/log/${LABEL}.err.log"
 APP_BUNDLE_ID="io.exo.EXO"
 
-# Resolve the invoking user's home, even when run via sudo.
-USER_HOME="$(eval echo "~${SUDO_USER:-$USER}")"
-EXO_DIR="$USER_HOME/.exo"
-
 # Colors for output
 RED='\033[0;31m'
 GREEN='\033[0;32m'
@@ -99,17 +69,11 @@ else
   echo_warn "LaunchDaemon plist not found (already removed?)"
 fi
 
-# Remove the script (current and legacy filenames) — backwards-compatible:
-# tolerate either, both, or neither being present.
-removed_any_script=0
-for script in "$CURRENT_SCRIPT_DEST" "$LEGACY_SCRIPT_DEST"; do
-  if [[ -f $script ]]; then
-    rm -f "$script"
-    echo_info "Removed network setup script: $script"
-    removed_any_script=1
-  fi
-done
-if [[ $removed_any_script -eq 0 ]]; then
+# Remove the script and parent directory
+if [[ -f $SCRIPT_DEST ]]; then
+  rm -f "$SCRIPT_DEST"
+  echo_info "Removed network setup script"
+else
   echo_warn "Network setup script not found (already removed?)"
 fi
 
@@ -151,22 +115,6 @@ if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge";
   echo_info "Re-enabled Thunderbolt Bridge"
 fi
 
-# Remove EXO data directory (~/.exo)
-EXO_DIR_REMOVED=""
-if [[ -d $EXO_DIR ]]; then
-  if [[ $KEEP_MODELS == "1" && -d "$EXO_DIR/models" ]]; then
-    find "$EXO_DIR" -mindepth 1 -maxdepth 1 ! -name models -exec rm -rf {} +
-    EXO_DIR_REMOVED="kept_models"
-    echo_info "Removed ~/.exo (preserved models/)"
-  else
-    rm -rf "$EXO_DIR"
-    EXO_DIR_REMOVED="full"
-    echo_info "Removed ~/.exo"
-  fi
-else
-  echo_warn "~/.exo not found (already removed?)"
-fi
-
 # Note about launch at login registration
 # SMAppService-based login items cannot be removed from a shell script.
 # They can only be unregistered from within the app itself or manually via System Settings.
@@ -196,10 +144,6 @@ echo "  • Network setup LaunchDaemon"
 echo "  • Network configuration script"
 echo "  • Log files"
 echo "  • 'exo' network location"
-case "$EXO_DIR_REMOVED" in
-full) echo "  • EXO data directory (~/.exo)" ;;
-kept_models) echo "  • EXO data directory (~/.exo, models preserved)" ;;
-esac
 echo ""
 echo "Your network has been restored to use the 'Automatic' location."
 echo "Thunderbolt Bridge has been re-enabled (if present)."

bench/eval_configs/models.toml

@@ -7,7 +7,7 @@
 #   name, patterns, reasoning
 #
 # Optional per-model overrides (CLI flags take priority over these):
-#   temperature, top_p, max_tokens, reasoning_effort, enable_thinking
+#   temperature, top_p, max_tokens, reasoning_effort
 #
 # Fallback defaults (when no per-model config):
 #   reasoning:     temperature=1.0, max_tokens=131072, reasoning_effort="high"
@@ -18,9 +18,10 @@
 
 # ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
 # Source: HuggingFace model cards (Qwen/Qwen3.5-*)
-# Model card recommends: temp=0.6, top_p=0.95, top_k=20
-# We omit top_k to match vllm eval (which doesn't set it).
-# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin).
+# 35B-A3B thinking general: temp=1.0, top_p=0.95, top_k=20
+# 397B thinking: temp=0.6, top_p=0.95, top_k=20
+# Non-thinking: temp=0.7, top_p=0.8, top_k=20
+# max_tokens: 32768 general, 81920 for complex math/code
 
 [[model]]
 name = "Qwen3.5 2B"
@@ -28,8 +29,7 @@ patterns = ["Qwen3.5-2B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 9B"
@@ -37,8 +37,7 @@ patterns = ["Qwen3.5-9B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 27B"
@@ -46,17 +45,15 @@ patterns = ["Qwen3.5-27B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 35B A3B"
 patterns = ["Qwen3.5-35B-A3B"]
 reasoning = true
-temperature = 0.6
+temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 122B A10B"
@@ -64,8 +61,7 @@ patterns = ["Qwen3.5-122B-A10B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 [[model]]
 name = "Qwen3.5 397B A17B"
@@ -73,14 +69,12 @@ patterns = ["Qwen3.5-397B-A17B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 81920
 
 # ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
 # Source: HuggingFace model cards (Qwen/Qwen3-*)
-# Model card recommends: temp=0.6, top_p=0.95, top_k=20
-# We omit top_k to match vllm eval (which doesn't set it).
-# Non-thinking: temp=0.7, top_p=0.8
+# Thinking: temp=0.6, top_p=0.95, top_k=20
+# Non-thinking: temp=0.7, top_p=0.8, top_k=20
 # max_tokens: 32768 general, 38912 for complex math/code
 
 [[model]]
@@ -89,7 +83,6 @@ patterns = ["Qwen3-0.6B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -98,7 +91,6 @@ patterns = ["Qwen3-30B-A3B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -107,7 +99,6 @@ patterns = ["Qwen3-235B-A22B"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -116,7 +107,6 @@ patterns = ["Qwen3-Next-80B-A3B-Thinking"]
 reasoning = true
 temperature = 0.6
 top_p = 0.95
-enable_thinking = true
 max_tokens = 38912
 
 [[model]]
@@ -139,9 +129,9 @@ max_tokens = 16384
 name = "Qwen3 Coder Next"
 patterns = ["Qwen3-Coder-Next"]
 reasoning = false
-temperature = 1.0
-top_p = 0.95
-max_tokens = 121072
+temperature = 0.7
+top_p = 0.8
+max_tokens = 16384
 
 # ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
 # Source: OpenAI GitHub README + HuggingFace discussion #21
@@ -175,38 +165,10 @@ patterns = ["DeepSeek-V3.1"]
 reasoning = true
 temperature = 0.0
 
-[[model]]
-name = "DeepSeek V3.2"
-patterns = ["DeepSeek-V3.2"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-
-# ─── NVIDIA Nemotron ───────────────────────────────────────────────────
-# Source: HuggingFace model cards
-# All variants: temp=1.0, top_p=0.95, enable_thinking=true
-
-[[model]]
-name = "Nemotron Cascade 2 30B A3B"
-patterns = ["Nemotron-Cascade-2-30B-A3B"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-
-[[model]]
-name = "Nemotron 3 Super 120B A12B"
-patterns = ["Nemotron-3-Super-120B-A12B", "NVIDIA-Nemotron-3-Super-120B-A12B"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-
 # ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
 # Source: HuggingFace model cards + generation_config.json + docs.z.ai
 # GLM 4.5+: temp=1.0, top_p=0.95
-# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin)
+# Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7
 
 [[model]]
 name = "GLM-5"
@@ -214,8 +176,7 @@ patterns = ["GLM-5"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 131072
 
 [[model]]
 name = "GLM 4.5 Air"
@@ -230,8 +191,7 @@ patterns = ["GLM-4.7-"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 131072
 # Note: matches both GLM-4.7 and GLM-4.7-Flash
 
 # ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
@@ -253,8 +213,7 @@ patterns = ["Kimi-K2.5"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
+max_tokens = 131072
 
 [[model]]
 name = "Kimi K2 Instruct"
@@ -264,17 +223,7 @@ temperature = 0.6
 
 # ─── MiniMax ─────────────────────────────────────────────────────────
 # Source: HuggingFace model cards + generation_config.json
-# All models: temp=1.0, top_p=0.95
-# max_tokens=90000 to match vllm eval (100000 context - 10000 safety margin)
-
-[[model]]
-name = "MiniMax M2.7"
-patterns = ["MiniMax-M2.7"]
-reasoning = true
-temperature = 1.0
-top_p = 0.95
-enable_thinking = true
-max_tokens = 90000
+# All models: temp=1.0, top_p=0.95, top_k=40
 
 [[model]]
 name = "MiniMax M2.5"
@@ -282,8 +231,6 @@ patterns = ["MiniMax-M2.5"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 90000
 
 [[model]]
 name = "MiniMax M2.1"
@@ -304,8 +251,6 @@ patterns = ["Step-3.5-Flash"]
 reasoning = true
 temperature = 1.0
 top_p = 0.95
-enable_thinking = true
-max_tokens = 121072
 
 # ─── Llama (Meta) ───────────────────────────────────────────────────
 # Source: generation_config.json + meta-llama/llama-models generation.py

bench/exo_bench.py

@@ -35,7 +35,6 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
-    find_existing_instance,
     instance_id_from_instance,
     node_ids_from_instance,
     nodes_used_in_instance,
@@ -80,7 +79,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
         model_path = Path(
             snapshot_download(
                 model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
+                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
             )
         )
 
@@ -278,72 +277,28 @@ def run_one_completion(
     prompt_sizer: PromptSizer,
     *,
     use_prefix_cache: bool = False,
-    stream: bool = False,
 ) -> tuple[dict[str, Any], int]:
     content, pp_tokens = prompt_sizer.build(pp_hint)
     payload: dict[str, Any] = {
         "model": model_id,
         "messages": [{"role": "user", "content": content}],
+        "stream": False,
         "max_tokens": tg,
         "logprobs": False,
         "use_prefix_cache": use_prefix_cache,
     }
 
-    if not stream:
-        payload["stream"] = False
-        t0 = time.perf_counter()
-        out = client.post_bench_chat_completions(payload)
-        elapsed = time.perf_counter() - t0
+    t0 = time.perf_counter()
+    out = client.post_bench_chat_completions(payload)
+    elapsed = time.perf_counter() - t0
 
-        stats = out.get("generation_stats")
-        choices = out.get("choices") or [{}]
-        message = choices[0].get("message", {}) if choices else {}
-        content = message.get("content") or ""
-        preview = content[:200] if content else ""
-    else:
-        tokens = 0
-        first_token_time = None
-        t0 = time.perf_counter()
-        text_parts: list[str] = []
-        stats = None
+    stats = out.get("generation_stats")
 
-        for raw_line in client.stream_bench_chat_completions(payload):
-            line = raw_line.strip()
-            if line.startswith(": generation_stats "):
-                with contextlib.suppress(json.JSONDecodeError):
-                    stats = json.loads(line[len(": generation_stats ") :])
-                continue
-            if not line.startswith("data: "):
-                continue
-            data = line[6:]
-            if data == "[DONE]":
-                break
-            try:
-                chunk = json.loads(data)
-                delta = chunk.get("choices", [{}])[0].get("delta", {})
-                if delta.get("content"):
-                    if first_token_time is None:
-                        first_token_time = time.perf_counter()
-                    tokens += 1
-                    text_parts.append(delta["content"])
-            except json.JSONDecodeError:
-                pass
-
-        elapsed = time.perf_counter() - t0
-        preview = "".join(text_parts)[:200]
-
-        if not stats:
-            ttft = (first_token_time - t0) if first_token_time else elapsed
-            gen_time = elapsed - ttft if tokens > 1 else elapsed
-            gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
-            prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
-            stats = {
-                "prompt_tokens": pp_tokens,
-                "generation_tokens": tokens,
-                "prompt_tps": round(prompt_tps, 2),
-                "generation_tps": round(gen_tps, 2),
-                "peak_memory_usage": {"inBytes": 0},
-            }
+    # Extract preview, handling None content (common for thinking models)
+    choices = out.get("choices") or [{}]
+    message = choices[0].get("message", {}) if choices else {}
+    content = message.get("content") or ""
+    preview = content[:200] if content else ""
 
     return {
         "elapsed_s": elapsed,
@@ -470,11 +425,6 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
-    ap.add_argument(
-        "--stream",
-        action="store_true",
-        help="Use /bench/chat/completions with streaming SSE response (bench=True still applies: no EOS detection, no KV cache).",
-    )
     ap.add_argument(
         "--no-system-metrics",
         action="store_true",
@@ -540,124 +490,81 @@ def main() -> int:
         logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
         raise
 
-    # Optionally reuse a running instance for this model
-    reused_instance_id: str | None = None
-    if args.reuse_instance:
-        existing = find_existing_instance(client, full_model_id)
-        if existing:
-            reused_instance_id = existing
-            logger.info(f"Reusing existing instance {reused_instance_id}")
-        else:
-            logger.warning(
-                "--reuse-instance: no existing instance found, creating a new one"
-            )
+    selected = settle_and_fetch_placements(
+        client, full_model_id, args, settle_timeout=args.settle_timeout
+    )
 
-    if reused_instance_id is not None:
-        # Use the existing instance directly — skip placement iteration
-        selected = []
-        download_duration_s = None
+    if not selected:
+        logger.error("No valid placements matched your filters.")
+        return 1
+
+    selected.sort(
+        key=lambda p: (
+            str(p.get("instance_meta", "")),
+            str(p.get("sharding", "")),
+            -nodes_used_in_instance(p["instance"]),
+        ),
+        reverse=True,
+    )
+
+    logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
+    logger.info(f"placements: {len(selected)}")
+    for p in selected:
+        logger.info(
+            f"  - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
+        )
+
+    if args.dry_run:
+        return 0
+
+    settle_deadline = (
+        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    )
+
+    logger.info("Planning phase: checking downloads...")
+    download_duration_s = run_planning_phase(
+        client,
+        full_model_id,
+        selected[0],
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    if download_duration_s is not None:
+        logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
     else:
-        selected = settle_and_fetch_placements(
-            client, full_model_id, args, settle_timeout=args.settle_timeout
-        )
-
-        if not selected:
-            logger.error("No valid placements matched your filters.")
-            return 1
-
-        selected.sort(
-            key=lambda p: (
-                str(p.get("instance_meta", "")),
-                str(p.get("sharding", "")),
-                nodes_used_in_instance(p["instance"]),
-            ),
-            reverse=True,
-        )
-
-        logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
-        logger.info(f"placements: {len(selected)}")
-        for p in selected:
-            logger.info(
-                f"  - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
-            )
-
-        if args.dry_run:
-            return 0
-
-        settle_deadline = (
-            time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-        )
-
-        logger.info("Planning phase: checking downloads...")
-        download_duration_s = run_planning_phase(
-            client,
-            full_model_id,
-            selected[0],
-            args.danger_delete_downloads,
-            args.timeout,
-            settle_deadline,
-        )
-        if download_duration_s is not None:
-            logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
-        else:
-            logger.info("Download: model already cached")
+        logger.info("Download: model already cached")
 
     cluster_snapshot = capture_cluster_snapshot(client)
     all_rows: list[dict[str, Any]] = []
     all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
 
-    # If reusing an existing instance, run a single benchmark pass against it
-    if reused_instance_id is not None:
-        selected = [None]
-
     for preview in selected:
-        created_instance = False
-        if preview is not None:
-            instance = preview["instance"]
-            instance_id = instance_id_from_instance(instance)
+        instance = preview["instance"]
+        instance_id = instance_id_from_instance(instance)
 
-            sharding = str(preview["sharding"])
-            instance_meta = str(preview["instance_meta"])
-            n_nodes = nodes_used_in_instance(instance)
+        sharding = str(preview["sharding"])
+        instance_meta = str(preview["instance_meta"])
+        n_nodes = nodes_used_in_instance(instance)
 
-            logger.info("=" * 80)
-            logger.info(
-                f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
-            )
+        logger.info("=" * 80)
+        logger.info(
+            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
+        )
 
-            # Delete any existing instances to free resources before placing
-            try:
-                state = client.request_json("GET", "/state")
-                for old_id in list(state.get("instances", {}).keys()):
-                    logger.info(f"Deleting stale instance {old_id}")
-                    with contextlib.suppress(ExoHttpError):
-                        client.request_json("DELETE", f"/instance/{old_id}")
-                if state.get("instances"):
-                    time.sleep(2)
-            except Exception as e:
-                logger.warning(f"Failed to clean up stale instances: {e}")
+        client.request_json("POST", "/instance", body={"instance": instance})
+        try:
+            wait_for_instance_ready(client, instance_id)
+        except (RuntimeError, TimeoutError) as e:
+            logger.error(f"Failed to initialize placement: {e}")
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            continue
 
-            client.request_json("POST", "/instance", body={"instance": instance})
-            try:
-                wait_for_instance_ready(client, instance_id)
-            except (RuntimeError, TimeoutError) as e:
-                logger.error(f"Failed to initialize placement: {e}")
-                with contextlib.suppress(ExoHttpError):
-                    client.request_json("DELETE", f"/instance/{instance_id}")
-                continue
-
-            time.sleep(1)
-            created_instance = True
-        else:
-            instance_id = reused_instance_id
-            sharding = "reused"
-            instance_meta = "reused"
-            n_nodes = 0
-            logger.info("=" * 80)
-            logger.info(f"Using existing instance {instance_id}")
+        time.sleep(1)
 
         sampler: SystemMetricsSampler | None = None
-        if not args.no_system_metrics and preview is not None:
+        if not args.no_system_metrics:
             nids = node_ids_from_instance(instance)
             sampler = SystemMetricsSampler(
                 ExoClient(args.host, args.port, timeout_s=30),
@@ -666,20 +573,16 @@ def main() -> int:
             )
             sampler.start()
 
-        def _do_one(c: ExoClient, pp: int, tg: int) -> tuple[dict[str, Any], int]:
-            return run_one_completion(
-                c,
-                full_model_id,
-                pp,
-                tg,
-                prompt_sizer,
-                use_prefix_cache=args.use_prefix_cache,
-                stream=args.stream,
-            )
-
         try:
             for i in range(args.warmup):
-                _do_one(client, pp_list[0], tg_list[0])
+                run_one_completion(
+                    client,
+                    full_model_id,
+                    pp_list[0],
+                    tg_list[0],
+                    prompt_sizer,
+                    use_prefix_cache=args.use_prefix_cache,
+                )
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
             # If pp and tg lists have same length, run in tandem (zip)
@@ -701,7 +604,14 @@ def main() -> int:
                             # Sequential: single request
                             try:
                                 inf_t0 = time.monotonic()
-                                row, actual_pp_tokens = _do_one(client, pp, tg)
+                                row, actual_pp_tokens = run_one_completion(
+                                    client,
+                                    full_model_id,
+                                    pp,
+                                    tg,
+                                    prompt_sizer,
+                                    use_prefix_cache=args.use_prefix_cache,
+                                )
                                 inference_windows.append((inf_t0, time.monotonic()))
                             except Exception as e:
                                 logger.error(e)
@@ -850,12 +760,10 @@ def main() -> int:
                         gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+                        peak = mean(
+                            x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+                        )
 
-                        def _peak_bytes(s: dict[str, Any]) -> float:
-                            pm = s["peak_memory_usage"]
-                            return pm.get("inBytes") or pm.get("in_bytes", 0)
-
-                        peak = mean(_peak_bytes(x["stats"]) for x in runs)
                         summary = (
                             f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                             f"prompt_tokens={ptok} gen_tokens={gtok}    "
@@ -880,16 +788,15 @@ def main() -> int:
                 if placement_metrics:
                     all_system_metrics.update(placement_metrics)
 
-            if created_instance and instance_id is not None:
-                try:
-                    client.request_json("DELETE", f"/instance/{instance_id}")
-                except ExoHttpError as e:
-                    if e.status != 404:
-                        raise
-                wait_for_instance_gone(client, instance_id)
-                logger.debug(f"Deleted instance {instance_id}")
+            try:
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            except ExoHttpError as e:
+                if e.status != 404:
+                    raise
+            wait_for_instance_gone(client, instance_id)
+            logger.debug(f"Deleted instance {instance_id}")
 
-                time.sleep(5)
+            time.sleep(5)
 
     output: dict[str, Any] = {"runs": all_rows}
     if cluster_snapshot:

bench/exo_eval.py

@@ -47,7 +47,6 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
-    find_existing_instance,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -63,15 +62,6 @@ from loguru import logger
 # ---------------------------------------------------------------------------
 
 MAX_RETRIES = 30
-INSTANCE_HEALTH_CHECK_AFTER = (
-    3  # Check instance health after this many consecutive failures
-)
-
-
-class InstanceFailedError(RuntimeError):
-    """Raised when the exo instance is detected as failed/gone."""
-
-
 DEFAULT_MAX_TOKENS = 16_384
 REASONING_MAX_TOKENS = 131_072
 TEMPERATURE_NON_REASONING = 0.0
@@ -281,7 +271,7 @@ def run_humaneval_test(
 
 @dataclass
 class QuestionResult:
-    question_id: int | str
+    question_id: int
     prompt: str
     response: str
     extracted_answer: str | None
@@ -291,11 +281,7 @@ class QuestionResult:
     prompt_tokens: int = 0
     completion_tokens: int = 0
     reasoning_tokens: int = 0
-    reasoning_content: str = ""
-    finish_reason: str = ""
     elapsed_s: float = 0.0
-    power_watts: float = 0.0
-    energy_joules: float = 0.0
 
 
 @dataclass
@@ -531,10 +517,6 @@ class ApiResult:
     prompt_tokens: int
     completion_tokens: int
     reasoning_tokens: int
-    reasoning_content: str = ""
-    finish_reason: str = ""
-    power_watts: float = 0.0
-    energy_joules: float = 0.0
 
 
 async def _call_api(
@@ -548,9 +530,6 @@ async def _call_api(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
-    top_k: int | None = None,
-    min_p: float | None = None,
-    enable_thinking: bool | None = None,
 ) -> ApiResult:
     messages = []
     if system_message:
@@ -567,12 +546,6 @@ async def _call_api(
         body["reasoning_effort"] = reasoning_effort
     if top_p is not None:
         body["top_p"] = top_p
-    if top_k is not None:
-        body["top_k"] = top_k
-    if min_p is not None:
-        body["min_p"] = min_p
-    if enable_thinking is not None:
-        body["enable_thinking"] = enable_thinking
 
     resp = await client.post(
         f"{base_url}/v1/chat/completions",
@@ -581,40 +554,19 @@ async def _call_api(
     )
     resp.raise_for_status()
     data = resp.json()
-    choice = data["choices"][0]
-    message = choice["message"]
-    content = message.get("content") or ""
-    reasoning_content = message.get("reasoning_content") or ""
-    finish_reason = choice.get("finish_reason") or ""
-
-    # For thinking models, empty content is expected when finish_reason is "length"
-    if not content.strip() and finish_reason != "length" and not reasoning_content:
+    content = data["choices"][0]["message"]["content"]
+    if not content or not content.strip():
         raise ValueError("Empty response from model")
     usage = data.get("usage", {})
     details = usage.get("completion_tokens_details", {})
-    power = data.get("power_usage") or {}
     return ApiResult(
         content=content,
         prompt_tokens=usage.get("prompt_tokens", 0),
         completion_tokens=usage.get("completion_tokens", 0),
         reasoning_tokens=details.get("reasoning_tokens", 0) if details else 0,
-        reasoning_content=reasoning_content,
-        finish_reason=finish_reason,
-        power_watts=power.get("total_avg_sys_power_watts", 0.0),
-        energy_joules=power.get("total_energy_joules", 0.0),
     )
 
 
-async def _check_instance_health(base_url: str) -> bool:
-    """Return True if the exo instance is still reachable."""
-    try:
-        async with httpx.AsyncClient() as c:
-            resp = await c.get(f"{base_url}/models", timeout=5.0)
-            return resp.status_code == 200
-    except Exception:
-        return False
-
-
 async def call_with_retries(
     client: httpx.AsyncClient,
     base_url: str,
@@ -626,14 +578,8 @@ async def call_with_retries(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
-    top_k: int | None = None,
-    min_p: float | None = None,
-    enable_thinking: bool | None = None,
-    instance_failed: asyncio.Event | None = None,
 ) -> ApiResult | None:
     for attempt in range(MAX_RETRIES):
-        if instance_failed and instance_failed.is_set():
-            raise InstanceFailedError("Instance already marked as failed")
         try:
             return await _call_api(
                 client,
@@ -646,30 +592,8 @@ async def call_with_retries(
                 system_message,
                 reasoning_effort,
                 top_p,
-                top_k,
-                min_p,
-                enable_thinking,
             )
         except Exception as e:
-            is_conn_error = isinstance(
-                e,
-                (
-                    httpx.ConnectError,
-                    httpx.RemoteProtocolError,
-                    ConnectionRefusedError,
-                    OSError,
-                ),
-            )
-            if (
-                is_conn_error
-                and attempt >= INSTANCE_HEALTH_CHECK_AFTER
-                and not await _check_instance_health(base_url)
-            ):
-                if instance_failed:
-                    instance_failed.set()
-                raise InstanceFailedError(
-                    f"Instance is down after {attempt + 1} failures: {e}"
-                ) from e
             if attempt < MAX_RETRIES - 1:
                 wait = min(2**attempt, 60)
                 logger.warning(
@@ -694,16 +618,10 @@ async def evaluate_benchmark(
     max_tokens: int,
     concurrency: int = 1,
     limit: int | None = None,
-    offset: int = 0,
     timeout: float | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
-    top_k: int | None = None,
-    min_p: float | None = None,
-    enable_thinking: bool | None = None,
     difficulty: str | None = None,
-    checkpoint_path: Path | None = None,
-    release_version: str | None = None,
 ) -> list[QuestionResult]:
     """Run a benchmark. Returns per-question results."""
     import datasets
@@ -734,21 +652,7 @@ async def evaluate_benchmark(
         ds = ds.filter(lambda x: x["difficulty"] == difficulty)
         logger.info(f"Filtered to {len(ds)} {difficulty} problems")
 
-    if release_version and "release_version" in ds.column_names:
-        ds = ds.filter(lambda x: x["release_version"] == release_version)
-        logger.info(
-            f"Filtered to {len(ds)} problems with release_version={release_version}"
-        )
-
-    # Sort by question_id to match LCB runner ordering (scenario_router.py:60).
-    # This ensures [offset:offset+limit] slices select the same problems as vllm.
-    if "question_id" in ds.column_names:
-        ds = ds.sort("question_id")
-
     total = len(ds)
-    if offset > 0:
-        ds = ds.select(range(min(offset, total), total))
-        total = len(ds)
     if limit and limit < total:
         ds = ds.select(range(limit))
         total = limit
@@ -756,13 +660,6 @@ async def evaluate_benchmark(
     logger.info(
         f"Evaluating {benchmark_name}: {total} questions, concurrency={concurrency}, "
         f"temperature={temperature}, max_tokens={max_tokens}"
-        + (f", top_k={top_k}" if top_k is not None else "")
-        + (f", min_p={min_p}" if min_p is not None else "")
-        + (
-            f", enable_thinking={enable_thinking}"
-            if enable_thinking is not None
-            else ""
-        )
     )
 
     if config.kind == "code":
@@ -770,64 +667,16 @@ async def evaluate_benchmark(
             "Code benchmarks execute model-generated code. Use a sandboxed environment."
         )
 
-    # Load checkpoint for resume
-    checkpoint_data: dict[str | int, dict[str, Any]] = {}
-    if checkpoint_path and checkpoint_path.exists():
-        with open(checkpoint_path) as f:
-            for line in f:
-                entry = json.loads(line)
-                checkpoint_data[entry["question_id"]] = entry
-        logger.info(f"Loaded {len(checkpoint_data)} checkpointed results")
-
     semaphore = asyncio.Semaphore(concurrency)
-    instance_failed = asyncio.Event()
     results: list[QuestionResult | None] = [None] * total
     completed = 0
     lock = asyncio.Lock()
 
-    def _get_question_id(idx: int, doc: dict) -> str | int:
-        """Get a stable question ID for checkpointing."""
-        if benchmark_name == "livecodebench":
-            return doc.get("question_id", idx)
-        elif benchmark_name == "humaneval":
-            return doc.get("task_id", idx)
-        return idx
-
     async def process_question(
         idx: int, doc: dict, http_client: httpx.AsyncClient
     ) -> None:
         nonlocal completed
         system_msg = None
-        question_id = _get_question_id(idx, doc)
-
-        # Bail out early if instance is already dead
-        if instance_failed.is_set():
-            return
-
-        # Check checkpoint
-        if question_id in checkpoint_data:
-            cached = checkpoint_data[question_id]
-            results[idx] = QuestionResult(
-                question_id=question_id,
-                prompt=cached.get("prompt", ""),
-                response=cached.get("response", ""),
-                extracted_answer=cached.get("extracted_answer"),
-                gold_answer=cached.get("gold_answer", ""),
-                correct=cached.get("correct", False),
-                error=cached.get("error"),
-                prompt_tokens=cached.get("prompt_tokens", 0),
-                completion_tokens=cached.get("completion_tokens", 0),
-                reasoning_tokens=cached.get("reasoning_tokens", 0),
-                reasoning_content=cached.get("reasoning_content", ""),
-                finish_reason=cached.get("finish_reason", ""),
-                elapsed_s=cached.get("elapsed_s", 0.0),
-                power_watts=cached.get("power_watts", 0.0),
-                energy_joules=cached.get("energy_joules", 0.0),
-            )
-            async with lock:
-                completed += 1
-            logger.info(f"  [{completed}/{total}] {question_id} (cached)")
-            return
 
         if benchmark_name == "gpqa_diamond":
             prompt, gold = format_gpqa_question(doc, idx)
@@ -848,50 +697,24 @@ async def evaluate_benchmark(
             raise ValueError(f"Unknown benchmark: {benchmark_name}")
 
         async with semaphore:
-            if instance_failed.is_set():
-                return
             t0 = time.monotonic()
-            try:
-                # Race the API call against the instance_failed event
-                api_task = asyncio.create_task(
-                    call_with_retries(
-                        http_client,
-                        base_url,
-                        model,
-                        prompt,
-                        temperature,
-                        max_tokens,
-                        timeout,
-                        system_message=system_msg,
-                        reasoning_effort=reasoning_effort,
-                        top_p=top_p,
-                        top_k=top_k,
-                        min_p=min_p,
-                        enable_thinking=enable_thinking,
-                        instance_failed=instance_failed,
-                    )
-                )
-                failed_waiter = asyncio.create_task(instance_failed.wait())
-                done, pending = await asyncio.wait(
-                    [api_task, failed_waiter],
-                    return_when=asyncio.FIRST_COMPLETED,
-                )
-                for p in pending:
-                    p.cancel()
-                    with contextlib.suppress(asyncio.CancelledError):
-                        await p
-                if instance_failed.is_set() and api_task not in done:
-                    logger.error(f"Instance failed, aborting {question_id}")
-                    return
-                api_result = api_task.result()
-            except InstanceFailedError:
-                logger.error(f"Instance failed, skipping {question_id}")
-                return
+            api_result = await call_with_retries(
+                http_client,
+                base_url,
+                model,
+                prompt,
+                temperature,
+                max_tokens,
+                timeout,
+                system_message=system_msg,
+                reasoning_effort=reasoning_effort,
+                top_p=top_p,
+            )
             elapsed = time.monotonic() - t0
 
         if api_result is None:
             result = QuestionResult(
-                question_id=question_id,
+                question_id=idx,
                 prompt=prompt,
                 response="",
                 extracted_answer=None,
@@ -906,17 +729,13 @@ async def evaluate_benchmark(
                 "prompt_tokens": api_result.prompt_tokens,
                 "completion_tokens": api_result.completion_tokens,
                 "reasoning_tokens": api_result.reasoning_tokens,
-                "reasoning_content": api_result.reasoning_content,
-                "finish_reason": api_result.finish_reason,
                 "elapsed_s": elapsed,
-                "power_watts": api_result.power_watts,
-                "energy_joules": api_result.energy_joules,
             }
 
             if config.kind == "mc":
                 extracted = extract_mc_answer(response, valid_letters)
                 result = QuestionResult(
-                    question_id=question_id,
+                    question_id=idx,
                     prompt=prompt,
                     response=response,
                     extracted_answer=extracted,
@@ -930,7 +749,7 @@ async def evaluate_benchmark(
                     check_aime_answer(extracted, int(gold)) if extracted else False
                 )
                 result = QuestionResult(
-                    question_id=question_id,
+                    question_id=idx,
                     prompt=prompt,
                     response=response,
                     extracted_answer=extracted,
@@ -944,7 +763,7 @@ async def evaluate_benchmark(
                 code = extract_code_block(response, preserve_indent=keep_indent)
                 if code is None:
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -959,7 +778,7 @@ async def evaluate_benchmark(
                         code,
                     )
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer="pass" if passed else "fail",
@@ -974,7 +793,7 @@ async def evaluate_benchmark(
                         exec_meta["sample"],
                     )
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer="pass" if passed else "fail",
@@ -985,7 +804,7 @@ async def evaluate_benchmark(
                     )
                 else:
                     result = QuestionResult(
-                        question_id=question_id,
+                        question_id=idx,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -996,7 +815,7 @@ async def evaluate_benchmark(
                     )
             else:
                 result = QuestionResult(
-                    question_id=question_id,
+                    question_id=idx,
                     prompt=prompt,
                     response=response,
                     extracted_answer=None,
@@ -1008,82 +827,24 @@ async def evaluate_benchmark(
 
         results[idx] = result
 
-        # Write checkpoint (skip infra failures so they get retried on resume,
-        # but keep wrong answers — they are legitimate results)
-        if checkpoint_path is not None and result.response:
-            _write_checkpoint(checkpoint_path, result)
-
         async with lock:
             completed += 1
             n = completed
-
-        # Log progress
-        thinking_info = ""
-        if result.reasoning_content:
-            thinking_info = f", {len(result.reasoning_content)} chars thinking"
-        logger.info(
-            f"  [{n}/{total}] {question_id}: {len(result.response)} chars{thinking_info}, "
-            f"tokens: {result.prompt_tokens}+{result.completion_tokens} "
-            f"[{result.finish_reason}]"
-            + (f" {result.extracted_answer}" if result.extracted_answer else "")
-        )
-
-    async def _health_monitor() -> None:
-        """Periodically check if the instance is still alive."""
-        # Wait a bit before first check to let things start
-        await asyncio.sleep(10)
-        while not instance_failed.is_set():
-            if not await _check_instance_health(base_url):
-                # Double-check to avoid false positives
-                await asyncio.sleep(2)
-                if not await _check_instance_health(base_url):
-                    logger.error("Health monitor: instance is down!")
-                    instance_failed.set()
-                    return
-            await asyncio.sleep(5)
+        if n % max(1, total // 20) == 0 or n == total:
+            correct_so_far = sum(1 for r in results if r is not None and r.correct)
+            answered = sum(1 for r in results if r is not None)
+            logger.info(
+                f"  [{n}/{total}] {correct_so_far}/{answered} correct "
+                f"({correct_so_far / max(answered, 1):.1%})"
+            )
 
     async with httpx.AsyncClient() as http_client:
-        monitor = asyncio.create_task(_health_monitor())
         tasks = [process_question(i, doc, http_client) for i, doc in enumerate(ds)]
         await asyncio.gather(*tasks)
-        monitor.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await monitor
-
-    if instance_failed.is_set():
-        completed_count = sum(1 for r in results if r is not None)
-        logger.error(
-            f"Instance failed! Completed {completed_count}/{total} problems. "
-            f"Checkpoint saved — restart to resume remaining problems."
-        )
-        raise InstanceFailedError("Instance failed during evaluation")
 
     return [r for r in results if r is not None]
 
 
-def _write_checkpoint(path: Path, result: QuestionResult) -> None:
-    """Append a single result to the JSONL checkpoint file."""
-    entry = {
-        "question_id": result.question_id,
-        "prompt": result.prompt,
-        "response": result.response,
-        "extracted_answer": result.extracted_answer,
-        "gold_answer": result.gold_answer,
-        "correct": result.correct,
-        "error": result.error,
-        "prompt_tokens": result.prompt_tokens,
-        "completion_tokens": result.completion_tokens,
-        "reasoning_tokens": result.reasoning_tokens,
-        "reasoning_content": result.reasoning_content,
-        "finish_reason": result.finish_reason,
-        "elapsed_s": round(result.elapsed_s, 2),
-        "power_watts": round(result.power_watts, 2),
-        "energy_joules": round(result.energy_joules, 2),
-    }
-    with open(path, "a") as f:
-        f.write(json.dumps(entry) + "\n")
-
-
 # ---------------------------------------------------------------------------
 # Results display
 # ---------------------------------------------------------------------------
@@ -1106,8 +867,6 @@ def print_results(
     total_elapsed = sum(r.elapsed_s for r in results)
     wall_clock = max(r.elapsed_s for r in results) if results else 0.0
     avg_gen_tps = total_completion_tokens / total_elapsed if total_elapsed > 0 else 0.0
-    total_energy = sum(r.energy_joules for r in results)
-    avg_power = sum(r.power_watts for r in results) / max(total, 1)
 
     label = f"[c={concurrency}] " if concurrency is not None else ""
     print(f"\n{label}{benchmark_name}: {correct}/{total} ({accuracy:.1%})")
@@ -1119,10 +878,6 @@ def print_results(
         f"  |  total time: {total_elapsed:.1f}s  wall clock: {wall_clock:.1f}s"
     )
     print(tok_line)
-    if total_energy > 0:
-        print(
-            f"  power: avg {avg_power:.1f}W  |  total energy: {total_energy:.1f}J ({total_energy / 3600:.2f}Wh)"
-        )
     if errors:
         print(f"  API errors: {errors}")
     if no_extract:
@@ -1141,8 +896,6 @@ def print_results(
         "total_elapsed_s": total_elapsed,
         "wall_clock_s": wall_clock,
         "avg_gen_tps": avg_gen_tps,
-        "avg_power_watts": avg_power,
-        "total_energy_joules": total_energy,
     }
 
 
@@ -1300,11 +1053,7 @@ def save_results(
                 "prompt_tokens": r.prompt_tokens,
                 "completion_tokens": r.completion_tokens,
                 "reasoning_tokens": r.reasoning_tokens,
-                "reasoning_content": r.reasoning_content,
-                "finish_reason": r.finish_reason,
                 "elapsed_s": round(r.elapsed_s, 2),
-                "power_watts": round(r.power_watts, 2),
-                "energy_joules": round(r.energy_joules, 2),
             }
             for r in results
         ],
@@ -1320,15 +1069,6 @@ def save_results(
 # ---------------------------------------------------------------------------
 
 
-def _checkpoint_path(
-    results_dir: str, benchmark: str, model: str, concurrency: int
-) -> Path:
-    """Return the JSONL checkpoint path for a benchmark run."""
-    out_dir = Path(results_dir) / model.replace("/", "_") / benchmark
-    out_dir.mkdir(parents=True, exist_ok=True)
-    return out_dir / f"c{concurrency}.checkpoint.jsonl"
-
-
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -1356,12 +1096,6 @@ def main() -> int:
         default=None,
         help="Max questions per benchmark (for fast iteration).",
     )
-    ap.add_argument(
-        "--offset",
-        type=int,
-        default=0,
-        help="Skip first N questions (0-based).",
-    )
 
     reasoning_group = ap.add_mutually_exclusive_group()
     reasoning_group.add_argument(
@@ -1381,8 +1115,6 @@ def main() -> int:
         "--temperature", type=float, default=None, help="Override temperature."
     )
     ap.add_argument("--top-p", type=float, default=None, help="Override top_p.")
-    ap.add_argument("--top-k", type=int, default=None, help="Override top_k.")
-    ap.add_argument("--min-p", type=float, default=None, help="Override min_p.")
     ap.add_argument(
         "--max-tokens", type=int, default=None, help="Override max output tokens."
     )
@@ -1416,31 +1148,15 @@ def main() -> int:
         choices=["easy", "medium", "hard"],
         help="Filter by difficulty (livecodebench only). E.g. --difficulty hard",
     )
-    ap.add_argument(
-        "--release-version",
-        default=None,
-        help="LCB dataset release version (livecodebench only). E.g. release_v5",
-    )
     ap.add_argument(
         "--results-dir",
         default="eval_results",
         help="Directory for result JSON files (default: eval_results).",
     )
     ap.add_argument(
-        "--enable-thinking",
-        type=lambda v: v.lower() in ("true", "1", "yes"),
-        default=None,
-        help="Enable thinking mode for models that support it.",
-    )
-    ap.add_argument(
-        "--force",
+        "--skip-instance-setup",
         action="store_true",
-        help="Discard any existing checkpoint and run from scratch.",
-    )
-    ap.add_argument(
-        "--keep-instance",
-        action="store_true",
-        help="Skip deleting the instance after eval (for chaining runs).",
+        help="Skip exo instance management (assumes model is already running).",
     )
 
     args, _ = ap.parse_known_args()
@@ -1461,26 +1177,13 @@ def main() -> int:
     # Instance management
     client = ExoClient(args.host, args.port, timeout_s=args.timeout)
     instance_id: str | None = None
-    created_instance = False
 
-    _short_id, full_model_id = resolve_model_short_id(
-        client,
-        args.model,
-        force_download=args.force_download,
-    )
-
-    # Optionally reuse a running instance for this model
-    if args.reuse_instance:
-        existing = find_existing_instance(client, full_model_id)
-        if existing:
-            instance_id = existing
-            logger.info(f"Reusing existing instance {instance_id}")
-        else:
-            logger.warning(
-                "--reuse-instance: no existing instance found, creating a new one"
-            )
-
-    if instance_id is None:
+    if not args.skip_instance_setup:
+        short_id, full_model_id = resolve_model_short_id(
+            client,
+            args.model,
+            force_download=args.force_download,
+        )
         selected = settle_and_fetch_placements(
             client,
             full_model_id,
@@ -1495,7 +1198,7 @@ def main() -> int:
             key=lambda p: (
                 str(p.get("instance_meta", "")),
                 str(p.get("sharding", "")),
-                nodes_used_in_instance(p["instance"]),
+                -nodes_used_in_instance(p["instance"]),
             ),
             reverse=True,
         )
@@ -1522,18 +1225,6 @@ def main() -> int:
         if download_duration is not None:
             logger.info(f"Download: {download_duration:.1f}s")
 
-        # Delete any existing instances to free resources before placing
-        try:
-            state = client.request_json("GET", "/state")
-            for old_id in list(state.get("instances", {}).keys()):
-                logger.info(f"Deleting stale instance {old_id}")
-                with contextlib.suppress(ExoHttpError):
-                    client.request_json("DELETE", f"/instance/{old_id}")
-            if state.get("instances"):
-                time.sleep(2)
-        except Exception as e:
-            logger.warning(f"Failed to clean up stale instances: {e}")
-
         client.request_json("POST", "/instance", body={"instance": instance})
         try:
             wait_for_instance_ready(client, instance_id)
@@ -1543,9 +1234,10 @@ def main() -> int:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             return 1
         time.sleep(1)
-        created_instance = True
-
-    cluster_snapshot = capture_cluster_snapshot(client)
+        cluster_snapshot = capture_cluster_snapshot(client)
+    else:
+        full_model_id = args.model
+        cluster_snapshot = None
 
     # Auto-detect reasoning from model config
     model_config = load_model_config(full_model_id)
@@ -1599,57 +1291,16 @@ def main() -> int:
         reasoning_effort = str(cfg["reasoning_effort"])
     else:
         reasoning_effort = "high" if is_reasoning else None
-
-    if args.top_k is not None:
-        top_k: int | None = args.top_k
-    elif "top_k" in cfg:
-        top_k = int(cfg["top_k"])
-    else:
-        top_k = None
-
-    if args.min_p is not None:
-        min_p: float | None = args.min_p
-    elif "min_p" in cfg:
-        min_p = float(cfg["min_p"])
-    else:
-        min_p = None
-
-    if args.enable_thinking is not None:
-        enable_thinking: bool | None = args.enable_thinking
-    elif "enable_thinking" in cfg:
-        enable_thinking = bool(cfg["enable_thinking"])
-    else:
-        enable_thinking = None
-
     base_url = f"http://{args.host}:{args.port}"
 
     logger.info(f"Model: {full_model_id}")
     logger.info(
         f"Settings: temperature={temperature}, max_tokens={max_tokens}, "
         + (f"top_p={top_p}, " if top_p is not None else "")
-        + (f"top_k={top_k}, " if top_k is not None else "")
-        + (f"min_p={min_p}, " if min_p is not None else "")
         + f"reasoning={'yes' if is_reasoning else 'no'}"
         + (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "")
-        + (
-            f", enable_thinking={enable_thinking}"
-            if enable_thinking is not None
-            else ""
-        )
     )
 
-    # Common kwargs for evaluate_benchmark
-    eval_kwargs: dict[str, Any] = {
-        "reasoning_effort": reasoning_effort,
-        "top_p": top_p,
-        "top_k": top_k,
-        "min_p": min_p,
-        "enable_thinking": enable_thinking,
-        "difficulty": args.difficulty,
-        "offset": args.offset,
-        "release_version": args.release_version,
-    }
-
     try:
         if args.compare_concurrency:
             concurrency_levels = parse_int_list(args.compare_concurrency)
@@ -1658,11 +1309,6 @@ def main() -> int:
                 for c in concurrency_levels:
                     logger.info(f"\n{'=' * 50}")
                     logger.info(f"Running {task_name} at concurrency={c}")
-                    checkpoint_path = _checkpoint_path(
-                        args.results_dir, task_name, full_model_id, c
-                    )
-                    if args.force and checkpoint_path.exists():
-                        checkpoint_path.unlink()
                     results = asyncio.run(
                         evaluate_benchmark(
                             task_name,
@@ -1673,8 +1319,9 @@ def main() -> int:
                             concurrency=c,
                             limit=args.limit,
                             timeout=args.request_timeout,
-                            checkpoint_path=checkpoint_path,
-                            **eval_kwargs,
+                            reasoning_effort=reasoning_effort,
+                            top_p=top_p,
+                            difficulty=args.difficulty,
                         )
                     )
                     if results:
@@ -1689,18 +1336,10 @@ def main() -> int:
                             cluster=cluster_snapshot,
                         )
                         results_by_c[c] = results
-                    # Clean up checkpoint on success
-                    if checkpoint_path.exists():
-                        checkpoint_path.unlink()
                 if len(results_by_c) >= 2:
                     print_comparison(task_name, results_by_c)
         else:
             for task_name in task_names:
-                checkpoint_path = _checkpoint_path(
-                    args.results_dir, task_name, full_model_id, args.num_concurrent
-                )
-                if args.force and checkpoint_path.exists():
-                    checkpoint_path.unlink()
                 results = asyncio.run(
                     evaluate_benchmark(
                         task_name,
@@ -1711,8 +1350,9 @@ def main() -> int:
                         concurrency=args.num_concurrent,
                         limit=args.limit,
                         timeout=args.request_timeout,
-                        checkpoint_path=checkpoint_path,
-                        **eval_kwargs,
+                        reasoning_effort=reasoning_effort,
+                        top_p=top_p,
+                        difficulty=args.difficulty,
                     )
                 )
                 if results:
@@ -1726,25 +1366,14 @@ def main() -> int:
                         scores,
                         cluster=cluster_snapshot,
                     )
-                # Clean up checkpoint on success
-                if checkpoint_path.exists():
-                    checkpoint_path.unlink()
     finally:
-        if created_instance and instance_id is not None:
-            if args.keep_instance:
-                logger.info(f"Keeping instance {instance_id} (--keep-instance)")
-            else:
-                try:
-                    client.request_json("DELETE", f"/instance/{instance_id}")
-                except ExoHttpError as e:
-                    if e.status != 404:
-                        raise
-                try:
-                    wait_for_instance_gone(client, instance_id)
-                except TimeoutError:
-                    logger.warning(
-                        f"Timed out waiting for instance {instance_id} to be deleted"
-                    )
+        if instance_id is not None:
+            try:
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            except ExoHttpError as e:
+                if e.status != 404:
+                    raise
+            wait_for_instance_gone(client, instance_id)
 
     return 0
 

bench/harness.py

@@ -6,7 +6,6 @@ import http.client
 import json
 import os
 import time
-from collections.abc import Iterator
 from typing import Any
 from urllib.parse import urlencode
 
@@ -70,30 +69,6 @@ class ExoClient:
     def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
         return self.request_json("POST", "/bench/chat/completions", body=payload)
 
-    def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
-        """POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
-        payload = {**payload, "stream": True}
-        data = json.dumps(payload).encode("utf-8")
-        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
-        try:
-            conn.request(
-                "POST",
-                "/bench/chat/completions",
-                body=data,
-                headers={
-                    "Content-Type": "application/json",
-                    "Accept": "text/event-stream",
-                },
-            )
-            resp = conn.getresponse()
-            if resp.status >= 400:
-                raw = resp.read().decode("utf-8", errors="replace")
-                raise ExoHttpError(resp.status, resp.reason, raw[:300])
-            for line in resp:
-                yield line.decode("utf-8", errors="replace")
-        finally:
-            conn.close()
-
     def get_state_path(self, path: str) -> Any:
         try:
             return self.request_json("GET", f"/state/{path}")
@@ -293,15 +268,11 @@ def sharding_filter(sharding: str, wanted: str) -> bool:
 
 
 def fetch_and_filter_placements(
-    client: ExoClient,
-    full_model_id: str,
-    args: argparse.Namespace,
-    node_id: str | None = None,
+    client: ExoClient, full_model_id: str, args: argparse.Namespace
 ) -> list[dict[str, Any]]:
-    params: dict[str, str] = {"model_id": full_model_id}
-    if node_id is not None:
-        params["node_ids"] = node_id
-    previews_resp = client.request_json("GET", "/instance/previews", params=params)
+    previews_resp = client.request_json(
+        "GET", "/instance/previews", params={"model_id": full_model_id}
+    )
     previews = previews_resp.get("previews") or []
 
     selected: list[dict[str, Any]] = []
@@ -361,9 +332,8 @@ def settle_and_fetch_placements(
     full_model_id: str,
     args: argparse.Namespace,
     settle_timeout: float = 0,
-    node_id: str | None = None,
 ) -> list[dict[str, Any]]:
-    selected = fetch_and_filter_placements(client, full_model_id, args, node_id=node_id)
+    selected = fetch_and_filter_placements(client, full_model_id, args)
 
     if not selected and settle_timeout > 0:
         backoff = _SETTLE_INITIAL_BACKOFF_S
@@ -376,9 +346,7 @@ def settle_and_fetch_placements(
             )
             time.sleep(min(backoff, remaining))
             backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            selected = fetch_and_filter_placements(
-                client, full_model_id, args, node_id=node_id
-            )
+            selected = fetch_and_filter_placements(client, full_model_id, args)
 
     return selected
 
@@ -494,8 +462,9 @@ def run_planning_phase(
         )
         logger.info(f"Started download on {node_id}")
 
-    # Wait for downloads (no timeout — poll until complete or failed)
-    while True:
+    # Wait for downloads
+    start = time.time()
+    while time.time() - start < timeout:
         all_done = True
         for node_id in node_ids:
             node_downloads = client.get_node_downloads(node_id) or []
@@ -545,24 +514,9 @@ def run_planning_phase(
             if download_t0 is not None:
                 return time.perf_counter() - download_t0
             return None
-        time.sleep(10)
+        time.sleep(1)
 
-
-def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
-    """Find an existing running instance for the given model."""
-    try:
-        state = client.request_json("GET", "/state")
-    except Exception:
-        return None
-    for inst_id, inst in state.get("instances", {}).items():
-        # Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
-        for _inst_type, inner in inst.items():
-            if not isinstance(inner, dict):
-                continue
-            sa = inner.get("shardAssignments", {})
-            if sa.get("modelId") == model_id:
-                return inst_id
-    return None
+    raise TimeoutError("Downloads did not complete in time")
 
 
 def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
@@ -618,8 +572,3 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Delete existing models from smallest to largest to make room for benchmark model.",
     )
-    ap.add_argument(
-        "--reuse-instance",
-        action="store_true",
-        help="Reuse an existing running instance for this model instead of creating a new one.",
-    )

bench/prefill-decode.toml

@@ -1,36 +0,0 @@
-# Prefill/Decode disaggregation benchmark config.
-#
-# Top-level keys are bench-wide. [prefill] and [decode] sections set per-side
-# placement filters and (optionally) per-side model.
-#
-# Example:
-#   uv run python bench/prefill_decode_bench.py --config bench/prefill-decode.toml
-
-host = "james"
-port = 52415
-timeout = 7200.0
-settle_timeout = 60.0
-
-# Workload
-pp = [4096]
-tg = [512]
-repeat = 1
-warmup = 0
-
-json_out = "bench/prefill_decode_results.json"
-
-[prefill]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "mike"
-instance_meta = "ring"
-sharding = "pipeline"
-min_nodes = 1
-max_nodes = 1
-
-[decode]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "james"
-instance_meta = "ring"
-sharding = "pipeline"
-min_nodes = 1
-max_nodes = 1

bench/prefill_decode_bench.py

@@ -1,785 +0,0 @@
-# type: ignore
-#!/usr/bin/env python3
-"""Disaggregated prefill-decode benchmark for exo (MLX → MLX).
-
-Spins up two MLX instances on the cluster, marks one as Prefill source and
-the other as Decode target via /v1/instance-links, then sends chat
-completions to the API. The master routes the request to the decode
-instance and stamps `prefill_endpoint` pointing at the prefill instance —
-the worker decides per-request whether to ship prefill remotely
-(uncached_count > REMOTE_PREFILL_MIN_TOKENS).
-
-Usage:
-    uv run python bench/prefill_decode_bench.py --model <id> --pp 2048,8192 --tg 128
-    uv run python bench/prefill_decode_bench.py --model <id> --pp 4096 --tg 128 --repeat 3
-    uv run python bench/prefill_decode_bench.py --model <id> --pp 2048 --tg 128 --dry-run
-"""
-
-from __future__ import annotations
-
-import argparse
-import contextlib
-import copy
-import itertools
-import json
-import sys
-import time
-import tomllib
-from pathlib import Path
-from statistics import mean
-from typing import Any
-
-from exo_bench import (
-    PromptSizer,
-    format_peak_memory,
-    load_tokenizer_for_bench,
-    parse_int_list,
-)
-from harness import (
-    ExoClient,
-    ExoHttpError,
-    add_common_instance_args,
-    instance_id_from_instance,
-    node_ids_from_instance,
-    nodes_used_in_instance,
-    resolve_model_short_id,
-    run_planning_phase,
-    settle_and_fetch_placements,
-    unwrap_instance,
-    wait_for_instance_gone,
-    wait_for_instance_ready,
-)
-from loguru import logger
-
-
-def _node_id_to_friendly(client: ExoClient) -> dict[str, str]:
-    identities = client.get_node_identities() or {}
-    out: dict[str, str] = {}
-    for node_id, identity in identities.items():
-        if isinstance(identity, dict):
-            name = identity.get("friendlyName") or identity.get("friendly_name")
-            if isinstance(name, str):
-                out[str(node_id)] = name
-    return out
-
-
-def _placement_node_friendly_names(
-    placement: dict[str, Any], id_to_friendly: dict[str, str]
-) -> list[str]:
-    instance = placement["instance"]
-    return [id_to_friendly.get(nid, nid) for nid in node_ids_from_instance(instance)]
-
-
-def _filter_by_node(
-    placements: list[dict[str, Any]],
-    friendly_name: str,
-    id_to_friendly: dict[str, str],
-) -> list[dict[str, Any]]:
-    target = friendly_name.lower()
-    matched: list[dict[str, Any]] = []
-    for p in placements:
-        names = [n.lower() for n in _placement_node_friendly_names(p, id_to_friendly)]
-        if any(target == n or target in n for n in names):
-            matched.append(p)
-    return matched
-
-
-def _node_id_by_friendly(id_to_friendly: dict[str, str], target: str) -> str | None:
-    target_lc = target.lower()
-    for nid, name in id_to_friendly.items():
-        if target_lc == name.lower() or target_lc in name.lower():
-            return nid
-    return None
-
-
-def _load_toml(path: str) -> dict[str, Any]:
-    with Path(path).open("rb") as f:
-        return tomllib.load(f)
-
-
-_TOP_LEVEL_TOML_KEYS = {
-    "host",
-    "port",
-    "timeout",
-    "settle_timeout",
-    "model",
-    "pp",
-    "tg",
-    "repeat",
-    "warmup",
-    "json_out",
-    "instance_meta",
-    "sharding",
-    "min_nodes",
-    "max_nodes",
-    "force_download",
-    "danger_delete_downloads",
-    "all_combinations",
-}
-
-
-def _inject_toml_into_argv() -> None:
-    """If --config X is in sys.argv, pre-load it and inject required CLI args
-    (--model, --pp, --tg) so argparse's required=True checks pass."""
-    argv = sys.argv
-    if "--config" not in argv:
-        return
-    idx = argv.index("--config")
-    if idx + 1 >= len(argv):
-        return
-    cfg_path = argv[idx + 1]
-    cfg = _load_toml(cfg_path)
-    decode = cfg.get("decode", {})
-
-    def _has(flag: str) -> bool:
-        return any(a == flag or a.startswith(flag + "=") for a in argv)
-
-    # --model: prefer top-level, then [decode].model
-    if not _has("--model"):
-        model = cfg.get("model") or decode.get("model")
-        if model:
-            argv += ["--model", str(model)]
-    if not _has("--pp"):
-        pp = cfg.get("pp")
-        if pp:
-            argv += (
-                ["--pp", *(str(x) for x in pp)]
-                if isinstance(pp, list)
-                else [
-                    "--pp",
-                    str(pp),
-                ]
-            )
-    if not _has("--tg"):
-        tg = cfg.get("tg")
-        if tg:
-            argv += (
-                ["--tg", *(str(x) for x in tg)]
-                if isinstance(tg, list)
-                else [
-                    "--tg",
-                    str(tg),
-                ]
-            )
-
-
-def _merge_toml_into_args(args: argparse.Namespace, cfg: dict[str, Any]) -> None:
-    """Apply top-level toml keys onto args namespace where args has a default."""
-    for key, value in cfg.items():
-        if key in {"prefill", "decode"}:
-            continue
-        if key not in _TOP_LEVEL_TOML_KEYS:
-            continue
-        attr = key
-        current = getattr(args, attr, None)
-        if current in (None, [], False):
-            setattr(args, attr, value)
-
-
-def _side_args(
-    base: argparse.Namespace, overrides: dict[str, Any]
-) -> argparse.Namespace:
-    out = copy.copy(base)
-    for k in (
-        "instance_meta",
-        "sharding",
-        "min_nodes",
-        "max_nodes",
-        "skip_pipeline_jaccl",
-        "skip_tensor_ring",
-    ):
-        if k in overrides:
-            setattr(out, k, overrides[k])
-    return out
-
-
-def _pick_two_distinct_placements(
-    placements: list[dict[str, Any]],
-) -> tuple[dict[str, Any], dict[str, Any]] | None:
-    if len(placements) < 2:
-        return None
-    seen_nodes: set[tuple[str, ...]] = set()
-    chosen: list[dict[str, Any]] = []
-    for p in placements:
-        nodes = tuple(sorted(str(n) for n in p.get("nodes", [])))
-        if nodes in seen_nodes:
-            continue
-        seen_nodes.add(nodes)
-        chosen.append(p)
-        if len(chosen) == 2:
-            return chosen[0], chosen[1]
-    return None
-
-
-def _create_instance_link(
-    client: ExoClient,
-    prefill_instance_id: str,
-    decode_instance_id: str,
-) -> str:
-    out = client.request_json(
-        "POST",
-        "/v1/instance-links",
-        body={
-            "prefill_instances": [prefill_instance_id],
-            "decode_instances": [decode_instance_id],
-        },
-    )
-    return str(out.get("commandId", ""))
-
-
-def _list_instance_links(client: ExoClient) -> list[dict[str, Any]]:
-    out = client.request_json("GET", "/v1/instance-links")
-    return out if isinstance(out, list) else []
-
-
-def _delete_instance_link(client: ExoClient, link_id: str) -> None:
-    client.request_json("DELETE", f"/v1/instance-links/{link_id}")
-
-
-def run_one(
-    client: ExoClient,
-    model_id: str,
-    pp_hint: int,
-    tg: int,
-    prompt_sizer: PromptSizer,
-) -> tuple[dict[str, Any], int]:
-    content, pp_tokens = prompt_sizer.build(pp_hint)
-    payload: dict[str, Any] = {
-        "model": model_id,
-        "messages": [{"role": "user", "content": content}],
-        "stream": False,
-        "max_tokens": tg,
-    }
-
-    t0 = time.perf_counter()
-    out = client.post_bench_chat_completions(payload)
-    elapsed = time.perf_counter() - t0
-
-    stats = out.get("generation_stats")
-    choices = out.get("choices") or [{}]
-    message = choices[0].get("message", {}) if choices else {}
-    text = message.get("content") or ""
-    preview = text[:200] if text else ""
-
-    return {
-        "elapsed_s": elapsed,
-        "output_text_preview": preview,
-        "stats": stats,
-    }, pp_tokens
-
-
-def _run_phase(
-    *,
-    client: ExoClient,
-    label: str,
-    pp_tg_pairs: list[tuple[int, int]],
-    model_id: str,
-    prompt_sizer: PromptSizer,
-    warmup: int,
-    repeat: int,
-    common_meta: dict[str, Any],
-) -> list[dict[str, Any]]:
-    logger.info(f"=== phase: {label} (model={model_id}) ===")
-    rows: list[dict[str, Any]] = []
-    for i in range(warmup):
-        run_one(client, model_id, pp_tg_pairs[0][0], pp_tg_pairs[0][1], prompt_sizer)
-        logger.debug(f"  warmup {i + 1}/{warmup} done")
-
-    for pp, tg in pp_tg_pairs:
-        logger.info(f"--- {label}: pp={pp} tg={tg} ---")
-        runs: list[dict[str, Any]] = []
-        for r in range(repeat):
-            time.sleep(2)
-            try:
-                row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
-            except Exception as e:
-                logger.error(e)
-                continue
-            row.update(common_meta)
-            row.update(
-                {
-                    "phase": label,
-                    "phase_model_id": model_id,
-                    "pp_tokens": actual_pp_tokens,
-                    "tg": tg,
-                    "repeat_index": r,
-                }
-            )
-            runs.append(row)
-            rows.append(row)
-
-        if runs:
-            prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-            gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
-            ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
-            gtok = mean(x["stats"]["generation_tokens"] for x in runs)
-            peak = mean(x["stats"]["peak_memory_usage"]["inBytes"] for x in runs)
-            avg_elapsed = mean(x["elapsed_s"] for x in runs)
-            logger.info(
-                f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
-                f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                f"peak_memory={format_peak_memory(peak)}    "
-                f"avg_elapsed={avg_elapsed:.2f}s"
-            )
-        time.sleep(2)
-    return rows
-
-
-def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, float]]:
-    grouped: dict[tuple[int, int], list[dict[str, Any]]] = {}
-    for r in rows:
-        key = (int(r["pp_tokens"]), int(r["tg"]))
-        grouped.setdefault(key, []).append(r)
-    out: dict[tuple[int, int], dict[str, float]] = {}
-    for key, runs in grouped.items():
-        out[key] = {
-            "prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
-            "gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
-            "elapsed_s": mean(x["elapsed_s"] for x in runs),
-        }
-    return out
-
-
-def _print_diff(
-    disagg_rows: list[dict[str, Any]],
-    decode_alone_rows: list[dict[str, Any]],
-    prefill_alone_rows: list[dict[str, Any]],
-) -> None:
-    disagg = _summarise(disagg_rows)
-    decode_alone = _summarise(decode_alone_rows)
-    prefill_alone = _summarise(prefill_alone_rows)
-    keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
-
-    width = 64
-    for key in sorted(keys):
-        pp, tg = key
-        logger.info("─" * width)
-        logger.info(f"  pp={pp}  tg={tg}")
-        logger.info("─" * width)
-        logger.info(
-            f"  {'phase':<16} {'elapsed':>10}  {'prompt_tps':>11}  {'gen_tps':>9}"
-        )
-        for label, summary in (
-            ("disaggregated", disagg.get(key)),
-            ("decode_alone", decode_alone.get(key)),
-            ("prefill_alone", prefill_alone.get(key)),
-        ):
-            if summary is None:
-                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
-                continue
-            logger.info(
-                f"  {label:<16} "
-                f"{summary['elapsed_s']:>9.2f}s  "
-                f"{summary['prompt_tps']:>11.1f}  "
-                f"{summary['gen_tps']:>9.2f}"
-            )
-
-        d = disagg.get(key)
-        da = decode_alone.get(key)
-        pa = prefill_alone.get(key)
-        if d and da and d["elapsed_s"] > 0:
-            logger.info(
-                f"  speedup vs decode_alone:  {da['elapsed_s'] / d['elapsed_s']:.2f}x"
-            )
-        if d and pa and d["elapsed_s"] > 0:
-            logger.info(
-                f"  speedup vs prefill_alone: {pa['elapsed_s'] / d['elapsed_s']:.2f}x"
-            )
-    logger.info("─" * width)
-
-
-def main() -> int:
-    _inject_toml_into_argv()
-    ap = argparse.ArgumentParser(
-        prog="prefill-decode-bench",
-        description="Benchmark MLX-MLX disaggregated prefill/decode via instance links.",
-    )
-    add_common_instance_args(ap)
-    ap.add_argument(
-        "--pp",
-        nargs="+",
-        required=True,
-        help="Prompt-size hints (ints, must be >1000). Accepts commas.",
-    )
-    ap.add_argument(
-        "--tg",
-        nargs="+",
-        required=True,
-        help="Generation lengths (ints). Accepts commas.",
-    )
-    ap.add_argument(
-        "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
-    )
-    ap.add_argument(
-        "--warmup",
-        type=int,
-        default=0,
-        help="Warmup runs (uses first pp/tg).",
-    )
-    ap.add_argument(
-        "--json-out",
-        default="bench/prefill_decode_results.json",
-        help="Write raw per-run results JSON to this path.",
-    )
-    ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
-    ap.add_argument(
-        "--dry-run", action="store_true", help="List selected placements and exit."
-    )
-    ap.add_argument(
-        "--all-combinations",
-        action="store_true",
-        help="Force all pp×tg combinations even when lists have equal length.",
-    )
-    ap.add_argument(
-        "--prefill-model",
-        default=None,
-        help="Model id for the prefill instance. Defaults to --model.",
-    )
-    ap.add_argument(
-        "--prefill-node",
-        default=None,
-        help="friendly_name of the node hosting the prefill instance.",
-    )
-    ap.add_argument(
-        "--decode-node",
-        default=None,
-        help="friendly_name of the node hosting the decode instance.",
-    )
-    ap.add_argument(
-        "--config",
-        default=None,
-        help="TOML config file. CLI flags override toml values.",
-    )
-    ap.add_argument(
-        "--compare-baseline",
-        action="store_true",
-        help="Also run each (pp,tg) pair without the prefill/decode link "
-        "(decode instance does its own prefill) and report the diff.",
-    )
-    args = ap.parse_args()
-    cfg = _load_toml(args.config) if args.config else {}
-    _merge_toml_into_args(args, cfg)
-    prefill_overrides = cfg.get("prefill", {}) if cfg else {}
-    decode_overrides = cfg.get("decode", {}) if cfg else {}
-    if args.prefill_model is None and "model" in prefill_overrides:
-        args.prefill_model = prefill_overrides["model"]
-    if args.prefill_node is None and "node" in prefill_overrides:
-        args.prefill_node = prefill_overrides["node"]
-    if args.decode_node is None and "node" in decode_overrides:
-        args.decode_node = decode_overrides["node"]
-    if "model" in decode_overrides and not args.model:
-        args.model = decode_overrides["model"]
-
-    pp_list = parse_int_list(args.pp)
-    tg_list = parse_int_list(args.tg)
-    if not pp_list or not tg_list:
-        logger.error("pp and tg lists must be non-empty")
-        return 2
-    for pp in pp_list:
-        if pp <= 1000:
-            logger.error(
-                f"pp={pp} must be >1000 (remote prefill triggers when uncached >1000)"
-            )
-            return 2
-    if args.repeat <= 0:
-        logger.error("--repeat must be >= 1")
-        return 2
-
-    use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
-    if use_combinations:
-        logger.info(
-            f"pp/tg mode: combinations (product) — {len(pp_list) * len(tg_list)} pairs"
-        )
-    else:
-        logger.info(f"pp/tg mode: tandem (zip) — {len(pp_list)} pairs")
-
-    client = ExoClient(args.host, args.port, timeout_s=args.timeout)
-
-    decode_short_id, decode_full_id = resolve_model_short_id(
-        client, args.model, force_download=args.force_download
-    )
-    if args.prefill_model:
-        prefill_short_id, prefill_full_id = resolve_model_short_id(
-            client, args.prefill_model, force_download=args.force_download
-        )
-    else:
-        prefill_short_id, prefill_full_id = decode_short_id, decode_full_id
-
-    tokenizer = load_tokenizer_for_bench(decode_full_id)
-    if tokenizer is None:
-        raise RuntimeError("[prefill-decode-bench] decode tokenizer load failed")
-    try:
-        decode_prompt_sizer = PromptSizer(tokenizer)
-    except Exception:
-        logger.error("[prefill-decode-bench] decode prompt sizing failed")
-        raise
-
-    if prefill_full_id == decode_full_id:
-        prefill_prompt_sizer = decode_prompt_sizer
-    else:
-        prefill_tokenizer = load_tokenizer_for_bench(prefill_full_id)
-        if prefill_tokenizer is None:
-            raise RuntimeError("[prefill-decode-bench] prefill tokenizer load failed")
-        prefill_prompt_sizer = PromptSizer(prefill_tokenizer)
-
-    id_to_friendly = _node_id_to_friendly(client)
-
-    prefill_args = _side_args(args, prefill_overrides)
-    decode_args = _side_args(args, decode_overrides)
-
-    if prefill_full_id == decode_full_id and prefill_overrides == decode_overrides:
-        placements = settle_and_fetch_placements(
-            client, decode_full_id, args, settle_timeout=args.settle_timeout
-        )
-        prefill_candidates = (
-            _filter_by_node(placements, args.prefill_node, id_to_friendly)
-            if args.prefill_node
-            else placements
-        )
-        decode_candidates = (
-            _filter_by_node(placements, args.decode_node, id_to_friendly)
-            if args.decode_node
-            else placements
-        )
-        if args.prefill_node and not prefill_candidates:
-            logger.error(f"No placement on prefill node {args.prefill_node!r}.")
-            return 1
-        if args.decode_node and not decode_candidates:
-            logger.error(f"No placement on decode node {args.decode_node!r}.")
-            return 1
-        if args.prefill_node and args.decode_node:
-            prefill_p = prefill_candidates[0]
-            decode_p = decode_candidates[0]
-        else:
-            pair = _pick_two_distinct_placements(placements)
-            if pair is None:
-                logger.error(
-                    "Need at least two distinct-node MLX placements for the same model."
-                )
-                return 1
-            prefill_p, decode_p = pair
-            if args.prefill_node:
-                prefill_p = prefill_candidates[0]
-            if args.decode_node:
-                decode_p = decode_candidates[0]
-    else:
-        prefill_node_id = (
-            _node_id_by_friendly(id_to_friendly, args.prefill_node)
-            if args.prefill_node
-            else None
-        )
-        decode_node_id = (
-            _node_id_by_friendly(id_to_friendly, args.decode_node)
-            if args.decode_node
-            else None
-        )
-        if args.prefill_node and prefill_node_id is None:
-            logger.error(f"Unknown node {args.prefill_node!r}.")
-            return 1
-        if args.decode_node and decode_node_id is None:
-            logger.error(f"Unknown node {args.decode_node!r}.")
-            return 1
-        prefill_placements = settle_and_fetch_placements(
-            client,
-            prefill_full_id,
-            prefill_args,
-            settle_timeout=args.settle_timeout,
-            node_id=prefill_node_id,
-        )
-        decode_placements = settle_and_fetch_placements(
-            client,
-            decode_full_id,
-            decode_args,
-            settle_timeout=args.settle_timeout,
-            node_id=decode_node_id,
-        )
-        if not prefill_placements:
-            logger.error(
-                f"No placement found for prefill model {prefill_full_id}"
-                f"{f' on node {args.prefill_node!r}' if args.prefill_node else ''}."
-            )
-            return 1
-        if not decode_placements:
-            logger.error(
-                f"No placement found for decode model {decode_full_id}"
-                f"{f' on node {args.decode_node!r}' if args.decode_node else ''}."
-            )
-            return 1
-        prefill_p = prefill_placements[0]
-        decode_p = decode_placements[0]
-
-    prefill_node_names = _placement_node_friendly_names(prefill_p, id_to_friendly)
-    decode_node_names = _placement_node_friendly_names(decode_p, id_to_friendly)
-    _ = unwrap_instance
-
-    prefill_instance = prefill_p["instance"]
-    decode_instance = decode_p["instance"]
-    prefill_id = instance_id_from_instance(prefill_instance)
-    decode_id = instance_id_from_instance(decode_instance)
-    prefill_meta = str(prefill_p.get("instance_meta", ""))
-    decode_meta = str(decode_p.get("instance_meta", ""))
-    prefill_nodes = nodes_used_in_instance(prefill_instance)
-    decode_nodes = nodes_used_in_instance(decode_instance)
-
-    logger.info("=" * 80)
-    logger.info(
-        f"PREFILL: {prefill_meta} / nodes={prefill_nodes} ({','.join(prefill_node_names)}) "
-        f"/ {prefill_short_id} ({prefill_full_id}) / instance_id={prefill_id}"
-    )
-    logger.info(
-        f"DECODE:  {decode_meta} / nodes={decode_nodes} ({','.join(decode_node_names)}) "
-        f"/ {decode_short_id} ({decode_full_id}) / instance_id={decode_id}"
-    )
-
-    if args.dry_run:
-        return 0
-
-    settle_deadline = (
-        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-    )
-
-    logger.info("Planning phase: prefill...")
-    run_planning_phase(
-        client,
-        prefill_full_id,
-        prefill_p,
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-    logger.info("Planning phase: decode...")
-    run_planning_phase(
-        client,
-        decode_full_id,
-        decode_p,
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-
-    if use_combinations:
-        pp_tg_pairs = list(itertools.product(pp_list, tg_list))
-    else:
-        pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
-
-    common_meta = {
-        "decode_model_short_id": decode_short_id,
-        "decode_model_id": decode_full_id,
-        "prefill_model_short_id": prefill_short_id,
-        "prefill_model_id": prefill_full_id,
-        "prefill_instance_id": prefill_id,
-        "prefill_instance_meta": prefill_meta,
-        "prefill_nodes": prefill_nodes,
-        "decode_instance_id": decode_id,
-        "decode_instance_meta": decode_meta,
-        "decode_nodes": decode_nodes,
-    }
-
-    all_rows: list[dict[str, Any]] = []
-    disagg_rows: list[dict[str, Any]] = []
-    decode_alone_rows: list[dict[str, Any]] = []
-    prefill_alone_rows: list[dict[str, Any]] = []
-    link_id = ""
-    prefill_alive = False
-    decode_alive = False
-    try:
-        logger.info("Creating prefill instance...")
-        client.request_json("POST", "/instance", body={"instance": prefill_instance})
-        wait_for_instance_ready(client, prefill_id)
-        prefill_alive = True
-        logger.info("Prefill instance ready")
-
-        if args.compare_baseline:
-            time.sleep(2)
-            prefill_alone_rows = _run_phase(
-                client=client,
-                label="prefill_alone",
-                pp_tg_pairs=pp_tg_pairs,
-                model_id=prefill_full_id,
-                prompt_sizer=prefill_prompt_sizer,
-                warmup=args.warmup,
-                repeat=args.repeat,
-                common_meta=common_meta,
-            )
-            all_rows.extend(prefill_alone_rows)
-
-        logger.info("Creating decode instance...")
-        client.request_json("POST", "/instance", body={"instance": decode_instance})
-        wait_for_instance_ready(client, decode_id)
-        decode_alive = True
-        logger.info("Decode instance ready")
-
-        logger.info("Linking instances (prefill → decode)...")
-        _create_instance_link(client, prefill_id, decode_id)
-        time.sleep(1)
-        links = _list_instance_links(client)
-        if not links:
-            logger.error("Link did not appear in state.")
-            return 1
-        link_id = str(links[-1].get("linkId") or links[-1].get("link_id") or "")
-        logger.info(f"Link created: {link_id}")
-        time.sleep(2)
-
-        disagg_rows = _run_phase(
-            client=client,
-            label="disaggregated",
-            pp_tg_pairs=pp_tg_pairs,
-            model_id=decode_full_id,
-            prompt_sizer=decode_prompt_sizer,
-            warmup=args.warmup,
-            repeat=args.repeat,
-            common_meta=common_meta,
-        )
-        all_rows.extend(disagg_rows)
-
-        if args.compare_baseline:
-            logger.info("Removing link and prefill instance to isolate decode_alone.")
-            with contextlib.suppress(ExoHttpError):
-                if link_id:
-                    _delete_instance_link(client, link_id)
-                    link_id = ""
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{prefill_id}")
-            wait_for_instance_gone(client, prefill_id)
-            prefill_alive = False
-            time.sleep(2)
-
-            decode_alone_rows = _run_phase(
-                client=client,
-                label="decode_alone",
-                pp_tg_pairs=pp_tg_pairs,
-                model_id=decode_full_id,
-                prompt_sizer=decode_prompt_sizer,
-                warmup=args.warmup,
-                repeat=args.repeat,
-                common_meta=common_meta,
-            )
-            all_rows.extend(decode_alone_rows)
-
-            _print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
-    finally:
-        with contextlib.suppress(ExoHttpError):
-            if link_id:
-                _delete_instance_link(client, link_id)
-        if decode_alive:
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{decode_id}")
-            wait_for_instance_gone(client, decode_id)
-        if prefill_alive:
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{prefill_id}")
-            wait_for_instance_gone(client, prefill_id)
-        logger.debug("Deleted both instances")
-
-    if args.stdout:
-        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
-    elif args.json_out:
-        with open(args.json_out, "w", encoding="utf-8") as f:
-            json.dump(all_rows, f, indent=2, ensure_ascii=False)
-        logger.debug(f"\nWrote results JSON: {args.json_out}")
-
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

dashboard/src/lib/components/HeaderNav.svelte

@@ -1,8 +1,5 @@
 <script lang="ts">
   import { browser } from "$app/environment";
-  import { featureFlags } from "$lib/stores/app.svelte";
-
-  const showAdvanced = $derived(featureFlags()["disaggregation"] === true);
 
   interface Props {
     showHome?: boolean;
@@ -300,28 +297,5 @@
       </svg>
       <span class="hidden sm:inline">Integrations</span>
     </a>
-    {#if showAdvanced}
-      <a
-        href="/#/advanced"
-        class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
-        title="Advanced cluster settings"
-      >
-        <svg
-          class="w-4 h-4"
-          viewBox="0 0 24 24"
-          fill="none"
-          stroke="currentColor"
-          stroke-width="2"
-          stroke-linecap="round"
-          stroke-linejoin="round"
-        >
-          <circle cx="12" cy="12" r="3" />
-          <path
-            d="M19.4 15a1.65 1.65 0 0 0 .33 1.82l.06.06a2 2 0 0 1 0 2.83 2 2 0 0 1-2.83 0l-.06-.06a1.65 1.65 0 0 0-1.82-.33 1.65 1.65 0 0 0-1 1.51V21a2 2 0 0 1-4 0v-.09A1.65 1.65 0 0 0 9 19.4a1.65 1.65 0 0 0-1.82.33l-.06.06a2 2 0 0 1-2.83 0 2 2 0 0 1 0-2.83l.06-.06a1.65 1.65 0 0 0 .33-1.82 1.65 1.65 0 0 0-1.51-1H3a2 2 0 0 1 0-4h.09A1.65 1.65 0 0 0 4.6 9a1.65 1.65 0 0 0-.33-1.82l-.06-.06a2 2 0 0 1 0-2.83 2 2 0 0 1 2.83 0l.06.06a1.65 1.65 0 0 0 1.82.33H9a1.65 1.65 0 0 0 1-1.51V3a2 2 0 0 1 4 0v.09a1.65 1.65 0 0 0 1 1.51 1.65 1.65 0 0 0 1.82-.33l.06-.06a2 2 0 0 1 2.83 0 2 2 0 0 1 0 2.83l-.06.06a1.65 1.65 0 0 0-.33 1.82V9a1.65 1.65 0 0 0 1.51 1H21a2 2 0 0 1 0 4h-.09a1.65 1.65 0 0 0-1.51 1z"
-          />
-        </svg>
-        <span class="hidden sm:inline">Advanced</span>
-      </a>
-    {/if}
   </nav>
 </header>

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -1,565 +0,0 @@
-<script lang="ts">
-  import { onMount, onDestroy } from "svelte";
-  import FamilyLogos from "$lib/components/FamilyLogos.svelte";
-  import {
-    instances,
-    instanceLinks,
-    nodeIdentities,
-    refreshState,
-    createInstanceLink,
-    updateInstanceLink,
-    deleteInstanceLink,
-    type Instance,
-  } from "$lib/stores/app.svelte";
-  import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
-
-  type InstanceWrapper = {
-    MlxRingInstance?: Instance;
-    MlxJacclInstance?: Instance;
-    VllmInstance?: Instance;
-  };
-
-  let interval: ReturnType<typeof setInterval> | null = null;
-
-  onMount(() => {
-    refreshState();
-    interval = setInterval(refreshState, 3000);
-  });
-  onDestroy(() => {
-    if (interval) clearInterval(interval);
-  });
-
-  type InstanceRow = {
-    id: string;
-    modelId: string;
-    family: string;
-    baseModel: string;
-    nodeNames: string[];
-    nodeCount: number;
-  };
-
-  const instanceRows = $derived.by<InstanceRow[]>(() => {
-    const rows: InstanceRow[] = [];
-    const ids = nodeIdentities();
-    for (const [id, raw] of Object.entries(instances())) {
-      const wrapper = raw as InstanceWrapper;
-      const inst =
-        wrapper.MlxRingInstance ??
-        wrapper.MlxJacclInstance ??
-        wrapper.VllmInstance;
-      const modelId = inst?.shardAssignments?.modelId ?? "";
-      const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
-      const nodeIds = Object.keys(nodeToRunner);
-      const nodeNames = nodeIds
-        .map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
-        .filter((name) => !!name);
-      rows.push({
-        id,
-        modelId,
-        family: deriveFamily(modelId),
-        baseModel: deriveBaseModel(modelId),
-        nodeNames,
-        nodeCount: nodeIds.length,
-      });
-    }
-    rows.sort((a, b) => a.modelId.localeCompare(b.modelId));
-    return rows;
-  });
-
-  const instanceById = $derived(
-    Object.fromEntries(instanceRows.map((r) => [r.id, r])),
-  );
-
-  type LinkRow = {
-    linkId: string;
-    prefill: string[];
-    decode: string[];
-    families: string[];
-    multiNode: boolean;
-  };
-
-  const linkRows = $derived.by<LinkRow[]>(() => {
-    const rows: LinkRow[] = [];
-    for (const [, link] of Object.entries(instanceLinks())) {
-      const fams = new Set<string>();
-      let multiNode = false;
-      for (const id of [...link.prefillInstances, ...link.decodeInstances]) {
-        const r = instanceById[id];
-        if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
-        if (r && r.nodeCount > 1) multiNode = true;
-      }
-      rows.push({
-        linkId: link.linkId,
-        prefill: link.prefillInstances,
-        decode: link.decodeInstances,
-        families: Array.from(fams),
-        multiNode,
-      });
-    }
-    return rows;
-  });
-
-  let editingLinkId = $state<string | null>(null);
-  let editingPrefill = $state<Set<string>>(new Set());
-  let editingDecode = $state<Set<string>>(new Set());
-  let saving = $state(false);
-  let errorMessage = $state<string | null>(null);
-
-  function startCreate() {
-    editingLinkId = "new";
-    editingPrefill = new Set();
-    editingDecode = new Set();
-    errorMessage = null;
-  }
-
-  function startEdit(row: LinkRow) {
-    editingLinkId = row.linkId;
-    editingPrefill = new Set(row.prefill);
-    editingDecode = new Set(row.decode);
-    errorMessage = null;
-  }
-
-  function cancelEdit() {
-    editingLinkId = null;
-    editingPrefill = new Set();
-    editingDecode = new Set();
-    errorMessage = null;
-  }
-
-  type Role = "prefill" | "decode" | "none";
-
-  function roleOf(id: string): Role {
-    if (editingPrefill.has(id)) return "prefill";
-    if (editingDecode.has(id)) return "decode";
-    return "none";
-  }
-
-  function setRole(id: string, role: Role) {
-    const p = new Set(editingPrefill);
-    const d = new Set(editingDecode);
-    p.delete(id);
-    d.delete(id);
-    if (role === "prefill") p.add(id);
-    if (role === "decode") d.add(id);
-    editingPrefill = p;
-    editingDecode = d;
-  }
-
-  const editingFamilies = $derived.by<string[]>(() => {
-    const fams = new Set<string>();
-    for (const id of [...editingPrefill, ...editingDecode]) {
-      const r = instanceById[id];
-      if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
-    }
-    return Array.from(fams);
-  });
-
-  const editingMultiNode = $derived.by<string[]>(() => {
-    const names: string[] = [];
-    for (const id of [...editingPrefill, ...editingDecode]) {
-      const r = instanceById[id];
-      if (r && r.nodeCount > 1) {
-        names.push(r.baseModel || r.modelId);
-      }
-    }
-    return names;
-  });
-
-  const editingMismatch = $derived(editingFamilies.length > 1);
-  const canSave = $derived(
-    editingLinkId !== null &&
-      editingPrefill.size > 0 &&
-      editingDecode.size > 0 &&
-      !saving,
-  );
-
-  async function save() {
-    if (editingLinkId === null) return;
-    saving = true;
-    errorMessage = null;
-    try {
-      const prefill = Array.from(editingPrefill);
-      const decode = Array.from(editingDecode);
-      if (editingLinkId === "new") {
-        await createInstanceLink(prefill, decode);
-      } else {
-        await updateInstanceLink(editingLinkId, prefill, decode);
-      }
-      cancelEdit();
-      await refreshState();
-    } catch (err) {
-      errorMessage = err instanceof Error ? err.message : String(err);
-    } finally {
-      saving = false;
-    }
-  }
-
-  async function remove(linkId: string) {
-    if (!confirm("Remove this routing?")) return;
-    try {
-      await deleteInstanceLink(linkId);
-      if (editingLinkId === linkId) cancelEdit();
-      await refreshState();
-    } catch (err) {
-      errorMessage = err instanceof Error ? err.message : String(err);
-    }
-  }
-</script>
-
-<div class="font-mono text-foreground">
-  <div class="mb-6 space-y-4">
-    <details open class="group [&_summary::-webkit-details-marker]:hidden">
-      <summary
-        class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
-      >
-        <span
-          class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
-          >▶</span
-        >
-        Prefill vs Decode
-      </summary>
-      <div class="mt-2 text-white/80 text-sm leading-relaxed">
-        Prefill is the compute-heavy pass that consumes the entire prompt and
-        builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
-        tokens sequentially from that cache. The two phases have very different
-        bottlenecks, so running them on different hardware can be substantially
-        faster than doing both on one node.
-      </div>
-    </details>
-    <details class="group [&_summary::-webkit-details-marker]:hidden">
-      <summary
-        class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
-      >
-        <span
-          class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
-          >▶</span
-        >
-        Linking Instances
-      </summary>
-      <div class="mt-2 text-white/80 text-sm leading-relaxed space-y-2">
-        <p>
-          A linked route here tells the cluster: when a request is sent to a
-          model in that cluster, the decode node (or the least active one if
-          there are multiple) will handle it. If it decides it must do a lot of
-          prefill not already cached in the prefix cache, it routes the request
-          to the prefill node over TCP IP. The prefill node streams the KV cache
-          back to the decode node which picks up from there.
-        </p>
-        <p>
-          Linked instances must be running the same model family — KV layouts
-          differ across architectures. More on the <a
-            class="text-exo-yellow underline underline-offset-2 hover:text-exo-yellow-darker transition-colors"
-            href="https://blog.exolabs.net/nvidia-dgx-spark/"
-            target="_blank"
-            rel="noreferrer noopener">blog</a
-          >.
-        </p>
-      </div>
-    </details>
-  </div>
-
-  {#if errorMessage}
-    <div
-      class="mb-4 px-4 py-3 bg-red-500/10 border border-red-500/40 text-red-300 text-sm"
-    >
-      {errorMessage}
-    </div>
-  {/if}
-
-  <section class="mt-12">
-    <h2
-      class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
-    >
-      Existing routes
-    </h2>
-
-    {#if linkRows.length === 0}
-      {#if editingLinkId === null}
-        <div class="flex items-center justify-between">
-          <p class="text-exo-light-gray italic text-sm m-0">
-            No routes yet. Create one to enable remote prefill.
-          </p>
-          <button
-            class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
-            onclick={startCreate}
-          >
-            + New route
-          </button>
-        </div>
-      {/if}
-    {:else}
-      {#if editingLinkId === null}
-        <div class="flex justify-end mb-3">
-          <button
-            class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
-            onclick={startCreate}
-          >
-            + New route
-          </button>
-        </div>
-      {/if}
-      <div
-        class="bg-exo-dark-gray/60 border border-exo-medium-gray/40 flex flex-col"
-      >
-        {#each linkRows as row (row.linkId)}
-          {#if editingLinkId !== row.linkId}
-            <article
-              class="p-4 border-b border-exo-light-gray/25 last:border-b-0"
-            >
-              {#if row.multiNode}
-                <div
-                  class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
-                >
-                  ⚠ Multi-node instance detected. Remote prefill currently only
-                  works on single-node (rank-0) instances. This route will not
-                  function until that's supported.
-                </div>
-              {/if}
-              {#if row.families.length > 1}
-                <div
-                  class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
-                >
-                  ⚠ Mixed model families: {row.families.join(", ")}
-                </div>
-              {/if}
-              <div
-                class="grid grid-cols-[1fr_auto_1fr_auto] items-center gap-x-3 gap-y-2"
-              >
-                <span
-                  class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-yellow/15 border border-exo-yellow/40 text-exo-yellow"
-                  >Prefill</span
-                >
-                <span></span>
-                <span
-                  class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-medium-gray/40 border border-exo-medium-gray/60 text-foreground"
-                  >Decode</span
-                >
-                <span></span>
-                <div class="min-w-0">
-                  <ul class="list-none p-0 m-0 flex flex-col gap-2">
-                    {#each row.prefill as id (id)}
-                      {@const r = instanceById[id]}
-                      {#if r}
-                        <li
-                          class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
-                        >
-                          <FamilyLogos family={r.family} />
-                          <div class="min-w-0 flex-1">
-                            <div
-                              class="text-exo-yellow text-xs font-mono truncate"
-                            >
-                              {r.baseModel || r.modelId}
-                            </div>
-                            <div
-                              class="text-exo-light-gray text-[11px] truncate"
-                            >
-                              {r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
-                                ? ` (${r.nodeCount} nodes)`
-                                : ""}
-                            </div>
-                            <div
-                              class="text-exo-light-gray/40 text-[10px] font-mono truncate"
-                              title={r.id}
-                            >
-                              {r.id.slice(0, 8)}
-                            </div>
-                          </div>
-                        </li>
-                      {/if}
-                    {/each}
-                  </ul>
-                </div>
-                <div class="text-exo-yellow/60 text-xl px-2" aria-hidden="true">
-                  →
-                </div>
-                <div class="min-w-0">
-                  <ul class="list-none p-0 m-0 flex flex-col gap-2">
-                    {#each row.decode as id (id)}
-                      {@const r = instanceById[id]}
-                      {#if r}
-                        <li
-                          class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
-                        >
-                          <FamilyLogos family={r.family} />
-                          <div class="min-w-0 flex-1">
-                            <div
-                              class="text-exo-yellow text-xs font-mono truncate"
-                            >
-                              {r.baseModel || r.modelId}
-                            </div>
-                            <div
-                              class="text-exo-light-gray text-[11px] truncate"
-                            >
-                              {r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
-                                ? ` (${r.nodeCount} nodes)`
-                                : ""}
-                            </div>
-                            <div
-                              class="text-exo-light-gray/40 text-[10px] font-mono truncate"
-                              title={r.id}
-                            >
-                              {r.id.slice(0, 8)}
-                            </div>
-                          </div>
-                        </li>
-                      {/if}
-                    {/each}
-                  </ul>
-                </div>
-                <div class="flex gap-2 pl-3">
-                  <button
-                    class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 hover:text-exo-yellow disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
-                    onclick={() => startEdit(row)}
-                    disabled={editingLinkId !== null}
-                  >
-                    Edit
-                  </button>
-                  <button
-                    class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-red-500/15 border border-red-500/40 rounded text-red-300 hover:bg-red-500/25 transition-colors"
-                    onclick={() => remove(row.linkId)}
-                  >
-                    Remove
-                  </button>
-                </div>
-              </div>
-            </article>
-          {/if}
-        {/each}
-      </div>
-    {/if}
-  </section>
-
-  {#if editingLinkId !== null && instanceRows.length === 0}
-    <section
-      class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 px-4 py-2.5 flex items-center justify-between gap-3"
-    >
-      <span class="text-exo-light-gray italic text-sm font-mono"
-        >No instances available.</span
-      >
-      <button
-        class="px-3 py-1 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 transition-colors"
-        onclick={cancelEdit}
-      >
-        Cancel
-      </button>
-    </section>
-  {:else if editingLinkId !== null}
-    <section class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 p-5">
-      <h2
-        class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
-      >
-        {editingLinkId === "new" ? "New route" : "Edit route"}
-      </h2>
-
-      {#if editingMismatch}
-        <div
-          class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
-        >
-          ⚠ Selected instances span multiple model families: <strong
-            >{editingFamilies.join(", ")}</strong
-          >. Linking across families produces a corrupt KV cache.
-        </div>
-      {/if}
-
-      {#if editingMultiNode.length > 0}
-        <div
-          class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
-        >
-          ⚠ Multi-node instance(s) selected: <strong
-            >{editingMultiNode.join(", ")}</strong
-          >. Remote prefill currently only works on single-node instances. This
-          route will not function until multi-node support lands.
-        </div>
-      {/if}
-
-      <p class="text-exo-light-gray text-xs mb-4">
-        Pick a role for each instance:
-        <span class="text-exo-yellow">Prefill</span>
-        serves KV cache,
-        <span class="text-foreground">Decode</span> consumes it.
-      </p>
-      <div
-        class="grid gap-2.5"
-        style="grid-template-columns: repeat(auto-fill, minmax(360px, 1fr));"
-      >
-        {#each instanceRows as row (row.id)}
-          {@const role = roleOf(row.id)}
-          <div
-            class="border p-3 flex flex-col gap-2.5 transition-colors {role ===
-            'prefill'
-              ? 'border-exo-yellow/60 bg-exo-dark-gray/60'
-              : role === 'decode'
-                ? 'border-exo-light-gray/60 bg-exo-dark-gray/60'
-                : 'border-exo-medium-gray/40 bg-exo-dark-gray/40'}"
-          >
-            <div class="flex items-center gap-2">
-              <FamilyLogos family={row.family} />
-              <div class="min-w-0 flex-1">
-                <div class="text-exo-yellow text-xs font-mono truncate">
-                  {row.baseModel || row.modelId}
-                </div>
-                <div class="text-exo-light-gray text-[11px] truncate">
-                  {row.nodeNames.join(", ") || "?"}{row.nodeCount > 1
-                    ? ` (${row.nodeCount} nodes)`
-                    : ""}
-                </div>
-                <div
-                  class="text-exo-light-gray/40 text-[10px] font-mono truncate"
-                  title={row.id}
-                >
-                  {row.id.slice(0, 8)}
-                </div>
-              </div>
-              {#if row.nodeCount > 1}
-                <span
-                  class="text-[9px] font-mono tracking-widest uppercase px-1.5 py-0.5 bg-red-500/15 border border-red-500/40 text-red-300"
-                  title="Multi-node instances are not supported by remote prefill yet."
-                  >Unsupported</span
-                >
-              {/if}
-            </div>
-            <div
-              class="flex rounded-md overflow-hidden border border-exo-light-gray/40 divide-x divide-exo-light-gray/40"
-            >
-              <button
-                class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
-                'prefill'
-                  ? 'bg-exo-yellow/20 text-exo-yellow'
-                  : 'bg-transparent text-white/80 hover:text-exo-yellow'}"
-                onclick={() =>
-                  setRole(row.id, role === "prefill" ? "none" : "prefill")}
-                >Prefill</button
-              >
-              <button
-                class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
-                'decode'
-                  ? 'bg-exo-medium-gray/50 text-foreground'
-                  : 'bg-transparent text-white/80 hover:text-foreground'}"
-                onclick={() =>
-                  setRole(row.id, role === "decode" ? "none" : "decode")}
-                >Decode</button
-              >
-            </div>
-          </div>
-        {/each}
-      </div>
-
-      <div class="flex gap-2 mt-5 justify-end">
-        <button
-          class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
-          onclick={save}
-          disabled={!canSave}
-        >
-          {saving ? "Saving..." : "Save route"}
-        </button>
-        <button
-          class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 text-foreground hover:border-exo-yellow/60 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
-          onclick={cancelEdit}
-          disabled={saving}
-        >
-          Cancel
-        </button>
-      </div>
-    </section>
-  {/if}
-</div>

dashboard/src/lib/stores/app.svelte.ts

@@ -74,12 +74,6 @@ export interface Instance {
   };
 }
 
-export interface RawInstanceLink {
-  linkId: string;
-  prefillInstances: string[];
-  decodeInstances: string[];
-}
-
 // Granular node state types from the new state structure
 interface RawNodeIdentity {
   modelId?: string;
@@ -229,7 +223,6 @@ interface RawStateResponse {
     }
   >;
   runners?: Record<string, unknown>;
-  instanceLinks?: Record<string, RawInstanceLink>;
   downloads?: Record<string, unknown[]>;
   // New granular node state fields
   nodeIdentities?: Record<string, RawNodeIdentity>;
@@ -548,8 +541,6 @@ class AppStore {
   topologyData = $state<TopologyData | null>(null);
   instances = $state<Record<string, unknown>>({});
   runners = $state<Record<string, unknown>>({});
-  instanceLinks = $state<Record<string, RawInstanceLink>>({});
-  featureFlags = $state<Record<string, boolean>>({});
   downloads = $state<Record<string, unknown[]>>({});
   nodeDisk = $state<
     Record<
@@ -1283,7 +1274,6 @@ class AppStore {
 
   startPolling() {
     this.fetchState();
-    this.fetchFeatureFlags();
     this.fetchInterval = setInterval(() => this.fetchState(), 1000);
   }
 
@@ -1295,16 +1285,6 @@ class AppStore {
     this.stopPreviewsPolling();
   }
 
-  async fetchFeatureFlags() {
-    try {
-      const response = await fetch("/v1/feature-flags");
-      if (!response.ok) return;
-      this.featureFlags = await response.json();
-    } catch {
-      // Silently ignore — defaults to all-disabled.
-    }
-  }
-
   async fetchState() {
     try {
       const response = await fetch("/state");
@@ -1330,11 +1310,6 @@ class AppStore {
       if (data.runners) {
         this.runners = data.runners;
       }
-      if (data.instanceLinks) {
-        this.instanceLinks = data.instanceLinks;
-      } else {
-        this.instanceLinks = {};
-      }
       if (data.downloads) {
         this.downloads = data.downloads;
       }
@@ -1695,15 +1670,7 @@ class AppStore {
               }
             }
           }
-          const out: {
-            role: string;
-            content: string;
-            reasoning_content?: string;
-          } = { role: m.role, content: msgContent };
-          if (m.role === "assistant" && m.thinking) {
-            out.reasoning_content = m.thinking;
-          }
-          return out;
+          return { role: m.role, content: msgContent };
         }),
       ];
 
@@ -1910,15 +1877,7 @@ class AppStore {
       const apiMessages = [
         systemPrompt,
         ...targetConversation.messages.slice(0, -1).map((m) => {
-          const out: {
-            role: string;
-            content: string;
-            reasoning_content?: string;
-          } = { role: m.role, content: m.content };
-          if (m.role === "assistant" && m.thinking) {
-            out.reasoning_content = m.thinking;
-          }
-          return out;
+          return { role: m.role, content: m.content };
         }),
       ];
 
@@ -2449,15 +2408,10 @@ class AppStore {
               contentParts.push({ type: "text", text: textContent });
             }
 
-            const out: {
-              role: string;
-              content: typeof contentParts;
-              reasoning_content?: string;
-            } = { role: m.role, content: contentParts };
-            if (m.role === "assistant" && m.thinking) {
-              out.reasoning_content = m.thinking;
-            }
-            return out;
+            return {
+              role: m.role,
+              content: contentParts,
+            };
           }
 
           // Text-only message (original path)
@@ -2475,15 +2429,10 @@ class AppStore {
             }
           }
 
-          const out: {
-            role: string;
-            content: string;
-            reasoning_content?: string;
-          } = { role: m.role, content: msgContent };
-          if (m.role === "assistant" && m.thinking) {
-            out.reasoning_content = m.thinking;
-          }
-          return out;
+          return {
+            role: m.role,
+            content: msgContent,
+          };
         }),
       ];
 
@@ -3332,60 +3281,6 @@ class AppStore {
     }
   }
 
-  async createInstanceLink(
-    prefillInstances: string[],
-    decodeInstances: string[],
-  ): Promise<void> {
-    const response = await fetch("/v1/instance-links", {
-      method: "POST",
-      headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({
-        prefill_instances: prefillInstances,
-        decode_instances: decodeInstances,
-      }),
-    });
-    if (!response.ok) {
-      throw new Error(
-        `Failed to create instance link: ${response.status} ${await response.text()}`,
-      );
-    }
-  }
-
-  async updateInstanceLink(
-    linkId: string,
-    prefillInstances: string[],
-    decodeInstances: string[],
-  ): Promise<void> {
-    const response = await fetch(
-      `/v1/instance-links/${encodeURIComponent(linkId)}`,
-      {
-        method: "PUT",
-        headers: { "Content-Type": "application/json" },
-        body: JSON.stringify({
-          prefill_instances: prefillInstances,
-          decode_instances: decodeInstances,
-        }),
-      },
-    );
-    if (!response.ok) {
-      throw new Error(
-        `Failed to update instance link: ${response.status} ${await response.text()}`,
-      );
-    }
-  }
-
-  async deleteInstanceLink(linkId: string): Promise<void> {
-    const response = await fetch(
-      `/v1/instance-links/${encodeURIComponent(linkId)}`,
-      { method: "DELETE" },
-    );
-    if (!response.ok) {
-      throw new Error(
-        `Failed to delete instance link: ${response.status} ${await response.text()}`,
-      );
-    }
-  }
-
   /**
    * Delete a downloaded model from a specific node
    */
@@ -3484,19 +3379,6 @@ export const prefillProgress = () => appStore.prefillProgress;
 export const topologyData = () => appStore.topologyData;
 export const instances = () => appStore.instances;
 export const runners = () => appStore.runners;
-export const instanceLinks = () => appStore.instanceLinks;
-export const featureFlags = () => appStore.featureFlags;
-export const createInstanceLink = (
-  prefillInstances: string[],
-  decodeInstances: string[],
-) => appStore.createInstanceLink(prefillInstances, decodeInstances);
-export const updateInstanceLink = (
-  linkId: string,
-  prefillInstances: string[],
-  decodeInstances: string[],
-) => appStore.updateInstanceLink(linkId, prefillInstances, decodeInstances);
-export const deleteInstanceLink = (linkId: string) =>
-  appStore.deleteInstanceLink(linkId);
 export const downloads = () => appStore.downloads;
 export const nodeDisk = () => appStore.nodeDisk;
 export const placementPreviews = () => appStore.placementPreviews;

dashboard/src/lib/utils/model_family.ts

@@ -1,44 +0,0 @@
-// Mirrors src/exo/shared/models/model_cards.py:derive_base_model
-const QUANT_SUFFIXES = new RegExp(
-  "[-_ ](?:MLX|MXFP[0-9]+|NVFP[0-9]+|GPTQ|AWQ|GGUF|fp16|bf16|fp8|int[0-9]+|[0-9]+(?:\\.[0-9]+)?bit|Q[0-9]+(?:_[A-Z0-9]+)?|gs[0-9]+)" +
-    "(?:[-_ ](?:MLX|Q[0-9]+|Int[0-9]+|[A-Z0-9]+|gs[0-9]+))*$",
-  "i",
-);
-
-function normalize(s: string): string {
-  return s
-    .replaceAll("-", " ")
-    .replaceAll("_", " ")
-    .replaceAll("  ", " ")
-    .trim();
-}
-
-export function deriveBaseModel(modelId: string): string {
-  const short = modelId.includes("/")
-    ? (modelId.split("/").pop() ?? modelId)
-    : modelId;
-  const stripped = short.replace(QUANT_SUFFIXES, "");
-  return normalize(stripped);
-}
-
-export function baseModelsCompatible(a: string, b: string): boolean {
-  return deriveBaseModel(a).toLowerCase() === deriveBaseModel(b).toLowerCase();
-}
-
-// Mirrors src/exo/shared/models/model_cards.py:derive_family
-export function deriveFamily(modelId: string): string {
-  const short = modelId.includes("/")
-    ? (modelId.split("/").pop() ?? modelId)
-    : modelId;
-  const stripped = short
-    .replace(QUANT_SUFFIXES, "")
-    .toLowerCase()
-    .replaceAll("_", "-");
-  const parts = stripped.split(/[-.]/);
-  const familyParts: string[] = [];
-  for (const p of parts) {
-    if (/^\d+$/.test(p) || /^\d+[bm]?$/i.test(p)) break;
-    familyParts.push(p);
-  }
-  return familyParts.length > 0 ? familyParts.join("-") : stripped;
-}

dashboard/src/routes/+page.svelte

@@ -3435,7 +3435,6 @@
             >
               <li>Connect nodes with TB5 cables</li>
               <li>Boot to Recovery (hold power 10s → Options)</li>
-              <li>Open Terminal from the Utilities menu</li>
               <li>
                 Run
                 <code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4823,7 +4822,6 @@
                 >
                   <li>Connect nodes with TB5 cables</li>
                   <li>Boot to Recovery (hold power 10s → Options)</li>
-                  <li>Open Terminal from the Utilities menu</li>
                   <li>
                     Run
                     <code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4970,7 +4968,6 @@
                   >
                     <li>Connect nodes with TB5 cables</li>
                     <li>Boot to Recovery (hold power 10s → Options)</li>
-                    <li>Open Terminal from the Utilities menu</li>
                     <li>
                       Run
                       <code

dashboard/src/routes/advanced/+page.svelte

@@ -1,81 +0,0 @@
-<script lang="ts">
-  import { browser } from "$app/environment";
-  import HeaderNav from "$lib/components/HeaderNav.svelte";
-  import PrefillDecodeDisaggregation from "$lib/components/PrefillDecodeDisaggregation.svelte";
-  import { featureFlags, refreshState } from "$lib/stores/app.svelte";
-  import { onMount } from "svelte";
-
-  type TabId = "prefill-decode";
-
-  const tabs: { id: TabId; label: string }[] = [
-    { id: "prefill-decode", label: "Prefill / Decode" },
-  ];
-
-  let activeTab = $state<TabId>(tabs[0].id);
-  let flagsLoaded = $state(false);
-
-  onMount(() => {
-    refreshState().finally(() => {
-      flagsLoaded = true;
-    });
-  });
-
-  const flags = $derived(featureFlags());
-  const enabled = $derived(flags["disaggregation"] === true);
-
-  $effect(() => {
-    if (browser && flagsLoaded && !enabled) {
-      // No advanced features enabled — bounce home.
-      window.location.hash = "/";
-    }
-  });
-</script>
-
-<div class="min-h-screen bg-exo-dark-gray flex flex-col">
-  <HeaderNav />
-
-  <main class="flex-1 max-w-[1100px] mx-auto w-full px-4 md:px-6 py-8">
-    {#if !flagsLoaded}
-      <div class="text-exo-light-gray/60 text-sm">Loading…</div>
-    {:else if !enabled}
-      <div class="text-exo-light-gray/60 text-sm">
-        No advanced features enabled. Set <code
-          class="text-exo-yellow font-mono">ENABLE_DISAGGREGATION=true</code
-        > on the cluster to access prefill/decode disaggregation.
-      </div>
-    {:else}
-      <div class="mb-4">
-        <h1
-          class="text-white text-xl md:text-2xl font-semibold tracking-wide mb-2"
-        >
-          Advanced
-        </h1>
-        <p class="text-exo-light-gray/60 text-sm">
-          Cluster-level configuration. Most users don't need anything here.
-        </p>
-      </div>
-
-      <div
-        class="flex flex-wrap gap-2 mb-6 border-b border-exo-light-gray/10 pb-3"
-      >
-        {#each tabs as tab (tab.id)}
-          <button
-            onclick={() => (activeTab = tab.id)}
-            class="px-3 py-1.5 text-xs rounded-md transition-all cursor-pointer
-              {activeTab === tab.id
-              ? 'bg-exo-yellow/15 text-exo-yellow border border-exo-yellow/30'
-              : 'text-exo-light-gray/60 hover:text-white/80 border border-transparent hover:border-exo-light-gray/20'}"
-          >
-            {tab.label}
-          </button>
-        {/each}
-      </div>
-
-      <div class="space-y-4">
-        {#if activeTab === "prefill-decode"}
-          <PrefillDecodeDisaggregation />
-        {/if}
-      </div>
-    {/if}
-  </main>
-</div>

dashboard/src/routes/integrations/+page.svelte

@@ -14,7 +14,6 @@
 
   let modelCapabilities = $state<Record<string, string[]>>({});
   let modelContextLengths = $state<Record<string, number>>({});
-  let modelReasoningDialects = $state<Record<string, string>>({});
 
   const runningModels = $derived.by(() => {
     const models: string[] = [];
@@ -133,7 +132,6 @@
     for (const modelId of runningModels) {
       const caps = modelCapabilities[modelId] || [];
       const ctxLen = modelContextLengths[modelId] || 0;
-      const dialect = modelReasoningDialects[modelId];
       const entry: Record<string, unknown> = { name: modelId };
       if (ctxLen > 0) {
         entry.limit = { context: ctxLen, output: Math.min(ctxLen, 16384) };
@@ -141,27 +139,6 @@
       if (caps.includes("vision")) {
         entry.modalities = { input: ["text", "image"], output: ["text"] };
       }
-      // Reasoning round-trip: opencode's `interleaved` field tells the
-      // openai-compatible adapter to send the assistant's prior
-      // reasoning_content back in subsequent turns. Emit it for dialects
-      // whose chat templates use prior reasoning:
-      //   - `tool_conditional` (DeepSeek V3.2 / V4): wrapper preserves all
-      //     reasoning when tools are present.
-      //   - `post_last_user` (Qwen3-Thinking, GLM 4.5+, MiniMax M2.x):
-      //     Jinja template reads reasoning_content for assistant turns since
-      //     the last user message — exactly the tool-chain window.
-      //   - `channel` (gpt-oss / Harmony): the model's Jinja template reads
-      //     `message.thinking` rather than `message.reasoning_content`, but
-      //     the server bridges `reasoning_content` → `thinking` before
-      //     rendering, so the round-trip works through the standard field.
-      // `suffix` (Kimi): reasoning lives in content; no separate field path.
-      if (
-        dialect === "tool_conditional" ||
-        dialect === "post_last_user" ||
-        dialect === "channel"
-      ) {
-        entry.interleaved = { field: "reasoning_content" };
-      }
       models[modelId] = entry;
     }
     if (Object.keys(models).length === 0) {
@@ -373,25 +350,16 @@
     try {
       const resp = await fetch("/v1/models");
       const data = (await resp.json()) as {
-        data: {
-          id: string;
-          capabilities: string[];
-          context_length: number;
-          reasoning_dialect?: string;
-        }[];
+        data: { id: string; capabilities: string[]; context_length: number }[];
       };
       const caps: Record<string, string[]> = {};
       const ctxs: Record<string, number> = {};
-      const dialects: Record<string, string> = {};
       for (const model of data.data) {
         caps[model.id] = model.capabilities || [];
         if (model.context_length > 0) ctxs[model.id] = model.context_length;
-        if (model.reasoning_dialect)
-          dialects[model.id] = model.reasoning_dialect;
       }
       modelCapabilities = caps;
       modelContextLengths = ctxs;
-      modelReasoningDialects = dialects;
     } catch {
       /* ignore */
     }

pyproject.toml

@@ -15,11 +15,11 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo-pyo3-bindings",                         # rust bindings
+  "exo-pyo3-bindings",                       # rust bindings
   "anyio==4.11.0",
   "mlx==0.31.2; sys_platform == 'darwin'",
   "mlx-lm; sys_platform=='darwin'",
-  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
@@ -28,12 +28,13 @@ dependencies = [
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
-  "mlx-vlm>=0.3.11; sys_platform == 'darwin'",
+  "mlx-vlm>=0.3.11",
   "transformers>=5.6.2",
 ]
 
 [project.scripts]
 exo = "exo.main:main"
+exo-reasoning-proxy = "exo.reasoning_proxy.main:main"
 
 # dependencies only required for development
 [dependency-groups]
@@ -52,21 +53,18 @@ cpu = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cpu==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
-  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda12 = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
-  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda13 = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
-  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 

resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "4bit"
 base_model = "DeepSeek V3.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
 
 context_length = 131072
 

resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V3.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "post_last_user"
 
 context_length = 131072
 

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-4bit.toml

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "4bit"
 base_model = "DeepSeek V3.2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
 
 context_length = 131072
 

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-8bit.toml

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V3.2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
 
 context_length = 131072
 

resources/inference_model_cards/mlx-community--DeepSeek-V4-Flash.toml

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V4 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
 
 context_length = 1048576
 

resources/inference_model_cards/mlx-community--DeepSeek-V4-Pro.toml

@@ -8,7 +8,6 @@ family = "deepseek"
 quantization = "8bit"
 base_model = "DeepSeek V4 Pro"
 capabilities = ["text", "thinking", "thinking_toggle"]
-reasoning_dialect = "tool_conditional"
 
 context_length = 1048576
 

src/exo/api/adapters/chat_completions.py

@@ -131,13 +131,9 @@ async def chat_request_to_text_generation(
                         multimodal_content.append({"type": "text", "text": part.text})
                     else:
                         multimodal_content.append({"type": "image"})
-                multimodal_msg: dict[str, Any] = {
-                    "role": msg.role,
-                    "content": multimodal_content,
-                }
-                if msg.reasoning_content is not None:
-                    multimodal_msg["reasoning_content"] = msg.reasoning_content
-                chat_template_messages.append(multimodal_msg)
+                chat_template_messages.append(
+                    {"role": msg.role, "content": multimodal_content}
+                )
                 continue
             msg_copy = msg.model_copy(update={"content": content})
 
@@ -172,8 +168,6 @@ async def chat_request_to_text_generation(
         min_p=request.min_p,
         repetition_penalty=request.repetition_penalty,
         repetition_context_size=request.repetition_context_size,
-        presence_penalty=request.presence_penalty,
-        frequency_penalty=request.frequency_penalty,
         images=images,
     )
 

src/exo/api/main.py

@@ -79,8 +79,6 @@ from exo.api.types import (
     ImageListItem,
     ImageListResponse,
     ImageSize,
-    InstanceLinkBody,
-    InstanceLinkResponse,
     ModelList,
     ModelListModel,
     PlaceInstanceParams,
@@ -121,12 +119,9 @@ from exo.api.types.openai_responses import (
 )
 from exo.master.image_store import ImageStore
 from exo.master.placement import place_instance as get_instance_placements
-from exo.routing.event_router import EventRouter
-from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.apply import apply
 from exo.shared.constants import (
     DASHBOARD_DIR,
-    ENABLE_DISAGGREGATION,
     EXO_CACHE_HOME,
     EXO_EVENT_LOG_DIR,
     EXO_IMAGE_CACHE_DIR,
@@ -159,22 +154,19 @@ from exo.shared.types.commands import (
     DeleteCustomModelCard,
     DeleteDownload,
     DeleteInstance,
-    DeleteInstanceLink,
     DownloadCommand,
     ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
-    RequestSnapshot,
     SendInputChunk,
-    SetInstanceLink,
     StartDownload,
     TaskCancelled,
     TaskFinished,
     TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SystemId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -182,9 +174,7 @@ from exo.shared.types.events import (
     InstanceDeleted,
     TracesMerged,
 )
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.memory import Memory
-from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -211,8 +201,6 @@ from exo.utils.task_group import TaskGroup
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
 ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
 
-_SNAPSHOT_FETCH_TIMEOUT_SECONDS = 30
-
 
 def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
     return f"image/{image_format or 'png'}"
@@ -227,27 +215,13 @@ def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
     return params
 
 
-def _require_disaggregation_enabled() -> None:
-    if not ENABLE_DISAGGREGATION:
-        raise HTTPException(
-            status_code=HTTPStatus.NOT_FOUND,
-            detail=(
-                "Prefill/decode disaggregation is disabled. "
-                "Set ENABLE_DISAGGREGATION=true to enable."
-            ),
-        )
-
-
 class API:
     def __init__(
         self,
         node_id: NodeId,
-        session_id: SessionId,
         *,
         port: int,
-        event_router: EventRouter,
         event_receiver: Receiver[IndexedEvent],
-        snapshot_chunk_receiver: Receiver[SnapshotChunk],
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
@@ -256,16 +230,14 @@ class API:
         self.state = State()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self._system_id = SystemId()
-        self.session_id = session_id
-        self.event_router = event_router
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
         self.event_receiver = event_receiver
-        self.snapshot_chunk_receiver = snapshot_chunk_receiver
         self.election_receiver = election_receiver
         self.node_id: NodeId = node_id
         self.last_completed_election: int = 0
         self.port = port
+        self._sent_image_hashes: set[str] = set()
 
         self.paused: bool = False
         self.paused_ev: anyio.Event = anyio.Event()
@@ -303,29 +275,19 @@ class API:
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup = TaskGroup()
 
-    def reset(
-        self,
-        result_clock: int,
-        session_id: SessionId,
-        event_router: EventRouter,
-        event_receiver: Receiver[IndexedEvent],
-        snapshot_chunk_receiver: Receiver[SnapshotChunk],
-    ):
+    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
         logger.info("Resetting API State")
         self._event_log.close()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self.state = State()
         self._system_id = SystemId()
-        self.session_id = session_id
-        self.event_router = event_router
         self._text_generation_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
         self.event_receiver.close()
         self.event_receiver = event_receiver
-        self.snapshot_chunk_receiver.close()
-        self.snapshot_chunk_receiver = snapshot_chunk_receiver
-        self._tg.start_soon(self._bootstrap_then_apply_state)
+        self._tg.start_soon(self._apply_state)
+        self._sent_image_hashes = set()
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -366,11 +328,6 @@ class API:
         self.app.get("/instance/previews")(self.get_placement_previews)
         self.app.get("/instance/{instance_id}")(self.get_instance)
         self.app.delete("/instance/{instance_id}")(self.delete_instance)
-        self.app.get("/v1/instance-links")(self.list_instance_links)
-        self.app.post("/v1/instance-links")(self.create_instance_link)
-        self.app.put("/v1/instance-links/{link_id}")(self.update_instance_link)
-        self.app.delete("/v1/instance-links/{link_id}")(self.delete_instance_link)
-        self.app.get("/v1/feature-flags")(self.get_feature_flags)
         self.app.get("/models")(self.get_models)
         self.app.get("/v1/models")(self.get_models)
         self.app.post("/models/add")(self.add_custom_model)
@@ -379,9 +336,7 @@ class API:
         self.app.post("/v1/chat/completions", response_model=None)(
             self.chat_completions
         )
-        self.app.post("/bench/chat/completions", response_model=None)(
-            self.bench_chat_completions
-        )
+        self.app.post("/bench/chat/completions")(self.bench_chat_completions)
         self.app.post("/v1/images/generations", response_model=None)(
             self.image_generations
         )
@@ -660,49 +615,6 @@ class API:
             instance_id=instance_id,
         )
 
-    async def get_feature_flags(self) -> dict[str, bool]:
-        return {"disaggregation": ENABLE_DISAGGREGATION}
-
-    async def list_instance_links(self) -> list[InstanceLink]:
-        if not ENABLE_DISAGGREGATION:
-            return []
-        return list(self.state.instance_links.values())
-
-    async def create_instance_link(
-        self, body: InstanceLinkBody
-    ) -> InstanceLinkResponse:
-        _require_disaggregation_enabled()
-        return await self._set_instance_link(InstanceLinkId(), body)
-
-    async def update_instance_link(
-        self, link_id: InstanceLinkId, body: InstanceLinkBody
-    ) -> InstanceLinkResponse:
-        _require_disaggregation_enabled()
-        return await self._set_instance_link(link_id, body)
-
-    async def _set_instance_link(
-        self, link_id: InstanceLinkId, body: InstanceLinkBody
-    ) -> InstanceLinkResponse:
-        command = SetInstanceLink(
-            link_id=link_id,
-            prefill_instances=list(body.prefill_instances),
-            decode_instances=list(body.decode_instances),
-        )
-        await self._send(command)
-        return InstanceLinkResponse(
-            message="Command received.", command_id=command.command_id
-        )
-
-    async def delete_instance_link(
-        self, link_id: InstanceLinkId
-    ) -> InstanceLinkResponse:
-        _require_disaggregation_enabled()
-        command = DeleteInstanceLink(link_id=link_id)
-        await self._send(command)
-        return InstanceLinkResponse(
-            message="Command received.", command_id=command.command_id
-        )
-
     async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
         """Cancel an active command by closing its stream and notifying workers."""
         sender = self._text_generation_queues.get(
@@ -847,8 +759,18 @@ class API:
         )
         command = TextGeneration(task_params=task_params)
 
+        new_images: list[tuple[int, str]] = []
+        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
+            if h not in self._sent_image_hashes:
+                self._sent_image_hashes.add(h)
+                new_images.append((idx, img))
+
+        if not new_images:
+            await self._send(command)
+            return command
+
         all_chunks: list[tuple[int, str]] = []
-        for img_idx, img_data in enumerate(images):
+        for img_idx, img_data in new_images:
             for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                 all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
 
@@ -907,7 +829,7 @@ class API:
 
     async def bench_chat_completions(
         self, payload: BenchChatCompletionRequest
-    ) -> BenchChatCompletionResponse | StreamingResponse:
+    ) -> BenchChatCompletionResponse:
         task_params = await chat_request_to_text_generation(payload)
         resolved_model = await self._resolve_and_validate_text_model(
             ModelId(task_params.model)
@@ -924,22 +846,6 @@ class API:
 
         command = await self._send_text_generation_with_images(task_params)
 
-        if payload.stream:
-            return StreamingResponse(
-                with_sse_keepalive(
-                    generate_chat_stream(
-                        command.command_id,
-                        self._token_chunk_stream(command.command_id),
-                    ),
-                ),
-                media_type="text/event-stream",
-                headers={
-                    "Cache-Control": "no-cache",
-                    "Connection": "close",
-                    "X-Accel-Buffering": "no",
-                },
-            )
-
         return await self._collect_text_generation_with_stats(command.command_id)
 
     async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
@@ -1859,7 +1765,7 @@ class API:
         try:
             async with self._tg as tg:
                 logger.info("Starting API")
-                tg.start_soon(self._bootstrap_then_apply_state)
+                tg.start_soon(self._apply_state)
                 tg.start_soon(self._pause_on_new_election)
                 tg.start_soon(self._cleanup_expired_images)
                 print_startup_banner(self.port)
@@ -1873,7 +1779,6 @@ class API:
             self._event_log.close()
             self.command_sender.close()
             self.event_receiver.close()
-            self.snapshot_chunk_receiver.close()
 
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
@@ -1889,43 +1794,9 @@ class API:
                 shutdown_trigger=ev.wait,
             )
 
-    async def _bootstrap_then_apply_state(self):
-        await self._fetch_snapshot()
-        await self._apply_state()
-
-    async def _fetch_snapshot(self) -> None:
-        receiver = SnapshotReceiver(self.node_id, self.session_id)
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=RequestSnapshot(requester_node_id=self.node_id),
-            )
-        )
-
-        with anyio.move_on_after(_SNAPSHOT_FETCH_TIMEOUT_SECONDS):
-            with self.snapshot_chunk_receiver as chunks:
-                async for chunk in chunks:
-                    received = receiver.ingest(chunk)
-                    if received is None:
-                        continue
-                    self.state = received.state
-                    self.event_router.set_buffer_start(
-                        received.last_event_applied_idx + 1
-                    )
-                    logger.info(
-                        f"API bootstrapped from snapshot at idx "
-                        f"{received.last_event_applied_idx}"
-                    )
-                    return
-        logger.info(
-            "API: no snapshot received before timeout; falling back to full event-log replay"
-        )
-
     async def _apply_state(self):
         with self.event_receiver as events:
             async for i_event in events:
-                if i_event.idx <= self.state.last_event_applied_idx:
-                    continue
                 self._event_log.append(i_event.event)
                 self.state = apply(self.state, i_event)
                 event = i_event.event

src/exo/api/tests/test_api_snapshot_bootstrap.py

@@ -1,136 +0,0 @@
-# pyright: reportPrivateUsage=false
-
-import hashlib
-
-import anyio
-import pytest
-import zstandard
-
-from exo.api.main import API
-from exo.routing.event_router import EventRouter
-from exo.shared.types.commands import ForwarderCommand, RequestSnapshot
-from exo.shared.types.common import NodeId, SessionId, SystemId
-from exo.shared.types.events import (
-    Event,
-    GlobalForwarderEvent,
-    IndexedEvent,
-    LocalForwarderEvent,
-    TestEvent,
-)
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-from exo.utils.channels import Receiver, Sender, channel
-
-
-class _FakeEventLog:
-    def __init__(self) -> None:
-        self.appended: list[Event] = []
-
-    def append(self, event: Event) -> None:
-        self.appended.append(event)
-
-
-def _snapshot_chunk(
-    state: State, *, requester_node_id: NodeId, session_id: SessionId
-) -> SnapshotChunk:
-    body = zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-    return SnapshotChunk.from_data(
-        data=body,
-        transfer_id=SnapshotTransferId("transfer-1"),
-        requester_node_id=requester_node_id,
-        session_id=session_id,
-        schema_version=state.schema_version,
-        last_event_applied_idx=state.last_event_applied_idx,
-        chunk_index=0,
-        total_chunks=1,
-        sha256_hex=hashlib.sha256(body).hexdigest(),
-    )
-
-
-def _api(
-    node_id: NodeId, session_id: SessionId
-) -> tuple[
-    API,
-    EventRouter,
-    Receiver[ForwarderCommand],
-    Sender[SnapshotChunk],
-    Sender[IndexedEvent],
-    _FakeEventLog,
-]:
-    router_command_sender, _router_command_receiver = channel[ForwarderCommand]()
-    _global_event_sender, global_event_receiver = channel[GlobalForwarderEvent]()
-    local_event_sender, _local_event_receiver = channel[LocalForwarderEvent]()
-    event_router = EventRouter(
-        session_id=session_id,
-        command_sender=router_command_sender,
-        external_inbound=global_event_receiver,
-        external_outbound=local_event_sender,
-    )
-
-    event_sender, event_receiver = channel[IndexedEvent]()
-    command_sender, command_receiver = channel[ForwarderCommand]()
-    snapshot_sender, snapshot_receiver = channel[SnapshotChunk]()
-
-    api = object.__new__(API)
-    api.node_id = node_id
-    api.session_id = session_id
-    api.event_router = event_router
-    api.event_receiver = event_receiver
-    api.snapshot_chunk_receiver = snapshot_receiver
-    api.command_sender = command_sender
-    api._system_id = SystemId("api-system")
-    api.state = State()
-    event_log = _FakeEventLog()
-    api._event_log = event_log  # pyright: ignore[reportAttributeAccessIssue]
-    api._image_generation_queues = {}
-    api._text_generation_queues = {}
-    return api, event_router, command_receiver, snapshot_sender, event_sender, event_log
-
-
-@pytest.mark.asyncio
-async def test_api_fetch_snapshot_applies_state_and_fast_forwards_router() -> None:
-    node_id = NodeId("api")
-    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
-    api, event_router, command_receiver, snapshot_sender, _event_sender, _event_log = (
-        _api(node_id, session_id)
-    )
-    state = State(last_event_applied_idx=7)
-
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(api._fetch_snapshot)
-        command = await command_receiver.receive()
-        assert isinstance(command.command, RequestSnapshot)
-        assert command.command.requester_node_id == node_id
-
-        await snapshot_sender.send(
-            _snapshot_chunk(state, requester_node_id=node_id, session_id=session_id)
-        )
-
-    assert api.state.last_event_applied_idx == 7
-    assert event_router.event_buffer.next_idx_to_release == 8
-
-
-@pytest.mark.asyncio
-async def test_api_apply_state_ignores_events_covered_by_snapshot() -> None:
-    node_id = NodeId("api")
-    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
-    (
-        api,
-        _event_router,
-        _command_receiver,
-        _snapshot_sender,
-        event_sender,
-        event_log,
-    ) = _api(node_id, session_id)
-    api.state = State(last_event_applied_idx=7)
-
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(api._apply_state)
-        await event_sender.send(IndexedEvent(idx=7, event=TestEvent()))
-        await event_sender.send(IndexedEvent(idx=8, event=TestEvent()))
-
-        while api.state.last_event_applied_idx != 8:
-            await anyio.sleep(0.001)
-        tg.cancel_scope.cancel()
-
-    assert len(event_log.appended) == 1

src/exo/api/types/__init__.py

@@ -34,8 +34,6 @@ from .api import ImageGenerationTaskParams as ImageGenerationTaskParams
 from .api import ImageListItem as ImageListItem
 from .api import ImageListResponse as ImageListResponse
 from .api import ImageSize as ImageSize
-from .api import InstanceLinkBody as InstanceLinkBody
-from .api import InstanceLinkResponse as InstanceLinkResponse
 from .api import Logprobs as Logprobs
 from .api import LogprobsContentItem as LogprobsContentItem
 from .api import ModelList as ModelList

src/exo/api/types/api.py

@@ -296,16 +296,6 @@ class CancelCommandResponse(BaseModel):
     command_id: CommandId
 
 
-class InstanceLinkBody(BaseModel):
-    prefill_instances: list[InstanceId]
-    decode_instances: list[InstanceId]
-
-
-class InstanceLinkResponse(BaseModel):
-    message: str
-    command_id: CommandId
-
-
 ImageSize = Literal[
     "auto",
     "512x512",

src/exo/download/download_utils.py

@@ -1,12 +1,11 @@
 import asyncio
 import hashlib
 import os
-import random
 import shutil
 import ssl
 import time
 import traceback
-from collections.abc import Awaitable, Mapping
+from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -56,36 +55,6 @@ class HuggingFaceAuthenticationError(Exception):
 class HuggingFaceRateLimitError(Exception):
     """429 Huggingface code"""
 
-    def __init__(self, msg: str, retry_after: float | None = None) -> None:
-        super().__init__(msg)
-        self.retry_after = retry_after
-
-
-def _parse_retry_after(headers: Mapping[str, str]) -> float | None:
-    """Parse seconds-to-reset from HF's RateLimit header.
-
-    HF sends e.g. ``ratelimit: "api";r=0;t=52`` on 429s; ``t`` is the wait.
-    Returns ``None`` if the header is missing or has no ``t`` field.
-    """
-    raw = headers.get("RateLimit") or headers.get("ratelimit")
-    if raw is None:
-        return None
-    for part in raw.split(";"):
-        key, _, val = part.strip().partition("=")
-        if key == "t":
-            try:
-                return float(val)
-            except ValueError:
-                return None
-    return None
-
-
-# reset window is 5 min
-_RATE_LIMIT_MAX_SLEEP_SECS = 300.0
-
-# 24h. Manually clear the cache (or `delete_model`) to force a refresh.
-_FILE_LIST_CACHE_TTL_SECS = 24 * 60 * 60
-
 
 async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
     token = await get_hf_token()
@@ -379,6 +348,9 @@ async def _build_file_list_from_local_directory(
     return None
 
 
+_fetched_file_lists_this_session: set[str] = set()
+
+
 async def fetch_file_list_with_cache(
     model_id: ModelId,
     revision: str = "main",
@@ -388,16 +360,13 @@ async def fetch_file_list_with_cache(
 ) -> list[FileListEntry]:
     target_dir = await ensure_cache_dir(model_id)
     cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
+    cache_key = f"{model_id.normalize()}--{revision}"
 
-    # cache survives process restarts so cold starts don't re-burst HF
-    if await aios.path.exists(cache_file):
-        try:
-            cache_age = time.time() - (await aios.stat(cache_file)).st_mtime
-        except OSError:
-            cache_age = float("inf")
-        if cache_age < _FILE_LIST_CACHE_TTL_SECS:
-            async with aiofiles.open(cache_file, "r") as f:
-                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+    if cache_key in _fetched_file_lists_this_session and await aios.path.exists(
+        cache_file
+    ):
+        async with aiofiles.open(cache_file, "r") as f:
+            return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
 
     if skip_internet:
         if await aios.path.exists(cache_file):
@@ -426,6 +395,7 @@ async def fetch_file_list_with_cache(
             await f.write(
                 TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
             )
+        _fetched_file_lists_this_session.add(cache_key)
         return file_list
     except Exception as e:
         logger.opt(exception=e).warning(
@@ -456,29 +426,17 @@ async def fetch_file_list_with_retry(
     recursive: bool = False,
     on_connection_lost: Callable[[], None] = lambda: None,
 ) -> list[FileListEntry]:
-    n_attempts = 5
+    n_attempts = 3
     for attempt in range(n_attempts):
         try:
             return await _fetch_file_list(model_id, revision, path, recursive)
         except HuggingFaceAuthenticationError:
             raise
-        except HuggingFaceRateLimitError as e:
-            if attempt == n_attempts - 1:
-                raise
-            sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
-            sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
-                0, 1
-            )
-            logger.warning(
-                f"Rate limited by HuggingFace fetching file list for {model_id}; "
-                f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
-            )
-            await asyncio.sleep(sleep_for)
         except Exception as e:
             on_connection_lost()
             if attempt == n_attempts - 1:
                 raise e
-            await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
+            await asyncio.sleep(2.0**attempt)
     raise Exception(
         f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
     )
@@ -489,9 +447,6 @@ async def _fetch_file_list(
 ) -> list[FileListEntry]:
     api_url = f"{get_hf_endpoint()}/api/models/{model_id}/tree/{revision}"
     url = f"{api_url}/{path}" if path else api_url
-    # ?recursive=true returns the whole subtree in one request
-    if recursive:
-        url = f"{url}?recursive=true"
 
     headers = await get_download_headers()
     async with (
@@ -503,8 +458,7 @@ async def _fetch_file_list(
             raise HuggingFaceAuthenticationError(msg)
         elif response.status == 429:
             raise HuggingFaceRateLimitError(
-                f"HuggingFace rate limit hit fetching file list for {model_id}",
-                retry_after=_parse_retry_after(response.headers),
+                f"Couldn't download {model_id} because of HuggingFace rate limit."
             )
         elif response.status == 200:
             data_json = await response.text()
@@ -514,14 +468,10 @@ async def _fetch_file_list(
                 if item.type == "file":
                     files.append(FileListEntry.model_validate(item))
                 elif item.type == "directory" and recursive:
-                    # already inlined by ?recursive=true
-                    continue
-            if recursive and len(data) >= 1000:
-                # HF tree endpoint paginates at 1000; we don't follow cursors
-                logger.warning(
-                    f"File list for {model_id} hit the 1000-entry page cap "
-                    "and may be truncated; cursor pagination is not implemented"
-                )
+                    subfiles = await _fetch_file_list(
+                        model_id, revision, item.path, recursive
+                    )
+                    files.extend(subfiles)
             return files
         else:
             raise Exception(f"Failed to fetch file list: {response.status}")
@@ -602,11 +552,6 @@ async def file_meta(
         if r.status in [401, 403]:
             msg = await _build_auth_error_message(r.status, model_id)
             raise HuggingFaceAuthenticationError(msg)
-        if r.status == 429:
-            raise HuggingFaceRateLimitError(
-                f"HuggingFace rate limit hit fetching metadata for {model_id}/{path}",
-                retry_after=_parse_retry_after(r.headers),
-            )
         content_length = int(
             r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
         )
@@ -626,7 +571,7 @@ async def download_file_with_retry(
     on_connection_lost: Callable[[], None] = lambda: None,
     skip_internet: bool = False,
 ) -> Path:
-    n_attempts = 5
+    n_attempts = 3
     for attempt in range(n_attempts):
         try:
             return await _download_file(
@@ -638,16 +583,12 @@ async def download_file_with_retry(
             raise
         except HuggingFaceRateLimitError as e:
             if attempt == n_attempts - 1:
-                raise
-            sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
-            sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
-                0, 1
+                raise e
+            logger.error(
+                f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
             )
-            logger.warning(
-                f"Rate limited by HuggingFace downloading {model_id}/{path}; "
-                f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
-            )
-            await asyncio.sleep(sleep_for)
+            logger.error(traceback.format_exc())
+            await asyncio.sleep(2.0**attempt)
         except Exception as e:
             if attempt == n_attempts - 1:
                 on_connection_lost()
@@ -656,7 +597,7 @@ async def download_file_with_retry(
                 f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
             )
             logger.error(traceback.format_exc())
-            await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
+            await asyncio.sleep(2.0**attempt)
     raise Exception(
         f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
     )
@@ -724,11 +665,6 @@ async def _download_file(
             if r.status in [401, 403]:
                 msg = await _build_auth_error_message(r.status, model_id)
                 raise HuggingFaceAuthenticationError(msg)
-            if r.status == 429:
-                raise HuggingFaceRateLimitError(
-                    f"HuggingFace rate limit hit downloading {model_id}/{path}",
-                    retry_after=_parse_retry_after(r.headers),
-                )
             assert r.status in [200, 206], (
                 f"Failed to download {path} from {url}: {r.status}"
             )

src/exo/download/tests/test_offline_mode.py

@@ -1,7 +1,5 @@
 """Tests for offline/air-gapped mode."""
 
-import os
-import time
 from collections.abc import AsyncIterator
 from pathlib import Path
 from unittest.mock import AsyncMock, patch
@@ -233,64 +231,3 @@ class TestFetchFileListOffline:
         raise FileNotFoundError."""
         with pytest.raises(FileNotFoundError, match="No internet"):
             await fetch_file_list_with_cache(model_id, "main", skip_internet=True)
-
-
-class TestFileListCacheTTL:
-    async def test_uses_fresh_cache_without_fetching(
-        self, model_id: ModelId, temp_models_dir: Path
-    ) -> None:
-        from pydantic import TypeAdapter
-
-        cache_dir = temp_models_dir / "caches" / model_id.normalize()
-        await aios.makedirs(cache_dir, exist_ok=True)
-
-        cached_list = [
-            FileListEntry(type="file", path="model.safetensors", size=1000),
-        ]
-        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
-        async with aiofiles.open(cache_file, "w") as f:
-            await f.write(
-                TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode()
-            )
-
-        with patch(
-            "exo.download.download_utils.fetch_file_list_with_retry",
-            new_callable=AsyncMock,
-        ) as mock_fetch:
-            result = await fetch_file_list_with_cache(model_id, "main")
-
-        assert result == cached_list
-        mock_fetch.assert_not_called()
-
-    async def test_refetches_when_cache_older_than_ttl(
-        self, model_id: ModelId, temp_models_dir: Path
-    ) -> None:
-        from pydantic import TypeAdapter
-
-        from exo.download.download_utils import (
-            _FILE_LIST_CACHE_TTL_SECS,  # pyright: ignore[reportPrivateUsage]
-        )
-
-        cache_dir = temp_models_dir / "caches" / model_id.normalize()
-        await aios.makedirs(cache_dir, exist_ok=True)
-
-        stale_list = [FileListEntry(type="file", path="stale.bin", size=1)]
-        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
-        async with aiofiles.open(cache_file, "w") as f:
-            await f.write(
-                TypeAdapter(list[FileListEntry]).dump_json(stale_list).decode()
-            )
-
-        old_mtime = time.time() - _FILE_LIST_CACHE_TTL_SECS - 60
-        os.utime(cache_file, (old_mtime, old_mtime))
-
-        fresh_list = [FileListEntry(type="file", path="fresh.bin", size=2)]
-        with patch(
-            "exo.download.download_utils.fetch_file_list_with_retry",
-            new_callable=AsyncMock,
-            return_value=fresh_list,
-        ) as mock_fetch:
-            result = await fetch_file_list_with_cache(model_id, "main")
-
-        assert result == fresh_list
-        mock_fetch.assert_called_once()

src/exo/download/tests/test_rate_limit_handling.py

@@ -1,355 +0,0 @@
-"""Tests for HuggingFace 429 rate-limit handling in download_utils."""
-
-from collections.abc import AsyncIterator
-from pathlib import Path
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import aiofiles.os as aios
-import pytest
-
-from exo.download.download_utils import (
-    HuggingFaceRateLimitError,
-    _download_file,  # pyright: ignore[reportPrivateUsage]
-    _fetch_file_list,  # pyright: ignore[reportPrivateUsage]
-    _parse_retry_after,  # pyright: ignore[reportPrivateUsage]
-    download_file_with_retry,
-    fetch_file_list_with_retry,
-    file_meta,
-)
-from exo.shared.types.common import ModelId
-
-# captured from a real HF 429 on 2026-04-30 (header is lowercased by Cloudfront)
-REAL_HF_429_HEADERS_2026_04_30 = {
-    "ratelimit": '"api";r=0;t=52',
-    "ratelimit-policy": '"fixed window";"api";q=500;w=300',
-}
-
-
-class TestParseRetryAfter:
-    def test_parses_documented_format(self) -> None:
-        assert _parse_retry_after({"RateLimit": '"api";r=0;t=243'}) == 243.0
-
-    def test_parses_real_hf_response(self) -> None:
-        assert _parse_retry_after(REAL_HF_429_HEADERS_2026_04_30) == 52.0
-
-    def test_parses_resolvers_bucket(self) -> None:
-        assert _parse_retry_after({"ratelimit": '"resolvers";r=0;t=120'}) == 120.0
-
-    def test_parses_pages_bucket(self) -> None:
-        assert _parse_retry_after({"ratelimit": '"pages";r=0;t=10'}) == 10.0
-
-    def test_returns_none_when_header_missing(self) -> None:
-        assert _parse_retry_after({}) is None
-
-    def test_returns_none_when_only_retry_after_present(self) -> None:
-        assert _parse_retry_after({"Retry-After": "60"}) is None
-
-    def test_returns_none_when_format_unrecognised(self) -> None:
-        assert _parse_retry_after({"ratelimit": "garbage"}) is None
-
-    def test_handles_extra_whitespace(self) -> None:
-        assert _parse_retry_after({"ratelimit": '"api"; r=0; t=42'}) == 42.0
-
-
-class TestFetchFileListRetry:
-    async def test_uses_retry_after_from_error(self) -> None:
-        sleeps: list[float] = []
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
-            if not sleeps:
-                raise HuggingFaceRateLimitError("rate limited", retry_after=2.0)
-            return []
-
-        with (
-            patch(
-                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-        ):
-            result = await fetch_file_list_with_retry(ModelId("test/model"))
-
-        assert result == []
-        assert len(sleeps) == 1
-        assert 2.0 <= sleeps[0] < 3.0  # retry_after + jitter[0,1)
-
-    async def test_falls_back_to_exp_backoff_when_no_retry_after(self) -> None:
-        sleeps: list[float] = []
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
-            if not sleeps:
-                raise HuggingFaceRateLimitError("rate limited", retry_after=None)
-            return []
-
-        with (
-            patch(
-                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-        ):
-            await fetch_file_list_with_retry(ModelId("test/model"))
-
-        assert len(sleeps) == 1
-        assert 1.0 <= sleeps[0] < 2.0  # 2**0 + jitter[0,1)
-
-    async def test_caps_sleep_at_max_window(self) -> None:
-        sleeps: list[float] = []
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
-            if not sleeps:
-                raise HuggingFaceRateLimitError("rate limited", retry_after=10_000.0)
-            return []
-
-        with (
-            patch(
-                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-        ):
-            await fetch_file_list_with_retry(ModelId("test/model"))
-
-        assert len(sleeps) == 1
-        assert 300.0 <= sleeps[0] < 301.0  # cap + jitter[0,1)
-
-    async def test_retries_up_to_five_times(self) -> None:
-        sleeps: list[float] = []
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
-            raise HuggingFaceRateLimitError("rate limited", retry_after=1.0)
-
-        with (
-            patch(
-                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-            pytest.raises(HuggingFaceRateLimitError),
-        ):
-            await fetch_file_list_with_retry(ModelId("test/model"))
-
-        assert len(sleeps) == 4  # 5 attempts -> 4 sleeps before giving up
-
-
-class TestDownloadFileRetry:
-    @pytest.fixture
-    async def target_dir(self, tmp_path: Path) -> AsyncIterator[Path]:
-        target = tmp_path / "downloads"
-        await aios.makedirs(target, exist_ok=True)
-        yield target
-
-    async def test_uses_retry_after_from_error(self, target_dir: Path) -> None:
-        sleeps: list[float] = []
-        results: list[Path] = [target_dir / "file.bin"]
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        async def fake_download(*args: object, **kwargs: object) -> Path:
-            if not sleeps:
-                raise HuggingFaceRateLimitError("rate limited", retry_after=5.0)
-            return results[0]
-
-        with (
-            patch(
-                "exo.download.download_utils._download_file",
-                side_effect=fake_download,
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-        ):
-            result = await download_file_with_retry(
-                ModelId("test/model"), "main", "file.bin", target_dir
-            )
-
-        assert result == results[0]
-        assert len(sleeps) == 1
-        assert 5.0 <= sleeps[0] < 6.0
-
-    async def test_caps_sleep_at_max_window(self, target_dir: Path) -> None:
-        sleeps: list[float] = []
-        results: list[Path] = [target_dir / "file.bin"]
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        async def fake_download(*args: object, **kwargs: object) -> Path:
-            if not sleeps:
-                raise HuggingFaceRateLimitError("rate limited", retry_after=99_999.0)
-            return results[0]
-
-        with (
-            patch(
-                "exo.download.download_utils._download_file",
-                side_effect=fake_download,
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-        ):
-            await download_file_with_retry(
-                ModelId("test/model"), "main", "file.bin", target_dir
-            )
-
-        assert len(sleeps) == 1
-        assert 300.0 <= sleeps[0] < 301.0
-
-    async def test_retries_up_to_five_times(self, target_dir: Path) -> None:
-        sleeps: list[float] = []
-
-        async def fake_sleep(seconds: float) -> None:
-            sleeps.append(seconds)
-
-        with (
-            patch(
-                "exo.download.download_utils._download_file",
-                new_callable=AsyncMock,
-                side_effect=HuggingFaceRateLimitError("rate limited", retry_after=1.0),
-            ),
-            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
-            pytest.raises(HuggingFaceRateLimitError),
-        ):
-            await download_file_with_retry(
-                ModelId("test/model"), "main", "file.bin", target_dir
-            )
-
-        assert len(sleeps) == 4
-
-
-def _make_mock_session_returning(
-    response_attrs: dict[str, object], method: str = "get"
-) -> MagicMock:
-    """Build a MagicMock that mimics ``create_http_session`` returning a
-    response whose ``status`` / ``headers`` are set from ``response_attrs``.
-
-    Mocks the chain ``create_http_session().__aenter__() -> session``, and
-    ``session.<method>().__aenter__() -> response``.
-    """
-    mock_response = MagicMock()
-    for k, v in response_attrs.items():
-        setattr(mock_response, k, v)
-
-    mock_session = MagicMock()
-    method_mock = getattr(mock_session, method)  # pyright: ignore[reportAny]
-    method_mock.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
-        return_value=mock_response
-    )
-    method_mock.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
-        return_value=None
-    )
-
-    mock_factory = MagicMock()
-    mock_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
-        return_value=mock_session
-    )
-    mock_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
-        return_value=None
-    )
-    return mock_factory
-
-
-REAL_HF_429_HEADER_DICT = {"ratelimit": '"api";r=0;t=52'}
-
-
-class TestRateLimitAtHttpCallSites:
-    """Verify each HF call site translates an HTTP 429 into a
-    ``HuggingFaceRateLimitError`` carrying the parsed ``retry_after``.
-
-    These tests would catch regressions where (a) the 429 branch is
-    deleted, (b) ``_parse_retry_after`` stops being called, or
-    (c) the wrong header object is passed to it.
-    """
-
-    async def test_fetch_file_list_maps_429_to_rate_limit_error(self) -> None:
-        mock_factory = _make_mock_session_returning(
-            {"status": 429, "headers": REAL_HF_429_HEADER_DICT}
-        )
-        with (
-            patch("exo.download.download_utils.create_http_session", mock_factory),
-            pytest.raises(HuggingFaceRateLimitError) as exc_info,
-        ):
-            await _fetch_file_list(ModelId("test/model"), "main")
-        assert exc_info.value.retry_after == 52.0
-
-    async def test_file_meta_maps_429_to_rate_limit_error(self) -> None:
-        mock_factory = _make_mock_session_returning(
-            {"status": 429, "headers": REAL_HF_429_HEADER_DICT}, method="head"
-        )
-        with (
-            patch("exo.download.download_utils.create_http_session", mock_factory),
-            pytest.raises(HuggingFaceRateLimitError) as exc_info,
-        ):
-            await file_meta(ModelId("test/model"), "main", "weights.safetensors")
-        assert exc_info.value.retry_after == 52.0
-
-    async def test_file_meta_maps_429_after_307_redirect(self) -> None:
-        """When the initial HEAD 307s and the redirected HEAD then 429s,
-        the 429 must still surface as ``HuggingFaceRateLimitError``."""
-        # First HEAD -> 307 with a Location header pointing somewhere new.
-        first_response = MagicMock()
-        first_response.status = 307
-        first_response.headers = {"location": "/redirected/url"}
-
-        # Second HEAD (the recursive call) -> 429 with the real-HF header.
-        second_response = MagicMock()
-        second_response.status = 429
-        second_response.headers = REAL_HF_429_HEADER_DICT
-
-        responses = iter([first_response, second_response])
-
-        mock_session = MagicMock()
-        mock_session.head.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
-            side_effect=lambda: next(responses)
-        )
-        mock_session.head.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
-            return_value=None
-        )
-
-        mock_factory = MagicMock()
-        mock_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
-            return_value=mock_session
-        )
-        mock_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
-            return_value=None
-        )
-
-        with (
-            patch("exo.download.download_utils.create_http_session", mock_factory),
-            pytest.raises(HuggingFaceRateLimitError) as exc_info,
-        ):
-            await file_meta(ModelId("test/model"), "main", "weights.safetensors")
-        assert exc_info.value.retry_after == 52.0
-
-    async def test_download_file_maps_429_to_rate_limit_error(
-        self, tmp_path: Path
-    ) -> None:
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-        # No local file -> _download_file goes straight to file_meta then GET.
-        # We need both calls to succeed enough to reach the GET branch:
-        #   - file_meta returns a non-429 (size, etag) so we proceed.
-        #   - the GET then 429s.
-        with (
-            patch(
-                "exo.download.download_utils.file_meta",
-                new_callable=AsyncMock,
-                return_value=(100, "abc123"),
-            ),
-            patch(
-                "exo.download.download_utils.create_http_session",
-                _make_mock_session_returning(
-                    {"status": 429, "headers": REAL_HF_429_HEADER_DICT}
-                ),
-            ),
-            pytest.raises(HuggingFaceRateLimitError) as exc_info,
-        ):
-            await _download_file(
-                ModelId("test/model"), "main", "weights.safetensors", target_dir
-            )
-        assert exc_info.value.retry_after == 52.0

src/exo/main.py

@@ -59,7 +59,6 @@ class Node:
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
         await router.register_topic(topics.DOWNLOAD_COMMANDS)
-        await router.register_topic(topics.SNAPSHOT_RESPONSES)
         event_router = EventRouter(
             session_id,
             command_sender=router.sender(topics.COMMANDS),
@@ -84,11 +83,8 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
-                session_id,
                 port=args.api_port,
-                event_router=event_router,
                 event_receiver=event_router.receiver(),
-                snapshot_chunk_receiver=router.receiver(topics.SNAPSHOT_RESPONSES),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -99,11 +95,8 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
-                session_id,
-                event_router=event_router,
                 event_receiver=event_router.receiver(),
                 event_sender=event_router.sender(),
-                snapshot_chunk_receiver=router.receiver(topics.SNAPSHOT_RESPONSES),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 api_port=args.api_port,
@@ -119,7 +112,6 @@ class Node:
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
-            snapshot_chunk_sender=router.sender(topics.SNAPSHOT_RESPONSES),
             download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
         )
 
@@ -218,9 +210,6 @@ class Node:
                         global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                         local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
                         command_receiver=self.router.receiver(topics.COMMANDS),
-                        snapshot_chunk_sender=self.router.sender(
-                            topics.SNAPSHOT_RESPONSES
-                        ),
                         download_command_sender=self.router.sender(
                             topics.DOWNLOAD_COMMANDS
                         ),
@@ -257,13 +246,8 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
-                            result.session_id,
-                            event_router=self.event_router,
                             event_receiver=self.event_router.receiver(),
                             event_sender=self.event_router.sender(),
-                            snapshot_chunk_receiver=self.router.receiver(
-                                topics.SNAPSHOT_RESPONSES
-                            ),
                             command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
@@ -272,13 +256,7 @@ class Node:
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(
-                            result.won_clock,
-                            result.session_id,
-                            self.event_router,
-                            self.event_router.receiver(),
-                            self.router.receiver(topics.SNAPSHOT_RESPONSES),
-                        )
+                        self.api.reset(result.won_clock, self.event_router.receiver())
                     self._tg.start_soon(self.event_router.run)
                 else:
                     if self.api:

src/exo/master/main.py

@@ -1,8 +1,6 @@
-import hashlib
 from datetime import datetime, timedelta, timezone
 
 import anyio
-from anyio import to_thread
 from loguru import logger
 
 from exo.master.placement import (
@@ -12,7 +10,6 @@ from exo.master.placement import (
     get_transition_events,
     place_instance,
 )
-from exo.master.placement_utils import find_ip_prioritised
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
@@ -20,16 +17,13 @@ from exo.shared.types.commands import (
     CreateInstance,
     DeleteCustomModelCard,
     DeleteInstance,
-    DeleteInstanceLink,
     ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
     RequestEventLog,
-    RequestSnapshot,
     SendInputChunk,
-    SetInstanceLink,
     TaskCancelled,
     TaskFinished,
     TestCommand,
@@ -44,8 +38,6 @@ from exo.shared.types.events import (
     IndexedEvent,
     InputChunkReceived,
     InstanceDeleted,
-    InstanceLinkCreated,
-    InstanceLinkDeleted,
     LocalForwarderEvent,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -56,8 +48,6 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
-from exo.shared.types.instance_link import InstanceLink
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -78,55 +68,6 @@ from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
-_SNAPSHOT_CHUNK_BYTES = 512 * 1024
-_MAX_EVENT_LOG_REPLAY_BATCH = 1000
-
-
-def _encode_state_for_transfer(state: State) -> bytes:
-    import zstandard
-
-    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-
-
-def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
-    decode = state.instances.get(decode_instance_id)
-    if decode is None:
-        return None
-    decode_node = next(iter(decode.shard_assignments.node_to_runner.keys()), None)
-    if decode_node is None:
-        return None
-
-    sources: set[InstanceId] = set()
-    for link in state.instance_links.values():
-        if decode_instance_id in link.decode_instances:
-            sources.update(link.prefill_instances)
-    sources.discard(decode_instance_id)
-
-    in_flight = {TaskStatus.Pending, TaskStatus.Running}
-    task_counts: dict[InstanceId, int] = {
-        src_id: sum(
-            1
-            for task in state.tasks.values()
-            if task.instance_id == src_id and task.task_status in in_flight
-        )
-        for src_id in sources
-    }
-    for src_id in sorted(sources, key=lambda sid: task_counts[sid]):
-        instance = state.instances.get(src_id)
-        if instance is None:
-            continue
-        for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
-            port = state.prefill_server_ports.get(runner_id)
-            if port is None:
-                continue
-            ip = find_ip_prioritised(
-                decode_node, node_id, state.topology, state.node_network, ring=True
-            )
-            if ip is None:
-                continue
-            return f"{ip}:{port}"
-    return None
-
 
 class Master:
     def __init__(
@@ -138,7 +79,6 @@ class Master:
         event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
-        snapshot_chunk_sender: Sender[SnapshotChunk],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
         self.node_id = node_id
@@ -149,7 +89,6 @@ class Master:
         self.command_receiver = command_receiver
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
-        self.snapshot_chunk_sender = snapshot_chunk_sender
         self.download_command_sender = download_command_sender
         self.event_sender = event_sender
         self._system_id = SystemId()
@@ -170,7 +109,6 @@ class Master:
             self._event_log.close()
             self.global_event_sender.close()
             self.local_event_receiver.close()
-            self.snapshot_chunk_sender.close()
             self.command_receiver.close()
 
     async def shutdown(self):
@@ -190,24 +128,15 @@ class Master:
                         case TestCommand():
                             pass
                         case TextGeneration():
-                            prefill_only: set[InstanceId] = set()
-                            for link in self.state.instance_links.values():
-                                prefill_only.update(link.prefill_instances)
-                            for link in self.state.instance_links.values():
-                                prefill_only.difference_update(link.decode_instances)
-
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
                                     == command.task_params.model
-                                    and instance.instance_id not in prefill_only
                                 ):
-                                    in_flight = {TaskStatus.Pending, TaskStatus.Running}
                                     task_count = sum(
                                         1
                                         for task in self.state.tasks.values()
                                         if task.instance_id == instance.instance_id
-                                        and task.task_status in in_flight
                                     )
                                     instance_task_counts[instance.instance_id] = (
                                         task_count
@@ -225,27 +154,20 @@ class Master:
                                 ],
                             )
 
-                            decode_instance_id = available_instance_ids[0]
                             task_id = TaskId()
-                            params = command.task_params.model_copy(
-                                update={
-                                    "prefill_endpoint": _prefill_endpoint_for(
-                                        self.state, decode_instance_id
-                                    ),
-                                }
-                            )
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
                                     task=TextGenerationTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
-                                        instance_id=decode_instance_id,
+                                        instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=params,
+                                        task_params=command.task_params,
                                     ),
                                 )
                             )
+
                             self.command_task_mapping[command.command_id] = task_id
                         case ImageGeneration():
                             for instance in self.state.instances.values():
@@ -253,12 +175,10 @@ class Master:
                                     instance.shard_assignments.model_id
                                     == command.task_params.model
                                 ):
-                                    in_flight = {TaskStatus.Pending, TaskStatus.Running}
                                     task_count = sum(
                                         1
                                         for task in self.state.tasks.values()
                                         if task.instance_id == instance.instance_id
-                                        and task.task_status in in_flight
                                     )
                                     instance_task_counts[instance.instance_id] = (
                                         task_count
@@ -309,12 +229,10 @@ class Master:
                                     instance.shard_assignments.model_id
                                     == command.task_params.model
                                 ):
-                                    in_flight = {TaskStatus.Pending, TaskStatus.Running}
                                     task_count = sum(
                                         1
                                         for task in self.state.tasks.values()
                                         if task.instance_id == instance.instance_id
-                                        and task.task_status in in_flight
                                     )
                                     instance_task_counts[instance.instance_id] = (
                                         task_count
@@ -439,37 +357,15 @@ class Master:
                             generated_events.append(
                                 CustomModelCardDeleted(model_id=command.model_id)
                             )
-                        case SetInstanceLink():
-                            link = InstanceLink(
-                                link_id=command.link_id,
-                                prefill_instances=list(
-                                    dict.fromkeys(command.prefill_instances)
-                                ),
-                                decode_instances=list(
-                                    dict.fromkeys(command.decode_instances)
-                                ),
-                            )
-                            generated_events.append(InstanceLinkCreated(link=link))
-                        case DeleteInstanceLink():
-                            generated_events.append(
-                                InstanceLinkDeleted(link_id=command.link_id)
-                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
-                            end = min(
-                                command.since_idx + _MAX_EVENT_LOG_REPLAY_BATCH,
-                                len(self._event_log),
-                            )
+                            end = min(command.since_idx + 1000, len(self._event_log))
                             for i, event in enumerate(
                                 self._event_log.read_range(command.since_idx, end),
                                 start=command.since_idx,
                             ):
                                 await self._send_event(IndexedEvent(idx=i, event=event))
-                        case RequestSnapshot():
-                            self._tg.start_soon(
-                                self._serve_snapshot, command.requester_node_id
-                            )
                     for event in generated_events:
                         await self.event_sender.send(event)
                 except ValueError as e:
@@ -529,42 +425,6 @@ class Master:
                     self._event_log.append(event)
                     await self._send_event(indexed)
 
-    async def _serve_snapshot(self, requester_node_id: NodeId) -> None:
-        state = self.state
-        if state.last_event_applied_idx < 0:
-            logger.info(
-                f"RequestSnapshot from {requester_node_id} but master has no events yet"
-            )
-            return
-
-        body = await to_thread.run_sync(_encode_state_for_transfer, state)
-        sha256 = hashlib.sha256(body).hexdigest()
-        chunks = [
-            body[i : i + _SNAPSHOT_CHUNK_BYTES]
-            for i in range(0, len(body), _SNAPSHOT_CHUNK_BYTES)
-        ] or [b""]
-        transfer_id = SnapshotTransferId()
-
-        logger.info(
-            f"Serving snapshot to {requester_node_id}: "
-            f"idx={state.last_event_applied_idx}, "
-            f"{len(chunks)} chunk(s), {len(body)} bytes total"
-        )
-        for index, chunk in enumerate(chunks):
-            await self.snapshot_chunk_sender.send(
-                SnapshotChunk.from_data(
-                    data=chunk,
-                    transfer_id=transfer_id,
-                    requester_node_id=requester_node_id,
-                    session_id=self.session_id,
-                    schema_version=state.schema_version,
-                    last_event_applied_idx=state.last_event_applied_idx,
-                    chunk_index=index,
-                    total_chunks=len(chunks),
-                    sha256_hex=sha256,
-                )
-            )
-
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly

src/exo/master/placement.py

@@ -1,3 +1,4 @@
+import random
 from collections.abc import Mapping
 from copy import deepcopy
 from typing import Sequence
@@ -45,7 +46,11 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
 )
 from exo.shared.types.worker.shards import Sharding
-from exo.utils.ports import random_ephemeral_port
+
+
+def random_ephemeral_port() -> int:
+    port = random.randint(49153, 65535)
+    return port - 1 if port <= 52415 else port
 
 
 def add_instance_to_placements(

src/exo/master/placement_utils.py

@@ -336,7 +336,7 @@ def _find_connection_ip(
             yield connection.sink_multiaddr.ip_address
 
 
-def find_ip_prioritised(
+def _find_ip_prioritised(
     node_id: NodeId,
     other_node_id: NodeId,
     cycle_digraph: Topology,
@@ -413,7 +413,7 @@ def get_mlx_ring_hosts_by_node(
                 hosts_for_node.append(Host(ip="198.51.100.1", port=0))
                 continue
 
-            connection_ip = find_ip_prioritised(
+            connection_ip = _find_ip_prioritised(
                 node_id, other_node_id, cycle_digraph, node_network, ring=True
             )
             if connection_ip is None:
@@ -445,7 +445,7 @@ def get_mlx_jaccl_coordinators(
         if n == coordinator:
             return "0.0.0.0"
 
-        ip = find_ip_prioritised(
+        ip = _find_ip_prioritised(
             n, coordinator, cycle_digraph, node_network, ring=False
         )
         if ip is not None:

src/exo/master/tests/test_master.py

@@ -7,14 +7,12 @@ from loguru import logger
 
 from exo.master.main import Master
 from exo.routing.router import get_node_id_keypair
-from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.commands import (
     CommandId,
     ForwarderCommand,
     ForwarderDownloadCommand,
     PlaceInstance,
-    RequestSnapshot,
     TextGeneration,
 )
 from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
@@ -31,7 +29,6 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryUsage,
 )
-from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.tasks import TaskStatus
 from exo.shared.types.tasks import TextGeneration as TextGenerationTask
 from exo.shared.types.text_generation import (
@@ -59,7 +56,6 @@ async def test_master():
     local_event_sender, le_receiver = channel[LocalForwarderEvent]()
     fcds, _fcdr = channel[ForwarderDownloadCommand]()
     ev_send, ev_recv = channel[Event]()
-    snapshot_chunk_send, _snapshot_chunk_recv = channel[SnapshotChunk]()
 
     async def mock_event_router():
         idx = 0
@@ -96,7 +92,6 @@ async def test_master():
         global_event_sender=ge_sender,
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
-        snapshot_chunk_sender=snapshot_chunk_send,
         download_command_sender=fcds,
     )
     logger.info("run the master")
@@ -234,52 +229,3 @@ async def test_master():
 
         ev_send.close()
         await master.shutdown()
-
-
-@pytest.mark.asyncio
-async def test_master_serves_snapshot_for_current_state():
-    node_id = NodeId("master")
-    requester_node_id = NodeId("worker")
-    session_id = SessionId(master_node_id=node_id, election_clock=0)
-
-    ge_sender, _global_event_receiver = channel[GlobalForwarderEvent]()
-    command_sender, command_receiver = channel[ForwarderCommand]()
-    _local_event_sender, local_event_receiver = channel[LocalForwarderEvent]()
-    download_command_sender, _download_command_receiver = channel[
-        ForwarderDownloadCommand
-    ]()
-    event_sender, _event_receiver = channel[Event]()
-    snapshot_chunk_sender, snapshot_chunk_receiver = channel[SnapshotChunk]()
-
-    master = Master(
-        node_id,
-        session_id,
-        event_sender=event_sender,
-        global_event_sender=ge_sender,
-        local_event_receiver=local_event_receiver,
-        command_receiver=command_receiver,
-        snapshot_chunk_sender=snapshot_chunk_sender,
-        download_command_sender=download_command_sender,
-    )
-    master.state = master.state.model_copy(update={"last_event_applied_idx": 12})
-
-    receiver = SnapshotReceiver(requester_node_id, session_id)
-    async with anyio.create_task_group() as tg:
-        tg.start_soon(master.run)
-        await command_sender.send(
-            ForwarderCommand(
-                origin=SystemId("api"),
-                command=RequestSnapshot(requester_node_id=requester_node_id),
-            )
-        )
-
-        received = None
-        while received is None:
-            chunk = await snapshot_chunk_receiver.receive()
-            received = receiver.ingest(chunk)
-
-        assert received.last_event_applied_idx == 12
-        assert received.state.last_event_applied_idx == 12
-
-        await master.shutdown()
-        tg.cancel_scope.cancel()

src/exo/reasoning_proxy/_helpers.py

@@ -0,0 +1,33 @@
+from typing import cast
+
+
+def as_str(value: object) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+def as_list(value: object) -> list[object] | None:
+    if isinstance(value, list):
+        return cast(list[object], value)
+    return None
+
+
+def as_dict(value: object) -> dict[str, object] | None:
+    if isinstance(value, dict):
+        return cast(dict[str, object], value)
+    return None
+
+
+def as_int(value: object, default: int = 0) -> int:
+    return value if isinstance(value, int) and not isinstance(value, bool) else default
+
+
+def dict_get_str(d: dict[str, object], key: str) -> str | None:
+    return as_str(d.get(key))
+
+
+def dict_get_list(d: dict[str, object], key: str) -> list[object] | None:
+    return as_list(d.get(key))
+
+
+def dict_get_dict(d: dict[str, object], key: str) -> dict[str, object] | None:
+    return as_dict(d.get(key))

src/exo/reasoning_proxy/accumulator.py

@@ -0,0 +1,261 @@
+"""Accumulators for capturing the emitted assistant shape from a streaming response.
+
+Both accumulators are fed raw SSE chunks (bytes) as they pass through. At stream
+end, they expose a canonical assistant-message shape suitable for hashing, plus
+the reasoning text that should be cached against that hash.
+"""
+
+import json
+import logging
+from typing import cast
+
+from exo.reasoning_proxy._helpers import (
+    as_dict,
+    as_str,
+    dict_get_dict,
+    dict_get_list,
+    dict_get_str,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIAccumulator:
+    """Captures content, tool_calls, and reasoning_content from OpenAI SSE chunks.
+
+    OpenAI can emit multiple choices per chunk; we only track choice index 0
+    (the common case for chat completions; n>1 is uncommon and re-hash misses
+    there degrade gracefully to no-op cache insert).
+    """
+
+    def __init__(self) -> None:
+        self._content_parts: list[str] = []
+        self._reasoning_parts: list[str] = []
+        self._tool_calls_by_index: dict[int, dict[str, object]] = {}
+        self._buffer = ""
+
+    def feed_bytes(self, chunk: bytes) -> None:
+        self._buffer += chunk.decode("utf-8", errors="replace")
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            line = line.strip()
+            if not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]" or not payload:
+                continue
+            try:
+                parsed = cast(object, json.loads(payload))
+            except json.JSONDecodeError:
+                continue
+            event = as_dict(parsed)
+            if event is None:
+                continue
+            self._consume_event(event)
+
+    def _consume_event(self, event: dict[str, object]) -> None:
+        choices = dict_get_list(event, "choices")
+        if not choices:
+            return
+        for choice_raw in choices:
+            choice = as_dict(choice_raw)
+            if choice is None:
+                continue
+            index_val = choice.get("index", 0)
+            if (
+                not (isinstance(index_val, int) and not isinstance(index_val, bool))
+                or index_val != 0
+            ):
+                continue
+            delta = dict_get_dict(choice, "delta")
+            if delta is None:
+                continue
+            content = dict_get_str(delta, "content")
+            if content is not None:
+                self._content_parts.append(content)
+            reasoning = dict_get_str(delta, "reasoning_content")
+            if reasoning is not None:
+                self._reasoning_parts.append(reasoning)
+            tool_calls = dict_get_list(delta, "tool_calls")
+            if tool_calls is not None:
+                self._merge_tool_calls(tool_calls)
+
+    def _merge_tool_calls(self, deltas: list[object]) -> None:
+        for raw in deltas:
+            d = as_dict(raw)
+            if d is None:
+                continue
+            index_val = d.get("index", 0)
+            if not (isinstance(index_val, int) and not isinstance(index_val, bool)):
+                continue
+            entry = self._tool_calls_by_index.setdefault(
+                index_val,
+                {
+                    "id": "",
+                    "type": "function",
+                    "function": {"name": "", "arguments": ""},
+                },
+            )
+            tc_id = dict_get_str(d, "id")
+            if tc_id is not None:
+                entry["id"] = tc_id
+            tc_type = dict_get_str(d, "type")
+            if tc_type is not None:
+                entry["type"] = tc_type
+            fn = dict_get_dict(d, "function")
+            if fn is not None:
+                entry_fn = entry.get("function")
+                if not isinstance(entry_fn, dict):
+                    entry_fn = {"name": "", "arguments": ""}
+                    entry["function"] = entry_fn
+                entry_fn_typed = cast(dict[str, object], entry_fn)
+                name = dict_get_str(fn, "name")
+                if name is not None:
+                    prev_name = as_str(entry_fn_typed.get("name")) or ""
+                    entry_fn_typed["name"] = prev_name + name
+                args = dict_get_str(fn, "arguments")
+                if args is not None:
+                    prev_args = as_str(entry_fn_typed.get("arguments")) or ""
+                    entry_fn_typed["arguments"] = prev_args + args
+
+    @property
+    def content(self) -> str | None:
+        joined = "".join(self._content_parts)
+        return joined if joined else None
+
+    @property
+    def tool_calls(self) -> list[dict[str, object]] | None:
+        if not self._tool_calls_by_index:
+            return None
+        ordered = [
+            self._tool_calls_by_index[i] for i in sorted(self._tool_calls_by_index)
+        ]
+        return ordered
+
+    @property
+    def reasoning(self) -> str:
+        return "".join(self._reasoning_parts)
+
+
+class ClaudeAccumulator:
+    """Captures Claude streaming content blocks.
+
+    Tracks per-index content blocks. At end, exposes the final `content_blocks`
+    list (excluding thinking blocks — those go into `reasoning` as joined text)
+    in a shape suitable for hashing.
+    """
+
+    def __init__(self) -> None:
+        self._blocks_by_index: dict[int, dict[str, object]] = {}
+        self._buffer = ""
+        self._current_event: str | None = None
+
+    def feed_bytes(self, chunk: bytes) -> None:
+        self._buffer += chunk.decode("utf-8", errors="replace")
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            line = line.rstrip("\r")
+            if not line:
+                self._current_event = None
+                continue
+            if line.startswith("event:"):
+                self._current_event = line[len("event:") :].strip()
+                continue
+            if line.startswith("data:"):
+                payload = line[len("data:") :].strip()
+                try:
+                    parsed = cast(object, json.loads(payload))
+                except json.JSONDecodeError:
+                    continue
+                event = as_dict(parsed)
+                if event is None:
+                    continue
+                self._consume_event(event)
+
+    def _consume_event(self, event: dict[str, object]) -> None:
+        event_type = dict_get_str(event, "type") or self._current_event
+        if event_type == "content_block_start":
+            index_val = event.get("index", 0)
+            if not (isinstance(index_val, int) and not isinstance(index_val, bool)):
+                return
+            block = dict_get_dict(event, "content_block")
+            if block is None:
+                return
+            btype = dict_get_str(block, "type")
+            if btype == "text":
+                self._blocks_by_index[index_val] = {"type": "text", "text": ""}
+            elif btype == "thinking":
+                self._blocks_by_index[index_val] = {
+                    "type": "thinking",
+                    "thinking": "",
+                }
+            elif btype == "tool_use":
+                self._blocks_by_index[index_val] = {
+                    "type": "tool_use",
+                    "id": dict_get_str(block, "id") or "",
+                    "name": dict_get_str(block, "name") or "",
+                    "input_json": "",
+                }
+        elif event_type == "content_block_delta":
+            index_val = event.get("index", 0)
+            if not (isinstance(index_val, int) and not isinstance(index_val, bool)):
+                return
+            delta = dict_get_dict(event, "delta")
+            if delta is None:
+                return
+            block = self._blocks_by_index.get(index_val)
+            if block is None:
+                return
+            dtype = dict_get_str(delta, "type")
+            if dtype == "text_delta":
+                text = dict_get_str(delta, "text")
+                if text is not None:
+                    prev = as_str(block.get("text")) or ""
+                    block["text"] = prev + text
+            elif dtype == "thinking_delta":
+                thinking = dict_get_str(delta, "thinking")
+                if thinking is not None:
+                    prev = as_str(block.get("thinking")) or ""
+                    block["thinking"] = prev + thinking
+            elif dtype == "input_json_delta":
+                partial = dict_get_str(delta, "partial_json")
+                if partial is not None:
+                    prev = as_str(block.get("input_json")) or ""
+                    block["input_json"] = prev + partial
+
+    @property
+    def content_blocks(self) -> list[dict[str, object]]:
+        """Public blocks (excludes thinking), with tool_use input parsed from JSON."""
+        public: list[dict[str, object]] = []
+        for index in sorted(self._blocks_by_index):
+            block = self._blocks_by_index[index]
+            if block.get("type") == "thinking":
+                continue
+            if block.get("type") == "tool_use":
+                input_json = as_str(block.get("input_json")) or "{}"
+                parsed_input_raw: object
+                try:
+                    parsed_input_raw = cast(object, json.loads(input_json))
+                except json.JSONDecodeError:
+                    parsed_input_raw = {}
+                parsed_input = as_dict(parsed_input_raw) or {}
+                public.append(
+                    {
+                        "type": "tool_use",
+                        "id": as_str(block.get("id")) or "",
+                        "name": as_str(block.get("name")) or "",
+                        "input": parsed_input,
+                    }
+                )
+            else:
+                public.append({k: v for k, v in block.items() if k != "input_json"})
+        return public
+
+    @property
+    def reasoning(self) -> str:
+        parts: list[str] = []
+        for index in sorted(self._blocks_by_index):
+            block = self._blocks_by_index[index]
+            if block.get("type") == "thinking":
+                parts.append(as_str(block.get("thinking")) or "")
+        return "".join(parts)

src/exo/reasoning_proxy/cache.py

@@ -0,0 +1,21 @@
+import threading
+
+
+class ReasoningCache:
+    def __init__(self) -> None:
+        self._store: dict[str, str] = {}
+        self._lock = threading.Lock()
+
+    def get(self, content_hash: str) -> str | None:
+        with self._lock:
+            return self._store.get(content_hash)
+
+    def put(self, content_hash: str, reasoning: str) -> None:
+        if not reasoning:
+            return
+        with self._lock:
+            self._store[content_hash] = reasoning
+
+    def size(self) -> int:
+        with self._lock:
+            return len(self._store)

src/exo/reasoning_proxy/dialects.py

@@ -0,0 +1,170 @@
+"""Dialect strategies for deciding which assistant history indices should receive
+cached reasoning on inbound requests.
+
+Each dialect inspects the message list and returns the set of indices where
+reasoning_content (OpenAI) or a thinking block (Claude) should be reattached if
+the cache has it. The dialect does not mutate messages — the caller does.
+
+Dialect selection is driven by the `reasoning_dialect` field on each model card,
+surfaced through /v1/models.
+"""
+
+from typing import Protocol
+
+from exo.reasoning_proxy._helpers import as_dict, as_list, dict_get_list, dict_get_str
+from exo.shared.types.text_generation import ReasoningDialect
+
+
+class Dialect(Protocol):
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]: ...
+
+
+def _is_assistant(msg: dict[str, object]) -> bool:
+    return msg.get("role") == "assistant"
+
+
+def _is_user(msg: dict[str, object]) -> bool:
+    return msg.get("role") == "user"
+
+
+class NoneDialect:
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        return set()
+
+
+class PostLastUserDialect:
+    """MiniMax / GLM / Qwen-thinking / V4-with-tools.
+
+    Preserve reasoning on every assistant message appearing after the last
+    non-tool-response user message. Tool-response user messages (role=tool, or
+    role=user with tool_call_id set, or Claude's tool_result block) don't count
+    as "real" user turns — they're part of the assistant's tool-calling chain.
+    """
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        last_user_index = -1
+        for i, msg in enumerate(messages):
+            if _is_user(msg) and not _is_tool_response(msg):
+                last_user_index = i
+        return {
+            i
+            for i, msg in enumerate(messages)
+            if i > last_user_index and _is_assistant(msg)
+        }
+
+
+class SuffixDialect:
+    """Kimi K2 Thinking / K2.6.
+
+    Preserve reasoning only on the tail run of tool-call-carrying assistant
+    messages (the current, unresolved tool-call chain). Walk backward: include
+    every assistant with tool_calls until we hit an assistant without tool_calls
+    or a non-assistant message.
+    """
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        indices: set[int] = set()
+        for i in range(len(messages) - 1, -1, -1):
+            msg = messages[i]
+            if not _is_assistant(msg):
+                if _is_tool_response(msg):
+                    continue
+                break
+            if not _has_tool_calls(msg):
+                break
+            indices.add(i)
+        return indices
+
+
+class ChannelDialect:
+    """GPT-OSS Harmony format.
+
+    Preserve analysis-channel content on assistant turns that follow the most
+    recent assistant message tagged with a "final" channel marker. If no prior
+    final exists, the whole conversation is one unresolved chain.
+    """
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        last_final_index = -1
+        for i, msg in enumerate(messages):
+            if _is_assistant(msg) and _has_final_channel(msg):
+                last_final_index = i
+        return {
+            i
+            for i, msg in enumerate(messages)
+            if i > last_final_index and _is_assistant(msg)
+        }
+
+
+class ToolConditionalDialect:
+    """DeepSeek V4 Flash.
+
+    If the request has tools, behave as PostLastUserDialect; otherwise passthrough.
+    """
+
+    def __init__(self) -> None:
+        self._inner = PostLastUserDialect()
+
+    def select_attach_indices(
+        self, messages: list[dict[str, object]], has_tools: bool
+    ) -> set[int]:
+        if not has_tools:
+            return set()
+        return self._inner.select_attach_indices(messages, has_tools)
+
+
+def _is_tool_response(msg: dict[str, object]) -> bool:
+    if msg.get("role") == "tool":
+        return True
+    if msg.get("role") == "user" and msg.get("tool_call_id"):
+        return True
+    content = as_list(msg.get("content"))
+    if content is not None:
+        for raw in content:
+            block = as_dict(raw)
+            if block is not None and block.get("type") == "tool_result":
+                return True
+    return False
+
+
+def _has_tool_calls(msg: dict[str, object]) -> bool:
+    tc = dict_get_list(msg, "tool_calls")
+    if tc:
+        return True
+    content = as_list(msg.get("content"))
+    if content is not None:
+        for raw in content:
+            block = as_dict(raw)
+            if block is not None and block.get("type") == "tool_use":
+                return True
+    return False
+
+
+def _has_final_channel(msg: dict[str, object]) -> bool:
+    if msg.get("channel") == "final":
+        return True
+    content = dict_get_str(msg, "content")
+    return bool(content and content.strip())
+
+
+_DIALECTS: dict[ReasoningDialect, Dialect] = {
+    "none": NoneDialect(),
+    "post_last_user": PostLastUserDialect(),
+    "suffix": SuffixDialect(),
+    "channel": ChannelDialect(),
+    "tool_conditional": ToolConditionalDialect(),
+}
+
+
+def get_dialect(name: ReasoningDialect) -> Dialect:
+    return _DIALECTS[name]

src/exo/reasoning_proxy/hashing.py

@@ -0,0 +1,75 @@
+import hashlib
+import json
+
+from exo.reasoning_proxy._helpers import as_dict, dict_get_str
+
+
+def _canonical_tool_calls(
+    tool_calls: list[dict[str, object]] | None,
+) -> list[dict[str, object]]:
+    if not tool_calls:
+        return []
+    result: list[dict[str, object]] = []
+    for tc in tool_calls:
+        entry: dict[str, object] = {}
+        if "id" in tc:
+            entry["id"] = tc["id"]
+        fn = as_dict(tc.get("function"))
+        if fn is not None:
+            entry["function"] = {
+                "name": dict_get_str(fn, "name") or "",
+                "arguments": dict_get_str(fn, "arguments") or "",
+            }
+        if "type" in tc:
+            entry["type"] = tc["type"]
+        result.append(entry)
+    return result
+
+
+def hash_openai_assistant(
+    content: str | list[object] | None,
+    tool_calls: list[dict[str, object]] | None,
+) -> str:
+    """Deterministic hash of an OpenAI assistant message's observable surface.
+
+    Canonicalizes None content to "" and tool_calls to a minimal id/function shape
+    so trivial shape differences between client render and our re-emit don't miss.
+    """
+    shape: dict[str, object] = {
+        "content": content if content is not None else "",
+        "tool_calls": _canonical_tool_calls(tool_calls),
+    }
+    payload = json.dumps(
+        shape, sort_keys=True, ensure_ascii=False, separators=(",", ":")
+    )
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def hash_claude_assistant(content_blocks: list[dict[str, object]]) -> str:
+    """Deterministic hash of a Claude assistant message's observable surface.
+
+    Skips thinking blocks (we're hashing what the *client sends back*, which typically
+    omits thinking) and normalizes tool_use blocks to id/name/input.
+    """
+    normalized: list[dict[str, object]] = []
+    for block in content_blocks:
+        btype = block.get("type")
+        if btype == "text":
+            normalized.append(
+                {"type": "text", "text": dict_get_str(block, "text") or ""}
+            )
+        elif btype == "tool_use":
+            normalized.append(
+                {
+                    "type": "tool_use",
+                    "id": dict_get_str(block, "id") or "",
+                    "name": dict_get_str(block, "name") or "",
+                    "input": block.get("input")
+                    if block.get("input") is not None
+                    else {},
+                }
+            )
+    payload = json.dumps(
+        normalized, sort_keys=True, ensure_ascii=False, separators=(",", ":")
+    )
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()

src/exo/reasoning_proxy/main.py

@@ -0,0 +1,70 @@
+import argparse
+import logging
+from contextlib import asynccontextmanager
+from typing import cast
+
+import httpx
+import uvicorn
+from fastapi import FastAPI
+
+from exo.reasoning_proxy.cache import ReasoningCache
+from exo.reasoning_proxy.registry import DialectRegistry
+from exo.reasoning_proxy.routes import register_routes
+
+logger = logging.getLogger("exo.reasoning_proxy")
+
+
+def build_app(upstream: str) -> FastAPI:
+    client = httpx.AsyncClient(timeout=httpx.Timeout(None, connect=10.0))
+    cache = ReasoningCache()
+    registry = DialectRegistry(upstream=upstream, client=client)
+
+    @asynccontextmanager
+    async def lifespan(_: FastAPI):
+        await registry.refresh()
+        try:
+            yield
+        finally:
+            await client.aclose()
+
+    app = FastAPI(lifespan=lifespan, title="exo-reasoning-proxy")
+    register_routes(
+        app=app,
+        client=client,
+        upstream=upstream.rstrip("/"),
+        cache=cache,
+        registry=registry,
+    )
+    return app
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(prog="exo-reasoning-proxy")
+    _ = parser.add_argument("--upstream", default="http://localhost:52415")
+    _ = parser.add_argument("--host", default="127.0.0.1")
+    _ = parser.add_argument("--port", type=int, default=52416)
+    _ = parser.add_argument("-v", "--verbose", action="count", default=0)
+    args = parser.parse_args()
+
+    verbose = cast(int, args.verbose)
+    upstream = cast(str, args.upstream)
+    host = cast(str, args.host)
+    port = cast(int, args.port)
+
+    level = logging.WARNING
+    if verbose == 1:
+        level = logging.INFO
+    elif verbose >= 2:
+        level = logging.DEBUG
+    logging.basicConfig(
+        level=level, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
+    )
+
+    logger.info("Starting exo-reasoning-proxy on %s:%d → %s", host, port, upstream)
+
+    app = build_app(upstream=upstream)
+    uvicorn.run(app, host=host, port=port, log_level=level)
+
+
+if __name__ == "__main__":
+    main()

src/exo/reasoning_proxy/registry.py

@@ -0,0 +1,75 @@
+import asyncio
+import logging
+from typing import cast, get_args
+
+import httpx
+
+from exo.reasoning_proxy._helpers import as_dict, as_list, dict_get_str
+from exo.shared.types.text_generation import ReasoningDialect
+
+logger = logging.getLogger(__name__)
+
+
+class DialectRegistry:
+    def __init__(self, upstream: str, client: httpx.AsyncClient) -> None:
+        self._upstream = upstream.rstrip("/")
+        self._client = client
+        self._by_model: dict[str, ReasoningDialect] = {}
+        self._unknown_logged: set[str] = set()
+        self._lock = asyncio.Lock()
+        self._initialized = False
+
+    async def refresh(self) -> None:
+        await self._fetch()
+
+    async def _fetch(self) -> None:
+        try:
+            resp = await self._client.get(f"{self._upstream}/v1/models", timeout=10.0)
+            resp.raise_for_status()
+            body = as_dict(cast(object, resp.json()))
+            if body is None:
+                return
+            data = as_list(body.get("data")) or []
+            updated: dict[str, ReasoningDialect] = {}
+            for entry_raw in data:
+                entry = as_dict(entry_raw)
+                if entry is None:
+                    continue
+                model_id = dict_get_str(entry, "id")
+                dialect_raw = entry.get("reasoning_dialect", "none")
+                if model_id is not None:
+                    updated[model_id] = _coerce_dialect(dialect_raw)
+            self._by_model = updated
+            self._initialized = True
+            logger.info(
+                "Loaded %d model dialects from %s", len(updated), self._upstream
+            )
+        except Exception as exc:
+            logger.warning(
+                "Failed to fetch /v1/models from %s: %s", self._upstream, exc
+            )
+
+    async def resolve(self, model_id: str) -> ReasoningDialect:
+        async with self._lock:
+            if not self._initialized:
+                await self._fetch()
+            if model_id in self._by_model:
+                return self._by_model[model_id]
+            await self._fetch()
+            if model_id in self._by_model:
+                return self._by_model[model_id]
+            if model_id not in self._unknown_logged:
+                logger.info(
+                    "No dialect declared for model %s; passing through", model_id
+                )
+                self._unknown_logged.add(model_id)
+            return "none"
+
+
+_VALID_DIALECTS: frozenset[str] = frozenset(get_args(ReasoningDialect))
+
+
+def _coerce_dialect(value: object) -> ReasoningDialect:
+    if isinstance(value, str) and value in _VALID_DIALECTS:
+        return cast(ReasoningDialect, value)
+    return "none"

src/exo/reasoning_proxy/routes.py

@@ -0,0 +1,384 @@
+"""FastAPI handlers for the reasoning proxy.
+
+Two handlers, one shape: read body → resolve dialect → reattach cached
+reasoning to designated history indices → forward → tee the response stream
+→ capture emitted reasoning → cache under the emitted assistant's hash.
+"""
+
+import json
+import logging
+from collections.abc import AsyncIterator
+from typing import cast
+
+import httpx
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.datastructures import Headers as StarletteHeaders
+
+from exo.reasoning_proxy._helpers import (
+    as_dict,
+    as_list,
+    dict_get_dict,
+    dict_get_list,
+    dict_get_str,
+)
+from exo.reasoning_proxy.accumulator import ClaudeAccumulator, OpenAIAccumulator
+from exo.reasoning_proxy.cache import ReasoningCache
+from exo.reasoning_proxy.dialects import get_dialect
+from exo.reasoning_proxy.hashing import hash_claude_assistant, hash_openai_assistant
+from exo.reasoning_proxy.registry import DialectRegistry
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_body(raw_body: bytes) -> dict[str, object] | None:
+    try:
+        parsed = cast(object, json.loads(raw_body))
+    except json.JSONDecodeError:
+        return None
+    return as_dict(parsed)
+
+
+def _content_for_hash(value: object) -> str | list[object] | None:
+    if isinstance(value, str):
+        return value
+    as_listed = as_list(value)
+    if as_listed is not None:
+        return as_listed
+    return None
+
+
+def _attach_openai_reasoning(
+    messages: list[dict[str, object]],
+    indices: set[int],
+    cache: ReasoningCache,
+) -> None:
+    for i in indices:
+        msg = messages[i]
+        existing = dict_get_str(msg, "reasoning_content")
+        if existing:
+            continue
+        content = _content_for_hash(msg.get("content"))
+        tool_calls_raw = dict_get_list(msg, "tool_calls") or []
+        tool_calls: list[dict[str, object]] = [
+            d for d in (as_dict(t) for t in tool_calls_raw) if d is not None
+        ]
+        h = hash_openai_assistant(content, tool_calls or None)
+        cached = cache.get(h)
+        if cached is not None:
+            msg["reasoning_content"] = cached
+
+
+def _attach_claude_reasoning(
+    messages: list[dict[str, object]],
+    indices: set[int],
+    cache: ReasoningCache,
+) -> None:
+    for i in indices:
+        msg = messages[i]
+        content = as_list(msg.get("content"))
+        if content is None:
+            continue
+        has_thinking = False
+        normalized: list[dict[str, object]] = []
+        for raw in content:
+            block = as_dict(raw)
+            if block is None:
+                continue
+            if block.get("type") == "thinking":
+                has_thinking = True
+            normalized.append(block)
+        if has_thinking:
+            continue
+        h = hash_claude_assistant(normalized)
+        cached = cache.get(h)
+        if cached is None:
+            continue
+        new_content: list[dict[str, object]] = [
+            {"type": "thinking", "thinking": cached}
+        ]
+        new_content.extend(normalized)
+        msg["content"] = new_content
+
+
+async def _stream_and_capture_openai(
+    upstream_resp: httpx.Response,
+    cache: ReasoningCache,
+) -> AsyncIterator[bytes]:
+    accumulator = OpenAIAccumulator()
+    try:
+        async for chunk in upstream_resp.aiter_raw():
+            accumulator.feed_bytes(chunk)
+            yield chunk
+    finally:
+        await upstream_resp.aclose()
+    reasoning = accumulator.reasoning
+    if not reasoning:
+        return
+    h = hash_openai_assistant(accumulator.content, accumulator.tool_calls)
+    cache.put(h, reasoning)
+
+
+async def _stream_and_capture_claude(
+    upstream_resp: httpx.Response,
+    cache: ReasoningCache,
+) -> AsyncIterator[bytes]:
+    accumulator = ClaudeAccumulator()
+    try:
+        async for chunk in upstream_resp.aiter_raw():
+            accumulator.feed_bytes(chunk)
+            yield chunk
+    finally:
+        await upstream_resp.aclose()
+    reasoning = accumulator.reasoning
+    if not reasoning:
+        return
+    h = hash_claude_assistant(accumulator.content_blocks)
+    cache.put(h, reasoning)
+
+
+def _capture_openai_nonstream(body_text: str, cache: ReasoningCache) -> None:
+    body = _parse_body(body_text.encode("utf-8"))
+    if body is None:
+        return
+    choices = dict_get_list(body, "choices")
+    if not choices:
+        return
+    first = as_dict(choices[0])
+    if first is None:
+        return
+    message = dict_get_dict(first, "message")
+    if message is None:
+        return
+    reasoning = dict_get_str(message, "reasoning_content")
+    if not reasoning:
+        return
+    content = _content_for_hash(message.get("content"))
+    tool_calls_raw = dict_get_list(message, "tool_calls") or []
+    tool_calls: list[dict[str, object]] = [
+        d for d in (as_dict(t) for t in tool_calls_raw) if d is not None
+    ]
+    h = hash_openai_assistant(content, tool_calls or None)
+    cache.put(h, reasoning)
+
+
+def _capture_claude_nonstream(body_text: str, cache: ReasoningCache) -> None:
+    body = _parse_body(body_text.encode("utf-8"))
+    if body is None:
+        return
+    content = as_list(body.get("content"))
+    if content is None:
+        return
+    reasoning_parts: list[str] = []
+    public_blocks: list[dict[str, object]] = []
+    for raw in content:
+        block = as_dict(raw)
+        if block is None:
+            continue
+        if block.get("type") == "thinking":
+            thinking = dict_get_str(block, "thinking")
+            if thinking is not None:
+                reasoning_parts.append(thinking)
+        else:
+            public_blocks.append(block)
+    reasoning = "".join(reasoning_parts)
+    if not reasoning:
+        return
+    h = hash_claude_assistant(public_blocks)
+    cache.put(h, reasoning)
+
+
+def _messages_from_body(body: dict[str, object]) -> list[dict[str, object]] | None:
+    raw = as_list(body.get("messages"))
+    if raw is None:
+        return None
+    result: list[dict[str, object]] = []
+    for item in raw:
+        m = as_dict(item)
+        if m is None:
+            return None
+        result.append(m)
+    return result
+
+
+def register_routes(
+    app: FastAPI,
+    client: httpx.AsyncClient,
+    upstream: str,
+    cache: ReasoningCache,
+    registry: DialectRegistry,
+) -> None:
+    async def handle_chat_completions(request: Request) -> Response:
+        raw_body = await request.body()
+        body = _parse_body(raw_body)
+        if body is None:
+            return _bad_request("invalid JSON body")
+
+        model_id = dict_get_str(body, "model")
+        if model_id is None:
+            return _bad_request("missing or invalid 'model' field")
+
+        dialect_name = await registry.resolve(model_id)
+        dialect = get_dialect(dialect_name)
+
+        messages = _messages_from_body(body)
+        if messages is not None:
+            has_tools = bool(body.get("tools"))
+            indices = dialect.select_attach_indices(messages, has_tools=has_tools)
+            if indices:
+                _attach_openai_reasoning(messages, indices, cache)
+                body["messages"] = messages
+
+        forward_body = json.dumps(body).encode("utf-8")
+        forward_headers = _copy_headers(request.headers)
+        forward_headers["content-length"] = str(len(forward_body))
+
+        is_stream = bool(body.get("stream"))
+
+        try:
+            if is_stream:
+                req = client.build_request(
+                    "POST",
+                    f"{upstream}/v1/chat/completions",
+                    content=forward_body,
+                    headers=forward_headers,
+                )
+                upstream_resp = await client.send(req, stream=True)
+                return StreamingResponse(
+                    _stream_and_capture_openai(upstream_resp, cache),
+                    status_code=upstream_resp.status_code,
+                    media_type=_media_type(upstream_resp.headers, "text/event-stream"),
+                    headers=_response_headers(upstream_resp.headers),
+                )
+            upstream_resp = await client.post(
+                f"{upstream}/v1/chat/completions",
+                content=forward_body,
+                headers=forward_headers,
+            )
+            text = upstream_resp.text
+            if upstream_resp.status_code == 200:
+                _capture_openai_nonstream(text, cache)
+            return Response(
+                content=text,
+                status_code=upstream_resp.status_code,
+                media_type=_media_type(upstream_resp.headers, "application/json"),
+                headers=_response_headers(upstream_resp.headers),
+            )
+        except httpx.RequestError as exc:
+            logger.warning("Upstream request failed: %s", exc)
+            return _bad_gateway(str(exc))
+
+    async def handle_claude_messages(request: Request) -> Response:
+        raw_body = await request.body()
+        body = _parse_body(raw_body)
+        if body is None:
+            return _bad_request("invalid JSON body")
+
+        model_id = dict_get_str(body, "model")
+        if model_id is None:
+            return _bad_request("missing or invalid 'model' field")
+
+        dialect_name = await registry.resolve(model_id)
+        dialect = get_dialect(dialect_name)
+
+        messages = _messages_from_body(body)
+        if messages is not None:
+            has_tools = bool(body.get("tools"))
+            indices = dialect.select_attach_indices(messages, has_tools=has_tools)
+            if indices:
+                _attach_claude_reasoning(messages, indices, cache)
+                body["messages"] = messages
+
+        forward_body = json.dumps(body).encode("utf-8")
+        forward_headers = _copy_headers(request.headers)
+        forward_headers["content-length"] = str(len(forward_body))
+
+        is_stream = bool(body.get("stream"))
+
+        try:
+            if is_stream:
+                req = client.build_request(
+                    "POST",
+                    f"{upstream}/v1/messages",
+                    content=forward_body,
+                    headers=forward_headers,
+                )
+                upstream_resp = await client.send(req, stream=True)
+                return StreamingResponse(
+                    _stream_and_capture_claude(upstream_resp, cache),
+                    status_code=upstream_resp.status_code,
+                    media_type=_media_type(upstream_resp.headers, "text/event-stream"),
+                    headers=_response_headers(upstream_resp.headers),
+                )
+            upstream_resp = await client.post(
+                f"{upstream}/v1/messages",
+                content=forward_body,
+                headers=forward_headers,
+            )
+            text = upstream_resp.text
+            if upstream_resp.status_code == 200:
+                _capture_claude_nonstream(text, cache)
+            return Response(
+                content=text,
+                status_code=upstream_resp.status_code,
+                media_type=_media_type(upstream_resp.headers, "application/json"),
+                headers=_response_headers(upstream_resp.headers),
+            )
+        except httpx.RequestError as exc:
+            logger.warning("Upstream request failed: %s", exc)
+            return _bad_gateway(str(exc))
+
+    async def health() -> dict[str, object]:
+        return {"status": "ok", "cache_entries": cache.size()}
+
+    _ = app.post("/v1/chat/completions")(handle_chat_completions)
+    _ = app.post("/v1/messages")(handle_claude_messages)
+    _ = app.get("/health")(health)
+
+
+_HOP_BY_HOP = {
+    "connection",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailers",
+    "transfer-encoding",
+    "upgrade",
+    "content-length",
+    "host",
+}
+
+
+def _copy_headers(headers: httpx.Headers | StarletteHeaders) -> dict[str, str]:
+    out: dict[str, str] = {}
+    for k, v in headers.items():
+        if k.lower() in _HOP_BY_HOP:
+            continue
+        out[k] = v
+    return out
+
+
+def _response_headers(headers: httpx.Headers) -> dict[str, str]:
+    return _copy_headers(headers)
+
+
+def _media_type(headers: httpx.Headers, default: str) -> str:
+    value = cast(object, headers.get("content-type", default))
+    return value if isinstance(value, str) else default
+
+
+def _bad_request(msg: str) -> JSONResponse:
+    return JSONResponse(
+        status_code=400,
+        content={"error": {"message": msg, "type": "invalid_request_error"}},
+    )
+
+
+def _bad_gateway(msg: str) -> JSONResponse:
+    return JSONResponse(
+        status_code=502,
+        content={
+            "error": {"message": f"upstream unreachable: {msg}", "type": "bad_gateway"}
+        },
+    )

src/exo/routing/event_router.py

@@ -80,9 +80,6 @@ class EventRouter:
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 
-    def set_buffer_start(self, idx: int) -> None:
-        self.event_buffer.fast_forward_to(idx)
-
     async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
         idx = 0
         with recv as events:
@@ -98,6 +95,7 @@ class EventRouter:
                 self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
 
     async def _run_ext_in(self):
+        buf = OrderedBuffer[Event]()
         with self.external_inbound as events:
             async for event in events:
                 if event.session != self.session_id:
@@ -105,12 +103,12 @@ class EventRouter:
                 if event.origin != self.session_id.master_node_id:
                     continue
 
-                self.event_buffer.ingest(event.origin_idx, event.event)
+                buf.ingest(event.origin_idx, event.event)
                 event_id = event.event.event_id
                 if event_id in self.out_for_delivery:
                     self.out_for_delivery.pop(event_id)
 
-                drained = self.event_buffer.drain_indexed()
+                drained = buf.drain_indexed()
                 if drained:
                     self._nack_attempts = 0
                     if self._nack_cancel_scope:
@@ -121,9 +119,7 @@ class EventRouter:
                     or self._nack_cancel_scope.cancel_called
                 ):
                     # Request the next index.
-                    self._tg.start_soon(
-                        self._nack_request, self.event_buffer.next_idx_to_release
-                    )
+                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
                     continue
 
                 for idx, event in drained:

src/exo/routing/snapshot_receiver.py

@@ -1,109 +0,0 @@
-"""Reassembles a snapshot from a stream of `SnapshotChunk`s.
-
-A receiver belongs to one node; it ignores chunks addressed to other
-requesters and chunks from prior sessions. Once a transfer's chunks have
-all been collected and the SHA-256 checks out, the snapshot is decoded into
-a `State`. Concurrent transfers (for the same requester) are tolerated:
-each is keyed by `transfer_id`.
-"""
-
-from __future__ import annotations
-
-import hashlib
-from dataclasses import dataclass, field
-from typing import final
-
-import zstandard
-from loguru import logger
-
-from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-
-
-@final
-@dataclass
-class _Assembly:
-    """Partial state for one in-flight snapshot transfer."""
-
-    total_chunks: int
-    sha256_hex: str
-    schema_version: int
-    last_event_applied_idx: int
-    chunks: dict[int, bytes] = field(default_factory=dict)
-
-    def is_complete(self) -> bool:
-        return len(self.chunks) == self.total_chunks
-
-    def assemble(self) -> bytes:
-        return b"".join(self.chunks[i] for i in range(self.total_chunks))
-
-
-@dataclass
-class ReceivedSnapshot:
-    last_event_applied_idx: int
-    state: State
-
-
-class SnapshotReceiver:
-    """Filters and reassembles inbound chunks into a `ReceivedSnapshot`.
-
-    Stateless w.r.t. delivery: callers feed `SnapshotChunk`s in via `ingest`
-    and check the return value for completion.
-    """
-
-    def __init__(self, my_node_id: NodeId, session_id: SessionId) -> None:
-        self._my_node_id = my_node_id
-        self._session_id = session_id
-        self._assemblies: dict[SnapshotTransferId, _Assembly] = {}
-
-    def ingest(self, chunk: SnapshotChunk) -> ReceivedSnapshot | None:
-        """Absorb a chunk; return the snapshot once a transfer completes.
-
-        Returns None for partial transfers, mismatched recipients, stale
-        sessions, version mismatches, or corrupt payloads.
-        """
-        if chunk.requester_node_id != self._my_node_id:
-            return None
-        if chunk.session_id != self._session_id:
-            return None
-
-        existing = self._assemblies.get(chunk.transfer_id)
-        if existing is None:
-            existing = _Assembly(
-                total_chunks=chunk.total_chunks,
-                sha256_hex=chunk.sha256_hex,
-                schema_version=chunk.schema_version,
-                last_event_applied_idx=chunk.last_event_applied_idx,
-            )
-            self._assemblies[chunk.transfer_id] = existing
-        existing.chunks[chunk.chunk_index] = chunk.data
-
-        if not existing.is_complete():
-            return None
-
-        # Transfer complete — finalise and remove from the in-flight map.
-        del self._assemblies[chunk.transfer_id]
-        body = existing.assemble()
-        if hashlib.sha256(body).hexdigest() != existing.sha256_hex:
-            logger.warning(f"Snapshot {chunk.transfer_id} failed checksum; discarding")
-            return None
-        try:
-            decompressed = zstandard.ZstdDecompressor().decompress(body)
-            state = State.model_validate_json(decompressed.decode("utf-8"))
-        except (zstandard.ZstdError, ValueError) as e:
-            logger.opt(exception=e).warning(
-                f"Snapshot {chunk.transfer_id} could not be decoded; discarding"
-            )
-            return None
-        if state.schema_version != existing.schema_version:
-            # Should not happen — the master writes schema_version into both
-            # the chunk meta and the State payload — but treat it as corrupt.
-            logger.warning(
-                f"Snapshot {chunk.transfer_id} schema version mismatch "
-                f"(chunk={existing.schema_version}, state={state.schema_version})"
-            )
-            return None
-        return ReceivedSnapshot(
-            last_event_applied_idx=existing.last_event_applied_idx, state=state
-        )

src/exo/routing/tests/test_event_buffer.py

@@ -141,28 +141,3 @@ async def test_drain_and_ingest_with_new_sequence(buffer: OrderedBuffer[Event]):
     assert [e[0] for e in drained] == [2]
     assert buffer.next_idx_to_release == 3
     assert 4 in buffer.store
-
-
-@pytest.mark.asyncio
-async def test_fast_forward_discards_buffered_stale_events(
-    buffer: OrderedBuffer[Event],
-):
-    buffer.ingest(*make_indexed_event(0))
-    buffer.ingest(*make_indexed_event(2))
-    buffer.ingest(*make_indexed_event(4))
-
-    buffer.fast_forward_to(3)
-
-    assert buffer.next_idx_to_release == 3
-    assert set(buffer.store) == {4}
-
-
-@pytest.mark.asyncio
-async def test_fast_forward_only_moves_forward(buffer: OrderedBuffer[Event]):
-    buffer.ingest(*make_indexed_event(0))
-    buffer.ingest(*make_indexed_event(1))
-    buffer.drain()
-
-    buffer.fast_forward_to(1)
-
-    assert buffer.next_idx_to_release == 2

src/exo/routing/tests/test_snapshot_receiver.py

@@ -1,151 +0,0 @@
-import hashlib
-
-import pytest
-import zstandard
-
-from exo.routing.snapshot_receiver import SnapshotReceiver
-from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-from exo.shared.types.state import State
-
-
-@pytest.fixture
-def session_id() -> SessionId:
-    return SessionId(master_node_id=NodeId("master"), election_clock=0)
-
-
-@pytest.fixture
-def my_node() -> NodeId:
-    return NodeId("worker-1")
-
-
-def _encode(state: State) -> bytes:
-    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
-
-
-def _make_chunks(
-    body: bytes,
-    *,
-    chunk_size: int,
-    requester_node_id: NodeId,
-    session_id: SessionId,
-    state: State,
-    transfer_id: SnapshotTransferId | None = None,
-) -> list[SnapshotChunk]:
-    sha256 = hashlib.sha256(body).hexdigest()
-    transfer_id = transfer_id or SnapshotTransferId()
-    pieces = [body[i : i + chunk_size] for i in range(0, len(body), chunk_size)] or [
-        b""
-    ]
-    return [
-        SnapshotChunk.from_data(
-            data=piece,
-            transfer_id=transfer_id,
-            requester_node_id=requester_node_id,
-            session_id=session_id,
-            schema_version=state.schema_version,
-            last_event_applied_idx=state.last_event_applied_idx,
-            chunk_index=i,
-            total_chunks=len(pieces),
-            sha256_hex=sha256,
-        )
-        for i, piece in enumerate(pieces)
-    ]
-
-
-def test_completes_on_full_transfer(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=42)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=my_node,
-        session_id=session_id,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-
-    received = None
-    for chunk in chunks:
-        received = receiver.ingest(chunk)
-    assert received is not None
-    assert received.last_event_applied_idx == 42
-    assert received.state.last_event_applied_idx == 42
-
-
-def test_handles_out_of_order_chunks(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=99)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=32,
-        requester_node_id=my_node,
-        session_id=session_id,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-
-    # Reverse them.
-    received = None
-    for chunk in reversed(chunks):
-        received = receiver.ingest(chunk)
-    assert received is not None
-    assert received.last_event_applied_idx == 99
-
-
-def test_ignores_chunks_for_other_recipients(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=1)
-    other = NodeId("worker-2")
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=other,
-        session_id=session_id,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-    for chunk in chunks:
-        assert receiver.ingest(chunk) is None
-
-
-def test_ignores_chunks_from_stale_session(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=1)
-    other_session = SessionId(master_node_id=NodeId("other-master"), election_clock=99)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=my_node,
-        session_id=other_session,
-        state=state,
-    )
-    receiver = SnapshotReceiver(my_node, session_id)
-    for chunk in chunks:
-        assert receiver.ingest(chunk) is None
-
-
-def test_discards_on_checksum_mismatch(my_node: NodeId, session_id: SessionId):
-    state = State(last_event_applied_idx=1)
-    chunks = _make_chunks(
-        _encode(state),
-        chunk_size=64,
-        requester_node_id=my_node,
-        session_id=session_id,
-        state=state,
-    )
-    # Corrupt the last byte of the last chunk.
-    original = chunks[-1]
-    chunks[-1] = SnapshotChunk.from_data(
-        data=original.data + b"\x00garbage",
-        transfer_id=original.transfer_id,
-        requester_node_id=original.requester_node_id,
-        session_id=original.session_id,
-        schema_version=original.schema_version,
-        last_event_applied_idx=original.last_event_applied_idx,
-        chunk_index=original.chunk_index,
-        total_chunks=original.total_chunks,
-        sha256_hex=original.sha256_hex,
-    )
-
-    receiver = SnapshotReceiver(my_node, session_id)
-    received = None
-    for chunk in chunks:
-        received = receiver.ingest(chunk)
-    assert received is None

src/exo/routing/tests/test_snapshot_routing_types.py

@@ -1,37 +0,0 @@
-from exo.routing import topics
-from exo.shared.types.commands import ForwarderCommand, RequestSnapshot
-from exo.shared.types.common import NodeId, SessionId, SystemId
-from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
-
-
-def test_request_snapshot_round_trips_through_forwarder_command() -> None:
-    command = ForwarderCommand(
-        origin=SystemId("system-1"),
-        command=RequestSnapshot(requester_node_id=NodeId("worker-1")),
-    )
-
-    restored = ForwarderCommand.model_validate_json(command.model_dump_json())
-
-    assert isinstance(restored.command, RequestSnapshot)
-    assert restored.command.requester_node_id == NodeId("worker-1")
-
-
-def test_snapshot_response_topic_round_trips_chunk() -> None:
-    chunk = SnapshotChunk.from_data(
-        data=b"snapshot-bytes",
-        transfer_id=SnapshotTransferId("transfer-1"),
-        requester_node_id=NodeId("worker-1"),
-        session_id=SessionId(master_node_id=NodeId("master"), election_clock=1),
-        schema_version=1,
-        last_event_applied_idx=42,
-        chunk_index=0,
-        total_chunks=1,
-        sha256_hex="unused",
-    )
-
-    restored = topics.SNAPSHOT_RESPONSES.deserialize(
-        topics.SNAPSHOT_RESPONSES.serialize(chunk)
-    )
-
-    assert restored == chunk
-    assert restored.data == b"snapshot-bytes"

src/exo/routing/topics.py

@@ -8,7 +8,6 @@ from exo.shared.types.events import (
     GlobalForwarderEvent,
     LocalForwarderEvent,
 )
-from exo.shared.types.snapshots import SnapshotChunk
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -50,6 +49,3 @@ CONNECTION_MESSAGES = TypedTopic(
 DOWNLOAD_COMMANDS = TypedTopic(
     "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
 )
-SNAPSHOT_RESPONSES = TypedTopic(
-    "snapshot_responses", PublishPolicy.Always, SnapshotChunk
-)

src/exo/shared/apply.py

@@ -14,8 +14,6 @@ from exo.shared.types.events import (
     InputChunkReceived,
     InstanceCreated,
     InstanceDeleted,
-    InstanceLinkCreated,
-    InstanceLinkDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -31,7 +29,6 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     NodeIdentity,
     NodeNetworkInfo,
@@ -40,23 +37,11 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    ImageEdits,
-    ImageGeneration,
-    Task,
-    TaskId,
-    TaskStatus,
-    TextGeneration,
-)
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
-from exo.shared.types.worker.runners import (
-    RunnerId,
-    RunnerReady,
-    RunnerShutdown,
-    RunnerStatus,
-)
+from exo.shared.types.worker.runners import RunnerId, RunnerShutdown, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import (
     MacmonMetrics,
     MacThunderboltConnections,
@@ -79,6 +64,7 @@ def event_apply(event: Event, state: State) -> State:
             TestEvent()
             | ChunkGenerated()
             | TaskAcknowledged()
+            | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
             | CustomModelCardAdded()
@@ -99,8 +85,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_runner_status_updated(event, state)
         case TaskCreated():
             return apply_task_created(event, state)
-        case InputChunkReceived():
-            return apply_input_chunk_received(event, state)
         case TaskDeleted():
             return apply_task_deleted(event, state)
         case TaskFailed():
@@ -111,10 +95,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_topology_edge_created(event, state)
         case TopologyEdgeDeleted():
             return apply_topology_edge_deleted(event, state)
-        case InstanceLinkCreated():
-            return apply_instance_link_created(event, state)
-        case InstanceLinkDeleted():
-            return apply_instance_link_deleted(event, state)
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -165,32 +145,10 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
     return state.model_copy(update={"tasks": new_tasks})
 
 
-def apply_input_chunk_received(event: InputChunkReceived, state: State) -> State:
-    command_chunks = {
-        **state.input_chunks.get(event.command_id, {}),
-        event.chunk.chunk_index: event.chunk,
-    }
-    return state.model_copy(
-        update={
-            "input_chunks": {**state.input_chunks, event.command_id: command_chunks}
-        }
-    )
-
-
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
-    task = state.tasks.get(event.task_id)
     new_tasks: Mapping[TaskId, Task] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
-    if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
-        new_input_chunks = {
-            command_id: chunks
-            for command_id, chunks in state.input_chunks.items()
-            if command_id != task.command_id
-        }
-        return state.model_copy(
-            update={"tasks": new_tasks, "input_chunks": new_input_chunks}
-        )
     return state.model_copy(update={"tasks": new_tasks})
 
 
@@ -236,38 +194,7 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     new_instances: Mapping[InstanceId, Instance] = {
         iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
     }
-    new_links: dict[InstanceLinkId, InstanceLink] = {}
-    for link_id, link in state.instance_links.items():
-        prefill = [i for i in link.prefill_instances if i != event.instance_id]
-        decode = [i for i in link.decode_instances if i != event.instance_id]
-        if not prefill or not decode:
-            continue
-        if prefill == list(link.prefill_instances) and decode == list(
-            link.decode_instances
-        ):
-            new_links[link_id] = link
-        else:
-            new_links[link_id] = link.model_copy(
-                update={"prefill_instances": prefill, "decode_instances": decode}
-            )
-    return state.model_copy(
-        update={"instances": new_instances, "instance_links": new_links}
-    )
-
-
-def apply_instance_link_created(event: InstanceLinkCreated, state: State) -> State:
-    new_links: Mapping[InstanceLinkId, InstanceLink] = {
-        **state.instance_links,
-        event.link.link_id: event.link,
-    }
-    return state.model_copy(update={"instance_links": new_links})
-
-
-def apply_instance_link_deleted(event: InstanceLinkDeleted, state: State) -> State:
-    new_links: Mapping[InstanceLinkId, InstanceLink] = {
-        lid: link for lid, link in state.instance_links.items() if lid != event.link_id
-    }
-    return state.model_copy(update={"instance_links": new_links})
+    return state.model_copy(update={"instances": new_instances})
 
 
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
@@ -275,28 +202,12 @@ def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> Sta
         new_runners: Mapping[RunnerId, RunnerStatus] = {
             rid: rs for rid, rs in state.runners.items() if rid != event.runner_id
         }
-        new_ports: Mapping[RunnerId, int] = {
-            rid: p
-            for rid, p in state.prefill_server_ports.items()
-            if rid != event.runner_id
-        }
-        return state.model_copy(
-            update={"runners": new_runners, "prefill_server_ports": new_ports}
-        )
+        return state.model_copy(update={"runners": new_runners})
     new_runners = {
         **state.runners,
         event.runner_id: event.runner_status,
     }
-    update: dict[str, object] = {"runners": new_runners}
-    if (
-        isinstance(event.runner_status, RunnerReady)
-        and event.runner_status.prefill_server_port is not None
-    ):
-        update["prefill_server_ports"] = {
-            **state.prefill_server_ports,
-            event.runner_id: event.runner_status.prefill_server_port,
-        }
-    return state.model_copy(update=update)
+    return state.model_copy(update={"runners": new_runners})
 
 
 def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:

src/exo/shared/constants.py

@@ -96,8 +96,6 @@ EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
 
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
 
-ENABLE_DISAGGREGATION = os.getenv("ENABLE_DISAGGREGATION", "false").lower() == "true"
-
 EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
 
 EXO_MAX_INSTANCE_RETRIES = 5

src/exo/shared/tests/test_apply/test_apply_input_chunks.py

@@ -1,85 +0,0 @@
-from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId
-from exo.shared.types.events import (
-    IndexedEvent,
-    InputChunkReceived,
-    TaskCreated,
-    TaskDeleted,
-)
-from exo.shared.types.state import State
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import (
-    InputMessage,
-    InputMessageContent,
-    TextGenerationTaskParams,
-)
-from exo.shared.types.worker.instances import InstanceId
-
-
-def test_apply_input_chunk_received_stores_chunk_in_state() -> None:
-    command_id = CommandId("command-1")
-    chunk = InputImageChunk(
-        model=ModelId("mlx-community/test-model"),
-        command_id=command_id,
-        data="abc",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-
-    state = apply(
-        State(),
-        IndexedEvent(
-            idx=0,
-            event=InputChunkReceived(command_id=command_id, chunk=chunk),
-        ),
-    )
-
-    assert state.input_chunks == {command_id: {0: chunk}}
-
-
-def test_apply_task_deleted_removes_chunks_for_generation_command() -> None:
-    command_id = CommandId("command-1")
-    task_id = TaskId("task-1")
-    chunk = InputImageChunk(
-        model=ModelId("mlx-community/test-model"),
-        command_id=command_id,
-        data="abc",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-    task = TextGeneration(
-        task_id=task_id,
-        instance_id=InstanceId("instance-1"),
-        task_status=TaskStatus.Pending,
-        command_id=command_id,
-        task_params=TextGenerationTaskParams(
-            model=ModelId("mlx-community/test-model"),
-            input=[
-                InputMessage(role="user", content=InputMessageContent("hello")),
-            ],
-        ),
-    )
-
-    state = State()
-    state = apply(
-        state,
-        IndexedEvent(
-            idx=0,
-            event=InputChunkReceived(command_id=command_id, chunk=chunk),
-        ),
-    )
-    state = apply(
-        state,
-        IndexedEvent(idx=1, event=TaskCreated(task_id=task_id, task=task)),
-    )
-    state = apply(
-        state,
-        IndexedEvent(idx=2, event=TaskDeleted(task_id=task_id)),
-    )
-
-    assert state.tasks == {}
-    assert state.input_chunks == {}

src/exo/shared/tests/test_apply/test_apply_instance_link.py

@@ -1,72 +0,0 @@
-from exo.shared.apply import (
-    apply_instance_deleted,
-    apply_instance_link_created,
-    apply_instance_link_deleted,
-)
-from exo.shared.types.events import (
-    InstanceDeleted,
-    InstanceLinkCreated,
-    InstanceLinkDeleted,
-)
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
-from exo.shared.types.state import State
-from exo.shared.types.worker.instances import InstanceId
-
-
-def _link(
-    prefill: list[InstanceId],
-    decode: list[InstanceId],
-    link_id: InstanceLinkId | None = None,
-) -> InstanceLink:
-    return InstanceLink(
-        link_id=link_id or InstanceLinkId(),
-        prefill_instances=prefill,
-        decode_instances=decode,
-    )
-
-
-def test_create_link() -> None:
-    state = State()
-    link = _link([InstanceId("a")], [InstanceId("b")])
-    new_state = apply_instance_link_created(InstanceLinkCreated(link=link), state)
-    assert new_state.instance_links == {link.link_id: link}
-
-
-def test_update_replaces_existing_link() -> None:
-    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
-    link = _link([a], [b])
-    state = State(instance_links={link.link_id: link})
-
-    updated = link.model_copy(update={"decode_instances": [b, c]})
-    new_state = apply_instance_link_created(InstanceLinkCreated(link=updated), state)
-    assert set(new_state.instance_links[link.link_id].decode_instances) == {b, c}
-
-
-def test_delete_link() -> None:
-    link = _link([InstanceId("a")], [InstanceId("b")])
-    state = State(instance_links={link.link_id: link})
-
-    new_state = apply_instance_link_deleted(
-        InstanceLinkDeleted(link_id=link.link_id), state
-    )
-    assert new_state.instance_links == {}
-
-
-def test_instance_deleted_strips_from_links() -> None:
-    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
-    link = _link([a, c], [b])
-    state = State(instance_links={link.link_id: link})
-
-    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
-    remaining = new_state.instance_links[link.link_id]
-    assert remaining.prefill_instances == [c]
-    assert remaining.decode_instances == [b]
-
-
-def test_instance_deleted_drops_link_when_role_empties() -> None:
-    a, b = InstanceId("a"), InstanceId("b")
-    link = _link([a], [b])
-    state = State(instance_links={link.link_id: link})
-
-    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
-    assert link.link_id not in new_state.instance_links

src/exo/shared/tests/test_state_serialization.py

@@ -25,7 +25,6 @@ def test_state_serialization_roundtrip() -> None:
     json_repr = state.model_dump_json()
     restored_state = State.model_validate_json(json_repr)
 
-    assert restored_state.schema_version == state.schema_version
     assert (
         state.topology.to_snapshot().nodes
         == restored_state.topology.to_snapshot().nodes

src/exo/shared/types/chunks.py

@@ -85,6 +85,6 @@ class PrefillProgressChunk(BaseChunk):
     total_tokens: int
 
 
-StatusChunk = PrefillProgressChunk
-GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
-Chunk = StatusChunk | GenerationChunk
+GenerationChunk = (
+    TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk | PrefillProgressChunk
+)

src/exo/shared/types/commands.py

@@ -7,7 +7,6 @@ from exo.api.types import (
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
-from exo.shared.types.instance_link import InstanceLinkId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -67,12 +66,6 @@ class RequestEventLog(BaseCommand):
     since_idx: int
 
 
-class RequestSnapshot(BaseCommand):
-    """Ask the current master to send a State snapshot to this node."""
-
-    requester_node_id: NodeId
-
-
 class StartDownload(BaseCommand):
     target_node_id: NodeId
     shard_metadata: ShardMetadata
@@ -96,23 +89,12 @@ class DeleteCustomModelCard(BaseCommand):
     model_id: ModelId
 
 
-class SetInstanceLink(BaseCommand):
-    link_id: InstanceLinkId
-    prefill_instances: list[InstanceId]
-    decode_instances: list[InstanceId]
-
-
-class DeleteInstanceLink(BaseCommand):
-    link_id: InstanceLinkId
-
-
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
 Command = (
     TestCommand
     | RequestEventLog
-    | RequestSnapshot
     | TextGeneration
     | ImageGeneration
     | ImageEdits
@@ -124,8 +106,6 @@ Command = (
     | SendInputChunk
     | AddCustomModelCard
     | DeleteCustomModelCard
-    | SetInstanceLink
-    | DeleteInstanceLink
 )
 
 

src/exo/shared/types/events.py

@@ -5,9 +5,8 @@ from pydantic import Field
 
 from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
-from exo.shared.types.chunks import Chunk, InputImageChunk
+from exo.shared.types.chunks import GenerationChunk, InputImageChunk
 from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -92,7 +91,7 @@ class NodeDownloadProgress(BaseEvent):
 
 class ChunkGenerated(BaseEvent):
     command_id: CommandId
-    chunk: Chunk
+    chunk: GenerationChunk
 
 
 class InputChunkReceived(BaseEvent):
@@ -138,14 +137,6 @@ class TracesMerged(BaseEvent):
     traces: list[TraceEventData]
 
 
-class InstanceLinkCreated(BaseEvent):
-    link: InstanceLink
-
-
-class InstanceLinkDeleted(BaseEvent):
-    link_id: InstanceLinkId
-
-
 Event = (
     TestEvent
     | TaskCreated
@@ -167,8 +158,6 @@ Event = (
     | TracesMerged
     | CustomModelCardAdded
     | CustomModelCardDeleted
-    | InstanceLinkCreated
-    | InstanceLinkDeleted
 )
 
 

src/exo/shared/types/instance_link.py

@@ -1,13 +0,0 @@
-from exo.shared.types.common import Id
-from exo.shared.types.worker.instances import InstanceId
-from exo.utils.pydantic_ext import FrozenModel
-
-
-class InstanceLinkId(Id):
-    pass
-
-
-class InstanceLink(FrozenModel):
-    link_id: InstanceLinkId
-    prefill_instances: list[InstanceId]
-    decode_instances: list[InstanceId]

src/exo/shared/types/mlx.py

@@ -11,16 +11,10 @@ from mlx_lm.models.cache import (
     QuantizedKVCache,
     RotatingKVCache,
 )
-from mlx_lm.models.deepseek_v4 import DeepseekV4Cache
 
 # This list contains one cache entry per transformer layer
 KVCacheType = Sequence[
-    KVCache
-    | RotatingKVCache
-    | QuantizedKVCache
-    | ArraysCache
-    | CacheList
-    | DeepseekV4Cache
+    KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
 ]
 
 

src/exo/shared/types/snapshots.py

@@ -1,54 +0,0 @@
-"""Wire types for snapshot transfer between master and a joining node.
-
-Snapshots can be tens of MB; the gossipsub message ceiling is around 1 MB.
-We slice the compressed snapshot body into chunks and publish each chunk on
-the SNAPSHOT_RESPONSES topic. The receiver collects chunks for its own
-`requester_node_id`, validates the SHA-256 of the reassembled body, and
-materialises the State.
-"""
-
-import base64
-
-from exo.shared.types.common import Id, NodeId, SessionId
-from exo.utils.pydantic_ext import FrozenModel
-
-
-class SnapshotTransferId(Id):
-    """Identifies a single snapshot transfer (one master response to one
-    `RequestSnapshot`). Distinct transfers may interleave; the id lets
-    receivers keep them apart."""
-
-
-class SnapshotChunk(FrozenModel):
-    """One slice of a snapshot in flight.
-
-    `data_b64` carries a base64-encoded slice of the zstd-compressed JSON
-    dump of State. Concatenating the *decoded* bytes of all chunks for a
-    `transfer_id` in order of `chunk_index` yields the full compressed
-    body; `sha256_hex` is the SHA-256 of that decoded blob.
-
-    We use base64 explicitly because the topic layer JSON-encodes messages,
-    and JSON can't carry raw binary. Helpers `from_data` / `data` keep the
-    base64 detail at the boundaries.
-    """
-
-    transfer_id: SnapshotTransferId
-    requester_node_id: NodeId
-    session_id: SessionId
-    schema_version: int
-    last_event_applied_idx: int
-    chunk_index: int
-    total_chunks: int
-    sha256_hex: str
-    data_b64: str
-
-    @classmethod
-    def from_data(cls, *, data: bytes, **kwargs: object) -> "SnapshotChunk":
-        return cls(data_b64=base64.b64encode(data).decode("ascii"), **kwargs)  # pyright: ignore[reportArgumentType]
-
-    @property
-    def data(self) -> bytes:
-        return base64.b64decode(self.data_b64)
-
-
-__all__ = ["SnapshotChunk", "SnapshotTransferId"]

src/exo/shared/types/state.py

@@ -6,9 +6,7 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
+from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
     DiskUsage,
     MemoryUsage,
@@ -42,16 +40,10 @@ class State(FrozenModel):
         strict=True,
         arbitrary_types_allowed=True,
     )
-    # Bump when a State change makes older snapshots unsafe to restore.
-    schema_version: int = Field(default=1, ge=1)
-
     instances: Mapping[InstanceId, Instance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
-    # Durable request input chunks for active image requests. Workers rebuild
-    # local image caches from this state instead of reading events directly.
-    input_chunks: Mapping[CommandId, Mapping[int, InputImageChunk]] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
@@ -69,9 +61,6 @@ class State(FrozenModel):
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
 
-    instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
-    prefill_server_ports: Mapping[RunnerId, int] = {}
-
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()

src/exo/shared/types/tasks.py

@@ -101,6 +101,3 @@ Task = (
     | ImageEdits
     | Shutdown
 )
-TextTask = TextGeneration
-ImageTask = ImageGeneration | ImageEdits
-GenerationTask = TextTask | ImageTask

src/exo/shared/types/text_generation.py

@@ -13,20 +13,6 @@ from exo.shared.types.common import ModelId, TruncatingString
 
 MessageRole = Literal["user", "assistant", "system", "developer", "tool"]
 ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
-# How a model wants prior-turn reasoning content handled. Drives both the
-# server-side encoder (drop vs keep) and the integration configs we emit
-# (e.g. opencode's per-model `interleaved` flag).
-#   - "none":            model has no reasoning channel.
-#   - "post_last_user":  reasoning is only meaningful for the latest assistant
-#                        turn; older turns can drop it (drop_thinking=True).
-#   - "suffix":          reasoning is embedded in the assistant content as a
-#                        suffix/prefix; round-tripping content already covers
-#                        it (no separate `reasoning_content` round-trip).
-#   - "channel":         reasoning lives on a dedicated channel (Harmony, etc.)
-#                        and must be sent back verbatim every turn.
-#   - "tool_conditional": always round-trip when the conversation has tools;
-#                        the model relies on prior reasoning to chain tool
-#                        calls (DeepSeek V3.2 / V4).
 ReasoningDialect = Literal[
     "none", "post_last_user", "suffix", "channel", "tool_conditional"
 ]
@@ -132,8 +118,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     images: list[Base64Image] = Field(default_factory=list)
     image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
 
-    prefill_endpoint: str | None = None
-
     def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
         from exo.shared.models.model_cards import get_card
 

src/exo/shared/types/worker/runner_response.py

@@ -16,6 +16,10 @@ class BaseRunnerResponse(TaggedModel):
     pass
 
 
+class TokenizedResponse(BaseRunnerResponse):
+    prompt_tokens: int
+
+
 class GenerationResponse(BaseRunnerResponse):
     text: str
     token: int
@@ -71,10 +75,6 @@ class ModelLoadingResponse(BaseRunnerResponse):
     total: int
 
 
-class CancelledResponse(BaseRunnerResponse):
-    pass
-
-
 class PrefillProgressResponse(BaseRunnerResponse):
     processed_tokens: int
     total_tokens: int

src/exo/shared/types/worker/runners.py

@@ -47,7 +47,7 @@ class RunnerWarmingUp(BaseRunnerStatus):
 
 
 class RunnerReady(BaseRunnerStatus):
-    prefill_server_port: int | None = None
+    pass
 
 
 class RunnerRunning(BaseRunnerStatus):

src/exo/utils/event_buffer.py

@@ -47,18 +47,6 @@ class OrderedBuffer[T]:
         logger.trace(f"Releasing event {ret}")
         return ret
 
-    def fast_forward_to(self, idx: int) -> None:
-        """Skip every event before idx.
-
-        Snapshot restore uses this after applying state that already includes
-        events before idx. Any buffered or future event below idx is stale.
-        """
-        if idx <= self.next_idx_to_release:
-            return
-        self.next_idx_to_release = idx
-        for stale_idx in [i for i in self.store if i < idx]:
-            del self.store[stale_idx]
-
 
 class MultiSourceBuffer[SourceId, T]:
     """

src/exo/utils/ports.py

@@ -1,6 +0,0 @@
-import random
-
-
-def random_ephemeral_port() -> int:
-    port = random.randint(49153, 65535)
-    return port - 1 if port <= 52415 else port

src/exo/worker/disaggregated/protocol.py

@@ -1,152 +0,0 @@
-from typing import BinaryIO, Literal
-
-import msgspec
-
-DType = Literal["bfloat16", "float16", "float32"]
-
-
-class ProtocolError(Exception):
-    pass
-
-
-class Header(msgspec.Struct):
-    request_id: str = ""
-    model_id: str = ""
-    num_layers: int = 0
-    dtype: DType = "bfloat16"
-    start_pos: int = 0
-
-
-class TensorBlob(msgspec.Struct):
-    dtype: DType
-    shape: tuple[int, ...]
-    data: bytes
-
-
-class KVChunk(msgspec.Struct, tag="kv_chunk"):
-    layer_idx: int
-    num_tokens: int
-    n_heads: int
-    head_dim: int
-    dtype: DType
-    keys: bytes
-    values: bytes
-
-    @property
-    def shape(self) -> tuple[int, int, int]:
-        return (self.num_tokens, self.n_heads, self.head_dim)
-
-
-class ArraysState(msgspec.Struct, tag="arrays_state"):
-    layer_idx: int
-    arrays: list[TensorBlob] = []
-
-
-class Done(msgspec.Struct, tag="done"):
-    total_tokens: int
-
-
-class ErrorMessage(msgspec.Struct, tag="error"):
-    code: int
-    message: str
-
-
-Message = KVChunk | ArraysState | Done | ErrorMessage
-
-_msg_encoder = msgspec.msgpack.Encoder()
-_msg_decoder: msgspec.msgpack.Decoder[Message] = msgspec.msgpack.Decoder(Message)
-_header_encoder = msgspec.msgpack.Encoder()
-_header_decoder: msgspec.msgpack.Decoder[Header] = msgspec.msgpack.Decoder(Header)
-
-
-def _read_exactly(stream: BinaryIO, n: int) -> bytes:
-    buf = bytearray()
-    while len(buf) < n:
-        chunk = stream.read(n - len(buf))
-        if not chunk:
-            if len(buf) == 0:
-                return b""
-            raise ConnectionError(f"Connection closed after {len(buf)}/{n} bytes")
-        buf.extend(chunk)
-    return bytes(buf)
-
-
-def write_frame(stream: BinaryIO, payload: bytes) -> None:
-    stream.write(len(payload).to_bytes(4, "big"))
-    stream.write(payload)
-    stream.flush()
-
-
-def read_frame(stream: BinaryIO) -> bytes:
-    raw = _read_exactly(stream, 4)
-    if not raw:
-        return b""
-    length = int.from_bytes(raw, "big")
-    return _read_exactly(stream, length)
-
-
-def write_header(stream: BinaryIO, header: Header) -> None:
-    write_frame(stream, _header_encoder.encode(header))
-
-
-def read_header(stream: BinaryIO) -> Header:
-    payload = read_frame(stream)
-    if not payload:
-        raise ConnectionError("No header received")
-    try:
-        return _header_decoder.decode(payload)
-    except msgspec.DecodeError as exc:
-        raise ProtocolError(f"Bad header: {exc}") from exc
-
-
-def write_message(stream: BinaryIO, msg: Message) -> None:
-    write_frame(stream, _msg_encoder.encode(msg))
-
-
-def read_message(stream: BinaryIO) -> Message | None:
-    payload = read_frame(stream)
-    if not payload:
-        return None
-    try:
-        return _msg_decoder.decode(payload)
-    except msgspec.DecodeError as exc:
-        raise ProtocolError(f"Bad message: {exc}") from exc
-
-
-def write_kv_chunk(
-    stream: BinaryIO,
-    *,
-    layer_idx: int,
-    num_tokens: int,
-    n_heads: int,
-    head_dim: int,
-    dtype: DType,
-    keys: bytes,
-    values: bytes,
-) -> None:
-    write_message(
-        stream,
-        KVChunk(
-            layer_idx=layer_idx,
-            num_tokens=num_tokens,
-            n_heads=n_heads,
-            head_dim=head_dim,
-            dtype=dtype,
-            keys=keys,
-            values=values,
-        ),
-    )
-
-
-def write_arrays_state(
-    stream: BinaryIO, layer_idx: int, arrays: list[TensorBlob]
-) -> None:
-    write_message(stream, ArraysState(layer_idx=layer_idx, arrays=arrays))
-
-
-def write_done(stream: BinaryIO, total_tokens: int) -> None:
-    write_message(stream, Done(total_tokens=total_tokens))
-
-
-def write_error(stream: BinaryIO, code: int, message: str) -> None:
-    write_message(stream, ErrorMessage(code=code, message=message))

src/exo/worker/disaggregated/server.py

@@ -1,105 +0,0 @@
-import socket
-import socketserver
-import threading
-from collections.abc import Callable
-from typing import BinaryIO, cast
-
-import msgspec
-from loguru import logger
-
-from exo.worker.disaggregated.protocol import (
-    Header,
-    read_frame,
-    write_error,
-    write_frame,
-    write_header,
-)
-
-
-class PrefillRequest(msgspec.Struct):
-    request_id: str = ""
-    model_id: str = ""
-    token_ids: list[int] = msgspec.field(default_factory=list)
-    start_pos: int = 0
-
-
-_request_encoder = msgspec.msgpack.Encoder()
-_request_decoder: msgspec.msgpack.Decoder[PrefillRequest] = msgspec.msgpack.Decoder(
-    PrefillRequest
-)
-
-
-def write_request(stream: BinaryIO, job: PrefillRequest) -> None:
-    write_frame(stream, _request_encoder.encode(job))
-
-
-def read_request(stream: BinaryIO) -> PrefillRequest:
-    payload = read_frame(stream)
-    if not payload:
-        raise ConnectionError("No request received")
-    return _request_decoder.decode(payload)
-
-
-ResolveHandler = Callable[[PrefillRequest, BinaryIO], bool]
-
-
-def _send_error(wfile: BinaryIO, code: int, message: str) -> None:
-    try:
-        write_header(wfile, Header(num_layers=0, dtype="float32"))
-        write_error(wfile, code=code, message=message)
-    except Exception:
-        pass
-
-
-class _PrefillHandler(socketserver.StreamRequestHandler):
-    def setup(self) -> None:
-        super().setup()
-        sock = cast(socket.socket, self.request)
-        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024)
-
-    def handle(self) -> None:
-        server = cast(PrefillServer, self.server)
-        wfile: BinaryIO = cast(BinaryIO, cast(object, self.wfile))
-        rfile: BinaryIO = cast(BinaryIO, cast(object, self.rfile))
-        try:
-            job = read_request(rfile)
-        except ConnectionError:
-            return
-        except (msgspec.DecodeError, ValueError) as exc:
-            _send_error(wfile, 400, f"Bad request: {exc}")
-            return
-        try:
-            picked_up = server.resolve(job, wfile)
-        except Exception as e:
-            logger.opt(exception=e).warning(
-                f"Prefill resolve error for request_id={job.request_id}"
-            )
-            _send_error(wfile, 500, str(e))
-            return
-        if not picked_up:
-            _send_error(
-                wfile, 503, f"Prefill not picked up for request_id={job.request_id!r}"
-            )
-
-
-class PrefillServer(socketserver.ThreadingTCPServer):
-    allow_reuse_address = True
-    daemon_threads = True
-    resolve: ResolveHandler
-
-    def __init__(self, resolve: ResolveHandler, host: str, port: int) -> None:
-        super().__init__((host, port), _PrefillHandler)
-        self.resolve = resolve
-        self._thread = threading.Thread(
-            target=self.serve_forever, name="prefill-server"
-        )
-        self._thread.start()
-        logger.info(f"Prefill server listening on {host}:{port}")
-
-    def stop(self) -> None:
-        self.shutdown()
-        self.server_close()
-        if self._thread is not None:
-            self._thread.join(timeout=5)
-            self._thread = None

src/exo/worker/engines/base.py

@@ -1,60 +0,0 @@
-from abc import ABC, abstractmethod
-from collections.abc import Generator, Iterable
-from typing import BinaryIO
-
-from exo.shared.types.chunks import Chunk
-from exo.shared.types.tasks import CANCEL_ALL_TASKS, GenerationTask, TaskId
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runner_response import (
-    CancelledResponse,
-    FinishedResponse,
-    ModelLoadingResponse,
-)
-from exo.worker.disaggregated.server import PrefillRequest
-
-
-class Engine(ABC):
-    _cancelled_tasks: set[TaskId]
-
-    def should_cancel(self, task_id: TaskId) -> bool:
-        return (
-            task_id in self._cancelled_tasks
-            or CANCEL_ALL_TASKS in self._cancelled_tasks
-        )
-
-    @abstractmethod
-    def warmup(self) -> None: ...
-
-    @abstractmethod
-    def submit(
-        self,
-        task: GenerationTask,
-    ) -> None: ...
-
-    @abstractmethod
-    def step(
-        self,
-    ) -> Iterable[tuple[TaskId, Chunk | CancelledResponse | FinishedResponse]]: ...
-
-    @abstractmethod
-    def close(self) -> None: ...
-
-    @abstractmethod
-    def serve_prefill(self, request: PrefillRequest, wfile: BinaryIO) -> None: ...
-
-
-class Builder(ABC):
-    @abstractmethod
-    def connect(self, bound_instance: BoundInstance) -> None: ...
-
-    @abstractmethod
-    def load(
-        self,
-        bound_instance: BoundInstance,
-    ) -> Generator[ModelLoadingResponse]: ...
-
-    @abstractmethod
-    def build(self) -> Engine: ...
-
-    @abstractmethod
-    def close(self) -> None: ...

src/exo/worker/engines/image/__init__.py

@@ -1,16 +1,12 @@
-from exo.worker.engines.image.builder import (
-    ImageEngine,
-    MfluxBuilder,
-)
 from exo.worker.engines.image.distributed_model import (
     DistributedImageModel,
+    initialize_image_model,
 )
 from exo.worker.engines.image.generate import generate_image, warmup_image_generator
 
 __all__ = [
-    "MfluxBuilder",
-    "ImageEngine",
     "DistributedImageModel",
     "generate_image",
+    "initialize_image_model",
     "warmup_image_generator",
 ]

src/exo/worker/engines/image/builder.py

@@ -1,219 +0,0 @@
-import contextlib
-from collections import deque
-from collections.abc import Generator, Iterable
-from dataclasses import dataclass, field
-from typing import BinaryIO
-
-import mlx.core as mx
-from loguru import logger
-
-from exo.api.types import ImageEditsTaskParams, ImageGenerationTaskParams
-from exo.shared.constants import EXO_TRACING_ENABLED
-from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
-from exo.shared.types.chunks import Chunk, ErrorChunk
-from exo.shared.types.events import (
-    Event,
-    TraceEventData,
-    TracesCollected,
-)
-from exo.shared.types.tasks import (
-    GenerationTask,
-    ImageEdits,
-    ImageGeneration,
-    ImageTask,
-    TaskId,
-)
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runner_response import (
-    CancelledResponse,
-    FinishedResponse,
-    ModelLoadingResponse,
-)
-from exo.shared.types.worker.shards import (
-    CfgShardMetadata,
-    PipelineShardMetadata,
-    ShardMetadata,
-)
-from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.disaggregated.server import PrefillRequest
-from exo.worker.engines.base import Builder, Engine
-from exo.worker.engines.image.distributed_model import (
-    DistributedImageModel,
-)
-from exo.worker.engines.image.generate import (
-    generate_image,
-    warmup_image_generator,
-)
-from exo.worker.engines.mlx.utils_mlx import (
-    initialize_mlx,
-)
-
-
-def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
-    """Check if this node is the primary output node for image generation.
-
-    For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
-    For non-CFG models: the last pipeline stage.
-    """
-    if isinstance(shard_metadata, CfgShardMetadata):
-        is_pipeline_last = (
-            shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
-        )
-        return is_pipeline_last and shard_metadata.cfg_rank == 0
-    elif isinstance(shard_metadata, PipelineShardMetadata):
-        return shard_metadata.device_rank == shard_metadata.world_size - 1
-    return False
-
-
-def _send_traces_if_enabled(
-    event_sender: MpSender[Event],
-    task_id: TaskId,
-    rank: int,
-) -> None:
-    if not EXO_TRACING_ENABLED:
-        return
-
-    traces = get_trace_buffer()
-    if traces:
-        trace_data = [
-            TraceEventData(
-                name=t.name,
-                start_us=t.start_us,
-                duration_us=t.duration_us,
-                rank=t.rank,
-                category=t.category,
-            )
-            for t in traces
-        ]
-        event_sender.send(
-            TracesCollected(
-                task_id=task_id,
-                rank=rank,
-                traces=trace_data,
-            )
-        )
-    clear_trace_buffer()
-
-
-@dataclass
-class MfluxBuilder(Builder):
-    event_sender: MpSender[Event]
-    cancel_receiver: MpReceiver[TaskId]
-    shard_metadata: ShardMetadata | None = None
-    image_model: DistributedImageModel | None = None
-    group: mx.distributed.Group | None = None
-
-    def connect(self, bound_instance: BoundInstance) -> None:
-        self.group = initialize_mlx(bound_instance)
-
-    def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
-        self.shard_metadata = bound_instance.bound_shard
-        self.image_model = DistributedImageModel.from_shard_metadata(
-            bound_instance.bound_shard, self.group
-        )
-        return
-        # very important!
-        yield
-
-    def close(self) -> None:
-        with contextlib.suppress(NameError, AttributeError):
-            del self.image_model, self.group
-
-    def build(
-        self,
-    ) -> Engine:
-        assert self.image_model
-        assert self.shard_metadata
-
-        return ImageEngine(
-            self.image_model,
-            self.shard_metadata,
-            self.event_sender,
-            self.cancel_receiver,
-        )
-
-
-@dataclass
-class ImageEngine(Engine):
-    image_model: DistributedImageModel
-    shard_metadata: ShardMetadata
-    event_sender: MpSender[Event]
-    cancel_receiver: MpReceiver[TaskId]
-    current_gen: (
-        Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]] | None
-    ) = field(init=False, default=None)
-    queue: deque[ImageTask] = field(init=False, default_factory=deque)
-
-    def warmup(self) -> None:
-        image = warmup_image_generator(model=self.image_model)
-        if image is not None:
-            logger.info(f"warmed up by generating {image.size} image")
-        else:
-            logger.info("warmup completed (non-primary node)")
-
-    def submit(
-        self,
-        task: GenerationTask,
-    ) -> None:
-        assert isinstance(task, (ImageGeneration, ImageEdits))
-        self.queue.append(task)
-
-    def step(
-        self,
-    ) -> Iterable[tuple[TaskId, Chunk | CancelledResponse | FinishedResponse]]:
-        resp = None
-        if self.current_gen is not None:
-            resp = next(self.current_gen, None)
-        if resp is None and len(self.queue) > 0:
-            task = self.queue.popleft()
-            self.current_gen = self._run_image_task(task.task_id, task.task_params)
-            resp = next(self.current_gen, None)
-        return (resp,) if resp is not None else ()
-
-    def close(self) -> None:
-        with contextlib.suppress(NameError, AttributeError):
-            del self.image_model
-
-    def serve_prefill(self, request: PrefillRequest, wfile: BinaryIO) -> None:
-        raise NotImplementedError() from None
-
-    def _run_image_task(
-        self,
-        task_id: TaskId,
-        task_params: ImageGenerationTaskParams | ImageEditsTaskParams,
-    ) -> Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]]:
-        assert self.image_model
-        logger.info(f"received image task: {str(task_params)[:500]}")
-
-        def cancel_checker() -> bool:
-            for cancel_id in self.cancel_receiver.collect():
-                self._cancelled_tasks.add(cancel_id)
-            return self.should_cancel(task_id)
-
-        try:
-            # todo: yield CancelledResponse properly
-            for response in generate_image(
-                model=self.image_model,
-                task=task_params,
-                cancel_checker=cancel_checker,
-            ):
-                if _is_primary_output_node(self.shard_metadata):
-                    yield (task_id, response)
-        except Exception as e:
-            if _is_primary_output_node(self.shard_metadata):
-                yield (
-                    task_id,
-                    ErrorChunk(
-                        model=self.shard_metadata.model_card.model_id,
-                        finish_reason="error",
-                        error_message=str(e),
-                    ),
-                )
-            raise
-        finally:
-            _send_traces_if_enabled(
-                self.event_sender, task_id, self.shard_metadata.device_rank
-            )
-            yield (task_id, FinishedResponse())
-
-        return

src/exo/worker/engines/image/distributed_model.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable, Generator
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, Optional
 
 import mlx.core as mx
 from mflux.models.common.config.config import Config
@@ -9,11 +9,8 @@ from PIL import Image
 from exo.api.types import AdvancedImageParams
 from exo.download.download_utils import build_model_path
 from exo.shared.types.common import ModelId
-from exo.shared.types.worker.shards import (
-    CfgShardMetadata,
-    PipelineShardMetadata,
-    ShardMetadata,
-)
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
     create_adapter_for_model,
@@ -21,7 +18,7 @@ from exo.worker.engines.image.models import (
 )
 from exo.worker.engines.image.models.base import ModelAdapter
 from exo.worker.engines.image.pipeline import DiffusionRunner
-from exo.worker.engines.mlx.utils_mlx import mx_barrier
+from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init, mx_barrier
 from exo.worker.runner.bootstrap import logger
 
 
@@ -36,7 +33,7 @@ class DistributedImageModel:
         model_id: ModelId,
         local_path: Path,
         shard_metadata: PipelineShardMetadata | CfgShardMetadata,
-        group: mx.distributed.Group | None,
+        group: Optional[mx.distributed.Group] = None,
         quantize: int | None = None,
     ):
         config = get_config_for_model(model_id)
@@ -79,21 +76,32 @@ class DistributedImageModel:
         self._runner = runner
 
     @classmethod
-    def from_shard_metadata(
-        cls, shard: ShardMetadata, group: mx.distributed.Group | None
+    def from_bound_instance(
+        cls, bound_instance: BoundInstance
     ) -> "DistributedImageModel":
-        model_id = shard.model_card.model_id
+        model_id = bound_instance.bound_shard.model_card.model_id
         model_path = build_model_path(model_id)
 
-        if not isinstance(shard, (PipelineShardMetadata, CfgShardMetadata)):
+        shard_metadata = bound_instance.bound_shard
+        if not isinstance(shard_metadata, (PipelineShardMetadata, CfgShardMetadata)):
             raise ValueError(
                 "Expected PipelineShardMetadata or CfgShardMetadata for image generation"
             )
 
+        is_distributed = (
+            len(bound_instance.instance.shard_assignments.node_to_runner) > 1
+        )
+
+        if is_distributed:
+            logger.info("Starting distributed init for image model")
+            group = mlx_distributed_init(bound_instance)
+        else:
+            group = None
+
         return cls(
             model_id=model_id,
             local_path=model_path,
-            shard_metadata=shard,
+            shard_metadata=shard_metadata,
             group=group,
         )
 
@@ -168,3 +176,7 @@ class DistributedImageModel:
             else:
                 logger.info("generated image")
                 yield result
+
+
+def initialize_image_model(bound_instance: BoundInstance) -> DistributedImageModel:
+    return DistributedImageModel.from_bound_instance(bound_instance)

src/exo/worker/engines/mlx/auto_parallel.py

@@ -918,7 +918,7 @@ class DeepseekV4ShardingStrategy(TensorParallelShardingStrategy):
             # Head-parallel attention with interleaved-per-group sharding.
             _shard_v4_attention_heads(layer.attn, self.N, self.group.rank())
             self.sharded_to_all_linear_in_place(layer.attn.wo_a)
-            layer.attn.wo_b = _AllSumLinear(layer.attn.wo_b, self.group)  # type: ignore
+            layer.attn.wo_b = _AllSumLinear(layer.attn.wo_b, self.group)  # type: ignore[assignment]
 
             ffn = layer.ffn
             if getattr(ffn, "shared_experts", None) is not None:
@@ -930,7 +930,7 @@ class DeepseekV4ShardingStrategy(TensorParallelShardingStrategy):
             self.all_to_sharded_linear_in_place(ffn.switch_mlp.up_proj)
             wrapped = ShardedMoEV4(ffn)
             wrapped.sharding_group = self.group
-            layer.ffn = wrapped  # type: ignore
+            layer.ffn = wrapped  # type: ignore[assignment]
 
             mx.eval(layer)
             mx.clear_cache()

src/exo/worker/engines/mlx/builder.py

@@ -1,113 +0,0 @@
-import contextlib
-import os
-from collections.abc import Generator
-from dataclasses import dataclass
-
-import mlx.core as mx
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-
-from exo.shared.types.common import ModelId
-from exo.shared.types.events import Event
-from exo.shared.types.tasks import TaskId
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runner_response import ModelLoadingResponse
-from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.base import Builder, Engine
-from exo.worker.runner.bootstrap import logger
-from exo.worker.runner.llm_inference.batch_generator import (
-    BatchGenerator,
-    SequentialGenerator,
-)
-from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
-
-from .cache import KVPrefixCache
-from .types import Model
-from .utils_mlx import (
-    initialize_mlx,
-    load_mlx_items,
-)
-from .vision import VisionProcessor
-
-
-@dataclass
-class MlxBuilder(Builder):
-    model_id: ModelId
-    event_sender: MpSender[Event]
-    cancel_receiver: MpReceiver[TaskId]
-    inference_model: Model | None = None
-    tokenizer: TokenizerWrapper | None = None
-    group: mx.distributed.Group | None = None
-    vision_processor: VisionProcessor | None = None
-
-    def connect(self, bound_instance: BoundInstance) -> None:
-        self.group = initialize_mlx(bound_instance)
-
-    def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
-        (
-            self.inference_model,
-            self.tokenizer,
-            self.vision_processor,
-        ) = yield from load_mlx_items(bound_instance, self.group)
-
-    def close(self) -> None:
-        with contextlib.suppress(NameError, AttributeError):
-            del self.inference_model
-        with contextlib.suppress(NameError, AttributeError):
-            del self.tokenizer
-        with contextlib.suppress(NameError, AttributeError):
-            del self.group
-
-    def build(
-        self,
-    ) -> Engine:
-        assert self.inference_model
-        assert self.tokenizer
-
-        vision_processor = self.vision_processor
-
-        tool_parser = None
-        logger.info(
-            f"model has_tool_calling={self.tokenizer.has_tool_calling} using tokens {self.tokenizer.tool_call_start}, {self.tokenizer.tool_call_end}"
-        )
-        if (
-            self.tokenizer.tool_call_start
-            and self.tokenizer.tool_call_end
-            and self.tokenizer.tool_parser  # type: ignore
-        ):
-            tool_parser = make_mlx_parser(
-                self.tokenizer.tool_call_start,
-                self.tokenizer.tool_call_end,
-                self.tokenizer.tool_parser,  # type: ignore
-            )
-
-        kv_prefix_cache = KVPrefixCache(self.group)
-
-        device_rank = 0 if self.group is None else self.group.rank()
-        if os.environ.get("EXO_NO_BATCH"):
-            logger.info("using SequentialGenerator (batching disabled)")
-            return SequentialGenerator(
-                model=self.inference_model,
-                tokenizer=self.tokenizer,
-                group=self.group,
-                tool_parser=tool_parser,
-                kv_prefix_cache=kv_prefix_cache,
-                model_id=self.model_id,
-                device_rank=device_rank,
-                cancel_receiver=self.cancel_receiver,
-                event_sender=self.event_sender,
-                vision_processor=vision_processor,
-            )
-        else:
-            logger.info("using BatchGenerator")
-            return BatchGenerator(
-                model=self.inference_model,
-                tokenizer=self.tokenizer,
-                group=self.group,
-                tool_parser=tool_parser,
-                kv_prefix_cache=kv_prefix_cache,
-                model_id=self.model_id,
-                device_rank=device_rank,
-                cancel_receiver=self.cancel_receiver,
-                event_sender=self.event_sender,
-                vision_processor=vision_processor,
-            )

src/exo/worker/engines/mlx/cache.py

@@ -13,17 +13,11 @@ from mlx_lm.models.cache import (
     QuantizedKVCache,
     RotatingKVCache,
 )
-from mlx_lm.models.deepseek_v4 import (
-    DeepseekV4Cache,
-)
-from mlx_lm.models.deepseek_v4 import (
-    _CompressorBranch as CompressorBranch,  # type: ignore
-)
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
-from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.runner.bootstrap import logger
 
 if TYPE_CHECKING:
@@ -53,9 +47,7 @@ class CacheSnapshot:
 
     def __init__(
         self,
-        states: list[
-            RotatingKVCache | ArraysCache | CacheList | DeepseekV4Cache | None
-        ],
+        states: list[RotatingKVCache | ArraysCache | CacheList | None],
         token_count: int,
     ):
         self.states = states
@@ -120,71 +112,21 @@ def _copy_cache_list(cl: CacheList) -> CacheList:
     return CacheList(*copied)
 
 
-def _detached_copy_or_none(a: mx.array | None) -> mx.array | None:
-    if a is None:
+def restore_snapshot_entry(
+    entry: ArraysCache | RotatingKVCache | CacheList | None,
+) -> ArraysCache | RotatingKVCache | CacheList | None:
+    if entry is None:
         return None
-    out = _detached_copy(a)
-    mx.eval(out)
-    return out
-
-
-def _copy_compressor_branch(b: CompressorBranch) -> CompressorBranch:
-    out = CompressorBranch.__new__(CompressorBranch)
-    out.buffer_kv = _detached_copy_or_none(b.buffer_kv)
-    out.buffer_gate = _detached_copy_or_none(b.buffer_gate)
-    out.prev_kv = _detached_copy_or_none(b.prev_kv)
-    out.prev_gate = _detached_copy_or_none(b.prev_gate)
-    out.pool = _detached_copy_or_none(b.pool)
-    out.buffer_lengths = deepcopy(b.buffer_lengths)
-    out.pool_lengths = deepcopy(b.pool_lengths)
-    out.buffer_count = deepcopy(b.buffer_count)
-    out._new_pool_lengths = deepcopy(b._new_pool_lengths)
-    return out
-
-
-def _copy_v4_cache(c: DeepseekV4Cache) -> DeepseekV4Cache:
-    snap = DeepseekV4Cache.__new__(DeepseekV4Cache)
-
-    local: RotatingKVCache = c.local
-    local_snap = copy_rotating_kv_cache(local)
-    if local_snap is None:
-        local_snap = RotatingKVCache.__new__(RotatingKVCache)
-        local_snap.keys = None
-        local_snap.values = None
-        local_snap.offset = local.offset
-        local_snap._idx = 0
-        local_snap.keep = local.keep
-        local_snap.max_size = local.max_size
-    snap.local = local_snap
-
-    snap._branches = {
-        key: _copy_compressor_branch(branch) for key, branch in c._branches.items()
-    }
-    snap._pending_lengths = deepcopy(c._pending_lengths)
-    return snap
-
-
-def copy_snapshot_entry(
-    entry: ArraysCache | RotatingKVCache | CacheList | DeepseekV4Cache | None,
-) -> ArraysCache | RotatingKVCache | CacheList | DeepseekV4Cache | None:
-    match entry:
-        case None:
-            return None
-        case RotatingKVCache():
-            snap = copy_rotating_kv_cache(entry)
-            return snap if snap is not None else deepcopy(entry)
-        case ArraysCache():
-            return _copy_arrays_cache(entry)
-        case CacheList():
-            return _copy_cache_list(entry)
-        case DeepseekV4Cache():
-            return _copy_v4_cache(entry)
+    if isinstance(entry, RotatingKVCache):
+        snap = copy_rotating_kv_cache(entry)
+        return snap if snap is not None else deepcopy(entry)
+    if isinstance(entry, ArraysCache):
+        return _copy_arrays_cache(entry)
+    return _copy_cache_list(entry)
 
 
 def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
-    states: list[
-        RotatingKVCache | ArraysCache | CacheList | DeepseekV4Cache | None
-    ] = []
+    states: list[ArraysCache | RotatingKVCache | CacheList | None] = []
     for c in cache:
         if isinstance(c, ArraysCache):
             states.append(_copy_arrays_cache(c))
@@ -192,8 +134,6 @@ def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
             states.append(copy_rotating_kv_cache(c))
         elif isinstance(c, CacheList) and not bool(c.is_trimmable()):  # type: ignore[reportUnknownMemberType]
             states.append(_copy_cache_list(c))
-        elif isinstance(c, DeepseekV4Cache):
-            states.append(_copy_v4_cache(c))
         else:
             states.append(None)
     token_count = cache_length(cache)
@@ -213,20 +153,14 @@ def _find_nearest_snapshot(
     return best
 
 
-def is_non_trimmable_cache_entry(c: object) -> bool:
-    """A cache entry is non-trimmable if `trim(n)` can't roll back its full
-    state — meaning the prefill +2 rollback must snapshot+restore it instead.
-    """
-    if isinstance(c, (ArraysCache, RotatingKVCache)):
-        return True
-    if isinstance(c, CacheList):
-        return not bool(c.is_trimmable())  # type: ignore[reportUnknownMemberType]
-    return isinstance(c, DeepseekV4Cache)
-
-
 def has_non_kv_caches(cache: KVCacheType) -> bool:
     """Check if a cache contains any ArraysCache (SSM) entries."""
-    return any(is_non_trimmable_cache_entry(c) for c in cache)
+    for c in cache:
+        if isinstance(c, CacheList):
+            return any(isinstance(_c, (ArraysCache, RotatingKVCache)) for _c in c)  # type: ignore[reportUnknownVariableType]
+        elif isinstance(c, (ArraysCache, RotatingKVCache)):
+            return True
+    return False
 
 
 class KVPrefixCache:
@@ -382,10 +316,6 @@ class KVPrefixCache:
             trim_cache(prompt_cache, tokens_to_trim, restore_snap)
             # Reset cache offset to match trimmed length
             for c in prompt_cache:
-                if isinstance(c, (ArraysCache, RotatingKVCache)):
-                    continue
-                if isinstance(c, DeepseekV4Cache):
-                    continue
                 if hasattr(c, "offset"):
                     c.offset = restore_pos
 
@@ -481,22 +411,16 @@ def trim_cache(
         )
         if non_trimmable:
             if snapshot is not None and snapshot.states[i] is not None:
-                restored = copy_snapshot_entry(snapshot.states[i])
+                restored = restore_snapshot_entry(snapshot.states[i])
                 if restored is not None:
                     cache[i] = restored  # type: ignore
             elif isinstance(c, (ArraysCache, RotatingKVCache)):
                 c.state = [None] * len(c.state)
-                if isinstance(c, RotatingKVCache):
-                    c.offset = 0
-                    c._idx = 0
             else:
                 # CacheList without a snapshot — zero each inner cache's state
                 for inner in c:  # type: ignore[reportUnknownVariableType]
                     if isinstance(inner, (ArraysCache, RotatingKVCache)):
                         inner.state = [None] * len(inner.state)
-                        if isinstance(inner, RotatingKVCache):
-                            inner.offset = 0
-                            inner._idx = 0
         else:
             c.trim(num_tokens)
 
@@ -514,12 +438,7 @@ def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
 
 
 def _entry_length(
-    c: KVCache
-    | RotatingKVCache
-    | QuantizedKVCache
-    | ArraysCache
-    | CacheList
-    | DeepseekV4Cache,
+    c: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList,
 ) -> int:
     # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache).
     if hasattr(c, "offset"):

src/exo/worker/engines/mlx/deepseek_v4_encoding.py

@@ -602,6 +602,7 @@ def encode_messages(
 
     prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
 
+    # Resolve drop_thinking: if any message has tools defined, don't drop thinking
     effective_drop_thinking = drop_thinking
     if any(m.get("tools") for m in full_messages):
         effective_drop_thinking = False

src/exo/worker/engines/mlx/disaggregated/adapter.py

@@ -1,233 +0,0 @@
-from typing import BinaryIO
-
-import mlx.core as mx
-import numpy as np
-from mlx_lm.models.cache import (
-    ArraysCache,
-    CacheList,
-    KVCache,
-    QuantizedKVCache,
-    RotatingKVCache,
-)
-from mlx_lm.models.deepseek_v4 import DeepseekV4Cache
-
-from exo.worker.disaggregated.protocol import (
-    DType,
-    Header,
-    KVChunk,
-    TensorBlob,
-    write_arrays_state,
-    write_done,
-    write_header,
-    write_kv_chunk,
-)
-from exo.worker.engines.mlx.types import KVCacheType
-from exo.worker.runner.bootstrap import logger
-
-_STR_TO_MX: dict[DType, mx.Dtype] = {
-    "bfloat16": mx.bfloat16,
-    "float16": mx.float16,
-    "float32": mx.float32,
-}
-
-_MX_TO_STR: dict[mx.Dtype, DType] = {v: k for k, v in _STR_TO_MX.items()}
-
-
-def mx_dtype_to_str(dtype: mx.Dtype) -> DType:
-    if dtype not in _MX_TO_STR:
-        raise ValueError(f"Unsupported mlx dtype on wire: {dtype}")
-    return _MX_TO_STR[dtype]
-
-
-def wire_dtype_from_cache(caches: KVCacheType) -> DType:
-    for c in caches:
-        keys: mx.array | None = getattr(c, "keys", None)
-        if keys is None:
-            continue
-        if keys.dtype in _MX_TO_STR:
-            return _MX_TO_STR[keys.dtype]
-        break
-    return "bfloat16"
-
-
-def str_to_mx_dtype(dtype: DType) -> mx.Dtype:
-    if dtype not in _STR_TO_MX:
-        raise ValueError(f"Unsupported wire dtype: {dtype!r}")
-    return _STR_TO_MX[dtype]
-
-
-def array_to_bytes(t: mx.array) -> bytes:
-    # bf16 has no native numpy dtype; bitcast through uint16.
-    if t.dtype == mx.bfloat16:
-        return np.asarray(t.view(mx.uint16)).tobytes()
-    if t.dtype in (mx.float16, mx.float32):
-        return np.asarray(t).tobytes()
-    raise ValueError(f"Unsupported mlx dtype for wire: {t.dtype}")
-
-
-def bytes_to_array(data: bytes, shape: tuple[int, ...], dtype: DType) -> mx.array:
-    match dtype:
-        case "bfloat16":
-            arr = np.frombuffer(data, dtype=np.uint16).reshape(shape).copy()
-            return mx.array(arr).view(mx.bfloat16)
-        case "float16":
-            arr = np.frombuffer(data, dtype=np.float16).reshape(shape).copy()
-            return mx.array(arr)
-        case "float32":
-            arr = np.frombuffer(data, dtype=np.float32).reshape(shape).copy()
-            return mx.array(arr)
-
-
-def bhsd_to_nhd(t: mx.array) -> mx.array:
-    if t.ndim != 4 or int(t.shape[0]) != 1:
-        raise ValueError(f"Expected BHSD with B=1, got shape={tuple(t.shape)}")
-    return mx.transpose(t[0], (1, 0, 2))
-
-
-def nhd_to_bhsd(t: mx.array) -> mx.array:
-    if t.ndim != 3:
-        raise ValueError(f"Expected NHD (3D), got shape={tuple(t.shape)}")
-    return mx.expand_dims(mx.transpose(t, (1, 0, 2)), 0)
-
-
-def send_mlx_kv_cache(
-    stream: BinaryIO,
-    caches: KVCacheType,
-    *,
-    dtype: DType,
-    start_pos: int = 0,
-    max_tokens: int | None = None,
-) -> int:
-    tokens_sent = 0
-    for layer_idx, c in enumerate(caches):
-        match c:
-            case QuantizedKVCache() | CacheList() | DeepseekV4Cache():
-                raise NotImplementedError
-            case KVCache() | RotatingKVCache():
-                keys = c.keys
-                values = c.values
-                if keys is None or values is None:
-                    continue
-                offset = int(c.offset)
-                if max_tokens is not None:
-                    offset = min(offset, max_tokens)
-                if offset <= start_pos:
-                    continue
-                with mx.stream(mx.Device(mx.cpu)):
-                    k = mx.array(keys[:, :, start_pos:offset, :])
-                    v = mx.array(values[:, :, start_pos:offset, :])
-                    k_nhd = bhsd_to_nhd(k)
-                    v_nhd = bhsd_to_nhd(v)
-                    mx.eval(k_nhd, v_nhd)
-                num_tokens = int(k_nhd.shape[0])
-                n_heads = int(k_nhd.shape[1])
-                head_dim = int(k_nhd.shape[2])
-                write_kv_chunk(
-                    stream,
-                    layer_idx=layer_idx,
-                    num_tokens=num_tokens,
-                    n_heads=n_heads,
-                    head_dim=head_dim,
-                    dtype=dtype,
-                    keys=array_to_bytes(k_nhd),
-                    values=array_to_bytes(v_nhd),
-                )
-                if tokens_sent != 0 and num_tokens != tokens_sent:
-                    logger.critical(
-                        f"Unexpected number of tokens sent {num_tokens} != {tokens_sent}"
-                    )
-                tokens_sent = num_tokens
-            case ArraysCache():
-                blobs: list[TensorBlob] = []
-                for a in c.state:
-                    if a is None:
-                        continue
-                    with mx.stream(mx.Device(mx.cpu)):
-                        a_cpu = mx.array(a)
-                        mx.eval(a_cpu)
-                    blobs.append(
-                        TensorBlob(
-                            dtype=mx_dtype_to_str(a_cpu.dtype),
-                            shape=tuple(int(d) for d in a_cpu.shape),
-                            data=array_to_bytes(a_cpu),
-                        )
-                    )
-                if blobs:
-                    write_arrays_state(stream, layer_idx, blobs)
-    return tokens_sent
-
-
-def chunk_to_mlx_nhd(chunk: KVChunk) -> tuple[mx.array, mx.array]:
-    shape = chunk.shape
-    return (
-        bytes_to_array(chunk.keys, shape, chunk.dtype),
-        bytes_to_array(chunk.values, shape, chunk.dtype),
-    )
-
-
-def blob_to_mlx(blob: TensorBlob) -> mx.array:
-    return bytes_to_array(blob.data, blob.shape, blob.dtype)
-
-
-def inject_kv_chunk(
-    cache: KVCache,
-    keys_nhd: mx.array,
-    values_nhd: mx.array,
-    offset: int,
-    *,
-    start_pos: int = 0,
-    existing_k: mx.array | None = None,
-    existing_v: mx.array | None = None,
-) -> None:
-    k_bhsd = nhd_to_bhsd(keys_nhd)
-    v_bhsd = nhd_to_bhsd(values_nhd)
-    if start_pos > 0 and existing_k is not None and existing_v is not None:
-        cache.keys = mx.concatenate([existing_k[:, :, :start_pos, :], k_bhsd], axis=2)
-        cache.values = mx.concatenate([existing_v[:, :, :start_pos, :], v_bhsd], axis=2)
-    else:
-        cache.keys = k_bhsd
-        cache.values = v_bhsd
-    cache.offset = offset
-
-
-def inject_rotating_kv_chunk(
-    cache: RotatingKVCache,
-    keys_nhd: mx.array,
-    values_nhd: mx.array,
-    offset: int,
-) -> None:
-    k_bhsd = nhd_to_bhsd(keys_nhd)
-    v_bhsd = nhd_to_bhsd(values_nhd)
-    cache.keys = k_bhsd
-    cache.values = v_bhsd
-    cache.offset = offset
-    cache._idx = int(k_bhsd.shape[2])
-
-
-def inject_arrays_cache(cache: ArraysCache, blobs: list[TensorBlob]) -> None:
-    cache.state = [blob_to_mlx(b) for b in blobs]
-
-
-def write_cache_to_wire(
-    wfile: BinaryIO,
-    cache: KVCacheType,
-    *,
-    request_id: str = "",
-    model_id: str = "",
-    start_pos: int = 0,
-) -> int:
-    dtype = wire_dtype_from_cache(cache)
-    write_header(
-        wfile,
-        Header(
-            request_id=request_id,
-            model_id=model_id,
-            num_layers=len(cache),
-            dtype=dtype,
-            start_pos=start_pos,
-        ),
-    )
-    tokens_sent = send_mlx_kv_cache(wfile, cache, dtype=dtype, start_pos=start_pos)
-    write_done(wfile, tokens_sent)
-    wfile.flush()
-    return tokens_sent

src/exo/worker/engines/mlx/disaggregated/client.py

@@ -1,147 +0,0 @@
-import socket
-from collections import defaultdict
-from collections.abc import Callable
-from dataclasses import dataclass, field
-from typing import BinaryIO, cast
-
-import mlx.core as mx
-from loguru import logger
-from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
-
-from exo.worker.disaggregated.protocol import (
-    ArraysState,
-    Done,
-    Header,
-    KVChunk,
-    TensorBlob,
-    read_header,
-    read_message,
-)
-from exo.worker.disaggregated.server import PrefillRequest, write_request
-from exo.worker.engines.mlx.disaggregated.adapter import (
-    chunk_to_mlx_nhd,
-    inject_arrays_cache,
-    inject_kv_chunk,
-    inject_rotating_kv_chunk,
-)
-
-_SOCKET_TIMEOUT_SECS = 60
-_RECV_BUFFER_BYTES = 4 * 1024 * 1024
-
-
-@dataclass
-class PrefillResult:
-    header: Header
-    kv_chunks: dict[int, list[KVChunk]] = field(
-        default_factory=dict[int, list[KVChunk]]
-    )
-    arrays: dict[int, list[TensorBlob]] = field(
-        default_factory=dict[int, list[TensorBlob]]
-    )
-    total_tokens: int = 0
-
-
-def _parse_endpoint(endpoint: str) -> tuple[str, int]:
-    if ":" in endpoint:
-        host, port_str = endpoint.rsplit(":", 1)
-        return host, int(port_str)
-    raise ValueError(f"Invalid endpoint {endpoint}")
-
-
-def remote_prefill_fetch(
-    endpoint: str,
-    request: PrefillRequest,
-    on_header: Callable[[Header], None] | None = None,
-    on_kv_chunk: Callable[[KVChunk, int], None] | None = None,
-    timeout_secs: float = _SOCKET_TIMEOUT_SECS,
-) -> PrefillResult:
-    host, port = _parse_endpoint(endpoint)
-    logger.info(
-        f"Connecting to prefill server at {host}:{port} "
-        f"({len(request.token_ids)} tokens, start_pos={request.start_pos})"
-    )
-
-    sock = socket.create_connection((host, port), timeout=timeout_secs)
-    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, _RECV_BUFFER_BYTES)
-    try:
-        wfile = sock.makefile("wb", buffering=256 * 1024)
-        wstream: BinaryIO = cast(BinaryIO, cast(object, wfile))
-        write_request(wstream, request)
-
-        raw_stream = sock.makefile("rb", buffering=256 * 1024)
-        stream: BinaryIO = cast(BinaryIO, cast(object, raw_stream))
-
-        header = read_header(stream)
-        if on_header is not None:
-            on_header(header)
-
-        result = PrefillResult(header=header)
-        kv_by_layer: dict[int, list[KVChunk]] = defaultdict(list)
-        chunks_received = 0
-
-        while True:
-            msg = read_message(stream)
-            if msg is None:
-                break
-            if isinstance(msg, KVChunk):
-                kv_by_layer[msg.layer_idx].append(msg)
-                chunks_received += 1
-                if on_kv_chunk is not None:
-                    on_kv_chunk(msg, chunks_received)
-            elif isinstance(msg, ArraysState):
-                result.arrays[msg.layer_idx] = msg.arrays
-            elif isinstance(msg, Done):
-                result.total_tokens = msg.total_tokens
-                break
-            else:
-                raise RuntimeError(f"Prefill server error [{msg.code}]: {msg.message}")
-
-        result.kv_chunks = dict(kv_by_layer)
-        return result
-    finally:
-        sock.close()
-
-
-def ingest_into_mlx_cache(
-    result: PrefillResult,
-    caches: list[KVCache | RotatingKVCache | ArraysCache],
-    *,
-    start_pos: int = 0,
-) -> int:
-    max_received = max(
-        (sum(c.num_tokens for c in chunks) for chunks in result.kv_chunks.values()),
-        default=0,
-    )
-    final_offset = start_pos + max_received
-
-    for i, cache in enumerate(caches):
-        if i in result.kv_chunks:
-            chunks = result.kv_chunks[i]
-            if len(chunks) == 1:
-                k_nhd, v_nhd = chunk_to_mlx_nhd(chunks[0])
-            else:
-                decoded = [chunk_to_mlx_nhd(c) for c in chunks]
-                k_nhd = mx.concatenate([k for k, _ in decoded], axis=0)
-                v_nhd = mx.concatenate([v for _, v in decoded], axis=0)
-
-            if isinstance(cache, RotatingKVCache):
-                inject_rotating_kv_chunk(cache, k_nhd, v_nhd, final_offset)
-            elif isinstance(cache, KVCache):
-                if start_pos > 0:
-                    inject_kv_chunk(
-                        cache,
-                        k_nhd,
-                        v_nhd,
-                        final_offset,
-                        start_pos=start_pos,
-                        existing_k=cache.keys,
-                        existing_v=cache.values,
-                    )
-                else:
-                    inject_kv_chunk(cache, k_nhd, v_nhd, final_offset)
-
-        if i in result.arrays and isinstance(cache, ArraysCache):
-            inject_arrays_cache(cache, result.arrays[i])
-
-    return final_offset

src/exo/worker/engines/mlx/disaggregated/serve.py

@@ -1,86 +0,0 @@
-import time
-
-import mlx.core as mx
-from mlx_lm.sample_utils import make_sampler
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-
-from exo.worker.disaggregated.server import PrefillRequest
-from exo.worker.engines.mlx.cache import (
-    KVPrefixCache,
-    cache_length,
-    make_kv_cache,
-    snapshot_ssm_states,
-)
-from exo.worker.engines.mlx.generator.generate import prefill as mlx_prefill
-from exo.worker.engines.mlx.types import KVCacheType, Model
-from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens
-from exo.worker.runner.bootstrap import logger
-
-
-def run_prefill_for_request(
-    *,
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    group: mx.distributed.Group | None,
-    kv_prefix_cache: KVPrefixCache | None,
-    request: PrefillRequest,
-) -> KVCacheType:
-    prompt_tokens = mx.array(request.token_ids)
-    prompt_tokens = fix_unmatched_think_end_tokens(prompt_tokens, tokenizer)
-    n_tokens = int(prompt_tokens.shape[0])
-    t0 = time.perf_counter()
-
-    matched_index: int | None = None
-    prefix_hit_length = 0
-    if kv_prefix_cache is not None:
-        cache, remaining, matched_index, _ = kv_prefix_cache.get_kv_cache(
-            model, prompt_tokens
-        )
-        prefix_hit_length = n_tokens - int(remaining.shape[0])
-    else:
-        cache = make_kv_cache(model)
-        remaining = prompt_tokens
-
-    target_offset = max(0, n_tokens - 2)
-    new_tokens = max(0, target_offset - prefix_hit_length)
-    prefill_input = remaining[:new_tokens]
-    if int(prefill_input.shape[0]) > 0:
-        sampler = make_sampler(temp=1.0)
-        _ = mlx_prefill(
-            model=model,
-            tokenizer=tokenizer,
-            sampler=sampler,
-            prompt_tokens=prefill_input,
-            cache=cache,
-            group=group,
-            on_prefill_progress=None,
-            distributed_prompt_progress_callback=None,
-        )
-
-    if kv_prefix_cache is not None:
-        try:
-            cache_snapshots = [snapshot_ssm_states(cache)]
-            hit_ratio = prefix_hit_length / n_tokens if n_tokens > 0 else 0.0
-            if matched_index is not None and hit_ratio >= 0.5:
-                kv_prefix_cache.update_kv_cache(
-                    matched_index,
-                    prompt_tokens,
-                    cache,
-                    cache_snapshots,
-                    restore_pos=prefix_hit_length,
-                )
-            else:
-                kv_prefix_cache.add_kv_cache(prompt_tokens, cache, cache_snapshots)
-        except Exception:
-            logger.opt(exception=True).warning(
-                "Failed to save prefix cache on prefill server"
-            )
-
-    elapsed = time.perf_counter() - t0
-    final_offset = cache_length(cache)
-    logger.info(
-        f"Prefill: request_id={request.request_id} "
-        f"{n_tokens} tokens (prefix_hit={prefix_hit_length}, "
-        f"final_offset={final_offset}) in {elapsed * 1000:.0f}ms"
-    )
-    return cache

src/exo/worker/engines/mlx/disaggregated/tests/test_end_to_end.py

@@ -1,167 +0,0 @@
-from typing import BinaryIO
-
-import mlx.core as mx
-import numpy as np
-import pytest
-from mlx_lm.models.cache import KVCache
-
-from exo.worker.disaggregated.protocol import Header, write_done, write_header
-from exo.worker.disaggregated.server import PrefillRequest, PrefillServer
-from exo.worker.engines.mlx.disaggregated.adapter import (
-    send_mlx_kv_cache,
-    wire_dtype_from_cache,
-)
-from exo.worker.engines.mlx.disaggregated.client import (
-    ingest_into_mlx_cache,
-    remote_prefill_fetch,
-)
-
-
-def _equal(a: mx.array, b: mx.array) -> bool:
-    if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
-        return False
-    if a.dtype == mx.bfloat16:
-        return bool(
-            np.array_equal(np.asarray(a.view(mx.uint16)), np.asarray(b.view(mx.uint16)))
-        )
-    return bool(np.array_equal(np.asarray(a), np.asarray(b)))
-
-
-def _make_cache(seq_len: int, n_heads: int, head_dim: int) -> KVCache:
-    mx.random.seed(0)
-    cache = KVCache()
-    with mx.stream(mx.Device(mx.cpu)):
-        cache.keys = (
-            mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10
-        ).astype(mx.bfloat16)
-        cache.values = (
-            mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10
-        ).astype(mx.bfloat16)
-        mx.eval(cache.keys, cache.values)
-    cache.offset = seq_len
-    return cache
-
-
-def _stream_cache(
-    wfile: BinaryIO, cache: KVCache, *, request_id: str, start_pos: int = 0
-) -> None:
-    dtype = wire_dtype_from_cache([cache])
-    write_header(
-        wfile,
-        Header(
-            request_id=request_id,
-            model_id="test-model",
-            num_layers=1,
-            dtype=dtype,
-            start_pos=start_pos,
-        ),
-    )
-    tokens_sent = send_mlx_kv_cache(wfile, [cache], dtype=dtype, start_pos=start_pos)
-    write_done(wfile, tokens_sent)
-    wfile.flush()
-
-
-@pytest.mark.slow
-def test_server_client_roundtrip() -> None:
-    seq_len = 5
-    n_heads = 2
-    head_dim = 4
-    gold = _make_cache(seq_len, n_heads, head_dim)
-
-    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
-        _stream_cache(wfile, gold, request_id=job.request_id)
-        return True
-
-    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52417)
-    try:
-        result = remote_prefill_fetch(
-            endpoint="127.0.0.1:52417",
-            request=PrefillRequest(
-                model_id="test-model",
-                token_ids=list(range(seq_len)),
-                request_id="req-1",
-            ),
-        )
-        assert result.total_tokens == seq_len
-        assert 0 in result.kv_chunks
-
-        dst = KVCache()
-        final_offset = ingest_into_mlx_cache(result, [dst])
-        assert final_offset == seq_len
-        assert dst.offset == seq_len
-        dst_k = dst.keys
-        dst_v = dst.values
-        gold_k = gold.keys
-        gold_v = gold.values
-        assert dst_k is not None and dst_v is not None
-        assert gold_k is not None and gold_v is not None
-        assert _equal(dst_k, gold_k)
-        assert _equal(dst_v, gold_v)
-    finally:
-        server.stop()
-
-
-@pytest.mark.slow
-def test_server_reports_pickup_failure() -> None:
-    def resolve(_job: PrefillRequest, _wfile: BinaryIO) -> bool:
-        return False
-
-    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52418)
-    try:
-        with pytest.raises(RuntimeError, match="not picked up"):
-            _ = remote_prefill_fetch(
-                endpoint="127.0.0.1:52418",
-                request=PrefillRequest(
-                    model_id="test-model",
-                    token_ids=[1, 2, 3],
-                    request_id="never-registered",
-                ),
-            )
-    finally:
-        server.stop()
-
-
-@pytest.mark.slow
-def test_server_client_roundtrip_with_start_pos() -> None:
-    seq_len = 8
-    start_pos = 5
-    n_heads = 2
-    head_dim = 4
-    gold = _make_cache(seq_len, n_heads, head_dim)
-
-    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
-        _stream_cache(wfile, gold, request_id=job.request_id, start_pos=start_pos)
-        return True
-
-    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52419)
-    try:
-        result = remote_prefill_fetch(
-            endpoint="127.0.0.1:52419",
-            request=PrefillRequest(
-                model_id="test-model",
-                token_ids=list(range(seq_len)),
-                request_id="req-1",
-                start_pos=start_pos,
-            ),
-        )
-        assert result.total_tokens == seq_len - start_pos
-        assert result.header.start_pos == start_pos
-
-        dst = KVCache()
-        gold_k = gold.keys
-        gold_v = gold.values
-        assert gold_k is not None and gold_v is not None
-        dst.keys = mx.array(gold_k[:, :, :start_pos, :])
-        dst.values = mx.array(gold_v[:, :, :start_pos, :])
-        dst.offset = start_pos
-
-        final_offset = ingest_into_mlx_cache(result, [dst], start_pos=start_pos)
-        assert final_offset == seq_len
-        assert dst.offset == seq_len
-        dst_k = dst.keys
-        dst_v = dst.values
-        assert dst_k is not None and dst_v is not None
-        assert _equal(dst_k, gold_k)
-        assert _equal(dst_v, gold_v)
-    finally:
-        server.stop()

src/exo/worker/engines/mlx/disaggregated/tests/test_mlx_adapter.py

@@ -1,270 +0,0 @@
-import io
-
-import mlx.core as mx
-import numpy as np
-from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
-
-from exo.worker.disaggregated.protocol import (
-    ArraysState,
-    Done,
-    Header,
-    KVChunk,
-    TensorBlob,
-    read_header,
-    read_message,
-    write_done,
-    write_header,
-)
-from exo.worker.engines.mlx.disaggregated.adapter import (
-    array_to_bytes,
-    bhsd_to_nhd,
-    bytes_to_array,
-    chunk_to_mlx_nhd,
-    inject_arrays_cache,
-    inject_kv_chunk,
-    inject_rotating_kv_chunk,
-    nhd_to_bhsd,
-    send_mlx_kv_cache,
-    wire_dtype_from_cache,
-)
-from exo.worker.engines.mlx.disaggregated.client import (
-    PrefillResult,
-    ingest_into_mlx_cache,
-)
-
-
-def _equal(a: mx.array, b: mx.array) -> bool:
-    if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
-        return False
-    if a.dtype == mx.bfloat16:
-        return bool(
-            np.array_equal(np.asarray(a.view(mx.uint16)), np.asarray(b.view(mx.uint16)))
-        )
-    return bool(np.array_equal(np.asarray(a), np.asarray(b)))
-
-
-def _rand(shape: tuple[int, ...], dtype: mx.Dtype) -> mx.array:
-    mx.random.seed(0)
-    return (mx.random.uniform(shape=shape) * 10).astype(dtype)
-
-
-def _make_kv_cache(seq_len: int, n_heads: int, head_dim: int) -> KVCache:
-    cache = KVCache()
-    cache.keys = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
-    cache.values = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
-    cache.offset = seq_len
-    return cache
-
-
-def test_bytes_roundtrip_bf16() -> None:
-    x = _rand((2, 3, 4), mx.bfloat16)
-    y = bytes_to_array(array_to_bytes(x), (2, 3, 4), "bfloat16")
-    assert _equal(x, y)
-
-
-def test_bytes_roundtrip_f16() -> None:
-    x = _rand((5,), mx.float16)
-    y = bytes_to_array(array_to_bytes(x), (5,), "float16")
-    assert _equal(x, y)
-
-
-def test_bytes_roundtrip_f32() -> None:
-    x = _rand((2, 2), mx.float32)
-    y = bytes_to_array(array_to_bytes(x), (2, 2), "float32")
-    assert _equal(x, y)
-
-
-def test_bhsd_nhd_roundtrip() -> None:
-    bhsd = _rand((1, 4, 7, 8), mx.float32)
-    nhd = bhsd_to_nhd(bhsd)
-    assert tuple(nhd.shape) == (7, 4, 8)
-    back = nhd_to_bhsd(nhd)
-    assert _equal(bhsd, back)
-
-
-def test_kv_cache_inject_roundtrip() -> None:
-    n_heads, seq_len, head_dim = 3, 5, 4
-    k_bhsd = _rand((1, n_heads, seq_len, head_dim), mx.float32)
-    v_bhsd = _rand((1, n_heads, seq_len, head_dim), mx.float32)
-    k_nhd = bhsd_to_nhd(k_bhsd)
-    v_nhd = bhsd_to_nhd(v_bhsd)
-
-    cache = KVCache()
-    inject_kv_chunk(cache, k_nhd, v_nhd, offset=seq_len)
-    assert cache.offset == seq_len
-    assert cache.keys is not None and cache.values is not None
-    assert _equal(cache.keys, k_bhsd)
-    assert _equal(cache.values, v_bhsd)
-
-
-def test_arrays_cache_inject() -> None:
-    a = _rand((3,), mx.float32)
-    b = _rand((2, 2), mx.bfloat16)
-    blobs = [
-        TensorBlob(dtype="float32", shape=(3,), data=array_to_bytes(a)),
-        TensorBlob(dtype="bfloat16", shape=(2, 2), data=array_to_bytes(b)),
-    ]
-    cache = ArraysCache(size=2)
-    inject_arrays_cache(cache, blobs)
-    s0 = cache.state[0]
-    s1 = cache.state[1]
-    assert s0 is not None and s1 is not None
-    assert _equal(s0, a)
-    assert _equal(s1, b)
-
-
-def test_send_mlx_cache_end_to_end() -> None:
-    n_heads, head_dim = 2, 4
-    seq_len = 3
-    src = _make_kv_cache(seq_len, n_heads, head_dim)
-    k_bhsd, v_bhsd = src.keys, src.values
-    assert k_bhsd is not None and v_bhsd is not None
-
-    buf = io.BytesIO()
-    write_header(buf, Header(num_layers=1, dtype="bfloat16"))
-    tokens = send_mlx_kv_cache(buf, [src], dtype="bfloat16")
-    write_done(buf, tokens)
-    buf.seek(0)
-
-    got_hdr = read_header(buf)
-    assert got_hdr.num_layers == 1
-
-    msg = read_message(buf)
-    assert isinstance(msg, KVChunk)
-    assert msg.num_tokens == seq_len
-    k_nhd, v_nhd = chunk_to_mlx_nhd(msg)
-    dst = KVCache()
-    inject_kv_chunk(dst, k_nhd, v_nhd, offset=msg.num_tokens)
-
-    done = read_message(buf)
-    assert isinstance(done, Done)
-    assert done.total_tokens == seq_len
-
-    assert dst.offset == seq_len
-    assert dst.keys is not None and dst.values is not None
-    assert _equal(dst.keys, k_bhsd)
-    assert _equal(dst.values, v_bhsd)
-    _ = ArraysState
-
-
-def test_send_with_start_pos_only_ships_suffix() -> None:
-    n_heads, head_dim = 2, 4
-    seq_len, start_pos = 6, 4
-    src = _make_kv_cache(seq_len, n_heads, head_dim)
-
-    buf = io.BytesIO()
-    write_header(buf, Header(num_layers=1, dtype="bfloat16", start_pos=start_pos))
-    tokens = send_mlx_kv_cache(buf, [src], dtype="bfloat16", start_pos=start_pos)
-    write_done(buf, tokens)
-    buf.seek(0)
-
-    _ = read_header(buf)
-    msg = read_message(buf)
-    assert isinstance(msg, KVChunk)
-    assert msg.num_tokens == seq_len - start_pos
-
-
-def test_send_skips_layer_when_offset_below_start_pos() -> None:
-    n_heads, head_dim = 2, 4
-    seq_len, start_pos = 3, 5
-    src = _make_kv_cache(seq_len, n_heads, head_dim)
-
-    buf = io.BytesIO()
-    write_header(buf, Header(num_layers=1, dtype="bfloat16", start_pos=start_pos))
-    tokens = send_mlx_kv_cache(buf, [src], dtype="bfloat16", start_pos=start_pos)
-    write_done(buf, tokens)
-    buf.seek(0)
-
-    _ = read_header(buf)
-    msg = read_message(buf)
-    assert isinstance(msg, Done)
-    assert msg.total_tokens == 0
-    assert tokens == 0
-
-
-def test_wire_dtype_from_cache() -> None:
-    src = _make_kv_cache(3, 2, 4)
-    assert wire_dtype_from_cache([src]) == "bfloat16"
-
-    f32 = KVCache()
-    f32.keys = _rand((1, 2, 3, 4), mx.float32)
-    f32.values = _rand((1, 2, 3, 4), mx.float32)
-    f32.offset = 3
-    assert wire_dtype_from_cache([f32]) == "float32"
-
-
-def _decode_payload(payload: bytes) -> PrefillResult:
-    buf = io.BytesIO(payload)
-    hdr = read_header(buf)
-    result = PrefillResult(header=hdr)
-    while True:
-        msg = read_message(buf)
-        if msg is None:
-            break
-        if isinstance(msg, KVChunk):
-            result.kv_chunks.setdefault(msg.layer_idx, []).append(msg)
-        elif isinstance(msg, ArraysState):
-            result.arrays[msg.layer_idx] = msg.arrays
-        elif isinstance(msg, Done):
-            result.total_tokens = msg.total_tokens
-            break
-    return result
-
-
-def test_mixed_cache_roundtrip() -> None:
-    n_heads, head_dim, seq_len = 2, 4, 6
-
-    src_kv = _make_kv_cache(seq_len, n_heads, head_dim)
-
-    src_rot = RotatingKVCache(max_size=16, keep=0)
-    src_rot.keys = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
-    src_rot.values = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
-    src_rot.offset = seq_len
-    src_rot._idx = seq_len
-
-    src_arr = ArraysCache(size=2)
-    arr_a = _rand((3,), mx.bfloat16)
-    arr_b = _rand((2, 4), mx.bfloat16)
-    src_arr.state = [arr_a, arr_b]
-
-    buf = io.BytesIO()
-    write_header(
-        buf,
-        Header(request_id="req", model_id="m", num_layers=3, dtype="bfloat16"),
-    )
-    tokens_sent = send_mlx_kv_cache(buf, [src_kv, src_rot, src_arr], dtype="bfloat16")
-    write_done(buf, tokens_sent)
-    result = _decode_payload(buf.getvalue())
-
-    assert result.header.num_layers == 3
-    assert result.total_tokens == seq_len
-
-    dst_kv = KVCache()
-    dst_rot = RotatingKVCache(max_size=16, keep=0)
-    dst_arr = ArraysCache(size=2)
-    final_offset = ingest_into_mlx_cache(result, [dst_kv, dst_rot, dst_arr])
-
-    assert final_offset == seq_len
-
-    assert dst_kv.offset == seq_len
-    assert dst_kv.keys is not None and dst_kv.values is not None
-    src_kv_k, src_kv_v = src_kv.keys, src_kv.values
-    assert src_kv_k is not None and src_kv_v is not None
-    assert _equal(dst_kv.keys, src_kv_k)
-    assert _equal(dst_kv.values, src_kv_v)
-
-    assert dst_rot.offset == seq_len
-    assert dst_rot.keys is not None and dst_rot.values is not None
-    src_rot_k, src_rot_v = src_rot.keys, src_rot.values
-    assert src_rot_k is not None and src_rot_v is not None
-    assert _equal(dst_rot.keys, src_rot_k)
-    assert _equal(dst_rot.values, src_rot_v)
-    assert dst_rot._idx == seq_len
-
-    assert len(dst_arr.state) == 2
-    s0, s1 = dst_arr.state[0], dst_arr.state[1]
-    assert s0 is not None and s1 is not None
-    assert _equal(s0, arr_a)
-    assert _equal(s1, arr_b)
-    _ = inject_rotating_kv_chunk
-    _ = nhd_to_bhsd

src/exo/worker/engines/mlx/disaggregated/tests/test_protocol_roundtrip.py

@@ -1,154 +0,0 @@
-import io
-
-import pytest
-
-from exo.worker.disaggregated.protocol import (
-    ArraysState,
-    Done,
-    ErrorMessage,
-    Header,
-    KVChunk,
-    ProtocolError,
-    TensorBlob,
-    read_header,
-    read_message,
-    write_arrays_state,
-    write_done,
-    write_error,
-    write_header,
-    write_kv_chunk,
-)
-
-
-def _mk_bytes(n: int) -> bytes:
-    return bytes(i & 0xFF for i in range(n))
-
-
-def test_header_roundtrip() -> None:
-    hdr = Header(
-        request_id="r",
-        model_id="m",
-        num_layers=32,
-        dtype="bfloat16",
-        start_pos=42,
-    )
-    buf = io.BytesIO()
-    write_header(buf, hdr)
-    buf.seek(0)
-    got = read_header(buf)
-    assert got == hdr
-    assert got.dtype == "bfloat16"
-    assert got.num_layers == 32
-    assert got.start_pos == 42
-
-
-def test_kv_chunk_roundtrip() -> None:
-    num_tokens, n_heads, head_dim = 7, 4, 8
-    n_bytes = num_tokens * n_heads * head_dim * 2
-    keys = _mk_bytes(n_bytes)
-    values = _mk_bytes(n_bytes)[::-1]
-
-    buf = io.BytesIO()
-    write_kv_chunk(
-        buf,
-        layer_idx=3,
-        num_tokens=num_tokens,
-        n_heads=n_heads,
-        head_dim=head_dim,
-        dtype="bfloat16",
-        keys=keys,
-        values=values,
-    )
-    buf.seek(0)
-    msg = read_message(buf)
-    assert isinstance(msg, KVChunk)
-    assert msg.layer_idx == 3
-    assert msg.shape == (num_tokens, n_heads, head_dim)
-    assert msg.dtype == "bfloat16"
-    assert msg.keys == keys
-    assert msg.values == values
-
-
-def test_arrays_state_roundtrip() -> None:
-    arrs = [
-        TensorBlob(dtype="float32", shape=(2, 3), data=_mk_bytes(2 * 3 * 4)),
-        TensorBlob(dtype="bfloat16", shape=(5,), data=_mk_bytes(5 * 2)),
-    ]
-    buf = io.BytesIO()
-    write_arrays_state(buf, layer_idx=9, arrays=arrs)
-    buf.seek(0)
-    msg = read_message(buf)
-    assert isinstance(msg, ArraysState)
-    assert msg.layer_idx == 9
-    assert len(msg.arrays) == 2
-    assert msg.arrays[0].dtype == "float32"
-    assert msg.arrays[0].shape == (2, 3)
-    assert msg.arrays[0].data == arrs[0].data
-    assert msg.arrays[1].dtype == "bfloat16"
-    assert msg.arrays[1].shape == (5,)
-    assert msg.arrays[1].data == arrs[1].data
-
-
-def test_done_roundtrip() -> None:
-    buf = io.BytesIO()
-    write_done(buf, 1234)
-    buf.seek(0)
-    msg = read_message(buf)
-    assert isinstance(msg, Done)
-    assert msg.total_tokens == 1234
-
-
-def test_error_roundtrip() -> None:
-    buf = io.BytesIO()
-    write_error(buf, code=42, message="boom")
-    buf.seek(0)
-    msg = read_message(buf)
-    assert isinstance(msg, ErrorMessage)
-    assert msg.code == 42
-    assert msg.message == "boom"
-
-
-def test_stream_of_messages() -> None:
-    hdr = Header(num_layers=2, dtype="float32")
-    buf = io.BytesIO()
-    write_header(buf, hdr)
-    write_kv_chunk(
-        buf,
-        layer_idx=0,
-        num_tokens=1,
-        n_heads=1,
-        head_dim=2,
-        dtype="float32",
-        keys=_mk_bytes(1 * 1 * 2 * 4),
-        values=_mk_bytes(1 * 1 * 2 * 4),
-    )
-    write_arrays_state(
-        buf,
-        layer_idx=1,
-        arrays=[TensorBlob(dtype="float32", shape=(1,), data=_mk_bytes(4))],
-    )
-    write_done(buf, total_tokens=1)
-    buf.seek(0)
-
-    got_hdr = read_header(buf)
-    assert got_hdr == hdr
-
-    m1 = read_message(buf)
-    m2 = read_message(buf)
-    m3 = read_message(buf)
-    m4 = read_message(buf)
-    assert isinstance(m1, KVChunk)
-    assert isinstance(m2, ArraysState)
-    assert isinstance(m3, Done)
-    assert m4 is None
-
-
-def test_corrupt_message_raises() -> None:
-    buf = io.BytesIO()
-    write_header(buf, Header(num_layers=1, dtype="float32"))
-    buf.write((5).to_bytes(4, "big"))
-    buf.write(b"\xff\xff\xff\xff\xff")
-    buf.seek(0)
-    _ = read_header(buf)
-    with pytest.raises(ProtocolError):
-        _ = read_message(buf)

src/exo/worker/engines/mlx/disaggregated/tests/test_server_drain.py

@@ -1,120 +0,0 @@
-"""Server thread receives request, runs resolve in another thread (mimicking
-runner main thread + work queue), streams cache bytes."""
-
-import queue
-import threading
-from typing import BinaryIO
-
-import mlx.core as mx
-import numpy as np
-import pytest
-from mlx_lm.models.cache import KVCache
-
-from exo.utils.ports import random_ephemeral_port
-from exo.worker.disaggregated.protocol import Header, write_done, write_header
-from exo.worker.disaggregated.server import PrefillRequest, PrefillServer
-from exo.worker.engines.mlx.disaggregated.adapter import (
-    send_mlx_kv_cache,
-    wire_dtype_from_cache,
-)
-from exo.worker.engines.mlx.disaggregated.client import (
-    PrefillResult,
-    ingest_into_mlx_cache,
-    remote_prefill_fetch,
-)
-
-
-def _equal(a: mx.array, b: mx.array) -> bool:
-    if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
-        return False
-    if a.dtype == mx.bfloat16:
-        return bool(
-            np.array_equal(np.asarray(a.view(mx.uint16)), np.asarray(b.view(mx.uint16)))
-        )
-    return bool(np.array_equal(np.asarray(a), np.asarray(b)))
-
-
-def _make_cache(seq_len: int, n_heads: int, head_dim: int) -> KVCache:
-    mx.random.seed(0)
-    cache = KVCache()
-    cache.keys = (mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10).astype(
-        mx.bfloat16
-    )
-    cache.values = (
-        mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10
-    ).astype(mx.bfloat16)
-    cache.offset = seq_len
-    return cache
-
-
-@pytest.mark.slow
-def test_server_drains_via_main_thread() -> None:
-    seq_len = 4
-    n_heads = 2
-    head_dim = 4
-    gold = _make_cache(seq_len, n_heads, head_dim)
-
-    request_queue: queue.Queue[tuple[PrefillRequest, BinaryIO, threading.Event]] = (
-        queue.Queue()
-    )
-
-    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
-        done = threading.Event()
-        request_queue.put((job, wfile, done))
-        return done.wait(timeout=5)
-
-    server = PrefillServer(
-        resolve=resolve, host="127.0.0.1", port=(port := random_ephemeral_port())
-    )
-
-    def serve_one(wfile: BinaryIO) -> None:
-        dtype = wire_dtype_from_cache([gold])
-        write_header(
-            wfile,
-            Header(request_id="req-1", model_id="m", num_layers=1, dtype=dtype),
-        )
-        tokens = send_mlx_kv_cache(wfile, [gold], dtype=dtype)
-        write_done(wfile, tokens)
-        wfile.flush()
-
-    drained_job: list[PrefillRequest] = []
-    fetch_result: list[PrefillResult] = []
-
-    def fetcher() -> None:
-        fetch_result.append(
-            remote_prefill_fetch(
-                endpoint=f"127.0.0.1:{port}",
-                request=PrefillRequest(
-                    model_id="m", token_ids=list(range(seq_len)), request_id="req-1"
-                ),
-            )
-        )
-
-    fetch = threading.Thread(target=fetcher, daemon=True)
-    fetch.start()
-    try:
-        job, wfile, done = request_queue.get(timeout=5)
-        drained_job.append(job)
-        try:
-            serve_one(wfile)
-        finally:
-            done.set()
-        fetch.join(timeout=5)
-        assert fetch_result, "fetcher did not return"
-        result = fetch_result[0]
-        assert drained_job[0].request_id == "req-1"
-        assert result.total_tokens == seq_len
-
-        dst = KVCache()
-        ingest_into_mlx_cache(result, [dst])
-        assert dst.offset == seq_len
-        dst_k = dst.keys
-        dst_v = dst.values
-        gold_k = gold.keys
-        gold_v = gold.values
-        assert dst_k is not None and dst_v is not None
-        assert gold_k is not None and gold_v is not None
-        assert _equal(dst_k, gold_k)
-        assert _equal(dst_v, gold_v)
-    finally:
-        server.stop()

src/exo/worker/engines/mlx/dsml_encoding.py

@@ -34,17 +34,11 @@ def encode_messages(
     add_default_bos_token: bool = True,
     tools: Any = None,  # pyright: ignore[reportAny]
 ) -> str:
-    # V3.2 (like V4) is `tool_conditional`: when tools are in play, prior-turn
-    # reasoning_content must be retained so multi-step tool chains stay
-    # coherent.
-    effective_drop_thinking = drop_thinking
-    if tools:
-        effective_drop_thinking = False
     prompt: str = deepseek_v32.encode_messages(
         messages,
         thinking_mode=thinking_mode,
         context=context,
-        drop_thinking=effective_drop_thinking,
+        drop_thinking=drop_thinking,
         add_default_bos_token=add_default_bos_token,
         tools=tools,
     )

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -1,6 +1,5 @@
 import contextlib
 import time
-import uuid
 from dataclasses import dataclass, field
 from typing import Callable, Literal, cast
 
@@ -24,6 +23,7 @@ from exo.api.types import (
     Usage,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import GenerationResponse
 from exo.worker.engines.mlx.cache import (
@@ -40,12 +40,10 @@ from exo.worker.engines.mlx.generator.generate import (
     patch_embed_tokens,
     prefill,
 )
-from exo.worker.engines.mlx.generator.remote_prefill import remote_prefill
 from exo.worker.engines.mlx.patches.opt_batch_gen import (
     set_needs_topk,
     take_ready_topk,
 )
-from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.engines.mlx.utils_mlx import (
     fix_unmatched_think_end_tokens,
     system_prompt_token_count,
@@ -59,7 +57,6 @@ from exo.worker.engines.mlx.vision import (
 from exo.worker.runner.bootstrap import logger
 
 _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
-REMOTE_PREFILL_MIN_TOKENS = 1000
 
 
 def _stop_sequences(task_params: TextGenerationTaskParams) -> list[str]:
@@ -202,45 +199,17 @@ class ExoBatchGenerator:
             if vision is not None
             else contextlib.nullcontext()
         )
-        uncached_count = len(prompt_tokens)
-        use_remote = (
-            uncached_count > REMOTE_PREFILL_MIN_TOKENS
-            and task_params.prefill_endpoint is not None
-        )
-
-        _prefill_tps: float = 0.0
-        _prefill_tokens: int = 0
-        cache_snapshots: list[CacheSnapshot] = []
-        remote_prefilled = False
         with vision_ctx:
-            if use_remote and task_params.prefill_endpoint is not None:
-                try:
-                    _prefill_tps, _prefill_tokens, cache_snapshots = remote_prefill(
-                        prompt_tokens[:-1],
-                        cache,
-                        on_prefill_progress,
-                        endpoint=task_params.prefill_endpoint,
-                        request_id=str(uuid.uuid4()),
-                        model_id=str(task_params.model),
-                        start_pos=prefix_hit_length,
-                    )
-                    remote_prefilled = True
-                except Exception:
-                    logger.opt(exception=True).warning(
-                        "Remote prefill failed, falling back to local prefill"
-                    )
-
-            if not remote_prefilled:
-                _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
-                    self.model,
-                    self.tokenizer,
-                    sampler,
-                    prompt_tokens[:-1],
-                    cache,
-                    self.group,
-                    on_prefill_progress,
-                    distributed_prompt_progress_callback,
-                )
+            _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+                self.model,
+                self.tokenizer,
+                sampler,
+                prompt_tokens[:-1],
+                cache,
+                self.group,
+                on_prefill_progress,
+                distributed_prompt_progress_callback,
+            )
 
         prefix_cache_hit: Literal["none", "partial", "exact"] = "none"
         if matched_index is not None and prefix_hit_length > 0:
@@ -488,7 +457,6 @@ class ExoBatchGenerator:
 
     def close(self) -> None:
         self._mlx_gen.close()
-        mx.clear_cache()
 
     def _save_prefix_cache(
         self,

src/exo/worker/engines/mlx/generator/generate.py

@@ -2,7 +2,6 @@ import contextlib
 import functools
 import math
 import time
-import uuid
 from typing import Callable, Generator, cast, get_args
 
 import mlx.core as mx
@@ -10,6 +9,7 @@ from mlx_lm.generate import (
     maybe_quantize_kv_cache,
     stream_generate,
 )
+from mlx_lm.models.cache import ArraysCache, CacheList, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
@@ -23,6 +23,7 @@ from exo.api.types import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.shared.types.text_generation import (
     InputMessage,
     InputMessageContent,
@@ -42,11 +43,10 @@ from exo.worker.engines.mlx.auto_parallel import (
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
-    copy_snapshot_entry,
     encode_prompt,
     has_non_kv_caches,
-    is_non_trimmable_cache_entry,
     make_kv_cache,
+    restore_snapshot_entry,
     snapshot_ssm_states,
 )
 from exo.worker.engines.mlx.constants import (
@@ -55,8 +55,6 @@ from exo.worker.engines.mlx.constants import (
     KV_GROUP_SIZE,
     MAX_TOKENS,
 )
-from exo.worker.engines.mlx.generator.remote_prefill import remote_prefill
-from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     fix_unmatched_think_end_tokens,
@@ -72,8 +70,6 @@ from exo.worker.engines.mlx.vision import (
 )
 from exo.worker.runner.bootstrap import logger
 
-REMOTE_PREFILL_MIN_TOKENS = 1000
-
 generation_stream = mx.new_stream(mx.default_device())
 
 _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
@@ -376,10 +372,12 @@ def prefill(
     # Because of needing to roll back arrays cache, we will generate on 2 tokens so trim 1 more.
     pre_gen = snapshots[-2] if has_ssm else None
     for i, c in enumerate(cache):
-        non_trimmable = is_non_trimmable_cache_entry(c)
+        non_trimmable = isinstance(c, (ArraysCache, RotatingKVCache)) or (
+            isinstance(c, CacheList) and not bool(c.is_trimmable())  # type: ignore[reportUnknownMemberType]
+        )
         if has_ssm and non_trimmable:
             assert pre_gen is not None
-            restored = copy_snapshot_entry(pre_gen.states[i])
+            restored = restore_snapshot_entry(pre_gen.states[i])
             if restored is not None:
                 cache[i] = restored  # type: ignore
         else:
@@ -637,42 +635,17 @@ def mlx_generate(
         if vision is not None
         else contextlib.nullcontext()
     )
-    use_remote = (
-        len(prompt_tokens) > REMOTE_PREFILL_MIN_TOKENS
-        and task.prefill_endpoint is not None
-    )
-    remote_prefilled = False
-    prefill_tps = 0.0
-    prefill_tokens = 0
-    ssm_snapshots_list: list[CacheSnapshot] = []
     with maybe_vision_ctx:
-        if use_remote and task.prefill_endpoint is not None:
-            try:
-                prefill_tps, prefill_tokens, ssm_snapshots_list = remote_prefill(
-                    prompt_tokens[:-1],
-                    caches,
-                    on_prefill_progress,
-                    endpoint=task.prefill_endpoint,
-                    request_id=str(uuid.uuid4()),
-                    model_id=str(task.model),
-                    start_pos=prefix_hit_length,
-                )
-                remote_prefilled = True
-            except Exception:
-                logger.opt(exception=True).warning(
-                    "Remote prefill failed, falling back to local prefill"
-                )
-        if not remote_prefilled:
-            prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
-                model,
-                tokenizer,
-                sampler,
-                prompt_tokens[:-1],
-                caches,
-                group,
-                on_prefill_progress,
-                distributed_prompt_progress_callback,
-            )
+        prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
+            model,
+            tokenizer,
+            sampler,
+            prompt_tokens[:-1],
+            caches,
+            group,
+            on_prefill_progress,
+            distributed_prompt_progress_callback,
+        )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
     if kv_prefix_cache is not None and matched_index is not None and is_exact_hit:

src/exo/worker/engines/mlx/generator/remote_prefill.py

@@ -1,72 +0,0 @@
-import time
-from collections.abc import Callable
-from typing import cast
-
-import mlx.core as mx
-from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
-
-from exo.worker.disaggregated.protocol import Header, KVChunk
-from exo.worker.disaggregated.server import PrefillRequest
-from exo.worker.engines.mlx.cache import CacheSnapshot, snapshot_ssm_states
-from exo.worker.engines.mlx.disaggregated.client import (
-    ingest_into_mlx_cache,
-    remote_prefill_fetch,
-)
-from exo.worker.engines.mlx.types import KVCacheType
-from exo.worker.runner.bootstrap import logger
-
-
-def remote_prefill(
-    prompt_tokens: mx.array,
-    cache: KVCacheType,
-    on_prefill_progress: Callable[[int, int], None] | None,
-    *,
-    endpoint: str,
-    request_id: str,
-    model_id: str,
-    start_pos: int = 0,
-) -> tuple[float, int, list[CacheSnapshot]]:
-    t0 = time.perf_counter()
-    total_prompt_tokens = int(prompt_tokens.shape[0])
-    num_layers: int = 0
-
-    def _on_header(header: Header) -> None:
-        nonlocal num_layers
-        num_layers = header.num_layers
-
-    def _on_chunk(_chunk: KVChunk, chunks_received: int) -> None:
-        nonlocal num_layers
-        if on_prefill_progress is None:
-            return
-        if num_layers > 0 and chunks_received % num_layers == 0:
-            tokens_so_far = chunks_received // num_layers
-            on_prefill_progress(
-                min(tokens_so_far, total_prompt_tokens),
-                total_prompt_tokens,
-            )
-
-    request = PrefillRequest(
-        model_id=model_id,
-        token_ids=cast(list[int], prompt_tokens.tolist()),
-        start_pos=start_pos,
-        request_id=request_id,
-    )
-    result = remote_prefill_fetch(
-        endpoint, request, on_header=_on_header, on_kv_chunk=_on_chunk
-    )
-    t_received = time.perf_counter()
-
-    caches = cast(list[KVCache | RotatingKVCache | ArraysCache], list(cache))
-    final_offset = ingest_into_mlx_cache(result, caches, start_pos=start_pos)
-    t_done = time.perf_counter()
-
-    num_tokens = final_offset - start_pos
-    tps = num_tokens / max(t_done - t0, 0.001)
-
-    logger.info(
-        f"Remote prefill: {num_tokens} tokens (start_pos={start_pos}, "
-        f"final_offset={final_offset}) at {tps:.0f} tok/s, "
-        f"transfer={(t_received - t0) * 1000:.0f}ms, "
-        f"inject={(t_done - t_received) * 1000:.0f}ms"
-    )
-    return tps, num_tokens, [snapshot_ssm_states(cache)]

src/exo/worker/engines/mlx/patches/__init__.py

@@ -1,5 +1,6 @@
 from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
 from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
+from exo.worker.engines.mlx.patches.v4_offset_sync import apply as apply_v4_offset_sync
 
 _applied = False
 
@@ -11,3 +12,4 @@ def apply_mlx_patches() -> None:
     _applied = True
     patch_yarn_rope()
     apply_batch_gen_patch()
+    apply_v4_offset_sync()

src/exo/worker/engines/mlx/patches/v4_offset_sync.py

@@ -0,0 +1,74 @@
+from typing import Callable, cast
+
+import mlx.core as mx
+from mlx_lm.models.deepseek_v4 import Compressor, DeepseekV4Model
+
+_current_int_offset: int | None = None
+_applied: bool = False
+
+
+def _extract_int_offset(cache: object | None) -> int | None:
+    if cache is None:
+        return None
+    for entry in cast(list[object], cache):
+        inner_caches = getattr(entry, "caches", None)
+        win = inner_caches[0] if inner_caches is not None else entry
+        int_off = getattr(win, "_offset", None)
+        if isinstance(int_off, int):
+            return int_off
+        maybe_int = getattr(win, "offset", None)
+        if isinstance(maybe_int, int):
+            return maybe_int
+    return None
+
+
+_ModelCall = Callable[[DeepseekV4Model, mx.array, list[object] | None], mx.array]
+_CompressorCall = Callable[
+    [Compressor, mx.array, object, object, int, int, int], mx.array | None
+]
+
+
+def apply() -> None:
+    global _applied
+    if _applied:
+        return
+    _applied = True
+
+    original_model_call = cast(_ModelCall, DeepseekV4Model.__call__)
+
+    def patched_model_call(
+        self: DeepseekV4Model,
+        inputs: mx.array,
+        cache: list[object] | None = None,
+    ) -> mx.array:
+        global _current_int_offset
+        prev = _current_int_offset
+        _current_int_offset = _extract_int_offset(cache)
+        try:
+            return original_model_call(self, inputs, cache)
+        finally:
+            _current_int_offset = prev
+
+    DeepseekV4Model.__call__ = patched_model_call
+
+    original_compressor_call = cast(_CompressorCall, Compressor.__call__)
+
+    def patched_compressor_call(
+        self: Compressor,
+        x: mx.array,
+        state: object,
+        offset: object,
+        slot_compressed: int,
+        slot_kv_state: int,
+        slot_score_state: int,
+    ) -> mx.array | None:
+        if isinstance(offset, mx.array) and _current_int_offset is not None:
+            offset = _current_int_offset
+        return original_compressor_call(
+            self, x, state, offset, slot_compressed, slot_kv_state, slot_score_state
+        )
+
+    Compressor.__call__ = patched_compressor_call
+
+
+apply()

src/exo/worker/engines/mlx/tests/test_batch_generate.py

@@ -24,9 +24,9 @@ from transformers import AutoTokenizer
 
 # Import batch_generate to activate the right-padding BatchKVCache patch
 import exo.worker.engines.mlx.generator.batch_generate  # noqa: F401
+from exo.shared.types.mlx import Model
 from exo.worker.engines.mlx.cache import encode_prompt, make_kv_cache
 from exo.worker.engines.mlx.generator.generate import prefill
-from exo.worker.engines.mlx.types import Model
 
 NUM_STEPS = 20