fix: detect MoE across all expert-count config spellings (#112)

[solderzzc]

Fixes #112.

Root cause

ModelProfiler decoded the routed-expert count from a single key:

case numExperts = "num_local_experts"   // ModelProfiler.swift:192
let isMoE = config.numExperts != nil && (config.numExperts ?? 0) > 1

num_local_experts is the Mixtral spelling. Qwen3/Qwen3.5 MoE use num_experts; DeepSeek V2/V3, GLM MoE and MiniMax use n_routed_experts — both already known elsewhere in this repo (DeepseekV3DFlash.swift:59, KimiLinearDFlash.swift:83).

So profile.isMoE was false for Qwen3.5-397B-A17B and Server.swift:347 silently dropped --stream-experts. The reporter's own log confirms it: qwen3_5_moe is not MoE.

The zsh: killed follows directly — with streaming off, modelConfig.lazyLoad is never set (Server.swift:355), so the full 397B model is materialised and the OS OOM-kills the process. That also explains why b648 (before the guard landed) worked.

Changes

• ModelProfiler.swift — decode all three spellings, plus the same keys nested under text_config for multimodal wrappers, resolved in ModelConfig.numExperts.
• modelTypeImpliesMoE() — last-resort fallback for families whose expert key we haven't seen (any model_type containing moe, plus mixtral/dbrx/grok). The asymmetry is deliberate: a false positive costs a weight-file walk, a false negative costs an OOM kill.
• The > 1 threshold is preserved, so a single-expert (dense FFN) config stays dense.
• Server.swift:347 — the rejection message now names the keys it looked for, so the next misdetection is diagnosable from the log alone.

Tests

New tests/SwiftLMTests/ModelProfilerMoEDetectionTests.swift — 9 tests writing real config.json fixtures through ModelProfiler.profile(): one per spelling (the Qwen3.5 case reproduces the issue's exact model), nested text_config, the unknown-key fallback, and negatives for dense / single-expert / missing config.

Verified red against the pre-fix logic — 7 of the 9 fail, including testQwen35MoEDetectedViaNumExperts. Green after: 9/9.

Full SwiftLMTests run: 86 tests across 8 suites, 0 failures. PromptCacheTests aborts under swift test with Failed to load the default metallib — reproduced at HEAD without these changes, so it is pre-existing and environmental.

Not covered here: the reporter should confirm the actual --stream-experts run on their 397B model, since the OOM path itself is hardware-dependent.

🤖 Generated with Claude Code

[Gwada]

Runtime validation on a MacBook Pro M4 Pro with 24 GB unified memory:

• Model: lmstudio-community/Qwen3.6-35B-A3B-MLX-8bit
• Before the fix, SwiftLM logged qwen3_5_moe is not MoE, ignored --stream-experts, and profiled the model at about 3 GB.
• With nested text_config expert detection applied, SwiftLM recognized the 37.7 GB MoE model, enabled SSD expert streaming, loaded all 40 layers, and started the OpenAI-compatible server successfully.
• POST /v1/chat/completions returned HTTP 200; warm generation measured about 4.04 tok/s.
• Targeted ModelProfiler regression tests passed: 2 tests, 0 failures.

Validated command:

.build/release/SwiftLM \
  --model lmstudio-community/Qwen3.6-35B-A3B-MLX-8bit \
  --stream-experts \
  --turbo-kv \
  --vision \
  --port 5413

[solderzzc]

Thanks @Gwada — that's a more demanding case than the one in the issue, and I've now pinned it in a test (95ba329).

I checked your config against the cached Qwen--Qwen3.6-35B-A3B and it matches your report exactly:

model_type:                  qwen3_5_moe          ← hence "qwen3_5_moe is not MoE"
top-level expert keys:       {}                   ← nothing at all
text_config.num_experts:     256
text_config.num_experts_per_tok: 8
text_config.num_hidden_layers:   40               ← your "loaded all 40 layers"

So this model needs both halves of the fix: the num_experts spelling and the nested text_config lookup. Every expert and dimension key lives under text_config because it's multimodal — the top level holds only vision_config and the quantization block.

Worth noting for anyone reading the diff: model_type here contains "moe", so the heuristic fallback alone would have flagged it as MoE while leaving numExperts/numActiveExperts nil — enough to enable --stream-experts, but with wrong parameter estimates and no MoE advisory. The nested lookup is what makes the 256/8 counts correct. The new test asserts the counts rather than just isMoE for that reason, and I verified it fails with the nested lookup removed.

The 37.7 GB model running at ~4 tok/s on 24 GB unified memory is a much better demonstration of the streaming path than anything I could produce locally. Much appreciated.

One clarification, since it affects what your result proves: was this PR's branch as-is, or your own patch? You mention "2 tests" and this branch carries 9 (now 10), so I want to be sure the validated code is what's up for merge here.

[solderzzc]

Reproduced @Gwada's runtime validation independently, on mlx-community/Qwen3.6-35B-A3B-4bit (same architecture and config shape as their 8-bit build — qwen3_5_moe, no top-level expert keys, text_config.num_experts: 256, 40 layers, vision_config present).

Same command, same binary, only ModelProfiler.swift differing between the two runs:

.build/release/SwiftLM --model <local snapshot> --stream-experts --turbo-kv --vision --port 5413

Before (main):

⚠️  Model does not support SSD expert streaming (qwen3_5_moe is not MoE). Ignoring --stream-experts flag.
✅ Memory strategy: FULL GPU (1.5GB model, 64.7GB available)
Config: … ssd_stream=disabled

After (this branch):

Enabled Async SSD Streaming on directory: 38740b84…
✅ Memory strategy: FULL GPU (20.4GB model, 64.7GB available)
💾 SSD Expert Streaming enabled (lazy load + layer-sync)
Config: … ssd_stream=enabled
ready: {"strategy":"ssd_streaming","total_layers":40,"gpu_layers":40,"model_weight_gb":20.4,"ssd_stream":true,"vision":true}

The bogus weight figure reproduces too — 1.5 GB for the 4-bit here against Gwada's ~3 GB for the 8-bit, which is the expected 2× ratio. It comes from measureWeightFiles() being skipped entirely for models profiled as dense, so the misdetection also corrupts the memory-strategy input, not just the streaming flag.

POST /v1/chat/completions → HTTP 200, coherent output, 16.95 tok/s sustained over 113 tokens.

Two caveats on my run, since they bound what it proves: this box has 68.7 GB of RAM, so streaming activates but is never memory-pressured the way Gwada's 24 GB machine exercised it — their throughput number is the meaningful one, not mine. And I could not test the exact 8-bit build: it is 37.7 GB against 39 GiB free here, so I used the 4-bit of the same model rather than fill the disk.

Detection, streaming activation, and generation all confirmed on a second machine.

[solderzzc]

Answering my own question above: #117 makes it clear the validation was of @Gwada's own patch, not this branch — which is exactly why their runtime result was useful, so no concern there. Recording how the two fixes relate, since #117 is now closed and the reasoning shouldn't disappear with it.

Config shape	#117	#114
num_local_experts (Mixtral, Phi-MoE)	✅	✅
text_config.num_experts (Qwen3.6 multimodal MoE)	✅	✅
top-level num_experts (Qwen3.5 text-only MoE)	❌	✅
n_routed_experts (DeepSeek V2/V3, GLM MoE, MiniMax)	❌	✅
model_type fallback for unknown spellings	❌	✅

The third row is the one worth calling out: #117 left config.numExperts bound to num_local_experts and added only the nested lookup, so a flat top-level num_experts still profiled as dense — and that is the config of Qwen3.5-397B-A17B, the model #112 was actually filed about. Two different real models, two different key placements, same symptom.

Credit where it is due: @Gwada found the nested text_config case independently and validated it on hardware I do not have. That report is what prompted me to check the real Qwen3.6-35B-A3B config, confirm every expert and dimension key sits under text_config because the model is multimodal, and pin that exact shape as a test here (95ba329) rather than the synthetic one I had.

For the record on this branch specifically: I ran main and this branch as separate release binaries against mlx-community/Qwen3.6-35B-A3B-4bit — same architecture and config shape as Gwada's 8-bit build. Before: qwen3_5_moe is not MoE, ssd_stream=disabled, 1.5 GB profile. After: 20.4 GB, ssd_streaming, 40 layers, HTTP 200 at 16.95 tok/s. Details in the comment above.

[Gwada]

Thanks for the detailed follow-up and for reproducing the result independently.

I originally came across SwiftLM because I wanted to try the project, and I was genuinely pleasantly surprised — it is a really good tool. I tested with a smaller model because my Mac cannot reasonably handle anything larger.

I only noticed #114 a few minutes after publishing my own PR. Otherwise, I would have added my findings here directly.

I will keep following the project, and I am really looking forward to future improvements, especially further performance gains. Thanks again for the work!