perf: stack per-expert tensors from borrows instead of copies

[inureyes]

Summary

stack_individual_experts_with_count called mlxcel_core::copy on every per-expert projection tensor before stacking it. The copies existed only to produce owned UniquePtr values for stack_owned and contributed nothing to the result, so the stack is now built from borrowed pointers instead. That removes one Copy graph node per per-expert tensor, 5376 of them for models/ling-lite-1.5.

Correction to the issue's cost model, and to this PR's first draft

Issue #948 states that the copies cost about 31 GB of redundant memory traffic at the first forward pass for Ling-lite, plus a per-projection peak-memory transient of roughly 352 MiB. That is wrong, and this PR originally repeated it. Checked against the pinned MLX checkout:

• mx::copy builds a Copy node (mlx/ops.cpp:303).
• Copy::eval_gpu (mlx/backend/gpu/primitives.cpp:64) and Copy::eval_cpu (mlx/backend/cpu/primitives.cpp:88) both just call Copy::eval.
• Copy::eval (mlx/backend/common/common.cpp:55) is out.copy_shared_buffer(inputs[0]), which repoints array_desc_->data at the input's buffer. It allocates nothing and writes nothing.

The materializing copy_gpu(..., CopyType::General) a few lines above in that same file belongs to Contiguous::eval_gpu, a different primitive, which is the likely source of the original misreading.

So there was never 31 GB of traffic to save and never a 352 MiB transient to remove. What this change actually saves is graph metadata and eval-scheduling overhead for 5376 nodes. The change is still correct and still worth landing (it is strictly less indirection for identical numerics), it is just a much smaller win than the issue claims. The acceptance criterion that would have caught this is the before/after measurement, which is the one still unticked.

The corrected reasoning is now in the source comments rather than only here, since a wrong rationale in a comment is worse than a wrong PR body: the next reader trusts it. Both comments carry an explicit warning not to reintroduce the traffic claim without re-reading the pinned MLX source.

What changed

src/models/switch_layers.rs

• stack_individual_experts_with_count collects *const MlxArray straight from weights.get(...) and calls mlxcel_core::stack instead of building Vec<UniquePtr<MlxArray>> via copy and calling stack_owned. No new borrowing entry point was needed: ops::stack already takes pointers and is public.
• The pointers are collected and consumed inside the live &WeightMap borrow, and the map is only read. This deliberately does not switch to the remove-based draining that src/models/olmoe.rs uses, because this helper is called from a &WeightMap context.
• New test stacked_experts_outlive_the_weight_map_they_were_borrowed_from drops the source map before evaluating the stacked array, then checks per-plane values, not just shapes. It is a tripwire rather than a proof, and its comment says so: freed memory usually reads back intact, so the real guarantee is the by-value arrs.push_back(a->inner) in the C++ stack shim (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:2577), which copies the refcounted mx::array handle so the stacked node holds its inputs independently of the map.
• The Used by: comment was stale, naming three families. It now lists the nineteen that reach the helper through the shared SwitchGLU / SwitchLinear plus DeepSeek v1's direct call at deepseek.rs:857, verified by grep, with a caveat that PhiMoE and OLMoE pre-stack in their own sanitize_weights and so normally take the pre-stacked branch instead.

src/models/phimoe.rs

Fixed in the same pass rather than deferred. It built expert_arrays with mlxcel_core::copy and then collected pointers from those copies, so it carried the identical redundant node over its own expert set. It is not a no-copy precedent, contrary to how it is sometimes described; only olmoe.rs is, because that one drains with weights.remove and never copies. A new inline sanitize_tests module covers the stacking path end to end (w1/w2/w3 to gate/down/up, source map dropped before eval, non-expert passthrough, per-expert keys gone), since phimoe_tests.rs is a sibling module and cannot reach the private sanitize_weights. This mirrors how deepseek_v3.rs tests its own.

Decision on the pre-stacked branch

Not included here. Follow-up #959 was filed for it and has now been closed, because the corrected cost model removes its justification.

The adjacent pre-stacked branch copies the already-joined [num_experts, ...] tensors, and removing that copy would need a handle-sharing bridge shim (std::make_unique<MlxArray>(a.inner), no primitive) that does not exist today, since mlxcel_core::copy is currently the only way to get an owned UniquePtr from a &MlxArray. That was worth scoping separately for two reasons: it adds cxx FFI surface, and it changes an invariant every existing copy caller gets for free, namely that the loaded layer holds a node independent of the WeightMap entry.

With copy now known to be buffer-sharing rather than materializing, the payoff for that shim is graph-node reduction only, across roughly thirteen family-local SwitchLinear implementations plus phimoe's two whole-map fallbacks. New FFI surface and an aliasing audit across thirteen model files is not a good trade for that, so #959 is closed with the evidence recorded rather than left open to mislead whoever picks it up. It can be reopened if someone measures a real cost.

Scope note

The affected-family list is not every MoE family. Fourteen files define their own local SwitchLinear / SwitchGLU and never reach this branch (Qwen3-MoE and Qwen3-VL-MoE, Qwen3-Next and Qwen3.5, DeepSeek v1/v2/V3/V3.2, GLM4-MoE and GLM4-MoE-Lite, ERNIE 4.5 MoE, ExaOne MoE, Hunyuan-MoE, Llama 4, Jamba, gpt-oss, Step3p5). None of them has a per-expert stacking branch at all, so an unstacked checkpoint for those families fails to load rather than taking a slow path. Verified by grep, not assumed.

Test plan

• cargo test --release --lib --features metal,accelerate models::switch_layers (22 passed, including the new ownership test)
• cargo test --release --lib --features metal,accelerate models::phimoe (5 passed, including the new sanitize_tests stacking test)
• Shared-loader family sweep, all green: models::{phimoe,olmoe,deepseek,mixtral,qwen2_moe,bailing_moe,dots1,gemma4,minimax,llada2_moe,granitemoehybrid,lfm2,mellum,solar_open,cohere2_moe,moondream3}
• models::deepseek still passes both direct stack_individual_experts tests, so the expert-count return value and the truncated-checkpoint error keep their contiguous-gather-until-first-gap semantics
• cargo clippy --release --lib --tests --features metal,accelerate clean apart from the two warnings already on main in src/multimodal/host_preprocessor.rs and host_preprocessor_export.rs
• cargo fmt --all -- --check
• Token-exact numerics on a real MoE checkpoint that takes the per-expert branch (models/ling-lite-1.5 is the reference case). Not run here: the release binaries are not built in this environment and the checkpoint is 33.6 GB. Left for the follow-up validation pass. Numerics should be identical by construction, since the graph goes from stack(copy(w_0..w_n)) to stack(w_0..w_n) and Copy aliases its input.
• Second real checkpoint from a different family on the same branch (a per-expert Mixtral or Qwen2-MoE export), same reason.
• First-forward wall time and peak memory before and after. Expect this to show little or nothing now that the cost model is corrected; peak in particular should be unchanged, not improved. Worth running precisely because it is the measurement that disproved the original claim.

Closes #948