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

Problem / Background

Current behavior

The shared Llama attention path still never reads rope_traditional from config.json, but the code around it has moved since this issue was filed. #930 (implementing #837, Kyutai Helium) has MERGED, as 1f21b1f4c, so on main today src/models/llama3.rs does carry ModelArgs::rope_traditional and Attention::rope_traditional, and both graph RoPE routes do receive the flag instead of a literal false:

// src/models/llama3.rs, Attention::forward (single-sequence prefill and decode)
let q = mlxcel_core::fast_rope(&q, self.rope_dims, self.rope_traditional, self.rope_base, 1.0, offset);
let k = mlxcel_core::fast_rope(&k, self.rope_dims, self.rope_traditional, self.rope_base, 1.0, offset);

// src/models/llama3.rs, Attention::forward_split_attention (batched decode, batched prefill, paged decode)
let q_batched = mlxcel_core::fast_rope_batched(
    &q_batched,
    self.rope_dims,
    self.rope_traditional,
    self.rope_base,
    1.0,
    &metadata.rope_offsets,
);
let k_batched = mlxcel_core::fast_rope_batched(
    &k_batched,
    self.rope_dims,
    self.rope_traditional,
    self.rope_base,
    1.0,
    &metadata.rope_offsets,
);

What #930 deliberately did not do is deserialize the field (src/models/llama3.rs:70-80):

/// Rotate interleaved channel pairs `(2i, 2i+1)` instead of the split-half
/// pairs `(i, i + dims/2)`.
///
/// Deliberately **not** deserialized. Every family that reaches this loader
/// today rotates split-half, and a checkpoint author who writes
/// `"rope_traditional": true` into a Llama `config.json` must not be able to
/// silently change how an existing checkpoint decodes. [...]
#[serde(skip)]
pub rope_traditional: bool,

The only writer is helium::ModelArgs::to_llama3_args (src/models/helium.rs:415), which sets it programmatically. So after #930 the plumbing exists on the graph paths, but a Llama or Qwen2 config.json that declares "rope_traditional": true is still silently ignored, and #930 pins that behavior with a test, a_llama_config_cannot_turn_on_traditional_rope_through_json (src/models/helium_tests.rs:296), which parses a Llama config carrying the key and asserts the parsed value is false. That was the right scoping call for a single model port. It leaves the general question open, which is what this issue is for. One consequence for whoever picks this up: that test asserts the exact opposite of the first acceptance criterion below, so the fix must replace it rather than keep it green.

Correct behavior

The reference implementation carries the key and honors it. In mlx_lm/models/llama.py, ModelArgs declares rope_traditional: bool = False and Attention.__init__ passes args.rope_traditional straight into initialize_rope:

self.rope = initialize_rope(
    self.head_dim,
    args.rope_theta,
    args.rope_traditional,
    args.rope_scaling,
    args.max_position_embeddings,
)

Any checkpoint whose config.json sets the key therefore rotates interleaved under the reference and split-half under mlxcel. mlxcel diverges from the reference for exactly that class of checkpoint.

Which families inherit the gap

Re-verified against main at a3c823611, after the Helium merge.

src/model_metadata.rs still registers exactly two model types against models::llama3::ModelArgs:

• Llama => { ... args: models::llama3::ModelArgs, weight_builder: models::Llama3Model::from_weights, ... } (model_metadata.rs:96)
• Qwen2 => { ... args: models::llama3::ModelArgs, weight_builder: models::Qwen2Model::from_weights, ... } (model_metadata.rs:100)

src/models/qwen2.rs is still a pure re-export (pub use super::llama3::{Attention, Llama3Model as Qwen2Model, MLP, ModelArgs, Quantization, RopeScaling, TransformerBlock};), and src/models/detection.rs maps "qwen2" => Ok(ModelType::Qwen2), which is the model_type string Qwen2.5 checkpoints ship. So Llama, Qwen2 and Qwen2.5 all share the affected attention. Helium's merge did not add a third entry to that list: model_metadata.rs:207 registers Helium against its own models::helium::ModelArgs, which reaches the shared decoder by conversion (to_llama3_args) rather than by a second registration.

The two registrations are not the whole surface, though. llama3::ModelArgs is deserialized from JSON in nine further places, and every one of them would begin honoring the key the moment it becomes deserializable:

The Llama-3.2-Vision (mllama) text decoder is the one reuser that would not inherit it: MllamaTextConfig::to_llama3_args (src/models/mllama/config.rs:77) builds a synthetic JSON object from a fixed key list that never contains rope_traditional, so it keeps parsing to false regardless of what the checkpoint declares.

The // Used by: inventory now lives on Attention::forward (llama3.rs:149-155), and forward_split_attention (:348-350) defers to it rather than repeating a shorter list:

Used by: Llama (llama / mistral checkpoints), Qwen2 / Qwen2.5 (which re-export this attention), Helium (which reuses it with rope_traditional set), the Llama-3.2-Vision (mllama) text decoder's self-attention layers, every VLM whose text backbone is Llama3Model or Qwen2Model (Pixtral, LLaVA, SmolVLM / Idefics3, Idefics2, InternVL, FastVLM, dots.ocr), the llama and mistral pipeline stage executors, and the tensor-parallel Llama runtime.

The tensor-parallel Llama runtime builds its per-rank config with local_llama_args in src/distributed/tensor_parallel/llama_runtime.rs:1878, which does let mut local = args.clone(); and mutates only the sharded dimensions. A deserialized flag would propagate to every rank for free, but that should be asserted rather than assumed. Related: tensor parallel is currently REFUSED for Helium for precisely this reason, recorded at src/distributed/tensor_parallel/inference.rs:251-261 ("a sharded Helium would silently rotate split-half pairs while the single-process path rotates interleaved pairs"). Closing this issue removes that blocker, so that comment and the runtime_kind_for gap it describes need revisiting in the same change.

Whether any local checkpoint triggers it

Re-scanned every config.json under models/ on main at a3c823611 (182 files, following symlinks): no checkpoint declares rope_traditional. 22 of those configs declare a model_type of llama or qwen2, so the affected families are well represented locally, but none of them sets the key. There is no reproducer in the local checkpoint set today, and validating a fix will need either a public checkpoint that ships the key or a synthetic config plus a reference comparison.

Why this matters: the failure is silent

Applying the split-half rotation to a model trained with the traditional (interleaved) convention, or the reverse, is not a crash and not a shape error. Both conventions consume [batch, heads, seq, head_dim] and produce the same shape, the KV cache stays consistent, attention runs, and the model emits fluent, plausible, grammatical text. Nothing throws, no assertion fires, and reading the output does not reveal it: the degradation looks like a slightly worse model, not like a bug.

This is invisible to shape tests, to cache-consistency tests, and to logits-shape tests. The only thing that catches it is a token-level comparison against a reference implementation on the same weights and the same prompt. That constraint dictates what an acceptance test for this fix has to be, so it is stated here rather than left to the implementer to rediscover. #930 added traditional_and_split_half_rope_are_different_rotations (src/models/helium_tests.rs), which demonstrates the same point at the fast_rope level: the two conventions agree on every shape and disagree on the values.

Technical Considerations

The two fused quantized launchers still cannot express the flag. In src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp, inside fused_qkv_project_split_rope (declared at line 3822):

q = mlx::core::fast::rope(q, rope_dims, false, rope_base, 1.0f, cache_offset);
k = mlx::core::fast::rope(k, rope_dims, false, rope_base, 1.0f, cache_offset);

and inside fused_causal_prefill_attention (declared at line 4028):

q = mlx::core::fast::rope(q, rope_dims, false, rope_base, 1.0f, 0);
k = mlx::core::fast::rope(k, rope_dims, false, rope_base, 1.0f, 0);

The third positional argument is MLX's traditional parameter, hardcoded false in both bodies, and neither the cxx bridge signature nor the Rust helper takes a flag to override it. Both launchers are reachable only for quantized weights and both are opt-in behind an environment variable (MLXCEL_ENABLE_FUSED_QKV_SPLIT_ROPE and MLXCEL_ENABLE_FUSED_CAUSAL_PREFILL_ATTENTION respectively). Two neighboring launchers in the same file, fused_qkv_project_and_rope and fused_qkv_project_split_norm_rope, hardcode false the same way; they are not on the Llama attention path today, but a bridge signature change should account for them.

#930 handled this by gating both fused branches on !self.rope_traditional (llama3.rs:198 and llama3.rs:235), so a traditional-RoPE model always falls back to the graph path that does receive the flag, and recorded the decision and its rationale in the # Traditional RoPE and the fused fast paths section of the Attention::forward doc comment (llama3.rs:157-183). That is acceptable for one family whose flag is fixed at load time and whose fast paths are off by default. It is not a general answer for a config-driven flag: it means any checkpoint that sets the key silently loses both quantized fast paths, and the loss is invisible at the call site.

Acceptance Criteria

• llama3::ModelArgs::rope_traditional is deserialized from config.json with a default of false, matching the reference ModelArgs default.
• The flag reaches every RoPE route the Llama attention can take: single-sequence prefill and decode (fast_rope), batched decode and batched or padded prefill (fast_rope_batched), and paged decode (which applies RoPE before the paged dispatch).
• The nine non-registration deserialization sites listed above (the generic Llama loader, six VLM text_config loaders, the two pipeline stage executors, and the tensor-parallel runtime) are each either confirmed to pick the flag up correctly or explicitly scoped out with a reason. The decision is recorded, not left implicit.
• The tensor-parallel Llama runtime propagates the flag to every rank, verified by a test rather than inferred from the args.clone() in local_llama_args. The Helium tensor-parallel refusal documented at src/distributed/tensor_parallel/inference.rs:251-261 is re-evaluated, and either lifted or kept with an updated rationale.
• A decision is made and implemented for the two fused quantized launchers: either the cxx bridge carries the flag, or the launchers stay bypassed when it is set. Whichever is chosen replaces the interim rationale already recorded in the Attention::forward doc comment (llama3.rs:157-183).
• A test asserts the chosen fused-path behavior directly, so that the assumption it rests on fails loudly if the launchers later change.
• Correctness is proved by a token-level comparison against a reference implementation, not by a shape assertion or a "runs without error" check. Same weights, same prompt, same sampling settings, comparing generated token ids (or per-step logits within a stated tolerance) against the reference for a checkpoint that sets rope_traditional. A shape-only or smoke test does not satisfy this criterion, because the wrong rotation is shape-correct and produces fluent output.
• Existing Llama, Qwen2 and Qwen2.5 checkpoints produce byte-identical output to the current build. The default is false and none of these checkpoints declares the key, so nothing about their graph or their generated tokens may change. Verify on at least one Llama checkpoint and one Qwen2 or Qwen2.5 checkpoint from the recommended test set.
• Helium (feat(models): add Kyutai Helium (helium) text model support #837 / feat(models): add Kyutai Helium (helium) text model support #930) keeps working: its loader still sets the flag programmatically through to_llama3_args, and every Helium test stays green except a_llama_config_cannot_turn_on_traditional_rope_through_json (src/models/helium_tests.rs:296), which asserts the current non-deserializing behavior and must be replaced by a test asserting the parse.
• A unit test covers the parse itself: a config that omits the key yields false, a config that sets it to true yields true.
• cargo fmt --all -- --check and cargo clippy --all-targets -- -D warnings are clean.

Open Question

Should the cxx bridge signature be extended so fused_qkv_project_split_rope and fused_causal_prefill_attention can express the rotation convention, or should they keep being bypassed when the flag is set?

• Extend the bridge. Correct in all cases and removes a silent performance cliff. Cost: it is an FFI signature change on a surface shared by every model family that reaches these launchers, so it touches the Rust helper, the C++ body, and every call site at once, and it has to be revalidated on the quantized paths for the families already using them. It also adds a parameter that is false for every caller except one, which is maintenance burden with no immediate user.
• Keep bypassing. Zero risk to existing families and no FFI churn. Cost: a checkpoint that sets the key silently loses both quantized fast paths, and nothing at the call site tells the user why. Since neither path is on by default today the cost is currently theoretical, but it becomes real the moment either one is promoted to default-on, and at that point the regression would appear only for traditional-RoPE checkpoints and would be hard to attribute.

A middle option worth weighing: keep the bypass, but emit a one-time warning at load when a checkpoint sets rope_traditional while a fused path was requested through its environment variable, so the fallback is at least observable. Note that tracing::warn! is a no-op in the CLI binary (see the implementation context below), so such a warning has to be eprintln!.

References

Originating work: #930 (implementation, merged as 1f21b1f4c) and #837 (the Helium port that surfaced this).