fix(internlm3): stop doubling RoPE positions under dynamic NTK scaling and add a shared DynamicNTK RoPE helper #1324

Summary

internlm3 passes scale = 2.0 to mlxcel_core::fast_rope whenever rope_scaling is absent or its rope_type is "dynamic". The checkpoint the family is validated on (mlx-community/internlm3-8b-instruct-4bit) ships rope_scaling: {"factor": 6.0, "rope_type": "dynamic"}, so every query and key is rotated at twice its true position on every request. The correct schedule for dynamic NTK is: positions unscaled (scale = 1.0), and only the rotary base grows once the sequence passes max_position_embeddings. This change fixes the scale, moves the dynamic NTK base computation into a shared src/models/dynamic_ntk_rope.rs helper that internlm2 will reuse, and re-pins the InternLM3 greedy parity reference, which was captured against the same wrong schedule.

Current behavior

src/models/internlm3.rs:108-119:

pub fn rope_scale(&self) -> f32 {
    self.rope_scaling
        .as_ref()
        .and_then(|s| {
            if s.rope_type == "linear" {
                Some(1.0 / s.factor)
            } else {
                None
            }
        })
        .unwrap_or(2.0)
}

rope_scaling == None and rope_type == "dynamic" both fall through to unwrap_or(2.0). The value is stored as Attention::rope_scale (src/models/internlm3.rs:297) and handed to both rotary calls at src/models/internlm3.rs:209-224:

let q = mlxcel_core::fast_rope(&q, self.head_dim, self.rope_traditional, rope_base, self.rope_scale, offset);
let k = mlxcel_core::fast_rope(&k, self.head_dim, self.rope_traditional, rope_base, self.rope_scale, offset);

fast_rope multiplies every position by scale before computing the rotation angle, so with scale = 2.0 the token at absolute position p is rotated as if it sat at 2p. For rope_theta = 5e7, head_dim = 128, position 7, frequency pair j = 64: the correct angle is 7 * 5e7^(-64/128) = 9.899e-4; the runtime applies 1.980e-3. Relative angles between any two tokens are doubled as well, so attention patterns differ from what the weights were trained on at every context length, not only past 32768.

The base computation itself is right. compute_dynamic_ntk_base (src/models/internlm3.rs:140-154) returns rope_theta while seq_len <= max_position_embeddings and otherwise rescales it, which is equivalent to clamping seq_len to max(seq_len, max_position_embeddings) before applying the formula.

The private RopeScaling struct at src/models/internlm3.rs:68-72 requires a rope_type key; a config that spells the key type fails to deserialize. src/models/config.rs:309-324 already has a shared RopeScaling with #[serde(rename = "type", alias = "rope_type")] scaling_type: Option<String> and factor: Option<f32>, but no RoPE family uses it for this purpose.

tests/causal_prefill_greedy_parity.rs:104-107 pins INTERNLM3_REF_OUT for the 56-id prompt at INTERNLM3_INPUT_IDS. Those 24 ids were produced by a reference that applied the same scale = 2.0, so they pin the defect. docs/supported-models.md:38 repeats the claim that the family is token-exact against that reference.

Expected behavior

For a [B, n_heads, L, head_dim] query or key tensor at cache offset offset, with dims = head_dim, base = rope_theta, max_pos = max_position_embeddings:

• rope_scaling absent, or scaling_type in {"default", None}: fast_rope(x, dims, traditional, base, 1.0, offset).
• scaling_type == "linear" with factor = f: fast_rope(x, dims, traditional, base, 1.0 / f, offset). Position p is rotated as p / f.
• scaling_type == "dynamic" with factor = f:
  • seq_len = max(L + offset, max_pos)
  • base_eff = base * ((f * seq_len / max_pos) - (f - 1)) ^ (dims / (dims - 2))
  • fast_rope(x, dims, traditional, base_eff, 1.0, offset). Positions are never scaled.
• Any other scaling_type (llama3, yarn, ...) is an error at load time for this family. A "linear" or "dynamic" entry without factor is also an error.

Worked values for the validation checkpoint (base = 5e7, f = 6.0, dims = 128, max_pos = 32768):

Angle for position p and frequency pair j (0 <= j < dims / 2): theta(p, j) = p * base_eff ^ (-2j / dims). At p = 7, j = 0: 7.0; at p = 7, j = 64: 9.899e-4.

Implementation plan

1. Add src/models/dynamic_ntk_rope.rs (register pub mod dynamic_ntk_rope; in src/models/mod.rs, and attach dynamic_ntk_rope_tests.rs with the usual #[cfg(test)] #[path = "dynamic_ntk_rope_tests.rs"] mod tests;). Contents:

   /// Used by: InternLM3, InternLM2
   #[derive(Debug, Clone, Copy, PartialEq)]
   pub enum DynamicNtkRopeMode {
       Default,
       Linear { factor: f32 },
       Dynamic { factor: f32 },
   }
   
   /// Used by: InternLM3, InternLM2
   #[derive(Debug, Clone)]
   pub struct DynamicNtkRope {
       pub dims: i32,
       pub base: f32,
       pub traditional: bool,
       pub max_position_embeddings: usize,
       pub mode: DynamicNtkRopeMode,
   }
   
   impl DynamicNtkRope {
       /// Build from the shared `config::RopeScaling`. `None` and
       /// `scaling_type` in {None, "default"} give `Default`; "linear" and
       /// "dynamic" require `factor` (Err otherwise); anything else is Err.
       pub fn from_scaling(dims: i32, base: f32, traditional: bool,
           max_position_embeddings: usize,
           scaling: Option<&crate::models::config::RopeScaling>) -> Result<Self, String>;
   
       /// 1.0 for Default and Dynamic, 1.0 / factor for Linear.
       pub fn position_scale(&self) -> f32;
   
       /// `base` for Default and Linear; the NTK-rescaled base for Dynamic,
       /// computed from `seq_len = max(seq_len, max_position_embeddings)`.
       pub fn base_for(&self, seq_len: i32) -> f32;
   
       /// `mlxcel_core::fast_rope(x, dims, traditional, base_for(seq_len), position_scale(), offset)`
       pub fn apply(&self, x: &MlxArray, offset: i32, seq_len: i32) -> UniquePtr<MlxArray>;
   }

   base_for is the existing compute_dynamic_ntk_base body with the if seq_len > max_pos branch replaced by the max clamp (same result, fewer branches). Use f32 throughout to match fast_rope's base: f32 parameter.

2. In src/models/internlm3.rs:

   • Delete the private RopeScaling struct (lines 68-72) and type rope_scaling as Option<crate::models::config::RopeScaling>. The shared struct's alias = "rope_type" keeps the existing rope_type key working.
   • Delete ModelArgs::rope_scale, ModelArgs::dynamic_ntk_factor and compute_dynamic_ntk_base.
   • Replace the rope_base, rope_scale, rope_traditional, max_position_embeddings, dynamic_ntk_factor fields of Attention with one rope: DynamicNtkRope, built in Attention::from_weights with DynamicNtkRope::from_scaling(head_dim, args.rope_theta, args.rope_traditional, args.max_position_embeddings, args.rope_scaling.as_ref())? (the constructor already returns Result<Self, String>).
   • In Attention::forward, replace the two fast_rope calls with self.rope.apply(&q, offset, l + offset) and self.rope.apply(&k, offset, l + offset).
   • Update the module doc comment at the top of the file (it describes the dynamic base but not the position scale).

3. Re-pin INTERNLM3_REF_OUT in tests/causal_prefill_greedy_parity.rs from an oracle that applies the corrected schedule (positions unscaled, base unchanged below 32768 tokens) to the same 4-bit checkpoint and the same 56 prompt ids, and update the comment above the constant to say which oracle produced it. Leave INTERNLM3_INPUT_IDS as is.

4. Update docs/supported-models.md:38: drop the claim of token-exactness against the old reference, describe the corrected schedule in one sentence (positions unscaled; base rescaled only past max_position_embeddings), and name the oracle the re-pinned ids came from.

Validation

(a) Unit tests, src/models/dynamic_ntk_rope_tests.rs:

• scale_is_one_when_rope_scaling_absent: from_scaling(128, 5e7, false, 32768, None) gives position_scale() == 1.0 and mode == Default.
• scale_is_one_for_dynamic: RopeScaling { scaling_type: Some("dynamic"), factor: Some(6.0), ..Default::default() } gives position_scale() == 1.0 and mode == Dynamic { factor: 6.0 }.
• scale_is_inverse_factor_for_linear: scaling_type: Some("linear"), factor: Some(4.0) gives position_scale() == 0.25 and base_for(100_000) == base (linear never rescales the base).
• rope_type_alias_is_accepted: deserialize {"factor": 6.0, "rope_type": "dynamic"} and {"factor": 2.0, "type": "dynamic"} into config::RopeScaling; both build Dynamic.
• dynamic_base_is_unchanged_at_or_below_max_position: base_for(100) and base_for(32768) both equal 5e7 exactly for the factor = 6.0, dims = 128 configuration.
• dynamic_base_grows_past_max_position: base_for(40000) is within 1e-3 relative of 117_777_118.66 and base_for(65536) within 1e-3 relative of 360_979_300.43.
• unknown_and_incomplete_scaling_are_rejected: scaling_type: Some("yarn") is Err; scaling_type: Some("dynamic"), factor: None is Err.
• apply_matches_fast_rope_with_unit_scale: for a [1, 2, 8, 16] float32 tensor of arange values and offset = 3, apply(x, 3, 11) under Dynamic { factor: 6.0 }, max_pos = 32768 is array_equal to mlxcel_core::fast_rope(x, 16, false, base, 1.0, 3); and under Linear { factor: 2.0 } it is array_equal to fast_rope(x, 16, false, base, 0.5, 3).

Also add src/models/internlm3_tests.rs (new, attached via #[path]) with internlm3_attention_uses_unit_scale_for_dynamic_config: build ModelArgs from the validation checkpoint's config JSON (rope_theta: 50000000, rope_scaling: {"factor": 6.0, "rope_type": "dynamic"}, max_position_embeddings: 32768) and assert the attention's rope.position_scale() == 1.0 and rope.base_for(56) == 50_000_000.0.

(b) Real checkpoint: mlx-community/internlm3-8b-instruct-4bit (local models/internlm3-8b-4bit).

./target/release/mlxcel generate -m models/internlm3-8b-4bit -p "The Industrial Revolution changed the way people lived and worked because" -n 64
cargo test --profile test-fast --features metal,accelerate --test causal_prefill_greedy_parity internlm3

Acceptance: the first 24 greedy tokens on the 56-id prompt are token-exact against the corrected oracle (re-pinned INTERNLM3_REF_OUT); internlm3_prefill_is_causal still passes at its existing tolerance; CLI output is finite and fluent. A second run with a prompt of at least 64 tokens must also be token-exact, since a 6-token prompt cannot distinguish the two schedules at small positions.

Acceptance criteria

• internlm3 passes scale = 1.0 to fast_rope for rope_scaling = None and for rope_type = "dynamic", and 1 / factor only for "linear"
• Dynamic NTK base still equals rope_theta for L + offset <= max_position_embeddings and matches the formula above past it (table values)
• src/models/dynamic_ntk_rope.rs exists with a // Used by: comment and internlm3 uses it; the private RopeScaling in internlm3.rs is gone in favor of config::RopeScaling
• INTERNLM3_REF_OUT re-pinned and tests/causal_prefill_greedy_parity.rs passes on models/internlm3-8b-4bit
• docs/supported-models.md InternLM entry updated
• cargo test --workspace --profile test-fast --features metal,accelerate passes
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Out of scope

• internlm2 (tracked separately; it will consume the helper introduced here).
• YaRN or Llama-3 style scaling for the InternLM families; no public checkpoint uses them.

Implementation note (PR #1389). Two deviations from the plan above, both deliberate:

1. The helper is built on rope_utils::RopeScalingSpec (fix(llama3): apply llama3 rope_scaling frequencies on the shared Llama path (Llama 3.1/3.2 text and Llama 3.2 Vision) #1355), not on crate::models::config::RopeScaling. That struct has no importers anywhere in the tree and uses the #[serde(rename = "type", alias = "rope_type")] construct that hard-errors with duplicate field on a config carrying both keys. It is left untouched; the private internlm3::RopeScaling is deleted as planned.
2. internlm2 is in scope after all. Its ModelArgs did not declare rope_scaling at all, so its checkpoints' live {"type": "dynamic", "factor": 2.0} was dropped at deserialization; shipping a helper "shaped for" a family that still ignores its own block was not worth the deferral. Its position scale was already correct, so everything inside max_position_embeddings is bit-identical and models/internlm2-7b-4bit generates byte-identical output before and after.

The workspace gate lines are checked against the local equivalents that were actually run (--lib for the mlxcel package plus the parity integration test, and clippy --lib --tests); the full --workspace sweep is the orchestrator's.

Note also that stock mlx-lm 0.31.3 carries this same defect, so the acceptance criterion "token-exact against the mlx-lm greedy reference" would have been satisfied by the buggy code. The re-pin comes from an mlx-lm whose rope module was corrected first, never from mlxcel's own output.