78 transformer blocks › 75 MoE blocks › Sparse latent attention › Sparse attention › Mix and project
Output projection
×learnedself.o_proj(attn_output)
Map 16,384 head features to the 6,144-feature update.
- in
- [B batch, T positions, 16384 heads × value]
- out
- [B batch, T positions, 6144 d]
- parameters
- 100,663,296 · × 75 MoE blocks = 7,549,747,200
Parameters
100,663,296 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Weight matrix | 100,663,296 | |
| Per-instance total | 100,663,296 | |
This projection has no bias parameter.
The reference module holds the weight output by input, [out, in].
Across 75 independent MoE blocks: parameters.
Derived from the configured matrix dimensions; the model total matches the checkpoint’s tensors.
Codemodeling_glm_moe_dsa.py 293–462
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) attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
attn_output = self.o_proj(attn_output)
return attn_output, attn_weights, topk_indicesSources2
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
Head outputs are flattened and projected back to the model width.
mla-output · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 460–462
Read code excerpt
attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
attn_output = self.o_proj(attn_output)
return attn_output, attn_weights, topk_indices