Warmup = 2

## Motivation

Process-isolated runner crashes and C-extension failures can write
directly to fd-level stdout/stderr, bypassing Python/loguru. We need to
capture that output per runner process without polluting the main
process or other workers, and without breaking operation when the parent
stdio is detached.

## Changes

- Added `AsyncProcess`, a spawn-only multiprocessing wrapper that
redirects child stdout/stderr to pipes and exposes them as in-memory
`Receiver[bytes]`s
- Replaced runner-supervisor's raw `multiprocessing.Process` usage with
`AsyncProcess`
- Added `--no-stdio`, redirecting stdin/stdout/stderr to `/dev/null`
after logging is configured
- Disabled verbose MLX
- Added tests covering stdio capture, child crashes, repeated bad
children, SIGTERM/SIGKILL shutdown escalation, stdio detachment, and
spawning captured children from a stdio-detached parent

## Why It Works

The parent can redirect its own stdio fds to `/dev/null`, while
`AsyncProcess` installs fresh pipe fds over fd 1 and 2 inside each
spawned child. That keeps stdio-detached parents quiet while preserving
per-runner stdout/stderr capture. Runner shutdown is still bounded:
SIGTERM grace first, then SIGKILL escalation if needed.

Next direction: the runner supervisor currently drains captured output
and logs it as stdout/debug and stderr/warning. This should be split
into more useful process-isolated error reporting instead of just log
forwarding (regex match on errors to obtain "reason" string, best
effort).

## Test Plan

### Manual Testing

Ran on 4 Mac Minis in a Thunderbolt 4 ring, can see that runner's
stdout/stderr contents are being captured.

### Automated Testing

- Added async-process tests for fd-level stdout/stderr capture, Python
traceback capture, bounded-buffer output, child `exit`/abort, parent
stdio preservation, fd leak checks, spawn-context mp channels, and
SIGTERM/SIGKILL shutdown behavior
- Added stdio-detach tests proving stdio detaches to `/dev/null`, a
stdio-detached parent can still spawn and capture a child, and the same
stdio-detached parent can spawn/capture multiple children sequentially
- Updated runner-supervisor tests for the new `AsyncProcess.exitcode`
path

.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/build-app.yml

@@ -32,7 +32,6 @@ jobs:
       SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
       SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
       SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
-      EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
       AWS_REGION: ${{ secrets.AWS_REGION }}
       EXO_BUILD_NUMBER: ${{ github.run_number }}
       EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
@@ -159,7 +158,7 @@ jobs:
           fi
 
       - name: Install Homebrew packages
-        run: brew install just awscli macmon
+        run: brew install just awscli
 
       - name: Install UV
         uses: astral-sh/setup-uv@v6
@@ -239,10 +238,92 @@ jobs:
           # Export keychain path for other steps
           echo "BUILD_KEYCHAIN_PATH=$KEYCHAIN_PATH" >> $GITHUB_ENV
 
+      # ============================================================
+      # Pre-flight credential / profile validation
+      # Runs BEFORE the ~16 min build so auth/expiry failures surface in <1 min.
+      # ============================================================
+
+      - name: Validate Apple notarization credentials
+        env:
+          APPLE_NOTARIZATION_USERNAME: ${{ secrets.APPLE_NOTARIZATION_USERNAME }}
+          APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
+          APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
+        run: |
+          # All-or-nothing: either all three creds are set, or none are.
+          CRED_COUNT=0
+          for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
+            [[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
+          done
+          if [[ "$CRED_COUNT" -eq 0 ]]; then
+            echo "No notarization credentials configured — skipping notarization for this build."
+            exit 0
+          fi
+          if [[ "$CRED_COUNT" -ne 3 ]]; then
+            echo "ERROR: partial notarization credentials set ($CRED_COUNT/3). Aborting before build."
+            exit 1
+          fi
+          # Cheap, ~5s, auth-only call. Fails instantly with a clear message if
+          # the app-specific password is stale, wrong team-id, etc.
+          echo "Verifying Apple notarization credentials via notarytool history..."
+          if ! xcrun notarytool history \
+            --apple-id "$APPLE_NOTARIZATION_USERNAME" \
+            --password "$APPLE_NOTARIZATION_PASSWORD" \
+            --team-id "$APPLE_NOTARIZATION_TEAM" >/dev/null; then
+            echo "ERROR: notarytool rejected the provided credentials. Fix before rerunning."
+            echo "Common causes: app-specific password expired/revoked, wrong team-id,"
+            echo "Apple ID not on the team, or 2FA not configured for this Apple ID."
+            exit 1
+          fi
+          echo "Apple notarization credentials OK."
+
+      - name: Validate provisioning profile expiry
+        run: |
+          PROFILE="$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles/EXO.provisionprofile"
+          if [[ ! -f "$PROFILE" ]]; then
+            echo "ERROR: provisioning profile not found at $PROFILE"
+            exit 1
+          fi
+          EXPIRY=$(security cms -D -i "$PROFILE" | plutil -extract ExpirationDate raw -o - - 2>/dev/null || true)
+          if [[ -z "$EXPIRY" ]]; then
+            echo "WARNING: could not read ExpirationDate from provisioning profile; skipping expiry check."
+            exit 0
+          fi
+          # Try a couple of known plutil date formats. If none parse, skip the check rather
+          # than risk a false-positive "expired" block on a format we didn't anticipate.
+          EXPIRY_EPOCH=""
+          for fmt in "%Y-%m-%dT%H:%M:%SZ" "%Y-%m-%d %H:%M:%S %z" "%Y-%m-%d %H:%M:%S +0000"; do
+            if parsed=$(date -j -f "$fmt" "$EXPIRY" +%s 2>/dev/null); then
+              EXPIRY_EPOCH="$parsed"
+              break
+            fi
+          done
+          if [[ -z "$EXPIRY_EPOCH" ]]; then
+            echo "WARNING: could not parse ExpirationDate '$EXPIRY'; skipping expiry check."
+            exit 0
+          fi
+          NOW_EPOCH=$(date +%s)
+          if [[ "$EXPIRY_EPOCH" -le "$NOW_EPOCH" ]]; then
+            echo "ERROR: provisioning profile expired on $EXPIRY. Regenerate it before rerunning."
+            exit 1
+          fi
+          DAYS_LEFT=$(( (EXPIRY_EPOCH - NOW_EPOCH) / 86400 ))
+          echo "Provisioning profile valid until $EXPIRY ($DAYS_LEFT days remaining)."
+          if [[ "$DAYS_LEFT" -lt 14 ]]; then
+            echo "WARNING: profile expires in under 14 days — regenerate soon."
+          fi
+
       # ============================================================
       # Build the bundle
       # ============================================================
 
+      - name: Add pinned macmon to PATH
+        run: |
+          MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
+          echo "Using macmon from: $MACMON_DIR"
+          echo "$MACMON_DIR" >> $GITHUB_PATH
+          # Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
+          brew uninstall macmon 2>/dev/null || true
+
       - name: Build PyInstaller bundle
         run: uv run pyinstaller packaging/pyinstaller/exo.spec
 
@@ -265,7 +346,6 @@ jobs:
             EXO_BUILD_COMMIT="$GITHUB_SHA" \
             SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
             SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
-            EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
             CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
             CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
           mkdir -p ../../output
@@ -298,11 +378,41 @@ jobs:
           APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
           APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
         run: |
+          set -o pipefail
           cd output
           security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH"
           SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}')
+
+          # Fail fast if notarization creds are partial. All-or-nothing.
+          CRED_COUNT=0
+          for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
+            [[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
+          done
+          if [[ "$CRED_COUNT" -ne 0 && "$CRED_COUNT" -ne 3 ]]; then
+            echo "ERROR: partial Apple notarization credentials set ($CRED_COUNT/3). Aborting."
+            exit 1
+          fi
+
           /usr/bin/codesign --deep --force --timestamp --options runtime \
             --sign "$SIGNING_IDENTITY" EXO.app
+
+          # Pre-flight: verify the signed app BEFORE building DMG and submitting to Apple.
+          # If this fails, notarization will fail too — cheap way to fail in seconds, not 15 minutes.
+          echo "===== codesign --verify EXO.app ====="
+          if ! /usr/bin/codesign --verify --deep --strict --verbose=2 EXO.app; then
+            echo "ERROR: EXO.app failed codesign verification. Dumping signing status of every executable:"
+            find EXO.app -type f \( -perm -111 -o -name "*.dylib" -o -name "*.so" -o -name "*.framework" \) -print0 |
+              while IFS= read -r -d '' f; do
+                printf -- '--- %s\n' "$f"
+                /usr/bin/codesign -dv --verbose=2 "$f" 2>&1 | sed 's/^/    /' || true
+              done
+            exit 1
+          fi
+
+          # Gatekeeper assessment. A failure here strongly predicts notarization rejection.
+          echo "===== spctl assessment (predicts notarization outcome) ====="
+          /usr/bin/spctl -a -vvv -t install EXO.app || echo "WARNING: spctl assessment failed — notarization is likely to fail too."
+
           mkdir -p dmg-root
           cp -R EXO.app dmg-root/
           ln -s /Applications dmg-root/Applications
@@ -310,12 +420,22 @@ jobs:
           hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME"
           /usr/bin/codesign --force --timestamp --options runtime \
             --sign "$SIGNING_IDENTITY" "$DMG_NAME"
+
+          echo "===== codesign --verify DMG ====="
+          if ! /usr/bin/codesign --verify --verbose=2 "$DMG_NAME"; then
+            echo "ERROR: DMG failed codesign verification."
+            exit 1
+          fi
+
           if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then
+            echo "===== notarytool submit ====="
+            # `|| true` so set -e doesn't abort before we can echo output / fetch the log.
+            # We rely on the parsed STATUS below to decide pass/fail.
             SUBMISSION_OUTPUT=$(xcrun notarytool submit "$DMG_NAME" \
               --apple-id "$APPLE_NOTARIZATION_USERNAME" \
               --password "$APPLE_NOTARIZATION_PASSWORD" \
               --team-id "$APPLE_NOTARIZATION_TEAM" \
-              --wait --timeout 15m 2>&1)
+              --wait --timeout 15m 2>&1) || true
             echo "$SUBMISSION_OUTPUT"
 
             SUBMISSION_ID=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="id:" && $2 ~ /^[0-9a-fA-F-]+$/ {print $2; exit}')

.github/workflows/pipeline.yml

@@ -91,9 +91,6 @@ jobs:
             nix build .#metal-toolchain
           fi
 
-          # Build mlx (depends on metal-toolchain)
-          nix build .#mlx
-
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '

.gitignore

@@ -38,3 +38,6 @@ bench/**/*.json
 
 # tmp
 tmp/models
+/build/exo
+/.claude/skills
+/.claude

.idea/exo-v2.iml

@@ -1,31 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<module type="EMPTY_MODULE" version="4">
-  <component name="FacetManager">
-    <facet type="Python" name="Python facet">
-      <configuration sdkName="Python 3.13 virtualenv at ~/Desktop/exo/.venv" />
-    </facet>
-  </component>
-  <component name="Go" enabled="true" />
-  <component name="NewModuleRootManager">
-    <content url="file://$MODULE_DIR$">
-      <sourceFolder url="file://$MODULE_DIR$/scripts/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/tests" isTestSource="true" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/util/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/examples" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/tests" isTestSource="true" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/system_custodian/src" isTestSource="false" />
-      <excludeFolder url="file://$MODULE_DIR$/.venv" />
-      <excludeFolder url="file://$MODULE_DIR$/.direnv" />
-      <excludeFolder url="file://$MODULE_DIR$/build" />
-      <excludeFolder url="file://$MODULE_DIR$/dist" />
-      <excludeFolder url="file://$MODULE_DIR$/.go_cache" />
-      <excludeFolder url="file://$MODULE_DIR$/rust/target" />
-    </content>
-    <orderEntry type="jdk" jdkName="Python 3.13 (exo)" jdkType="Python SDK" />
-    <orderEntry type="sourceFolder" forTests="false" />
-    <orderEntry type="library" name="Python 3.13 virtualenv at ~/Desktop/exo/.venv interpreter library" level="application" />
-  </component>
-</module>

.idea/externalDependencies.xml

@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalDependencies">
-    <plugin id="systems.fehn.intellijdirenv" />
+    <plugin id="al.aoli.intellijdirenv" />
   </component>
 </project>

.idea/inspectionProfiles/Project_Default.xml

@@ -1,14 +0,0 @@
-<component name="InspectionProjectProfileManager">
-  <profile version="1.0">
-    <option name="myName" value="Project Default" />
-    <inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
-      <option name="ourVersions">
-        <value>
-          <list size="1">
-            <item index="0" class="java.lang.String" itemvalue="3.14" />
-          </list>
-        </value>
-      </option>
-    </inspection_tool>
-  </profile>
-</component>

.idea/misc.xml

@@ -4,7 +4,4 @@
     <option name="sdkName" value="Python 3.13 (exo)" />
   </component>
   <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (exo)" project-jdk-type="Python SDK" />
-  <component name="PythonCompatibilityInspectionAdvertiser">
-    <option name="version" value="3" />
-  </component>
 </project>

.mlx_typings/mlx/core/__init__.pyi

@@ -1767,12 +1767,12 @@ def clip(
         array: The clipped array.
     """
 
-def compile(
-    fun: Callable,
+def compile[F: Callable[..., object]](
+    fun: F,
     inputs: object | None = ...,
     outputs: object | None = ...,
     shapeless: bool = ...,
-) -> Callable:
+) -> F:
     """
     Returns a compiled function which produces the same output as ``fun``.
 
@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.
@@ -2915,8 +2915,8 @@ def gather_mm(
     a: array,
     b: array,
     /,
-    lhs_indices: array,
-    rhs_indices: array,
+    lhs_indices: array | None = ...,
+    rhs_indices: array | None = ...,
     *,
     sorted_indices: bool = ...,
     stream: Stream | Device | None = ...,
@@ -4707,6 +4707,7 @@ def softmax(
     /,
     axis: int | Sequence[int] | None = ...,
     *,
+    precise: bool = ...,
     stream: Stream | Device | None = ...,
 ) -> array:
     """

.mlx_typings/mlx/nn/__init__.pyi

@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *
 
 from . import init as init
 from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,10 +53,14 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    __call__: Callable
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 
+    def __getitem__(self, key: str) -> mx.array | Module: ...
+    def get(
+        self, key: str, default: mx.array | Module | None = ...
+    ) -> mx.array | Module | None: ...
     @property
     def training(self):  # -> bool:
         """Boolean indicating if the model is in training mode."""

.mlx_typings/mlx/nn/layers/quantized.pyi

@@ -2,7 +2,7 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import mlx.core as mx
 from base import Module
@@ -13,8 +13,10 @@ def quantize(
     bits: int = ...,
     *,
     mode: str = ...,
-    class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
-):  # -> None:
+    class_predicate: Optional[
+        Callable[[str, Module], Union[bool, dict[str, Any]]]
+    ] = ...,
+) -> None:
     """Quantize the sub-modules of a module according to a predicate.
 
     By default all layers that define a ``to_quantized(group_size, bits)``

.mlx_typings/mlx_lm/_version.pyi

@@ -0,0 +1,5 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+__version__ = ...

.mlx_typings/mlx_lm/generate.pyi

@@ -3,13 +3,12 @@ This type stub file was generated by pyright.
 """
 
 import contextlib
-from dataclasses import dataclass
-from typing import Any, Callable, Generator, List, Optional, Tuple, Union
-
 import mlx.core as mx
 import mlx.nn as nn
+from dataclasses import dataclass
+from collections import deque
+from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
 from transformers import PreTrainedTokenizer
-
 from .tokenizer_utils import TokenizerWrapper
 
 DEFAULT_PROMPT = ...
@@ -29,8 +28,9 @@ def str2bool(string):  # -> bool:
     ...
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
+    ...
 
-generation_stream = ...
+generation_stream: mx.Stream
 
 @contextlib.contextmanager
 def wired_limit(
@@ -43,6 +43,7 @@ def wired_limit(
     async eval could be running pass in the streams to synchronize with prior
     to exiting the context manager.
     """
+    ...
 @dataclass
 class GenerationResponse:
     """
@@ -91,7 +92,7 @@ def generate_step(
     kv_bits: Optional[int] = ...,
     kv_group_size: int = ...,
     quantized_kv_start: int = ...,
-    prompt_progress_callback: Optional[Callable[[int], int]] = ...,
+    prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
     input_embeddings: Optional[mx.array] = ...,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -117,7 +118,7 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[[int], int]): A call-back which takes the
+        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
@@ -125,6 +126,7 @@ def generate_step(
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
+    ...
 
 def speculative_generate_step(
     prompt: mx.array,
@@ -170,6 +172,7 @@ def speculative_generate_step(
         Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
           and a bool indicating if the token was generated by the draft model
     """
+    ...
 
 def stream_generate(
     model: nn.Module,
@@ -177,7 +180,7 @@ def stream_generate(
     prompt: Union[str, mx.array, List[int]],
     max_tokens: int = ...,
     draft_model: Optional[nn.Module] = ...,
-    **kwargs: object,
+    **kwargs: Any,
 ) -> Generator[GenerationResponse, None, None]:
     """
     A generator producing text based on the given prompt from the model.
@@ -199,6 +202,7 @@ def stream_generate(
         GenerationResponse: An instance containing the generated text segment and
             associated metadata. See :class:`GenerationResponse` for details.
     """
+    ...
 
 def generate(
     model: nn.Module,
@@ -219,6 +223,9 @@ def generate(
        kwargs: The remaining options get passed to :func:`stream_generate`.
           See :func:`stream_generate` for more details.
     """
+    ...
+
+def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
 @dataclass
 class BatchStats:
     """
@@ -242,10 +249,263 @@ class BatchStats:
     generation_time: float = ...
     peak_memory: float = ...
 
+class SequenceStateMachine:
+    """A state machine that uses one Aho-Corasick trie per state to efficiently
+    track state across a generated sequence.
+
+    The transitions are provided as state -> [(sequence, new_state)].
+
+    Example:
+
+        sm = SequenceStateMachine(
+            transitions={
+                "normal": [
+                    (think_start_tokens, "reasoning"),
+                    (tool_start_tokens, "tool"),
+                    (eos, None),
+                ],
+                "reasoning": [
+                    (think_end_tokens, "normal"),
+                    (eos, None),
+                ],
+                "tool": [
+                    (tool_end_tokens, None),
+                    (eos, None)
+                ],
+            },
+            initial="normal"
+        )
+    """
+    def __init__(self, transitions=..., initial=...) -> None: ...
+    def __deepcopy__(self, memo):  # -> SequenceStateMachine:
+        ...
+    def make_state(self):  # -> tuple[str, Any, dict[Any, Any]]:
+        ...
+    @staticmethod
+    def match(state, x):  # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
+        ...
+
+class PromptProcessingBatch:
+    """
+    A batch processor for prompt tokens with support for incremental processing.
+
+    This class handles batched prompt processing, managing KV caches and preparing
+    tokens for generation. It supports extending, filtering, and splitting batches.
+    """
+    @dataclass
+    class Response:
+        uid: int
+        progress: tuple
+        end_of_segment: bool
+        end_of_prompt: bool
+        ...
+
+    def __init__(
+        self,
+        model: nn.Module,
+        uids: List[int],
+        caches: List[List[Any]],
+        tokens: Optional[List[List[int]]] = ...,
+        prefill_step_size: int = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+        max_tokens: Optional[List[int]] = ...,
+    ) -> None: ...
+    def __len__(self):  # -> int:
+        ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
+    def extend(self, batch):  # -> None:
+        ...
+    def split(self, indices: List[int]):  # -> Self:
+        ...
+    def filter(self, keep: List[int]):  # -> None:
+        ...
+    def prompt(self, tokens: List[List[int]]):  # -> None:
+        """
+        Process prompt tokens through the model.
+
+        Args:
+            tokens: List of token sequences to process.
+        """
+        ...
+
+    def generate(self, tokens: List[List[int]]):  # -> GenerationBatch:
+        """
+        Transition from prompt processing to generation.
+
+        Args:
+            tokens: Final tokens for each sequence to start generation.
+
+        Returns:
+            A GenerationBatch ready for token generation.
+        """
+        ...
+
+    @classmethod
+    def empty(
+        cls,
+        model: nn.Module,
+        fallback_sampler: Callable[[mx.array], mx.array],
+        prefill_step_size: int = ...,
+    ):  # -> Self:
+        ...
+
+class GenerationBatch:
+    """
+    A batched token generator that manages multiple sequences in parallel.
+
+    This class handles the generation phase after prompt processing, managing
+    KV caches, sampling, and stop sequence detection for multiple sequences.
+    """
+    @dataclass
+    class Response:
+        uid: int
+        token: int
+        logprobs: mx.array
+        finish_reason: Optional[str]
+        current_state: Optional[str]
+        match_sequence: Optional[List[int]]
+        prompt_cache: Optional[List[Any]]
+        all_tokens: Optional[List[int]]
+        ...
+
+    model: nn.Module
+    uids: List[int]
+    prompt_cache: List[Any]
+    tokens: List[List[int]]
+    samplers: Optional[List[Callable[[mx.array], mx.array]]]
+    fallback_sampler: Callable[[mx.array], mx.array]
+    logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
+    state_machines: List[SequenceStateMachine]
+    max_tokens: List[int]
+    _current_tokens: Optional[mx.array]
+    _current_logprobs: mx.array | List[mx.array]
+    _next_tokens: Optional[mx.array]
+    _next_logprobs: mx.array | List[mx.array]
+    _token_context: List[Any]
+    _num_tokens: List[int]
+    _matcher_states: List[Any]
+
+    def __init__(
+        self,
+        model: nn.Module,
+        uids: List[int],
+        inputs: mx.array,
+        prompt_cache: List[Any],
+        tokens: List[List[int]],
+        samplers: Optional[List[Callable[[mx.array], mx.array]]],
+        fallback_sampler: Callable[[mx.array], mx.array],
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ],
+        state_machines: List[SequenceStateMachine],
+        max_tokens: List[int],
+    ) -> None: ...
+    def __len__(self) -> int: ...
+    def extend(self, batch: GenerationBatch) -> None: ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
+    def filter(self, keep: List[int]) -> None: ...
+    def _step(self) -> Tuple[List[int], List[mx.array]]: ...
+    def next(self) -> List[Response]:
+        """
+        Generate the next batch of tokens.
+
+        Returns:
+            List of Response objects for each sequence in the batch.
+        """
+        ...
+
+    @classmethod
+    def empty(
+        cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
+    ):  # -> Self:
+        ...
+
+class BatchGenerator:
+    """
+    A batch generator implements continuous batching.
+
+    This class provides automatic management of prompt processing and generation
+    batches, handling the transition between the two.
+
+    It also allows for segmented prompt processing which guarantees that the
+    generator will stop at these boundaries when processing an input.
+    """
+    def __init__(
+        self,
+        model: nn.Module,
+        max_tokens: int = ...,
+        stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
+        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        logits_processors: Optional[
+            List[Callable[[mx.array, mx.array], mx.array]]
+        ] = ...,
+        completion_batch_size: int = ...,
+        prefill_batch_size: int = ...,
+        prefill_step_size: int = ...,
+    ) -> None: ...
+    def close(self) -> None: ...
+    def __del__(self):  # -> None:
+        ...
+    @contextlib.contextmanager
+    def stats(self, stats=...):  # -> Generator[Any | BatchStats, Any, None]:
+        ...
+    _unprocessed_sequences: deque[tuple[Any, ...]]
+    _prompt_batch: PromptProcessingBatch
+    _generation_batch: GenerationBatch
+    _currently_processing: list[Any]
+    _gen_tokens_counter: int
+    _steps_counter: int
+    def _next(
+        self,
+    ) -> tuple[
+        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+    ]: ...
+    def insert(
+        self,
+        prompts: List[List[int]],
+        max_tokens: Optional[List[int]] = ...,
+        caches: Optional[List[List[Any]]] = ...,
+        all_tokens: Optional[List[List[int]]] = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+    ) -> List[int]: ...
+    def insert_segments(
+        self,
+        segments: List[List[List[int]]],
+        max_tokens: Optional[List[int]] = ...,
+        caches: Optional[List[List[Any]]] = ...,
+        all_tokens: Optional[List[List[int]]] = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+    ) -> List[int]: ...
+    def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
+    def remove(
+        self, uids: List[int], return_prompt_caches: bool = ...
+    ) -> dict[int, Any]: ...
+    @property
+    def prompt_cache_nbytes(self) -> int: ...
+    def next(
+        self,
+    ) -> tuple[
+        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+    ]: ...
+    def next_generated(self) -> List[GenerationBatch.Response]: ...
+
 @dataclass
 class BatchResponse:
     """
-    An data object to hold a batch generation response.
+    A data object to hold a batch generation response.
 
     Args:
         texts: (List[str]): The generated text for each prompt.
@@ -255,93 +515,17 @@ class BatchResponse:
     texts: List[str]
     stats: BatchStats
     caches: Optional[List[List[Any]]]
-
-def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _make_cache(
-    model: Any, left_padding: Any, max_kv_size: Optional[int]
-) -> List[Any]: ...
-def _merge_caches(caches: Any) -> List[Any]: ...
-@dataclass
-class Batch:
-    uids: List[int]
-    y: mx.array
-    logprobs: mx.array
-    max_tokens: List[int]
-    num_tokens: List[int]
-    cache: List[Any]
-    samplers: List[Any]
-    logits_processors: List[Any]
-    tokens: List[mx.array]
-    def __len__(self) -> int: ...
-    def filter(self, keep_idx: List[int]) -> None: ...
-    def extend(self, other: "Batch") -> None: ...
-    def extract_cache(self, idx: int) -> List[Any]: ...
-
-class BatchGenerator:
-    model: Any
-    max_kv_size: Optional[int]
-    prefill_step_size: int
-    unprocessed_prompts: List[Any]
-    active_batch: Optional[Batch]
-    prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
-    _stats: BatchStats
-
-    @dataclass
-    class Response:
-        uid: int
-        token: int
-        logprobs: mx.array
-        finish_reason: Optional[str]
-        prompt_cache: Any
-
-    def __init__(
-        self,
-        model: Any,
-        max_tokens: int = ...,
-        stop_tokens: Optional[set[int]] = ...,
-        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
-        logits_processors: Optional[
-            List[Callable[[mx.array, mx.array], mx.array]]
-        ] = ...,
-        completion_batch_size: int = ...,
-        prefill_batch_size: int = ...,
-        prefill_step_size: int = ...,
-        prompt_progress_callback: Optional[
-            Callable[[List[Tuple[int, int, int]]], None]
-        ] = ...,
-        max_kv_size: Optional[int] = ...,
-    ) -> None: ...
-    def close(self) -> None: ...
-    def insert(
-        self,
-        prompts: Any,
-        max_tokens: Union[List[int], int, None] = ...,
-        caches: Any = ...,
-        samplers: Optional[List[Any]] = ...,
-        logits_processors: Optional[List[Any]] = ...,
-    ) -> List[int]: ...
-    def remove(
-        self, uids: List[int], return_prompt_caches: bool = ...
-    ) -> Optional[dict[int, List[Any]]]: ...
-    def stats(self) -> BatchStats: ...
-    def next(self) -> List[Response]: ...
-    def _process_prompts(self, prompts: List[Any]) -> Batch: ...
-    def _step(
-        self,
-        input_tokens: mx.array,
-        prompt_cache: List[Any],
-        samplers: Optional[List[Any]],
-        logits_processors: Optional[List[Any]],
-        tokens: List[mx.array],
-    ) -> Tuple[mx.array, List[mx.array]]: ...
+    ...
 
 def batch_generate(
     model,
     tokenizer,
-    prompts: List[int],
+    prompts: List[List[int]],
+    prompt_caches: Optional[List[List[Any]]] = ...,
     max_tokens: Union[int, List[int]] = ...,
     verbose: bool = ...,
+    return_prompt_caches: bool = ...,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -350,14 +534,22 @@ def batch_generate(
     Args:
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
-       prompt (List[List[int]]): The input prompts.
+       prompts (List[List[int]]): The input prompts.
+       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
+          for each input prompt. Note, unlike ``generate_step``, the caches
+          won't be updated in-place.
        verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
        max_tokens (Union[int, List[int]): Maximum number of output tokens. This
           can be per prompt if a list is provided.
+       return_prompt_caches (bool): Return the prompt caches in the batch
+          responses. Default: ``False``.
+       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
+    ...
 
 def main():  # -> None:
     ...

.mlx_typings/mlx_lm/models/activations.pyi

@@ -0,0 +1,19 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+from functools import partial
+
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, x): ...
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):  # -> array:
+    ...
+
+class XieLU(nn.Module):
+    def __init__(
+        self, alpha_p_init=..., alpha_n_init=..., beta=..., eps=...
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...

.mlx_typings/mlx_lm/models/base.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 import mlx.core as mx
 
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
     bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries,
-    keys,
-    values,
-    cache,
+    queries: mx.array,
+    keys: mx.array,
+    values: mx.array,
+    cache: Optional[Any],
     scale: float,
     mask: Optional[mx.array],
     sinks: Optional[mx.array] = ...,

.mlx_typings/mlx_lm/models/cache.pyi

@@ -88,8 +88,8 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array
-    values: mx.array
+    keys: mx.array | None
+    values: mx.array | None
     offset: int
     @property
     def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,29 +268,14 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step = ...
-    def __init__(self, left_padding: List[int]) -> None:
-        """
-        The BatchKV cache expects inputs to be left-padded.
-
-        E.g. the following prompts:
-
-            [1, 3, 5]
-            [7]
-            [2, 6, 8, 9]
-
-        Should be padded like so:
-
-            [0, 1, 3, 5]
-            [0, 0, 0, 7]
-            [2, 6, 8, 9]
-
-        And ``left_padding`` specifies the amount of padding for each.
-        In this case, ``left_padding = [1, 3, 0]``.
-        """
-
-    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    _idx: int
+    def __init__(self, left_padding: List[int]) -> None: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,
@@ -316,12 +301,21 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step = ...
-    def __init__(self, max_size, left_padding: List[int]) -> None: ...
-    def update_and_fetch(
-        self, keys, values
-    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    max_size: int
+    _idx: int
+    _offset: int
+    rotated: bool
+    _lengths: array | None
+    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/deepseek_v4.pyi

@@ -0,0 +1,280 @@
+"""Type stubs for mlx_lm.models.deepseek_v4"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import ArraysCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: Optional[int]
+    n_routed_experts: int
+    num_experts_per_tok: int
+    head_dim: int
+    qk_rope_head_dim: int
+    q_lora_rank: int
+    o_lora_rank: int
+    o_groups: int
+    sliding_window: int
+    hc_mult: int
+    hc_sinkhorn_iters: int
+    hc_eps: float
+    compress_ratios: Optional[List[int]]
+    compress_rope_theta: float
+    rope_theta: float
+    rope_scaling: Optional[Dict[str, Any]]
+    rms_norm_eps: float
+    swiglu_limit: float
+    attention_bias: bool
+    max_position_embeddings: int
+
+class DeepseekV4RoPE(nn.Module):
+    dims: int
+    freqs: mx.array
+
+    def __init__(
+        self,
+        dims: int,
+        base: float,
+        scaling_config: Optional[Dict[str, Any]] = None,
+    ) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        offset: int = 0,
+        inverse: bool = False,
+    ) -> mx.array: ...
+
+class HyperConnection(nn.Module):
+    dim: int
+    hc_mult: int
+    norm_eps: float
+
+    def __init__(
+        self,
+        dim: int,
+        hc_mult: int,
+        norm_eps: float,
+        sinkhorn_iters: int,
+        hc_eps: float,
+    ) -> None: ...
+
+class HyperHead(nn.Module):
+    dim: int
+    hc_mult: int
+
+    def __init__(
+        self,
+        dim: int,
+        hc_mult: int,
+        norm_eps: float,
+        hc_eps: float,
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Compressor(nn.Module):
+    dim: int
+    head_dim: int
+    rope_head_dim: int
+    compress_ratio: int
+    overlap: bool
+    wkv_gate: nn.Linear
+    ape: mx.array
+    norm: nn.RMSNorm
+    rope: DeepseekV4RoPE
+
+    def __init__(
+        self,
+        dim: int,
+        compress_ratio: int,
+        head_dim: int,
+        rope_head_dim: int,
+        rms_norm_eps: float,
+        rope: DeepseekV4RoPE,
+    ) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        cache: "DeepseekV4Cache",
+        offset: Any,
+        key: str = ...,
+    ) -> mx.array: ...
+
+class Indexer(nn.Module):
+    def __init__(
+        self,
+        args: ModelArgs,
+        compress_ratio: int,
+        rope: DeepseekV4RoPE,
+    ) -> None: ...
+
+class _CompressorBranch:
+    buffer_kv: Optional[mx.array]
+    buffer_gate: Optional[mx.array]
+    prev_kv: Optional[mx.array]
+    prev_gate: Optional[mx.array]
+    pool: Optional[mx.array]
+    buffer_lengths: Optional[List[int]]
+    pool_lengths: Optional[List[int]]
+    buffer_count: int
+    _new_pool_lengths: Optional[List[int]]
+
+    def __init__(self) -> None: ...
+
+class DeepseekV4Cache:
+    local: RotatingKVCache
+    offset: int
+    keys: Optional[mx.array]
+    values: Optional[mx.array]
+    state: Any
+    meta_state: Any
+    nbytes: int
+    _branches: Dict[str, _CompressorBranch]
+    _pending_lengths: Optional[List[int]]
+
+    def __init__(self, sliding_window: int) -> None: ...
+    def update_and_fetch(
+        self, keys: mx.array, values: mx.array
+    ) -> tuple[mx.array, mx.array]: ...
+    def is_trimmable(self) -> bool: ...
+    def trim(self, n: int) -> int: ...
+    def empty(self) -> bool: ...
+    def size(self) -> int: ...
+    def prepare(
+        self,
+        *,
+        left_padding: Optional[List[int]] = None,
+        lengths: Optional[List[int]] = None,
+        right_padding: Optional[List[int]] = None,
+    ) -> None: ...
+    def finalize(self) -> None: ...
+    def filter(self, batch_indices: mx.array) -> None: ...
+    def extend(self, other: "DeepseekV4Cache") -> None: ...
+    def extract(self, idx: int) -> "DeepseekV4Cache": ...
+    @classmethod
+    def merge(cls, caches: List["DeepseekV4Cache"]) -> "DeepseekV4Cache": ...
+
+class V4Attention(nn.Module):
+    args: ModelArgs
+    layer_id: int
+    dim: int
+    n_heads: int
+    head_dim: int
+    rope_head_dim: int
+    nope_head_dim: int
+    n_groups: int
+    q_lora_rank: int
+    o_lora_rank: int
+    window: int
+    eps: float
+    scale: float
+    compress_ratio: int
+    wqkv_a: nn.Linear
+    q_norm: nn.RMSNorm
+    wq_b: nn.Linear
+    kv_norm: nn.RMSNorm
+    attn_sink: mx.array
+    wo_a: nn.Linear
+    wo_b: nn.Linear
+    rope: DeepseekV4RoPE
+    compressor: Compressor
+    indexer: Indexer
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DeepseekV4MLP(nn.Module):
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        swiglu_limit: float = 0.0,
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MoEGate(nn.Module):
+    weight: mx.array
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(
+        self, x: mx.array, input_ids: mx.array
+    ) -> tuple[mx.array, mx.array]: ...
+
+class DeepseekV4MoE(nn.Module):
+    num_experts_per_tok: int
+    switch_mlp: SwitchGLU
+    gate: MoEGate
+    shared_experts: DeepseekV4MLP
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: ...
+
+class DeepseekV4Block(nn.Module):
+    attn_norm: nn.RMSNorm
+    attn: V4Attention
+    hc_attn: HyperConnection
+    ffn_norm: nn.RMSNorm
+    ffn: DeepseekV4MoE
+    hc_ffn: HyperConnection
+
+    def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
+    def __call__(
+        self,
+        h: mx.array,
+        cache: Optional[Any],
+        input_ids: mx.array,
+    ) -> mx.array: ...
+
+class DeepseekV4Model(nn.Module):
+    args: ModelArgs
+    vocab_size: int
+    embed_tokens: nn.Embedding
+    layers: list[DeepseekV4Block]
+    norm: nn.RMSNorm
+    hc_head: HyperHead
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: DeepseekV4Model
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    def make_cache(self) -> list[RotatingKVCache | DeepseekV4Cache]: ...
+    @property
+    def layers(self) -> list[DeepseekV4Block]: ...

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+    use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mlx_lm/models/gemma4.pyi

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from . import gemma4_text
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: Optional[dict[str, Any]]
+    vocab_size: int
+
+    def __post_init__(self) -> None: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: gemma4_text.Model
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[gemma4_text.DecoderLayer]: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...

.mlx_typings/mlx_lm/models/gemma4_text.pyi

@@ -0,0 +1,179 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    head_dim: int
+    global_head_dim: int
+    global_partial_rotary_factor: float
+    rms_norm_eps: float
+    vocab_size: int
+    vocab_size_per_layer_input: int
+    num_key_value_heads: int
+    num_global_key_value_heads: Optional[int]
+    num_kv_shared_layers: int
+    pad_token_id: int
+    hidden_size_per_layer_input: int
+    rope_traditional: bool
+    partial_rotary_factor: float
+    rope_parameters: Optional[Dict[str, Any]]
+    sliding_window: int
+    sliding_window_pattern: int
+    max_position_embeddings: int
+    attention_k_eq_v: bool
+    final_logit_softcapping: float
+    use_double_wide_mlp: bool
+    enable_moe_block: bool
+    num_experts: Optional[int]
+    top_k_experts: Optional[int]
+    moe_intermediate_size: Optional[int]
+    layer_types: Optional[List[str]]
+    tie_word_embeddings: bool
+
+    def __post_init__(self) -> None: ...
+
+class MLP(nn.Module):
+    gate_proj: nn.Linear
+    down_proj: nn.Linear
+    up_proj: nn.Linear
+
+    def __init__(self, config: ModelArgs, layer_idx: int = 0) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Router(nn.Module):
+    proj: nn.Linear
+    scale: mx.array
+    per_expert_scale: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class Experts(nn.Module):
+    switch_glu: SwitchGLU
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
+    ) -> mx.array: ...
+
+class Attention(nn.Module):
+    layer_idx: int
+    layer_type: str
+    is_sliding: bool
+    head_dim: int
+    n_heads: int
+    n_kv_heads: int
+    use_k_eq_v: bool
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    v_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+class DecoderLayer(nn.Module):
+    layer_idx: int
+    layer_type: str
+    self_attn: Attention
+    mlp: MLP
+    enable_moe: bool
+    router: Router
+    experts: Experts
+    input_layernorm: nn.Module
+    post_attention_layernorm: nn.Module
+    pre_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm_1: nn.Module
+    post_feedforward_layernorm_2: nn.Module
+    pre_feedforward_layernorm_2: nn.Module
+    hidden_size_per_layer_input: int
+    per_layer_input_gate: Optional[nn.Linear]
+    per_layer_projection: Optional[nn.Linear]
+    post_per_layer_input_norm: Optional[nn.Module]
+    layer_scalar: mx.array
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = ...,
+        cache: Optional[Any] = ...,
+        per_layer_input: Optional[mx.array] = ...,
+        shared_kv: Optional[tuple[mx.array, mx.array]] = ...,
+        offset: Optional[mx.array] = ...,
+    ) -> tuple[mx.array, tuple[mx.array, mx.array], mx.array]: ...
+
+class Gemma4TextModel(nn.Module):
+    config: ModelArgs
+    vocab_size: int
+    window_size: int
+    sliding_window_pattern: int
+    num_hidden_layers: int
+    embed_tokens: nn.Embedding
+    embed_scale: float
+    layers: list[DecoderLayer]
+    norm: nn.Module
+    hidden_size_per_layer_input: int
+    embed_tokens_per_layer: Optional[nn.Embedding]
+    per_layer_model_projection: Optional[nn.Linear]
+    per_layer_projection_norm: Optional[nn.Module]
+    previous_kvs: list[int]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: Optional[mx.array] = ...,
+        cache: Optional[list[Any]] = ...,
+        input_embeddings: Optional[mx.array] = ...,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _get_per_layer_inputs(
+        self,
+        input_ids: Optional[mx.array],
+        input_embeddings: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _project_per_layer_inputs(
+        self,
+        input_embeddings: mx.array,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _make_masks(self, h: mx.array, cache: list[Any]) -> list[Any]: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: Gemma4TextModel
+    final_logit_softcapping: float
+    tie_word_embeddings: bool
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    @property
+    def head_dim(self) -> int: ...
+    @property
+    def n_kv_heads(self) -> int: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...

.mlx_typings/mlx_lm/models/gpt_oss.pyi

@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    vocab_size: int
+    rms_norm_eps: float
+    sliding_window: int
+    layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_key_value_groups: int
+    sinks: mx.array
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    sm_scale: float
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+    self_attn: AttentionBlock
+    mlp: MLPBlock
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+    hidden_size: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    experts: SwitchGLU
+    router: nn.Linear
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+    embed_tokens: nn.Embedding
+    norm: nn.RMSNorm
+    layer_types: List[str]
+    layers: list[TransformerBlock]
+    window_size: int
+    swa_idx: int
+    ga_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: GptOssMoeModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[nn.Module]: ...
+    def make_cache(self) -> list[KVCache]: ...

.mlx_typings/mlx_lm/models/minimax.pyi

@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+    num_heads: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+    num_experts_per_tok: int
+    gate: nn.Linear
+    switch_mlp: SwitchGLU
+    e_score_correction_bias: mx.array
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+    self_attn: MiniMaxAttention
+    block_sparse_moe: MiniMaxSparseMoeBlock
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[MiniMaxDecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: MiniMaxModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[MiniMaxDecoderLayer]: ...

.mlx_typings/mlx_lm/models/nemotron_h.pyi

@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
         cache: Optional[KVCache] = None,
     ) -> mx.array: ...
 
+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    weight: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
 class NemotronHMLP(nn.Module):
     up_proj: nn.Linear
     down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
     def __call__(self, x: mx.array) -> mx.array: ...
 
 class NemotronHMoE(nn.Module):
+    config: ModelArgs
     num_experts_per_tok: int
+    moe_latent_size: Optional[int]
     switch_mlp: SwitchMLP
+    gate: MoEGate
     shared_experts: NemotronHMLP
+    fc1_latent_proj: nn.Linear
+    fc2_latent_proj: nn.Linear
 
     def __init__(self, config: ModelArgs) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
 class Qwen3NextSparseMoeBlock(nn.Module):
     norm_topk_prob: bool
     num_experts: int
+    num_experts_per_tok: int
     top_k: int
     gate: nn.Linear
     switch_mlp: SwitchGLU

.mlx_typings/mlx_lm/models/rope_utils.pyi

@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        beta_fast: float = ...,
+        beta_slow: float = ...,
+        mscale: float = ...,
+        mscale_all_dim: float = ...,
+    ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        low_freq_factor: float = ...,
+        high_freq_factor: float = ...,
+    ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        short_factor: Any = ...,
+        long_factor: Any = ...,
+        original_max_position_embeddings: int = ...,
+    ) -> None: ...
+
+def initialize_rope(
+    dims: int,
+    base: float = ...,
+    traditional: bool = ...,
+    scaling_config: Optional[dict[str, Any]] = ...,
+    max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...

.mlx_typings/mlx_lm/sample_utils.pyi

@@ -48,6 +48,10 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = ...,
     repetition_penalty: Optional[float] = ...,
     repetition_context_size: Optional[int] = ...,
+    presence_penalty: Optional[float] = ...,
+    presence_context_size: Optional[int] = ...,
+    frequency_penalty: Optional[float] = ...,
+    frequency_context_size: Optional[int] = ...,
 ) -> list[Callable[[mx.array, mx.array], mx.array]]:
     """
     Make logits processors for use with ``generate_step``.

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -117,6 +117,8 @@ class TokenizerWrapper:
     think_end: str | None
     think_start_id: int | None
     think_end_id: int | None
+    think_start_tokens: list[int] | None
+    think_end_tokens: list[int] | None
 
     def __init__(
         self,

.mlx_typings/mlx_lm/tuner/dora.pyi

@@ -0,0 +1,50 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class DoRALinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> DoRALinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def set_linear(self, linear):  # -> None:
+        """
+        Set the self.linear layer and recompute self.m.
+        """
+        ...
+
+    def __call__(self, x): ...
+
+class DoRAEmbedding(nn.Module):
+    def from_base(
+        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> DoRAEmbedding:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> Embedding:
+        ...
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+    ) -> None: ...
+    def set_embedding(self, embedding: nn.Module):  # -> None:
+        ...
+    def __call__(self, x): ...
+    def as_linear(self, x): ...

.mlx_typings/mlx_lm/tuner/lora.pyi

@@ -0,0 +1,66 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class LoRALinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRALinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def __call__(self, x):  # -> array:
+        ...
+
+class LoRASwitchLinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Module, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRASwitchLinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedSwitchLinear | SwitchLinear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        num_experts: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def __call__(self, x, indices, sorted_indices=...): ...
+
+class LoRAEmbedding(nn.Module):
+    @staticmethod
+    def from_base(
+        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRAEmbedding:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedEmbedding | Embedding:
+        ...
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+    ) -> None: ...
+    def __call__(self, x):  # -> array:
+        ...
+    def as_linear(self, x):  # -> array:
+        ...

.mlx_typings/mlx_lm/tuner/utils.pyi

@@ -0,0 +1,57 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+from typing import Dict
+
+def build_schedule(schedule_config: Dict):  # -> Any:
+    """
+    Build a learning rate schedule from the given config.
+    """
+    ...
+
+def linear_to_lora_layers(
+    model: nn.Module, num_layers: int, config: Dict, use_dora: bool = ...
+):  # -> None:
+    """
+    Convert some of the models linear layers to lora layers.
+
+    Args:
+        model (nn.Module): The neural network model.
+        num_layers (int): The number of blocks to convert to lora layers
+        starting from the last layer.
+        config (dict): More configuration parameters for LoRA, including the
+          rank, scale, and optional layer keys.
+        use_dora (bool): If True, uses DoRA instead of LoRA.
+          Default: ``False``
+    """
+    ...
+
+def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
+    """
+    Load any fine-tuned adapters / layers.
+
+    Args:
+        model (nn.Module): The neural network model.
+        adapter_path (str): Path to the adapter configuration file.
+
+    Returns:
+        nn.Module: The updated model with LoRA layers applied.
+    """
+    ...
+
+def remove_lora_layers(model: nn.Module) -> nn.Module:
+    """
+    Remove the LoRA layers from the model.
+
+    Args:
+        model (nn.Module): The model with LoRA layers.
+
+    Returns:
+        nn.Module: The model without LoRA layers.
+    """
+    ...
+
+def print_trainable_parameters(model):  # -> None:
+    ...

.mlx_typings/mlx_vlm/prompt_utils.pyi

@@ -0,0 +1,12 @@
+from typing import Any
+
+def get_message_json(
+    model_name: str,
+    prompt: str,
+    role: str = "user",
+    skip_image_token: bool = False,
+    skip_audio_token: bool = False,
+    num_images: int = 0,
+    num_audios: int = 0,
+    **kwargs: Any,
+) -> dict[str, Any]: ...

.mlx_typings/mlx_vlm/utils.pyi

@@ -0,0 +1,15 @@
+from pathlib import Path
+from typing import Any
+
+class ImageProcessor:
+    def preprocess(
+        self, images: list[dict[str, Any]], **kwargs: Any
+    ) -> dict[str, Any]: ...
+    def __call__(self, **kwargs: Any) -> dict[str, Any]: ...
+
+def load_image_processor(
+    model_path: str | Path, **kwargs: Any
+) -> ImageProcessor | None: ...
+def load_processor(
+    model_path: str | Path, add_detokenizer: bool = ..., **kwargs: Any
+) -> ImageProcessor: ...

.mlx_typings/safetensors/__init__.pyi

@@ -0,0 +1,8 @@
+from typing import Any, Self
+
+class safe_open:
+    def __init__(self, filename: str, framework: str = "pt") -> None: ...
+    def __enter__(self) -> Self: ...
+    def __exit__(self, *args: Any) -> None: ...
+    def keys(self) -> list[str]: ...
+    def get_tensor(self, name: str) -> Any: ...

.vscode/extensions.json

@@ -1,7 +1,8 @@
 {
     "recommendations": [
         "detachhead.basedpyright",
-        "ms-python.python"
+        "ms-python.python",
+        "jnoortheen.nix-ide"
     ],
     "unwantedRecommendations": [
         "ms-python.vscode-pylance",

.vscode/settings.json

@@ -1,3 +1,32 @@
 {
-    "basedpyright.importStrategy": "fromEnvironment"
+    "files.associations": {
+        "*.nix": "nix",
+    },
+    "nix.enableLanguageServer": true,
+    "nix.serverPath": "nixd",
+    "nix.serverSettings": {
+        "nixd": {
+        "formatting": {
+            "command": ["nixpkgs-fmt"]
+        },
+        "nixpkgs": {
+            "expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.config._module.args.pkgs"
+        },
+        "options": {
+            "flake-parts": {
+                "expr": "(builtins.getFlake \"path:${workspaceFolder}\").debug.options"
+            },
+            "flake-parts-perSystem": {
+                "expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.options"
+            }
+        }
+      }
+    },
+    "[nix]": {
+        "editor.defaultFormatter": "jnoortheen.nix-ide"
+    },
+    
+    "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
+    "basedpyright.analysis.configFilePath": "${workspaceFolder}/pyproject.toml",
+    "basedpyright.importStrategy": "fromEnvironment",
 }

AGENTS.md

@@ -120,6 +120,81 @@ From .cursorrules:
 
 Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests.
 
+Integration tests live in `tests/` (root) and are opt-in via `--ignore=tests` in the default pytest addopts. They require an `eco`-managed cluster:
+
+```bash
+uv run pytest tests/ -v                    # constraint-driven host pick
+uv run pytest tests/ -v --hosts s4         # explicit host override
+```
+
+## Benchmarking
+
+Benchmarks live in `bench/`. The framework is a CLI with subcommands; each benchmark is a small library module under `bench/lib/<name>.py` plus a CLI front-end under `bench/cli/<name>.py`.
+
+```
+bench/
+├── lib/                       # composable, typed building blocks
+│   ├── prompt.py              # PromptSizer, load_tokenizer_for_bench
+│   ├── completion.py          # run_one_completion + typed payloads
+│   ├── session.py             # BenchSession (cluster + client + instance)
+│   ├── results.py             # RunMetadata, ResultsBundle, JSON schema
+│   ├── model_meta.py          # HF API: total weights size, max context, layers
+│   ├── cluster.py             # managed_cluster + managed_instance ctx-managers
+│   └── context_scaling.py     # prompt-TPS / decode-TPS vs context-size sweep
+├── cli/                       # CLI subcommands
+│   ├── _common.py             # shared argparse args + SharedOptions
+│   ├── context_scaling.py     # `python -m bench.cli context-scaling …`
+│   └── __main__.py            # subcommand dispatcher
+└── exo_bench.py, prefill_decode_bench.py
+                                # legacy CLI scripts; PromptSizer / run_one_completion
+                                # / load_tokenizer_for_bench are re-exports of bench.lib.
+```
+
+Run a benchmark:
+
+```bash
+# Defaults assume a multi-node, Thunderbolt-connected cluster with tensor
+# parallelism + JACCL: --sharding Tensor --comm MlxJaccl --thunderbolt a2a.
+# Memory + disk minimums are auto-derived from HF metadata.
+uv run python -m bench.cli context-scaling \
+  --model mlx-community/Qwen3-30B-A3B-4bit --nodes 2 --num-steps 32
+
+# Single-node smoke: opt out of TB / tensor / jaccl
+uv run python -m bench.cli context-scaling --hosts s4 \
+  --model mlx-community/Llama-3.2-1B-Instruct-4bit --num-steps 4 \
+  --sharding Pipeline --comm MlxRing --thunderbolt none
+
+# From a TOML config (CLI flags override config values)
+uv run python -m bench.cli context-scaling \
+  --config bench/configs/context_scaling.example.toml --hosts s4,s9
+```
+
+Shared CLI flags (every subcommand inherits these via `bench/cli/_common.py`):
+- `--config <path>.toml` — load run parameters from a TOML file
+- `--model`, `--sharding {Pipeline,Tensor}` (default Tensor), `--comm {MlxRing,MlxJaccl}` (default MlxJaccl), `--min-nodes` — placement
+- `--hosts`, `--nodes` (number of cluster hosts; distinct from `--min-nodes`), `--thunderbolt {a2a,ring,none}` (default a2a), `--chip` — host pool
+- `--min-memory-gb`, `--max-memory-gb`, `--min-disk-gb`, `--max-disk-gb` (minimums auto-derived from HF model size when not supplied)
+- `--evict-downloads` (default on; auto-evicts smallest-first when disk is short)
+- `--cleanup-instance` (default on; deletes the instance on exit)
+- `--output-dir`, `--tag key=value` (repeatable)
+
+Run a multi-run campaign from a single TOML file (each `[[runs]]` = its own cluster deploy + bench + teardown; `[defaults]` is shared, per-run keys override; `[plot]` triggers a comparison PNG):
+
+```bash
+uv run python -m bench.cli campaign bench/configs/llama-family-smoke.toml
+```
+
+Plot any results JSON to a PNG (auto-detects benchmark type from `metadata.benchmark`):
+
+```bash
+uv run python -m bench.cli plot bench/results/context_scaling/latest.json
+uv run python -m bench.cli plot a.json b.json --label-tag operator   # multi-run comparison
+```
+
+Adding a new benchmark = (1) write a `bench/lib/<name>.py` exposing a typed `run(session, params, bundle)` callable; (2) add a `bench/cli/<name>.py` with `add_subparser(...)` + `run(args) -> Path`; (3) register the imports in `bench/cli/__main__.py`. To enable plotting for the new benchmark, add a `render_<name>(inputs)` function in `bench/lib/plotting.py` and a dispatch entry in `bench/cli/plot.py::run`.
+
+Results land at `bench/results/<benchmark>/<run_id>.json` (with a `latest.json` symlink alongside) containing metadata (exo SHA, hostname, platform, ISO timestamps, methodology version, user tags), full cluster snapshot, the resolved + derived params, per-step rows, optional cold-control rows, and any derived summaries (e.g. `t_cum_seconds[]` for context-scaling).
+
 ## Dashboard UI Testing & Screenshots
 
 ### Building and Running the Dashboard

CONTRIBUTING.md

@@ -11,9 +11,18 @@ To run EXO from source:
   ```bash
   brew install uv
   ```
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
   ```bash
-  brew install macmon
+  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+  rustup toolchain install nightly
+  ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+  Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  ```bash
+  cargo install --git https://github.com/vladkens/macmon \
+    --rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
+    macmon \
+    --force
   ```
 
 ```bash

Cargo.lock

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

README.md

@@ -95,11 +95,10 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
   ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)
 
   ```bash
-  brew install uv macmon node
+  brew install uv node
   ```
 - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
 
@@ -107,6 +106,17 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
   rustup toolchain install nightly
   ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+
+  Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  Homebrew `macmon 0.6.1` still crashes on Apple M5.
+
+  ```bash
+  cargo install --git https://github.com/vladkens/macmon \
+    --rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
+    macmon \
+    --force
+  ```
 
 Clone the repo, build the dashboard, and run exo:
 
@@ -285,8 +295,9 @@ exo supports several environment variables for configuration:
 
 | Variable | Description | Default |
 |----------|-------------|---------|
-| `EXO_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None |
-| `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
+| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
 | `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
 | `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
 | `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
@@ -296,8 +307,11 @@ exo supports several environment variables for configuration:
 **Example usage:**
 
 ```bash
-# Use pre-downloaded models from NFS mount
-EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-models uv run exo
+# Use pre-downloaded models from NFS mount (read-only)
+EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
+
+# Download models to an external SSD (falls back to default dir if full)
+EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
 
 # Run in offline mode
 EXO_OFFLINE=true uv run exo
@@ -536,6 +550,120 @@ uv run bench/exo_bench.py \
 
 The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
 
+### Composable benchmarks (CLI)
+
+For benchmarks that need an `eco`-managed cluster and a stable JSON result format, exo ships a CLI under `bench/cli/`. The CLI handles cluster lifecycle, instance placement, model-metadata resolution (HuggingFace), and result capture; benchmark logic lives in `bench/lib/` so each new benchmark is a small library module + a CLI subcommand.
+
+**Run the prompt-TPS / decode-TPS vs context-size sweep:**
+
+The defaults assume a multi-node, Thunderbolt-connected cluster with tensor parallelism + JACCL — the typical exo benchmarking setup:
+
+```bash
+# Defaults: --sharding Tensor --comm MlxJaccl --thunderbolt a2a, with
+# memory/disk minimums auto-derived from the HF model size. eco picks
+# `--nodes` hosts from its inventory that form a TB clique and satisfy
+# those constraints.
+uv run python -m bench.cli context-scaling \
+  --model mlx-community/Qwen3-30B-A3B-4bit --nodes 2 --num-steps 32
+
+# Pin to specific hosts (defaults still apply for sharding/comm/topology)
+uv run python -m bench.cli context-scaling --hosts s4,s9 \
+  --model mlx-community/Qwen3-30B-A3B-4bit --num-steps 16
+
+# Single-node smoke test: explicit single-node placement overrides
+uv run python -m bench.cli context-scaling --hosts s4 \
+  --model mlx-community/Llama-3.2-1B-Instruct-4bit --num-steps 4 \
+  --sharding Pipeline --comm MlxRing --thunderbolt none
+
+# Override the auto-derived ramp / cold controls
+uv run python -m bench.cli context-scaling --hosts s4,s9 --model X \
+  --pp-step 4096 --num-steps 32 --cold-controls 8192,32768,65536,131072
+
+# Custom output dir + tags
+uv run python -m bench.cli context-scaling --hosts s4,s9 --model X \
+  --output-dir bench/results/2026-05-10/ --tag operator=$USER --tag run=full
+
+# Run from a TOML config (CLI flags override values from the file)
+uv run python -m bench.cli context-scaling \
+  --config bench/configs/context_scaling.example.toml
+```
+
+**Shared flags (every benchmark subcommand has these):**
+
+- `--model` — HuggingFace model id (required)
+- `--config <path>.toml` — load run parameters from a TOML file
+- `--sharding {Pipeline,Tensor}` (default **Tensor**) — sharding mode
+- `--comm {MlxRing,MlxJaccl}` (default **MlxJaccl**) — inter-node comm mode
+- `--min-nodes N` (default 1) — minimum nodes for the placement
+- `--hosts s4,s9` — pin to specific hosts; bypasses constraint search
+- `--nodes N` (default 1) — number of cluster hosts to reserve when `--hosts` is unset (distinct from `--min-nodes`, which controls the model's instance placement)
+- `--thunderbolt {a2a,ring,none}` (default **a2a**) — required Thunderbolt topology
+- `--chip "M3 Ultra"` — required chip (substring match; comment to allow any)
+- `--min-memory-gb`, `--max-memory-gb`, `--min-disk-gb`, `--max-disk-gb` — host RAM / disk constraints. The minimums are auto-derived from the HF model size (×1.30 + 1 GiB for memory, ×1.10 + 1 GiB for disk) when not supplied; explicit values always win.
+- `--evict-downloads` (default **on**) — auto-evict existing models smallest-first on disk-full; pass `--no-evict-downloads` to keep
+- `--cleanup-instance` (default **on**) — delete the placed instance after exit; pass `--no-cleanup-instance` to leave it running for debugging
+- `--output-dir bench/results` — base directory for JSON results (subcommands add their own subfolder)
+- `--tag key=value` — append to `metadata.tags` (repeatable)
+
+**Context-scaling-specific flags:**
+
+- `--num-steps N` — number of equally-spaced ramp points (default 32)
+- `--pp-step Δ` — explicit Δ in tokens (overrides auto-derivation from `max_position_embeddings`)
+- `--fraction-of-max F` — when Δ is auto-derived, use `F × max_context` as the upper bound
+- `--tg` — tokens generated per step (default 64)
+- `--warmup` — warmup requests at `pp=Δ` (default 1)
+- `--cold-controls auto` (4 evenly-spaced points across the ramp) **or** `--cold-controls 8192,32768,…` (explicit pp values). Default: no cold controls.
+
+**Output:** each run writes `bench/results/<benchmark>/<run_id>.json` plus a `latest.json` symlink. The JSON contains metadata (exo SHA, hostname, platform, user tags), the full cluster snapshot at run start, the resolved + derived params, per-step rows, cold-control rows, and derived summaries (`t_cum_seconds`, `control_gaps`).
+
+**Multi-run campaigns** — `bench campaign` runs a list of bench invocations from a single TOML file. Each `[[runs]]` entry is its own cluster deploy + bench + teardown, with a shared `[defaults]` table for DRY config:
+
+```toml
+# bench/configs/llama-family-smoke.toml
+[defaults]
+nodes = 4
+num_steps = 8
+fraction_of_max = 0.5
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.2-3b-4bit"
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.1-8b-4bit"
+
+[plot]
+label_tag = "model_short"
+```
+
+```bash
+uv run python -m bench.cli campaign bench/configs/llama-family-smoke.toml
+```
+
+After all runs finish, an optional `[plot]` table triggers a comparison plot per benchmark group (one PNG per benchmark type with ≥2 runs).
+
+**Plotting** — `bench plot` renders any results JSON to a 2-panel PNG (prompt_tps + generation_tps vs pp_tokens, cold controls overlaid as 'x' markers):
+
+```bash
+# Plot the most recent run next to its JSON
+uv run python -m bench.cli plot bench/results/context_scaling/latest.json
+
+# Compare multiple runs (one line per run; legend label = the chosen tag)
+uv run python -m bench.cli plot run_a.json run_b.json --label-tag operator
+
+# Custom output path + title
+uv run python -m bench.cli plot run.json --output /tmp/scaling.png --title "30B 4-node sweep"
+```
+
+The benchmark type is detected from each JSON's `metadata.benchmark`, so the same `plot` command will work for future benchmarks once their renderer is registered in `bench/lib/plotting.py`.
+
+Methodology for the context-scaling benchmark is documented in detail in `bench/lib/context_scaling.py`'s module docstring and in `bench/METHODOLOGY.md`.
+
 ---
 
 ## Hardware Accelerator Support

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,118 +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) {
-                    Text("What's the issue? (optional)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    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: {
@@ -711,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"
@@ -848,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/ExoProcessController.swift

@@ -4,9 +4,30 @@ import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
+private let hfEndpointKey = "EXOHFEndpoint"
 private let enableImageModelsKey = "EXOEnableImageModels"
 private let offlineModeKey = "EXOOfflineMode"
+private let fastSynchEnabledKey = "EXOFastSynchEnabled"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
+private let defaultModelsDirKey = "EXODefaultModelsDir"
+private let additionalModelsDirsKey = "EXOAdditionalModelsDirs"
+private let readOnlyModelsDirsKey = "EXOReadOnlyModelsDirs"
+private let customEnvironmentVariablesKey = "EXOCustomEnvironmentVariables"
+
+/// A user-defined environment variable that is injected into the exo child
+/// process at launch. Used as an escape hatch for env vars that don't have
+/// first-class typed UI in Settings.
+struct CustomEnvironmentVariable: Codable, Identifiable, Equatable {
+    var id: UUID
+    var key: String
+    var value: String
+
+    init(id: UUID = UUID(), key: String = "", value: String = "") {
+        self.id = id
+        self.key = key
+        self.value = value
+    }
+}
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -53,6 +74,14 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
         }
     }
+    @Published var hfEndpoint: String = {
+        return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
+        }
+    }
     @Published var enableImageModels: Bool = {
         return UserDefaults.standard.bool(forKey: enableImageModelsKey)
     }()
@@ -69,6 +98,60 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
         }
     }
+    @Published var fastSynchEnabled: Bool = {
+        if UserDefaults.standard.object(forKey: fastSynchEnabledKey) == nil {
+            return true
+        }
+        return UserDefaults.standard.bool(forKey: fastSynchEnabledKey)
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(fastSynchEnabled, forKey: fastSynchEnabledKey)
+        }
+    }
+    @Published var defaultModelsDir: String = {
+        return UserDefaults.standard.string(forKey: defaultModelsDirKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(defaultModelsDir, forKey: defaultModelsDirKey)
+        }
+    }
+    @Published var additionalModelsDirs: String = {
+        return UserDefaults.standard.string(forKey: additionalModelsDirsKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(additionalModelsDirs, forKey: additionalModelsDirsKey)
+        }
+    }
+    @Published var readOnlyModelsDirs: String = {
+        return UserDefaults.standard.string(forKey: readOnlyModelsDirsKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(readOnlyModelsDirs, forKey: readOnlyModelsDirsKey)
+        }
+    }
+    @Published var customEnvironmentVariables: [CustomEnvironmentVariable] = {
+        guard
+            let data = UserDefaults.standard.data(forKey: customEnvironmentVariablesKey),
+            let decoded = try? JSONDecoder().decode(
+                [CustomEnvironmentVariable].self, from: data
+            )
+        else {
+            return []
+        }
+        return decoded
+    }()
+    {
+        didSet {
+            guard let data = try? JSONEncoder().encode(customEnvironmentVariables) else {
+                return
+            }
+            UserDefaults.standard.set(data, forKey: customEnvironmentVariablesKey)
+        }
+    }
 
     /// Fires once when EXO transitions to `.running` for the very first time (fresh install).
     @Published private(set) var isFirstLaunchReady = false
@@ -273,12 +356,16 @@ final class ExoProcessController: ObservableObject {
         if !hfToken.isEmpty {
             environment["HF_TOKEN"] = hfToken
         }
+        if !hfEndpoint.isEmpty {
+            environment["HF_ENDPOINT"] = hfEndpoint
+        }
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
         if offlineMode {
             environment["EXO_OFFLINE"] = "true"
         }
+        environment["EXO_FAST_SYNCH"] = fastSynchEnabled ? "true" : "false"
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {
@@ -303,6 +390,29 @@ final class ExoProcessController: ObservableObject {
         }
 
         environment["PATH"] = paths.joined(separator: ":")
+
+        let trimmedDefaultModelsDir = defaultModelsDir.trimmingCharacters(in: .whitespaces)
+        if !trimmedDefaultModelsDir.isEmpty {
+            environment["EXO_DEFAULT_MODELS_DIR"] = trimmedDefaultModelsDir
+        }
+        let trimmedAdditionalModelsDirs = additionalModelsDirs.trimmingCharacters(in: .whitespaces)
+        if !trimmedAdditionalModelsDirs.isEmpty {
+            environment["EXO_MODELS_DIRS"] = trimmedAdditionalModelsDirs
+        }
+        let trimmedReadOnlyModelsDirs = readOnlyModelsDirs.trimmingCharacters(in: .whitespaces)
+        if !trimmedReadOnlyModelsDirs.isEmpty {
+            environment["EXO_MODELS_READ_ONLY_DIRS"] = trimmedReadOnlyModelsDirs
+        }
+
+        // Apply user-defined arbitrary environment variables last so that
+        // power users can override any of the typed fields above when
+        // necessary. Empty keys are ignored.
+        for variable in customEnvironmentVariables {
+            let trimmedKey = variable.key.trimmingCharacters(in: .whitespaces)
+            guard !trimmedKey.isEmpty else { continue }
+            environment[trimmedKey] = variable.value
+        }
+
         return environment
     }
 

app/EXO/EXO/Info.plist

@@ -9,7 +9,7 @@
 	<key>EXOBuildCommit</key>
 	<string>$(EXO_BUILD_COMMIT)</string>
 	<key>EXOBugReportPresignedUrlEndpoint</key>
-	<string>$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)</string>
+	<string>https://reports.exolabs.net/presigned-urls</string>
 	<key>NSLocalNetworkUsageDescription</key>
 	<string>EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference.</string>
 	<key>NSBonjourServices</key>

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

@@ -12,11 +12,15 @@ struct SettingsView: View {
 
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
+    @State private var pendingHFEndpoint: String = ""
     @State private var pendingEnableImageModels = false
     @State private var pendingOfflineMode = false
+    @State private var pendingFastSynchEnabled = false
+    @State private var pendingDefaultModelsDir: String = ""
+    @State private var pendingAdditionalModelsDirs: String = ""
+    @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 {
@@ -33,17 +37,27 @@ struct SettingsView: View {
                 .tabItem {
                     Label("Advanced", systemImage: "wrench.and.screwdriver")
                 }
+            environmentTab
+                .tabItem {
+                    Label("Environment", systemImage: "terminal")
+                }
             aboutTab
                 .tabItem {
                     Label("About", systemImage: "info.circle")
                 }
         }
-        .frame(width: 450, height: 400)
+        .frame(width: 640, height: 560)
         .onAppear {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
+            pendingHFEndpoint = controller.hfEndpoint
             pendingEnableImageModels = controller.enableImageModels
             pendingOfflineMode = controller.offlineMode
+            pendingFastSynchEnabled = controller.fastSynchEnabled
+            pendingDefaultModelsDir = controller.defaultModelsDir
+            pendingAdditionalModelsDirs = controller.additionalModelsDirs
+            pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
+            pendingCustomEnvironmentVariables = controller.customEnvironmentVariables
             needsRestart = false
         }
     }
@@ -54,9 +68,9 @@ struct SettingsView: View {
         Form {
             Section {
                 LabeledContent("Cluster Namespace") {
-                    TextField("default", text: $pendingNamespace)
+                    TextField("", text: $pendingNamespace, prompt: Text("default"))
                         .textFieldStyle(.roundedBorder)
-                        .frame(width: 200)
+                        .frame(width: 260)
                 }
                 Text("Nodes with the same namespace form a cluster. Leave empty for default.")
                     .font(.caption)
@@ -65,15 +79,26 @@ struct SettingsView: View {
 
             Section {
                 LabeledContent("HuggingFace Token") {
-                    SecureField("optional", text: $pendingHFToken)
+                    SecureField("", text: $pendingHFToken, prompt: Text("optional"))
                         .textFieldStyle(.roundedBorder)
-                        .frame(width: 200)
+                        .frame(width: 260)
                 }
                 Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
                     .font(.caption)
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("HuggingFace Endpoint") {
+                    TextField("", text: $pendingHFEndpoint, prompt: Text("default"))
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 260)
+                }
+                Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+            }
+
             Section {
                 Toggle("Offline Mode", isOn: $pendingOfflineMode)
                 Text("Skip internet checks and use only locally available models.")
@@ -124,6 +149,23 @@ struct SettingsView: View {
 
     private var advancedTab: some View {
         Form {
+            Section("Performance") {
+                Toggle("Fast Synch Enabled", isOn: $pendingFastSynchEnabled)
+                Text(
+                    "Experimental: enables fast CPU to GPU synchronization. Can sometimes cause a \"GPU lock\" where inference hangs for ~10 seconds before starting. Necessary for low latency with RDMA and Tensor Parallelism."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+
+                HStack {
+                    Spacer()
+                    Button("Save & Restart") {
+                        applyAdvancedSettings()
+                    }
+                    .disabled(!hasAdvancedChanges)
+                }
+            }
+
             Section("Onboarding") {
                 HStack {
                     VStack(alignment: .leading) {
@@ -158,8 +200,6 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
-
-                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -180,6 +220,128 @@ struct SettingsView: View {
         .padding()
     }
 
+    // MARK: - Environment Tab
+
+    private var environmentTab: some View {
+        Form {
+            Section("Models Directories") {
+                LabeledContent("Default Models Directory") {
+                    TextField(
+                        "",
+                        text: $pendingDefaultModelsDir,
+                        prompt: Text("~/.exo/models")
+                    )
+                    .textFieldStyle(.roundedBorder)
+                    .font(.system(.body, design: .monospaced))
+                    .frame(width: 260)
+                }
+                Text("Sets EXO_DEFAULT_MODELS_DIR. Where models are downloaded.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+
+                LabeledContent("Additional Directories") {
+                    TextField(
+                        "",
+                        text: $pendingAdditionalModelsDirs,
+                        prompt: Text("optional, colon-separated")
+                    )
+                    .textFieldStyle(.roundedBorder)
+                    .font(.system(.body, design: .monospaced))
+                    .frame(width: 260)
+                }
+                Text("Sets EXO_MODELS_DIRS. Extra writable model directories.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+
+                LabeledContent("Read-Only Directories") {
+                    TextField(
+                        "",
+                        text: $pendingReadOnlyModelsDirs,
+                        prompt: Text("optional, colon-separated")
+                    )
+                    .textFieldStyle(.roundedBorder)
+                    .font(.system(.body, design: .monospaced))
+                    .frame(width: 260)
+                }
+                Text("Sets EXO_MODELS_READ_ONLY_DIRS. Never written to.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+            }
+
+            Section("Custom Environment Variables") {
+                Text(
+                    "Escape hatch for env vars that don't have typed fields above. "
+                        + "Values here override the typed fields on conflict."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+
+                if pendingCustomEnvironmentVariables.isEmpty {
+                    Text("No custom variables.")
+                        .font(.caption)
+                        .foregroundColor(.secondary)
+                } else {
+                    ForEach($pendingCustomEnvironmentVariables) { $variable in
+                        HStack(alignment: .center, spacing: 8) {
+                            VStack(spacing: 4) {
+                                TextField("key", text: $variable.key)
+                                    .labelsHidden()
+                                    .textFieldStyle(.roundedBorder)
+                                    .font(.system(.body, design: .monospaced))
+                                TextField("value", text: $variable.value)
+                                    .labelsHidden()
+                                    .textFieldStyle(.roundedBorder)
+                                    .font(.system(.body, design: .monospaced))
+                            }
+                            VStack(spacing: 4) {
+                                Button {
+                                    pendingCustomEnvironmentVariables.removeAll {
+                                        $0.id == variable.id
+                                    }
+                                } label: {
+                                    Image(systemName: "minus.circle")
+                                }
+                                .buttonStyle(.borderless)
+                                .help("Remove variable")
+                                if !isValidEnvironmentVariableName(variable.key) {
+                                    Image(systemName: "exclamationmark.triangle.fill")
+                                        .foregroundColor(.orange)
+                                        .help(
+                                            "Invalid environment variable name. "
+                                                + "Must match [A-Za-z_][A-Za-z0-9_]*."
+                                        )
+                                }
+                            }
+                        }
+                    }
+                }
+
+                HStack {
+                    Button {
+                        pendingCustomEnvironmentVariables.append(
+                            CustomEnvironmentVariable()
+                        )
+                    } label: {
+                        Label("Add Variable", systemImage: "plus")
+                    }
+                    Spacer()
+                }
+            }
+
+            Section {
+                HStack {
+                    Spacer()
+                    Button("Save & Restart") {
+                        applyEnvironmentSettings()
+                    }
+                    .disabled(!hasEnvironmentChanges)
+                }
+            }
+        }
+        .formStyle(.grouped)
+        .padding()
+    }
+
     // MARK: - About Tab
 
     private var aboutTab: some View {
@@ -338,63 +500,30 @@ struct SettingsView: View {
         }
     }
 
-    private var sendBugReportButton: some View {
-        VStack(alignment: .leading, spacing: 4) {
-            Button {
-                Task {
-                    await sendBugReport()
-                }
-            } label: {
-                HStack {
-                    if bugReportInFlight {
-                        ProgressView()
-                            .scaleEffect(0.6)
-                    }
-                    Text("Send Bug Report")
-                        .font(.caption)
-                        .fontWeight(.semibold)
-                    Spacer()
-                }
-            }
-            .disabled(bugReportInFlight)
-
-            if let message = bugReportMessage {
-                Text(message)
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                    .fixedSize(horizontal: false, vertical: true)
-            }
-        }
-    }
-
     // MARK: - Actions
 
-    private func sendBugReport() async {
-        bugReportInFlight = true
-        bugReportMessage = "Collecting logs..."
-        let service = BugReportService()
-        do {
-            let outcome = try await service.sendReport(isManual: true)
-            bugReportMessage = outcome.message
-        } catch {
-            bugReportMessage = error.localizedDescription
-        }
-        bugReportInFlight = false
-    }
-
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"
         alert.informativeText = """
-            This will remove EXO and all its system components:
+            This will remove EXO and all its components:
 
             • Network configuration daemon
             • Launch at login registration
             • EXO network location
+            • EXO data directory (~/.exo)
 
             The app will be moved to Trash.
             """
         alert.alertStyle = .warning
+
+        let checkbox = NSButton(
+            checkboxWithTitle: "Keep downloaded models (~/.exo/models)",
+            target: nil, action: nil)
+        checkbox.state = .off
+        checkbox.sizeToFit()
+        alert.accessoryView = checkbox
+
         alert.addButton(withTitle: "Uninstall")
         alert.addButton(withTitle: "Cancel")
 
@@ -404,11 +533,11 @@ struct SettingsView: View {
 
         let response = alert.runModal()
         if response == .alertFirstButtonReturn {
-            performUninstall()
+            performUninstall(keepModels: checkbox.state == .on)
         }
     }
 
-    private func performUninstall() {
+    private func performUninstall(keepModels: Bool) {
         uninstallInProgress = true
 
         controller.cancelPendingLaunch()
@@ -418,6 +547,7 @@ struct SettingsView: View {
         DispatchQueue.global(qos: .utility).async {
             do {
                 try NetworkSetupHelper.uninstall()
+                try Self.removeExoDirectory(keepModels: keepModels)
 
                 DispatchQueue.main.async {
                     LaunchAtLoginHelper.disable()
@@ -441,6 +571,23 @@ struct SettingsView: View {
         }
     }
 
+    private static func removeExoDirectory(keepModels: Bool) throws {
+        let fm = FileManager.default
+        let exoDir = ExoProcessController.exoDirectoryURL
+        guard fm.fileExists(atPath: exoDir.path) else { return }
+
+        if !keepModels {
+            try fm.removeItem(at: exoDir)
+            return
+        }
+
+        let contents = try fm.contentsOfDirectory(
+            at: exoDir, includingPropertiesForKeys: nil, options: [])
+        for entry in contents where entry.lastPathComponent != "models" {
+            try? fm.removeItem(at: entry)
+        }
+    }
+
     private func moveAppToTrash() {
         guard let appURL = Bundle.main.bundleURL as URL? else { return }
         do {
@@ -454,6 +601,7 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
+            || pendingHFEndpoint != controller.hfEndpoint
             || pendingOfflineMode != controller.offlineMode
     }
 
@@ -461,9 +609,21 @@ struct SettingsView: View {
         pendingEnableImageModels != controller.enableImageModels
     }
 
+    private var hasAdvancedChanges: Bool {
+        pendingFastSynchEnabled != controller.fastSynchEnabled
+    }
+
+    private var hasEnvironmentChanges: Bool {
+        pendingDefaultModelsDir != controller.defaultModelsDir
+            || pendingAdditionalModelsDirs != controller.additionalModelsDirs
+            || pendingReadOnlyModelsDirs != controller.readOnlyModelsDirs
+            || pendingCustomEnvironmentVariables != controller.customEnvironmentVariables
+    }
+
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
+        controller.hfEndpoint = pendingHFEndpoint
         controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
@@ -473,6 +633,75 @@ struct SettingsView: View {
         restartIfRunning()
     }
 
+    private func applyAdvancedSettings() {
+        controller.fastSynchEnabled = pendingFastSynchEnabled
+        restartIfRunning()
+    }
+
+    private func applyEnvironmentSettings() {
+        controller.defaultModelsDir = pendingDefaultModelsDir.trimmingCharacters(
+            in: .whitespaces)
+        controller.additionalModelsDirs = pendingAdditionalModelsDirs.trimmingCharacters(
+            in: .whitespaces)
+        controller.readOnlyModelsDirs = pendingReadOnlyModelsDirs.trimmingCharacters(
+            in: .whitespaces)
+
+        pendingDefaultModelsDir = controller.defaultModelsDir
+        pendingAdditionalModelsDirs = controller.additionalModelsDirs
+        pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
+
+        // Trim whitespace from keys and drop empty ones so that the stored
+        // form matches what is actually injected into the child process and
+        // hasEnvironmentChanges doesn't show a stale diff after save.
+        let trimmed: [CustomEnvironmentVariable] =
+            pendingCustomEnvironmentVariables.compactMap { variable in
+                let key = variable.key.trimmingCharacters(in: .whitespaces)
+                guard !key.isEmpty else { return nil }
+                return CustomEnvironmentVariable(
+                    id: variable.id, key: key, value: variable.value
+                )
+            }
+
+        // De-duplicate keys, keeping the last occurrence. This matches the
+        // effective semantics of the dictionary assignment in
+        // ExoProcessController.makeEnvironment and avoids silently losing
+        // visible rows after save.
+        var seenKeys = Set<String>()
+        var deduplicatedReversed: [CustomEnvironmentVariable] = []
+        for variable in trimmed.reversed() {
+            if seenKeys.insert(variable.key).inserted {
+                deduplicatedReversed.append(variable)
+            }
+        }
+        let sanitized = Array(deduplicatedReversed.reversed())
+
+        pendingCustomEnvironmentVariables = sanitized
+        controller.customEnvironmentVariables = sanitized
+
+        restartIfRunning()
+    }
+
+    /// Validates a POSIX-style environment variable name:
+    /// `[A-Za-z_][A-Za-z0-9_]*`. Uses an ASCII-only charset so that
+    /// Unicode letters (e.g. `ñ`, Cyrillic) are rejected in line with what
+    /// the help tooltip advertises. Empty strings are treated as valid
+    /// here so that a freshly added blank row does not immediately look
+    /// broken; the save step filters empty keys out instead.
+    private func isValidEnvironmentVariableName(_ key: String) -> Bool {
+        if key.isEmpty { return true }
+        let headAllowed = CharacterSet(
+            charactersIn: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_"
+        )
+        let tailAllowed = headAllowed.union(CharacterSet(charactersIn: "0123456789"))
+        guard let first = key.unicodeScalars.first, headAllowed.contains(first) else {
+            return false
+        }
+        for scalar in key.unicodeScalars.dropFirst() {
+            if !tailAllowed.contains(scalar) { return false }
+        }
+        return true
+    }
+
     private func restartIfRunning() {
         if controller.status == .running || controller.status == .starting {
             controller.restart()

app/EXO/EXO/Views/SettingsWindowController.swift

@@ -30,7 +30,7 @@ final class SettingsWindowController: ObservableObject {
         let hostingView = NSHostingView(rootView: settingsView)
 
         let newWindow = NSWindow(
-            contentRect: NSRect(x: 0, y: 0, width: 450, height: 400),
+            contentRect: NSRect(x: 0, y: 0, width: 640, height: 560),
             styleMask: [.titled, .closable],
             backing: .buffered,
             defer: false

app/EXO/uninstall-exo.sh

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

bench/METHODOLOGY.md

@@ -0,0 +1,160 @@
+# Exo-Bench — Methodology
+
+exo bench measures inference throughput and resource consumption of an exo cluster under controlled conditions. It sends prompts to the `/bench/chat/completions` endpoint, collects server-reported timing statistics, and records system-level metrics (power, GPU utilisation, temperature) throughout each run.
+
+The goal is to have accurate, transparent and reproducible numbers to compare speed and scaling across different models and different setups, and to be able to track these results as optimizations and features are added to EXO.
+
+Below is a technical summary of how Exo-Bench works. While the methodology and benchmark may change over time, this document will be kept up to date whenever this happens. If you find an issue with the methodology, or would like a  feature to be added, please open a GitHub issue!
+
+---
+
+## Prompt Construction
+
+Benchmarks need prompts of an exact token length. Unfortunately, we do not have direct access to the model but just the chat completion endpoint. To get around this fact, we create a request that will tokenise to a certain prompt length.
+
+This is achieved by:
+
+1. Tokenising a sample message through the model's `apply_chat_template()` to measure overhead (system tokens, special tokens, chat formatting).
+2. Binary-searching over a repeated atom string (default `"a "`) to find the content length that produces exactly the target number of tokens after template expansion.
+3. Returning both the content string and the verified token count.
+
+The actual token count is recorded in every result row as `pp_tokens`, so downstream analysis can confirm the prompt hit its target.
+
+Chat template formatting means that it may be impossible to attain very small pp benchmarks. e.g. pp=32 may not work. This tradeoff was made because the result of such a small prompt does not seem very interesting or useful for any real-world use cases.
+
+---
+
+## Bench Endpoint
+
+When a request reaches the server via the `/bench/chat/completions` endpoint, three things change compared to a normal chat completion:
+
+- **KV prefix cache is disabled by default**. Every request starts from a cold cache, ensuring prefill timing is not affected by prior requests. See [Prefix Cache Mode](#prefix-cache-mode) for the `--use-prefix-cache` option.
+- **EOS tokens are banned**. A logits processor suppresses all end-of-sequence tokens, forcing the model to generate exactly `max_tokens` tokens. This guarantees consistent generation length for fair TPS comparison — the model cannot short-circuit a run by stopping early.
+- **No model output parsing**. The bench collection path concatenates raw token text without any model-specific post-processing (thinking tag extraction, structured output handling, etc.). This is to avoid model outputs such as tool parsing or any structural mistakes from breaking the benchmark - we are testing for speed; see Exo-Eval for performance metrics.
+
+---
+
+## Timing
+
+### Prefill TPS
+
+Measured server-side per task.
+
+```
+prefill_tps = num_prompt_tokens / prefill_wall_seconds
+```
+
+### Generation TPS
+
+Measured server-side per task. Each task records wall-clock timestamps as tokens arrive:
+
+- First generated token: timestamp recorded
+- Every subsequent token: timestamp updated
+
+When generation completes:
+
+```
+gen_span = last_token_time - first_token_time
+generation_tps = (completion_tokens - 1) / gen_span
+```
+
+The first token is excluded from the numerator because the rate measures inter-token throughput — the time between the first and last token divided by the number of intervals.
+
+This does mean that tg=1 will not work.
+
+---
+
+## Concurrency
+
+### Single Request
+
+The client records wall-clock `elapsed_s` around the HTTP round-trip (network latency + server prefill + generation + response serialisation). This is a convenience metric for end-to-end latency. The authoritative TPS numbers come from the server-side per-task timing in the `generation_stats` response.
+
+### Concurrent Requests
+
+When `--concurrency N` is set with N > 1, all N requests must hit the server at the same instant. The mechanism:
+
+1. The prompt is built once and shared across all threads.
+2. Each thread gets its own HTTP connection.
+3. A thread barrier blocks all threads until every thread is ready.
+4. The first thread past the barrier records the batch start time and signals the others.
+5. All threads use the same start time as their reference, then fire their HTTP request.
+6. Each thread's `elapsed_s` is measured from the shared start time to its own response completion.
+
+**Batch wall time** is the maximum `elapsed_s` across all N requests — the time until the last request finishes.
+
+### Aggregate TPS
+
+```
+per_req_tps = max(generation_tps across N concurrent requests)
+agg_gen_tps = per_req_tps * concurrency
+```
+
+`max` is used instead of `mean` because all requests run in parallel against the same model. The fastest request's generation rate represents the system's per-stream throughput capacity; multiplying by concurrency gives aggregate throughput.
+
+---
+
+## Prefix Cache Mode
+
+When `--use-prefix-cache` is passed, the KV prefix cache remains active during benchmarking. This speeds up repeated runs by skipping redundant prefill work, which is useful when prompt processing is not the focus of the benchmark (e.g. when measuring generation throughput or power consumption across many configurations).
+
+Each response includes a `prefix_cache_hit` field (`"none"`, `"partial"`, or `"exact"`):
+
+- **none**: Cold prefill — no cached KV state was available. The reported `prompt_tps` is the real prefill throughput.
+- **partial**: A prefix of the prompt was found in cache. Only the remaining tokens were prefilled. The reported `prompt_tps` reflects the real throughput on the uncached portion. This occurs when multiple ascending `--pp` values share a common prefix (e.g. `--pp 1000,5000` — the 5000-token prompt reuses the 1000-token cache entry and prefills the remaining 4000 tokens).
+- **exact**: The entire prompt was found in cache (e.g. same `--pp` value on a `--repeat`). No prefill work was done. The reported `prompt_tps` is the TPS from when the cache entry was originally created, not a new measurement.
+
+**Prompt TPS is approximate in this mode.** Exact-hit runs report the stored TPS from the original cold/partial prefill rather than a freshly measured value. For accurate cold prefill numbers, run without `--use-prefix-cache`.
+
+Ascending `--pp` order (e.g. `--pp 1000,5000,10000`) gives the most useful data: each size gets a meaningful partial hit except the first which is cold. Descending order produces exact hits with approximate TPS from a longer prompt's original run.
+
+---
+
+## Warmup
+
+Before measurement begins, `--warmup N` (default: 0) discarded requests are sent using the first pp/tg pair. Warmup results are not included in the output.
+
+---
+
+## System Metrics
+
+A background thread polls each node at 1 Hz, collecting:
+
+- GPU utilisation (%)
+- Temperature (C)
+- System power draw (W)
+- CPU cluster usage (performance and efficiency cores)
+
+**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
+
+---
+
+## Output Format
+
+Results are written as JSON with three top-level keys:
+
+- **`runs`**: Array of per-request result objects, each containing:
+  - `elapsed_s`, `output_text_preview` (first 200 chars)
+  - `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
+  - Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
+  - Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
+  - `download_duration_s` (if model was freshly downloaded)
+- **`cluster`**: Cluster state snapshot at time of benchmarking.
+- **`system_metrics`**: Per-node time-series samples (GPU, power, temperature).
+
+---
+
+## Reproducing Results
+
+```bash
+cd bench && uv run python exo_bench.py \
+  --model "mlx-community/Qwen3.5-27B-4bit" \
+  --instance-meta jaccl \
+  --sharding tensor \
+  --min-nodes 2 --max-nodes 2 \
+  --pp 512 4096 --tg 128 \
+  --repeat 3 \
+  --warmup 1
+```
+
+Run --help for all the available flags.

bench/cli/__init__.py

@@ -0,0 +1,14 @@
+"""CLI front-ends for bench library benchmarks.
+
+Each benchmark is a sub-package / module with two pieces:
+
+  - a ``run(...)`` callable in ``bench.lib.<name>`` that does the actual
+    measurement (no argparse, no eco, no I/O)
+  - an ``add_subparser(subparsers)`` helper here that wires CLI args to a
+    handler invoking the lib
+
+The main entry point dispatches to the requested subcommand:
+
+    uv run python -m bench.cli context-scaling --hosts s4 \\
+        --model mlx-community/Qwen3-30B-A3B-4bit
+"""

bench/cli/__main__.py

@@ -0,0 +1,56 @@
+"""``python -m bench.cli`` — dispatcher for benchmark subcommands.
+
+To add a new benchmark:
+
+  1. Implement the methodology in ``bench.lib.<name>`` exposing a typed
+     ``run(session, params, bundle)`` callable (no argparse, no eco I/O).
+  2. Implement a ``bench.cli.<name>`` module with an ``add_subparser`` and
+     a ``run(args) -> Path`` handler.
+  3. Add an ``import + add_subparser(subparsers)`` line below.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from typing import cast
+
+from bench.cli import campaign, context_scaling, plot
+from bench.cli._common import expand_config_in_argv
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m bench.cli",
+        description=(
+            "Composable, eco-managed benchmarks for exo. "
+            "Pick a subcommand and pass its model / cluster options."
+        ),
+    )
+    subparsers = parser.add_subparsers(
+        dest="subcommand",
+        required=True,
+        metavar="SUBCOMMAND",
+    )
+    context_scaling.add_subparser(subparsers)
+    plot.add_subparser(subparsers)
+    campaign.add_subparser(subparsers)
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    raw_argv = list(argv if argv is not None else sys.argv[1:])
+    expanded = expand_config_in_argv(raw_argv)
+    args = _build_parser().parse_args(expanded)
+    handler = getattr(args, "handler", None)
+    if not callable(handler):
+        subcommand = getattr(args, "subcommand", "<unknown>")
+        raise SystemExit(f"subcommand {subcommand!r} did not register a handler")
+    cast("Callable[[argparse.Namespace], Path]", handler)(args)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

bench/cli/_common.py

@@ -0,0 +1,345 @@
+"""Shared CLI argument parsing for the bench command-line interface.
+
+Every benchmark subcommand inherits the same model / cluster / output
+arguments via :func:`add_shared_args` and consumes them through
+:class:`SharedOptions`. argparse's ``Namespace.<attr>`` is fundamentally
+typed ``Any``; the :func:`get_arg` / :func:`get_arg_optional` helpers are
+the single boundary where we coerce to typed values.
+
+A ``--config <path>.toml`` flag lets the caller capture a run definition
+in a TOML file. :func:`expand_config_in_argv` rewrites argv in place,
+substituting the config's keys as CLI flags placed *before* any explicit
+user args so that explicit CLI flags always win.
+"""
+
+from __future__ import annotations
+
+import argparse
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, TypeVar
+
+from exo_tools.cluster import Chip, Thunderbolt
+from exo_tools.harness import Comm, Sharding
+
+_T = TypeVar("_T")
+
+
+def get_arg(args: argparse.Namespace, name: str, type_: type[_T]) -> _T:
+    """Return ``args.<name>``, asserting it's an instance of ``type_``.
+
+    For ``int`` and ``float`` we additionally accept inputs that ``int(.)`` /
+    ``float(.)`` would parse, since argparse's ``type=int`` already coerces
+    cleanly on input but post-`set_defaults` callers may pass raw values.
+    """
+    raw: Any = getattr(args, name)  # type: ignore[reportAny]
+    if isinstance(raw, type_):
+        return raw
+    if type_ is int and isinstance(raw, (int, str)):
+        return int(raw)  # type: ignore[return-value]
+    if type_ is float and isinstance(raw, (int, float, str)):
+        return float(raw)  # type: ignore[return-value]
+    raise TypeError(
+        f"argparse field {name!r} expected {type_.__name__}, got {type(raw).__name__}"  # type: ignore[reportUnknownArgumentType]
+    )
+
+
+def get_arg_optional(args: argparse.Namespace, name: str, type_: type[_T]) -> _T | None:
+    """Like :func:`get_arg` but allows the field to be missing or None."""
+    raw = getattr(args, name, None)
+    if raw is None:
+        return None
+    if isinstance(raw, type_):
+        return raw
+    if type_ is int and isinstance(raw, (int, str)):
+        return int(raw)  # type: ignore[return-value]
+    if type_ is float and isinstance(raw, (int, float, str)):
+        return float(raw)  # type: ignore[return-value]
+    raise TypeError(
+        f"argparse field {name!r} expected {type_.__name__} or None, "
+        f"got {type(raw).__name__}"  # type: ignore[reportUnknownArgumentType]
+    )
+
+
+@dataclass(frozen=True)
+class SharedOptions:
+    """Parsed shared CLI options for any benchmark."""
+
+    model: str
+    hosts: tuple[str, ...]
+    nodes: int
+    thunderbolt: Thunderbolt | None
+    chip: Chip | None
+    min_memory_gb: float | None
+    max_memory_gb: float | None
+    min_disk_gb: float | None
+    max_disk_gb: float | None
+    evict_downloads: bool
+    sharding: Sharding
+    comm: Comm
+    min_nodes: int
+    output_dir: Path
+    tags: dict[str, str]
+    cleanup_instance: bool
+    user_prefix: str
+
+    @classmethod
+    def from_namespace(cls, args: argparse.Namespace) -> SharedOptions:
+        hosts_raw = get_arg_optional(args, "hosts", str)
+        thunderbolt_raw = get_arg_optional(args, "thunderbolt", str)
+        chip_raw = get_arg_optional(args, "chip", str)
+        tag_list_raw: object = getattr(args, "tag", None) or []
+        if isinstance(tag_list_raw, list):
+            tag_list: list[str] = [
+                str(t)  # type: ignore[reportUnknownArgumentType]
+                for t in tag_list_raw  # type: ignore[reportUnknownVariableType]
+            ]
+        else:
+            tag_list = []
+        return cls(
+            model=get_arg(args, "model", str),
+            hosts=tuple(_parse_csv(hosts_raw)) if hosts_raw else (),
+            nodes=get_arg(args, "nodes", int),
+            thunderbolt=Thunderbolt(thunderbolt_raw) if thunderbolt_raw else None,
+            chip=Chip(chip_raw) if chip_raw else None,
+            min_memory_gb=get_arg_optional(args, "min_memory_gb", float),
+            max_memory_gb=get_arg_optional(args, "max_memory_gb", float),
+            min_disk_gb=get_arg_optional(args, "min_disk_gb", float),
+            max_disk_gb=get_arg_optional(args, "max_disk_gb", float),
+            evict_downloads=get_arg(args, "evict_downloads", bool),
+            sharding=Sharding(get_arg(args, "sharding", str)),
+            comm=Comm(get_arg(args, "comm", str)),
+            min_nodes=get_arg(args, "min_nodes", int),
+            output_dir=Path(get_arg(args, "output_dir", str)),
+            tags=_parse_tags(tag_list),
+            cleanup_instance=get_arg(args, "cleanup_instance", bool),
+            user_prefix=get_arg(args, "eco_user_prefix", str),
+        )
+
+
+def add_shared_args(parser: argparse.ArgumentParser) -> None:
+    """Register the shared-arg group on ``parser``.
+
+    The bool flags (``--auto-constrain``, ``--evict-downloads``,
+    ``--cleanup-instance``) all default to True and use
+    :class:`argparse.BooleanOptionalAction` so callers opt out via the
+    ``--no-X`` form (or set ``X = false`` in a TOML config).
+    """
+    g_config = parser.add_argument_group("config file")
+    g_config.add_argument(
+        "--config",
+        default=None,
+        help="TOML file with run parameters. CLI flags placed after --config "
+        "override values from the file.",
+    )
+
+    g_model = parser.add_argument_group("model")
+    g_model.add_argument(
+        "--model",
+        required=True,
+        help="HuggingFace model id. To run multiple models in one go, use "
+        "the 'campaign' subcommand with a TOML file listing each as a "
+        "separate [[runs]] entry.",
+    )
+    g_model.add_argument(
+        "--sharding",
+        default=Sharding.TENSOR.value,
+        choices=[s.value for s in Sharding],
+        help="Sharding mode for the placed instance. Default 'Tensor' (splits "
+        "layers within nodes; pairs with --comm MlxJaccl for high throughput "
+        "on TB-connected clusters). Use 'Pipeline' for layer-per-node sharding "
+        "(typical for single-node smoke tests).",
+    )
+    g_model.add_argument(
+        "--comm",
+        default=Comm.JACCL.value,
+        choices=[c.value for c in Comm],
+        help="Inter-node communication mode. Default 'MlxJaccl' (RDMA over "
+        "Thunderbolt; pairs with --sharding Tensor and --thunderbolt a2a). "
+        "Use 'MlxRing' for ring all-reduce over the regular network.",
+    )
+    g_model.add_argument("--min-nodes", type=int, default=1)
+
+    g_cluster = parser.add_argument_group("cluster")
+    g_cluster.add_argument(
+        "--hosts",
+        default=None,
+        help="Comma-separated host list (e.g. s4,s9). Bypasses constraint search.",
+    )
+    g_cluster.add_argument(
+        "--nodes",
+        type=int,
+        default=1,
+        help="Number of cluster nodes (hosts) to deploy on. "
+        "Distinct from --min-nodes which controls the model's instance placement.",
+    )
+    g_cluster.add_argument(
+        "--thunderbolt",
+        default=Thunderbolt.A2A.value,
+        choices=[t.value for t in Thunderbolt],
+        help="Thunderbolt topology required: 'a2a' (clique, default; needed "
+        "for tensor parallelism + JACCL), 'ring' (cycle; for pipeline + JACCL), "
+        "or 'none' (exclude TB-connected hosts; pair with --sharding Pipeline "
+        "--comm MlxRing).",
+    )
+    g_cluster.add_argument(
+        "--chip",
+        default=None,
+        choices=[c.value for c in Chip],
+        help="Chip required (e.g. 'M3 Ultra')",
+    )
+    g_cluster.add_argument(
+        "--min-memory-gb",
+        type=float,
+        default=None,
+        help="Min RAM (GB) on each host. If unset, auto-derived from the HF "
+        "model size (×1.30 + 1 GiB).",
+    )
+    g_cluster.add_argument(
+        "--max-memory-gb",
+        type=float,
+        default=None,
+        help="Max RAM (GB) on each host. Useful to leave bigger machines "
+        "free for other workloads.",
+    )
+    g_cluster.add_argument(
+        "--min-disk-gb",
+        type=float,
+        default=None,
+        help="Min free disk (GB) on each host. If unset, auto-derived from "
+        "the HF model size (×1.10 + 1 GiB).",
+    )
+    g_cluster.add_argument(
+        "--max-disk-gb",
+        type=float,
+        default=None,
+        help="Max disk (GB) on each host.",
+    )
+
+    g_runtime = parser.add_argument_group("runtime")
+    g_runtime.add_argument(
+        "--evict-downloads",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Auto-evict existing models (smallest first) when disk is short to "
+        "make room for the bench model. Default on; pass --no-evict-downloads "
+        "to keep existing downloads.",
+    )
+    g_runtime.add_argument(
+        "--cleanup-instance",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Clean up the placed instance after the benchmark exits. "
+        "Default on; pass --no-cleanup-instance to leave it running for debugging.",
+    )
+    g_runtime.add_argument(
+        "--eco-user-prefix",
+        default="bench",
+        help="USER prefix for the eco session (default: 'bench').",
+    )
+
+    g_output = parser.add_argument_group("output")
+    g_output.add_argument(
+        "--output-dir",
+        default="bench/results",
+        help="Base directory for JSON results. Subcommands may add a sub-folder.",
+    )
+    g_output.add_argument(
+        "--tag",
+        action="append",
+        default=[],
+        help="Add a 'key=value' tag to metadata.tags (repeatable).",
+    )
+
+
+# ---------------------------------------------------------------------------
+# TOML config expansion
+# ---------------------------------------------------------------------------
+
+
+def expand_config_in_argv(argv: list[str]) -> list[str]:
+    """If ``--config <path>`` appears in ``argv``, splice the TOML's contents in.
+
+    The TOML file's keys are converted to CLI flags (``foo_bar`` →
+    ``--foo-bar``) and inserted *before* the user's other args, so explicit
+    CLI flags always override the config. The ``--config <path>`` pair
+    itself is removed from argv. The first arg (the subcommand name) is
+    preserved at index 0.
+
+    Special handling:
+      - ``[tags]`` table → repeated ``--tag key=value`` occurrences
+      - lists → joined as a comma-separated value (matches the parser's
+        CSV handling for ``--hosts``)
+      - bool true/false → ``--key`` / ``--no-key`` (assumes the underlying
+        flag uses :class:`argparse.BooleanOptionalAction`)
+    """
+    if "--config" not in argv:
+        return list(argv)
+
+    idx = argv.index("--config")
+    if idx + 1 >= len(argv):
+        raise ValueError("--config requires a path argument")
+    config_path = Path(argv[idx + 1])
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    with config_path.open("rb") as f:
+        config_data: dict[str, Any] = tomllib.load(f)
+
+    expanded = _config_to_argv(config_data)
+    stripped = list(argv[:idx]) + list(argv[idx + 2 :])
+    if not stripped:
+        return expanded
+    # The subcommand name must come first; insert config-derived args
+    # right after it so that the user's later explicit args override.
+    return [stripped[0]] + expanded + stripped[1:]
+
+
+def _config_to_argv(data: dict[str, Any]) -> list[str]:
+    """Convert a TOML-loaded dict to a list of argv-style CLI flags."""
+    out: list[str] = []
+    for key in data:
+        value: Any = data[key]  # type: ignore[reportAny]
+        if key == "tags" and isinstance(value, dict):
+            for tag_key, tag_value in value.items():  # type: ignore[reportUnknownVariableType]
+                out.extend(["--tag", f"{tag_key}={tag_value}"])
+            continue
+        if value is None:
+            continue
+        flag = "--" + key.replace("_", "-")
+        if isinstance(value, bool):
+            out.append(flag if value else f"--no-{key.replace('_', '-')}")
+        elif isinstance(value, list):
+            joined = ",".join(
+                str(x)  # type: ignore[reportUnknownArgumentType]
+                for x in value  # type: ignore[reportUnknownVariableType]
+            )
+            out.extend([flag, joined])
+        else:
+            out.extend([flag, str(value)])  # type: ignore[reportAny]
+    return out
+
+
+def _parse_csv(raw: str) -> list[str]:
+    return [s.strip() for s in raw.split(",") if s.strip()]
+
+
+def _parse_tags(raw: list[str]) -> dict[str, str]:
+    out: dict[str, str] = {}
+    for entry in raw:
+        if "=" not in entry:
+            raise argparse.ArgumentTypeError(
+                f"--tag must be 'key=value', got {entry!r}"
+            )
+        k, v = entry.split("=", 1)
+        out[k.strip()] = v.strip()
+    return out
+
+
+@dataclass
+class CommandResult:
+    """Return value from a benchmark CLI handler."""
+
+    output_path: Path | None = None
+    extra: dict[str, str] = field(default_factory=dict)

bench/cli/campaign.py

@@ -0,0 +1,220 @@
+"""Run a campaign of bench invocations from a single TOML file.
+
+A campaign config has a ``[defaults]`` table (applied to every run) and a
+list of ``[[runs]]`` entries (each a fully-formed invocation with its own
+``subcommand``). The campaign runner merges defaults with each run's
+overrides, dispatches to the matching subcommand handler, and collects
+the output JSON paths.
+
+Each run gets its own cluster — the deploy / teardown happens per-run.
+After all runs finish, an optional ``[plot]`` table triggers a comparison
+plot per benchmark group.
+
+Schema::
+
+    [defaults]
+    nodes = 4
+    num_steps = 8
+
+    [[runs]]
+    subcommand = "context-scaling"
+    model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+    [runs.tags]
+    model_short = "llama-3.2-3b"
+
+    [[runs]]
+    subcommand = "context-scaling"
+    model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
+    [runs.tags]
+    model_short = "llama-3.1-8b"
+
+    [plot]
+    label_tag = "model_short"
+"""
+
+from __future__ import annotations
+
+import argparse
+import tomllib
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, cast
+
+from loguru import logger
+
+from bench.cli import context_scaling
+from bench.cli._common import (
+    _config_to_argv,  # type: ignore[reportPrivateUsage]
+    get_arg,
+)
+from bench.lib.plotting import PlotInputs, render_context_scaling
+
+# Each subcommand exposes its argparse via add_subparser. The campaign
+# runner builds a one-off parser per run with only the chosen subcommand
+# registered, parses the run-derived argv, and invokes the handler.
+_SUBCOMMAND_PARSERS: dict[
+    str,
+    Callable[
+        [Any], None
+    ],  # subparsers action — argparse private; Any-typed at boundary
+] = {
+    "context-scaling": context_scaling.add_subparser,
+}
+
+
+def add_subparser(
+    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],  # type: ignore[type-arg]
+) -> None:
+    parser = subparsers.add_parser(
+        "campaign",
+        help="Run a list of bench invocations from a single TOML config.",
+        description=__doc__,
+    )
+    parser.add_argument(
+        "config",
+        type=str,
+        help="TOML campaign file (with [defaults] + [[runs]] tables).",
+    )
+    parser.add_argument(
+        "--no-plot",
+        action="store_true",
+        help="Skip the optional comparison plot at the end of the campaign.",
+    )
+    parser.set_defaults(handler=run)
+
+
+# ---------------------------------------------------------------------------
+
+
+def run(args: argparse.Namespace) -> Path | None:
+    config_path = Path(get_arg(args, "config", str))
+    if not config_path.is_file():
+        raise SystemExit(f"campaign: file not found: {config_path}")
+
+    with config_path.open("rb") as f:
+        raw = tomllib.load(f)
+
+    defaults = _table(raw, "defaults")
+    runs_obj = raw.get("runs")
+    if not isinstance(runs_obj, list) or not runs_obj:
+        raise SystemExit(f"campaign: {config_path}: missing or empty [[runs]] list")
+    runs_raw: list[Any] = cast("list[Any]", runs_obj)
+    plot_cfg = _table(raw, "plot")
+
+    n_runs = len(runs_raw)
+    output_paths: dict[str, list[Path]] = {}
+    for i, run_obj in enumerate(runs_raw):  # type: ignore[reportAny]
+        if not isinstance(run_obj, dict):
+            raise SystemExit(
+                f"campaign: run #{i + 1}: expected a TOML table, "
+                f"got {type(run_obj).__name__}"  # type: ignore[reportUnknownArgumentType]
+            )
+        run_cfg = cast("dict[str, Any]", run_obj)
+
+        merged = _merge(defaults, run_cfg)
+        subcommand_obj: Any = merged.pop("subcommand", None)  # type: ignore[reportAny]
+        if not isinstance(subcommand_obj, str):
+            raise SystemExit(
+                f"campaign: run #{i + 1}: 'subcommand' field is required (str)"
+            )
+        if subcommand_obj not in _SUBCOMMAND_PARSERS:
+            raise SystemExit(
+                f"campaign: run #{i + 1}: unknown subcommand "
+                f"{subcommand_obj!r} (have {sorted(_SUBCOMMAND_PARSERS)})"
+            )
+
+        argv_for_run = _config_to_argv(merged)
+        sub_args = _parse_for_subcommand(subcommand_obj, argv_for_run)
+        handler = getattr(sub_args, "handler", None)
+        if not callable(handler):
+            raise SystemExit(f"campaign: subcommand {subcommand_obj!r} has no handler")
+
+        logger.info(
+            f"campaign: starting run {i + 1}/{n_runs} "
+            f"({subcommand_obj}; {len(merged)} flags)"
+        )
+        out = cast("Callable[[argparse.Namespace], Path]", handler)(sub_args)
+        output_paths.setdefault(subcommand_obj, []).append(Path(out))
+        logger.info(f"campaign: finished run {i + 1}/{n_runs} → {out}")
+
+    last_path: Path | None = None
+    for paths in output_paths.values():
+        if paths:
+            last_path = paths[-1]
+
+    if get_arg(args, "no_plot", bool):
+        return last_path
+
+    comparison = _render_comparisons(output_paths, plot_cfg)
+    return comparison or last_path
+
+
+# ---------------------------------------------------------------------------
+
+
+def _table(data: dict[str, Any], key: str) -> dict[str, Any]:
+    """Return ``data[key]`` if it's a table, else an empty dict."""
+    val: Any = data.get(key)
+    return cast("dict[str, Any]", val) if isinstance(val, dict) else {}
+
+
+def _merge(defaults: dict[str, Any], run: dict[str, Any]) -> dict[str, Any]:
+    """Shallow-merge ``defaults`` with ``run``; ``run`` wins on conflict.
+
+    The ``tags`` table is deep-merged (defaults' tags + run's tags) so a
+    campaign-level operator tag and a per-run model_short tag both survive.
+    """
+    merged: dict[str, Any] = {**defaults, **run}
+    default_tags = _table(defaults, "tags")
+    run_tags = _table(run, "tags")
+    if default_tags or run_tags:
+        merged["tags"] = {**default_tags, **run_tags}
+    return merged
+
+
+def _parse_for_subcommand(
+    subcommand: str, argv_for_run: list[str]
+) -> argparse.Namespace:
+    """Build a one-off parser with ``subcommand`` registered + parse argv."""
+    parser = argparse.ArgumentParser(prog=f"bench campaign:{subcommand}")
+    subparsers = parser.add_subparsers(dest="subcommand", required=True)
+    _SUBCOMMAND_PARSERS[subcommand](subparsers)
+    return parser.parse_args([subcommand] + argv_for_run)
+
+
+def _render_comparisons(
+    output_paths: dict[str, list[Path]],
+    plot_cfg: dict[str, Any],
+) -> Path | None:
+    """Render one comparison plot per benchmark group with ≥2 outputs."""
+    label_tag = _str_or_none(plot_cfg.get("label_tag"))
+    title = _str_or_none(plot_cfg.get("title"))
+
+    last: Path | None = None
+    for subcommand, paths in output_paths.items():
+        if len(paths) < 2:
+            continue
+        if subcommand != "context-scaling":
+            logger.warning(
+                f"campaign: no comparison renderer registered for {subcommand!r}; "
+                "skipping comparison plot"
+            )
+            continue
+        out = paths[0].with_name(f"campaign_{subcommand}_compare.png")
+        last = render_context_scaling(
+            PlotInputs(
+                results=paths,
+                output=out,
+                label_tag=label_tag,
+                title=title,
+            )
+        )
+        logger.info(f"campaign: wrote comparison plot {last}")
+    return last
+
+
+def _str_or_none(value: Any) -> str | None:  # type: ignore[reportAny]
+    return value if isinstance(value, str) else None
+
+
+__all__ = ["add_subparser", "run"]

bench/cli/context_scaling.py

@@ -0,0 +1,270 @@
+"""Context-scaling benchmark — CLI subcommand.
+
+Wraps :func:`bench.lib.context_scaling.run` with:
+
+  - HF model-metadata resolution
+  - Auto-derived constraints (memory, disk) and context ramp (Δ, K)
+  - eco cluster + instance lifecycle (managed_cluster + managed_instance)
+  - Cold-control isolation (delete sweep instance before controls)
+  - JSON results + ``latest.json`` symlink under ``<output-dir>/context_scaling/``
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+from exo_tools.cluster import EcoSession
+from loguru import logger
+
+from bench.cli._common import (
+    SharedOptions,
+    add_shared_args,
+    get_arg,
+    get_arg_optional,
+)
+from bench.lib import context_scaling
+from bench.lib.cluster import managed_cluster, managed_instance
+from bench.lib.context_scaling import (
+    ContextScalingParams,
+    make_cold_control_factory,
+)
+from bench.lib.model_meta import (
+    ModelMeta,
+    derive_cold_controls,
+    derive_context_ramp,
+    fetch_model_meta,
+)
+from bench.lib.results import ResultsBundle, RunMetadata, find_repo_root
+
+
+def add_subparser(
+    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],  # type: ignore[type-arg]
+) -> None:
+    parser = subparsers.add_parser(
+        "context-scaling",
+        help="Prompt-TPS / decode-TPS vs context-size sweep",
+        description=__doc__,
+    )
+    add_shared_args(parser)
+    g = parser.add_argument_group("context-scaling")
+    g.add_argument(
+        "--num-steps",
+        type=int,
+        default=32,
+        help="Number of equally-spaced PP points in the ramp (K).",
+    )
+    g.add_argument(
+        "--pp-step",
+        type=int,
+        default=None,
+        help="Δ (token step). If unset, derived from the model's max context.",
+    )
+    g.add_argument(
+        "--fraction-of-max",
+        type=float,
+        default=1.0,
+        help="When Δ is auto-derived, use this fraction of the model's "
+        "max_position_embeddings as the ramp's upper bound (0 < f ≤ 1).",
+    )
+    g.add_argument(
+        "--tg",
+        type=int,
+        default=64,
+        help="Tokens to generate per step (decode duration; constant across ramp).",
+    )
+    g.add_argument(
+        "--warmup",
+        type=int,
+        default=2,
+        help="Warmup requests at pp=Δ before the measured ramp. "
+        "First warmup is cache-disabled (kernel JIT only); subsequent "
+        "warmups are cache-enabled (the second is the one that primes "
+        "the cache entry with a hot-kernel rate). Default 2 is the "
+        "sweet spot: warmup=0 leaves JIT cost in step 0; warmup=1 has "
+        "step 0 as a 'none' hit (still hot-kernel cold prefill, just "
+        "classified differently).",
+    )
+    g.add_argument(
+        "--cold-controls",
+        type=str,
+        default=None,
+        help="Cold-control pp values to take after the cached sweep. Either "
+        "'auto' (4 evenly-spaced points across the ramp) or a comma-separated "
+        "list of explicit pp values (e.g. '8192,32768,65536'). "
+        "Default: no cold controls.",
+    )
+    g.add_argument(
+        "--sleep-between-s",
+        type=float,
+        default=1.0,
+        help="Seconds to sleep between consecutive sweep requests.",
+    )
+    parser.set_defaults(handler=run)
+
+
+# ---------------------------------------------------------------------------
+
+
+def run(args: argparse.Namespace) -> Path:
+    """Execute the context-scaling benchmark per the parsed args.
+
+    Returns the path of the JSON results file.
+    """
+    shared = SharedOptions.from_namespace(args)
+    repo_root = find_repo_root()
+
+    # 1. Fetch HF metadata up-front; everything else can be derived from it.
+    logger.info(f"fetching HuggingFace metadata for {shared.model}")
+    meta = fetch_model_meta(shared.model)
+    logger.info(
+        f"  weights: {meta.total_weight_gb:.1f}GB; "
+        f"max context: {meta.max_position_embeddings} tokens; "
+        f"layers: {meta.num_hidden_layers}"
+    )
+
+    # 2. Derive constraints (user values always win; otherwise fall back to
+    # ModelMeta heuristics for the *minimums*).
+    min_memory_gb = (
+        shared.min_memory_gb
+        if shared.min_memory_gb is not None
+        else meta.memory_constraint_gb
+    )
+    min_disk_gb = (
+        shared.min_disk_gb
+        if shared.min_disk_gb is not None
+        else meta.disk_constraint_gb
+    )
+    logger.info(f"  cluster constraint: min memory {min_memory_gb:.1f}GB")
+    logger.info(f"  cluster constraint: min disk {min_disk_gb:.1f}GB")
+    if shared.max_memory_gb is not None:
+        logger.info(f"  cluster constraint: max memory {shared.max_memory_gb:.1f}GB")
+    if shared.max_disk_gb is not None:
+        logger.info(f"  cluster constraint: max disk {shared.max_disk_gb:.1f}GB")
+
+    explicit_pp_step = get_arg_optional(args, "pp_step", int)
+    num_steps = get_arg(args, "num_steps", int)
+    if explicit_pp_step is not None:
+        pp_step = explicit_pp_step
+    else:
+        pp_step, num_steps = derive_context_ramp(
+            meta,
+            num_steps=num_steps,
+            fraction_of_max=get_arg(args, "fraction_of_max", float),
+        )
+        logger.info(
+            f"  derived ramp: Δ={pp_step} × K={num_steps} "
+            f"= {pp_step * num_steps} tokens (max {meta.max_position_embeddings})"
+        )
+
+    cold_controls = _resolve_cold_controls(
+        args, meta, pp_step=pp_step, num_steps=num_steps
+    )
+    if cold_controls:
+        logger.info(f"  cold controls: {list(cold_controls)}")
+
+    # 3. Spin up cluster + instance + run.
+    eco = EcoSession(user_prefix=shared.user_prefix)
+    output_dir = (shared.output_dir / "context_scaling").resolve()
+    metadata = RunMetadata.new(
+        benchmark="context_scaling",
+        repo_root=repo_root,
+        tags={**shared.tags, "host_pool": ",".join(shared.hosts) or "<auto>"},
+    )
+    bundle = ResultsBundle(metadata=metadata)
+
+    with (
+        managed_cluster(
+            eco,
+            hosts=list(shared.hosts) or None,
+            count=shared.nodes,
+            thunderbolt=shared.thunderbolt,
+            chip=shared.chip,
+            min_memory_gb=min_memory_gb,
+            max_memory_gb=shared.max_memory_gb,
+            min_disk_gb=min_disk_gb,
+            max_disk_gb=shared.max_disk_gb,
+        ) as cluster,
+        managed_instance(
+            cluster,
+            eco,
+            shared.model,
+            sharding=shared.sharding,
+            comm=shared.comm,
+            min_nodes=shared.min_nodes,
+            evict_downloads=shared.evict_downloads,
+            cleanup_on_exit=shared.cleanup_instance,
+        ) as session,
+    ):
+        params = ContextScalingParams(
+            pp_step=pp_step,
+            num_steps=num_steps,
+            tg=get_arg(args, "tg", int),
+            warmup=get_arg(args, "warmup", int),
+            cold_controls=cold_controls,
+            sleep_between_s=get_arg(args, "sleep_between_s", float),
+        )
+        factory = (
+            make_cold_control_factory(
+                session, shared.sharding, shared.comm, shared.min_nodes
+            )
+            if cold_controls
+            else None
+        )
+        context_scaling.run(session, params, bundle, cold_control_factory=factory)
+
+    out_path = bundle.write_json(output_dir)
+    _update_latest_symlink(out_path)
+    logger.info(f"wrote results → {out_path}")
+
+    _validate_partial_hits(bundle)
+    return out_path
+
+
+# ---------------------------------------------------------------------------
+
+
+def _resolve_cold_controls(
+    args: argparse.Namespace,
+    meta: ModelMeta,
+    *,
+    pp_step: int,
+    num_steps: int,
+) -> tuple[int, ...]:
+    raw = get_arg_optional(args, "cold_controls", str)
+    if raw is None or not raw.strip():
+        return ()
+    if raw.strip().lower() == "auto":
+        return derive_cold_controls(meta, pp_step=pp_step, num_steps=num_steps, count=4)
+    return tuple(int(s.strip()) for s in raw.split(",") if s.strip())
+
+
+def _update_latest_symlink(out_path: Path) -> None:
+    """Update ``<dir>/latest.json`` to point at the newly-written file."""
+    link = out_path.parent / "latest.json"
+    try:
+        if link.is_symlink() or link.exists():
+            link.unlink()
+        os.symlink(out_path.name, link)
+    except OSError as e:
+        logger.warning(f"could not update latest.json symlink: {e}")
+
+
+def _validate_partial_hits(bundle: ResultsBundle) -> None:
+    """Hard-fail if the cached sweep didn't see ``partial`` on every step ≥ 1.
+
+    Step 0 is allowed to be ``exact`` (warmup primed the cache at pp=Δ); a
+    later ``exact`` means Δ was effectively absorbed into the cache and the
+    cold-rate measurement is meaningless. ``none`` means the cache was
+    discarded mid-sweep and ``T_cum`` is unreliable.
+    """
+    cached = [r for r in bundle.runs if r.get("phase") == "cached_sweep"]
+    bad = [r for r in cached[1:] if r.get("prefix_cache_hit") != "partial"]
+    if bad:
+        bad_summary = [(r["step_index"], r["prefix_cache_hit"]) for r in bad]
+        raise RuntimeError(
+            f"{len(bad)} cached-sweep step(s) reported "
+            f"prefix_cache_hit != 'partial': {bad_summary!r}; "
+            "T_cum is unreliable."
+        )

bench/cli/plot.py

@@ -0,0 +1,125 @@
+"""Plot benchmark results — CLI subcommand.
+
+    uv run python -m bench.cli plot bench/results/context_scaling/latest.json
+    uv run python -m bench.cli plot run_a.json run_b.json --label-tag operator
+    uv run python -m bench.cli plot latest.json --output /tmp/scaling.png
+
+The benchmark type is detected from each JSON's ``metadata.benchmark`` —
+all input files must share the same benchmark.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from loguru import logger
+
+from bench.cli._common import get_arg_optional
+from bench.lib.plotting import PlotInputs, render_context_scaling
+
+
+def add_subparser(
+    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],  # type: ignore[type-arg]
+) -> None:
+    parser = subparsers.add_parser(
+        "plot",
+        help="Render benchmark JSON result(s) as a PNG.",
+        description=__doc__,
+    )
+    parser.add_argument(
+        "paths",
+        nargs="+",
+        type=str,
+        help="One or more bench results JSON files. Multiple files are "
+        "rendered as a comparison plot (one line per file).",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="PNG output path. Default: replace the first JSON's '.json' "
+        "suffix with '.png' (or '.compare.png' when multiple inputs).",
+    )
+    parser.add_argument(
+        "--label-tag",
+        type=str,
+        default=None,
+        help="Use metadata.tags[<KEY>] as the legend label for each run "
+        "(falls back to run_id if unset or missing).",
+    )
+    parser.add_argument(
+        "--title",
+        type=str,
+        default=None,
+        help="Override the auto-generated figure title.",
+    )
+    parser.set_defaults(handler=run)
+
+
+def run(args: argparse.Namespace) -> Path:
+    paths_raw = getattr(args, "paths", None)
+    if not isinstance(paths_raw, list) or not paths_raw:
+        raise SystemExit("plot: at least one JSON path is required")
+    paths = [
+        Path(str(p))  # type: ignore[reportUnknownArgumentType]
+        for p in cast("list[Any]", paths_raw)  # type: ignore[reportAny]
+    ]
+    for p in paths:
+        if not p.is_file():
+            raise SystemExit(f"plot: file not found: {p}")
+
+    benchmarks = {_benchmark_for(p) for p in paths}
+    if len(benchmarks) != 1:
+        raise SystemExit(
+            f"plot: all input JSONs must share the same benchmark, got {benchmarks!r}"
+        )
+    benchmark = next(iter(benchmarks))
+
+    output_arg = get_arg_optional(args, "output", str)
+    output = Path(output_arg) if output_arg is not None else _default_output(paths)
+    inputs = PlotInputs(
+        results=paths,
+        output=output,
+        label_tag=get_arg_optional(args, "label_tag", str),
+        title=get_arg_optional(args, "title", str),
+    )
+
+    if benchmark == "context_scaling":
+        out_path = render_context_scaling(inputs)
+    else:
+        raise SystemExit(f"plot: no renderer registered for benchmark {benchmark!r}")
+
+    logger.info(f"plot: wrote {out_path}")
+    return out_path
+
+
+# ---------------------------------------------------------------------------
+
+
+def _benchmark_for(path: Path) -> str:
+    with path.open() as f:
+        loaded: Any = json.load(f)  # type: ignore[reportAny]
+    if not isinstance(loaded, dict):
+        raise SystemExit(f"plot: {path}: expected top-level JSON object")
+    metadata: Any = loaded.get("metadata", {})  # type: ignore[reportAny]
+    if not isinstance(metadata, dict):
+        raise SystemExit(f"plot: {path}: metadata is not an object")
+    benchmark: Any = metadata.get("benchmark")  # type: ignore[reportAny]
+    if not isinstance(benchmark, str):
+        raise SystemExit(f"plot: {path}: metadata.benchmark missing or not a string")
+    return benchmark
+
+
+def _default_output(paths: list[Path]) -> Path:
+    """Auto-derive a PNG path next to the first JSON.
+
+    Single input → ``<path>.png`` (replaces ``.json``).
+    Multiple inputs → ``<path>.compare.png`` next to the first JSON.
+    """
+    first = paths[0]
+    if len(paths) == 1:
+        return first.with_suffix(".png")
+    return first.with_name(first.stem + ".compare.png")

bench/cli/tests/test_campaign.py

@@ -0,0 +1,109 @@
+"""Unit tests for ``bench.cli.campaign``.
+
+The pure helpers (defaults+run merge, table-lookup, str-or-none) are
+tested here. End-to-end campaign execution requires a real eco cluster
+and is exercised manually via ``bench campaign <toml>``.
+"""
+
+from __future__ import annotations
+
+from bench.cli.campaign import (
+    _merge,  # type: ignore[reportPrivateUsage]
+    _str_or_none,  # type: ignore[reportPrivateUsage]
+    _table,  # type: ignore[reportPrivateUsage]
+)
+
+# ---------------------------------------------------------------------------
+# _table
+# ---------------------------------------------------------------------------
+
+
+class TestTable:
+    def test_present_table(self) -> None:
+        data = {"defaults": {"nodes": 4}}
+        assert _table(data, "defaults") == {"nodes": 4}
+
+    def test_missing_key_returns_empty(self) -> None:
+        assert _table({}, "absent") == {}
+
+    def test_non_table_value_returns_empty(self) -> None:
+        # `nodes = 4` is an int at top level, not a table; treat as empty.
+        assert _table({"nodes": 4}, "nodes") == {}
+
+    def test_list_value_returns_empty(self) -> None:
+        assert _table({"runs": [{"a": 1}]}, "runs") == {}
+
+
+# ---------------------------------------------------------------------------
+# _merge
+# ---------------------------------------------------------------------------
+
+
+class TestMerge:
+    def test_run_wins_on_conflict(self) -> None:
+        defaults = {"nodes": 4, "tg": 64}
+        run = {"nodes": 2}
+        assert _merge(defaults, run) == {"nodes": 2, "tg": 64}
+
+    def test_disjoint_keys(self) -> None:
+        defaults = {"nodes": 4}
+        run = {"model": "test/foo"}
+        assert _merge(defaults, run) == {"nodes": 4, "model": "test/foo"}
+
+    def test_run_only(self) -> None:
+        assert _merge({}, {"a": 1, "b": 2}) == {"a": 1, "b": 2}
+
+    def test_defaults_only(self) -> None:
+        assert _merge({"a": 1}, {}) == {"a": 1}
+
+    def test_tags_deep_merged_defaults_only(self) -> None:
+        defaults = {"tags": {"operator": "ciaranbor"}}
+        run = {"model": "test/foo"}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {"operator": "ciaranbor"}
+
+    def test_tags_deep_merged_run_only(self) -> None:
+        defaults = {"nodes": 4}
+        run = {"tags": {"model_short": "llama-3b"}}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {"model_short": "llama-3b"}
+
+    def test_tags_deep_merged_both(self) -> None:
+        defaults = {"tags": {"operator": "ciaranbor", "campaign": "smoke"}}
+        run = {"tags": {"model_short": "llama-3b"}}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {
+            "operator": "ciaranbor",
+            "campaign": "smoke",
+            "model_short": "llama-3b",
+        }
+
+    def test_run_tags_override_defaults_tags(self) -> None:
+        defaults = {"tags": {"operator": "ciaranbor"}}
+        run = {"tags": {"operator": "alice"}}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {"operator": "alice"}
+
+    def test_no_tags_table_means_no_tags_key(self) -> None:
+        # When neither side has tags, we don't synthesise an empty dict.
+        merged = _merge({"nodes": 4}, {"model": "test/foo"})
+        assert "tags" not in merged
+
+
+# ---------------------------------------------------------------------------
+# _str_or_none
+# ---------------------------------------------------------------------------
+
+
+class TestStrOrNone:
+    def test_str_passes_through(self) -> None:
+        assert _str_or_none("hello") == "hello"
+
+    def test_none_returns_none(self) -> None:
+        assert _str_or_none(None) is None
+
+    def test_int_returns_none(self) -> None:
+        assert _str_or_none(42) is None
+
+    def test_list_returns_none(self) -> None:
+        assert _str_or_none(["a", "b"]) is None

bench/cli/tests/test_common.py

@@ -0,0 +1,275 @@
+"""Unit tests for the argparse boundary helpers in ``bench.cli._common``."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import pytest
+
+from bench.cli._common import (
+    _config_to_argv,  # type: ignore[reportPrivateUsage]
+    _parse_csv,  # type: ignore[reportPrivateUsage]
+    _parse_tags,  # type: ignore[reportPrivateUsage]
+    expand_config_in_argv,
+    get_arg,
+    get_arg_optional,
+)
+
+# ---------------------------------------------------------------------------
+# _parse_csv
+# ---------------------------------------------------------------------------
+
+
+class TestParseCsv:
+    def test_simple_list(self) -> None:
+        assert _parse_csv("a,b,c") == ["a", "b", "c"]
+
+    def test_strips_whitespace(self) -> None:
+        assert _parse_csv("  a , b , c  ") == ["a", "b", "c"]
+
+    def test_skips_empty_entries(self) -> None:
+        assert _parse_csv("a,,b,") == ["a", "b"]
+        assert _parse_csv(",,,") == []
+
+    def test_empty_string_returns_empty(self) -> None:
+        assert _parse_csv("") == []
+
+    def test_single_value(self) -> None:
+        assert _parse_csv("only") == ["only"]
+
+
+# ---------------------------------------------------------------------------
+# _parse_tags
+# ---------------------------------------------------------------------------
+
+
+class TestParseTags:
+    def test_empty_input_returns_empty_dict(self) -> None:
+        assert _parse_tags([]) == {}
+
+    def test_single_tag(self) -> None:
+        assert _parse_tags(["operator=ciaranbor"]) == {"operator": "ciaranbor"}
+
+    def test_multiple_tags(self) -> None:
+        assert _parse_tags(["a=1", "b=2", "c=3"]) == {"a": "1", "b": "2", "c": "3"}
+
+    def test_strips_whitespace_around_key_and_value(self) -> None:
+        assert _parse_tags(["  key  =  value  "]) == {"key": "value"}
+
+    def test_value_can_contain_equals(self) -> None:
+        assert _parse_tags(["url=http://example.com/?a=b"]) == {
+            "url": "http://example.com/?a=b"
+        }
+
+    def test_later_duplicate_key_wins(self) -> None:
+        # Standard dict behaviour; explicit so we notice if it changes.
+        assert _parse_tags(["k=v1", "k=v2"]) == {"k": "v2"}
+
+    def test_missing_equals_raises(self) -> None:
+        with pytest.raises(argparse.ArgumentTypeError, match="key=value"):
+            _ = _parse_tags(["malformed"])
+
+    def test_one_malformed_in_list_raises(self) -> None:
+        with pytest.raises(argparse.ArgumentTypeError):
+            _ = _parse_tags(["good=1", "bad", "alsogood=2"])
+
+
+# ---------------------------------------------------------------------------
+# get_arg / get_arg_optional
+# ---------------------------------------------------------------------------
+
+
+class TestGetArg:
+    def test_str_passes_through(self) -> None:
+        ns = argparse.Namespace(name="hello")
+        assert get_arg(ns, "name", str) == "hello"
+
+    def test_int_passes_through(self) -> None:
+        ns = argparse.Namespace(count=42)
+        assert get_arg(ns, "count", int) == 42
+
+    def test_int_coerces_from_string(self) -> None:
+        ns = argparse.Namespace(count="42")
+        assert get_arg(ns, "count", int) == 42
+
+    def test_float_passes_through(self) -> None:
+        ns = argparse.Namespace(rate=3.14)
+        assert get_arg(ns, "rate", float) == 3.14
+
+    def test_float_coerces_from_int(self) -> None:
+        ns = argparse.Namespace(rate=3)
+        assert get_arg(ns, "rate", float) == 3.0
+
+    def test_float_coerces_from_string(self) -> None:
+        ns = argparse.Namespace(rate="3.14")
+        assert get_arg(ns, "rate", float) == 3.14
+
+    def test_bool_passes_through(self) -> None:
+        ns = argparse.Namespace(flag=True)
+        assert get_arg(ns, "flag", bool) is True
+
+    def test_wrong_type_raises(self) -> None:
+        ns = argparse.Namespace(name=42)
+        with pytest.raises(TypeError, match="expected str"):
+            _ = get_arg(ns, "name", str)
+
+    def test_missing_attribute_raises(self) -> None:
+        ns = argparse.Namespace()
+        with pytest.raises(AttributeError):
+            _ = get_arg(ns, "missing", str)
+
+
+class TestGetArgOptional:
+    def test_missing_returns_none(self) -> None:
+        ns = argparse.Namespace()
+        assert get_arg_optional(ns, "missing", str) is None
+
+    def test_explicit_none_returns_none(self) -> None:
+        ns = argparse.Namespace(value=None)
+        assert get_arg_optional(ns, "value", str) is None
+
+    def test_present_value_returns_typed(self) -> None:
+        ns = argparse.Namespace(value="present")
+        assert get_arg_optional(ns, "value", str) == "present"
+
+    def test_int_coerces_from_string(self) -> None:
+        ns = argparse.Namespace(value="42")
+        assert get_arg_optional(ns, "value", int) == 42
+
+    def test_float_coerces_from_int(self) -> None:
+        ns = argparse.Namespace(value=42)
+        assert get_arg_optional(ns, "value", float) == 42.0
+
+    def test_wrong_type_raises(self) -> None:
+        ns = argparse.Namespace(value=[1, 2, 3])
+        with pytest.raises(TypeError, match="expected str or None"):
+            _ = get_arg_optional(ns, "value", str)
+
+
+# ---------------------------------------------------------------------------
+# _config_to_argv
+# ---------------------------------------------------------------------------
+
+
+class TestConfigToArgv:
+    def test_empty(self) -> None:
+        assert _config_to_argv({}) == []
+
+    def test_string_value(self) -> None:
+        assert _config_to_argv({"model": "mlx/foo"}) == ["--model", "mlx/foo"]
+
+    def test_int_and_float_values(self) -> None:
+        out = _config_to_argv({"num_steps": 32, "fraction_of_max": 0.5})
+        assert out == ["--num-steps", "32", "--fraction-of-max", "0.5"]
+
+    def test_underscore_keys_become_hyphenated_flags(self) -> None:
+        out = _config_to_argv({"min_memory_gb": 21.0})
+        assert out == ["--min-memory-gb", "21.0"]
+
+    def test_bool_true_emits_flag(self) -> None:
+        assert _config_to_argv({"auto_constrain": True}) == ["--auto-constrain"]
+
+    def test_bool_false_emits_no_form(self) -> None:
+        assert _config_to_argv({"auto_constrain": False}) == ["--no-auto-constrain"]
+
+    def test_none_value_skipped(self) -> None:
+        assert _config_to_argv({"chip": None, "model": "foo"}) == [
+            "--model",
+            "foo",
+        ]
+
+    def test_list_joined_as_csv(self) -> None:
+        out = _config_to_argv({"hosts": ["s4", "s9"], "cold_controls": [1024, 2048]})
+        assert out == [
+            "--hosts",
+            "s4,s9",
+            "--cold-controls",
+            "1024,2048",
+        ]
+
+    def test_tags_table_expands_to_repeated_tag_args(self) -> None:
+        out = _config_to_argv({"tags": {"operator": "ciaranbor", "run": "full"}})
+        # Order within a TOML table is preserved by tomllib
+        assert out == [
+            "--tag",
+            "operator=ciaranbor",
+            "--tag",
+            "run=full",
+        ]
+
+
+# ---------------------------------------------------------------------------
+# expand_config_in_argv
+# ---------------------------------------------------------------------------
+
+
+class TestExpandConfigInArgv:
+    def test_no_config_flag_passthrough(self) -> None:
+        argv = ["context-scaling", "--model", "foo"]
+        assert expand_config_in_argv(argv) == argv
+
+    def test_config_at_end(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text('model = "from_config"\nnum_steps = 16\n')
+        argv = ["context-scaling", "--config", str(cfg)]
+        # Config flags are inserted right after the subcommand
+        assert expand_config_in_argv(argv) == [
+            "context-scaling",
+            "--model",
+            "from_config",
+            "--num-steps",
+            "16",
+        ]
+
+    def test_explicit_cli_overrides_config(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text('model = "from_config"\nnum_steps = 16\n')
+        # User overrides --num-steps explicitly. Argparse takes the last
+        # occurrence for non-append actions, so the user's 32 wins.
+        argv = ["context-scaling", "--config", str(cfg), "--num-steps", "32"]
+        out = expand_config_in_argv(argv)
+        assert out == [
+            "context-scaling",
+            "--model",
+            "from_config",
+            "--num-steps",
+            "16",
+            "--num-steps",
+            "32",
+        ]
+
+    def test_missing_path_arg_raises(self) -> None:
+        with pytest.raises(ValueError, match="--config requires a path"):
+            _ = expand_config_in_argv(["context-scaling", "--config"])
+
+    def test_nonexistent_file_raises(self, tmp_path: Path) -> None:
+        with pytest.raises(FileNotFoundError, match="Config file not found"):
+            _ = expand_config_in_argv(
+                ["context-scaling", "--config", str(tmp_path / "missing.toml")]
+            )
+
+    def test_bool_false_in_config(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text("auto_constrain = false\n")
+        argv = ["context-scaling", "--config", str(cfg)]
+        assert expand_config_in_argv(argv) == [
+            "context-scaling",
+            "--no-auto-constrain",
+        ]
+
+    def test_tags_table(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text(
+            'model = "foo"\n[tags]\noperator = "ciaranbor"\nrun = "full"\n'
+        )
+        argv = ["context-scaling", "--config", str(cfg)]
+        assert expand_config_in_argv(argv) == [
+            "context-scaling",
+            "--model",
+            "foo",
+            "--tag",
+            "operator=ciaranbor",
+            "--tag",
+            "run=full",
+        ]

bench/configs/context_scaling.example.toml

@@ -0,0 +1,70 @@
+# Example context-scaling run configuration.
+#
+# Use it like this:
+#
+#   uv run python -m bench.cli context-scaling --config bench/configs/context_scaling.example.toml
+#
+# CLI flags placed after `--config` override individual values.
+#
+# All shared and subcommand-specific flags can appear here. Keys mirror the
+# CLI flag names with hyphens replaced by underscores. Boolean keys map to
+# `--key` / `--no-key`; lists are joined as CSV; the `[tags]` table maps to
+# repeated `--tag key=value` flags.
+#
+# NOTE: TOML scoping — once a `[table]` header is opened, all subsequent
+# top-level-looking assignments belong to that table until the next header.
+# Keep tables (like `[tags]`) at the END of the file.
+
+# ---- Model + placement ----
+model = "mlx-community/Qwen3-30B-A3B-4bit"
+# sharding = "Tensor"      # default: "Tensor"; pairs with --comm MlxJaccl + --thunderbolt a2a
+# comm = "MlxJaccl"        # default: "MlxJaccl" (RDMA over Thunderbolt)
+# min_nodes = 1
+
+# ---- Cluster ----
+# Either pin to specific hosts...
+# hosts = ["s4"]
+# ...or let eco pick hosts that satisfy the constraints below.
+# nodes = 1
+# chip = "M3 Ultra"        # eco chip name (case-insensitive substring); comment to allow any
+# thunderbolt = "a2a"      # default: "a2a" (clique, for Tensor+JACCL)
+# "ring" (cycle; for Pipeline+JACCL)
+# "none" (exclude TB; pair with sharding=Pipeline + comm=MlxRing for non-TB hosts)
+
+# Memory + disk minimums are auto-derived from the HF model size
+# (×1.30 + 1 GiB for memory, ×1.10 + 1 GiB for disk). Set any of these
+# explicitly to override the auto-derived value.
+# min_memory_gb = 96.0
+# max_memory_gb = 256.0    # leave bigger machines free for other workloads
+# min_disk_gb = 24.0
+# max_disk_gb = 4000.0
+
+# ---- Runtime ----
+# evict_downloads is true by default — frees disk smallest-first to fit
+# the bench model. Set to false to keep existing downloads.
+# evict_downloads = false
+
+# cleanup_instance is true by default — deletes the placed instance on exit.
+# Set to false to leave it running for debugging.
+# cleanup_instance = false
+
+# ---- Output ----
+output_dir = "bench/results"
+
+# ---- Context-scaling sweep ----
+num_steps = 32 # K — number of equally-spaced ramp points
+# pp_step = 1024              # Δ — explicit override; otherwise auto-derived
+# fraction_of_max = 1.0       # use this fraction of max_position_embeddings
+tg = 64        # tokens generated per step
+# warmup = 2   # default: 2 (1 cache-disabled JIT warmup + 1 cache-priming warmup)
+# cold_controls = "auto"      # 4 evenly-spaced controls across the ramp, or:
+# cold_controls = "8192,16384,32768,40960"   # explicit pp values
+sleep_between_s = 1.0
+
+# ---- Tags ----
+# Survive into metadata.tags in the output JSON; useful for filtering or
+# grouping runs across SHAs / hosts / configs. `$USER` is NOT expanded
+# (TOML is literal); pass `--tag operator=$USER` on the CLI for shell expansion.
+# Must be the LAST table in the file (see TOML scoping note above).
+[tags]
+run = "full"

bench/configs/llama-family-smoke.toml

@@ -0,0 +1,34 @@
+# 4-node smoke campaign: two small/medium Llama models, abbreviated ramps,
+# auto-everything else (TB a2a + tensor + JACCL + auto-derived constraints).
+#
+# Run with:
+#   uv run python -m bench.cli campaign bench/configs/llama-family-smoke.toml
+#
+# Each [[runs]] gets its own cluster (deploy + bench + teardown). After
+# both runs finish, a side-by-side comparison plot is written next to the
+# JSONs.
+
+[defaults]
+nodes = 4
+num_steps = 8
+fraction_of_max = 0.5
+
+[defaults.tags]
+campaign = "llama-family-smoke"
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.2-3b-4bit"
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.1-8b-4bit"
+
+# Final comparison plot (one PNG per benchmark group with ≥2 runs).
+[plot]
+label_tag = "model_short"
+title = "Llama 3 family — 4-node tensor + JACCL smoke"

bench/eval_configs/models.toml

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

bench/eval_tool_calls.py

@@ -3,20 +3,22 @@ from __future__ import annotations
 
 import argparse
 import contextlib
+import io
 import json
 import os
 import sys
 import time
 import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
 
 import httpx
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -208,7 +210,7 @@ def _openai_build_request(
         "model": model,
         "messages": messages,
         "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     return "/v1/chat/completions", body
@@ -275,7 +277,7 @@ def _openai_build_followup(
         "model": model,
         "messages": followup_messages,
         "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     return "/v1/chat/completions", body
@@ -378,7 +380,7 @@ def _claude_build_request(
         "model": model,
         "messages": claude_messages,
         "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     if system_content is not None:
@@ -488,7 +490,7 @@ def _claude_build_followup(
         "model": model,
         "messages": claude_messages,
         "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
         "temperature": 0.0,
     }
     if system_content is not None:
@@ -912,6 +914,12 @@ Examples:
         default=1,
         help="Repeat each scenario N times (default: 1)",
     )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help="Run up to N scenarios in parallel against the same instance (default: 1)",
+    )
     parser.add_argument(
         "--scenarios",
         nargs="*",
@@ -934,6 +942,13 @@ Examples:
     )
     args = parser.parse_args()
 
+    if args.concurrency < 1:
+        print(
+            f"--concurrency must be >= 1 (got {args.concurrency})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
     all_scenarios = load_scenarios(SCENARIOS_PATH)
     if args.scenarios:
         scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1006,44 +1021,75 @@ Examples:
         sys.exit(1)
 
     time.sleep(1)
+    cluster_snapshot = capture_cluster_snapshot(exo)
     all_results: list[ScenarioResult] = []
 
+    tasks: list[tuple[int, Scenario, ApiName]] = [
+        (run_idx, scenario, api_name)
+        for run_idx in range(args.repeat)
+        for scenario in scenarios
+        for api_name in api_names
+    ]
+
+    def _run_one(
+        http_client: httpx.Client,
+        task: tuple[int, Scenario, ApiName],
+    ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+        run_idx, scenario, api_name = task
+        buf = io.StringIO()
+        run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+        print(
+            f"\n  {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+            file=buf,
+        )
+        scenario_results = run_scenario(
+            http_client,
+            args.host,
+            args.port,
+            full_model_id,
+            scenario,
+            api_name,
+            args.timeout,
+            args.verbose,
+        )
+        for r in scenario_results:
+            status = "PASS" if r.passed else "FAIL"
+            print(
+                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
+                file=buf,
+            )
+            for check_name, check_ok in r.checks.items():
+                mark = "+" if check_ok else "-"
+                print(f"      {mark} {check_name}", file=buf)
+            if r.error:
+                print(f"      ! {r.error}", file=buf)
+        return task, scenario_results, buf.getvalue()
+
     try:
         with httpx.Client() as http_client:
-            for run_idx in range(args.repeat):
-                if args.repeat > 1:
-                    print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
-                for scenario in scenarios:
-                    for api_name in api_names:
-                        print(
-                            f"\n  [{api_name:>9}] {scenario.name}: {scenario.description}",
-                            file=log,
-                        )
-
-                        scenario_results = run_scenario(
-                            http_client,
-                            args.host,
-                            args.port,
-                            full_model_id,
-                            scenario,
-                            api_name,
-                            args.timeout,
-                            args.verbose,
-                        )
+            if args.concurrency == 1:
+                current_run = -1
+                for task in tasks:
+                    run_idx = task[0]
+                    if args.repeat > 1 and run_idx != current_run:
+                        print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+                        current_run = run_idx
+                    _, scenario_results, buffered = _run_one(http_client, task)
+                    all_results.extend(scenario_results)
+                    log.write(buffered)
+                    log.flush()
+            else:
+                print(
+                    f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+                    file=log,
+                )
+                with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+                    futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+                    for fut in as_completed(futures):
+                        _, scenario_results, buffered = fut.result()
                         all_results.extend(scenario_results)
-
-                        for r in scenario_results:
-                            status = "PASS" if r.passed else "FAIL"
-                            print(
-                                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
-                                file=log,
-                            )
-                            for check_name, check_ok in r.checks.items():
-                                mark = "+" if check_ok else "-"
-                                print(f"      {mark} {check_name}", file=log)
-                            if r.error:
-                                print(f"      ! {r.error}", file=log)
+                        log.write(buffered)
+                        log.flush()
     finally:
         try:
             exo.request_json("DELETE", f"/instance/{instance_id}")
@@ -1084,16 +1130,19 @@ Examples:
                 print(f"  - {r.name} [{r.api}/{r.phase}]: {r.error}", file=log)
 
     json_results = [result_to_dict(r) for r in all_results]
+    output: dict[str, Any] = {"results": json_results}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
 
     if args.stdout:
-        print(json.dumps(json_results, indent=2))
+        print(json.dumps(output, indent=2))
     else:
         json_path = args.json_out
         parent = os.path.dirname(json_path)
         if parent:
             os.makedirs(parent, exist_ok=True)
         with open(json_path, "w") as f:
-            json.dump(json_results, f, indent=2)
+            json.dump(output, f, indent=2)
             f.write("\n")
         print(f"\nJSON results written to {json_path}", file=log)
 

bench/exo_bench.py

@@ -22,18 +22,19 @@ import contextlib
 import itertools
 import json
 import sys
+import threading
 import time
-from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from pathlib import Path
 from statistics import mean
 from typing import Any
 
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
+    capture_cluster_snapshot,
+    find_existing_instance,
     instance_id_from_instance,
+    node_ids_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
@@ -42,85 +43,44 @@ from harness import (
     wait_for_instance_ready,
 )
 from loguru import logger
-from transformers import AutoTokenizer
 
-# Monkey-patch for transformers 5.x compatibility
-# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
-# which was moved in transformers 5.0.0rc2
-try:
-    import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
-    from transformers.convert_slow_tokenizer import bytes_to_unicode
+# PromptSizer / run_one_completion / load_tokenizer_for_bench are the
+# canonical, fully-typed implementations under bench/lib/. They are
+# re-exported here for backwards compatibility with prefill_decode_bench.py
+# and any other consumers of `from exo_bench import …`.
+from bench.lib.completion import run_one_completion as _lib_run_one_completion
+from bench.lib.prompt import (
+    PromptSizer as _LibPromptSizer,
+)
+from bench.lib.prompt import (
+    load_tokenizer_for_bench as _lib_load_tokenizer_for_bench,
+)
 
-    if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
-        gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[attr-defined]
-except ImportError:
-    pass  # transformers < 5.0 or bytes_to_unicode not available
+PromptSizer = _LibPromptSizer
+load_tokenizer_for_bench = _lib_load_tokenizer_for_bench
 
 
-def load_tokenizer_for_bench(model_id: str) -> Any:
-    """
-    Load tokenizer for benchmarking, with special handling for Kimi models.
-
-    Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer.
-    This function replicates the logic from utils_mlx.py for bench compatibility.
-    """
-    model_id_lower = model_id.lower()
-
-    if "kimi-k2" in model_id_lower:
-        import importlib.util
-        import types
-
-        from huggingface_hub import snapshot_download
-
-        # Download/get the model path
-        model_path = Path(
-            snapshot_download(
-                model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
-            )
-        )
-
-        sys.path.insert(0, str(model_path))
-
-        # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
-        tool_decl_path = model_path / "tool_declaration_ts.py"
-        if tool_decl_path.exists():
-            spec = importlib.util.spec_from_file_location(
-                "tool_declaration_ts", tool_decl_path
-            )
-            if spec and spec.loader:
-                tool_decl_module = importlib.util.module_from_spec(spec)
-                sys.modules["tool_declaration_ts"] = tool_decl_module
-                spec.loader.exec_module(tool_decl_module)
-
-        # Load tokenization_kimi with patched source (convert relative to absolute import)
-        tok_path = model_path / "tokenization_kimi.py"
-        source = tok_path.read_text()
-        source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
-        spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
-        if spec:
-            tok_module = types.ModuleType("tokenization_kimi")
-            tok_module.__file__ = str(tok_path)
-            sys.modules["tokenization_kimi"] = tok_module
-            exec(compile(source, tok_path, "exec"), tok_module.__dict__)  # noqa: S102
-            TikTokenTokenizer = tok_module.TikTokenTokenizer  # noqa: N806
-        else:
-            from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
-
-        hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)
-
-        # Patch encode to use internal tiktoken model directly
-        # transformers 5.x has a bug in the encode->pad path for slow tokenizers
-        def _patched_encode(text: str, **kwargs: object) -> list[int]:
-            # Pass allowed_special="all" to handle special tokens like <|im_user|>
-            return list(hf_tokenizer.model.encode(text, allowed_special="all"))
-
-        hf_tokenizer.encode = _patched_encode
-
-        return hf_tokenizer
-
-    # Default: use AutoTokenizer
-    return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+def run_one_completion(
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool = False,
+    stream: bool = False,
+) -> tuple[dict[str, Any], int]:
+    """Backwards-compatible shim returning a plain ``dict`` row."""
+    row, pp_tokens = _lib_run_one_completion(
+        client,
+        model_id,
+        pp_hint,
+        tg,
+        prompt_sizer,
+        use_prefix_cache=use_prefix_cache,
+        stream=stream,
+    )
+    return dict(row), pp_tokens
 
 
 def format_peak_memory(b: float) -> str:
@@ -131,6 +91,91 @@ def format_peak_memory(b: float) -> str:
     raise ValueError("You're using petabytes of memory. Something went wrong...")
 
 
+_SAMPLER_METRICS = ("gpuUsage", "temp", "sysPower", "pcpuUsage", "ecpuUsage")
+
+
+class SystemMetricsSampler:
+    def __init__(self, client: ExoClient, node_ids: list[str], interval_s: float = 1.0):
+        self._client = client
+        self._node_ids = node_ids
+        self._interval_s = interval_s
+        self._samples: dict[str, list[tuple[float, dict[str, float]]]] = {
+            nid: [] for nid in node_ids
+        }
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._stop.clear()
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=5)
+
+    def _poll_loop(self) -> None:
+        while not self._stop.is_set():
+            t = time.monotonic()
+            for nid in self._node_ids:
+                try:
+                    data = self._client.get_node_system(nid)
+                    if data:
+                        self._samples[nid].append(
+                            (t, {k: data.get(k, 0.0) for k in _SAMPLER_METRICS})
+                        )
+                except Exception:
+                    pass
+            self._stop.wait(self._interval_s)
+
+    def energy_between(self, t0: float, t1: float) -> float:
+        total_joules = 0.0
+        for _nid, samples in self._samples.items():
+            window = [(t, s["sysPower"]) for t, s in samples if t0 <= t <= t1]
+            if len(window) >= 2:
+                for i in range(1, len(window)):
+                    dt = window[i][0] - window[i - 1][0]
+                    avg_power = (window[i][1] + window[i - 1][1]) / 2
+                    total_joules += avg_power * dt
+            elif len(window) == 1:
+                total_joules += window[0][1] * (t1 - t0)
+        return total_joules
+
+    def summarize(self) -> dict[str, dict[str, dict[str, float]]]:
+        result: dict[str, dict[str, dict[str, float]]] = {}
+        for nid, samples in self._samples.items():
+            if not samples:
+                continue
+            metrics: dict[str, dict[str, float]] = {}
+            for key in _SAMPLER_METRICS:
+                values = [s[key] for t, s in samples]
+                metrics[key] = {
+                    "min": round(min(values), 2),
+                    "max": round(max(values), 2),
+                    "mean": round(mean(values), 2),
+                    "samples": len(values),
+                }
+            result[nid] = metrics
+        return result
+
+    def print_summary(self, placement_label: str) -> None:
+        summary = self.summarize()
+        if not summary:
+            return
+        logger.info(f"--- System Metrics ({placement_label}) ---")
+        for nid, metrics in summary.items():
+            gpu = metrics.get("gpuUsage", {})
+            temp = metrics.get("temp", {})
+            power = metrics.get("sysPower", {})
+            logger.info(
+                f"  {nid}:  "
+                f"GPU {gpu.get('mean', 0) * 100:.0f}% avg ({gpu.get('min', 0) * 100:.0f}–{gpu.get('max', 0) * 100:.0f}%)  |  "
+                f"{temp.get('mean', 0):.1f}°C avg  |  "
+                f"{power.get('mean', 0):.1f}W avg"
+            )
+
+
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -141,97 +186,6 @@ def parse_int_list(values: list[str]) -> list[int]:
     return items
 
 
-def run_one_completion(
-    client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer
-) -> tuple[dict[str, Any], int]:
-    content, pp_tokens = prompt_sizer.build(pp_hint)
-    payload: dict[str, Any] = {
-        "model": model_id,
-        "messages": [{"role": "user", "content": content}],
-        "stream": False,
-        "max_tokens": tg,
-    }
-
-    t0 = time.perf_counter()
-    out = client.post_bench_chat_completions(payload)
-    elapsed = time.perf_counter() - t0
-
-    stats = out.get("generation_stats")
-
-    # Extract preview, handling None content (common for thinking models)
-    choices = out.get("choices") or [{}]
-    message = choices[0].get("message", {}) if choices else {}
-    content = message.get("content") or ""
-    preview = content[:200] if content else ""
-
-    return {
-        "elapsed_s": elapsed,
-        "output_text_preview": preview,
-        "stats": stats,
-    }, pp_tokens
-
-
-class PromptSizer:
-    def __init__(self, tokenizer: Any, atom: str = "a "):
-        self.tokenizer = tokenizer
-        self.atom = atom
-        self.count_fn = PromptSizer._make_counter(tokenizer)
-        self.base_tokens = self.count_fn("")
-
-    @staticmethod
-    def _make_counter(tokenizer: Any) -> Callable[[str], int]:
-        def count_fn(user_content: str) -> int:
-            messages = [{"role": "user", "content": user_content}]
-            ids = tokenizer.apply_chat_template(
-                messages, tokenize=True, add_generation_prompt=True
-            )
-            # Fix for transformers 5.x
-            if hasattr(ids, "input_ids"):
-                ids = ids.input_ids
-            return int(len(ids))
-
-        return count_fn
-
-    def build(self, target_prompt_tokens: int) -> tuple[str, int]:
-        target = int(target_prompt_tokens)
-        if target < self.base_tokens:
-            raise RuntimeError(
-                f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
-            )
-
-        # Estimate tokens per atom using a sample
-        sample_count = 100
-        sample_content = self.atom * sample_count
-        sample_tokens = self.count_fn(sample_content) - self.base_tokens
-        tokens_per_atom = sample_tokens / sample_count
-
-        # Estimate starting point
-        needed_tokens = target - self.base_tokens
-        estimated_atoms = int(needed_tokens / tokens_per_atom)
-
-        # Binary search to find exact atom count
-        low, high = 0, estimated_atoms * 2 + 100
-        while low < high:
-            mid = (low + high) // 2
-            tok = self.count_fn(self.atom * mid)
-            if tok < target:
-                low = mid + 1
-            else:
-                high = mid
-
-        content = self.atom * low
-        tok = self.count_fn(content)
-        logger.info(f"{tok=}")
-
-        if tok != target:
-            raise RuntimeError(
-                f"Overshot: got {tok} tokens (target {target}). "
-                f"Pick a different atom (try ' a' or '\\n' or '0 ')."
-            )
-
-        return content, tok
-
-
 def main() -> int:
     ap = argparse.ArgumentParser(
         prog="exo-bench",
@@ -279,6 +233,27 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
+    ap.add_argument(
+        "--stream",
+        action="store_true",
+        help="Use /bench/chat/completions with streaming SSE response (bench=True still applies: no EOS detection, no KV cache).",
+    )
+    ap.add_argument(
+        "--no-system-metrics",
+        action="store_true",
+        help="Disable GPU utilization, temperature, and power collection during inference.",
+    )
+    ap.add_argument(
+        "--metrics-interval",
+        type=float,
+        default=1.0,
+        help="System metrics polling interval in seconds (default: 1.0).",
+    )
+    ap.add_argument(
+        "--use-prefix-cache",
+        action="store_true",
+        help="Enable KV prefix cache during bench (default: disabled for cold-cache measurements).",
+    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -294,6 +269,15 @@ def main() -> int:
         logger.error("--concurrency values must be >= 1")
         return 2
 
+    if args.use_prefix_cache:
+        logger.warning(
+            "--use-prefix-cache: prompt TPS will be approximate. See METHODOLOGY.md for details."
+        )
+        if pp_list != sorted(pp_list):
+            logger.warning(
+                "--pp values are not in ascending order: prompt TPS will be less accurate. Use ascending --pp for best results."
+            )
+
     # Log pairing mode
     use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
     if use_combinations:
@@ -319,82 +303,146 @@ def main() -> int:
         logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
         raise
 
-    selected = settle_and_fetch_placements(
-        client, full_model_id, args, settle_timeout=args.settle_timeout
-    )
+    # Optionally reuse a running instance for this model
+    reused_instance_id: str | None = None
+    if args.reuse_instance:
+        existing = find_existing_instance(client, full_model_id)
+        if existing:
+            reused_instance_id = existing
+            logger.info(f"Reusing existing instance {reused_instance_id}")
+        else:
+            logger.warning(
+                "--reuse-instance: no existing instance found, creating a new one"
+            )
 
-    if not selected:
-        logger.error("No valid placements matched your filters.")
-        return 1
-
-    selected.sort(
-        key=lambda p: (
-            str(p.get("instance_meta", "")),
-            str(p.get("sharding", "")),
-            -nodes_used_in_instance(p["instance"]),
-        ),
-        reverse=True,
-    )
-
-    logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
-    logger.info(f"placements: {len(selected)}")
-    for p in selected:
-        logger.info(
-            f"  - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
+    if reused_instance_id is not None:
+        # Use the existing instance directly — skip placement iteration
+        selected = []
+        download_duration_s = None
+    else:
+        selected = settle_and_fetch_placements(
+            client, full_model_id, args, settle_timeout=args.settle_timeout
         )
 
-    if args.dry_run:
-        return 0
+        if not selected:
+            logger.error("No valid placements matched your filters.")
+            return 1
 
-    settle_deadline = (
-        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
-    )
+        selected.sort(
+            key=lambda p: (
+                str(p.get("instance_meta", "")),
+                str(p.get("sharding", "")),
+                nodes_used_in_instance(p["instance"]),
+            ),
+            reverse=True,
+        )
 
-    logger.info("Planning phase: checking downloads...")
-    download_duration_s = run_planning_phase(
-        client,
-        full_model_id,
-        selected[0],
-        args.danger_delete_downloads,
-        args.timeout,
-        settle_deadline,
-    )
-    if download_duration_s is not None:
-        logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
-    else:
-        logger.info("Download: model already cached")
+        logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
+        logger.info(f"placements: {len(selected)}")
+        for p in selected:
+            logger.info(
+                f"  - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
+            )
 
+        if args.dry_run:
+            return 0
+
+        settle_deadline = (
+            time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+        )
+
+        logger.info("Planning phase: checking downloads...")
+        download_duration_s = run_planning_phase(
+            client,
+            full_model_id,
+            selected[0],
+            args.danger_delete_downloads,
+            args.timeout,
+            settle_deadline,
+        )
+        if download_duration_s is not None:
+            logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
+        else:
+            logger.info("Download: model already cached")
+
+    cluster_snapshot = capture_cluster_snapshot(client)
     all_rows: list[dict[str, Any]] = []
+    all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
+
+    # If reusing an existing instance, run a single benchmark pass against it
+    if reused_instance_id is not None:
+        selected = [None]
 
     for preview in selected:
-        instance = preview["instance"]
-        instance_id = instance_id_from_instance(instance)
+        created_instance = False
+        if preview is not None:
+            instance = preview["instance"]
+            instance_id = instance_id_from_instance(instance)
 
-        sharding = str(preview["sharding"])
-        instance_meta = str(preview["instance_meta"])
-        n_nodes = nodes_used_in_instance(instance)
+            sharding = str(preview["sharding"])
+            instance_meta = str(preview["instance_meta"])
+            n_nodes = nodes_used_in_instance(instance)
 
-        logger.info("=" * 80)
-        logger.info(
-            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
-        )
+            logger.info("=" * 80)
+            logger.info(
+                f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
+            )
 
-        client.request_json("POST", "/instance", body={"instance": instance})
-        try:
-            wait_for_instance_ready(client, instance_id)
-        except (RuntimeError, TimeoutError) as e:
-            logger.error(f"Failed to initialize placement: {e}")
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{instance_id}")
-            continue
+            # Delete any existing instances to free resources before placing
+            try:
+                state = client.request_json("GET", "/state")
+                for old_id in list(state.get("instances", {}).keys()):
+                    logger.info(f"Deleting stale instance {old_id}")
+                    with contextlib.suppress(ExoHttpError):
+                        client.request_json("DELETE", f"/instance/{old_id}")
+                if state.get("instances"):
+                    time.sleep(2)
+            except Exception as e:
+                logger.warning(f"Failed to clean up stale instances: {e}")
 
-        time.sleep(1)
+            client.request_json("POST", "/instance", body={"instance": instance})
+            try:
+                wait_for_instance_ready(client, instance_id)
+            except (RuntimeError, TimeoutError) as e:
+                logger.error(f"Failed to initialize placement: {e}")
+                with contextlib.suppress(ExoHttpError):
+                    client.request_json("DELETE", f"/instance/{instance_id}")
+                continue
+
+            time.sleep(1)
+            created_instance = True
+        else:
+            instance_id = reused_instance_id
+            sharding = "reused"
+            instance_meta = "reused"
+            n_nodes = 0
+            logger.info("=" * 80)
+            logger.info(f"Using existing instance {instance_id}")
+
+        sampler: SystemMetricsSampler | None = None
+        if not args.no_system_metrics and preview is not None:
+            nids = node_ids_from_instance(instance)
+            sampler = SystemMetricsSampler(
+                ExoClient(args.host, args.port, timeout_s=30),
+                nids,
+                interval_s=args.metrics_interval,
+            )
+            sampler.start()
+
+        def _do_one(c: ExoClient, pp: int, tg: int) -> tuple[dict[str, Any], int]:
+            return run_one_completion(
+                c,
+                full_model_id,
+                pp,
+                tg,
+                prompt_sizer,
+                use_prefix_cache=args.use_prefix_cache,
+                stream=args.stream,
+            )
 
         try:
             for i in range(args.warmup):
-                run_one_completion(
-                    client, full_model_id, pp_list[0], tg_list[0], prompt_sizer
-                )
+                _do_one(client, pp_list[0], tg_list[0])
                 logger.debug(f"  warmup {i + 1}/{args.warmup} done")
 
             # If pp and tg lists have same length, run in tandem (zip)
@@ -408,15 +456,16 @@ def main() -> int:
                 for concurrency in concurrency_list:
                     logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---")
                     runs: list[dict[str, Any]] = []
+                    inference_windows: list[tuple[float, float]] = []
                     for r in range(args.repeat):
                         time.sleep(3)
 
                         if concurrency <= 1:
                             # Sequential: single request
                             try:
-                                row, actual_pp_tokens = run_one_completion(
-                                    client, full_model_id, pp, tg, prompt_sizer
-                                )
+                                inf_t0 = time.monotonic()
+                                row, actual_pp_tokens = _do_one(client, pp, tg)
+                                inference_windows.append((inf_t0, time.monotonic()))
                             except Exception as e:
                                 logger.error(e)
                                 continue
@@ -443,20 +492,54 @@ def main() -> int:
                             all_rows.append(row)
                         else:
                             # Concurrent: fire N requests in parallel
-                            # Each thread gets its own ExoClient (separate HTTP connection)
+                            # Pre-build prompt once, barrier ensures simultaneous dispatch
+                            content, actual_pp = prompt_sizer.build(pp)
+                            pre_built_payload: dict[str, Any] = {
+                                "model": full_model_id,
+                                "messages": [{"role": "user", "content": content}],
+                                "stream": False,
+                                "max_tokens": tg,
+                                "logprobs": False,
+                                "use_prefix_cache": args.use_prefix_cache,
+                            }
+                            barrier = threading.Barrier(concurrency)
+                            batch_start = threading.Event()
+                            batch_t0: float = 0.0
                             batch_results: list[tuple[dict[str, Any], int]] = []
                             batch_errors = 0
 
                             def _run_concurrent(
-                                idx: int, *, _pp: int = pp, _tg: int = tg
+                                idx: int,
+                                _barrier: threading.Barrier = barrier,
+                                _batch_start: threading.Event = batch_start,
+                                _payload: dict[str, Any] = pre_built_payload,
+                                _actual_pp: int = actual_pp,
                             ) -> tuple[dict[str, Any], int]:
+                                nonlocal batch_t0
                                 c = ExoClient(
                                     args.host, args.port, timeout_s=args.timeout
                                 )
-                                return run_one_completion(
-                                    c, full_model_id, _pp, _tg, prompt_sizer
+                                if _barrier.wait() == 0:
+                                    batch_t0 = time.perf_counter()
+                                    _batch_start.set()
+                                else:
+                                    _batch_start.wait()
+                                t0 = batch_t0
+                                out = c.post_bench_chat_completions(_payload)
+                                elapsed = time.perf_counter() - t0
+                                stats = out.get("generation_stats")
+                                choices = out.get("choices") or [{}]
+                                message = (
+                                    choices[0].get("message", {}) if choices else {}
                                 )
+                                text = message.get("content") or ""
+                                return {
+                                    "elapsed_s": elapsed,
+                                    "output_text_preview": text[:200],
+                                    "stats": stats,
+                                }, _actual_pp
 
+                            inf_t0 = time.monotonic()
                             with ThreadPoolExecutor(max_workers=concurrency) as pool:
                                 futures = {
                                     pool.submit(_run_concurrent, i): i
@@ -468,6 +551,12 @@ def main() -> int:
                                     except Exception as e:
                                         logger.error(f"Concurrent request failed: {e}")
                                         batch_errors += 1
+                            batch_wall_s = (
+                                max(x["elapsed_s"] for x, _ in batch_results)
+                                if batch_results
+                                else time.perf_counter() - batch_t0
+                            )
+                            inference_windows.append((inf_t0, time.monotonic()))
 
                             for idx, (row, actual_pp_tokens) in enumerate(
                                 batch_results
@@ -501,51 +590,81 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                agg_gen_tps = (
-                                    mean(valid_gen_tps) if valid_gen_tps else 0.0
+                                per_req_tps = (
+                                    max(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
-                                gen_tps = agg_gen_tps / concurrency
+                                agg_gen_tps = per_req_tps * concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
-                                    f"gen_tps={gen_tps:.2f}  "
+                                    f"per_req_tps={per_req_tps:.2f}  "
+                                    f"wall_s={batch_wall_s:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(
-                            x["stats"]["generation_tps"] / x["concurrency"]
+                        valid_gen = [
+                            x["stats"]["generation_tps"]
                             for x in runs
-                        )
+                            if x["stats"]["generation_tps"] > 0
+                        ]
+                        per_req_tps = max(valid_gen) if valid_gen else 0.0
+                        gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
-                        peak = mean(
-                            x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
-                        )
 
-                        logger.info(
+                        def _peak_bytes(s: dict[str, Any]) -> float:
+                            pm = s["peak_memory_usage"]
+                            return pm.get("inBytes") or pm.get("in_bytes", 0)
+
+                        peak = mean(_peak_bytes(x["stats"]) for x in runs)
+                        summary = (
                             f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                             f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                            f"peak_memory={format_peak_memory(peak)}\n"
+                            f"peak_memory={format_peak_memory(peak)}"
                         )
+                        if sampler 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
+                            summary += f"    energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
+                        logger.info(f"{summary}\n")
                     time.sleep(2)
         finally:
-            try:
-                client.request_json("DELETE", f"/instance/{instance_id}")
-            except ExoHttpError as e:
-                if e.status != 404:
-                    raise
-            wait_for_instance_gone(client, instance_id)
-            logger.debug(f"Deleted instance {instance_id}")
+            if sampler:
+                sampler.stop()
+                placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
+                sampler.print_summary(placement_label)
+                placement_metrics = sampler.summarize()
+                if placement_metrics:
+                    all_system_metrics.update(placement_metrics)
 
-            time.sleep(5)
+            if created_instance and instance_id is not None:
+                try:
+                    client.request_json("DELETE", f"/instance/{instance_id}")
+                except ExoHttpError as e:
+                    if e.status != 404:
+                        raise
+                wait_for_instance_gone(client, instance_id)
+                logger.debug(f"Deleted instance {instance_id}")
+
+                time.sleep(5)
+
+    output: dict[str, Any] = {"runs": all_rows}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
+    if all_system_metrics:
+        output["system_metrics"] = all_system_metrics
 
     if args.stdout:
-        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
+        json.dump(output, sys.stdout, indent=2, ensure_ascii=False)
     elif args.json_out:
         with open(args.json_out, "w", encoding="utf-8") as f:
-            json.dump(all_rows, f, indent=2, ensure_ascii=False)
+            json.dump(output, f, indent=2, ensure_ascii=False)
         logger.debug(f"\nWrote results JSON: {args.json_out}")
 
     return 0

bench/exo_eval.py

@@ -42,10 +42,11 @@ from pathlib import Path
 from typing import Any
 
 import httpx
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
+    capture_cluster_snapshot,
+    find_existing_instance,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -61,6 +62,15 @@ from loguru import logger
 # ---------------------------------------------------------------------------
 
 MAX_RETRIES = 30
+INSTANCE_HEALTH_CHECK_AFTER = (
+    3  # Check instance health after this many consecutive failures
+)
+
+
+class InstanceFailedError(RuntimeError):
+    """Raised when the exo instance is detected as failed/gone."""
+
+
 DEFAULT_MAX_TOKENS = 16_384
 REASONING_MAX_TOKENS = 131_072
 TEMPERATURE_NON_REASONING = 0.0
@@ -270,7 +280,7 @@ def run_humaneval_test(
 
 @dataclass
 class QuestionResult:
-    question_id: int
+    question_id: int | str
     prompt: str
     response: str
     extracted_answer: str | None
@@ -280,7 +290,11 @@ class QuestionResult:
     prompt_tokens: int = 0
     completion_tokens: int = 0
     reasoning_tokens: int = 0
+    reasoning_content: str = ""
+    finish_reason: str = ""
     elapsed_s: float = 0.0
+    power_watts: float = 0.0
+    energy_joules: float = 0.0
 
 
 @dataclass
@@ -516,6 +530,10 @@ class ApiResult:
     prompt_tokens: int
     completion_tokens: int
     reasoning_tokens: int
+    reasoning_content: str = ""
+    finish_reason: str = ""
+    power_watts: float = 0.0
+    energy_joules: float = 0.0
 
 
 async def _call_api(
@@ -529,6 +547,9 @@ async def _call_api(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    enable_thinking: bool | None = None,
 ) -> ApiResult:
     messages = []
     if system_message:
@@ -545,6 +566,12 @@ async def _call_api(
         body["reasoning_effort"] = reasoning_effort
     if top_p is not None:
         body["top_p"] = top_p
+    if top_k is not None:
+        body["top_k"] = top_k
+    if min_p is not None:
+        body["min_p"] = min_p
+    if enable_thinking is not None:
+        body["enable_thinking"] = enable_thinking
 
     resp = await client.post(
         f"{base_url}/v1/chat/completions",
@@ -553,19 +580,40 @@ async def _call_api(
     )
     resp.raise_for_status()
     data = resp.json()
-    content = data["choices"][0]["message"]["content"]
-    if not content or not content.strip():
+    choice = data["choices"][0]
+    message = choice["message"]
+    content = message.get("content") or ""
+    reasoning_content = message.get("reasoning_content") or ""
+    finish_reason = choice.get("finish_reason") or ""
+
+    # For thinking models, empty content is expected when finish_reason is "length"
+    if not content.strip() and finish_reason != "length" and not reasoning_content:
         raise ValueError("Empty response from model")
     usage = data.get("usage", {})
     details = usage.get("completion_tokens_details", {})
+    power = data.get("power_usage") or {}
     return ApiResult(
         content=content,
         prompt_tokens=usage.get("prompt_tokens", 0),
         completion_tokens=usage.get("completion_tokens", 0),
         reasoning_tokens=details.get("reasoning_tokens", 0) if details else 0,
+        reasoning_content=reasoning_content,
+        finish_reason=finish_reason,
+        power_watts=power.get("total_avg_sys_power_watts", 0.0),
+        energy_joules=power.get("total_energy_joules", 0.0),
     )
 
 
+async def _check_instance_health(base_url: str) -> bool:
+    """Return True if the exo instance is still reachable."""
+    try:
+        async with httpx.AsyncClient() as c:
+            resp = await c.get(f"{base_url}/models", timeout=5.0)
+            return resp.status_code == 200
+    except Exception:
+        return False
+
+
 async def call_with_retries(
     client: httpx.AsyncClient,
     base_url: str,
@@ -577,8 +625,14 @@ async def call_with_retries(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    enable_thinking: bool | None = None,
+    instance_failed: asyncio.Event | None = None,
 ) -> ApiResult | None:
     for attempt in range(MAX_RETRIES):
+        if instance_failed and instance_failed.is_set():
+            raise InstanceFailedError("Instance already marked as failed")
         try:
             return await _call_api(
                 client,
@@ -591,8 +645,30 @@ async def call_with_retries(
                 system_message,
                 reasoning_effort,
                 top_p,
+                top_k,
+                min_p,
+                enable_thinking,
             )
         except Exception as e:
+            is_conn_error = isinstance(
+                e,
+                (
+                    httpx.ConnectError,
+                    httpx.RemoteProtocolError,
+                    ConnectionRefusedError,
+                    OSError,
+                ),
+            )
+            if (
+                is_conn_error
+                and attempt >= INSTANCE_HEALTH_CHECK_AFTER
+                and not await _check_instance_health(base_url)
+            ):
+                if instance_failed:
+                    instance_failed.set()
+                raise InstanceFailedError(
+                    f"Instance is down after {attempt + 1} failures: {e}"
+                ) from e
             if attempt < MAX_RETRIES - 1:
                 wait = min(2**attempt, 60)
                 logger.warning(
@@ -617,10 +693,16 @@ async def evaluate_benchmark(
     max_tokens: int,
     concurrency: int = 1,
     limit: int | None = None,
+    offset: int = 0,
     timeout: float | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    enable_thinking: bool | None = None,
     difficulty: str | None = None,
+    checkpoint_path: Path | None = None,
+    release_version: str | None = None,
 ) -> list[QuestionResult]:
     """Run a benchmark. Returns per-question results."""
     import datasets
@@ -651,7 +733,21 @@ async def evaluate_benchmark(
         ds = ds.filter(lambda x: x["difficulty"] == difficulty)
         logger.info(f"Filtered to {len(ds)} {difficulty} problems")
 
+    if release_version and "release_version" in ds.column_names:
+        ds = ds.filter(lambda x: x["release_version"] == release_version)
+        logger.info(
+            f"Filtered to {len(ds)} problems with release_version={release_version}"
+        )
+
+    # Sort by question_id to match LCB runner ordering (scenario_router.py:60).
+    # This ensures [offset:offset+limit] slices select the same problems as vllm.
+    if "question_id" in ds.column_names:
+        ds = ds.sort("question_id")
+
     total = len(ds)
+    if offset > 0:
+        ds = ds.select(range(min(offset, total), total))
+        total = len(ds)
     if limit and limit < total:
         ds = ds.select(range(limit))
         total = limit
@@ -659,6 +755,13 @@ async def evaluate_benchmark(
     logger.info(
         f"Evaluating {benchmark_name}: {total} questions, concurrency={concurrency}, "
         f"temperature={temperature}, max_tokens={max_tokens}"
+        + (f", top_k={top_k}" if top_k is not None else "")
+        + (f", min_p={min_p}" if min_p is not None else "")
+        + (
+            f", enable_thinking={enable_thinking}"
+            if enable_thinking is not None
+            else ""
+        )
     )
 
     if config.kind == "code":
@@ -666,16 +769,64 @@ async def evaluate_benchmark(
             "Code benchmarks execute model-generated code. Use a sandboxed environment."
         )
 
+    # Load checkpoint for resume
+    checkpoint_data: dict[str | int, dict[str, Any]] = {}
+    if checkpoint_path and checkpoint_path.exists():
+        with open(checkpoint_path) as f:
+            for line in f:
+                entry = json.loads(line)
+                checkpoint_data[entry["question_id"]] = entry
+        logger.info(f"Loaded {len(checkpoint_data)} checkpointed results")
+
     semaphore = asyncio.Semaphore(concurrency)
+    instance_failed = asyncio.Event()
     results: list[QuestionResult | None] = [None] * total
     completed = 0
     lock = asyncio.Lock()
 
+    def _get_question_id(idx: int, doc: dict) -> str | int:
+        """Get a stable question ID for checkpointing."""
+        if benchmark_name == "livecodebench":
+            return doc.get("question_id", idx)
+        elif benchmark_name == "humaneval":
+            return doc.get("task_id", idx)
+        return idx
+
     async def process_question(
         idx: int, doc: dict, http_client: httpx.AsyncClient
     ) -> None:
         nonlocal completed
         system_msg = None
+        question_id = _get_question_id(idx, doc)
+
+        # Bail out early if instance is already dead
+        if instance_failed.is_set():
+            return
+
+        # Check checkpoint
+        if question_id in checkpoint_data:
+            cached = checkpoint_data[question_id]
+            results[idx] = QuestionResult(
+                question_id=question_id,
+                prompt=cached.get("prompt", ""),
+                response=cached.get("response", ""),
+                extracted_answer=cached.get("extracted_answer"),
+                gold_answer=cached.get("gold_answer", ""),
+                correct=cached.get("correct", False),
+                error=cached.get("error"),
+                prompt_tokens=cached.get("prompt_tokens", 0),
+                completion_tokens=cached.get("completion_tokens", 0),
+                reasoning_tokens=cached.get("reasoning_tokens", 0),
+                reasoning_content=cached.get("reasoning_content", ""),
+                finish_reason=cached.get("finish_reason", ""),
+                elapsed_s=cached.get("elapsed_s", 0.0),
+                power_watts=cached.get("power_watts", 0.0),
+                energy_joules=cached.get("energy_joules", 0.0),
+            )
+            async with lock:
+                completed += 1
+            logger.info(f"  [{completed}/{total}] {question_id} (cached)")
+            return
 
         if benchmark_name == "gpqa_diamond":
             prompt, gold = format_gpqa_question(doc, idx)
@@ -696,24 +847,50 @@ async def evaluate_benchmark(
             raise ValueError(f"Unknown benchmark: {benchmark_name}")
 
         async with semaphore:
+            if instance_failed.is_set():
+                return
             t0 = time.monotonic()
-            api_result = await call_with_retries(
-                http_client,
-                base_url,
-                model,
-                prompt,
-                temperature,
-                max_tokens,
-                timeout,
-                system_message=system_msg,
-                reasoning_effort=reasoning_effort,
-                top_p=top_p,
-            )
+            try:
+                # Race the API call against the instance_failed event
+                api_task = asyncio.create_task(
+                    call_with_retries(
+                        http_client,
+                        base_url,
+                        model,
+                        prompt,
+                        temperature,
+                        max_tokens,
+                        timeout,
+                        system_message=system_msg,
+                        reasoning_effort=reasoning_effort,
+                        top_p=top_p,
+                        top_k=top_k,
+                        min_p=min_p,
+                        enable_thinking=enable_thinking,
+                        instance_failed=instance_failed,
+                    )
+                )
+                failed_waiter = asyncio.create_task(instance_failed.wait())
+                done, pending = await asyncio.wait(
+                    [api_task, failed_waiter],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                for p in pending:
+                    p.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await p
+                if instance_failed.is_set() and api_task not in done:
+                    logger.error(f"Instance failed, aborting {question_id}")
+                    return
+                api_result = api_task.result()
+            except InstanceFailedError:
+                logger.error(f"Instance failed, skipping {question_id}")
+                return
             elapsed = time.monotonic() - t0
 
         if api_result is None:
             result = QuestionResult(
-                question_id=idx,
+                question_id=question_id,
                 prompt=prompt,
                 response="",
                 extracted_answer=None,
@@ -728,13 +905,17 @@ async def evaluate_benchmark(
                 "prompt_tokens": api_result.prompt_tokens,
                 "completion_tokens": api_result.completion_tokens,
                 "reasoning_tokens": api_result.reasoning_tokens,
+                "reasoning_content": api_result.reasoning_content,
+                "finish_reason": api_result.finish_reason,
                 "elapsed_s": elapsed,
+                "power_watts": api_result.power_watts,
+                "energy_joules": api_result.energy_joules,
             }
 
             if config.kind == "mc":
                 extracted = extract_mc_answer(response, valid_letters)
                 result = QuestionResult(
-                    question_id=idx,
+                    question_id=question_id,
                     prompt=prompt,
                     response=response,
                     extracted_answer=extracted,
@@ -748,7 +929,7 @@ async def evaluate_benchmark(
                     check_aime_answer(extracted, int(gold)) if extracted else False
                 )
                 result = QuestionResult(
-                    question_id=idx,
+                    question_id=question_id,
                     prompt=prompt,
                     response=response,
                     extracted_answer=extracted,
@@ -762,7 +943,7 @@ async def evaluate_benchmark(
                 code = extract_code_block(response, preserve_indent=keep_indent)
                 if code is None:
                     result = QuestionResult(
-                        question_id=idx,
+                        question_id=question_id,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -777,7 +958,7 @@ async def evaluate_benchmark(
                         code,
                     )
                     result = QuestionResult(
-                        question_id=idx,
+                        question_id=question_id,
                         prompt=prompt,
                         response=response,
                         extracted_answer="pass" if passed else "fail",
@@ -792,7 +973,7 @@ async def evaluate_benchmark(
                         exec_meta["sample"],
                     )
                     result = QuestionResult(
-                        question_id=idx,
+                        question_id=question_id,
                         prompt=prompt,
                         response=response,
                         extracted_answer="pass" if passed else "fail",
@@ -803,7 +984,7 @@ async def evaluate_benchmark(
                     )
                 else:
                     result = QuestionResult(
-                        question_id=idx,
+                        question_id=question_id,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -814,7 +995,7 @@ async def evaluate_benchmark(
                     )
             else:
                 result = QuestionResult(
-                    question_id=idx,
+                    question_id=question_id,
                     prompt=prompt,
                     response=response,
                     extracted_answer=None,
@@ -826,24 +1007,82 @@ async def evaluate_benchmark(
 
         results[idx] = result
 
+        # Write checkpoint (skip infra failures so they get retried on resume,
+        # but keep wrong answers — they are legitimate results)
+        if checkpoint_path is not None and result.response:
+            _write_checkpoint(checkpoint_path, result)
+
         async with lock:
             completed += 1
             n = completed
-        if n % max(1, total // 20) == 0 or n == total:
-            correct_so_far = sum(1 for r in results if r is not None and r.correct)
-            answered = sum(1 for r in results if r is not None)
-            logger.info(
-                f"  [{n}/{total}] {correct_so_far}/{answered} correct "
-                f"({correct_so_far / max(answered, 1):.1%})"
-            )
+
+        # Log progress
+        thinking_info = ""
+        if result.reasoning_content:
+            thinking_info = f", {len(result.reasoning_content)} chars thinking"
+        logger.info(
+            f"  [{n}/{total}] {question_id}: {len(result.response)} chars{thinking_info}, "
+            f"tokens: {result.prompt_tokens}+{result.completion_tokens} "
+            f"[{result.finish_reason}]"
+            + (f" {result.extracted_answer}" if result.extracted_answer else "")
+        )
+
+    async def _health_monitor() -> None:
+        """Periodically check if the instance is still alive."""
+        # Wait a bit before first check to let things start
+        await asyncio.sleep(10)
+        while not instance_failed.is_set():
+            if not await _check_instance_health(base_url):
+                # Double-check to avoid false positives
+                await asyncio.sleep(2)
+                if not await _check_instance_health(base_url):
+                    logger.error("Health monitor: instance is down!")
+                    instance_failed.set()
+                    return
+            await asyncio.sleep(5)
 
     async with httpx.AsyncClient() as http_client:
+        monitor = asyncio.create_task(_health_monitor())
         tasks = [process_question(i, doc, http_client) for i, doc in enumerate(ds)]
         await asyncio.gather(*tasks)
+        monitor.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await monitor
+
+    if instance_failed.is_set():
+        completed_count = sum(1 for r in results if r is not None)
+        logger.error(
+            f"Instance failed! Completed {completed_count}/{total} problems. "
+            f"Checkpoint saved — restart to resume remaining problems."
+        )
+        raise InstanceFailedError("Instance failed during evaluation")
 
     return [r for r in results if r is not None]
 
 
+def _write_checkpoint(path: Path, result: QuestionResult) -> None:
+    """Append a single result to the JSONL checkpoint file."""
+    entry = {
+        "question_id": result.question_id,
+        "prompt": result.prompt,
+        "response": result.response,
+        "extracted_answer": result.extracted_answer,
+        "gold_answer": result.gold_answer,
+        "correct": result.correct,
+        "error": result.error,
+        "prompt_tokens": result.prompt_tokens,
+        "completion_tokens": result.completion_tokens,
+        "reasoning_tokens": result.reasoning_tokens,
+        "reasoning_content": result.reasoning_content,
+        "finish_reason": result.finish_reason,
+        "elapsed_s": round(result.elapsed_s, 2),
+        "power_watts": round(result.power_watts, 2),
+        "energy_joules": round(result.energy_joules, 2),
+    }
+    with open(path, "a") as f:
+        f.write(json.dumps(entry) + "\n")
+
+
 # ---------------------------------------------------------------------------
 # Results display
 # ---------------------------------------------------------------------------
@@ -866,6 +1105,8 @@ def print_results(
     total_elapsed = sum(r.elapsed_s for r in results)
     wall_clock = max(r.elapsed_s for r in results) if results else 0.0
     avg_gen_tps = total_completion_tokens / total_elapsed if total_elapsed > 0 else 0.0
+    total_energy = sum(r.energy_joules for r in results)
+    avg_power = sum(r.power_watts for r in results) / max(total, 1)
 
     label = f"[c={concurrency}] " if concurrency is not None else ""
     print(f"\n{label}{benchmark_name}: {correct}/{total} ({accuracy:.1%})")
@@ -877,6 +1118,10 @@ def print_results(
         f"  |  total time: {total_elapsed:.1f}s  wall clock: {wall_clock:.1f}s"
     )
     print(tok_line)
+    if total_energy > 0:
+        print(
+            f"  power: avg {avg_power:.1f}W  |  total energy: {total_energy:.1f}J ({total_energy / 3600:.2f}Wh)"
+        )
     if errors:
         print(f"  API errors: {errors}")
     if no_extract:
@@ -895,6 +1140,8 @@ def print_results(
         "total_elapsed_s": total_elapsed,
         "wall_clock_s": wall_clock,
         "avg_gen_tps": avg_gen_tps,
+        "avg_power_watts": avg_power,
+        "total_energy_joules": total_energy,
     }
 
 
@@ -1027,16 +1274,18 @@ def save_results(
     concurrency: int,
     results: list[QuestionResult],
     scores: dict[str, Any],
+    cluster: dict[str, Any] | None = None,
 ) -> Path:
     out_dir = Path(results_dir) / model.replace("/", "_") / benchmark_name
     out_dir.mkdir(parents=True, exist_ok=True)
     ts = time.strftime("%Y%m%d_%H%M%S")
     path = out_dir / f"c{concurrency}_{ts}.json"
 
-    data = {
+    data: dict[str, Any] = {
         "benchmark": benchmark_name,
         "model": model,
         "concurrency": concurrency,
+        **({"cluster": cluster} if cluster else {}),
         "scores": scores,
         "results": [
             {
@@ -1050,7 +1299,11 @@ def save_results(
                 "prompt_tokens": r.prompt_tokens,
                 "completion_tokens": r.completion_tokens,
                 "reasoning_tokens": r.reasoning_tokens,
+                "reasoning_content": r.reasoning_content,
+                "finish_reason": r.finish_reason,
                 "elapsed_s": round(r.elapsed_s, 2),
+                "power_watts": round(r.power_watts, 2),
+                "energy_joules": round(r.energy_joules, 2),
             }
             for r in results
         ],
@@ -1066,6 +1319,15 @@ def save_results(
 # ---------------------------------------------------------------------------
 
 
+def _checkpoint_path(
+    results_dir: str, benchmark: str, model: str, concurrency: int
+) -> Path:
+    """Return the JSONL checkpoint path for a benchmark run."""
+    out_dir = Path(results_dir) / model.replace("/", "_") / benchmark
+    out_dir.mkdir(parents=True, exist_ok=True)
+    return out_dir / f"c{concurrency}.checkpoint.jsonl"
+
+
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -1093,6 +1355,12 @@ def main() -> int:
         default=None,
         help="Max questions per benchmark (for fast iteration).",
     )
+    ap.add_argument(
+        "--offset",
+        type=int,
+        default=0,
+        help="Skip first N questions (0-based).",
+    )
 
     reasoning_group = ap.add_mutually_exclusive_group()
     reasoning_group.add_argument(
@@ -1112,6 +1380,8 @@ def main() -> int:
         "--temperature", type=float, default=None, help="Override temperature."
     )
     ap.add_argument("--top-p", type=float, default=None, help="Override top_p.")
+    ap.add_argument("--top-k", type=int, default=None, help="Override top_k.")
+    ap.add_argument("--min-p", type=float, default=None, help="Override min_p.")
     ap.add_argument(
         "--max-tokens", type=int, default=None, help="Override max output tokens."
     )
@@ -1145,15 +1415,31 @@ def main() -> int:
         choices=["easy", "medium", "hard"],
         help="Filter by difficulty (livecodebench only). E.g. --difficulty hard",
     )
+    ap.add_argument(
+        "--release-version",
+        default=None,
+        help="LCB dataset release version (livecodebench only). E.g. release_v5",
+    )
     ap.add_argument(
         "--results-dir",
         default="eval_results",
         help="Directory for result JSON files (default: eval_results).",
     )
     ap.add_argument(
-        "--skip-instance-setup",
+        "--enable-thinking",
+        type=lambda v: v.lower() in ("true", "1", "yes"),
+        default=None,
+        help="Enable thinking mode for models that support it.",
+    )
+    ap.add_argument(
+        "--force",
         action="store_true",
-        help="Skip exo instance management (assumes model is already running).",
+        help="Discard any existing checkpoint and run from scratch.",
+    )
+    ap.add_argument(
+        "--keep-instance",
+        action="store_true",
+        help="Skip deleting the instance after eval (for chaining runs).",
     )
 
     args, _ = ap.parse_known_args()
@@ -1174,13 +1460,26 @@ def main() -> int:
     # Instance management
     client = ExoClient(args.host, args.port, timeout_s=args.timeout)
     instance_id: str | None = None
+    created_instance = False
 
-    if not args.skip_instance_setup:
-        short_id, full_model_id = resolve_model_short_id(
-            client,
-            args.model,
-            force_download=args.force_download,
-        )
+    _short_id, full_model_id = resolve_model_short_id(
+        client,
+        args.model,
+        force_download=args.force_download,
+    )
+
+    # Optionally reuse a running instance for this model
+    if args.reuse_instance:
+        existing = find_existing_instance(client, full_model_id)
+        if existing:
+            instance_id = existing
+            logger.info(f"Reusing existing instance {instance_id}")
+        else:
+            logger.warning(
+                "--reuse-instance: no existing instance found, creating a new one"
+            )
+
+    if instance_id is None:
         selected = settle_and_fetch_placements(
             client,
             full_model_id,
@@ -1195,7 +1494,7 @@ def main() -> int:
             key=lambda p: (
                 str(p.get("instance_meta", "")),
                 str(p.get("sharding", "")),
-                -nodes_used_in_instance(p["instance"]),
+                nodes_used_in_instance(p["instance"]),
             ),
             reverse=True,
         )
@@ -1222,6 +1521,18 @@ def main() -> int:
         if download_duration is not None:
             logger.info(f"Download: {download_duration:.1f}s")
 
+        # Delete any existing instances to free resources before placing
+        try:
+            state = client.request_json("GET", "/state")
+            for old_id in list(state.get("instances", {}).keys()):
+                logger.info(f"Deleting stale instance {old_id}")
+                with contextlib.suppress(ExoHttpError):
+                    client.request_json("DELETE", f"/instance/{old_id}")
+            if state.get("instances"):
+                time.sleep(2)
+        except Exception as e:
+            logger.warning(f"Failed to clean up stale instances: {e}")
+
         client.request_json("POST", "/instance", body={"instance": instance})
         try:
             wait_for_instance_ready(client, instance_id)
@@ -1231,8 +1542,9 @@ def main() -> int:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             return 1
         time.sleep(1)
-    else:
-        full_model_id = args.model
+        created_instance = True
+
+    cluster_snapshot = capture_cluster_snapshot(client)
 
     # Auto-detect reasoning from model config
     model_config = load_model_config(full_model_id)
@@ -1286,16 +1598,57 @@ def main() -> int:
         reasoning_effort = str(cfg["reasoning_effort"])
     else:
         reasoning_effort = "high" if is_reasoning else None
+
+    if args.top_k is not None:
+        top_k: int | None = args.top_k
+    elif "top_k" in cfg:
+        top_k = int(cfg["top_k"])
+    else:
+        top_k = None
+
+    if args.min_p is not None:
+        min_p: float | None = args.min_p
+    elif "min_p" in cfg:
+        min_p = float(cfg["min_p"])
+    else:
+        min_p = None
+
+    if args.enable_thinking is not None:
+        enable_thinking: bool | None = args.enable_thinking
+    elif "enable_thinking" in cfg:
+        enable_thinking = bool(cfg["enable_thinking"])
+    else:
+        enable_thinking = None
+
     base_url = f"http://{args.host}:{args.port}"
 
     logger.info(f"Model: {full_model_id}")
     logger.info(
         f"Settings: temperature={temperature}, max_tokens={max_tokens}, "
         + (f"top_p={top_p}, " if top_p is not None else "")
+        + (f"top_k={top_k}, " if top_k is not None else "")
+        + (f"min_p={min_p}, " if min_p is not None else "")
         + f"reasoning={'yes' if is_reasoning else 'no'}"
         + (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "")
+        + (
+            f", enable_thinking={enable_thinking}"
+            if enable_thinking is not None
+            else ""
+        )
     )
 
+    # Common kwargs for evaluate_benchmark
+    eval_kwargs: dict[str, Any] = {
+        "reasoning_effort": reasoning_effort,
+        "top_p": top_p,
+        "top_k": top_k,
+        "min_p": min_p,
+        "enable_thinking": enable_thinking,
+        "difficulty": args.difficulty,
+        "offset": args.offset,
+        "release_version": args.release_version,
+    }
+
     try:
         if args.compare_concurrency:
             concurrency_levels = parse_int_list(args.compare_concurrency)
@@ -1304,6 +1657,11 @@ def main() -> int:
                 for c in concurrency_levels:
                     logger.info(f"\n{'=' * 50}")
                     logger.info(f"Running {task_name} at concurrency={c}")
+                    checkpoint_path = _checkpoint_path(
+                        args.results_dir, task_name, full_model_id, c
+                    )
+                    if args.force and checkpoint_path.exists():
+                        checkpoint_path.unlink()
                     results = asyncio.run(
                         evaluate_benchmark(
                             task_name,
@@ -1314,9 +1672,8 @@ def main() -> int:
                             concurrency=c,
                             limit=args.limit,
                             timeout=args.request_timeout,
-                            reasoning_effort=reasoning_effort,
-                            top_p=top_p,
-                            difficulty=args.difficulty,
+                            checkpoint_path=checkpoint_path,
+                            **eval_kwargs,
                         )
                     )
                     if results:
@@ -1328,12 +1685,21 @@ def main() -> int:
                             c,
                             results,
                             scores,
+                            cluster=cluster_snapshot,
                         )
                         results_by_c[c] = results
+                    # Clean up checkpoint on success
+                    if checkpoint_path.exists():
+                        checkpoint_path.unlink()
                 if len(results_by_c) >= 2:
                     print_comparison(task_name, results_by_c)
         else:
             for task_name in task_names:
+                checkpoint_path = _checkpoint_path(
+                    args.results_dir, task_name, full_model_id, args.num_concurrent
+                )
+                if args.force and checkpoint_path.exists():
+                    checkpoint_path.unlink()
                 results = asyncio.run(
                     evaluate_benchmark(
                         task_name,
@@ -1344,9 +1710,8 @@ def main() -> int:
                         concurrency=args.num_concurrent,
                         limit=args.limit,
                         timeout=args.request_timeout,
-                        reasoning_effort=reasoning_effort,
-                        top_p=top_p,
-                        difficulty=args.difficulty,
+                        checkpoint_path=checkpoint_path,
+                        **eval_kwargs,
                     )
                 )
                 if results:
@@ -1358,15 +1723,27 @@ def main() -> int:
                         args.num_concurrent,
                         results,
                         scores,
+                        cluster=cluster_snapshot,
                     )
+                # Clean up checkpoint on success
+                if checkpoint_path.exists():
+                    checkpoint_path.unlink()
     finally:
-        if instance_id is not None:
-            try:
-                client.request_json("DELETE", f"/instance/{instance_id}")
-            except ExoHttpError as e:
-                if e.status != 404:
-                    raise
-            wait_for_instance_gone(client, instance_id)
+        if created_instance and instance_id is not None:
+            if args.keep_instance:
+                logger.info(f"Keeping instance {instance_id} (--keep-instance)")
+            else:
+                try:
+                    client.request_json("DELETE", f"/instance/{instance_id}")
+                except ExoHttpError as e:
+                    if e.status != 404:
+                        raise
+                try:
+                    wait_for_instance_gone(client, instance_id)
+                except TimeoutError:
+                    logger.warning(
+                        f"Timed out waiting for instance {instance_id} to be deleted"
+                    )
 
     return 0
 

bench/lib/__init__.py

@@ -0,0 +1,18 @@
+"""Composable bench library for exo.
+
+Provides reusable building blocks for benchmarks:
+
+- :class:`bench.lib.session.BenchSession` — cluster + instance + client wrapper
+- :class:`bench.lib.results.ResultsBundle` — structured results + JSON writer
+- :func:`bench.lib.cluster.managed_cluster` /
+  :func:`bench.lib.cluster.managed_instance` — eco-managed lifecycle ctx-managers
+- :func:`bench.lib.model_meta.fetch_model_meta` — HF metadata fetcher driving
+  cluster constraints + auto-derived context ramps
+- :mod:`bench.lib.context_scaling` — prompt-TPS / decode-TPS vs context-size sweep
+
+CLI entrypoints under ``bench/cli/`` consume this library via
+``python -m bench.cli <subcommand>``. Adding a new benchmark = (i) write
+``bench/lib/<name>.py`` exposing a typed ``run(session, params, bundle)``
+callable, (ii) write ``bench/cli/<name>.py`` with an ``add_subparser`` and
+a handler, (iii) register it in ``_REGISTRY`` in ``bench/cli/__main__.py``.
+"""

bench/lib/cluster.py

@@ -0,0 +1,215 @@
+"""Eco-managed cluster + instance lifecycle helpers for the bench CLI.
+
+Two context managers:
+
+- :func:`managed_cluster` deploys exo on the requested hosts (or via
+  constraint-based reservation) and tears it down on exit.
+- :func:`managed_instance` resolves the model on the cluster, optionally
+  frees disk via ``--danger-delete-downloads`` (default on for benches),
+  places the instance, and deletes it on exit.
+
+The library never reaches for global state — every call takes an
+explicit :class:`EcoSession`. Callers are expected to instantiate one
+session per CLI invocation and use it across both context managers.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any, cast
+
+from exo_tools.client import ExoClient
+from exo_tools.cluster import Chip, ClusterInfo, EcoSession, Thunderbolt
+from exo_tools.harness import (
+    Comm,
+    Sharding,
+    cleanup_all_instances,
+    place_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+)
+from loguru import logger
+
+from .session import BenchSession
+
+
+@contextmanager
+def managed_cluster(
+    eco: EcoSession,
+    *,
+    hosts: list[str] | None = None,
+    count: int = 1,
+    thunderbolt: Thunderbolt | None = None,
+    chip: Chip | None = None,
+    min_memory_gb: float | None = None,
+    max_memory_gb: float | None = None,
+    min_disk_gb: float | None = None,
+    max_disk_gb: float | None = None,
+    deploy_timeout_s: int = 600,
+) -> Iterator[ClusterInfo]:
+    """Deploy exo for the duration of the ``with`` block, then ``eco stop``.
+
+    If ``hosts`` is given, deploys on exactly those hosts (constraint flags
+    are ignored — eco doesn't re-validate the explicit list). Otherwise eco
+    reserves any matching hosts that satisfy all of:
+
+      - ``count`` (number of hosts)
+      - ``thunderbolt`` topology (``A2A``, ``RING``, or ``NONE`` to
+        exclude TB-connected hosts)
+      - ``chip`` (substring match against eco's chip names)
+      - memory bounds (``min_memory_gb`` / ``max_memory_gb``)
+      - disk bounds (``min_disk_gb`` / ``max_disk_gb``)
+    """
+    if hosts:
+        cluster = eco.start_deploy(
+            hosts=hosts[:count],
+            wait=True,
+            timeout=deploy_timeout_s,
+        )
+    else:
+        cluster = eco.start_deploy(
+            count=count,
+            thunderbolt=thunderbolt,
+            chip=chip,
+            min_memory_gb=min_memory_gb,
+            max_memory_gb=max_memory_gb,
+            min_disk_gb=min_disk_gb,
+            max_disk_gb=max_disk_gb,
+            wait=True,
+            timeout=deploy_timeout_s,
+        )
+    logger.info(
+        f"cluster deployed: {len(cluster.hosts)} host(s) "
+        f"({', '.join(cluster.hosts)}); namespace={cluster.namespace}"
+    )
+    try:
+        yield cluster
+    finally:
+        with contextlib.suppress(Exception):
+            eco.stop(cluster.hosts)
+        logger.info("cluster stopped")
+
+
+@contextmanager
+def managed_instance(
+    cluster: ClusterInfo,
+    eco: EcoSession,
+    model_id: str,
+    *,
+    sharding: Sharding = Sharding.PIPELINE,
+    comm: Comm = Comm.RING,
+    min_nodes: int = 1,
+    evict_downloads: bool = True,
+    cleanup_on_exit: bool = True,
+    instance_timeout_s: float = 7200.0,
+    settle_timeout_s: float = 60.0,
+) -> Iterator[BenchSession]:
+    """Resolve the model on the cluster, place an instance, yield a session.
+
+    Steps on entry:
+      1. Resolve ``model_id`` to ``(short_id, full_id)`` against the cluster's
+         ``/models`` endpoint (auto-adds from HuggingFace if missing).
+      2. Run the harness's planning phase: validates each node has enough
+         disk for the model and starts the download (or reuses an existing
+         download). When ``evict_downloads=True`` (the default for benches),
+         this also evicts smaller existing models if disk is short.
+      3. Place the instance, wait for it to be ``RunnerReady``.
+      4. Yield a :class:`BenchSession` pointing at the cluster's primary API.
+
+    On exit: deletes the placed instance (and any other lingering
+    instances) so the cluster is clean for the next benchmark.
+    """
+    client = cluster.make_client(timeout_s=instance_timeout_s)
+
+    short_id, full_id = resolve_model_short_id(client, model_id, force_download=True)
+    logger.info(f"resolved model: short_id={short_id} full_id={full_id}")
+
+    # The planning phase needs a concrete preview (instance + runner-to-shard
+    # mapping) to know which nodes to download to. Pull the placements API
+    # directly and take the first valid one — bench cares about disk +
+    # download, not the specific shard mapping.
+    preview = _first_valid_preview(client, full_id, settle_timeout_s)
+    if preview is None:
+        raise RuntimeError(
+            f"No placement available for {full_id} on cluster {cluster.hosts}"
+        )
+
+    duration = run_planning_phase(
+        client,
+        full_id,
+        preview,
+        danger_delete=evict_downloads,
+        timeout=instance_timeout_s,
+        settle_deadline=None,
+    )
+    if duration is not None:
+        logger.info(f"download: {duration:.1f}s (freshly downloaded)")
+    else:
+        logger.info("download: model already cached on all nodes")
+
+    instance_id = place_instance(
+        client,
+        model_id,
+        sharding=sharding,
+        comm=comm,
+        min_nodes=min_nodes,
+        timeout=instance_timeout_s,
+    )
+    logger.info(f"placed instance {instance_id} ({sharding.value}/{comm.value})")
+
+    sess = BenchSession(
+        cluster=cluster,
+        eco=eco,
+        instance_id=instance_id,
+        model_id=short_id,
+        full_model_id=full_id,
+    )
+    try:
+        yield sess
+    finally:
+        if cleanup_on_exit:
+            with contextlib.suppress(Exception):
+                cleanup_all_instances(sess.client)
+        else:
+            logger.info(
+                f"cleanup_on_exit=False: leaving instance(s) on {cluster.hosts}"
+            )
+
+
+def _first_valid_preview(
+    client: ExoClient, full_model_id: str, settle_timeout_s: float
+) -> dict[str, Any] | None:
+    """Poll ``/instance/previews`` until at least one valid preview comes back."""
+    deadline = time.monotonic() + settle_timeout_s
+    backoff_s = 1.0
+    while True:
+        resp_obj: Any = client.request_json(  # type: ignore[reportAny]
+            "GET", "/instance/previews", params={"model_id": full_model_id}
+        )
+        resp: dict[str, Any] = (
+            cast("dict[str, Any]", resp_obj) if isinstance(resp_obj, dict) else {}
+        )
+        previews_raw: object = resp.get("previews") or []
+        previews: list[Any] = (
+            cast("list[Any]", previews_raw) if isinstance(previews_raw, list) else []
+        )
+        for raw in previews:  # type: ignore[reportAny]
+            if not isinstance(raw, dict):
+                continue
+            entry = cast("dict[str, Any]", raw)
+            if entry.get("error") is not None:
+                continue
+            instance = entry.get("instance")
+            if isinstance(instance, dict):
+                return entry
+        if time.monotonic() >= deadline:
+            return None
+        logger.info(
+            f"waiting for placement to appear for {full_model_id} "
+            f"({deadline - time.monotonic():.0f}s remaining)..."
+        )
+        time.sleep(min(backoff_s, max(0.0, deadline - time.monotonic())))
+        backoff_s = min(backoff_s * 2, 30.0)

bench/lib/completion.py

@@ -0,0 +1,194 @@
+"""Typed wrapper around ``/bench/chat/completions`` for benchmarks.
+
+The bench endpoint disables EOS suppression and KV prefix caching by
+default (see ``bench/METHODOLOGY.md``). This module exposes a single
+function :func:`run_one_completion` that:
+
+  1. Builds an exact-token-length prompt via :class:`PromptSizer`.
+  2. POSTs to ``/bench/chat/completions``.
+  3. Returns a ``(BenchRow, prompt_tokens)`` pair where ``BenchRow`` is a
+     :class:`typing.TypedDict` with the fields the caller needs.
+
+Streaming is supported but rarely needed for context-scaling — the
+non-streaming path is the default.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from typing import Any, Literal, NotRequired, TypedDict, cast
+
+from exo_tools.client import ExoClient
+
+from .prompt import PromptSizer
+
+PrefixCacheHit = Literal["none", "partial", "exact"]
+
+
+class GenerationStats(TypedDict, total=False):
+    """Server-reported per-task timing stats."""
+
+    prompt_tps: float
+    generation_tps: float
+    prompt_tokens: int
+    generation_tokens: int
+    peak_memory_usage: dict[str, int]
+    prefix_cache_hit: PrefixCacheHit
+
+
+class BenchRow(TypedDict):
+    """Per-request result row returned to callers."""
+
+    elapsed_s: float
+    output_text_preview: str
+    stats: GenerationStats
+    error: NotRequired[str]
+
+
+def _as_dict(value: Any) -> dict[str, Any]:  # type: ignore[reportAny]
+    """Narrow an arbitrary JSON value to a typed ``dict[str, Any]``."""
+    if isinstance(value, dict):
+        return cast("dict[str, Any]", value)
+    return {}
+
+
+def _as_list(value: Any) -> list[Any]:  # type: ignore[reportAny]
+    if isinstance(value, list):
+        return cast("list[Any]", value)
+    return []
+
+
+def _extract_stats(raw_response: dict[str, Any]) -> GenerationStats:
+    stats_obj = raw_response.get("generation_stats")
+    if not isinstance(stats_obj, dict):
+        return {}
+    return cast("GenerationStats", cast("object", stats_obj))
+
+
+def _extract_preview(raw_response: dict[str, Any], limit: int = 200) -> str:
+    choices = _as_list(raw_response.get("choices"))
+    if not choices:
+        return ""
+    first = _as_dict(choices[0])
+    message = _as_dict(first.get("message"))
+    content_obj = message.get("content")
+    if isinstance(content_obj, str):
+        return content_obj[:limit]
+    return ""
+
+
+def run_one_completion(
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool = False,
+    stream: bool = False,
+) -> tuple[BenchRow, int]:
+    """Send one request to ``/bench/chat/completions`` and return its row.
+
+    ``pp_hint`` is the *target* prompt-token count; the actual prompt is
+    sized via :class:`PromptSizer` and the verified value is returned as
+    the second element of the tuple.
+    """
+    content, pp_tokens = prompt_sizer.build(pp_hint)
+    payload: dict[str, Any] = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": content}],
+        "max_tokens": tg,
+        "logprobs": False,
+        "use_prefix_cache": use_prefix_cache,
+    }
+
+    if not stream:
+        payload["stream"] = False
+        t0 = time.perf_counter()
+        raw_obj = client.post_bench_chat_completions(payload)
+        elapsed = time.perf_counter() - t0
+        raw = _as_dict(raw_obj)
+        return (
+            BenchRow(
+                elapsed_s=elapsed,
+                output_text_preview=_extract_preview(raw),
+                stats=_extract_stats(raw),
+            ),
+            pp_tokens,
+        )
+
+    return _run_streaming(client, payload, pp_tokens)
+
+
+def _run_streaming(
+    client: ExoClient,
+    payload: dict[str, Any],
+    pp_tokens: int,
+) -> tuple[BenchRow, int]:
+    """Streaming variant: parse SSE lines, recover ``GenerationStats``."""
+    payload = {**payload, "stream": True}
+
+    tokens = 0
+    first_token_time: float | None = None
+    t0 = time.perf_counter()
+    text_parts: list[str] = []
+    stats: GenerationStats = {}
+
+    for raw_line in client.stream_bench_chat_completions(payload):
+        line = raw_line.strip()
+        if line.startswith(": generation_stats "):
+            with contextlib.suppress(json.JSONDecodeError):
+                parsed_obj: Any = json.loads(  # type: ignore[reportAny]
+                    line[len(": generation_stats ") :]
+                )
+                if isinstance(parsed_obj, dict):
+                    stats = cast("GenerationStats", cast("object", parsed_obj))
+            continue
+        if not line.startswith("data: "):
+            continue
+        data = line[6:]
+        if data == "[DONE]":
+            break
+        try:
+            chunk_obj: Any = json.loads(data)  # type: ignore[reportAny]
+        except json.JSONDecodeError:
+            continue
+        chunk = _as_dict(chunk_obj)
+        choices = _as_list(chunk.get("choices"))
+        if not choices:
+            continue
+        first = _as_dict(choices[0])
+        delta = _as_dict(first.get("delta"))
+        delta_content_obj = delta.get("content")
+        if isinstance(delta_content_obj, str) and delta_content_obj:
+            if first_token_time is None:
+                first_token_time = time.perf_counter()
+            tokens += 1
+            text_parts.append(delta_content_obj)
+
+    elapsed = time.perf_counter() - t0
+    preview = "".join(text_parts)[:200]
+
+    if not stats:
+        ttft = (first_token_time - t0) if first_token_time is not None else elapsed
+        gen_time = elapsed - ttft if tokens > 1 else elapsed
+        gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
+        prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
+        stats = GenerationStats(
+            prompt_tokens=pp_tokens,
+            generation_tokens=tokens,
+            prompt_tps=round(prompt_tps, 2),
+            generation_tps=round(gen_tps, 2),
+            peak_memory_usage={"inBytes": 0},
+        )
+
+    return (
+        BenchRow(
+            elapsed_s=elapsed,
+            output_text_preview=preview,
+            stats=stats,
+        ),
+        pp_tokens,
+    )

bench/lib/context_scaling.py

@@ -0,0 +1,428 @@
+"""Prompt-TPS / decode-TPS vs context-size sweep.
+
+Methodology (see also ``bench/METHODOLOGY.md``):
+
+  Run a single ascending ramp of equally-spaced prompt lengths
+  ``pp ∈ {Δ, 2Δ, …, K·Δ}`` with ``prefix_cache=enabled``, ``repeat=1``,
+  ``concurrency=1`` and one warmup at ``pp=Δ``.
+
+  Because each step's prefix is exactly what the previous step left in
+  the cache, every step beyond the first is a *partial* hit and the
+  server-reported ``prompt_tps`` reflects the true cold rate over the
+  fresh ``Δ``-token suffix. We accept the warmup's reported rate as the
+  cold equivalent for ``pp=Δ`` (the warmup itself is the cold prefill).
+
+  ``decode TPS`` is independent of prefill mechanics — every step's
+  ``generation_tps`` is a real decode-rate-at-N data point.
+
+  Cumulative cold-prefill upper bound:
+      ``T_cum(pp_k) = Σ_{i=1..k} (Δ_i / prompt_tps_i)``
+
+  Optional cold-control points (``prefix_cache=disabled``) validate the
+  approximation; the gap quantifies per-task overhead. To preserve the
+  ``none`` cache-hit classification AND ensure the request actually
+  hits a freshly-placed runner (the master picks the instance with the
+  lowest in-flight task count, which is non-deterministic when multiple
+  same-model instances exist), :func:`run` deletes the sweep instance
+  *before* invoking the cold-control factory. The factory itself places
+  a fresh instance per control and deletes it on exit; the
+  :func:`bench.lib.cluster.managed_instance` ctx-manager calls
+  ``cleanup_all_instances`` on exit as a final safety net.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from collections.abc import Callable, Iterator
+from contextlib import AbstractContextManager, contextmanager
+from dataclasses import asdict, dataclass
+from typing import Any
+
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
+    Comm,
+    Sharding,
+    place_instance,
+    wait_for_instance_gone,
+)
+from loguru import logger
+
+from .completion import GenerationStats, PrefixCacheHit, run_one_completion
+from .prompt import PromptSizer
+from .results import ResultsBundle
+from .session import BenchSession
+
+
+@dataclass(frozen=True)
+class ContextScalingParams:
+    """Inputs for a single context-scaling sweep."""
+
+    pp_step: int
+    num_steps: int
+    tg: int
+    warmup: int = 1
+    cold_controls: tuple[int, ...] = ()
+    sleep_between_s: float = 1.0
+
+
+@dataclass
+class StepResult:
+    pp_tokens: int
+    delta_tokens: int
+    prompt_tps: float
+    generation_tps: float
+    prefix_cache_hit: PrefixCacheHit | str
+    prompt_tokens: int
+    generation_tokens: int
+    elapsed_s: float
+    peak_memory_bytes: int = 0
+    output_text_preview: str = ""
+
+
+def _peak_bytes(stats: GenerationStats) -> int:
+    pm = stats.get("peak_memory_usage") or {}
+    return int(pm.get("inBytes") or pm.get("in_bytes") or 0)
+
+
+def _build_step_result(
+    pp_tokens: int,
+    delta_tokens: int,
+    elapsed_s: float,
+    output_text_preview: str,
+    stats: GenerationStats,
+) -> StepResult:
+    return StepResult(
+        pp_tokens=pp_tokens,
+        delta_tokens=delta_tokens,
+        prompt_tps=float(stats.get("prompt_tps") or 0.0),
+        generation_tps=float(stats.get("generation_tps") or 0.0),
+        prefix_cache_hit=stats.get("prefix_cache_hit") or "unknown",
+        prompt_tokens=int(stats.get("prompt_tokens") or pp_tokens),
+        generation_tokens=int(stats.get("generation_tokens") or 0),
+        elapsed_s=elapsed_s,
+        peak_memory_bytes=_peak_bytes(stats),
+        output_text_preview=output_text_preview[:200],
+    )
+
+
+def _run_request(
+    client: ExoClient,
+    full_model_id: str,
+    pp: int,
+    tg: int,
+    sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool,
+) -> tuple[StepResult, int]:
+    """Send one request and return ``(StepResult, actual_pp_tokens)``."""
+    row, actual_pp = run_one_completion(
+        client,
+        full_model_id,
+        pp,
+        tg,
+        sizer,
+        use_prefix_cache=use_prefix_cache,
+        stream=False,
+    )
+    step = _build_step_result(
+        pp_tokens=actual_pp,
+        delta_tokens=actual_pp,  # caller overrides for cached sweep
+        elapsed_s=row["elapsed_s"],
+        output_text_preview=row["output_text_preview"],
+        stats=row["stats"],
+    )
+    return step, actual_pp
+
+
+def _compute_t_cum(steps: list[StepResult]) -> list[float]:
+    t_cum = 0.0
+    out: list[float] = []
+    for s in steps:
+        if s.prompt_tps > 0 and s.delta_tokens > 0:
+            t_cum += s.delta_tokens / s.prompt_tps
+        out.append(round(t_cum, 6))
+    return out
+
+
+def run_cached_sweep(
+    session: BenchSession,
+    params: ContextScalingParams,
+    bundle: ResultsBundle,
+) -> list[StepResult]:
+    """Run the ascending PP sweep with ``prefix_cache=enabled``.
+
+    Mutates ``bundle.runs`` in place and returns the typed step list.
+    """
+    if session.full_model_id is None:
+        raise RuntimeError(
+            "BenchSession.full_model_id must be set for context-scaling."
+        )
+
+    sizer = session.get_prompt_sizer()
+    client = session.client
+    pp_targets = [params.pp_step * i for i in range(1, params.num_steps + 1)]
+    logger.info(
+        f"context-scaling: K={params.num_steps} steps, Δ={params.pp_step} tokens, "
+        f"tg={params.tg}, warmup={params.warmup}, cached"
+    )
+
+    # Warmup discipline:
+    #   - First warmup runs with the prefix cache DISABLED. This triggers
+    #     the MLX kernel JIT compile + KV-buffer alloc for this exact
+    #     (Δ, dtype, batch) shape, but does NOT write a cache entry — so
+    #     the cold-with-JIT rate isn't fossilised.
+    #   - Subsequent warmups run with the prefix cache ENABLED. The
+    #     second one finds an empty cache, does a real cold prefill with
+    #     a HOT kernel, and writes the resulting rate into the cache
+    #     entry at pp=Δ.
+    #   - Step 0 (also cache-enabled) is then an exact hit on that entry
+    #     and reports the hot rate.
+    # Default warmup=2 gives both effects; warmup=1 still does the JIT
+    # warmup but leaves step 0 as a "none" hit (cold prefill at the hot
+    # kernel, creates the cache entry on the way through).
+    for w in range(params.warmup):
+        is_jit_warmup = w == 0
+        kind = "JIT warmup" if is_jit_warmup else "cache-prime warmup"
+        logger.info(
+            f"  warmup {w + 1}/{params.warmup} ({kind}, pp={params.pp_step})"
+        )
+        _run_request(
+            client,
+            session.full_model_id,
+            params.pp_step,
+            params.tg,
+            sizer,
+            use_prefix_cache=not is_jit_warmup,
+        )
+
+    steps: list[StepResult] = []
+    prev_pp = 0
+    for i, pp in enumerate(pp_targets):
+        time.sleep(params.sleep_between_s)
+        try:
+            step, actual_pp = _run_request(
+                client,
+                session.full_model_id,
+                pp,
+                params.tg,
+                sizer,
+                use_prefix_cache=True,
+            )
+        except Exception as e:
+            logger.error(f"step {i + 1}/{params.num_steps} (pp={pp}) failed: {e}")
+            raise
+
+        step.delta_tokens = actual_pp - prev_pp
+        steps.append(step)
+        bundle.runs.append({"step_index": i, "phase": "cached_sweep", **asdict(step)})
+        logger.info(
+            f"  step {i + 1}/{params.num_steps} pp={actual_pp} Δ={step.delta_tokens} "
+            f"prompt_tps={step.prompt_tps:.1f} gen_tps={step.generation_tps:.2f} "
+            f"hit={step.prefix_cache_hit}"
+        )
+        prev_pp = actual_pp
+
+    return steps
+
+
+def run_cold_controls(
+    factory: Callable[[], AbstractContextManager[ExoClient]],
+    session: BenchSession,
+    params: ContextScalingParams,
+    bundle: ResultsBundle,
+) -> list[StepResult]:
+    """Run cold-control points on a fresh instance to preserve ``none`` hits.
+
+    A cold control is a single request at ``pp=N`` with
+    ``prefix_cache=disabled``, executed against a freshly-placed instance
+    (and with no other same-model instance live, so the master's task
+    routing is deterministic). The caller is expected to delete the
+    sweep instance before invoking this — see :func:`run`.
+    """
+    if not params.cold_controls:
+        return []
+    if session.full_model_id is None:
+        raise RuntimeError("BenchSession.full_model_id must be set for cold controls.")
+
+    sizer = session.get_prompt_sizer()
+    out: list[StepResult] = []
+    for control_pp in params.cold_controls:
+        logger.info(f"cold control: pp={control_pp} (fresh instance, cache disabled)")
+        with factory() as fresh_client:
+            step, actual_pp = _run_request(
+                fresh_client,
+                session.full_model_id,
+                control_pp,
+                params.tg,
+                sizer,
+                use_prefix_cache=False,
+            )
+        step.delta_tokens = actual_pp
+        out.append(step)
+        bundle.cold_controls.append({"phase": "cold_control", **asdict(step)})
+        logger.info(
+            f"  cold pp={actual_pp} prompt_tps={step.prompt_tps:.1f} "
+            f"gen_tps={step.generation_tps:.2f} hit={step.prefix_cache_hit}"
+        )
+        if step.prefix_cache_hit != "none":
+            logger.warning(
+                f"cold control at pp={actual_pp} reported "
+                f"prefix_cache_hit={step.prefix_cache_hit!r}; "
+                f"control may not be cold."
+            )
+    return out
+
+
+def derive_summary(
+    steps: list[StepResult],
+    cold_controls: list[StepResult],
+) -> dict[str, Any]:
+    """Compute the cumulative cold-prefill upper bound + control gaps."""
+    t_cum = _compute_t_cum(steps)
+    bracketed = sorted(
+        ((s.pp_tokens, t) for s, t in zip(steps, t_cum, strict=True)),
+        key=lambda x: x[0],
+    )
+    control_gaps: list[dict[str, float]] = []
+    for ctrl in cold_controls:
+        cold_t = ctrl.pp_tokens / ctrl.prompt_tps if ctrl.prompt_tps > 0 else 0.0
+        cum_t = _interp(bracketed, ctrl.pp_tokens)
+        gap = cum_t - cold_t
+        control_gaps.append(
+            {
+                "pp_tokens": ctrl.pp_tokens,
+                "cold_t_seconds": round(cold_t, 4),
+                "t_cum_seconds_at_pp": round(cum_t, 4),
+                "gap_seconds": round(gap, 4),
+                "gap_fraction": round(gap / cold_t, 4) if cold_t > 0 else 0.0,
+            }
+        )
+    return {
+        "t_cum_seconds": t_cum,
+        "control_gaps": control_gaps,
+    }
+
+
+def _interp(points: list[tuple[int, float]], x: int) -> float:
+    """Linear interpolate y at x, given sorted ``(x, y)`` points."""
+    if not points:
+        return 0.0
+    if x <= points[0][0]:
+        return points[0][1]
+    if x >= points[-1][0]:
+        return points[-1][1]
+    for i in range(1, len(points)):
+        x0, y0 = points[i - 1]
+        x1, y1 = points[i]
+        if x0 <= x <= x1 and x1 != x0:
+            return y0 + (y1 - y0) * (x - x0) / (x1 - x0)
+    return points[-1][1]
+
+
+# ---------------------------------------------------------------------------
+# Cold-control instance factory
+# ---------------------------------------------------------------------------
+
+
+def make_cold_control_factory(
+    session: BenchSession,
+    sharding: Sharding,
+    comm: Comm,
+    min_nodes: int,
+    instance_timeout_s: float = 1800.0,
+) -> Callable[[], AbstractContextManager[ExoClient]]:
+    """Return a callable yielding a context manager that places a fresh instance.
+
+    Each ``with factory() as client:`` block places a brand-new instance,
+    yields its client, then deletes the instance on exit. Used to isolate
+    cold-control runs.
+
+    The caller is responsible for ensuring no other same-model instance is
+    live during the ``with`` block — otherwise master routing is
+    non-deterministic and the cold control may be served by a stale runner.
+    See :func:`run` for the orchestration.
+    """
+
+    @contextmanager
+    def factory() -> Iterator[ExoClient]:
+        if session.full_model_id is None:
+            raise RuntimeError("session.full_model_id is unset")
+        client = session.client
+        instance_id = place_instance(
+            client,
+            session.full_model_id,
+            sharding=sharding,
+            comm=comm,
+            min_nodes=min_nodes,
+            timeout=instance_timeout_s,
+        )
+        try:
+            yield client
+        finally:
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            with contextlib.suppress(Exception):
+                wait_for_instance_gone(client, instance_id, timeout=60.0)
+
+    return factory
+
+
+def _delete_instance(client: ExoClient, instance_id: str) -> None:
+    """Best-effort delete of a placed instance."""
+    with contextlib.suppress(ExoHttpError):
+        client.request_json("DELETE", f"/instance/{instance_id}")
+    with contextlib.suppress(Exception):
+        wait_for_instance_gone(client, instance_id, timeout=60.0)
+
+
+def run(
+    session: BenchSession,
+    params: ContextScalingParams,
+    bundle: ResultsBundle,
+    *,
+    cold_control_factory: Callable[[], AbstractContextManager[ExoClient]] | None = None,
+) -> ResultsBundle:
+    """End-to-end: cached sweep + optional cold controls + derived summary.
+
+    To make cold controls truly isolated from the sweep instance, we
+    delete the sweep instance *before* running the controls (otherwise
+    the master might route a control's request to the stale sweep
+    instance, since both match the same ``model_id``). The controls then
+    each place their own fresh instance via the factory.
+    """
+    bundle.params.update(
+        {
+            "pp_step": params.pp_step,
+            "num_steps": params.num_steps,
+            "tg": params.tg,
+            "warmup": params.warmup,
+            "cold_controls": list(params.cold_controls),
+            "sleep_between_s": params.sleep_between_s,
+            "model_id": session.model_id,
+            "full_model_id": session.full_model_id,
+        }
+    )
+    bundle.capture_cluster(session.client)
+
+    cached_steps = run_cached_sweep(session, params, bundle)
+
+    cold_steps: list[StepResult] = []
+    if params.cold_controls and cold_control_factory is not None:
+        # Delete the sweep instance so the cold-control fresh instance is
+        # the only same-model instance live for the duration of the controls.
+        if session.instance_id is not None:
+            logger.info(
+                f"cold controls: deleting sweep instance {session.instance_id} "
+                "to isolate fresh instance routing"
+            )
+            _delete_instance(session.client, session.instance_id)
+            session.instance_id = None
+        cold_steps = run_cold_controls(cold_control_factory, session, params, bundle)
+    elif params.cold_controls and cold_control_factory is None:
+        logger.warning(
+            "Cold controls requested but no cold_control_factory supplied; skipping."
+        )
+
+    bundle.derived.update(derive_summary(cached_steps, cold_steps))
+    return bundle

bench/lib/model_meta.py

@@ -0,0 +1,183 @@
+"""Fetch HuggingFace model metadata for benchmark planning.
+
+Two pieces of metadata drive every benchmark we run:
+
+  1. **Total weight size** — used to derive ``min-memory`` and ``min-disk``
+     constraints when picking a host. We sum the sizes of all
+     ``.safetensors`` (or ``.bin``) shards from the repo's file listing.
+  2. **Max position embeddings** — the model's training context length.
+     Used to bound a context-scaling sweep at the model's max context, and
+     to derive a sensible Δ given a target step count.
+
+The fetcher uses the ``huggingface_hub`` python API, which talks to the
+public HF Hub HTTPS endpoints — no exo cluster required, no download
+of weights.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+# Files that count toward the on-disk weight footprint.
+_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".gguf", ".pt", ".npz")
+
+
+@dataclass(frozen=True)
+class ModelMeta:
+    """Subset of HF metadata that a benchmark needs."""
+
+    model_id: str
+    total_weight_bytes: int
+    max_position_embeddings: int
+    num_hidden_layers: int
+    raw_config: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def total_weight_gb(self) -> float:
+        return self.total_weight_bytes / (1024**3)
+
+    @property
+    def memory_constraint_gb(self) -> float:
+        """Estimated minimum host memory to hold weights + overhead.
+
+        Picks the model size + 30 % headroom (KV cache, activations,
+        framework bookkeeping). Rounded up to the next whole GiB.
+        """
+        return float(int(self.total_weight_gb * 1.30) + 1)
+
+    @property
+    def disk_constraint_gb(self) -> float:
+        """Disk space the host must have free for the download."""
+        return float(int(self.total_weight_gb * 1.10) + 1)
+
+
+def _read_config_json(model_id: str) -> dict[str, Any]:
+    from huggingface_hub import (
+        hf_hub_download,  # type: ignore[reportUnknownVariableType]
+    )
+
+    raw_path = hf_hub_download(repo_id=model_id, filename="config.json", dry_run=False)
+    with open(raw_path) as f:
+        loaded: Any = json.load(f)  # type: ignore[reportAny]
+    return cast("dict[str, Any]", loaded) if isinstance(loaded, dict) else {}
+
+
+def _sum_weight_sizes(model_id: str) -> int:
+    """Sum sizes of all weight-shard files in the repo's file listing."""
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+    info = api.model_info(repo_id=model_id, files_metadata=True)
+    siblings = info.siblings or []
+    total = 0
+    for sib in siblings:
+        rfilename = getattr(sib, "rfilename", None)
+        size = getattr(sib, "size", None)
+        if not isinstance(rfilename, str) or not isinstance(size, int):
+            continue
+        if any(rfilename.endswith(suf) for suf in _WEIGHT_SUFFIXES):
+            total += size
+    return total
+
+
+def _first_int(config: dict[str, Any], *keys: str) -> int:
+    """Return the first key from ``config`` that holds a usable positive int."""
+    for key in keys:
+        value = config.get(key)
+        if isinstance(value, int) and value > 0:
+            return value
+        if isinstance(value, str):
+            try:
+                parsed = int(value)
+            except ValueError:
+                continue
+            if parsed > 0:
+                return parsed
+    return 0
+
+
+def fetch_model_meta(model_id: str) -> ModelMeta:
+    """Fetch the metadata our benchmarks care about for ``model_id``.
+
+    Args:
+        model_id: HuggingFace repo id, e.g. ``mlx-community/Qwen3-30B-A3B-4bit``.
+
+    Returns:
+        Populated :class:`ModelMeta`.
+
+    Raises:
+        Exception: any HTTP / parse error from ``huggingface_hub`` propagates.
+    """
+    config = _read_config_json(model_id)
+    return ModelMeta(
+        model_id=model_id,
+        total_weight_bytes=_sum_weight_sizes(model_id),
+        max_position_embeddings=_first_int(
+            config,
+            "max_position_embeddings",
+            "max_seq_len",
+            "model_max_length",
+            "n_positions",
+        ),
+        num_hidden_layers=_first_int(
+            config,
+            "num_hidden_layers",
+            "num_layers",
+            "n_layer",
+            "n_layers",
+            "num_decoder_layers",
+        ),
+        raw_config=config,
+    )
+
+
+def derive_context_ramp(
+    meta: ModelMeta,
+    *,
+    num_steps: int,
+    fraction_of_max: float = 1.0,
+    min_pp_step: int = 256,
+    round_to: int = 256,
+) -> tuple[int, int]:
+    """Pick ``(pp_step, num_steps)`` covering ``fraction_of_max`` of the context.
+
+    Δ is rounded down to the nearest ``round_to`` so the per-step prompt is a
+    clean number, and clamped to ``min_pp_step`` for tiny-context models.
+    """
+    if meta.max_position_embeddings <= 0:
+        raise ValueError(
+            f"{meta.model_id} reports max_position_embeddings=0 in config.json"
+        )
+    if not (0.0 < fraction_of_max <= 1.0):
+        raise ValueError(f"fraction_of_max must be in (0, 1], got {fraction_of_max}")
+    if num_steps <= 0:
+        raise ValueError(f"num_steps must be >0, got {num_steps}")
+
+    target_max = int(meta.max_position_embeddings * fraction_of_max)
+    raw_step = max(min_pp_step, target_max // num_steps)
+    pp_step = (raw_step // round_to) * round_to or round_to
+    return pp_step, num_steps
+
+
+def derive_cold_controls(
+    meta: ModelMeta,
+    *,
+    pp_step: int,
+    num_steps: int,
+    count: int = 4,
+) -> tuple[int, ...]:
+    """Pick ``count`` evenly-spaced cold-control points across the ramp.
+
+    Always includes the largest ramp point (``pp_step * num_steps``).
+    Returns control pp values in ascending order, deduped.
+    """
+    if count <= 0:
+        return ()
+    max_pp = pp_step * num_steps
+    if count == 1:
+        return (max_pp,)
+    spaced = sorted({(max_pp * (i + 1)) // count for i in range(count)})
+    # Filter out anything below pp_step (a control at <Δ is meaningless).
+    return tuple(p for p in spaced if p >= pp_step)

bench/lib/plotting.py

@@ -0,0 +1,308 @@
+"""Typed matplotlib renderers for benchmark JSON results.
+
+This module owns the *visualisation* of bench results, mirroring how
+``bench/lib/<name>.py`` owns the methodology and ``bench/cli/<name>.py``
+owns the orchestration. Adding plotting for a new benchmark = a new
+``render_<name>`` function here + a dispatch entry in ``bench/cli/plot.py``.
+
+Functions take typed inputs (``Path`` lists, options) and write a PNG.
+They never touch argparse or stdout — that's the CLI's job.
+
+matplotlib's type stubs are thin (most return values are ``Any``), so all
+calls into ``pyplot`` are concentrated at the bottom of this file with
+targeted ``# type: ignore[reportUnknownMemberType, reportAny]`` per line.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+# Tab10 cycle from matplotlib's default; we pick colours by index ourselves
+# instead of fishing them out of `Line2D.get_color()` so the strict-type
+# fallout stays small and predictable.
+_COLOR_CYCLE: tuple[str, ...] = (
+    "C0",
+    "C1",
+    "C2",
+    "C3",
+    "C4",
+    "C5",
+    "C6",
+    "C7",
+    "C8",
+    "C9",
+)
+
+
+@dataclass(frozen=True)
+class PlotInputs:
+    """Inputs for any benchmark renderer.
+
+    Attributes:
+        results: One or more bench JSON files. The first is used to
+            auto-derive the title when ``title`` is unset.
+        output: Path to write the PNG to.
+        label_tag: When set, use ``metadata.tags[label_tag]`` as the
+            legend label for each run; otherwise use the run id.
+        title: Override for the figure title.
+    """
+
+    results: list[Path]
+    output: Path
+    label_tag: str | None = None
+    title: str | None = None
+
+
+@dataclass(frozen=True)
+class _RunSeries:
+    """Pre-extracted plot data for one results JSON.
+
+    ``cached_prefill_seconds`` is the cumulative cold-prefill estimate
+    (``T_cum`` from the methodology — read from ``derived.t_cum_seconds``).
+    ``control_prefill_seconds`` is the actual cold prefill time per
+    control (``pp_tokens / prompt_tps`` from the cold-control row).
+    """
+
+    label: str
+    cached_pp: list[int]
+    cached_prefill_seconds: list[float]
+    cached_gen_tps: list[float]
+    control_pp: list[int]
+    control_prefill_seconds: list[float]
+
+
+# ---------------------------------------------------------------------------
+# Pure data extraction (strict-typed, no matplotlib)
+# ---------------------------------------------------------------------------
+
+
+def _load(path: Path) -> dict[str, Any]:
+    """Read a bench JSON file and assert top-level shape."""
+    with path.open() as f:
+        loaded: Any = json.load(f)  # type: ignore[reportAny]
+    if not isinstance(loaded, dict):
+        raise ValueError(f"{path}: expected top-level JSON object")
+    return cast("dict[str, Any]", loaded)
+
+
+# dict[str, Any].get(...) returns Any. The five _get_* helpers below
+# concentrate the Any boundary so the rest of the module can be strict.
+
+
+def _get_dict(d: dict[str, Any], key: str) -> dict[str, Any]:
+    val: Any = d.get(key)
+    return cast("dict[str, Any]", val) if isinstance(val, dict) else {}
+
+
+def _get_list(d: dict[str, Any], key: str) -> list[Any]:
+    val: Any = d.get(key)
+    return cast("list[Any]", val) if isinstance(val, list) else []
+
+
+def _get_str(d: dict[str, Any], key: str, default: str = "") -> str:
+    val: Any = d.get(key, default)  # type: ignore[reportAny]
+    return val if isinstance(val, str) else default
+
+
+def _get_int(row: dict[str, Any], key: str) -> int:
+    val: Any = row.get(key, 0)  # type: ignore[reportAny]
+    if isinstance(val, bool):  # bool is int; reject explicitly
+        return 0
+    if isinstance(val, (int, float)):
+        return int(val)
+    if isinstance(val, str):
+        try:
+            return int(float(val))
+        except ValueError:
+            return 0
+    return 0
+
+
+def _get_float(row: dict[str, Any], key: str) -> float:
+    val: Any = row.get(key, 0.0)  # type: ignore[reportAny]
+    if isinstance(val, bool):
+        return 0.0
+    if isinstance(val, (int, float)):
+        return float(val)
+    if isinstance(val, str):
+        try:
+            return float(val)
+        except ValueError:
+            return 0.0
+    return 0.0
+
+
+def _label_for(data: dict[str, Any], label_tag: str | None) -> str:
+    if label_tag is not None:
+        tags = _get_dict(_get_dict(data, "metadata"), "tags")
+        if label_tag in tags:
+            return _get_str(tags, label_tag, "(unnamed)")
+    return _get_str(_get_dict(data, "metadata"), "run_id", "(unnamed)")
+
+
+def _extract_series(data: dict[str, Any], label: str) -> _RunSeries:
+    """Pre-extract typed lists from a context-scaling bench JSON."""
+    cached_pp: list[int] = []
+    cached_gen_tps: list[float] = []
+    for raw in _get_list(data, "runs"):  # type: ignore[reportAny]
+        if not isinstance(raw, dict):
+            continue
+        row = cast("dict[str, Any]", raw)
+        if _get_str(row, "phase") != "cached_sweep":
+            continue
+        cached_pp.append(_get_int(row, "pp_tokens"))
+        cached_gen_tps.append(_get_float(row, "generation_tps"))
+
+    # Cumulative cold-prefill estimate is computed in derive_summary and
+    # written to derived.t_cum_seconds (parallel to the cached steps).
+    derived = _get_dict(data, "derived")
+    t_cum_raw = _get_list(derived, "t_cum_seconds")
+    cached_prefill_seconds: list[float] = []
+    for raw in t_cum_raw:  # type: ignore[reportAny]
+        if isinstance(raw, (int, float)) and not isinstance(raw, bool):
+            cached_prefill_seconds.append(float(raw))
+
+    # Cold controls give us the actual cold prefill time directly:
+    # pp_tokens / prompt_tps. Skip rows with zero/missing prompt_tps.
+    control_pp: list[int] = []
+    control_prefill_seconds: list[float] = []
+    for raw in _get_list(data, "cold_controls"):  # type: ignore[reportAny]
+        if not isinstance(raw, dict):
+            continue
+        row = cast("dict[str, Any]", raw)
+        pp = _get_int(row, "pp_tokens")
+        tps = _get_float(row, "prompt_tps")
+        if pp > 0 and tps > 0:
+            control_pp.append(pp)
+            control_prefill_seconds.append(pp / tps)
+
+    return _RunSeries(
+        label=label,
+        cached_pp=cached_pp,
+        cached_prefill_seconds=cached_prefill_seconds,
+        cached_gen_tps=cached_gen_tps,
+        control_pp=control_pp,
+        control_prefill_seconds=control_prefill_seconds,
+    )
+
+
+def _auto_title(data: dict[str, Any]) -> str:
+    metadata = _get_dict(data, "metadata")
+    params = _get_dict(data, "params")
+    model = (
+        _get_str(params, "full_model_id") or _get_str(params, "model_id") or "(unknown)"
+    )
+    sha = _get_str(metadata, "exo_sha") or "(no-sha)"
+    host = _get_str(metadata, "hostname") or "(no-host)"
+    return f"{model}\n{sha} on {host}"
+
+
+# ---------------------------------------------------------------------------
+# Matplotlib boundary — each call site has a narrow, justified ignore.
+# ---------------------------------------------------------------------------
+
+
+def render_context_scaling(inputs: PlotInputs) -> Path:
+    """Render a 2-panel context-scaling plot.
+
+    Top:    pp_tokens vs prompt_tps (line per run; cold controls as 'x' scatter)
+    Bottom: pp_tokens vs generation_tps (line per run)
+
+    Each line is a separate result file. Multi-file mode is for comparing
+    runs across exo SHAs / hosts / configs; the title is taken from the
+    first file's metadata unless ``inputs.title`` is set.
+    """
+    if not inputs.results:
+        raise ValueError("at least one results JSON path is required")
+
+    # Validate + extract first so any data-shape error surfaces before we
+    # even import matplotlib.
+    first_data: dict[str, Any] | None = None
+    series: list[_RunSeries] = []
+    for path in inputs.results:
+        data = _load(path)
+        if first_data is None:
+            first_data = data
+        benchmark = _get_str(_get_dict(data, "metadata"), "benchmark")
+        if benchmark != "context_scaling":
+            raise ValueError(
+                f"{path}: expected benchmark=='context_scaling', got {benchmark!r}"
+            )
+        series.append(_extract_series(data, _label_for(data, inputs.label_tag)))
+
+    title = inputs.title
+    if title is None and first_data is not None:
+        title = _auto_title(first_data)
+        if len(inputs.results) > 1:
+            title = f"{title}\n(comparison of {len(inputs.results)} runs)"
+
+    inputs.output.parent.mkdir(parents=True, exist_ok=True)
+    _draw(series, inputs.output, title=title)
+    return inputs.output
+
+
+def _draw(series: list[_RunSeries], output: Path, *, title: str | None) -> None:
+    """Concentrated matplotlib boundary."""
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, axes = plt.subplots(  # type: ignore[reportUnknownMemberType]
+        2, 1, figsize=(10, 8), sharex=True
+    )
+    top: Any = axes[0]  # type: ignore[reportAny]
+    bottom: Any = axes[1]  # type: ignore[reportAny]
+
+    for i, run in enumerate(series):
+        color = _COLOR_CYCLE[i % len(_COLOR_CYCLE)]
+        # Cumulative cold-prefill estimate (T_cum). Only plot points where
+        # we have a t_cum value — skip if derived was empty for this run.
+        n = min(len(run.cached_pp), len(run.cached_prefill_seconds))
+        if n > 0:
+            top.plot(  # type: ignore[reportAny, reportUnknownMemberType]
+                run.cached_pp[:n],
+                run.cached_prefill_seconds[:n],
+                "-o",
+                color=color,
+                label=run.label,
+            )
+        if run.control_pp:
+            top.scatter(  # type: ignore[reportAny, reportUnknownMemberType]
+                run.control_pp,
+                run.control_prefill_seconds,
+                marker="x",
+                s=80,
+                color=color,
+                label=f"{run.label} (cold one-shot)",
+            )
+        bottom.plot(  # type: ignore[reportAny, reportUnknownMemberType]
+            run.cached_pp,
+            run.cached_gen_tps,
+            "-o",
+            color=color,
+            label=run.label,
+        )
+
+    top.set_ylabel("prefill time (s)")  # type: ignore[reportAny, reportUnknownMemberType]
+    top.set_title(  # type: ignore[reportAny, reportUnknownMemberType]
+        "cumulative cold-prefill time vs context size "
+        "(line: T_cum estimate; ✕: cold one-shot control)"
+    )
+    top.grid(True, alpha=0.3)  # type: ignore[reportAny, reportUnknownMemberType]
+    top.legend(loc="best", fontsize=8)  # type: ignore[reportAny, reportUnknownMemberType]
+
+    bottom.set_xlabel("pp_tokens")  # type: ignore[reportAny, reportUnknownMemberType]
+    bottom.set_ylabel("generation_tps (tok/s)")  # type: ignore[reportAny, reportUnknownMemberType]
+    bottom.set_title("decode throughput vs context size")  # type: ignore[reportAny, reportUnknownMemberType]
+    bottom.grid(True, alpha=0.3)  # type: ignore[reportAny, reportUnknownMemberType]
+
+    if title is not None:
+        fig.suptitle(title, fontsize=10)  # type: ignore[reportUnknownMemberType]
+
+    fig.tight_layout()
+    fig.savefig(output, dpi=120, bbox_inches="tight")  # type: ignore[reportUnknownMemberType]
+    plt.close(fig)

bench/lib/prompt.py

@@ -0,0 +1,269 @@
+"""Typed prompt-sizing utilities for benchmarks.
+
+Wraps the HuggingFace ``transformers`` tokenizer (a fundamentally dynamic
+object — different models return different types from
+``apply_chat_template``) behind a small typed API so the rest of the bench
+library can stay strict-typed.
+
+``PromptSizer.build(target)`` returns a ``(content, exact_token_count)``
+pair. Internally it:
+  1. Tokenises the empty user message to learn the chat-template overhead
+     (``base_tokens``).
+  2. Estimates tokens-per-atom from a 100-atom sample.
+  3. Binary-searches over the atom count so the resulting message
+     tokenises to *exactly* ``target`` tokens.
+
+Callers downstream (``run_one_completion`` etc.) receive the verified
+token count, so analysis can confirm the prompt hit its target.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+import types
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Final, cast
+
+
+def _coerce_token_ids(raw: object) -> list[int]:
+    """Normalise ``apply_chat_template`` output to a flat list of token ids.
+
+    transformers' ``apply_chat_template`` may return:
+      - ``list[int]`` (slow tokenizers, ``tokenize=True``)
+      - a ``BatchEncoding`` with ``.input_ids`` (fast tokenizers)
+      - a tensor wrapped object (some models)
+
+    We only need ``len(.)`` of the result, so we just need to flatten to a
+    list and return it.
+    """
+    if isinstance(raw, list):
+        return cast("list[int]", raw)
+    input_ids = getattr(raw, "input_ids", None)
+    if isinstance(input_ids, list):
+        return cast("list[int]", input_ids)
+    raise TypeError(
+        f"Unsupported tokenizer output type {type(raw).__name__}; "
+        "expected list[int] or BatchEncoding-like with .input_ids."
+    )
+
+
+def _build_token_counter(tokenizer: object) -> Callable[[str], int]:
+    """Return a closure that counts tokens for a user message.
+
+    Tries ``apply_chat_template`` first; falls back to the DeepSeek-V4
+    Python encoder for models that don't ship a Jinja chat template.
+    """
+    apply_chat_template = cast(
+        Callable[..., object],
+        tokenizer.apply_chat_template,  # type: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+    )
+    encode = cast(
+        Callable[..., list[int]],
+        tokenizer.encode,  # type: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+    )
+
+    def count_fn(user_content: str) -> int:
+        messages = [{"role": "user", "content": user_content}]
+        try:
+            raw = apply_chat_template(
+                messages, tokenize=True, add_generation_prompt=True
+            )
+        except ValueError:
+            # Models without a Jinja chat template (e.g. DeepSeek V4 which
+            # ships its own Python encoder). Use the exo-side V4 encoder.
+            from exo.worker.engines.mlx.vendor.deepseek_v4_encoding import (  # type: ignore[reportMissingTypeStubs]
+                encode_messages as encode_v4,
+            )
+
+            prompt = cast(str, encode_v4(messages, thinking_mode="thinking"))  # type: ignore[reportUnknownArgumentType]
+            raw = encode(prompt, add_special_tokens=False)
+        return len(_coerce_token_ids(raw))
+
+    return count_fn
+
+
+class PromptSizer:
+    """Build a chat-completion content string of an exact token length."""
+
+    DEFAULT_ATOM: Final[str] = "a "
+
+    def __init__(self, tokenizer: object, atom: str = DEFAULT_ATOM):
+        self._tokenizer = tokenizer
+        self.atom = atom
+        self._count_fn = _build_token_counter(tokenizer)
+        self.base_tokens = self._count_fn("")
+
+    def count(self, content: str) -> int:
+        """Return the token count for ``content`` after chat-template expansion."""
+        return self._count_fn(content)
+
+    def build(self, target_prompt_tokens: int) -> tuple[str, int]:
+        """Return ``(content, exact_token_count)`` summing to ``target``.
+
+        Raises ``RuntimeError`` if the chosen ``atom`` overshoots the target
+        (try a different atom — see ``DEFAULT_ATOM``).
+        """
+        target = int(target_prompt_tokens)
+        if target < self.base_tokens:
+            raise RuntimeError(
+                f"Target ({target}) is smaller than template overhead "
+                f"({self.base_tokens})."
+            )
+
+        # Estimate tokens per atom using a sample.
+        sample_count = 100
+        sample_tokens = self._count_fn(self.atom * sample_count) - self.base_tokens
+        tokens_per_atom = sample_tokens / sample_count
+        needed_tokens = target - self.base_tokens
+        estimated_atoms = int(needed_tokens / tokens_per_atom)
+
+        # Binary search to find exact atom count.
+        low, high = 0, estimated_atoms * 2 + 100
+        while low < high:
+            mid = (low + high) // 2
+            if self._count_fn(self.atom * mid) < target:
+                low = mid + 1
+            else:
+                high = mid
+
+        content = self.atom * low
+        actual = self._count_fn(content)
+        if actual != target:
+            raise RuntimeError(
+                f"Overshot: got {actual} tokens (target {target}). "
+                f"Pick a different atom (try ' a' or '\\n' or '0 ')."
+            )
+        return content, actual
+
+
+def _load_kimi_tokenizer(model_id: str) -> object:
+    """Special-case Kimi K2's custom TikTokenTokenizer (transformers 5.x quirk)."""
+    from huggingface_hub import (
+        snapshot_download,  # type: ignore[reportUnknownVariableType]
+    )
+
+    raw_path = snapshot_download(
+        model_id,
+        allow_patterns=[
+            "*.json",
+            "*.py",
+            "*.tiktoken",
+            "*.model",
+            "*.jinja",
+        ],
+        dry_run=False,
+    )
+    model_path = Path(raw_path)
+    sys.path.insert(0, str(model_path))
+
+    tool_decl_path = model_path / "tool_declaration_ts.py"
+    if tool_decl_path.exists():
+        spec = importlib.util.spec_from_file_location(
+            "tool_declaration_ts", tool_decl_path
+        )
+        if spec is not None and spec.loader is not None:
+            tool_decl_module = importlib.util.module_from_spec(spec)
+            sys.modules["tool_declaration_ts"] = tool_decl_module
+            spec.loader.exec_module(tool_decl_module)
+
+    tok_path = model_path / "tokenization_kimi.py"
+    source = tok_path.read_text().replace(
+        "from .tool_declaration_ts", "from tool_declaration_ts"
+    )
+    tok_module = types.ModuleType("tokenization_kimi")
+    tok_module.__file__ = str(tok_path)
+    sys.modules["tokenization_kimi"] = tok_module
+    exec(compile(source, str(tok_path), "exec"), tok_module.__dict__)  # noqa: S102
+
+    tik_token_cls = cast(Any, tok_module).TikTokenTokenizer  # type: ignore[reportAny]
+    hf_tokenizer = cast(Any, tik_token_cls.from_pretrained(model_path))  # type: ignore[reportAny]
+
+    # Patch encode to use internal tiktoken model directly (transformers 5.x
+    # bug in the encode→pad path for slow tokenizers).
+    def _patched_encode(text: str, **_kwargs: object) -> list[int]:
+        return list(
+            hf_tokenizer.model.encode(text, allowed_special="all")  # type: ignore[reportAny, reportUnknownMemberType]
+        )
+
+    hf_tokenizer.encode = _patched_encode
+    return cast(object, hf_tokenizer)
+
+
+def load_tokenizer_for_bench(model_id: str) -> object:
+    """Load a HuggingFace tokenizer with bench-specific compatibility shims.
+
+    Returns the tokenizer as ``object`` because transformers' types are
+    fundamentally dynamic (concrete class depends on the model). Callers
+    should pass the result straight to :class:`PromptSizer`.
+    """
+    # Monkey-patch for transformers 5.x: Kimi's tokenization_kimi.py imports
+    # bytes_to_unicode from gpt2_tokenization which moved.
+    try:
+        import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
+        from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+        if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
+            gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[reportAttributeAccessIssue]
+    except ImportError:
+        pass
+
+    if "kimi-k2" in model_id.lower():
+        return _load_kimi_tokenizer(model_id)
+
+    from transformers import AutoTokenizer
+
+    try:
+        return cast(
+            object,
+            AutoTokenizer.from_pretrained(model_id, trust_remote_code=True),  # type: ignore[reportUnknownMemberType]
+        )
+    except (AttributeError, ValueError):
+        # Some models ship a Jinja template / encoder that AutoTokenizer
+        # can't introspect from HF directly — download artefacts and load
+        # from the local snapshot path.
+        from huggingface_hub import (
+            snapshot_download,  # type: ignore[reportUnknownVariableType]
+        )
+        from transformers import PretrainedConfig
+
+        raw_full_path = snapshot_download(
+            model_id,
+            allow_patterns=[
+                "*.json",
+                "*.py",
+                "tokenizer.model",
+                "*.tiktoken",
+                "tiktoken.model",
+                "*.txt",
+                "*.jsonl",
+                "*.jinja",
+            ],
+            dry_run=False,
+        )
+        model_path = Path(raw_full_path)
+        stub_kwargs: dict[str, Any] = {}
+        config_file = model_path / "config.json"
+        if config_file.exists():
+            with config_file.open() as f:
+                raw_config: dict[str, Any] = json.load(f)  # type: ignore[reportAny]
+            for key in (
+                "model_type",
+                "max_position_embeddings",
+                "vocab_size",
+                "bos_token_id",
+                "eos_token_id",
+                "pad_token_id",
+            ):
+                if key in raw_config:
+                    stub_kwargs[key] = raw_config[key]
+        return cast(
+            object,
+            AutoTokenizer.from_pretrained(  # type: ignore[reportUnknownMemberType]
+                str(model_path),
+                config=PretrainedConfig(**stub_kwargs),  # type: ignore[reportArgumentType, reportAny]
+                trust_remote_code=True,
+            ),
+        )

bench/lib/results.py

@@ -0,0 +1,139 @@
+"""Structured benchmark results — metadata capture + JSON output.
+
+Every benchmark run produces a single JSON file with a stable schema:
+
+- ``metadata``: exo SHA, ISO timestamps, hostnames, and any user-supplied
+  tags identifying the run.
+- ``cluster``: snapshot from the API (node identities, topology, memory).
+- ``params``: the benchmark's input parameters (sweep config, etc).
+- ``runs``: per-request result rows.
+- ``derived``: any computed summaries (``t_cum_seconds`` for context scaling).
+
+The format is intentionally additive so downstream tooling (plot scripts,
+dashboards) can rely on optional fields being absent rather than malformed.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import socket
+import subprocess
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from exo_tools.client import ExoClient
+from exo_tools.harness import capture_cluster_snapshot
+
+
+def _git_describe(repo_root: Path) -> str | None:
+    """Return ``<short-sha>[-dirty]`` for the repo at ``repo_root`` or None."""
+    try:
+        sha = subprocess.run(
+            ["git", "rev-parse", "--short=12", "HEAD"],
+            cwd=str(repo_root),
+            capture_output=True,
+            text=True,
+            timeout=5,
+            check=True,
+        ).stdout.strip()
+    except (
+        subprocess.CalledProcessError,
+        FileNotFoundError,
+        subprocess.TimeoutExpired,
+    ):
+        return None
+    try:
+        dirty = subprocess.run(
+            ["git", "status", "--porcelain"],
+            cwd=str(repo_root),
+            capture_output=True,
+            text=True,
+            timeout=5,
+            check=True,
+        ).stdout.strip()
+        return f"{sha}-dirty" if dirty else sha
+    except (
+        subprocess.CalledProcessError,
+        FileNotFoundError,
+        subprocess.TimeoutExpired,
+    ):
+        return sha
+
+
+@dataclass
+class RunMetadata:
+    """Identifies a single bench run."""
+
+    run_id: str
+    benchmark: str
+    started_at: str
+    finished_at: str | None = None
+    exo_sha: str | None = None
+    hostname: str = ""
+    platform: str = ""
+    tags: dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def new(
+        cls,
+        benchmark: str,
+        repo_root: Path,
+        *,
+        tags: dict[str, str] | None = None,
+    ) -> RunMetadata:
+        now = datetime.now(timezone.utc)
+        run_id = f"{benchmark}_{now.strftime('%Y%m%dT%H%M%SZ')}_{os.getpid()}"
+        return cls(
+            run_id=run_id,
+            benchmark=benchmark,
+            started_at=now.isoformat(),
+            exo_sha=_git_describe(repo_root),
+            hostname=socket.gethostname(),
+            platform=f"{platform.system()} {platform.release()} ({platform.machine()})",
+            tags=dict(tags or {}),
+        )
+
+
+@dataclass
+class ResultsBundle:
+    """Container for a single benchmark's results, before being written."""
+
+    metadata: RunMetadata
+    params: dict[str, Any] = field(default_factory=dict)
+    cluster: dict[str, Any] = field(default_factory=dict)
+    runs: list[dict[str, Any]] = field(default_factory=list)
+    cold_controls: list[dict[str, Any]] = field(default_factory=list)
+    derived: dict[str, Any] = field(default_factory=dict)
+
+    def capture_cluster(self, client: ExoClient) -> None:
+        """Snapshot the cluster state into ``self.cluster``."""
+        try:
+            snapshot = capture_cluster_snapshot(client)
+            if snapshot:
+                self.cluster.update(snapshot)
+        except Exception:
+            # Non-fatal: a benchmark without cluster snapshot is still valid
+            pass
+
+    def write_json(self, output_dir: Path) -> Path:
+        """Write the bundle as ``<output_dir>/<run_id>.json`` and return the path."""
+        if self.metadata.finished_at is None:
+            self.metadata.finished_at = datetime.now(timezone.utc).isoformat()
+        output_dir.mkdir(parents=True, exist_ok=True)
+        path = output_dir / f"{self.metadata.run_id}.json"
+        with path.open("w", encoding="utf-8") as f:
+            json.dump(asdict(self), f, indent=2, ensure_ascii=False)
+        return path
+
+
+def find_repo_root(start: Path | None = None) -> Path:
+    """Walk upwards from ``start`` (or this file) until a ``.git`` dir is found."""
+    cur = (start or Path(__file__)).resolve()
+    for parent in (cur, *cur.parents):
+        if (parent / ".git").is_dir() or (parent / ".git").is_file():
+            return parent
+    raise RuntimeError(f"Could not locate repo root above {cur}")

bench/lib/session.py

@@ -0,0 +1,65 @@
+"""BenchSession — wires together cluster + client + instance + tokenizer.
+
+Holds the ``EcoSession``, a deployed ``ClusterInfo``, an ``ExoClient`` for
+the cluster's primary endpoint, and (for benchmarks that need exact-token
+prompts) a lazily-constructed :class:`PromptSizer`.
+
+Benchmarks consume this via :func:`bench.lib.cluster.managed_instance`,
+which yields a populated ``BenchSession``. Library helpers (e.g.
+``context_scaling.run``) take a ``BenchSession`` and never reach for
+global state.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+from exo_tools.client import ExoClient
+from exo_tools.cluster import ClusterInfo, EcoSession, make_client_from_url
+
+from .prompt import PromptSizer, load_tokenizer_for_bench
+
+
+@dataclass
+class BenchSession:
+    """Bundle of cluster + client + (optional) instance for benchmarks."""
+
+    cluster: ClusterInfo
+    eco: EcoSession
+    instance_id: str | None = None
+    model_id: str | None = None
+    full_model_id: str | None = None
+    _prompt_sizer: PromptSizer | None = field(default=None, repr=False)
+
+    @property
+    def client(self) -> ExoClient:
+        return make_client_from_url(self.cluster.api_url)
+
+    def state(self) -> dict[str, Any]:
+        raw: Any = self.client.request_json("GET", "/state")  # type: ignore[reportAny]
+        if isinstance(raw, dict):
+            return cast("dict[str, Any]", raw)
+        return {}
+
+    def instances(self) -> dict[str, Any]:
+        result: Any = self.state().get("instances", {})  # type: ignore[reportAny]
+        if isinstance(result, dict):
+            return cast("dict[str, Any]", result)
+        return {}
+
+    def get_prompt_sizer(self) -> PromptSizer:
+        """Return a cached :class:`PromptSizer` for ``self.full_model_id``.
+
+        Loaded lazily because tokenizer load is expensive and not every
+        benchmark needs prompt sizing.
+        """
+        if self._prompt_sizer is not None:
+            return self._prompt_sizer
+        if self.full_model_id is None:
+            raise RuntimeError(
+                "BenchSession.full_model_id is not set; cannot build a PromptSizer."
+            )
+        tokenizer = load_tokenizer_for_bench(self.full_model_id)
+        self._prompt_sizer = PromptSizer(tokenizer)
+        return self._prompt_sizer

bench/lib/tests/test_context_scaling.py

@@ -0,0 +1,186 @@
+"""Unit tests for the pure helpers in ``bench.lib.context_scaling``.
+
+The orchestration entry points (``run``, ``run_cached_sweep``,
+``run_cold_controls``, ``make_cold_control_factory``) need a real
+``BenchSession`` and exo cluster, so they're exercised end-to-end via
+``python -m bench.cli context-scaling``. This module covers the
+underscore-prefixed pure helpers via direct private-symbol access (the
+private prefix discourages library users; tests for those helpers are
+the explicit exception).
+"""
+
+from __future__ import annotations
+
+import math
+from typing import cast
+
+from bench.lib.context_scaling import (
+    StepResult,
+    _compute_t_cum,  # type: ignore[reportPrivateUsage]
+    _interp,  # type: ignore[reportPrivateUsage]
+    derive_summary,
+)
+
+
+def _step(
+    *,
+    pp: int,
+    delta: int,
+    prompt_tps: float,
+    generation_tps: float = 100.0,
+    hit: str = "partial",
+) -> StepResult:
+    return StepResult(
+        pp_tokens=pp,
+        delta_tokens=delta,
+        prompt_tps=prompt_tps,
+        generation_tps=generation_tps,
+        prefix_cache_hit=hit,
+        prompt_tokens=pp,
+        generation_tokens=32,
+        elapsed_s=delta / prompt_tps if prompt_tps else 0.0,
+    )
+
+
+def _close(actual: float, expected: float, abs_tol: float = 1e-3) -> bool:
+    return math.isclose(actual, expected, abs_tol=abs_tol)
+
+
+# ---------------------------------------------------------------------------
+# _compute_t_cum
+# ---------------------------------------------------------------------------
+
+
+class TestComputeTCum:
+    def test_empty_returns_empty(self) -> None:
+        assert _compute_t_cum([]) == []
+
+    def test_single_step(self) -> None:
+        # 256 tokens at 1024 tps -> 0.25s
+        out = _compute_t_cum([_step(pp=256, delta=256, prompt_tps=1024.0)])
+        assert len(out) == 1
+        assert _close(out[0], 0.25)
+
+    def test_cumulative_sum_across_three_steps(self) -> None:
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1000.0),  # 0.256s
+            _step(pp=512, delta=256, prompt_tps=2000.0),  # +0.128s = 0.384s
+            _step(pp=768, delta=256, prompt_tps=512.0),  # +0.500s = 0.884s
+        ]
+        out = _compute_t_cum(steps)
+        assert _close(out[0], 0.256)
+        assert _close(out[1], 0.384)
+        assert _close(out[2], 0.884)
+        # Monotonically non-decreasing
+        assert out == sorted(out)
+
+    def test_zero_tps_step_skipped(self) -> None:
+        # A row with prompt_tps == 0 contributes nothing to the cumulative sum
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1024.0),  # +0.25s
+            _step(pp=512, delta=256, prompt_tps=0.0),  # +0
+            _step(pp=768, delta=256, prompt_tps=512.0),  # +0.5s
+        ]
+        out = _compute_t_cum(steps)
+        assert _close(out[0], 0.25)
+        assert _close(out[1], 0.25)  # unchanged
+        assert _close(out[2], 0.75)
+
+    def test_zero_delta_step_skipped(self) -> None:
+        # Defensive: a Δ=0 row would otherwise add zero anyway, but we
+        # explicitly guard against negative delta + 0/0.
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1000.0),
+            _step(pp=256, delta=0, prompt_tps=1000.0),  # explicit Δ=0
+        ]
+        out = _compute_t_cum(steps)
+        assert _close(out[0], 0.256)
+        assert _close(out[1], 0.256)
+
+
+# ---------------------------------------------------------------------------
+# _interp
+# ---------------------------------------------------------------------------
+
+
+class TestInterp:
+    def test_empty_points_returns_zero(self) -> None:
+        assert _interp([], 100) == 0.0
+
+    def test_single_point_returns_y(self) -> None:
+        assert _interp([(100, 1.5)], 50) == 1.5
+        assert _interp([(100, 1.5)], 100) == 1.5
+        assert _interp([(100, 1.5)], 200) == 1.5
+
+    def test_clamps_below_first(self) -> None:
+        points = [(100, 0.1), (200, 0.3), (300, 0.6)]
+        assert _interp(points, 0) == 0.1
+        assert _interp(points, 50) == 0.1
+        assert _interp(points, 100) == 0.1
+
+    def test_clamps_above_last(self) -> None:
+        points = [(100, 0.1), (200, 0.3), (300, 0.6)]
+        assert _interp(points, 300) == 0.6
+        assert _interp(points, 500) == 0.6
+        assert _interp(points, 1_000_000) == 0.6
+
+    def test_mid_bracket_linear_interpolation(self) -> None:
+        points = [(100, 0.0), (200, 1.0)]
+        assert _close(_interp(points, 150), 0.5)
+        assert _close(_interp(points, 175), 0.75)
+
+    def test_multi_segment_linear_interpolation(self) -> None:
+        # Two adjacent segments, x=250 falls in the second one
+        points = [(100, 0.1), (200, 0.3), (300, 0.6)]
+        # 200..300: 0.3 + (0.6-0.3) * (250-200)/(300-200) = 0.3 + 0.15 = 0.45
+        assert _close(_interp(points, 250), 0.45)
+
+
+# ---------------------------------------------------------------------------
+# derive_summary
+# ---------------------------------------------------------------------------
+
+
+def _gap_at(summary: dict[str, object], index: int) -> dict[str, float]:
+    """Cast ``summary['control_gaps'][index]`` into the typed shape we expect."""
+    raw = summary["control_gaps"]
+    assert isinstance(raw, list)
+    entry = cast("dict[str, float]", raw[index])
+    return entry
+
+
+class TestDeriveSummary:
+    def test_no_controls_only_t_cum(self) -> None:
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1024.0),
+            _step(pp=512, delta=256, prompt_tps=1024.0),
+        ]
+        summary = derive_summary(steps, [])
+        t_cum = cast("list[float]", summary["t_cum_seconds"])
+        assert _close(t_cum[0], 0.25)
+        assert _close(t_cum[1], 0.5)
+        assert summary["control_gaps"] == []
+
+    def test_control_gap_at_known_pp(self) -> None:
+        # Sweep: 0.25s @ pp=256, 0.5s @ pp=512
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1024.0),
+            _step(pp=512, delta=256, prompt_tps=1024.0),
+        ]
+        # Cold control at pp=512, 2x faster than the per-step rate -> 0.25s
+        controls = [_step(pp=512, delta=512, prompt_tps=2048.0, hit="none")]
+        summary = derive_summary(steps, controls)
+        gap = _gap_at(summary, 0)
+        assert gap["pp_tokens"] == 512
+        assert _close(gap["cold_t_seconds"], 0.25, abs_tol=0.01)
+        assert _close(gap["t_cum_seconds_at_pp"], 0.5, abs_tol=0.01)
+        assert _close(gap["gap_seconds"], 0.25, abs_tol=0.01)
+        # gap_fraction = 0.25 / 0.25 = 1.0
+        assert _close(gap["gap_fraction"], 1.0, abs_tol=0.01)
+
+    def test_control_gap_zero_cold_tps_yields_zero_fraction(self) -> None:
+        steps = [_step(pp=256, delta=256, prompt_tps=1000.0)]
+        controls = [_step(pp=256, delta=256, prompt_tps=0.0, hit="none")]
+        gap = _gap_at(derive_summary(steps, controls), 0)
+        assert gap["cold_t_seconds"] == 0.0
+        assert gap["gap_fraction"] == 0.0

bench/lib/tests/test_model_meta.py

@@ -0,0 +1,171 @@
+"""Unit tests for ``bench.lib.model_meta``.
+
+These exercise the pure derivation helpers (no HF round-trip). The HTTP
+fetchers (``fetch_model_meta``, ``_read_config_json``, ``_sum_weight_sizes``)
+hit the public hub and aren't covered here.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from bench.lib.model_meta import (
+    ModelMeta,
+    derive_cold_controls,
+    derive_context_ramp,
+)
+
+
+def _meta(
+    *,
+    weight_bytes: int = 0,
+    max_pos: int = 4096,
+    layers: int = 32,
+) -> ModelMeta:
+    return ModelMeta(
+        model_id="test/model",
+        total_weight_bytes=weight_bytes,
+        max_position_embeddings=max_pos,
+        num_hidden_layers=layers,
+    )
+
+
+# ---------------------------------------------------------------------------
+# ModelMeta properties
+# ---------------------------------------------------------------------------
+
+
+class TestModelMetaConstraints:
+    def test_zero_weight_yields_one_gib_floor(self) -> None:
+        meta = _meta(weight_bytes=0)
+        # int(0 * 1.30) + 1 == 1; int(0 * 1.10) + 1 == 1
+        assert meta.memory_constraint_gb == 1.0
+        assert meta.disk_constraint_gb == 1.0
+
+    def test_one_gib_weight_rounds_up(self) -> None:
+        meta = _meta(weight_bytes=1 * (1024**3))
+        # int(1.0 * 1.30) + 1 = 2; int(1.0 * 1.10) + 1 = 2
+        assert meta.memory_constraint_gb == 2.0
+        assert meta.disk_constraint_gb == 2.0
+
+    def test_sixteen_gib_weight_uses_30pct_memory_10pct_disk(self) -> None:
+        meta = _meta(weight_bytes=16 * (1024**3))
+        # memory: int(16 * 1.30) + 1 = 21; disk: int(16 * 1.10) + 1 = 18
+        assert meta.memory_constraint_gb == 21.0
+        assert meta.disk_constraint_gb == 18.0
+
+    def test_total_weight_gb_property(self) -> None:
+        meta = _meta(weight_bytes=2_147_483_648)  # 2 GiB exactly
+        assert math.isclose(meta.total_weight_gb, 2.0)
+
+
+# ---------------------------------------------------------------------------
+# derive_context_ramp
+# ---------------------------------------------------------------------------
+
+
+class TestDeriveContextRamp:
+    def test_full_max_evenly_divides_round_to(self) -> None:
+        meta = _meta(max_pos=131072)  # 128k
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=32)
+        # 131072 // 32 = 4096; rounded down to multiple of 256 = 4096
+        assert pp_step == 4096
+        assert num_steps == 32
+        # Top of ramp == max
+        assert pp_step * num_steps == 131072
+
+    def test_qwen30b_a3b_ramp(self) -> None:
+        meta = _meta(max_pos=40960)  # Qwen3-30B-A3B
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=32)
+        # 40960 // 32 = 1280; multiple of 256
+        assert pp_step == 1280
+        assert pp_step * num_steps == 40960
+
+    def test_fraction_of_max_half(self) -> None:
+        meta = _meta(max_pos=131072)
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=8, fraction_of_max=0.5)
+        # half = 65536; 65536 // 8 = 8192
+        assert pp_step == 8192
+        assert num_steps == 8
+
+    def test_min_pp_step_floor(self) -> None:
+        meta = _meta(max_pos=512)
+        # 512 // 32 = 16, but min_pp_step=256 floors it; rounded to 256
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=32)
+        assert pp_step == 256
+        assert num_steps == 32
+
+    def test_round_to_truncates_down(self) -> None:
+        meta = _meta(max_pos=10000)
+        pp_step, _ = derive_context_ramp(meta, num_steps=32, round_to=256)
+        # 10000 // 32 = 312; (312 // 256) * 256 = 256
+        assert pp_step == 256
+
+    def test_round_to_zero_step_falls_back_to_round_to(self) -> None:
+        # Pathological: huge round_to relative to per-step size
+        meta = _meta(max_pos=1024)
+        pp_step, _ = derive_context_ramp(meta, num_steps=8, round_to=1024)
+        # 1024 // 8 = 128, but min_pp_step=256 → 256; (256 // 1024) * 1024 = 0;
+        # `or round_to` rescues to 1024.
+        assert pp_step == 1024
+
+    def test_max_pos_zero_raises(self) -> None:
+        meta = _meta(max_pos=0)
+        with pytest.raises(ValueError, match="max_position_embeddings=0"):
+            _ = derive_context_ramp(meta, num_steps=32)
+
+    @pytest.mark.parametrize("fraction", [0.0, -0.1, 1.5, 2.0])
+    def test_fraction_outside_unit_interval_raises(self, fraction: float) -> None:
+        meta = _meta(max_pos=4096)
+        with pytest.raises(ValueError, match="fraction_of_max"):
+            _ = derive_context_ramp(meta, num_steps=4, fraction_of_max=fraction)
+
+    @pytest.mark.parametrize("steps", [0, -1, -100])
+    def test_num_steps_must_be_positive(self, steps: int) -> None:
+        meta = _meta(max_pos=4096)
+        with pytest.raises(ValueError, match="num_steps"):
+            _ = derive_context_ramp(meta, num_steps=steps)
+
+
+# ---------------------------------------------------------------------------
+# derive_cold_controls
+# ---------------------------------------------------------------------------
+
+
+class TestDeriveColdControls:
+    def test_count_zero_returns_empty_tuple(self) -> None:
+        meta = _meta()
+        assert derive_cold_controls(meta, pp_step=4096, num_steps=32, count=0) == ()
+
+    def test_count_one_returns_top_only(self) -> None:
+        meta = _meta()
+        assert derive_cold_controls(meta, pp_step=4096, num_steps=32, count=1) == (
+            131072,
+        )
+
+    def test_evenly_spaced_four(self) -> None:
+        meta = _meta()
+        out = derive_cold_controls(meta, pp_step=4096, num_steps=32, count=4)
+        # max_pp = 131072; (131072 * (i+1)) // 4 for i in {0,1,2,3}
+        # = {32768, 65536, 98304, 131072}
+        assert out == (32768, 65536, 98304, 131072)
+
+    def test_filters_below_pp_step(self) -> None:
+        meta = _meta()
+        out = derive_cold_controls(meta, pp_step=8192, num_steps=2, count=4)
+        # max_pp = 16384; spaced points = {4096, 8192, 12288, 16384};
+        # 4096 < pp_step=8192 → dropped.
+        assert out == (8192, 12288, 16384)
+
+    def test_dedups_at_low_count_high_step(self) -> None:
+        meta = _meta()
+        # max_pp = 1024; count=2 → spaced = {512, 1024}; 512 < pp_step? No (=).
+        out = derive_cold_controls(meta, pp_step=512, num_steps=2, count=2)
+        assert out == (512, 1024)
+
+    def test_returned_in_ascending_order(self) -> None:
+        meta = _meta()
+        out = derive_cold_controls(meta, pp_step=1024, num_steps=8, count=4)
+        assert list(out) == sorted(out)

bench/lib/tests/test_plotting.py

@@ -0,0 +1,189 @@
+"""Smoke tests for ``bench.lib.plotting``.
+
+Renders a synthetic benchmark JSON to a tmp PNG and verifies the file is
+non-empty. We deliberately don't assert on pixel values — matplotlib
+output isn't byte-stable across versions — but a non-empty PNG with a
+valid header is a strong signal the renderer didn't throw.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+import pytest
+
+from bench.lib.plotting import PlotInputs, render_context_scaling
+
+
+def _write_synthetic_run(path: Path, *, run_id: str, model: str = "test/model") -> None:
+    """Write a minimal context-scaling-shaped JSON for plotting tests."""
+    payload = {
+        "metadata": {
+            "run_id": run_id,
+            "benchmark": "context_scaling",
+            "started_at": "2026-05-10T00:00:00Z",
+            "exo_sha": "deadbeef",
+            "hostname": "test-host",
+            "platform": "Linux 6.0 (x86_64)",
+            "tags": {"operator": "tester"},
+        },
+        "params": {
+            "pp_step": 256,
+            "num_steps": 4,
+            "tg": 32,
+            "warmup": 1,
+            "full_model_id": model,
+        },
+        "cluster": {},
+        "runs": [
+            {
+                "step_index": 0,
+                "phase": "cached_sweep",
+                "pp_tokens": 256,
+                "delta_tokens": 256,
+                "prompt_tps": 1800.0,
+                "generation_tps": 410.0,
+                "prefix_cache_hit": "exact",
+                "prompt_tokens": 256,
+                "generation_tokens": 32,
+                "elapsed_s": 0.14,
+                "peak_memory_bytes": 1_000_000_000,
+                "output_text_preview": "",
+            },
+            {
+                "step_index": 1,
+                "phase": "cached_sweep",
+                "pp_tokens": 512,
+                "delta_tokens": 256,
+                "prompt_tps": 2000.0,
+                "generation_tps": 395.0,
+                "prefix_cache_hit": "partial",
+                "prompt_tokens": 512,
+                "generation_tokens": 32,
+                "elapsed_s": 0.13,
+                "peak_memory_bytes": 1_100_000_000,
+                "output_text_preview": "",
+            },
+            {
+                "step_index": 2,
+                "phase": "cached_sweep",
+                "pp_tokens": 768,
+                "delta_tokens": 256,
+                "prompt_tps": 2200.0,
+                "generation_tps": 378.0,
+                "prefix_cache_hit": "partial",
+                "prompt_tokens": 768,
+                "generation_tokens": 32,
+                "elapsed_s": 0.12,
+                "peak_memory_bytes": 1_200_000_000,
+                "output_text_preview": "",
+            },
+        ],
+        "cold_controls": [
+            {
+                "phase": "cold_control",
+                "pp_tokens": 512,
+                "delta_tokens": 512,
+                "prompt_tps": 3200.0,
+                "generation_tps": 400.0,
+                "prefix_cache_hit": "none",
+                "prompt_tokens": 512,
+                "generation_tokens": 32,
+                "elapsed_s": 0.16,
+                "peak_memory_bytes": 1_500_000_000,
+                "output_text_preview": "",
+            },
+        ],
+        "derived": {
+            "t_cum_seconds": [0.14, 0.27, 0.39],
+            "control_gaps": [],
+        },
+    }
+    _ = path.write_text(json.dumps(payload))
+
+
+def _png_is_valid(path: Path) -> bool:
+    """A PNG file starts with the 8-byte magic ``\\x89PNG\\r\\n\\x1a\\n``."""
+    if not path.is_file():
+        return False
+    if path.stat().st_size < 100:
+        return False
+    head = path.read_bytes()[:8]
+    return head == b"\x89PNG\r\n\x1a\n"
+
+
+# ---------------------------------------------------------------------------
+
+
+class TestRenderContextScaling:
+    def test_single_run(self, tmp_path: Path) -> None:
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "out.png"
+
+        returned = render_context_scaling(PlotInputs(results=[json_path], output=out))
+        assert returned == out
+        assert _png_is_valid(out)
+
+    def test_creates_output_parent_dir(self, tmp_path: Path) -> None:
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "nested" / "deep" / "out.png"
+
+        _ = render_context_scaling(PlotInputs(results=[json_path], output=out))
+        assert _png_is_valid(out)
+
+    def test_comparison_two_runs(self, tmp_path: Path) -> None:
+        a = tmp_path / "a.json"
+        b = tmp_path / "b.json"
+        _write_synthetic_run(a, run_id="run-a", model="test/model-a")
+        _write_synthetic_run(b, run_id="run-b", model="test/model-b")
+        out = tmp_path / "compare.png"
+
+        _ = render_context_scaling(PlotInputs(results=[a, b], output=out))
+        assert _png_is_valid(out)
+
+    def test_label_tag_uses_metadata_tag(self, tmp_path: Path) -> None:
+        # Smoke test: just confirm passing label_tag doesn't throw and the
+        # PNG renders. Label content is too matplotlib-internal to inspect.
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "out.png"
+
+        _ = render_context_scaling(
+            PlotInputs(results=[json_path], output=out, label_tag="operator")
+        )
+        assert _png_is_valid(out)
+
+    def test_explicit_title(self, tmp_path: Path) -> None:
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "out.png"
+
+        _ = render_context_scaling(
+            PlotInputs(results=[json_path], output=out, title="Custom Title")
+        )
+        assert _png_is_valid(out)
+
+    def test_empty_results_raises(self, tmp_path: Path) -> None:
+        with pytest.raises(ValueError, match="at least one"):
+            _ = render_context_scaling(
+                PlotInputs(results=[], output=tmp_path / "out.png")
+            )
+
+    def test_wrong_benchmark_raises(self, tmp_path: Path) -> None:
+        # Same shape but with the wrong metadata.benchmark
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        raw_loaded: Any = json.loads(json_path.read_text())  # type: ignore[reportAny]
+        assert isinstance(raw_loaded, dict)
+        data = cast("dict[str, dict[str, str]]", raw_loaded)
+        data["metadata"]["benchmark"] = "something_else"
+        _ = json_path.write_text(json.dumps(data))
+
+        with pytest.raises(ValueError, match="context_scaling"):
+            _ = render_context_scaling(
+                PlotInputs(results=[json_path], output=tmp_path / "out.png")
+            )

bench/prefill-decode.toml

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

bench/prefill_decode_bench.py

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

bench/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
   "lm-eval[api,math]>=0.4.0",
   "human-eval>=1.0.3",
   "numpy>=1.24.0",
+  "matplotlib>=3.8",
 ]
 
 [build-system]

dashboard/package-lock.json

@@ -11,7 +11,8 @@
 				"highlight.js": "^11.11.1",
 				"katex": "^0.16.27",
 				"marked": "^17.0.1",
-				"mode-watcher": "^1.1.0"
+				"mode-watcher": "^1.1.0",
+				"pdfjs-dist": "^5.6.205"
 			},
 			"devDependencies": {
 				"@sveltejs/adapter-static": "^3.0.10",
@@ -518,6 +519,256 @@
 				"@jridgewell/sourcemap-codec": "^1.4.14"
 			}
 		},
+		"node_modules/@napi-rs/canvas": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas/-/canvas-0.1.97.tgz",
+			"integrity": "sha512-8cFniXvrIEnVwuNSRCW9wirRZbHvrD3JVujdS2P5n5xiJZNZMOZcfOvJ1pb66c7jXMKHHglJEDVJGbm8XWFcXQ==",
+			"license": "MIT",
+			"optional": true,
+			"workspaces": [
+				"e2e/*"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			},
+			"optionalDependencies": {
+				"@napi-rs/canvas-android-arm64": "0.1.97",
+				"@napi-rs/canvas-darwin-arm64": "0.1.97",
+				"@napi-rs/canvas-darwin-x64": "0.1.97",
+				"@napi-rs/canvas-linux-arm-gnueabihf": "0.1.97",
+				"@napi-rs/canvas-linux-arm64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-arm64-musl": "0.1.97",
+				"@napi-rs/canvas-linux-riscv64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-x64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-x64-musl": "0.1.97",
+				"@napi-rs/canvas-win32-arm64-msvc": "0.1.97",
+				"@napi-rs/canvas-win32-x64-msvc": "0.1.97"
+			}
+		},
+		"node_modules/@napi-rs/canvas-android-arm64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-0.1.97.tgz",
+			"integrity": "sha512-V1c/WVw+NzH8vk7ZK/O8/nyBSCQimU8sfMsB/9qeSvdkGKNU7+mxy/bIF0gTgeBFmHpj30S4E9WHMSrxXGQuVQ==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"android"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-darwin-arm64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-arm64/-/canvas-darwin-arm64-0.1.97.tgz",
+			"integrity": "sha512-ok+SCEF4YejcxuJ9Rm+WWunHHpf2HmiPxfz6z1a/NFQECGXtsY7A4B8XocK1LmT1D7P174MzwPF9Wy3AUAwEPw==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"darwin"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-darwin-x64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-0.1.97.tgz",
+			"integrity": "sha512-PUP6e6/UGlclUvAQNnuXCcnkpdUou6VYZfQOQxExLp86epOylmiwLkqXIvpFmjoTEDmPmXrI+coL/9EFU1gKPA==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"darwin"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm-gnueabihf": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-0.1.97.tgz",
+			"integrity": "sha512-XyXH2L/cic8eTNtbrXCcvqHtMX/nEOxN18+7rMrAM2XtLYC/EB5s0wnO1FsLMWmK+04ZSLN9FBGipo7kpIkcOw==",
+			"cpu": [
+				"arm"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-0.1.97.tgz",
+			"integrity": "sha512-Kuq/M3djq0K8ktgz6nPlK7Ne5d4uWeDxPpyKWOjWDK2RIOhHVtLtyLiJw2fuldw7Vn4mhw05EZXCEr4Q76rs9w==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm64-musl": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-0.1.97.tgz",
+			"integrity": "sha512-kKmSkQVnWeqg7qdsiXvYxKhAFuHz3tkBjW/zyQv5YKUPhotpaVhpBGv5LqCngzyuRV85SXoe+OFj+Tv0a0QXkQ==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-riscv64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-0.1.97.tgz",
+			"integrity": "sha512-Jc7I3A51jnEOIAXeLsN/M/+Z28LUeakcsXs07FLq9prXc0eYOtVwsDEv913Gr+06IRo34gJJVgT0TXvmz+N2VA==",
+			"cpu": [
+				"riscv64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-x64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-0.1.97.tgz",
+			"integrity": "sha512-iDUBe7AilfuBSRbSa8/IGX38Mf+iCSBqoVKLSQ5XaY2JLOaqz1TVyPFEyIck7wT6mRQhQt5sN6ogfjIDfi74tg==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-x64-musl": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-0.1.97.tgz",
+			"integrity": "sha512-AKLFd/v0Z5fvgqBDqhvqtAdx+fHMJ5t9JcUNKq4FIZ5WH+iegGm8HPdj00NFlCSnm83Fp3Ln8I2f7uq1aIiWaA==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-win32-arm64-msvc": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-0.1.97.tgz",
+			"integrity": "sha512-u883Yr6A6fO7Vpsy9YE4FVCIxzzo5sO+7pIUjjoDLjS3vQaNMkVzx5bdIpEL+ob+gU88WDK4VcxYMZ6nmnoX9A==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"win32"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-win32-x64-msvc": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-0.1.97.tgz",
+			"integrity": "sha512-sWtD2EE3fV0IzN+iiQUqr/Q1SwqWhs2O1FKItFlxtdDkikpEj5g7DKQpY3x55H/MAOnL8iomnlk3mcEeGiUMoQ==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"win32"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
 		"node_modules/@polka/url": {
 			"version": "1.0.0-next.29",
 			"resolved": "https://registry.npmjs.org/@polka/url/-/url-1.0.0-next.29.tgz",
@@ -2635,6 +2886,26 @@
 				"node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1"
 			}
 		},
+		"node_modules/node-readable-to-web-readable-stream": {
+			"version": "0.4.2",
+			"resolved": "https://registry.npmjs.org/node-readable-to-web-readable-stream/-/node-readable-to-web-readable-stream-0.4.2.tgz",
+			"integrity": "sha512-/cMZNI34v//jUTrI+UIo4ieHAB5EZRY/+7OmXZgBxaWBMcW2tGdceIw06RFxWxrKZ5Jp3sI2i5TsRo+CBhtVLQ==",
+			"license": "MIT",
+			"optional": true
+		},
+		"node_modules/pdfjs-dist": {
+			"version": "5.6.205",
+			"resolved": "https://registry.npmjs.org/pdfjs-dist/-/pdfjs-dist-5.6.205.tgz",
+			"integrity": "sha512-tlUj+2IDa7G1SbvBNN74UHRLJybZDWYom+k6p5KIZl7huBvsA4APi6mKL+zCxd3tLjN5hOOEE9Tv7VdzO88pfg==",
+			"license": "Apache-2.0",
+			"engines": {
+				"node": ">=20.19.0 || >=22.13.0 || >=24"
+			},
+			"optionalDependencies": {
+				"@napi-rs/canvas": "^0.1.96",
+				"node-readable-to-web-readable-stream": "^0.4.2"
+			}
+		},
 		"node_modules/picocolors": {
 			"version": "1.1.1",
 			"resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.1.1.tgz",

dashboard/package.json

@@ -31,6 +31,7 @@
 		"highlight.js": "^11.11.1",
 		"katex": "^0.16.27",
 		"marked": "^17.0.1",
-		"mode-watcher": "^1.1.0"
+		"mode-watcher": "^1.1.0",
+		"pdfjs-dist": "^5.6.205"
 	}
 }

dashboard/src/lib/components/ChatMessages.svelte

@@ -139,6 +139,8 @@
         return "🖼";
       case "text":
         return "📄";
+      case "pdf":
+        return "📑";
       default:
         return "📎";
     }

dashboard/src/lib/components/FamilyLogos.svelte

@@ -88,6 +88,18 @@
       d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
     />
   </svg>
+{:else if family === "nemotron"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M8.948 8.798v-1.43a6.7 6.7 0 0 1 .424-.018c3.922-.124 6.493 3.374 6.493 3.374s-2.774 3.851-5.75 3.851c-.398 0-.787-.062-1.158-.185v-4.346c1.528.185 1.837.857 2.747 2.385l2.04-1.714s-1.492-1.952-4-1.952a6.016 6.016 0 0 0-.796.035m0-4.735v2.138l.424-.027c5.45-.185 9.01 4.47 9.01 4.47s-4.08 4.964-8.33 4.964c-.37 0-.733-.035-1.095-.097v1.325c.3.035.61.062.91.062 3.957 0 6.82-2.023 9.593-4.408.459.371 2.34 1.263 2.73 1.652-2.633 2.208-8.772 3.984-12.253 3.984-.335 0-.653-.018-.971-.053v1.864H24V4.063zm0 10.326v1.131c-3.657-.654-4.673-4.46-4.673-4.46s1.758-1.944 4.673-2.262v1.237H8.94c-1.528-.186-2.73 1.245-2.73 1.245s.68 2.412 2.739 3.11M2.456 10.9s2.164-3.197 6.5-3.533V6.201C4.153 6.59 0 10.653 0 10.653s2.35 6.802 8.948 7.42v-1.237c-4.84-.6-6.492-5.936-6.492-5.936z"
+    />
+  </svg>
+{:else if family === "gemma"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M12.48 10.92v3.28h7.84c-.24 1.84-.853 3.187-1.787 4.133-1.147 1.147-2.933 2.4-6.053 2.4-4.827 0-8.6-3.893-8.6-8.72s3.773-8.72 8.6-8.72c2.6 0 4.507 1.027 5.907 2.347l2.307-2.307C18.747 1.44 16.133 0 12.48 0 5.867 0 .307 5.387.307 12s5.56 12 12.173 12c3.573 0 6.267-1.173 8.373-3.36 2.16-2.16 2.84-5.213 2.84-7.667 0-.76-.053-1.467-.173-2.053H12.48z"
+    />
+  </svg>
 {:else}
   <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
     <path

dashboard/src/lib/components/FamilySidebar.svelte

@@ -31,6 +31,8 @@
     kimi: "Kimi",
     flux: "FLUX",
     "qwen-image": "Qwen Img",
+    nemotron: "NVIDIA",
+    gemma: "Google",
   };
 
   function getFamilyName(family: string): string {
@@ -41,31 +43,20 @@
 </script>
 
 <div
-  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[72px] sm:min-w-[64px] overflow-y-auto scrollbar-hide"
+  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[80px] sm:min-w-[72px] overflow-y-auto scrollbar-hide"
 >
   <!-- All models (no filter) -->
   <button
     type="button"
     onclick={() => onSelect(null)}
-    class="group flex flex-col items-center justify-center p-2 sm:p-2 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
+    class="group flex items-center justify-center px-3 py-2.5 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
     null
       ? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
       : 'hover:bg-white/5 border-l-2 border-transparent'}"
     title="All models"
   >
-    <svg
-      class="w-5 h-5 {selectedFamily === null
-        ? 'text-exo-yellow'
-        : 'text-white/50 group-hover:text-white/70'}"
-      viewBox="0 0 24 24"
-      fill="currentColor"
-    >
-      <path
-        d="M4 8h4V4H4v4zm6 12h4v-4h-4v4zm-6 0h4v-4H4v4zm0-6h4v-4H4v4zm6 0h4v-4h-4v4zm6-10v4h4V4h-4zm-6 4h4V4h-4v4zm6 6h4v-4h-4v4zm0 6h4v-4h-4v4z"
-      />
-    </svg>
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === null
+      class="text-[12px] font-mono font-medium {selectedFamily === null
         ? 'text-exo-yellow'
         : 'text-white/40 group-hover:text-white/60'}">All</span
     >
@@ -89,7 +80,7 @@
           : "text-white/50 group-hover:text-amber-400/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'favorites'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'favorites'
           ? 'text-amber-400'
           : 'text-white/40 group-hover:text-white/60'}">Faves</span
       >
@@ -114,7 +105,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'recents'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'recents'
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}">Recent</span
       >
@@ -138,7 +129,7 @@
         : "text-white/50 group-hover:text-orange-400/70"}
     />
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === 'huggingface'
+      class="text-[11px] font-mono mt-0.5 {selectedFamily === 'huggingface'
         ? 'text-orange-400'
         : 'text-white/40 group-hover:text-white/60'}">Hub</span
     >
@@ -164,7 +155,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
+        class="text-[11px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
         family
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}"

dashboard/src/lib/components/HeaderNav.svelte

@@ -1,21 +1,41 @@
 <script lang="ts">
   import { browser } from "$app/environment";
+  import { featureFlags } from "$lib/stores/app.svelte";
 
-  export let showHome = true;
-  export let onHome: (() => void) | null = null;
-  export let showSidebarToggle = false;
-  export let sidebarVisible = true;
-  export let onToggleSidebar: (() => void) | null = null;
-  export let showMobileMenuToggle = false;
-  export let mobileMenuOpen = false;
-  export let onToggleMobileMenu: (() => void) | null = null;
-  export let showMobileRightToggle = false;
-  export let mobileRightOpen = false;
-  export let onToggleMobileRight: (() => void) | null = null;
-  export let downloadProgress: {
-    count: number;
-    percentage: number;
-  } | null = null;
+  const showAdvanced = $derived(featureFlags()["disaggregation"] === true);
+
+  interface Props {
+    showHome?: boolean;
+    onHome?: (() => void) | null;
+    showSidebarToggle?: boolean;
+    sidebarVisible?: boolean;
+    onToggleSidebar?: (() => void) | null;
+    showMobileMenuToggle?: boolean;
+    mobileMenuOpen?: boolean;
+    onToggleMobileMenu?: (() => void) | null;
+    showMobileRightToggle?: boolean;
+    mobileRightOpen?: boolean;
+    onToggleMobileRight?: (() => void) | null;
+    downloadProgress?: {
+      count: number;
+      percentage: number;
+    } | null;
+  }
+
+  let {
+    showHome = true,
+    onHome = null,
+    showSidebarToggle = false,
+    sidebarVisible = true,
+    onToggleSidebar = null,
+    showMobileMenuToggle = false,
+    mobileMenuOpen = false,
+    onToggleMobileMenu = null,
+    showMobileRightToggle = false,
+    mobileRightOpen = false,
+    onToggleMobileRight = null,
+    downloadProgress = null,
+  }: Props = $props();
 
   function handleHome(): void {
     if (onHome) {
@@ -259,5 +279,49 @@
       {/if}
       <span class="hidden sm:inline">Downloads</span>
     </a>
+    <a
+      href="/#/integrations"
+      class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
+      title="Integration configs for external tools"
+    >
+      <svg
+        class="w-4 h-4"
+        viewBox="0 0 24 24"
+        fill="none"
+        stroke="currentColor"
+        stroke-width="2"
+        stroke-linecap="round"
+        stroke-linejoin="round"
+      >
+        <path d="M10 13a5 5 0 0 0 7.54.54l3-3a5 5 0 0 0-7.07-7.07l-1.72 1.71" />
+        <path
+          d="M14 11a5 5 0 0 0-7.54-.54l-3 3a5 5 0 0 0 7.07 7.07l1.71-1.71"
+        />
+      </svg>
+      <span class="hidden sm:inline">Integrations</span>
+    </a>
+    {#if showAdvanced}
+      <a
+        href="/#/advanced"
+        class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
+        title="Advanced cluster settings"
+      >
+        <svg
+          class="w-4 h-4"
+          viewBox="0 0 24 24"
+          fill="none"
+          stroke="currentColor"
+          stroke-width="2"
+          stroke-linecap="round"
+          stroke-linejoin="round"
+        >
+          <circle cx="12" cy="12" r="3" />
+          <path
+            d="M19.4 15a1.65 1.65 0 0 0 .33 1.82l.06.06a2 2 0 0 1 0 2.83 2 2 0 0 1-2.83 0l-.06-.06a1.65 1.65 0 0 0-1.82-.33 1.65 1.65 0 0 0-1 1.51V21a2 2 0 0 1-4 0v-.09A1.65 1.65 0 0 0 9 19.4a1.65 1.65 0 0 0-1.82.33l-.06.06a2 2 0 0 1-2.83 0 2 2 0 0 1 0-2.83l.06-.06a1.65 1.65 0 0 0 .33-1.82 1.65 1.65 0 0 0-1.51-1H3a2 2 0 0 1 0-4h.09A1.65 1.65 0 0 0 4.6 9a1.65 1.65 0 0 0-.33-1.82l-.06-.06a2 2 0 0 1 0-2.83 2 2 0 0 1 2.83 0l.06.06a1.65 1.65 0 0 0 1.82.33H9a1.65 1.65 0 0 0 1-1.51V3a2 2 0 0 1 4 0v.09a1.65 1.65 0 0 0 1 1.51 1.65 1.65 0 0 0 1.82-.33l.06-.06a2 2 0 0 1 2.83 0 2 2 0 0 1 0 2.83l-.06.06a1.65 1.65 0 0 0-.33 1.82V9a1.65 1.65 0 0 0 1.51 1H21a2 2 0 0 1 0 4h-.09a1.65 1.65 0 0 0-1.51 1z"
+          />
+        </svg>
+        <span class="hidden sm:inline">Advanced</span>
+      </a>
+    {/if}
   </nav>
 </header>

dashboard/src/lib/components/IntegrationCard.svelte

@@ -0,0 +1,62 @@
+<script lang="ts">
+  import { copyText } from "$lib/utils/clipboard";
+
+  interface Props {
+    title: string;
+    subtitle: string;
+    config: string;
+    description?: string;
+    language?: "json" | "bash";
+  }
+
+  let {
+    title,
+    subtitle,
+    config,
+    description = "",
+    language = "json",
+  }: Props = $props();
+
+  let copied = $state(false);
+  let failed = $state(false);
+
+  async function copyToClipboard() {
+    const ok = await copyText(config);
+    if (ok) {
+      copied = true;
+      setTimeout(() => (copied = false), 2000);
+    } else {
+      failed = true;
+      setTimeout(() => (failed = false), 2000);
+    }
+  }
+</script>
+
+<div
+  class="border border-exo-light-gray/20 rounded-lg bg-exo-medium-gray/20 overflow-hidden"
+>
+  <div class="flex items-center justify-between px-5 py-4">
+    <div>
+      <h3 class="text-white text-sm font-semibold tracking-wide">{title}</h3>
+      <p class="text-exo-light-gray/60 text-xs mt-0.5 font-mono">{subtitle}</p>
+    </div>
+    <button
+      onclick={copyToClipboard}
+      class="px-3 py-1.5 text-xs rounded border transition-all duration-200 cursor-pointer
+        {copied
+        ? 'border-green-500/50 text-green-400 bg-green-500/10'
+        : failed
+          ? 'border-red-500/50 text-red-400 bg-red-500/10'
+          : 'border-exo-light-gray/30 text-exo-light-gray hover:border-exo-yellow/50 hover:text-exo-yellow'}"
+    >
+      {copied ? "Copied!" : failed ? "Copy failed" : "Copy"}
+    </button>
+  </div>
+  {#if description}
+    <p class="text-exo-light-gray/70 text-xs px-5 pb-3">{description}</p>
+  {/if}
+  <div class="bg-black/30 border-t border-exo-light-gray/10">
+    <pre
+      class="text-xs text-exo-light-gray/90 font-mono p-4 overflow-x-auto whitespace-pre">{config}</pre>
+  </div>
+</div>

dashboard/src/lib/components/ModelCard.svelte

@@ -16,7 +16,9 @@
       perNode?: Array<{
         nodeId: string;
         nodeName: string;
-        progress: DownloadProgress;
+        status: "completed" | "partial" | "pending" | "downloading";
+        percentage: number;
+        progress: DownloadProgress | null;
       }>;
     } | null;
     nodes?: Record<string, NodeInfo>;
@@ -145,10 +147,7 @@
     return `${s}s`;
   }
 
-  const isDownloading = $derived(downloadStatus?.isDownloading ?? false);
-  const progress = $derived(downloadStatus?.progress);
-  const percentage = $derived(progress?.percentage ?? 0);
-  let expandedNodes = $state<Set<string>>(new Set());
+  const perNode = $derived(downloadStatus?.perNode ?? []);
 
   function toggleNodeDetails(nodeId: string): void {
     const next = new Set(expandedNodes);
@@ -587,23 +586,49 @@
       </span>
     </div>
 
-    <!-- Download Status -->
-    {#if isDownloading && progress}
+    <!-- Download Status (per-node) -->
+    {#if perNode.length > 0}
       <div class="mb-2 space-y-1">
-        <div class="flex items-center justify-between text-xs font-mono">
-          <span class="text-blue-400 tracking-wider uppercase">Downloading</span
-          >
-          <span class="text-white/60"
-            >{percentage.toFixed(1)}% &middot; {formatSpeed(progress.speed)}
-            &middot; {formatEta(progress.etaMs)}</span
-          >
-        </div>
-        <div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
-          <div
-            class="h-full bg-blue-500/70 transition-all duration-300"
-            style="width: {percentage}%"
-          ></div>
+        <div
+          class="text-[10px] font-mono text-white/20 tracking-widest uppercase"
+        >
+          Download progress
         </div>
+        {#each perNode as node}
+          <div class="flex items-center gap-2 text-xs font-mono">
+            <span class="text-white/40 w-20 truncate" title={node.nodeId}
+              >{node.nodeName}</span
+            >
+            <div
+              class="flex-1 h-1 bg-exo-medium-gray/30 rounded overflow-hidden"
+            >
+              <div
+                class="h-full transition-all duration-300 {node.status ===
+                'downloading'
+                  ? 'bg-blue-500/70'
+                  : node.status === 'completed'
+                    ? 'bg-exo-yellow/40'
+                    : 'bg-white/20'}"
+                style="width: {node.percentage}%"
+              ></div>
+            </div>
+            <span
+              class="text-right {node.status === 'completed'
+                ? 'text-exo-yellow/60'
+                : node.status === 'downloading'
+                  ? 'text-blue-400/60'
+                  : 'text-white/30'}"
+            >
+              {#if node.status === "downloading" && node.progress}
+                {Math.round(node.percentage)}% {formatSpeed(
+                  node.progress.speed,
+                )}
+              {:else}
+                {node.percentage > 0 ? `${Math.round(node.percentage)}%` : "0%"}
+              {/if}
+            </span>
+          </div>
+        {/each}
       </div>
     {/if}
 
@@ -662,15 +687,7 @@
             {@const allConnections =
               isDebugMode && usedNodes.length > 1
                 ? (() => {
-                    const conns: Array<{
-                      ip: string;
-                      iface: string | null;
-                      from: string;
-                      to: string;
-                      midX: number;
-                      midY: number;
-                      arrow: string;
-                    }> = [];
+                    const conns: Array = [];
                     for (let i = 0; i < usedNodes.length; i++) {
                       for (let j = i + 1; j < usedNodes.length; j++) {
                         const n1 = usedNodes[i];
@@ -682,7 +699,12 @@
                           const toPos = nodePositions[c.to];
                           const arrow =
                             fromPos && toPos ? getArrow(fromPos, toPos) : "→";
-                          conns.push({ ...c, midX, midY, arrow });
+                          conns.push({
+                            ...c,
+                            midX,
+                            midY,
+                            arrow,
+                          });
                         }
                       }
                     }

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -73,7 +73,7 @@
 
 <!-- svelte-ignore a11y_no_static_element_interactions -->
 <div
-  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-10"
+  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-20"
   transition:fly={{ y: -10, duration: 200, easing: cubicOut }}
   onclick={(e) => e.stopPropagation()}
   role="dialog"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -457,8 +457,10 @@
       "deepseek",
       "gpt-oss",
       "llama",
+      "gemma",
       "flux",
       "qwen-image",
+      "nemotron",
     ];
     return Array.from(families).sort((a, b) => {
       const aIdx = familyOrder.indexOf(a);

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

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

dashboard/src/lib/components/index.ts

@@ -13,3 +13,4 @@ export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
 export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
 export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
 export { default as ChatModelSelector } from "./ChatModelSelector.svelte";
+export { default as IntegrationCard } from "./IntegrationCard.svelte";

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

@@ -74,6 +74,12 @@ export interface Instance {
   };
 }
 
+export interface RawInstanceLink {
+  linkId: string;
+  prefillInstances: string[];
+  decodeInstances: string[];
+}
+
 // Granular node state types from the new state structure
 interface RawNodeIdentity {
   modelId?: string;
@@ -223,6 +229,7 @@ interface RawStateResponse {
     }
   >;
   runners?: Record<string, unknown>;
+  instanceLinks?: Record<string, RawInstanceLink>;
   downloads?: Record<string, unknown[]>;
   // New granular node state fields
   nodeIdentities?: Record<string, RawNodeIdentity>;
@@ -257,11 +264,12 @@ interface RawStateResponse {
 }
 
 export interface MessageAttachment {
-  type: "image" | "text" | "file" | "generated-image";
+  type: "image" | "text" | "file" | "generated-image" | "pdf";
   name: string;
   content?: string;
   preview?: string;
   mimeType?: string;
+  pageImages?: string[];
 }
 
 export interface TopLogprob {
@@ -540,6 +548,8 @@ class AppStore {
   topologyData = $state<TopologyData | null>(null);
   instances = $state<Record<string, unknown>>({});
   runners = $state<Record<string, unknown>>({});
+  instanceLinks = $state<Record<string, RawInstanceLink>>({});
+  featureFlags = $state<Record<string, boolean>>({});
   downloads = $state<Record<string, unknown[]>>({});
   nodeDisk = $state<
     Record<
@@ -1273,6 +1283,7 @@ class AppStore {
 
   startPolling() {
     this.fetchState();
+    this.fetchFeatureFlags();
     this.fetchInterval = setInterval(() => this.fetchState(), 1000);
   }
 
@@ -1284,6 +1295,16 @@ class AppStore {
     this.stopPreviewsPolling();
   }
 
+  async fetchFeatureFlags() {
+    try {
+      const response = await fetch("/v1/feature-flags");
+      if (!response.ok) return;
+      this.featureFlags = await response.json();
+    } catch {
+      // Silently ignore — defaults to all-disabled.
+    }
+  }
+
   async fetchState() {
     try {
       const response = await fetch("/state");
@@ -1309,6 +1330,11 @@ class AppStore {
       if (data.runners) {
         this.runners = data.runners;
       }
+      if (data.instanceLinks) {
+        this.instanceLinks = data.instanceLinks;
+      } else {
+        this.instanceLinks = {};
+      }
       if (data.downloads) {
         this.downloads = data.downloads;
       }
@@ -1669,7 +1695,15 @@ class AppStore {
               }
             }
           }
-          return { role: m.role, content: msgContent };
+          const out: {
+            role: string;
+            content: string;
+            reasoning_content?: string;
+          } = { role: m.role, content: msgContent };
+          if (m.role === "assistant" && m.thinking) {
+            out.reasoning_content = m.thinking;
+          }
+          return out;
         }),
       ];
 
@@ -1793,6 +1827,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final update
@@ -1868,7 +1910,15 @@ class AppStore {
       const apiMessages = [
         systemPrompt,
         ...targetConversation.messages.slice(0, -1).map((m) => {
-          return { role: m.role, content: m.content };
+          const out: {
+            role: string;
+            content: string;
+            reasoning_content?: string;
+          } = { role: m.role, content: m.content };
+          if (m.role === "assistant" && m.thinking) {
+            out.reasoning_content = m.thinking;
+          }
+          return out;
         }),
       ];
 
@@ -1990,6 +2040,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final cleanup of the message (if conversation still exists)
@@ -2226,6 +2284,7 @@ class AppStore {
       type: string;
       textContent?: string;
       preview?: string;
+      pageImages?: string[];
     }[],
     enableThinking?: boolean | null,
   ): Promise<void> {
@@ -2261,6 +2320,20 @@ class AppStore {
             preview: file.preview,
             mimeType: file.type,
           });
+        } else if (
+          file.pageImages ||
+          (file.textContent && file.type === "application/pdf")
+        ) {
+          attachments.push({
+            type: "pdf",
+            name: file.name,
+            content: file.textContent,
+            pageImages: file.pageImages,
+            mimeType: file.type,
+          });
+          if (file.textContent) {
+            fileContext += `\n\n[File: ${file.name}]\n\`\`\`\n${file.textContent}\n\`\`\``;
+          }
         } else if (file.textContent) {
           attachments.push({
             type: "text",
@@ -2328,22 +2401,89 @@ class AppStore {
       const apiMessages = [
         systemPrompt,
         ...targetConversation.messages.slice(0, -1).map((m) => {
-          // Build content including any text file attachments
+          // Check if this message has image or PDF attachments
+          const visualAttachments = m.attachments?.filter(
+            (a) =>
+              (a.type === "image" && a.preview) ||
+              (a.type === "pdf" && a.pageImages?.length),
+          );
+
+          if (visualAttachments && visualAttachments.length > 0) {
+            // Build multimodal content array (OpenAI vision format)
+            const contentParts: Array<
+              | { type: "text"; text: string }
+              | { type: "image_url"; image_url: { url: string } }
+            > = [];
+
+            // Add image parts first
+            for (const att of visualAttachments) {
+              if (att.type === "image" && att.preview) {
+                contentParts.push({
+                  type: "image_url",
+                  image_url: { url: att.preview },
+                });
+              } else if (att.type === "pdf" && att.pageImages) {
+                for (const pageImg of att.pageImages) {
+                  contentParts.push({
+                    type: "image_url",
+                    image_url: { url: pageImg },
+                  });
+                }
+              }
+            }
+
+            // Build text content including any text/pdf file attachments
+            let textContent = m.content;
+            if (m.attachments) {
+              for (const attachment of m.attachments) {
+                if (
+                  (attachment.type === "text" || attachment.type === "pdf") &&
+                  attachment.content
+                ) {
+                  textContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``;
+                }
+              }
+            }
+
+            if (textContent) {
+              contentParts.push({ type: "text", text: textContent });
+            }
+
+            const out: {
+              role: string;
+              content: typeof contentParts;
+              reasoning_content?: string;
+            } = { role: m.role, content: contentParts };
+            if (m.role === "assistant" && m.thinking) {
+              out.reasoning_content = m.thinking;
+            }
+            return out;
+          }
+
+          // Text-only message (original path)
           let msgContent = m.content;
 
-          // Add text attachments as context
+          // Add text/pdf attachments as context
           if (m.attachments) {
             for (const attachment of m.attachments) {
-              if (attachment.type === "text" && attachment.content) {
+              if (
+                (attachment.type === "text" || attachment.type === "pdf") &&
+                attachment.content
+              ) {
                 msgContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``;
               }
             }
           }
 
-          return {
-            role: m.role,
-            content: msgContent,
-          };
+          const out: {
+            role: string;
+            content: string;
+            reasoning_content?: string;
+          } = { role: m.role, content: msgContent };
+          if (m.role === "assistant" && m.thinking) {
+            out.reasoning_content = m.thinking;
+          }
+          return out;
         }),
       ];
 
@@ -2397,7 +2537,7 @@ class AppStore {
 
       let streamedContent = "";
       let streamedThinking = "";
-
+      let serverTpsReceived = false;
       interface ChatCompletionChunk {
         choices?: Array<{
           delta?: { content?: string; reasoning_content?: string };
@@ -2462,7 +2602,6 @@ class AppStore {
             tokenCount += 1;
             this.totalTokens = tokenCount;
 
-            // Update real-time TPS during streaming
             if (firstTokenTime !== null && tokenCount > 1) {
               const elapsed = performance.now() - firstTokenTime;
               this.tps = (tokenCount / elapsed) * 1000;
@@ -2513,16 +2652,24 @@ class AppStore {
               startedAt: this.prefillProgress?.startedAt ?? performance.now(),
             };
           },
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+              serverTpsReceived = true;
+            }
+          },
         },
       );
 
       // Clear prefill progress after stream ends
       this.prefillProgress = null;
 
-      // Calculate final TPS
-      if (firstTokenTime !== null && tokenCount > 1) {
+      // Use server-side TPS if available, otherwise fall back to client-side
+      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
+        this.tps = (tokenCount / totalGenerationTime) * 1000;
       }
 
       // Final cleanup of the message (if conversation still exists)
@@ -2627,6 +2774,9 @@ class AppStore {
     this.syncActiveMessagesIfNeeded(targetConversationId);
     this.saveConversationsToStorage();
 
+    const abortController = new AbortController();
+    this.currentAbortController = abortController;
+
     try {
       // Determine the model to use
       const model = this.getModelForRequest(modelId);
@@ -2681,6 +2831,7 @@ class AppStore {
           "Content-Type": "application/json",
         },
         body: JSON.stringify(requestBody),
+        signal: abortController.signal,
       });
 
       if (!response.ok) {
@@ -2820,14 +2971,27 @@ class AppStore {
         );
       }
     } catch (error) {
-      console.error("Error generating image:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to generate image",
-      );
+      if (abortController.signal.aborted) {
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "Cancelled";
+            msg.attachments = [];
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        console.error("Error generating image:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to generate image",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
       this.isLoading = false;
       this.saveConversationsToStorage();
     }
@@ -2891,6 +3055,9 @@ class AppStore {
     // Clear editing state
     this.editingImage = null;
 
+    const abortController = new AbortController();
+    this.currentAbortController = abortController;
+
     try {
       // Determine the model to use
       const model = this.getModelForRequest(modelId);
@@ -2952,6 +3119,7 @@ class AppStore {
       const apiResponse = await fetch("/v1/images/edits", {
         method: "POST",
         body: formData,
+        signal: abortController.signal,
       });
 
       if (!apiResponse.ok) {
@@ -3052,14 +3220,27 @@ class AppStore {
         );
       }
     } catch (error) {
-      console.error("Error editing image:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to edit image",
-      );
+      if (abortController.signal.aborted) {
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "Cancelled";
+            msg.attachments = [];
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        console.error("Error editing image:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to edit image",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
       this.isLoading = false;
       this.saveConversationsToStorage();
     }
@@ -3126,6 +3307,85 @@ class AppStore {
     }
   }
 
+  /**
+   * Cancel/pause an active download on a specific node
+   */
+  async cancelDownload(nodeId: string, modelId: string): Promise<void> {
+    try {
+      const response = await fetch("/download/cancel", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          targetNodeId: nodeId,
+          modelId: modelId,
+        }),
+      });
+      if (!response.ok) {
+        const errorText = await response.text();
+        throw new Error(
+          `Failed to cancel download: ${response.status} - ${errorText}`,
+        );
+      }
+    } catch (error) {
+      console.error("Error cancelling download:", error);
+      throw error;
+    }
+  }
+
+  async createInstanceLink(
+    prefillInstances: string[],
+    decodeInstances: string[],
+  ): Promise<void> {
+    const response = await fetch("/v1/instance-links", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({
+        prefill_instances: prefillInstances,
+        decode_instances: decodeInstances,
+      }),
+    });
+    if (!response.ok) {
+      throw new Error(
+        `Failed to create instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
+  async updateInstanceLink(
+    linkId: string,
+    prefillInstances: string[],
+    decodeInstances: string[],
+  ): Promise<void> {
+    const response = await fetch(
+      `/v1/instance-links/${encodeURIComponent(linkId)}`,
+      {
+        method: "PUT",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          prefill_instances: prefillInstances,
+          decode_instances: decodeInstances,
+        }),
+      },
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to update instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
+  async deleteInstanceLink(linkId: string): Promise<void> {
+    const response = await fetch(
+      `/v1/instance-links/${encodeURIComponent(linkId)}`,
+      { method: "DELETE" },
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to delete instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
   /**
    * Delete a downloaded model from a specific node
    */
@@ -3224,6 +3484,19 @@ export const prefillProgress = () => appStore.prefillProgress;
 export const topologyData = () => appStore.topologyData;
 export const instances = () => appStore.instances;
 export const runners = () => appStore.runners;
+export const instanceLinks = () => appStore.instanceLinks;
+export const featureFlags = () => appStore.featureFlags;
+export const createInstanceLink = (
+  prefillInstances: string[],
+  decodeInstances: string[],
+) => appStore.createInstanceLink(prefillInstances, decodeInstances);
+export const updateInstanceLink = (
+  linkId: string,
+  prefillInstances: string[],
+  decodeInstances: string[],
+) => appStore.updateInstanceLink(linkId, prefillInstances, decodeInstances);
+export const deleteInstanceLink = (linkId: string) =>
+  appStore.deleteInstanceLink(linkId);
 export const downloads = () => appStore.downloads;
 export const nodeDisk = () => appStore.nodeDisk;
 export const placementPreviews = () => appStore.placementPreviews;
@@ -3248,6 +3521,7 @@ export const sendMessage = (
     type: string;
     textContent?: string;
     preview?: string;
+    pageImages?: string[];
   }[],
   enableThinking?: boolean | null,
 ) => appStore.sendMessage(content, files, enableThinking);
@@ -3346,6 +3620,8 @@ export const resetImageGenerationParams = () =>
 // Download actions
 export const startDownload = (nodeId: string, shardMetadata: object) =>
   appStore.startDownload(nodeId, shardMetadata);
+export const cancelDownload = (nodeId: string, modelId: string) =>
+  appStore.cancelDownload(nodeId, modelId);
 export const deleteDownload = (nodeId: string, modelId: string) =>
   appStore.deleteDownload(nodeId, modelId);
 

dashboard/src/lib/types/files.ts

@@ -2,6 +2,41 @@
  * File attachment types for the chat interface
  */
 
+import { getDocument, GlobalWorkerOptions, version } from "pdfjs-dist";
+import type { DocumentInitParameters } from "pdfjs-dist/types/src/display/api";
+
+// Safari (through at least 18/26) does not implement
+// ReadableStream.prototype[Symbol.asyncIterator], which pdfjs-dist uses
+// internally in getTextContent(). Without this polyfill, `for await (const n
+// of stream)` throws "undefined is not a function" and PDF processing fails.
+if (
+  typeof ReadableStream !== "undefined" &&
+  !(ReadableStream.prototype as unknown as Record<symbol, unknown>)[
+    Symbol.asyncIterator
+  ]
+) {
+  (ReadableStream.prototype as unknown as Record<symbol, unknown>)[
+    Symbol.asyncIterator
+  ] = async function* (this: ReadableStream<unknown>) {
+    const reader = this.getReader();
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) return;
+        yield value;
+      }
+    } finally {
+      reader.releaseLock();
+    }
+  };
+}
+
+GlobalWorkerOptions.workerSrc = `https://cdn.jsdelivr.net/npm/pdfjs-dist@${version}/build/pdf.worker.mjs`;
+
+const PDF_PAGE_SCALE = 2.0;
+const PDF_MAX_PAGES = 20;
+const PDF_MAX_TEXT_CHARS = 100_000;
+
 export interface ChatUploadedFile {
   id: string;
   name: string;
@@ -10,6 +45,7 @@ export interface ChatUploadedFile {
   file: File;
   preview?: string;
   textContent?: string;
+  pageImages?: string[];
 }
 
 export interface ChatAttachment {
@@ -194,6 +230,58 @@ export function readFileAsText(file: File): Promise<string> {
   });
 }
 
+async function extractPdfContent(
+  file: File,
+): Promise<{ text: string; pageImages: string[] }> {
+  const arrayBuffer = await file.arrayBuffer();
+  const pdf = await getDocument({
+    data: new Uint8Array(arrayBuffer),
+    useSystemFonts: true,
+  } as DocumentInitParameters).promise;
+
+  const numPages = Math.min(pdf.numPages, PDF_MAX_PAGES);
+  const pageTexts: string[] = [];
+  const pageImages: string[] = [];
+
+  for (let i = 1; i <= numPages; i++) {
+    const page = await pdf.getPage(i);
+
+    const content = await page.getTextContent();
+    const strings = content.items
+      .filter((item: any) => "str" in item)
+      .map((item: any) => item.str as string);
+    pageTexts.push(strings.join(" "));
+
+    const viewport = page.getViewport({ scale: PDF_PAGE_SCALE });
+    const canvas = new OffscreenCanvas(viewport.width, viewport.height);
+    const ctx = canvas.getContext("2d");
+    if (ctx) {
+      await page.render({ canvasContext: ctx as any, viewport }).promise;
+      const blob = await canvas.convertToBlob({
+        type: "image/jpeg",
+        quality: 0.8,
+      });
+      const reader = new FileReader();
+      const dataUrl = await new Promise<string>((resolve, reject) => {
+        reader.onload = () => resolve(reader.result as string);
+        reader.onerror = () => reject(reader.error);
+        reader.readAsDataURL(blob);
+      });
+      pageImages.push(dataUrl);
+    }
+  }
+
+  let text = pageTexts.join("\n\n").trim();
+  if (text.length > PDF_MAX_TEXT_CHARS) {
+    text = text.slice(0, PDF_MAX_TEXT_CHARS) + "\n\n[truncated]";
+  }
+  if (pdf.numPages > PDF_MAX_PAGES) {
+    text += `\n\n[showing ${PDF_MAX_PAGES} of ${pdf.numPages} pages]`;
+  }
+
+  return { text, pageImages };
+}
+
 /**
  * Process uploaded files into ChatUploadedFile format
  */
@@ -223,7 +311,12 @@ export async function processUploadedFiles(
         const textContent = await readFileAsText(file);
         results.push({ ...base, textContent });
       } else if (category === "pdf") {
-        results.push(base);
+        const { text, pageImages } = await extractPdfContent(file);
+        results.push({
+          ...base,
+          textContent: text || undefined,
+          pageImages: pageImages.length > 0 ? pageImages : undefined,
+        });
       } else if (category === "audio") {
         const preview = await readFileAsDataURL(file);
         results.push({ ...base, preview });

dashboard/src/lib/utils/clipboard.ts

@@ -0,0 +1,55 @@
+export async function copyText(text: string): Promise<boolean> {
+  if (
+    typeof window !== "undefined" &&
+    window.isSecureContext &&
+    navigator.clipboard?.writeText
+  ) {
+    try {
+      await navigator.clipboard.writeText(text);
+      return true;
+    } catch {
+      // fall through to execCommand fallback
+    }
+  }
+
+  if (typeof document === "undefined") {
+    return false;
+  }
+
+  const textarea = document.createElement("textarea");
+  textarea.value = text;
+  textarea.setAttribute("readonly", "");
+  textarea.style.position = "fixed";
+  textarea.style.top = "0";
+  textarea.style.left = "0";
+  textarea.style.width = "1px";
+  textarea.style.height = "1px";
+  textarea.style.padding = "0";
+  textarea.style.border = "none";
+  textarea.style.outline = "none";
+  textarea.style.boxShadow = "none";
+  textarea.style.background = "transparent";
+  textarea.style.opacity = "0";
+  document.body.appendChild(textarea);
+
+  const previousSelection = document.getSelection();
+  const previousRange =
+    previousSelection && previousSelection.rangeCount > 0
+      ? previousSelection.getRangeAt(0)
+      : null;
+
+  try {
+    textarea.focus();
+    textarea.select();
+    textarea.setSelectionRange(0, text.length);
+    return document.execCommand("copy");
+  } catch {
+    return false;
+  } finally {
+    document.body.removeChild(textarea);
+    if (previousRange && previousSelection) {
+      previousSelection.removeAllRanges();
+      previousSelection.addRange(previousRange);
+    }
+  }
+}

dashboard/src/lib/utils/model_family.ts

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

dashboard/src/routes/+page.svelte

@@ -42,6 +42,7 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
+    thinkingEnabled,
     generateImage,
     editImage,
     editingImage,
@@ -264,6 +265,7 @@
 
   let mounted = $state(false);
   let localNodeId = $state<string | null>(null);
+  let pendingFirefoxQuery = $state<string | null>(null); // ?q= param deferred until state loads
 
   // ── Onboarding wizard state ──
   const ONBOARDING_COMPLETE_KEY = "exo-onboarding-complete";
@@ -852,7 +854,7 @@
   ) {
     const model = selectedChatModel();
     if (!model) {
-      sendMessage(content, files, null);
+      sendMessage(content, files, thinkingEnabled());
       return;
     }
 
@@ -880,7 +882,7 @@
     }
 
     // Default: text chat
-    sendMessage(content, files, null);
+    sendMessage(content, files, thinkingEnabled());
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
@@ -1308,6 +1310,20 @@
       return;
     }
 
+    // Firefox AI sidebar integration: handle ?q= query parameter
+    // Firefox's built-in AI sidebar (about:config: browser.ml.chat.enabled) sends
+    // the user's prompt as ?q=<URL-encoded prompt> to the configured provider URL.
+    // See: https://support.mozilla.org/en-US/kb/ai-chatbot
+    const queryParam = params.get("q");
+    if (queryParam) {
+      // Clean up the URL to prevent re-submission on page refresh
+      window.history.replaceState({}, "", window.location.pathname);
+      // Defer the auto-send until cluster state is loaded (topologyData,
+      // instances, availableMemory) so that model auto-selection works
+      // correctly. The $effect below will pick this up once data is ready.
+      pendingFirefoxQuery = queryParam;
+    }
+
     // Check server-side onboarding state (persisted in ~/.exo)
     try {
       const res = await fetch("/onboarding");
@@ -1328,6 +1344,18 @@
     }
   });
 
+  // Deferred Firefox AI sidebar auto-send: wait for cluster state and model
+  // list before submitting. Both data (from /state polling) and models (from
+  // the async /models fetch in onMount) must be loaded for handleAutoSend to
+  // correctly auto-select a model.
+  $effect(() => {
+    if (pendingFirefoxQuery && data && models.length > 0) {
+      const query = pendingFirefoxQuery;
+      pendingFirefoxQuery = null;
+      handleChatSend(query);
+    }
+  });
+
   async function fetchModels() {
     try {
       const response = await fetch("/models");
@@ -1535,34 +1563,44 @@
   }
 
   // Helper to get download status for a model (checks all downloads for matching model ID)
-  function getModelDownloadStatus(modelId: string): {
+  type NodeDownloadStatus = {
+    nodeId: string;
+    nodeName: string;
+    status: "completed" | "partial" | "pending" | "downloading";
+    percentage: number;
+    progress: DownloadProgress | null;
+  };
+
+  // Shared helper: collect per-node download status for a model across a set of nodes.
+  // Handles deduplication, entry parsing, and aggregation in one place.
+  function collectDownloadStatus(
+    modelId: string,
+    nodeIds?: string[],
+  ): {
     isDownloading: boolean;
     progress: DownloadProgress | null;
-    perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }>;
+    perNode: NodeDownloadStatus[];
+    failedError: string | null;
   } {
+    const empty = {
+      isDownloading: false,
+      progress: null,
+      perNode: [] as NodeDownloadStatus[],
+      failedError: null,
+    };
+
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
-      return { isDownloading: false, progress: null, perNode: [] };
+      return empty;
     }
 
-    let totalBytes = 0;
-    let downloadedBytes = 0;
-    let totalSpeed = 0;
-    let completedFiles = 0;
-    let totalFiles = 0;
-    let isDownloading = false;
-    const allFiles: DownloadProgress["files"] = [];
-    const perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }> = [];
+    // Deduplicate by nodeId — a node can have multiple entries for the same model
+    // (e.g. PipelineShardMetadata + TensorShardMetadata). Keep the last entry,
+    // which is the most recently applied event.
+    const perNodeMap = new Map<string, NodeDownloadStatus>();
 
-    // Check all nodes for downloads matching this model
+    const nodeIdSet = nodeIds ? new Set(nodeIds) : null;
     for (const [nodeId, nodeDownloads] of Object.entries(downloadsData)) {
+      if (nodeIdSet && !nodeIdSet.has(nodeId)) continue;
       if (!Array.isArray(nodeDownloads)) continue;
 
       for (const downloadWrapped of nodeDownloads) {
@@ -1575,29 +1613,45 @@
         const downloadPayload = (downloadWrapped as Record<string, unknown>)[
           downloadKind
         ] as Record<string, unknown>;
-
-        if (
-          downloadKind !== "DownloadOngoing" &&
-          downloadKind !== "DownloadPending"
-        )
-          continue;
         if (!downloadPayload) continue;
 
         const downloadModelId = extractModelIdFromDownload(downloadPayload);
+        if (!downloadModelId || downloadModelId !== modelId) continue;
 
-        // Match if the model ID contains or equals the requested model
-        // (handles cases like "mlx-community/Meta-Llama..." matching)
-        if (
-          !downloadModelId ||
-          !downloadModelId.includes(modelId.split("/").pop() || modelId)
-        ) {
-          // Try exact match or partial match
-          if (downloadModelId !== modelId) continue;
+        // DownloadFailed — return with any data collected so far
+        if (downloadKind === "DownloadFailed") {
+          return {
+            isDownloading: false,
+            progress: null,
+            perNode: Array.from(perNodeMap.values()),
+            failedError:
+              (downloadPayload.errorMessage as string) ||
+              (downloadPayload.error_message as string) ||
+              "Download failed",
+          };
+        }
+
+        if (
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending" &&
+          downloadKind !== "DownloadCompleted"
+        )
+          continue;
+
+        const nodeName =
+          data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+
+        if (downloadKind === "DownloadCompleted") {
+          perNodeMap.set(nodeId, {
+            nodeId,
+            nodeName,
+            status: "completed",
+            percentage: 100,
+            progress: null,
+          });
+          continue;
         }
 
-        // For DownloadPending with partial bytes (paused/resumed downloads),
-        // synthesize a progress object from the top-level downloaded/total fields
-        let progress: DownloadProgress | null;
         if (downloadKind === "DownloadPending") {
           const pendingDownloaded = getBytes(
             downloadPayload.downloaded ??
@@ -1610,44 +1664,67 @@
               downloadPayload.totalBytes,
           );
           if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
-          isDownloading = true;
-          progress = {
-            totalBytes: pendingTotal,
-            downloadedBytes: pendingDownloaded,
-            speed: 0,
-            etaMs: 0,
-            percentage:
-              pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
-            completedFiles: 0,
-            totalFiles: 0,
-            files: [],
-          };
-        } else {
-          isDownloading = true;
-          progress = parseDownloadProgress(downloadPayload);
+          const pct =
+            pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0;
+          perNodeMap.set(nodeId, {
+            nodeId,
+            nodeName,
+            status: pendingDownloaded > 0 ? "partial" : "pending",
+            percentage: pct,
+            progress: null,
+          });
+          continue;
         }
 
-        if (progress) {
-          // Sum all values across nodes - each node downloads independently
-          totalBytes += progress.totalBytes;
-          downloadedBytes += progress.downloadedBytes;
-          totalSpeed += progress.speed;
-          completedFiles += progress.completedFiles;
-          totalFiles += progress.totalFiles;
-          allFiles.push(...progress.files);
+        // DownloadOngoing
+        const progress = parseDownloadProgress(downloadPayload);
+        if (
+          !progress ||
+          (progress.downloadedBytes <= 0 && progress.totalBytes <= 0)
+        )
+          continue;
 
-          const nodeName =
-            data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
-          perNode.push({ nodeId, nodeName, progress });
-        }
+        perNodeMap.set(nodeId, {
+          nodeId,
+          nodeName,
+          status: "downloading",
+          percentage: progress.percentage,
+          progress,
+        });
+      }
+    }
+
+    // Aggregate from deduplicated per-node entries
+    const perNode = Array.from(perNodeMap.values());
+    let totalBytes = 0;
+    let downloadedBytes = 0;
+    let totalSpeed = 0;
+    let completedFiles = 0;
+    let totalFiles = 0;
+    let isDownloading = false;
+    const allFiles: DownloadProgress["files"] = [];
+
+    for (const node of perNode) {
+      if (node.status === "downloading" && node.progress) {
+        isDownloading = true;
+        totalBytes += node.progress.totalBytes;
+        downloadedBytes += node.progress.downloadedBytes;
+        totalSpeed += node.progress.speed;
+        completedFiles += node.progress.completedFiles;
+        totalFiles += node.progress.totalFiles;
+        allFiles.push(...node.progress.files);
       }
     }
 
     if (!isDownloading) {
-      return { isDownloading: false, progress: null, perNode: [] };
+      return {
+        isDownloading: false,
+        progress: null,
+        perNode,
+        failedError: null,
+      };
     }
 
-    // ETA = total remaining bytes / total speed across all nodes
     const remainingBytes = totalBytes - downloadedBytes;
     const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
 
@@ -1664,9 +1741,21 @@
         files: allFiles,
       },
       perNode,
+      failedError: null,
     };
   }
 
+  function getModelDownloadStatus(
+    modelId: string,
+    nodeIds?: string[],
+  ): {
+    isDownloading: boolean;
+    progress: DownloadProgress | null;
+    perNode: NodeDownloadStatus[];
+  } {
+    return collectDownloadStatus(modelId, nodeIds);
+  }
+
   // Helper to get download status for an instance
   function getInstanceDownloadStatus(
     instanceId: string,
@@ -1677,26 +1766,9 @@
     errorMessage: string | null;
     progress: DownloadProgress | null;
     statusText: string;
-    perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }>;
+    perNode: NodeDownloadStatus[];
   } {
-    if (!downloadsData || Object.keys(downloadsData).length === 0) {
-      // No download data yet — defer to runner status instead of assuming RUNNING
-      const statusInfo = deriveInstanceStatus(instanceWrapped);
-      return {
-        isDownloading: false,
-        isFailed: false,
-        errorMessage: null,
-        progress: null,
-        statusText: statusInfo.statusText,
-        perNode: [],
-      };
-    }
-
-    // Unwrap the instance
+    // Unwrap the instance to get shard assignments
     const [instanceTag, instance] = getTagged(instanceWrapped);
     if (!instance || typeof instance !== "object") {
       return {
@@ -1716,132 +1788,9 @@
         modelId?: string;
       };
     };
-    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
-    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
     const instanceModelId = inst.shardAssignments?.modelId;
 
-    // Build reverse mapping: runnerId -> nodeId
-    const runnerToNode: Record<string, string> = {};
-    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
-      runnerToNode[runnerId] = nodeId;
-    }
-
-    let totalBytes = 0;
-    let downloadedBytes = 0;
-    let totalSpeed = 0;
-    let completedFiles = 0;
-    let totalFiles = 0;
-    let isDownloading = false;
-    const allFiles: DownloadProgress["files"] = [];
-    const perNode: Array<{
-      nodeId: string;
-      nodeName: string;
-      progress: DownloadProgress;
-    }> = [];
-
-    // Check downloads for nodes that are part of this instance
-    for (const runnerId of Object.keys(runnerToShard)) {
-      const nodeId = runnerToNode[runnerId];
-      if (!nodeId) continue;
-
-      const nodeDownloads = downloadsData[nodeId];
-      if (!Array.isArray(nodeDownloads)) continue;
-
-      for (const downloadWrapped of nodeDownloads) {
-        if (!downloadWrapped || typeof downloadWrapped !== "object") continue;
-
-        const keys = Object.keys(downloadWrapped as Record<string, unknown>);
-        if (keys.length !== 1) continue;
-
-        const downloadKind = keys[0];
-        const downloadPayload = (downloadWrapped as Record<string, unknown>)[
-          downloadKind
-        ] as Record<string, unknown>;
-
-        // Handle DownloadFailed - return immediately with error info
-        if (downloadKind === "DownloadFailed") {
-          const downloadModelId = extractModelIdFromDownload(downloadPayload);
-          if (
-            instanceModelId &&
-            downloadModelId &&
-            downloadModelId === instanceModelId
-          ) {
-            return {
-              isDownloading: false,
-              isFailed: true,
-              errorMessage:
-                (downloadPayload.errorMessage as string) || "Download failed",
-              progress: null,
-              statusText: "FAILED",
-              perNode: [],
-            };
-          }
-        }
-
-        if (
-          downloadKind !== "DownloadOngoing" &&
-          downloadKind !== "DownloadPending"
-        )
-          continue;
-        if (!downloadPayload) continue;
-
-        // Check if this download is for this instance's model
-        const downloadModelId = extractModelIdFromDownload(downloadPayload);
-        if (
-          instanceModelId &&
-          downloadModelId &&
-          downloadModelId === instanceModelId
-        ) {
-          // For DownloadPending with partial bytes, synthesize progress
-          let progress: DownloadProgress | null;
-          if (downloadKind === "DownloadPending") {
-            const pendingDownloaded = getBytes(
-              downloadPayload.downloaded ??
-                downloadPayload.downloaded_bytes ??
-                downloadPayload.downloadedBytes,
-            );
-            const pendingTotal = getBytes(
-              downloadPayload.total ??
-                downloadPayload.total_bytes ??
-                downloadPayload.totalBytes,
-            );
-            if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
-            isDownloading = true;
-            progress = {
-              totalBytes: pendingTotal,
-              downloadedBytes: pendingDownloaded,
-              speed: 0,
-              etaMs: 0,
-              percentage:
-                pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
-              completedFiles: 0,
-              totalFiles: 0,
-              files: [],
-            };
-          } else {
-            isDownloading = true;
-            progress = parseDownloadProgress(downloadPayload);
-          }
-
-          if (progress) {
-            // Sum all values across nodes - each node downloads independently
-            totalBytes += progress.totalBytes;
-            downloadedBytes += progress.downloadedBytes;
-            totalSpeed += progress.speed;
-            completedFiles += progress.completedFiles;
-            totalFiles += progress.totalFiles;
-            allFiles.push(...progress.files);
-
-            const nodeName =
-              data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
-            perNode.push({ nodeId, nodeName, progress });
-          }
-        }
-      }
-    }
-
-    if (!isDownloading) {
-      // Check runner status for other states
+    if (!instanceModelId) {
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
         isDownloading: false,
@@ -1853,26 +1802,49 @@
       };
     }
 
-    // ETA = total remaining bytes / total speed across all nodes
-    const remainingBytes = totalBytes - downloadedBytes;
-    const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
+    // Get node IDs assigned to this instance
+    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
+    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const runnerToNode: Record<string, string> = {};
+    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
+      runnerToNode[runnerId] = nodeId;
+    }
+    const instanceNodeIds = Object.keys(runnerToShard)
+      .map((runnerId) => runnerToNode[runnerId])
+      .filter(Boolean);
+
+    const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
+
+    if (result.failedError) {
+      return {
+        isDownloading: false,
+        isFailed: true,
+        errorMessage: result.failedError,
+        progress: null,
+        statusText: "FAILED",
+        perNode: [],
+      };
+    }
+
+    if (!result.isDownloading) {
+      const statusInfo = deriveInstanceStatus(instanceWrapped);
+      return {
+        isDownloading: false,
+        isFailed: statusInfo.statusText === "FAILED",
+        errorMessage: null,
+        progress: null,
+        statusText: statusInfo.statusText,
+        perNode: result.perNode,
+      };
+    }
 
     return {
       isDownloading: true,
       isFailed: false,
       errorMessage: null,
-      progress: {
-        totalBytes,
-        downloadedBytes,
-        speed: totalSpeed,
-        etaMs,
-        percentage: totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0,
-        completedFiles,
-        totalFiles,
-        files: allFiles,
-      },
+      progress: result.progress,
       statusText: "DOWNLOADING",
-      perNode,
+      perNode: result.perNode,
     };
   }
 
@@ -3463,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"
@@ -4630,7 +4603,7 @@
                 type="button"
                 onclick={() => {
                   completeOnboarding();
-                  sendMessage(chip);
+                  sendMessage(chip, undefined, thinkingEnabled());
                 }}
                 class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
               >
@@ -4850,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"
@@ -4996,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
@@ -5369,10 +5344,10 @@
                             <div
                               class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                             >
-                              {#each downloadInfo.perNode as nodeProg}
+                              {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
                                 {@const nodePercent = Math.min(
                                   100,
-                                  Math.max(0, nodeProg.progress.percentage),
+                                  Math.max(0, nodeProg.percentage),
                                 )}
                                 {@const isExpanded =
                                   instanceDownloadExpandedNodes.has(
@@ -5428,15 +5403,17 @@
                                     >
                                       <span
                                         >{formatBytes(
-                                          nodeProg.progress.downloadedBytes,
+                                          nodeProg.progress?.downloadedBytes ??
+                                            0,
                                         )} / {formatBytes(
-                                          nodeProg.progress.totalBytes,
+                                          nodeProg.progress?.totalBytes ?? 0,
                                         )}</span
                                       >
                                       <span
-                                        >{formatSpeed(nodeProg.progress.speed)} •
-                                        ETA {formatEta(
-                                          nodeProg.progress.etaMs,
+                                        >{formatSpeed(
+                                          nodeProg.progress?.speed ?? 0,
+                                        )} • ETA {formatEta(
+                                          nodeProg.progress?.etaMs ?? 0,
                                         )}</span
                                       >
                                     </div>
@@ -5444,14 +5421,14 @@
 
                                   {#if isExpanded}
                                     <div class="mt-2 space-y-1.5">
-                                      {#if nodeProg.progress.files.length === 0}
+                                      {#if nodeProg.progress?.files ?? [].length === 0}
                                         <div
                                           class="text-[11px] font-mono text-exo-light-gray/70"
                                         >
                                           No file details reported.
                                         </div>
                                       {:else}
-                                        {#each nodeProg.progress.files as f}
+                                        {#each nodeProg.progress?.files ?? [] as f}
                                           {@const filePercent = Math.min(
                                             100,
                                             Math.max(0, f.percentage ?? 0),
@@ -5927,12 +5904,15 @@
                 )}
                 {@const allPreviews = filteredPreviews()}
                 {#if selectedModel && allPreviews.length > 0}
-                  {@const downloadStatus = getModelDownloadStatus(
-                    selectedModel.id,
-                  )}
                   {@const tags = modelTags()[selectedModel.id] || []}
                   <div class="space-y-3">
                     {#each allPreviews as apiPreview, i}
+                      {@const downloadStatus = getModelDownloadStatus(
+                        selectedModel.id,
+                        apiPreview.memory_delta_by_node
+                          ? Object.keys(apiPreview.memory_delta_by_node)
+                          : undefined,
+                      )}
                       <div
                         role="group"
                         onmouseenter={() => {
@@ -6120,7 +6100,7 @@
                           onclick={() => {
                             chatLaunchState = "idle";
                             selectedChatCategory = null;
-                            sendMessage(prompt);
+                            sendMessage(prompt, undefined, thinkingEnabled());
                           }}
                           class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
                         >
@@ -6503,10 +6483,10 @@
                               <div
                                 class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                               >
-                                {#each downloadInfo.perNode as nodeProg}
+                                {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
                                   {@const nodePercent = Math.min(
                                     100,
-                                    Math.max(0, nodeProg.progress.percentage),
+                                    Math.max(0, nodeProg.percentage),
                                   )}
                                   {@const isExpanded =
                                     instanceDownloadExpandedNodes.has(
@@ -6565,16 +6545,17 @@
                                       >
                                         <span
                                           >{formatBytes(
-                                            nodeProg.progress.downloadedBytes,
+                                            nodeProg.progress
+                                              ?.downloadedBytes ?? 0,
                                           )} / {formatBytes(
-                                            nodeProg.progress.totalBytes,
+                                            nodeProg.progress?.totalBytes ?? 0,
                                           )}</span
                                         >
                                         <span
                                           >{formatSpeed(
-                                            nodeProg.progress.speed,
+                                            nodeProg.progress?.speed ?? 0,
                                           )} • ETA {formatEta(
-                                            nodeProg.progress.etaMs,
+                                            nodeProg.progress?.etaMs ?? 0,
                                           )}</span
                                         >
                                       </div>
@@ -6582,14 +6563,14 @@
 
                                     {#if isExpanded}
                                       <div class="mt-2 space-y-1.5">
-                                        {#if nodeProg.progress.files.length === 0}
+                                        {#if nodeProg.progress?.files ?? [].length === 0}
                                           <div
                                             class="text-[11px] font-mono text-exo-light-gray/70"
                                           >
                                             No file details reported.
                                           </div>
                                         {:else}
-                                          {#each nodeProg.progress.files as f}
+                                          {#each nodeProg.progress?.files ?? [] as f}
                                             {@const filePercent = Math.min(
                                               100,
                                               Math.max(0, f.percentage ?? 0),