fix(models): flatten mistral4 MoE tokens to 2D before SwitchGLU routing (#425)

[inureyes]

Summary

Mistral4MoE::forward passed 3D tensors to the shared SwitchGLU::forward, which expects 2D inputs. This violated the SwitchGLU shape contract and aborted every real mistral4 prompt, every chunked-prefill continuation chunk, and any image-bearing multimodal prefill. The fix flattens tokens to [n_tokens, hidden] before routing and reshapes the result back, mirroring qwen3_moe's proven pattern.

The bug: 3D-vs-2D SwitchGLU contract violation

SwitchGLU::forward (src/models/switch_layers.rs:324) expects 2D inputs: x is [n_tokens, hidden] and the routing indices are [n_tokens, top_k]. It reads indices_shape[0] as n_tokens and indices_shape[1] as top_k, then decides the sort path with do_sort = n_tokens * top_k >= 64.

Mistral4MoE::forward instead passed a 3D x = [batch, seq, hidden] and 3D inds = [batch, seq, top_k]. With [1, seq, top_k] indices, SwitchGLU misread n_tokens = 1 and top_k = seq:

• For seq < 64, do_sort = false limped along on the wrong shapes.
• For seq >= 64 (any real prompt, a 512-token prefill chunk, or an image-bearing prefill), do_sort = true entered gather_sort with the wrong dims and aborted with [reshape] Cannot reshape array of size seq*hidden into (1,1,1).

This is the abort tracked in #425. It blocked:

• All real text prompts on mlx-community/Mistral-Small-4-119B-2603-4bit.
• Chunked-prefill continuation under --max-kv-size 1024 --prefill-chunk-size 512 (reshape ... of size 2097152 = 512 x 4096 into (1,1,1)).
• Image+text multimodal prefill (reshape ... of size 12480512 = 3047 x 4096 into (1,1,1)).

It stayed latent because mistral4 was only ever run on a single-token warmup before #424 made the family loadable; the warmup's [1, 1, hidden] MoE forward (n_tokens=1, top_k=1, total < 64) happens to succeed.

The fix: qwen3_moe-mirrored flatten / reshape-back

Following SparseMoeBlock::forward in src/models/qwen3_moe.rs:

1. Capture orig_shape and hidden_dim.
2. Flatten x to [n_tokens, hidden] (n_tokens = product of leading dims): prefill [1, l, hidden] -> [l, hidden], decode [1, 1, hidden] -> [1, hidden]. The <= 2D path copies through unchanged.
3. Run the existing routing math (softmax, argpartition, slice_axis, take_along_axis, the norm_topk_prob block with the 1e-20 eps, and the routed_scaling_factor multiply) on the flat tensor, so indices and scores come out correctly shaped [n_tokens, top_k]. Ops and eps are unchanged, so per-token routing is numerically identical, just correctly shaped.
4. Dispatch experts and shared experts on the flat tensor.
5. Reshape the result back to the original rank.

SwitchGLU, moe_weighted_sum, and all other models are untouched; only Mistral4MoE::forward changes. No fused single-token kernel path was added (correctness-only; a fused mistral4 MoE optimization can be a separate change).

This also corrects the decode path, which the old 3D shape was misrouting. Decode output may differ slightly from the pre-fix buggy path; that is the intended correction, not a regression.

What changed

• src/models/mistral4.rs: rewrote Mistral4MoE::forward to flatten tokens to 2D before SwitchGLU routing and reshape the result back, with a comment documenting the 2D SwitchGLU contract and the >= 64-token do_sort crash.

Test plan

A Mistral4MoE::forward unit test needs real expert weights (a SwitchGLU built from SwitchLinear per-expert weights), which is not CI-feasible, and there is no existing checkpoint-free MoE shape-test pattern in the tree to mirror, so no synthetic test was forced. Real-model 119B validation (chunked-prefill text plus image+text) is performed by the orchestrator.

• cargo fmt --check
• cargo check --lib --features metal,accelerate
• cargo clippy --lib --tests --features metal,accelerate -- -D warnings

Closes #425

[inureyes]

Real-model validation (Mistral-Small-4-119B-2603-4bit) — mistral4 now runs normally

Before this fix, every mistral4 forward with >= 64 tokens aborted ([reshape] ... size seq*hidden into (1,1,1)): the MoE passed 3D [batch, seq, hidden] / [batch, seq, top_k] to SwitchGLU::forward, which expects 2D [n_tokens, hidden] / [n_tokens, top_k] and takes the gather_sort path once n_tokens*top_k >= 64. After flattening to 2D before routing (mirroring qwen3_moe) and reshaping back:

• Text, chunked prefill (--prefill-chunk-size 512, a ~600-token prompt whose first chunk is 512 >= 64, i.e. the exact crash repro): HTTP 200, coherent: "The main theme is the incremental evolution of computing technology through successive innovations." (concise, correct top_k routing).
• Image + text (the original fix(models): mistral4-VLM image+text path aborts in multimodal prefill (reshape size 12480512 into (1,1,1)) #425 multimodal symptom, /v1/chat/completions with an image): HTTP 200, accurate vision: "The image is a bar chart comparing the performance of various internal models, categorized by size (Small, Medium, Large) and type (Instruct and Reasoning), across different text and vision benchmarks..." (correct description of the benchmark chart).
• Trim (--max-kv-size 512 --prefill-chunk-size 256, moderate trim): HTTP 200, coherent: "The main theme is the incremental and foundational nature of computing history." This exercises the fix(models): complete offset->live_len prefill-mask migration for dense-KVCache models under --max-kv-size (mistral4/nemotron_nas/qwen*-vl) #421 live_len mask and the fix(models): mistral4-VLM image+text path aborts in multimodal prefill (reshape size 12480512 into (1,1,1)) #425 MoE flatten together with no abort. (An aggressive --max-kv-size 1024 on a 2632-token prompt returns empty output, i.e. immediate EOS from losing ~60% of context on a full-attention model, a truncation artifact, not a crash.)

No reshape/broadcast/abort signals in any run. This completes the mistral4 bring-up chain (#423 load routing -> embed_tokens -> #425 MoE flatten): mistral4 (Mistral Small 4) now loads and runs text and image+text inference.

[inureyes]

Security and performance review

Reviewed src/models/mistral4.rs (the only changed file, +39/-8) for security/DoS and performance. No new findings; auto-fix made no changes.

Security / DoS (primary value): resolved, no new surface

• The change eliminates a remote-client DoS. The pre-fix path passed 3D [batch, seq, hidden] tensors and 3D [batch, seq, top_k] indices to the 2D-only SwitchGLU::forward, which misreads n_tokens=batch, top_k=seq. For seq >= 64 (do_sort = n_tokens*top_k >= 64) it entered gather_sort with wrong dims and aborted via an uncaught C++ std::invalid_argument (reshape ... into (1,1,1)). That throw crosses the cxx boundary on a non-Result FFI, so it is uncatchable by the worker catch_unwind and terminates the whole server process. Any prompt >= 64 tokens (every real prompt, a 512-token prefill chunk, an image-bearing prefill) triggered it.
• The reshape sizes are derived solely from the model's own hidden-state tensor shape (orig_shape), not from request content beyond sequence length, which is now handled correctly. No new untrusted-input path.
• No unsafe, no unwrap/expect, no new panic surface introduced (diff scanned). The orig_shape[len-1] indexing is inherited verbatim from the proven qwen3_moe pattern and is unreachable for sub-2D inputs (the MoE always receives the 3D residual stream from post_attention_layernorm).

Performance: no new hot-path cost

• The MoE forward runs once per layer per forward, unchanged. The added work is shape-only: reshape([1, l, hidden] -> [l, hidden]) and the reshape-back are O(1) metadata views on the contiguous residual stream (dropping/adding a leading-1 dim), no data copy.
• Decode ([1, 1, hidden]) takes the same reshape path as qwen3_moe, no extra data copy. The <= 2D copy(x) branch is unreachable for mistral4 (inputs are always 3D) and exists only to mirror qwen3_moe.

Correctness-adjacent: reshape-back is size-preserving

• n = product of leading dims = batch*seq. SwitchGLU + moe_weighted_sum return [n_tokens, hidden]; reshape to orig_shape holds because n_tokens*hidden == batch*seq*hidden. Shared-expert output and moe_weighted_sum both yield [n_tokens, hidden], so the final residual add(&h, &mlp_out) is shape-correct. Routing ops and the 1e-20 eps are unchanged, so per-token routing is numerically identical, just correctly shaped. The implementation matches SparseMoeBlock::forward in src/models/qwen3_moe.rs line for line.

Verification

• cargo check --lib --features metal,accelerate: clean.
• cargo clippy --lib --tests --features metal,accelerate -- -D warnings: clean.

No CRITICAL/HIGH/MEDIUM/LOW findings to escalate. The PR removes a CRITICAL-severity process-abort DoS without introducing new risk.

[inureyes]

PR Finalization

Checks run

• cargo check --lib: pass
• cargo clippy --lib --tests -- -D warnings: pass, no findings
• cargo fmt --check: pass

Code review

src/models/mistral4.rs is clean: no debug statements, no dead code. The Mistral4MoE::forward flatten/reshape block (lines 449-524) follows the same pattern as qwen3_moe and the comment is accurate. No // TODO or // FIXME markers remain.

src/models/mod.rs:595 registers ModelType::Mistral4 => ("Mistral 4 (MLA)", "Mistral"), so mlxcel arch already lists the family; no registration gap.

Docs

docs/supported-models.md was already updated in commit 088b31328. The Mistral 4 bullet sits after the dots.llm1 entry, which is correct placement. The arch values (q_lora_rank 1024, kv_lora_rank 256, 128 routed + 1 shared experts, 4 active per token, norm_topk_prob, Llama-4 attention scaling) are consistent with the config struct and the validated checkpoint. The Mistral 3 VLM note at line 77 correctly states image-plus-text is validated on both backbones including Mistral Small 4 119B. No em dashes. No slop phrases. No style issues.

Tests

Mistral4MoE::forward has no unit test. The mirrored qwen3_moe has no companion unit test either. A meaningful test for the MoE routing path requires real expert weight tensors, which are not CI-feasible. Synthetic coverage would only verify that the reshape arithmetic is consistent with itself, not that the routing produces correct outputs. No synthetic test was added.

Summary

No code or doc changes were required. The PR is ready to merge as-is.