feat: dense arch pack (Seed-OSS, MiMo, InternLM3, ExaOne) in mlxcel-xla (#499)

[inureyes]

Summary

Adds four dense Llama-family architectures to the OpenXLA (mlxcel-xla) emitter as config / naming deltas over the two already-proven forwards (Llama and Qwen2), part of epic #493. The emitter's forward pass (emitter/model.rs) is untouched: each family was verified from its modeling code to use the emitter's half-split RoPE, cat(freqs, freqs) table and head_dim^-0.5 scaling, RMSNorm and SwiGLU, so its emitted graph is a proven graph rather than a new one.

Families added

• Seed-OSS (seed_oss): q/k/v projection bias (from attention_bias), untied embeddings, rope_type = "default" served as plain RoPE. The proven Qwen2 bias forward with standard names.
• MiMo (mimo): MiMoForCausalLM subclasses Qwen2ForCausalLM and reuses Qwen2Attention / Qwen2MLP / Qwen2RMSNorm verbatim (q/k/v bias, untied, plain RoPE); its multi-token-prediction heads are not loaded and its config sliding_window is served globally (None), as for Qwen2.
• InternLM3 (internlm3): standard names, untied, rope_type = "dynamic" served as plain RoPE (dynamic NTK is identity within the original context; long-context rescale is a follow-up).
• ExaOne 3.x (exaone): llama3 RoPE, tied, the GPT-2-style transformer.h.{i}... tensor names via a new WeightScheme::Exaone, and the num_layers / layer_norm_epsilon alternate config fields. Verified against the checkpoint's modeling_exaone.py (gated MLP c_proj(act(c_fc_0(x)) * c_fc_1(x)), so c_fc_0 is the gate and c_fc_1 the up projection).

What changed

• emitter/config.rs: generalizes the q/k/v bias source (Qwen2 hard-coded, else attention_bias / qkv_bias), widens the accepted rope types (default / in-context dynamic -> plain), reads the alternate field names, and adds a WeightScheme. Out-of-scope deltas are rejected with clear errors: an attention output bias, an MLP bias, a non-SwiGLU activation, an unsupported rope type (e.g. yarn), and interleaved (GPT-J) RoPE (e.g. ERNIE-4.5).
• weight_names.rs (new): the emitter's weight arg order moves to a pure-Rust module so ExaOne's GPT-2-style remap is unit-tested without the iree feature; iree.rs now calls it. The Llama scheme reproduces the previous names byte-for-byte, so Llama / Qwen2 / Gemma2 loading is unchanged.
• validation.rs: registers each family as a byte-exact structural fixture (small synthetic configs keep the goldens small) and adds dense_pack_families_reuse_proven_graphs, which asserts each family emits StableHLO byte-for-byte identical to a proven Llama / Qwen2 reference across the single, prefill, and ragged graph kinds.
• assets/{seed_oss,mimo,internlm3,exaone}/: fixture config.json + frozen decode.mlir / prefill.mlir + a README documenting each delta and its validation.
• spike/openxla/synthetic_arch_check.py (new): builds a tiny random SeedOssForCausalLM and matches the emitter's prefill logits through IREE to HF; spike/openxla/arch_execution_check.py (new): drives a real-checkpoint continuation against an HF fp32 oracle. Neither builds xla-iree.

Validation

• cargo test -p mlxcel-xla --lib: 72 passed (byte-exact fixtures, the reuse-proven-graphs equivalence, config-parse and weight-name unit tests, and the existing Llama/Qwen2/Gemma2 gates all green; the llama-3.2-1b goldens are unchanged).
• Forward parity: synthetic_arch_check.py matches Seed-OSS to HF eager to max|logit diff| = 3.7e-9. The arch_execution_check.py harness ran end to end on the smallest checkpoint and correctly rejected ERNIE-4.5 (interleaved RoPE diverged from HF by ~8, vs ~1e-9 for a supported family), which is why it is deferred.
• cargo clippy -p mlxcel-xla --lib --tests -- -D warnings: clean; cargo fmt applied.
• The heavy token-exact / serve gate needs the native IREE build and runs post-merge via scripts/xla/validate_arch.sh on the real checkpoints.

Deferred (documented follow-ups)

Out of scope for a dense half-split pack, each needing a distinct emitter subsystem: ERNIE-4.5 (interleaved GPT-J RoPE), Mistral / Ministral (YaRN rope + VLM language_model. prefix), ExaOne4 (QK-norm + post-norm placement), InternLM2 (fused wqkv loader surgery), GLM4-flash and Solar-Open (MoE + MLA, not dense), Baichuan-M1 (conv + differentiated sliding-window heads), Apertus (xIELU activation), and Hunyuan-dense (cross-layer attention). The config.json parser rejects each with a clear, specific error.

Test plan

• cargo test -p mlxcel-xla --lib (72 passed)
• cargo clippy -p mlxcel-xla --lib --tests -- -D warnings (clean)
• spike/openxla/synthetic_arch_check.py (Seed-OSS forward parity ~1e-9)
• Post-merge: scripts/xla/validate_arch.sh --model <ckpt> token-exact / serve gate per family (needs the native IREE build)

Closes #499

[inureyes]

Merge note: conflicts with #558 (complementary arch sets)

This PR was branched from e7ffc53. Sibling epic-#493 unit #558 ("dense arch pack for qk-norm and Gemma family") landed on main afterwards and restructured the same shared files, so GitHub reports this PR as conflicting. The two arch sets are disjoint and complementary, not duplicated:

• This PR (feat: dense arch pack for remaining Llama-family architectures (Mistral/Ministral, ExaOne, InternLM, Baichuan, GLM4, and more) #499): seed_oss, mimo, internlm3, exaone (config / naming deltas over the proven Llama and Qwen2 forwards; no emitter/model.rs change).
• feat: dense arch pack for qk-norm and Gemma family (Qwen3, Gemma1, Gemma3, SmolLM3, OLMo2/3) #558: qwen3, gemma1, gemma3, smollm3, olmo2, olmo3 (QK-norm / norm-placement / NoPE, with model.rs + rope.rs changes).

The merge needs to re-base this PR's additions onto #558's redesigned Config, which are otherwise orthogonal:

• emitter/config.rs::from_json_str: add the four model_types, the rope-type widening (default / in-context dynamic -> plain), the generalized q/k/v bias source (attention_bias / qkv_bias), the ExaOne alternate fields (num_layers / layer_norm_epsilon), the WeightScheme field, and the interleaved-RoPE (ERNIE-4.5) rejection. feat: dense arch pack for qk-norm and Gemma family (Qwen3, Gemma1, Gemma3, SmolLM3, OLMo2/3) #558 replaced the gemma2: bool with NormStyle / QkNorm / embed_scale / norm_one_plus / mlp_geglu flags, so the Config literals here (weight_scheme) need to coexist with feat: dense arch pack for qk-norm and Gemma family (Qwen3, Gemma1, Gemma3, SmolLM3, OLMo2/3) #558's new fields.
• Tensor naming: this PR extracted weight_names into a pure-Rust weight_names module (so ExaOne's GPT-2-style remap is unit-tested without the iree feature); feat: dense arch pack for qk-norm and Gemma family (Qwen3, Gemma1, Gemma3, SmolLM3, OLMo2/3) #558 kept weight_names inline in iree.rs and extended it (q/k norms, has_input_norm). Reconcile to one location, folding in the WeightScheme::Exaone remap.
• validation.rs: union the fixture registrations and the dense_pack_families_reuse_proven_graphs cases.

Per the workflow, no merge was performed here.