fix: bounds-check topk_group inside group_mask_scores #947

Problem / Background

group_mask_scores (src/models/switch_layers.rs:681) is the shared grouped-expert routing mask behind every DeepSeek-V3-style MoE gate. It takes n_group and topk_group as plain i32 and never checks either one against the other or against the score row:

let k = n_group - topk_group;
let group_idx = mlxcel_core::argpartition(&group_scores, k - 1, -2); // [n, n_group, 1]
let group_idx = slice_axis(&group_idx, -2, 0, k);                    // [n, k, 1]
let grouped = mlxcel_core::put_along_axis(&grouped, &group_idx, &zero, -2);

slice_axis (src/lib/mlxcel-core/src/utils.rs:43) implements Python slice semantics: end == -1 means "to the end of the axis", and any other negative end is taken relative to the axis length. A negative k therefore does not raise, it selects a different set of groups.

The silent case: topk_group == n_group + 1

k is -1. argpartition(kth = k - 1 = -2, axis = -2) is fine because MLX normalizes a negative kth (mlx/ops.cpp:2757 at the pinned commit b7c3dd6d, int kth_ = kth < 0 ? kth + a.shape(axis) : kth;), giving n_group - 2. Then slice_axis(group_idx, -2, 0, -1) selects the whole group axis, so group_idx is a full permutation of the group indices and put_along_axis writes zero into every group of the row.

What the model emits is not zero output. Every caller masks a selection copy of the scores and gathers the combine weights from the untouched orig_scores, so the damage lands entirely on selection: with an all-zero (hence all-equal) selection row, argpartition returns a tie-break order that has nothing to do with the router, so every token routes to the same arbitrary expert set, and the combine weights are the genuine unbiased sigmoid/softmax scores at those arbitrary experts, renormalized and scaled as usual. The output is finite, in range, and plausible. The router has simply been switched off. Nothing throws and nothing logs.

The neighboring regimes

• n_group + 2 <= topk_group <= 2 * n_group - 1: k <= -2, slice_axis resolves end to n_group + k, so the wrong number of groups is zeroed. Also silent, also wrong routing.
• topk_group >= 2 * n_group: the normalized kth goes below zero and MLX throws std::invalid_argument. fn argpartition is declared non-Result in the bridge (src/lib/mlxcel-core/src/lib.rs:924), so the exception is an uncatchable abort at the first forward pass rather than a load error.
• topk_group == 0: k == n_group, so the slice takes the whole group axis and every group is zeroed. Same silent-degenerate-routing shape as the n_group + 1 case.
• num_experts % n_group != 0 fails the same way one step earlier: experts_per_group is computed by truncating integer division, so reshape(scores, [n, n_group, n_experts / n_group]) throws on the size mismatch and aborts.

Caller inventory

Eight production call sites, not six. Only Bailing MoE validates:

All eight already gate the call on self.n_group > 1, so the n_group <= 1 regime is covered by every current caller. That is a call-site convention rather than a property of the function, which is the same shape of exposure the rest of this issue describes: it holds until a ninth caller forgets it.

validate_routing (src/models/bailing_moe.rs:583) is the model to generalize from. On the grouped branch it checks num_experts % n_group == 0 (the reshape precondition), experts_per_group >= 2 (group scoring sums each group's top two, so a one-expert group is meaningless), and 1 <= topk_group <= n_group - 1, and it runs at load so a bad config is a load error with a message rather than wrong tokens.

The function's own // Used by: comment lists six callers ("BailingMoe, DeepSeekV3, DeepSeekV32, GLM4Moe, GLM4MoeLite, ExaOneMoe") and is missing SolarOpen and Dots1. docs/code-guidelines.md makes that comment the discovery mechanism for "what breaks if I change this", so the stale list is part of the problem here and is very likely why the caller count was under-counted in the first place.

Both the function and the seven unvalidated call sites predate #946 / #838. The Bailing MoE work is only where this surfaced; nothing in that PR introduced it.

Proposed Solution

Put the bounds check inside group_mask_scores so all eight call sites are covered at once, instead of copying validate_routing's topk_group clause into seven more model files (and having to remember it for the ninth caller).

The function has an unmasked identity available, so it can be made total: when the mask is not well defined, return the scores unmasked rather than computing a negative k. That covers every case above:

• n_group <= 1, or topk_group >= n_group: mask nothing. This is what the math already means (you cannot keep more groups than exist) and it matches the n_group != topk_group guard dots1.rs:378 already carries.
• topk_group == 0: today k == n_group, which zeroes every group. Same silent-degenerate-routing shape, so it belongs in the same guard.
• n_experts % n_group != 0 or n_experts / n_group < 2: mask nothing instead of letting reshape or argpartition throw across a non-Result bridge and abort the process.

So that an out-of-range config is not merely silently ignored, pair the guard with a one-time diagnostic (a warning naming the model, n_group, and topk_group) on the skip path. Model families that want a hard failure can keep or add a load-time check like validate_routing; the in-function guard is the floor that stops wrong output and aborts, not a replacement for load-time validation. Note that tracing::warn! is a no-op in the mlxcel CLI binary (see the implementation context below), so the diagnostic has to be eprintln! to be visible on the CLI path.

Acceptance Criteria

• group_mask_scores validates n_group, topk_group, and the score row width before computing k, and never reaches slice_axis or argpartition with a negative k or a negative normalized kth.
• An out-of-range topk_group can no longer produce a mask that zeroes every group, and can no longer reach an MLX exception through the non-Result bridge.
• Regression test that fails on main: call group_mask_scores directly with n_group = 4, topk_group = 5 and assert the returned row still carries the input scores. On main that call returns an all-zero row.
• Regression test for the abort regime: n_group = 4, topk_group = 8 returns a value instead of terminating the process.
• Regression test for topk_group == 0.
• In-range behavior is bit-identical: for every valid (n_group, topk_group) pair the masked output matches main exactly, including the existing group_mask_scores_keeps_the_groups_with_the_largest_top_two_sum and the n_group = 4, topk_group = 3 case in src/models/bailing_moe_tests.rs.
• Existing Bailing MoE behavior is unchanged: validate_routing still rejects the same configs at load with the same messages, and cargo test --release --lib models::bailing_moe passes with no test modified.
• Real-checkpoint smoke run on at least one MoE family whose gate calls this function (for example models/ling-lite-1.5 or models/solar-open-100b-4bit), token-exact against main. Deferred to the orchestrator: the only local checkpoint in these families is models/ling-lite-1.5 at 33.6 GB, 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.
• The // Used by: comment above group_mask_scores lists all eight callers.

Technical Considerations

Upstream divergence, offered as context and not as a required change. Bailing MoE rejects topk_group == n_group, and its error text says argpartition(kth = n_group - topk_group - 1) "is out of range as soon as topk_group reaches n_group". That reasoning does not hold against the pinned MLX: mlx/ops.cpp:2757 normalizes a negative kth, so kth = -1 becomes n_group - 1, which is in range, and k == 0 then makes the slice empty so no group is masked. Upstream mlx-lm accepts topk_group == n_group for exactly that reason. If the guard proposed above skips the mask for topk_group >= n_group, relaxing Bailing's check from >= n_group to > n_group would be bit-exact, because topk_group == n_group masks zero groups either way. Whoever picks this up should decide whether to relax it here or leave it to a separate change; either way the error message's stated rationale needs correcting.

Hot path. The guard runs once per MoE layer per forward, on scalars already held in the gate struct, so it is arithmetic on i32s with no array work and no measurable cost. Do not add an eval or a device round-trip to implement it.

Scalar casts. Several callers reach the function through args.topk_group as i32 from a usize config field (deepseek_v3.rs:592, glm4_moe_lite.rs:720, and the inline casts in the table above). A usize above i32::MAX wraps to a negative i32 there, before group_mask_scores ever sees it. That makes k larger than n_group rather than negative, which zeroes every group via the end.min(dim_size) clamp and then hands argpartition a kth at or above n_group, so it lands in the abort regime. bailing_moe.rs:1568 already saturates with i32::try_from(args.topk_group).unwrap_or(i32::MAX); the in-function guard should be written so a negative topk_group also reaches the mask-nothing path.