From 6db32f16861de19c8e1b7d03d7f32cb6dbd50395 Mon Sep 17 00:00:00 2001 From: Mark Smeltzer Date: Wed, 9 Sep 2026 22:13:12 -0700 Subject: [PATCH] bugfix: Support Qwen3.8 processor and cache Add Q8/BF16 Metal cards, registration coverage, component regressions and public request replay instructions. Carry the ArraysCache.advance dependency correction through Nix without changing dependency pins. Adapt the cache metadata dependency mechanism from Pierre Lamy's mlx-lm commit bb615ebdb5aff33eb931ac0627cae142bd7adbaa (advance hunk only). Implemented with GitHub Copilot under Mark Smeltzer's direction; human code review remains pending. Validate on macOS: Nix runtime/dashboard/test builds, Q8/BF16 real loaders, cache regression, 471 native tests, strict typecheck, Ruff and nix fmt. Historical four-node qualification is not a fresh validation of this upstream port. --- nix/arrays-cache-metadata.patch | 18 +++ nix/tests/README.qwen38.md | 128 ++++++++++++++++++ nix/tests/arrays-cache-metadata-regression.py | 30 ++++ nix/tests/qwen-requests.py | 103 ++++++++++++++ nix/tests/qwen-vision-regression.py | 14 ++ python/parts.nix | 3 + .../mlx-community--Qwen3.8-27B-8bit.toml | 15 ++ .../mlx-community--Qwen3.8-27B-bf16.toml | 15 ++ .../shared/tests/test_qwen38_model_cards.py | 31 +++++ src/exo/worker/engines/mlx/vision.py | 9 +- 10 files changed, 364 insertions(+), 2 deletions(-) create mode 100644 nix/arrays-cache-metadata.patch create mode 100644 nix/tests/README.qwen38.md create mode 100644 nix/tests/arrays-cache-metadata-regression.py create mode 100644 nix/tests/qwen-requests.py create mode 100644 nix/tests/qwen-vision-regression.py create mode 100644 resources/inference_model_cards/mlx-community--Qwen3.8-27B-8bit.toml create mode 100644 resources/inference_model_cards/mlx-community--Qwen3.8-27B-bf16.toml create mode 100644 src/exo/shared/tests/test_qwen38_model_cards.py diff --git a/nix/arrays-cache-metadata.patch b/nix/arrays-cache-metadata.patch new file mode 100644 index 0000000000..87ebde2fae --- /dev/null +++ b/nix/arrays-cache-metadata.patch @@ -0,0 +1,18 @@ +diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py +--- a/mlx_lm/models/cache.py ++++ b/mlx_lm/models/cache.py +@@ -689,6 +689,14 @@ + self.lengths -= N + if self.left_padding is not None: + self.left_padding -= N ++ metadata = tuple( ++ value for value in (self.lengths, self.left_padding) if value is not None ++ ) ++ if metadata: ++ for index, value in enumerate(self.cache): ++ if value is not None: ++ self.cache[index] = mx.depends(value, metadata) ++ break + + def make_mask(self, N: int): + if self.left_padding is not None: \ No newline at end of file diff --git a/nix/tests/README.qwen38.md b/nix/tests/README.qwen38.md new file mode 100644 index 0000000000..cc663dfbf3 --- /dev/null +++ b/nix/tests/README.qwen38.md @@ -0,0 +1,128 @@ +# Qwen3.8-27B Reproduction + +This contribution targets Q8 and BF16 text generation on Apple Silicon. The cards +do not advertise vision, CUDA or CPU qualification. No weights are included. +The processor and cache corrections address separate failures; the Nix override +applies the cache correction to EXO's actual MLX-LM dependency. `uv run exo` does +not apply this Nix-only patch. Dependency pins are unchanged in `uv.lock`. + +The cache patch adapts only `ArraysCache.advance()` from Pierre Lamy's +[bb615eb patch](https://github.com/pierre427/mlx-lm/commit/bb615ebdb5aff33eb931ac0627cae142bd7adbaa), +not its `extract()` change. Related upstream work includes +[MLX-LM #1632](https://github.com/ml-explore/mlx-lm/pull/1632) and +[#1845](https://github.com/ml-explore/mlx-lm/issues/1845). A companion MLX-LM PR +has not yet been published. These references do not imply endorsement. + +## Component Checks + +Run from this checkout on macOS with EXO's documented Nix prerequisites. +Use Python without `-O` or `PYTHONOPTIMIZE`; the scripts include assertions. + +```bash +set -e +EXO_TEST_PYTHON_ENV="$(nix build --no-link --print-out-paths path:.#exo.venv)" +EXO_TEST_DASHBOARD="$(nix build --no-link --print-out-paths path:.#dashboard)" +"$EXO_TEST_PYTHON_ENV/bin/python" nix/tests/arrays-cache-metadata-regression.py +``` + +Expected: `arrays-cache metadata regression: PASS`. No weights or EXO service +are needed. This checks metadata graph bounds after 256 cache advances and the +final numerical values, not distributed inference or long-run resource usage. + +Download the complete model repositories at these revisions using Hugging Face: + +| Model | Revision | +| --- | --- | +| `mlx-community/Qwen3.8-27B-8bit` | `815b83c0df8ffd1d1b5244cf75fd6ef14fca9ef9` | +| `mlx-community/Qwen3.8-27B-bf16` | `6f265714824f3c38d4452baa1628aef3d9b9aae9` | + +Set `MODEL_ROOT` to a directory containing the complete downloads named +`mlx-community--Qwen3.8-27B-8bit` and `mlx-community--Qwen3.8-27B-bf16`. + +```bash +MODEL_ROOT="/absolute/path/to/local-models" +( + set -e + export HOME="$(mktemp -d)" + unset EXO_HOME EXO_MODELS_READ_ONLY_DIRS PYTHONPATH PYTHONOPTIMIZE + export EXO_DEFAULT_MODELS_DIR="$HOME/models" + export EXO_MODELS_DIRS="$MODEL_ROOT" + export EXO_RESOURCES_DIR="$PWD/resources" + export EXO_DASHBOARD_DIR="$EXO_TEST_DASHBOARD" + "$EXO_TEST_PYTHON_ENV/bin/python" nix/tests/qwen-vision-regression.py mlx-community/Qwen3.8-27B-8bit + "$EXO_TEST_PYTHON_ENV/bin/python" nix/tests/qwen-vision-regression.py mlx-community/Qwen3.8-27B-bf16 +) +``` + +Each invocation must print `Qwen vision loader regression: PASS (MODEL_ID)`. +It loads real vision weights but does not perform image understanding. The +temporary home is retained for inspection; no services or model files are changed. + +Registration and native source tests, following the macOS CI approach: + +```bash +TEST_ENV="$(nix build --no-link --print-out-paths path:.#exo-test-env)" +( + export HOME="$(mktemp -d)" + unset EXO_HOME + export PYTHONPATH="$PWD/src" + export EXO_RESOURCES_DIR="$PWD/resources" + export EXO_DASHBOARD_DIR="$PWD/dashboard" + "$TEST_ENV/bin/python" -m pytest src -m "not slow" --import-mode=importlib +) +"$TEST_ENV/bin/basedpyright" --pythonpath "$TEST_ENV/bin/python" +"$TEST_ENV/bin/ruff" check +nix fmt +``` + +## Distributed Request Replay + +This is a manual, opt-in workload, not part of pytest. Build `nix build path:.#exo` +and run that build on all four nodes using EXO's documented JACCL/RDMA setup. +Do not mix builds. Place one variant at a time. Inspect `/state`: require one +matching `MlxJacclInstance`, four Tensor shards with `worldSize` four, and all +four assigned runners in `RunnerReady`. Recheck before and after each request. +Neither the generator nor curl places models, restarts services or verifies this +topology for you. Other serving modes are not equivalent to this qualification. + +```bash +MODEL="mlx-community/Qwen3.8-27B-8bit" +EXO_URL="http://localhost:52415" +"$EXO_TEST_PYTHON_ENV/bin/python" nix/tests/qwen-requests.py "$MODEL" sustained > sustained.json +curl --fail-with-body --max-time 1800 -N -H 'Content-Type: application/json' \ + --data-binary @sustained.json "$EXO_URL/v1/chat/completions" > sustained.sse +"$EXO_TEST_PYTHON_ENV/bin/python" nix/tests/qwen-requests.py "$MODEL" recall > recall.json +curl --fail-with-body --max-time 1800 -N -H 'Content-Type: application/json' \ + --data-binary @recall.json "$EXO_URL/v1/chat/completions" > recall.sse +``` + +For the cached extension, concatenate `choices[].delta.content` from the recall +SSE JSON messages, preserving the answer exactly, into `recall-answer.txt`. +Do not include reasoning, SSE framing, or an added newline. Then: + +```bash +"$EXO_TEST_PYTHON_ENV/bin/python" nix/tests/qwen-requests.py "$MODEL" generation \ + --previous-answer recall-answer.txt > generation.json +curl --fail-with-body --max-time 1800 -N -H 'Content-Type: application/json' \ + --data-binary @generation.json "$EXO_URL/v1/chat/completions" > generation.sse +``` + +Repeat separately with the BF16 ID. HTTP success alone is insufficient: inspect +SSE for errors, usage, nonempty output, `finish_reason: stop` and `[DONE]`. Recall +must return `copper-41, harbor-72, spruce-93`; generation must begin with the +updated middle value `anchor-84`. Check actual prompt/cached token counts, not +the nominal workload name. The original requests produced 210,668 prompt tokens +for recall and 243,192 for generation with 210,666 cached; token counts can change +with dependencies or answers. A short output does not test sustained decoding. +If any rank fails or curl times out, stop the workload and inspect the server; +client timeout does not prove the server stopped generation. Do not treat errors, +truncation or a missing marker as passing. Generated C# has not been compiled. + +## Evidence Limits + +On September 9, 2026, a development build with both corrections completed Q8 +(17,992 output tokens) and BF16 (19,211 output tokens) sustained workloads and +the cached 243k conversations on four M3 Ultra nodes. These are historical +compatibility observations, not a four-node validation of this upstream port. +The public-source component checks ran on macOS 26.6.1 (25G76). Full-window cold +prefill, concurrency, overnight stability and multimodal inference remain unqualified. \ No newline at end of file diff --git a/nix/tests/arrays-cache-metadata-regression.py b/nix/tests/arrays-cache-metadata-regression.py new file mode 100644 index 0000000000..7681642f0d --- /dev/null +++ b/nix/tests/arrays-cache-metadata-regression.py @@ -0,0 +1,30 @@ +import pathlib +import tempfile + +import mlx.core as mx +from mlx_lm.models.cache import ArraysCache + +cache = ArraysCache(2, left_padding=[2]) +cache.prepare(lengths=[3]) +cache[0] = mx.array([0]) +cache[1] = mx.array([0]) +for _step in range(256): + cache[0] = cache[0] + 1 + cache.advance(1) + mx.eval(cache[0]) + +with tempfile.TemporaryDirectory() as directory: + for name, metadata in ( + ("lengths", cache.lengths), + ("left-padding", cache.left_padding), + ): + path = pathlib.Path(directory, f"arrays-cache-{name}.dot") + mx.export_to_dot(str(path), metadata) + edges = path.read_text(encoding="utf-8").count("->") + if edges > 8: + raise AssertionError(f"{name} graph has {edges} edges, expected <= 8") + +assert cache[0].item() == 256 +assert cache.lengths.item() == 3 - 256 +assert cache.left_padding.item() == 2 - 256 +print("arrays-cache metadata regression: PASS") diff --git a/nix/tests/qwen-requests.py b/nix/tests/qwen-requests.py new file mode 100644 index 0000000000..92b8f14e44 --- /dev/null +++ b/nix/tests/qwen-requests.py @@ -0,0 +1,103 @@ +import argparse +import json +from pathlib import Path + + +def recall_prompt() -> str: + records = [ + f"Record {index:04d}: depot inventory has 17 units reserved and 23 units " + "available for the next delivery." + for index in range(1, 7801) + ] + records[0] += " START_CODE=copper-41." + records[3899] += " MIDDLE_CODE=harbor-72." + records[7799] += " END_CODE=spruce-93." + return ( + "Read the following records. Return exactly the START_CODE, MIDDLE_CODE " + "and END_CODE values, in that order, comma separated, without explanation.\n" + + "\n".join(records) + ) + + +def generation_prompt() -> str: + records = [ + f"Record {index:04d}: depot inventory has 17 units reserved and 23 units " + "available for the next delivery." + for index in range(7801, 9001) + ] + return ( + "Append these records to the previous ledger.\n" + + "\n".join(records) + + "\nCorrection: MIDDLE_CODE is now anchor-84. The other two codes are " + "unchanged. Begin your answer with the three current code values in order. " + "Then write a complete C# streaming ledger parser with a runnable example, " + "bounded memory, cancellation, validation of record IDs and numeric " + "quantities, and aggregation of reserved/available totals. Include focused " + "tests for malformed records and cancellation. Work from the entire " + "conversation, not just this last message." + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Emit the Qwen qualification requests." + ) + parser.add_argument( + "model", + choices=["mlx-community/Qwen3.8-27B-8bit", "mlx-community/Qwen3.8-27B-bf16"], + ) + parser.add_argument("scenario", choices=["sustained", "recall", "generation"]) + parser.add_argument("--previous-answer", type=Path) + arguments = parser.parse_args() + if arguments.scenario == "generation" and arguments.previous_answer is None: + parser.error( + "generation requires --previous-answer with the exact recall answer" + ) + if arguments.scenario == "sustained": + messages = [ + {"role": "user", "content": "ping"}, + {"role": "assistant", "content": "pong"}, + { + "role": "user", + "content": "Write a complete C# implementation of a high-throughput, " + "reliable UDP protocol that uses NACK-based retransmission. Include " + "packet framing, sequence tracking, loss detection, retransmission " + "queues, bounded memory, cancellation, and a runnable example. " + "Include every source file and project file without ellipses, then " + "a complete deterministic packet-loss/reordering simulator with at " + "least 12 explicitly implemented validation scenarios, and explain " + "the invariants each scenario checks.", + }, + ] + else: + messages = [{"role": "user", "content": recall_prompt()}] + if arguments.scenario == "generation": + messages.extend( + [ + { + "role": "assistant", + "content": arguments.previous_answer.read_text( + encoding="utf-8" + ), + }, + {"role": "user", "content": generation_prompt()}, + ] + ) + request = { + "model": arguments.model, + "stream": True, + "temperature": 0.7 if arguments.scenario == "sustained" else 0, + "enable_thinking": arguments.scenario != "recall", + "logprobs": True, + "top_logprobs": 5, + "messages": messages, + } + if arguments.scenario == "sustained": + request["reasoning_effort"] = "medium" + else: + request["max_tokens"] = 1024 if arguments.scenario == "recall" else 8192 + print(json.dumps(request, ensure_ascii=True)) + + +if __name__ == "__main__": + main() diff --git a/nix/tests/qwen-vision-regression.py b/nix/tests/qwen-vision-regression.py new file mode 100644 index 0000000000..86c74a02d2 --- /dev/null +++ b/nix/tests/qwen-vision-regression.py @@ -0,0 +1,14 @@ +import sys + +from exo.shared.models.model_cards import detect_vision_from_config +from exo.shared.types.common import ModelId +from exo.worker.engines.mlx.vision import VisionEncoder + +assert len(sys.argv) == 2, "Usage: qwen-vision-regression.py MODEL_ID" +model_id = ModelId(sys.argv[1]) +vision_config = detect_vision_from_config(model_id) +assert vision_config is not None, "Expected local model vision metadata" +assert vision_config.model_type == "qwen3_5", "Expected the Qwen3.8 architecture" +encoder = VisionEncoder(vision_config, model_id) +encoder.ensure_loaded() +print(f"Qwen vision loader regression: PASS ({model_id})") diff --git a/python/parts.nix b/python/parts.nix index 26318c4a28..7c1452f23a 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -44,6 +44,9 @@ let paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]); }; exoOverlay = final: prev: { + mlx-lm = prev.mlx-lm.overrideAttrs (old: { + patches = (old.patches or [ ]) ++ [ ../nix/arrays-cache-metadata.patch ]; + }); # Replace workspace exo_rs with Nix-built wheel. # Preserve passthru so mkVirtualEnv can resolve dependency groups. # Copy .pyi stub + py.typed marker so basedpyright can find the types. diff --git a/resources/inference_model_cards/mlx-community--Qwen3.8-27B-8bit.toml b/resources/inference_model_cards/mlx-community--Qwen3.8-27B-8bit.toml new file mode 100644 index 0000000000..0e5e42cc04 --- /dev/null +++ b/resources/inference_model_cards/mlx-community--Qwen3.8-27B-8bit.toml @@ -0,0 +1,15 @@ +model_id = "mlx-community/Qwen3.8-27B-8bit" +n_layers = 64 +hidden_size = 5120 +num_key_value_heads = 4 +supports_tensor = true +tasks = ["TextGeneration"] +family = "qwen" +quantization = "8bit" +base_model = "Qwen3.8 27B" +capabilities = ["text", "thinking", "thinking_toggle"] +reasoning_dialect = "post_last_user" +context_length = 262144 +backends = ["MlxMetal"] +[storage_size] +in_bytes = 29500938720 diff --git a/resources/inference_model_cards/mlx-community--Qwen3.8-27B-bf16.toml b/resources/inference_model_cards/mlx-community--Qwen3.8-27B-bf16.toml new file mode 100644 index 0000000000..2ac38245b3 --- /dev/null +++ b/resources/inference_model_cards/mlx-community--Qwen3.8-27B-bf16.toml @@ -0,0 +1,15 @@ +model_id = "mlx-community/Qwen3.8-27B-bf16" +n_layers = 64 +hidden_size = 5120 +num_key_value_heads = 4 +supports_tensor = true +tasks = ["TextGeneration"] +family = "qwen" +quantization = "bf16" +base_model = "Qwen3.8 27B" +capabilities = ["text", "thinking", "thinking_toggle"] +reasoning_dialect = "post_last_user" +context_length = 262144 +backends = ["MlxMetal"] +[storage_size] +in_bytes = 54713457120 diff --git a/src/exo/shared/tests/test_qwen38_model_cards.py b/src/exo/shared/tests/test_qwen38_model_cards.py new file mode 100644 index 0000000000..b273a8ede7 --- /dev/null +++ b/src/exo/shared/tests/test_qwen38_model_cards.py @@ -0,0 +1,31 @@ +import pytest + +from exo.shared.models.model_cards import ModelTask, card_cache +from exo.shared.types.backends import Backend +from exo.shared.types.common import ModelId + + +@pytest.mark.parametrize( + ("quantization", "storage_bytes"), + [("8bit", 29500938720), ("bf16", 54713457120)], +) +async def test_qwen38_builtin_registration( + quantization: str, storage_bytes: int +) -> None: + cache = type(card_cache)() + await cache.refresh() + model_id = ModelId(f"mlx-community/Qwen3.8-27B-{quantization}") + card = cache.get(model_id) + assert card is not None + assert not card.is_custom + assert card.n_layers == 64 + assert card.hidden_size == 5120 + assert card.num_key_value_heads == 4 + assert card.context_length == 262144 + assert card.supports_tensor + assert card.tasks == [ModelTask.TextGeneration] + assert card.backends == [Backend.MlxMetal] + assert card.quantization == quantization + assert card.storage_size.in_bytes == storage_bytes + assert card.reasoning_dialect == "post_last_user" + assert card.capabilities == ["text", "thinking", "thinking_toggle"] diff --git a/src/exo/worker/engines/mlx/vision.py b/src/exo/worker/engines/mlx/vision.py index a79925bc75..2685b47c54 100644 --- a/src/exo/worker/engines/mlx/vision.py +++ b/src/exo/worker/engines/mlx/vision.py @@ -20,7 +20,7 @@ from mlx_vlm.utils import load_image_processor from PIL import Image from safetensors import safe_open -from transformers import AutoImageProcessor +from transformers import AutoConfig, AutoImageProcessor, PreTrainedConfig from exo.download.download_utils import build_model_path from exo.shared.models.model_cards import VisionCardConfig @@ -349,8 +349,13 @@ def _load_weights(self) -> None: if image_proc is not None: self._processor = image_proc else: + load_config = cast( + Callable[..., PreTrainedConfig], AutoConfig.from_pretrained + ) self._processor = AutoImageProcessor.from_pretrained( # type: ignore - repo, trust_remote_code=True + repo, + config=load_config(repo, trust_remote_code=True), + trust_remote_code=True, ) if processor_repo: self._merge_kernel_size = vision_cfg.get("merge_kernel_size", [2, 2]) # type: ignore