fix(speculative): read rope_theta from rope_parameters first

CHANGELOG.rst

@@ -41,6 +41,7 @@ Changelog
 
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
+- Fix the DFlash draft inheriting the wrong RoPE base from a Transformers 5 target. Such a config can carry both a ``rope_parameters`` dict holding the model's real base and a top-level ``rope_theta`` left at the config-class default, and ModelOpt read the flat field first — so a Qwen3-8B draft trained and exported with ``rope_theta`` 10000 against a target using 1000000. Retrain and re-export any DFlash-family draft (DFlash, Domino, DSpark) built against such a target.
 
 0.46 (2026-08-17)
 ^^^^^^^^^^^^^^^^^

modelopt/torch/export/plugins/hf_spec_export.py

@@ -31,18 +31,26 @@
 
 
 def _get_rope_theta(config, default=None):
-    """Get RoPE theta from either legacy or Transformers 5 config fields."""
-    rope_theta = getattr(config, "rope_theta", None)
-    if rope_theta is not None:
-        return rope_theta
-
+    """Get RoPE theta from either legacy or Transformers 5 config fields.
+
+    ``rope_parameters`` is checked FIRST. A config can carry both fields with
+    different values: Transformers 5 stores the real base under
+    ``rope_parameters`` while the class default (10000.0 for Qwen3) may still be
+    visible as a top-level ``rope_theta``. Reading ``rope_theta`` first silently
+    exports a draft whose RoPE base is 100x off the target's, which breaks
+    serving because DFlash injects the target's KV into every draft layer.
+    """
     # Transformers 5 stores this under rope_parameters (and exposes the same
     # data through rope_scaling for backwards compatibility).
     for attr in ("rope_parameters", "rope_scaling"):
         rope_config = getattr(config, attr, None)
         if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None:
             return rope_config["rope_theta"]
 
+    rope_theta = getattr(config, "rope_theta", None)
+    if rope_theta is not None:
+        return rope_theta
+
     return default
 
 

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -419,10 +419,21 @@ def modify(self, config):
         # overwrite any user value and warn. (rope_scaling is intentionally NOT inherited:
         # DFlash uses standard Qwen3 RotaryEmbedding; the long-context YaRN scaling is
         # added only at export via dflash_export_rope_scaling.)
+        # A config can carry BOTH a top-level rope_theta and a rope_parameters dict
+        # with different values: Transformers 5 keeps the real base in
+        # rope_parameters while the class default (10000.0 for Qwen3) stays visible
+        # as rope_theta. rope_parameters wins, otherwise the draft trains against a
+        # RoPE base 100x off the target's.
+        base_rope_params = getattr(base_config, "rope_parameters", None)
+        if not isinstance(base_rope_params, dict):
+            base_rope_params = {}
         for attr in ("rope_theta", "rope_type", "rope_interleaved"):
-            if not hasattr(base_config, attr):
+            if attr in base_rope_params:
+                base_val = base_rope_params[attr]
+            elif hasattr(base_config, attr):
+                base_val = getattr(base_config, attr)
+            else:
                 continue
-            base_val = getattr(base_config, attr)
             user_val = getattr(self.dflash_config, attr, None)
             if user_val is not None and user_val != base_val:
                 logger.warning(
@@ -434,6 +445,12 @@ def modify(self, config):
                     base_val,
                 )
             setattr(self.dflash_config, attr, base_val)
+            # Qwen3Config populates rope_parameters at construction, so a later
+            # setattr on the flat field alone would leave the dict — which is what
+            # the rotary module reads — holding the stale value.
+            draft_rope_params = getattr(self.dflash_config, "rope_parameters", None)
+            if isinstance(draft_rope_params, dict) and attr in draft_rope_params:
+                draft_rope_params[attr] = base_val
 
         self.dflash_config.head_dim = getattr(
             self.dflash_config,

tests/unit/torch/export/test_hf_spec_rope_export.py

@@ -20,7 +20,11 @@
 
 import torch
 
-from modelopt.torch.export.plugins.hf_spec_export import DFlashExporter, EagleExporter
+from modelopt.torch.export.plugins.hf_spec_export import (
+    DFlashExporter,
+    EagleExporter,
+    _get_rope_theta,
+)
 
 DEFAULT_ROPE_SCALING = {
     "rope_type": "yarn",
@@ -152,3 +156,30 @@ def test_dflash_rope_theta_inherits_base_rope_parameters():
     config = exporter._export_config()
 
     assert config["rope_theta"] == 5000000.0
+
+
+def test_get_rope_theta_prefers_rope_parameters_over_flat_field():
+    """rope_parameters wins when a config carries both fields with different values.
+
+    A real Transformers 5 Qwen3-8B config keeps the true base (1e6) in
+    rope_parameters while the Qwen3Config class default (1e4) stays visible as a
+    top-level rope_theta. Reading the flat field first yields a draft whose RoPE
+    base is 100x off the target's.
+    """
+    config = SimpleNamespace(
+        rope_theta=10000.0,  # class default, NOT the model's real base
+        rope_parameters={"rope_theta": 1000000, "rope_type": "default"},
+    )
+    assert _get_rope_theta(config) == 1000000
+
+
+def test_get_rope_theta_falls_back_to_flat_field():
+    """Legacy configs that only have the flat field still resolve."""
+    config = SimpleNamespace(rope_theta=500000.0, rope_parameters=None)
+    assert _get_rope_theta(config) == 500000.0
+
+
+def test_get_rope_theta_default_when_absent():
+    """No rope information anywhere returns the caller's default."""
+    config = SimpleNamespace(rope_theta=None, rope_parameters=None, rope_scaling=None)
+    assert _get_rope_theta(config, default=1234) == 1234