78 transformer blocks › 75 MoE blocks › Sparse latent attention › Sparse attention › Attention weights
Apply sparse mask
×no parametersattn_weights + attention_mask
Restrict each query to the positions the indexer selected.
- in
- [B batch, 64 heads, T queries, T keys] [B batch, T queries, T keys]
- out
- [B batch, 64 heads, T queries, T keys]
Codemodeling_glm_moe_dsa.py 268–290, 404–458
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,
)Sources1
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,
)