SIGN IN SIGN UP

fix(speculative): align the GLM-5.2 DSpark draft RoPE with its target (#3267)

* fix(speculative): align the GLM-5.2 DSpark draft RoPE with its target

The GLM-5.2 DSpark draft trains on the target's hidden states, so any
divergence in its positional encoding or attention temperature is a silent
acceptance loss: the run still converges. Three such divergences existed.

* rope_theta fell back to 10000.0 when it was not under config.rope_parameters,
  while GlmMoeDsaModel reads rope_parameters["rope_theta"] and otherwise the
  top-level field. A draft trained at theta=10000 against a theta=1e6 target
  looks healthy in the loss curve. There is now no default: a config carrying
  no rope_theta anywhere raises.
* The frequency table was built with rope_scaling=None, so a YaRN-scaled target
  produced an unscaled draft. It now forwards the same rope_scaling expression
  GlmMoeDsaModel uses.
* The softmax scale omitted the YaRN mscale**2 correction that
  GlmMoeDsaAttention applies, so the draft attended at a different temperature
  than the target on any extended-context checkpoint.

The mscale derivation was already duplicated per model, so it moves into
mla_softmax_scale next to yarn_get_mscale in deepseek_v3/rope_utils.py, and the
GLM target attention and the draft now share it: the pair whose drift this
change is about cannot separate again. The identical copies in the DeepSeek V3
and V3.2 attention layers are left alone here (their tests patch the imported
yarn_get_mscale symbol); adopting the helper there is a mechanical follow-up.

qk_head_dim also follows the target's precedence (an explicit config field wins
over the nope+rope sum), so the draft's q_b_proj width and scale are derived the
same way on both sides.

Testing: added unit tests for the shared helper and for the draft's scale, rope
table, rope_theta source, and missing-rope_theta failure; ran the speculative,
glm_moe_dsa, deepseek_v3 and deepseek_v32 unit tests on CPU. No GPU run is
needed: released GLM-5.2 ships rope_type "default", so this is a no-op for the
current checkpoint and only changes behavior on a YaRN-scaled or theta-relocated
config.

Signed-off-by: khazic <khazzz1c@gmail.com>

* fix(speculative): move mla_softmax_scale into glm_moe_dsa

The helper is only used by the GLM MoE DSA attention and its DSpark draft,
so it lives next to them in glm_moe_dsa/rope_utils.py instead of in the
DeepSeek-V3 rope utilities. This also leaves the DeepSeek-V3 module and its
tests untouched by this PR.

Signed-off-by: khazic <khazzz1c@gmail.com>

* chore: retrigger CI

Signed-off-by: khazic <khazzz1c@gmail.com>

---------

Signed-off-by: khazic <khazzz1c@gmail.com>
K
khazzz1c committed
7e6a9fb15da0bf575f0f7c095edf161ccae0aa77
Parent: 8f4fdfe
Committed by GitHub <noreply@github.com> on 8/10/2026, 5:18:40 PM