78 transformer blocks › 3 dense blocks › Sparse latent attention
learnedkey_states
Reconstruct 64 key heads of 256 features: per-head content from the latent and one rotated key shared by all heads.
Why
The authors of latent attention explain that a key has two parts because rotation by position does not go together with the compression. If the rebuilt keys were rotated, the rotation would stand between the query matrix and the key’s up-projection, and the keys of all earlier tokens would have to be recomputed for every new token. So position is carried by a separate rotary part, one per position and shared by all heads, joined to each head’s content part. DeepSeek-V2 §2.1.3
All 5 steps
- 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,
)