Skip to content

fix(mlx): interleave prefill with active decode - #2312

Open
STRML wants to merge 2 commits into
exo-explore:mainfrom
STRML:fix/interleave-prefill-decode
Open

STRML wants to merge 2 commits into
exo-explore:mainfrom
STRML:fix/interleave-prefill-decode

Conversation

@STRML

@STRML STRML commented Sep 15, 2026

Copy link
Copy Markdown

Addresses #1908.

Summary

  • admit ordinary text prompts into mlx-lm continuous batching when another request is active
  • run active decode first, then process one bounded prompt chunk per scheduler iteration
  • serialize competing prefills so simultaneous arrivals cannot multiply the decode pause
  • default to 256-token chunks, configurable with EXO_INTERLEAVED_PREFILL_STEP_SIZE
  • preserve prefix-cache updates, hybrid SSM snapshots, prefill progress, and distributed cancellation
  • retain the existing synchronous path for pipeline-parallel, vision, and remote-prefill requests

The first request on an idle engine still uses the high-throughput synchronous prefill path. Only newly arriving prompts compete with existing work. This implements the scheduler-level chunking direction proposed in #1908 without requiring mixed prefill/decode kernels.

Self-review notes

  • prefill_batch_size is deliberately 1, preserving the issue stated one-prefill-at-a-time invariant and making the chunk size a cluster-wide latency budget rather than a per-request budget
  • interleaved scheduler chunks are 256 tokens, while hybrid-cache state snapshots retain a 4096-token cadence plus the final reusable state; this avoids a 16x snapshot-memory increase
  • rotating caches are clamped before mlx-lm merges them into either path
  • exact and partial prefix-cache metadata, progress callbacks, and cancellation agreement are carried across the asynchronous prompt lifecycle

Automated testing

  • ruff check .
  • basedpyright
  • pytest src -q: 470 passed, 3 skipped, 190 deselected
  • the new regression is part of the normal CI set, uses a self-contained tiny tokenizer/model, and runs locally in about 2.4 seconds
  • it verifies active decode continues, only one competing prefill advances, chunked prefill reaches decode, progress is monotonic, prefix-cache save and exact reuse work, and distributed cancellation removes queued prefill
  • existing slow B=2 tests require the source-built EXO MLX fork; with the stock prebuilt MLX wheel they fail identically on untouched main and this branch

Nix is not installed on the local test machine, so nix fmt was not run. No Nix files changed.

Scope and follow-ups

This patch targets single-node and tensor-parallel text generation. Pipeline-parallel and vision prefills need specialized incremental handling and remain on the existing synchronous path. Remote prefill retains its existing synchronous transfer/cache-ingestion lifecycle. Hardware tuning can adjust EXO_INTERLEAVED_PREFILL_STEP_SIZE without changing scheduler behavior.

@STRML
STRML marked this pull request as ready for review September 15, 2026 15:08
@STRML

STRML commented Sep 15, 2026

Copy link
Copy Markdown
Author

@rltakashige this is ready for review and implements the scheduler-level chunked-prefill direction from #1908. The regression now runs in the normal CI set and covers competing prefill serialization, active decode progress, prefix-cache reuse, and cancellation. GitHub is holding the fork workflow at action_required pending maintainer approval.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant