fix(models): honor rope_traditional from config.json on the shared Llama path

[inureyes]

Summary

The shared Llama attention path never read rope_traditional from config.json, while the reference ModelArgs in mlx_lm/models/llama.py declares rope_traditional: bool = False and passes it straight into initialize_rope. A Llama, Qwen2 or Qwen2.5 checkpoint declaring the key was rotated split-half where upstream rotates interleaved. The failure is silent: both conventions consume and produce [batch, heads, seq, head_dim], the KV cache stays consistent, and the model emits fluent text out of a mis-rotated attention. No shape test and no reading of the output can catch it.

#930 threaded the flag from ModelArgs through every RoPE route but deliberately left it #[serde(skip)]. This removes the skip.

What changed

• src/models/llama3.rs — ModelArgs::rope_traditional deserializes with a default of false. An explicit null maps to false rather than erroring, via a small deserialize_with: the field was #[serde(skip)], which ignored whatever the key held, so a config carrying "rope_traditional": null loaded fine and must keep loading. A non-boolean is still rejected. Added FUSED_CAUSAL_PREFILL_ENV / FUSED_QKV_SPLIT_ROPE_ENV / FUSED_ROPE_ENV_VARS so the gate, the notice and the tests read one list, and report_fused_rope_bypass_once(), called from Attention::from_weights.
• src/models/llama3_tests.rs (new) — the parse contract, the two VLM config shapes, the sanitizer, the flag reaching Attention, the logits-level oracle, the batched route, and the fused-prefill bypass on a genuinely quantized block.
• src/models/mllama/config.rs — to_llama3_args now forwards rope_traditional, plus a test.
• src/models/helium_tests.rs — replaced the test that asserted the old behavior; renamed attention_carries_the_traditional_flag_only_for_helium.
• src/distributed/tensor_parallel/inference.rs — rewrote the Helium TP refusal rationale.
• src/distributed/tensor_parallel/llama_runtime_tests.rs — new tensor_parallel_llama_propagates_rope_traditional_to_every_rank.
• docs/environment-variables.md, docs/supported-models.md — the two fused-path rows now state the bypass; the Helium entry's TP claim was wrong after this change and is corrected.

The nine deserialization sites

Each was read, not assumed. All pick the key up from wherever the checkpoint declares it and default to false when it is absent.

Site	Shape it parses	Behavior with the default
src/loading/mod.rs:304	full config, string-sanitized first	unchanged
src/loading/vlm_pixtral.rs:236	text_config sub-object	unchanged
src/loading/vlm_llava.rs:173	text_config sub-object	unchanged
src/loading/vlm_smolvlm.rs:189	text_config sub-object	unchanged
src/loading/vlm_idefics2.rs:170	text_config sub-object	unchanged
src/loading/vlm_internvl.rs:125	text_config sub-object	unchanged
src/loading/vlm_fastvlm.rs:169	full config (FastVLM keeps text fields at the root)	unchanged
src/distributed/pipeline/stage_executor/{llama,mistral}.rs	full config, string-sanitized first	unchanged
src/distributed/tensor_parallel/llama_runtime.rs:70	full config, string-sanitized first	unchanged

Two facts that mattered and are now pinned by tests. First, no site injects the key: the six VLM loaders copy only quantization into text_config, so a wrapper-level rope_traditional cannot leak into a nested text backbone, which is correct because in a VLM config the top level describes the multimodal wrapper and not the decoder. Second, sanitize_config_json, which four of these sites run before parsing, only rewrites Infinity and NaN literals by string substitution and leaves the key intact.

MllamaTextConfig::to_llama3_args: a gap, not a scoping decision

It builds a synthetic config from a fixed key list, so a key it does not name cannot reach the shared decoder. MllamaTextConfig has always deserialized rope_traditional with #[serde(default)], and the conversion dropped it, so an mllama checkpoint declaring the key rotated split-half in its self-attention layers against what its own config said. Fixed by naming the key. Cross-attention layers are unaffected: they attend to vision features and apply no RoPE.

Decision on the fused launchers: keep bypassing

fused_qkv_project_split_rope and fused_causal_prefill_attention still hardcode traditional = false in C++, and Attention::forward still routes a traditional-RoPE model around both. The #930 rationale comment is replaced, not supplemented. Reasons, in order of weight:

1. Neither is a fused kernel. Reading both C++ bodies, each is quantized_matmul, three slices, reshape, transpose, fast::rope, which is the same MLX graph the Rust fallback builds op for op. What they save is roughly eleven cxx crossings per layer per forward. The "silent performance cliff" argument for extending the bridge therefore describes FFI call overhead, not throughput, and it is invisible next to the quantized matmuls on either side of it.
2. Both are opt-in and off by default, so nothing shipped today loses anything.
3. The signature is shared. fused_qkv_project_and_rope and fused_qkv_project_split_norm_rope hardcode the same constant and serve other families. Changing two of four leaves an inconsistent surface; changing four adds a parameter that is false at every call site in the tree and pulls families outside this fix into the blast radius.
4. It could not be validated where it matters. No checkpoint pairs rope_traditional with the quantized Llama path, so an extended bridge would ship on synthetic evidence only.

What the bypass actually lacked was observability, so that is fixed rather than argued away: Attention::from_weights prints a one-time stderr notice when a traditional-RoPE checkpoint is loaded while either variable asked for a fused path, naming the variables and saying the graph path is used instead. eprintln! because tracing has no subscriber in the mlxcel CLI binary. Revisit if either launcher is promoted to default-on or grows a real kernel; the notice makes both visible.

a_traditional_rope_block_is_routed_around_the_fused_prefill_launcher asserts the chosen behavior on a genuinely quantized block, in both halves: the flag makes the launcher request a no-op, and the same request on the same weights with the flag off produces a visibly different rotation, so the equality is not vacuous. the_fused_qkv_rope_launcher_cannot_express_traditional_rope (from #930) continues to pin the other launcher.

Tensor parallelism: refusal kept, rationale corrected

The old reason ("rope_traditional is deliberately not deserializable") is no longer true. The flag now reaches every rank: the TP runtime parses config.json into the shared args and local_llama_args clones them, and tensor_parallel_llama_propagates_rope_traditional_to_every_rank asserts the flag on every rank's every layer rather than inferring it from the args.clone().

The refusal survives for a different reason. Helium's convention is fixed in upstream code, not in its config: upstream builds nn.RoPE(..., traditional=True) and the published config.json carries no rope_traditional key at all, which helium_still_needs_the_conversion_because_its_config_omits_the_key pins. The TP runtime parses config.json directly and never goes through to_llama3_args, so a sharded Helium would still parse false. Lifting the refusal means routing the rank-local config through the conversion, which is a change to how the runtime builds its config and not something deserializing a key can deliver.

The #930 test that had to go

a_llama_config_cannot_turn_on_traditional_rope_through_json parsed a Llama config carrying the key and asserted the parsed value was false, which is the exact inverse of this issue's first acceptance criterion. It was replaced (not made to pass) by helium_still_needs_the_conversion_because_its_config_omits_the_key, which asserts the thing that is actually still true and is the reason to_llama3_args is not now redundant. attention_carries_the_traditional_flag_only_for_helium was renamed to ..._from_its_args, since "only for Helium" stopped being accurate.

On the token-level acceptance criterion

The issue asks for a token-level comparison against a reference on a checkpoint that sets rope_traditional. No such public Llama, Qwen2 or Qwen2.5 checkpoint exists, and none of the 182 local configs declares the key, so there is nothing to generate from. The strongest available substitute is implemented instead, and it is a logits comparison rather than a shape check: a model built from a JSON config carrying the key is bit-identical (max_abs_diff == 0.0) to one with the flag set programmatically, and separated from the split-half model by both an absolute and a relative floor. That proves the key selects exactly the interleaved rotation and nothing else. What a real-checkpoint run would add is that mlxcel's interleaved rotation matches the reference's, and that equivalence is fast_rope(traditional = true) itself, which #930 validated token-exactly against mlx-lm on a real Helium checkpoint.

Byte-identical output for existing Llama / Qwen2 / Qwen2.5 checkpoints follows from the default: the parsed args are identical, so the graph is identical. Real-model regression across the local checkpoints is left to the orchestrator.

Test plan

• cargo check --release --lib --tests --features metal,accelerate clean (two warnings already on main)
• cargo test --release --lib --features metal,accelerate models::llama3 — 11 passed
• cargo test --release --lib --features metal,accelerate models::helium — 27 passed
• cargo test --release --lib --features metal,accelerate models::mllama — new test passes; ragged_real_tile_rows_match_reference_masked_full_rows fails on main too (fix(mllama): ragged cross-attention test fails on main, and no CI job runs cargo test #939)
• cargo test --release --lib --features metal,accelerate loading:: — 220 passed
• cargo test --release --lib --features metal,accelerate distributed:: — 1266 passed
• cargo test --release --lib --features metal,accelerate server::startup — 54 passed
• cargo clippy --release --lib --tests --features metal,accelerate clean
• cargo fmt --all -- --check clean
• python3 scripts/ci/check_cross_repo_refs.py clean

Closes #931