Optimizations 6

.envrc

@@ -1,8 +1 @@
 use flake
-
-# creates .venv if doesn't exist and loads its environment
-export VIRTUAL_ENV=".venv"
-if ! [ -d "./$VIRTUAL_ENV" ]; then
-  uv venv
-fi
-layout python

.gitignore

@@ -40,4 +40,3 @@ bench/**/*.json
 tmp/models
 /build/exo
 /.claude/skills
-/.claude

.mlx_typings/mlx_lm/models/cache.pyi

@@ -191,13 +191,10 @@ class RotatingKVCache(_BaseCache):
     def state(self, v):  # -> None:
         ...
     @property
-    def meta_state(self):  # -> tuple[str, ...]:
-        ...
+    def meta_state(self) -> tuple[str, ...]: ...
     @meta_state.setter
-    def meta_state(self, v):  # -> None:
-        ...
-    def is_trimmable(self):  # -> bool:
-        ...
+    def meta_state(self, v: tuple[str, ...]) -> None: ...
+    def is_trimmable(self) -> bool: ...
     def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...

Cargo.lock

@@ -916,13 +916,11 @@ dependencies = [
  "libp2p",
  "log",
  "networking",
- "pidfile-rs",
  "pin-project",
  "pyo3",
  "pyo3-async-runtimes",
  "pyo3-log",
  "pyo3-stub-gen",
- "thiserror 2.0.17",
  "tokio",
  "util",
 ]
@@ -966,16 +964,6 @@ version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3a3076410a55c90011c298b04d0cfa770b00fa04e1e3c97d3f6c9de105a03844"
 
-[[package]]
-name = "flopen"
-version = "0.1.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fbfb8b5fbd1f27929f216650081a07b6ceb0741f0542c8c43ff7ef8e93a35a5d"
-dependencies = [
- "libc",
- "nix 0.31.2",
-]
-
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -1801,9 +1789,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
 
 [[package]]
 name = "libc"
-version = "0.2.186"
+version = "0.2.178"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
+checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091"
 
 [[package]]
 name = "libp2p"
@@ -2819,18 +2807,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "nix"
-version = "0.31.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3"
-dependencies = [
- "bitflags 2.10.0",
- "cfg-if",
- "cfg_aliases",
- "libc",
-]
-
 [[package]]
 name = "nohash-hasher"
 version = "0.2.0"
@@ -3084,18 +3060,6 @@ dependencies = [
  "siphasher",
 ]
 
-[[package]]
-name = "pidfile-rs"
-version = "0.3.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d1a8aa9a30b1b65ef48b333931b80f2324a14e00208eb2b8f5788f1180791bcc"
-dependencies = [
- "flopen",
- "libc",
- "log",
- "thiserror 1.0.69",
-]
-
 [[package]]
 name = "pin-project"
 version = "1.1.10"
@@ -3704,7 +3668,7 @@ dependencies = [
  "netlink-packet-utils",
  "netlink-proto",
  "netlink-sys",
- "nix 0.26.4",
+ "nix",
  "thiserror 1.0.69",
  "tokio",
 ]

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,8 +504,50 @@ 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"

bench/eval_tool_calls.py

@@ -15,8 +15,9 @@ from pathlib import Path
 from typing import Any, Literal
 
 import httpx
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
     instance_id_from_instance,

bench/exo_bench.py

@@ -30,8 +30,9 @@ from pathlib import Path
 from statistics import mean
 from typing import Any
 
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,

bench/exo_eval.py

@@ -42,8 +42,9 @@ from pathlib import Path
 from typing import Any
 
 import httpx
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,

bench/harness.py

@@ -1,39 +1,129 @@
 # type: ignore
-"""Instance lifecycle helpers for exo clusters.
-
-Provides utilities for placing instances, waiting for readiness,
-managing downloads, filtering placements, and common CLI arguments.
-"""
-
 from __future__ import annotations
 
 import argparse
-import contextlib
+import http.client
+import json
 import os
 import time
-from enum import Enum
+from collections.abc import Iterator
 from typing import Any
+from urllib.parse import urlencode
 
 from loguru import logger
 
-from .client import ExoClient, ExoHttpError
-
-
-class Sharding(str, Enum):
-    PIPELINE = "Pipeline"  # layers split across nodes
-    TENSOR = "Tensor"  # layers split within (across nodes)
-
-
-class Comm(str, Enum):
-    RING = "MlxRing"  # ring all-reduce over network
-    JACCL = "MlxJaccl"  # RDMA over Thunderbolt
-
-
 _SETTLE_INITIAL_BACKOFF_S = 1.0
 _SETTLE_MAX_BACKOFF_S = 60.0
 _SETTLE_BACKOFF_MULTIPLIER = 2.0
 
 
+class ExoHttpError(RuntimeError):
+    def __init__(self, status: int, reason: str, body_preview: str):
+        super().__init__(f"HTTP {status} {reason}: {body_preview}")
+        self.status = status
+
+
+class ExoClient:
+    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
+        self.host = host
+        self.port = port
+        self.timeout_s = timeout_s
+
+    def request_json(
+        self,
+        method: str,
+        path: str,
+        params: dict[str, Any] | None = None,
+        body: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Any:
+        if not path.startswith("/"):
+            path = "/" + path
+        if params:
+            path = path + "?" + urlencode(params)
+
+        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
+        try:
+            payload: bytes | None = None
+            hdrs: dict[str, str] = {"Accept": "application/json"}
+
+            if body is not None:
+                payload = json.dumps(body).encode("utf-8")
+                hdrs["Content-Type"] = "application/json"
+            if headers:
+                hdrs.update(headers)
+
+            conn.request(method.upper(), path, body=payload, headers=hdrs)
+            resp = conn.getresponse()
+            raw = resp.read()
+            text = raw.decode("utf-8", errors="replace") if raw else ""
+
+            if resp.status >= 400:
+                raise ExoHttpError(resp.status, resp.reason, text[:300])
+
+            if not text:
+                return None
+            return json.loads(text)
+        finally:
+            conn.close()
+
+    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}")
+        except ExoHttpError as e:
+            if e.status == 404:
+                return None
+            raise
+
+    def get_instance(self, instance_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"instances/{instance_id}")
+
+    def get_runner(self, runner_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"runners/{runner_id}")
+
+    def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
+        return self.get_state_path(f"downloads/{node_id}")
+
+    def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeDisk/{node_id}")
+
+    def get_node_system(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeSystem/{node_id}")
+
+    def get_node_identities(self) -> dict[str, Any] | None:
+        return self.get_state_path("nodeIdentities")
+
+    def get_topology(self) -> dict[str, Any] | None:
+        return self.get_state_path("topology")
+
+
 def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
     if len(instance) != 1:
         raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
@@ -465,6 +555,7 @@ def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
     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
@@ -498,7 +589,9 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         help="Only consider placements using >= this many nodes.",
     )
     ap.add_argument(
-        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
+        "--instance-meta",
+        choices=["ring", "jaccl", "vllm", "both"],
+        default="both",
     )
     ap.add_argument(
         "--sharding", choices=["pipeline", "tensor", "both"], default="both"
@@ -532,112 +625,3 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Reuse an existing running instance for this model instead of creating a new one.",
     )
-
-
-# ---------------------------------------------------------------------------
-# Cluster/instance orchestration helpers (used by tests, bench, eval)
-# ---------------------------------------------------------------------------
-
-
-def get_instance_ids(client: ExoClient) -> set[str]:
-    """Return the set of current instance IDs from cluster state."""
-    state = client.request_json("GET", "/state") or {}
-    result: set[str] = set()
-    for instance in state.get("instances", {}).values():
-        with contextlib.suppress(Exception):
-            result.add(instance_id_from_instance(instance))
-    return result
-
-
-def wait_for_cluster_ready(
-    client: ExoClient, expected_nodes: int = 1, timeout: float = 120.0
-) -> None:
-    """Wait until the cluster has all expected nodes visible and reporting memory.
-
-    Placement requires nodeMemory for all nodes in a cycle. This polls until
-    both nodeIdentities and nodeMemory have at least `expected_nodes` entries.
-    """
-    start = time.time()
-    while time.time() - start < timeout:
-        try:
-            state = client.request_json("GET", "/state") or {}
-            if (
-                len(state.get("nodeIdentities", {})) >= expected_nodes
-                and len(state.get("nodeMemory", {})) >= expected_nodes
-            ):
-                return
-        except Exception:
-            pass
-        time.sleep(1.0)
-    raise TimeoutError(f"Cluster not ready: expected {expected_nodes} nodes")
-
-
-def place_instance(
-    client: ExoClient,
-    model_id: str,
-    *,
-    sharding: Sharding = Sharding.PIPELINE,
-    comm: Comm = Comm.RING,
-    min_nodes: int = 1,
-    timeout: float = 600.0,
-    placement_retries: int = 10,
-    placement_retry_delay: float = 10.0,
-) -> str:
-    """Place an instance and wait for it to be ready. Returns the instance_id.
-
-    The /place_instance API returns a command_id, but instances are stored
-    under a separately-generated instance_id. This polls cluster state for the
-    new instance, retrying placement if the cluster is still settling.
-    """
-    wait_for_cluster_ready(client, expected_nodes=min_nodes)
-
-    body = {
-        "model_id": model_id,
-        "sharding": sharding.value,
-        "instance_meta": comm.value,
-        "min_nodes": min_nodes,
-    }
-
-    instance_id: str | None = None
-    for attempt in range(placement_retries):
-        before_ids = get_instance_ids(client)
-        client.request_json("POST", "/place_instance", body=body)
-
-        poll_deadline = time.time() + 30.0
-        while time.time() < poll_deadline:
-            new_ids = get_instance_ids(client) - before_ids
-            if new_ids:
-                instance_id = next(iter(new_ids))
-                break
-            time.sleep(1.0)
-
-        if instance_id is not None:
-            break
-
-        if attempt < placement_retries - 1:
-            time.sleep(placement_retry_delay)
-
-    if instance_id is None:
-        raise TimeoutError(
-            f"Placement failed after {placement_retries} attempts "
-            f"({sharding.value}/{comm.value} for {model_id})"
-        )
-
-    wait_for_instance_ready(client, instance_id, timeout=timeout)
-    return instance_id
-
-
-def cleanup_all_instances(client: ExoClient) -> None:
-    """Remove all running instances from the cluster."""
-    state = client.request_json("GET", "/state") or {}
-    for instance in state.get("instances", {}).values():
-        with contextlib.suppress(Exception):
-            iid = instance_id_from_instance(instance)
-            client.request_json("DELETE", f"/instance/{iid}")
-            wait_for_instance_gone(client, iid, timeout=30.0)
-
-
-def is_model_downloaded(client: ExoClient, model_id: str) -> bool:
-    response = client.request_json("GET", "/models", params={"status": "downloaded"})
-    data = (response or {}).get("data", [])
-    return all(model.get("id") == model_id for model in data)

bench/prefill-decode.toml

@@ -12,24 +12,24 @@ timeout = 7200.0
 settle_timeout = 60.0
 
 # Workload
-pp = [4096]
-tg = [512]
+pp = [4096, 8192]
+tg = [128]
 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"
+model = "sakamakismile/Qwen3.6-27B-NVFP4"
+node = "gx10-de89"
+instance_meta = "vllm"
 sharding = "pipeline"
 min_nodes = 1
 max_nodes = 1
 
 [decode]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "james"
+model = "mlx-community/Qwen3.6-27B-4bit"
+node = "Ryuichi’s MacBook Pro"
 instance_meta = "ring"
 sharding = "pipeline"
 min_nodes = 1

bench/prefill_decode_bench.py

@@ -31,12 +31,14 @@ from typing import Any
 
 from exo_bench import (
     PromptSizer,
+    SystemMetricsSampler,
     format_peak_memory,
     load_tokenizer_for_bench,
     parse_int_list,
 )
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     instance_id_from_instance,
     node_ids_from_instance,
@@ -277,6 +279,7 @@ def _run_phase(
     warmup: int,
     repeat: int,
     common_meta: dict[str, Any],
+    sampler: SystemMetricsSampler | None = None,
 ) -> list[dict[str, Any]]:
     logger.info(f"=== phase: {label} (model={model_id}) ===")
     rows: list[dict[str, Any]] = []
@@ -287,10 +290,13 @@ def _run_phase(
     for pp, tg in pp_tg_pairs:
         logger.info(f"--- {label}: pp={pp} tg={tg} ---")
         runs: list[dict[str, Any]] = []
+        inference_windows: list[tuple[float, float]] = []
         for r in range(repeat):
             time.sleep(2)
             try:
+                inf_t0 = time.monotonic()
                 row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
+                inference_windows.append((inf_t0, time.monotonic()))
             except Exception as e:
                 logger.error(e)
                 continue
@@ -314,11 +320,26 @@ def _run_phase(
             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)
+            energy_str = ""
+            if sampler is not None and inference_windows:
+                joules = sum(
+                    sampler.energy_between(t0, t1) for t0, t1 in inference_windows
+                )
+                inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
+                avg_watts = joules / inf_seconds if inf_seconds > 0 else 0.0
+                energy_per_run = joules / len(runs) if runs else 0.0
+                energy_str = (
+                    f"    energy={joules:.1f}J ({avg_watts:.1f}W avg over "
+                    f"{inf_seconds:.1f}s inference, {energy_per_run:.1f}J/run)"
+                )
+                for run_row, (t0, t1) in zip(runs, inference_windows, strict=False):
+                    run_row["energy_joules"] = sampler.energy_between(t0, t1)
+                    run_row["inference_window_s"] = t1 - t0
             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"
+                f"avg_elapsed={avg_elapsed:.2f}s{energy_str}"
             )
         time.sleep(2)
     return rows
@@ -331,14 +352,36 @@ def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, fl
         grouped.setdefault(key, []).append(r)
     out: dict[tuple[int, int], dict[str, float]] = {}
     for key, runs in grouped.items():
+        energy_runs = [x.get("energy_joules") for x in runs if "energy_joules" in x]
+        window_runs = [
+            x.get("inference_window_s") for x in runs if "inference_window_s" in x
+        ]
         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),
+            "prompt_tokens": mean(x["stats"]["prompt_tokens"] for x in runs),
+            "gen_tokens": mean(x["stats"]["generation_tokens"] for x in runs),
+            "energy_j": mean(energy_runs) if energy_runs else 0.0,
+            "inference_window_s": mean(window_runs) if window_runs else 0.0,
         }
     return out
 
 
+def _normalised_seconds(summary: dict[str, float], pp: int, tg: int) -> float | None:
+    """Wall-clock time implied by reported tps for the *configured* pp/tg.
+
+    elapsed_s is not comparable across phases when models EOS at different
+    lengths. This formula reconstructs "what would this phase take to do
+    pp prompt tokens + tg generation tokens" using its own reported rates.
+    """
+    p_tps = summary.get("prompt_tps", 0.0)
+    g_tps = summary.get("gen_tps", 0.0)
+    if p_tps <= 0 or g_tps <= 0:
+        return None
+    return pp / p_tps + tg / g_tps
+
+
 def _print_diff(
     disagg_rows: list[dict[str, Any]],
     decode_alone_rows: list[dict[str, Any]],
@@ -349,14 +392,17 @@ def _print_diff(
     prefill_alone = _summarise(prefill_alone_rows)
     keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
 
-    width = 64
+    width = 110
     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}"
+            f"  {'phase':<16} {'elapsed':>9}  {'norm':>9}  "
+            f"{'prompt_tps':>11}  {'gen_tps':>8}  "
+            f"{'p_tok':>6}  {'g_tok':>6}  "
+            f"{'energy':>9}  {'avg_W':>7}"
         )
         for label, summary in (
             ("disaggregated", disagg.get(key)),
@@ -364,26 +410,51 @@ def _print_diff(
             ("prefill_alone", prefill_alone.get(key)),
         ):
             if summary is None:
-                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
+                logger.info(
+                    f"  {label:<16} {'—':>9}  {'—':>9}  "
+                    f"{'—':>11}  {'—':>8}  {'—':>6}  {'—':>6}  "
+                    f"{'—':>9}  {'—':>7}"
+                )
                 continue
+            norm = _normalised_seconds(summary, pp, tg)
+            norm_str = f"{norm:>8.2f}s" if norm is not None else f"{'—':>9}"
+            energy = summary.get("energy_j", 0.0)
+            window = summary.get("inference_window_s", 0.0)
+            energy_str = f"{energy:>8.1f}J" if energy > 0 else f"{'—':>9}"
+            avg_w = energy / window if window > 0 else 0.0
+            avg_w_str = f"{avg_w:>6.1f}W" if avg_w > 0 else f"{'—':>7}"
             logger.info(
                 f"  {label:<16} "
-                f"{summary['elapsed_s']:>9.2f}s  "
+                f"{summary['elapsed_s']:>8.2f}s  "
+                f"{norm_str}  "
                 f"{summary['prompt_tps']:>11.1f}  "
-                f"{summary['gen_tps']:>9.2f}"
+                f"{summary['gen_tps']:>8.2f}  "
+                f"{summary['prompt_tokens']:>6.0f}  "
+                f"{summary['gen_tokens']:>6.0f}  "
+                f"{energy_str}  "
+                f"{avg_w_str}"
             )
 
         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"
-            )
+        d_norm = _normalised_seconds(d, pp, tg) if d else None
+        if d_norm and da:
+            da_norm = _normalised_seconds(da, pp, tg)
+            if da_norm:
+                logger.info(
+                    f"  norm speedup vs decode_alone:  {da_norm / d_norm:.2f}x  "
+                    f"(prefill {d['prompt_tps'] / da['prompt_tps']:.2f}x, "
+                    f"decode {d['gen_tps'] / da['gen_tps']:.2f}x)"
+                )
+        if d_norm and pa:
+            pa_norm = _normalised_seconds(pa, pp, tg)
+            if pa_norm:
+                logger.info(
+                    f"  norm speedup vs prefill_alone: {pa_norm / d_norm:.2f}x  "
+                    f"(prefill {d['prompt_tps'] / pa['prompt_tps']:.2f}x, "
+                    f"decode {d['gen_tps'] / pa['gen_tps']:.2f}x)"
+                )
     logger.info("─" * width)
 
 
@@ -681,6 +752,16 @@ def main() -> int:
     link_id = ""
     prefill_alive = False
     decode_alive = False
+    sampler_nodes = sorted(
+        {
+            *node_ids_from_instance(prefill_instance),
+            *node_ids_from_instance(decode_instance),
+        }
+    )
+    sampler = SystemMetricsSampler(
+        ExoClient(args.host, args.port, timeout_s=30), sampler_nodes
+    )
+    sampler.start()
     try:
         logger.info("Creating prefill instance...")
         client.request_json("POST", "/instance", body={"instance": prefill_instance})
@@ -699,6 +780,7 @@ def main() -> int:
                 warmup=args.warmup,
                 repeat=args.repeat,
                 common_meta=common_meta,
+                sampler=sampler,
             )
             all_rows.extend(prefill_alone_rows)
 
@@ -728,6 +810,7 @@ def main() -> int:
             warmup=args.warmup,
             repeat=args.repeat,
             common_meta=common_meta,
+            sampler=sampler,
         )
         all_rows.extend(disagg_rows)
 
@@ -752,11 +835,13 @@ def main() -> int:
                 warmup=args.warmup,
                 repeat=args.repeat,
                 common_meta=common_meta,
+                sampler=sampler,
             )
             all_rows.extend(decode_alone_rows)
 
             _print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
     finally:
+        sampler.stop()
         with contextlib.suppress(ExoHttpError):
             if link_id:
                 _delete_instance_link(client, link_id)

dashboard/src/lib/components/ChatSidebar.svelte

@@ -202,6 +202,7 @@
     let instanceType: string | null = null;
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
+    else if (instanceTag === "VllmInstance") instanceType = "vLLM";
 
     let sharding: string | null = null;
     const inst = instance as {

dashboard/src/lib/components/DeviceIcon.svelte

@@ -9,7 +9,7 @@
    */
 
   interface Props {
-    /** "macbook pro" | "mac studio" | "mac mini" etc. */
+    /** "macbook pro" | "mac studio" | "mac mini" | "dgx spark" | "linux" etc. */
     deviceType: string;
     /** Center X coordinate in SVG space */
     cx: number;
@@ -38,10 +38,43 @@
   const LOGO_NATIVE_WIDTH = 814;
   const LOGO_NATIVE_HEIGHT = 1000;
 
+  // NVIDIA logo SVG path
+  const NVIDIA_LOGO_PATH =
+    "M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
+
   const wireColor = "rgba(179,179,179,0.8)";
   const strokeWidth = 1.5;
 
   const modelLower = $derived(deviceType.toLowerCase());
+  const isSpark = $derived(
+    modelLower.includes("dgx") || modelLower.includes("gx10"),
+  );
+  const isLinux = $derived(!isSpark && modelLower.startsWith("linux"));
+  const isLinuxLaptop = $derived(isLinux && modelLower.includes("laptop"));
+
+  // ── DGX Spark dimensions ──
+  const dgxW = $derived(size * 1.55);
+  const dgxH = $derived(size * 0.58);
+  const dgxX = $derived(cx - dgxW / 2);
+  const dgxY = $derived(cy - dgxH / 2);
+  const dgxChassisX = $derived(dgxX - dgxW * 0.03);
+  const dgxChassisW = $derived(dgxW * 1.05);
+  const dgxHandleW = $derived(dgxW * 0.27);
+  const dgxHandleGap = $derived(dgxH * 0.05);
+  const dgxHandleH = $derived(dgxH - dgxHandleGap * 2);
+  const dgxHandleY = $derived(dgxY + dgxHandleGap);
+  const dgxInnerHandleW = $derived(dgxW * 0.12);
+  const dgxInnerHandleH = $derived(dgxHandleH - dgxH * 0.06);
+  const dgxLeftHandleX = $derived(dgxX + 4);
+  const dgxRightHandleX = $derived(dgxX + dgxW - dgxHandleW - 4);
+  const dgxClipId = $derived(`di-dgx-${uid}`);
+  const dgxTextureId = $derived(`di-dgx-tex-${uid}`);
+
+  // ── Linux Desktop dimensions (reuses Mac Studio proportions) ──
+  const linuxDesktopClipId = $derived(`di-linux-desktop-${uid}`);
+
+  // ── Linux Laptop dimensions (reuses MacBook proportions) ──
+  const linuxScreenClipId = $derived(`di-linux-screen-${uid}`);
 
   // ── Mac Studio dimensions (same ratios as TopologyGraph) ──
   const studioW = $derived(size * 1.25);
@@ -114,7 +147,264 @@
   const studioClipId = $derived(`di-studio-${uid}`);
 </script>
 
-{#if modelLower === "mac studio" || modelLower === "mac mini"}
+{#if isSpark}
+  <!-- DGX Spark -->
+  <defs>
+    <clipPath id={dgxClipId}>
+      <rect x={dgxX} y={dgxY} width={dgxW} height={dgxH} rx="3" />
+    </clipPath>
+    <pattern
+      id={dgxTextureId}
+      patternUnits="userSpaceOnUse"
+      width="8"
+      height="8"
+    >
+      <rect width="8" height="8" fill="#6f6248" />
+      <circle cx="2" cy="2" r="1" fill="#5a4f3b" opacity="0.5" />
+      <circle cx="6" cy="6" r="1" fill="#4a4232" opacity="0.45" />
+    </pattern>
+  </defs>
+
+  <!-- Main body -->
+  <rect
+    x={dgxChassisX}
+    y={dgxY}
+    width={dgxChassisW}
+    height={dgxH}
+    rx="3"
+    fill="url(#{dgxTextureId})"
+    stroke={wireColor}
+    stroke-width={strokeWidth}
+  />
+
+  <!-- Side border accents -->
+  <rect
+    x={dgxChassisX}
+    y={dgxY}
+    width={dgxW * 0.02}
+    height={dgxH}
+    fill="#8a7a56"
+  />
+  <rect
+    x={dgxChassisX + dgxChassisW - dgxW * 0.02}
+    y={dgxY}
+    width={dgxW * 0.02}
+    height={dgxH}
+    fill="#8a7a56"
+  />
+
+  <!-- Memory fill -->
+  {#if ramPercent > 0}
+    <rect
+      x={dgxX}
+      y={dgxY + dgxH - (ramPercent / 100) * dgxH}
+      width={dgxW}
+      height={(ramPercent / 100) * dgxH}
+      fill="rgba(255,215,0,0.45)"
+      clip-path="url(#{dgxClipId})"
+    />
+  {/if}
+
+  <!-- Left handle -->
+  <rect
+    x={dgxLeftHandleX}
+    y={dgxHandleY}
+    width={dgxHandleW}
+    height={dgxHandleH}
+    rx="2.4"
+    fill="#b3a170"
+    stroke="#403723"
+    stroke-width="0.7"
+  />
+  <rect
+    x={dgxLeftHandleX + dgxHandleW * 0.06}
+    y={dgxHandleY + dgxH * 0.03}
+    width={dgxInnerHandleW}
+    height={dgxInnerHandleH}
+    rx="1.6"
+    fill="#8a7a56"
+  />
+
+  <!-- Right handle -->
+  <rect
+    x={dgxRightHandleX}
+    y={dgxHandleY}
+    width={dgxHandleW}
+    height={dgxHandleH}
+    rx="2.4"
+    fill="#b3a170"
+    stroke="#403723"
+    stroke-width="0.7"
+  />
+  <rect
+    x={dgxRightHandleX + dgxHandleW - dgxInnerHandleW - dgxHandleW * 0.08}
+    y={dgxHandleY + dgxH * 0.03}
+    width={dgxInnerHandleW}
+    height={dgxInnerHandleH}
+    rx="1.6"
+    fill="#8a7a56"
+  />
+
+  <!-- NVIDIA logo (rotated 90deg on left handle) -->
+  {@const badgeW = dgxW * 0.09}
+  {@const badgeH = dgxHandleH * 0.5}
+  {@const badgeX = dgxLeftHandleX + dgxHandleW - badgeW - dgxHandleW * 0.06}
+  {@const badgeYPos = dgxHandleY + (dgxHandleH - badgeH) / 2}
+  {@const textSz = badgeW * 0.58}
+  {@const logoW = textSz * 1.2}
+  {@const logoH = logoW * (1.438 / 2.174)}
+  {@const ctrX = badgeX + badgeW / 2 - badgeW * 0.03}
+  {@const ctrY = badgeYPos + badgeH / 2}
+  {@const labelGap = badgeW * 0.15}
+  {@const totalW = logoW + labelGap + textSz * 3.6}
+  <g transform="rotate(90 {ctrX} {ctrY})">
+    <svg
+      x={ctrX - totalW / 2}
+      y={ctrY - logoH / 2}
+      width={logoW}
+      height={logoH}
+      viewBox="0 0 2.174 1.438"
+    >
+      <path d={NVIDIA_LOGO_PATH} fill="#76b900" />
+    </svg>
+    <text
+      x={ctrX - totalW / 2 + logoW + labelGap}
+      y={ctrY}
+      text-anchor="start"
+      dominant-baseline="middle"
+      fill="#8a7a56"
+      font-size={textSz}
+      font-family="monospace"
+      font-weight="700">NVIDIA</text
+    >
+  </g>
+{:else if isLinuxLaptop}
+  <!-- Linux Laptop — MacBook shape with Tux logo -->
+  <defs>
+    <clipPath id={linuxScreenClipId}>
+      <rect
+        x={mbScreenX + mbBezel}
+        y={mbY + mbBezel}
+        width={mbScreenW - mbBezel * 2}
+        height={mbScreenH - mbBezel * 2}
+        rx="2"
+      />
+    </clipPath>
+  </defs>
+
+  <rect
+    x={mbScreenX}
+    y={mbY}
+    width={mbScreenW}
+    height={mbScreenH}
+    rx="3"
+    fill="#1a1a1a"
+    stroke={wireColor}
+    stroke-width={strokeWidth}
+  />
+  <rect
+    x={mbScreenX + mbBezel}
+    y={mbY + mbBezel}
+    width={mbScreenW - mbBezel * 2}
+    height={mbScreenH - mbBezel * 2}
+    rx="2"
+    fill="#0a0a12"
+  />
+  {#if ramPercent > 0}
+    <rect
+      x={mbScreenX + mbBezel}
+      y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
+      width={mbScreenW - mbBezel * 2}
+      height={mbMemH}
+      fill="rgba(255,215,0,0.85)"
+      clip-path="url(#{linuxScreenClipId})"
+    />
+  {/if}
+
+  <!-- Terminal prompt on screen -->
+  <text
+    x={cx}
+    y={mbY + mbScreenH / 2}
+    text-anchor="middle"
+    dominant-baseline="middle"
+    fill="#FFFFFF"
+    opacity="0.9"
+    font-size={mbScreenH * 0.25}
+    font-family="SF Mono, Monaco, monospace"
+    font-weight="700">{">_"}</text
+  >
+
+  <path
+    d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
+      mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
+      mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
+    fill="#2c2c2c"
+    stroke={wireColor}
+    stroke-width="1"
+  />
+  <rect
+    x={mbKbX}
+    y={mbKbY}
+    width={mbKbW}
+    height={mbKbH}
+    fill="rgba(0,0,0,0.2)"
+    rx="2"
+  />
+  <rect
+    x={mbTpX}
+    y={mbTpY}
+    width={mbTpW}
+    height={mbTpH}
+    fill="rgba(255,255,255,0.08)"
+    rx="2"
+  />
+{:else if isLinux}
+  <!-- Linux Desktop — Mac Studio shape with Tux logo -->
+  <defs>
+    <clipPath id={linuxDesktopClipId}>
+      <rect
+        x={studioX}
+        y={studioY + studioTopH}
+        width={studioW}
+        height={studioH - studioTopH}
+        rx={studioCorner - 1}
+      />
+    </clipPath>
+  </defs>
+
+  <rect
+    x={studioX}
+    y={studioY}
+    width={studioW}
+    height={studioH}
+    rx={studioCorner}
+    fill="#1a1a1a"
+    stroke={wireColor}
+    stroke-width={strokeWidth}
+  />
+  {#if ramPercent > 0}
+    <rect
+      x={studioX}
+      y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
+      width={studioW}
+      height={studioMemH}
+      fill="rgba(255,215,0,0.75)"
+      clip-path="url(#{linuxDesktopClipId})"
+    />
+  {/if}
+
+  <!-- Terminal prompt on front face -->
+  <text
+    x={cx}
+    y={studioY + studioTopH + (studioH - studioTopH) / 2}
+    text-anchor="middle"
+    dominant-baseline="middle"
+    fill="rgba(255,255,255,0.5)"
+    font-size={(studioH - studioTopH) * 0.4}
+    font-family="SF Mono, Monaco, monospace"
+    font-weight="700">{">_"}</text
+  >
+{:else if modelLower === "mac studio" || modelLower === "mac mini"}
   <!-- Mac Studio / Mac Mini -->
   <defs>
     <clipPath id={studioClipId}>

dashboard/src/lib/components/ModelCard.svelte

@@ -23,7 +23,7 @@
     } | null;
     nodes?: Record<string, NodeInfo>;
     sharding?: "Pipeline" | "Tensor";
-    runtime?: "MlxRing" | "MlxJaccl";
+    runtime?: "MlxRing" | "MlxJaccl" | "Vllm";
     onLaunch?: () => void;
     tags?: string[];
     apiPreview?: PlacementPreview | null;
@@ -168,8 +168,10 @@
 
   function getDeviceType(
     name: string,
-  ): "macbook" | "studio" | "mini" | "unknown" {
+  ): "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown" {
     const lower = name.toLowerCase();
+    if (lower.includes("dgx") || lower.includes("gx10")) return "dgx";
+    if (lower.includes("linux")) return "linux";
     if (lower.includes("macbook")) return "macbook";
     if (lower.includes("studio")) return "studio";
     if (lower.includes("mini")) return "mini";
@@ -576,13 +578,17 @@
         class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
         title={runtime === "MlxRing"
           ? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
-          : "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
+          : runtime === "MlxJaccl"
+            ? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
+            : "vLLM: NVIDIA CUDA inference engine."}
       >
         {runtime === "MlxRing"
           ? "MLX Ring"
           : runtime === "MlxJaccl"
             ? "MLX RDMA"
-            : runtime}
+            : runtime === "Vllm"
+              ? "vLLM"
+              : runtime}
       </span>
     </div>
 
@@ -990,6 +996,81 @@
                     />
                   {/if}
                 </g>
+              {:else if node.deviceType === "dgx"}
+                <!-- DGX Spark icon -->
+                {@const s = node.iconSize}
+                {@const dgxW = s * 1.4}
+                {@const dgxH = s * 0.52}
+                <g transform="translate({-dgxW / 2}, {-dgxH / 2})">
+                  <!-- Chassis -->
+                  <rect
+                    x="0"
+                    y="0"
+                    width={dgxW}
+                    height={dgxH}
+                    rx="2"
+                    fill="#6f6248"
+                    stroke={node.isUsed ? "#FFD700" : "#4B5563"}
+                    stroke-width="1.5"
+                  />
+                  <!-- Side accents -->
+                  <rect
+                    x="0"
+                    y="0"
+                    width={dgxW * 0.02}
+                    height={dgxH}
+                    fill="#8a7a56"
+                  />
+                  <rect
+                    x={dgxW - dgxW * 0.02}
+                    y="0"
+                    width={dgxW * 0.02}
+                    height={dgxH}
+                    fill="#8a7a56"
+                  />
+                  <!-- Left handle -->
+                  <rect
+                    x={dgxW * 0.04}
+                    y={dgxH * 0.08}
+                    width={dgxW * 0.22}
+                    height={dgxH * 0.84}
+                    rx="2"
+                    fill="#b3a170"
+                    stroke="#403723"
+                    stroke-width="0.5"
+                  />
+                  <!-- Right handle -->
+                  <rect
+                    x={dgxW - dgxW * 0.04 - dgxW * 0.22}
+                    y={dgxH * 0.08}
+                    width={dgxW * 0.22}
+                    height={dgxH * 0.84}
+                    rx="2"
+                    fill="#b3a170"
+                    stroke="#403723"
+                    stroke-width="0.5"
+                  />
+                  <!-- Memory fill -->
+                  <rect
+                    x="2"
+                    y={dgxH - dgxH * (node.currentPercent / 100)}
+                    width={dgxW - 4}
+                    height={dgxH * (node.currentPercent / 100)}
+                    fill="rgba(255,215,0,0.35)"
+                  />
+                  {#if node.modelUsageGB > 0 && node.isUsed}
+                    <rect
+                      x="2"
+                      y={dgxH - dgxH * (node.newPercent / 100)}
+                      width={dgxW - 4}
+                      height={dgxH *
+                        ((node.newPercent - node.currentPercent) / 100)}
+                      fill="#FFD700"
+                      filter="url(#memGlow-{filterId})"
+                      class="animate-pulse-slow"
+                    />
+                  {/if}
+                </g>
               {:else}
                 <!-- Unknown device - hexagon -->
                 <g

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -9,6 +9,7 @@
     capabilities?: string[];
     family?: string;
     is_custom?: boolean;
+    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -19,6 +20,7 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
+    requiresVllm: boolean;
   }
 
   type DownloadAvailability = {
@@ -213,6 +215,14 @@
         <span class="font-mono text-sm text-white truncate">
           {group.name}
         </span>
+        {#if group.requiresVllm}
+          <span
+            class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
+            title="Requires vLLM runtime"
+          >
+            vLLM
+          </span>
+        {/if}
         <!-- Capability icons -->
         {#each group.capabilities.filter((c) => c !== "text") as cap}
           {#if cap === "thinking"}
@@ -523,6 +533,15 @@
             {variant.quantization || "default"}
           </span>
 
+          {#if variant.requires_vllm}
+            <span
+              class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
+              title="Requires vLLM runtime"
+            >
+              vLLM
+            </span>
+          {/if}
+
           <!-- Size -->
           <span
             class="text-xs font-mono flex-1 {getSizeClassForFitStatus(
@@ -628,6 +647,7 @@
                 variants: [variant],
                 smallestVariant: variant,
                 hasMultipleVariants: false,
+                requiresVllm: variant.requires_vllm === true,
               });
             }}
             title="View variant details"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -22,6 +22,7 @@
     is_custom?: boolean;
     tasks?: string[];
     hugging_face_id?: string;
+    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -32,6 +33,7 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
+    requiresVllm: boolean;
   }
 
   interface FilterState {
@@ -396,6 +398,7 @@
           variants: [],
           smallestVariant: model,
           hasMultipleVariants: false,
+          requiresVllm: true,
         });
       }
 
@@ -430,6 +433,7 @@
           (a.storage_size_megabytes || 0) - (b.storage_size_megabytes || 0),
       );
       group.hasMultipleVariants = group.variants.length > 1;
+      group.requiresVllm = group.variants.every((v) => v.requires_vllm);
     }
 
     // Convert to array and sort by smallest variant size (biggest first)
@@ -587,6 +591,7 @@
           variants: [model],
           smallestVariant: model,
           hasMultipleVariants: false,
+          requiresVllm: model.requires_vllm === true,
         });
       }
     }
@@ -1165,6 +1170,17 @@
           <span class="text-white/40">Variants:</span>
           <span class="text-white/70">{infoGroup.variants.length}</span>
         </div>
+        {#if infoGroup.requiresVllm}
+          <div class="flex items-center gap-2">
+            <span class="text-white/40">Runtime:</span>
+            <span
+              class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 tracking-wider uppercase"
+            >
+              vLLM
+            </span>
+            <span class="text-white/40 text-[11px]">required</span>
+          </div>
+        {/if}
         {#if infoGroup.variants.length > 0}
           <div class="mt-3 pt-3 border-t border-exo-yellow/10">
             <span class="text-white/40">Available quantizations:</span>

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -219,7 +219,7 @@
         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
+        Prefill is the compute-bound 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

dashboard/src/lib/components/TopologyGraph.svelte

@@ -117,6 +117,10 @@
   const LOGO_NATIVE_WIDTH = 814;
   const LOGO_NATIVE_HEIGHT = 1000;
 
+  // NVIDIA logo SVG path (from exo-nvidia)
+  const NVIDIA_LOGO_PATH =
+    "M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
+
   function formatBytes(bytes: number, decimals = 1): string {
     if (!bytes || bytes === 0) return "0B";
     const k = 1024;
@@ -554,6 +558,13 @@
       const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
 
       const modelLower = modelId.toLowerCase();
+      const identity = identitiesData[nodeInfo.id];
+      const nameLower = (friendlyName || "").toLowerCase();
+      const isSpark = modelLower.includes("dgx") || modelLower.includes("gx10");
+      const isLinux =
+        !isSpark &&
+        (modelLower.startsWith("linux") || identity?.osVersion === "Linux");
+      const isLinuxLaptop = isLinux && modelLower.includes("laptop");
 
       // Check node states for styling
       const isHighlighted = highlightedNodes.has(nodeInfo.id);
@@ -623,7 +634,382 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      if (modelLower === "mac studio") {
+      if (isSpark) {
+        // NVIDIA DGX Spark — gold chassis with textured front, side handles, and NVIDIA badge
+        iconBaseWidth = nodeRadius * 1.55;
+        iconBaseHeight = nodeRadius * 0.58;
+        const x = nodeInfo.x - iconBaseWidth / 2;
+        const y = nodeInfo.y - iconBaseHeight / 2;
+        const chassisX = x - iconBaseWidth * 0.03;
+        const chassisWidth = iconBaseWidth * 1.05;
+        const cornerRadius = 3;
+
+        const dgxClipId = `dgx-clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("clipPath")
+          .attr("id", dgxClipId)
+          .append("rect")
+          .attr("x", x)
+          .attr("y", y)
+          .attr("width", iconBaseWidth)
+          .attr("height", iconBaseHeight)
+          .attr("rx", cornerRadius);
+
+        // Chassis texture pattern
+        const textureId = `chassis-texture-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("pattern")
+          .attr("id", textureId)
+          .attr("patternUnits", "userSpaceOnUse")
+          .attr("width", 8)
+          .attr("height", 8);
+        const texturePattern = defs.select(`#${textureId}`);
+        texturePattern
+          .append("rect")
+          .attr("width", 8)
+          .attr("height", 8)
+          .attr("fill", "#6f6248");
+        texturePattern
+          .append("circle")
+          .attr("cx", 2)
+          .attr("cy", 2)
+          .attr("r", 1)
+          .attr("fill", "#5a4f3b")
+          .attr("opacity", 0.5);
+        texturePattern
+          .append("circle")
+          .attr("cx", 6)
+          .attr("cy", 6)
+          .attr("r", 1)
+          .attr("fill", "#4a4232")
+          .attr("opacity", 0.45);
+
+        // Main body
+        nodeG
+          .append("rect")
+          .attr("class", "node-outline")
+          .attr("x", chassisX)
+          .attr("y", y)
+          .attr("width", chassisWidth)
+          .attr("height", iconBaseHeight)
+          .attr("rx", cornerRadius)
+          .attr("fill", `url(#${textureId})`)
+          .attr("stroke", wireColor)
+          .attr("stroke-width", strokeWidth);
+
+        // Side border accents
+        const sideThickness = iconBaseWidth * 0.02;
+        nodeG
+          .append("rect")
+          .attr("x", chassisX)
+          .attr("y", y)
+          .attr("width", sideThickness)
+          .attr("height", iconBaseHeight)
+          .attr("fill", "#8a7a56");
+        nodeG
+          .append("rect")
+          .attr("x", chassisX + chassisWidth - sideThickness)
+          .attr("y", y)
+          .attr("width", sideThickness)
+          .attr("height", iconBaseHeight)
+          .attr("fill", "#8a7a56");
+
+        // Memory fill (bottom up)
+        if (ramUsagePercent > 0) {
+          const memFillHeight = (ramUsagePercent / 100) * iconBaseHeight;
+          nodeG
+            .append("rect")
+            .attr("x", x)
+            .attr("y", y + iconBaseHeight - memFillHeight)
+            .attr("width", iconBaseWidth)
+            .attr("height", memFillHeight)
+            .attr("fill", "rgba(255,215,0,0.45)")
+            .attr("clip-path", `url(#${dgxClipId})`);
+        }
+
+        // Side handles with inner recess
+        const handleWidth = iconBaseWidth * 0.27;
+        const handleGap = iconBaseHeight * 0.05;
+        const handleHeight = iconBaseHeight - handleGap * 2;
+        const handleY = y + handleGap;
+        const innerHandleWidth = iconBaseWidth * 0.12;
+        const innerHandleHeight = handleHeight - iconBaseHeight * 0.06;
+        const leftHandleX = x + 4;
+        const rightHandleX = x + iconBaseWidth - handleWidth - 4;
+
+        // Left handle
+        nodeG
+          .append("rect")
+          .attr("x", leftHandleX)
+          .attr("y", handleY)
+          .attr("width", handleWidth)
+          .attr("height", handleHeight)
+          .attr("rx", 2.4)
+          .attr("fill", "#b3a170")
+          .attr("stroke", "#403723")
+          .attr("stroke-width", 0.7);
+        nodeG
+          .append("rect")
+          .attr("x", leftHandleX + handleWidth * 0.06)
+          .attr("y", handleY + iconBaseHeight * 0.03)
+          .attr("width", innerHandleWidth)
+          .attr("height", innerHandleHeight)
+          .attr("rx", 1.6)
+          .attr("fill", "#8a7a56");
+
+        // Right handle
+        nodeG
+          .append("rect")
+          .attr("x", rightHandleX)
+          .attr("y", handleY)
+          .attr("width", handleWidth)
+          .attr("height", handleHeight)
+          .attr("rx", 2.4)
+          .attr("fill", "#b3a170")
+          .attr("stroke", "#403723")
+          .attr("stroke-width", 0.7);
+        nodeG
+          .append("rect")
+          .attr(
+            "x",
+            rightHandleX + handleWidth - innerHandleWidth - handleWidth * 0.08,
+          )
+          .attr("y", handleY + iconBaseHeight * 0.03)
+          .attr("width", innerHandleWidth)
+          .attr("height", innerHandleHeight)
+          .attr("rx", 1.6)
+          .attr("fill", "#8a7a56");
+
+        // NVIDIA logo + text label (rotated 90 deg on left handle)
+        const badgeWidth = iconBaseWidth * 0.09;
+        const badgeHeight = handleHeight * 0.5;
+        const badgeX =
+          leftHandleX + handleWidth - badgeWidth - handleWidth * 0.06;
+        const badgeY = handleY + (handleHeight - badgeHeight) / 2;
+        const textSize = badgeWidth * 0.58;
+        const logoWidth = textSize * 1.2;
+        const logoHeight = logoWidth * (1.438 / 2.174);
+        const centerX = badgeX + badgeWidth / 2 - badgeWidth * 0.03;
+        const centerY = badgeY + badgeHeight / 2;
+        const gap = badgeWidth * 0.15;
+        const totalWidth = logoWidth + gap + textSize * 3.6;
+
+        const labelGroup = nodeG
+          .append("g")
+          .attr("transform", `rotate(90 ${centerX} ${centerY})`);
+
+        labelGroup
+          .append("svg")
+          .attr("x", centerX - totalWidth / 2)
+          .attr("y", centerY - logoHeight / 2)
+          .attr("width", logoWidth)
+          .attr("height", logoHeight)
+          .attr("viewBox", "0 0 2.174 1.438")
+          .append("path")
+          .attr("d", NVIDIA_LOGO_PATH)
+          .attr("fill", "#76b900");
+
+        labelGroup
+          .append("text")
+          .attr("x", centerX - totalWidth / 2 + logoWidth + gap)
+          .attr("y", centerY)
+          .attr("text-anchor", "start")
+          .attr("dominant-baseline", "middle")
+          .attr("fill", "#8a7a56")
+          .attr("font-size", textSize)
+          .attr("font-family", "monospace")
+          .attr("font-weight", "700")
+          .text("NVIDIA");
+      } else if (isLinuxLaptop) {
+        // Linux Laptop — same shape as MacBook but with Tux logo
+        iconBaseWidth = nodeRadius * 1.6;
+        iconBaseHeight = nodeRadius * 1.15;
+        const x = nodeInfo.x - iconBaseWidth / 2;
+        const y = nodeInfo.y - iconBaseHeight / 2;
+
+        const screenHeight = iconBaseHeight * 0.7;
+        const baseHeight = iconBaseHeight * 0.3;
+        const screenWidth = iconBaseWidth * 0.85;
+        const screenX = nodeInfo.x - screenWidth / 2;
+        const screenBezel = 3;
+
+        const linuxScreenClipId = `linux-screen-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("clipPath")
+          .attr("id", linuxScreenClipId)
+          .append("rect")
+          .attr("x", screenX + screenBezel)
+          .attr("y", y + screenBezel)
+          .attr("width", screenWidth - screenBezel * 2)
+          .attr("height", screenHeight - screenBezel * 2)
+          .attr("rx", 2);
+
+        // Screen outer frame
+        nodeG
+          .append("rect")
+          .attr("class", "node-outline")
+          .attr("x", screenX)
+          .attr("y", y)
+          .attr("width", screenWidth)
+          .attr("height", screenHeight)
+          .attr("rx", 3)
+          .attr("fill", "#1a1a1a")
+          .attr("stroke", wireColor)
+          .attr("stroke-width", strokeWidth);
+
+        // Screen inner
+        nodeG
+          .append("rect")
+          .attr("x", screenX + screenBezel)
+          .attr("y", y + screenBezel)
+          .attr("width", screenWidth - screenBezel * 2)
+          .attr("height", screenHeight - screenBezel * 2)
+          .attr("rx", 2)
+          .attr("fill", "#0a0a12");
+
+        // Memory fill on screen
+        if (ramUsagePercent > 0) {
+          const memFillTotalHeight = screenHeight - screenBezel * 2;
+          const memFillActualHeight =
+            (ramUsagePercent / 100) * memFillTotalHeight;
+          nodeG
+            .append("rect")
+            .attr("x", screenX + screenBezel)
+            .attr(
+              "y",
+              y + screenBezel + (memFillTotalHeight - memFillActualHeight),
+            )
+            .attr("width", screenWidth - screenBezel * 2)
+            .attr("height", memFillActualHeight)
+            .attr("fill", "rgba(255,215,0,0.85)")
+            .attr("clip-path", `url(#${linuxScreenClipId})`);
+        }
+
+        // Terminal prompt on screen
+        nodeG
+          .append("text")
+          .attr("x", nodeInfo.x)
+          .attr("y", y + screenHeight / 2)
+          .attr("text-anchor", "middle")
+          .attr("dominant-baseline", "middle")
+          .attr("fill", "#FFFFFF")
+          .attr("opacity", 0.9)
+          .attr("font-size", screenHeight * 0.25)
+          .attr("font-family", "SF Mono, Monaco, monospace")
+          .attr("font-weight", "700")
+          .text(">_");
+
+        // Keyboard base (trapezoidal)
+        const baseY = y + screenHeight;
+        const baseTopWidth = screenWidth;
+        const baseBottomWidth = iconBaseWidth;
+        const baseTopX = nodeInfo.x - baseTopWidth / 2;
+        const baseBottomX = nodeInfo.x - baseBottomWidth / 2;
+
+        nodeG
+          .append("path")
+          .attr(
+            "d",
+            `M ${baseTopX} ${baseY} L ${baseTopX + baseTopWidth} ${baseY} L ${baseBottomX + baseBottomWidth} ${baseY + baseHeight} L ${baseBottomX} ${baseY + baseHeight} Z`,
+          )
+          .attr("fill", "#2c2c2c")
+          .attr("stroke", wireColor)
+          .attr("stroke-width", 1);
+
+        // Keyboard area
+        const keyboardX = baseTopX + 6;
+        const keyboardY = baseY + 3;
+        const keyboardWidth = baseTopWidth - 12;
+        const keyboardHeight = baseHeight * 0.55;
+        nodeG
+          .append("rect")
+          .attr("x", keyboardX)
+          .attr("y", keyboardY)
+          .attr("width", keyboardWidth)
+          .attr("height", keyboardHeight)
+          .attr("fill", "rgba(0,0,0,0.2)")
+          .attr("rx", 2);
+
+        // Trackpad
+        const trackpadWidth = baseTopWidth * 0.4;
+        const trackpadX = nodeInfo.x - trackpadWidth / 2;
+        const trackpadY = baseY + keyboardHeight + 5;
+        const trackpadHeight = baseHeight * 0.3;
+        nodeG
+          .append("rect")
+          .attr("x", trackpadX)
+          .attr("y", trackpadY)
+          .attr("width", trackpadWidth)
+          .attr("height", trackpadHeight)
+          .attr("fill", "rgba(255,255,255,0.08)")
+          .attr("rx", 2);
+      } else if (isLinux) {
+        // Linux Desktop — same shape as Mac Studio but with Tux logo
+        iconBaseWidth = nodeRadius * 1.25;
+        iconBaseHeight = nodeRadius * 0.85;
+        const x = nodeInfo.x - iconBaseWidth / 2;
+        const y = nodeInfo.y - iconBaseHeight / 2;
+        const cornerRadius = 4;
+        const topSurfaceHeight = iconBaseHeight * 0.15;
+
+        const linuxDesktopClipId = `linux-desktop-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("clipPath")
+          .attr("id", linuxDesktopClipId)
+          .append("rect")
+          .attr("x", x)
+          .attr("y", y + topSurfaceHeight)
+          .attr("width", iconBaseWidth)
+          .attr("height", iconBaseHeight - topSurfaceHeight)
+          .attr("rx", cornerRadius - 1);
+
+        // Main body
+        nodeG
+          .append("rect")
+          .attr("class", "node-outline")
+          .attr("x", x)
+          .attr("y", y)
+          .attr("width", iconBaseWidth)
+          .attr("height", iconBaseHeight)
+          .attr("rx", cornerRadius)
+          .attr("fill", "#1a1a1a")
+          .attr("stroke", wireColor)
+          .attr("stroke-width", strokeWidth);
+
+        // Memory fill
+        if (ramUsagePercent > 0) {
+          const memFillTotalHeight = iconBaseHeight - topSurfaceHeight;
+          const memFillActualHeight =
+            (ramUsagePercent / 100) * memFillTotalHeight;
+          nodeG
+            .append("rect")
+            .attr("x", x)
+            .attr(
+              "y",
+              y + topSurfaceHeight + (memFillTotalHeight - memFillActualHeight),
+            )
+            .attr("width", iconBaseWidth)
+            .attr("height", memFillActualHeight)
+            .attr("fill", "rgba(255,215,0,0.75)")
+            .attr("clip-path", `url(#${linuxDesktopClipId})`);
+        }
+
+        // Terminal prompt on front face
+        nodeG
+          .append("text")
+          .attr("x", nodeInfo.x)
+          .attr(
+            "y",
+            y + topSurfaceHeight + (iconBaseHeight - topSurfaceHeight) / 2,
+          )
+          .attr("text-anchor", "middle")
+          .attr("dominant-baseline", "middle")
+          .attr("fill", "rgba(255,255,255,0.5)")
+          .attr("font-size", (iconBaseHeight - topSurfaceHeight) * 0.4)
+          .attr("font-family", "SF Mono, Monaco, monospace")
+          .attr("font-weight", "700")
+          .text(">_");
+      } else if (modelLower === "mac studio") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
         iconBaseHeight = nodeRadius * 0.85;
@@ -1182,8 +1568,12 @@
           debugLabelY += debugLineHeight;
         }
 
-        const identity = identitiesData[nodeInfo.id];
-        if (identity?.osVersion) {
+        const dbgIdentity = identitiesData[nodeInfo.id];
+        if (dbgIdentity?.osVersion) {
+          const osLabel =
+            dbgIdentity.osVersion === "Linux"
+              ? "Linux"
+              : `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
           nodeG
             .append("text")
             .attr("x", nodeInfo.x)
@@ -1192,9 +1582,7 @@
             .attr("fill", "rgba(179,179,179,0.7)")
             .attr("font-size", debugFontSize)
             .attr("font-family", "SF Mono, Monaco, monospace")
-            .text(
-              `macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
-            );
+            .text(osLabel);
         }
       }
     });

dashboard/src/routes/+page.svelte

@@ -65,6 +65,7 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
+    featureFlags,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -702,7 +703,10 @@
       ? Object.keys(topologyData()!.nodes).length
       : 1;
     const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
-    const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
+    const instanceType =
+      nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
+        ? "MlxRing"
+        : selectedInstanceType;
     try {
       const placementResponse = await fetch(
         `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
@@ -783,6 +787,7 @@
       quantization?: string;
       base_model?: string;
       capabilities?: string[];
+      requires_vllm?: boolean;
     }>
   >([]);
   type ModelMemoryFitStatus =
@@ -886,7 +891,7 @@
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
-  type InstanceMeta = "MlxRing" | "MlxJaccl";
+  type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
@@ -932,7 +937,12 @@
     // Apply sharding and instance type unconditionally
     selectedSharding = defaults.sharding;
     selectedInstanceType =
-      defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
+      defaults.instanceType === "MlxRing"
+        ? "MlxRing"
+        : defaults.instanceType === "Vllm"
+          ? "Vllm"
+          : "MlxJaccl";
+    userPickedInstanceType = true;
 
     // Apply minNodes if valid (between 1 and maxNodes)
     if (
@@ -954,6 +964,23 @@
   }
 
   let selectedInstanceType = $state<InstanceMeta>("MlxRing");
+  let userPickedInstanceType = $state(false);
+  $effect(() => {
+    if (!userPickedInstanceType && featureFlags()["vllm_available"]) {
+      selectedInstanceType = "Vllm";
+    }
+  });
+  const selectedModelRequiresVllm = $derived.by((): boolean => {
+    const id = selectedPreviewModelId();
+    if (!id) return false;
+    const model = models.find((m) => m.id === id);
+    return model?.requires_vllm === true;
+  });
+  $effect(() => {
+    if (selectedModelRequiresVllm) {
+      selectedInstanceType = "Vllm";
+    }
+  });
   let selectedMinNodes = $state<number>(1);
   let minNodesInitialized = $state(false);
   let launchingModelId = $state<string | null>(null);
@@ -1146,9 +1173,7 @@
   }
 
   const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
-    selectedInstanceType === "MlxRing"
-      ? runtime === "MlxRing"
-      : runtime === "MlxJaccl";
+    runtime === selectedInstanceType;
 
   // Helper to check if a model can be launched (has valid placement with >= minNodes)
   function canModelFit(modelId: string): boolean {
@@ -2063,6 +2088,7 @@
     let instanceType = "Unknown";
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
+    else if (instanceTag === "VllmInstance") instanceType = "vLLM";
 
     const inst = instance as {
       shardAssignments?: {
@@ -3435,7 +3461,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 +4848,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 +4994,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
@@ -5772,14 +5795,18 @@
                     </div>
                     <div class="flex gap-2">
                       <button
+                        disabled={selectedModelRequiresVllm}
                         onclick={() => {
+                          if (selectedModelRequiresVllm) return;
                           selectedInstanceType = "MlxRing";
+                          userPickedInstanceType = true;
                           saveLaunchDefaults();
                         }}
-                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
-                        'MlxRing'
-                          ? 'bg-transparent text-exo-yellow border-exo-yellow'
-                          : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
+                          ? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
+                          : selectedInstanceType === 'MlxRing'
+                            ? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
+                            : 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5795,14 +5822,18 @@
                         TCP/IP
                       </button>
                       <button
+                        disabled={selectedModelRequiresVllm}
                         onclick={() => {
+                          if (selectedModelRequiresVllm) return;
                           selectedInstanceType = "MlxJaccl";
+                          userPickedInstanceType = true;
                           saveLaunchDefaults();
                         }}
-                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
-                        'MlxJaccl'
-                          ? 'bg-transparent text-exo-yellow border-exo-yellow'
-                          : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
+                          ? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
+                          : selectedInstanceType === 'MlxJaccl'
+                            ? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
+                            : 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5817,7 +5848,41 @@
                         </span>
                         RDMA (Fast)
                       </button>
+                      {#if featureFlags()["vllm_available"] || selectedModelRequiresVllm}
+                        <button
+                          onclick={() => {
+                            selectedInstanceType = "Vllm";
+                            userPickedInstanceType = true;
+                            saveLaunchDefaults();
+                          }}
+                          class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
+                          'Vllm'
+                            ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                            : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        >
+                          <span
+                            class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
+                            'Vllm'
+                              ? 'border-exo-yellow'
+                              : 'border-exo-medium-gray'}"
+                          >
+                            {#if selectedInstanceType === "Vllm"}
+                              <span
+                                class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
+                              ></span>
+                            {/if}
+                          </span>
+                          vLLM (CUDA)
+                        </button>
+                      {/if}
                     </div>
+                    {#if selectedModelRequiresVllm}
+                      <div
+                        class="mt-2 text-[11px] font-mono text-orange-300/80"
+                      >
+                        This model requires vLLM.
+                      </div>
+                    {/if}
                   </div>
 
                   <!-- Minimum Devices -->

flake.nix

@@ -146,7 +146,7 @@
                 config.treefmt.build.wrapper
 
                 # PYTHON
-                self'.packages.editableVenv
+                self'.packages.exo.passthru.evenv
                 uv
 
                 # RUST

justfile

@@ -40,6 +40,19 @@ build-app: rust-rebuild sync-clean package
     xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
     @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
 
+sync-cuda:
+    #!/usr/bin/env bash
+    set -euo pipefail
+    uv sync --extra vllm-cuda13 --extra mlx-cpu --no-install-package vllm
+    dest=".venv/lib/python3.13/site-packages"
+    [[ -d $dest/vllm ]] || {
+        nix build .#exo-cuda-13.passthru.evenv
+        # will also grab vllm-0.19.1-distinfo
+        cp -aL result/lib/python3.13/site-packages/vllm* .venv/lib/python3.13/site-packages
+        chmod -R u+rwX .venv/lib/python3.13/site-packages/vllm*
+        rm result
+    }
+
 clean:
     rm -rf **/__pycache__
     rm -rf target/

nix/vllm-setuppy-cmake.patch

@@ -0,0 +1,26 @@
+diff --git a/setup.py b/setup.py
+index 6dc2ed028..bdcc6354a 100644
+--- a/setup.py
++++ b/setup.py
+@@ -18,6 +18,13 @@ from setuptools import Extension, setup
+ from setuptools.command.build_ext import build_ext
+ 
+ 
++if "NIX_ATTRS_JSON_FILE" in os.environ:
++    with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
++        NIX_ATTRS = json.load(f)
++else:
++    NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
++
++
+ def load_module_from_path(module_name, path):
+     spec = importlib.util.spec_from_file_location(module_name, path)
+     module = importlib.util.module_from_spec(spec)
+@@ -184,6 +191,7 @@ class cmake_build_ext(build_ext):
+         cmake_args = [
+             "-DCMAKE_BUILD_TYPE={}".format(cfg),
+             "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
++            *NIX_ATTRS["cmakeFlags"],
+         ]
+ 
+         verbose = envs.VERBOSE

pyproject.toml

@@ -15,21 +15,18 @@ 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",
   "tomlkit>=0.14.0",
-  "mflux==0.17.2; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
-  "mlx-vlm>=0.3.11; sys_platform == 'darwin'",
   "transformers>=5.6.2",
+  "nvidia-ml-py>=13.595.45",
 ]
 
 [project.scripts]
@@ -40,7 +37,6 @@ exo = "exo.main:main"
 dev = [
   "basedpyright>=1.29.0",
   "pyinstaller>=6.17.0",
-  "playwright>=1.52.0",
   "pytest>=8.4.0",
   "pytest-asyncio>=1.0.0",
   "pytest-env",
@@ -49,26 +45,30 @@ dev = [
 
 [project.optional-dependencies]
 build = ["nanobind"]
-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'",
+mlx-none = ["anyio"]
+mlx = [
+  "mlx==0.31.2",
+  "mlx-lm",
+  "mlx-vlm>=0.3.11",
+  "mflux==0.17.5",
+  # pinning vllms versions for consistency.
+  "torch==2.10.0; sys_platform == 'darwin'",
+  "torch==2.10.0; sys_platform == 'linux'",
+  "torchaudio==2.10.0; sys_platform == 'darwin'",
+  "torchaudio==2.10.0; sys_platform == 'linux'",
+  "torchvision==0.25.0; sys_platform == 'darwin'",
+  "torchvision==0.25.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'",
+mlx-cpu = ["exo[mlx]", "mlx-cpu==0.31.2; sys_platform == 'linux'"]
+mlx-cuda12 = ["exo[mlx]", "mlx-cuda-12==0.31.1; sys_platform == 'linux'"]
+mlx-cuda13 = ["exo[mlx]", "mlx-cuda-13==0.31.1; sys_platform == 'linux'"]
+vllm-none = ["anyio"]
+vllm-cuda13 = [
+  "vllm[cuda13, fastsafetensors]; sys_platform == 'linux'",
+  "torch==2.10.0; sys_platform == 'linux'",
+  "torchaudio==2.10.0; sys_platform == 'linux'",
+  "torchvision==0.25.0; sys_platform == 'linux'",
 ]
 
 ###
@@ -76,18 +76,29 @@ cuda13 = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench", "tools"]
+members = ["rust/exo_pyo3_bindings", "bench"]
 
 [tool.uv.sources]
 exo-pyo3-bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
+mflux = { git = "http://github.com/evanev7/mflux", branch = "exo" }
+vllm = { git = "http://github.com/evanev7/vllm", branch = "exo2" }
 torch = [
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
-  { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
-  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
+]
+torchvision = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
+]
+torchaudio = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
 ]
-vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 
 [[tool.uv.index]]
 name = "pytorch-cu130"
@@ -95,8 +106,8 @@ url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 [[tool.uv.index]]
-name = "pytorch-cu120"
-url = "https://download.pytorch.org/whl/cu120"
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
 explicit = true
 
 [[tool.uv.index]]
@@ -113,7 +124,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench", "tools"]
+include = ["src", "bench"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -147,13 +158,6 @@ reportMissingModuleSource = false
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
 
-[[tool.basedpyright.executionEnvironments]]
-root = "bench"
-extraPaths = ["tools/src"]
-
-[[tool.basedpyright.executionEnvironments]]
-root = "tools/src"
-
 
 ###
 # uv configuration
@@ -164,11 +168,19 @@ root = "tools/src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
-conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
-constraint-dependencies = ["transformers>=5.6.2"]
-override-dependencies = [
-  "mlx==0.31.1; sys_platform=='linux'",
-  "mlx; sys_platform=='darwin'",
+override-dependencies = ["opencv-python; python_version < '0'"]
+conflicts = [
+  [
+    { extra = "mlx-cuda13" },
+    { extra = "mlx-cuda12" },
+    { extra = "mlx-cpu" },
+    { extra = "mlx-none" },
+  ],
+  [
+    { extra = "vllm-cuda13" },
+    { extra = "mlx-cuda12" },
+    { extra = "vllm-none" },
+  ],
 ]
 
 [tool.uv.extra-build-dependencies]
@@ -183,6 +195,7 @@ mlx = [
   "ninja",
 ]
 mlx-lm = ["setuptools"]
+mflux = ["uv_build"]
 xgrammar = [
   "nanobind",
   "setuptools",
@@ -228,5 +241,5 @@ pythonpath = "."
 asyncio_mode = "auto"
 markers = ["slow: marks tests as slow (deselected by default)"]
 env = ["EXO_TESTS=1"]
-addopts = "-m 'not slow' --ignore=tests"
+addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
 filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

python/parts.nix

@@ -10,8 +10,10 @@ let
       inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
       inherit (pkgs.config) cudaSupport;
       inherit (pkgs) cudaPackages;
-      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
-      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
+      libmlx_source =
+        if (builtins.elem "mlx-cuda13" members.exo or [ ]) then "mlx-cuda-13"
+        else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
+        else "mlx-cpu";
       python = pkgs.python313;
       cudaLibs = with cudaPackages; [
         cuda_cudart
@@ -113,37 +115,213 @@ let
               });
           } // lib.optionalAttrs isLinux {
           mlx = prev.mlx.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ lib.optionals cudaSupport [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
-            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
             postInstall = ''
               cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
             '';
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
         } // lib.optionalAttrs cudaSupport {
           "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
             autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-            buildInputs = old.buildInputs ++ [ cudaLibs ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = old.buildInputs ++ cudaLibs;
           });
           torch = prev.torch.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
             autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
+          torchaudio = prev.torchaudio.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = old.buildInputs ++ [ cudaPackages.cuda_cudart ];
+            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
+          });
+          torchvision = prev.torchvision.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
+          });
+
+          torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
+          });
+          # Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
+          vllm = prev.vllm.overrideAttrs (old:
+            let
+              cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
+                mkdir -p $out/include
+                ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
+              '';
+
+              cudaRoot = pkgs.symlinkJoin {
+                name = "cuda-merged-exo";
+                paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
+              };
+
+              cutlass = pkgs.fetchFromGitHub {
+                name = "cutlass-source";
+                owner = "NVIDIA";
+                repo = "cutlass";
+                tag = "v4.2.1";
+                hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
+              };
+              triton-kernels = pkgs.fetchFromGitHub {
+                owner = "triton-lang";
+                repo = "triton";
+                tag = "v3.6.0";
+                hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
+              };
+
+              cutlass-flashmla = pkgs.fetchFromGitHub {
+                owner = "NVIDIA";
+                repo = "cutlass";
+                rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
+                hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
+              };
+
+              flashmla = pkgs.stdenv.mkDerivation {
+                pname = "flashmla";
+                version = "1.0.0";
+
+                src = pkgs.fetchFromGitHub {
+                  name = "FlashMLA-source";
+                  owner = "vllm-project";
+                  repo = "FlashMLA";
+                  rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
+                  hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
+                };
+
+                dontConfigure = true;
+
+                buildPhase = ''
+                  rm -rf csrc/cutlass
+                  ln -sf ${cutlass-flashmla} csrc/cutlass
+                '';
+
+                installPhase = ''
+                  cp -rva . $out
+                '';
+              };
+              qutlass = pkgs.fetchFromGitHub {
+                name = "qutlass-source";
+                owner = "IST-DASLab";
+                repo = "qutlass";
+                rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
+                hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
+              };
+              vllm-flash-attn = pkgs.stdenv.mkDerivation {
+                pname = "vllm-flash-attn";
+                version = "2.7.2.post1";
+
+                src = pkgs.fetchFromGitHub {
+                  name = "flash-attention-source";
+                  owner = "vllm-project";
+                  repo = "flash-attention";
+                  rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
+                  hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
+                };
+
+                patches = [
+                  (pkgs.fetchpatch {
+                    url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
+                    hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
+                  })
+                  (pkgs.fetchpatch {
+                    url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
+                    hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
+                  })
+                ];
+
+                dontConfigure = true;
+
+                buildPhase = ''
+                  rm -rf csrc/cutlass
+                  ln -sf ${cutlass} csrc/cutlass
+                '';
+
+                installPhase = ''
+                  cp -rva . $out
+                '';
+              };
+            in
+            {
+              patches = (old.patches or [ ]) ++ [ ../nix/vllm-setuppy-cmake.patch ];
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+                pkgs.cmake
+                pkgs.ninja
+                pkgs.autoAddDriverRunpath
+              ] ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_nvcc
+              ];
+              # TODO: vllm rocm/cpu
+              VLLM_TARGET_DEVICE = "empty";
+              preConfigure = ''
+                export MAX_JOBS="$NIX_BUILD_CORES"
+              '';
+
+              # TODO: vllm non cuda13 support, more arch's, etc.
+            } // lib.optionalAttrs cudaSupport {
+              buildInputs = cudaLibs ++ [ cudaRoot ];
+
+              VLLM_CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
+              CUDA_HOME = "${cudaRoot}";
+              CUDAToolkit_ROOT = "${cudaRoot}";
+              CUDACXX = "${cudaRoot}/bin/nvcc";
+              VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
+              VLLM_TARGET_DEVICE = "cuda";
+              TORCH_CUDA_ARCH_LIST = "12.0;12.1";
+              TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
+              FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
+              QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
+              VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
+              CAFFE2_USE_CUDNN = "ON";
+              CAFFE2_USE_CUFILE = "ON";
+              CUTLASS_ENABLE_CUBLAS = "ON";
+              CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
+
+              cmakeFlags = [
+                (lib.cmakeBool "CMAKE_SKIP_INSTALL_RPATH" true)
+                (lib.cmakeBool "CMAKE_BUILD_WITH_INSTALL_RPATH" true)
+                (lib.cmakeFeature "CUDA_HOME" "${cudaRoot}")
+                (lib.cmakeFeature "CUDAToolkit_ROOT" "${cudaRoot}")
+                (lib.cmakeFeature "CMAKE_CUDA_COMPILER" "${cudaRoot}/bin/nvcc")
+                (lib.cmakeFeature "CMAKE_PREFIX_PATH" "${cudaRoot}")
+                (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}")
+                (lib.cmakeFeature "FLASH_MLA_SRC_DIR" "${lib.getDev flashmla}")
+                (lib.cmakeFeature "VLLM_FLASH_ATTN_SRC_DIR" "${lib.getDev vllm-flash-attn}")
+                (lib.cmakeFeature "QUTLASS_SRC_DIR" "${lib.getDev qutlass}")
+                (lib.cmakeFeature "TORCH_CUDA_ARCH_LIST" "12.0;12.1")
+                (lib.cmakeFeature "CUTLASS_NVCC_ARCHS_ENABLED" "${cudaPackages.flags.cmakeCudaArchitecturesString}")
+                (lib.cmakeFeature "CUDA_TOOLKIT_ROOT_DIR" "${cudaRoot}")
+                (lib.cmakeFeature "CAFFE2_USE_CUDNN" "ON")
+                (lib.cmakeFeature "CAFFE2_USE_CUFILE" "ON")
+                (lib.cmakeFeature "CUTLASS_ENABLE_CUBLAS" "ON")
+              ];
+            });
+
+        } // lib.optionalAttrs (cudaSupport && isx86_64) {
+          numba = prev.numba.overrideAttrs (old: {
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
+          });
         };
       pyprojectOverlay = workspace.mkPyprojectOverlay {
         sourcePreference = "wheel";
@@ -164,24 +342,28 @@ let
           buildSystemsOverlay
         ]
       );
-      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
-      mkApp = cmd: name: pkgs.writeShellApplication {
+      # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
+      venv = name: (pythonSet.mkVirtualEnv "${name}-venv" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
+      mkApp = text: name: pkgs.writeShellApplication {
         inherit name;
+        text = "exec " + lib.optionalString cudaSupport "nixglhost " + text;
         runtimeEnv = {
           EXO_DASHBOARD_DIR = self'.packages.dashboard;
           EXO_RESOURCES_DIR = inputs.self + /resources;
         };
         runtimeInputs = [
-          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
           (venv name)
+          pkgs.nix-gl-host
         ]
         ++ lib.optionals isDarwin [ pkgs.macmon ];
-        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
+        passthru = {
+          venv = venv name;
+          evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
+        };
       };
     in
     {
       inherit venv;
-      editablePythonSet = pythonSet.overrideScope editableOverlay;
       mkPythonScript = path: mkApp ''python ${path} "$@"'';
       mkExo = mkApp ''exo "$@"'';
     };
@@ -191,18 +373,18 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
 
       # Virtual environment with dev dependencies for testing
       testVenv = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+        exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).venv "exo-test";
 
       mkBenchScript = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "cpu" ];
+        exo = [ "mlx-cpu" "vllm-none" ];
         exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).mkPythonScript;
@@ -212,12 +394,12 @@ in
         runtimeInputs = [ pkgs.python313 ];
         text = ''exec python ${path} "$@"'';
       };
-
+      cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" "vllm-none" ]; }; };
+      cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cpu" "vllm-cuda13" ]; }; };
     in
     {
       packages = {
         exo = mkExo "exo";
-        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
         # for running tests in ci
         exo-test-env = testVenv;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
@@ -226,8 +408,8 @@ in
         # used by ./tests/run_exo_on.sh
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
       } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
-        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
+        exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
+        exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
       };
 
       checks = {

resources/inference_model_cards/2imi9--gpt-oss-20B-NVFP4A16-BF16.toml

@@ -0,0 +1,21 @@
+model_id = "2imi9/gpt-oss-20B-NVFP4A16-BF16"
+n_layers = 24
+hidden_size = 2880
+num_key_value_heads = 8
+supports_tensor = false
+tasks = ["TextGeneration"]
+family = "gpt-oss"
+quantization = "nvfp4"
+base_model = "GPT-OSS 20B"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "channel"
+context_length = 131072
+requires_vllm = true
+
+[storage_size]
+in_bytes = 41829514752
+
+[sampling_defaults]
+temperature = 1.0
+top_p = 1.0
+top_k = 0

resources/inference_model_cards/nvidia--Qwen3-30B-A3B-NVFP4.toml

@@ -0,0 +1,27 @@
+model_id = "nvidia/Qwen3-30B-A3B-NVFP4"
+n_layers = 48
+hidden_size = 2048
+num_key_value_heads = 4
+supports_tensor = false
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "nvfp4"
+base_model = "Qwen3 30B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+context_length = 32768
+requires_vllm = true
+
+[storage_size]
+in_bytes = 18087458688
+
+[sampling_defaults]
+temperature = 0.6
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0

resources/inference_model_cards/openai--gpt-oss-120b.toml

@@ -0,0 +1,20 @@
+model_id = "openai/gpt-oss-120b"
+n_layers = 36
+hidden_size = 2880
+num_key_value_heads = 8
+supports_tensor = false
+tasks = ["TextGeneration"]
+family = "gpt-oss"
+quantization = "mxfp4"
+base_model = "GPT-OSS 120B"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "channel"
+context_length = 131072
+
+[storage_size]
+in_bytes = 65248815744
+
+[sampling_defaults]
+temperature = 1.0
+top_p = 1.0
+top_k = 0

resources/inference_model_cards/sakamakismile--Qwen3.6-27B-NVFP4.toml

@@ -0,0 +1,32 @@
+model_id = "sakamakismile/Qwen3.6-27B-NVFP4"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = false
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "nvfp4"
+base_model = "Qwen3.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+reasoning_dialect = "post_last_user"
+context_length = 262144
+requires_vllm = true
+
+[storage_size]
+in_bytes = 16703361232
+
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5

rust/exo_pyo3_bindings/Cargo.toml

@@ -46,12 +46,9 @@ pyo3-async-runtimes = { version = "0.27.0", features = [
 ] }
 pyo3-log = "0.13.2"
 
-pidfile-rs = "0.3"
-
 # macro dependencies
 extend = { workspace = true }
 delegate = { workspace = true }
-thiserror = "2.0"
 
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -2,8 +2,6 @@
 # ruff: noqa: E501, F401
 
 import builtins
-import os
-import pathlib
 import typing
 
 @typing.final
@@ -71,48 +69,6 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-@typing.final
-class Pidfile:
-    r"""
-    A PID file protected with a lock.
-    
-    An instance of `Pidfile` can be used to manage a PID file: create it,
-    lock it, detect already running daemons. It is backed by [`pidfile`][]
-    functions of `libbsd`/`libutil` which use `flopen` to lock the PID
-    file.
-    
-    When a PID file is created, the process ID of the current process is
-    *not* written there, making it possible to lock the PID file before
-    forking and only write the ID of the forked process when it is ready.
-    
-    The PID file is deleted automatically when the `Pidfile` comes out of
-    the scope. To close the PID file without deleting it, for example, in
-    the parent process of a forked daemon, call `close()`.
-    
-    [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
-    [`pidfile`]: https://linux.die.net/man/3/pidfile
-    [`daemon`(3)]: https://linux.die.net/man/3/daemon
-    """
-    def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
-        r"""
-        Creates a new PID file and locks it.
-        
-        If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
-        a PID of the already running process, or `None` if no PID has been written to
-        the PID file yet.
-        """
-    def write(self) -> None:
-        r"""
-        Writes the current process ID to the PID file.
-        
-        The file is truncated before writing.
-        """
-
-@typing.final
-class PidfileError(builtins.Exception):
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
 class PyFromSwarm:
     @typing.final
     class Connection(PyFromSwarm):

rust/exo_pyo3_bindings/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.2.2"
+version = "0.2.1"
 description = "Add your description here"
 readme = "README.md"
 authors = [

rust/exo_pyo3_bindings/src/lib.rs

@@ -7,11 +7,9 @@
 mod allow_threading;
 mod ident;
 mod networking;
-mod pidfile;
 
 use crate::ident::PyKeypair;
 use crate::networking::networking_submodule;
-use crate::pidfile::pidfile_submodule;
 use pyo3::prelude::PyModule;
 use pyo3::types::PyModuleMethods;
 use pyo3::{Bound, PyResult, pyclass, pymodule};
@@ -166,7 +164,6 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     //       too many importing issues...
     m.add_class::<PyKeypair>()?;
     networking_submodule(m)?;
-    pidfile_submodule(m)?;
 
     // top-level constructs
     // TODO: ...

rust/exo_pyo3_bindings/src/pidfile.rs

@@ -1,87 +0,0 @@
-use pidfile_rs::{Pidfile, PidfileError};
-use pyo3::exceptions::PyException;
-use pyo3::prelude::{PyModule, PyModuleMethods};
-use pyo3::{Bound, PyErr, PyResult, Python, pyclass, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
-use std::fs::Permissions;
-use std::os::unix::prelude::PermissionsExt;
-use std::path::PathBuf;
-
-#[gen_stub_pyclass]
-#[pyclass(frozen, extends=PyException, name="PidfileError")]
-pub struct PyPidfileError(PidfileError);
-
-impl PyPidfileError {
-    // TODO: I actually like this pattern a LOT more but how to abstract??
-    fn into_pyerr(self, py: Python) -> PyErr {
-        match Bound::new(py, self) {
-            Ok(err) => PyErr::from_value(err.into_any()),
-            Err(err) => err,
-        }
-    }
-}
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl PyPidfileError {
-    fn __repr__(&self) -> String {
-        format!("PidfileError(\"{}\")", self.0)
-    }
-
-    fn __str__(&self) -> String {
-        self.0.to_string()
-    }
-}
-
-/// A PID file protected with a lock.
-///
-/// An instance of `Pidfile` can be used to manage a PID file: create it,
-/// lock it, detect already running daemons. It is backed by [`pidfile`][]
-/// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
-/// file.
-///
-/// When a PID file is created, the process ID of the current process is
-/// *not* written there, making it possible to lock the PID file before
-/// forking and only write the ID of the forked process when it is ready.
-///
-/// The PID file is deleted automatically when the `Pidfile` comes out of
-/// the scope. To close the PID file without deleting it, for example, in
-/// the parent process of a forked daemon, call `close()`.
-///
-/// [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
-/// [`pidfile`]: https://linux.die.net/man/3/pidfile
-/// [`daemon`(3)]: https://linux.die.net/man/3/daemon
-#[gen_stub_pyclass]
-#[pyclass(name = "Pidfile")]
-pub struct PyPidfile(Pidfile);
-
-#[gen_stub_pymethods]
-#[pymethods]
-impl PyPidfile {
-    /// Creates a new PID file and locks it.
-    ///
-    /// If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
-    /// a PID of the already running process, or `None` if no PID has been written to
-    /// the PID file yet.
-    #[new]
-    fn py_new(py: Python, path: PathBuf, mode: u32) -> PyResult<Self> {
-        Ok(Self(
-            Pidfile::new(&path, Permissions::from_mode(mode))
-                .map_err(|e| PyPidfileError(e).into_pyerr(py))?,
-        ))
-    }
-
-    /// Writes the current process ID to the PID file.
-    ///
-    /// The file is truncated before writing.
-    fn write<'py>(&mut self, py: Python<'py>) -> PyResult<()> {
-        self.0.write().map_err(|e| PyPidfileError(e).into_pyerr(py))
-    }
-}
-
-pub fn pidfile_submodule(m: &Bound<PyModule>) -> PyResult<()> {
-    m.add_class::<PyPidfileError>()?;
-    m.add_class::<PyPidfile>()?;
-
-    Ok(())
-}

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,12 +1,10 @@
 import asyncio
 
 import pytest
-from _pytest.capture import CaptureFixture
 from exo_pyo3_bindings import (
     Keypair,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
-    Pidfile,
     PyFromSwarm,
 )
 
@@ -28,13 +26,6 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-def test_pidfile(capsys: CaptureFixture[str]):
-    with capsys.disabled():
-        print("\nbefore python")
-        scoped_lock_file()
-        print("after python")
-
-
 async def _await_recv(h: NetworkingHandle):
     while True:
         event = await h.recv()
@@ -43,7 +34,3 @@ async def _await_recv(h: NetworkingHandle):
                 print(f"PYTHON: connection update: {c}")
             case PyFromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
-
-
-def scoped_lock_file():
-    a = Pidfile("/tmp/lock.pid", 0o0600)

scripts/check_serve_prefill.py

@@ -0,0 +1,222 @@
+#!/usr/bin/env python
+"""Standalone smoke test for VllmEngine.serve_prefill.
+
+Loads a real vLLM engine, runs serve_prefill against an in-memory buffer
+twice in a row with the same prompt, and verifies both runs produce a
+well-formed wire stream (header -> KV chunks -> Done).
+
+The second run is the regression guard: with vLLM APC enabled this would
+trip the chunked-prefill + APC + custom kv-connector CUDA assert
+(`vectorized_gather_kernel: ind >= ind_dim_size`) and the server would
+close the socket before the Done frame.
+
+Usage on the Spark (gx10-de89):
+
+    cd /home/larry/exo
+    /nix/store/2b82iz9ac0pxqafrgxmgdkq8sr2hwlx6-exo-cuda-13-venv/bin/python \\
+        scripts/check_serve_prefill.py Qwen/Qwen3-0.6B
+
+Exits 0 on success, non-zero with a diagnostic on failure.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import os
+import sys
+import traceback
+from pathlib import Path
+from typing import cast
+
+
+def _ensure_repo_on_path() -> None:
+    repo = Path(__file__).resolve().parent.parent
+    src = repo / "src"
+    if str(src) not in sys.path:
+        sys.path.insert(0, str(src))
+
+
+_ensure_repo_on_path()
+
+from exo.shared.types.common import ModelId  # noqa: E402
+from exo.worker.disaggregated.protocol import (  # noqa: E402
+    ArraysState,
+    Done,
+    ErrorMessage,
+    KVChunk,
+    read_header,
+    read_message,
+)
+from exo.worker.disaggregated.server import PrefillRequest  # noqa: E402
+
+
+def _make_token_ids(n: int) -> list[int]:
+    return [(i * 1009 + 17) % 30000 + 100 for i in range(n)]
+
+
+def _decode(
+    payload: bytes,
+) -> tuple[list[KVChunk], list[ArraysState], Done | None, ErrorMessage | None]:
+    buf = io.BytesIO(payload)
+    _ = read_header(buf)
+    chunks: list[KVChunk] = []
+    arrays: list[ArraysState] = []
+    done: Done | None = None
+    error: ErrorMessage | None = None
+    while True:
+        msg = read_message(buf)
+        if msg is None:
+            break
+        if isinstance(msg, KVChunk):
+            chunks.append(msg)
+        elif isinstance(msg, ArraysState):
+            arrays.append(msg)
+        elif isinstance(msg, Done):
+            done = msg
+            break
+        elif isinstance(msg, ErrorMessage):
+            error = msg
+            break
+    return chunks, arrays, done, error
+
+
+def _build_engine(model_id: ModelId) -> object:
+    from exo.worker.engines.vllm.engine import VllmEngine
+    from exo.worker.engines.vllm.generator import VllmBatchEngine, load_vllm_engine
+    from exo.worker.engines.vllm.kv_connector import (
+        ExoKVProducerConnector,
+        _patch_gdn_capture,
+        _patch_vllm_for_connector,
+    )
+
+    _patch_vllm_for_connector(ExoKVProducerConnector)
+    _patch_gdn_capture()
+
+    llm_engine, tool_parser = load_vllm_engine(
+        model_id=model_id,
+        trust_remote_code=False,
+        n_layers=1,
+        kv_connector_cls=ExoKVProducerConnector,
+    )
+    gen = VllmBatchEngine(engine=llm_engine, model_id=model_id)
+
+    class _S:
+        def send(self, _: object) -> None: ...
+
+    class _R:
+        def collect(self) -> list[object]:
+            return []
+
+    return VllmEngine(
+        tool_parser=tool_parser,
+        model_id=model_id,
+        cancel_receiver=cast("object", _R()),  # pyright: ignore[reportArgumentType]
+        event_sender=cast("object", _S()),  # pyright: ignore[reportArgumentType]
+        _gen=gen,
+        max_concurrent_requests=1,
+    )
+
+
+def _run_one(engine: object, n_tokens: int, label: str) -> int:
+    request = PrefillRequest(
+        request_id=f"check-{label}-{os.getpid()}",
+        model_id="ignored",
+        token_ids=_make_token_ids(n_tokens),
+        start_pos=0,
+        use_prefix_cache=True,
+    )
+    buf = io.BytesIO()
+    engine.serve_prefill(request, buf)  # pyright: ignore[reportAttributeAccessIssue]
+    payload = buf.getvalue()
+    if not payload:
+        raise AssertionError(f"{label}: server wrote nothing")
+
+    chunks, arrays, done, error = _decode(payload)
+    if error is not None:
+        raise AssertionError(
+            f"{label}: server returned ErrorMessage [{error.code}]: {error.message}"
+        )
+    if done is None:
+        raise AssertionError(
+            f"{label}: stream did not end with Done "
+            f"({len(chunks)} kv chunks, {len(arrays)} arrays)"
+        )
+    if done.total_tokens <= 0:
+        raise AssertionError(f"{label}: Done reported {done.total_tokens} tokens")
+    if not chunks:
+        raise AssertionError(f"{label}: no KV chunks shipped")
+
+    expected = max(0, n_tokens - 2)
+    if done.total_tokens < expected - 64:
+        raise AssertionError(
+            f"{label}: got {done.total_tokens} tokens, expected ~{expected}"
+        )
+    print(
+        f"  [{label}] OK: tokens={done.total_tokens} "
+        f"kv_chunks={len(chunks)} arrays={len(arrays)}"
+    )
+    return done.total_tokens
+
+
+def main(argv: list[str]) -> int:
+    if len(argv) < 2:
+        print(__doc__)
+        return 2
+    model_id = ModelId(argv[1])
+
+    from exo.download.download_utils import build_model_path
+
+    model_path = build_model_path(model_id)
+    if not model_path.exists():
+        print(f"FAIL: model {model_id} not found at {model_path}")
+        return 1
+    print(f"Loading vLLM engine for {model_id} ({model_path}) ...")
+
+    engine = _build_engine(model_id)
+    failures: list[str] = []
+    try:
+        try:
+            t1 = _run_one(engine, n_tokens=512, label="run1-fresh")
+        except AssertionError as e:
+            failures.append(f"run1: {e}")
+            t1 = 0
+        try:
+            t2 = _run_one(engine, n_tokens=512, label="run2-same-prompt")
+        except AssertionError as e:
+            failures.append(f"run2: {e}")
+            t2 = 0
+        if t1 and t2 and t1 != t2:
+            failures.append(
+                f"run1 returned {t1} tokens but run2 returned {t2} (should match)"
+            )
+        try:
+            ta = _run_one(engine, n_tokens=256, label="run3-shorter")
+            tb = _run_one(engine, n_tokens=768, label="run4-longer")
+            if ta and tb and tb <= ta:
+                failures.append(
+                    f"longer prompt should produce more tokens: 256->{ta} 768->{tb}"
+                )
+        except AssertionError as e:
+            failures.append(f"length-variation: {e}")
+    finally:
+        with contextlib.suppress(Exception):
+            engine.close()  # pyright: ignore[reportAttributeAccessIssue]
+
+    if failures:
+        print()
+        print("FAIL")
+        for f in failures:
+            print(f"  - {f}")
+        return 1
+    print()
+    print("PASS")
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main(sys.argv))
+    except Exception:
+        traceback.print_exc()
+        sys.exit(1)

setup_linklocal_dgx.sh

@@ -0,0 +1,124 @@
+#!/usr/bin/env bash
+set -Eeuo pipefail
+
+SELF_IP="169.254.100.1"
+PEER_IP="169.254.100.2"
+PREFIX="16"
+IFACE="enP7s7"
+USE_NM="auto"
+DRY_RUN=0
+
+usage() {
+  cat <<EOF
+Usage: sudo $(basename "$0") [options]
+
+Configure a Linux Ethernet interface with a static IPv4 for a host-to-host
+link to a Mac peer.
+
+Defaults: this host = ${SELF_IP}/${PREFIX}, peer = ${PEER_IP}, iface = ${IFACE}.
+
+Options:
+  --iface IFACE     Default: ${IFACE}
+  --self-ip IP      Default: ${SELF_IP}
+  --peer-ip IP      For verification ping. Default: ${PEER_IP}
+  --prefix N        Default: ${PREFIX}
+  --no-nm           Use 'ip addr' directly (transient, no NetworkManager).
+  --dry-run         Print actions without applying.
+  -h, --help        Show this help.
+EOF
+}
+
+while (($#)); do
+  case "$1" in
+  --iface)
+    shift
+    IFACE="${1:?}"
+    ;;
+  --self-ip)
+    shift
+    SELF_IP="${1:?}"
+    ;;
+  --peer-ip)
+    shift
+    PEER_IP="${1:?}"
+    ;;
+  --prefix)
+    shift
+    PREFIX="${1:?}"
+    ;;
+  --no-nm) USE_NM=no ;;
+  --dry-run) DRY_RUN=1 ;;
+  -h | --help)
+    usage
+    exit 0
+    ;;
+  *)
+    echo "Unknown arg: $1" >&2
+    usage >&2
+    exit 1
+    ;;
+  esac
+  shift
+done
+
+[[ $EUID -eq 0 ]] || {
+  echo "Run as root." >&2
+  exit 1
+}
+
+run() {
+  printf '+'
+  printf ' %q' "$@"
+  printf '\n'
+  ((DRY_RUN)) || "$@"
+}
+
+ip link show "$IFACE" >/dev/null 2>&1 || {
+  echo "Interface $IFACE does not exist." >&2
+  exit 1
+}
+
+if [[ $USE_NM == "auto" ]]; then
+  if command -v nmcli >/dev/null 2>&1 && systemctl is-active --quiet NetworkManager 2>/dev/null; then
+    USE_NM=yes
+  else
+    USE_NM=no
+  fi
+fi
+
+if [[ $USE_NM == "yes" ]]; then
+  CONN="$(nmcli -g GENERAL.CONNECTION device show "$IFACE" 2>/dev/null | head -n1 || true)"
+  if [[ -z $CONN || $CONN == "--" ]]; then
+    CONN="static-${IFACE}"
+    run nmcli connection add type ethernet ifname "$IFACE" con-name "$CONN"
+  fi
+  run nmcli connection modify "$CONN" \
+    connection.interface-name "$IFACE" \
+    connection.autoconnect yes \
+    connection.autoconnect-priority 100 \
+    ipv4.method manual \
+    ipv4.addresses "${SELF_IP}/${PREFIX}" \
+    ipv4.gateway "" \
+    ipv4.dns "" \
+    ipv4.never-default yes \
+    ipv6.method link-local \
+    ipv6.addr-gen-mode stable-privacy
+  run nmcli connection up "$CONN"
+else
+  run ip link set "$IFACE" up
+  run ip addr flush dev "$IFACE"
+  run ip addr add "${SELF_IP}/${PREFIX}" dev "$IFACE"
+fi
+
+if ((!DRY_RUN)); then
+  printf '\n'
+  ip -br addr show "$IFACE"
+  printf '\n'
+  if ping -c2 -W2 "$PEER_IP" >/dev/null 2>&1; then
+    echo "OK: $PEER_IP reachable on $IFACE."
+  else
+    echo "WARN: $PEER_IP not reachable yet."
+    echo "      Verify the peer is configured (run setup_linklocal_mac.sh on the Mac)."
+    echo "      ip neigh show dev $IFACE  # check for the peer MAC"
+  fi
+fi

setup_linklocal_mac.sh

@@ -0,0 +1,170 @@
+#!/usr/bin/env bash
+set -Eeuo pipefail
+
+SELF_IP="169.254.100.2"
+PEER_IP="169.254.100.1"
+NETMASK="255.255.0.0"
+IFACE=""
+DRY_RUN=0
+
+usage() {
+  cat <<EOF
+Usage: sudo $(basename "$0") [options]
+
+Configure a Mac Ethernet interface with a static IPv4 for a host-to-host link
+to the DGX/GX10 peer.
+
+Defaults: this Mac = ${SELF_IP}, peer = ${PEER_IP}, mask = ${NETMASK}.
+
+Options:
+  --iface IFACE     Interface (e.g. en12). Default: auto-detect.
+  --self-ip IP      This Mac's address. Default: ${SELF_IP}.
+  --peer-ip IP      Peer for verification ping. Default: ${PEER_IP}.
+  --netmask MASK    Default: ${NETMASK}.
+  --dry-run         Print actions without applying.
+  -h, --help        Show this help.
+EOF
+}
+
+while (($#)); do
+  case "$1" in
+  --iface)
+    shift
+    IFACE="${1:?}"
+    ;;
+  --self-ip)
+    shift
+    SELF_IP="${1:?}"
+    ;;
+  --peer-ip)
+    shift
+    PEER_IP="${1:?}"
+    ;;
+  --netmask)
+    shift
+    NETMASK="${1:?}"
+    ;;
+  --dry-run) DRY_RUN=1 ;;
+  -h | --help)
+    usage
+    exit 0
+    ;;
+  *)
+    echo "Unknown arg: $1" >&2
+    usage >&2
+    exit 1
+    ;;
+  esac
+  shift
+done
+
+[[ $EUID -eq 0 ]] || {
+  echo "Run with sudo." >&2
+  exit 1
+}
+
+run() {
+  printf '+'
+  printf ' %q' "$@"
+  printf '\n'
+  ((DRY_RUN)) || "$@"
+}
+
+target_subnet_prefix() {
+  local ip="$1"
+  printf '%s.' "${ip%.*}"
+}
+
+iface_score() {
+  local iface="$1" info subnet
+  info="$(ifconfig "$iface" 2>/dev/null || true)"
+  [[ -n $info ]] || {
+    echo 0
+    return
+  }
+  grep -q 'status: active' <<<"$info" || {
+    echo 0
+    return
+  }
+  subnet="$(target_subnet_prefix "$SELF_IP")"
+  if grep -qE "inet ${subnet//./\\.}" <<<"$info"; then
+    echo 100
+    return
+  fi
+  if grep -qE 'inet 169\.254\.' <<<"$info"; then
+    echo 80
+    return
+  fi
+  if ! grep -qE '^[[:space:]]*inet ' <<<"$info"; then
+    echo 60
+    return
+  fi
+  echo 10
+}
+
+detect_iface() {
+  local best="" best_score=0 iface score
+  for iface in $(ifconfig -l); do
+    [[ $iface =~ ^en[0-9]+$ ]] || continue
+    score="$(iface_score "$iface")"
+    if ((score > best_score)); then
+      best="$iface"
+      best_score="$score"
+    fi
+  done
+  ((best_score >= 60)) || return 1
+  printf '%s\n' "$best"
+}
+
+iface_to_service() {
+  local iface="$1" line port=""
+  while IFS= read -r line; do
+    if [[ $line == "Hardware Port: "* ]]; then
+      port="${line#Hardware Port: }"
+    elif [[ $line == "Device: $iface" ]]; then
+      printf '%s\n' "$port"
+      return 0
+    fi
+  done < <(networksetup -listallhardwareports)
+  return 1
+}
+
+if [[ -z $IFACE ]]; then
+  IFACE="$(detect_iface || true)"
+  [[ -n $IFACE ]] || {
+    echo "Could not auto-detect a wired interface. Pass --iface enX." >&2
+    echo "Active interfaces:" >&2
+    ifconfig -l | tr ' ' '\n' | grep -E '^en[0-9]+$' | while read -r i; do
+      printf '  %-6s %s\n' "$i" "$(ifconfig "$i" | grep -E 'status:|inet ' | tr '\n' ' ')" >&2
+    done
+    exit 1
+  }
+  echo "Auto-detected interface: $IFACE"
+fi
+
+ifconfig "$IFACE" >/dev/null 2>&1 || {
+  echo "Interface $IFACE does not exist." >&2
+  exit 1
+}
+
+SERVICE="$(iface_to_service "$IFACE" || true)"
+[[ -n $SERVICE ]] || {
+  echo "No network service maps to $IFACE. Check System Settings -> Network." >&2
+  exit 1
+}
+echo "Network service: $SERVICE"
+
+run networksetup -setmanual "$SERVICE" "$SELF_IP" "$NETMASK" ""
+
+if ((!DRY_RUN)); then
+  printf '\n'
+  ifconfig "$IFACE" | grep -E 'inet |status:'
+  printf '\n'
+  if ping -c2 -t3 "$PEER_IP" >/dev/null 2>&1; then
+    echo "OK: $PEER_IP reachable on $IFACE."
+  else
+    echo "WARN: $PEER_IP not reachable yet."
+    echo "      Verify the peer is configured (run setup_linklocal_dgx.sh on the GX10)."
+    echo "      arp -an -i $IFACE  # check for the peer MAC"
+  fi
+fi

src/exo/api/main.py

@@ -20,7 +20,6 @@ from fastapi.staticfiles import StaticFiles
 from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType]
 from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
-from hypercorn.utils import LifespanTimeoutError
 from loguru import logger
 
 from exo.api.adapters.chat_completions import (
@@ -134,10 +133,12 @@ from exo.shared.constants import (
 )
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
-from exo.shared.models import model_cards
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
+    add_to_card_cache,
+    get_card,
+    get_model_cards,
 )
 from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.chunks import (
@@ -480,7 +481,6 @@ class API:
                 topology=self.state.topology,
                 current_instances=self.state.instances,
                 download_status=self.state.downloads,
-                node_rdma_ctl=self.state.node_rdma_ctl,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -526,8 +526,8 @@ class API:
                         )
                     ]
                 )
-        # TODO: PDD
-        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
+        if any(self.state.node_vllm.values()):
+            instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
             try:
@@ -544,7 +544,6 @@ class API:
                     current_instances=self.state.instances,
                     required_nodes=required_nodes,
                     download_status=self.state.downloads,
-                    node_rdma_ctl=self.state.node_rdma_ctl,
                 )
             except ValueError as exc:
                 if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -641,7 +640,10 @@ class API:
         )
 
     async def get_feature_flags(self) -> dict[str, bool]:
-        return {"disaggregation": ENABLE_DISAGGREGATION}
+        return {
+            "disaggregation": ENABLE_DISAGGREGATION,
+            "vllm_available": any(self.state.node_vllm.values()),
+        }
 
     async def list_instance_links(self) -> list[InstanceLink]:
         if not ENABLE_DISAGGREGATION:
@@ -1634,16 +1636,17 @@ class API:
     async def ollama_tags(self) -> OllamaTagsResponse:
         """Returns list of models in Ollama tags format. We return the downloaded ones only."""
 
-        downloaded_model_ids: set[ModelId] = set()
+        def none_if_empty(value: str) -> str | None:
+            return value or None
+
+        downloaded_model_ids: set[str] = set()
         for node_downloads in self.state.downloads.values():
             for dl in node_downloads:
                 if isinstance(dl, DownloadCompleted):
                     downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
 
         cards = [
-            c
-            for c in await model_cards.card_cache.list_all()
-            if c.model_id in downloaded_model_ids
+            c for c in await get_model_cards() if c.model_id in downloaded_model_ids
         ]
 
         now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -1656,8 +1659,8 @@ class API:
                     size=card.storage_size.in_bytes,
                     digest="sha256:000000000000",
                     details=OllamaModelDetails(
-                        family=card.family or None,
-                        quantization_level=card.quantization or None,
+                        family=none_if_empty(card.family),
+                        quantization_level=none_if_empty(card.quantization),
                     ),
                 )
                 for card in cards
@@ -1720,7 +1723,7 @@ class API:
 
     async def get_models(self, status: str | None = Query(default=None)) -> ModelList:
         """Returns list of available models, optionally filtered by being downloaded."""
-        cards = await model_cards.card_cache.list_all()
+        cards = await get_model_cards()
 
         if status == "downloaded":
             downloaded_model_ids: set[str] = set()
@@ -1748,6 +1751,7 @@ class API:
                     capabilities=card.capabilities,
                     reasoning_dialect=card.reasoning_dialect,
                     context_length=card.context_length,
+                    requires_vllm=card.requires_vllm,
                 )
                 for card in cards
             ]
@@ -1771,7 +1775,7 @@ class API:
 
         # Immediately update the local cache so the subsequent GET /models
         # returns the new model without waiting for the event round-trip.
-        model_cards.card_cache.cc[card.model_id] = card
+        add_to_card_cache(card)
 
         return ModelListModel(
             id=card.model_id,
@@ -1787,7 +1791,7 @@ class API:
 
     async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
         """Delete a user-added custom model card and sync deletion across the cluster."""
-        card = model_cards.card_cache.get(model_id)
+        card = get_card(model_id)
         if card is None or not card.is_custom:
             raise HTTPException(status_code=404, detail="Custom model card not found")
 
@@ -1857,21 +1861,12 @@ class API:
                     await anyio.sleep_forever()
                 finally:
                     with anyio.CancelScope(shield=True):
-                        # IMPORTANT: when new queues are added, update this (for proper shutdown semantics)
-                        self._shutdown_queues(self._text_generation_queues)
-                        self._shutdown_queues(self._image_generation_queues)
-
                         shutdown_ev.set()
         finally:
             self._event_log.close()
             self.command_sender.close()
             self.event_receiver.close()
 
-    @staticmethod
-    def _shutdown_queues[K, V](queues: dict[K, Sender[V]]):
-        for v in queues.values():
-            v.close()
-
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
         cfg.bind = [f"0.0.0.0:{self.port}"]
@@ -1879,23 +1874,12 @@ class API:
         cfg.accesslog = None
         cfg.errorlog = "-"
         cfg.logger_class = InterceptLogger
-
-        # prevents hangs when mid-request and connection refuses to close
-        cfg.graceful_timeout = 2  # seconds
-        cfg.shutdown_timeout = 3  # seconds
-
         with anyio.CancelScope(shield=True):
-            try:
-                await serve(
-                    cast(ASGIFramework, self.app),
-                    cfg,
-                    shutdown_trigger=ev.wait,
-                )
-            except LifespanTimeoutError as e:
-                logger.warning(
-                    "Graceful server shutdown timed out, some connections forcebly closed"
-                )
-                logger.opt(exception=e).debug("")
+            await serve(
+                cast(ASGIFramework, self.app),
+                cfg,
+                shutdown_trigger=ev.wait,
+            )
 
     async def _apply_state(self):
         with self.event_receiver as events:

src/exo/api/types/api.py

@@ -49,6 +49,7 @@ class ModelListModel(BaseModel):
     base_model: str = Field(default="")
     capabilities: list[str] = Field(default_factory=list)
     reasoning_dialect: ReasoningDialect = "none"
+    requires_vllm: bool = Field(default=False)
 
 
 class ModelList(BaseModel):

src/exo/download/coordinator.py

@@ -16,8 +16,7 @@ from exo.download.download_utils import (
 )
 from exo.download.shard_downloader import ShardDownloader
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
-from exo.shared.models import model_cards
-from exo.shared.models.model_cards import ModelId
+from exo.shared.models.model_cards import ModelId, get_model_cards
 from exo.shared.types.commands import (
     CancelDownload,
     DeleteDownload,
@@ -423,7 +422,7 @@ class DownloadCoordinator:
                     )
                 # Scan read-only directories for pre-downloaded models
                 if EXO_MODELS_READ_ONLY_DIRS:
-                    for card in await model_cards.card_cache.list_all():
+                    for card in await get_model_cards():
                         mid = card.model_id
                         if mid in self.active_downloads:
                             continue

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/impl_shard_downloader.py

@@ -11,11 +11,11 @@ from exo.download.download_utils import (
     download_shard,
 )
 from exo.download.shard_downloader import ShardDownloader
-from exo.shared.models import model_cards
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
     ModelTask,
+    get_model_cards,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.shards import (
@@ -258,7 +258,7 @@ class ResumableShardDownloader(ShardDownloader):
 
         tasks = [
             create_task(download_with_semaphore(model_card))
-            for model_card in await model_cards.card_cache.list_all()
+            for model_card in await get_model_cards()
         ]
 
         for task in asyncio.as_completed(tasks):

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

@@ -3,7 +3,6 @@ import multiprocessing as mp
 import os
 import resource
 import signal
-import sys
 from dataclasses import dataclass, field
 from typing import Self
 
@@ -23,8 +22,6 @@ from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
-from exo.utils.daemon import detach_stdio_to_devnull
-from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
 from exo.utils.pydantic_ext import FrozenModel
 from exo.utils.task_group import TaskGroup
 from exo.worker.main import Worker
@@ -267,26 +264,14 @@ class Node:
 
 
 def main():
-    # Exit early if no PID file (not compatible with double-for daemonization yet)
-    try:
-        pidfile = acquire_exo_pidfile()
-    except PidfileLockError as exception:
-        print(exception, file=sys.stderr)
-        raise SystemExit(1) from exception
-
     args = Args.parse()
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
     target = min(max(soft, 65535), hard)
     resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
 
     mp.set_start_method("spawn", force=True)
-
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
-    if args.no_stdio:
-        detach_stdio_to_devnull()
-        logger.info("Detached stdio to /dev/null")
-
     logger.info(f"{'=' * 40}")
     logger.info(f"Starting EXO | pid={os.getpid()}")
     logger.info(f"{'=' * 40}")
@@ -321,7 +306,6 @@ def main():
     finally:
         logger.info("EXO Shutdown complete")
         logger_cleanup()
-        del pidfile
 
 
 class Args(FrozenModel):
@@ -335,7 +319,6 @@ class Args(FrozenModel):
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    no_stdio: bool = False
     bootstrap_peers: list[str] = []
     libp2p_port: int
 
@@ -395,11 +378,6 @@ class Args(FrozenModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
-        parser.add_argument(
-            "--no-stdio",
-            action="store_true",
-            help="Detach stdin/stdout/stderr to /dev/null after logging is configured",
-        )
         parser.add_argument(
             "--bootstrap-peers",
             type=lambda s: [p for p in s.split(",") if p],

src/exo/master/main.py

@@ -180,10 +180,31 @@ class Master:
                             for link in self.state.instance_links.values():
                                 prefill_only.difference_update(link.decode_instances)
 
+                            # If the user typed a prefill-only model id (e.g.
+                            # the vLLM-side producer of a P/D pair), the
+                            # candidate decode side is whatever it's linked
+                            # to. Expand the requested model id to also
+                            # include those linked decode instances.
+                            requested_model = command.task_params.model
+                            linked_decode_ids: set[InstanceId] = set()
+                            for link in self.state.instance_links.values():
+                                if any(
+                                    self.state.instances.get(pid) is not None
+                                    and self.state.instances[
+                                        pid
+                                    ].shard_assignments.model_id
+                                    == requested_model
+                                    for pid in link.prefill_instances
+                                ):
+                                    linked_decode_ids.update(link.decode_instances)
+
                             for instance in self.state.instances.values():
-                                if (
+                                model_match = (
                                     instance.shard_assignments.model_id
-                                    == command.task_params.model
+                                    == requested_model
+                                ) or (instance.instance_id in linked_decode_ids)
+                                if (
+                                    model_match
                                     and instance.instance_id not in prefill_only
                                 ):
                                     in_flight = {TaskStatus.Pending, TaskStatus.Running}
@@ -365,7 +386,6 @@ class Master:
                                 self.state.node_memory,
                                 self.state.node_network,
                                 download_status=self.state.downloads,
-                                node_rdma_ctl=self.state.node_rdma_ctl,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks

src/exo/master/placement.py

@@ -28,7 +28,7 @@ from exo.shared.types.events import (
     TaskStatusUpdated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo, NodeRdmaCtlStatus
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
@@ -43,6 +43,7 @@ from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.ports import random_ephemeral_port
@@ -105,7 +106,6 @@ def place_instance(
     node_network: Mapping[NodeId, NodeNetworkInfo],
     required_nodes: set[NodeId] | None = None,
     download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
-    node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -167,18 +167,8 @@ def place_instance(
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
-    rdma_ctl_status = node_rdma_ctl or {}
-
-    def _all_rdma_ctl_enabled(cycle: Cycle) -> bool:
-        return all(
-            ((status := rdma_ctl_status.get(node_id)) is not None and status.enabled)
-            for node_id in cycle
-        )
-
     smallest_rdma_cycles = [
-        cycle
-        for cycle in smallest_cycles
-        if topology.is_rdma_cycle(cycle) and _all_rdma_ctl_enabled(cycle)
+        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
     if command.instance_meta == InstanceMeta.MlxJaccl:
@@ -213,7 +203,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
+    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
         command = command.model_copy(
             update={
                 "instance_meta": InstanceMeta.MlxRing,
@@ -277,6 +267,11 @@ def place_instance(
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
             )
+        case InstanceMeta.Vllm:
+            target_instances[instance_id] = VllmInstance(
+                instance_id=instance_id,
+                shard_assignments=shard_assignments,
+            )
 
     return target_instances
 

src/exo/master/placement_utils.py

@@ -375,7 +375,13 @@ def find_ip_prioritised(
             "maybe_ethernet": 3,
             "thunderbolt": 4,
         }
-    return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
+
+    def _key(ip: str) -> tuple[int, int]:
+        link_local = 0 if ip.startswith("169.254.") else 1
+        type_pri = priority.get(ip_to_type.get(ip, "unknown"), 2)
+        return (link_local, type_pri)
+
+    return min(ips, key=_key)
 
 
 def get_mlx_ring_hosts_by_node(

src/exo/master/tests/test_placement.py

@@ -21,11 +21,7 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import (
-    NetworkInterfaceInfo,
-    NodeNetworkInfo,
-    NodeRdmaCtlStatus,
-)
+from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import (
     InputMessage,
@@ -443,21 +439,8 @@ def test_tensor_rdma_backend_connectivity_matrix(
         min_nodes=1,
     )
 
-    node_rdma_ctl = {
-        node_a: NodeRdmaCtlStatus(enabled=True),
-        node_b: NodeRdmaCtlStatus(enabled=True),
-        node_c: NodeRdmaCtlStatus(enabled=True),
-    }
-
     # act
-    placements = place_instance(
-        cic,
-        topology,
-        {},
-        node_memory,
-        node_network,
-        node_rdma_ctl=node_rdma_ctl,
-    )
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
 
     # assert
     assert len(placements) == 1
@@ -499,131 +482,6 @@ def test_tensor_rdma_backend_connectivity_matrix(
             assert len(ip_part.split(".")) == 4
 
 
-def _build_three_node_rdma_topology() -> tuple[
-    Topology, NodeId, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]
-]:
-    topology = Topology()
-    node_a = NodeId()
-    node_b = NodeId()
-    node_c = NodeId()
-
-    ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
-    ethernet_conn = SocketConnection(
-        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
-    )
-    node_network = {
-        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
-        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
-        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
-    }
-
-    for n in (node_a, node_b, node_c):
-        topology.add_node(n)
-
-    rdma_pairs = [
-        (node_a, node_b, 3),
-        (node_b, node_a, 3),
-        (node_b, node_c, 4),
-        (node_c, node_b, 4),
-        (node_a, node_c, 5),
-        (node_c, node_a, 5),
-    ]
-    for src, sink, iface in rdma_pairs:
-        topology.add_connection(
-            Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
-        )
-
-    socket_pairs = [
-        (node_a, node_b),
-        (node_b, node_c),
-        (node_c, node_a),
-        (node_a, node_c),
-        (node_b, node_a),
-        (node_c, node_b),
-    ]
-    for src, sink in socket_pairs:
-        topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))
-
-    return topology, node_a, node_b, node_c, node_network
-
-
-def test_place_mlx_jaccl_rejects_when_a_node_has_rdma_ctl_disabled(
-    model_card: ModelCard,
-):
-    # arrange
-    model_card = model_card.model_copy(
-        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
-    )
-    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
-    node_memory = {
-        node_a: create_node_memory(500),
-        node_b: create_node_memory(500),
-        node_c: create_node_memory(500),
-    }
-    node_rdma_ctl = {
-        node_a: NodeRdmaCtlStatus(enabled=True),
-        node_b: NodeRdmaCtlStatus(enabled=True),
-        node_c: NodeRdmaCtlStatus(enabled=False),
-    }
-    cic = PlaceInstance(
-        sharding=Sharding.Tensor,
-        instance_meta=InstanceMeta.MlxJaccl,
-        command_id=CommandId(),
-        model_card=model_card,
-        min_nodes=3,
-    )
-
-    # act / assert
-    with pytest.raises(
-        ValueError, match="Requested RDMA \\(MlxJaccl\\) but no RDMA-connected cycles"
-    ):
-        place_instance(
-            cic,
-            topology,
-            {},
-            node_memory,
-            node_network,
-            node_rdma_ctl=node_rdma_ctl,
-        )
-
-
-def test_place_mlx_jaccl_rejects_when_node_rdma_ctl_missing(model_card: ModelCard):
-    """A node with no observed rdma_ctl status must not participate in RDMA placement."""
-    # arrange
-    model_card = model_card.model_copy(
-        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
-    )
-    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
-    node_memory = {
-        node_a: create_node_memory(500),
-        node_b: create_node_memory(500),
-        node_c: create_node_memory(500),
-    }
-    # node_c has no rdma_ctl entry at all
-    node_rdma_ctl = {
-        node_a: NodeRdmaCtlStatus(enabled=True),
-        node_b: NodeRdmaCtlStatus(enabled=True),
-    }
-    cic = PlaceInstance(
-        sharding=Sharding.Tensor,
-        instance_meta=InstanceMeta.MlxJaccl,
-        command_id=CommandId(),
-        model_card=model_card,
-        min_nodes=3,
-    )
-
-    # act / assert
-    with pytest.raises(ValueError):
-        place_instance(
-            cic,
-            topology,
-            {},
-            node_memory,
-            node_network,
-            node_rdma_ctl=node_rdma_ctl,
-        )
-
-
 def _make_task(
     instance_id: InstanceId,
     status: TaskStatus = TaskStatus.Running,

src/exo/shared/apply.py

@@ -4,8 +4,7 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.models.model_cards import ModelCard
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
     CustomModelCardAdded,
@@ -60,24 +59,14 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeDiskUsage,
     NodeNetworkInterfaces,
+    NvmlMetrics,
     RdmaCtlStatus,
     StaticNodeInformation,
     ThunderboltBridgeInfo,
+    VllmCapability,
 )
 
 
-def _is_rdma_ctl_enabled(
-    node_id: NodeId, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus]
-) -> bool:
-    """A node is RDMA-capable only if rdma_ctl status has been observed as enabled.
-
-    Missing entries default to ``False`` — if we have not yet observed (or the node
-    cannot run) ``rdma_ctl``, it must not participate in an RDMA-backed instance.
-    """
-    status = node_rdma_ctl.get(node_id)
-    return status is not None and status.enabled
-
-
 def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
@@ -88,12 +77,10 @@ def event_apply(event: Event, state: State) -> State:
             | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
+            | CustomModelCardAdded()
+            | CustomModelCardDeleted()
         ):  # Pass-through events that don't modify state
             return state
-        case CustomModelCardAdded():
-            return apply_custom_model_card_added(event, state)
-        case CustomModelCardDeleted():
-            return apply_custom_model_card_deleted(event, state)
         case InstanceCreated():
             return apply_instance_created(event, state)
         case InstanceDeleted():
@@ -319,6 +306,9 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     node_rdma_ctl = {
         key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
     }
+    node_vllm = {
+        key: value for key, value in state.node_vllm.items() if key != event.node_id
+    }
     # Only recompute cycles if the leaving node had TB bridge enabled
     leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
     leaving_node_had_tb_enabled = (
@@ -341,6 +331,7 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "node_thunderbolt": node_thunderbolt,
             "node_thunderbolt_bridge": node_thunderbolt_bridge,
             "node_rdma_ctl": node_rdma_ctl,
+            "node_vllm": node_vllm,
             "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
         }
     )
@@ -367,6 +358,11 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 event.node_id: info.system_profile,
             }
             update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
+        case NvmlMetrics():
+            update["node_system"] = {
+                **state.node_system,
+                event.node_id: info.system_profile,
+            }
         case MemoryUsage():
             update["node_memory"] = {**state.node_memory, event.node_id: info}
         case NodeDiskUsage():
@@ -412,9 +408,6 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 for nid in state.node_thunderbolt
                 for tb_ident in state.node_thunderbolt[nid].interfaces
             }
-            source_is_rdma_enabled = _is_rdma_ctl_enabled(
-                event.node_id, state.node_rdma_ctl
-            )
             as_rdma_conns = [
                 Connection(
                     source=event.node_id,
@@ -427,10 +420,6 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 for tb_conn in info.conns
                 if tb_conn.source_uuid in conn_map
                 if tb_conn.sink_uuid in conn_map
-                if source_is_rdma_enabled
-                and _is_rdma_ctl_enabled(
-                    conn_map[tb_conn.sink_uuid][0], state.node_rdma_ctl
-                )
             ]
             topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
         case ThunderboltBridgeInfo():
@@ -454,12 +443,11 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
-            # If RDMA just got disabled on this node, drop any RDMA edges touching it
-            # so placement / topology consumers cannot pick a disabled node for an
-            # RDMA-backed instance. (Edges will repopulate on the next
-            # MacThunderboltConnections poll once both endpoints are enabled again.)
-            if not info.enabled:
-                topology.remove_all_rdma_connections_touching(event.node_id)
+        case VllmCapability():
+            update["node_vllm"] = {
+                **state.node_vllm,
+                event.node_id: info.available,
+            }
 
     return state.model_copy(update=update)
 
@@ -475,22 +463,3 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
     topology.remove_connection(event.conn)
     # TODO: Clean up removing the reverse connection
     return state.model_copy(update={"topology": topology})
-
-
-def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
-    new_cards: Mapping[ModelId, ModelCard] = {
-        **state.custom_model_cards,
-        event.model_card.model_id: event.model_card,
-    }
-    return state.model_copy(update={"custom_model_cards": new_cards})
-
-
-def apply_custom_model_card_deleted(
-    event: CustomModelCardDeleted, state: State
-) -> State:
-    new_cards: Mapping[ModelId, ModelCard] = {
-        model_id: card
-        for model_id, card in state.custom_model_cards.items()
-        if model_id != event.model_id
-    }
-    return state.model_copy(update={"custom_model_cards": new_cards})

src/exo/shared/constants.py

@@ -68,12 +68,7 @@ DASHBOARD_DIR = (
 # Log files (data/logs or cache)
 EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
 EXO_LOG = EXO_LOG_DIR / "exo.log"
-EXO_RUNNER_LOG_DIR = EXO_LOG_DIR / "runner_log"
-EXO_RUNNER_STDOUT_LOG = EXO_RUNNER_LOG_DIR / "stdout.log"
-EXO_RUNNER_STDERR_LOG = EXO_RUNNER_LOG_DIR / "stderr.log"
-
 EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
-EXO_PID_FILE = EXO_CACHE_HOME / "exo.pid"
 
 # Identity (config)
 EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"

src/exo/shared/models/model_cards.py

@@ -39,57 +39,7 @@ _BUILTIN_CARD_DIRS = [
     Path(RESOURCES_DIR) / "image_model_cards",
 ]
 
-
-class _CardCache:
-    def __init__(self):
-        self.cc: dict[ModelId, "ModelCard"] = {}
-
-    def get(self, model_id: ModelId) -> "ModelCard | None":
-        return self.cc.get(model_id)
-
-    async def save(self, card: "ModelCard"):
-        self.cc[card.model_id] = card
-        try:
-            await card.save_to_custom_dir()
-        except OSError as e:
-            logger.warning(f"failed to save custom model card ({e.strerror})")
-
-    async def pop(self, model_id: ModelId) -> "ModelCard | None":
-        """Delete a user-added custom model card. Returns True if deleted."""
-        card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
-        try:
-            if await card_path.exists():
-                await card_path.unlink()
-                return self.cc.pop(model_id, None)
-        except OSError as e:
-            logger.warning(f"failed to delete custom model card ({e.strerror})")
-
-    async def list_all(self) -> list["ModelCard"]:
-        if len(self.cc) == 0:
-            await self.refresh()
-        if EXO_ENABLE_IMAGE_MODELS:
-            return list(self.cc.values())
-        return [c for c in self.cc.values() if not _is_image_card(c)]
-
-    async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
-        """Load all TOML model cards from a directory into the cache."""
-        async for toml_file in directory.rglob("*.toml"):
-            try:
-                card = await ModelCard.load_from_path(toml_file)
-                if is_custom:
-                    card = card.model_copy(update={"is_custom": True})
-                if self.get(card.model_id) is None:
-                    self.cc[card.model_id] = card
-            except (ValidationError, TOMLKitError):
-                pass
-
-    async def refresh(self) -> None:
-        for path in _BUILTIN_CARD_DIRS:
-            await self._load_cards_from_dir(path, is_custom=False)
-        await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
-
-
-card_cache = _CardCache()
+_card_cache: dict[ModelId, "ModelCard"] = {}
 
 
 def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -109,10 +59,42 @@ def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
     return None
 
 
+async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
+    """Load all TOML model cards from a directory into the cache."""
+    async for toml_file in directory.rglob("*.toml"):
+        try:
+            card = await ModelCard.load_from_path(toml_file)
+            if is_custom:
+                card = card.model_copy(update={"is_custom": True})
+            if card.model_id not in _card_cache:
+                _card_cache[card.model_id] = card
+        except (ValidationError, TOMLKitError):
+            pass
+
+
+async def _refresh_card_cache() -> None:
+    for path in _BUILTIN_CARD_DIRS:
+        await _load_cards_from_dir(path, is_custom=False)
+    await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
+
+
 def _is_image_card(card: "ModelCard") -> bool:
     return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
 
 
+def get_card(model_id: ModelId) -> "ModelCard | None":
+    """Look up a single model card from the cache by ID."""
+    return _card_cache.get(model_id)
+
+
+async def get_model_cards() -> list["ModelCard"]:
+    if len(_card_cache) == 0:
+        await _refresh_card_cache()
+    if EXO_ENABLE_IMAGE_MODELS:
+        return list(_card_cache.values())
+    return [c for c in _card_cache.values() if not _is_image_card(c)]
+
+
 class ModelTask(str, Enum):
     TextGeneration = "TextGeneration"
     TextToImage = "TextToImage"
@@ -168,6 +150,7 @@ class ModelCard(FrozenModel):
     context_length: int = 0
     uses_cfg: bool = False
     trust_remote_code: bool = True
+    requires_vllm: bool = False
     is_custom: bool = False
     vision: VisionCardConfig | None = None
     sampling_defaults: SamplingDefaults = Field(default_factory=SamplingDefaults)
@@ -214,13 +197,14 @@ class ModelCard(FrozenModel):
     # Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
     @staticmethod
     async def load(model_id: ModelId) -> "ModelCard":
-        if card_cache.get(model_id) is None:
-            await card_cache.refresh()
-        if (mc := card_cache.get(model_id)) is not None:
+        if model_id not in _card_cache:
+            await _refresh_card_cache()
+        if (mc := _card_cache.get(model_id)) is not None:
             return mc
 
         mc = await ModelCard.fetch_from_hf(model_id)
         await mc.save_to_custom_dir()
+        _card_cache[model_id] = mc
         return mc
 
     @staticmethod
@@ -250,6 +234,21 @@ class ModelCard(FrozenModel):
         )
 
 
+def add_to_card_cache(card: "ModelCard") -> None:
+    """Add or update a model card in the in-memory cache."""
+    _card_cache[card.model_id] = card
+
+
+async def delete_custom_card(model_id: ModelId) -> bool:
+    """Delete a user-added custom model card. Returns True if deleted."""
+    card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
+    if await card_path.exists():
+        await card_path.unlink()
+        _card_cache.pop(model_id, None)
+        return True
+    return False
+
+
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 
@@ -351,7 +350,11 @@ async def fetch_config_data(model_id: ModelId) -> ConfigData:
 
 
 async def fetch_safetensors_size(model_id: ModelId) -> Memory:
-    """Gets model size from safetensors index or falls back to HF API."""
+    """Gets model size from safetensors index or falls back to HF API.
+
+    Single-shard repos don't have a `model.safetensors.index.json`; fall back
+    to the HF API for those.
+    """
     from exo.download.download_utils import (
         download_file_with_retry,
         resolve_model_dir,
@@ -359,21 +362,25 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
     from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
     target_dir = await resolve_model_dir(model_id)
-    index_path = await download_file_with_retry(
-        model_id,
-        "main",
-        "model.safetensors.index.json",
-        target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
-            f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
-        ),
-    )
-    async with aiofiles.open(index_path, "r") as f:
-        index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+    try:
+        index_path = await download_file_with_retry(
+            model_id,
+            "main",
+            "model.safetensors.index.json",
+            target_dir,
+            lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+                f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
+            ),
+        )
+    except FileNotFoundError:
+        index_path = None
 
-    metadata = index_data.metadata
-    if metadata is not None and metadata.total_size is not None:
-        return Memory.from_bytes(metadata.total_size)
+    if index_path is not None:
+        async with aiofiles.open(index_path, "r") as f:
+            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+        metadata = index_data.metadata
+        if metadata is not None and metadata.total_size is not None:
+            return Memory.from_bytes(metadata.total_size)
 
     info = model_info(model_id)
     if info.safetensors is None:

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

@@ -1,44 +0,0 @@
-from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.common import ModelId
-from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
-    IndexedEvent,
-)
-from exo.shared.types.memory import Memory
-from exo.shared.types.state import State
-
-
-def _model_card(model_id: ModelId) -> ModelCard:
-    return ModelCard(
-        model_id=model_id,
-        n_layers=1,
-        storage_size=Memory.from_bytes(1),
-        hidden_size=1,
-        supports_tensor=True,
-        tasks=[ModelTask.TextGeneration],
-    )
-
-
-def test_custom_model_card_added_is_reduced_into_state() -> None:
-    card = _model_card(ModelId("custom/model"))
-
-    state = apply(
-        State(),
-        IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
-    )
-
-    assert state.custom_model_cards == {card.model_id: card}
-
-
-def test_custom_model_card_deleted_removes_card_from_state() -> None:
-    card = _model_card(ModelId("custom/model"))
-    state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
-
-    state = apply(
-        state,
-        IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
-    )
-
-    assert state.custom_model_cards == {}

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

@@ -1,231 +0,0 @@
-from datetime import datetime, timezone
-
-from exo.shared.apply import apply_node_gathered_info
-from exo.shared.topology import Topology
-from exo.shared.types.common import NodeId
-from exo.shared.types.events import NodeGatheredInfo
-from exo.shared.types.profiling import (
-    NodeRdmaCtlStatus,
-    NodeThunderboltInfo,
-)
-from exo.shared.types.state import State
-from exo.shared.types.thunderbolt import ThunderboltConnection, ThunderboltIdentifier
-from exo.shared.types.topology import RDMAConnection
-from exo.utils.info_gatherer.info_gatherer import (
-    MacThunderboltConnections,
-    RdmaCtlStatus,
-)
-
-
-def _now() -> str:
-    return datetime.now(timezone.utc).isoformat()
-
-
-def _make_state_with_thunderbolt_idents(
-    *node_ids_and_uuids: tuple[NodeId, str, str],
-    rdma_ctl: dict[NodeId, NodeRdmaCtlStatus] | None = None,
-) -> State:
-    """Build a State with Thunderbolt identifiers per node so the apply MacThunderboltConnections
-    case can resolve uuid -> (node, iface)."""
-    node_thunderbolt = {
-        nid: NodeThunderboltInfo(
-            interfaces=[ThunderboltIdentifier(rdma_interface=iface, domain_uuid=uuid)]
-        )
-        for nid, uuid, iface in node_ids_and_uuids
-    }
-    return State(
-        node_thunderbolt=node_thunderbolt,
-        node_rdma_ctl=rdma_ctl or {},
-    )
-
-
-def _has_rdma_edge(topology: Topology, source: NodeId, sink: NodeId) -> bool:
-    return any(
-        isinstance(edge, RDMAConnection)
-        for edge in topology.get_all_connections_between(source, sink)
-    )
-
-
-def test_mac_thunderbolt_connections_emits_rdma_when_both_endpoints_enabled():
-    node_a = NodeId()
-    node_b = NodeId()
-    state = _make_state_with_thunderbolt_idents(
-        (node_a, "uuid-a", "rdma_en1"),
-        (node_b, "uuid-b", "rdma_en1"),
-        rdma_ctl={
-            node_a: NodeRdmaCtlStatus(enabled=True),
-            node_b: NodeRdmaCtlStatus(enabled=True),
-        },
-    )
-
-    event = NodeGatheredInfo(
-        node_id=node_a,
-        when=_now(),
-        info=MacThunderboltConnections(
-            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
-        ),
-    )
-
-    new_state = apply_node_gathered_info(event, state)
-
-    assert _has_rdma_edge(new_state.topology, node_a, node_b)
-
-
-def test_mac_thunderbolt_connections_skips_rdma_when_source_rdma_ctl_disabled():
-    node_a = NodeId()
-    node_b = NodeId()
-    state = _make_state_with_thunderbolt_idents(
-        (node_a, "uuid-a", "rdma_en1"),
-        (node_b, "uuid-b", "rdma_en1"),
-        rdma_ctl={
-            node_a: NodeRdmaCtlStatus(enabled=False),
-            node_b: NodeRdmaCtlStatus(enabled=True),
-        },
-    )
-
-    event = NodeGatheredInfo(
-        node_id=node_a,
-        when=_now(),
-        info=MacThunderboltConnections(
-            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
-        ),
-    )
-
-    new_state = apply_node_gathered_info(event, state)
-
-    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
-
-
-def test_mac_thunderbolt_connections_skips_rdma_when_sink_rdma_ctl_disabled():
-    node_a = NodeId()
-    node_b = NodeId()
-    state = _make_state_with_thunderbolt_idents(
-        (node_a, "uuid-a", "rdma_en1"),
-        (node_b, "uuid-b", "rdma_en1"),
-        rdma_ctl={
-            node_a: NodeRdmaCtlStatus(enabled=True),
-            node_b: NodeRdmaCtlStatus(enabled=False),
-        },
-    )
-
-    event = NodeGatheredInfo(
-        node_id=node_a,
-        when=_now(),
-        info=MacThunderboltConnections(
-            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
-        ),
-    )
-
-    new_state = apply_node_gathered_info(event, state)
-
-    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
-
-
-def test_mac_thunderbolt_connections_skips_rdma_when_rdma_ctl_status_missing():
-    """Missing rdma_ctl status defaults to not-enabled — node is RDMA-incapable."""
-    node_a = NodeId()
-    node_b = NodeId()
-    state = _make_state_with_thunderbolt_idents(
-        (node_a, "uuid-a", "rdma_en1"),
-        (node_b, "uuid-b", "rdma_en1"),
-        rdma_ctl={
-            node_a: NodeRdmaCtlStatus(enabled=True),
-            # node_b intentionally absent
-        },
-    )
-
-    event = NodeGatheredInfo(
-        node_id=node_a,
-        when=_now(),
-        info=MacThunderboltConnections(
-            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
-        ),
-    )
-
-    new_state = apply_node_gathered_info(event, state)
-
-    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
-
-
-def test_rdma_ctl_status_disabled_purges_existing_rdma_edges():
-    """When a node reports rdma_ctl disabled, all RDMA edges touching it must be removed."""
-    node_a = NodeId()
-    node_b = NodeId()
-
-    # Start with both nodes RDMA-enabled and existing RDMA edges in the topology.
-    state = _make_state_with_thunderbolt_idents(
-        (node_a, "uuid-a", "rdma_en1"),
-        (node_b, "uuid-b", "rdma_en1"),
-        rdma_ctl={
-            node_a: NodeRdmaCtlStatus(enabled=True),
-            node_b: NodeRdmaCtlStatus(enabled=True),
-        },
-    )
-    state = apply_node_gathered_info(
-        NodeGatheredInfo(
-            node_id=node_a,
-            when=_now(),
-            info=MacThunderboltConnections(
-                conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
-            ),
-        ),
-        state,
-    )
-    state = apply_node_gathered_info(
-        NodeGatheredInfo(
-            node_id=node_b,
-            when=_now(),
-            info=MacThunderboltConnections(
-                conns=[ThunderboltConnection(source_uuid="uuid-b", sink_uuid="uuid-a")]
-            ),
-        ),
-        state,
-    )
-    assert _has_rdma_edge(state.topology, node_a, node_b)
-    assert _has_rdma_edge(state.topology, node_b, node_a)
-
-    # Now node_a flips to rdma_ctl disabled — both directions of RDMA edge must drop.
-    state = apply_node_gathered_info(
-        NodeGatheredInfo(
-            node_id=node_a, when=_now(), info=RdmaCtlStatus(enabled=False)
-        ),
-        state,
-    )
-
-    assert not _has_rdma_edge(state.topology, node_a, node_b)
-    assert not _has_rdma_edge(state.topology, node_b, node_a)
-    assert state.node_rdma_ctl[node_a].enabled is False
-
-
-def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
-    """Purging RDMA edges for a disabled node must not affect non-RDMA edges."""
-    from exo.shared.types.multiaddr import Multiaddr
-    from exo.shared.types.topology import Connection, SocketConnection
-
-    topology = Topology()
-    node_a = NodeId()
-    node_b = NodeId()
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_connection(
-        Connection(
-            source=node_a,
-            sink=node_b,
-            edge=RDMAConnection(
-                source_rdma_iface="rdma_en1", sink_rdma_iface="rdma_en1"
-            ),
-        )
-    )
-    socket_edge = SocketConnection(
-        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
-    )
-    topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
-
-    topology.remove_all_rdma_connections_touching(node_a)
-
-    assert not _has_rdma_edge(topology, node_a, node_b)
-    # Socket edge survives.
-    assert any(
-        isinstance(edge, SocketConnection)
-        for edge in topology.get_all_connections_between(node_a, node_b)
-    )

src/exo/shared/topology.py

@@ -169,22 +169,6 @@ class Topology:
         for conn in new_connections:
             self.add_connection(conn)
 
-    def remove_all_rdma_connections_touching(self, node_id: NodeId) -> None:
-        """Remove every RDMA edge incident to ``node_id`` (incoming or outgoing)."""
-        if node_id not in self._vertex_indices:
-            return
-        rx_idx = self._vertex_indices[node_id]
-        rdma_edge_idxs = [
-            edge_idx
-            for edge_idx in (
-                *self._graph.out_edge_indices(rx_idx),
-                *self._graph.in_edge_indices(rx_idx),
-            )
-            if isinstance(self._graph.get_edge_data_by_index(edge_idx), RDMAConnection)
-        ]
-        for edge_idx in rdma_edge_idxs:
-            self._graph.remove_edge_from_index(edge_idx)
-
     def remove_connection(self, conn: Connection) -> None:
         if (
             conn.source not in self._vertex_indices

src/exo/shared/types/state.py

@@ -5,9 +5,8 @@ from typing import Any, cast
 from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
-from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -59,6 +58,7 @@ class State(FrozenModel):
     node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
     node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
+    node_vllm: Mapping[NodeId, bool] = {}
 
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
@@ -66,9 +66,6 @@ class State(FrozenModel):
     instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
     prefill_server_ports: Mapping[RunnerId, int] = {}
 
-    # User-added model cards. Workers can reconcile their on-disk custom card cache
-    custom_model_cards: Mapping[ModelId, ModelCard] = {}
-
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()

src/exo/shared/types/text_generation.py

@@ -135,9 +135,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     prefill_endpoint: str | None = None
 
     def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
-        from exo.shared.models import model_cards
+        from exo.shared.models.model_cards import get_card
 
-        card = model_cards.card_cache.get(self.model)
+        card = get_card(self.model)
         if card is None:
             return self
 

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

@@ -15,6 +15,7 @@ class InstanceId(Id):
 class InstanceMeta(str, Enum):
     MlxRing = "MlxRing"
     MlxJaccl = "MlxJaccl"
+    Vllm = "Vllm"
 
 
 class BaseInstance(TaggedModel):
@@ -35,8 +36,12 @@ class MlxJacclInstance(BaseInstance):
     jaccl_coordinators: dict[NodeId, str]
 
 
+class VllmInstance(BaseInstance):
+    pass
+
+
 # TODO: Single node instance
-Instance = MlxRingInstance | MlxJacclInstance
+Instance = MlxRingInstance | MlxJacclInstance | VllmInstance
 
 
 class BoundInstance(FrozenModel):

src/exo/utils/async_process.py

@@ -1,290 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import faulthandler
-import multiprocessing as mp
-import os
-import sys
-from collections.abc import Callable, Iterable, Mapping
-from multiprocessing.process import BaseProcess
-from multiprocessing.resource_sharer import DupFd
-from typing import final
-
-from anyio import (
-    TASK_STATUS_IGNORED,
-    BrokenResourceError,
-    CancelScope,
-    ClosedResourceError,
-    Event,
-    create_task_group,
-    move_on_after,
-    sleep,
-    wait_readable,
-)
-from anyio.abc import TaskStatus
-from loguru import logger
-
-from exo.utils.channels import Receiver, Sender, channel
-
-_STDOUT_FD = 1
-_STDERR_FD = 2
-_READ_CHUNK_SIZE = 64 * 1024
-_TERMINATE_GRACE_SECONDS = 10.0
-_TERMINATE_RETRY_GRACE_SECONDS = 2.0
-_TERMINATE_ATTEMPTS = 10
-_KILL_GRACE_SECONDS = 5.0
-
-
-@final
-class AsyncProcess:
-    def __init__(
-        self,
-        target: Callable[..., object] | None = None,
-        name: str | None = None,
-        args: Iterable[object] = (),
-        kwargs: Mapping[str, object] | None = None,
-        *,
-        daemon: bool | None = None,
-    ) -> None:
-        # setup state
-        self._target = target
-        self._name = name
-        self._args = args
-        self._kwargs = kwargs
-        self._daemon = daemon
-
-        # lifecycle state
-        self._process: BaseProcess | None = None
-        self._pid: int | None = None
-        self._stdout_tx, self._stdout_rx = channel[bytes]()
-        self._stderr_tx, self._stderr_rx = channel[bytes]()
-        self._started = Event()
-        self._done = Event()
-        self._run_cancel_scope: CancelScope | None = None
-        self._start_error: BaseException | None = None
-        self._exitcode: int | None = None
-
-    async def run(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) -> None:
-        if self._run_cancel_scope is not None or self._done.is_set():
-            raise RuntimeError("process has already been started")
-
-        stdout_read_fd: int | None = None
-        stdout_write_fd: int | None = None
-        stderr_read_fd: int | None = None
-        stderr_write_fd: int | None = None
-
-        def cleanup_stdio_fd() -> None:
-            nonlocal stdout_read_fd, stdout_write_fd, stderr_read_fd, stderr_write_fd
-            stdout_read_fd = _close_fd(stdout_read_fd)
-            stdout_write_fd = _close_fd(stdout_write_fd)
-            stderr_read_fd = _close_fd(stderr_read_fd)
-            stderr_write_fd = _close_fd(stderr_write_fd)
-
-        try:
-            with CancelScope() as run_cancel_scope:
-                self._run_cancel_scope = run_cancel_scope
-                stdout_read_fd, stdout_write_fd = os.pipe()
-                stderr_read_fd, stderr_write_fd = os.pipe()
-
-                process = mp.Process(
-                    target=_run_with_captured_stdio,
-                    name=self._name,
-                    args=(
-                        DupFd(stdout_write_fd),
-                        DupFd(stderr_write_fd),
-                        self._target,
-                        *self._args,
-                    ),
-                    kwargs={} if self._kwargs is None else self._kwargs,
-                    daemon=self._daemon,
-                )
-                process.start()
-                pid = process.pid
-                if pid is None:
-                    raise RuntimeError("started process has no pid")
-
-                # important to close parent write-side FD to prevent hangs
-                stdout_write_fd = _close_fd(stdout_write_fd)
-                stderr_write_fd = _close_fd(stderr_write_fd)
-
-                self._process = process
-                self._pid = pid
-                self._started.set()
-
-                async with create_task_group() as tg:
-                    tg.start_soon(_drain_fd, stdout_read_fd, self._stdout_tx)
-                    stdout_read_fd = None
-                    tg.start_soon(_drain_fd, stderr_read_fd, self._stderr_tx)
-                    stderr_read_fd = None
-                    task_status.started()
-                    await self.wait()
-        except BaseException as exc:
-            if not self._started.is_set():
-                self._start_error = exc
-                self._started.set()
-            raise
-        finally:
-            try:
-                with CancelScope(shield=True):
-                    await self._terminate_if_still_alive()
-            finally:
-                cleanup_stdio_fd()
-                for tx in (self._stdout_tx, self._stderr_tx):
-                    with contextlib.suppress(Exception):
-                        await tx.aclose()
-                if self._process is not None:
-                    with contextlib.suppress(ValueError):
-                        self._process.close()
-                self._run_cancel_scope = None
-                self._done.set()
-
-    async def stop(self) -> None:
-        if self._run_cancel_scope is None and not self._done.is_set():
-            raise RuntimeError("process has not been started")
-        if self._run_cancel_scope is not None:
-            self._run_cancel_scope.cancel()
-        await self._done.wait()
-
-    async def aclose(self) -> None:
-        await self.stop()
-
-    async def wait(self) -> int:
-        if self._exitcode is not None:
-            return self._exitcode
-
-        await self._started.wait()
-        if self._start_error is not None:
-            raise self._start_error
-        assert self._process is not None
-
-        while True:
-            exitcode = self.exitcode
-            if exitcode is not None:
-                return exitcode
-            await sleep(0.01)
-
-    @property
-    def pid(self) -> int:
-        if self._pid is None:
-            raise RuntimeError("process has not been started")
-        return self._pid
-
-    @property
-    def exitcode(self) -> int | None:
-        if self._exitcode is not None:
-            return self._exitcode
-        if self._process is None:
-            return None
-
-        with contextlib.suppress(ValueError):
-            exitcode = self._process.exitcode
-            if exitcode is not None:
-                self._exitcode = exitcode
-            return exitcode
-        return None
-
-    def is_alive(self) -> bool:
-        if self._process is None:
-            return False
-
-        with contextlib.suppress(ValueError):
-            return self._process.is_alive()
-        return False
-
-    # TODO: maybe in the future if needed, create stdin that is also installed,
-    #       and a ByteSendStream handle is provided for it :)
-
-    @property
-    def stdout(self) -> Receiver[bytes]:
-        return self._stdout_rx
-
-    @property
-    def stderr(self) -> Receiver[bytes]:
-        return self._stderr_rx
-
-    async def _terminate_if_still_alive(self) -> None:
-        process = self._process
-        if process is None:
-            return
-
-        if self.exitcode is not None:
-            return
-
-        with contextlib.suppress(ValueError):
-            if not process.is_alive():
-                return
-
-            logger.warning("Child process didn't shut down successfully, terminating")
-            process.terminate()
-            with move_on_after(_TERMINATE_GRACE_SECONDS):
-                await self.wait()
-
-            if self.exitcode is not None or not process.is_alive():
-                logger.warning("Terminated nicely in the first attempt!")
-                return
-
-            for attempt in range(2, _TERMINATE_ATTEMPTS + 1):
-                process.terminate()
-                with move_on_after(_TERMINATE_RETRY_GRACE_SECONDS):
-                    await self.wait()
-
-                if self.exitcode is not None or not process.is_alive():
-                    logger.warning(f"That took {attempt} attempts :)")
-                    return
-
-            logger.critical("Child process didn't respond to SIGTERM, killing")
-            j = 0
-            while True:
-                process.kill()
-                with move_on_after(_KILL_GRACE_SECONDS):
-                    await self.wait()
-                j += 1
-                if self.exitcode is not None or not process.is_alive():
-                    break
-            logger.warning(f"That took {j} attempts :(")
-
-
-# Spawn-mode multiprocessing requires a module-level target that can be pickled.
-def _run_with_captured_stdio(
-    stdout: DupFd,
-    stderr: DupFd,
-    target: Callable[..., object] | None,
-    *target_args: object,
-    **target_kwargs: object,
-) -> None:
-    stdout_fd = stdout.detach()
-    stderr_fd = stderr.detach()
-
-    try:
-        os.dup2(stdout_fd, _STDOUT_FD)
-        os.dup2(stderr_fd, _STDERR_FD)
-    finally:
-        for fd in (stdout_fd, stderr_fd):
-            if fd not in (_STDOUT_FD, _STDERR_FD):
-                _close_fd(fd)
-
-    faulthandler.enable(file=sys.stderr, all_threads=True)
-    if target is not None:
-        target(*target_args, **target_kwargs)
-
-
-async def _drain_fd(fd: int, tx: Sender[bytes]) -> None:
-    try:
-        while True:
-            await wait_readable(fd)
-            chunk = os.read(fd, _READ_CHUNK_SIZE)
-            if not chunk:
-                return
-            await tx.send(chunk)
-    except (BrokenPipeError, BrokenResourceError, ClosedResourceError):
-        pass
-    finally:
-        _close_fd(fd)
-        await tx.aclose()
-
-
-def _close_fd(fd: int | None) -> None:
-    if fd is None:
-        return
-    with contextlib.suppress(OSError):
-        os.close(fd)

src/exo/utils/banner.py

@@ -25,12 +25,12 @@ def print_startup_banner(port: int) -> None:
     banner = f"""
 ╔═══════════════════════════════════════════════════════════════════════╗
 ║                                                                       ║
-║   ███████╗██╗  ██╗ ██████╗                                            ║
-║   ██╔════╝╚██╗██╔╝██╔═══██╗                                           ║
-║   █████╗   ╚███╔╝ ██║   ██║                                           ║
-║   ██╔══╝   ██╔██╗ ██║   ██║                                           ║
-║   ███████╗██╔╝ ██╗╚██████╔╝                                           ║
-║   ╚══════╝╚═╝  ╚═╝ ╚═════╝                                            ║
+║   ███████╗██╗  ██╗ ██████╗               ██████╗  ██████╗ ██╗  ██╗    ║
+║   ██╔════╝╚██╗██╔╝██╔═══██╗      ██      ██╔══██╗██╔════╝ ╚██╗██╔╝    ║
+║   █████╗   ╚███╔╝ ██║   ██║    ██████╗   ██║  ██║██║  ███╗ ╚███╔╝     ║
+║   ██╔══╝   ██╔██╗ ██║   ██║    ╚═██╔═╝   ██║  ██║██║   ██║ ██╔██╗     ║
+║   ███████╗██╔╝ ██╗╚██████╔╝      ╚═╝     ██████╔╝╚██████╔╝██╔╝ ██╗    ║
+║   ╚══════╝╚═╝  ╚═╝ ╚═════╝               ╚═════╝  ╚═════╝ ╚═╝  ╚═╝    ║
 ║                                                                       ║
 ║   Distributed AI Inference Cluster                                    ║
 ║                                                                       ║

src/exo/utils/daemon.py

@@ -1,28 +0,0 @@
-import os
-import sys
-
-_STDIN_FD = 0
-_STDOUT_FD = 1
-_STDERR_FD = 2
-
-
-def detach_stdio_to_devnull() -> None:
-    """Redirect process stdio file descriptors to /dev/null."""
-
-    for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
-        if stream is not None:
-            stream.flush()
-
-    stdin_fd = os.open(os.devnull, os.O_RDONLY)
-    stdout_fd = os.open(os.devnull, os.O_WRONLY)
-    stderr_fd = os.open(os.devnull, os.O_WRONLY)
-
-    try:
-        # dup2 closes the target fd first, but leaves the source fd open.
-        os.dup2(stdin_fd, _STDIN_FD)
-        os.dup2(stdout_fd, _STDOUT_FD)
-        os.dup2(stderr_fd, _STDERR_FD)
-    finally:
-        for fd in (stdin_fd, stdout_fd, stderr_fd):
-            if fd not in (_STDIN_FD, _STDOUT_FD, _STDERR_FD):
-                os.close(fd)

src/exo/utils/info_gatherer/info_gatherer.py

@@ -31,6 +31,7 @@ from exo.utils.pydantic_ext import TaggedModel
 from exo.utils.task_group import TaskGroup
 
 from .macmon import MacmonMetrics
+from .nvml import NvmlMetrics, gather_nvidia_metrics, has_nvml
 from .system_info import (
     get_friendly_name,
     get_model_and_chip,
@@ -353,6 +354,24 @@ async def _gather_iface_map() -> dict[str, str] | None:
     return ports
 
 
+class VllmCapability(TaggedModel):
+    available: bool
+    version: str | None = None
+
+    @classmethod
+    async def gather(cls) -> Self:
+        try:
+            import importlib
+
+            vllm = importlib.import_module("vllm")
+            return cls(
+                available=True,
+                version=cast(str | None, getattr(vllm, "__version__", None)),
+            )
+        except ImportError:
+            return cls(available=False)
+
+
 GatheredInfo = (
     MacmonMetrics
     | MemoryUsage
@@ -361,6 +380,8 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
+    | NvmlMetrics
+    | VllmCapability
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -419,6 +440,8 @@ class InfoGatherer:
                 tg.start_soon(self._monitor_rdma_ctl_status, 10)
             if not IS_DARWIN:
                 tg.start_soon(self._monitor_memory_usage, 1)
+                if has_nvml():
+                    tg.start_soon(self._monitor_nvml_metrics, 1)
             tg.start_soon(self._watch_system_info, 10)
             tg.start_soon(self._monitor_misc, 60)
             tg.start_soon(self._monitor_static_info, 60)
@@ -427,6 +450,10 @@ class InfoGatherer:
             nc = await NodeConfig.gather()
             if nc is not None:
                 await self.info_sender.send(nc)
+            try:
+                await self.info_sender.send(await VllmCapability.gather())
+            except Exception as e:
+                logger.warning(f"Error gathering vLLM capability: {e}")
 
     def shutdown(self):
         self._tg.cancel_tasks()
@@ -475,6 +502,16 @@ class InfoGatherer:
                 logger.opt(exception=e).warning("Error gathering Thunderbolt data")
             await anyio.sleep(system_profiler_interval)
 
+    async def _monitor_nvml_metrics(self, nvml_poll_rate: float):
+        while True:
+            try:
+                metrics = gather_nvidia_metrics()
+                if metrics is not None:
+                    await self.info_sender.send(metrics)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering NVML metrics")
+            await anyio.sleep(nvml_poll_rate)
+
     async def _monitor_memory_usage(self, memory_poll_rate: float):
         if self._psutil_enabled:
             return

src/exo/utils/info_gatherer/nvml.py

@@ -0,0 +1,70 @@
+from exo.shared.types.profiling import SystemPerformanceProfile
+from exo.utils.pydantic_ext import TaggedModel
+
+try:
+    import pynvml as nvml
+except ImportError:
+    nvml = None
+
+_CPU_POWER_IDLE = 20.0
+_CPU_POWER_MAX = 100.0
+_GPU_POWER_MAX = 120.0
+
+
+class NvmlMetrics(TaggedModel):
+    system_profile: SystemPerformanceProfile
+
+
+def has_nvml() -> bool:
+    if nvml is None:
+        return False
+    try:
+        nvml.nvmlInit()
+        count = nvml.nvmlDeviceGetCount()
+        nvml.nvmlShutdown()
+        return count > 0
+    except Exception:
+        return False
+
+
+def gather_nvidia_metrics() -> NvmlMetrics | None:
+    if nvml is None:
+        return None
+
+    is_init = False
+    try:
+        nvml.nvmlInit()
+        is_init = True
+        count = nvml.nvmlDeviceGetCount()
+        if count == 0:
+            return None
+
+        total_gpu_util = 0.0
+        total_temp = 0.0
+        total_gpu_power = 0.0
+        for i in range(count):
+            handle = nvml.nvmlDeviceGetHandleByIndex(i)
+            util = nvml.nvmlDeviceGetUtilizationRates(handle)
+            total_gpu_util += float(util.gpu)
+            total_temp += float(
+                nvml.nvmlDeviceGetTemperatureV(handle, nvml.NVML_TEMPERATURE_GPU)
+            )
+            total_gpu_power += float(nvml.nvmlDeviceGetPowerUsage(handle)) / 1000.0
+
+        gpu_load_fraction = min(total_gpu_power / _GPU_POWER_MAX, 1.0)
+        estimated_cpu_power = (
+            _CPU_POWER_IDLE + (_CPU_POWER_MAX - _CPU_POWER_IDLE) * gpu_load_fraction
+        )
+
+        return NvmlMetrics(
+            system_profile=SystemPerformanceProfile(
+                gpu_usage=total_gpu_util / count / 100.0,
+                temp=total_temp / count,
+                sys_power=total_gpu_power + estimated_cpu_power,
+            ),
+        )
+    except Exception:
+        return None
+    finally:
+        if is_init:
+            nvml.nvmlShutdown()

src/exo/utils/info_gatherer/system_info.py

@@ -1,6 +1,7 @@
 import platform
 import socket
 import sys
+from pathlib import Path
 from subprocess import CalledProcessError
 
 import psutil
@@ -117,12 +118,90 @@ async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     return interfaces_info
 
 
+def _read_dmi_field(name: str) -> str | None:
+    try:
+        path = Path(f"/sys/class/dmi/id/{name}")
+        if path.exists():
+            return path.read_text().strip()
+    except (OSError, PermissionError):
+        pass
+    return None
+
+
+async def _get_linux_model_and_chip() -> tuple[str, str]:
+    model = "Linux"
+    chip = "Unknown Chip"
+
+    product_name = _read_dmi_field("product_name")
+    sys_vendor = _read_dmi_field("sys_vendor")
+
+    # DGX Spark: DMI product_name may be "DGX_Spark" or "gx10" variant
+    product_lower = (product_name or "").lower()
+    if product_name and ("dgx" in product_lower or "gx10" in product_lower):
+        model = "DGX Spark"
+        try:
+            process = await run_process(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
+            )
+            gpu_name = process.stdout.decode().strip().split("\n")[0]
+            chip = gpu_name if gpu_name and gpu_name != "[N/A]" else "NVIDIA GB10"
+        except (CalledProcessError, FileNotFoundError):
+            chip = "NVIDIA GB10"
+        return (model, chip)
+
+    # Other NVIDIA systems (sys_vendor contains "NVIDIA")
+    if sys_vendor and "NVIDIA" in sys_vendor:
+        model = product_name.replace("_", " ") if product_name else "NVIDIA System"
+        try:
+            process = await run_process(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
+            )
+            gpu_name = process.stdout.decode().strip().split("\n")[0]
+            if gpu_name and gpu_name != "[N/A]":
+                chip = gpu_name
+        except (CalledProcessError, FileNotFoundError):
+            pass
+        return (model, chip)
+
+    # Generic Linux — detect laptop vs desktop via chassis_type
+    # SMBIOS chassis types: 8,9,10,14,31,32 = portable/laptop
+    chassis_type = _read_dmi_field("chassis_type")
+    laptop_chassis_types = {"8", "9", "10", "14", "31", "32"}
+    if chassis_type in laptop_chassis_types:
+        model = "Linux Laptop"
+    elif chassis_type is not None:
+        model = "Linux Desktop"
+
+    # Also check for battery as a fallback laptop indicator
+    if model == "Linux" and Path("/sys/class/power_supply/BAT0").exists():
+        model = "Linux Laptop"
+
+    # Use /proc/cpuinfo for chip
+    cpuinfo_path = Path("/proc/cpuinfo")
+    if cpuinfo_path.exists():
+        try:
+            for line in cpuinfo_path.read_text().splitlines():
+                if line.startswith("model name"):
+                    chip = line.split(":", 1)[1].strip()
+                    break
+        except OSError:
+            pass
+
+    return (model, chip)
+
+
 async def get_model_and_chip() -> tuple[str, str]:
-    """Get Mac system information using system_profiler."""
+    """Get system model and chip information.
+
+    On macOS, uses ``system_profiler``.  On Linux, reads DMI data from
+    sysfs and CPU info from ``/proc/cpuinfo``.
+    """
     model = "Unknown Model"
     chip = "Unknown Chip"
 
-    # TODO: better non mac support
+    if sys.platform == "linux":
+        return await _get_linux_model_and_chip()
+
     if sys.platform != "darwin":
         return (model, chip)
 

src/exo/utils/pidfile.py

@@ -1,28 +0,0 @@
-from __future__ import annotations
-
-import os
-from typing import Final
-
-from exo_pyo3_bindings import Pidfile, PidfileError
-
-from exo.shared.constants import EXO_PID_FILE
-
-_PIDFILE_MODE: Final = 0o600
-
-
-class PidfileLockError(RuntimeError):
-    pass
-
-
-def acquire_exo_pidfile() -> Pidfile:
-    path = EXO_PID_FILE
-    os.makedirs(os.path.dirname(path), exist_ok=True)
-    try:
-        pidfile = Pidfile(path, _PIDFILE_MODE)
-        pidfile.write()
-    except (OSError, PidfileError) as exception:
-        raise PidfileLockError(
-            f"Failed to acquire EXO pidfile at {path}: {exception}"
-        ) from exception
-
-    return pidfile

src/exo/utils/power_sampler.py

@@ -19,21 +19,19 @@ class PowerSampler:
     ):
         self._get_node_system = get_node_system
         self._interval = interval
-        self._samples: defaultdict[
-            NodeId, list[tuple[float, SystemPerformanceProfile]]
-        ] = defaultdict(list)
+        self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
+            defaultdict(list)
+        )
         self._start_time: float | None = None
         self._stopped = False
 
-    def _take_sample(self, t_rel: float | None = None) -> None:
-        assert self._start_time is not None
-        ts = t_rel if t_rel is not None else time.perf_counter() - self._start_time
+    def _take_sample(self) -> None:
         for node_id, profile in self._get_node_system().items():
-            self._samples[node_id].append((ts, profile))
+            self._samples[node_id].append(profile)
 
     async def run(self) -> None:
         self._start_time = time.perf_counter()
-        self._take_sample(t_rel=0.0)
+        self._take_sample()
         while not self._stopped:
             await anyio.sleep(self._interval)
             self._take_sample()
@@ -41,51 +39,26 @@ class PowerSampler:
     def result(self) -> PowerUsage:
         self._stopped = True
         assert self._start_time is not None, "result() called before run()"
+        self._take_sample()
         elapsed = time.perf_counter() - self._start_time
-        self._take_sample(t_rel=elapsed)
 
         node_stats: list[NodePowerStats] = []
-        total_energy_j = 0.0
-        for node_id, ts_profiles in self._samples.items():
-            n = len(ts_profiles)
+        for node_id, profiles in self._samples.items():
+            n = len(profiles)
             if n == 0:
                 continue
-            node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
-            avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
-            total_energy_j += node_energy_j
             node_stats.append(
                 NodePowerStats(
                     node_id=node_id,
                     samples=n,
-                    avg_sys_power=avg_power_w,
+                    avg_sys_power=sum(p.sys_power for p in profiles) / n,
                 )
             )
 
-        total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
+        total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
         return PowerUsage(
             elapsed_seconds=elapsed,
             nodes=node_stats,
-            total_avg_sys_power_watts=total_avg_sys_w,
-            total_energy_joules=total_energy_j,
+            total_avg_sys_power_watts=total_avg_sys,
+            total_energy_joules=total_avg_sys * elapsed,
         )
-
-
-def trapezoidal_energy(
-    ts_profiles: list[tuple[float, SystemPerformanceProfile]],
-    elapsed: float,
-) -> float:
-    """Integrate sys_power(t) over the sample window using the trapezoidal rule.
-    First sample is anchored at t=0 and last at t=elapsed (set by `run` /
-    `result`), so the integral spans the full request interval. Falls back to
-    power * elapsed when only one sample exists (constant-power assumption)."""
-    if len(ts_profiles) == 1:
-        return ts_profiles[0][1].sys_power * elapsed
-    energy_j = 0.0
-    for i in range(1, len(ts_profiles)):
-        t_prev, p_prev = ts_profiles[i - 1]
-        t_cur, p_cur = ts_profiles[i]
-        dt = t_cur - t_prev
-        if dt <= 0:
-            continue
-        energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
-    return energy_j

src/exo/utils/tests/conftest.py

@@ -1,8 +0,0 @@
-import multiprocessing as mp
-
-import pytest
-
-
-@pytest.fixture(scope="session", autouse=True)
-def mp_force_spawn():
-    mp.set_start_method("spawn", force=True)

src/exo/utils/tests/test_async_process.py

@@ -1,515 +0,0 @@
-import contextlib
-import os
-import signal
-import sys
-import time
-from collections.abc import AsyncIterator, Callable
-from types import FrameType
-
-import mlx.core as mx
-import pytest
-from _pytest.capture import CaptureFixture
-from anyio import EndOfStream, create_task_group, fail_after
-from pytest import MonkeyPatch
-
-import exo.utils.async_process as async_process
-from exo.utils.async_process import (
-    AsyncProcess,
-)
-from exo.utils.channels import MpSender, Receiver, mp_channel
-
-
-def _write_to_stdio(prefix: str, *, stderr_suffix: str) -> None:
-    print(f"{prefix}: python stdout")
-    print(f"{prefix}: python stderr {stderr_suffix}", file=sys.stderr)
-    os.write(1, f"{prefix}: fd stdout\n".encode())
-    os.write(2, f"{prefix}: fd stderr {stderr_suffix}\n".encode())
-
-
-def _write_large_output() -> None:
-    os.write(1, b"stdout-0123456789")
-    os.write(2, b"stderr-0123456789")
-
-
-def _write_all(fd: int, data: bytes) -> None:
-    remaining = memoryview(data)
-    while remaining:
-        written = os.write(fd, remaining)
-        remaining = remaining[written:]
-
-
-def _write_large_exact_output(size: int) -> None:
-    _write_all(1, b"stdout:" + (b"x" * size))
-    _write_all(2, b"stderr:" + (b"y" * size))
-
-
-def _raise_after_stderr_write() -> None:
-    os.write(2, b"stderr before exception\n")
-    raise RuntimeError("child boom")
-
-
-def _exit_after_stdio_write(prefix: str, exitcode: int) -> None:
-    os.write(1, f"{prefix}: stdout before _exit\n".encode())
-    os.write(2, f"{prefix}: stderr before _exit\n".encode())
-    os._exit(exitcode)
-
-
-def _abort_after_stdio_write(prefix: str) -> None:
-    os.write(1, f"{prefix}: stdout before abort\n".encode())
-    os.write(2, f"{prefix}: stderr before abort\n".encode())
-    os.abort()
-
-
-def _close_stdio_and_exit() -> None:
-    os.close(1)
-    os.close(2)
-    os._exit(0)
-
-
-def _exit_on_sigterm(exitcode: int) -> None:
-    def handle_sigterm(_signum: int, _frame: FrameType | None) -> None:
-        os._exit(exitcode)
-
-    signal.signal(signal.SIGTERM, handle_sigterm)
-    os.write(1, b"sigterm-ready\n")
-    while True:
-        time.sleep(0.1)
-
-
-def _exit_after_repeated_sigterm(required_count: int, exitcode: int) -> None:
-    sigterm_count = 0
-
-    def handle_sigterm(_signum: int, _frame: FrameType | None) -> None:
-        nonlocal sigterm_count
-        sigterm_count += 1
-        if sigterm_count >= required_count:
-            os._exit(exitcode)
-
-    signal.signal(signal.SIGTERM, handle_sigterm)
-    os.write(1, b"sigterm-ready\n")
-    while True:
-        time.sleep(0.1)
-
-
-def _ignore_sigterm_forever() -> None:
-    signal.signal(signal.SIGTERM, signal.SIG_IGN)
-    os.write(1, b"sigterm-ready\n")
-    while True:
-        time.sleep(0.1)
-
-
-def _sleep_forever() -> None:
-    while True:
-        time.sleep(0.1)
-
-
-def _send_over_mp_channel(send: MpSender[str]) -> None:
-    send.send("hello from child")
-    send.close()
-
-
-def _mlx_force_oom(size: int = 40_000) -> None:
-    """
-    Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
-    """
-    print("CHILD: start")
-
-    mx.set_default_device(mx.gpu)
-    a = mx.random.uniform(shape=(size, size), dtype=mx.float32)
-    b = mx.random.uniform(shape=(size, size), dtype=mx.float32)
-    mx.eval(a, b)
-    c = mx.matmul(a, b)
-    d = mx.matmul(a, c)
-    e = mx.matmul(b, c)
-    f = mx.sigmoid(d + e)
-    mx.eval(f)
-
-    print("CHILD: end")
-
-
-async def _collect_stream(
-    stream: Receiver[bytes],
-    output: bytearray,
-) -> None:
-    while True:
-        try:
-            output.extend(await stream.receive())
-        except EndOfStream:
-            return
-
-
-async def _collect_process_output(
-    process: AsyncProcess,
-) -> tuple[int, bytes, bytes]:
-    stdout = bytearray()
-    stderr = bytearray()
-    exitcodes: list[int] = []
-
-    async with create_task_group() as task_group:
-        task_group.start_soon(_collect_stream, process.stdout, stdout)
-        task_group.start_soon(_collect_stream, process.stderr, stderr)
-        exitcodes.append(await process.wait())
-
-    if not exitcodes:
-        raise RuntimeError("process exited without a return code")
-    return exitcodes[0], bytes(stdout), bytes(stderr)
-
-
-def _fd_identity(fd: int) -> tuple[int, int]:
-    fd_stat = os.fstat(fd)
-    return fd_stat.st_dev, fd_stat.st_ino
-
-
-def _fd_count() -> int | None:
-    for fd_dir in ("/proc/self/fd", "/dev/fd"):
-        with contextlib.suppress(OSError):
-            return len(os.listdir(fd_dir))
-    return None
-
-
-@contextlib.asynccontextmanager
-async def _started_process(process: AsyncProcess) -> AsyncIterator[None]:
-    async with create_task_group() as task_group:
-        await task_group.start(process.run)
-        try:
-            yield
-        finally:
-            await process.stop()
-
-
-async def _run_and_collect(
-    target: Callable[..., object] | None,
-    *,
-    args: tuple[object, ...] = (),
-    kwargs: dict[str, object] | None = None,
-) -> tuple[int, bytes, bytes]:
-    process = AsyncProcess(
-        target,
-        args=args,
-        kwargs=kwargs,
-    )
-    async with _started_process(process):
-        return await _collect_process_output(process)
-
-
-@pytest.mark.anyio
-async def test_spawn_process_captures_stdout_and_stderr_separately(
-    capfd: CaptureFixture[str],
-) -> None:
-    process = AsyncProcess(
-        _write_to_stdio,
-        args=("child",),
-        kwargs={"stderr_suffix": "error"},
-    )
-    async with _started_process(process):
-        exitcode, stdout_bytes, stderr_bytes = await _collect_process_output(process)
-
-    parent_output = capfd.readouterr()
-    stdout = stdout_bytes.decode("utf-8", errors="replace")
-    stderr = stderr_bytes.decode("utf-8", errors="replace")
-
-    assert exitcode == 0
-    assert "child: python stdout" in stdout
-    assert "child: fd stdout" in stdout
-    assert "child: python stderr error" in stderr
-    assert "child: fd stderr error" in stderr
-    assert "child:" not in parent_output.out
-    assert "child:" not in parent_output.err
-
-
-@pytest.mark.anyio
-async def test_process_with_no_target_exits_successfully() -> None:
-    exitcode, stdout, stderr = await _run_and_collect(None)
-
-    assert exitcode == 0
-    assert stdout == b""
-    assert stderr == b""
-
-
-@pytest.mark.anyio
-async def test_output_receivers_and_wait_are_safe_immediately_after_run_starts() -> (
-    None
-):
-    process = AsyncProcess(
-        _write_to_stdio,
-        args=("immediate",),
-        kwargs={"stderr_suffix": "error"},
-    )
-    result: tuple[int, bytes, bytes] | None = None
-
-    async with create_task_group() as task_group:
-        await task_group.start(process.run)
-        try:
-            result = await _collect_process_output(process)
-        finally:
-            await process.stop()
-
-    assert result is not None
-    exitcode, stdout, stderr = result
-    assert exitcode == 0
-    assert b"immediate: fd stdout\n" in stdout
-    assert b"immediate: fd stderr error\n" in stderr
-
-
-@pytest.mark.anyio
-async def test_stop_before_run_raises() -> None:
-    process = AsyncProcess(
-        _write_to_stdio,
-        args=("never",),
-        kwargs={"stderr_suffix": "run"},
-    )
-
-    assert not process.is_alive()
-    with pytest.raises(RuntimeError, match="process has not been started"):
-        await process.stop()
-
-
-@pytest.mark.anyio
-async def test_process_run_is_one_shot() -> None:
-    process = AsyncProcess(None)
-
-    await process.run()
-
-    with pytest.raises(RuntimeError, match="process has already been started"):
-        await process.run()
-
-
-@pytest.mark.anyio
-async def test_process_started_with_task_group_start_can_stop_immediately() -> None:
-    process = AsyncProcess(_sleep_forever)
-
-    async with create_task_group() as task_group:
-        await task_group.start(process.run)
-        assert process.is_alive()
-        with fail_after(2):
-            await process.stop()
-
-    assert not process.is_alive()
-
-
-@pytest.mark.anyio
-async def test_stdout_receiver_yields_bytes_chunks() -> None:
-    process = AsyncProcess(_write_large_output)
-
-    async with _started_process(process):
-        first_stdout = await process.stdout.receive()
-        exitcode, remaining_stdout, stderr = await _collect_process_output(process)
-
-    assert exitcode == 0
-    assert first_stdout + remaining_stdout == b"stdout-0123456789"
-    assert stderr == b"stderr-0123456789"
-
-
-@pytest.mark.anyio
-async def test_output_can_be_read_after_process_exits() -> None:
-    process = AsyncProcess(_write_large_output)
-
-    async with create_task_group() as task_group:
-        await task_group.start(process.run)
-        assert await process.wait() == 0
-
-    assert await process.stdout.receive() == b"stdout-0123456789"
-    assert await process.stderr.receive() == b"stderr-0123456789"
-    with pytest.raises(EndOfStream):
-        await process.stdout.receive()
-    with pytest.raises(EndOfStream):
-        await process.stderr.receive()
-
-
-@pytest.mark.anyio
-async def test_large_stdout_and_stderr_are_not_lost() -> None:
-    size = 1024 * 1024
-    exitcode, stdout, stderr = await _run_and_collect(
-        _write_large_exact_output,
-        args=(size,),
-    )
-
-    assert exitcode == 0
-    assert stdout == b"stdout:" + (b"x" * size)
-    assert stderr == b"stderr:" + (b"y" * size)
-
-
-@pytest.mark.anyio
-async def test_child_exception_traceback_is_captured_from_stderr() -> None:
-    process = AsyncProcess(_raise_after_stderr_write)
-
-    async with _started_process(process):
-        exitcode, _, stderr_bytes = await _collect_process_output(process)
-
-    assert exitcode == 1
-    stderr = stderr_bytes.decode("utf-8", errors="replace")
-    assert "stderr before exception" in stderr
-    assert "RuntimeError: child boom" in stderr
-
-
-@pytest.mark.anyio
-async def test_repeated_bad_children_do_not_pollute_or_replace_parent_stdio(
-    capfd: CaptureFixture[str],
-) -> None:
-    stdout_object = sys.stdout
-    stderr_object = sys.stderr
-    stdout_identity = _fd_identity(1)
-    stderr_identity = _fd_identity(2)
-
-    cases: tuple[tuple[Callable[..., object], tuple[object, ...]], ...] = (
-        (_raise_after_stderr_write, ()),
-        (_exit_after_stdio_write, ("exit-child", 17)),
-        (_abort_after_stdio_write, ("abort-child",)),
-    )
-
-    for iteration in range(3):
-        for target, args in cases:
-            exitcode, stdout, stderr = await _run_and_collect(
-                target,
-                args=args,
-            )
-
-            assert exitcode != 0
-            if target is _exit_after_stdio_write:
-                assert stdout == b"exit-child: stdout before _exit\n"
-                assert stderr == b"exit-child: stderr before _exit\n"
-            elif target is _abort_after_stdio_write:
-                assert b"abort-child: stdout before abort\n" in stdout
-                assert b"abort-child: stderr before abort\n" in stderr
-                assert exitcode == -signal.SIGABRT
-            else:
-                assert stdout == b""
-                assert b"stderr before exception\n" in stderr
-                assert b"RuntimeError: child boom" in stderr
-
-        print(f"parent stdout still works {iteration}")
-        print(f"parent stderr still works {iteration}", file=sys.stderr)
-
-    parent_output = capfd.readouterr()
-
-    assert sys.stdout is stdout_object
-    assert sys.stderr is stderr_object
-    assert _fd_identity(1) == stdout_identity
-    assert _fd_identity(2) == stderr_identity
-    assert "parent stdout still works 0" in parent_output.out
-    assert "parent stdout still works 2" in parent_output.out
-    assert "parent stderr still works 0" in parent_output.err
-    assert "parent stderr still works 2" in parent_output.err
-    assert "exit-child:" not in parent_output.out
-    assert "exit-child:" not in parent_output.err
-    assert "abort-child:" not in parent_output.out
-    assert "abort-child:" not in parent_output.err
-    assert "child boom" not in parent_output.err
-
-
-@pytest.mark.anyio
-async def test_child_can_close_stdio_without_corrupting_parent_stdio(
-    capfd: CaptureFixture[str],
-) -> None:
-    stdout_identity = _fd_identity(1)
-    stderr_identity = _fd_identity(2)
-
-    exitcode, stdout, stderr = await _run_and_collect(_close_stdio_and_exit)
-    os.write(1, b"parent stdout after child closed stdio\n")
-    os.write(2, b"parent stderr after child closed stdio\n")
-    parent_output = capfd.readouterr()
-
-    assert exitcode == 0
-    assert stdout == b""
-    assert stderr == b""
-    assert _fd_identity(1) == stdout_identity
-    assert _fd_identity(2) == stderr_identity
-    assert "parent stdout after child closed stdio" in parent_output.out
-    assert "parent stderr after child closed stdio" in parent_output.err
-
-
-@pytest.mark.anyio
-async def test_repeated_crashing_children_do_not_grow_parent_fd_table() -> None:
-    await _run_and_collect(_exit_after_stdio_write, args=("warmup", 23))
-    before = _fd_count()
-    if before is None:
-        pytest.skip("fd table count is not available on this platform")
-
-    for iteration in range(20):
-        exitcode, stdout, stderr = await _run_and_collect(
-            _exit_after_stdio_write,
-            args=(f"fd-child-{iteration}", 31),
-        )
-
-        assert exitcode == 31
-        assert stdout == f"fd-child-{iteration}: stdout before _exit\n".encode()
-        assert stderr == f"fd-child-{iteration}: stderr before _exit\n".encode()
-
-    after = _fd_count()
-    assert after is not None
-    assert after <= before + 2
-
-
-@pytest.mark.anyio
-async def test_stop_allows_child_to_exit_after_sigterm() -> None:
-    process = AsyncProcess(_exit_on_sigterm, args=(43,))
-
-    async with _started_process(process):
-        assert await process.stdout.receive() == b"sigterm-ready\n"
-
-        with fail_after(2):
-            await process.stop()
-
-    assert process.exitcode == 43
-
-
-@pytest.mark.anyio
-async def test_stop_retries_sigterm_before_sigkill(monkeypatch: MonkeyPatch) -> None:
-    monkeypatch.setattr(async_process, "_TERMINATE_GRACE_SECONDS", 0.01)
-    monkeypatch.setattr(async_process, "_TERMINATE_RETRY_GRACE_SECONDS", 0.01)
-    process = AsyncProcess(_exit_after_repeated_sigterm, args=(3, 44))
-
-    async with _started_process(process):
-        assert await process.stdout.receive() == b"sigterm-ready\n"
-
-        with fail_after(2):
-            await process.stop()
-
-    assert process.exitcode == 44
-
-
-@pytest.mark.anyio
-async def test_stop_escalates_to_sigkill_when_child_ignores_sigterm(
-    monkeypatch: MonkeyPatch,
-) -> None:
-    monkeypatch.setattr(async_process, "_TERMINATE_GRACE_SECONDS", 0.1)
-    monkeypatch.setattr(async_process, "_TERMINATE_RETRY_GRACE_SECONDS", 0.01)
-    process = AsyncProcess(_ignore_sigterm_forever)
-
-    async with _started_process(process):
-        assert await process.stdout.receive() == b"sigterm-ready\n"
-
-        with fail_after(3):
-            await process.stop()
-
-    assert process.exitcode == -signal.SIGKILL
-
-
-@pytest.mark.anyio
-async def test_process_can_use_mp_channel_with_global_spawn_context() -> None:
-    send, recv = mp_channel[str]()
-    process = AsyncProcess(_send_over_mp_channel, args=(send,))
-
-    async with _started_process(process):
-        with fail_after(2):
-            assert await recv.receive_async() == "hello from child"
-            assert await process.wait() == 0
-
-    with contextlib.suppress(Exception):
-        recv.close()
-
-
-@pytest.mark.anyio
-@pytest.mark.skip(reason="manual MLX OOM isolation check")
-async def test_death(capsys: CaptureFixture[str]) -> None:
-    with capsys.disabled():
-        process = AsyncProcess(_mlx_force_oom)
-        stdout = b""
-        stderr = b""
-        async with _started_process(process):
-            _, stdout, stderr = await _collect_process_output(process)
-
-        print("PARENT: done")
-
-        print("CHILD out:", stdout.decode("utf-8", errors="replace"))
-        print("CHILD err:", stderr.decode("utf-8", errors="replace"), "hello :)")

src/exo/utils/tests/test_daemon.py

@@ -1,168 +0,0 @@
-import contextlib
-import os
-from collections.abc import AsyncIterator
-
-import anyio
-import pytest
-from anyio import EndOfStream, create_task_group, fail_after
-
-from exo.utils.async_process import AsyncProcess
-from exo.utils.channels import MpReceiver, MpSender, Receiver, mp_channel
-from exo.utils.daemon import detach_stdio_to_devnull
-
-
-def _write_before_and_after_detach() -> None:
-    os.write(1, b"before stdout\n")
-    os.write(2, b"before stderr\n")
-    detach_stdio_to_devnull()
-    os.write(1, b"after stdout\n")
-    os.write(2, b"after stderr\n")
-
-
-def _write_grandchild_stdio(label: str) -> None:
-    os.write(1, f"{label} stdout\n".encode())
-    os.write(2, f"{label} stderr\n".encode())
-
-
-async def _spawn_grandchild_and_report(
-    result_sender: MpSender[tuple[int, bytes, bytes]],
-    label: str,
-) -> None:
-    result_sender.send(await _collect_spawned_child(label))
-    result_sender.close()
-
-
-async def _collect_spawned_child(label: str) -> tuple[int, bytes, bytes]:
-    process = AsyncProcess(_write_grandchild_stdio, args=(label,))
-    async with _started_process(process):
-        return await _collect_process_output(process)
-
-
-def _detach_stdio_then_spawn_captured_child(
-    result_sender: MpSender[tuple[int, bytes, bytes]],
-) -> None:
-    detach_stdio_to_devnull()
-    anyio.run(_spawn_grandchild_and_report, result_sender, "grandchild")
-
-
-def _detach_stdio_then_spawn_captured_children_sequentially(
-    result_sender: MpSender[list[tuple[int, bytes, bytes]]],
-) -> None:
-    async def run_children() -> list[tuple[int, bytes, bytes]]:
-        results: list[tuple[int, bytes, bytes]] = []
-        for index in range(5):
-            results.append(await _collect_spawned_child(f"grandchild-{index}"))
-        return results
-
-    detach_stdio_to_devnull()
-    result_sender.send(anyio.run(run_children))
-    result_sender.close()
-
-
-async def _collect_stream(stream: Receiver[bytes], output: bytearray) -> None:
-    while True:
-        try:
-            output.extend(await stream.receive())
-        except EndOfStream:
-            return
-
-
-async def _collect_process_output(
-    process: AsyncProcess,
-) -> tuple[int, bytes, bytes]:
-    stdout = bytearray()
-    stderr = bytearray()
-    exitcodes: list[int] = []
-
-    async with create_task_group() as collect_group:
-        collect_group.start_soon(_collect_stream, process.stdout, stdout)
-        collect_group.start_soon(_collect_stream, process.stderr, stderr)
-        exitcodes.append(await process.wait())
-
-    if not exitcodes:
-        raise RuntimeError("process exited without a return code")
-    return exitcodes[0], bytes(stdout), bytes(stderr)
-
-
-@contextlib.asynccontextmanager
-async def _started_process(process: AsyncProcess) -> AsyncIterator[None]:
-    async with create_task_group() as task_group:
-        await task_group.start(process.run)
-        try:
-            yield
-        finally:
-            await process.stop()
-
-
-async def _run_process_and_receive[T](
-    process: AsyncProcess,
-    recv: MpReceiver[T],
-    *,
-    timeout: float,
-) -> tuple[int, T]:
-    async with _started_process(process):
-        with fail_after(timeout):
-            result = await recv.receive_async()
-            exitcode = await process.wait()
-
-    return exitcode, result
-
-
-@pytest.mark.anyio
-async def test_detach_stdio_to_devnull_redirects_stdio_away_from_capture() -> None:
-    process = AsyncProcess(_write_before_and_after_detach)
-
-    async with _started_process(process):
-        exitcode, stdout, stderr = await _collect_process_output(process)
-
-    assert exitcode == 0
-    assert stdout == b"before stdout\n"
-    assert stderr == b"before stderr\n"
-
-
-@pytest.mark.anyio
-async def test_detached_stdio_process_can_spawn_and_capture_child_stdio() -> None:
-    send, recv = mp_channel[tuple[int, bytes, bytes]]()
-    process = AsyncProcess(_detach_stdio_then_spawn_captured_child, args=(send,))
-
-    try:
-        daemonized_parent_exitcode, result = await _run_process_and_receive(
-            process, recv, timeout=5
-        )
-    finally:
-        recv.close()
-
-    child_exitcode, child_stdout, child_stderr = result
-
-    assert daemonized_parent_exitcode == 0
-    assert child_exitcode == 0
-    assert child_stdout == b"grandchild stdout\n"
-    assert child_stderr == b"grandchild stderr\n"
-
-
-@pytest.mark.anyio
-async def test_detached_stdio_process_can_spawn_captured_children_sequentially() -> (
-    None
-):
-    send, recv = mp_channel[list[tuple[int, bytes, bytes]]]()
-    process = AsyncProcess(
-        _detach_stdio_then_spawn_captured_children_sequentially,
-        args=(send,),
-    )
-
-    try:
-        daemonized_parent_exitcode, results = await _run_process_and_receive(
-            process, recv, timeout=10
-        )
-    finally:
-        recv.close()
-
-    assert daemonized_parent_exitcode == 0
-    assert results == [
-        (
-            0,
-            f"grandchild-{index} stdout\n".encode(),
-            f"grandchild-{index} stderr\n".encode(),
-        )
-        for index in range(5)
-    ]

src/exo/utils/tests/test_pidfile.py

@@ -1,84 +0,0 @@
-from __future__ import annotations
-
-import gc
-import os
-import subprocess
-import sys
-import textwrap
-from pathlib import Path
-from typing import Final
-
-import pytest
-
-import exo.utils.pidfile as pidfile
-from exo.utils.pidfile import acquire_exo_pidfile
-
-_CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
-    """
-    import sys
-    from pathlib import Path
-    from unittest.mock import patch
-
-    import exo.utils.pidfile as pidfile
-    from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
-
-    with patch.object(pidfile, "EXO_PID_FILE", Path(sys.argv[1])):
-        try:
-            handle = acquire_exo_pidfile()
-        except PidfileLockError as exception:
-            print(str(exception))
-            raise SystemExit(73) from exception
-
-        del handle
-    """
-)
-
-
-def _use_pidfile_path(monkeypatch: pytest.MonkeyPatch, path: Path) -> None:
-    monkeypatch.setattr(pidfile, "EXO_PID_FILE", path)
-
-
-def _run_child_acquire_pidfile(path: Path) -> subprocess.CompletedProcess[str]:
-    return subprocess.run(
-        [sys.executable, "-c", _CHILD_ACQUIRE_PIDFILE_SCRIPT, str(path)],
-        check=False,
-        capture_output=True,
-        text=True,
-    )
-
-
-def test_acquire_exo_pidfile_writes_current_pid_and_removes_on_drop(
-    tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    path = tmp_path / "exo.pid"
-    _use_pidfile_path(monkeypatch, path)
-
-    handle = acquire_exo_pidfile()
-    assert path.read_text() == str(os.getpid())
-
-    del handle
-    gc.collect()
-
-    assert not path.exists()
-
-
-def test_acquire_exo_pidfile_rejects_second_process(
-    tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    path = tmp_path / "exo.pid"
-    _use_pidfile_path(monkeypatch, path)
-
-    handle = acquire_exo_pidfile()
-    try:
-        blocked_child = _run_child_acquire_pidfile(path)
-        assert blocked_child.returncode == 73
-        assert "Failed to acquire EXO pidfile" in blocked_child.stdout
-    finally:
-        del handle
-        gc.collect()
-
-    unblocked_child = _run_child_acquire_pidfile(path)
-    assert unblocked_child.returncode == 0
-    assert unblocked_child.stdout == ""

src/exo/utils/tests/test_power_sampler.py

@@ -111,36 +111,6 @@ async def test_empty_state() -> None:
     assert result.total_energy_joules == 0.0
 
 
-def test_trapezoidal_unit_dt_weighting() -> None:
-    """Pure unit test on the integration helper. Crafted samples where the
-    arithmetic mean is wildly wrong vs the time-weighted result."""
-    from exo.utils.power_sampler import trapezoidal_energy
-
-    # 5 s window. Power = 10 W for the first 4.9 s, then 100 W for the last 0.1 s.
-    # Three samples: t=0 W=10, t=4.9 W=10, t=5.0 W=100.
-    samples = [
-        (0.0, _make_profile(10.0)),
-        (4.9, _make_profile(10.0)),
-        (5.0, _make_profile(100.0)),
-    ]
-    energy = trapezoidal_energy(samples, elapsed=5.0)
-    # (10+10)/2 * 4.9 + (10+100)/2 * 0.1 = 49 + 5.5 = 54.5 J
-    assert abs(energy - 54.5) < 1e-9
-    avg = energy / 5.0  # 10.9 W
-    # Arithmetic mean of the three samples would be (10+10+100)/3 ≈ 40 W.
-    # Trapezoidal correctly weights each segment by its dt.
-    assert abs(avg - 10.9) < 1e-9
-
-
-def test_trapezoidal_unit_single_sample() -> None:
-    """One sample: no window to integrate over, so fall back to constant power
-    over the elapsed duration."""
-    from exo.utils.power_sampler import trapezoidal_energy
-
-    samples = [(0.0, _make_profile(42.0))]
-    assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
-
-
 async def test_result_stops_sampling() -> None:
     """Calling result() should stop the sampler's run loop."""
     state: dict[NodeId, SystemPerformanceProfile] = {

src/exo/worker/disaggregated/protocol.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from typing import BinaryIO, Literal
 
 import msgspec
@@ -23,7 +24,28 @@ class TensorBlob(msgspec.Struct):
     data: bytes
 
 
-class KVChunk(msgspec.Struct, tag="kv_chunk"):
+class _KVChunkHeader(msgspec.Struct, tag="kv_chunk"):
+    """Wire-side KV chunk metadata. Raw `keys` then `values` bytes follow on
+    the stream, lengths given by `keys_len` / `values_len`. Splitting them out
+    of the msgpack frame lets the producer pass tensor buffers via the buffer
+    protocol straight into the socket (one host-side memcpy total).
+    """
+
+    layer_idx: int
+    num_tokens: int
+    n_heads: int
+    head_dim: int
+    dtype: DType
+    keys_len: int
+    values_len: int
+
+
+@dataclass(frozen=True)
+class KVChunk:
+    """In-memory KV chunk reconstructed by `read_message` from
+    `_KVChunkHeader` + the raw bytes that follow on the wire.
+    """
+
     layer_idx: int
     num_tokens: int
     n_heads: int
@@ -51,10 +73,13 @@ class ErrorMessage(msgspec.Struct, tag="error"):
     message: str
 
 
+_WireMessage = _KVChunkHeader | ArraysState | Done | ErrorMessage
 Message = KVChunk | ArraysState | Done | ErrorMessage
 
 _msg_encoder = msgspec.msgpack.Encoder()
-_msg_decoder: msgspec.msgpack.Decoder[Message] = msgspec.msgpack.Decoder(Message)
+_msg_decoder: msgspec.msgpack.Decoder[_WireMessage] = msgspec.msgpack.Decoder(
+    _WireMessage
+)
 _header_encoder = msgspec.msgpack.Encoder()
 _header_decoder: msgspec.msgpack.Decoder[Header] = msgspec.msgpack.Decoder(Header)
 
@@ -99,7 +124,7 @@ def read_header(stream: BinaryIO) -> Header:
         raise ProtocolError(f"Bad header: {exc}") from exc
 
 
-def write_message(stream: BinaryIO, msg: Message) -> None:
+def write_message(stream: BinaryIO, msg: _WireMessage) -> None:
     write_frame(stream, _msg_encoder.encode(msg))
 
 
@@ -108,9 +133,22 @@ def read_message(stream: BinaryIO) -> Message | None:
     if not payload:
         return None
     try:
-        return _msg_decoder.decode(payload)
+        msg = _msg_decoder.decode(payload)
     except msgspec.DecodeError as exc:
         raise ProtocolError(f"Bad message: {exc}") from exc
+    if isinstance(msg, _KVChunkHeader):
+        keys = _read_exactly(stream, msg.keys_len)
+        values = _read_exactly(stream, msg.values_len)
+        return KVChunk(
+            layer_idx=msg.layer_idx,
+            num_tokens=msg.num_tokens,
+            n_heads=msg.n_heads,
+            head_dim=msg.head_dim,
+            dtype=msg.dtype,
+            keys=keys,
+            values=values,
+        )
+    return msg
 
 
 def write_kv_chunk(
@@ -121,21 +159,35 @@ def write_kv_chunk(
     n_heads: int,
     head_dim: int,
     dtype: DType,
-    keys: bytes,
-    values: bytes,
+    keys: "bytes | memoryview",
+    values: "bytes | memoryview",
 ) -> None:
-    write_message(
-        stream,
-        KVChunk(
+    """Stream KV chunk metadata + raw key/value bytes to the wire.
+
+    `keys` / `values` may be bytes-like (bytes, bytearray, memoryview) — the
+    raw payload is written directly to the buffered stream after the
+    msgpack-framed header, avoiding a memcpy through the msgpack encoder.
+    """
+    keys_len = len(keys)
+    values_len = len(values)
+    header_payload = _msg_encoder.encode(
+        _KVChunkHeader(
             layer_idx=layer_idx,
             num_tokens=num_tokens,
             n_heads=n_heads,
             head_dim=head_dim,
             dtype=dtype,
-            keys=keys,
-            values=values,
-        ),
+            keys_len=keys_len,
+            values_len=values_len,
+        )
     )
+    stream.write(len(header_payload).to_bytes(4, "big"))
+    stream.write(header_payload)
+    stream.write(keys)
+    stream.write(values)
+    # No per-chunk flush: the K/V payload is far larger than the
+    # BufferedWriter's internal buffer so it bypasses to the socket directly.
+    # The trailing `Done` frame's `write_frame` flushes once at the end.
 
 
 def write_arrays_state(

src/exo/worker/disaggregated/server.py

@@ -21,6 +21,7 @@ class PrefillRequest(msgspec.Struct):
     model_id: str = ""
     token_ids: list[int] = msgspec.field(default_factory=list)
     start_pos: int = 0
+    use_prefix_cache: bool = True
 
 
 _request_encoder = msgspec.msgpack.Encoder()
@@ -56,7 +57,10 @@ class _PrefillHandler(socketserver.StreamRequestHandler):
         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)
+        # 64MB send buffer: K/V chunks are ~33MB each; a small SNDBUF
+        # back-pressures the writer thread between chunks and serializes
+        # network with compute.
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 64 * 1024 * 1024)
 
     def handle(self) -> None:
         server = cast(PrefillServer, self.server)

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

@@ -143,7 +143,6 @@ class ImageEngine(Engine):
         Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]] | None
     ) = field(init=False, default=None)
     queue: deque[ImageTask] = field(init=False, default_factory=deque)
-    _cancelled_tasks: set[TaskId] = field(init=False, default_factory=set)
 
     def warmup(self) -> None:
         image = warmup_image_generator(model=self.image_model)
@@ -169,11 +168,7 @@ class ImageEngine(Engine):
             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 and _is_primary_output_node(self.shard_metadata)
-            else ()
-        )
+        return (resp,) if resp is not None else ()
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):

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

@@ -71,48 +71,14 @@ if TYPE_CHECKING:
 
 _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []
 
-_last_dist_op: mx.array | None = None
-
-
-def _link(out: mx.array) -> mx.array:
-    global _last_dist_op
-    if _last_dist_op is not None:
-        out = mx.depends(out, _last_dist_op)
-    mx.async_eval(out)
-    _last_dist_op = out
-    return out
-
-
-def send(x: mx.array, dst: int, group: mx.distributed.Group) -> mx.array:
-    return _link(mx.distributed.send(x, dst, group=group, stream=mx.Device(mx.cpu)))
-
-
-def recv_like(x: mx.array, src: int, group: mx.distributed.Group) -> mx.array:
-    return _link(
-        mx.distributed.recv_like(x, src, group=group, stream=mx.Device(mx.cpu))
-    )
-
-
-def all_gather(x: mx.array, group: mx.distributed.Group) -> mx.array:
-    return _link(mx.distributed.all_gather(x, group=group, stream=mx.Device(mx.cpu)))
-
 
 def flush_prefill_sends() -> None:
     for output, dst, group in _pending_prefill_sends:
-        send(output, dst, group)
+        sent = mx.distributed.send(output, dst, group=group)
+        mx.async_eval(sent)
     _pending_prefill_sends.clear()
 
 
-def reset_chain_head() -> None:
-    global _last_dist_op
-    _last_dist_op = None
-
-
-def dist_chain_head() -> mx.array:
-    assert _last_dist_op is not None
-    return _last_dist_op
-
-
 def clear_prefill_sends() -> None:
     # Discard pending sends (e.g. on cancellation).
     _pending_prefill_sends.clear()
@@ -165,7 +131,11 @@ class PipelineFirstLayer(CustomMlxLayer):
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         if self.r != 0:
-            x = recv_like(x, (self.r - 1), group=self.group)
+            # We want to avoid GPU timeout errors by evalling the distributed operation
+            # so that it stays on CPU, which does not have a timeout.
+            mx.eval(x)
+            x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
+            mx.eval(x)
         return self.original_layer(x, *args, **kwargs)
 
 
@@ -192,31 +162,36 @@ class PipelineLastLayer(CustomMlxLayer):
 
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
+        # Eval layer output to materialize it before send — this splits the graph
+        # so the send is isolated and the receiving rank's recv can complete.
+        mx.eval(output)
+
         if self.r != self.s - 1:
             if self.queue_sends:
                 _pending_prefill_sends.append(
                     (output, (self.r + 1) % self.s, self.group)
                 )
-                sent_output = output
             else:
-                sent_output = send(output, (self.r + 1) % self.s, self.group)
-                if cache is not None:
-                    # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
-                    # doesn't have .keys directly; access via first sub-cache.
-                    _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                    if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
-                        _cache.keys = mx.depends(_cache.keys, sent_output)  # type: ignore
-        else:
-            sent_output = output
+                output = mx.distributed.send(
+                    output, (self.r + 1) % self.s, group=self.group
+                )
+            if cache is not None:
+                # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
+                # doesn't have .keys directly; access via first sub-cache.
+                _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
+                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
+                    _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+            mx.eval(output)
+            if cache is not None and hasattr(_cache, "keys"):  # type: ignore
+                mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
-            all_gathered_output = all_gather(sent_output, group=self.group)[
+            output = mx.distributed.all_gather(output, group=self.group)[
                 -output.shape[0] :
             ]
-        else:
-            all_gathered_output = sent_output
+            mx.eval(output)
 
-        return all_gathered_output
+        return output
 
 
 def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
@@ -417,6 +392,64 @@ def pipeline_auto_parallel(
         "Expected a list of layers after auto-parallel initialisation"
     )
 
+    return patch_pipeline_model(model, group)
+
+
+def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
+    # Patch __call__ on the model's class
+    cls = model.__class__
+    original_call = cls.__call__  # type :ignore
+    call_signature = signature(original_call)  # type :ignore
+
+    def patched_call(
+        self: T,
+        *args: object,
+        **kwargs: object,
+    ) -> mx.array:
+        logits: mx.array = original_call(self, *args, **kwargs)  # type: ignore
+        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
+            "cache", None
+        )
+
+        # Add dependency to last cache entry to ensure distributed ops are evaluated
+        if cache is not None and len(cache) > 0:  # type: ignore
+            last = cache[-1]  # type: ignore
+            dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
+            if hasattr(dep_cache, "keys") and dep_cache.keys is not None:  # type: ignore
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
+
+        return logits
+
+    cls.__call__ = patched_call
+    return model
+
+
+def patch_tensor_model[T](model: T) -> T:
+    """Patch model's __call__ to ensure distributed ops sync during inference."""
+    cls = model.__class__
+    original_call = cls.__call__
+    call_signature = signature(original_call)
+
+    def patched_call(
+        self: T,
+        *args: object,
+        **kwargs: object,
+    ) -> mx.array:
+        logits: mx.array = original_call(self, *args, **kwargs)  # pyright: ignore[reportAny]
+        cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
+            "cache", None
+        )
+
+        # Add dependency to last cache entry to ensure distributed ops are evaluated
+        if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
+            last = cache[-1]  # pyright: ignore[reportAny]
+            dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
+            if hasattr(dep_cache, "keys"):  # type: ignore
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny]
+
+        return logits
+
+    cls.__call__ = patched_call
     return model
 
 
@@ -566,7 +599,7 @@ def tensor_auto_parallel(
         raise ValueError(f"Unsupported model type: {type(model)}")
 
     model = yield from tensor_parallel_sharding_strategy.shard_model(model)
-    return model
+    return patch_tensor_model(model)
 
 
 class TensorParallelShardingStrategy(ABC):
@@ -730,10 +763,8 @@ class ShardedMoE(CustomMlxLayer):
             x = sum_gradients(self.sharding_group)(x)
         y = self.original_layer.__call__(x)
         if self.sharding_group is not None:
-            z = mx.distributed.all_sum(y, group=self.sharding_group)
-        else:
-            z = y
-        return z
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+        return y
 
 
 class ShardedMoEV4(CustomMlxLayer):
@@ -749,10 +780,8 @@ class ShardedMoEV4(CustomMlxLayer):
             x = sum_gradients(self.sharding_group)(x)
         y = self._v4_inner(x, input_ids)
         if self.sharding_group is not None:
-            z = mx.distributed.all_sum(y, group=self.sharding_group)
-        else:
-            z = y
-        return z
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+        return y
 
 
 def _shard_quantized_rows(

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

@@ -34,7 +34,7 @@ class MlxBuilder(Builder):
     model_id: ModelId
     event_sender: MpSender[Event]
     cancel_receiver: MpReceiver[TaskId]
-    inference_model: Model | None = None
+    model: Model | None = None
     tokenizer: TokenizerWrapper | None = None
     group: mx.distributed.Group | None = None
     vision_processor: VisionProcessor | None = None
@@ -44,14 +44,14 @@ class MlxBuilder(Builder):
 
     def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
         (
-            self.inference_model,
+            self.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
+            del self.model
         with contextlib.suppress(NameError, AttributeError):
             del self.tokenizer
         with contextlib.suppress(NameError, AttributeError):
@@ -60,7 +60,7 @@ class MlxBuilder(Builder):
     def build(
         self,
     ) -> Engine:
-        assert self.inference_model
+        assert self.model
         assert self.tokenizer
 
         vision_processor = self.vision_processor
@@ -86,7 +86,7 @@ class MlxBuilder(Builder):
         if os.environ.get("EXO_NO_BATCH"):
             logger.info("using SequentialGenerator (batching disabled)")
             return SequentialGenerator(
-                model=self.inference_model,
+                model=self.model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,
@@ -100,7 +100,7 @@ class MlxBuilder(Builder):
         else:
             logger.info("using BatchGenerator")
             return BatchGenerator(
-                model=self.inference_model,
+                model=self.model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,

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

@@ -1,6 +1,5 @@
 import gc
 import os
-import time
 from copy import deepcopy
 from typing import TYPE_CHECKING
 
@@ -94,17 +93,6 @@ def copy_rotating_kv_cache(cache: RotatingKVCache) -> RotatingKVCache | None:
     return snap
 
 
-def copy_kv_cache(cache: KVCache) -> KVCache:
-    snap = KVCache()
-    snap.offset = cache.offset
-    if cache.keys is not None:
-        assert cache.values is not None
-        snap.keys = _detached_copy(cache.keys)
-        snap.values = _detached_copy(cache.values)
-        mx.eval(snap.keys, snap.values)
-    return snap
-
-
 def _copy_arrays_cache(ac: ArraysCache) -> ArraysCache:
     entries: list[mx.array | None] = []
     for entry in ac.cache:  # type: ignore[reportUnknownMemberType]
@@ -193,41 +181,6 @@ def copy_snapshot_entry(
             return _copy_v4_cache(entry)
 
 
-def copy_cache_entry(
-    entry: KVCache
-    | RotatingKVCache
-    | QuantizedKVCache
-    | ArraysCache
-    | CacheList
-    | DeepseekV4Cache,
-) -> (
-    KVCache
-    | RotatingKVCache
-    | QuantizedKVCache
-    | ArraysCache
-    | CacheList
-    | DeepseekV4Cache
-):
-    match entry:
-        case KVCache():
-            return copy_kv_cache(entry)
-        case RotatingKVCache():
-            snap = copy_rotating_kv_cache(entry)
-            return snap if snap is not None else deepcopy(entry)
-        case QuantizedKVCache():
-            return deepcopy(entry)
-        case ArraysCache():
-            return _copy_arrays_cache(entry)
-        case CacheList():
-            return _copy_cache_list(entry)
-        case DeepseekV4Cache():
-            return _copy_v4_cache(entry)
-
-
-def copy_kv_cache_state(cache: KVCacheType) -> KVCacheType:
-    return [copy_cache_entry(entry) for entry in cache]
-
-
 def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
     states: list[
         RotatingKVCache | ArraysCache | CacheList | DeepseekV4Cache | None
@@ -307,7 +260,7 @@ class KVPrefixCache:
         """Add a new cache entry. Evicts LRU entries if memory is high."""
         self._evict_if_needed()
         self.prompts.append(prompt_tokens)
-        self.caches.append(copy_kv_cache_state(cache))
+        self.caches.append(deepcopy(cache))
         self._snapshots.append(ssm_snapshots)
         self._media_regions.append(media_regions or [])
         self.prefill_tps.append(prefill_tps)
@@ -334,7 +287,7 @@ class KVPrefixCache:
             merged.extend(snapshots)
 
         self.prompts[index] = prompt_tokens
-        self.caches[index] = copy_kv_cache_state(cache)
+        self.caches[index] = deepcopy(cache)
         self._snapshots[index] = merged or None
         self._media_regions[index] = media_regions or []
         self.prefill_tps[index] = prefill_tps
@@ -423,7 +376,7 @@ class KVPrefixCache:
         if restore_snap is None and has_ssm:
             return make_kv_cache(model), prompt_tokens, None, False
 
-        prompt_cache = copy_kv_cache_state(self.caches[best_index])
+        prompt_cache = deepcopy(self.caches[best_index])
         tokens_to_trim = cached_length - restore_pos
         if tokens_to_trim > 0:
             trim_cache(prompt_cache, tokens_to_trim, restore_snap)
@@ -503,7 +456,6 @@ class KVPrefixCache:
             mx.clear_cache()
 
     def get_memory_used_percentage(self) -> float:
-        t0 = time.perf_counter()
         local_pressure: float = get_memory_used_percentage()
 
         if self._group is None:
@@ -512,14 +464,9 @@ class KVPrefixCache:
         all_pressure = mx.distributed.all_gather(
             mx.array([local_pressure], dtype=mx.float32),
             group=self._group,
-            stream=mx.Device(mx.cpu),
         )
         # .item() evals.
         max_pressure = float(mx.max(all_pressure).item())
-        elapsed_ms = (time.perf_counter() - t0) * 1000
-        logger.info(
-            f"get_memory_used_percentage took {elapsed_ms:.2f}ms (group size {self._group.size()})"
-        )
         return max_pressure
 
 

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

@@ -22,7 +22,6 @@ from exo.worker.disaggregated.protocol import (
     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,
@@ -90,6 +89,18 @@ def nhd_to_bhsd(t: mx.array) -> mx.array:
     return mx.expand_dims(mx.transpose(t, (1, 0, 2)), 0)
 
 
+def _rotating_to_temporal(buf: mx.array, idx: int, offset: int, keep: int) -> mx.array:
+    seq = int(buf.shape[2])
+    if idx == seq:
+        return buf
+    if idx < offset:
+        return mx.concatenate(
+            [buf[..., :keep, :], buf[..., idx:, :], buf[..., keep:idx, :]],
+            axis=2,
+        )
+    return buf[..., :idx, :]
+
+
 def send_mlx_kv_cache(
     stream: BinaryIO,
     caches: KVCacheType,
@@ -103,7 +114,7 @@ def send_mlx_kv_cache(
         match c:
             case QuantizedKVCache() | CacheList() | DeepseekV4Cache():
                 raise NotImplementedError
-            case KVCache() | RotatingKVCache():
+            case KVCache():
                 keys = c.keys
                 values = c.values
                 if keys is None or values is None:
@@ -132,11 +143,39 @@ def send_mlx_kv_cache(
                     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
+                tokens_sent = max(tokens_sent, num_tokens)
+            case RotatingKVCache():
+                keys = c.keys
+                values = c.values
+                if keys is None or values is None:
+                    continue
+                offset = int(c.offset)
+                if offset <= 0:
+                    continue
+                idx = int(c._idx)
+                keep = int(c.keep)
+                with mx.stream(mx.Device(mx.cpu)):
+                    k_temporal = _rotating_to_temporal(keys, idx, offset, keep)
+                    v_temporal = _rotating_to_temporal(values, idx, offset, keep)
+                    k = mx.array(k_temporal)
+                    v = mx.array(v_temporal)
+                    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),
+                )
+                tokens_sent = max(tokens_sent, offset)
             case ArraysCache():
                 blobs: list[TensorBlob] = []
                 for a in c.state:

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

@@ -79,6 +79,7 @@ def remote_prefill_fetch(
         result = PrefillResult(header=header)
         kv_by_layer: dict[int, list[KVChunk]] = defaultdict(list)
         chunks_received = 0
+        done_seen = False
 
         while True:
             msg = read_message(stream)
@@ -93,10 +94,17 @@ def remote_prefill_fetch(
                 result.arrays[msg.layer_idx] = msg.arrays
             elif isinstance(msg, Done):
                 result.total_tokens = msg.total_tokens
+                done_seen = True
                 break
             else:
                 raise RuntimeError(f"Prefill server error [{msg.code}]: {msg.message}")
 
+        if not done_seen:
+            raise ConnectionError(
+                "Prefill server closed before Done frame "
+                f"(received {chunks_received} kv chunks, {len(result.arrays)} arrays)"
+            )
+
         result.kv_chunks = dict(kv_by_layer)
         return result
     finally:

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

@@ -49,7 +49,6 @@ 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,
-    mx_barrier
 )
 from exo.worker.engines.mlx.vision import (
     MediaRegion,
@@ -106,7 +105,7 @@ class ExoBatchGenerator:
         self._mlx_gen = MlxBatchGenerator(
             model=self.model,
             stop_tokens=[[t] for t in eos_ids_from_tokenizer(self.tokenizer)],
-            prefill_step_size=1,
+            prefill_step_size=4096,
         )
         self._step_count = 0
 
@@ -216,14 +215,18 @@ class ExoBatchGenerator:
         with vision_ctx:
             if use_remote and task_params.prefill_endpoint is not None:
                 try:
+                    # Send full prompt; producer's vLLM APC handles the prefix
+                    # match. `start_pos` aligns the writer's skip_tokens with
+                    # the consumer's locally-cached prefix.
                     _prefill_tps, _prefill_tokens, cache_snapshots = remote_prefill(
-                        prompt_tokens[:-1],
+                        all_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,
+                        use_prefix_cache=not is_bench or task_params.use_prefix_cache,
                     )
                     remote_prefilled = True
                 except Exception:
@@ -232,7 +235,6 @@ class ExoBatchGenerator:
                     )
 
             if not remote_prefilled:
-                mx_barrier(self.group)
                 _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
                     self.model,
                     self.tokenizer,
@@ -326,8 +328,6 @@ class ExoBatchGenerator:
             media_regions=media_regions,
         )
 
-        mx_barrier(self.group)
-
         return uid
 
     def step(self) -> list[tuple[int, GenerationResponse]]:

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

@@ -35,9 +35,7 @@ from exo.worker.engines.mlx.auto_parallel import (
     PipelineFirstLayer,
     PipelineLastLayer,
     clear_prefill_sends,
-    dist_chain_head,
     flush_prefill_sends,
-    reset_chain_head,
     set_pipeline_prefill,
     set_pipeline_queue_sends,
 )
@@ -194,8 +192,7 @@ def pipeline_parallel_prefill(
     This function is designed to match mlx_lm's stream_generate exactly in terms of
     side effects (given the same prefill step size)
     """
-    world_size = group.size()
-    prefill_step_size = prefill_step_size // min(4, world_size)
+    prefill_step_size = prefill_step_size // min(4, group.size())
 
     quantize_cache_fn: Callable[..., None] = functools.partial(
         maybe_quantize_kv_cache,
@@ -206,18 +203,14 @@ def pipeline_parallel_prefill(
 
     _prompt_cache: KVCacheType = prompt_cache
     rank = group.rank()
+    world_size = group.size()
 
-    # Build list of real prompt chunk sizes.
-    # For pipeline parallel to overlap stages, we need at least world_size real chunks;
-    # otherwise a single-chunk prompt serializes across ranks (each rank's recv blocks).
+    # Build list of real prompt chunk sizes
     total = len(prompt)
-    remaining = total - 1
-    chunk_size = min(
-        prefill_step_size, max(1, (remaining + world_size - 1) // world_size)
-    )
     real_chunk_sizes: list[int] = []
+    remaining = total - 1
     while remaining:
-        n = min(chunk_size, remaining)
+        n = min(prefill_step_size, remaining)
         real_chunk_sizes.append(n)
         remaining -= n
     n_real = len(real_chunk_sizes)
@@ -239,77 +232,43 @@ def pipeline_parallel_prefill(
 
     try:
         with mx.stream(generation_stream):
-            t_leading = time.perf_counter()
             for _ in range(n_leading):
                 if distributed_prompt_progress_callback is not None:
                     distributed_prompt_progress_callback()
-            logger.info(
-                f"[R{rank}] leading dummies ({n_leading}) took {(time.perf_counter() - t_leading) * 1000:.1f}ms"
-            )
 
             for i in range(n_real):
-                t_iter = time.perf_counter()
                 chunk_size = real_chunk_sizes[i]
-                t_model = time.perf_counter()
                 model(
                     prompt[processed : processed + chunk_size][None],
                     cache=_prompt_cache,
                 )
                 quantize_cache_fn(_prompt_cache)
-                t_after_model = time.perf_counter()
                 processed += chunk_size
 
                 if distributed_prompt_progress_callback is not None:
                     distributed_prompt_progress_callback()
-                t_after_cb = time.perf_counter()
 
                 flush_prefill_sends()
-                t_after_flush = time.perf_counter()
-                mx.eval([c.state for c in _prompt_cache])  # type: ignore
-                t_after_eval = time.perf_counter()
 
                 prompt_progress_callback(processed, total)
-                logger.info(
-                    f"[R{rank}] iter {i}/{n_real} ({chunk_size} tok): "
-                    f"model+quant {(t_after_model - t_model) * 1000:.1f}ms, "
-                    f"cb {(t_after_cb - t_after_model) * 1000:.1f}ms, "
-                    f"flush {(t_after_flush - t_after_cb) * 1000:.1f}ms, "
-                    f"eval {(t_after_eval - t_after_flush) * 1000:.1f}ms, "
-                    f"total {(time.perf_counter() - t_iter) * 1000:.1f}ms"
-                )
 
-            t_trailing = time.perf_counter()
             for _ in range(n_trailing):
                 if distributed_prompt_progress_callback is not None:
                     distributed_prompt_progress_callback()
-            logger.info(
-                f"[R{rank}] trailing dummies ({n_trailing}) took {(time.perf_counter() - t_trailing) * 1000:.1f}ms"
-            )
 
     finally:
         clear_prefill_sends()
 
     # Post-loop: process remaining 1 token + add +1 entry to match stream_generate.
-    for j in range(2):
-        t_post = time.perf_counter()
+    for _ in range(2):
         with mx.stream(generation_stream):
             model(prompt[-1:][None], cache=_prompt_cache)
             quantize_cache_fn(_prompt_cache)
         flush_prefill_sends()
-        logger.info(
-            f"[R{rank}] post-loop iter {j} took {(time.perf_counter() - t_post) * 1000:.1f}ms"
-        )
 
     assert _prompt_cache is not None
-    t_final = time.perf_counter()
     with mx.stream(generation_stream):
-        mx.eval([c.state for c in _prompt_cache], dist_chain_head())  # type: ignore
-    mx.synchronize(mx.default_stream(mx.Device(mx.cpu)))
-    mx.synchronize(generation_stream)
-    reset_chain_head()
-    logger.info(
-        f"[R{rank}] final eval+chain drain took {(time.perf_counter() - t_final) * 1000:.1f}ms"
-    )
+        mx.eval([c.state for c in _prompt_cache])  # type: ignore
 
     # Final callback matching generate_step
     prompt_progress_callback(total, total)
@@ -343,8 +302,6 @@ def prefill(
         return 0.0, 0, []
 
     logger.debug(f"Prefilling {num_tokens} tokens...")
-
-    reset_chain_head()
     start_time = time.perf_counter()
     has_ssm = has_non_kv_caches(cache)
     snapshots: list[CacheSnapshot] = []
@@ -377,10 +334,7 @@ def prefill(
     prefill_step_size = 4096
 
     try:
-        if is_pipeline:
-            logger.info(
-                f"prefill path: pipeline_parallel_prefill ({num_tokens} tokens)"
-            )
+        if is_pipeline and num_tokens >= prefill_step_size:
             set_pipeline_queue_sends(model, queue_sends=True)
             assert group is not None, "Pipeline prefill requires a distributed group"
             pipeline_parallel_prefill(
@@ -395,10 +349,6 @@ def prefill(
                 group=group,
             )
         else:
-            logger.info(
-                f"prefill path: stream_generate ({num_tokens} tokens, is_pipeline={is_pipeline})"
-            )
-            t_stream = time.perf_counter()
             # Use max_tokens=1 because max_tokens=0 does not work.
             # We just throw away the generated token - we only care about filling the cache
             for _ in stream_generate(
@@ -414,9 +364,6 @@ def prefill(
                 prompt_progress_callback=combined_progress_callback,
             ):
                 break  # Stop after first iteration - cache is now filled
-            logger.info(
-                f"stream_generate prefill took {(time.perf_counter() - t_stream) * 1000:.1f}ms"
-            )
     except PrefillCancelled:
         set_pipeline_queue_sends(model, queue_sends=False)
         set_pipeline_prefill(model, is_prefill=False)
@@ -596,7 +543,6 @@ def mlx_generate(
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
-    reset_chain_head()
     # TODO: Randomise task seed and set in taskparams, instead of hard coding as 42.
     seed = task.seed or 42
     mx.random.seed(seed)
@@ -702,14 +648,20 @@ def mlx_generate(
     with maybe_vision_ctx:
         if use_remote and task.prefill_endpoint is not None:
             try:
+                # Send the FULL prompt to the producer (not the cache-stripped
+                # suffix). vLLM's APC handles the prefix match internally;
+                # `start_pos` tells our extractor / wire writer how much of the
+                # producer-side capture corresponds to tokens the consumer
+                # already has, so the writer's skip_tokens math aligns.
                 prefill_tps, prefill_tokens, ssm_snapshots_list = remote_prefill(
-                    prompt_tokens[:-1],
+                    all_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,
+                    use_prefix_cache=not is_bench or task.use_prefix_cache,
                 )
                 remote_prefilled = True
             except Exception:
@@ -770,7 +722,6 @@ def mlx_generate(
     usage: Usage | None = None
     logger.info("Starting decode")
     mx_barrier(group)
-    reset_chain_head()
 
     for completion_tokens, out in enumerate(
         stream_generate(

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

@@ -25,23 +25,25 @@ def remote_prefill(
     request_id: str,
     model_id: str,
     start_pos: int = 0,
+    use_prefix_cache: bool = True,
 ) -> tuple[float, int, list[CacheSnapshot]]:
     t0 = time.perf_counter()
     total_prompt_tokens = int(prompt_tokens.shape[0])
     num_layers: int = 0
+    tokens_received_total: 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
+    def _on_chunk(chunk: KVChunk, chunks_received: int) -> None:
+        nonlocal num_layers, tokens_received_total
+        tokens_received_total += chunk.num_tokens
         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),
+                min(tokens_received_total // num_layers, total_prompt_tokens),
                 total_prompt_tokens,
             )
 
@@ -50,6 +52,7 @@ def remote_prefill(
         token_ids=cast(list[int], prompt_tokens.tolist()),
         start_pos=start_pos,
         request_id=request_id,
+        use_prefix_cache=use_prefix_cache,
     )
     result = remote_prefill_fetch(
         endpoint, request, on_header=_on_header, on_kv_chunk=_on_chunk
@@ -61,6 +64,28 @@ def remote_prefill(
     t_done = time.perf_counter()
 
     num_tokens = final_offset - start_pos
+    # The producer strips the last 2 tokens of the prompt (consumer warm-starts
+    # decode from those locally). Anything within `producer_strip` of the full
+    # suffix is the expected outcome, not a bug.
+    producer_strip = 2 if total_prompt_tokens > 2 else 0
+    expected_min = max(0, total_prompt_tokens - start_pos - producer_strip)
+    expected_max = max(0, total_prompt_tokens - start_pos)
+    if num_tokens <= 0:
+        raise RuntimeError(
+            f"Remote prefill returned no KV (start_pos={start_pos}, "
+            f"final_offset={final_offset}, expected={expected_min}, "
+            f"transfer={(t_received - t0) * 1000:.0f}ms)"
+        )
+    if num_tokens < expected_min:
+        logger.warning(
+            f"Remote prefill returned {num_tokens} tokens, expected at least "
+            f"{expected_min} (start_pos={start_pos}, final_offset={final_offset})"
+        )
+    elif num_tokens > expected_max:
+        logger.warning(
+            f"Remote prefill returned {num_tokens} tokens, expected at most "
+            f"{expected_max} (start_pos={start_pos}, final_offset={final_offset})"
+        )
     tps = num_tokens / max(t_done - t0, 0.001)
 
     logger.info(

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

@@ -50,6 +50,7 @@ from exo.shared.types.worker.instances import (
     BoundInstance,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.runner_response import ModelLoadingResponse
 from exo.shared.types.worker.shards import (
@@ -115,8 +116,7 @@ def mlx_distributed_init(
 
                 os.environ["MLX_HOSTFILE"] = coordination_file
                 os.environ["MLX_RANK"] = str(rank)
-                # os.environ["MLX_RING_VERBOSE"] = "1"  # NOTE: we don't use it enough to care (turn on again if need to)
-
+                os.environ["MLX_RING_VERBOSE"] = "1"
                 group = mx.distributed.init(backend="ring", strict=True)
 
             case MlxJacclInstance(
@@ -141,6 +141,8 @@ def mlx_distributed_init(
                 os.environ["MLX_RANK"] = str(rank)
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
+            case VllmInstance():
+                raise ValueError("loaded VllmInstance in MLX engine")
 
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 
@@ -311,11 +313,12 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
     model_id_lower = model_id.lower()
     if "kimi-k2" in model_id_lower:
         return [163586]
-    elif "glm-5" in model_id_lower:
+    elif "glm-5" in model_id_lower or "glm-4.7" in model_id_lower:
+        # For GLM-5 and GLM-4.7
         # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|>
         return [154820, 154827, 154829]
     elif "glm" in model_id_lower:
-        # For GLM-4.7 and older
+        # For GLM-4.5 and older
         return [151336, 151329, 151338]
     elif "gpt-oss" in model_id_lower:
         return [200002, 200012]

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

@@ -0,0 +1,88 @@
+import contextlib
+import os
+from collections.abc import Generator
+from dataclasses import dataclass
+
+from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
+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.engines.vllm.engine import VllmEngine
+from exo.worker.engines.vllm.generator import VllmBatchEngine, load_vllm_engine
+from exo.worker.runner.bootstrap import logger
+
+
+@dataclass
+class VllmBuilder(Builder):
+    model_id: ModelId
+    event_sender: MpSender[Event]
+    cancel_receiver: MpReceiver[TaskId]
+
+    def connect(self, bound_instance: BoundInstance) -> None:
+        raise NotImplementedError(
+            "Multiple node VLLM instances are not supported at the moment!"
+        )
+
+    def load(
+        self,
+        bound_instance: BoundInstance,
+    ) -> Generator[ModelLoadingResponse]:
+        from exo.worker.engines.vllm.kv_connector import (
+            ExoKVProducerConnector,
+            _patch_gdn_capture,
+            _patch_vllm_for_connector,
+        )
+
+        # Apply bypass patches before vLLM init reads its connector registry
+        # and the unifier touches hybrid kv-cache specs.
+        _patch_vllm_for_connector(ExoKVProducerConnector)
+        _patch_gdn_capture()
+
+        kv_connector_cls: type[object] | None = ExoKVProducerConnector
+        # overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+
+        def on_layer_loaded(loaded: int, total: int) -> None:
+            pass
+
+        self._bound_runner_id = bound_instance.bound_runner_id
+        self._engine, self._tool_parser = load_vllm_engine(
+            model_id=self.model_id,
+            trust_remote_code=bound_instance.bound_shard.model_card.trust_remote_code,
+            n_layers=bound_instance.bound_shard.model_card.n_layers,
+            on_layer_loaded=on_layer_loaded,
+            kv_connector_cls=kv_connector_cls,
+        )
+        return
+        yield
+
+    def build(self) -> Engine:
+        gen = VllmBatchEngine(
+            engine=self._engine,
+            model_id=self.model_id,
+        )
+        try:
+            max_concurrent = (
+                1
+                if bool(os.getenv("EXO_NO_BATCH", False))
+                else EXO_MAX_CONCURRENT_REQUESTS
+            )
+        except Exception:
+            max_concurrent = EXO_MAX_CONCURRENT_REQUESTS
+
+        logger.info(f"using VllmEngine (max_concurrent={max_concurrent})")
+        return VllmEngine(
+            tool_parser=self._tool_parser,
+            model_id=self.model_id,
+            cancel_receiver=self.cancel_receiver,
+            event_sender=self.event_sender,
+            _gen=gen,
+            max_concurrent_requests=max_concurrent,
+        )
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self._engine, self._tool_parser

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

@@ -0,0 +1,238 @@
+"""vLLM-side disaggregation adapter.
+
+Mirrors `engines/mlx/disaggregated/adapter.py` for the vLLM engine: owns
+torch dtype ↔ wire dtype, byte (de)serialization, layout conversion (vLLM's
+paged block storage → NHD per-token), and the wire-write helpers used by
+the producer connector + serve_prefill flow.
+
+Wire format is `engines/.../disaggregated/protocol.py` (msgpack), shared with
+the MLX side.
+"""
+
+import os
+from typing import BinaryIO
+
+import torch
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from exo.worker.disaggregated.protocol import (
+    DType,
+    Header,
+    TensorBlob,
+    write_arrays_state,
+    write_done,
+    write_header,
+    write_kv_chunk,
+)
+
+_TORCH_TO_WIRE: dict[torch.dtype, DType] = {
+    torch.bfloat16: "bfloat16",
+    torch.float16: "float16",
+    torch.float32: "float32",
+}
+_WIRE_TO_TORCH: dict[DType, torch.dtype] = {v: k for k, v in _TORCH_TO_WIRE.items()}
+
+
+def torch_dtype_to_wire(dtype: torch.dtype) -> DType:
+    if dtype not in _TORCH_TO_WIRE:
+        raise ValueError(f"Unsupported torch dtype on wire: {dtype}")
+    return _TORCH_TO_WIRE[dtype]
+
+
+def wire_to_torch_dtype(dtype: DType) -> torch.dtype:
+    return _WIRE_TO_TORCH[dtype]
+
+
+def tensor_to_wire_bytes(t: torch.Tensor) -> bytes:
+    """Serialize an NHD-laid-out tensor to wire bytes.
+
+    bfloat16 has no native numpy dtype — bitcast through uint16.
+    """
+    t = t.detach().contiguous().cpu()
+    if t.dtype == torch.bfloat16:
+        return bytes(t.view(torch.uint16).numpy().tobytes())
+    return bytes(t.numpy().tobytes())
+
+
+def to_nhd(t: torch.Tensor) -> torch.Tensor:
+    """Permute HND → NHD when vLLM's KV cache layout is HND."""
+    if os.environ.get("VLLM_KV_CACHE_LAYOUT", "HND") == "HND" and t.dim() == 3:
+        return t.permute(1, 0, 2)
+    return t
+
+
+def to_bf16(t: torch.Tensor) -> torch.Tensor:
+    """Coerce to bfloat16, dequantizing fp8 / uint8-encoded fp8 if needed."""
+    if t.dtype == torch.uint8:
+        t = t.view(torch.float8_e4m3fn)
+    if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        return t.to(torch.float32).to(torch.bfloat16)
+    if t.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return t
+    return t.to(torch.bfloat16)
+
+
+def extract_kv_via_slot_mapping(
+    kv_layer: torch.Tensor, slot_mapping: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pull (keys, values) for the fresh tokens of one layer using slot_mapping.
+
+    `kv_layer` is vLLM's per-layer paged storage. Layout depends on attention
+    backend: either `[2, num_blocks, block_size, H, D]` or `[num_blocks, 2,
+    block_size, H, D]`. NHD is enforced via `VLLM_KV_CACHE_LAYOUT=NHD`.
+    `slot_mapping` is the per-token slot index; entries `< 0` are padding.
+
+    Returned tensors stay on the GPU — the D2H copy is deferred to the
+    writer thread (`tensor_to_wire_bytes` calls `.cpu()`) so it doesn't
+    block forward of subsequent layers.
+    """
+    if kv_layer.shape[0] == 2:
+        k_all = to_nhd(kv_layer[0])
+        v_all = to_nhd(kv_layer[1])
+    else:
+        k_all = to_nhd(kv_layer[:, 0])
+        v_all = to_nhd(kv_layer[:, 1])
+    k_flat = k_all.reshape(-1, *k_all.shape[-2:])
+    v_flat = v_all.reshape(-1, *v_all.shape[-2:])
+    valid = slot_mapping >= 0
+    safe_sm = slot_mapping.clamp(min=0)
+    keys = to_bf16(k_flat[safe_sm][valid])
+    values = to_bf16(v_flat[safe_sm][valid])
+    return keys, values
+
+
+def write_kv_layer_chunk(
+    wfile: BinaryIO,
+    layer_idx: int,
+    keys: torch.Tensor,
+    values: torch.Tensor,
+) -> None:
+    """Serialize one layer's NHD-shaped K/V to a `KVChunk` on the wire."""
+    if keys.dim() == 4:
+        keys = keys.reshape(-1, keys.shape[-2], keys.shape[-1])
+        values = values.reshape(-1, values.shape[-2], values.shape[-1])
+    num_tokens = int(keys.shape[0])
+    n_heads = int(keys.shape[1])
+    head_dim = int(keys.shape[2])
+    write_kv_chunk(
+        wfile,
+        layer_idx=layer_idx,
+        num_tokens=num_tokens,
+        n_heads=n_heads,
+        head_dim=head_dim,
+        dtype=torch_dtype_to_wire(keys.dtype),
+        keys=tensor_to_wire_bytes(keys),
+        values=tensor_to_wire_bytes(values),
+    )
+
+
+def arrays_to_blobs(arrays: list[torch.Tensor]) -> list[TensorBlob]:
+    """Convert torch tensors (CPU or GPU) to wire-ready `TensorBlob`s."""
+    return [
+        TensorBlob(
+            dtype=torch_dtype_to_wire(arr.dtype),
+            shape=tuple(int(d) for d in arr.shape),
+            data=tensor_to_wire_bytes(arr),
+        )
+        for arr in arrays
+    ]
+
+
+def write_layer_arrays_blobs(
+    wfile: BinaryIO,
+    layer_idx: int,
+    blobs: list[TensorBlob],
+) -> None:
+    write_arrays_state(wfile, layer_idx, blobs)
+
+
+def write_layer_arrays(
+    wfile: BinaryIO,
+    layer_idx: int,
+    arrays: list[torch.Tensor],
+) -> None:
+    """Serialize a layer's auxiliary state (SSM/conv) as `ArraysState`."""
+    write_layer_arrays_blobs(wfile, layer_idx, arrays_to_blobs(arrays))
+
+
+def write_prefill_header(
+    wfile: BinaryIO,
+    *,
+    request_id: str,
+    model_id: str,
+    num_layers: int,
+    dtype: DType = "bfloat16",
+    start_pos: int = 0,
+) -> None:
+    write_header(
+        wfile,
+        Header(
+            request_id=request_id,
+            model_id=model_id,
+            num_layers=num_layers,
+            dtype=dtype,
+            start_pos=start_pos,
+        ),
+    )
+
+
+def write_prefill_done(wfile: BinaryIO, total_tokens: int) -> None:
+    write_done(wfile, total_tokens)
+
+
+def build_layer_to_group(kv_cache_config: KVCacheConfig) -> list[int]:
+    """Map each layer index (model_runner.kv_caches order) to its kv_cache group.
+
+    vLLM's hybrid models split layers across multiple KV cache groups (e.g.
+    full attention vs sliding-window attention). `request_finished_all_groups`
+    returns block_ids per group; we need this map to look up the right group
+    when reading a layer's blocks.
+    """
+    group_lookup: dict[str, int] = {}
+    for group_idx, group_spec in enumerate(kv_cache_config.kv_cache_groups):
+        for layer_name in group_spec.layer_names:
+            group_lookup[layer_name] = group_idx
+
+    layer_to_group: list[int] = []
+    for tensor_spec in kv_cache_config.kv_cache_tensors:
+        for name in tensor_spec.shared_by:
+            layer_to_group.append(group_lookup[name])
+    return layer_to_group
+
+
+def gather_layer_kv_from_blocks(
+    layer_kv: torch.Tensor,
+    block_ids: list[int],
+    num_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Read K and V for `num_tokens` from a layer's paged block storage.
+
+    Captures APC-cached blocks identically to freshly-computed ones — the
+    block pool doesn't distinguish.
+
+    `layer_kv` shapes (NHD, set via `VLLM_KV_CACHE_LAYOUT=NHD`):
+        - `[2, num_pool_blocks, block_size, n_kv_heads, head_dim]`, or
+        - `[num_pool_blocks, 2, block_size, n_kv_heads, head_dim]`.
+    Returns NHD-shaped K and V of shape `[num_tokens, n_kv_heads, head_dim]`
+    on the same CUDA device as `layer_kv`. The caller is responsible for
+    issuing the D2H copy on a side stream so the scheduler thread isn't
+    blocked.
+    """
+    if not block_ids:
+        return torch.empty(0, device=layer_kv.device), torch.empty(
+            0, device=layer_kv.device
+        )
+    block_idx_tensor = torch.tensor(block_ids, dtype=torch.long, device=layer_kv.device)
+    if layer_kv.shape[0] == 2:
+        # [2, blocks, block, H, D]
+        gathered_k = layer_kv[0][block_idx_tensor]
+        gathered_v = layer_kv[1][block_idx_tensor]
+    else:
+        # [blocks, 2, block, H, D]
+        gathered = layer_kv[block_idx_tensor]
+        gathered_k = gathered[:, 0]
+        gathered_v = gathered[:, 1]
+    # gathered_k/v: [num_blocks, block_size, H, D]. Concat blocks along seq.
+    keys = gathered_k.reshape(-1, *gathered_k.shape[-2:])[:num_tokens]
+    values = gathered_v.reshape(-1, *gathered_v.shape[-2:])[:num_tokens]
+    return to_bf16(keys), to_bf16(values)

src/exo/worker/engines/vllm/disaggregated/tests/conftest.py

@@ -0,0 +1,10 @@
+import pytest
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--model-id",
+        action="store",
+        default=None,
+        help="HuggingFace-style model id (e.g. Qwen/Qwen3-0.6B) — must be downloaded",
+    )

src/exo/worker/engines/vllm/disaggregated/tests/test_serve_prefill_integration.py

@@ -0,0 +1,201 @@
+"""End-to-end test for VllmEngine.serve_prefill.
+
+Boots a real vLLM engine with a small model, calls serve_prefill twice in a
+row against an in-memory wire buffer, and verifies both runs produce a
+well-formed stream (header -> KV chunks -> Done).
+
+The second run is the regression case: with vLLM APC enabled this would hit
+the chunked-prefill + APC + custom kv-connector CUDA assert
+(`vectorized_gather_kernel: ind >= ind_dim_size`) and the server would close
+the socket without a Done frame. With APC disabled at engine creation time
+each request runs a full forward pass and the stream is well-formed.
+
+Run on Spark (gx10-de89):
+    cd /home/larry/exo
+    uv run pytest -q -s -m "" \\
+        src/exo/worker/engines/vllm/disaggregated/tests/test_serve_prefill_integration.py \\
+        --model-id Qwen/Qwen3-0.6B
+
+The test is gated on `--model-id` being passed; the model must already be
+present at `EXO_DEFAULT_MODELS_DIR/<id-with-/-as--->` (the standard exo
+download layout). On machines without CUDA / vLLM the test is skipped.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+from collections.abc import Iterator
+from typing import cast
+
+import pytest
+
+from exo.shared.types.common import ModelId
+from exo.worker.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    ErrorMessage,
+    KVChunk,
+    read_header,
+    read_message,
+)
+from exo.worker.disaggregated.server import PrefillRequest
+from exo.worker.engines.base import Engine
+from exo.worker.engines.vllm.engine import VllmEngine
+
+
+def _has_cuda() -> bool:
+    try:
+        import torch
+    except ImportError:
+        return False
+    return bool(torch.cuda.is_available())
+
+
+def _make_token_ids(n: int) -> list[int]:
+    # Deterministic synthetic tokens. Vocab >= ~30k for the Qwen tokenizers we
+    # care about, so 100..30099 is safe.
+    return [(i * 1009 + 17) % 30000 + 100 for i in range(n)]
+
+
+def _decode_stream(
+    payload: bytes,
+) -> tuple[list[KVChunk], list[ArraysState], Done | None, ErrorMessage | None]:
+    buf = io.BytesIO(payload)
+    _ = read_header(buf)
+    chunks: list[KVChunk] = []
+    arrays: list[ArraysState] = []
+    done: Done | None = None
+    error: ErrorMessage | None = None
+    while True:
+        match msg := read_message(buf):
+            case None:
+                break
+            case KVChunk():
+                chunks.append(msg)
+            case ArraysState():
+                arrays.append(msg)
+            case Done():
+                done = msg
+                break
+            case ErrorMessage():
+                error = msg
+                break
+    return chunks, arrays, done, error
+
+
+@pytest.fixture(scope="module")
+def vllm_engine(request: pytest.FixtureRequest) -> Iterator[object]:
+    """Build a real VllmEngine pointed at a downloaded HF model."""
+    if not _has_cuda():
+        pytest.skip("CUDA not available")
+    model_id_str = cast(str, request.config.getoption("--model-id"))
+    if not model_id_str:
+        pytest.skip("pass --model-id <hf-id> to run this test")
+
+    model_id = ModelId(model_id_str)
+
+    from exo.download.download_utils import build_model_path
+
+    if not build_model_path(model_id).exists():
+        pytest.skip(f"model {model_id} not downloaded locally")
+
+    from exo.worker.engines.vllm.generator import (
+        VllmBatchEngine,
+        load_vllm_engine,
+    )
+    from exo.worker.engines.vllm.kv_connector import (
+        ExoKVProducerConnector,
+        _patch_gdn_capture,
+        _patch_vllm_for_connector,
+    )
+
+    # Mirror VllmBuilder.load() — patches must run before LLMEngine init.
+    _patch_vllm_for_connector(ExoKVProducerConnector)
+    _patch_gdn_capture()
+
+    llm_engine, tool_parser = load_vllm_engine(
+        model_id=model_id,
+        trust_remote_code=False,
+        n_layers=1,
+        kv_connector_cls=ExoKVProducerConnector,
+    )
+    gen = VllmBatchEngine(engine=llm_engine, model_id=model_id)
+
+    # serve_prefill only touches self._gen.engine; the channel fields exist
+    # for the (unused-here) generation path.
+    class _DummySender:
+        def send(self, _: object) -> None: ...
+
+    class _DummyReceiver:
+        def collect(self) -> list[object]:
+            return []
+
+    engine = VllmEngine(
+        tool_parser=tool_parser,
+        model_id=model_id,
+        cancel_receiver=_DummyReceiver(),  # pyright: ignore[reportArgumentType]
+        event_sender=_DummySender(),  # pyright: ignore[reportArgumentType]
+        _gen=gen,
+        max_concurrent_requests=1,
+    )
+    try:
+        yield engine
+    finally:
+        with contextlib.suppress(Exception):
+            engine.close()
+
+
+def _run_one(engine: Engine, n_tokens: int, label: str) -> Done:
+    request = PrefillRequest(
+        request_id=f"itest-{label}",
+        model_id="ignored",
+        token_ids=_make_token_ids(n_tokens),
+        start_pos=0,
+        use_prefix_cache=True,
+    )
+    buf = io.BytesIO()
+    engine.serve_prefill(request, buf)
+
+    payload = buf.getvalue()
+    assert payload, f"{label}: server wrote nothing"
+
+    chunks, arrays, done, error = _decode_stream(payload)
+    if error is not None:
+        pytest.fail(
+            f"{label}: server returned ErrorMessage [{error.code}]: {error.message}"
+        )
+    assert done is not None, (
+        f"{label}: stream did not end with Done "
+        f"(received {len(chunks)} kv chunks, {len(arrays)} arrays)"
+    )
+    expected = max(0, n_tokens - 2)  # serve_prefill drops the last 2 tokens
+    assert done.total_tokens > 0, f"{label}: Done reported 0 tokens"
+    assert done.total_tokens >= expected - 64, (
+        f"{label}: got {done.total_tokens} tokens, expected ~{expected}"
+    )
+    assert chunks, f"{label}: no KV chunks shipped"
+    return done
+
+
+def test_serve_prefill_two_runs_no_apc_assert(vllm_engine: VllmEngine) -> None:
+    """Two consecutive prefills against the same engine must both succeed.
+
+    Before the fix, the second call hit a CUDA assert (vLLM APC + chunked
+    prefill + custom kv-connector). With APC off, each request runs a full
+    forward and the stream is well-formed both times.
+    """
+    first = _run_one(vllm_engine, n_tokens=512, label="run1")
+    second = _run_one(vllm_engine, n_tokens=512, label="run2-same-prompt")
+    assert second.total_tokens == first.total_tokens, (
+        f"run2 returned {second.total_tokens}, run1 returned {first.total_tokens}"
+    )
+
+
+def test_serve_prefill_different_lengths(vllm_engine: VllmEngine) -> None:
+    """A second prefill with a different prompt length still succeeds."""
+    a = _run_one(vllm_engine, n_tokens=256, label="run-256")
+    b = _run_one(vllm_engine, n_tokens=768, label="run-768")
+    assert b.total_tokens > a.total_tokens, (
+        f"longer prompt should ship more tokens: 256->{a.total_tokens} 768->{b.total_tokens}"
+    )

src/exo/worker/engines/vllm/engine.py

@@ -0,0 +1,623 @@
+import contextlib
+import itertools
+import threading
+import time
+from collections import deque
+from collections.abc import Generator, Iterator
+from dataclasses import dataclass, field
+from typing import BinaryIO
+
+import torch
+from vllm import SamplingParams
+from vllm.outputs import RequestOutput
+
+from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
+from exo.shared.types.chunks import ErrorChunk, GenerationChunk, PrefillProgressChunk
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import ChunkGenerated, Event
+from exo.shared.types.tasks import (
+    CANCEL_ALL_TASKS,
+    GenerationTask,
+    TaskId,
+    TextGeneration,
+)
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import (
+    CancelledResponse,
+    FinishedResponse,
+    GenerationResponse,
+)
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.disaggregated.protocol import write_error, write_kv_chunk
+from exo.worker.disaggregated.server import PrefillRequest
+from exo.worker.engines.base import Engine
+from exo.worker.engines.vllm.disaggregated.adapter import (
+    arrays_to_blobs,
+    tensor_to_wire_bytes,
+    torch_dtype_to_wire,
+    write_layer_arrays_blobs,
+    write_prefill_done,
+    write_prefill_header,
+)
+from exo.worker.engines.vllm.generator import VllmBatchEngine
+from exo.worker.engines.vllm.growable_cache import get_model_runner
+from exo.worker.engines.vllm.kv_connector import (
+    get_arrays_queue,
+    get_gdn_shipped,
+    get_gdn_states,
+    get_kv_queue,
+    get_save_kv_layer_diag,
+    init_gdn_layer_order,
+    reset_capture_state,
+)
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.model_output_parsers import (
+    apply_all_parsers,
+    map_responses_to_chunks,
+)
+from exo.worker.runner.llm_inference.tool_parsers import ToolParser
+
+
+class GeneratorQueue[T]:
+    def __init__(self) -> None:
+        self._q = deque[T]()
+
+    def push(self, t: T) -> None:
+        self._q.append(t)
+
+    def gen(self) -> Generator[T | None]:
+        while True:
+            if len(self._q) == 0:
+                yield None
+            else:
+                yield self._q.popleft()
+
+
+EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
+EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
+
+
+def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
+    """Keep the cheap debug prompt hooks without importing the MLX engine."""
+    if len(task_params.input) == 0:
+        return
+    prompt = task_params.input[0].content
+    if not prompt:
+        return
+    if EXO_RUNNER_MUST_FAIL in prompt:
+        raise Exception("Artificial runner exception - for testing purposes only.")
+    if EXO_RUNNER_MUST_TIMEOUT in prompt:
+        time.sleep(100)
+
+
+@dataclass(eq=False)
+class VllmEngine(Engine):
+    """Single-node vLLM implementation of the exo Engine interface.
+
+    This intentionally duplicates the local orchestration from the MLX
+    BatchGenerator instead of trying to share a batch abstraction too early.
+    The vLLM-specific tokenization/sampling/stepping remains inside
+    VllmBatchEngine.
+    """
+
+    tool_parser: ToolParser | None
+    model_id: ModelId
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    _gen: VllmBatchEngine
+    max_concurrent_requests: int = EXO_MAX_CONCURRENT_REQUESTS
+    check_for_cancel_every: int = 50
+
+    _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
+    _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
+    _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
+    _active_tasks: dict[
+        TaskId,
+        tuple[
+            TextGeneration,
+            GeneratorQueue[GenerationResponse],
+            Iterator[GenerationChunk | None],
+        ],
+    ] = field(default_factory=dict, init=False)
+
+    def warmup(self) -> None:
+        self.check_for_cancel_every = self._gen.warmup()
+
+    def submit(self, task: GenerationTask) -> None:
+        assert isinstance(task, TextGeneration)
+        self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
+        self._all_tasks[task.task_id] = task
+        self._queue.append(task)
+
+    def step(
+        self,
+    ) -> Iterator[
+        tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
+    ]:
+        self._collect_cancellations()
+        output: list[
+            tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
+        ] = list(self._apply_cancellations())
+
+        while self._queue and len(self._active_tasks) < self.max_concurrent_requests:
+            task = self._queue.popleft()
+            if self.should_cancel(task.task_id):
+                output.append((task.task_id, CancelledResponse()))
+                self._all_tasks.pop(task.task_id, None)
+                continue
+
+            try:
+                task_id, queue, output_generator = self._start_task(task)
+            except Exception as e:
+                self._send_error(task, e)
+                self._all_tasks.pop(task.task_id, None)
+                raise
+
+            self._active_tasks[task_id] = (task, queue, output_generator)
+
+        if not self._gen.has_work:
+            return iter(output)
+
+        results = self._gen.step()
+        for task_id, response in results:
+            if task_id not in self._active_tasks:
+                logger.warning(f"{task_id=} not found in active vLLM tasks")
+                continue
+
+            task, queue, output_generator = self._active_tasks[task_id]
+            queue.push(response)
+            while (parsed := next(output_generator, None)) is not None:
+                output.append((task.task_id, parsed))
+
+            if response.finish_reason is not None:
+                output.append((task.task_id, FinishedResponse()))
+                del self._active_tasks[task_id]
+                self._all_tasks.pop(task.task_id, None)
+
+        return itertools.chain(output, self._apply_cancellations())
+
+    def close(self) -> None:
+        self._gen.close()
+
+    def serve_prefill(self, request: PrefillRequest, wfile: BinaryIO) -> None:
+        engine = self._gen.engine
+        if engine.has_unfinished_requests():
+            logger.warning("serve_prefill: engine busy, refusing prefill request")
+            write_error(wfile, code=503, message="engine busy")
+            return
+
+        model_runner = get_model_runner()
+        if model_runner is None:
+            logger.warning("serve_prefill: model runner not initialized")
+            write_error(wfile, code=503, message="model runner not initialized")
+            return
+
+        init_gdn_layer_order(model_runner.kv_caches)
+
+        prefill_token_ids = (
+            request.token_ids[:-2]
+            if len(request.token_ids) > 2
+            else list(request.token_ids)
+        )
+        n_layers = len(model_runner.kv_caches)
+
+        reset_capture_state()
+        arrays_queue = get_arrays_queue()
+        kv_queue = get_kv_queue()
+
+        # We strip the trailing 2 tokens because the consumer warm-starts
+        # decode from them locally.
+        sp = SamplingParams(max_tokens=2, temperature=0.0, detokenize=False)
+        engine.add_request(
+            request.request_id,
+            {"prompt_token_ids": prefill_token_ids},
+            sp,
+        )
+
+        write_prefill_header(
+            wfile,
+            request_id=request.request_id,
+            model_id=request.model_id,
+            num_layers=n_layers,
+            start_pos=request.start_pos,
+        )
+
+        skip_tokens = request.start_pos
+        chunks_sent = 0
+        arrays_streamed = 0
+        layer_token_counts: dict[int, int] = {}
+        # Both writer threads serialize through this lock — BufferedWriter
+        # is not thread-safe and we don't want partial frames interleaved.
+        wfile_lock = threading.Lock()
+        # Diag for end-of-request bandwidth report.
+        writer_stats = {
+            "bytes_shipped": 0,
+            "wait_event_secs": 0.0,
+            "socket_secs": 0.0,
+            "first_byte_t": 0.0,
+            "last_byte_t": 0.0,
+            "started_t": 0.0,
+        }
+
+        def writer_loop() -> None:
+            nonlocal chunks_sent
+            writer_stats["started_t"] = time.perf_counter()
+            last_hb = time.perf_counter()
+            try:
+                while True:
+                    try:
+                        item = kv_queue.get(timeout=3.0)
+                    except Exception:
+                        item = ...  # sentinel for "no item yet"
+                    if item is ...:
+                        now = time.perf_counter()
+                        if now - last_hb > 3.0:
+                            logger.info(
+                                f"serve_prefill writer idle: "
+                                f"chunks_sent={chunks_sent} "
+                                f"kv_queue_size={kv_queue.qsize()} "
+                                f"arrays_queue_size={arrays_queue.qsize()}"
+                            )
+                            last_hb = now
+                        continue
+                    if item is None:
+                        break
+                    layer_idx, count, keys, values, copy_event = item
+                    # Wait for the side-stream D2H to finish populating the
+                    # pinned host buffers. CPU-side wait, doesn't block GPU.
+                    t_wait = time.perf_counter()
+                    copy_event.synchronize()
+                    writer_stats["wait_event_secs"] += time.perf_counter() - t_wait
+                    previous = layer_token_counts.get(layer_idx, 0)
+                    new_total = previous + count
+                    layer_token_counts[layer_idx] = new_total
+
+                    if new_total <= skip_tokens:
+                        continue
+                    # Reshape paged 4-D layouts to per-token 3-D up front so
+                    # the trim slice operates on the token axis.
+                    if keys.dim() == 4:
+                        keys = keys.reshape(-1, keys.shape[-2], keys.shape[-1])
+                        values = values.reshape(-1, values.shape[-2], values.shape[-1])
+                    # Slice keys/values to exactly `count` tokens — the source
+                    # tensor may be larger if shape disagrees with logical
+                    # token count (e.g. paged storage gathered over more
+                    # blocks than tokens consumed).
+                    if int(keys.shape[0]) > count:
+                        keys = keys[:count]
+                        values = values[:count]
+                    if previous < skip_tokens:
+                        trim = skip_tokens - previous
+                        keys = keys[trim:]
+                        values = values[trim:]
+                    num_tokens = int(keys.shape[0])
+                    n_heads = int(keys.shape[1])
+                    head_dim = int(keys.shape[2])
+                    dtype_w = torch_dtype_to_wire(keys.dtype)
+                    keys_bytes = tensor_to_wire_bytes(keys)
+                    values_bytes = tensor_to_wire_bytes(values)
+                    payload_bytes = len(keys_bytes) + len(values_bytes)
+                    if chunks_sent == 0:
+                        writer_stats["first_byte_t"] = time.perf_counter()
+                        logger.info(
+                            f"First KV chunk: layer={layer_idx} keys={keys.shape} "
+                            f"keys.dtype={keys.dtype} values.dtype={values.dtype}"
+                        )
+                    t_sock = time.perf_counter()
+                    with wfile_lock:
+                        write_kv_chunk(
+                            wfile,
+                            layer_idx=layer_idx,
+                            num_tokens=num_tokens,
+                            n_heads=n_heads,
+                            head_dim=head_dim,
+                            dtype=dtype_w,
+                            keys=keys_bytes,
+                            values=values_bytes,
+                        )
+                    writer_stats["socket_secs"] += time.perf_counter() - t_sock
+                    writer_stats["bytes_shipped"] += payload_bytes
+                    writer_stats["last_byte_t"] = time.perf_counter()
+                    chunks_sent += 1
+            except Exception:
+                logger.opt(exception=True).warning(
+                    "serve_prefill writer thread crashed"
+                )
+
+        def arrays_writer_loop() -> None:
+            nonlocal arrays_streamed
+            try:
+                while True:
+                    item = arrays_queue.get()
+                    if item is None:
+                        break
+                    layer_idx, arrays, copy_event = item
+                    if copy_event is not None:
+                        copy_event.synchronize()
+                    with wfile_lock:
+                        write_layer_arrays_blobs(
+                            wfile, layer_idx, arrays_to_blobs(arrays)
+                        )
+                    arrays_streamed += 1
+            except Exception:
+                logger.opt(exception=True).warning(
+                    "serve_prefill arrays writer thread crashed"
+                )
+
+        writer_thread = threading.Thread(target=writer_loop, daemon=True)
+        writer_thread.start()
+        arrays_writer_thread = threading.Thread(target=arrays_writer_loop, daemon=True)
+        arrays_writer_thread.start()
+
+        t0 = time.perf_counter()
+        forward_error: Exception | None = None
+        step_count = 0
+        last_step_log = time.perf_counter()
+        first_output_logged = False
+        try:
+            while engine.has_unfinished_requests():
+                outputs = engine.step()
+                step_count += 1
+                now = time.perf_counter()
+                if now - last_step_log > 3.0:
+                    logger.info(
+                        f"serve_prefill {request.request_id}: "
+                        f"step #{step_count} (kv_queue={kv_queue.qsize()})"
+                    )
+                    last_step_log = now
+                aborted = False
+                for output in outputs:
+                    if not first_output_logged:
+                        first_output_logged = True
+                        logger.info(
+                            f"serve_prefill {request.request_id}: first output "
+                            f"id={output.request_id!r} "
+                            f"tokens={len(output.outputs[0].token_ids) if isinstance(output, RequestOutput) and output.outputs else 0}"
+                        )
+                    # Match either the external id we passed or the
+                    # internal-suffixed id vLLM may surface ('-XXXXXXXX').
+                    if (
+                        isinstance(output, RequestOutput)
+                        and (
+                            output.request_id == request.request_id
+                            or output.request_id.startswith(request.request_id)
+                        )
+                        and output.outputs
+                        and output.outputs[0].token_ids
+                    ):
+                        engine.abort_request([request.request_id])
+                        aborted = True
+                        break
+                if aborted:
+                    break
+            # Post-abort drain. Bail out hard after 5s — if the request
+            # didn't finish by then something is wrong upstream and we'd
+            # otherwise spin forever in a no-op step loop.
+            drain_deadline = time.perf_counter() + 5.0
+            while engine.has_unfinished_requests():
+                if time.perf_counter() > drain_deadline:
+                    logger.warning(
+                        f"serve_prefill {request.request_id}: post-abort drain "
+                        f"timeout, force-aborting and breaking out"
+                    )
+                    with contextlib.suppress(Exception):
+                        engine.abort_request([request.request_id])
+                    break
+                _ = engine.step()
+                step_count += 1
+        except Exception as exc:
+            forward_error = exc
+            with contextlib.suppress(Exception):
+                engine.abort_request([request.request_id])
+        finally:
+            logger.info(
+                f"serve_prefill {request.request_id}: "
+                f"kv_queue={kv_queue.qsize()} arrays_queue={arrays_queue.qsize()}"
+            )
+            kv_queue.put(None)
+            arrays_queue.put(None)
+            writer_thread.join(timeout=30)
+            arrays_writer_thread.join(timeout=30)
+            if writer_thread.is_alive():
+                logger.warning("serve_prefill: kv writer thread did not exit")
+            if arrays_writer_thread.is_alive():
+                logger.warning("serve_prefill: arrays writer thread did not exit")
+
+        if forward_error is not None:
+            logger.opt(exception=forward_error).error(
+                f"serve_prefill {request.request_id}: engine.step() raised"
+            )
+            with contextlib.suppress(Exception):
+                write_error(wfile, code=500, message=f"engine.step: {forward_error!r}")
+            return
+
+        # The K/V writer and arrays writer both drained their queues during
+        # forward (see writer_loop / arrays_writer_loop above). What remains
+        # here is the fallback for any GDN layer whose conv+ssm pair never
+        # reached `_try_ship_gdn` — e.g., ssm captured but not conv. We skip
+        # layers already shipped by the streaming path.
+        gdn = get_gdn_states()
+        gdn_shipped = get_gdn_shipped()
+        unshipped = [li for li in sorted(gdn.keys()) if li not in gdn_shipped]
+        arrays_layers = arrays_streamed
+        if unshipped:
+            torch.cuda.synchronize()
+            for layer_idx in unshipped:
+                state = gdn[layer_idx]
+                arrs: list[torch.Tensor] = []
+                if "conv" in state:
+                    arrs.append(state["conv"])
+                if "ssm" in state:
+                    arrs.append(state["ssm"])
+                if arrs:
+                    write_layer_arrays_blobs(wfile, layer_idx, arrays_to_blobs(arrs))
+                    arrays_layers += 1
+
+        forwarded_per_layer = max(layer_token_counts.values(), default=0)
+        tokens_sent = max(0, forwarded_per_layer - skip_tokens)
+        write_prefill_done(wfile, tokens_sent)
+        elapsed = time.perf_counter() - t0
+        diag = get_save_kv_layer_diag()
+        diag_summary = ", ".join(
+            f"L{li}:{','.join(str(s) for s in sizes)}"
+            for li, sizes in sorted(diag.items())
+        )
+        logger.info(
+            f"serve_prefill {request.request_id}: save_kv_layer calls per layer "
+            f"(positive=non-list/tuple kv, negative=list/tuple kv) → {diag_summary}"
+        )
+        logger.info(
+            f"serve_prefill {request.request_id}: layer_token_counts="
+            f"{dict(sorted(layer_token_counts.items()))}"
+        )
+        # Bandwidth + per-stage breakdown for the writer thread.
+        bytes_shipped = writer_stats["bytes_shipped"]
+        wait_secs = writer_stats["wait_event_secs"]
+        sock_secs = writer_stats["socket_secs"]
+        first_byte_dt = (
+            writer_stats["first_byte_t"] - t0 if writer_stats["first_byte_t"] else 0.0
+        )
+        ship_secs = (
+            writer_stats["last_byte_t"] - writer_stats["first_byte_t"]
+            if writer_stats["last_byte_t"]
+            else 0.0
+        )
+        eff_bw_mbps = (bytes_shipped / 1e6 / ship_secs) if ship_secs > 0 else 0.0
+        peak_bw_mbps = (bytes_shipped / 1e6 / sock_secs) if sock_secs > 0 else 0.0
+        logger.info(
+            f"serve_prefill {request.request_id}: "
+            f"streamed_chunks={chunks_sent} arrays_layers={arrays_layers} "
+            f"tokens={tokens_sent} elapsed_ms={elapsed * 1000:.0f} "
+            f"bytes={bytes_shipped / 1e6:.0f}MB ttfb_ms={first_byte_dt * 1000:.0f} "
+            f"ship_ms={ship_secs * 1000:.0f} "
+            f"wait_event_ms={wait_secs * 1000:.0f} sock_ms={sock_secs * 1000:.0f} "
+            f"eff_bw={eff_bw_mbps:.0f}MB/s peak_bw={peak_bw_mbps:.0f}MB/s"
+        )
+
+    def _start_task(
+        self, task: TextGeneration
+    ) -> tuple[
+        TaskId,
+        GeneratorQueue[GenerationResponse],
+        Iterator[GenerationChunk | None],
+    ]:
+        from exo.worker.engines.vllm.prompt_format import format_vllm_prompt
+
+        _check_for_debug_prompts(task.task_params)
+
+        token_ids, prompt_text, _ = format_vllm_prompt(
+            self._gen.engine, task.task_params
+        )
+
+        queue = GeneratorQueue[GenerationResponse]()
+        if task.task_params.bench:
+            output_generator: Iterator[GenerationChunk | None] = map(
+                lambda r: map_responses_to_chunks(r, self.model_id), queue.gen()
+            )
+        else:
+            from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+            output_generator = apply_all_parsers(
+                queue.gen(),
+                prompt_text,
+                self.tool_parser,
+                TokenizerWrapper(self._gen.engine.get_tokenizer()),
+                self.model_id,
+                task.task_params.tools,
+            )
+
+        check_for_cancel_every = max(self.check_for_cancel_every, 1)
+        tokens_since_cancel_check = check_for_cancel_every
+
+        def on_prefill_progress(processed: int, total: int) -> None:
+            self._collect_cancellations()
+            if self.should_cancel(task.task_id):
+                self._cancelled_tasks.add(task.task_id)
+            self.event_sender.send(
+                ChunkGenerated(
+                    command_id=task.command_id,
+                    chunk=PrefillProgressChunk(
+                        model=self.model_id,
+                        processed_tokens=processed,
+                        total_tokens=total,
+                    ),
+                )
+            )
+
+        def on_generation_token() -> None:
+            nonlocal tokens_since_cancel_check
+            tokens_since_cancel_check += 1
+            if tokens_since_cancel_check >= check_for_cancel_every:
+                tokens_since_cancel_check = 0
+                self._collect_cancellations()
+                if self.should_cancel(task.task_id):
+                    self._cancelled_tasks.add(task.task_id)
+
+        task_id = self._gen.submit(
+            task_id=task.task_id,
+            task_params=task.task_params,
+            on_prefill_progress=on_prefill_progress,
+            on_generation_token=on_generation_token,
+            token_ids=token_ids,
+        )
+        return task_id, queue, output_generator
+
+    def _collect_cancellations(self) -> None:
+        for task_id in self.cancel_receiver.collect():
+            if task_id == CANCEL_ALL_TASKS:
+                self._cancelled_tasks.add(CANCEL_ALL_TASKS)
+            elif task_id in self._all_tasks:
+                self._cancelled_tasks.add(task_id)
+
+    def _apply_cancellations(self) -> Iterator[tuple[TaskId, CancelledResponse]]:
+        if not self._cancelled_tasks:
+            return iter([])
+
+        cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks
+        results: list[tuple[TaskId, CancelledResponse]] = []
+        task_ids_to_abort: list[TaskId] = []
+
+        for task_id, (task, _, _) in list(self._active_tasks.items()):
+            if cancel_all or task.task_id in self._cancelled_tasks:
+                task_ids_to_abort.append(task_id)
+                results.append((task.task_id, CancelledResponse()))
+                del self._active_tasks[task_id]
+                self._all_tasks.pop(task.task_id, None)
+
+        if self._queue:
+            kept_queue: deque[TextGeneration] = deque()
+            for task in self._queue:
+                if cancel_all or task.task_id in self._cancelled_tasks:
+                    results.append((task.task_id, CancelledResponse()))
+                    self._all_tasks.pop(task.task_id, None)
+                else:
+                    kept_queue.append(task)
+            self._queue = kept_queue
+
+        if task_ids_to_abort:
+            self._gen.cancel(task_ids_to_abort)
+
+        already_cancelled = {task_id for task_id, _ in results}
+        for task_id in self._cancelled_tasks:
+            if (
+                task_id != CANCEL_ALL_TASKS
+                and task_id in self._all_tasks
+                and task_id not in already_cancelled
+            ):
+                results.append((task_id, CancelledResponse()))
+                self._all_tasks.pop(task_id, None)
+
+        self._cancelled_tasks.clear()
+        return iter(results)
+
+    def _send_error(self, task: TextGeneration, e: Exception) -> None:
+        self.event_sender.send(
+            ChunkGenerated(
+                command_id=task.command_id,
+                chunk=ErrorChunk(
+                    model=self.model_id,
+                    finish_reason="error",
+                    error_message=str(e),
+                ),
+            )
+        )

src/exo/worker/engines/vllm/generator.py

@@ -0,0 +1,490 @@
+import gc
+import json
+import math
+import re
+import sys
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass, field
+from typing import cast
+
+import torch
+from vllm.config import CompilationConfig
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
+from vllm.engine.arg_utils import EngineArgs
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionMessageParam,
+    CustomChatCompletionMessageParam,
+)
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import SamplingParams
+from vllm.tokenizers import TokenizerLike
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.engine.llm_engine import LLMEngine
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from exo.api.types import (
+    CompletionTokensDetails,
+    GenerationStats,
+    PromptTokensDetails,
+    Usage,
+)
+from exo.download.download_utils import build_model_path
+from exo.shared.types.common import ModelId
+from exo.shared.types.memory import Memory
+from exo.shared.types.tasks import TaskId
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import GenerationResponse
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
+
+
+@dataclass
+class _EngineRequest:
+    request_id: str
+    prompt_token_count: int
+    prefill_done: bool = False
+    prefill_steps: int = 0
+    prev_text: str = ""
+    prev_token_count: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+    first_token_time: float | None = None
+    on_generation_token: Callable[[], None] | None = None
+    on_prefill_progress: Callable[[int, int], None] | None = None
+
+
+def _stop_token_ids(tokenizer: TokenizerLike, model_id: ModelId) -> set[int]:
+    from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
+
+    ids: set[int] = set()
+    eos_id = getattr(tokenizer, "eos_token_id", None)
+    if eos_id is not None:
+        ids.add(eos_id)  # pyright: ignore[reportAny]
+    extra = get_eos_token_ids_for_model(model_id)
+    if extra:
+        ids.update(extra)
+    return ids
+
+
+def _build_generation_response(
+    tokenizer: TokenizerLike,
+    token_id: int,
+    finish_reason: str | None,
+    prompt_token_count: int,
+    completion_tokens: int,
+    start_time: float,
+    first_token_time: float | None,
+    suppress_text: bool = False,
+) -> GenerationResponse:
+    token_text: str = "" if suppress_text else tokenizer.decode([token_id])
+    finish_usage: Usage | None = None
+    finish_stats: GenerationStats | None = None
+    mapped_finish_reason: str | None = None
+    if finish_reason:
+        now = time.perf_counter()
+        prefill_elapsed = (first_token_time or now) - start_time
+        decode_elapsed = now - (first_token_time or now)
+        finish_usage = Usage(
+            prompt_tokens=prompt_token_count,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_token_count + completion_tokens,
+            prompt_tokens_details=PromptTokensDetails(),
+            completion_tokens_details=CompletionTokensDetails(),
+        )
+        finish_stats = GenerationStats(
+            prompt_tps=prompt_token_count / prefill_elapsed
+            if prefill_elapsed > 0
+            else 0.0,
+            generation_tps=completion_tokens / decode_elapsed
+            if decode_elapsed > 0
+            else 0.0,
+            prompt_tokens=prompt_token_count,
+            generation_tokens=completion_tokens,
+            peak_memory_usage=Memory.from_bytes(torch.cuda.max_memory_allocated()),
+        )
+        mapped_finish_reason = (
+            finish_reason
+            if finish_reason in ("stop", "length", "content_filter")
+            else "stop"
+        )
+    return GenerationResponse(
+        text=token_text,
+        token=token_id,
+        finish_reason=mapped_finish_reason,
+        usage=finish_usage,
+        stats=finish_stats,
+    )
+
+
+def warmup_vllm_engine(engine: LLMEngine) -> int:
+    tokenizer = engine.get_tokenizer()
+    messages = [
+        cast(
+            ChatCompletionMessageParam,
+            CustomChatCompletionMessageParam(
+                role="user",
+                content="Prompt to warm up the inference engine. Repeat this.",
+            ),
+        )
+    ]
+    prompt_text: str | list[int] = tokenizer.apply_chat_template(  # pyright: ignore[reportUnknownMemberType]
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    if isinstance(prompt_text, list):
+        token_ids = prompt_text
+    else:
+        token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)
+
+    params = SamplingParams(max_tokens=50, detokenize=False)
+    engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
+    t = time.monotonic()
+    tokens_generated = 0
+    while engine.has_unfinished_requests():
+        engine.step()
+        tokens_generated += 1
+    elapsed = max(time.monotonic() - t, 0.001)
+    check_for_cancel_every = min(math.ceil(tokens_generated / elapsed), 100)
+    logger.info(
+        f"vLLM warmup complete, check_for_cancel_every={check_for_cancel_every}"
+    )
+    return check_for_cancel_every
+
+
+@dataclass(eq=False)
+class VllmBatchEngine:
+    engine: LLMEngine
+    model_id: ModelId
+
+    _active: dict[TaskId, _EngineRequest] = field(default_factory=dict, init=False)
+
+    def warmup(self) -> int:
+        return warmup_vllm_engine(self.engine)
+
+    @property
+    def has_work(self) -> bool:
+        return bool(self._active) or self.engine.has_unfinished_requests()
+
+    def submit(
+        self,
+        task_id: TaskId,
+        task_params: TextGenerationTaskParams,
+        token_ids: list[int],
+        on_prefill_progress: Callable[[int, int], None] | None = None,
+        on_generation_token: Callable[[], None] | None = None,
+    ) -> TaskId:
+        from exo.worker.engines.vllm.prompt_format import make_vllm_sampling_params
+
+        sampling_params = make_vllm_sampling_params(
+            self.engine, task_params, self.model_id
+        )
+        self.engine.add_request(
+            task_id, {"prompt_token_ids": token_ids}, sampling_params
+        )
+        self._active[task_id] = _EngineRequest(
+            request_id=task_id,
+            prompt_token_count=len(token_ids),
+            on_generation_token=on_generation_token,
+            on_prefill_progress=on_prefill_progress,
+        )
+        return task_id
+
+    def step(self) -> list[tuple[TaskId, GenerationResponse]]:
+        if not self.has_work:
+            return []
+
+        outputs = self.engine.step()
+        tokenizer = self.engine.get_tokenizer()
+        stop_ids = _stop_token_ids(tokenizer, self.model_id)
+        max_batch_tokens: int = (
+            getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048
+        )
+        results: list[tuple[TaskId, GenerationResponse]] = []
+
+        for output in outputs:
+            # todo: PoolingRequestOutputs
+            assert isinstance(output, RequestOutput)
+            task_id = TaskId(output.request_id)
+            if task_id not in self._active:
+                continue
+            req = self._active[task_id]
+            completion = output.outputs[0]
+            new_token_count = len(completion.token_ids)
+            new_tokens = completion.token_ids[req.prev_token_count :]
+            finish_reason = completion.finish_reason
+            req.prev_token_count = new_token_count
+
+            if not req.prefill_done and not new_tokens:
+                req.prefill_steps += 1
+                if req.on_prefill_progress:
+                    req.on_prefill_progress(
+                        min(
+                            req.prefill_steps * max_batch_tokens, req.prompt_token_count
+                        ),
+                        req.prompt_token_count,
+                    )
+                continue
+
+            if not req.prefill_done and new_tokens:
+                req.first_token_time = time.perf_counter()
+                req.prefill_done = True
+
+            for i, token_id in enumerate(new_tokens):
+                is_last = i == len(new_tokens) - 1
+                is_final_stop = is_last and finish_reason and token_id in stop_ids
+                if req.on_generation_token:
+                    req.on_generation_token()
+                results.append(
+                    (
+                        task_id,
+                        _build_generation_response(
+                            tokenizer,
+                            token_id,
+                            finish_reason if is_last and finish_reason else None,
+                            req.prompt_token_count,
+                            new_token_count,
+                            req.start_time,
+                            req.first_token_time,
+                            suppress_text=bool(is_final_stop),
+                        ),
+                    )
+                )
+
+            if finish_reason:
+                del self._active[task_id]
+
+        for req in self._active.values():
+            if not req.prefill_done:
+                req.prefill_steps += 1
+                if req.on_prefill_progress:
+                    req.on_prefill_progress(
+                        min(
+                            req.prefill_steps * max_batch_tokens, req.prompt_token_count
+                        ),
+                        req.prompt_token_count,
+                    )
+
+        return results
+
+    def cancel(self, task_ids: list[TaskId]) -> None:
+        to_abort = [str(tid) for tid in task_ids if tid in self._active]
+        if to_abort:
+            self.engine.abort_request(to_abort)
+        for tid in task_ids:
+            self._active.pop(tid, None)
+
+    def close(self) -> None:
+        if not hasattr(self, "engine"):
+            return
+        rids = [req.request_id for req in self._active.values()]
+        if rids:
+            self.engine.abort_request(rids)
+        self._active.clear()
+        del self.engine
+        gc.collect()
+        torch.cuda.empty_cache()
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+
+
+_weight_loading_callback: Callable[[int, int], None] | None = None
+_weight_loading_patched = False
+
+
+def get_weight_loading_callback() -> Callable[[int, int], None] | None:
+    return _weight_loading_callback
+
+
+def set_weight_loading_callback(cb: Callable[[int, int], None] | None) -> None:
+    global _weight_loading_callback
+    _weight_loading_callback = cb
+
+
+_LAYER_INDEX_PATTERN = re.compile(r"\.layers\.(\d+)\.")
+_n_layers: int = 1
+
+
+def get_n_layers() -> int:
+    return _n_layers
+
+
+def set_n_layers(n: int) -> None:
+    global _n_layers
+    _n_layers = n
+
+
+def _wrap_weights_iterator(
+    original: Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]],
+) -> Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]]:
+    def patched(
+        hf_weights_files: list[str], *args: object, **kwargs: object
+    ) -> Generator[tuple[str, "torch.Tensor"], None, None]:
+        callback = get_weight_loading_callback()
+        if callback is not None and hf_weights_files:
+            total_layers = get_n_layers()
+            seen_layers: set[int] = set()
+            last_reported = 0
+            for name, tensor in original(hf_weights_files, *args, **kwargs):
+                yield name, tensor
+                match = _LAYER_INDEX_PATTERN.search(name)
+                if match:
+                    seen_layers.add(int(match.group(1)))
+                current = len(seen_layers)
+                if current > last_reported:
+                    callback(current, total_layers)
+                    last_reported = current
+            callback(total_layers, total_layers)
+        else:
+            yield from original(hf_weights_files, *args, **kwargs)
+
+    return patched
+
+
+def _monkey_patch_iterator(weight_utils: object, attr_name: str) -> None:
+    original = getattr(weight_utils, attr_name, None)
+    if original is None:
+        return
+    patched = _wrap_weights_iterator(original)  # pyright: ignore[reportAny]
+    setattr(weight_utils, attr_name, patched)
+    for mod in list(sys.modules.values()):
+        if mod is weight_utils:
+            continue
+        for name in list(vars(mod)):
+            if vars(mod)[name] is original:
+                setattr(mod, name, patched)
+
+
+def _patch_weight_loading_progress() -> None:
+    global _weight_loading_patched
+    if _weight_loading_patched:
+        return
+    _weight_loading_patched = True
+
+    from vllm.model_executor.model_loader import (
+        weight_utils,
+    )
+
+    _monkey_patch_iterator(weight_utils, "safetensors_weights_iterator")
+    _monkey_patch_iterator(weight_utils, "fastsafetensors_weights_iterator")
+
+    import huggingface_hub
+
+    def _noop_metadata(*_a: object, **_kw: object) -> None:
+        pass
+
+    original_metadata = huggingface_hub.get_safetensors_metadata
+    huggingface_hub.get_safetensors_metadata = _noop_metadata
+    for mod in list(sys.modules.values()):
+        if mod is huggingface_hub:
+            continue
+        for attr in list(vars(mod)):
+            if vars(mod)[attr] is original_metadata:
+                setattr(mod, attr, _noop_metadata)
+
+
+def build_layer_groups(kv_cache_config: KVCacheConfig) -> list[int]:
+    group_lookup: dict[str, int] = {}
+    for group_idx, group_spec in enumerate(kv_cache_config.kv_cache_groups):
+        for layer_name in group_spec.layer_names:
+            group_lookup[layer_name] = group_idx
+
+    layer_to_group: list[int] = []
+    for tensor_spec in kv_cache_config.kv_cache_tensors:
+        for name in tensor_spec.shared_by:
+            layer_to_group.append(group_lookup[name])
+    return layer_to_group
+
+
+def load_vllm_engine(
+    model_id: ModelId,
+    trust_remote_code: bool,
+    n_layers: int = 1,
+    on_layer_loaded: Callable[[int, int], None] | None = None,
+    kv_connector_cls: type[object] | None = None,
+) -> tuple[LLMEngine, ToolParser | None]:
+    model_path = build_model_path(model_id)
+    _patch_weight_loading_progress()
+
+    set_n_layers(n_layers)
+
+    # Use the dict-with-colon form the original branch used. The typed
+    # `KVTransferConfig` object goes through a different vLLM code path
+    # and (with `kv_load_failure_policy="recompute"`) trips the
+    # APC/hybrid/chunked-prefill kv-cache nesting bug.
+    kv_transfer_config: dict[str, str] | None = None
+    if kv_connector_cls is not None:
+        kv_transfer_config = {
+            "kv_connector": (
+                f"{kv_connector_cls.__module__}:{kv_connector_cls.__name__}"
+            ),
+            "kv_role": "kv_both",
+        }
+
+    has_mamba = False
+    try:
+        with open(model_path / "config.json") as f:
+            model_config = json.load(f)  # pyright: ignore[reportAny]
+        text_config = model_config.get("text_config", model_config)  # pyright: ignore[reportAny]
+        has_mamba = "mamba_ssm_dtype" in text_config or "linear_attention" in (
+            text_config.get("layer_types") or []  # pyright: ignore[reportAny]
+        )
+    except Exception:
+        pass
+    if has_mamba:
+        backends = [AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TRITON_ATTN]
+    else:
+        backends = [
+            AttentionBackendEnum.FLASHINFER,
+            AttentionBackendEnum.FLASH_ATTN,
+            AttentionBackendEnum.TRITON_ATTN,
+        ]
+
+    engine: LLMEngine | None = None
+    for backend in backends:
+        try:
+            engine_args = EngineArgs(
+                model=str(model_path.expanduser().resolve()),
+                served_model_name=str(model_id),
+                gpu_memory_utilization=0.05,
+                trust_remote_code=trust_remote_code,
+                load_format="fastsafetensors",
+                enable_prefix_caching=True,
+                attention_backend=backend,
+                compilation_config=CompilationConfig(
+                    mode=CompilationMode.NONE,
+                    cudagraph_mode=CUDAGraphMode.NONE,
+                ),
+                disable_log_stats=True,
+                max_num_batched_tokens=4096,
+                kv_transfer_config=kv_transfer_config, # pyright: ignore[reportArgumentType]
+                disable_hybrid_kv_cache_manager=False,
+                kv_cache_dtype="auto",
+            )
+
+            set_weight_loading_callback(on_layer_loaded)
+            engine = LLMEngine.from_engine_args(engine_args)
+            logger.info(f"vLLM engine using attention backend: {backend}")
+            break
+        except (ValueError, RuntimeError, NotImplementedError) as e:
+            logger.warning(f"Attention backend {backend} failed: {e}, trying next")
+            engine = None
+
+            gc.collect()
+            torch.cuda.empty_cache()
+            continue
+
+    if engine is None:
+        raise RuntimeError(f"No attention backend worked for {model_id}")
+
+    tool_parser: ToolParser | None = None
+    tokenizer = engine.get_tokenizer()
+    chat_template = getattr(tokenizer, "chat_template", None)
+    if isinstance(chat_template, str):
+        tool_parser = infer_tool_parser(chat_template)
+        if tool_parser:
+            logger.info(
+                f"inferred tool parser: {tool_parser.start_parsing} / {tool_parser.end_parsing}"
+            )
+
+    logger.info(f"vLLM engine loaded for {model_id}")
+
+    return engine, tool_parser

src/exo/worker/engines/vllm/growable_cache.py

@@ -0,0 +1,480 @@
+# pyright: reportPrivateUsage=false, reportAttributeAccessIssue=false
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.request import Request
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+from exo.shared.logging import logger
+
+INITIAL_FRACTION = 0.05
+GROWTH_HEADROOM_BYTES = 512 * 1024 * 1024
+MIN_GROWTH_BLOCKS = 16
+
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+
+_patched = False
+_model_runner: GPUModelRunner | None = None
+
+
+def get_model_runner() -> GPUModelRunner | None:
+    return _model_runner
+
+
+def set_model_runner(runner: GPUModelRunner | None) -> None:
+    global _model_runner
+    _model_runner = runner
+
+
+def patch_vllm() -> None:
+    global _patched
+    if _patched:
+        return
+    _patched = True
+
+    _patch_nogds()
+    _patch_determine_available_memory()
+    _patch_check_enough_kv_cache_memory()
+    _patch_initialize_kv_cache_tensors()
+    _patch_initialize_from_config()
+    _patch_kv_cache_manager_init()
+    _patch_allocate_slots()
+    _patch_moe_sum()
+    _patch_marlin_w2_thread_config()
+    logger.info("vLLM growable KV cache patch applied")
+
+
+def _patch_nogds() -> None:
+    from vllm.model_executor.model_loader import weight_utils
+
+    original = weight_utils._init_fastsafetensors_loader
+
+    def patched(
+        pg: torch.distributed.ProcessGroup,
+        device: torch.device,
+        f_list: list[str],
+        *,
+        nogds: bool = False,
+    ) -> object:
+        return original(pg, device, f_list, nogds=True)
+
+    weight_utils._init_fastsafetensors_loader = patched
+
+
+def _patch_determine_available_memory() -> None:
+    from vllm.v1.worker.gpu_worker import Worker
+
+    # original = Worker.determine_available_memory
+
+    @torch.inference_mode()
+    def patched(self: Worker) -> int:
+        import pathlib
+        import shutil
+
+        compile_cache = pathlib.Path.home() / ".cache" / "vllm" / "torch_compile_cache"
+        if compile_cache.exists():
+            shutil.rmtree(compile_cache, ignore_errors=True)
+
+        free_bytes, _ = torch.cuda.mem_get_info()
+        # vLLM's get_kv_cache_configs computes per-group block counts via
+        # `tensor.size // num_blocks_old` and asserts the result divides
+        # evenly. With a small `available_kv_cache_memory_bytes` and
+        # multi-MiB-per-slot Mamba/hybrid groups, num_blocks_old can come
+        # back as 0 → ZeroDivisionError. Floor the initial budget so each
+        # group lands at least one block at init; growth picks up from
+        # there.
+        min_initial = 1024 * 1024 * 1024  # 1 GiB
+        if free_bytes < min_initial:
+            raise RuntimeError(
+                f"Insufficient GPU memory for KV cache initialization: "
+                f"{free_bytes / (1024**3):.2f} GiB free, need at least "
+                f"{min_initial / (1024**3):.2f} GiB. Stop other GPU "
+                f"processes (check `nvidia-smi`)."
+            )
+        initial = max(int(free_bytes * INITIAL_FRACTION), min_initial)
+        self._growable_max_kv_bytes = free_bytes
+        self.available_kv_cache_memory_bytes = initial
+        logger.info(
+            f"Growable KV cache: initial {initial / (1024**3):.2f} GiB "
+            f"(max {free_bytes / (1024**3):.2f} GiB)"
+        )
+        return initial
+
+    Worker.determine_available_memory = patched
+
+
+def _patch_check_enough_kv_cache_memory() -> None:
+    from vllm.v1.core import kv_cache_utils
+
+    def noop(*_args: object, **_kwargs: object) -> None:
+        pass
+
+    kv_cache_utils._check_enough_kv_cache_memory = noop
+
+
+def _patch_initialize_kv_cache_tensors() -> None:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+    original_alloc = GPUModelRunner._allocate_kv_cache_tensors
+
+    def patched_alloc(
+        self: GPUModelRunner, kv_cache_config: KVCacheConfig
+    ) -> dict[str, torch.Tensor]:
+        raw_tensors = original_alloc(self, kv_cache_config)
+        self._growable_raw_tensors = {name: t for name, t in raw_tensors.items()}
+        return raw_tensors
+
+    GPUModelRunner._allocate_kv_cache_tensors = patched_alloc
+
+    original_init_tensors = GPUModelRunner.initialize_kv_cache_tensors
+
+    def patched_init_tensors(
+        self: GPUModelRunner,
+        kv_cache_config: KVCacheConfig,
+        kernel_block_sizes: list[int],
+    ) -> dict[str, torch.Tensor]:
+        self._growable_kv_cache_config = kv_cache_config
+        self._growable_kernel_block_sizes = kernel_block_sizes
+        return original_init_tensors(self, kv_cache_config, kernel_block_sizes)
+
+    GPUModelRunner.initialize_kv_cache_tensors = patched_init_tensors
+
+
+def _patch_initialize_from_config() -> None:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+    from vllm.v1.worker.gpu_worker import Worker
+
+    original_init_attn = GPUModelRunner.initialize_attn_backend
+
+    def clear_and_reinit_attn(
+        self: GPUModelRunner,
+        kv_cache_config: KVCacheConfig,
+    ) -> None:
+        self.attn_groups.clear()
+        original_init_attn(self, kv_cache_config)
+
+    GPUModelRunner.initialize_attn_backend = clear_and_reinit_attn
+
+    original = Worker.initialize_from_config
+
+    def patched(self: Worker, kv_cache_config: KVCacheConfig) -> None:
+        original(self, kv_cache_config)
+        set_model_runner(self.model_runner)
+
+    Worker.initialize_from_config = patched
+
+
+def _patch_kv_cache_manager_init() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+    original_init = KVCacheManager.__init__
+
+    def patched_init(
+        self: KVCacheManager,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        hash_block_size: int,
+        enable_caching: bool = True,
+        use_eagle: bool = False,
+        log_stats: bool = False,
+        enable_kv_cache_events: bool = False,
+        dcp_world_size: int = 1,
+        pcp_world_size: int = 1,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+    ) -> None:
+        original_init(
+            self,
+            kv_cache_config,
+            max_model_len,
+            hash_block_size,
+            enable_caching,
+            use_eagle,
+            log_stats,
+            enable_kv_cache_events,
+            dcp_world_size,
+            pcp_world_size,
+            metrics_collector,
+        )
+        self._growable_model_runner = get_model_runner()
+
+    KVCacheManager.__init__ = patched_init
+
+
+def _patch_allocate_slots() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+
+    original = KVCacheManager.allocate_slots
+
+    def patched(
+        self: KVCacheManager,
+        request: Request,
+        num_new_tokens: int,
+        num_new_computed_tokens: int = 0,
+        new_computed_blocks: KVCacheBlocks | None = None,
+        num_lookahead_tokens: int = 0,
+        num_external_computed_tokens: int = 0,
+        delay_cache_blocks: bool = False,
+        num_encoder_tokens: int = 0,
+    ) -> KVCacheBlocks | None:
+        result = original(
+            self,
+            request,
+            num_new_tokens,
+            num_new_computed_tokens,
+            new_computed_blocks,
+            num_lookahead_tokens,
+            num_external_computed_tokens,
+            delay_cache_blocks,
+            num_encoder_tokens,
+        )
+        while result is None and _try_grow_cache(self):
+            result = original(
+                self,
+                request,
+                num_new_tokens,
+                num_new_computed_tokens,
+                new_computed_blocks,
+                num_lookahead_tokens,
+                num_external_computed_tokens,
+                delay_cache_blocks,
+                num_encoder_tokens,
+            )
+
+        return result
+
+    KVCacheManager.allocate_slots = patched
+
+    if hasattr(KVCacheManager, "can_fit_full_sequence"):
+        original_can_fit = cast(
+            Callable[..., bool],
+            KVCacheManager.can_fit_full_sequence,
+        )
+
+        def patched_can_fit(
+            self: KVCacheManager,
+            request: Request,
+            num_new_computed_tokens: int = 0,
+            new_computed_blocks: KVCacheBlocks | None = None,
+            num_external_computed_tokens: int = 0,
+            num_encoder_tokens: int = 0,
+        ) -> bool:
+            result: bool = original_can_fit(
+                self,
+                request,
+                num_new_computed_tokens,
+                new_computed_blocks,
+                num_external_computed_tokens,
+                num_encoder_tokens,
+            )
+            while not result and _try_grow_cache(self):
+                result = original_can_fit(
+                    self,
+                    request,
+                    num_new_computed_tokens,
+                    new_computed_blocks,
+                    num_external_computed_tokens,
+                    num_encoder_tokens,
+                )
+            return result
+
+        KVCacheManager.can_fit_full_sequence = patched_can_fit
+
+
+def _try_grow_cache(kv_cache_manager: "KVCacheManager") -> bool:
+    block_pool = kv_cache_manager.block_pool
+    model_runner = cast(GPUModelRunner | None, kv_cache_manager._growable_model_runner)
+
+    if model_runner is None:
+        return False
+
+    free_bytes, _ = torch.cuda.mem_get_info()
+    if free_bytes < GROWTH_HEADROOM_BYTES:
+        return False
+
+    kv_cache_config = cast(KVCacheConfig, model_runner._growable_kv_cache_config)
+    old_num_blocks: int = kv_cache_config.num_blocks
+
+    total_tensor_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
+    per_block_bytes = total_tensor_bytes // old_num_blocks
+
+    usable_bytes = int(free_bytes * 0.8)
+    growth_blocks = min(usable_bytes // per_block_bytes, old_num_blocks)
+
+    if growth_blocks < MIN_GROWTH_BLOCKS:
+        return False
+
+    new_num_blocks = old_num_blocks + growth_blocks
+
+    logger.info(
+        f"Growing KV cache: {old_num_blocks} → {new_num_blocks} blocks "
+        f"(+{growth_blocks * per_block_bytes / (1024**3):.2f} GiB)"
+    )
+
+    try:
+        kv_cache_config.num_blocks = new_num_blocks
+        for tensor_spec in kv_cache_config.kv_cache_tensors:
+            tensor_spec.size = int(tensor_spec.size * new_num_blocks / old_num_blocks)
+        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
+        _grow_block_pool(block_pool, old_num_blocks, new_num_blocks)
+        logger.info(f"KV cache grown successfully to {new_num_blocks} blocks")
+        return True
+    except Exception:
+        logger.opt(exception=True).error("Failed to grow KV cache")
+        return False
+
+
+def _grow_tensors(
+    model_runner: GPUModelRunner,
+    kv_cache_config: KVCacheConfig,
+    old_num_blocks: int,
+    new_num_blocks: int,
+) -> None:
+    raw_tensors: dict[str, torch.Tensor] = cast(
+        dict[str, torch.Tensor], model_runner._growable_raw_tensors
+    )
+    ratio = new_num_blocks / old_num_blocks
+
+    already_grown: dict[int, torch.Tensor] = {}
+    new_raw_tensors: dict[str, torch.Tensor] = {}
+
+    for layer_name, old_raw in raw_tensors.items():
+        storage_id = old_raw.data_ptr()
+        if storage_id in already_grown:
+            new_raw_tensors[layer_name] = already_grown[storage_id]
+            continue
+
+        old_size = old_raw.numel()
+        new_size = int(old_size * ratio)
+        new_raw = torch.zeros(new_size, dtype=torch.int8, device=old_raw.device)
+        new_raw[:old_size] = old_raw
+        already_grown[storage_id] = new_raw
+        new_raw_tensors[layer_name] = new_raw
+
+    model_runner._growable_raw_tensors = new_raw_tensors
+
+    kernel_block_sizes: list[int] = cast(
+        list[int], model_runner._growable_kernel_block_sizes
+    )
+    new_kv_caches: dict[str, torch.Tensor] = model_runner._reshape_kv_cache_tensors(
+        kv_cache_config,
+        new_raw_tensors,
+        kernel_block_sizes,
+    )
+
+    forward_context: dict[str, Any] = (
+        model_runner.compilation_config.static_forward_context
+    )
+    runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches
+
+    from collections import defaultdict
+
+    from vllm.model_executor.models.utils import extract_layer_index
+
+    num_attn_module = 1
+    hf_config = getattr(getattr(model_runner, "model_config", None), "hf_config", None)
+    if getattr(hf_config, "model_type", "") == "longcat_flash":
+        num_attn_module = 2
+
+    index2name: dict[int, list[str]] = defaultdict(list)
+    for ln in new_kv_caches:
+        index2name[extract_layer_index(ln, num_attn_module)].append(ln)
+
+    new_ordered: list[torch.Tensor] = []
+    for layer_index in sorted(index2name.keys()):
+        for ln in index2name[layer_index]:
+            new_ordered.append(new_kv_caches[ln])
+
+    for i, new_kv in enumerate(new_ordered):
+        if i < len(runner_kv_caches):
+            runner_kv_caches[i] = new_kv
+        else:
+            runner_kv_caches.append(new_kv)
+
+    new_kv_typed = cast(dict[str, torch.Tensor | list[torch.Tensor]], new_kv_caches)
+    for layer_name, new_kv in new_kv_typed.items():
+        # vLLM uses different shapes per layer kind (gpu_model_runner.py:5852):
+        #   - full / sliding-window attention: `attn.kv_cache: torch.Tensor`
+        #     (paged storage with K/V stacked along dim 0; consumers call
+        #     `.unbind(0)` so it MUST be a Tensor, not a list)
+        #   - Mamba / hybrid:                  `attn.kv_cache: list[Tensor]`
+        #     ([conv_state, ssm_state])
+        # Preserve that distinction here. In-place .set_() keeps the existing
+        # tensor identities valid for any captured refs (torch.compile graph,
+        # layer module attrs); we only fall back to assignment on first
+        # install or a shape mismatch.
+        old_kv = cast(
+            list[Any] | list[torch.Tensor] | torch.Tensor,
+            forward_context[layer_name].kv_cache,
+        )
+
+        if isinstance(new_kv, list):
+            if (
+                isinstance(old_kv, list)
+                and len(old_kv) == len(new_kv)
+                and all(isinstance(t, torch.Tensor) for t in old_kv)
+            ):
+                for old_t, new_t in zip(old_kv, new_kv, strict=True):
+                    old_t.set_(
+                        new_t.storage(),
+                        new_t.storage_offset(),
+                        new_t.shape,
+                        new_t.stride(),
+                    )
+            else:
+                forward_context[layer_name].kv_cache = new_kv
+        else:
+            if isinstance(old_kv, torch.Tensor) and old_kv.numel() > 0:
+                old_kv.set_(
+                    new_kv.storage(),
+                    new_kv.storage_offset(),
+                    new_kv.shape,
+                    new_kv.stride(),
+                )
+            else:
+                forward_context[layer_name].kv_cache = new_kv
+
+
+def _grow_block_pool(
+    block_pool: BlockPool, old_num_blocks: int, new_num_blocks: int
+) -> None:
+    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+
+    new_blocks: list[KVCacheBlock] = []
+    for idx in range(old_num_blocks, new_num_blocks):
+        block = KVCacheBlock(idx)
+        block_pool.blocks.append(block)
+        new_blocks.append(block)
+
+    block_pool.free_block_queue.append_n(new_blocks)
+    block_pool.num_gpu_blocks = new_num_blocks
+
+
+def _patch_moe_sum() -> None:
+    import vllm._custom_ops as ops
+
+    def moe_sum_f32(x: torch.Tensor, output: torch.Tensor) -> None:
+        output[:] = x.to(torch.float32).sum(dim=1).to(output.dtype)
+
+    ops.moe_sum = moe_sum_f32
+
+
+def _patch_marlin_w2_thread_config() -> None:
+    try:
+        import vllm._custom_ops as ops
+    except ImportError:
+        return
+
+    original_gemm = cast(Callable[..., object], ops.moe_wna16_marlin_gemm)
+
+    def patched_gemm(*args: object, **kwargs: object) -> object:
+        kwargs["thread_k"] = 64
+        kwargs["thread_n"] = 128
+        return original_gemm(*args, **kwargs)
+
+    ops.moe_wna16_marlin_gemm = patched_gemm

src/exo/worker/engines/vllm/kv_connector.py

@@ -0,0 +1,787 @@
+# pyright: reportAny = false
+import contextlib
+import queue
+import re
+from dataclasses import dataclass
+from typing import Any, cast
+
+import torch
+from vllm.config import VllmConfig
+from vllm.config.kv_transfer import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.request import Request
+
+from exo.worker.engines.vllm.disaggregated.adapter import (
+    extract_kv_via_slot_mapping,
+    to_bf16,
+)
+from exo.worker.runner.bootstrap import logger
+
+_LAYER_RE = re.compile(r"layers\.(\d+)\.")
+
+
+# Module-level shared state. Populated by the connector's hooks (running inside
+# vLLM's scheduler/worker, same process since V1 multiprocessing is off);
+# drained by the producer engine in `serve_prefill` after the request finishes.
+#
+# `_kv_queue` is the original streaming-connector path ported into this module.
+# We defer prefix reuse to vLLM APC and do not keep a separate TorchKVCache.
+# 4-tuple: (layer_idx, keys_host_pinned, values_host_pinned, copy_done_event)
+# The writer thread does `event.synchronize()` (CPU-side, doesn't block GPU)
+# before reading the pinned host bytes.
+# 5-tuple: (layer_idx, num_tokens, keys_host_pinned, values_host_pinned, copy_done_event)
+# `num_tokens` is the authoritative token count for this item. The writer uses
+# it for skip_tokens accounting *and* to slice the keys/values tensors before
+# writing to wire — never trusts `keys.shape[0]`, since shape can disagree with
+# token count when the source path packs/reshapes (e.g. NVFP4 layouts).
+_kv_queue: queue.Queue[
+    tuple[int, int, torch.Tensor, torch.Tensor, torch.cuda.Event] | None
+] = queue.Queue()
+# 3-tuple: (layer_idx, arrays_host_or_gpu, copy_done_event_or_none)
+# - From save_kv_layer hybrid path: tensors are GPU, event=None (writer .cpu()s)
+# - From GDN capture (after both conv+ssm ready): tensors are pinned host,
+#   event is a CUDA event the writer must synchronize on before reading
+_arrays_queue: queue.Queue[
+    tuple[int, list[torch.Tensor], torch.cuda.Event | None] | None
+] = queue.Queue()
+# Per-layer tracking of which layers' GDN states have been shipped via the
+# async pipeline. Entries here are excluded from the post-writer fallback drain.
+_gdn_shipped: set[int] = set()
+_captured_layers: dict[int, dict[str, torch.Tensor]] = {}
+_captured_arrays: dict[int, list[torch.Tensor]] = {}
+# Hybrid-model SSM/conv state captured via causal_conv1d + delta-rule patches.
+_gdn_states: dict[int, dict[str, torch.Tensor]] = {}
+_gdn_layer_order: list[int] = []
+_gdn_call_idx: list[int] = [0]
+_ssm_call_idx: list[int] = [0]
+# Per-layer save_kv_layer call diagnostics: list of slot_mapping sizes seen.
+_save_kv_layer_diag: dict[int, list[int]] = {}
+# Side CUDA stream for K/V extract + async D2H, so vLLM's compute stream
+# isn't blocked on D2H/extract during forward.
+_save_stream: torch.cuda.Stream | None = None
+# Holds a reference to the set tracked by patched_schedule so
+# `reset_capture_state` can clear it between requests.
+_apc_extracted_set_ref: dict[str, set[str]] = {}
+# request_id → actual APC hit token count (captured at the moment vLLM's
+# kv_cache_manager.get_computed_blocks runs, before scheduler chunks the
+# remaining tokens). Used by patched_schedule to pre-extract exactly the
+# matched portion, not the matched+about-to-forward portion.
+_apc_hit_tokens: dict[str, int] = {}
+
+
+def _get_save_stream() -> torch.cuda.Stream:
+    global _save_stream
+    if _save_stream is None:
+        _save_stream = torch.cuda.Stream()
+    return _save_stream
+
+
+def get_kv_queue() -> queue.Queue[
+    tuple[int, int, torch.Tensor, torch.Tensor, torch.cuda.Event] | None
+]:
+    return _kv_queue
+
+
+def get_arrays_queue() -> queue.Queue[
+    tuple[int, list[torch.Tensor], torch.cuda.Event | None] | None
+]:
+    return _arrays_queue
+
+
+def get_gdn_states() -> dict[int, dict[str, torch.Tensor]]:
+    return _gdn_states
+
+
+def get_gdn_shipped() -> set[int]:
+    return _gdn_shipped
+
+
+def _try_ship_gdn(layer_idx: int) -> None:
+    """If both conv and ssm have been captured for `layer_idx`, kick off an
+    async pinned D2H on the side stream and enqueue an arrays-state item so
+    the writer thread can ship the bytes during forward instead of after.
+
+    Called from BOTH the conv and ssm capture patches. Conv always fires
+    before ssm in a Mamba layer's forward, so this is a no-op after conv
+    (state lacks ssm) and ships once after ssm. For chunked prefill the
+    pair fires once per chunk: we ship every time, and the consumer's
+    `arrays[layer_idx] = ...` last-write-wins keeps the final-chunk state
+    (Mamba state is cumulative, only the final state matters).
+    """
+    state = _gdn_states.get(layer_idx)
+    if state is None or "conv" not in state or "ssm" not in state:
+        return
+    conv_gpu = state["conv"]
+    ssm_gpu = state["ssm"]
+    side_stream = _get_save_stream()
+    side_stream.wait_stream(torch.cuda.current_stream())  # pyright: ignore[reportUnknownMemberType]
+    with torch.cuda.stream(side_stream):
+        conv_host = torch.empty(conv_gpu.shape, dtype=conv_gpu.dtype, pin_memory=True)
+        ssm_host = torch.empty(ssm_gpu.shape, dtype=ssm_gpu.dtype, pin_memory=True)
+        conv_host.copy_(conv_gpu, non_blocking=True)
+        ssm_host.copy_(ssm_gpu, non_blocking=True)
+    event = torch.cuda.Event()
+    event.record(side_stream)
+    _arrays_queue.put((layer_idx, [conv_host, ssm_host], event))
+    _gdn_shipped.add(layer_idx)
+
+
+def get_save_kv_layer_diag() -> dict[int, list[int]]:
+    return _save_kv_layer_diag
+
+
+def get_captured_layers() -> dict[int, dict[str, torch.Tensor]]:
+    return _captured_layers
+
+
+def get_captured_arrays() -> dict[int, list[torch.Tensor]]:
+    return _captured_arrays
+
+
+def reset_capture_state() -> None:
+    while not _kv_queue.empty():
+        try:
+            _kv_queue.get_nowait()
+        except queue.Empty:
+            break
+    while not _arrays_queue.empty():
+        try:
+            _arrays_queue.get_nowait()
+        except queue.Empty:
+            break
+    _captured_layers.clear()
+    _captured_arrays.clear()
+    _gdn_states.clear()
+    _gdn_shipped.clear()
+    _gdn_call_idx[0] = 0
+    _ssm_call_idx[0] = 0
+    _save_kv_layer_diag.clear()
+    _apc_hit_tokens.clear()
+    apc_set = _apc_extracted_set_ref.get("set")
+    if apc_set is not None:
+        apc_set.clear()
+
+
+@dataclass
+class StreamingConnectorMetadata(KVConnectorMetadata):
+    pass
+
+
+@dataclass
+class BatchConnectorMetadata(KVConnectorMetadata):
+    pass
+
+
+class StreamingConnector(KVConnectorBase_V1, SupportsHMA):
+    """Original streaming producer connector, kept under the new server abstraction."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: KVCacheConfig | None = None,
+    ) -> None:
+        super().__init__(vllm_config, role, kv_cache_config)
+        self._save_count = 0
+
+    # =========================================================================
+    # Worker-side hooks (the only ones we actually use)
+    # =========================================================================
+
+    def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:
+        return
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        return
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: Any,
+        attn_metadata: Any,
+        **kwargs: Any,
+    ) -> None:
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)
+        m = _LAYER_RE.search(layer_name)
+        layer_idx_for_diag = int(m.group(1)) if m else -1
+        slot_size = int(slot_mapping.shape[0]) if slot_mapping is not None else -1
+        is_list_kv = isinstance(kv_layer, (list, tuple))
+        # Tag list/tuple as negative so the diag log distinguishes hybrid from
+        # non-hybrid even when slot_size is the same.
+        _save_kv_layer_diag.setdefault(layer_idx_for_diag, []).append(
+            -slot_size if is_list_kv else slot_size
+        )
+
+        # Skip decode-step saves (small slot mapping); we only want prefill.
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:
+            return
+        if m is None:
+            return
+        layer_idx = int(m.group(1))
+
+        # Hybrid (Mamba+attention) layers: kv_layer is a list/tuple of state
+        # tensors (conv + ssm). Send them straight to the arrays queue —
+        # they don't live in the paged KV cache. Stay on GPU; the writer
+        # thread does the D2H copy via `tensor_to_wire_bytes`.
+        if isinstance(kv_layer, (list, tuple)):
+            arrays = [
+                to_bf16(t)
+                for t in cast(list[torch.Tensor] | tuple[torch.Tensor, ...], kv_layer)
+            ]
+            _arrays_queue.put((layer_idx, arrays, None))
+            return
+
+        # Standard attention layers (full or sliding-window): extract K/V
+        # via slot_mapping, which points to where vLLM is *writing* this
+        # forward step's tokens. Capturing here, before sliding-window
+        # eviction in the block pool, is the only way to ship every prompt
+        # token's K/V regardless of attention type.
+        #
+        # All of this work — gather + bf16 cast + D2H — runs on a side
+        # CUDA stream into pinned host memory. vLLM's compute stream is
+        # never blocked: it only has to record-event for our side stream
+        # to wait on, then it continues into the next layer's forward.
+        # The writer thread later waits on the CUDA event (CPU-side wait,
+        # doesn't block GPU) and ships the already-on-host bytes.
+        if slot_mapping is not None:
+            try:
+                save_stream = _get_save_stream()
+                save_stream.wait_stream(torch.cuda.current_stream())  # pyright: ignore[reportUnknownMemberType] # TODO: stub
+                with torch.cuda.stream(save_stream):
+                    keys_gpu, values_gpu = extract_kv_via_slot_mapping(
+                        kv_layer, slot_mapping
+                    )
+                    keys_host = torch.empty(
+                        keys_gpu.shape, dtype=keys_gpu.dtype, pin_memory=True
+                    )
+                    values_host = torch.empty(
+                        values_gpu.shape, dtype=values_gpu.dtype, pin_memory=True
+                    )
+                    keys_host.copy_(keys_gpu, non_blocking=True)
+                    values_host.copy_(values_gpu, non_blocking=True)
+                    num_tokens = int(keys_gpu.shape[0])
+                event = torch.cuda.Event()
+                event.record(save_stream)
+            except Exception as exc:
+                logger.warning(
+                    f"save_kv_layer extract failed layer={layer_idx} "
+                    f"kv_layer.shape={getattr(kv_layer, 'shape', None)} "
+                    f"slot_mapping.shape={slot_mapping.shape}: {exc!r}"
+                )
+                return
+            _kv_queue.put((layer_idx, num_tokens, keys_host, values_host, event))
+
+    def wait_for_save(self) -> None:
+        return
+
+    # =========================================================================
+    # Scheduler-side hooks (no-ops; we don't load and don't track allocs)
+    # =========================================================================
+
+    def get_num_new_matched_tokens(
+        self, request: Any, num_computed_tokens: int
+    ) -> tuple[int, bool]:
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: Any, blocks: Any, num_external_tokens: int
+    ) -> None:
+        return
+
+    def build_connector_meta(self, scheduler_output: Any) -> StreamingConnectorMetadata:
+        return StreamingConnectorMetadata()
+
+    def request_finished(
+        self, request: Any, block_ids: list[int]
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return False, None
+
+    def request_finished_all_groups(
+        self, request: Any, block_ids: tuple[list[int], ...]
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return False, None
+
+
+class BatchConnector(KVConnectorBase_V1, SupportsHMA):
+    """Original batch producer connector, ported for parity with the old branch."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: KVCacheConfig | None = None,
+    ) -> None:
+        super().__init__(vllm_config, role, kv_cache_config)
+
+    def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:
+        return
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        return
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: Any,
+        attn_metadata: Any,
+        **kwargs: Any,
+    ) -> None:
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:
+            return
+
+        m = _LAYER_RE.search(layer_name)
+        if m is None:
+            return
+        layer_idx = int(m.group(1))
+
+        if isinstance(kv_layer, (list, tuple)):
+            _captured_arrays[layer_idx] = [
+                to_bf16(t).cpu()
+                for t in cast(list[torch.Tensor] | tuple[torch.Tensor, ...], kv_layer)
+            ]
+            return
+
+        if slot_mapping is None:
+            return
+        keys, values = extract_kv_via_slot_mapping(kv_layer, slot_mapping)
+        prev = _captured_layers.get(layer_idx)
+        if prev is None:
+            _captured_layers[layer_idx] = {"keys": keys, "values": values}
+        else:
+            _captured_layers[layer_idx] = {
+                "keys": torch.cat([prev["keys"], keys], dim=0),
+                "values": torch.cat([prev["values"], values], dim=0),
+            }
+
+    def wait_for_save(self) -> None:
+        return
+
+    def request_finished(
+        self, request: Any, block_ids: list[int]
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return False, None
+
+    def request_finished_all_groups(
+        self, request: Any, block_ids: tuple[list[int], ...]
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return False, None
+
+    def get_num_new_matched_tokens(
+        self, request: Any, num_computed_tokens: int
+    ) -> tuple[int, bool]:
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: Any, blocks: Any, num_external_tokens: int
+    ) -> None:
+        return
+
+    def build_connector_meta(self, scheduler_output: Any) -> BatchConnectorMetadata:
+        return BatchConnectorMetadata()
+
+
+ExoKVProducerConnector = StreamingConnector
+
+
+# =============================================================================
+# Bypass patches — necessary to make our connector usable inside vLLM 1.x.
+# Ported from the original branch's prefill_server.py:_patch_vllm_for_connector.
+# =============================================================================
+
+_connector_patched = False
+
+
+def _patch_vllm_for_connector(connector_class: type[Any]) -> None:
+    """Three patches that make a custom save-only connector cooperate with vLLM.
+
+    1. Suppress `unify_hybrid_kv_cache_specs` ValueError on hybrid (Mamba +
+       attention) models — the unifier complains about mixed cache specs we
+       don't need to actually unify for save-only operation.
+    2. Override `Scheduler._connector_finished` to short-circuit the
+       async-save state machine. We're synchronous on the producer side.
+    3. Make `KVConnectorFactory._get_connector_class_with_compat` recognize
+       our class name and return our class directly, bypassing vLLM's
+       registry of built-in connectors.
+    """
+    global _connector_patched
+    if _connector_patched:
+        return
+    _connector_patched = True
+
+    from vllm.v1.core import kv_cache_utils
+
+    original_unify = kv_cache_utils.unify_hybrid_kv_cache_specs
+
+    def patched_unify(kv_cache_spec: Any) -> None:
+        with contextlib.suppress(ValueError):
+            original_unify(kv_cache_spec)
+
+    kv_cache_utils.unify_hybrid_kv_cache_specs = patched_unify
+
+    from vllm.v1.core.sched import scheduler as sched_mod
+
+    def patched_connector_finished(
+        self: sched_mod.Scheduler, request: Request
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return False, None
+
+    sched_mod.Scheduler._connector_finished = patched_connector_finished  # pyright: ignore[reportPrivateUsage]
+
+    from vllm.distributed.kv_transfer.kv_connector import factory
+
+    original_get = factory.KVConnectorFactory._get_connector_class_with_compat  # pyright: ignore[reportPrivateUsage]
+
+    def patched_get(kv_transfer_config: KVTransferConfig) -> tuple[Any, Any]:
+        kv_conn = kv_transfer_config.kv_connector or ""
+        kv_conn_lower = kv_conn.lower()
+        if (
+            kv_conn
+            in {
+                connector_class.__name__,
+                f"{connector_class.__module__}:{connector_class.__name__}",
+                "ExoKVProducerConnector",
+                f"{__name__}:ExoKVProducerConnector",
+                "StreamingConnector",
+                f"{__name__}:StreamingConnector",
+            }
+            or "streaming_connector" in kv_conn_lower
+        ):
+            return connector_class, None
+        if "batch_connector" in kv_conn_lower:
+            return BatchConnector, None
+        return original_get(kv_transfer_config)
+
+    factory.KVConnectorFactory._get_connector_class_with_compat = patched_get  # pyright: ignore[reportPrivateUsage]
+
+    # Patch KVCacheManager.get_computed_blocks so we capture the actual APC-hit
+    # token count for each request at the moment vLLM looks it up — *before*
+    # the scheduler bumps `req.num_computed_tokens` with the chunked-prefill
+    # first-chunk size. Reading `req.num_computed_tokens` post-schedule yields
+    # `apc_hit + first_chunk` and would cause us to extract bytes from blocks
+    # that haven't been written yet for the first-chunk tail.
+    try:
+        from vllm.v1.core.kv_cache_manager import (  # pyright: ignore[reportMissingImports]
+            KVCacheManager,
+        )
+    except ImportError:
+        KVCacheManager = None  # noqa: N806
+
+    if KVCacheManager is not None:
+        original_get_computed_blocks = KVCacheManager.get_computed_blocks
+
+        def patched_get_computed_blocks(self: Any, request: Any) -> Any:
+            result = original_get_computed_blocks(self, request)
+            try:
+                req_id = getattr(request, "request_id", None)
+                if req_id is not None:
+                    if isinstance(result, tuple) and len(result) >= 2:  # pyright: ignore[reportUnknownArgumentType]
+                        num = int(result[1])  # pyright: ignore[reportUnknownArgumentType]
+                    else:
+                        num = 0
+                    total = int(getattr(request, "num_tokens", 0) or 0)
+                    logger.info(
+                        f"APC get_computed_blocks: req={req_id} hit={num} total={total}"
+                    )
+                    if num > 0:
+                        _apc_hit_tokens[req_id] = num
+            except Exception:
+                logger.opt(exception=True).warning(
+                    "patched_get_computed_blocks: capture failed"
+                )
+            return result
+
+        KVCacheManager.get_computed_blocks = patched_get_computed_blocks  # pyright: ignore[reportAttributeAccessIssue]
+
+    # Patch Scheduler.schedule so APC-cached prefix blocks are extracted out
+    # of the paged pool and pushed to _kv_queue at scheduling time — BEFORE
+    # forward runs. Forward will only execute the suffix (vLLM's own APC
+    # behavior). save_kv_layer fires for the suffix as usual. The writer
+    # thread sees: prefix items from this hook + suffix items from save_kv_layer
+    # and ships them in arrival order (prefix before suffix per layer).
+    original_schedule = sched_mod.Scheduler.schedule
+    _scheduled_apc_extracted: set[str] = set()
+
+    def patched_schedule(self: sched_mod.Scheduler) -> Any:
+        scheduler_output = original_schedule(self)
+        try:
+            new_reqs = getattr(scheduler_output, "scheduled_new_reqs", None) or []
+            if not new_reqs:
+                return scheduler_output
+            from exo.worker.engines.vllm.disaggregated.adapter import (
+                build_layer_to_group,
+                gather_layer_kv_from_blocks,
+            )
+            from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+            mr = get_model_runner()
+            if mr is None:
+                return scheduler_output
+            cfg = getattr(mr, "_growable_kv_cache_config", None)
+            if cfg is None:
+                return scheduler_output
+            layer_to_group = build_layer_to_group(cfg)
+            n_layers = len(mr.kv_caches)
+
+            for new_req in new_reqs:
+                req_id = getattr(new_req, "req_id", None)
+                if req_id is None or req_id in _scheduled_apc_extracted:
+                    continue
+                pre_layers_shipped = 0
+                pre_bytes_shipped = 0
+                req = self.requests.get(req_id)
+                if req is None:
+                    continue
+                # Use the count captured by patched_get_computed_blocks (the
+                # actual APC hit), NOT req.num_computed_tokens — that field has
+                # already been bumped by the scheduler with the first chunk's
+                # about-to-forward token count and would over-extract.
+                num_apc = _apc_hit_tokens.get(req_id, 0)
+                req_total = int(getattr(req, "num_tokens", 0) or 0)
+                req_computed = int(getattr(req, "num_computed_tokens", 0) or 0)
+                logger.info(
+                    f"APC patched_schedule: req={req_id} apc_hit={num_apc} "
+                    f"req.num_computed_tokens={req_computed} req.num_tokens={req_total}"
+                )
+                if num_apc <= 0:
+                    _scheduled_apc_extracted.add(req_id)
+                    continue
+                # Pull the request's full per-group block list from
+                # scheduler_output.scheduled_new_reqs[i].block_ids — that field
+                # includes APC-cached prefix blocks. The KVCacheManager's
+                # `req_to_blocks` only tracks newly-allocated blocks for this
+                # step's suffix, so reading from there misses the prefix and
+                # makes gather return ~bock_count_suffix tokens of garbage.
+                req_block_ids_per_group: tuple[list[int], ...] | None = getattr(
+                    new_req, "block_ids", None
+                )
+                if not req_block_ids_per_group:
+                    logger.warning(
+                        f"APC pre-extract: new_req.block_ids missing for {req_id}"
+                    )
+                    _scheduled_apc_extracted.add(req_id)
+                    continue
+                save_stream = _get_save_stream()
+                save_stream.wait_stream(torch.cuda.current_stream())  # pyright: ignore[reportUnknownMemberType]
+                first_log_done = False
+                # Run the entire gather + cast + pinned alloc + D2H on the
+                # side stream — scheduler thread only issues kernel launches
+                # and records an event per layer. Compute stream is untouched.
+                with torch.cuda.stream(save_stream):
+                    for layer_idx in range(n_layers):
+                        kv_layer = mr.kv_caches[layer_idx]
+                        if isinstance(kv_layer, (list, tuple)):
+                            continue
+                        gi = (
+                            layer_to_group[layer_idx]
+                            if layer_idx < len(layer_to_group)
+                            else 0
+                        )
+                        if gi >= len(req_block_ids_per_group):
+                            continue
+                        block_ids = list(req_block_ids_per_group[gi])
+                        if not block_ids:
+                            continue
+                        keys_gpu, values_gpu = gather_layer_kv_from_blocks(
+                            kv_layer, block_ids, num_apc
+                        )
+                        if not first_log_done:
+                            first_log_done = True
+                            logger.info(
+                                f"APC pre-extract layer={layer_idx}: "
+                                f"kv_layer.shape={tuple(kv_layer.shape)} "
+                                f"kv_layer.dtype={kv_layer.dtype} "
+                                f"len(block_ids)={len(block_ids)} num_apc={num_apc} "
+                                f"keys_gpu.shape={tuple(keys_gpu.shape)} "
+                                f"keys_gpu.dtype={keys_gpu.dtype}"
+                            )
+                        if keys_gpu.numel() == 0:
+                            continue
+                        keys_host = torch.empty(
+                            keys_gpu.shape,
+                            dtype=keys_gpu.dtype,
+                            pin_memory=True,
+                        )
+                        values_host = torch.empty(
+                            values_gpu.shape,
+                            dtype=values_gpu.dtype,
+                            pin_memory=True,
+                        )
+                        keys_host.copy_(keys_gpu, non_blocking=True)
+                        values_host.copy_(values_gpu, non_blocking=True)
+                        event = torch.cuda.Event()
+                        event.record(save_stream)
+                        _kv_queue.put(
+                            (layer_idx, num_apc, keys_host, values_host, event)
+                        )
+                        pre_layers_shipped += 1
+                        pre_bytes_shipped += (
+                            keys_host.numel() * keys_host.element_size()
+                            + values_host.numel() * values_host.element_size()
+                        )
+                logger.info(
+                    f"APC pre-extract done: req={req_id} layers={pre_layers_shipped} "
+                    f"tokens={num_apc} bytes={pre_bytes_shipped}"
+                )
+                _scheduled_apc_extracted.add(req_id)
+        except Exception:
+            logger.opt(exception=True).warning(
+                "patched_schedule: APC pre-extract failed; continuing"
+            )
+        return scheduler_output
+
+    sched_mod.Scheduler.schedule = patched_schedule
+    # Reset the per-request-extracted set when reset_capture_state runs.
+    _apc_extracted_set_ref["set"] = _scheduled_apc_extracted
+    logger.info("Installed vLLM connector bypass patches")
+
+
+# =============================================================================
+# Hybrid-model GDN state capture (Qwen3.5/3.6 etc.).
+# Patches the conv1d kernel + delta-rule fns to grab conv/ssm states per layer.
+# =============================================================================
+
+_gdn_patched = False
+
+
+def _patch_gdn_capture() -> None:
+    global _gdn_patched
+    if _gdn_patched:
+        return
+    _gdn_patched = True
+
+    try:
+        import vllm.model_executor.layers.mamba.ops.causal_conv1d as cc_mod
+        from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+            causal_conv1d_fn as orig_fn,
+        )
+    except ImportError:
+        return
+
+    def patched_fn(
+        *args: Any, conv_states: Any = None, cache_indices: Any = None, **kwargs: Any
+    ) -> Any:
+        result = orig_fn(
+            *args, conv_states=conv_states, cache_indices=cache_indices, **kwargs
+        )
+        if conv_states is not None and cache_indices is not None:
+            x = args[0] if args else None
+            if x is not None and x.shape[0] <= 100:
+                return result
+            ci: int = cache_indices[0].item() if cache_indices.numel() > 0 else 0
+            idx = _gdn_call_idx[0]
+            if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                # `.contiguous()` decouples the slice from the underlying
+                # buffer; D2H is deferred to the writer thread.
+                conv_at_ci = conv_states[ci : ci + 1].transpose(-1, -2).contiguous()
+                _gdn_states.setdefault(layer_idx, {})["conv"] = conv_at_ci
+                _gdn_states[layer_idx]["ci"] = ci
+                # Don't ship from here: conv fires before ssm in a Mamba
+                # forward, so state["ssm"] is either missing (chunk 1) or
+                # stale from the previous chunk (chunk N>=2). Shipping here
+                # would emit a mismatched (conv_N, ssm_{N-1}) pair that the
+                # ssm patch's later ship would overwrite. Just wait for ssm.
+            _gdn_call_idx[0] += 1
+        return result
+
+    cc_mod.causal_conv1d_fn = patched_fn
+    import sys
+
+    for mod in list(sys.modules.values()):
+        if mod is cc_mod:
+            continue
+        # transformers' image_processing_* shims have a lazy __getattr__
+        # that emits a noisy deprecation warning on every attribute probe.
+        # They never use causal_conv1d_fn, so skip them.
+        mod_name = getattr(mod, "__name__", "") or ""
+        if mod_name.startswith("transformers."):
+            continue
+        if (
+            mod.__dict__.get("causal_conv1d_fn") is orig_fn
+            if hasattr(mod, "__dict__")
+            else False
+        ):
+            mod.causal_conv1d_fn = patched_fn
+    logger.info("Patched causal_conv1d_fn for GDN conv-state capture")
+
+    # The GDN delta-rule functions live in `mamba/gdn_linear_attn` (defined or
+    # re-imported there) and may also be re-exported by model modules. Patch
+    # all candidate modules + propagate to anywhere they're imported.
+    candidate_modules = [
+        "vllm.model_executor.layers.mamba.gdn_linear_attn",
+        "vllm.model_executor.models.qwen3_next",
+        "vllm.model_executor.models.qwen3_5",
+    ]
+    fn_names = ("fi_chunk_gated_delta_rule", "fla_chunk_gated_delta_rule")
+    patched_targets: list[str] = []
+
+    for mod_path in candidate_modules:
+        try:
+            mod = __import__(mod_path, fromlist=["*"])
+        except ImportError:
+            continue
+        for fn_name in fn_names:
+            orig = getattr(mod, fn_name, None)
+            if orig is None:
+                continue
+
+            def make_patched(orig_fn_inner: Any) -> Any:
+                def patched_chunk(*args: Any, **kwargs: Any) -> Any:
+                    result = orig_fn_inner(*args, **kwargs)
+                    output_final_state = kwargs.get("output_final_state", False)
+                    if (
+                        output_final_state
+                        and isinstance(result, tuple)
+                        and len(result) == 2  # pyright: ignore[reportUnknownArgumentType]
+                    ):
+                        _, ssm_state = result  # pyright: ignore[reportUnknownVariableType]
+                        idx = _ssm_call_idx[0]
+                        if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                            layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                            _gdn_states.setdefault(layer_idx, {})["ssm"] = ssm_state
+                            _try_ship_gdn(layer_idx)
+                        _ssm_call_idx[0] += 1
+                    return result  # pyright: ignore[reportUnknownVariableType]
+
+                return patched_chunk
+
+            patched_fn = make_patched(orig)
+            setattr(mod, fn_name, patched_fn)
+            patched_targets.append(f"{mod_path}.{fn_name}")
+            # Propagate to any module that imported the original function.
+            import sys as _sys
+
+            for other in list(_sys.modules.values()):
+                if other is mod:
+                    continue
+                other_name = getattr(other, "__name__", "") or ""
+                # Skip transformers — see causal_conv1d_fn loop above.
+                if other_name.startswith("transformers."):
+                    continue
+                if other.__dict__.get(fn_name) is orig:
+                    setattr(other, fn_name, patched_fn)
+                    patched_targets.append(f"{other.__name__}.{fn_name} (propagated)")
+    if patched_targets:
+        logger.info(f"Patched delta-rule fns for SSM capture: {patched_targets}")
+    else:
+        logger.warning(
+            "GDN SSM-capture patch installed no targets — hybrid models may miss ssm state"
+        )
+
+
+def init_gdn_layer_order(kv_caches: Any) -> None:
+    """Identify hybrid layers (those with list/tuple kv_cache entries)."""
+    _gdn_layer_order.clear()
+    for li in range(len(kv_caches)):
+        kv = kv_caches[li]
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:  # pyright: ignore[reportUnknownArgumentType]
+            _gdn_layer_order.append(li)
+    if _gdn_layer_order:
+        logger.info(f"GDN layer order: {len(_gdn_layer_order)} hybrid layers detected")

src/exo/worker/engines/vllm/prompt_format.py

@@ -0,0 +1,63 @@
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    get_eos_token_ids_for_model,
+)
+
+
+def format_vllm_prompt(
+    engine: LLMEngine, params: TextGenerationTaskParams
+) -> tuple[list[int], str, int]:
+    # we should have our own wrapper
+    # (instead of abusing mlx's TokenizerWrapper, use tokenizers Tokenizer)
+    tokenizer = TokenizerWrapper(engine.get_tokenizer())
+    prompt_text = apply_chat_template(tokenizer, params)
+    token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)
+    return token_ids, prompt_text, len(token_ids)
+
+
+def make_vllm_sampling_params(
+    engine: LLMEngine,
+    params: TextGenerationTaskParams,
+    model_id: ModelId | None = None,
+) -> SamplingParams:
+    kwargs: SamplingParams = SamplingParams()
+
+    if params.max_output_tokens is not None:
+        kwargs.max_tokens = params.max_output_tokens
+    else:
+        kwargs.max_tokens = min(engine.model_config.max_model_len, 32168)
+    if params.temperature is not None:
+        kwargs.temperature = params.temperature
+    if params.top_p is not None:
+        kwargs.top_p = params.top_p
+    if params.top_k is not None:
+        kwargs.top_k = params.top_k
+    if params.min_p is not None:
+        kwargs.min_p = params.min_p
+    if params.stop is not None:
+        kwargs.stop = params.stop
+    if params.seed is not None:
+        kwargs.seed = params.seed
+    if params.repetition_penalty is not None:
+        kwargs.repetition_penalty = params.repetition_penalty
+    if params.logprobs:
+        kwargs.logprobs = params.top_logprobs or 1
+
+    if model_id is not None:
+        extra_stop = get_eos_token_ids_for_model(model_id)
+        if extra_stop:
+            kwargs.stop_token_ids = extra_stop
+
+    if params.bench:
+        kwargs.ignore_eos = True
+        kwargs.min_tokens = kwargs.max_tokens
+        if not params.use_prefix_cache:
+            kwargs.skip_reading_prefix_cache = True
+
+    return kwargs