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

.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 }}
@@ -347,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

.gitignore

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

.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``.
 
@@ -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/layers/base.pyi

@@ -57,6 +57,10 @@ class Module(dict):
     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_lm/generate.pyi

@@ -383,11 +383,12 @@ class GenerationBatch:
     state_machines: List[SequenceStateMachine]
     max_tokens: List[int]
     _current_tokens: Optional[mx.array]
-    _current_logprobs: List[mx.array]
-    _next_tokens: mx.array
-    _next_logprobs: List[mx.array]
-    _token_context: List[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,

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

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

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

@@ -550,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/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

@@ -21,8 +21,6 @@ struct SettingsView: View {
     @State private var pendingReadOnlyModelsDirs: String = ""
     @State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
     @State private var needsRestart = false
-    @State private var bugReportInFlight = false
-    @State private var bugReportMessage: String?
     @State private var uninstallInProgress = false
 
     var body: some View {
@@ -202,8 +200,6 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
-
-                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -504,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")
 
@@ -570,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()
@@ -584,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()
@@ -607,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 {

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/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,19 +3,20 @@ 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,
@@ -209,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
@@ -276,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
@@ -379,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:
@@ -489,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:
@@ -913,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="*",
@@ -935,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]
@@ -1010,42 +1024,72 @@ Examples:
     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}")

bench/exo_bench.py

@@ -24,17 +24,15 @@ 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,
@@ -45,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:
@@ -229,105 +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,
-    *,
-    use_prefix_cache: bool = False,
-) -> tuple[dict[str, Any], int]:
-    content, pp_tokens = prompt_sizer.build(pp_hint)
-    payload: dict[str, Any] = {
-        "model": model_id,
-        "messages": [{"role": "user", "content": content}],
-        "stream": False,
-        "max_tokens": tg,
-        "logprobs": False,
-        "use_prefix_cache": use_prefix_cache,
-    }
-
-    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",
@@ -375,6 +233,11 @@ 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",
@@ -440,81 +303,124 @@ 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:
+        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),
@@ -523,16 +429,20 @@ def main() -> int:
             )
             sampler.start()
 
+        def _do_one(c: ExoClient, pp: int, tg: int) -> tuple[dict[str, Any], int]:
+            return run_one_completion(
+                c,
+                full_model_id,
+                pp,
+                tg,
+                prompt_sizer,
+                use_prefix_cache=args.use_prefix_cache,
+                stream=args.stream,
+            )
+
         try:
             for i in range(args.warmup):
-                run_one_completion(
-                    client,
-                    full_model_id,
-                    pp_list[0],
-                    tg_list[0],
-                    prompt_sizer,
-                    use_prefix_cache=args.use_prefix_cache,
-                )
+                _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)
@@ -554,14 +464,7 @@ def main() -> int:
                             # Sequential: single request
                             try:
                                 inf_t0 = time.monotonic()
-                                row, actual_pp_tokens = run_one_completion(
-                                    client,
-                                    full_model_id,
-                                    pp,
-                                    tg,
-                                    prompt_sizer,
-                                    use_prefix_cache=args.use_prefix_cache,
-                                )
+                                row, actual_pp_tokens = _do_one(client, pp, tg)
                                 inference_windows.append((inf_t0, time.monotonic()))
                             except Exception as e:
                                 logger.error(e)
@@ -710,10 +613,12 @@ def main() -> int:
                         gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
-                        peak = mean(
-                            x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
-                        )
 
+                        def _peak_bytes(s: dict[str, Any]) -> float:
+                            pm = s["peak_memory_usage"]
+                            return pm.get("inBytes") or pm.get("in_bytes", 0)
+
+                        peak = mean(_peak_bytes(x["stats"]) for x in runs)
                         summary = (
                             f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                             f"prompt_tokens={ptok} gen_tokens={gtok}    "
@@ -738,15 +643,16 @@ def main() -> int:
                 if placement_metrics:
                     all_system_metrics.update(placement_metrics)
 
-            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 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)
+                time.sleep(5)
 
     output: dict[str, Any] = {"runs": all_rows}
     if cluster_snapshot:

bench/exo_eval.py

@@ -42,11 +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,
@@ -62,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
@@ -271,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
@@ -281,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
@@ -517,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(
@@ -530,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:
@@ -546,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",
@@ -554,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,
@@ -578,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,
@@ -592,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(
@@ -618,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
@@ -652,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
@@ -660,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":
@@ -667,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)
@@ -697,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,
@@ -729,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,
@@ -749,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,
@@ -763,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,
@@ -778,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",
@@ -793,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",
@@ -804,7 +984,7 @@ async def evaluate_benchmark(
                     )
                 else:
                     result = QuestionResult(
-                        question_id=idx,
+                        question_id=question_id,
                         prompt=prompt,
                         response=response,
                         extracted_answer=None,
@@ -815,7 +995,7 @@ async def evaluate_benchmark(
                     )
             else:
                 result = QuestionResult(
-                    question_id=idx,
+                    question_id=question_id,
                     prompt=prompt,
                     response=response,
                     extracted_answer=None,
@@ -827,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
 # ---------------------------------------------------------------------------
@@ -867,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%})")
@@ -878,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:
@@ -896,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,
     }
 
 
@@ -1053,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
         ],
@@ -1069,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:
@@ -1096,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(
@@ -1115,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."
     )
@@ -1148,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()
@@ -1177,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,
@@ -1198,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,
         )
@@ -1225,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)
@@ -1234,10 +1542,9 @@ def main() -> int:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             return 1
         time.sleep(1)
-        cluster_snapshot = capture_cluster_snapshot(client)
-    else:
-        full_model_id = args.model
-        cluster_snapshot = None
+        created_instance = True
+
+    cluster_snapshot = capture_cluster_snapshot(client)
 
     # Auto-detect reasoning from model config
     model_config = load_model_config(full_model_id)
@@ -1291,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)
@@ -1309,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,
@@ -1319,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:
@@ -1336,10 +1688,18 @@ def main() -> int:
                             cluster=cluster_snapshot,
                         )
                         results_by_c[c] = results
+                    # Clean up checkpoint on success
+                    if checkpoint_path.exists():
+                        checkpoint_path.unlink()
                 if len(results_by_c) >= 2:
                     print_comparison(task_name, results_by_c)
         else:
             for task_name in task_names:
+                checkpoint_path = _checkpoint_path(
+                    args.results_dir, task_name, full_model_id, args.num_concurrent
+                )
+                if args.force and checkpoint_path.exists():
+                    checkpoint_path.unlink()
                 results = asyncio.run(
                     evaluate_benchmark(
                         task_name,
@@ -1350,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:
@@ -1366,14 +1725,25 @@ def main() -> int:
                         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/src/lib/components/HeaderNav.svelte

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

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -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/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>;
@@ -541,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<
@@ -1274,6 +1283,7 @@ class AppStore {
 
   startPolling() {
     this.fetchState();
+    this.fetchFeatureFlags();
     this.fetchInterval = setInterval(() => this.fetchState(), 1000);
   }
 
@@ -1285,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");
@@ -1310,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;
       }
@@ -1670,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;
         }),
       ];
 
@@ -1877,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;
         }),
       ];
 
@@ -2408,10 +2449,15 @@ class AppStore {
               contentParts.push({ type: "text", text: textContent });
             }
 
-            return {
-              role: m.role,
-              content: contentParts,
-            };
+            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)
@@ -2429,10 +2475,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;
         }),
       ];
 
@@ -3281,6 +3332,60 @@ class AppStore {
     }
   }
 
+  async createInstanceLink(
+    prefillInstances: string[],
+    decodeInstances: string[],
+  ): Promise<void> {
+    const response = await fetch("/v1/instance-links", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({
+        prefill_instances: prefillInstances,
+        decode_instances: decodeInstances,
+      }),
+    });
+    if (!response.ok) {
+      throw new Error(
+        `Failed to create instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
+  async updateInstanceLink(
+    linkId: string,
+    prefillInstances: string[],
+    decodeInstances: string[],
+  ): Promise<void> {
+    const response = await fetch(
+      `/v1/instance-links/${encodeURIComponent(linkId)}`,
+      {
+        method: "PUT",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          prefill_instances: prefillInstances,
+          decode_instances: decodeInstances,
+        }),
+      },
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to update instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
+  async deleteInstanceLink(linkId: string): Promise<void> {
+    const response = await fetch(
+      `/v1/instance-links/${encodeURIComponent(linkId)}`,
+      { method: "DELETE" },
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to delete instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
   /**
    * Delete a downloaded model from a specific node
    */
@@ -3379,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;

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

@@ -3435,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"
@@ -4822,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"
@@ -4968,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

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

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

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

@@ -14,6 +14,7 @@
 
   let modelCapabilities = $state<Record<string, string[]>>({});
   let modelContextLengths = $state<Record<string, number>>({});
+  let modelReasoningDialects = $state<Record<string, string>>({});
 
   const runningModels = $derived.by(() => {
     const models: string[] = [];
@@ -88,10 +89,12 @@
   let codexModel = $state("");
   let codexMcpPath = $state("/Users/username");
   let openClawModel = $state("");
+  let piModel = $state("");
   $effect(() => {
     const def = modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id";
     codexModel = def;
     openClawModel = def;
+    piModel = def;
   });
 
   const claudeShellCommand = $derived(
@@ -130,6 +133,7 @@
     for (const modelId of runningModels) {
       const caps = modelCapabilities[modelId] || [];
       const ctxLen = modelContextLengths[modelId] || 0;
+      const dialect = modelReasoningDialects[modelId];
       const entry: Record<string, unknown> = { name: modelId };
       if (ctxLen > 0) {
         entry.limit = { context: ctxLen, output: Math.min(ctxLen, 16384) };
@@ -137,6 +141,27 @@
       if (caps.includes("vision")) {
         entry.modalities = { input: ["text", "image"], output: ["text"] };
       }
+      // Reasoning round-trip: opencode's `interleaved` field tells the
+      // openai-compatible adapter to send the assistant's prior
+      // reasoning_content back in subsequent turns. Emit it for dialects
+      // whose chat templates use prior reasoning:
+      //   - `tool_conditional` (DeepSeek V3.2 / V4): wrapper preserves all
+      //     reasoning when tools are present.
+      //   - `post_last_user` (Qwen3-Thinking, GLM 4.5+, MiniMax M2.x):
+      //     Jinja template reads reasoning_content for assistant turns since
+      //     the last user message — exactly the tool-chain window.
+      //   - `channel` (gpt-oss / Harmony): the model's Jinja template reads
+      //     `message.thinking` rather than `message.reasoning_content`, but
+      //     the server bridges `reasoning_content` → `thinking` before
+      //     rendering, so the round-trip works through the standard field.
+      // `suffix` (Kimi): reasoning lives in content; no separate field path.
+      if (
+        dialect === "tool_conditional" ||
+        dialect === "post_last_user" ||
+        dialect === "channel"
+      ) {
+        entry.interleaved = { field: "reasoning_content" };
+      }
       models[modelId] = entry;
     }
     if (Object.keys(models).length === 0) {
@@ -218,6 +243,55 @@
     ),
   );
 
+  const piModelsJson = $derived.by(() => {
+    const models: Record<string, unknown>[] = [];
+    for (const modelId of runningModels) {
+      const caps = modelCapabilities[modelId] || [];
+      const ctxLen = modelContextLengths[modelId] || 0;
+      const entry: Record<string, unknown> = { id: modelId };
+      if (caps.includes("vision")) {
+        entry.input = ["text", "image"];
+      }
+      // Mark thinking-capable models so pi surfaces its thinking-level selector
+      // for them. exo capability strings: "thinking" (model emits reasoning
+      // content) and "thinking_toggle" (user can turn it on/off).
+      if (caps.includes("thinking") || caps.includes("thinking_toggle")) {
+        entry.reasoning = true;
+      }
+      if (ctxLen > 0) {
+        entry.contextWindow = ctxLen;
+      }
+      models.push(entry);
+    }
+    if (models.length === 0) {
+      models.push({ id: "your-model-id" });
+    }
+    return JSON.stringify(
+      {
+        providers: {
+          exo: {
+            baseUrl: `${apiUrl}/v1`,
+            api: "openai-completions",
+            apiKey: "exo",
+            compat: {
+              supportsDeveloperRole: false,
+              // exo's OpenAI surface takes a boolean `enable_thinking` toggle,
+              // not graded effort levels, so disable pi's `reasoning_effort`
+              // parameter and use the matching top-level-boolean format.
+              supportsReasoningEffort: false,
+              thinkingFormat: "qwen",
+            },
+            models,
+          },
+        },
+      },
+      null,
+      2,
+    );
+  });
+
+  const piShellCommand = $derived(`pi --provider exo --model ${piModel}`);
+
   const ollamaCommand = $derived(
     `OLLAMA_HOST=${apiUrl}/ollama ollama run ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"}`,
   );
@@ -277,6 +351,7 @@
     "OpenCode",
     "Codex",
     "OpenClaw",
+    "Pi",
     "Open WebUI",
     "n8n",
     "Firefox",
@@ -298,16 +373,25 @@
     try {
       const resp = await fetch("/v1/models");
       const data = (await resp.json()) as {
-        data: { id: string; capabilities: string[]; context_length: number }[];
+        data: {
+          id: string;
+          capabilities: string[];
+          context_length: number;
+          reasoning_dialect?: string;
+        }[];
       };
       const caps: Record<string, string[]> = {};
       const ctxs: Record<string, number> = {};
+      const dialects: Record<string, string> = {};
       for (const model of data.data) {
         caps[model.id] = model.capabilities || [];
         if (model.context_length > 0) ctxs[model.id] = model.context_length;
+        if (model.reasoning_dialect)
+          dialects[model.id] = model.reasoning_dialect;
       }
       modelCapabilities = caps;
       modelContextLengths = ctxs;
+      modelReasoningDialects = dialects;
     } catch {
       /* ignore */
     }
@@ -515,6 +599,33 @@
           config={`openclaw doctor --fix${(modelCapabilities[openClawModel] || []).includes("vision") ? `\nopenclaw models set-image exo/${openClawModel}` : ""}\nopenclaw gateway &\nopenclaw dashboard`}
           language="bash"
         />
+      {:else if activeTab === "Pi"}
+        {#if runningModels.length > 1}
+          <div class="text-xs">
+            <span
+              class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+              >Model</span
+            >
+            <select bind:value={piModel} class={selectClass}>
+              {#each runningModels as model}
+                <option value={model}>{model.split("/").pop()}</option>
+              {/each}
+            </select>
+          </div>
+        {/if}
+        <IntegrationCard
+          title="Models Config"
+          subtitle="~/.pi/agent/models.json"
+          description="Register exo as a custom provider in pi. Create or edit this file, then run pi and pick an exo model via /model. Install pi with: npm install -g @mariozechner/pi-coding-agent"
+          config={piModelsJson}
+        />
+        <IntegrationCard
+          title="Shell Command"
+          subtitle="Run in terminal"
+          description="Launch pi directly with the exo provider and model selected."
+          config={piShellCommand}
+          language="bash"
+        />
       {:else if activeTab === "Open WebUI"}
         <IntegrationCard
           title="1. Start Open WebUI"

package-lock.json

@@ -0,0 +1,6 @@
+{
+  "name": "exo",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}

pyproject.toml

@@ -3,7 +3,7 @@ name = "exo"
 version = "0.3.70"
 description = "Exo"
 readme = "README.md"
-requires-python = ">=3.13"
+requires-python = "==3.13.*"
 dependencies = [
   "aiofiles>=24.1.0",
   "aiohttp>=3.12.14",
@@ -15,11 +15,11 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo-pyo3-bindings",                       # rust bindings
+  "exo-pyo3-bindings",                         # rust bindings
   "anyio==4.11.0",
-  "mlx==0.31.1; sys_platform == 'darwin'",
-  "mlx-lm",
-  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
+  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
@@ -28,8 +28,8 @@ dependencies = [
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
-  "mlx-vlm>=0.3.11",
-  "transformers>=5.0.0,<5.4.0",
+  "mlx-vlm>=0.3.11; sys_platform == 'darwin'",
+  "transformers>=5.6.2",
 ]
 
 [project.scripts]
@@ -40,6 +40,7 @@ exo = "exo.main:main"
 dev = [
   "basedpyright>=1.29.0",
   "pyinstaller>=6.17.0",
+  "playwright>=1.52.0",
   "pytest>=8.4.0",
   "pytest-asyncio>=1.0.0",
   "pytest-env",
@@ -49,15 +50,24 @@ dev = [
 [project.optional-dependencies]
 build = ["nanobind"]
 cpu = [
+  "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cpu==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda12 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda13 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
   "torch>=2.10.0; sys_platform == 'linux'",
 ]
 
@@ -66,16 +76,16 @@ cuda13 = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench"]
+members = ["rust/exo_pyo3_bindings", "bench", "tools"]
 
 [tool.uv.sources]
 exo-pyo3-bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
 torch = [
   { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
   { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'cpu' and extra != 'cuda12' and extra != 'cuda13'" },
+  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
 ]
 vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 
@@ -103,7 +113,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench"]
+include = ["src", "bench", "tools"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -137,6 +147,15 @@ reportMissingModuleSource = false
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
 
+[[tool.basedpyright.executionEnvironments]]
+root = "bench"
+# `.` keeps `from bench.lib.X import …` resolvable (pytest adds the project
+# root to sys.path; we want type-checking to agree with runtime).
+extraPaths = ["tools/src", "."]
+
+[[tool.basedpyright.executionEnvironments]]
+root = "tools/src"
+
 
 ###
 # uv configuration
@@ -148,7 +167,11 @@ required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
-constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+constraint-dependencies = ["transformers>=5.6.2"]
+override-dependencies = [
+  "mlx==0.31.1; sys_platform=='linux'",
+  "mlx; sys_platform=='darwin'",
+]
 
 [tool.uv.extra-build-dependencies]
 miniaudio = ["setuptools", "cffi", "pycparser"]
@@ -203,9 +226,12 @@ extend-exclude = [
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
 
 [tool.pytest.ini_options]
-pythonpath = "."
+pythonpath = ["."]
 asyncio_mode = "auto"
 markers = ["slow: marks tests as slow (deselected by default)"]
 env = ["EXO_TESTS=1"]
-addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
+# `tests/` requires an eco cluster (opt-in). `tmp/` holds throwaway scripts
+# that run a top-level `sys.exit(...)` at import time, which otherwise blows
+# up the default pytest collection.
+addopts = "-m 'not slow' --ignore=tests --ignore=tmp"
 filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

python/parts.nix

@@ -5,14 +5,13 @@ let
     workspaceRoot = ../.;
   };
 
-  mkPythonSet = { pkgs, lib, self' }:
+  mkPythonSet = { pkgs, lib, self', members }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
       inherit (pkgs.config) cudaSupport;
       inherit (pkgs) cudaPackages;
       cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
       libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
-      uv_extra = if cuda13Support then "cuda13" else if cudaSupport then "cuda12" else "cpu";
       python = pkgs.python313;
       cudaLibs = with cudaPackages; [
         cuda_cudart
@@ -51,7 +50,7 @@ let
           '';
         };
       };
-      buildSystemsOverlay = final: prev: { } //
+      buildSystemsOverlay = final: prev:
         lib.optionalAttrs isDarwin
           {
             mlx = prev.mlx.overrideAttrs (old:
@@ -81,7 +80,7 @@ let
                 nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
                 # TODO: non-sdk_26 support
                 buildInputs = (old.buildInputs or [ ])
-                  ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
+                ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
                 patches = [
                   (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
                     sdkVersion = pkgs.apple-sdk_26.version;
@@ -113,42 +112,42 @@ let
                 MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
               });
           } // lib.optionalAttrs isLinux {
-        mlx = prev.mlx.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
-          autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
-          postInstall = (old.postInstall or "") + ''
-            cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
-          '';
-        });
-      } // lib.optionalAttrs cudaSupport {
-        "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ cudaLibs ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        torch = prev.torch.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-      };
+          mlx = prev.mlx.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
+            postInstall = ''
+              cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
+            '';
+          });
+        } // lib.optionalAttrs cudaSupport {
+          "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          torch = prev.torch.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+        };
       pyprojectOverlay = workspace.mkPyprojectOverlay {
         sourcePreference = "wheel";
-        dependencies = { exo = [ uv_extra ]; exo-bench = [ ]; };
+        dependencies = members;
       };
       editableOverlay = workspace.mkEditablePyprojectOverlay {
         # Use environment variable pointing to editable root directory
@@ -165,8 +164,8 @@ let
           buildSystemsOverlay
         ]
       );
-
-      mkApp = cmd: name: members: pkgs.writeShellApplication {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: name: pkgs.writeShellApplication {
         inherit name;
         runtimeEnv = {
           EXO_DASHBOARD_DIR = self'.packages.dashboard;
@@ -174,17 +173,17 @@ let
         };
         runtimeInputs = [
           # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
-          ((pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; }))
+          (venv name)
         ]
         ++ lib.optionals isDarwin [ pkgs.macmon ];
         text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
       };
     in
     {
-      inherit pythonSet;
+      inherit venv;
       editablePythonSet = pythonSet.overrideScope editableOverlay;
-      mkPythonScript = members: name: path: mkApp ''python ${path} "$@"'' name members;
-      mkExo = name: members: mkApp ''exo "$@"'' name members;
+      mkPythonScript = path: mkApp ''python ${path} "$@"'';
+      mkExo = mkApp ''exo "$@"'';
     };
 in
 {
@@ -192,16 +191,21 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; }) pythonSet editablePythonSet mkPythonScript mkExo;
-
-      exoVenv = pythonSet.mkVirtualEnv "exo-env" { exo = lib.optionals isLinux [ "cpu" ]; };
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
 
       # Virtual environment with dev dependencies for testing
-      testVenv = pythonSet.mkVirtualEnv "exo-test-env" {
-        exo = [ "dev" ] ++ lib.optionals isLinux [ "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+      testVenv = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
       };
+      }).venv "exo-test";
 
-      mkBenchScript = mkPythonScript { exo-bench = [ ]; };
+      mkBenchScript = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "cpu" ];
+        exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
+      };
+      }).mkPythonScript;
 
       mkSimplePythonScript = name: path: pkgs.writeShellApplication {
         inherit name;
@@ -212,9 +216,7 @@ in
     in
     {
       packages = {
-        exo = mkExo "exo" { exo = lib.optionals isLinux [ "cpu" ]; };
-        # for devShell
-        exo-venv = exoVenv;
+        exo = mkExo "exo";
         editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
         # for running tests in ci
         exo-test-env = testVenv;
@@ -224,8 +226,8 @@ in
         # used by ./tests/run_exo_on.sh
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
       } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; }).mkExo "exo-cuda-12" { exo = [ "cuda12" ]; };
-        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; }).mkExo "exo-cuda-13" { exo = [ "cuda13" ]; };
+        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
+        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
       };
 
       checks = {
@@ -235,19 +237,11 @@ in
           touch $out
         '';
 
-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python --project ${inputs.self}/pyproject.toml
-            touch $out
-          '';
+        typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv ]; } ''
+          cd ${inputs.self}
+          basedpyright
+          touch $out
+        '';
       };
     };
 }

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,21 @@
+model_id = "mlx-community/DeepSeek-V4-Flash"
+n_layers = 43
+hidden_size = 4096
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "8bit"
+base_model = "DeepSeek V4 Flash"
+capabilities = ["text", "thinking", "thinking_toggle"]
+reasoning_dialect = "tool_conditional"
+
+context_length = 1048576
+
+[storage_size]
+in_bytes = 155095760030
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
+[sampling_defaults]
+temperature = 1.0
+top_p = 1.0

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

@@ -0,0 +1,21 @@
+model_id = "mlx-community/DeepSeek-V4-Pro"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "8bit"
+base_model = "DeepSeek V4 Pro"
+capabilities = ["text", "thinking", "thinking_toggle"]
+reasoning_dialect = "tool_conditional"
+
+context_length = 1048576
+
+[storage_size]
+in_bytes = 849681803879
+
+# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro
+[sampling_defaults]
+temperature = 1.0
+top_p = 1.0

resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM 4.5 Air"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 131072
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "4bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "5bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "6bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM 4.7 Flash"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-8bit.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "8bit"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "MXFP4-Q8"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5-bf16.toml

@@ -8,7 +8,7 @@ family = "glm"
 quantization = "bf16"
 base_model = "GLM-5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 202752
 
 [storage_size]

resources/inference_model_cards/mlx-community--GLM-5.1-DQ4plus-q8.toml

@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-DQ4plus-q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "8bit"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "post_last_user"
+context_length = 202752
+
+[storage_size]
+in_bytes = 465173655552
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-5.1-MXFP4-Q8.toml

@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-MXFP4-Q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "MXFP4-Q8"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "post_last_user"
+context_length = 202752
+
+[storage_size]
+in_bytes = 405480321024
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--GLM-5.1.toml

@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "bf16"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+reasoning_dialect = "post_last_user"
+context_length = 202752
+
+[storage_size]
+in_bytes = 1487822475264
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95

resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "suffix"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2.5.toml

@@ -8,7 +8,7 @@ family = "kimi"
 quantization = ""
 base_model = "Kimi K2.5"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "suffix"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Kimi-K2.6-mlx-DQ3_K_M-q8.toml

@@ -0,0 +1,33 @@
+model_id = "mlx-community/Kimi-K2.6-mlx-DQ3_K_M-q8"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = "3bit"
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+reasoning_dialect = "suffix"
+context_length = 262144
+
+[storage_size]
+in_bytes = 470628683776
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01

resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "3bit"
 base_model = "MiniMax M2.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.1"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "6bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.5"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-4bit-mxfp4.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit-mxfp4"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-4bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "4bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-5bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "5bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-6bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "6bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7-8bit.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "8bit"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--MiniMax-M2.7.toml

@@ -8,7 +8,7 @@ family = "minimax"
 quantization = "bf16"
 base_model = "MiniMax M2.7"
 capabilities = ["text", "thinking"]
-
+reasoning_dialect = "post_last_user"
 context_length = 196608
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3 Next 80B"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "8bit"
 base_model = "Qwen3 Next 80B"
 capabilities = ["text", "thinking", "thinking_toggle"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-4bit.toml

@@ -8,7 +8,7 @@ family = "qwen"
 quantization = "4bit"
 base_model = "Qwen3.5 122B A10B"
 capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-
+reasoning_dialect = "post_last_user"
 context_length = 262144
 
 [storage_size]