fix: extend MLX thread-affinity fix (#886) to the Qwen3 LLM backend - #909
Open
shsunmoonlee wants to merge 4 commits into
Open
fix: extend MLX thread-affinity fix (#886) to the Qwen3 LLM backend#909shsunmoonlee wants to merge 4 commits into
shsunmoonlee wants to merge 4 commits into
Conversation
Two related bugs on Apple Silicon, both stemming from the same root cause (mlx-lm/mlx-audio never getting installed correctly by `just setup-python`), plus one cosmetic health-endpoint bug found along the way: 1. `backend/backends/mlx_backend.py`: pin every MLX model load/generate/ transcribe call to a single dedicated worker thread instead of asyncio.to_thread()'s default pool. MLX's Metal command encoder/stream is thread-local — if model load and generation land on different pool threads (which asyncio.to_thread() does not guarantee against), MLX raises "There is no Stream(gpu, N) in current thread." Fixes jamiepine#699. 2. `justfile`: `just setup-python` installed requirements-mlx.txt (mlx + miniaudio) on Apple Silicon and stopped — it never ran the `pip install --no-deps mlx-lm==0.31.1` / `mlx-audio==0.4.1` steps that .github/workflows/release.yml already does for packaged builds. Without mlx-audio installed, get_backend_type() falls back to "pytorch" silently: no error, no log line calling this out, just a backend that's much slower than intended and never actually uses MLX. Fixes jamiepine#823. Installing mlx-lm without --no-deps (the obvious thing to try once you notice jamiepine#823) reintroduces a second problem: mlx-lm>=0.30.0 requires transformers>=5.0.0, which conflicts with the transformers<=4.57.6 cap in requirements.txt. transformers>=5.0 changed `check_model_inputs` from a decorator factory (`@check_model_inputs()`) to a plain decorator (`@check_model_inputs`, called as `check_model_inputs(func)`) — qwen_tts's CustomVoice tokenizer still uses the old calling convention, so loading Qwen CustomVoice raises `TypeError: check_model_inputs() missing 1 required positional argument: 'func'` once transformers gets silently upgraded. --no-deps avoids this entirely since mlx-lm's own transformers requirement is never consulted; documented this in requirements-mlx.txt so the next person doesn't hit the same trap. 3. `backend/routes/health.py`: `/health`'s `backend_variant` field never had an `mlx` branch, so it always reported "cpu" regardless of whether MLX was active or working — cosmetic, but it's what a couple of the linked issues (jamiepine#706, jamiepine#723) used to conclude MLX wasn't running on the GPU when `backend_type` said otherwise. Added the missing branch so `backend_variant` reports "mps" when the mlx backend is selected. All three verified on real Apple Silicon hardware (M4): repeated generations with no thread crash, `just setup-python` now installs a working MLX backend from a clean venv, Qwen CustomVoice downloads and generates successfully post-fix, and `/health` reports `backend_variant: "mps"` when MLX is active. Fixes jamiepine#699, fixes jamiepine#823
…e Silicon get_torch_device() already supports allow_mps, and other engines (luxtts_backend.py, qwen_llm_backend.py) already pass allow_mps=True — but PyTorchTTSBackend, PyTorchSTTBackend, and QwenCustomVoiceBackend never did, so any Apple Silicon install that ends up on the PyTorch backend (e.g. MLX unavailable, or before this PR's justfile fix) silently ran these three on plain CPU instead of falling back to PyTorch+MPS. This turned out to matter more than expected: benchmarking on M4 hardware, PyTorch+MPS generation (RTF ~1.7x) was faster than this same PR's MLX thread-fix (RTF ~3.4-4.3x) for the same qwen engine and model size, and faster than a CUDA RTX 2080 Ti (RTF ~2.3-3.1x) too. I'm not proposing changing the default backend selection here — one engine, one model size, one machine is not enough data for that, and MLX-Whisper is reportedly faster than PyTorch-CPU for STT — but MPS should clearly be an option rather than silently absent. Verified on M4: PyTorchTTSBackend()._get_device() and QwenCustomVoiceBackend()._get_device() both now return "mps" (were "cpu").
…ad thread-safety Two follow-ups from the automated review on this PR: 1. backend/backends/base.py: empty_device_cache() only handled "cuda" and "xpu", never "mps" — so once the previous commit enabled MPS for the PyTorch Qwen3-TTS/CustomVoice engines, switching model sizes on Apple Silicon would leak MPS-cached memory (torch.mps.empty_cache() never got called). Added the missing branch. 2. backend/backends/qwen_custom_voice_backend.py: unload_model() had its own inline `if torch.cuda.is_available(): torch.cuda.empty_cache()` instead of using the shared empty_device_cache() helper — meant it never freed MPS memory either, same bug as jamiepine#1 but duplicated instead of shared. Switched to the shared helper (matches pytorch_backend.py's existing pattern). 3. backend/backends/mlx_backend.py: load_model_async()'s unload-on- model-size-change call went straight to self.unload_model() instead of routing through _run_on_mlx_thread like every other MLX call in this file (the whole point of the earlier commit in this PR). Fixed for consistency with the thread-pinning fix. Verified on M4: 0.6B -> generate still works post-changes (regression check), torch.mps.empty_cache() confirmed callable without error in this environment.
The dedicated-MLX-thread fix (jamiepine#886) covers the TTS and STT backends, but MLXQwenLLMBackend still dispatches load/generate through asyncio.to_thread(), so /llm/generate (personality compose/rewrite, dictation refinement) crashes with the same 'There is no Stream(gpu, 0) in current thread.' — raised from mlx_lm.generate's wired_limit on exit — whenever load and generate land on different pool threads. Route the MLX LLM backend's unload/load/generate through the same _run_on_mlx_thread helper so every MLX call in the process shares one worker thread. Verified on Apple M3 Pro (macOS 26.5, mlx 0.32.0, mlx-lm 0.31.1, mlx-audio 0.4.1): /llm/generate failed 100% before, succeeds after; TTS + STT unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (8)
📝 WalkthroughWalkthroughApple Silicon support now includes MPS cache clearing and device selection, dedicated-thread execution for MLX operations, corrected MLX health reporting, and explicit MLX runtime installation steps. ChangesApple Silicon backend support
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Confirming we hit exactly this: on Voicebox 0.5.0 / Apple Silicon (macOS 26.5.1, arm64), dictation refinement through the Qwen3 LLM backend fails 100% with TTS/STT aside, this is the path that breaks refinement for us (repro'd with |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Builds on #886 (its three commits are included here) — extends the dedicated-MLX-thread fix to the one backend it missed:
MLXQwenLLMBackend.#886 pins TTS and STT MLX calls to a single worker thread, but the Qwen3 LLM backend (
backend/backends/qwen_llm_backend.py) still dispatchesload_model/generatethroughasyncio.to_thread(). Result: with #886 applied, TTS and STT work, butPOST /llm/generate(personality Compose/Rewrite, dictation refinement,voicebox.speakwithpersonality: true) still fails 100% with the same error, raised frommlx_lm.generate'swired_limitcontext manager on exit:Fix
Route the MLX LLM backend's unload/load/generate through the same
_run_on_mlx_threadhelper #886 introduces, so every MLX call in the process (TTS, STT, and LLM) shares the one dedicated worker thread. Sharing a single serialized worker also matches the app's "one LLM, one model cache, one GPU-memory footprint" design.Verification
Apple M3 Pro, macOS 26.5 (Darwin 25.5.0), Python 3.12.8,
mlx 0.32.0/mlx-lm 0.31.1/mlx-audio 0.4.1, Qwen3 0.6B (4-bit MLX):POST /llm/generate→ 500, traceback above, 100% repro from a fresh backend start.POST /llm/generatereturns generated text; repeated calls stable./generateand Whisper Turbo/transcribestill work after the change (round-trip verified: generated audio transcribes back to the input text).Happy to rebase/fold into #886 if @AlexStephenLytton prefers to pick the commit into that branch instead.
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Chores