78 transformer blocks › 75 MoE blocks › Sparse latent attention
KV cache
×state · kept between callspast_key_values.update(k_pass, k_rot, self.layer_idx)
The normalized key–value latent and the rotated key of every position so far, kept between forward calls.
- shape
- c_KV [B batch, 1001 cached positions, 512 kv latent] k_R [B batch, 1001 cached positions, 64 rotary]
This call writes the new token’s latent and rotated key; the cache returns all 1,001 positions, and the key and value heads are expanded from every one of them again. The rotated key is made inside the key heads, where its write and read are drawn. 576 values per position, against 32,768 for the expanded keys and values. The cache has no fixed size; the slider stops at the configuration’s 202,752 positions (max_position_embeddings).
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)