Compare commits

...
Author SHA1 Message Date
Ryuichi Leo Takashige 84db569167 Optimizations 6 2026-04-29 13:29:36 +01:00
Ryuichi Leo Takashige b83d5e6a6f Optimizations 5 2026-04-29 13:29:36 +01:00
Ryuichi Leo Takashige 9e67e89862 Optimizations 4 2026-04-29 13:29:36 +01:00
Evan 92ea4ed0a4 banner update 2026-04-29 13:29:28 +01:00
Evan 9f37340b89 cleanup 2026-04-29 08:01:53 +01:00
Evan 344381fd74 snailed it 2026-04-29 08:00:50 +01:00
Ryuichi Leo Takashige 701c9b1cf6 Optimizations 3 2026-04-29 00:53:48 +01:00
Ryuichi Leo Takashige dc709e933a Optimizations 2 2026-04-29 00:02:18 +01:00
Ryuichi Leo Takashige 0a736d7eaf Optimizations 2026-04-28 20:50:47 +01:00
Ryuichi Leo Takashige 94b1813f76 tmp 2 2026-04-28 20:40:56 +01:00
Ryuichi Leo Takashige 8774513367 tmp 2026-04-28 17:13:42 +01:00
Ryuichi Leo Takashige 35e3335d6d Select VLLM instances 2026-04-28 15:58:09 +01:00
Ryuichi Leo Takashige c2b35f4d9e Fix linux CI 2026-04-28 14:59:31 +01:00
Ryuichi Leo Takashige d96f8379ce Add Linux dashboard 2026-04-28 14:52:57 +01:00
Ryuichi Leo Takashige c1eca8d026 Fix pyproject for Macs 2026-04-28 14:36:21 +01:00
Evan dbc736c845 vllm support 2026-04-28 02:08:10 +01:00
667a3bb0e5 feat: keep-models option when uninstalling EXO (#1997)
## Summary

- Adds a **Keep downloaded models (~/.exo/models)** checkbox to the
macOS uninstall confirmation dialog (Settings → Advanced → Danger Zone).
The full `~/.exo` directory is now removed on uninstall by default; if
the checkbox is checked, `~/.exo/models` is preserved.
- The standalone `app/EXO/uninstall-exo.sh` gains a matching
`--keep-models` flag and the same `~/.exo` cleanup so GUI and CLI flows
stay in sync. Resolves the user home via `$SUDO_USER` since the script
runs under `sudo`.

Previously, "Uninstall EXO" only cleaned up system-level components
(LaunchDaemon, network location, logs, app bundle) and left the entire
`~/.exo` directory behind. Now uninstalling actually removes EXO's user
data, with a one-click opt-out for the (potentially many GB) of
downloaded models.

![Uninstall dialog with new
checkbox](https://raw.githubusercontent.com/exo-explore/exo/703b7fbbf13441217ad2903bb199f07e92af4490/uninstall-dialog.png)

> Note: the rendered icon in the screenshot above is the generic system
folder icon because it was captured from a small standalone Swift binary
(no app bundle / icon resource). When triggered from the actual EXO.app,
the EXO app icon is shown.

## Test plan

- [ ] Build EXO.app locally; open Settings → Advanced → Danger Zone →
Uninstall EXO; confirm the new "Keep downloaded models (~/.exo/models)"
checkbox is present and unchecked by default.
- [ ] Uninstall with the checkbox **checked** → `~/.exo/models/`
survives, all other entries under `~/.exo` are gone, system components
removed, app moved to Trash.
- [ ] Uninstall with the checkbox **unchecked** → `~/.exo` is fully
removed.
- [ ] `sudo app/EXO/uninstall-exo.sh --keep-models` → `~/.exo/models/`
is preserved, the rest of `~/.exo` is removed.
- [ ] `sudo app/EXO/uninstall-exo.sh` (no flag) → `~/.exo` is fully
removed.
- [ ] `app/EXO/uninstall-exo.sh --help` prints usage and exits 0;
unknown args exit 2 with a usage hint.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Evan <evanev7@gmail.com>
2026-04-28 01:06:02 +00:00
Evan Quiney c80b10c013 implement engine abstraction for mlx and mflux (#2000)
refactor for future versions.
2026-04-28 00:58:17 +00:00
Alex CheemaandClaude Opus 4.7 18ffe1df23 fix: uninstall-exo.sh removes both current and legacy bridge scripts (#1998)
## Summary

The standalone `app/EXO/uninstall-exo.sh` only knew about the legacy
filename `disable_bridge_enable_dhcp.sh`. On machines installed with
newer EXO versions, the current `/Library/Application
Support/EXO/disable_bridge.sh` was left behind, and the script then
reported `EXO support directory not empty, leaving in place`.

This PR makes the script try both filenames, removing whichever ones
exist. Tolerates **either**, **both**, or **neither** being present
without erroring.

The Swift `NetworkSetupHelper.makeUninstallScript()` already handles
both paths correctly, so the GUI uninstall flow is unaffected — this is
a script-only fix.

Caught while running an end-to-end uninstall on a real machine for
#1997.

## Test plan

Verified the new block in isolation against all four states:

- [x] both `disable_bridge.sh` and `disable_bridge_enable_dhcp.sh`
present → both removed
- [x] only `disable_bridge.sh` present → removed cleanly
- [x] only `disable_bridge_enable_dhcp.sh` present → removed cleanly
(legacy install)
- [x] neither present → prints the existing "already removed?" warning,
exits 0

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-28 00:28:12 +00:00
rltakashige f0d1371d89 MLX P/D (#1993)
## Motivation

MLX only prefill server for Apple Silicon
2026-04-28 00:12:42 +00:00
5d10188d3a fix: route by in-flight tasks only — completed tasks were skewing load balance (#1989)
The load balancer counted ALL tasks (Complete, Cancelled, TimedOut,
Failed) instead of only Pending/Running ones. With 138 accumulated tasks
and only 7 active, routing decisions were based on historical
distribution, causing one node to appear permanently 'busier' and
starving the other of work.

Co-authored-by: Adam Durham <adam@example.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-27 16:03:12 +00:00
ciaranbor f2a0db4e23 Extend bench/eval tooling (#1905)
## Motivation

Extend bench/eval tooling with robustness features, streaming support,
and align model configs with vllm eval for reproducible comparisons.

## Changes

- **exo_eval**: Checkpoint/resume (JSONL), instance health monitoring +
early abort, `top_k`/`min_p`/`enable_thinking` params, LCB
`--release-version`/`--offset`
- **exo_bench**: Streaming SSE (`--stream`), Kimi tokenizer fix for
transformers 5.x
- **Both tools**: Auto-detect running instances instead of requiring
`--skip-instance-setup`; `--fresh-instance` to override
- **harness**: SSE streaming client, `find_existing_instance()` shared
helper, removed download timeout, settle-timeout default 0→7200s
- **models.toml**: Added `enable_thinking`, aligned `max_tokens`/temps
with vllm, added new models
- **API**: Streaming SSE for `/bench/chat/completions`

## Why It Works

- Checkpoint/resume uses append-only JSONL + skip-on-load so interrupted
evals resume without re-running completed questions
- Health monitoring races an `asyncio.Event` against API calls for fast
abort when the instance dies
- Auto-detection queries `/state` for existing instances matching the
model ID before attempting placement
- Streaming reuses the existing `generate_chat_stream` infrastructure
from the regular chat endpoint
2026-04-27 16:53:43 +01:00
rltakashigeandEvan 37f6f4f6c2 Add DeepSeek V4 Flash/Pro (#1978)
Wait for upstream merge.

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-04-27 15:20:50 +01:00
Adam DurhamandAdam Durham 48a922fd5c fix: map presence_penalty and frequency_penalty from ChatCompletionRequest (#1991)
Upstream PR #1947 added `presence_penalty` and `frequency_penalty` to
`TextGenerationTaskParams` and the mlx-lm generator call sites, but
missed wiring them up in the API adapter so they were silently dropped
from incoming requests. This fixes the API mapping.

Co-authored-by: Adam Durham <adam@example.com>
2026-04-27 08:58:59 +00:00
184 changed files with 18399 additions and 1660 deletions

No files matched your search

-7
View File
@@ -1,8 +1 @@
use flake
# creates .venv if doesn't exist and loads its environment
export VIRTUAL_ENV=".venv"
if ! [ -d "./$VIRTUAL_ENV" ]; then
uv venv
fi
layout python
+5 -4
View File
@@ -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,
+3 -6
View File
@@ -191,13 +191,10 @@ class RotatingKVCache(_BaseCache):
def state(self, v): # -> None:
...
@property
def meta_state(self): # -> tuple[str, ...]:
...
def meta_state(self) -> tuple[str, ...]: ...
@meta_state.setter
def meta_state(self, v): # -> None:
...
def is_trimmable(self): # -> bool:
...
def meta_state(self, v: tuple[str, ...]) -> None: ...
def is_trimmable(self) -> bool: ...
def trim(self, n: int) -> int: ...
def to_quantized(
self, group_size: int = ..., bits: int = ...
+280
View File
@@ -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]: ...
File diff suppressed because it is too large. Load diff
+30 -3
View File
@@ -552,15 +552,24 @@ struct SettingsView: View {
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 +579,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 +593,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 +617,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 {
+63 -7
View File
@@ -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)."
+77 -22
View File
@@ -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
+244 -101
View File
@@ -35,6 +35,7 @@ from harness import (
ExoHttpError,
add_common_instance_args,
capture_cluster_snapshot,
find_existing_instance,
instance_id_from_instance,
node_ids_from_instance,
nodes_used_in_instance,
@@ -79,7 +80,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
model_path = Path(
snapshot_download(
model_id,
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
)
)
@@ -122,8 +123,48 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
return hf_tokenizer
# Default: use AutoTokenizer
return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
# TODO: Change back to using only transformers
try:
return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
except (AttributeError, ValueError):
from huggingface_hub import snapshot_download
from transformers import PretrainedConfig
model_path = Path(
snapshot_download(
model_id,
allow_patterns=[
"*.json",
"*.py",
"tokenizer.model",
"*.tiktoken",
"tiktoken.model",
"*.txt",
"*.jsonl",
"*.jinja",
],
)
)
stub_kwargs: dict[str, Any] = {}
config_file = model_path / "config.json"
if config_file.exists():
with open(config_file) as f:
raw = json.load(f)
for key in (
"model_type",
"max_position_embeddings",
"vocab_size",
"bos_token_id",
"eos_token_id",
"pad_token_id",
):
if key in raw:
stub_kwargs[key] = raw[key]
return AutoTokenizer.from_pretrained(
str(model_path),
config=PretrainedConfig(**stub_kwargs),
trust_remote_code=True,
)
def format_peak_memory(b: float) -> str:
@@ -237,28 +278,72 @@ def run_one_completion(
prompt_sizer: PromptSizer,
*,
use_prefix_cache: bool = False,
stream: bool = False,
) -> tuple[dict[str, Any], int]:
content, pp_tokens = prompt_sizer.build(pp_hint)
payload: dict[str, Any] = {
"model": model_id,
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
"logprobs": False,
"use_prefix_cache": use_prefix_cache,
}
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
if not stream:
payload["stream"] = False
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
stats = out.get("generation_stats")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
preview = content[:200] if content else ""
else:
tokens = 0
first_token_time = None
t0 = time.perf_counter()
text_parts: list[str] = []
stats = None
# 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 ""
for raw_line in client.stream_bench_chat_completions(payload):
line = raw_line.strip()
if line.startswith(": generation_stats "):
with contextlib.suppress(json.JSONDecodeError):
stats = json.loads(line[len(": generation_stats ") :])
continue
if not line.startswith("data: "):
continue
data = line[6:]
if data == "[DONE]":
break
try:
chunk = json.loads(data)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if delta.get("content"):
if first_token_time is None:
first_token_time = time.perf_counter()
tokens += 1
text_parts.append(delta["content"])
except json.JSONDecodeError:
pass
elapsed = time.perf_counter() - t0
preview = "".join(text_parts)[:200]
if not stats:
ttft = (first_token_time - t0) if first_token_time else elapsed
gen_time = elapsed - ttft if tokens > 1 else elapsed
gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
stats = {
"prompt_tokens": pp_tokens,
"generation_tokens": tokens,
"prompt_tps": round(prompt_tps, 2),
"generation_tps": round(gen_tps, 2),
"peak_memory_usage": {"inBytes": 0},
}
return {
"elapsed_s": elapsed,
@@ -278,9 +363,19 @@ class PromptSizer:
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
)
try:
ids = tokenizer.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.deepseek_v4_encoding import (
encode_messages as encode_v4,
)
prompt = encode_v4(messages, thinking_mode="thinking")
ids = tokenizer.encode(prompt, add_special_tokens=False)
# Fix for transformers 5.x
if hasattr(ids, "input_ids"):
ids = ids.input_ids
@@ -375,6 +470,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 +540,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 +666,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 +701,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 +850,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 +880,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:
+427 -56
View File
@@ -47,6 +47,7 @@ from harness import (
ExoHttpError,
add_common_instance_args,
capture_cluster_snapshot,
find_existing_instance,
instance_id_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
@@ -62,6 +63,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 +281,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 +291,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 +531,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 +548,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 +567,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 +581,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 +626,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 +646,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 +694,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 +734,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 +756,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 +770,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 +848,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 +906,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 +930,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 +944,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 +959,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 +974,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 +985,7 @@ async def evaluate_benchmark(
)
else:
result = QuestionResult(
question_id=idx,
question_id=question_id,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -815,7 +996,7 @@ async def evaluate_benchmark(
)
else:
result = QuestionResult(
question_id=idx,
question_id=question_id,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -827,24 +1008,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 +1106,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 +1119,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 +1141,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 +1300,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 +1320,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 +1356,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 +1381,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 +1416,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 +1461,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 +1495,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 +1522,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 +1543,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 +1599,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 +1658,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 +1673,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 +1689,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 +1711,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 +1726,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
+65 -12
View File
@@ -6,6 +6,7 @@ import http.client
import json
import os
import time
from collections.abc import Iterator
from typing import Any
from urllib.parse import urlencode
@@ -69,6 +70,30 @@ class ExoClient:
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
return self.request_json("POST", "/bench/chat/completions", body=payload)
def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
"""POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
payload = {**payload, "stream": True}
data = json.dumps(payload).encode("utf-8")
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
try:
conn.request(
"POST",
"/bench/chat/completions",
body=data,
headers={
"Content-Type": "application/json",
"Accept": "text/event-stream",
},
)
resp = conn.getresponse()
if resp.status >= 400:
raw = resp.read().decode("utf-8", errors="replace")
raise ExoHttpError(resp.status, resp.reason, raw[:300])
for line in resp:
yield line.decode("utf-8", errors="replace")
finally:
conn.close()
def get_state_path(self, path: str) -> Any:
try:
return self.request_json("GET", f"/state/{path}")
@@ -268,11 +293,15 @@ def sharding_filter(sharding: str, wanted: str) -> bool:
def fetch_and_filter_placements(
client: ExoClient, full_model_id: str, args: argparse.Namespace
client: ExoClient,
full_model_id: str,
args: argparse.Namespace,
node_id: str | None = None,
) -> list[dict[str, Any]]:
previews_resp = client.request_json(
"GET", "/instance/previews", params={"model_id": full_model_id}
)
params: dict[str, str] = {"model_id": full_model_id}
if node_id is not None:
params["node_ids"] = node_id
previews_resp = client.request_json("GET", "/instance/previews", params=params)
previews = previews_resp.get("previews") or []
selected: list[dict[str, Any]] = []
@@ -332,8 +361,9 @@ def settle_and_fetch_placements(
full_model_id: str,
args: argparse.Namespace,
settle_timeout: float = 0,
node_id: str | None = None,
) -> list[dict[str, Any]]:
selected = fetch_and_filter_placements(client, full_model_id, args)
selected = fetch_and_filter_placements(client, full_model_id, args, node_id=node_id)
if not selected and settle_timeout > 0:
backoff = _SETTLE_INITIAL_BACKOFF_S
@@ -346,7 +376,9 @@ def settle_and_fetch_placements(
)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
selected = fetch_and_filter_placements(client, full_model_id, args)
selected = fetch_and_filter_placements(
client, full_model_id, args, node_id=node_id
)
return selected
@@ -462,9 +494,8 @@ def run_planning_phase(
)
logger.info(f"Started download on {node_id}")
# Wait for downloads
start = time.time()
while time.time() - start < timeout:
# Wait for downloads (no timeout — poll until complete or failed)
while True:
all_done = True
for node_id in node_ids:
node_downloads = client.get_node_downloads(node_id) or []
@@ -514,9 +545,24 @@ def run_planning_phase(
if download_t0 is not None:
return time.perf_counter() - download_t0
return None
time.sleep(1)
time.sleep(10)
raise TimeoutError("Downloads did not complete in time")
def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
"""Find an existing running instance for the given model."""
try:
state = client.request_json("GET", "/state")
except Exception:
return None
for inst_id, inst in state.get("instances", {}).items():
# Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
for _inst_type, inner in inst.items():
if not isinstance(inner, dict):
continue
sa = inner.get("shardAssignments", {})
if sa.get("modelId") == model_id:
return inst_id
return None
def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
@@ -543,7 +589,9 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
help="Only consider placements using >= this many nodes.",
)
ap.add_argument(
"--instance-meta", choices=["ring", "jaccl", "both"], default="both"
"--instance-meta",
choices=["ring", "jaccl", "vllm", "both"],
default="both",
)
ap.add_argument(
"--sharding", choices=["pipeline", "tensor", "both"], default="both"
@@ -572,3 +620,8 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
action="store_true",
help="Delete existing models from smallest to largest to make room for benchmark model.",
)
ap.add_argument(
"--reuse-instance",
action="store_true",
help="Reuse an existing running instance for this model instead of creating a new one.",
)
+36
View File
@@ -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, 8192]
tg = [128]
repeat = 1
warmup = 0
json_out = "bench/prefill_decode_results.json"
[prefill]
model = "sakamakismile/Qwen3.6-27B-NVFP4"
node = "gx10-de89"
instance_meta = "vllm"
sharding = "pipeline"
min_nodes = 1
max_nodes = 1
[decode]
model = "mlx-community/Qwen3.6-27B-4bit"
node = "Ryuichis MacBook Pro"
instance_meta = "ring"
sharding = "pipeline"
min_nodes = 1
max_nodes = 1
+869
View File
@@ -0,0 +1,869 @@
# 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,
SystemMetricsSampler,
format_peak_memory,
load_tokenizer_for_bench,
parse_int_list,
)
from harness import (
ExoClient,
ExoHttpError,
add_common_instance_args,
instance_id_from_instance,
node_ids_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
run_planning_phase,
settle_and_fetch_placements,
unwrap_instance,
wait_for_instance_gone,
wait_for_instance_ready,
)
from loguru import logger
def _node_id_to_friendly(client: ExoClient) -> dict[str, str]:
identities = client.get_node_identities() or {}
out: dict[str, str] = {}
for node_id, identity in identities.items():
if isinstance(identity, dict):
name = identity.get("friendlyName") or identity.get("friendly_name")
if isinstance(name, str):
out[str(node_id)] = name
return out
def _placement_node_friendly_names(
placement: dict[str, Any], id_to_friendly: dict[str, str]
) -> list[str]:
instance = placement["instance"]
return [id_to_friendly.get(nid, nid) for nid in node_ids_from_instance(instance)]
def _filter_by_node(
placements: list[dict[str, Any]],
friendly_name: str,
id_to_friendly: dict[str, str],
) -> list[dict[str, Any]]:
target = friendly_name.lower()
matched: list[dict[str, Any]] = []
for p in placements:
names = [n.lower() for n in _placement_node_friendly_names(p, id_to_friendly)]
if any(target == n or target in n for n in names):
matched.append(p)
return matched
def _node_id_by_friendly(id_to_friendly: dict[str, str], target: str) -> str | None:
target_lc = target.lower()
for nid, name in id_to_friendly.items():
if target_lc == name.lower() or target_lc in name.lower():
return nid
return None
def _load_toml(path: str) -> dict[str, Any]:
with Path(path).open("rb") as f:
return tomllib.load(f)
_TOP_LEVEL_TOML_KEYS = {
"host",
"port",
"timeout",
"settle_timeout",
"model",
"pp",
"tg",
"repeat",
"warmup",
"json_out",
"instance_meta",
"sharding",
"min_nodes",
"max_nodes",
"force_download",
"danger_delete_downloads",
"all_combinations",
}
def _inject_toml_into_argv() -> None:
"""If --config X is in sys.argv, pre-load it and inject required CLI args
(--model, --pp, --tg) so argparse's required=True checks pass."""
argv = sys.argv
if "--config" not in argv:
return
idx = argv.index("--config")
if idx + 1 >= len(argv):
return
cfg_path = argv[idx + 1]
cfg = _load_toml(cfg_path)
decode = cfg.get("decode", {})
def _has(flag: str) -> bool:
return any(a == flag or a.startswith(flag + "=") for a in argv)
# --model: prefer top-level, then [decode].model
if not _has("--model"):
model = cfg.get("model") or decode.get("model")
if model:
argv += ["--model", str(model)]
if not _has("--pp"):
pp = cfg.get("pp")
if pp:
argv += (
["--pp", *(str(x) for x in pp)]
if isinstance(pp, list)
else [
"--pp",
str(pp),
]
)
if not _has("--tg"):
tg = cfg.get("tg")
if tg:
argv += (
["--tg", *(str(x) for x in tg)]
if isinstance(tg, list)
else [
"--tg",
str(tg),
]
)
def _merge_toml_into_args(args: argparse.Namespace, cfg: dict[str, Any]) -> None:
"""Apply top-level toml keys onto args namespace where args has a default."""
for key, value in cfg.items():
if key in {"prefill", "decode"}:
continue
if key not in _TOP_LEVEL_TOML_KEYS:
continue
attr = key
current = getattr(args, attr, None)
if current in (None, [], False):
setattr(args, attr, value)
def _side_args(
base: argparse.Namespace, overrides: dict[str, Any]
) -> argparse.Namespace:
out = copy.copy(base)
for k in (
"instance_meta",
"sharding",
"min_nodes",
"max_nodes",
"skip_pipeline_jaccl",
"skip_tensor_ring",
):
if k in overrides:
setattr(out, k, overrides[k])
return out
def _pick_two_distinct_placements(
placements: list[dict[str, Any]],
) -> tuple[dict[str, Any], dict[str, Any]] | None:
if len(placements) < 2:
return None
seen_nodes: set[tuple[str, ...]] = set()
chosen: list[dict[str, Any]] = []
for p in placements:
nodes = tuple(sorted(str(n) for n in p.get("nodes", [])))
if nodes in seen_nodes:
continue
seen_nodes.add(nodes)
chosen.append(p)
if len(chosen) == 2:
return chosen[0], chosen[1]
return None
def _create_instance_link(
client: ExoClient,
prefill_instance_id: str,
decode_instance_id: str,
) -> str:
out = client.request_json(
"POST",
"/v1/instance-links",
body={
"prefill_instances": [prefill_instance_id],
"decode_instances": [decode_instance_id],
},
)
return str(out.get("commandId", ""))
def _list_instance_links(client: ExoClient) -> list[dict[str, Any]]:
out = client.request_json("GET", "/v1/instance-links")
return out if isinstance(out, list) else []
def _delete_instance_link(client: ExoClient, link_id: str) -> None:
client.request_json("DELETE", f"/v1/instance-links/{link_id}")
def run_one(
client: ExoClient,
model_id: str,
pp_hint: int,
tg: int,
prompt_sizer: PromptSizer,
) -> tuple[dict[str, Any], int]:
content, pp_tokens = prompt_sizer.build(pp_hint)
payload: dict[str, Any] = {
"model": model_id,
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
}
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
text = message.get("content") or ""
preview = text[:200] if text else ""
return {
"elapsed_s": elapsed,
"output_text_preview": preview,
"stats": stats,
}, pp_tokens
def _run_phase(
*,
client: ExoClient,
label: str,
pp_tg_pairs: list[tuple[int, int]],
model_id: str,
prompt_sizer: PromptSizer,
warmup: int,
repeat: int,
common_meta: dict[str, Any],
sampler: SystemMetricsSampler | None = None,
) -> 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]] = []
inference_windows: list[tuple[float, float]] = []
for r in range(repeat):
time.sleep(2)
try:
inf_t0 = time.monotonic()
row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
inference_windows.append((inf_t0, time.monotonic()))
except Exception as e:
logger.error(e)
continue
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)
energy_str = ""
if sampler is not None and inference_windows:
joules = sum(
sampler.energy_between(t0, t1) for t0, t1 in inference_windows
)
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0.0
energy_per_run = joules / len(runs) if runs else 0.0
energy_str = (
f" energy={joules:.1f}J ({avg_watts:.1f}W avg over "
f"{inf_seconds:.1f}s inference, {energy_per_run:.1f}J/run)"
)
for run_row, (t0, t1) in zip(runs, inference_windows, strict=False):
run_row["energy_joules"] = sampler.energy_between(t0, t1)
run_row["inference_window_s"] = t1 - t0
logger.info(
f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
f"prompt_tokens={ptok} gen_tokens={gtok} "
f"peak_memory={format_peak_memory(peak)} "
f"avg_elapsed={avg_elapsed:.2f}s{energy_str}"
)
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():
energy_runs = [x.get("energy_joules") for x in runs if "energy_joules" in x]
window_runs = [
x.get("inference_window_s") for x in runs if "inference_window_s" in x
]
out[key] = {
"prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
"gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
"elapsed_s": mean(x["elapsed_s"] for x in runs),
"prompt_tokens": mean(x["stats"]["prompt_tokens"] for x in runs),
"gen_tokens": mean(x["stats"]["generation_tokens"] for x in runs),
"energy_j": mean(energy_runs) if energy_runs else 0.0,
"inference_window_s": mean(window_runs) if window_runs else 0.0,
}
return out
def _normalised_seconds(summary: dict[str, float], pp: int, tg: int) -> float | None:
"""Wall-clock time implied by reported tps for the *configured* pp/tg.
elapsed_s is not comparable across phases when models EOS at different
lengths. This formula reconstructs "what would this phase take to do
pp prompt tokens + tg generation tokens" using its own reported rates.
"""
p_tps = summary.get("prompt_tps", 0.0)
g_tps = summary.get("gen_tps", 0.0)
if p_tps <= 0 or g_tps <= 0:
return None
return pp / p_tps + tg / g_tps
def _print_diff(
disagg_rows: list[dict[str, Any]],
decode_alone_rows: list[dict[str, Any]],
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 = 110
for key in sorted(keys):
pp, tg = key
logger.info("" * width)
logger.info(f" pp={pp} tg={tg}")
logger.info("" * width)
logger.info(
f" {'phase':<16} {'elapsed':>9} {'norm':>9} "
f"{'prompt_tps':>11} {'gen_tps':>8} "
f"{'p_tok':>6} {'g_tok':>6} "
f"{'energy':>9} {'avg_W':>7}"
)
for label, summary in (
("disaggregated", disagg.get(key)),
("decode_alone", decode_alone.get(key)),
("prefill_alone", prefill_alone.get(key)),
):
if summary is None:
logger.info(
f" {label:<16} {'':>9} {'':>9} "
f"{'':>11} {'':>8} {'':>6} {'':>6} "
f"{'':>9} {'':>7}"
)
continue
norm = _normalised_seconds(summary, pp, tg)
norm_str = f"{norm:>8.2f}s" if norm is not None else f"{'':>9}"
energy = summary.get("energy_j", 0.0)
window = summary.get("inference_window_s", 0.0)
energy_str = f"{energy:>8.1f}J" if energy > 0 else f"{'':>9}"
avg_w = energy / window if window > 0 else 0.0
avg_w_str = f"{avg_w:>6.1f}W" if avg_w > 0 else f"{'':>7}"
logger.info(
f" {label:<16} "
f"{summary['elapsed_s']:>8.2f}s "
f"{norm_str} "
f"{summary['prompt_tps']:>11.1f} "
f"{summary['gen_tps']:>8.2f} "
f"{summary['prompt_tokens']:>6.0f} "
f"{summary['gen_tokens']:>6.0f} "
f"{energy_str} "
f"{avg_w_str}"
)
d = disagg.get(key)
da = decode_alone.get(key)
pa = prefill_alone.get(key)
d_norm = _normalised_seconds(d, pp, tg) if d else None
if d_norm and da:
da_norm = _normalised_seconds(da, pp, tg)
if da_norm:
logger.info(
f" norm speedup vs decode_alone: {da_norm / d_norm:.2f}x "
f"(prefill {d['prompt_tps'] / da['prompt_tps']:.2f}x, "
f"decode {d['gen_tps'] / da['gen_tps']:.2f}x)"
)
if d_norm and pa:
pa_norm = _normalised_seconds(pa, pp, tg)
if pa_norm:
logger.info(
f" norm speedup vs prefill_alone: {pa_norm / d_norm:.2f}x "
f"(prefill {d['prompt_tps'] / pa['prompt_tps']:.2f}x, "
f"decode {d['gen_tps'] / pa['gen_tps']:.2f}x)"
)
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
sampler_nodes = sorted(
{
*node_ids_from_instance(prefill_instance),
*node_ids_from_instance(decode_instance),
}
)
sampler = SystemMetricsSampler(
ExoClient(args.host, args.port, timeout_s=30), sampler_nodes
)
sampler.start()
try:
logger.info("Creating prefill instance...")
client.request_json("POST", "/instance", body={"instance": prefill_instance})
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,
sampler=sampler,
)
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,
sampler=sampler,
)
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,
sampler=sampler,
)
all_rows.extend(decode_alone_rows)
_print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
finally:
sampler.stop()
with contextlib.suppress(ExoHttpError):
if link_id:
_delete_instance_link(client, link_id)
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())
@@ -202,6 +202,7 @@
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
else if (instanceTag === "VllmInstance") instanceType = "vLLM";
let sharding: string | null = null;
const inst = instance as {
+292 -2
View File
@@ -9,7 +9,7 @@
*/
interface Props {
/** "macbook pro" | "mac studio" | "mac mini" etc. */
/** "macbook pro" | "mac studio" | "mac mini" | "dgx spark" | "linux" etc. */
deviceType: string;
/** Center X coordinate in SVG space */
cx: number;
@@ -38,10 +38,43 @@
const LOGO_NATIVE_WIDTH = 814;
const LOGO_NATIVE_HEIGHT = 1000;
// NVIDIA logo SVG path
const NVIDIA_LOGO_PATH =
"M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
const wireColor = "rgba(179,179,179,0.8)";
const strokeWidth = 1.5;
const modelLower = $derived(deviceType.toLowerCase());
const isSpark = $derived(
modelLower.includes("dgx") || modelLower.includes("gx10"),
);
const isLinux = $derived(!isSpark && modelLower.startsWith("linux"));
const isLinuxLaptop = $derived(isLinux && modelLower.includes("laptop"));
// ── DGX Spark dimensions ──
const dgxW = $derived(size * 1.55);
const dgxH = $derived(size * 0.58);
const dgxX = $derived(cx - dgxW / 2);
const dgxY = $derived(cy - dgxH / 2);
const dgxChassisX = $derived(dgxX - dgxW * 0.03);
const dgxChassisW = $derived(dgxW * 1.05);
const dgxHandleW = $derived(dgxW * 0.27);
const dgxHandleGap = $derived(dgxH * 0.05);
const dgxHandleH = $derived(dgxH - dgxHandleGap * 2);
const dgxHandleY = $derived(dgxY + dgxHandleGap);
const dgxInnerHandleW = $derived(dgxW * 0.12);
const dgxInnerHandleH = $derived(dgxHandleH - dgxH * 0.06);
const dgxLeftHandleX = $derived(dgxX + 4);
const dgxRightHandleX = $derived(dgxX + dgxW - dgxHandleW - 4);
const dgxClipId = $derived(`di-dgx-${uid}`);
const dgxTextureId = $derived(`di-dgx-tex-${uid}`);
// ── Linux Desktop dimensions (reuses Mac Studio proportions) ──
const linuxDesktopClipId = $derived(`di-linux-desktop-${uid}`);
// ── Linux Laptop dimensions (reuses MacBook proportions) ──
const linuxScreenClipId = $derived(`di-linux-screen-${uid}`);
// ── Mac Studio dimensions (same ratios as TopologyGraph) ──
const studioW = $derived(size * 1.25);
@@ -114,7 +147,264 @@
const studioClipId = $derived(`di-studio-${uid}`);
</script>
{#if modelLower === "mac studio" || modelLower === "mac mini"}
{#if isSpark}
<!-- DGX Spark -->
<defs>
<clipPath id={dgxClipId}>
<rect x={dgxX} y={dgxY} width={dgxW} height={dgxH} rx="3" />
</clipPath>
<pattern
id={dgxTextureId}
patternUnits="userSpaceOnUse"
width="8"
height="8"
>
<rect width="8" height="8" fill="#6f6248" />
<circle cx="2" cy="2" r="1" fill="#5a4f3b" opacity="0.5" />
<circle cx="6" cy="6" r="1" fill="#4a4232" opacity="0.45" />
</pattern>
</defs>
<!-- Main body -->
<rect
x={dgxChassisX}
y={dgxY}
width={dgxChassisW}
height={dgxH}
rx="3"
fill="url(#{dgxTextureId})"
stroke={wireColor}
stroke-width={strokeWidth}
/>
<!-- Side border accents -->
<rect
x={dgxChassisX}
y={dgxY}
width={dgxW * 0.02}
height={dgxH}
fill="#8a7a56"
/>
<rect
x={dgxChassisX + dgxChassisW - dgxW * 0.02}
y={dgxY}
width={dgxW * 0.02}
height={dgxH}
fill="#8a7a56"
/>
<!-- Memory fill -->
{#if ramPercent > 0}
<rect
x={dgxX}
y={dgxY + dgxH - (ramPercent / 100) * dgxH}
width={dgxW}
height={(ramPercent / 100) * dgxH}
fill="rgba(255,215,0,0.45)"
clip-path="url(#{dgxClipId})"
/>
{/if}
<!-- Left handle -->
<rect
x={dgxLeftHandleX}
y={dgxHandleY}
width={dgxHandleW}
height={dgxHandleH}
rx="2.4"
fill="#b3a170"
stroke="#403723"
stroke-width="0.7"
/>
<rect
x={dgxLeftHandleX + dgxHandleW * 0.06}
y={dgxHandleY + dgxH * 0.03}
width={dgxInnerHandleW}
height={dgxInnerHandleH}
rx="1.6"
fill="#8a7a56"
/>
<!-- Right handle -->
<rect
x={dgxRightHandleX}
y={dgxHandleY}
width={dgxHandleW}
height={dgxHandleH}
rx="2.4"
fill="#b3a170"
stroke="#403723"
stroke-width="0.7"
/>
<rect
x={dgxRightHandleX + dgxHandleW - dgxInnerHandleW - dgxHandleW * 0.08}
y={dgxHandleY + dgxH * 0.03}
width={dgxInnerHandleW}
height={dgxInnerHandleH}
rx="1.6"
fill="#8a7a56"
/>
<!-- NVIDIA logo (rotated 90deg on left handle) -->
{@const badgeW = dgxW * 0.09}
{@const badgeH = dgxHandleH * 0.5}
{@const badgeX = dgxLeftHandleX + dgxHandleW - badgeW - dgxHandleW * 0.06}
{@const badgeYPos = dgxHandleY + (dgxHandleH - badgeH) / 2}
{@const textSz = badgeW * 0.58}
{@const logoW = textSz * 1.2}
{@const logoH = logoW * (1.438 / 2.174)}
{@const ctrX = badgeX + badgeW / 2 - badgeW * 0.03}
{@const ctrY = badgeYPos + badgeH / 2}
{@const labelGap = badgeW * 0.15}
{@const totalW = logoW + labelGap + textSz * 3.6}
<g transform="rotate(90 {ctrX} {ctrY})">
<svg
x={ctrX - totalW / 2}
y={ctrY - logoH / 2}
width={logoW}
height={logoH}
viewBox="0 0 2.174 1.438"
>
<path d={NVIDIA_LOGO_PATH} fill="#76b900" />
</svg>
<text
x={ctrX - totalW / 2 + logoW + labelGap}
y={ctrY}
text-anchor="start"
dominant-baseline="middle"
fill="#8a7a56"
font-size={textSz}
font-family="monospace"
font-weight="700">NVIDIA</text
>
</g>
{:else if isLinuxLaptop}
<!-- Linux Laptop — MacBook shape with Tux logo -->
<defs>
<clipPath id={linuxScreenClipId}>
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel}
width={mbScreenW - mbBezel * 2}
height={mbScreenH - mbBezel * 2}
rx="2"
/>
</clipPath>
</defs>
<rect
x={mbScreenX}
y={mbY}
width={mbScreenW}
height={mbScreenH}
rx="3"
fill="#1a1a1a"
stroke={wireColor}
stroke-width={strokeWidth}
/>
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel}
width={mbScreenW - mbBezel * 2}
height={mbScreenH - mbBezel * 2}
rx="2"
fill="#0a0a12"
/>
{#if ramPercent > 0}
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
width={mbScreenW - mbBezel * 2}
height={mbMemH}
fill="rgba(255,215,0,0.85)"
clip-path="url(#{linuxScreenClipId})"
/>
{/if}
<!-- Terminal prompt on screen -->
<text
x={cx}
y={mbY + mbScreenH / 2}
text-anchor="middle"
dominant-baseline="middle"
fill="#FFFFFF"
opacity="0.9"
font-size={mbScreenH * 0.25}
font-family="SF Mono, Monaco, monospace"
font-weight="700">{">_"}</text
>
<path
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
fill="#2c2c2c"
stroke={wireColor}
stroke-width="1"
/>
<rect
x={mbKbX}
y={mbKbY}
width={mbKbW}
height={mbKbH}
fill="rgba(0,0,0,0.2)"
rx="2"
/>
<rect
x={mbTpX}
y={mbTpY}
width={mbTpW}
height={mbTpH}
fill="rgba(255,255,255,0.08)"
rx="2"
/>
{:else if isLinux}
<!-- Linux Desktop — Mac Studio shape with Tux logo -->
<defs>
<clipPath id={linuxDesktopClipId}>
<rect
x={studioX}
y={studioY + studioTopH}
width={studioW}
height={studioH - studioTopH}
rx={studioCorner - 1}
/>
</clipPath>
</defs>
<rect
x={studioX}
y={studioY}
width={studioW}
height={studioH}
rx={studioCorner}
fill="#1a1a1a"
stroke={wireColor}
stroke-width={strokeWidth}
/>
{#if ramPercent > 0}
<rect
x={studioX}
y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
width={studioW}
height={studioMemH}
fill="rgba(255,215,0,0.75)"
clip-path="url(#{linuxDesktopClipId})"
/>
{/if}
<!-- Terminal prompt on front face -->
<text
x={cx}
y={studioY + studioTopH + (studioH - studioTopH) / 2}
text-anchor="middle"
dominant-baseline="middle"
fill="rgba(255,255,255,0.5)"
font-size={(studioH - studioTopH) * 0.4}
font-family="SF Mono, Monaco, monospace"
font-weight="700">{">_"}</text
>
{:else if modelLower === "mac studio" || modelLower === "mac mini"}
<!-- Mac Studio / Mac Mini -->
<defs>
<clipPath id={studioClipId}>
@@ -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>
+85 -4
View File
@@ -23,7 +23,7 @@
} | null;
nodes?: Record<string, NodeInfo>;
sharding?: "Pipeline" | "Tensor";
runtime?: "MlxRing" | "MlxJaccl";
runtime?: "MlxRing" | "MlxJaccl" | "Vllm";
onLaunch?: () => void;
tags?: string[];
apiPreview?: PlacementPreview | null;
@@ -168,8 +168,10 @@
function getDeviceType(
name: string,
): "macbook" | "studio" | "mini" | "unknown" {
): "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown" {
const lower = name.toLowerCase();
if (lower.includes("dgx") || lower.includes("gx10")) return "dgx";
if (lower.includes("linux")) return "linux";
if (lower.includes("macbook")) return "macbook";
if (lower.includes("studio")) return "studio";
if (lower.includes("mini")) return "mini";
@@ -576,13 +578,17 @@
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
title={runtime === "MlxRing"
? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
: runtime === "MlxJaccl"
? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
: "vLLM: NVIDIA CUDA inference engine."}
>
{runtime === "MlxRing"
? "MLX Ring"
: runtime === "MlxJaccl"
? "MLX RDMA"
: runtime}
: runtime === "Vllm"
? "vLLM"
: runtime}
</span>
</div>
@@ -990,6 +996,81 @@
/>
{/if}
</g>
{:else if node.deviceType === "dgx"}
<!-- DGX Spark icon -->
{@const s = node.iconSize}
{@const dgxW = s * 1.4}
{@const dgxH = s * 0.52}
<g transform="translate({-dgxW / 2}, {-dgxH / 2})">
<!-- Chassis -->
<rect
x="0"
y="0"
width={dgxW}
height={dgxH}
rx="2"
fill="#6f6248"
stroke={node.isUsed ? "#FFD700" : "#4B5563"}
stroke-width="1.5"
/>
<!-- Side accents -->
<rect
x="0"
y="0"
width={dgxW * 0.02}
height={dgxH}
fill="#8a7a56"
/>
<rect
x={dgxW - dgxW * 0.02}
y="0"
width={dgxW * 0.02}
height={dgxH}
fill="#8a7a56"
/>
<!-- Left handle -->
<rect
x={dgxW * 0.04}
y={dgxH * 0.08}
width={dgxW * 0.22}
height={dgxH * 0.84}
rx="2"
fill="#b3a170"
stroke="#403723"
stroke-width="0.5"
/>
<!-- Right handle -->
<rect
x={dgxW - dgxW * 0.04 - dgxW * 0.22}
y={dgxH * 0.08}
width={dgxW * 0.22}
height={dgxH * 0.84}
rx="2"
fill="#b3a170"
stroke="#403723"
stroke-width="0.5"
/>
<!-- Memory fill -->
<rect
x="2"
y={dgxH - dgxH * (node.currentPercent / 100)}
width={dgxW - 4}
height={dgxH * (node.currentPercent / 100)}
fill="rgba(255,215,0,0.35)"
/>
{#if node.modelUsageGB > 0 && node.isUsed}
<rect
x="2"
y={dgxH - dgxH * (node.newPercent / 100)}
width={dgxW - 4}
height={dgxH *
((node.newPercent - node.currentPercent) / 100)}
fill="#FFD700"
filter="url(#memGlow-{filterId})"
class="animate-pulse-slow"
/>
{/if}
</g>
{:else}
<!-- Unknown device - hexagon -->
<g
@@ -9,6 +9,7 @@
capabilities?: string[];
family?: string;
is_custom?: boolean;
requires_vllm?: boolean;
}
interface ModelGroup {
@@ -19,6 +20,7 @@
variants: ModelInfo[];
smallestVariant: ModelInfo;
hasMultipleVariants: boolean;
requiresVllm: boolean;
}
type DownloadAvailability = {
@@ -213,6 +215,14 @@
<span class="font-mono text-sm text-white truncate">
{group.name}
</span>
{#if group.requiresVllm}
<span
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
title="Requires vLLM runtime"
>
vLLM
</span>
{/if}
<!-- Capability icons -->
{#each group.capabilities.filter((c) => c !== "text") as cap}
{#if cap === "thinking"}
@@ -523,6 +533,15 @@
{variant.quantization || "default"}
</span>
{#if variant.requires_vllm}
<span
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
title="Requires vLLM runtime"
>
vLLM
</span>
{/if}
<!-- Size -->
<span
class="text-xs font-mono flex-1 {getSizeClassForFitStatus(
@@ -628,6 +647,7 @@
variants: [variant],
smallestVariant: variant,
hasMultipleVariants: false,
requiresVllm: variant.requires_vllm === true,
});
}}
title="View variant details"
@@ -22,6 +22,7 @@
is_custom?: boolean;
tasks?: string[];
hugging_face_id?: string;
requires_vllm?: boolean;
}
interface ModelGroup {
@@ -32,6 +33,7 @@
variants: ModelInfo[];
smallestVariant: ModelInfo;
hasMultipleVariants: boolean;
requiresVllm: boolean;
}
interface FilterState {
@@ -396,6 +398,7 @@
variants: [],
smallestVariant: model,
hasMultipleVariants: false,
requiresVllm: true,
});
}
@@ -430,6 +433,7 @@
(a.storage_size_megabytes || 0) - (b.storage_size_megabytes || 0),
);
group.hasMultipleVariants = group.variants.length > 1;
group.requiresVllm = group.variants.every((v) => v.requires_vllm);
}
// Convert to array and sort by smallest variant size (biggest first)
@@ -587,6 +591,7 @@
variants: [model],
smallestVariant: model,
hasMultipleVariants: false,
requiresVllm: model.requires_vllm === true,
});
}
}
@@ -1165,6 +1170,17 @@
<span class="text-white/40">Variants:</span>
<span class="text-white/70">{infoGroup.variants.length}</span>
</div>
{#if infoGroup.requiresVllm}
<div class="flex items-center gap-2">
<span class="text-white/40">Runtime:</span>
<span
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 tracking-wider uppercase"
>
vLLM
</span>
<span class="text-white/40 text-[11px]">required</span>
</div>
{/if}
{#if infoGroup.variants.length > 0}
<div class="mt-3 pt-3 border-t border-exo-yellow/10">
<span class="text-white/40">Available quantizations:</span>
@@ -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-bound pass that consumes the entire prompt and
builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
tokens sequentially from that cache. The two phases have very different
bottlenecks, so running them on different hardware can be substantially
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>
@@ -117,6 +117,10 @@
const LOGO_NATIVE_WIDTH = 814;
const LOGO_NATIVE_HEIGHT = 1000;
// NVIDIA logo SVG path (from exo-nvidia)
const NVIDIA_LOGO_PATH =
"M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
function formatBytes(bytes: number, decimals = 1): string {
if (!bytes || bytes === 0) return "0B";
const k = 1024;
@@ -554,6 +558,13 @@
const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
const modelLower = modelId.toLowerCase();
const identity = identitiesData[nodeInfo.id];
const nameLower = (friendlyName || "").toLowerCase();
const isSpark = modelLower.includes("dgx") || modelLower.includes("gx10");
const isLinux =
!isSpark &&
(modelLower.startsWith("linux") || identity?.osVersion === "Linux");
const isLinuxLaptop = isLinux && modelLower.includes("laptop");
// Check node states for styling
const isHighlighted = highlightedNodes.has(nodeInfo.id);
@@ -623,7 +634,382 @@
`${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
);
if (modelLower === "mac studio") {
if (isSpark) {
// NVIDIA DGX Spark — gold chassis with textured front, side handles, and NVIDIA badge
iconBaseWidth = nodeRadius * 1.55;
iconBaseHeight = nodeRadius * 0.58;
const x = nodeInfo.x - iconBaseWidth / 2;
const y = nodeInfo.y - iconBaseHeight / 2;
const chassisX = x - iconBaseWidth * 0.03;
const chassisWidth = iconBaseWidth * 1.05;
const cornerRadius = 3;
const dgxClipId = `dgx-clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
defs
.append("clipPath")
.attr("id", dgxClipId)
.append("rect")
.attr("x", x)
.attr("y", y)
.attr("width", iconBaseWidth)
.attr("height", iconBaseHeight)
.attr("rx", cornerRadius);
// Chassis texture pattern
const textureId = `chassis-texture-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
defs
.append("pattern")
.attr("id", textureId)
.attr("patternUnits", "userSpaceOnUse")
.attr("width", 8)
.attr("height", 8);
const texturePattern = defs.select(`#${textureId}`);
texturePattern
.append("rect")
.attr("width", 8)
.attr("height", 8)
.attr("fill", "#6f6248");
texturePattern
.append("circle")
.attr("cx", 2)
.attr("cy", 2)
.attr("r", 1)
.attr("fill", "#5a4f3b")
.attr("opacity", 0.5);
texturePattern
.append("circle")
.attr("cx", 6)
.attr("cy", 6)
.attr("r", 1)
.attr("fill", "#4a4232")
.attr("opacity", 0.45);
// Main body
nodeG
.append("rect")
.attr("class", "node-outline")
.attr("x", chassisX)
.attr("y", y)
.attr("width", chassisWidth)
.attr("height", iconBaseHeight)
.attr("rx", cornerRadius)
.attr("fill", `url(#${textureId})`)
.attr("stroke", wireColor)
.attr("stroke-width", strokeWidth);
// Side border accents
const sideThickness = iconBaseWidth * 0.02;
nodeG
.append("rect")
.attr("x", chassisX)
.attr("y", y)
.attr("width", sideThickness)
.attr("height", iconBaseHeight)
.attr("fill", "#8a7a56");
nodeG
.append("rect")
.attr("x", chassisX + chassisWidth - sideThickness)
.attr("y", y)
.attr("width", sideThickness)
.attr("height", iconBaseHeight)
.attr("fill", "#8a7a56");
// Memory fill (bottom up)
if (ramUsagePercent > 0) {
const memFillHeight = (ramUsagePercent / 100) * iconBaseHeight;
nodeG
.append("rect")
.attr("x", x)
.attr("y", y + iconBaseHeight - memFillHeight)
.attr("width", iconBaseWidth)
.attr("height", memFillHeight)
.attr("fill", "rgba(255,215,0,0.45)")
.attr("clip-path", `url(#${dgxClipId})`);
}
// Side handles with inner recess
const handleWidth = iconBaseWidth * 0.27;
const handleGap = iconBaseHeight * 0.05;
const handleHeight = iconBaseHeight - handleGap * 2;
const handleY = y + handleGap;
const innerHandleWidth = iconBaseWidth * 0.12;
const innerHandleHeight = handleHeight - iconBaseHeight * 0.06;
const leftHandleX = x + 4;
const rightHandleX = x + iconBaseWidth - handleWidth - 4;
// Left handle
nodeG
.append("rect")
.attr("x", leftHandleX)
.attr("y", handleY)
.attr("width", handleWidth)
.attr("height", handleHeight)
.attr("rx", 2.4)
.attr("fill", "#b3a170")
.attr("stroke", "#403723")
.attr("stroke-width", 0.7);
nodeG
.append("rect")
.attr("x", leftHandleX + handleWidth * 0.06)
.attr("y", handleY + iconBaseHeight * 0.03)
.attr("width", innerHandleWidth)
.attr("height", innerHandleHeight)
.attr("rx", 1.6)
.attr("fill", "#8a7a56");
// Right handle
nodeG
.append("rect")
.attr("x", rightHandleX)
.attr("y", handleY)
.attr("width", handleWidth)
.attr("height", handleHeight)
.attr("rx", 2.4)
.attr("fill", "#b3a170")
.attr("stroke", "#403723")
.attr("stroke-width", 0.7);
nodeG
.append("rect")
.attr(
"x",
rightHandleX + handleWidth - innerHandleWidth - handleWidth * 0.08,
)
.attr("y", handleY + iconBaseHeight * 0.03)
.attr("width", innerHandleWidth)
.attr("height", innerHandleHeight)
.attr("rx", 1.6)
.attr("fill", "#8a7a56");
// NVIDIA logo + text label (rotated 90 deg on left handle)
const badgeWidth = iconBaseWidth * 0.09;
const badgeHeight = handleHeight * 0.5;
const badgeX =
leftHandleX + handleWidth - badgeWidth - handleWidth * 0.06;
const badgeY = handleY + (handleHeight - badgeHeight) / 2;
const textSize = badgeWidth * 0.58;
const logoWidth = textSize * 1.2;
const logoHeight = logoWidth * (1.438 / 2.174);
const centerX = badgeX + badgeWidth / 2 - badgeWidth * 0.03;
const centerY = badgeY + badgeHeight / 2;
const gap = badgeWidth * 0.15;
const totalWidth = logoWidth + gap + textSize * 3.6;
const labelGroup = nodeG
.append("g")
.attr("transform", `rotate(90 ${centerX} ${centerY})`);
labelGroup
.append("svg")
.attr("x", centerX - totalWidth / 2)
.attr("y", centerY - logoHeight / 2)
.attr("width", logoWidth)
.attr("height", logoHeight)
.attr("viewBox", "0 0 2.174 1.438")
.append("path")
.attr("d", NVIDIA_LOGO_PATH)
.attr("fill", "#76b900");
labelGroup
.append("text")
.attr("x", centerX - totalWidth / 2 + logoWidth + gap)
.attr("y", centerY)
.attr("text-anchor", "start")
.attr("dominant-baseline", "middle")
.attr("fill", "#8a7a56")
.attr("font-size", textSize)
.attr("font-family", "monospace")
.attr("font-weight", "700")
.text("NVIDIA");
} else if (isLinuxLaptop) {
// Linux Laptop — same shape as MacBook but with Tux logo
iconBaseWidth = nodeRadius * 1.6;
iconBaseHeight = nodeRadius * 1.15;
const x = nodeInfo.x - iconBaseWidth / 2;
const y = nodeInfo.y - iconBaseHeight / 2;
const screenHeight = iconBaseHeight * 0.7;
const baseHeight = iconBaseHeight * 0.3;
const screenWidth = iconBaseWidth * 0.85;
const screenX = nodeInfo.x - screenWidth / 2;
const screenBezel = 3;
const linuxScreenClipId = `linux-screen-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
defs
.append("clipPath")
.attr("id", linuxScreenClipId)
.append("rect")
.attr("x", screenX + screenBezel)
.attr("y", y + screenBezel)
.attr("width", screenWidth - screenBezel * 2)
.attr("height", screenHeight - screenBezel * 2)
.attr("rx", 2);
// Screen outer frame
nodeG
.append("rect")
.attr("class", "node-outline")
.attr("x", screenX)
.attr("y", y)
.attr("width", screenWidth)
.attr("height", screenHeight)
.attr("rx", 3)
.attr("fill", "#1a1a1a")
.attr("stroke", wireColor)
.attr("stroke-width", strokeWidth);
// Screen inner
nodeG
.append("rect")
.attr("x", screenX + screenBezel)
.attr("y", y + screenBezel)
.attr("width", screenWidth - screenBezel * 2)
.attr("height", screenHeight - screenBezel * 2)
.attr("rx", 2)
.attr("fill", "#0a0a12");
// Memory fill on screen
if (ramUsagePercent > 0) {
const memFillTotalHeight = screenHeight - screenBezel * 2;
const memFillActualHeight =
(ramUsagePercent / 100) * memFillTotalHeight;
nodeG
.append("rect")
.attr("x", screenX + screenBezel)
.attr(
"y",
y + screenBezel + (memFillTotalHeight - memFillActualHeight),
)
.attr("width", screenWidth - screenBezel * 2)
.attr("height", memFillActualHeight)
.attr("fill", "rgba(255,215,0,0.85)")
.attr("clip-path", `url(#${linuxScreenClipId})`);
}
// Terminal prompt on screen
nodeG
.append("text")
.attr("x", nodeInfo.x)
.attr("y", y + screenHeight / 2)
.attr("text-anchor", "middle")
.attr("dominant-baseline", "middle")
.attr("fill", "#FFFFFF")
.attr("opacity", 0.9)
.attr("font-size", screenHeight * 0.25)
.attr("font-family", "SF Mono, Monaco, monospace")
.attr("font-weight", "700")
.text(">_");
// Keyboard base (trapezoidal)
const baseY = y + screenHeight;
const baseTopWidth = screenWidth;
const baseBottomWidth = iconBaseWidth;
const baseTopX = nodeInfo.x - baseTopWidth / 2;
const baseBottomX = nodeInfo.x - baseBottomWidth / 2;
nodeG
.append("path")
.attr(
"d",
`M ${baseTopX} ${baseY} L ${baseTopX + baseTopWidth} ${baseY} L ${baseBottomX + baseBottomWidth} ${baseY + baseHeight} L ${baseBottomX} ${baseY + baseHeight} Z`,
)
.attr("fill", "#2c2c2c")
.attr("stroke", wireColor)
.attr("stroke-width", 1);
// Keyboard area
const keyboardX = baseTopX + 6;
const keyboardY = baseY + 3;
const keyboardWidth = baseTopWidth - 12;
const keyboardHeight = baseHeight * 0.55;
nodeG
.append("rect")
.attr("x", keyboardX)
.attr("y", keyboardY)
.attr("width", keyboardWidth)
.attr("height", keyboardHeight)
.attr("fill", "rgba(0,0,0,0.2)")
.attr("rx", 2);
// Trackpad
const trackpadWidth = baseTopWidth * 0.4;
const trackpadX = nodeInfo.x - trackpadWidth / 2;
const trackpadY = baseY + keyboardHeight + 5;
const trackpadHeight = baseHeight * 0.3;
nodeG
.append("rect")
.attr("x", trackpadX)
.attr("y", trackpadY)
.attr("width", trackpadWidth)
.attr("height", trackpadHeight)
.attr("fill", "rgba(255,255,255,0.08)")
.attr("rx", 2);
} else if (isLinux) {
// Linux Desktop — same shape as Mac Studio but with Tux logo
iconBaseWidth = nodeRadius * 1.25;
iconBaseHeight = nodeRadius * 0.85;
const x = nodeInfo.x - iconBaseWidth / 2;
const y = nodeInfo.y - iconBaseHeight / 2;
const cornerRadius = 4;
const topSurfaceHeight = iconBaseHeight * 0.15;
const linuxDesktopClipId = `linux-desktop-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
defs
.append("clipPath")
.attr("id", linuxDesktopClipId)
.append("rect")
.attr("x", x)
.attr("y", y + topSurfaceHeight)
.attr("width", iconBaseWidth)
.attr("height", iconBaseHeight - topSurfaceHeight)
.attr("rx", cornerRadius - 1);
// Main body
nodeG
.append("rect")
.attr("class", "node-outline")
.attr("x", x)
.attr("y", y)
.attr("width", iconBaseWidth)
.attr("height", iconBaseHeight)
.attr("rx", cornerRadius)
.attr("fill", "#1a1a1a")
.attr("stroke", wireColor)
.attr("stroke-width", strokeWidth);
// Memory fill
if (ramUsagePercent > 0) {
const memFillTotalHeight = iconBaseHeight - topSurfaceHeight;
const memFillActualHeight =
(ramUsagePercent / 100) * memFillTotalHeight;
nodeG
.append("rect")
.attr("x", x)
.attr(
"y",
y + topSurfaceHeight + (memFillTotalHeight - memFillActualHeight),
)
.attr("width", iconBaseWidth)
.attr("height", memFillActualHeight)
.attr("fill", "rgba(255,215,0,0.75)")
.attr("clip-path", `url(#${linuxDesktopClipId})`);
}
// Terminal prompt on front face
nodeG
.append("text")
.attr("x", nodeInfo.x)
.attr(
"y",
y + topSurfaceHeight + (iconBaseHeight - topSurfaceHeight) / 2,
)
.attr("text-anchor", "middle")
.attr("dominant-baseline", "middle")
.attr("fill", "rgba(255,255,255,0.5)")
.attr("font-size", (iconBaseHeight - topSurfaceHeight) * 0.4)
.attr("font-family", "SF Mono, Monaco, monospace")
.attr("font-weight", "700")
.text(">_");
} else if (modelLower === "mac studio") {
// Mac Studio - classic cube with memory fill
iconBaseWidth = nodeRadius * 1.25;
iconBaseHeight = nodeRadius * 0.85;
@@ -1182,8 +1568,12 @@
debugLabelY += debugLineHeight;
}
const identity = identitiesData[nodeInfo.id];
if (identity?.osVersion) {
const dbgIdentity = identitiesData[nodeInfo.id];
if (dbgIdentity?.osVersion) {
const osLabel =
dbgIdentity.osVersion === "Linux"
? "Linux"
: `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
nodeG
.append("text")
.attr("x", nodeInfo.x)
@@ -1192,9 +1582,7 @@
.attr("fill", "rgba(179,179,179,0.7)")
.attr("font-size", debugFontSize)
.attr("font-family", "SF Mono, Monaco, monospace")
.text(
`macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
);
.text(osLabel);
}
}
});
+128 -10
View File
@@ -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;
+44
View File
@@ -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;
}
+82 -14
View File
@@ -65,6 +65,7 @@
nodeThunderboltBridge,
nodeIdentities,
isConnected,
featureFlags,
type DownloadProgress,
type PlacementPreview,
} from "$lib/stores/app.svelte";
@@ -702,7 +703,10 @@
? Object.keys(topologyData()!.nodes).length
: 1;
const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
const instanceType =
nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
? "MlxRing"
: selectedInstanceType;
try {
const placementResponse = await fetch(
`/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
@@ -783,6 +787,7 @@
quantization?: string;
base_model?: string;
capabilities?: string[];
requires_vllm?: boolean;
}>
>([]);
type ModelMemoryFitStatus =
@@ -886,7 +891,7 @@
}
let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
type InstanceMeta = "MlxRing" | "MlxJaccl";
type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
// Launch defaults persistence
const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
@@ -932,7 +937,12 @@
// Apply sharding and instance type unconditionally
selectedSharding = defaults.sharding;
selectedInstanceType =
defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
defaults.instanceType === "MlxRing"
? "MlxRing"
: defaults.instanceType === "Vllm"
? "Vllm"
: "MlxJaccl";
userPickedInstanceType = true;
// Apply minNodes if valid (between 1 and maxNodes)
if (
@@ -954,6 +964,23 @@
}
let selectedInstanceType = $state<InstanceMeta>("MlxRing");
let userPickedInstanceType = $state(false);
$effect(() => {
if (!userPickedInstanceType && featureFlags()["vllm_available"]) {
selectedInstanceType = "Vllm";
}
});
const selectedModelRequiresVllm = $derived.by((): boolean => {
const id = selectedPreviewModelId();
if (!id) return false;
const model = models.find((m) => m.id === id);
return model?.requires_vllm === true;
});
$effect(() => {
if (selectedModelRequiresVllm) {
selectedInstanceType = "Vllm";
}
});
let selectedMinNodes = $state<number>(1);
let minNodesInitialized = $state(false);
let launchingModelId = $state<string | null>(null);
@@ -1146,9 +1173,7 @@
}
const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
selectedInstanceType === "MlxRing"
? runtime === "MlxRing"
: runtime === "MlxJaccl";
runtime === selectedInstanceType;
// Helper to check if a model can be launched (has valid placement with >= minNodes)
function canModelFit(modelId: string): boolean {
@@ -2063,6 +2088,7 @@
let instanceType = "Unknown";
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
else if (instanceTag === "VllmInstance") instanceType = "vLLM";
const inst = instance as {
shardAssignments?: {
@@ -5769,14 +5795,18 @@
</div>
<div class="flex gap-2">
<button
disabled={selectedModelRequiresVllm}
onclick={() => {
if (selectedModelRequiresVllm) return;
selectedInstanceType = "MlxRing";
userPickedInstanceType = true;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
'MlxRing'
? 'bg-transparent text-exo-yellow border-exo-yellow'
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
: selectedInstanceType === 'MlxRing'
? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
: 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
>
<span
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5792,14 +5822,18 @@
TCP/IP
</button>
<button
disabled={selectedModelRequiresVllm}
onclick={() => {
if (selectedModelRequiresVllm) return;
selectedInstanceType = "MlxJaccl";
userPickedInstanceType = true;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
'MlxJaccl'
? 'bg-transparent text-exo-yellow border-exo-yellow'
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
: selectedInstanceType === 'MlxJaccl'
? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
: 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
>
<span
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5814,7 +5848,41 @@
</span>
RDMA (Fast)
</button>
{#if featureFlags()["vllm_available"] || selectedModelRequiresVllm}
<button
onclick={() => {
selectedInstanceType = "Vllm";
userPickedInstanceType = true;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
'Vllm'
? 'bg-transparent text-exo-yellow border-exo-yellow'
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
>
<span
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
'Vllm'
? 'border-exo-yellow'
: 'border-exo-medium-gray'}"
>
{#if selectedInstanceType === "Vllm"}
<span
class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
></span>
{/if}
</span>
vLLM (CUDA)
</button>
{/if}
</div>
{#if selectedModelRequiresVllm}
<div
class="mt-2 text-[11px] font-mono text-orange-300/80"
>
This model requires vLLM.
</div>
{/if}
</div>
<!-- Minimum Devices -->
@@ -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>
+33 -1
View File
@@ -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[] = [];
@@ -132,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) };
@@ -139,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) {
@@ -350,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 */
}
+1 -1
View File
@@ -146,7 +146,7 @@
config.treefmt.build.wrapper
# PYTHON
self'.packages.editableVenv
self'.packages.exo.passthru.evenv
uv
# RUST
+13
View File
@@ -40,6 +40,19 @@ build-app: rust-rebuild sync-clean package
xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
@echo "\nBuild complete. Run with:\n open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
sync-cuda:
#!/usr/bin/env bash
set -euo pipefail
uv sync --extra vllm-cuda13 --extra mlx-cpu --no-install-package vllm
dest=".venv/lib/python3.13/site-packages"
[[ -d $dest/vllm ]] || {
nix build .#exo-cuda-13.passthru.evenv
# will also grab vllm-0.19.1-distinfo
cp -aL result/lib/python3.13/site-packages/vllm* .venv/lib/python3.13/site-packages
chmod -R u+rwX .venv/lib/python3.13/site-packages/vllm*
rm result
}
clean:
rm -rf **/__pycache__
rm -rf target/
+26
View File
@@ -0,0 +1,26 @@
diff --git a/setup.py b/setup.py
index 6dc2ed028..bdcc6354a 100644
--- a/setup.py
+++ b/setup.py
@@ -18,6 +18,13 @@ from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
+if "NIX_ATTRS_JSON_FILE" in os.environ:
+ with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
+ NIX_ATTRS = json.load(f)
+else:
+ NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
+
+
def load_module_from_path(module_name, path):
spec = importlib.util.spec_from_file_location(module_name, path)
module = importlib.util.module_from_spec(spec)
@@ -184,6 +191,7 @@ class cmake_build_ext(build_ext):
cmake_args = [
"-DCMAKE_BUILD_TYPE={}".format(cfg),
"-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
+ *NIX_ATTRS["cmakeFlags"],
]
verbose = envs.VERBOSE
+59 -35
View File
@@ -15,21 +15,18 @@ dependencies = [
"huggingface-hub>=1.8.0",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo-pyo3-bindings", # rust bindings
"exo-pyo3-bindings", # rust bindings
"anyio==4.11.0",
"mlx==0.31.2; sys_platform == 'darwin'",
"mlx-lm; sys_platform=='darwin'",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
"openai-harmony>=0.0.8",
"httpx>=0.28.1",
"tomlkit>=0.14.0",
"mflux==0.17.2; sys_platform == 'darwin'",
"python-multipart>=0.0.21",
"msgspec>=0.19.0",
"zstandard>=0.23.0",
"mlx-vlm>=0.3.11",
"transformers>=5.0.0,<5.4.0",
"transformers>=5.6.2",
"nvidia-ml-py>=13.595.45",
]
[project.scripts]
@@ -48,23 +45,30 @@ dev = [
[project.optional-dependencies]
build = ["nanobind"]
cpu = [
"mlx==0.31.1; sys_platform == 'linux'",
"mlx-cpu==0.31.1; sys_platform == 'linux'",
"mlx-lm; sys_platform == 'linux'",
"torch>=2.10.0; sys_platform == 'linux'",
mlx-none = ["anyio"]
mlx = [
"mlx==0.31.2",
"mlx-lm",
"mlx-vlm>=0.3.11",
"mflux==0.17.5",
# pinning vllms versions for consistency.
"torch==2.10.0; sys_platform == 'darwin'",
"torch==2.10.0; sys_platform == 'linux'",
"torchaudio==2.10.0; sys_platform == 'darwin'",
"torchaudio==2.10.0; sys_platform == 'linux'",
"torchvision==0.25.0; sys_platform == 'darwin'",
"torchvision==0.25.0; sys_platform == 'linux'",
]
cuda12 = [
"mlx==0.31.1; sys_platform == 'linux'",
"mlx-cuda-12==0.31.1; sys_platform == 'linux'",
"mlx-lm; sys_platform == 'linux'",
"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'",
"torch>=2.10.0; sys_platform == 'linux'",
mlx-cpu = ["exo[mlx]", "mlx-cpu==0.31.2; sys_platform == 'linux'"]
mlx-cuda12 = ["exo[mlx]", "mlx-cuda-12==0.31.1; sys_platform == 'linux'"]
mlx-cuda13 = ["exo[mlx]", "mlx-cuda-13==0.31.1; sys_platform == 'linux'"]
vllm-none = ["anyio"]
vllm-cuda13 = [
"vllm[cuda13, fastsafetensors]; sys_platform == 'linux'",
"torch==2.10.0; sys_platform == 'linux'",
"torchaudio==2.10.0; sys_platform == 'linux'",
"torchvision==0.25.0; sys_platform == 'linux'",
]
###
@@ -77,13 +81,24 @@ members = ["rust/exo_pyo3_bindings", "bench"]
[tool.uv.sources]
exo-pyo3-bindings = { workspace = true }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
mflux = { git = "http://github.com/evanev7/mflux", branch = "exo" }
vllm = { git = "http://github.com/evanev7/vllm", branch = "exo2" }
torch = [
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
{ index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
{ index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
]
torchaudio = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
]
vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
[[tool.uv.index]]
name = "pytorch-cu130"
@@ -91,8 +106,8 @@ url = "https://download.pytorch.org/whl/cu130"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu120"
url = "https://download.pytorch.org/whl/cu120"
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
[[tool.uv.index]]
@@ -153,11 +168,19 @@ root = "src"
required-version = ">=0.8.6"
prerelease = "allow"
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
override-dependencies = [
"mlx==0.31.1; sys_platform=='linux'",
"mlx; sys_platform=='darwin'",
override-dependencies = ["opencv-python; python_version < '0'"]
conflicts = [
[
{ extra = "mlx-cuda13" },
{ extra = "mlx-cuda12" },
{ extra = "mlx-cpu" },
{ extra = "mlx-none" },
],
[
{ extra = "vllm-cuda13" },
{ extra = "mlx-cuda12" },
{ extra = "vllm-none" },
],
]
[tool.uv.extra-build-dependencies]
@@ -172,6 +195,7 @@ mlx = [
"ninja",
]
mlx-lm = ["setuptools"]
mflux = ["uv_build"]
xgrammar = [
"nanobind",
"setuptools",
+202 -20
View File
@@ -10,8 +10,10 @@ let
inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
inherit (pkgs.config) cudaSupport;
inherit (pkgs) cudaPackages;
cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
libmlx_source =
if (builtins.elem "mlx-cuda13" members.exo or [ ]) then "mlx-cuda-13"
else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
else "mlx-cpu";
python = pkgs.python313;
cudaLibs = with cudaPackages; [
cuda_cudart
@@ -113,37 +115,213 @@ let
});
} // lib.optionalAttrs isLinux {
mlx = prev.mlx.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ lib.optionals cudaSupport [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
postInstall = ''
cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
'';
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
} // lib.optionalAttrs cudaSupport {
"${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
buildInputs = old.buildInputs ++ [ cudaLibs ];
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
});
torch = prev.torch.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
torchaudio = prev.torchaudio.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ [ cudaPackages.cuda_cudart ];
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
});
torchvision = prev.torchvision.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
});
torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old: {
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
});
# Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
vllm = prev.vllm.overrideAttrs (old:
let
cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
mkdir -p $out/include
ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
'';
cudaRoot = pkgs.symlinkJoin {
name = "cuda-merged-exo";
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
};
cutlass = pkgs.fetchFromGitHub {
name = "cutlass-source";
owner = "NVIDIA";
repo = "cutlass";
tag = "v4.2.1";
hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
};
triton-kernels = pkgs.fetchFromGitHub {
owner = "triton-lang";
repo = "triton";
tag = "v3.6.0";
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
};
cutlass-flashmla = pkgs.fetchFromGitHub {
owner = "NVIDIA";
repo = "cutlass";
rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
};
flashmla = pkgs.stdenv.mkDerivation {
pname = "flashmla";
version = "1.0.0";
src = pkgs.fetchFromGitHub {
name = "FlashMLA-source";
owner = "vllm-project";
repo = "FlashMLA";
rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
};
dontConfigure = true;
buildPhase = ''
rm -rf csrc/cutlass
ln -sf ${cutlass-flashmla} csrc/cutlass
'';
installPhase = ''
cp -rva . $out
'';
};
qutlass = pkgs.fetchFromGitHub {
name = "qutlass-source";
owner = "IST-DASLab";
repo = "qutlass";
rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
};
vllm-flash-attn = pkgs.stdenv.mkDerivation {
pname = "vllm-flash-attn";
version = "2.7.2.post1";
src = pkgs.fetchFromGitHub {
name = "flash-attention-source";
owner = "vllm-project";
repo = "flash-attention";
rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
};
patches = [
(pkgs.fetchpatch {
url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
})
(pkgs.fetchpatch {
url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
})
];
dontConfigure = true;
buildPhase = ''
rm -rf csrc/cutlass
ln -sf ${cutlass} csrc/cutlass
'';
installPhase = ''
cp -rva . $out
'';
};
in
{
patches = (old.patches or [ ]) ++ [ ../nix/vllm-setuppy-cmake.patch ];
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
pkgs.cmake
pkgs.ninja
pkgs.autoAddDriverRunpath
] ++ lib.optionals cudaSupport [
cudaPackages.cuda_nvcc
];
# TODO: vllm rocm/cpu
VLLM_TARGET_DEVICE = "empty";
preConfigure = ''
export MAX_JOBS="$NIX_BUILD_CORES"
'';
# TODO: vllm non cuda13 support, more arch's, etc.
} // lib.optionalAttrs cudaSupport {
buildInputs = cudaLibs ++ [ cudaRoot ];
VLLM_CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
CUDA_HOME = "${cudaRoot}";
CUDAToolkit_ROOT = "${cudaRoot}";
CUDACXX = "${cudaRoot}/bin/nvcc";
VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
VLLM_TARGET_DEVICE = "cuda";
TORCH_CUDA_ARCH_LIST = "12.0;12.1";
TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
CAFFE2_USE_CUDNN = "ON";
CAFFE2_USE_CUFILE = "ON";
CUTLASS_ENABLE_CUBLAS = "ON";
CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
cmakeFlags = [
(lib.cmakeBool "CMAKE_SKIP_INSTALL_RPATH" true)
(lib.cmakeBool "CMAKE_BUILD_WITH_INSTALL_RPATH" true)
(lib.cmakeFeature "CUDA_HOME" "${cudaRoot}")
(lib.cmakeFeature "CUDAToolkit_ROOT" "${cudaRoot}")
(lib.cmakeFeature "CMAKE_CUDA_COMPILER" "${cudaRoot}/bin/nvcc")
(lib.cmakeFeature "CMAKE_PREFIX_PATH" "${cudaRoot}")
(lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}")
(lib.cmakeFeature "FLASH_MLA_SRC_DIR" "${lib.getDev flashmla}")
(lib.cmakeFeature "VLLM_FLASH_ATTN_SRC_DIR" "${lib.getDev vllm-flash-attn}")
(lib.cmakeFeature "QUTLASS_SRC_DIR" "${lib.getDev qutlass}")
(lib.cmakeFeature "TORCH_CUDA_ARCH_LIST" "12.0;12.1")
(lib.cmakeFeature "CUTLASS_NVCC_ARCHS_ENABLED" "${cudaPackages.flags.cmakeCudaArchitecturesString}")
(lib.cmakeFeature "CUDA_TOOLKIT_ROOT_DIR" "${cudaRoot}")
(lib.cmakeFeature "CAFFE2_USE_CUDNN" "ON")
(lib.cmakeFeature "CAFFE2_USE_CUFILE" "ON")
(lib.cmakeFeature "CUTLASS_ENABLE_CUBLAS" "ON")
];
});
} // lib.optionalAttrs (cudaSupport && isx86_64) {
numba = prev.numba.overrideAttrs (old: {
buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
});
};
pyprojectOverlay = workspace.mkPyprojectOverlay {
sourcePreference = "wheel";
@@ -164,24 +342,28 @@ let
buildSystemsOverlay
]
);
venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
mkApp = cmd: name: pkgs.writeShellApplication {
# mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
venv = name: (pythonSet.mkVirtualEnv "${name}-venv" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
mkApp = text: name: pkgs.writeShellApplication {
inherit name;
text = "exec " + lib.optionalString cudaSupport "nixglhost " + text;
runtimeEnv = {
EXO_DASHBOARD_DIR = self'.packages.dashboard;
EXO_RESOURCES_DIR = inputs.self + /resources;
};
runtimeInputs = [
# mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
(venv name)
pkgs.nix-gl-host
]
++ lib.optionals isDarwin [ pkgs.macmon ];
text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
passthru = {
venv = venv name;
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
};
};
in
{
inherit venv;
editablePythonSet = pythonSet.overrideScope editableOverlay;
mkPythonScript = path: mkApp ''python ${path} "$@"'';
mkExo = mkApp ''exo "$@"'';
};
@@ -191,18 +373,18 @@ in
{ self', pkgs, unfreePkgs, lib, ... }:
let
inherit (pkgs.stdenv.hostPlatform) isLinux;
inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
# Virtual environment with dev dependencies for testing
testVenv = (mkPythonSet {
inherit self' pkgs lib; members = {
exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
};
}).venv "exo-test";
mkBenchScript = (mkPythonSet {
inherit self' pkgs lib; members = {
exo = [ "cpu" ];
exo = [ "mlx-cpu" "vllm-none" ];
exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
};
}).mkPythonScript;
@@ -212,12 +394,12 @@ in
runtimeInputs = [ pkgs.python313 ];
text = ''exec python ${path} "$@"'';
};
cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" "vllm-none" ]; }; };
cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cpu" "vllm-cuda13" ]; }; };
in
{
packages = {
exo = mkExo "exo";
editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
# for running tests in ci
exo-test-env = testVenv;
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
@@ -226,8 +408,8 @@ in
# used by ./tests/run_exo_on.sh
exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
} // lib.optionalAttrs isLinux {
exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
};
checks = {
@@ -0,0 +1,21 @@
model_id = "2imi9/gpt-oss-20B-NVFP4A16-BF16"
n_layers = 24
hidden_size = 2880
num_key_value_heads = 8
supports_tensor = false
tasks = ["TextGeneration"]
family = "gpt-oss"
quantization = "nvfp4"
base_model = "GPT-OSS 20B"
capabilities = ["text", "thinking"]
reasoning_dialect = "channel"
context_length = 131072
requires_vllm = true
[storage_size]
in_bytes = 41829514752
[sampling_defaults]
temperature = 1.0
top_p = 1.0
top_k = 0
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -8,7 +8,7 @@ family = "glm"
quantization = "8bit"
base_model = "GLM-5.1"
capabilities = ["text", "thinking"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
@@ -8,7 +8,7 @@ family = "glm"
quantization = "MXFP4-Q8"
base_model = "GLM-5.1"
capabilities = ["text", "thinking"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
@@ -8,7 +8,7 @@ family = "glm"
quantization = "bf16"
base_model = "GLM-5.1"
capabilities = ["text", "thinking"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
@@ -8,7 +8,7 @@ family = "kimi"
quantization = ""
base_model = "Kimi K2"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "suffix"
context_length = 262144
[storage_size]
@@ -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]
@@ -8,7 +8,7 @@ family = "kimi"
quantization = "3bit"
base_model = "Kimi K2.6"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "suffix"
context_length = 262144
[storage_size]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -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]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "6bit"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "bf16"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 27B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 27B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 2B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 397B A17B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "6bit"
base_model = "Qwen3.5 397B A17B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 397B A17B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 9B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 9B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "4bit"
base_model = "Qwen3.6 27B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.6 27B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "bf16"
base_model = "Qwen3.6 27B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "4bit"
base_model = "Qwen3.6 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "5bit"
base_model = "Qwen3.6 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "8bit"
base_model = "Qwen3.6 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "qwen"
quantization = "bf16"
base_model = "Qwen3.6 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
[storage_size]
@@ -8,7 +8,7 @@ family = "gpt-oss"
quantization = "MXFP4-Q8"
base_model = "GPT-OSS 120B"
capabilities = ["text", "thinking"]
reasoning_dialect = "channel"
context_length = 131072
[storage_size]
@@ -8,7 +8,7 @@ family = "gpt-oss"
quantization = "MXFP4-Q8"
base_model = "GPT-OSS 20B"
capabilities = ["text", "thinking"]
reasoning_dialect = "channel"
context_length = 131072
[storage_size]
@@ -0,0 +1,27 @@
model_id = "nvidia/Qwen3-30B-A3B-NVFP4"
n_layers = 48
hidden_size = 2048
num_key_value_heads = 4
supports_tensor = false
tasks = ["TextGeneration"]
family = "qwen"
quantization = "nvfp4"
base_model = "Qwen3 30B"
capabilities = ["text", "thinking", "thinking_toggle"]
context_length = 32768
requires_vllm = true
[storage_size]
in_bytes = 18087458688
[sampling_defaults]
temperature = 0.6
top_p = 0.95
top_k = 20
min_p = 0.0
[sampling_defaults.non_thinking]
temperature = 0.7
top_p = 0.8
top_k = 20
min_p = 0.0
@@ -0,0 +1,20 @@
model_id = "openai/gpt-oss-120b"
n_layers = 36
hidden_size = 2880
num_key_value_heads = 8
supports_tensor = false
tasks = ["TextGeneration"]
family = "gpt-oss"
quantization = "mxfp4"
base_model = "GPT-OSS 120B"
capabilities = ["text", "thinking"]
reasoning_dialect = "channel"
context_length = 131072
[storage_size]
in_bytes = 65248815744
[sampling_defaults]
temperature = 1.0
top_p = 1.0
top_k = 0
@@ -0,0 +1,32 @@
model_id = "sakamakismile/Qwen3.6-27B-NVFP4"
n_layers = 64
hidden_size = 5120
num_key_value_heads = 4
supports_tensor = false
tasks = ["TextGeneration"]
family = "qwen"
quantization = "nvfp4"
base_model = "Qwen3.6 27B"
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
reasoning_dialect = "post_last_user"
context_length = 262144
requires_vllm = true
[storage_size]
in_bytes = 16703361232
[sampling_defaults]
temperature = 1.0
top_p = 0.95
top_k = 20
min_p = 0.0
repetition_penalty = 1.0
presence_penalty = 1.5
[sampling_defaults.non_thinking]
temperature = 0.7
top_p = 0.8
top_k = 20
min_p = 0.0
repetition_penalty = 1.0
presence_penalty = 1.5
+222
View File
@@ -0,0 +1,222 @@
#!/usr/bin/env python
"""Standalone smoke test for VllmEngine.serve_prefill.
Loads a real vLLM engine, runs serve_prefill against an in-memory buffer
twice in a row with the same prompt, and verifies both runs produce a
well-formed wire stream (header -> KV chunks -> Done).
The second run is the regression guard: with vLLM APC enabled this would
trip the chunked-prefill + APC + custom kv-connector CUDA assert
(`vectorized_gather_kernel: ind >= ind_dim_size`) and the server would
close the socket before the Done frame.
Usage on the Spark (gx10-de89):
cd /home/larry/exo
/nix/store/2b82iz9ac0pxqafrgxmgdkq8sr2hwlx6-exo-cuda-13-venv/bin/python \\
scripts/check_serve_prefill.py Qwen/Qwen3-0.6B
Exits 0 on success, non-zero with a diagnostic on failure.
"""
from __future__ import annotations
import contextlib
import io
import os
import sys
import traceback
from pathlib import Path
from typing import cast
def _ensure_repo_on_path() -> None:
repo = Path(__file__).resolve().parent.parent
src = repo / "src"
if str(src) not in sys.path:
sys.path.insert(0, str(src))
_ensure_repo_on_path()
from exo.shared.types.common import ModelId # noqa: E402
from exo.worker.disaggregated.protocol import ( # noqa: E402
ArraysState,
Done,
ErrorMessage,
KVChunk,
read_header,
read_message,
)
from exo.worker.disaggregated.server import PrefillRequest # noqa: E402
def _make_token_ids(n: int) -> list[int]:
return [(i * 1009 + 17) % 30000 + 100 for i in range(n)]
def _decode(
payload: bytes,
) -> tuple[list[KVChunk], list[ArraysState], Done | None, ErrorMessage | None]:
buf = io.BytesIO(payload)
_ = read_header(buf)
chunks: list[KVChunk] = []
arrays: list[ArraysState] = []
done: Done | None = None
error: ErrorMessage | None = None
while True:
msg = read_message(buf)
if msg is None:
break
if isinstance(msg, KVChunk):
chunks.append(msg)
elif isinstance(msg, ArraysState):
arrays.append(msg)
elif isinstance(msg, Done):
done = msg
break
elif isinstance(msg, ErrorMessage):
error = msg
break
return chunks, arrays, done, error
def _build_engine(model_id: ModelId) -> object:
from exo.worker.engines.vllm.engine import VllmEngine
from exo.worker.engines.vllm.generator import VllmBatchEngine, load_vllm_engine
from exo.worker.engines.vllm.kv_connector import (
ExoKVProducerConnector,
_patch_gdn_capture,
_patch_vllm_for_connector,
)
_patch_vllm_for_connector(ExoKVProducerConnector)
_patch_gdn_capture()
llm_engine, tool_parser = load_vllm_engine(
model_id=model_id,
trust_remote_code=False,
n_layers=1,
kv_connector_cls=ExoKVProducerConnector,
)
gen = VllmBatchEngine(engine=llm_engine, model_id=model_id)
class _S:
def send(self, _: object) -> None: ...
class _R:
def collect(self) -> list[object]:
return []
return VllmEngine(
tool_parser=tool_parser,
model_id=model_id,
cancel_receiver=cast("object", _R()), # pyright: ignore[reportArgumentType]
event_sender=cast("object", _S()), # pyright: ignore[reportArgumentType]
_gen=gen,
max_concurrent_requests=1,
)
def _run_one(engine: object, n_tokens: int, label: str) -> int:
request = PrefillRequest(
request_id=f"check-{label}-{os.getpid()}",
model_id="ignored",
token_ids=_make_token_ids(n_tokens),
start_pos=0,
use_prefix_cache=True,
)
buf = io.BytesIO()
engine.serve_prefill(request, buf) # pyright: ignore[reportAttributeAccessIssue]
payload = buf.getvalue()
if not payload:
raise AssertionError(f"{label}: server wrote nothing")
chunks, arrays, done, error = _decode(payload)
if error is not None:
raise AssertionError(
f"{label}: server returned ErrorMessage [{error.code}]: {error.message}"
)
if done is None:
raise AssertionError(
f"{label}: stream did not end with Done "
f"({len(chunks)} kv chunks, {len(arrays)} arrays)"
)
if done.total_tokens <= 0:
raise AssertionError(f"{label}: Done reported {done.total_tokens} tokens")
if not chunks:
raise AssertionError(f"{label}: no KV chunks shipped")
expected = max(0, n_tokens - 2)
if done.total_tokens < expected - 64:
raise AssertionError(
f"{label}: got {done.total_tokens} tokens, expected ~{expected}"
)
print(
f" [{label}] OK: tokens={done.total_tokens} "
f"kv_chunks={len(chunks)} arrays={len(arrays)}"
)
return done.total_tokens
def main(argv: list[str]) -> int:
if len(argv) < 2:
print(__doc__)
return 2
model_id = ModelId(argv[1])
from exo.download.download_utils import build_model_path
model_path = build_model_path(model_id)
if not model_path.exists():
print(f"FAIL: model {model_id} not found at {model_path}")
return 1
print(f"Loading vLLM engine for {model_id} ({model_path}) ...")
engine = _build_engine(model_id)
failures: list[str] = []
try:
try:
t1 = _run_one(engine, n_tokens=512, label="run1-fresh")
except AssertionError as e:
failures.append(f"run1: {e}")
t1 = 0
try:
t2 = _run_one(engine, n_tokens=512, label="run2-same-prompt")
except AssertionError as e:
failures.append(f"run2: {e}")
t2 = 0
if t1 and t2 and t1 != t2:
failures.append(
f"run1 returned {t1} tokens but run2 returned {t2} (should match)"
)
try:
ta = _run_one(engine, n_tokens=256, label="run3-shorter")
tb = _run_one(engine, n_tokens=768, label="run4-longer")
if ta and tb and tb <= ta:
failures.append(
f"longer prompt should produce more tokens: 256->{ta} 768->{tb}"
)
except AssertionError as e:
failures.append(f"length-variation: {e}")
finally:
with contextlib.suppress(Exception):
engine.close() # pyright: ignore[reportAttributeAccessIssue]
if failures:
print()
print("FAIL")
for f in failures:
print(f" - {f}")
return 1
print()
print("PASS")
return 0
if __name__ == "__main__":
try:
sys.exit(main(sys.argv))
except Exception:
traceback.print_exc()
sys.exit(1)
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env bash
set -Eeuo pipefail
SELF_IP="169.254.100.1"
PEER_IP="169.254.100.2"
PREFIX="16"
IFACE="enP7s7"
USE_NM="auto"
DRY_RUN=0
usage() {
cat <<EOF
Usage: sudo $(basename "$0") [options]
Configure a Linux Ethernet interface with a static IPv4 for a host-to-host
link to a Mac peer.
Defaults: this host = ${SELF_IP}/${PREFIX}, peer = ${PEER_IP}, iface = ${IFACE}.
Options:
--iface IFACE Default: ${IFACE}
--self-ip IP Default: ${SELF_IP}
--peer-ip IP For verification ping. Default: ${PEER_IP}
--prefix N Default: ${PREFIX}
--no-nm Use 'ip addr' directly (transient, no NetworkManager).
--dry-run Print actions without applying.
-h, --help Show this help.
EOF
}
while (($#)); do
case "$1" in
--iface)
shift
IFACE="${1:?}"
;;
--self-ip)
shift
SELF_IP="${1:?}"
;;
--peer-ip)
shift
PEER_IP="${1:?}"
;;
--prefix)
shift
PREFIX="${1:?}"
;;
--no-nm) USE_NM=no ;;
--dry-run) DRY_RUN=1 ;;
-h | --help)
usage
exit 0
;;
*)
echo "Unknown arg: $1" >&2
usage >&2
exit 1
;;
esac
shift
done
[[ $EUID -eq 0 ]] || {
echo "Run as root." >&2
exit 1
}
run() {
printf '+'
printf ' %q' "$@"
printf '\n'
((DRY_RUN)) || "$@"
}
ip link show "$IFACE" >/dev/null 2>&1 || {
echo "Interface $IFACE does not exist." >&2
exit 1
}
if [[ $USE_NM == "auto" ]]; then
if command -v nmcli >/dev/null 2>&1 && systemctl is-active --quiet NetworkManager 2>/dev/null; then
USE_NM=yes
else
USE_NM=no
fi
fi
if [[ $USE_NM == "yes" ]]; then
CONN="$(nmcli -g GENERAL.CONNECTION device show "$IFACE" 2>/dev/null | head -n1 || true)"
if [[ -z $CONN || $CONN == "--" ]]; then
CONN="static-${IFACE}"
run nmcli connection add type ethernet ifname "$IFACE" con-name "$CONN"
fi
run nmcli connection modify "$CONN" \
connection.interface-name "$IFACE" \
connection.autoconnect yes \
connection.autoconnect-priority 100 \
ipv4.method manual \
ipv4.addresses "${SELF_IP}/${PREFIX}" \
ipv4.gateway "" \
ipv4.dns "" \
ipv4.never-default yes \
ipv6.method link-local \
ipv6.addr-gen-mode stable-privacy
run nmcli connection up "$CONN"
else
run ip link set "$IFACE" up
run ip addr flush dev "$IFACE"
run ip addr add "${SELF_IP}/${PREFIX}" dev "$IFACE"
fi
if ((!DRY_RUN)); then
printf '\n'
ip -br addr show "$IFACE"
printf '\n'
if ping -c2 -W2 "$PEER_IP" >/dev/null 2>&1; then
echo "OK: $PEER_IP reachable on $IFACE."
else
echo "WARN: $PEER_IP not reachable yet."
echo " Verify the peer is configured (run setup_linklocal_mac.sh on the Mac)."
echo " ip neigh show dev $IFACE # check for the peer MAC"
fi
fi
+170
View File
@@ -0,0 +1,170 @@
#!/usr/bin/env bash
set -Eeuo pipefail
SELF_IP="169.254.100.2"
PEER_IP="169.254.100.1"
NETMASK="255.255.0.0"
IFACE=""
DRY_RUN=0
usage() {
cat <<EOF
Usage: sudo $(basename "$0") [options]
Configure a Mac Ethernet interface with a static IPv4 for a host-to-host link
to the DGX/GX10 peer.
Defaults: this Mac = ${SELF_IP}, peer = ${PEER_IP}, mask = ${NETMASK}.
Options:
--iface IFACE Interface (e.g. en12). Default: auto-detect.
--self-ip IP This Mac's address. Default: ${SELF_IP}.
--peer-ip IP Peer for verification ping. Default: ${PEER_IP}.
--netmask MASK Default: ${NETMASK}.
--dry-run Print actions without applying.
-h, --help Show this help.
EOF
}
while (($#)); do
case "$1" in
--iface)
shift
IFACE="${1:?}"
;;
--self-ip)
shift
SELF_IP="${1:?}"
;;
--peer-ip)
shift
PEER_IP="${1:?}"
;;
--netmask)
shift
NETMASK="${1:?}"
;;
--dry-run) DRY_RUN=1 ;;
-h | --help)
usage
exit 0
;;
*)
echo "Unknown arg: $1" >&2
usage >&2
exit 1
;;
esac
shift
done
[[ $EUID -eq 0 ]] || {
echo "Run with sudo." >&2
exit 1
}
run() {
printf '+'
printf ' %q' "$@"
printf '\n'
((DRY_RUN)) || "$@"
}
target_subnet_prefix() {
local ip="$1"
printf '%s.' "${ip%.*}"
}
iface_score() {
local iface="$1" info subnet
info="$(ifconfig "$iface" 2>/dev/null || true)"
[[ -n $info ]] || {
echo 0
return
}
grep -q 'status: active' <<<"$info" || {
echo 0
return
}
subnet="$(target_subnet_prefix "$SELF_IP")"
if grep -qE "inet ${subnet//./\\.}" <<<"$info"; then
echo 100
return
fi
if grep -qE 'inet 169\.254\.' <<<"$info"; then
echo 80
return
fi
if ! grep -qE '^[[:space:]]*inet ' <<<"$info"; then
echo 60
return
fi
echo 10
}
detect_iface() {
local best="" best_score=0 iface score
for iface in $(ifconfig -l); do
[[ $iface =~ ^en[0-9]+$ ]] || continue
score="$(iface_score "$iface")"
if ((score > best_score)); then
best="$iface"
best_score="$score"
fi
done
((best_score >= 60)) || return 1
printf '%s\n' "$best"
}
iface_to_service() {
local iface="$1" line port=""
while IFS= read -r line; do
if [[ $line == "Hardware Port: "* ]]; then
port="${line#Hardware Port: }"
elif [[ $line == "Device: $iface" ]]; then
printf '%s\n' "$port"
return 0
fi
done < <(networksetup -listallhardwareports)
return 1
}
if [[ -z $IFACE ]]; then
IFACE="$(detect_iface || true)"
[[ -n $IFACE ]] || {
echo "Could not auto-detect a wired interface. Pass --iface enX." >&2
echo "Active interfaces:" >&2
ifconfig -l | tr ' ' '\n' | grep -E '^en[0-9]+$' | while read -r i; do
printf ' %-6s %s\n' "$i" "$(ifconfig "$i" | grep -E 'status:|inet ' | tr '\n' ' ')" >&2
done
exit 1
}
echo "Auto-detected interface: $IFACE"
fi
ifconfig "$IFACE" >/dev/null 2>&1 || {
echo "Interface $IFACE does not exist." >&2
exit 1
}
SERVICE="$(iface_to_service "$IFACE" || true)"
[[ -n $SERVICE ]] || {
echo "No network service maps to $IFACE. Check System Settings -> Network." >&2
exit 1
}
echo "Network service: $SERVICE"
run networksetup -setmanual "$SERVICE" "$SELF_IP" "$NETMASK" ""
if ((!DRY_RUN)); then
printf '\n'
ifconfig "$IFACE" | grep -E 'inet |status:'
printf '\n'
if ping -c2 -t3 "$PEER_IP" >/dev/null 2>&1; then
echo "OK: $PEER_IP reachable on $IFACE."
else
echo "WARN: $PEER_IP not reachable yet."
echo " Verify the peer is configured (run setup_linklocal_dgx.sh on the GX10)."
echo " arp -an -i $IFACE # check for the peer MAC"
fi
fi
+9 -3
View File
@@ -131,9 +131,13 @@ async def chat_request_to_text_generation(
multimodal_content.append({"type": "text", "text": part.text})
else:
multimodal_content.append({"type": "image"})
chat_template_messages.append(
{"role": msg.role, "content": multimodal_content}
)
multimodal_msg: dict[str, Any] = {
"role": msg.role,
"content": multimodal_content,
}
if msg.reasoning_content is not None:
multimodal_msg["reasoning_content"] = msg.reasoning_content
chat_template_messages.append(multimodal_msg)
continue
msg_copy = msg.model_copy(update={"content": content})
@@ -168,6 +172,8 @@ async def chat_request_to_text_generation(
min_p=request.min_p,
repetition_penalty=request.repetition_penalty,
repetition_context_size=request.repetition_context_size,
presence_penalty=request.presence_penalty,
frequency_penalty=request.frequency_penalty,
images=images,
)
+92 -4
View File
@@ -79,6 +79,8 @@ from exo.api.types import (
ImageListItem,
ImageListResponse,
ImageSize,
InstanceLinkBody,
InstanceLinkResponse,
ModelList,
ModelListModel,
PlaceInstanceParams,
@@ -122,6 +124,7 @@ from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
from exo.shared.constants import (
DASHBOARD_DIR,
ENABLE_DISAGGREGATION,
EXO_CACHE_HOME,
EXO_EVENT_LOG_DIR,
EXO_IMAGE_CACHE_DIR,
@@ -154,6 +157,7 @@ from exo.shared.types.commands import (
DeleteCustomModelCard,
DeleteDownload,
DeleteInstance,
DeleteInstanceLink,
DownloadCommand,
ForwarderCommand,
ForwarderDownloadCommand,
@@ -161,6 +165,7 @@ from exo.shared.types.commands import (
ImageGeneration,
PlaceInstance,
SendInputChunk,
SetInstanceLink,
StartDownload,
TaskCancelled,
TaskFinished,
@@ -174,6 +179,7 @@ from exo.shared.types.events import (
InstanceDeleted,
TracesMerged,
)
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
from exo.shared.types.memory import Memory
from exo.shared.types.state import State
from exo.shared.types.tasks import (
@@ -215,6 +221,17 @@ def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
return params
def _require_disaggregation_enabled() -> None:
if not ENABLE_DISAGGREGATION:
raise HTTPException(
status_code=HTTPStatus.NOT_FOUND,
detail=(
"Prefill/decode disaggregation is disabled. "
"Set ENABLE_DISAGGREGATION=true to enable."
),
)
class API:
def __init__(
self,
@@ -328,6 +345,11 @@ class API:
self.app.get("/instance/previews")(self.get_placement_previews)
self.app.get("/instance/{instance_id}")(self.get_instance)
self.app.delete("/instance/{instance_id}")(self.delete_instance)
self.app.get("/v1/instance-links")(self.list_instance_links)
self.app.post("/v1/instance-links")(self.create_instance_link)
self.app.put("/v1/instance-links/{link_id}")(self.update_instance_link)
self.app.delete("/v1/instance-links/{link_id}")(self.delete_instance_link)
self.app.get("/v1/feature-flags")(self.get_feature_flags)
self.app.get("/models")(self.get_models)
self.app.get("/v1/models")(self.get_models)
self.app.post("/models/add")(self.add_custom_model)
@@ -336,7 +358,9 @@ class API:
self.app.post("/v1/chat/completions", response_model=None)(
self.chat_completions
)
self.app.post("/bench/chat/completions")(self.bench_chat_completions)
self.app.post("/bench/chat/completions", response_model=None)(
self.bench_chat_completions
)
self.app.post("/v1/images/generations", response_model=None)(
self.image_generations
)
@@ -502,8 +526,8 @@ class API:
)
]
)
# TODO: PDD
# instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
if any(self.state.node_vllm.values()):
instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
for sharding, instance_meta, min_nodes in instance_combinations:
try:
@@ -615,6 +639,52 @@ class API:
instance_id=instance_id,
)
async def get_feature_flags(self) -> dict[str, bool]:
return {
"disaggregation": ENABLE_DISAGGREGATION,
"vllm_available": any(self.state.node_vllm.values()),
}
async def list_instance_links(self) -> list[InstanceLink]:
if not ENABLE_DISAGGREGATION:
return []
return list(self.state.instance_links.values())
async def create_instance_link(
self, body: InstanceLinkBody
) -> InstanceLinkResponse:
_require_disaggregation_enabled()
return await self._set_instance_link(InstanceLinkId(), body)
async def update_instance_link(
self, link_id: InstanceLinkId, body: InstanceLinkBody
) -> InstanceLinkResponse:
_require_disaggregation_enabled()
return await self._set_instance_link(link_id, body)
async def _set_instance_link(
self, link_id: InstanceLinkId, body: InstanceLinkBody
) -> InstanceLinkResponse:
command = SetInstanceLink(
link_id=link_id,
prefill_instances=list(body.prefill_instances),
decode_instances=list(body.decode_instances),
)
await self._send(command)
return InstanceLinkResponse(
message="Command received.", command_id=command.command_id
)
async def delete_instance_link(
self, link_id: InstanceLinkId
) -> InstanceLinkResponse:
_require_disaggregation_enabled()
command = DeleteInstanceLink(link_id=link_id)
await self._send(command)
return InstanceLinkResponse(
message="Command received.", command_id=command.command_id
)
async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
"""Cancel an active command by closing its stream and notifying workers."""
sender = self._text_generation_queues.get(
@@ -829,7 +899,7 @@ class API:
async def bench_chat_completions(
self, payload: BenchChatCompletionRequest
) -> BenchChatCompletionResponse:
) -> BenchChatCompletionResponse | StreamingResponse:
task_params = await chat_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
@@ -846,6 +916,22 @@ class API:
command = await self._send_text_generation_with_images(task_params)
if payload.stream:
return StreamingResponse(
with_sse_keepalive(
generate_chat_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
),
),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
return await self._collect_text_generation_with_stats(command.command_id)
async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
@@ -1663,7 +1749,9 @@ class API:
quantization=card.quantization,
base_model=card.base_model,
capabilities=card.capabilities,
reasoning_dialect=card.reasoning_dialect,
context_length=card.context_length,
requires_vllm=card.requires_vllm,
)
for card in cards
]
Loaded 100 of 184 files, more files were not shown because too many files have changed in this diff. Show more