feat(models): add MiniMax-M3 text model (hybrid dense/MoE with block-sparse attention)

[inureyes]

Summary

Adds the MiniMax-M3 text architecture (config model_type: "minimax_m3"): a hybrid dense/MoE decoder with Gemma-style RMSNorm (weight+1) everywhere, clamp-SwiGLU (swigluoai) experts and dense MLPs, partial RoPE, a sigmoid router with a selection-only routing bias plus one shared expert, and a block-sparse "MSA" attention indexer on the sparse layers. This is the text decoder only; the VL wrapper (#764) and the MTP head are out of scope and build on this work. The loader was reconciled against verbatim tensor names and shapes from the real MiniMaxAI/MiniMax-M3 safetensors index.

Architecture implemented

Layer plan: the per-layer dense/MoE split is driven by moe_layer_freq (the real checkpoint runs 3 leading dense layers, then MoE), and the per-layer sparse-attention gate by sparse_attention_config.sparse_attention_freq (dense attention aligns with the leading dense-MLP layers). Dense layers load a plain mlp; MoE layers load block_sparse_moe.

Attention: GQA (64 query heads, 4 KV heads, head_dim 128; q_proj [8192, 6144], k/v_proj [512, 6144]) with per-head Gemma Q/K RMSNorm (qk_norm_type: "per_head", a single [128] gamma shared across heads) and partial RoPE (rotary_dim 64 of head_dim 128, rope_theta 5e6). Load-time output-dim checks on q/k/v fail with a clear error on a GQA layout mismatch.

Router and experts (block_sparse_moe): sigmoid scoring where the routing bias affects the top-k selection only and the returned mixture weights are the unbiased sigmoids (normalized, then scaled by routed_scaling_factor 2.0). Routed experts are stored under the Mixtral convention (block_sparse_moe.experts.{0..127}.w1/w2/w3.weight, w1=gate_proj, w3=up_proj, w2=down_proj) and stacked through the shared switch_layers loader with detected proj names (falling back to gate_proj/up_proj/down_proj). The single shared expert is a SEPARATE MLP (block_sparse_moe.shared_experts.{gate_proj,up_proj,down_proj}), loaded when present and added to the routed mixture; it is never packed into the switch tensors. The clamp-SwiGLU activation reuses mlxcel_core::compiled_gpt_oss_swiglu_activation, and the experts reuse switch_layers::SwitchLinear.

Block-sparse indexer (minimax_m3_indexer.rs, MQA, modeled on the DSA lightning indexer): index_q_proj [512, 6144] produces 4 query heads of sparse_index_dim 128; index_k_proj [128, 6144] produces a single shared index-key stream that all query heads score against; index_q_norm/index_k_norm are [128] Gemma norms over the index vectors, with RoPE on both. The single-head index key is cached alongside the regular K via a head-axis concat (guarded by sparse_index_dim == head_dim). Token scores are reduced to block scores by max over sparse_block_size, top-sparse_topk_blocks blocks are selected with forced initial and per-query local blocks, and the choice expands to an additive block mask on top of the causal mask. Decode stays dense while the selected window still covers at least half the live cache; only beyond that does the block mask apply. Load-time output-dim checks on index_q_proj/index_k_proj catch an MHA-vs-MQA layout mismatch.

Deviations from the issue spec, all reconciled against the real checkpoint: there is no standalone text-only repo, so the config parses a flat text_config shape and is factored so the VL wrapper can construct it from the nested block; sparse_init_block is 0 (the issue said 1); the layer plan is driven by moe_layer_freq. The indexer is MQA (single shared index-key head), so the old sparse_num_index_heads == num_key_value_heads constraint is gone; the side cache instead requires sparse_index_dim == head_dim and disables (dense fallback) otherwise or when index weights are absent (dense-attention layers).

What changed

• src/models/minimax_m3.rs (model, decoder layer, load/from_weights, weight sanitizer, LanguageModel), minimax_m3_config.rs (serde ModelArgs / SparseAttentionConfig with checkpoint-keyed defaults), minimax_m3_layers.rs (attention with head-axis index-key side cache and q/k/v shape checks, clamp-SwiGLU, dense MLP), minimax_m3_moe.rs (router, block_sparse_moe Mixtral-named experts, separate shared expert), minimax_m3_indexer.rs (MQA block-sparse indexer, shape checks), minimax_m3_tests.rs (unit tests). Files split to stay under the 500-line limit.
• Weight sanitizer normalizes language_model. prefixes to the flat model. layout, drops MTP / next-N tensors, and drops vision_tower. / multi_modal_projector. / patch_merge_mlp. tensors so a text-only load of the VL checkpoint ignores the vision front-end.
• Registered in src/models/detection.rs, src/model_metadata.rs, src/loaded_model.rs (enum + dispatch), src/distributed/tensor_parallel/inference.rs (arch string), src/models/mod.rs (ModelType, ALL_MODEL_TYPES, metadata, arch-coverage macro), and docs/supported-models.md.

Validation

The only public checkpoint is the 427B MiniMaxAI/MiniMax-M3 VL repo (top-level model_type: "minimax_m3_vl"), which exceeds this machine's 128 GB unified memory, so it cannot be downloaded or loaded here. Validation is synthetic reduced-config unit tests plus a real-text_config parse test, with the tensor-layout assumptions cross-checked against the real safetensors index (verbatim tensor names and shapes). The acceptance items "real checkpoint loads/generates" and "server chat completion" carry over to a follow-up runtime validation when a fitting quantized conversion exists. The block-sparse path is realized as a masked dense attention (selecting top-k blocks via an additive mask); it reproduces the selection semantics but does not yet realize a native sparse kernel's compute/memory savings.

Test plan

• cargo test --release --features cuda --lib models::minimax_m3 -- --test-threads=1 (9 passed): config parse and derived layer plan; sanitizer prefix rewrite plus vision/MTP drop on verbatim checkpoint keys; block_sparse_moe Mixtral (w1/w2/w3) load with a separate shared expert; MQA indexer shapes (4 query heads, 1 key head); dense-layer index absence; router bias-affects-selection-only; partial RoPE leaving the non-rotary tail untouched; per-head Q/K norm independence; indexer degeneration (full block coverage is bit-identical to dense).
• cargo test --release --features cuda --lib -- models::metadata_tests models::detection minimax (80 passed): the arch-coverage test confirms MiniMaxM3 is wired through every registry.
• cargo fmt --all -- --check clean.

Closes #763

[inureyes]

Implementation Review Summary

Intent

Add the MiniMax-M3 text decoder (model_type: "minimax_m3"): a hybrid dense/MoE decoder with Gemma RMSNorm, clamp-SwiGLU experts, a selection-only-bias sigmoid router with one shared expert, partial RoPE, and a block-sparse "MSA" indexer, wired through the normal load path. Closes #763.

What is correct and validated

The checkpoint-free surface is sound and matches project conventions:

• Router: bias steers top-k selection only; mixture weights are the unbiased sigmoids, normalized then scaled by routed_scaling_factor. Unit-tested.
• Activation: swigluoai reuses compiled_gpt_oss_swiglu_activation with the same (x_linear=up, x_glu=gate) mapping and argument order as gpt_oss.rs. Correct clamp-SwiGLU.
• Partial RoPE (rotary_dim of head_dim), per-head Gemma Q/K norm (last-axis, weight+1), and the block-mask degeneration-to-dense property are each unit-tested and correct.
• Sanitizer: language_model. prefix collapse and MTP/next-N drops look correct and do not touch scales/biases.
• Registration is complete and consistent across detection.rs, model_metadata.rs, loaded_model.rs, TP inference.rs, mod.rs (ModelType, ALL_MODEL_TYPES, summary, arch-coverage macro), and docs/supported-models.md. The model is reachable through the normal detection to from_weights flow; ModelArgs is pub/Deserialize with checkpoint-keyed defaults, so feat(vlm): add MiniMax-M3-VL multimodal support #764 can construct it from a nested text_config.
• Purely additive change; no unintended restructuring; files under the 500-line limit.

Findings (all are risks against the deferred "real checkpoint loads/generates" criterion, not defects in the validated surface)

Two independent reviews converge on the checkpoint weight-layout assumptions as the top risk, and each reaches a different conclusion about the actual tensor layout. Neither can load the 427B checkpoint here, so these items must be reconciled against the real weight map (the safetensors index) before the load path can be considered done.

• HIGH: MoE weight layout is unverified. The loader only accepts model.layers.N.mlp.switch_mlp.{gate,up,down}_proj (pre-stacked) or model.layers.N.mlp.experts.{i}.{gate,up,down}_proj (per-expert). If the released checkpoint instead stores MoE under block_sparse_moe.experts.{i}.{w1,w3,w2} (Mixtral-style), the first MoE layer fails to load. Reconcile the prefix and leaf names against the checkpoint; if needed, use SwitchGLU::from_weights_with_proj_names.
• HIGH: Shared-expert storage is unverified. shared_expert_is_packed() returns true whenever shared_intermediate_size == intermediate_size and then routes index num_local_experts (128). Every other MoE model in the tree (DeepSeek V2/V3/V32, MiniMax-M2) stores the shared expert as a separate shared_experts MLP. If the real checkpoint does the same, the packed path gathers at an out-of-bounds expert row and silently ignores the separate shared weights. Recommend gating "packed" on the actual presence of a num_local_experts + 1-row switch tensor, and adding a load-time row-count assertion so a mismatch errors clearly instead of corrupting the forward.
• HIGH: Index-key head count is unverified. keys() reshapes index_k_proj output to [b, s, sparse_num_index_heads, index_dim]. If the checkpoint's index_k_proj emits a single shared index-key head (DSA lightning-indexer style), this reshape fails on the first sparse forward. Validate the index_k_proj output width against num_index_heads * index_dim.
• MEDIUM: Sparse block selection is combined across index heads. token_scores takes mean over the index-head axis and produces one [b,1,s,kv_len] mask applied to all query heads. If the reference selects top blocks independently per KV group and repeats within the group, this global mask lets KV groups attend to blocks selected for other groups. Note: mean vs sum does not change selection, but per-group vs global does.
• MEDIUM: The should_apply_sparse heuristic (kv_len > 2 * topk_blocks * block_size) keeps full dense attention for 2049 to 4096 tokens, past the point (kv_len > topk_blocks * block_size, 2048) where the trained model would already drop unselected blocks. This changes long-prompt outputs relative to strict reference behavior. It is documented as an intentional dense-stay heuristic; confirm it matches the reference once a checkpoint is available.
• LOW: build_block_drop_mask is called with cache.offset (monotonic) for query_block, but the causal mask and the physical key-block layout use live_len() (live-window length). These are equal for untrimmed caches (all validated paths), so there is no effect today; they diverge only under cache trimming (live_start > 0) with sparse active. Passing live_before instead of offset aligns the two coordinate systems.
• LOW (perf): SwitchGluOai::forward always uses the non-sorted gather path, unlike the shared SwitchGLU which switches to gather_sort for large token batches. Prefill over MoE layers may be slower than the shared path. Correctness is unaffected.

Disposition

No auto-fix was applied. Every finding targets the real-checkpoint load path, which cannot be exercised here (the only public checkpoint is the 427B VL repo and the network is unavailable), and the two expert reviews propose different "correct" layouts. Editing the deliberate, config-informed loading logic to match an unverified layout would trade one unverifiable assumption for another with no test to adjudicate, and risks regressing the passing synthetic surface. The recommended next step is to fetch the checkpoint's safetensors weight map, reconcile the three HIGH items, add defensive load-time shape/row-count checks, then run the deferred real-weights load/generate validation.

Verification

• [~] All stated requirements implemented (text-decoder logic implemented and unit-tested; real-checkpoint load/generate deferred per the PR)
• No placeholder/mock code remaining (the block-sparse path is a real masked-dense implementation, documented as not a native sparse kernel)
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable (SwitchLinear, gpt-oss activation, GemmaRMSNorm)
• No unintended structural changes
• Tests pass (6/6 targeted + 80 registry/detection, per the provided run; not re-run since no code was changed)

[inureyes]

Re-review: HIGH findings resolved (commit aa2304a)

Verified the delta against the real checkpoint layout. All three HIGH items from the previous review are fixed, and I re-ran cargo test --release --features cuda --lib models::minimax_m3 -- --test-threads=1 (9 passed, 0 failed).

• HIGH: MoE weight layout. Now loads block_sparse_moe.experts.{i}.w1/w2/w3 (w1=gate, w3=up, w2=down) with auto-detection of the Mixtral vs gate/up/down_proj naming, router at block_sparse_moe.gate and bias at block_sparse_moe.e_score_correction_bias. Covered by a new load-plus-forward test.
• HIGH: Shared-expert storage. The packing heuristic is removed entirely; the shared expert is a separate block_sparse_moe.shared_experts.{gate,up,down}_proj MLP added to the routed mixture with weight 1.0 after routed_scaling_factor is applied to the routed weights only. This matches the DeepSeek convention already used in the tree.
• HIGH: Index-key head count. The indexer is reworked to MQA: index_q_proj gives 4 query heads of 128, index_k_proj a single shared key head of 128. The single key head rides on the regular K buffer via a head-axis concat guarded by index_dim == head_dim, so it trims in lockstep with K; the key head broadcasts across the query heads in scoring. Confirmed the default Fp16 KVCache handles the 5-head-K / 4-head-V asymmetry (keys and values are sliced by independent shapes). Covered by a new MQA-shape test.

Also verified: check_out_dim defensive load-time checks on q/k/v and index projections, the vision/multimodal prefix drops and lm_head rewrite for text-only VL loads (new verbatim-key sanitizer test), and no dead references to the removed packed-expert code.

One earlier MEDIUM is now moot: with a single ground-truth index-key head, the block mask is architecturally a single stream, so the previous per-KV-group-selection concern does not apply.

Remaining (non-blocking, unchanged by this commit, for runtime validation)

• MEDIUM: should_apply_sparse keeps dense attention until kv_len > 2 * topk_blocks * block_size, past the topk_blocks * block_size point where the trained model would already drop blocks. Documented as an intentional dense-stay heuristic; confirm against the reference at runtime.
• LOW: build_block_drop_mask receives cache.offset (monotonic) for query_block; it should use the live-window length to match the causal mask and physical key blocks. Equal for untrimmed caches (all current paths), diverges only under cache trimming with sparse active.
• LOW (perf): SwitchGluOai::forward always uses the non-sorted gather path, unlike the shared SwitchGLU which sorts for large token batches. Correctness is unaffected.
• LOW (edge): quantized/Turbo KV-cache modes would quantize the index-key head along with K; the default Fp16 path used by this model is unaffected.

Verdict: no remaining CRITICAL or HIGH. The load path now matches the verified checkpoint layout and the reduced-config load/forward paths are exercised by tests; a real-weights load/generate run remains the final validation step.