- in
- [B batch, T positions, 512 kv latent] [B batch, T positions, 6144 d] [T positions, 32 pairs]
- out
- [B batch, 64 heads, T positions, 256 content + rotary]
- parameters
- 6,684,672 · × 3 dense blocks = 20,054,016
Parameters
6,684,672 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Key content up-projection | 6,291,456 | |
| Rotary key projection | 393,216 | |
| Per-instance total | 6,684,672 | |
Across 3 independent dense blocks: parameters.
Derived from the pinned configuration and transformers implementation, and equal to the count of the checkpoint’s tensors read from its safetensors headers, less the multi-token-prediction layer the implementation does not load. The total includes the routers’ 19,200 correction-bias values, which the implementation holds as buffers.
Explanation
One matrix, kv_b_proj, produces the content keys and the values together; the implementation splits them after the product (modeling_glm_moe_dsa.py 371–372).
When decoding, the implementation rebuilds every head’s keys from the cached latent at each step (modeling_glm_moe_dsa.py 407–413).
Codemodeling_glm_moe_dsa.py 268–458
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)def eager_attention_forward(
module: nn.Module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None,
scaling: float,
dropout: float = 0.0,
**kwargs: Unpack[TransformersKwargs],
):
key_states = repeat_kv(key, module.num_key_value_groups)
value_states = repeat_kv(value, module.num_key_value_groups)
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).contiguous()
return attn_output, attn_weights cos, sin = position_embeddings
q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
# Cache read / write is performed while latent KV is still compressed
if past_key_values is not None:
k_pass, k_rot = past_key_values.update(k_pass, k_rot, self.layer_idx)
query_states = torch.cat((q_pass, q_rot), dim=-1)
key_states, value_states = self.expand_kv(k_pass, k_rot)
# DSA: select this layer's top-k tokens, or reuse the previous full layer's on `"shared"` layers.
if self.indexer is not None:
topk_indices = self.indexer(
hidden_states,
q_resid,
position_embeddings,
attention_mask[:, 0, :, :],
position_ids, # Kept for BC
past_key_values=past_key_values,
) # [B, S, topk]
else:
if prev_topk_indices is None:
raise ValueError("Shared DSA layers require top-k indices from a previous full indexer layer.")
topk_indices = prev_topk_indices
sparse_indices = None
if self.config._attn_implementation in ("eager", "sdpa"):
index_mask = (
topk_indices.new_ones((batch_size, seq_length, key_states.shape[2]), dtype=torch.bool)
.scatter(-1, topk_indices.long(), False)
.unsqueeze(1)
)
if attention_mask.dtype == torch.bool:
attention_mask = attention_mask & ~index_mask
else:
attention_mask = attention_mask.masked_fill(index_mask, torch.finfo(hidden_states.dtype).min)
else:
sparse_indices = topk_indices
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
self.config._attn_implementation, eager_attention_forward
)
attn_output, attn_weights = attention_interface(
self,
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.attention_dropout,
scaling=self.scaling,
indices=sparse_indices, # consumed by flash_mla_with_kvcache; ignored by eager / SDPA
**kwargs,
)Sources3
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)Direct evidence
The rotary parts are rotated; the normalized latent and the rotated key are what the cache stores; each head joins its content and rotary parts; the indexer’s selection becomes a mask that sets unselected positions to the dtype’s most negative value on top of the causal mask; attention is scaled dot-product attention with a float32 softmax.
mla-prefill · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 268–458
Read code excerpt
def eager_attention_forward(
module: nn.Module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None,
scaling: float,
dropout: float = 0.0,
**kwargs: Unpack[TransformersKwargs],
):
key_states = repeat_kv(key, module.num_key_value_groups)
value_states = repeat_kv(value, module.num_key_value_groups)
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).contiguous()
return attn_output, attn_weights… lines 291–403 omitted …
cos, sin = position_embeddings
q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
# Cache read / write is performed while latent KV is still compressed
if past_key_values is not None:
k_pass, k_rot = past_key_values.update(k_pass, k_rot, self.layer_idx)
query_states = torch.cat((q_pass, q_rot), dim=-1)
key_states, value_states = self.expand_kv(k_pass, k_rot)
# DSA: select this layer's top-k tokens, or reuse the previous full layer's on `"shared"` layers.
if self.indexer is not None:
topk_indices = self.indexer(
hidden_states,
q_resid,
position_embeddings,
attention_mask[:, 0, :, :],
position_ids, # Kept for BC
past_key_values=past_key_values,
) # [B, S, topk]
else:
if prev_topk_indices is None:
raise ValueError("Shared DSA layers require top-k indices from a previous full indexer layer.")
topk_indices = prev_topk_indices
sparse_indices = None
if self.config._attn_implementation in ("eager", "sdpa"):
index_mask = (
topk_indices.new_ones((batch_size, seq_length, key_states.shape[2]), dtype=torch.bool)
.scatter(-1, topk_indices.long(), False)
.unsqueeze(1)
)
if attention_mask.dtype == torch.bool:
attention_mask = attention_mask & ~index_mask
else:
attention_mask = attention_mask.masked_fill(index_mask, torch.finfo(hidden_states.dtype).min)
else:
sparse_indices = topk_indices
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
self.config._attn_implementation, eager_attention_forward
)
attn_output, attn_weights = attention_interface(
self,
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.attention_dropout,
scaling=self.scaling,
indices=sparse_indices, # consumed by flash_mla_with_kvcache; ignored by eager / SDPA
**kwargs,
)