Mixed per-layer quantization (MTP head bits override) crashes quantized_matmul with --stream-experts on qwen3_5_moe checkpoints

[bauntery]

Summary

Loading a qwen3_5_moe checkpoint that uses per-layer mixed-precision
quantization (a standard MLX quantization/quantization_config dict with
per-tensor bits overrides, including a separately-quantized MTP head)
crashes on the very first inference request with a quantized_matmul
shape-mismatch fatal error, when run with --stream-experts --mtp.

Environment

• MacBook Pro, Apple M5 Pro, 24 GB unified memory, macOS 26.6.1
• SwiftLM prebuilt release b698 (macos-arm64 tarball)
• Model: a Qwen3.6-35B-A3B-derived MoE checkpoint (hybrid linear-attention /
  full-attention layers, model_type: qwen3_5_moe,
  architectures: ["Qwen3_5MoeForConditionalGeneration"]), MLX base
  quantization with per-layer overrides (4/5/6-bit on the main model, uniform
  8-bit on the MTP head). ~22.6 GB weight files, 40 layers, 8/256 experts
  active per layer.

Repro

./SwiftLM --model /path/to/model --stream-experts --mtp --port 5413
# server starts fine, logs:
#   💾 Memory strategy: SSD STREAMING (page-cache managed, 16GB RAM budget, no swap)
#   ...mtp=enabled (3 tokens/round)
#   ✅ Ready. Listening on http://127.0.0.1:5413

curl -s http://127.0.0.1:5413/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d {"model":"m","messages":[{"role":"user","content":"Say hello."}],"max_tokens":64}

Observed crash

srv  slot_launch: id 0 | prompt=19t | thinking=false | prefilling...
MLX/ErrorHandler.swift:345: Fatal error: [quantized_matmul] The shapes of the weight and scales are incompatible based on bits and group_size. w.shape() == (2048,1024) and scales.shape() == (2048,64) with group_size=64 and bits=4 at .../mlx-swift/Source/Cmlx/mlx-c/mlx/c/ops.cpp:2567

Root cause (confirmed) and fix

Traced this to mlx-swift-lm's Libraries/MLXLMCommon/Load.swift: the
Swift module tree indexes MTP prediction layers as mtp.<depth>.layers....
(e.g. mtp.0.layers....), but checkpoints declare their per-layer
quantization overrides keyed as mtp.layers.... (no depth index). The two
forms never matched in the per-layer quantization lookup, so every MTP-head
module silently fell back to the top-level default quantization instead of
its own (in this case 8-bit) override — causing the shape mismatch above.

Opened a fix with a normalized .mtp.<N>. -> .mtp. lookup candidate,
mirroring an identical normalization that already exists a few lines below
in the same file for weight-key remapping. Verified end-to-end against the
checkpoint described above: no crash, real completions generated via
--stream-experts --mtp.

PR: SharpAI/mlx-swift-lm#56

Impact

Any qwen3_5_moe-family checkpoint using mixed-precision per-layer
quantization together with a separately-quantized MTP head (a fairly
standard mlx_lm.quantize output pattern, not a proprietary format)
currently cannot be served with --stream-experts --mtp.

[solderzzc]

Thanks for the detailed repro and root-cause trace — really appreciate you pinning it down to the exact normalization gap in Load.swift and opening the fix. Reviewing mlx-swift-lm#56 now; will link back here once it's merged and the submodule is bumped.

[solderzzc]

Fix merged: mlx-swift-lm#56. Submodule bump PR up here: #154 — will close this issue once that's merged and CI is green.

[solderzzc]

Resolved by SwiftLM#154, which bumped the mlx-swift-lm submodule to pick up mlx-swift-lm#56 — thanks again @bauntery for the clean root-cause and fix. Also added regression coverage for this exact bug shape in mlx-swift-lm#57, since nothing previously exercised a populated MTP head's quantization override lookup.