[presentyourlove]

Description

This PR resolves item 7 in TODO.md: "Solve the problem of in continuous batching when a new prompt comes in, it will block decode of the current batch until the prefill is complete."

Background & Problem

In continuous batching with BatchGenerator and ExoBatchGenerator, whenever a new request arrived while existing requests were actively decoding, the batch loop synchronously executed the entire prefill phase for the new request before advancing the generation loop:

1. Long prompts (e.g. 2,000–8,000 tokens) take hundreds of milliseconds to several seconds to prefill the KV cache.
2. During this entire prefill duration, _gen.step() was blocked, completely freezing token emission for all currently active streams (causing noticeable jitter, stutter, and degraded perceived latency for active users).
3. If an incoming request was cancelled while queued or prefilling, there was no granular mechanism to halt intermediate prefill execution without running the whole prompt.

Solution: Chunked Prefill Continuous Batching

This PR introduces Chunked Prefill to interleave prompt prefilling in configurable chunk budgets (default: 512 tokens) with active decode steps.

1. Non-blocking Multi-step Prefill Session (src/exo/worker/engines/mlx/generator/batch_generate.py):

   • Introduced PrefillSession dataclass tracking sliced prompt tokens, KV cache accumulator, media regions, progress callbacks, and completion status.
   • Added prepare_prefill(...) -> PrefillSession: sets up prompt tokenization, checks prefix KV cache hits, and creates a prefill session without executing forward computation.
   • Added step_prefill(session, chunk_size=512) -> tuple[bool, int | None]: evaluates one chunk of tokens against the model, updates rotating KV cache, notifies real-time prefill progress (PrefillProgressChunk), and registers the generation task once prefill is complete.
   • Maintained full backward compatibility for submit(...) by wrapping prepare_prefill + step_prefill.

2. Interleaved Decode and Prefill Loop (src/exo/worker/runner/llm_inference/batch_generator.py):

   • Added _prefilling_task tracking in BatchGenerator.
   • On each step:
     • Drains active generation tokens via _gen.step() so active streams emit tokens every step (~30ms) without any perceptible stutter.
     • Advances _prefilling_task by one chunk (default: 512 tokens when active decode streams exist; opportunistic fast-path 4,096 tokens when idle).
     • When prefill completes, seamlessly joins the generation batch in _active_tasks.
   • Enhanced cancellation handling (_apply_cancellations): promptly drops prefilling sessions if cancelled before completing prefill.

3. Cross-Platform Compatibility (src/exo/utils/channels.py, src/exo/worker/runner/bootstrap.py):

   • Added _MemoryObjectStreamState fallback in channels.py for varying AnyIO environments.
   • Made Unix-specific resource import optional in bootstrap.py for cross-platform robustness.

4. Comprehensive Unit Tests (src/exo/worker/tests/unittests/test_runner/test_chunked_prefill.py):

   • test_chunked_prefill_does_not_block_active_decode_batch: Verifies that Task A decodes a token at every single step while Task B prefills across multiple 512-token chunks, after which both decode concurrently.
   • test_cancellation_during_chunked_prefill: Verifies that a task cancelled during intermediate prefilling chunks cleanly terminates without corrupting the batch or stalling existing tasks.
   • test_legacy_engine_fallback_without_chunking: Verifies backward compatibility for legacy engines lacking chunked prefill support.

Verification

• Unit Tests:

  .venv/Scripts/python.exe -m pytest src/exo/worker/tests/unittests/test_runner/test_chunked_prefill.py -v
  # 3 passed in 0.57s

• Code Linting & Formatting:

  .venv/Scripts/python.exe -m ruff check src/exo/utils/channels.py src/exo/worker/engines/mlx/generator/batch_generate.py src/exo/worker/runner/bootstrap.py src/exo/worker/runner/llm_inference/batch_generator.py src/exo/worker/tests/unittests/test_runner/test_chunked_prefill.py
  .venv/Scripts/python.exe -m ruff format src/exo/utils/channels.py src/exo/worker/engines/mlx/generator/batch_generate.py src/exo/worker/runner/bootstrap.py src/exo/worker/runner/llm_inference/batch_generator.py src/exo/worker/tests/unittests/test_runner/test_chunked_prefill.py
  # All checks passed!

Related Issues

• Resolves: Item 7 in TODO.md ("Solve the problem of in continuous batching when a new prompt comes in, it will block decode of the current batch until the prefill is complete.")