fix: resolve rope_theta from wherever the config stores it

tests/test_draft_rope.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Tests that a draft's rotary base is the one its config declares.
+
+transformers 5.x moved `rope_theta` from a top-level config attribute into
+`rope_parameters`, which silently reduced every affected draft to the library
+default of 10000 while its config advertised something else. The sweep over
+`configs/draft_models/` is the part that fails if a future release moves the
+field again.
+"""
+
+import glob
+import json
+import os
+import unittest
+
+from transformers import LlamaConfig
+from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
+
+from torchspec.config.utils import resolve_rope_theta
+from torchspec.models.draft.auto import AutoDraftModelConfig
+from torchspec.models.draft.deepseek_eagle import DeepSeekMLAAttention
+from torchspec.models.draft.dflash import DFlashAttention, DFlashConfig
+from torchspec.models.draft.llama3_eagle import LlamaAttention
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DRAFT_CONFIG_DIR = os.path.join(REPO_ROOT, "configs", "draft_models")
+
+YARN = {
+    "rope_type": "yarn",
+    "factor": 32.0,
+    "original_max_position_embeddings": 64,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+}
+
+
+def _declared_theta(raw):
+    """The base a config file asks for, from wherever the file happens to put it."""
+    block = raw.get("rope_parameters") or raw.get("rope_scaling") or {}
+    nested = block.get("rope_theta")
+    return nested if nested is not None else raw.get("rope_theta")
+
+
+class TestResolveRopeTheta(unittest.TestCase):
+    def test_prefers_top_level_attribute(self):
+        cfg = DFlashConfig(rope_theta=50000.0)
+        self.assertEqual(resolve_rope_theta(cfg), 50000.0)
+
+    def test_falls_back_to_rope_parameters(self):
+        # transformers 5.x absorbs a top-level rope_theta into rope_parameters,
+        # leaving no attribute for the naive read to find.
+        cfg = LlamaConfig(rope_theta=1000000.0)
+        self.assertIsNone(getattr(cfg, "rope_theta", None))
+        self.assertEqual(resolve_rope_theta(cfg), 1000000.0)
+
+    def test_default_when_declared_nowhere(self):
+        cfg = DFlashConfig()
+        del cfg.rope_theta
+        self.assertEqual(resolve_rope_theta(cfg), 10000.0)
+        self.assertIsNone(resolve_rope_theta(cfg, default=None))
+
+    def test_conflicting_values_are_refused(self):
+        cfg = LlamaConfig(rope_theta=1000000.0)
+        cfg.rope_theta = 10000.0  # stale legacy value left beside the block
+        with self.assertRaisesRegex(ValueError, "Conflicting rope_theta"):
+            resolve_rope_theta(cfg)
+
+
+class TestShippedDraftConfigs(unittest.TestCase):
+    """Every config in configs/draft_models/ must resolve to the base it declares."""
+
+    def test_declared_theta_is_what_resolves(self):
+        paths = sorted(glob.glob(os.path.join(DRAFT_CONFIG_DIR, "*.json")))
+        self.assertGreater(len(paths), 0, "no draft configs found")
+        for path in paths:
+            with self.subTest(config=os.path.basename(path)):
+                declared = _declared_theta(json.load(open(path)))
+                if declared is None:
+                    self.skipTest("config declares no rope_theta")
+                cfg = AutoDraftModelConfig.from_file(path)
+                self.assertEqual(resolve_rope_theta(cfg), float(declared))
+
+
+class TestAttentionUsesResolvedTheta(unittest.TestCase):
+    """Each rotary construction path must reach the resolved base, not the default."""
+
+    def test_llama_yarn_path(self):
+        cfg = LlamaConfig(
+            hidden_size=64,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            max_position_embeddings=2048,
+            rope_theta=1000000.0,
+            rope_parameters=dict(YARN),
+        )
+        cfg.target_hidden_size = 64
+        self.assertEqual(LlamaAttention(cfg).rotary_emb.base, 1000000.0)
+
+    def test_deepseek_mla_yarn_path(self):
+        cfg = DeepseekV3Config(
+            hidden_size=64,
+            num_attention_heads=4,
+            q_lora_rank=32,
+            kv_lora_rank=16,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=8,
+            v_head_dim=16,
+            max_position_embeddings=2048,
+            rope_theta=50000.0,
+            rope_parameters=dict(YARN),
+        )
+        self.assertEqual(DeepSeekMLAAttention(cfg).rotary_emb.base, 50000.0)
+
+    def test_dflash_path(self):
+        cfg = DFlashConfig(
+            hidden_size=64,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            max_position_embeddings=512,
+            rope_theta=1000000.0,
+        )
+        self.assertEqual(DFlashAttention(cfg).rotary_emb.base, 1000000.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

torchspec/config/utils.py

@@ -22,6 +22,7 @@
 import json
 import logging
 import warnings
+from typing import Optional
 
 import torch
 from transformers import AutoConfig, AutoTokenizer
@@ -58,6 +59,39 @@ def _normalize_rope_scaling(rope_scaling):
     return normalized
 
 
+def resolve_rope_theta(config, default: Optional[float] = 10000.0) -> Optional[float]:
+    """Return the RoPE base a config asks for, wherever it happens to be stored.
+
+    transformers 5.x keeps `rope_theta` inside `rope_parameters` instead of as a
+    top-level attribute, so reading the attribute alone silently yields `default`
+    for any config loaded from such a checkpoint — a draft then trains against
+    frequencies for 10000 while its config declares something else entirely.
+
+    Configs written by this repo (`DFlashConfig` and friends) declare the
+    attribute themselves and have no `rope_parameters`, so the attribute is
+    preferred and those paths are unaffected. The two disagree only if a stale
+    legacy value was left behind next to an updated block, which is ambiguous
+    enough to refuse rather than silently resolve.
+
+    Pass `default=None` to distinguish "declared nowhere" from a real value.
+    """
+    top_level = getattr(config, "rope_theta", None)
+    params = getattr(config, "rope_parameters", None)
+    nested = params.get("rope_theta") if isinstance(params, dict) else None
+
+    if top_level is not None and nested is not None and float(top_level) != float(nested):
+        raise ValueError(
+            f"Conflicting rope_theta: attribute is {top_level} but rope_parameters "
+            f"carries {nested}. Remove the stale top-level value, or set both to the "
+            "base the checkpoint was trained with."
+        )
+    if top_level is not None:
+        return float(top_level)
+    if nested is not None:
+        return float(nested)
+    return None if default is None else float(default)
+
+
 def generate_draft_model_config(
     target_model_path: str, template_config_path: str = None, cache_dir: str = None
 ):
@@ -139,6 +173,14 @@ def generate_draft_model_config(
             value = _normalize_rope_scaling(value)
         draft_config[draft_param] = value
 
+    # The hasattr copy above cannot see a rope_theta that lives only inside the
+    # target's rope_parameters, which is where transformers 5.x puts it.
+    for source in (text_config, target_config):
+        theta = resolve_rope_theta(source, default=None)
+        if theta is not None:
+            draft_config["rope_theta"] = theta
+            break
+
     draft_config["num_hidden_layers"] = 1
     draft_config["tie_word_embeddings"] = False
     draft_config["use_cache"] = True

torchspec/models/draft/deepseek_eagle.py

@@ -34,19 +34,16 @@
 
 # TODO: Extract shared components into a common module to reduce duplication:
 # - LlamaMLP, LlamaRMSNorm, RoPE classes → torchspec/models/draft/modules.py
-# - _init_rope() is near-identical to LlamaAttention._init_rope (~60 lines)
 # - DecoderLayer.forward() is line-for-line identical to LlamaDecoderLayer.forward()
 # - embed_input_ids/project_hidden_states/compute_logits/backbone are identical
 #   to LlamaForCausalLMEagle3 and could live in Eagle3DraftModel base class
 # - Suffix attention loop could be batched (einsum instead of Python loop) for
 #   both this file and llama3_eagle.py
 from torchspec.models.draft.llama3_eagle import (
-    LlamaDynamicNTKScalingRotaryEmbedding,
-    LlamaLinearScalingRotaryEmbedding,
     LlamaMLP,
     LlamaRMSNorm,
-    LlamaRotaryEmbedding,
-    LlamaYarnRotaryEmbedding,
+    build_rotary_embedding,
+    rope_config_get,
     yarn_get_mscale,
 )
 from torchspec.models.ops.flex_attention import (
@@ -55,16 +52,6 @@
 )
 from torchspec.utils.logging import logger, print_with_rank
 
-
-def _rope_config_get(rope_scaling, key, default=None):
-    """Get a value from rope_scaling config (dict or object)."""
-    if rope_scaling is None:
-        return default
-    if isinstance(rope_scaling, dict):
-        return rope_scaling.get(key, default)
-    return getattr(rope_scaling, key, default)
-
-
 # ── Interleaved RoPE (DeepSeek convention) ────────────────────────────────
 #
 # DeepSeek MLA uses interleaved-pair rotation where consecutive dimension
@@ -157,70 +144,16 @@ def __init__(self, config: DeepseekV3Config):
         self.softmax_scale = self._compute_softmax_scale()
 
     def _init_rope(self):
-        """Initialize rotary embeddings with qk_rope_head_dim as the dimension."""
-        rope_dim = self.qk_rope_head_dim
-        rope_scaling = self.config.rope_scaling
-        rope_theta = getattr(self.config, "rope_theta", 10000)
-        rget = partial(_rope_config_get, rope_scaling)
-        scaling_type = rget("rope_type", rget("type"))
-
-        if rope_scaling is None or scaling_type == "default":
-            self.rotary_emb = LlamaRotaryEmbedding(
-                rope_dim,
-                max_position_embeddings=self.max_position_embeddings,
-                base=rope_theta,
-            )
-        else:
-            scaling_factor = rget("factor")
-
-            if scaling_type in (None, "default"):
-                self.rotary_emb = LlamaRotaryEmbedding(
-                    rope_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    base=rope_theta,
-                )
-            elif scaling_type == "linear":
-                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
-                    rope_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    scaling_factor=scaling_factor,
-                )
-            elif scaling_type == "dynamic":
-                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
-                    rope_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    scaling_factor=scaling_factor,
-                )
-            elif scaling_type == "llama3":
-                self.rotary_emb = LlamaRotaryEmbedding(
-                    rope_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    base=rope_theta,
-                    scaling_factor=scaling_factor if scaling_factor is not None else 1.0,
-                    low_freq_factor=rget("low_freq_factor"),
-                    high_freq_factor=rget("high_freq_factor"),
-                    orig_max_position=rget("original_max_position_embeddings"),
-                )
-            elif scaling_type == "yarn":
-                self.rotary_emb = LlamaYarnRotaryEmbedding(
-                    rope_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    base=rope_theta,
-                    original_max_position_embeddings=rget("original_max_position_embeddings"),
-                    scaling_factor=scaling_factor,
-                    beta_fast=rget("beta_fast"),
-                    beta_slow=rget("beta_slow"),
-                    mscale=rget("mscale"),
-                    mscale_all_dim=rget("mscale_all_dim"),
-                )
-            else:
-                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+        """Rotate only the rope side dims, so the cache is built at qk_rope_head_dim."""
+        self.rotary_emb = build_rotary_embedding(
+            self.config, self.qk_rope_head_dim, self.max_position_embeddings
+        )
 
     def _compute_softmax_scale(self) -> float:
         """Compute softmax scale, incorporating YaRN mscale if applicable."""
         rope_scaling = self.config.rope_scaling
         if rope_scaling is not None:
-            rget = partial(_rope_config_get, rope_scaling)
+            rget = partial(rope_config_get, rope_scaling)
             scaling_type = rget("rope_type", rget("type"))
             if scaling_type == "yarn":
                 factor = rget("factor", 1.0)

torchspec/models/draft/dflash.py

@@ -38,6 +38,8 @@
 from safetensors import safe_open
 from transformers import PretrainedConfig, PreTrainedModel
 
+from torchspec.config.utils import resolve_rope_theta
+
 
 class DFlashConfig(PretrainedConfig):
     """Configuration for DFlash draft model."""
@@ -189,7 +191,7 @@ def __init__(self, config: PretrainedConfig):
         self.rotary_emb = DFlashRotaryEmbedding(
             self.head_dim,
             max_position_embeddings=self.max_position_embeddings,
-            base=getattr(config, "rope_theta", 10000.0),
+            base=resolve_rope_theta(config),
         )
 
     def forward(