nemo_automodel.components.models.glm_moe_dsa.rope_utils

Rotary helpers shared by the GLM MoE DSA model and its DSpark draft.

Module Contents

Functions

Name	Description
`mla_softmax_scale`	MLA attention scale `qk_head_dim ** -0.5`, with the YaRN `mscale` correction.

Data

__all__

API

nemo_automodel.components.models.glm_moe_dsa.rope_utils.mla_softmax_scale(
    config,
    qk_head_dim: int
) -> float

MLA attention scale qk_head_dim ** -0.5, with the YaRN mscale correction.

When the rotary parameters carry a full YaRN spec (factor / mscale / original_max_position_embeddings) and the context was extended past the original window, the scale is multiplied by mscale ** 2 — the DeepSeek-V3 MLA convention that GLM-5.2 inherits. Shared so the DSpark draft, trained on this model’s hidden states, cannot drift from the target’s attention temperature.

Parameters:

config

Model config exposing rope_parameters (or the older rope_scaling) and max_position_embeddings.

qk_head_dim

int

Query/key head dimension the scale is derived from.

Returns: float

The softmax scale to pass to the attention kernel.

nemo_automodel.components.models.glm_moe_dsa.rope_utils.__all__ = ['mla_softmax_scale']