dont push to cachix

## Why

Workers currently update their custom model-card cache by reacting to
`CustomModelCardAdded` / `CustomModelCardDeleted` events directly. That
is another snapshot footgun: a worker restored from State may never see
the historical add/delete event, so the durable State must include the
desired custom-card set.

## How

- Add `State.custom_model_cards`, keyed by `ModelId`.
- Reduce `CustomModelCardAdded` into State.
- Reduce `CustomModelCardDeleted` into State.
- Add focused reducer tests for add and delete.

This PR only makes custom cards durable in State. A follow-up PR will
make workers reconcile their on-disk custom-card cache from this state
instead of relying on those events directly.

## Tests

- `uv run pytest
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py
src/exo/shared/tests/test_state_serialization.py`
- `uv run pytest`
- `uv run ruff check src/exo/shared/types/state.py
src/exo/shared/apply.py
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py`
- `uv run basedpyright`
- `nix fmt`

.envrc

@@ -1 +1,8 @@
 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

.github/workflows/pipeline.yml

@@ -36,6 +36,7 @@ jobs:
         with:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+          pushFilter: '-docker-image.tar.gz$'
 
       - name: Build Metal packages (macOS only)
         if: runner.os == 'macOS'

.mlx_typings/mlx_lm/models/cache.pyi

@@ -191,10 +191,13 @@ 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: tuple[str, ...]) -> None: ...
-    def is_trimmable(self) -> bool: ...
+    def meta_state(self, v):  # -> None:
+        ...
+    def is_trimmable(self):  # -> bool:
+        ...
     def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...

app/EXO/EXO/ContentView.swift

@@ -16,22 +16,13 @@ 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 = ""
@@ -294,6 +285,13 @@ 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()
@@ -477,40 +475,6 @@ 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
@@ -559,127 +523,6 @@ 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: {
@@ -720,61 +563,6 @@ 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"
@@ -857,13 +645,6 @@ 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,6 +22,7 @@ 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)
@@ -46,6 +47,7 @@ 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()
@@ -66,6 +68,7 @@ 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 = .shared
+        session: URLSession = ClusterStateService.makeNonCachingSession()
     ) {
         self.baseURL = baseURL
         self.endpoint = baseURL.appendingPathComponent("state")
@@ -27,6 +27,23 @@ 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

@@ -0,0 +1,242 @@
+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,8 +21,6 @@ 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 {
@@ -202,8 +200,6 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
-
-                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -504,50 +500,8 @@ 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/harness.py

@@ -589,9 +589,7 @@ 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", "vllm", "both"],
-        default="both",
+        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
     )
     ap.add_argument(
         "--sharding", choices=["pipeline", "tensor", "both"], default="both"

bench/prefill-decode.toml

@@ -12,24 +12,24 @@ timeout = 7200.0
 settle_timeout = 60.0
 
 # Workload
-pp = [4096, 8192]
-tg = [128]
+pp = [4096]
+tg = [512]
 repeat = 1
 warmup = 0
 
 json_out = "bench/prefill_decode_results.json"
 
 [prefill]
-model = "sakamakismile/Qwen3.6-27B-NVFP4"
-node = "gx10-de89"
-instance_meta = "vllm"
+model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+node = "mike"
+instance_meta = "ring"
 sharding = "pipeline"
 min_nodes = 1
 max_nodes = 1
 
 [decode]
-model = "mlx-community/Qwen3.6-27B-4bit"
-node = "Ryuichi’s MacBook Pro"
+model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+node = "james"
 instance_meta = "ring"
 sharding = "pipeline"
 min_nodes = 1

bench/prefill_decode_bench.py

@@ -31,7 +31,6 @@ from typing import Any
 
 from exo_bench import (
     PromptSizer,
-    SystemMetricsSampler,
     format_peak_memory,
     load_tokenizer_for_bench,
     parse_int_list,
@@ -279,7 +278,6 @@ 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]] = []
@@ -290,13 +288,10 @@ 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
@@ -320,26 +315,11 @@ 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{energy_str}"
+                f"avg_elapsed={avg_elapsed:.2f}s"
             )
         time.sleep(2)
     return rows
@@ -352,36 +332,14 @@ 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]],
@@ -392,17 +350,14 @@ def _print_diff(
     prefill_alone = _summarise(prefill_alone_rows)
     keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
 
-    width = 110
+    width = 64
     for key in sorted(keys):
         pp, tg = key
         logger.info("─" * width)
         logger.info(f"  pp={pp}  tg={tg}")
         logger.info("─" * width)
         logger.info(
-            f"  {'phase':<16} {'elapsed':>9}  {'norm':>9}  "
-            f"{'prompt_tps':>11}  {'gen_tps':>8}  "
-            f"{'p_tok':>6}  {'g_tok':>6}  "
-            f"{'energy':>9}  {'avg_W':>7}"
+            f"  {'phase':<16} {'elapsed':>10}  {'prompt_tps':>11}  {'gen_tps':>9}"
         )
         for label, summary in (
             ("disaggregated", disagg.get(key)),
@@ -410,51 +365,26 @@ def _print_diff(
             ("prefill_alone", prefill_alone.get(key)),
         ):
             if summary is None:
-                logger.info(
-                    f"  {label:<16} {'—':>9}  {'—':>9}  "
-                    f"{'—':>11}  {'—':>8}  {'—':>6}  {'—':>6}  "
-                    f"{'—':>9}  {'—':>7}"
-                )
+                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
                 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']:>8.2f}s  "
-                f"{norm_str}  "
+                f"{summary['elapsed_s']:>9.2f}s  "
                 f"{summary['prompt_tps']:>11.1f}  "
-                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}"
+                f"{summary['gen_tps']:>9.2f}"
             )
 
         d = disagg.get(key)
         da = decode_alone.get(key)
         pa = prefill_alone.get(key)
-        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)"
-                )
+        if d and da and d["elapsed_s"] > 0:
+            logger.info(
+                f"  speedup vs decode_alone:  {da['elapsed_s'] / d['elapsed_s']:.2f}x"
+            )
+        if d and pa and d["elapsed_s"] > 0:
+            logger.info(
+                f"  speedup vs prefill_alone: {pa['elapsed_s'] / d['elapsed_s']:.2f}x"
+            )
     logger.info("─" * width)
 
 
@@ -752,16 +682,6 @@ 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})
@@ -780,7 +700,6 @@ def main() -> int:
                 warmup=args.warmup,
                 repeat=args.repeat,
                 common_meta=common_meta,
-                sampler=sampler,
             )
             all_rows.extend(prefill_alone_rows)
 
@@ -810,7 +729,6 @@ def main() -> int:
             warmup=args.warmup,
             repeat=args.repeat,
             common_meta=common_meta,
-            sampler=sampler,
         )
         all_rows.extend(disagg_rows)
 
@@ -835,13 +753,11 @@ 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,7 +202,6 @@
     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" | "dgx spark" | "linux" etc. */
+    /** "macbook pro" | "mac studio" | "mac mini" etc. */
     deviceType: string;
     /** Center X coordinate in SVG space */
     cx: number;
@@ -38,43 +38,10 @@
   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);
@@ -147,264 +114,7 @@
   const studioClipId = $derived(`di-studio-${uid}`);
 </script>
 
-{#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"}
+{#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" | "Vllm";
+    runtime?: "MlxRing" | "MlxJaccl";
     onLaunch?: () => void;
     tags?: string[];
     apiPreview?: PlacementPreview | null;
@@ -168,10 +168,8 @@
 
   function getDeviceType(
     name: string,
-  ): "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown" {
+  ): "macbook" | "studio" | "mini" | "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";
@@ -578,17 +576,13 @@
         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)."
-          : runtime === "MlxJaccl"
-            ? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
-            : "vLLM: NVIDIA CUDA inference engine."}
+          : "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
       >
         {runtime === "MlxRing"
           ? "MLX Ring"
           : runtime === "MlxJaccl"
             ? "MLX RDMA"
-            : runtime === "Vllm"
-              ? "vLLM"
-              : runtime}
+            : runtime}
       </span>
     </div>
 
@@ -996,81 +990,6 @@
                     />
                   {/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,7 +9,6 @@
     capabilities?: string[];
     family?: string;
     is_custom?: boolean;
-    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -20,7 +19,6 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
-    requiresVllm: boolean;
   }
 
   type DownloadAvailability = {
@@ -215,14 +213,6 @@
         <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"}
@@ -533,15 +523,6 @@
             {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(
@@ -647,7 +628,6 @@
                 variants: [variant],
                 smallestVariant: variant,
                 hasMultipleVariants: false,
-                requiresVllm: variant.requires_vllm === true,
               });
             }}
             title="View variant details"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -22,7 +22,6 @@
     is_custom?: boolean;
     tasks?: string[];
     hugging_face_id?: string;
-    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -33,7 +32,6 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
-    requiresVllm: boolean;
   }
 
   interface FilterState {
@@ -398,7 +396,6 @@
           variants: [],
           smallestVariant: model,
           hasMultipleVariants: false,
-          requiresVllm: true,
         });
       }
 
@@ -433,7 +430,6 @@
           (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)
@@ -591,7 +587,6 @@
           variants: [model],
           smallestVariant: model,
           hasMultipleVariants: false,
-          requiresVllm: model.requires_vllm === true,
         });
       }
     }
@@ -1170,17 +1165,6 @@
           <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-bound pass that consumes the entire prompt and
+        Prefill is the compute-heavy pass that consumes the entire prompt and
         builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
         tokens sequentially from that cache. The two phases have very different
         bottlenecks, so running them on different hardware can be substantially

dashboard/src/lib/components/TopologyGraph.svelte

@@ -117,10 +117,6 @@
   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;
@@ -558,13 +554,6 @@
       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);
@@ -634,382 +623,7 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      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") {
+      if (modelLower === "mac studio") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
         iconBaseHeight = nodeRadius * 0.85;
@@ -1568,12 +1182,8 @@
           debugLabelY += debugLineHeight;
         }
 
-        const dbgIdentity = identitiesData[nodeInfo.id];
-        if (dbgIdentity?.osVersion) {
-          const osLabel =
-            dbgIdentity.osVersion === "Linux"
-              ? "Linux"
-              : `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
+        const identity = identitiesData[nodeInfo.id];
+        if (identity?.osVersion) {
           nodeG
             .append("text")
             .attr("x", nodeInfo.x)
@@ -1582,7 +1192,9 @@
             .attr("fill", "rgba(179,179,179,0.7)")
             .attr("font-size", debugFontSize)
             .attr("font-family", "SF Mono, Monaco, monospace")
-            .text(osLabel);
+            .text(
+              `macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
+            );
         }
       }
     });

dashboard/src/routes/+page.svelte

@@ -65,7 +65,6 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
-    featureFlags,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -703,10 +702,7 @@
       ? Object.keys(topologyData()!.nodes).length
       : 1;
     const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
-    const instanceType =
-      nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
-        ? "MlxRing"
-        : selectedInstanceType;
+    const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
     try {
       const placementResponse = await fetch(
         `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
@@ -787,7 +783,6 @@
       quantization?: string;
       base_model?: string;
       capabilities?: string[];
-      requires_vllm?: boolean;
     }>
   >([]);
   type ModelMemoryFitStatus =
@@ -891,7 +886,7 @@
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
-  type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
+  type InstanceMeta = "MlxRing" | "MlxJaccl";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
@@ -937,12 +932,7 @@
     // Apply sharding and instance type unconditionally
     selectedSharding = defaults.sharding;
     selectedInstanceType =
-      defaults.instanceType === "MlxRing"
-        ? "MlxRing"
-        : defaults.instanceType === "Vllm"
-          ? "Vllm"
-          : "MlxJaccl";
-    userPickedInstanceType = true;
+      defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
 
     // Apply minNodes if valid (between 1 and maxNodes)
     if (
@@ -964,23 +954,6 @@
   }
 
   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);
@@ -1173,7 +1146,9 @@
   }
 
   const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
-    runtime === selectedInstanceType;
+    selectedInstanceType === "MlxRing"
+      ? runtime === "MlxRing"
+      : runtime === "MlxJaccl";
 
   // Helper to check if a model can be launched (has valid placement with >= minNodes)
   function canModelFit(modelId: string): boolean {
@@ -2088,7 +2063,6 @@
     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?: {
@@ -3461,6 +3435,7 @@
             >
               <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"
@@ -4848,6 +4823,7 @@
                 >
                   <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"
@@ -4994,6 +4970,7 @@
                   >
                     <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
@@ -5795,18 +5772,14 @@
                     </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 {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'}"
+                        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'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5822,18 +5795,14 @@
                         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 {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'}"
+                        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'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5848,41 +5817,7 @@
                         </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.exo.passthru.evenv
+                self'.packages.editableVenv
                 uv
 
                 # RUST

justfile

@@ -40,19 +40,6 @@ 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

@@ -1,26 +0,0 @@
-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,18 +15,21 @@ 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",
-  "tiktoken>=0.12.0",                # required for kimi k2 tokenizer
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
+  "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]
@@ -45,30 +48,26 @@ dev = [
 
 [project.optional-dependencies]
 build = ["nanobind"]
-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'",
-
+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-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'",
+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'",
 ]
 
 ###
@@ -82,23 +81,12 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 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-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')" },
+  { 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'" },
 ]
+vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 
 [[tool.uv.index]]
 name = "pytorch-cu130"
@@ -106,8 +94,8 @@ url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 [[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
+name = "pytorch-cu120"
+url = "https://download.pytorch.org/whl/cu120"
 explicit = true
 
 [[tool.uv.index]]
@@ -168,19 +156,11 @@ root = "src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
-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" },
-  ],
+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'",
 ]
 
 [tool.uv.extra-build-dependencies]
@@ -195,7 +175,6 @@ mlx = [
   "ninja",
 ]
 mlx-lm = ["setuptools"]
-mflux = ["uv_build"]
 xgrammar = [
   "nanobind",
   "setuptools",

python/parts.nix

@@ -10,10 +10,8 @@ let
       inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
       inherit (pkgs.config) cudaSupport;
       inherit (pkgs) cudaPackages;
-      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";
+      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
+      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
       python = pkgs.python313;
       cudaLibs = with cudaPackages; [
         cuda_cudart
@@ -115,213 +113,37 @@ 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: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
-            buildInputs = old.buildInputs ++ cudaLibs;
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           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";
@@ -342,30 +164,43 @@ let
           buildSystemsOverlay
         ]
       );
-      # 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 {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: 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 ];
-        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" ]; });
-        };
+        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
       };
     in
     {
       inherit venv;
+      editablePythonSet = pythonSet.overrideScope editableOverlay;
       mkPythonScript = path: mkApp ''python ${path} "$@"'';
-      mkExo = mkApp ''exo "$@"'';
+      mkOutputs = name:
+        let package = mkApp ''exo "$@"'' name;
+        in {
+          ${name} = package;
+          "${name}-docker-image" = pkgs.dockerTools.buildLayeredImage {
+            name = "${name}-docker-image";
+            config = {
+              Entrypoint = [ (lib.getExe package) ];
+              Env = [
+                "EXO_HOME=/var/lib/${name}"
+              ];
+            };
+
+            extraCommands = ''
+              mkdir -p  var/lib/${name}
+            '';
+          };
+        };
     };
 in
 {
@@ -373,18 +208,18 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
+      inherit (mkPythonSet { inherit self' pkgs lib; members.exo = [ "cpu" ]; }) editablePythonSet mkOutputs;
 
       # Virtual environment with dev dependencies for testing
       testVenv = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).venv "exo-test";
 
       mkBenchScript = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "mlx-cpu" "vllm-none" ];
+        exo = [ "cpu" ];
         exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).mkPythonScript;
@@ -394,12 +229,15 @@ 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" ]; }; };
+
+      defaultOutputs = mkOutputs "exo";
+      cuda12Outputs = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members.exo=["cuda12"]; }).mkOutputs "exo-cuda-12";
+      cuda13Outputs = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members.exo=["cuda13"]; }).mkOutputs "exo-cuda-13";
     in
     {
       packages = {
-        exo = mkExo "exo";
+        # for devShell
+        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);
@@ -407,10 +245,9 @@ in
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         # 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 = cuda12Set.mkExo "exo-cuda-12";
-        exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
-      };
+      } // defaultOutputs
+      // lib.optionalAttrs isLinux cuda12Outputs
+      // lib.optionalAttrs isLinux cuda13Outputs;
 
       checks = {
         lint = pkgs.runCommand "ruff-lint" { } ''

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

@@ -1,21 +0,0 @@
-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

@@ -1,27 +0,0 @@
-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

@@ -1,20 +0,0 @@
-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

@@ -1,32 +0,0 @@
-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

scripts/check_serve_prefill.py

@@ -1,222 +0,0 @@
-#!/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

@@ -1,124 +0,0 @@
-#!/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

@@ -1,170 +0,0 @@
-#!/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

@@ -133,12 +133,10 @@ 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 (
@@ -481,6 +479,7 @@ 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 +525,8 @@ class API:
                         )
                     ]
                 )
-        if any(self.state.node_vllm.values()):
-            instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
+        # TODO: PDD
+        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
             try:
@@ -544,6 +543,7 @@ 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:
@@ -640,10 +640,7 @@ class API:
         )
 
     async def get_feature_flags(self) -> dict[str, bool]:
-        return {
-            "disaggregation": ENABLE_DISAGGREGATION,
-            "vllm_available": any(self.state.node_vllm.values()),
-        }
+        return {"disaggregation": ENABLE_DISAGGREGATION}
 
     async def list_instance_links(self) -> list[InstanceLink]:
         if not ENABLE_DISAGGREGATION:
@@ -1636,17 +1633,16 @@ class API:
     async def ollama_tags(self) -> OllamaTagsResponse:
         """Returns list of models in Ollama tags format. We return the downloaded ones only."""
 
-        def none_if_empty(value: str) -> str | None:
-            return value or None
-
-        downloaded_model_ids: set[str] = set()
+        downloaded_model_ids: set[ModelId] = 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 get_model_cards() if c.model_id in downloaded_model_ids
+            c
+            for c in await model_cards.card_cache.list_all()
+            if c.model_id in downloaded_model_ids
         ]
 
         now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -1659,8 +1655,8 @@ class API:
                     size=card.storage_size.in_bytes,
                     digest="sha256:000000000000",
                     details=OllamaModelDetails(
-                        family=none_if_empty(card.family),
-                        quantization_level=none_if_empty(card.quantization),
+                        family=card.family or None,
+                        quantization_level=card.quantization or None,
                     ),
                 )
                 for card in cards
@@ -1723,7 +1719,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 get_model_cards()
+        cards = await model_cards.card_cache.list_all()
 
         if status == "downloaded":
             downloaded_model_ids: set[str] = set()
@@ -1751,7 +1747,6 @@ class API:
                     capabilities=card.capabilities,
                     reasoning_dialect=card.reasoning_dialect,
                     context_length=card.context_length,
-                    requires_vllm=card.requires_vllm,
                 )
                 for card in cards
             ]
@@ -1775,7 +1770,7 @@ class API:
 
         # Immediately update the local cache so the subsequent GET /models
         # returns the new model without waiting for the event round-trip.
-        add_to_card_cache(card)
+        model_cards.card_cache.cc[card.model_id] = card
 
         return ModelListModel(
             id=card.model_id,
@@ -1791,7 +1786,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 = get_card(model_id)
+        card = model_cards.card_cache.get(model_id)
         if card is None or not card.is_custom:
             raise HTTPException(status_code=404, detail="Custom model card not found")
 

src/exo/api/types/api.py

@@ -49,7 +49,6 @@ 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,7 +16,8 @@ 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.model_cards import ModelId, get_model_cards
+from exo.shared.models import model_cards
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import (
     CancelDownload,
     DeleteDownload,
@@ -422,7 +423,7 @@ class DownloadCoordinator:
                     )
                 # Scan read-only directories for pre-downloaded models
                 if EXO_MODELS_READ_ONLY_DIRS:
-                    for card in await get_model_cards():
+                    for card in await model_cards.card_cache.list_all():
                         mid = card.model_id
                         if mid in self.active_downloads:
                             continue

src/exo/download/download_utils.py

@@ -1,11 +1,12 @@
 import asyncio
 import hashlib
 import os
+import random
 import shutil
 import ssl
 import time
 import traceback
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Mapping
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -55,6 +56,36 @@ 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()
@@ -348,9 +379,6 @@ 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",
@@ -360,13 +388,16 @@ 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}"
 
-    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())
+    # 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 skip_internet:
         if await aios.path.exists(cache_file):
@@ -395,7 +426,6 @@ 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(
@@ -426,17 +456,29 @@ async def fetch_file_list_with_retry(
     recursive: bool = False,
     on_connection_lost: Callable[[], None] = lambda: None,
 ) -> list[FileListEntry]:
-    n_attempts = 3
+    n_attempts = 5
     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)
+            await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
     raise Exception(
         f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
     )
@@ -447,6 +489,9 @@ 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 (
@@ -458,7 +503,8 @@ async def _fetch_file_list(
             raise HuggingFaceAuthenticationError(msg)
         elif response.status == 429:
             raise HuggingFaceRateLimitError(
-                f"Couldn't download {model_id} because of HuggingFace rate limit."
+                f"HuggingFace rate limit hit fetching file list for {model_id}",
+                retry_after=_parse_retry_after(response.headers),
             )
         elif response.status == 200:
             data_json = await response.text()
@@ -468,10 +514,14 @@ async def _fetch_file_list(
                 if item.type == "file":
                     files.append(FileListEntry.model_validate(item))
                 elif item.type == "directory" and recursive:
-                    subfiles = await _fetch_file_list(
-                        model_id, revision, item.path, recursive
-                    )
-                    files.extend(subfiles)
+                    # 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"
+                )
             return files
         else:
             raise Exception(f"Failed to fetch file list: {response.status}")
@@ -552,6 +602,11 @@ 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
         )
@@ -571,7 +626,7 @@ async def download_file_with_retry(
     on_connection_lost: Callable[[], None] = lambda: None,
     skip_internet: bool = False,
 ) -> Path:
-    n_attempts = 3
+    n_attempts = 5
     for attempt in range(n_attempts):
         try:
             return await _download_file(
@@ -583,12 +638,16 @@ async def download_file_with_retry(
             raise
         except HuggingFaceRateLimitError as e:
             if attempt == n_attempts - 1:
-                raise e
-            logger.error(
-                f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
+                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.error(traceback.format_exc())
-            await asyncio.sleep(2.0**attempt)
+            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)
         except Exception as e:
             if attempt == n_attempts - 1:
                 on_connection_lost()
@@ -597,7 +656,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)
+            await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
     raise Exception(
         f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
     )
@@ -665,6 +724,11 @@ 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 get_model_cards()
+            for model_card in await model_cards.card_cache.list_all()
         ]
 
         for task in asyncio.as_completed(tasks):

src/exo/download/tests/test_offline_mode.py

@@ -1,5 +1,7 @@
 """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
@@ -231,3 +233,64 @@ 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

@@ -0,0 +1,355 @@
+"""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/master/main.py

@@ -180,31 +180,10 @@ 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():
-                                model_match = (
-                                    instance.shard_assignments.model_id
-                                    == requested_model
-                                ) or (instance.instance_id in linked_decode_ids)
                                 if (
-                                    model_match
+                                    instance.shard_assignments.model_id
+                                    == command.task_params.model
                                     and instance.instance_id not in prefill_only
                                 ):
                                     in_flight = {TaskStatus.Pending, TaskStatus.Running}
@@ -386,6 +365,7 @@ 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
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo, NodeRdmaCtlStatus
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
@@ -43,7 +43,6 @@ 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
@@ -106,6 +105,7 @@ 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,8 +167,18 @@ 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)
+        cycle
+        for cycle in smallest_cycles
+        if topology.is_rdma_cycle(cycle) and _all_rdma_ctl_enabled(cycle)
     ]
 
     if command.instance_meta == InstanceMeta.MlxJaccl:
@@ -203,7 +213,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
+    if len(selected_cycle) == 1:
         command = command.model_copy(
             update={
                 "instance_meta": InstanceMeta.MlxRing,
@@ -267,11 +277,6 @@ 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,13 +375,7 @@ def find_ip_prioritised(
             "maybe_ethernet": 3,
             "thunderbolt": 4,
         }
-
-    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)
+    return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
 
 
 def get_mlx_ring_hosts_by_node(

src/exo/master/tests/test_placement.py

@@ -21,7 +21,11 @@ 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
+from exo.shared.types.profiling import (
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
+    NodeRdmaCtlStatus,
+)
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import (
     InputMessage,
@@ -439,8 +443,21 @@ 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)
+    placements = place_instance(
+        cic,
+        topology,
+        {},
+        node_memory,
+        node_network,
+        node_rdma_ctl=node_rdma_ctl,
+    )
 
     # assert
     assert len(placements) == 1
@@ -482,6 +499,131 @@ 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,7 +4,8 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.types.common import NodeId
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
     CustomModelCardAdded,
@@ -59,14 +60,24 @@ 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:
@@ -77,10 +88,12 @@ 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():
@@ -306,9 +319,6 @@ 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 = (
@@ -331,7 +341,6 @@ 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,
         }
     )
@@ -358,11 +367,6 @@ 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():
@@ -408,6 +412,9 @@ 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,
@@ -420,6 +427,10 @@ 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():
@@ -443,11 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
-        case VllmCapability():
-            update["node_vllm"] = {
-                **state.node_vllm,
-                event.node_id: info.available,
-            }
+            # 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)
 
     return state.model_copy(update=update)
 
@@ -463,3 +475,22 @@ 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

@@ -4,13 +4,13 @@ from pathlib import Path
 
 from exo.utils.dashboard_path import find_dashboard, find_resources
 
-_EXO_HOME_ENV = os.environ.get("EXO_HOME", None)
+_EXO_HOME_ENV = os.environ.get("EXO_HOME", "")
 
 
 def _get_xdg_dir(env_var: str, fallback: str) -> Path:
     """Get XDG directory, prioritising EXO_HOME environment variable if its set. On non-Linux platforms, default to ~/.exo."""
 
-    if _EXO_HOME_ENV is not None:
+    if _EXO_HOME_ENV != "":
         return Path.home() / _EXO_HOME_ENV
 
     if sys.platform != "linux":

src/exo/shared/models/model_cards.py

@@ -39,7 +39,57 @@ _BUILTIN_CARD_DIRS = [
     Path(RESOURCES_DIR) / "image_model_cards",
 ]
 
-_card_cache: dict[ModelId, "ModelCard"] = {}
+
+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()
 
 
 def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -59,42 +109,10 @@ 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"
@@ -150,7 +168,6 @@ 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)
@@ -197,14 +214,13 @@ 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 model_id not in _card_cache:
-            await _refresh_card_cache()
-        if (mc := _card_cache.get(model_id)) is not None:
+        if card_cache.get(model_id) is None:
+            await card_cache.refresh()
+        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
@@ -234,21 +250,6 @@ 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
 
@@ -350,11 +351,7 @@ 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.
-
-    Single-shard repos don't have a `model.safetensors.index.json`; fall back
-    to the HF API for those.
-    """
+    """Gets model size from safetensors index or falls back to HF API."""
     from exo.download.download_utils import (
         download_file_with_retry,
         resolve_model_dir,
@@ -362,25 +359,21 @@ 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)
-    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
+    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())
 
-    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)
+    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

@@ -0,0 +1,44 @@
+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

@@ -0,0 +1,231 @@
+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,6 +169,22 @@ 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,8 +5,9 @@ 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 NodeId
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -58,7 +59,6 @@ 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,6 +66,9 @@ 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.model_cards import get_card
+        from exo.shared.models import model_cards
 
-        card = get_card(self.model)
+        card = model_cards.card_cache.get(self.model)
         if card is None:
             return self
 

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

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

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

@@ -31,7 +31,6 @@ 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,
@@ -354,24 +353,6 @@ 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
@@ -380,8 +361,6 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
-    | NvmlMetrics
-    | VllmCapability
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -440,8 +419,6 @@ 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)
@@ -450,10 +427,6 @@ 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()
@@ -502,16 +475,6 @@ 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

@@ -1,70 +0,0 @@
-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,7 +1,6 @@
 import platform
 import socket
 import sys
-from pathlib import Path
 from subprocess import CalledProcessError
 
 import psutil
@@ -118,90 +117,12 @@ 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 system model and chip information.
-
-    On macOS, uses ``system_profiler``.  On Linux, reads DMI data from
-    sysfs and CPU info from ``/proc/cpuinfo``.
-    """
+    """Get Mac system information using system_profiler."""
     model = "Unknown Model"
     chip = "Unknown Chip"
 
-    if sys.platform == "linux":
-        return await _get_linux_model_and_chip()
-
+    # TODO: better non mac support
     if sys.platform != "darwin":
         return (model, chip)
 

src/exo/worker/disaggregated/protocol.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass
 from typing import BinaryIO, Literal
 
 import msgspec
@@ -24,28 +23,7 @@ class TensorBlob(msgspec.Struct):
     data: bytes
 
 
-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.
-    """
-
+class KVChunk(msgspec.Struct, tag="kv_chunk"):
     layer_idx: int
     num_tokens: int
     n_heads: int
@@ -73,13 +51,10 @@ 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[_WireMessage] = msgspec.msgpack.Decoder(
-    _WireMessage
-)
+_msg_decoder: msgspec.msgpack.Decoder[Message] = msgspec.msgpack.Decoder(Message)
 _header_encoder = msgspec.msgpack.Encoder()
 _header_decoder: msgspec.msgpack.Decoder[Header] = msgspec.msgpack.Decoder(Header)
 
@@ -124,7 +99,7 @@ def read_header(stream: BinaryIO) -> Header:
         raise ProtocolError(f"Bad header: {exc}") from exc
 
 
-def write_message(stream: BinaryIO, msg: _WireMessage) -> None:
+def write_message(stream: BinaryIO, msg: Message) -> None:
     write_frame(stream, _msg_encoder.encode(msg))
 
 
@@ -133,22 +108,9 @@ def read_message(stream: BinaryIO) -> Message | None:
     if not payload:
         return None
     try:
-        msg = _msg_decoder.decode(payload)
+        return _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(
@@ -159,35 +121,21 @@ def write_kv_chunk(
     n_heads: int,
     head_dim: int,
     dtype: DType,
-    keys: "bytes | memoryview",
-    values: "bytes | memoryview",
+    keys: bytes,
+    values: bytes,
 ) -> None:
-    """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(
+    write_message(
+        stream,
+        KVChunk(
             layer_idx=layer_idx,
             num_tokens=num_tokens,
             n_heads=n_heads,
             head_dim=head_dim,
             dtype=dtype,
-            keys_len=keys_len,
-            values_len=values_len,
-        )
+            keys=keys,
+            values=values,
+        ),
     )
-    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,7 +21,6 @@ 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()
@@ -57,10 +56,7 @@ class _PrefillHandler(socketserver.StreamRequestHandler):
         super().setup()
         sock = cast(socket.socket, self.request)
         sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-        # 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)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024)
 
     def handle(self) -> None:
         server = cast(PrefillServer, self.server)

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

@@ -143,6 +143,7 @@ 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)
@@ -168,7 +169,11 @@ 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 else ()
+        return (
+            (resp,)
+            if resp is not None and _is_primary_output_node(self.shard_metadata)
+            else ()
+        )
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):

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]
-    model: Model | None = None
+    inference_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.model,
+            self.inference_model,
             self.tokenizer,
             self.vision_processor,
         ) = yield from load_mlx_items(bound_instance, self.group)
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):
-            del self.model
+            del self.inference_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.model
+        assert self.inference_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.model,
+                model=self.inference_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.model,
+                model=self.inference_model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,

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

@@ -22,6 +22,7 @@ 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,
@@ -89,18 +90,6 @@ 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,
@@ -114,7 +103,7 @@ def send_mlx_kv_cache(
         match c:
             case QuantizedKVCache() | CacheList() | DeepseekV4Cache():
                 raise NotImplementedError
-            case KVCache():
+            case KVCache() | RotatingKVCache():
                 keys = c.keys
                 values = c.values
                 if keys is None or values is None:
@@ -143,39 +132,11 @@ def send_mlx_kv_cache(
                     keys=array_to_bytes(k_nhd),
                     values=array_to_bytes(v_nhd),
                 )
-                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)
+                if tokens_sent != 0 and num_tokens != tokens_sent:
+                    logger.critical(
+                        f"Unexpected number of tokens sent {num_tokens} != {tokens_sent}"
+                    )
+                tokens_sent = num_tokens
             case ArraysCache():
                 blobs: list[TensorBlob] = []
                 for a in c.state:

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

@@ -79,7 +79,6 @@ 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)
@@ -94,17 +93,10 @@ 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

@@ -215,18 +215,14 @@ 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(
-                        all_prompt_tokens[:-1],
+                        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:

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

@@ -648,20 +648,14 @@ 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(
-                    all_prompt_tokens[:-1],
+                    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:

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

@@ -25,25 +25,23 @@ 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, tokens_received_total
-        tokens_received_total += chunk.num_tokens
+    def _on_chunk(_chunk: KVChunk, chunks_received: int) -> None:
+        nonlocal num_layers
         if on_prefill_progress is None:
             return
         if num_layers > 0 and chunks_received % num_layers == 0:
+            tokens_so_far = chunks_received // num_layers
             on_prefill_progress(
-                min(tokens_received_total // num_layers, total_prompt_tokens),
+                min(tokens_so_far, total_prompt_tokens),
                 total_prompt_tokens,
             )
 
@@ -52,7 +50,6 @@ 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
@@ -64,28 +61,6 @@ 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,7 +50,6 @@ 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 (
@@ -141,8 +140,6 @@ 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")
 

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

@@ -1,88 +0,0 @@
-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

@@ -1,238 +0,0 @@
-"""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

@@ -1,10 +0,0 @@
-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

@@ -1,201 +0,0 @@
-"""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

@@ -1,623 +0,0 @@
-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

@@ -1,490 +0,0 @@
-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

@@ -1,480 +0,0 @@
-# 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

@@ -1,787 +0,0 @@
-# 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

@@ -1,63 +0,0 @@
-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

src/exo/worker/main.py

@@ -10,7 +10,7 @@ from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
-from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
+from exo.shared.models.model_cards import ModelId, card_cache
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
@@ -20,8 +20,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
@@ -110,6 +108,8 @@ class Worker:
                 tg.start_soon(self.plan_step)
                 tg.start_soon(self._event_applier)
                 tg.start_soon(self._poll_connection_updates)
+                tg.start_soon(self._reconcile_custom_cards)
+
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
@@ -151,7 +151,6 @@ class Worker:
                     self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
                         event.chunk
                     )
-
                     if (
                         len(self.input_chunk_buffer[cmd_id])
                         == self.input_chunk_counts[cmd_id]
@@ -172,12 +171,18 @@ class Worker:
                                 )
                             ] = img
 
-                if isinstance(event, CustomModelCardAdded):
-                    await event.model_card.save_to_custom_dir()
-                    add_to_card_cache(event.model_card)
+    async def _reconcile_custom_cards(self) -> None:
+        while True:
+            await anyio.sleep(1)
+            target = dict(self.state.custom_model_cards)
+            for model_id, card in target.items():
+                if card_cache.get(model_id) == card:
+                    continue
+                await card_cache.save(card)
 
-                if isinstance(event, CustomModelCardDeleted):
-                    await delete_custom_card(event.model_id)
+            for card in await card_cache.list_all():
+                if card.model_id not in target:
+                    await card_cache.pop(card.model_id)
 
     async def plan_step(self):
         while True:

src/exo/worker/runner/bootstrap.py

@@ -5,7 +5,7 @@ import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance, VllmInstance
+from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 from exo.worker.engines.base import Builder
@@ -46,21 +46,6 @@ def entrypoint(
             builder = MfluxBuilder(
                 event_sender, cancel_receiver, bound_instance.bound_shard
             )
-        elif isinstance(bound_instance.instance, VllmInstance):
-            from exo.worker.engines.vllm.builder import VllmBuilder
-            from exo.worker.engines.vllm.growable_cache import patch_vllm
-
-            os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
-            os.environ["VLLM_BATCH_INVARIANT"] = "1"
-            os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
-
-            patch_vllm()
-            builder = VllmBuilder(
-                bound_instance.bound_shard.model_card.model_id,
-                event_sender,
-                cancel_receiver,
-            )
         else:
             from exo.worker.engines.mlx.patches import apply_mlx_patches
 

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -138,8 +138,10 @@ class SequentialGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -197,9 +199,14 @@ class SequentialGenerator(Engine):
             self._active = None
             raise
 
-        return itertools.chain(
-            output,
-            map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
+        return filter(
+            lambda chunk: (
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
+            ),
+            itertools.chain(
+                output,
+                map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
+            ),
         )
 
     def _start_next(self) -> None:
@@ -221,6 +228,7 @@ class SequentialGenerator(Engine):
                 apply_chat_template(self.tokenizer, task.task_params),
                 self.tool_parser,
                 self.tokenizer,
+                type(self.model),
                 self.model_id,
                 task.task_params.tools,
             )
@@ -367,8 +375,10 @@ class BatchGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -417,6 +427,7 @@ class BatchGenerator(Engine):
                     apply_chat_template(self.tokenizer, task.task_params),
                     self.tool_parser,
                     self.tokenizer,
+                    type(self.model),
                     self.model_id,
                     task.task_params.tools,
                 )
@@ -447,7 +458,12 @@ class BatchGenerator(Engine):
                 output.append((task.task_id, FinishedResponse()))
                 del self._active_tasks[uid]
 
-        return itertools.chain(output, self._apply_cancellations())
+        return filter(
+            lambda chunk: (
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
+            ),
+            itertools.chain(output, self._apply_cancellations()),
+        )
 
     def _apply_cancellations(
         self,

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -2,6 +2,9 @@ from collections.abc import Callable, Generator, Iterator
 from functools import cache
 from typing import Any
 
+from mlx_lm.models.deepseek_v4 import Model as DeepseekV4Model
+from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     HarmonyEncodingName,
@@ -20,6 +23,7 @@ from exo.shared.types.chunks import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     detect_thinking_prompt_suffix,
 )
@@ -65,15 +69,16 @@ def apply_all_parsers(
     prompt: str,
     tool_parser: ToolParser | None,
     tokenizer: TokenizerWrapper,
+    model_type: type[Model],
     model_id: ModelId,
     tools: list[dict[str, Any]] | None,
 ) -> Iterator[GenerationChunk | None]:
     generator = receiver
 
-    normalized_id = model_id.short().lower()
-    if "gpt-oss" in normalized_id:
+    normalized_id = model_id.normalize().lower()
+    if issubclass(model_type, GptOssModel):
         generator = parse_gpt_oss(generator)
-    elif "deepseek-v3.2" in normalized_id:
+    elif issubclass(model_type, DeepseekV32Model) and "deepseek" in normalized_id:
         if tokenizer.has_thinking:
             generator = parse_thinking_models(
                 generator,
@@ -82,7 +87,7 @@ def apply_all_parsers(
                 starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
             )
         generator = parse_deepseek_v32(generator)
-    elif "deepseek-v4" in normalized_id:
+    elif issubclass(model_type, DeepseekV4Model) and "deepseek-v4" in normalized_id:
         if tokenizer.has_thinking:
             generator = parse_thinking_models(
                 generator,

src/exo/worker/runner/runner.py

@@ -178,10 +178,6 @@ class Runner:
 
     def _serve_prefill(self, req: PrefillTask) -> None:
         req.started.set()
-        nested = isinstance(self.current_status, RunnerRunning)
-        if not nested:
-            self.update_status(RunnerRunning())
-            logger.info("runner running")
         try:
             assert isinstance(self.generator, Engine)
             self.generator.serve_prefill(req.request, req.wfile)
@@ -191,11 +187,6 @@ class Runner:
             )
         finally:
             req.done.set()
-            if not nested:
-                self.update_status(
-                    RunnerReady(prefill_server_port=self._prefill_server_port)
-                )
-                logger.info("runner ready")
 
     def update_status(self, status: RunnerStatus):
         self.current_status = status
@@ -399,5 +390,5 @@ class Runner:
         chunk: Chunk,
         command_id: CommandId,
     ):
-        if self.device_rank == 0:
-            self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
+        assert isinstance(self.generator, Engine)
+        self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))

src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py

@@ -16,7 +16,7 @@ from exo.download.download_utils import (
     fetch_file_list_with_cache,
     resolve_model_dir,
 )
-from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
+from exo.shared.models.model_cards import ModelCard, ModelId, card_cache
 from exo.worker.engines.mlx.utils_mlx import (
     get_eos_token_ids_for_model,
     load_tokenizer_for_model_id,
@@ -76,7 +76,7 @@ def get_test_models() -> list[ModelCard]:
     """Get a representative sample of models to test."""
     # Pick one model from each family to test
     families: dict[str, ModelCard] = {}
-    for card in asyncio.run(get_model_cards()):
+    for card in asyncio.run(card_cache.list_all()):
         # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
         parts = card.model_id.short().split("-")
         family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
@@ -298,7 +298,7 @@ async def test_tokenizer_special_tokens(model_card: ModelCard) -> None:
 async def test_kimi_tokenizer_specifically():
     """Test Kimi tokenizer with its specific patches and quirks."""
     kimi_models = [
-        card for card in await get_model_cards() if "kimi" in card.model_id.lower()
+        card for card in await card_cache.list_all() if "kimi" in card.model_id.lower()
     ]
 
     if not kimi_models:
@@ -350,7 +350,7 @@ async def test_glm_tokenizer_specifically():
 
     glm_model_cards = [
         card
-        for card in await get_model_cards()
+        for card in await card_cache.list_all()
         if contains(card, "glm")
         and not contains(card, "-5")
         and not contains(card, "4.7")