fix: run BatchedEngine MLLM on dedicated MLXWorkerThread to prevent cross-thread stream errors

[xykong]

Summary

Fixes RuntimeError: There is no Stream(gpu, N) in current thread when running BatchedEngine (--continuous-batching) with MLX >= 0.31.2 on Apple Silicon.

Problem

MLX 0.31.2 made Metal CommandEncoders thread-local. When BatchedEngine._prepare_mllm_model() loads the model on the event-loop thread but MLLMScheduler._process_loop() runs step() on a different thread, the runtime raises:

RuntimeError: There is no Stream(gpu, N) in current thread

This happens because MLX arrays created during model loading are tagged with the loading thread's stream, which doesn't exist in the inference thread's TLS.

Fix

1. mlx_streams.py: Add MLXWorkerThread — a persistent single-thread executor that guarantees model load + inference share the same thread-local stream context.

2. engine/batched.py: Make _prepare_mllm_model() async and run model loading on MLXWorkerThread. Pass the worker to MLLMScheduler.

3. mllm_scheduler.py: Accept mlx_worker parameter. When provided, submit step() and preprocessing to the worker thread. Falls back to the legacy event-loop path (with bind_generation_streams()) when no worker is provided.

4. mllm_batch_generator.py: Remove MLLMBatchGenerator._stream (class-level mx.new_stream()) in favor of the worker thread's default stream, eliminating the last source of cross-thread stream references.

Testing

Tested on M4 Max 128GB with mlx-community/gemma-4-26b-a4b-it-4bit:

Metric	Result
Single request (short)	93 tok/s
Single request (long)	82 tok/s
4x concurrent aggregate	179 tok/s
TTFT (short)	35ms

Previously this configuration crashed immediately on the first request.

Backward Compatibility

• The legacy event-loop path (no MLXWorkerThread) is preserved as fallback when mlx_worker=None
• SimpleEngine path is unaffected
• LLM-only BatchedEngine path is unaffected

[janhilgard]

The overall approach is solid — model load + inference on the same persistent MLXWorkerThread is the right fix for MLX 0.31.2 thread-local CommandEncoders. The legacy fallback path is a nice touch for backward compatibility. No streaming regression here (unlike #478 for SimpleEngine).

A few issues to address:

---

1. if True: replaces with mx.stream(...) — should dedent instead

if True:  # use thread-default stream
    logits = self.language_model(last_token, cache=request_cache)

This appears 4 times. The if True: block is dead code that exists only to preserve indentation. Just remove the with mx.stream(...) wrapper and dedent the body. A # noqa or a comment is fine, but a no-op if True: is a code smell that will confuse future readers.

2. Preprocessing unnecessarily serialized on the MLX worker thread

The old code ran text-only preprocessing in the default thread-pool executor (loop.run_in_executor(None, ...)):

await loop.run_in_executor(None, bg._preprocess_request, req)

The new worker-thread path submits preprocessing to the same single MLX worker:

await self._mlx_worker.submit(loop, bg._preprocess_request, req)

The comment says "Preprocessing creates MLX arrays (input_ids) so it must run on the same worker thread" — but _preprocess_request does Jinja2 template rendering + tokenizer.encode(), which produces Python lists, not mx.array. The mx.array() conversion happens inside _process_prompts() called from step().

Running preprocessing on the single MLX worker thread serializes it with step(), which means a 30-second preprocessing of a 40K+ token conversation blocks ALL inference for 30 seconds. The old approach of offloading to the thread pool kept inference unblocked during preprocessing. Consider keeping run_in_executor for preprocessing.

3. Lost slow-preprocessing logging and dynamic n_yields

The old code logged slow preprocessing:

elapsed = time.perf_counter() - tic
if elapsed > 1.0:
    logger.info(f"Preprocessing {req.request_id[:12]}: {n_tok} tokens in {elapsed:.2f}s")

And had dynamic yield counts:

n_yields = 10 if elapsed > 1.0 else 5

The worker-thread path drops both — no timing, always 5 yields. Please preserve the slow-preprocessing logging (it's very useful for debugging production issues) and the dynamic yield count.

4. mx.new_thread_local_stream(mx.gpu) should live in MLXWorkerThread._run(), not in _load_on_worker

Currently the thread-local stream is initialized inside _load_on_worker():

def _load_on_worker():
    mx.new_thread_local_stream(mx.gpu)
    inst = MLXMultimodalLM(...)

But the MLXWorkerThread docstring promises "mx.new_stream() is called exactly once (during thread init)." If someone submits a task before model loading, there's no stream. Move mx.new_thread_local_stream(mx.gpu) into MLXWorkerThread._run() at the top, before the task loop.

5. Overlap with PR #478 on mlx_streams.py

Both #478 and #479 add MLXWorkerThread to mlx_streams.py. If both land, there will be a merge conflict. Consider coordinating — either land #478 first and rebase #479 on top, or extract MLXWorkerThread into a shared PR that both depend on.

6. Comments stripped from legacy path

Several explanatory comments were removed from the legacy path (early preprocessing phase rationale, health-check yield explanation). These comments document why the code works the way it does and are valuable for maintainability. Please keep them.

---

Summary

The MLXWorkerThread + BatchedEngine integration is architecturally sound. Main blockers: (1) preprocessing should stay on the thread pool to avoid blocking inference, (2) if True: should be proper dedent, (3) preserve the slow-step diagnostics logging. The rest are minor cleanups.

[CarloArpini]

Thanks for the great work on this PR which builds on top of #478; the MLLM path fix is solid and the two-phase preprocessing split is a nice touch.

I wanted to flag that the LLM-only BatchedEngine path (--continuous-batching with a text-only model, i.e. mllm=False) is still crashing with the same RuntimeError: There is no Stream(gpu, N) in current thread on MLX 0.31.2. The crash occurs in scheduler.step() when it calls self.batch_generator.next() on a thread that doesn't own the GPU stream created during BatchGenerator initialization.

Full traceback:

File "vllm_mlx/scheduler.py", line 2455, in step
    result = self.batch_generator.next()
  File "mlx_lm/generate.py", line 1855, in next
    return self._next()
  File "mlx_lm/generate.py", line 1841, in _next
    self._prompt_batch.prompt(prompts)
  File "mlx_lm/generate.py", line 1161, in prompt
    mx.eval([c.state for c in self.prompt_cache])
RuntimeError: There is no Stream(gpu, 1) in current thread.

Server startup confirms LLM-only path: BatchedEngine loaded: <model> (mllm=False).

TBH I don't have enough expertise to confirm this, but maybe the fix can follow the same pattern as this PR: the SchedulerCore (or wherever step() is dispatched) needs to submit batch_generator.next() to the same persistent MLXWorkerThread on which the BatchGenerator was created, mirroring what you've done here for MLLMScheduler._process_loop().

Environment: Apple M3 Ultra 512GB, macOS, MLX 0.31.2, mlx-lm 0.31.3, DeepSeek-R1-Distill-Llama-8B (text-only). Reproduced consistently across the original waybarrios/main, adriangalilea fork (PR #478), and this branch.

[waybarrios]

Hi @xykong, thanks for chasing the cross-thread stream bug. The MLXWorkerThread approach looks solid and the test results on Gemma look great. CI is red on three small issues; happy to push a follow-up commit to your branch if you'd prefer, otherwise here's what to fix.

The merge conflict is in vllm_mlx/mlx_streams.py and it's documentation only. main extended the bind_generation_streams docstring with a paragraph about admission/ownership semantics. Your PR added a .. deprecated:: block. Both texts are complementary, so the cleanest resolution after git fetch origin main && git merge origin/main is to keep both:

    MLX streams are thread-local. If a model is loaded on one thread and
    generation runs on another, module-level generation streams created during
    import can point at a stream that does not exist in the worker thread.

    This intentionally creates a fresh stream for the current worker call and
    replaces module-level generation_stream handles under a process-local lock.
    It is an admission/ownership fix, not a batching optimization; callers
    should invoke it at worker-entry boundaries rather than inside token loops.

    .. deprecated::
        Prefer ``MLXWorkerThread`` which keeps model load and inference on the
        same persistent thread, avoiding the need to rebind streams entirely.
    """

Then the lint failure on vllm_mlx/mllm_batch_generator.py:787. The from mlx_vlm.utils import prepare_inputs is unused after your preprocessing refactor in a157942 (the function now calls tokenizer(...) directly). Just delete that import; ruff check --fix will do it for you.

Finally, the failing test in tests/test_mllm_continuous_batching.py:1259. The mock FakeMLLMScheduler.__init__ doesn't accept the new mlx_worker kwarg that BatchedEngine._start_mllm now passes. One line:

        class FakeMLLMScheduler:
            def __init__(self, model, processor, config, mlx_worker=None):
                captured["scheduler_config"] = config
                captured["mlx_worker"] = mlx_worker

Capturing mlx_worker matches the existing capture style for config_kwargs and leaves the door open for a future assertion that the worker is wired correctly. This same change fixes the 3.11, 3.13, and tests aggregator jobs.

Let me know which path you want and I'll re-run CI when it lands.

[waybarrios]

Closing as superseded, same story as #478. Main's design keeps MLLM step() on the event-loop thread (see mllm_scheduler) and prepare_for_start runs inline since 3396981, so the dedicated worker thread approach moves in the opposite direction from what landed.