[Bug] DSpark draft layers beyond the target's compress_ratios silently fall back to yarn-scaled rope, contradicting the in-tree plain-rope contract

[fdaluiso]

Target repository: https://github.com/jasl/vllm (fork). Verified against tree at commit b5c0d43b.

Environment

• vLLM: jasl fork (github.com/jasl/vllm), DeepSeek-V4 / DSpark enablement branch, tree verified at b5c0d43b.
• Model: hybrid DeepSeek-V4 checkpoint, target num_hidden_layers = 43, target compress_ratios has 44 entries; DSpark speculative drafter of 3 layers, runtime layer ids 43/44/45 (base = num_hidden_layers).
• Draft config: compress_ratios has 46 entries with trailing 0s — i.e. under the draft's own config all three draft layers would resolve to plain (unscaled) rope.
• Speculative decoding: DSpark, vLLM V1 engine.
• Hardware: NVIDIA GB10 (sm_121a) — bug is hardware-independent.

Summary

resolve_layer_compress_ratio returns (1, False) — i.e. yarn-scaled rope — for any layer_id beyond the end of compress_ratios. Because DeepSeekV4DSparkLayer builds its attention from the target's hf_config, and the draft layers' runtime ids start at num_hidden_layers, draft layers can fall past the end of the target's compress_ratios list and silently get yarn-scaled rope (factor 16, original_max_position_embeddings 65536). NVIDIA's own in-tree comment states DSpark draft layers are trained with plain rope. In our configuration 2 of the 3 draft layers run with a positional encoding never seen in training.

Honesty note up front: this bug is proven by code reading and by the tree's own documentation, but our single-boot A/B measurement of its impact was inconclusive (within run-to-run variance). Details below.

Details, with file:line references (verified in-tree)

1. The resolver — vllm/models/deepseek_v4/attention.py:102 (resolve_layer_compress_ratio), body at :113-118:

if layer_id < config.num_hidden_layers:
    return max(1, config.compress_ratios[layer_id]), False
if layer_id < len(config.compress_ratios):
    raw_compress_ratio = config.compress_ratios[layer_id]
    return max(1, raw_compress_ratio), raw_compress_ratio == 0
return 1, False        # <-- out-of-range: use_unscaled_rope=False => yarn

The final fallback (attention.py:118) maps any layer beyond compress_ratios to use_unscaled_rope=False.

2. Which config is consulted — vllm/models/deepseek_v4/nvidia/dspark.py:137 and :146:

config = vllm_config.model_config.hf_config        # <-- the TARGET's hf_config
...
runtime_layer_idx = config.num_hidden_layers + dspark_layer_idx

DeepSeekV4DSparkLayer constructs its attention with the target's hf_config, so the rope decision for draft layers is made against the target's compress_ratios, not the draft's.

3. The arithmetic for our hybrid checkpoint:

• target: num_hidden_layers = 43, len(compress_ratios) = 44
• draft runtime ids: 43, 44, 45
• id 43: in range, compress_ratios[43] == 0 → (1, True) → plain rope ✔
• ids 44, 45: out of range → (1, False) → yarn-scaled rope ✘

4. What rope the "False" branch produces — vllm/models/deepseek_v4/common/rope.py: with use_unscaled_rope=False the rope_type stays "deepseek_yarn" (factor 16, original_max 65536), i.e. remapped frequencies.

5. The tree's own contract — vllm/models/deepseek_v4/common/rope.py:25-28:

if use_unscaled_rope:
    # The MTP draft layer of DSpark-style checkpoints (compress_ratios
    # entry 0) is trained with plain rope, not the yarn-scaled variant.
    rope_parameters["rope_type"] = "default"

NVIDIA documents DSpark-style draft layers as trained with plain rope. The draft checkpoint's own config agrees: its compress_ratios has 46 entries with trailing 0s, under which all three draft layers would be plain.

Root cause

Two independent conventions collide:

1. Draft layers are addressed with runtime ids num_hidden_layers + k but resolved against the target's config (dspark.py:137) — the same root as three other bugs we found in this enablement (dspark_block_size read as 0 from the target config; the KV zeroer packed-stride corruption; the aux-hidden-state off-by-one filed as a companion report). The target config has no dspark_* keys and its compress_ratios was never guaranteed to cover draft layer ids.
2. The out-of-range fallback in resolve_layer_compress_ratio (attention.py:118) chooses (1, False) = yarn — the exact opposite of what the tree itself documents for DSpark draft layers (common/rope.py:25-28). Official checkpoints dodge this only because their compress_ratios happens to be long enough.

Effect: 2 of the 3 draft blocks attend with a positional encoding never seen in training (within a 128-token speculation window, distances appear compressed on the ramped dimensions). The first draft layer is correct, which plausibly bounds the damage.

Proposed fix (minimal)

Make the out-of-range fallback honor the documented DSpark contract — layers beyond the target are draft layers and must get plain rope:

--- a/vllm/models/deepseek_v4/attention.py
+++ b/vllm/models/deepseek_v4/attention.py
@@ def resolve_layer_compress_ratio(config, layer_id: int) -> tuple[int, bool]:
     if layer_id < config.num_hidden_layers:
         return max(1, config.compress_ratios[layer_id]), False
     if layer_id < len(config.compress_ratios):
         raw_compress_ratio = config.compress_ratios[layer_id]
         return max(1, raw_compress_ratio), raw_compress_ratio == 0
-    return 1, False
+    # layer_id >= num_hidden_layers and beyond compress_ratios: this is a
+    # DSpark/MTP draft layer. DSpark drafters are trained with plain rope
+    # (see common/rope.py), not the yarn-scaled variant — do not fall back
+    # to yarn just because the TARGET's compress_ratios is too short.
+    return 1, True

Target layers (layer_id < num_hidden_layers) are untouched; official checkpoints with sufficiently long compress_ratios never reach the changed line. A deeper fix would be to resolve draft-layer rope against the draft's config (which carries the authoritative 46-entry list) instead of the target's — that would also close the whole dspark_*-keys-on-the-wrong-config bug family.

We applied this fix locally as an env-gated site-patch on the installed wheel and measured the impact below.

Measured impact — honest disclosure

• The bug is proven by static reading (the three code sites above), by the tree's own documentation (common/rope.py:25-28), and by the draft checkpoint's config (46 entries, trailing zeros ⇒ all-plain).
• The measured end-to-end impact is inconclusive. Single-boot-per-arm A/B, same probes: aux-fix only (rope bug present) → acceptance pos0 0.780 / mean accepted length 3.19; aux-fix + rope-fix → 0.774 / 3.16. The difference is within run-to-run variance for one boot per arm; we do not claim the fix improves (or worsens) acceptance on our checkpoint. Plausible mitigations of the impact: the first draft layer is correct, and speculation windows are short (≤128 tokens), where yarn ramp distortion is modest.
• We report it anyway because it is a correctness bug against the documented training contract, it silently depends on the accidental length of the target's compress_ratios, and other checkpoints/longer speculation windows may be hit harder.

Reproducibility

1. Static (no GPU): call resolve_layer_compress_ratio(target_config, layer_id) for layer_id in {num_hidden_layers, +1, +2} with a target whose compress_ratios has num_hidden_layers + 1 entries → observe (1, True) for the first and (1, False) (yarn) for the other two.
2. A CPU-only unit test does exactly that, plus an anti-vacuity arm (a faithful transcription of the native function must reproduce the yarn fallback for the out-of-range draft layers) and a non-regression sweep over all 43 target layers — passes, available on request.
3. End-to-end A/B requires multiple boots per arm to beat variance (we flagged our single-boot measurement as inconclusive).

Related

• Companion report: aux hidden-state off-by-one in the V1 runner for dspark_target_layer_ids (filed separately, measured gain 0.756→0.780 pos0).
• Same-root previously found: dspark_block_size read as 0 from the target hf_config (nvidia/dspark.py:145); KV zeroer packed-stride corruption.

---

This report was prepared with AI assistance (Claude, Anthropic); all file:line references were verified by hand against the tree at commit b5c0d43b.