fix: bounds-check topk_group inside group_mask_scores

[inureyes]

Summary

group_mask_scores (src/models/switch_layers.rs) computed k = n_group - topk_group and handed it to slice_axis and argpartition with no check of topk_group against n_group or against the score row. Seven of its eight callers pass the value straight from config.json. The bounds check now lives inside the function, so all eight call sites are covered at once instead of seven more model files each growing their own copy of it.

Why the silent regime matters

topk_group == n_group + 1 gives k == -1. slice_axis implements Python slice semantics, where end == -1 means "to the end of the axis", so the bottom-k slice selected the whole group axis and put_along_axis zeroed every group. That did not produce zero output. Every caller applies the mask to a selection copy and gathers its combine weights from an untouched orig_scores, so an all-zero (hence all-equal) selection row hands routing to argpartition's tie-break order, which has nothing to do with the router, while the combine weights remain the genuine unbiased scores at those arbitrary experts. Every token routes to the same arbitrary expert set, the output stays finite and in range, and nothing throws or logs. The router is simply off.

n_group + 2 <= topk_group <= 2 * n_group - 1 resolves end to n_group + k and zeroes the wrong number of groups, equally silently, and topk_group == 0 gives k == n_group and zeroes every group.

The remaining regimes abort instead. topk_group >= 2 * n_group and any negative topk_group drive argpartition's normalized kth out of range; num_experts % n_group != 0 and a group holding one expert fail one step earlier, in the reshape and in the top-2 argpartition respectively. mlxcel_core::argpartition is declared non-Result in the bridge, so an MLX C++ throw is an uncatchable std::terminate at the first forward pass rather than a load error: the model loads cleanly and the server dies on its first request. A negative topk_group is reachable without a negative config value, because several callers cast a usize config field through args.topk_group as i32, and a usize above i32::MAX wraps there before group_mask_scores ever sees it.

What changed

• src/models/switch_layers.rs: new group_mask_plan(shape, n_group, topk_group) -> Result<i32, GroupMaskSkip> classifies the triple before any array work and returns experts_per_group when the mask is well defined. group_mask_scores returns the scores unmasked on every skip reason. The guard is pure i32 arithmetic on values already held in the gate struct: no array work, no device round-trip, no eval, so the hot path is unaffected.
• GroupMaskSkip separates the two ordinary reasons to mask nothing (n_group <= 1, and topk_group == n_group, which keeps every group and is what upstream mlx-lm means by it) from the five broken-config reasons (InvalidGroupCount, TopkGroupOutOfRange, UnevenGroups, GroupTooSmall, BadScoreShape). Only the broken ones are reported.
• report_group_mask_skipped writes one eprintln! line per process naming the reason, n_group, topk_group, and the score-row shape. tracing::warn! is a no-op in the mlxcel CLI binary, where nothing installs a subscriber, so a tracing call would have been invisible on the CLI path. Every MoE layer of a model shares one (n_group, topk_group) pair, so one line is the right granularity; firing per layer per token would bury it. Verified: eight distinct bad-config calls in one process emit exactly one line.
• Check ordering puts n_group <= 0 ahead of the n_experts / n_group division so a zero group count cannot divide by zero, and puts the group-width check after divisibility so an empty score row (which passes 0 % n_group == 0) is still caught.
• The // Used by: comment listed six callers. A grep finds eight; it now names Dots1 and SolarOpen as well, per docs/code-guidelines.md. The upstream reference was also converted from a bare deepseek_v3.py mention to the web-accessible GitHub URL form the repository guidelines require.
• Eleven new tests in the models::switch_layers module, including one pure test of group_mask_plan that exercises the whole classification table without touching MLX.

Caller inventory, verified by grep

Eight production call sites, all of which currently gate the call on self.n_group > 1. The guard covers that regime too, so it no longer depends on every future caller repeating the convention.

Call site	Passes	Load-time validation
src/models/bailing_moe.rs:1491	self.topk_group	validate_routing
src/models/deepseek_v3.rs:542	self.topk_group	none
src/models/deepseek_v32.rs:613	self.topk_group as i32	none
src/models/glm4_moe.rs:778	self.topk_group as i32	none
src/models/glm4_moe_lite.rs:640	self.topk_group	none
src/models/exaone_moe.rs:340	self.topk_group as i32	none
src/models/solar_open.rs:617	self.topk_group as i32	none
src/models/dots1.rs:379	self.topk_group as i32	partial, the call-site guard skips the mask when n_group == topk_group

Deliberately not changed

Bailing MoE's validate_routing still rejects topk_group == n_group where upstream mlx-lm accepts it, and its error text still explains that with an argpartition range claim that does not hold against the pinned MLX (a negative kth is normalized, so kth == -1 becomes n_group - 1). The issue frames relaxing this as optional context, and the acceptance criteria require validate_routing to keep rejecting the same configs with the same messages, so it is untouched here. Correcting that rationale, and deciding whether to relax the bound, belongs in its own change.

src/models/llada2_moe/mod.rs:480 carries a private near-duplicate, group_limited_mask, which fills non-kept groups with a large negative constant rather than zero and has the same unguarded k = n_group - topk_group arithmetic. It is not a caller of group_mask_scores, so it is out of scope for this issue, but it is the same defect and should get its own issue.

Test plan

• DEVELOPER_DIR=... cargo test --release --lib --features metal,accelerate models::switch_layers passes 19/19.
• Each new regression test was run individually against the unmodified function first. The silent regimes fail on an assertion, the loud ones kill the binary with SIGABRT: topk_group = 5 returned [0, 0, ... 0] instead of the input row; topk_group = 6 returned a row with the wrong two groups zeroed; topk_group = 0 returned all zeros; topk_group = 8 aborted with [argpartition] Received invalid kth -5 along axis -2 for array with shape: (1,4,1); topk_group = -1 aborted with [argpartition] Received invalid kth 4; n_group = 3 over 16 experts aborted with [reshape] Cannot reshape array of size 16 into shape (1,3,5); n_group = 16 over 16 experts aborted with [argpartition] Received invalid kth 1 along axis -1; and n_group = 0 panicked on integer division by zero.
• In-range behavior is bit-identical. The n_group = 4, topk_group = 2 and topk_group = 3 cases, the topk_group == n_group identity, and a second n_group = 2, topk_group = 1 case all produce the same rows before and after the guard.
• DEVELOPER_DIR=... cargo test --release --lib --features metal,accelerate models::bailing_moe passes 42/42 with no test modified.
• The other caller families pass together, 61/61: models::deepseek_v3, models::deepseek_v32, models::glm4_moe, models::exaone_moe, models::solar_open, models::dots1, models::llada2_moe.
• DEVELOPER_DIR=... cargo clippy --release --lib --tests --features metal,accelerate exits 0 with nothing new; the only two warnings are the pre-existing src/multimodal/host_preprocessor.rs and host_preprocessor_export.rs ones already on main.
• cargo fmt --all -- --check clean.
• python3 scripts/ci/check_cross_repo_refs.py clean.
• Real-checkpoint smoke run on an MoE family whose gate calls this function, token-exact against main. Not run here: the available local checkpoint is models/ling-lite-1.5 at 33.6 GB, which is outside this unit's budget. Every published checkpoint in the eight families is in range, so the guard is a no-op for them by construction, and the in-range bit-identity tests above cover the path they take.

Closes #947

[inureyes]

The red cargo-fmt check is a pre-existing CI breakage, not a formatting problem in this branch. The job fails at the toolchain-install step, before cargo fmt ever runs:

info: syncing channel updates for 1.100.0-x86_64-unknown-linux-gnu
error: could not download nonexistent rust version `1.100.0-x86_64-unknown-linux-gnu`: ... 404

bdad23287 bumped dtolnay/rust-toolchain from 1.93.1 to 1.100.0. That action's tag is the Rust toolchain version, and Rust 1.100.0 does not exist, so the pin 404s. The same job failed on main in runs 30251560172 and 30253948313; the later green main run only passed because that commit touched no Rust and the changes filter skipped the job. cargo-deny, Detect changes and cross-repo refs all pass here.

Formatting was verified locally with cargo fmt --all -- --check (rustfmt 1.8.0-stable, Rust 1.93.1), which is clean. The CI pin needs its own fix.