78 transformer blocks › 75 MoE blocks › Sparse latent attention › Key–value compression
Key–value latent RMSNorm
×learnedself.kv_a_layernorm(kv_pass)
Rescale each position’s features using their root mean square, then apply learned gain.
- in
- [B batch, T positions, 512 kv latent]
- out
- [B batch, T positions, 512 kv latent]
- parameters
- 512 · × 75 MoE blocks = 38,400
Normalized in float32, cast back to the input type, then multiplied by the gain (modeling_glm_moe_dsa.py 58–62).
Parameters
512 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Learned gain | 512 | |
| Per-instance total | 512 | |
Each normalization instance owns its learned gain vector.
Across 75 independent MoE blocks: parameters.
Derived from the configured feature width; the model total matches the checkpoint’s tensors.
Codemodeling_glm_moe_dsa.py 48–403
class GlmMoeDsaRMSNorm(nn.Module):
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
"""
GlmMoeDsaRMSNorm is equivalent to T5LayerNorm
"""
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)def yarn_get_mscale(scale=1, mscale=1):
if scale <= 1:
return 1.0
return 0.1 * mscale * math.log(scale) + 1.0
def yarn_apply_mscale(rope_parameters, scaling):
if rope_parameters.get("rope_type", "default") != "default":
mscale_all_dim = rope_parameters.get("mscale_all_dim", 0)
scaling_factor = rope_parameters["factor"]
if mscale_all_dim:
mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
scaling = scaling * mscale * mscale
return scaling self.q_lora_rank = config.q_lora_rank
self.qk_rope_head_dim = config.qk_rope_head_dim
self.kv_lora_rank = config.kv_lora_rank
self.v_head_dim = config.v_head_dim
self.qk_nope_head_dim = config.qk_nope_head_dim
self.qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
self.is_causal = True
self.q_a_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=config.attention_bias)
self.q_a_layernorm = GlmMoeDsaRMSNorm(self.q_lora_rank)
self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
self.kv_a_proj_with_mqa = nn.Linear(
self.hidden_size,
config.kv_lora_rank + config.qk_rope_head_dim,
bias=config.attention_bias,
)
self.kv_a_layernorm = GlmMoeDsaRMSNorm(config.kv_lora_rank)
self.kv_b_proj = nn.Linear(
config.kv_lora_rank,
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
bias=False,
)
self.o_proj = nn.Linear(
self.num_heads * self.v_head_dim,
self.hidden_size,
bias=config.attention_bias,
)
self.scaling = yarn_apply_mscale(config.rope_parameters, self.qk_head_dim ** (-0.5))
# Refer: https://arxiv.org/abs/2603.12201 for more details.
self.skip_topk = config.indexer_types[layer_idx] == "shared"
self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx) def expand_kv(self, kv_nope: torch.Tensor, k_rot: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Expands the compressed latents into key and value states. Args:
- kv_nope: key + value without positional encoding, shape [batch_size, 1, seqlen, self.kv_lora_rank]
- k_rot: shared key with positional encoding, shape [batch_size, 1, seqlen, self.qk_rope_head_dim]
Returns the key and value states, two tensors of shape [batch, num_heads, seq, (k or v)_head_dim].
"""
batch_size, _, seq_length, _ = kv_nope.shape
key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
kv_nope = self.kv_b_proj(kv_nope).view(key_shape).transpose(1, 2)
k_nope, value_states = torch.split(kv_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
k_rot = k_rot.expand(-1, k_nope.shape[1], -1, -1) # does not affect the underlying storage
# Concatenate k_nope and k_rot in a new tensor
key_states = kv_nope.new_empty(*kv_nope.shape[:-1], self.qk_nope_head_dim + self.qk_rope_head_dim)
key_states[..., : self.qk_nope_head_dim].copy_(k_nope)
key_states[..., self.qk_nope_head_dim :].copy_(k_rot)
return key_states, value_states compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
kv_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
# Both latents are viewed as single-head, 4D tensors, as expected by `expand_kv`
k_pass = self.kv_a_layernorm(kv_pass).view(batch_size, 1, seq_length, self.kv_lora_rank)
k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)Sources3
Direct evidence
RMSNorm owns one gain per feature; it normalizes in float32, casts back to the input precision and multiplies by the gain.
rmsnorm · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 48–62
Read code excerpt
class GlmMoeDsaRMSNorm(nn.Module):
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
"""
GlmMoeDsaRMSNorm is equivalent to T5LayerNorm
"""
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)Direct evidence
MLA owns a query down-projection, query RMSNorm and up-projection, a joint key–value down-projection that also emits the shared rotary key, a latent RMSNorm, a key–value up-projection and an output projection; the latent RMSNorms are built without ε and take the default 10⁻⁶; the softmax scale is 1/√256, with no YaRN correction under default RoPE; every layer builds its own indexer.
mla-init · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 293–360
Read code excerpt
def yarn_get_mscale(scale=1, mscale=1):
if scale <= 1:
return 1.0
return 0.1 * mscale * math.log(scale) + 1.0
def yarn_apply_mscale(rope_parameters, scaling):
if rope_parameters.get("rope_type", "default") != "default":
mscale_all_dim = rope_parameters.get("mscale_all_dim", 0)
scaling_factor = rope_parameters["factor"]
if mscale_all_dim:
mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
scaling = scaling * mscale * mscale
return scaling… lines 307–325 omitted …
self.q_lora_rank = config.q_lora_rank
self.qk_rope_head_dim = config.qk_rope_head_dim
self.kv_lora_rank = config.kv_lora_rank
self.v_head_dim = config.v_head_dim
self.qk_nope_head_dim = config.qk_nope_head_dim
self.qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
self.is_causal = True
self.q_a_proj = nn.Linear(self.hidden_size, self.q_lora_rank, bias=config.attention_bias)
self.q_a_layernorm = GlmMoeDsaRMSNorm(self.q_lora_rank)
self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
self.kv_a_proj_with_mqa = nn.Linear(
self.hidden_size,
config.kv_lora_rank + config.qk_rope_head_dim,
bias=config.attention_bias,
)
self.kv_a_layernorm = GlmMoeDsaRMSNorm(config.kv_lora_rank)
self.kv_b_proj = nn.Linear(
config.kv_lora_rank,
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
bias=False,
)
self.o_proj = nn.Linear(
self.num_heads * self.v_head_dim,
self.hidden_size,
bias=config.attention_bias,
)
self.scaling = yarn_apply_mscale(config.rope_parameters, self.qk_head_dim ** (-0.5))
# Refer: https://arxiv.org/abs/2603.12201 for more details.
self.skip_topk = config.indexer_types[layer_idx] == "shared"
self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx)Direct evidence
The joint down-projection yields the key–value latent and one rotary key shared by all heads; the latent is normalized, and kv_b_proj reconstructs per-head content keys and values from it.
mla-kv · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 362–403
Read code excerpt
def expand_kv(self, kv_nope: torch.Tensor, k_rot: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Expands the compressed latents into key and value states. Args:
- kv_nope: key + value without positional encoding, shape [batch_size, 1, seqlen, self.kv_lora_rank]
- k_rot: shared key with positional encoding, shape [batch_size, 1, seqlen, self.qk_rope_head_dim]
Returns the key and value states, two tensors of shape [batch, num_heads, seq, (k or v)_head_dim].
"""
batch_size, _, seq_length, _ = kv_nope.shape
key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
kv_nope = self.kv_b_proj(kv_nope).view(key_shape).transpose(1, 2)
k_nope, value_states = torch.split(kv_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
k_rot = k_rot.expand(-1, k_nope.shape[1], -1, -1) # does not affect the underlying storage
# Concatenate k_nope and k_rot in a new tensor
key_states = kv_nope.new_empty(*kv_nope.shape[:-1], self.qk_nope_head_dim + self.qk_rope_head_dim)
key_states[..., : self.qk_nope_head_dim].copy_(k_nope)
key_states[..., self.qk_nope_head_dim :].copy_(k_rot)
return key_states, value_states… lines 380–397 omitted …
compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
kv_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
# Both latents are viewed as single-head, 4D tensors, as expected by `expand_kv`
k_pass = self.kv_a_layernorm(kv_pass).view(batch_size, 1, seq_length, self.kv_lora_rank)
k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)