78 transformer blocks › 75 MoE blocks › Sparse latent attention › Key heads
KV cache · rotated keys
×state · kept between callspast_key_values.update(k_pass, k_rot, self.layer_idx)
The rotated-key half of the layer’s cache: this call’s key is written, every cached one is read back.
- shape
- k_R [B batch, 1001 cached positions, 64 rotary]
Codemodeling_glm_moe_dsa.py 362–379, 404–413
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 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)Sources1
Contextual evidence
The cache is updated while the key–value latent is still compressed: the normalized latent and the rotated key go in, every cached latent and rotated key come back, and the key and value heads are expanded from all of them on every call.
kv-cache · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 362–413
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–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)