vendoring mlx-lm componenets - #880
Open
Lazarus-931 wants to merge 22 commits into
Open
Conversation
Each supporting ASR model now accepts a uniform hotwords: list[str] in its own generate() and folds it into its native prompt field: qwen3_asr->system_prompt, whisper->initial_prompt, vibevoice_asr->context, moss_transcribe_diarize->prompt. fun_asr_nano already handles hotwords natively. Backends without a hook ignore the field (silent drop, per maintainer preference). Shared merge_hotwords helper in stt/utils.py; models keep their native fields for backward compatibility.
Scaffolding for vendoring mlx-lm's transformer machinery into mlx-audio. The lm package stays import-cheap (no transformers at import time); parity tests live outside the CI test matrix and pin mlx-lm to the exact version being vendored from. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Trimmed copy of mlx_lm/models/base.py: drops the quantized SDPA path since no cache in mlx-audio exposes .bits. Swaps the 13 import sites. Parity tests assert bitwise equality with mlx-lm 0.31.3 across the mask cross-product, and that create_attention_mask returns the same *type* (str/None/array) so SDPA dispatch cannot silently change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Trimmed copy of mlx_lm/models/cache.py keeping the caches mlx-audio uses: KVCache, RotatingKVCache, BatchKVCache, ArraysCache, make_prompt_cache and the create_causal_mask re-export. Drops QuantizedKVCache, ConcatenateKVCache, ChunkedKVCache, CacheList, BatchRotatingKVCache, the prompt-cache save/load/trim helpers and the prompt trie. Extracted by AST so retained code is byte-identical to upstream. Swaps the 33 import sites. Parity tests cover RotatingKVCache wraparound across max_size/keep/sequence shapes comparing state after every step, BatchKVCache.merge in the shape continuous batching actually uses, and KVCache growth across the step-256 boundary. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Verbatim copy of mlx_lm/sample_utils.py (zero internal deps). Swaps the 23 import sites. Parity tests handle the mx.compile/RNG-state trap: these kernels are decorated with inputs=mx.random.state, so every comparison re-seeds before each call and also asserts the post-call RNG state matches, otherwise later sampler calls would diverge. Two upstream failure modes are pinned as parity evidence rather than hidden - min_p=0 hits math.log(0) and min_tokens_to_keep hits a put_along_axis signature mismatch in mlx-lm 0.31.3; both raise identically. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Upstream passes a Python bool to mx.put_along_axis, so any caller with min_p > 0 and min_tokens_to_keep > 1 raises TypeError. That path is reachable from vibevoice_asr, qwen3_asr, fun_asr_nano, glmasr and moss_music, all of which expose both knobs. Now that mlx-audio owns this code, fix it rather than pin the crash: pass a bool array, so min_tokens_to_keep is honoured (verified exactly N survivors under aggressive filtering) while the min_tokens_to_keep=1 path stays bitwise-identical to upstream. Also moves the parity CI job below its own comment so the modular job's comment is no longer orphaned. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
rope_utils.py and switch_layers.py copied verbatim; activations.py keeps only swiglu (xielu/XieLU is used solely by the apertus backbone, which mlx-audio does not vendor). Swaps the one SwitchGLU import in zonos2. This is the support layer the transformer backbones need. Parity tests cover initialize_rope across linear/llama3/yarn scaling with traditional on/off and non-zero offsets, SwitchLinear/SwitchGLU forward passes with shared weights, and the quantized switch path. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds mlx_audio/lm/{generate,load,convert}.py and flips the remaining call
sites: generate_step/stream_generate/generate across 13 STT and 4 TTS models,
the conversion helpers used by mlx_audio/convert.py and the per-model convert
scripts, and dramabox's text-encoder loader. Retargets the fun_asr_nano
monkeypatch, which patched mlx_lm.generate by module name and would otherwise
have passed vacuously.
Two bugs fixed in the vendored conversion helpers:
- save_model's donate_model path passed flattened dotted keys to
model.update(), which expects a nested tree; use tree_map over
model.parameters() as upstream does. This broke every quantized-checkpoint
round trip (caught by the cohere_asr 4-bit and 8-bit generate tests).
- quantize_model dropped upstream's fine-grained branch, so re-quantizing an
already-quantized config recorded no per-layer parameters.
Only sts/voice_pipeline.py still imports mlx-lm, by design.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
mlx-lm moves to an optional [llm] extra, needed only by the speech-to-speech pipeline's in-process LLM responder - the one place that loads an arbitrary chat model rather than a speech model's own transformer. Its imports move inside LocalLLMResponseEngine, with an actionable ImportError. This also fixes a live bug: voice_pipeline imported mlx_lm at module scope, so 'pip install mlx-audio[sts]' could not import VoicePipeline at all. The modular CI job only passed because the [all] install happened to pull mlx-lm in. Adds test_no_mlx_lm.py as a permanent guard: an AST walk rejecting mlx_lm imports outside the allowed file, a regex sweep for the dynamic importlib/__import__ form, and a check that mlx-lm never returns to [project.dependencies]. Verified with mlx-lm uninstalled: 1237 passed, 40 skipped, no collection errors; parity modules skip cleanly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
utils.py injects config["model_path"] for every model, but from_dict forwarded **config_dict straight into the dataclass, so loading LFM2.5-Audio raised TypeError before any model code ran. Filter to declared fields, matching how BaseModelArgs.from_dict behaves. Pre-existing and unrelated to the mlx-lm vendoring; found while verifying models on an M4 mini. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The generate parity test used a stub model, so it exercised none of the risky paths. Replace it with two structurally identical llama backbones sharing weights, compared across: greedy and sampled streams, logits processors, prefill chunk sizes including a single-step and an oversized one, max_kv_size (which forces RotatingKVCache wraparound), the max_tokens boundary, and input_embeddings. Adds a cross control running the vendored loop against the upstream model, so a loop bug is distinguishable from a backbone bug. Also documents the import table in the contributing guide and makes the guard test's failure message point at mlx_audio.lm, since several in-flight PRs add new mlx_lm imports and would otherwise hit an unexplained failure. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
generate.py: the EOS break happened before last_token was assigned, so the terminal GenerationResponse could never report finish_reason='stop' and re-emitted the previous token instead of the stop token. That matters for the four TTS consumers that append response.token to input_ids: the appended value was a duplicate audio code rather than the EOS the parsers filter, shifting codec framing at the tail of an utterance. Also restores a final response when the first sampled token is EOS, makes max_tokens=-1 unbounded again, and restores upstream's wired-limit warning. convert.py: save_model wrote a single monolithic safetensors file and donated the model only after serializing. Restore 5 GB sharding with the index map, and release references before writing so shards free as they go. mixed-quant now gives lm_head high bits unconditionally, as upstream does - the previous gate dropped it to low bits on models with fewer than eight layers. load.py: load weights strictly, so a mismatched text encoder raises instead of silently loading partly random weights. pyproject: [sts] and [all] now include mlx-lm. Dropping it from core also dropped it from those extras, which would have made the default S2S pipeline raise at startup. Adds parity coverage for the conversion helpers (all four mixed-quant recipes across 4/8/32 layers, which is what caught the lm_head gate) and permanent EOS-semantics tests that survive the parity teardown. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Local isort 7.0.0 wraps single-name imports in parentheses; the version CI pins keeps them inline. Run pre-commit rather than bare black/isort so the hook versions match. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Lazarus-931
marked this pull request as ready for review
August 10, 2026 13:32
Lazarus-931
requested review from
Blaizzy and
lucasnewman
and removed request for
lucasnewman
August 10, 2026 13:38
lucasnewman
reviewed
Aug 10, 2026
| @@ -0,0 +1,15 @@ | |||
| # Copyright © 2023-2024 Apple Inc. | |||
| # Vendored from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), | |||
Collaborator
There was a problem hiding this comment.
I don't think the commit hashes and implementation notes are particularly useful here, unless you think so?
lucasnewman
reviewed
Aug 10, 2026
| from mlx_lm.utils import load as load_llm | ||
| except ImportError as exc: | ||
| raise ImportError( | ||
| "The in-process LLM responder needs mlx-lm: pip install 'mlx-audio[llm]'. " |
Collaborator
There was a problem hiding this comment.
Can we make this a little bit more straightforward for the user? "Please run pip install -U mlx-lm to run this script" or something like that.
lucasnewman
approved these changes
Aug 10, 2026
lucasnewman
left a comment
lucasnewman
left a comment
Collaborator
There was a problem hiding this comment.
@Lazarus-931 Overall this looks great. I had a couple of minor comments inline.
Signed-off-by: Alazer Manakelew <alazermanakelewb@gmail.com>
Signed-off-by: Alazer Manakelew <alazermanakelewb@gmail.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
2 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.
Moves the transformer code mlx-audio's speech models rely on (KV caches, masks, samplers, the decode loop, and nine backbones) out of mlx-lm and into
mlx_audio/lm/, so mlx-lm is no longer a core dependency. Files are copied from mlx-lm v0.31.3 with a provenance header, laid out in the same directory shape so re-syncing stays easy. The one exception is the S2S pipeline's chat model, which loads an arbitrary LLM — that still uses mlx-lm, lazily, behind an optional[llm]extra.Verification: 1,531 tests pass, 1,237 with mlx-lm uninstalled entirely. 286 differential tests assert bitwise equality against mlx-lm 0.31.3. Found and fixed a few bugs on the way, including one where
stream_generatere-emitted the previous token instead of EOS, which shifted codec framing at the end of TTS output.