Compare commits

..

119 Commits

Author SHA1 Message Date
rltakashige fd17a3bc09 vllm loads! 2026-03-27 12:02:36 +00:00
rltakashige 7ca9edaf42 wohohoh 2026-03-25 10:35:46 +00:00
Evan d1f13f21a1 time to fiddle 2026-03-24 11:54:10 +00:00
Evan a6cdc93d3c disable 2026-03-24 11:54:04 +00:00
Evan a9c7b1c68a urgk 2026-03-24 11:54:04 +00:00
Evan 6355f3d8fb man i gotta commit 2026-03-23 16:00:57 +00:00
Evan 8cd6191c70 gotta move the cache soon 2026-03-21 22:22:25 +00:00
Evan 532a8f0b07 we got it to BUILD 2026-03-20 15:38:16 +00:00
Evan 581a1fcd79 move chunks to core 2026-03-19 16:18:40 +00:00
Evan ac7194dd90 mupdate 2026-03-19 16:01:35 +00:00
Evan f538b44211 shmovin 2026-03-19 14:51:08 +00:00
Evan a3ce437fd4 first 2026-03-19 11:10:58 +00:00
Ryuichi Leo Takashige be731d3a85 Merge main 2026-03-17 19:05:55 +00:00
Ryuichi Leo Takashige 655185cfe7 Address comments 4 - defer to the warmup into the exo batch generator and vllm batch engine and don't store model on the generators. 2026-03-17 19:00:12 +00:00
Ryuichi Leo Takashige 1dd9c28842 Address comments 3, mainly refactors 2026-03-17 18:31:41 +00:00
Ryuichi Leo Takashige cacd26e63c Type error lol 2026-03-17 18:00:39 +00:00
Ryuichi Leo Takashige 6a3eb2f37d close() 2026-03-17 17:31:15 +00:00
Ryuichi Leo Takashige e1df77bc4c No more future annotations 2026-03-17 17:07:32 +00:00
ciaranbor a6519ba006 Update mflux to 0.16.9 (#1751)
Prevents malformed output from Qwen-Image
2026-03-17 16:58:23 +00:00
Ryuichi Leo Takashige e78e53df6e Address comments 2 including vllm capability in state 2026-03-17 16:55:08 +00:00
Ryuichi Leo Takashige c70d9006e8 Address comments including task id interface 2026-03-17 15:31:55 +00:00
Ryuichi Leo Takashige 72cd8552ae Merge branch 'main' into leo/dgx-spark-integrations 2026-03-17 13:58:49 +00:00
rltakashige b713889f73 Fix exo bench again again (#1750)
Mb premature auto merge
2026-03-17 13:47:24 +00:00
rltakashige 6ee673147d Fix exo bench prefill and decode tps (#1749) 2026-03-17 13:36:44 +00:00
ciaranbor ff4d20eed9 Fix image models through dashboard (#1746)
## Motivation

Image generation/editing from the dashboard was broken. ChatForm
bypassed the parent's model-launch logic, so image requests fell through
to text chat.

## Changes

- Moved image routing logic from ChatForm.svelte to +page.svelte
(routeMessage())
- ChatForm now always delegates to parent via onAutoSend (made required)
- Fixed missing updateActiveConversation() call on message retry

## Why It Works

All sends now go through the parent's launch-then-route path, so image
models get launched before the request is dispatched to the correct
endpoint.
2026-03-17 11:22:01 +00:00
Ryuichi Leo Takashige 8cd1308336 Tidy pass 1 2026-03-17 00:38:16 +00:00
Ryuichi Leo Takashige 04dcdbd127 Merge main 2026-03-16 23:01:48 +00:00
Ryuichi Leo Takashige ec5d62f935 Strip vllm generator 2026-03-16 22:11:35 +00:00
Ryuichi Leo Takashige e96f084051 Distributed callbacks 2026-03-16 21:17:14 +00:00
Ryuichi Leo Takashige dc68ddbac0 Prompt formatting 2026-03-16 20:37:38 +00:00
Ryuichi Leo Takashige 073f8c1690 add batching 2026-03-16 19:25:45 +00:00
Ryuichi Leo Takashige 3c29d0dd4c test prefix caching 2026-03-16 16:43:51 +00:00
rltakashige 594ed99734 new uv lock for fastsafetensors 2026-03-13 17:08:49 +00:00
Evan Quiney 7ed4639540 use structured concurrency in download coordinator (#1722)
identical to #1721 but a little safer imo. 


## manual testing
ran a few downloads, cancelled non-master, cancelled master -- no errors
reported.
2026-03-13 14:55:37 +00:00
Mazin Sharaf 29d4165fe2 Add step logo condition to FamilyLogos component (#1676)
## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->
This was to fix a small issue which was that the StepFun logo was not
included in the sidebar, and I also noticed it from an open issue:
- #1662 

## Changes

<!-- Describe what you changed in detail -->
I added a condition to the FamilyLogos where if the family is "step"
then the logo will be included in the sidebar.

## Test Plan
Open exo, go choose a model, and scroll down the sidebar until you see
the step logo, which should be there.

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
Check for the step logo in the sidebar.
2026-03-13 13:31:46 +00:00
Ryuichi Leo Takashige e9e23e556e Have loading progress 2026-03-13 12:58:34 +00:00
ecohash-co 12af7c9586 fix: shield macmon cleanup from cancellation to prevent orphaned process (#1714)
## Summary

Fixes a bug where macmon metrics (GPU usage, temperature, power, RAM)
freeze permanently in the dashboard while non-macmon metrics (disk
usage) continue updating normally.

**Root cause: asyncio subprocess pipe transport flow control stall**

Observed on a 2-node Mac Studio M3 Ultra cluster (Python 3.13.12, anyio
4.11.0, macOS with kqueue) running EXO.app for ~36 hours. Diagnostics
confirmed:

1. **Only macmon monitoring is stuck** — disk metrics from the same
InfoGatherer continue updating, proving the Worker, EventRouter, and API
pipelines are healthy
2. **macmon IS producing data** — its stdout pipe is full at exactly
65536 bytes (64KB OS buffer), and macmon is blocked on write at 0% CPU
3. **The pipe read-end FD is still open** — the exo process holds it,
but asyncio isn't reading from it
4. **The stale GPU value (0.21) is wrong** — should be ~1.0 (matching
the other node under identical load)

This is consistent with asyncio's subprocess pipe transport getting
stuck in flow control: `pause_reading()` is called when the internal
buffer exceeds the high-water mark, but `resume_reading()` is never
called, permanently deregistering the FD from kqueue. The `receive()`
coroutine waits forever for data asyncio will never deliver.

```
$ lsof -p <macmon_pid>
macmon  74691  FD=1  PIPE  SIZE=65536  ->0x5ae78ecf376ccd0e  # full pipe

$ lsof -p <exo_pid> | grep <pipe_id>
exo     74689  FD=47  PIPE  SIZE=65536  ->0xd129da474c1340f7  # read end held open, not consumed
```

Note: anyio 4.11's `Process.aclose()` already uses
`CancelScope(shield=True)` for cleanup during cancellation — this is NOT
an election cleanup issue (confirmed by @ciaranbor's testing).

**Fix:** Replace the `async for` iteration with an explicit `receive()`
inside `fail_after()`. If macmon produces no output for 10× its
configured interval (minimum 30s), `TimeoutError` fires, `open_process`
cleanup kills macmon and tears down the stale transport, and the loop
restarts with a fresh subprocess and fresh asyncio transport.

## Test plan

- [x] All 246 existing tests pass
- [ ] Verify macmon restarts after simulated pipe stall (e.g., `SIGSTOP`
macmon, wait for timeout, confirm restart and metrics resume)
- [ ] Long-running soak test on multi-node cluster to confirm the fix
prevents recurrence
2026-03-13 12:54:30 +00:00
ciaranbor f28b2fd037 Extract mlx revision from uv lock (#1715)
## Motivation

The MLX version and git revision in nix/mlx.nix were hardcoded and had
to be manually kept in sync with uv.lock

## Changes

- flake.nix: Extract MLX git rev from uv.lock's source.git URL and pass
as uvLockMlxRev
- nix/mlx.nix: Use uvLockMlxVersion and uvLockMlxRev instead of
hardcoded values; remove version mismatch assertion

## Why It Works

uv.lock is already the source of truth — now Nix reads both version and
rev from it directly. The pinned fetchFromGitHub hash still guards
against unexpected changes.
2026-03-13 12:34:54 +00:00
Ryuichi Leo Takashige 169ea2a5e8 Fix GPT OSS by not retokenizing prompts 2026-03-12 18:02:43 +00:00
Ryuichi Leo Takashige 4a7901c548 Fix patches 2026-03-12 17:42:02 +00:00
Ryuichi Leo Takashige 7bb5cb4fc7 Allow memory profiling to be unstable 2026-03-12 17:34:53 +00:00
Ryuichi Leo Takashige 493e342f83 Skip impossible shardings 2026-03-12 17:33:12 +00:00
Ryuichi Leo Takashige 283b1809c9 Move VLLM runner into VLLM engine 2026-03-12 17:20:30 +00:00
Ryuichi Leo Takashige 35030119e3 ExoBench and ExoEval for CUDA 2026-03-12 17:09:51 +00:00
rltakashige 585dfe3549 Merge branch 'main' into leo/dgx-spark-integrations 2026-03-12 15:41:29 +00:00
Ryuichi Leo Takashige 5d7a005a13 Destroy process group on Keyboard Interrupt 2026-03-12 15:33:32 +00:00
Ryuichi Leo Takashige 87b7c5ef8b Pass CI 2026-03-12 15:24:49 +00:00
Ryuichi Leo Takashige 957ebbd21f Set max token length as max context length if no max tokens set 2026-03-12 15:15:39 +00:00
Mustafa Alp Yılmaz ea18a62581 fix: guard against ZeroDivisionError in mlx_lm stats (#1707)
## Problem

Running short completions (like `max_tokens=1` health check probes) can
finish so fast that `mlx_lm`'s internal `generation_time` rounds to
zero. When that happens, `BatchGenerator.stats()` in
`mlx_lm/generate.py` divides `generation_tokens / generation_time` and
throws a `ZeroDivisionError`, which kills the runner process.

EXO already handles this on its side — lines 289-293 in
`batch_generate.py` guard the TPS calculation with `if
generation_elapsed > 0`. But the call to `self._exo_gen.stats()` on line
294 goes into mlx_lm's *separate* timing code, which doesn't have the
same guard. Two different timers, only one is protected.

In my case, this was triggered by health check probes (content: `"a"`,
`max_tokens=1`). The generation completed in sub-microsecond time,
`generation_time` was exactly `0`, and the runner crashed. Since the
health check command stays in the queue and retries after recovery, it
created an infinite crash loop — every ~15 seconds the runner would load
the model, get the same health check, and die again.

## Fix

Wrap `self._exo_gen.stats()` in a `try/except ZeroDivisionError`. If it
throws, set `mlx_stats` to `None` and fall back to `0.0` for
`prompt_tps`. The only fields EXO reads from `mlx_stats` are
`prompt_tps` and `prompt_time` — losing them on a sub-microsecond
generation has no practical impact.

## Traceback

```
File "batch_generate.py", line 294, in step
    mlx_stats = self._exo_gen.stats()
File "mlx_lm/generate.py", line 1224, in stats
    self._stats.generation_tokens / self._stats.generation_time
ZeroDivisionError: division by zero
```
2026-03-12 14:48:36 +00:00
rltakashige 4b6dd7588f lockgit status 2026-03-12 14:41:33 +00:00
Ryuichi Leo Takashige 3f4f7c9ba6 Only do for aarch64 linux 2026-03-12 13:42:23 +00:00
Ryuichi Leo Takashige 1331465ba0 Add missing runner features 2026-03-12 10:51:37 +00:00
Miguel Cruz 0782d90ec5 fix: show partial download progress on initial dashboard load (#1706)
## Summary

- Dashboard now shows partial download progress for models that were
partially downloaded in a previous session, instead of showing 0%
- Both `getModelDownloadStatus()` and `getInstanceDownloadStatus()` now
handle `DownloadPending` entries that carry non-zero
`downloaded`/`total` bytes

Fixes #1042

## Root cause

When exo restarts with partially downloaded models, the
`DownloadCoordinator` emits `DownloadPending` events (because
`downloaded_this_session` is 0, even though real bytes exist on disk).
The main page dashboard only checked for `DownloadOngoing` entries, so
these partially downloaded models showed as 0%.

The dedicated `/downloads` page already handled this correctly — it
renders `DownloadPending` entries with a progress bar when `downloaded >
0`. This fix brings the same behavior to the main page.

## Changes

For both `getModelDownloadStatus()` and `getInstanceDownloadStatus()` in
`+page.svelte`:
- Accept `DownloadPending` in addition to `DownloadOngoing`
- For `DownloadPending` entries with `downloaded > 0` or `total > 0`,
synthesize a `DownloadProgress` object from the top-level fields (with
`speed: 0` and `etaMs: 0` since no active download is in progress)
- Skip `DownloadPending` entries where both `downloaded` and `total` are
0 (truly pending, not yet started)

## Test plan

- [ ] Partially download a model, quit exo, relaunch — dashboard should
show partial progress instead of 0%
- [ ] Fully downloaded models still show as complete
- [ ] Active downloads still show real-time progress with speed/ETA
- [ ] Models never downloaded show as not started (not falsely showing
progress)
- [ ] Dashboard builds without errors (`cd dashboard && npm run build`)

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-12 10:39:34 +00:00
Ryuichi Leo Takashige 8f94727f14 Ignore missing modules if type stubs exist 2026-03-12 00:20:28 +00:00
Ryuichi Leo Takashige 9ee23ee0d3 Make vllm inference runner closer to the normal inference runner 2026-03-11 23:41:34 +00:00
Ryuichi Leo Takashige f75d36cbe0 Fix cache patch 2026-03-11 22:28:45 +00:00
Ryuichi Leo Takashige 2683ac7b61 Add Torch typings 2026-03-11 21:44:10 +00:00
Ryuichi Leo Takashige 404b9769ac Download models without model.safetensors.index 2026-03-11 21:32:07 +00:00
Ryuichi Leo Takashige 3e097f7243 only download a single copy of the model. 2026-03-11 20:57:29 +00:00
Ryuichi Leo Takashige 0c8615f25c Only import VLLM once.. 2026-03-11 20:49:19 +00:00
Ryuichi Leo Takashige ba35a4ba13 Move VLLM into the runner and add type stubs 2026-03-11 19:15:55 +00:00
Ryuichi Leo Takashige 9a83fa6cdf Patch VLLM to load multiple models dynamically 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 659c1bc737 Progress: Run EXO-CUDA through nix! 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 34df811b92 Vibe coding design baby 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige ca5870a2e8 Some Linux Laptop/Desktop detection and goodbye penguin 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 6be6ea5fd2 Fix placement preview 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 5e9d27b753 Show Sparks and Linux in topology 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige cfc8f09004 Fast direct USB connectivity 2026-03-11 18:17:59 +00:00
rltakashige f221a6c85c Normalise Responses API tool call format (#1704)
## Motivation

The responses API often does not provide tool args nested under a
"function" field. Since we follow the chat completions format of tools
in the backend (for MLX chat templates), we need to normalise to this
format.

## Test Plan

### Manual Testing
Works on n8n!
<img width="3442" height="2076" alt="image"
src="https://github.com/user-attachments/assets/9e11d679-0102-4d83-9a8e-b0a7a5898708"
/>
2026-03-11 18:10:12 +00:00
Mustafa Alp Yılmaz 2994b41089 fix: validate num_key_value_heads in tensor sharding placement (#1669)
## Problem

Models with fewer KV heads than nodes crash during tensor parallelism.
For example, Qwen3.5 MoE models have only 2 KV heads — trying to shard
across 4 nodes produces empty tensors and a reshape error at runtime.

The placement system already validates `hidden_size % num_nodes == 0`
but doesn't check KV heads, so it creates configurations that look valid
but blow up when the worker tries to split the attention heads.

Affected models include Qwen3.5-35B-A3B, Qwen3.5-122B-A10B,
Qwen3.5-397B-A17B, Qwen3-Next-80B-A3B, and Qwen3-Coder-Next (all have 2
KV heads).

## Changes

**Placement validation** (`src/exo/master/placement.py`):
- Combined KV heads divisibility check with the existing hidden_size
filter in a single pass
- Cycles where `num_key_value_heads % len(cycle) != 0` are now excluded
for tensor sharding
- Error message includes both constraints when no valid cycle is found

**Model card schema** (`src/exo/shared/models/model_cards.py`):
- Added optional `num_key_value_heads` field to `ModelCard` and
`ConfigData`
- Extracted from HuggingFace `config.json` (handles both top-level and
`text_config` nesting)
- Passed through in `fetch_from_hf()` for dynamically fetched cards

**All 68 inference model cards**
(`resources/inference_model_cards/*.toml`):
- Populated `num_key_value_heads` from each model's HuggingFace config

**Utility script** (`scripts/fetch_kv_heads.py`):
- Fetches `num_key_value_heads` from HuggingFace and updates TOML cards
- `--missing`: only fills in cards that don't have the field yet
- `--all`: re-fetches and overwrites everything
- Uses tomlkit for safe TOML editing and ThreadPoolExecutor for parallel
fetches

## Behavior

- Instance previews no longer show tensor options for models that can't
split their KV heads across the cluster size
- `place_instance()` rejects with a clear error instead of crash-looping
- Pipeline parallelism is unaffected
- 2-node tensor still works for 2-KV-head models (2 ÷ 2 = 1)
- Field is optional — existing custom cards without it continue to work
(validation is skipped when `None`)
2026-03-11 13:46:33 +00:00
Mustafa Alp Yılmaz 38f0c09175 fix: use StreamingDetokenizer in batch generator to fix emoji/UTF-8 corruption (#1691)
## Problem

Emojis and other multi-byte UTF-8 characters are rendered as `\ufffd`
(Unicode Replacement Character) in batch streaming responses.

Byte-level BPE tokenizers (like Qwen's) can split multi-byte UTF-8
characters (e.g. a 4-byte emoji) across multiple tokens. The batch
generator decodes each token independently with
`tokenizer.decode([token_id])`, which produces U+FFFD for partial byte
sequences.

The sequential generator doesn't have this problem — it uses
`stream_generate()` from mlx_lm, which internally uses
`StreamingDetokenizer` to buffer incomplete bytes.

## Fix

Use `StreamingDetokenizer` in the batch generator, matching the
sequential path:

- Each `_EngineTask` gets its own `StreamingDetokenizer` instance
- `add_token()` buffers tokens, holding back incomplete UTF-8 byte
sequences until they form valid characters
- `last_segment` returns only complete, valid text
- `finalize()` flushes any remaining buffered bytes when generation
completes

Empty text segments during buffering are harmless — they're already
handled correctly by the downstream streaming pipeline.

---------

Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
2026-03-10 16:10:22 +00:00
ciaranbor f36fd56c38 Include power usage in bench responses (#1692)
## Motivation

Add power/energy usage tracking to /bench API responses.

## Changes

- New PowerSampler class that periodically samples sys_power from each
node during inference
- New PowerUsage / NodePowerStats API types
- Integrated into both bench chat completion and bench image generation
endpoints

## Why It Works

Runs concurrently via anyio.create_task_group, reads from existing
state.node_system heartbeat data — no new plumbing needed. Energy =
avg_power × elapsed_time.

## Test Plan

### Manual Testing

Run a /bench request and verify power_usage appears in the response.

### Automated Testing

7 async tests covering single/multi-node, energy math, dynamic power
changes, empty state, and lifecycle.
2026-03-10 16:02:55 +00:00
rltakashige 82c54dd6d6 Add support for Nemotron sharding (#1693)
### Automated Testing
tested logits match
2026-03-10 15:51:07 +00:00
ciaranbor 3536161f15 Fix stale state.runners (#1684)
## Motivation

When a runner shuts down, there's no mechanism to remove it from
state.runners, causing stale state to accumulate.

## Changes

- apply.py: apply_runner_status_updated now handles RunnerShutdown
status by removing the runner from state instead of adding/updating it.
Removed the separate apply_runner_deleted function entirely.
- events.py: Removed the RunnerDeleted event type — runner cleanup is
now handled via RunnerStatusUpdated with a RunnerShutdown status.
- Tests: New test_apply_runner_deleted.py with 2 tests: verifying
RunnerShutdown removes the runner, and that a normal status update adds
it.

## Why It Works

Instead of a separate RunnerDeleted event and coordinating its emission
from master/placement, runner cleanup is handled naturally through the
existing RunnerStatusUpdated path — a RunnerShutdown status signals
removal. This is simpler and requires no changes to master or placement
logic.

## Test Plan

### Automated Testing

2 new unit tests covering RunnerShutdown removal and normal runner
status addition
2026-03-10 09:32:25 +00:00
ArvidSU a6aa07ed83 feat(dashboard): add mobile drawer support for sidebars (#1677)
## Summary

- Adds mobile-friendly drawer UI for the chat sidebar and model family
sidebar
- Introduces responsive header navigation with mobile-specific toggle
buttons
- Uses slide-in drawer overlay pattern on small screens instead of
collapsible sidebars
- Stores mobile drawer state in app store for consistent state
management

## Testing

Tested on desktop (responsive mode) and mobile viewport widths. Sidebars
now slide in as drawers on small screens with proper z-index layering
and close-on-outside-click behavior.
2026-03-10 04:18:11 +00:00
rltakashige 131ad0ff36 Implement continuous batching (#1642)
## Motivation

Following the changes made in #1632 !
Closes #1020

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->

---------

Co-authored-by: Evan Quiney <evanev7@gmail.com>
2026-03-09 15:04:45 +00:00
ciaranbor d01636100a Ciaran/gpt oss tool call finish reason (#1673)
## Motivation

When gpt-oss truncates a response mid-tool-call (e.g. hitting max
tokens), the finish_reason was swallowed because the parser was in
tool-accumulation mode and never yielded it back to the caller. This
caused the API to hang or fail to signal completion.

## Changes

In parse_gpt_oss, when accumulating tool call arguments and a
finish_reason arrives, yield the response (with empty text) so the
finish reason propagates

## Why It Works

The parser now checks for finish_reason on every chunk while inside a
tool call. When the model stops generating (truncation), the finish
signal is forwarded immediately rather than being silently consumed.

## Test Plan

### Automated Testing

Added tests covering truncated tool calls (finish_reason="length") and
plain text truncation
2026-03-06 21:09:20 +00:00
ciaranbor 79c8dbeacd Ciaran/minor download bugs (#1674)
## Motivation

Two minor bugs in the download coordinator: cancelling a download didn't
reset its status to pending, and a single error in
_emit_existing_download_progress would kill the loop permanently.

## Changes

- Cancel emits pending status: When a download is cancelled, emit a
DownloadPending event so the UI/cluster state reflects that the model is
no longer actively downloading.
- Resilient progress emission loop: Move the try/except inside the while
loop so transient errors are logged but the periodic emission continues.

## Why It Works

- The cancel fix ensures the download status is consistent — without it,
a cancelled download would still appear as in-progress.
- The loop fix prevents a single exception from permanently stopping the
60-second periodic progress broadcast.
2026-03-06 17:46:18 +00:00
Evan Quiney 0096159728 up gossipsub limit (#1671)
bump gossipsub limit to 8MB + add warning
this is a bandaid solution for large messages while we figure out the
point to point protocol
2026-03-06 14:20:51 +00:00
Andrei Onel 7a36d3968d docs: Update documentation for v1.0.68 release (#1667)
## Motivation

Updated documentation for v1.0.68 release

## Changes

**docs/api.md:**
- Added documentation for new API endpoints: Claude Messages API
(`/v1/messages`), OpenAI Responses API (`/v1/responses`), and Ollama API
compatibility endpoints
- Documented custom model management endpoints (`POST /models/add`,
`DELETE /models/custom/{model_id}`)
- Added `enable_thinking` parameter documentation for thinking-capable
models (DeepSeek V3.1, Qwen3, GLM-4.7)
- Documented usage statistics in responses (prompt_tokens,
completion_tokens, total_tokens)
- Added streaming event format documentation for all API types
- Updated image generation section with FLUX.1-Kontext-dev support and
new dimensions (1024x1365, 1365x1024)
- Added request cancellation documentation
- Updated complete endpoint summary with all new endpoints
- Added security notes about trust_remote_code being opt-in

**README.md:**
- Updated Features section to highlight multiple API compatibility
options
- Added Environment Variables section documenting all configuration
options (EXO_MODELS_PATH, EXO_OFFLINE, EXO_ENABLE_IMAGE_MODELS,
EXO_LIBP2P_NAMESPACE, EXO_FAST_SYNCH, EXO_TRACING_ENABLED)
- Expanded "Using the API" section with examples for Claude Messages
API, OpenAI Responses API, and Ollama API
- Added custom model loading documentation with security notes
- Updated file locations to include log files and custom model cards
paths

**CONTRIBUTING.md:** 
- Added documentation for TOML model cards format and the API adapter
pattern

**docs/architecture.md:**
- Documented the adapter architecture introduced in PR #1167

Closes #1653

---------

Co-authored-by: askmanu[bot] <192355599+askmanu[bot]@users.noreply.github.com>
Co-authored-by: Evan Quiney <evanev7@gmail.com>
2026-03-06 11:32:46 +00:00
Mustafa Alp Yılmaz eee3432738 feat(mlx): add repetition_penalty and repetition_context_size to chat completions (#1665)
## Problem

Models running through exo can get stuck in repetition loops —
generating the same text over and over until hitting the token limit. It
happens more with quantized models where probability distributions can
become degenerate.

`mlx_lm` has `make_logits_processors(repetition_penalty,
repetition_context_size)` that handles this, but it is never called
anywhere in the pipeline.

## Solution

Wire `repetition_penalty` and `repetition_context_size` through the
stack:

- `ChatCompletionRequest` → accept the params from the client
- `TextGenerationTaskParams` → carry them through the pipeline
- `chat_completions.py` adapter → map to internal params
- `generate.py` → call `make_logits_processors()` and merge into the
`logits_processors` list

When `repetition_penalty` is `None` (default), `make_logits_processors`
returns `[]` so existing behavior is unchanged.

## Usage

```json
{
  "model": "mlx-community/...",
  "messages": [...],
  "repetition_penalty": 1.15,
  "repetition_context_size": 64
}
```
2026-03-06 10:51:25 +00:00
Alex Cheema e8c3a873a6 feat(dashboard): improve model picker feedback and HuggingFace search (#1661)
## Motivation

Fixes #1648. Users adding custom models from HuggingFace got no clear
feedback that the model was successfully added — the model didn't appear
prominently in the list. Additionally, searches were limited to
`mlx-community` models with no fallback.

## Changes

- **Success feedback**: After adding a custom model, a toast
notification appears, the view auto-switches to "All Models", and the
newly added model is scrolled into view with a green highlight that
fades over 4 seconds.
- **HuggingFace search fallback**: The `/models/search` endpoint now
searches `mlx-community` first; if no results are found, it falls back
to searching all of HuggingFace.
- **Inline HF results**: When the main search bar finds no local
matches, HuggingFace search results appear inline with "+ Add" buttons
and a "See all results on Hub" link.
- **Full repo ID for non-mlx models**: Non-mlx-community models now
display the full repo ID (e.g., `meta-llama/Llama-3.1-8B`) instead of
just the short name.

## Why It Works

The toast + scroll + highlight gives immediate, unambiguous feedback
that the model was added and where it lives in the list. The HF search
fallback broadens discoverability while still prioritising mlx-community
models. Inline results in the main search bar mean users don't need to
navigate to the HF tab to discover new models.

## Test Plan

### Manual Testing
- Open model picker → HuggingFace Hub tab → search for a model → click
"+ Add"
- Verify: toast appears, view switches to All Models, model highlighted
with green glow
- Search for a term with no mlx-community results → verify fallback to
full HF search
- Non-mlx-community results show full repo ID (e.g., `org/model-name`)
- On All Models tab, search for a model not in the local list → verify
"From HuggingFace" section appears

### Automated Testing
- Existing tests cover the model catalog and API; no new automated tests
needed for UI behaviour changes

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-06 10:23:08 +00:00
Michael Harrigan afab3095b0 fix: KVPrefixCache Regression (#1668) 2026-03-06 10:11:08 +00:00
ciaranbor b9d40e8e35 Ciaran/re download bug (#1658)
## Motivation

After deleting a model and re-downloading it, the CachedShardDownloader
returns the stale cached path, so ensure_shard short-circuits and no
download actually happens.

## Changes

- Added invalidate(model_id) method to the ShardDownloader ABC and all
implementations
- CachedShardDownloader.invalidate evicts cache entries matching the
model ID and delegates down
- DownloadCoordinator.delete_model calls invalidate after cancelling
active downloads, before deleting files
- Added end-to-end test that downloads, deletes, and re-downloads a
model through the coordinator

## Why It Works

The cache is cleared when a model is deleted, so the next ensure_shard
call performs a fresh download instead of returning the stale path.

## Test Plan

## Automated Testing

New test_re_download_after_delete_completes exercises the full download
→ delete → re-download flow through DownloadCoordinator with
CachedShardDownloader + SingletonShardDownloader wrappers matching
production.
2026-03-05 14:18:17 +00:00
wysie 3a4d635d0c Fix copy code button not working in dashboard (#1659)
Fixes #1657

## Motivation

The copy button on code blocks in the dashboard does nothing when
clicked. This affects all code blocks in assistant responses.

## Root Cause

In `MarkdownContent.svelte`, click event listeners are bound to
`.copy-code-btn` elements via `setupCopyButtons()` inside a Svelte 5
`$effect`. However, the effect fires before the DOM has been updated
with the new HTML, so `querySelectorAll(".copy-code-btn")` finds zero
buttons.

Additionally, during streaming, the `content` prop updates on every
token, causing the entire `{@html processedHtml}` to be re-rendered.
This destroys all previously bound event listeners, even if they were
successfully attached.

## Changes

Replaced the per-button `addEventListener` approach with **event
delegation** — a single click listener on the container element that
catches clicks bubbling up from any `.copy-code-btn` or
`.copy-math-btn`. This:

- Eliminates the timing issue (the listener exists before the buttons
are rendered)
- Survives HTML re-renders during streaming (no need to re-bind)
- Removes the need for `setupCopyButtons()` and the `data-listenerBound`
tracking

## Testing

1. Load any model
2. Prompt it to generate a code block (e.g. "write a hello world in
Python")
3. Click the copy button on the code block
4. Paste — the code is copied correctly
5. Verified the button also works during streaming (before generation
completes)

Co-authored-by: Wysie <wysie@users.noreply.github.com>
2026-03-05 12:41:17 +00:00
Miguel Miranda Dias 8485805042 fix(worker): emit error chunks when a runner dies mid-command (#1645)
Closes #1586

## Summary
- track in-flight tasks in `RunnerSupervisor` (not only unacknowledged
pending tasks)
- when `_check_runner()` detects a crashed runner, emit
`ChunkGenerated(ErrorChunk)` for each in-flight command task
(`TextGeneration`, `ImageGeneration`, `ImageEdits`)
- keep existing `RunnerStatusUpdated(RunnerFailed)` emission so
planner/state still transition correctly
- add a unit test for supervisor crash path to ensure an error chunk is
emitted before failed runner status

## Why
`#1586` reports streams that can hang forever when runners crash during
warmup/loading. This keeps failure signaling at the runner-supervisor
layer, matching maintainer guidance in the issue thread.

## Validation
- attempted: `uv run pytest
src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py`
- blocked locally by environment disk exhaustion while uv tried to
materialize heavy CUDA wheels (`No space left on device` during
`nvidia-cudnn-cu13` extraction)

I kept the change scoped and added a targeted unit test for the failure
path.

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-04 17:15:58 +00:00
ciaranbor 4de8f801c7 #Add reasoning parms to chat completion and responses APIs (#1654)
## Motivation

Adds reasoning_effort parameter support to both the Chat Completions and
Responses APIs, aligning with the OpenAI spec and enabling thinking
control for gpt-oss models

## Changes

- Added ReasoningEffort literal type ("none" | "minimal" | "low" |
"medium" | "high" | "xhigh") and a resolve_reasoning_params() helper
that cross-derives reasoning_effort ↔ enable_thinking when only one is
provided
- Added reasoning_effort field to ChatCompletionRequest and reasoning
(with Reasoning model) + enable_thinking to ResponsesRequest
- Both adapters now call resolve_reasoning_params() before building
TextGenerationTaskParams
- reasoning_effort is passed through to the MLX chat template as a
template variable

## Why It Works

resolve_reasoning_params is a pure function that normalises the two
overlapping knobs (reasoning_effort and enable_thinking) into a
consistent pair, so downstream code always has both values regardless of
which the caller supplied.

## Test Plan

### Automated Testing

Added test_resolve_reasoning_params.py with 10 test cases covering:
both-None, both-set passthrough, enable_thinking → effort derivation,
and effort → enable_thinking derivation for every ReasoningEffort
variant.
2026-03-04 14:27:49 +00:00
Owleksiy 5777bf3c39 fix: coerce tool-call argument types from tool schema (#1651)
Apply schema-aware coercion to parsed tool-call arguments so
Hermes-style toolcalls can still return typed JSON (e.g. integer ids).

 - pass request tools into parse_tool_calls
 - coerce parsed argument values by function parameters schema
 - add unit tests for coercion and unknown-tool passthrough

## Motivation

Models that use Hermes-based toolcall syntax (Qwen3.5) can't reliably
call tools with non-string parameters
Example tool:
```json
    {
      "type": "function",
      "function": {
        "name": "process",
        "description": "Manage background processes",
        "parameters": {
          "type": "object",
          "properties": {
            "action": {
              "type": "string",
              "enum": ["spawn", "output", "kill", "list"]
            },
            "id": {
              "type": "integer",
              "description": "Process id"
            },
            "command": {
              "type": "string",
              "description": "Command to run for spawn"
            }
          },
          "required": ["action"],
          "additionalProperties": false
        }
      }
    }
```
Model transcript:
```
<tool_call>
<function=process>
<parameter=action>
output
</parameter>
<parameter=id>
0
</parameter>
</function>
</tool_call>
```
And the API returns:

`{"id":"a8f11689-d840-4ca5-ab1d-ead3678a11a9","name":"process","arguments":"{\"action\":
\"output\", \"id\": \"0\"}"}}`

Tool definition declared `id` as `integer`, the model output is
type-agnostic, and the translation layer treats everything as a string.

The same Qwen3.5-27B on OpenRouter and GPT-4.1-mini on openai obey the
function signature and emit correct call:
```
{"name":"process","arguments":"{\"action\": \"output\", \"id\": 0}"}}
```

Steps to reproduce:
```
❯ curl -sS -v http://localhost:52415/v1/chat/completions \
  -H 'Content-Type: application/json' \  -d @- <<'JSON'
{
  "model": "mlx-community/Qwen3.5-27B-4bit",
  "stream": false,
  "temperature": 0,
  "messages": [
    {
      "role": "user",
      "content": "Call the process tool with action=output and id=0. Do not explain anything. Just make the tool call."
    }
  ],
  "tools": [
    {
      "type": "function",
      "function": {
        "name": "process",
        "description": "Manage background processes",
        "parameters": {
          "type": "object",
          "properties": {
            "action": {
              "type": "string",
              "enum": ["spawn", "output", "kill", "list"]
            },
            "id": {
              "type": "integer",
              "description": "Process id"
            },
            "command": {
              "type": "string",
              "description": "Command to run for spawn"
            }
          },
          "required": ["action"],
          "additionalProperties": false
        }
      }
    }
  ],
  "tool_choice": {
    "type": "function",
    "function": {
      "name": "process"
    }
  }
}
JSON
```
Look for type of `id` function call

## Changes

Function call parameters are now converted to the types that the
function declaration has

## Why It Works

We now explicitly convert types where we know it (and skip if we don't)

## Test Plan

### Manual Testing
Create Qwen3.5-<ANY> instance, send the curl command above. Check that
'id' is now serialized as a number

### Automated Testing
Unit tests to cover basic type conversions

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-04 12:22:53 +00:00
Evan Quiney 886192f1e6 ignore closed resource errors when trying to cancel a task (#1652)
fixes an occasional crash during the shutdown of a failed instance.
2026-03-03 16:48:43 +00:00
Evan Quiney d914acd64e check if we have a task before we delete it (#1634)
caused a crash we should instead be logging
2026-03-03 15:32:12 +00:00
rltakashige 37296c8249 Refactor runner for implementing batching (#1632)
## Motivation

Batching will require us to send tasks concurrently and queue them up.
Our current infrastructure cannot handle that all. This PR gets us
closer to this by allowing multiple tasks to be sent in parallel and
then queuing up tasks.

## Changes

Change Plan logic
Make runner main into a class
Add a "BatchGenerator" to which tasks can be submitted (although tasks
are handled sequentially) and sent back through an MpSender.
Refactor runner to accept tasks during generation
Keep the generator threading
Separate the runner into several files for better readability

## Test Plan

### Manual Testing
Tested manually, needs a lot more automated testing. Cancellation still
works on a single device. Needs checking on multiple devices.

### Automated Testing

---------

Co-authored-by: Evan Quiney <evanev7@gmail.com>
2026-03-03 14:38:55 +00:00
Daiz 28817d3ee3 Add support for Qwen3.5 (#1644)
## Motivation

Qwen3.5 MoE models (e.g., `Qwen3.5-397B-A17B-6bit`) are now supported by
`mlx-lm` via `qwen3_5_moe` model type, but exo lacks tensor parallel
sharding support for this architecture. This prevents running large
Qwen3.5 models across multiple nodes.

Qwen3.5 uses a GatedDeltaNet hybrid attention mechanism similar to
Qwen3-Next, but with a different projection layout — separate
`in_proj_qkv`, `in_proj_z`, `in_proj_b`, `in_proj_a` instead of
Qwen3-Next's combined `in_proj_qkvz` and `in_proj_ba`. This requires
architecture-aware sharding logic.

## Changes (evan summary)

- enable qwen3_5 dense + moe tensor parallelism from config
- defensively skip evalling _cache.keys if it doesn't exist
- ignore kwargs in qwen35 pipeline masking and ensure pipeline segments match global model parameters for mask creation
- add sharding for qwen3_5 moe linear attention
- added another 6 million model cards

## Why It Works

Qwen3.5's GatedDeltaNet has an `in_proj_qkv` linear layer with three
concatenated sections: `[q(key_dim), k(key_dim), v(value_dim)]`. A naive
contiguous split (`segments=1`) would slice across section boundaries,
corrupting q/k/v values and producing garbled output.

By passing `segments=[key_dim, key_dim + key_dim]` to `shard_linear()`,
each section is split independently before distributing across devices.
This ensures every rank receives correctly aligned q, k, and v
components.

The remaining separate projections (`in_proj_z`, `in_proj_b`,
`in_proj_a`) and the MoE layers follow the same `all_to_sharded` /
`sharded_to_all` pattern already used for Qwen3-Next.

Some pipeline splits didn't include an ssm layer or a linear layer resulting in a subset of the model acting like it shouldn't create the appropriate masks for the next layer - we patch the model to manually create such masks.

## Test Plan

tensor sharded 2,3,4 models & pipeline sharded 2,3,4 with simple eval.

---------

Co-authored-by: hw <hw@hwStudio1.local>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
Co-authored-by: Evan <evanev7@gmail.com>
2026-03-03 14:31:57 +00:00
Alex Cheema 0e1b9501c3 fix: mini topology sidebar navigates home on click (#1616)
## Motivation

Clicking on the mini topology view in the right sidebar during chat was
selecting/filtering nodes instead of navigating back to the home view.
The entire mini topology section is wrapped in a button that calls
`handleGoHome()`, but individual node clicks were intercepting the event
via `stopPropagation()`.

## Changes

- Removed `onNodeClick={togglePreviewNodeFilter}` from the mini
sidebar's `TopologyGraph` component
- Added `pointer-events-none` to the topology graph container so all
clicks pass through to the parent button

## Why It Works

`TopologyGraph` only registers click/hover handlers when `onNodeClick`
is provided. Removing it disables node interaction entirely in the mini
view. The `pointer-events-none` class ensures no SVG element can
intercept clicks — they all bubble up to the parent `<button
onclick={handleGoHome}>`, which resets the chat state and returns to the
main topology view.

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Start exo and open the dashboard
- Start a chat to trigger the mini topology sidebar
- Click on the mini topology view → should navigate back to home
- Confirm nodes in the mini topology are not selectable or hoverable

### Automated Testing
- Dashboard builds successfully (`npm run build`)
2026-03-03 10:46:52 +00:00
Mustafa Alp Yılmaz f0d4ccbeb3 feat: add POST /v1/cancel/{command_id} endpoint (#1579)
## Summary

- Adds explicit `POST /v1/cancel/{command_id}` REST endpoint to cancel
in-flight text and image generation commands by ID
- Previously, cancellation only worked via HTTP disconnect (client
closes SSE connection, triggering `anyio.get_cancelled_exc_class()`).
Non-streaming clients and external consumers had no way to cleanly
cancel active generation
- The endpoint looks up the command in active generation queues, sends
`TaskCancelled` to notify workers, and closes the sender stream. Returns
404 (OpenAI error format) if the command is not found or already
completed

## Changes

| File | Change |
|------|--------|
| `src/exo/shared/types/api.py` | Add `CancelCommandResponse` model |
| `src/exo/master/api.py` | Import, route registration, `cancel_command`
handler |
| `src/exo/master/tests/test_cancel_command.py` | 3 test cases (404,
text cancel, image cancel) |
| `docs/api.md` | Document new endpoint + update summary table |

## Design Decisions

- Uses `self._send()` (not raw `command_sender.send()`) — respects API
pause state during elections
- Uses `raise HTTPException` — feeds into exo's centralized OpenAI-style
error handler
- Returns typed `CancelCommandResponse` — consistent with
`CreateInstanceResponse` / `DeleteInstanceResponse` patterns
- `sender.close()` is idempotent so concurrent cancel requests for the
same command are harmless

## Test Plan

- [x] `test_cancel_nonexistent_command_returns_404` — verifies 404 with
OpenAI error format
- [x] `test_cancel_active_text_generation` — verifies 200,
`sender.close()` called, `TaskCancelled` sent with correct
`cancelled_command_id`
- [x] `test_cancel_active_image_generation` — same verification for
image queue
- [x] `basedpyright` — 0 new errors
- [x] `ruff check` — all checks passed
- [x] `pytest` — 3/3 new tests pass, no regressions

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-03 10:38:52 +00:00
wysie 858dc808df fix: replace Master event_sender after EventRouter recreation (#1630) (#1637)
## Motivation
Fix for https://github.com/exo-explore/exo/issues/1630.
When loading a model that spans 2 nodes as the first action after
startup, Exo crashes with anyio.BrokenResourceError in
Master._command_processor. The receive end of the Master's event_sender
channel is already closed (open_receive_channels=0 ) when the Master
tries to send an InstanceCreated event.

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## Changes
In _elect_loop, when is_new_master is True and the EventRouter is
replaced, the Master's event_sender is now swapped in-place with a fresh
sender from the new EventRouter. This mirrors the existing pattern where
the Worker and DownloadCoordinator are already recreated with new
channels from the new router.
<!-- Describe what you changed in detail -->

## Why It Works
When is_new_master fires (which happens on every startup when a second
node joins), the old EventRouter is shut down and a new one is created.
Shutting down the old router cancels its _ingest coroutine, which closes
the receive end of every channel created by event_router.sender(). The
Worker and DownloadCoordinator are already recreated with senders from
the new router, but the Master was not — its event_sender still pointed
to the dead channel from the old router.
This patch replaces self.master.event_sender with a new sender from the
new EventRouter, so events flow through a live channel. We swap in-place
rather than recreating the Master to preserve its DiskEventLog, internal
state, and running tasks.
The reason loading a single-node model first appeared to work around the
issue is a timing race: if the model loads fast enough, the Master sends
its events before the election fires and kills the old channel. With a
2-node model, placement and RDMA setup take longer, so the election
completes first and the channel is already dead by the time the Master
tries to send.
<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
2 x Mac Studio M4 Max 128GB, RDMA
<!-- What you did: -->
Loaded a model that requires 2 nodes and it works properly now.
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->

Co-authored-by: Wysie <wysie@users.noreply.github.com>
2026-03-02 10:05:47 +00:00
ciaranbor 635118ef24 Support trace deletion in dashboard (#1628)
## Motivation

Enable using the dashboard to delete traces

## Changes

- Backend (src/exo/master/api.py): Added POST /v1/traces/delete endpoint
that deletes trace files by task ID, with path traversal protection
- API types (src/exo/shared/types/api.py): Added DeleteTracesRequest and
DeleteTracesResponse models
  - Dashboard store (app.svelte.ts): Added deleteTraces() method
- Traces page (+page.svelte): Added multi-select UI (click to select,
select all/deselect all) with a bulk delete button and confirmation
dialog

## Why It Works

- Uses existing _get_trace_path helper (now hardened against path
traversal) to resolve and delete trace files
- Dashboard selection state is reactive via Svelte 5 runes; deleted
traces refresh the list automatically

## Test Plan

### Manual Testing

- Select individual traces, select all, delete with confirmation dialog
  - Verify traces are removed from disk and list refreshes
2026-02-27 17:29:57 +00:00
Jake Hillion dc0bb5e13b fmt: add taplo TOML formatter to treefmt configuration 2026-02-27 17:19:48 +00:00
rltakashige 152a27ea5d Fix pipeline mismatched send after 1587 (#1629)
## Motivation

Tests caught a bug. It was a real bug.
2026-02-26 16:48:34 +00:00
rltakashige db36bd5ac6 Add custom prefill for pipeline (#1587)
## Motivation

Since we need to do distributed communications between prefill step
sizes, the out-of-the-box stream_generate that we currently use prevents
pipeline parallel models from doing overlapped computation. While this
was technically a regression, this communication is necessary for
cancellation, and we will need various distributed communications in the
future (e.g. for coordinating batching).

500 lines are for one testing file, so the diffs aren't as bad as they
look!

## Changes

Added a special prefill function for pipeline parallel models
Edited the model to handle 
Added a test to verify this new prefill and the original prefill produce
identical results
Improved type stubs to remove some type: ignores 

## Why It Works
<img width="768" height="1246" alt="image"
src="https://github.com/user-attachments/assets/8986ff17-ac23-4a02-9bd7-e6253a0ca799"
/>

## Test Plan

### Manual Testing
Needs more testing, but seems good so far.

### Automated Testing
Passes CI, considerable speedup seen in benchmarks (up to 1.98x) on
prefill speed.

Before:
<img width="3280" height="1238" alt="image"
src="https://github.com/user-attachments/assets/9abc1cbc-ecdb-4e48-a675-2c4cb04a32a0"
/>


After:
<img width="3344" height="1236" alt="image"
src="https://github.com/user-attachments/assets/e03c7987-41b4-4950-9ac3-2840e774ce30"
/>
2026-02-26 16:00:38 +00:00
Evan Quiney 639243aa09 event router (#1572)
replace the nack & resend logic in the worker/download coordinator with
a dedicated subsystem in front of the topic router. this centralizes
that logic (and the concept of system ids) to reduce replication in the
codebase. each system reading or writing events now gets a clean stream
of events in and can trust written events will be retried until
acknowledged.
2026-02-26 14:17:02 +00:00
Evan Quiney db73c4fd5d move messaging into rust (#1549)
the main body of the rust refactor. fixes the tokio panic on shutdown.
simplifies the networking module significantly. doesn't touch lp2p
behaviour
2026-02-26 13:58:22 +00:00
ciaranbor eaed92952c Use tmpdir for coordination file (#1624)
## Motivation

Coordination files for MLX distributed init were written to the current
working directory (./hosts_*.json)

## Changes

- Move coordination file creation to a tempfile.TemporaryDirectory(),
which auto-cleans on context manager exit
2026-02-26 10:59:36 +00:00
Evan Quiney ba611f9cd0 Revert "report macmon failures more aggressively (#1618)" (#1625)
this pr broke macmon - revert it
2026-02-25 19:15:55 +00:00
Evan Quiney eab3e0b456 report macmon failures more aggressively (#1618)
macmon appears to be going silent. time it out and restart it after
3*interval (Default 3s) and report warning.

Co-authored-by: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
2026-02-25 17:49:10 +00:00
Evan Quiney c4e874e97d skip nan logprobs on tokens (#1622)
sometimes we generate NaN logprobs for tokens, this causes pydantic
validation errors on the receiving end. in this case we should just not
send the logprob items
2026-02-25 17:44:15 +00:00
rltakashige e23c3a3026 Address Mac Mini pipeline GPU timeouts (#1620)
## Motivation
Users were reporting GPU timeout errors on Mac Minis, which we never saw
on testing with Mac Studios. It also seems to only happen with large
models.

## Changes
Eval specific distributed operations.

## Why It Works

As I wrote in a Slack message:
Basically, prefill is too slow for pipeline communications. If there are
both communications and GPU operations as part of an mlx graph, the
communications become subject to the GPU's 5 second command buffer
timeout.

For normal generation, I added evals to the communications (only during
prefill, as it slows down decode) to do this, fixing GPU timeouts.

But we don't do this during warmup, as the prompt is absolutely tiny.
This is still too slow on an M4 Pro on some models that it causes a GPU
timeout during warmup...


----------------------
This was one of the issues. However, there is another issue:

mx.all_gather sometimes reads stale data with FAST_SYNCH enabled. I'm
still investigating the root cause, but the code as it is now works on
Mac Minis.



## Test Plan

### Manual Testing
<img width="2762" height="1808" alt="image"
src="https://github.com/user-attachments/assets/27c88542-606c-4551-8f7c-bd2c0471f54e"
/>

<img width="2820" height="1898" alt="image"
src="https://github.com/user-attachments/assets/0ba3478c-ee39-438d-902c-92893db23d05"
/>


### Automated Testing
Needs a bunch on mac minis
2026-02-25 17:37:32 +00:00
Alex Cheema 190e63e56d fix: log exceptions causing silent node shutdown (#1621)
## Motivation

Nodes silently shut down mid-inference with no exception logged. The
logs show a clean-looking shutdown cascade (unsubscribe all topics →
stop worker → runner communication closed) but no error explaining
*why*. This makes debugging cluster issues extremely difficult.

## Changes

**`src/exo/routing/router.py`** — Added `try/except Exception` around
`_networking_recv()` and `_networking_recv_connection_messages()` loops.
Logs the root cause at ERROR level via
`logger.opt(exception=...).error(...)` before re-raising. Uses
`Exception` (not `BaseException`) so clean SIGTERM cancellation is
unaffected.

**`src/exo/main.py`** — Wrapped `anyio.run(node.run)` in
`try/except/finally`:
- `except BaseException`: logs the fatal exception at CRITICAL level
through loguru
- `finally`: ensures `logger.info("EXO Shutdown complete")` and
`logger_cleanup()` always run, guaranteeing the async log queue is
flushed before process exit

## Why It Works

The Router's recv loops (`_networking_recv`,
`_networking_recv_connection_messages`) are infinite `while True` loops
calling Rust bindings with **no exception handling**. When the Rust
networking layer raises `ConnectionError` (e.g. channel closed), the
exception cascades silently through anyio task groups: Router → Node →
all components shut down. The exception was never logged because (1) no
try-except anywhere in the cascade, and (2) `logger_cleanup()` was
skipped when `anyio.run()` raised, so loguru's async queue was never
flushed.

## Test Plan

### Manual Testing
- Run `uv run exo`, send SIGTERM → should see clean shutdown with "EXO
Shutdown complete", no error/critical logs
- Next time the silent shutdown reproduces, logs should now show the
actual exception with full traceback at ERROR level from the recv loop,
plus CRITICAL level from main()

### Automated Testing
- All 222 existing tests pass (1 pre-existing failure in
`test_python.py::test_sleep_on_multiple_items` unrelated to this change)
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — all checks passed
- `nix fmt` — applied

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-25 08:24:18 -08:00
Alex Cheema 7660893538 fix: replace flaky internet checks with explicit offline mode (#1615)
## Motivation

The socket-based internet connectivity probes (connecting to
1.1.1.1/8.8.8.8/1.0.0.1 every 10 seconds) are flaky — they produce false
negatives on some networks/ISPs, causing downloads to silently fail or
get stuck. Instead of dynamically detecting connectivity, this replaces
it with an explicit offline mode that the user opts into.

## Changes

- **Removed internet check infrastructure**: Deleted
`_test_internet_connection()` (socket probes) and
`_check_internet_connection()` (10s polling loop) from
`DownloadCoordinator`
- **Renamed `internet_connection` → `offline`** on `ShardDownloader`
base class and all subclasses (`SingletonShardDownloader`,
`CachedShardDownloader`, `ResumableShardDownloader`)
- **Removed `on_connection_lost` callbacks** from download calls — no
longer needed without dynamic connectivity detection
- **Added `EXO_OFFLINE` env var support** in `constants.py` and `Args`
default in `main.py`
- **Added Offline Mode toggle** to macOS app Settings → General tab,
which sets `EXO_OFFLINE=true` in the process environment and restarts

## Why It Works

The download behavior (`skip_internet` conditionals in
`download_utils.py`) is unchanged — we just changed what drives it.
Instead of a flaky socket probe setting
`internet_connection=True/False`, the user explicitly sets offline mode
via:
1. `--offline` CLI flag
2. `EXO_OFFLINE=true` environment variable
3. macOS app Settings → General → Offline Mode toggle

This is more reliable and predictable. Users on flaky networks no longer
get intermittent download failures.

## Test Plan

### Manual Testing
- `uv run exo --offline` starts without internet checks, rejects
downloads for unavailable models
- `EXO_OFFLINE=true uv run exo` behaves identically
- macOS app Settings toggle persists across restarts and passes env var
to the exo process

### Automated Testing
- All existing tests pass (222 passed), including 8 offline-mode
specific tests
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — all checks passed
- `nix fmt` — 0 files changed

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 13:36:19 -08:00
Evan Quiney 9a2d2a4a7c bump (#1608)
should be release for 1.0.68 - let's synch our py version with our app
version - 0.3.68.
next minor release should be 1.4.0 and 0.4.0 respectively.

Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-24 19:14:15 +00:00
Alex Cheema bea64fe85e fix: prevent stale loading state and conversation loss when switching chats (#1613)
## Motivation

When navigating back to a previous conversation that used a different
model than the one currently loading, the UI incorrectly displayed the
other model's loading/download progress bar instead of the conversation
messages. Additionally, attempting to continue that old conversation by
sending a message would create an entirely new chat rather than
continuing the existing one.

## Changes

All changes in `dashboard/src/routes/+page.svelte`:

1. **Added fallthrough reset in the `chatLaunchState` restore
`$effect`**: When switching to a conversation whose model has no active
instance (not running, not downloading, not loading), the effect now
resets `chatLaunchState` to `"idle"` instead of leaving stale state from
a different model.

2. **Added `skipCreate` parameter to `launchModelForChat()`**: When
continuing an existing conversation (has messages),
`createConversation()` is skipped so the old conversation is preserved
rather than replaced with a new empty one.

3. **Reordered view conditional**: Progress views
(downloading/loading/launching) take priority, then conversation
messages (when idle with messages or model ready), then model selector
(idle with no messages). This ensures:
   - Old conversations display normally when their model isn't running
- Download progress shows when relaunching a model for any conversation
   - Model selector only appears when there are no messages

## Why It Works

- The `$effect` fallthrough prevents `chatLaunchState` from retaining a
stale value (e.g., `"downloading"`) from Model A when the user switches
to a conversation using Model B that has no active state.
- The `skipCreate` flag ensures `launchModelForChat` can be reused for
both new conversations (from model picker) and continuing existing ones
(from chat input) without always creating a new conversation.
- The reordered view logic ensures each state maps to the correct UI:
active launch → progress, existing messages → chat view, nothing → model
selector.

## Test Plan

### Manual Testing
- Load an old conversation while a different model is downloading →
should see conversation messages, not the other model's loading bar
- Send a message from that old conversation → model should
launch/download with progress shown → conversation continues with the
response appended
- Select a new model from the picker → should still create a new
conversation as before
- New conversation with model download → progress bar shows correctly

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 19:08:31 +00:00
rltakashige 14526d281a update mlx 2 (#1611)
## Motivation

GPU locks because of prompt progress callbacks taking time. Current
solution: Don't fix it, make the symptom better

## Changes
Shortened timeout by 2x
Get event leak fixes from latest upstream
2026-02-24 18:30:48 +00:00
Evan Quiney 73e50df827 fix glm5 tool calling (#1612)
glm5 is a deepseekv32 model, so was parsing dsml style tool calls
instead of glm style tool calls. fix it!!
2026-02-24 18:21:18 +00:00
Alex Cheema 2b417f28ff fix: sync model selectors between sidebar and chat input (#1610)
## Motivation

The "Load Model" dropdown in the right sidebar and the model selector
above the chat input were operating on independent state
(`selectedPreviewModelId` vs `selectedChatModel`). This caused several
UX issues:
- Selecting a model in one selector didn't update the other
- After going home from a chat, the chat selector showed "SELECT MODEL"
while the sidebar still showed the model
- On page refresh, only the sidebar restored the saved model
- Sending a chat message with a running model briefly showed the
recommended models view instead of starting the chat

## Changes

- **`handleModelPickerSelect`**: Also sets `selectedChatModel` when a
model is picked from the sidebar dropdown
- **`handleChatPickerSelect`**: Also sets `selectedPreviewModelId` when
a model is picked from the chat selector
- **`handleGoHome`**: Restores `selectedChatModel` from the sidebar's
`selectedModelId` instead of clearing it
- **`applyLaunchDefaults`**: Syncs `selectedChatModel` when restoring
saved defaults on page load
- **`handleChatSend`**: Sets `chatLaunchState = "ready"` before calling
`sendMessage` when the model is already running, ensuring the chat view
renders correctly on view transition

## Why It Works

The two selectors were backed by different store properties that were
never kept in sync. By updating both properties in every selection path
(sidebar pick, chat pick, go-home, page restore), they always reflect
the same model. The `chatLaunchState = "ready"` fix mirrors what
`launchModelForChat` already does (line 2649) and prevents the chat
state from being "idle" when transitioning from the welcome view to the
chat view.

## Test Plan

### Manual Testing
<!-- Hardware: MacBook Pro M4 Max 48GB -->
- Select a model from sidebar "Load Model" dropdown → chat input
selector updates to match
- Select a model from chat input selector → sidebar dropdown updates to
match
- Launch a model, chat with it, click "Go Home" → both selectors show
the same model
- Refresh the page → both selectors show the previously selected model
- With a running model, type and send a message from the welcome view →
chat starts directly without flashing the recommended models view

### Automated Testing
No new automated tests — this is a UI state synchronization fix in the
Svelte dashboard with no backend changes.

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 17:31:20 +00:00
Alex Cheema b65982ddd7 fix: improve text contrast on HOME and DOWNLOADS nav links (#1609)
## Motivation

Follow-up to #1601 (downloads page contrast fix). The HOME and DOWNLOADS
navigation links in the top-right header use `text-exo-light-gray`
(`oklch(0.6 0 0)`) which is too dim against the dark header background.

## Changes

Changed both nav links in `HeaderNav.svelte` from `text-exo-light-gray`
to `text-white/70` for better visibility. Hover state
(`text-exo-yellow`) is unchanged.

## Why It Works

`text-white/70` provides noticeably better contrast against
`bg-exo-dark-gray` while still looking subdued relative to the yellow
accent color on hover. This is consistent with the approach used in
#1601.

## Test Plan

### Manual Testing
- Verified both links are clearly readable on the home page and
downloads page
- Hover state still transitions to yellow as expected

### Automated Testing
- Dashboard builds successfully

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 13:48:40 +00:00
rltakashige 2fe689315b download .model files in exo bench (#1607)
## Motivation

failed again for kimi on a machine that had never downloaded it.

## Test Plan

### Manual Testing
it worked this time
2026-02-24 11:13:04 +00:00
Alex Cheema 644c5573ce fix: improve text contrast on downloads page (#1601)
## Summary
- Bumps opacity on dark-grey text/icons on the downloads page that were
nearly invisible against the dark background
- Informational text (GB sizes, speeds, disk free, model IDs) → full
`text-exo-light-gray`
- Interactive icons (delete, resume, retry) → `/70` at rest
- Hover-only elements (download button) → `/60`

## Test plan
- [ ] Open http://localhost:52415/downloads with models downloaded
- [ ] Verify GB downloaded amounts are clearly visible
- [ ] Verify delete trash icons are visible (not just on hover)
- [ ] Verify download speed text is readable
- [ ] Verify "paused" labels and resume buttons are visible
- [ ] Verify "X GB free" disk labels in column headers are readable

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 10:42:49 +00:00
rltakashige 12c3015f52 fix qwen moe tensor sharding (#1604)
## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
2026-02-23 21:23:27 +00:00
rltakashige 365dd68d9a Final fixes for release (#1603)
## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
2026-02-23 21:10:15 +00:00
Alex Cheema d3d129581e test: verify instance deletion cancels ongoing tasks (#1508)
## Summary
- The cancellation logic for issue #1215 already exists in
`get_transition_events()` (`src/exo/master/placement.py:208-227`) — when
an instance is deleted, `TaskStatusUpdated(Cancelled)` events are
emitted for all Pending/Running tasks on that instance
- Combined with PR #1276's token-boundary cancellation in runners, the
full pipeline works end-to-end
- However, the existing test
`test_get_transition_events_delete_instance` passed `{}` for tasks, so
this path was never exercised
- This PR adds 4 tests covering the cancellation behavior:
  - Running tasks are cancelled on instance deletion
  - Pending tasks are cancelled on instance deletion
  - Completed/Failed/TimedOut/Cancelled tasks are left alone
  - Only tasks matching the deleted instance are cancelled

Closes #1215

## Test plan
- [x] `uv run pytest src/exo/master/tests/test_placement.py -v` — all 15
tests pass
- [x] `uv run basedpyright` — 0 errors
- [x] `uv run ruff check` — all checks passed

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 20:12:23 +00:00
597 changed files with 21464 additions and 5278 deletions
+20
View File
@@ -0,0 +1,20 @@
from enum import Enum
class HarmonyEncodingName(Enum):
HARMONY_GPT_OSS = ...
class HarmonyEncoding: ...
class HarmonyError(Exception): ...
class Role(Enum):
ASSISTANT = ...
class StreamableParser:
last_content_delta: str
current_channel: str | None
current_recipient: str | None
def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
def process(self, token_id: int) -> None: ...
def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...
+17
View File
@@ -0,0 +1,17 @@
class NvmlMemoryInfo:
used: int
total: int
free: int
class NvmlUtilizationRates:
gpu: int
memory: int
def nvmlInit() -> None: ...
def nvmlShutdown() -> None: ...
def nvmlDeviceGetCount() -> int: ...
def nvmlDeviceGetHandleByIndex(index: int) -> object: ...
def nvmlDeviceGetUtilizationRates(handle: object) -> NvmlUtilizationRates: ...
def nvmlDeviceGetTemperature(handle: object, sensor_type: int) -> int: ...
def nvmlDeviceGetPowerUsage(handle: object) -> int: ...
def nvmlDeviceGetMemoryInfo(handle: object) -> NvmlMemoryInfo: ...
+61
View File
@@ -0,0 +1,61 @@
from typing import Any, Sequence
from torch import backends as backends
from torch import cuda as cuda
from torch import distributed as distributed
__version__: str
class version:
cuda: str
class dtype: ...
bfloat16: dtype
float16: dtype
float32: dtype
int8: dtype
int32: dtype
int64: dtype
long: dtype
float8_e4m3fn: dtype
class Tensor:
shape: Sequence[int]
dtype: dtype
def __getitem__(self, key: Any) -> Tensor: ...
def __setitem__(self, key: Any, value: Any) -> None: ...
def to(self, *args: Any, **kwargs: Any) -> Tensor: ...
def cpu(self) -> Tensor: ...
def detach(self) -> Tensor: ...
def clone(self) -> Tensor: ...
def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Tensor: ...
def view(self, *shape: Any) -> Tensor: ...
def squeeze(self, dim: int = ...) -> Tensor: ...
def unsqueeze(self, dim: int) -> Tensor: ...
def permute(self, *dims: int) -> Tensor: ...
def float(self) -> Tensor: ...
def numpy(self) -> Any: ...
def numel(self) -> int: ...
def nelement(self) -> int: ...
@property
def is_cuda(self) -> bool: ...
@property
def device(self) -> device: ...
def __len__(self) -> int: ...
def data_ptr(self) -> int: ...
def tolist(self) -> Any: ...
def abs(self) -> Tensor: ...
def max(self) -> Tensor: ...
def mean(self) -> Tensor: ...
def sum(self, dim: int = ...) -> Tensor: ...
def item(self) -> float: ...
def tensor(data: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def zeros(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def empty(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def from_numpy(ndarray: Any) -> Tensor: ...
def inference_mode() -> Any: ...
class device:
def __init__(self, type: str, index: int = ...) -> None: ...
@@ -0,0 +1 @@
from torch.backends import cuda as cuda
@@ -0,0 +1 @@
def is_built() -> bool: ...
+10
View File
@@ -0,0 +1,10 @@
class _DeviceProperties:
total_memory: int
def is_available() -> bool: ...
def get_device_name(device: int) -> str: ...
def get_device_properties(device: int) -> _DeviceProperties: ...
def empty_cache() -> None: ...
def mem_get_info() -> tuple[int, int]: ...
def synchronize() -> None: ...
def max_memory_allocated() -> int: ...
@@ -0,0 +1,2 @@
def is_initialized() -> bool: ...
def destroy_process_group() -> None: ...
+1
View File
@@ -0,0 +1 @@
__version__: str
+2
View File
@@ -0,0 +1,2 @@
class ModelConfig:
max_model_len: int
+18
View File
@@ -0,0 +1,18 @@
from dataclasses import dataclass
@dataclass
class EngineArgs:
model: str = ...
served_model_name: str | list[str] | None = ...
tokenizer: str | None = ...
trust_remote_code: bool = ...
dtype: str = ...
seed: int = ...
max_model_len: int | None = ...
gpu_memory_utilization: float = ...
enforce_eager: bool = ...
tensor_parallel_size: int = ...
pipeline_parallel_size: int = ...
quantization: str | None = ...
load_format: str = ...
enable_sleep_mode: bool = ...
+17
View File
@@ -0,0 +1,17 @@
class CompletionOutput:
index: int
text: str
token_ids: list[int]
cumulative_logprob: float | None
logprobs: object | None
finish_reason: str | None
stop_reason: int | str | None
def finished(self) -> bool: ...
class RequestOutput:
request_id: str
prompt: str | None
prompt_token_ids: list[int] | None
outputs: list[CompletionOutput]
finished: bool
+11
View File
@@ -0,0 +1,11 @@
class SamplingParams:
n: int
temperature: float
top_p: float
top_k: int
min_p: float
seed: int | None
stop: str | list[str] | None
max_tokens: int | None
logprobs: int | None
repetition_penalty: float
@@ -0,0 +1,3 @@
from vllm.tokenizers.protocol import TokenizerLike
__all__ = ["TokenizerLike"]
@@ -0,0 +1,15 @@
from typing import Protocol
class TokenizerLike(Protocol):
@property
def eos_token_id(self) -> int: ...
@property
def vocab_size(self) -> int: ...
def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]: ...
def decode(self, ids: list[int] | int, skip_special_tokens: bool = ...) -> str: ...
def apply_chat_template(
self,
messages: list[dict[str, str]],
tools: list[dict[str, object]] | None = ...,
**kwargs: object,
) -> str | list[int]: ...
+1
View File
@@ -0,0 +1 @@
+1
View File
@@ -0,0 +1 @@
@@ -0,0 +1,24 @@
from collections.abc import Sequence
from vllm.v1.core.kv_cache_utils import BlockPool, KVCacheBlock
from vllm.v1.kv_cache_interface import KVCacheConfig
class KVCacheBlocks:
blocks: tuple[Sequence[KVCacheBlock], ...]
def __init__(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None: ...
def get_block_ids(self) -> tuple[list[int], ...]: ...
class KVCacheManager:
block_pool: BlockPool
kv_cache_config: KVCacheConfig
enable_caching: bool
num_kv_cache_groups: int
coordinator: object
def __init__(self, *args: object, **kwargs: object) -> None: ...
def allocate_slots(
self, request: object, num_new_tokens: int, *args: object, **kwargs: object
) -> KVCacheBlocks | None: ...
def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
def create_kv_cache_blocks(
self, blocks: tuple[list[KVCacheBlock], ...]
) -> KVCacheBlocks: ...
@@ -0,0 +1,16 @@
class KVCacheBlock:
block_id: int
ref_cnt: int
def __init__(self, block_id: int) -> None: ...
class FreeKVCacheBlockQueue:
def append_n(self, blocks: list[KVCacheBlock]) -> None: ...
def popleft_n(self, n: int) -> list[KVCacheBlock]: ...
class BlockPool:
blocks: list[KVCacheBlock]
free_block_queue: FreeKVCacheBlockQueue
num_gpu_blocks: int
enable_caching: bool
def get_num_free_blocks(self) -> int: ...
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ...
@@ -0,0 +1,22 @@
from vllm.config import ModelConfig
from vllm.engine.arg_utils import EngineArgs
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.tokenizers import TokenizerLike
class LLMEngine:
tokenizer: TokenizerLike | None
model_config: ModelConfig
@classmethod
def from_engine_args(cls, engine_args: EngineArgs) -> LLMEngine: ...
def add_request(
self,
request_id: str,
prompt: str,
params: SamplingParams,
arrival_time: float | None = ...,
) -> None: ...
def step(self) -> list[RequestOutput]: ...
def has_unfinished_requests(self) -> bool: ...
def get_tokenizer(self) -> TokenizerLike: ...
@@ -0,0 +1,23 @@
from dataclasses import dataclass
@dataclass
class KVCacheSpec:
block_size: int
num_kv_heads: int
head_size: int
@dataclass
class KVCacheGroupSpec:
layer_names: list[str]
kv_cache_spec: KVCacheSpec
@dataclass
class KVCacheTensorSpec:
shared_by: list[str]
size: int
@dataclass
class KVCacheConfig:
num_blocks: int
kv_cache_groups: list[KVCacheGroupSpec]
kv_cache_tensors: list[KVCacheTensorSpec]
+6
View File
@@ -0,0 +1,6 @@
class Request:
request_id: str
prompt_token_ids: list[int] | None
num_prompt_tokens: int
num_computed_tokens: int
num_tokens: int
@@ -0,0 +1 @@
@@ -0,0 +1,24 @@
import torch
class _CompilationConfig:
static_forward_context: dict[str, object]
class _ModelConfig:
hf_config: object
class GPUModelRunner:
kv_caches: list[torch.Tensor]
compilation_config: _CompilationConfig
model_config: _ModelConfig | None
def _allocate_kv_cache_tensors(
self, kv_cache_config: object
) -> dict[str, torch.Tensor]: ...
def initialize_kv_cache_tensors(
self, kv_cache_config: object, kernel_block_sizes: list[int]
) -> dict[str, torch.Tensor]: ...
def _reshape_kv_cache_tensors(
self,
kv_cache_config: object,
raw_tensors: dict[str, torch.Tensor],
kernel_block_sizes: list[int],
) -> dict[str, torch.Tensor]: ...
@@ -0,0 +1,6 @@
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
class Worker:
model_runner: GPUModelRunner
def determine_available_memory(self) -> int: ...
def initialize_from_config(self, kv_cache_config: object) -> None: ...
+1
View File
@@ -0,0 +1 @@
def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...
-20
View File
@@ -1,20 +0,0 @@
"""
This type stub file was generated by pyright.
"""
from activations import *
from base import *
from containers import *
from convolution import *
from convolution_transpose import *
from distributed import *
from dropout import *
from embedding import *
from linear import *
from normalization import *
from pooling import *
from positional_encoding import *
from quantized import *
from recurrent import *
from transformer import *
from upsample import *
View File
@@ -2,10 +2,11 @@
This type stub file was generated by pyright.
"""
from typing import Protocol
import mlx.core as mx
import PIL.Image
import tqdm
from typing import Protocol
from mflux.models.common.config.config import Config
class BeforeLoopCallback(Protocol):
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.callbacks.callback import (
AfterLoopCallback,
BeforeLoopCallback,
@@ -2,10 +2,11 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
import PIL.Image
import tqdm
from typing import TYPE_CHECKING
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.models.common.config.config import Config
@@ -2,11 +2,12 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from pathlib import Path
from typing import Any
from tqdm import tqdm
import mlx.core as mx
from mflux.models.common.config.model_config import ModelConfig
from tqdm import tqdm
logger = ...
@@ -2,10 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from functools import lru_cache
from typing import Literal
import mlx.core as mx
class ModelConfig:
precision: mx.Dtype = ...
def __init__(
@@ -2,10 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from pathlib import Path
from typing import TYPE_CHECKING, TypeAlias
from mlx import nn
import mlx.core as mx
from mflux.models.common.vae.tiling_config import TilingConfig
from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
@@ -13,6 +13,7 @@ from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreat
from mflux.models.z_image.latent_creator.z_image_latent_creator import (
ZImageLatentCreator,
)
from mlx import nn
if TYPE_CHECKING:
LatentCreatorType: TypeAlias = type[
@@ -2,8 +2,8 @@
This type stub file was generated by pyright.
"""
from mlx import nn
from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
from mlx import nn
class FusedLoRALinear(nn.Module):
def __init__(
@@ -2,10 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
import mlx.nn as nn
from collections.abc import Callable
from dataclasses import dataclass
import mlx.core as mx
import mlx.nn as nn
from mflux.models.common.lora.mapping.lora_mapping import LoRATarget
@dataclass
@@ -2,11 +2,12 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from collections.abc import Callable
from dataclasses import dataclass
from typing import List, Protocol
import mlx.core as mx
@dataclass
class LoRATarget:
model_path: str
@@ -36,4 +36,3 @@ class Rule(NamedTuple):
name: str
check: str
action: QuantizationAction | PathAction | LoraAction | ConfigAction
...
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.config.model_config import ModelConfig
if TYPE_CHECKING: ...
@@ -2,9 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from abc import ABC, abstractmethod
import mlx.core as mx
class BaseScheduler(ABC):
@property
@abstractmethod
@@ -2,8 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from typing import TYPE_CHECKING
import mlx.core as mx
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -2,8 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from typing import TYPE_CHECKING
import mlx.core as mx
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -2,8 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from typing import TYPE_CHECKING
import mlx.core as mx
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -4,9 +4,10 @@ This type stub file was generated by pyright.
from abc import ABC, abstractmethod
from typing import Protocol, runtime_checkable
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
from PIL import Image
from transformers import PreTrainedTokenizer
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
"""
This type stub file was generated by pyright.
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from mflux.models.common.weights.loading.weight_definition import TokenizerDefinition
@@ -2,9 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
import mlx.core as mx
"""
This type stub file was generated by pyright.
"""
@@ -2,9 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from typing import Callable
import mlx.core as mx
class VAETiler:
@staticmethod
def encode_image_tiled(
@@ -3,8 +3,8 @@ This type stub file was generated by pyright.
"""
import mlx.core as mx
from mlx import nn
from mflux.models.common.vae.tiling_config import TilingConfig
from mlx import nn
class VAEUtil:
@staticmethod
@@ -2,8 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.nn as nn
from typing import TYPE_CHECKING
import mlx.nn as nn
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
@@ -2,11 +2,12 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from typing import Callable, List, TYPE_CHECKING, TypeAlias
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
from typing import TYPE_CHECKING, Callable, List, TypeAlias
import mlx.core as mx
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
from mflux.models.depth_pro.weights.depth_pro_weight_definition import (
DepthProWeightDefinition,
)
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
@@ -2,8 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from typing import Dict, List, Optional
import mlx.core as mx
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
class WeightMapper:
@@ -2,10 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from typing import Callable, List, Optional, Protocol
import mlx.core as mx
"""
This type stub file was generated by pyright.
"""
@@ -2,7 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import Any, TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType
if TYPE_CHECKING: ...
@@ -2,9 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from pathlib import Path
import mlx.core as mx
from PIL import Image
@dataclass
@@ -13,7 +14,6 @@ class DepthResult:
depth_array: mx.array
min_depth: float
max_depth: float
...
class DepthPro:
def __init__(self, quantize: int | None = ...) -> None: ...
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
TokenizerDefinition,
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import (
WeightMapping,
WeightTarget,
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
TokenizerDefinition,
@@ -3,6 +3,7 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import (
WeightMapping,
WeightTarget,

Some files were not shown because too many files have changed in this diff Show More