Warn on and normalize partial_rotary_factor declarations the DSpark modeling does not implement

deepspec/modeling/dspark/common.py

@@ -1,3 +1,4 @@
+import warnings
 from dataclasses import dataclass
 from typing import Optional
 
@@ -56,6 +57,54 @@ def extract_context_feature(hidden_states, layer_ids):
     )
 
 
+def declared_partial_rotary_factor(config) -> float:
+    for source in (
+        getattr(config, "rope_parameters", None),
+        getattr(config, "rope_scaling", None),
+    ):
+        if isinstance(source, dict) and source.get("partial_rotary_factor") is not None:
+            return float(source["partial_rotary_factor"])
+    value = getattr(config, "partial_rotary_factor", None)
+    return 1.0 if value is None else float(value)
+
+
+def check_partial_rotary_factor(config, normalize: bool = False):
+    """DSpark draft attention applies RoPE across the full head dim, so a config
+    declaring partial_rotary_factor != 1.0 (e.g. inherited from a Qwen3.5-style
+    target via copy.deepcopy, or from a DFlash warm-start checkpoint) is not what
+    training actually does. Serving engines that honor the field (e.g. vLLM's
+    dflash/dspark path) would then rotate fewer dims than training did, silently
+    collapsing the acceptance length. Warn, and optionally normalize the
+    declaration to 1.0 so saved checkpoints describe the trained geometry."""
+    declared = declared_partial_rotary_factor(config)
+    if declared == 1.0:
+        return config
+    if normalize:
+        for source in (
+            getattr(config, "rope_parameters", None),
+            getattr(config, "rope_scaling", None),
+        ):
+            if isinstance(source, dict) and "partial_rotary_factor" in source:
+                source["partial_rotary_factor"] = 1.0
+        if getattr(config, "partial_rotary_factor", None) is not None:
+            config.partial_rotary_factor = 1.0
+        warnings.warn(
+            f"Draft config declared partial_rotary_factor={declared}, but DSpark "
+            "modeling applies RoPE across the full head dim; normalized the "
+            "declaration to 1.0 so the saved checkpoint matches the trained "
+            "geometry."
+        )
+    else:
+        warnings.warn(
+            f"Config declares partial_rotary_factor={declared}, but DSpark "
+            "modeling applies RoPE across the full head dim and ignores it. "
+            "Serving engines that honor this field (e.g. vLLM) will apply "
+            "different attention geometry than training; set it to 1.0 in the "
+            "checkpoint config before deploying."
+        )
+    return config
+
+
 def validate_target_layer_ids(layer_ids, num_target_layers: int):
     layer_ids = [int(layer_id) for layer_id in layer_ids]
     assert layer_ids, "target_layer_ids must not be empty."

deepspec/modeling/dspark/qwen3/config.py

@@ -1,6 +1,9 @@
 import copy
 
-from deepspec.modeling.dspark.common import validate_target_layer_ids
+from deepspec.modeling.dspark.common import (
+    check_partial_rotary_factor,
+    validate_target_layer_ids,
+)
 
 
 TRAIN_ATTN_IMPLEMENTATION = "flex_attention"
@@ -35,6 +38,10 @@ def build_draft_config(
         )
 
     draft_config = copy.deepcopy(target_config)
+    # The target config may declare partial RoPE (e.g. Qwen3.5 targets), which
+    # the DSpark draft modeling does not implement; normalize the declaration so
+    # saved checkpoints describe the geometry that is actually trained.
+    check_partial_rotary_factor(draft_config, normalize=True)
     draft_config.architectures = ["Qwen3DSparkModel"]
     draft_config.num_target_layers = num_target_layers
     draft_config.num_hidden_layers = num_draft_layers

deepspec/modeling/dspark/qwen3/modeling.py

@@ -21,6 +21,7 @@
     AcceptRatePredictor,
     DSparkForwardOutput,
     build_eval_mask,
+    check_partial_rotary_factor,
     create_dspark_attention_mask,
     create_noise_embed,
     create_position_ids,
@@ -236,6 +237,10 @@ def __init__(self, config) -> None:
             ]
         )
         self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Warn on checkpoints whose config declares partial RoPE: this modeling
+        # rotates the full head dim, so the declaration is untrue and breaks
+        # serving engines that honor it (e.g. vLLM's dflash/dspark path).
+        check_partial_rotary_factor(config)
         self.rotary_emb = Qwen3RotaryEmbedding(config)
         self.fc = nn.Linear(
             len(self.target_layer_ids) * config.hidden_size,