Coverage for transformer_lens/model_bridge/generalized_components/glm_moe_dsa_attention.py: 72%
125 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""GLM-MoE-DSA attention bridge component."""
2from __future__ import annotations
4from typing import Any, Dict, Optional
6import torch
7import torch.nn.functional as F
9from transformer_lens.hook_points import HookPoint
10from transformer_lens.model_bridge.generalized_components.base import (
11 GeneralizedComponent,
12)
13from transformer_lens.model_bridge.generalized_components.mla_attention import (
14 MLAAttentionBridge,
15)
18def _apply_rotary_pos_emb_single(
19 x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int
20) -> torch.Tensor:
21 """Apply interleaved-pair rotary position embeddings (transformers >= 5.13).
23 HF 5.13 switched GLM-MoE-DSA from split-half NeoX-style RoPE to interleaved-
24 pair rotation (``apply_rotary_pos_emb_interleave``). Even-dimension elements
25 are paired with the following odd dimension: (d0,d1), (d2,d3), …
26 """
27 cos = cos[..., : cos.shape[-1] // 2].unsqueeze(unsqueeze_dim)
28 sin = sin[..., : sin.shape[-1] // 2].unsqueeze(unsqueeze_dim)
29 x1, x2 = x[..., 0::2], x[..., 1::2]
30 return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
33def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
34 batch, num_key_value_heads, slen, head_dim = hidden_states.shape
35 if n_rep == 1: 35 ↛ 37line 35 didn't jump to line 37 because the condition on line 35 was always true
36 return hidden_states
37 hidden_states = hidden_states[:, :, None, :, :].expand(
38 batch, num_key_value_heads, n_rep, slen, head_dim
39 )
40 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
43class GlmMoeDsaAttentionBridge(MLAAttentionBridge):
44 """Bridge for GLM-5 DeepSeek Sparse Attention.
46 GLM-MoE-DSA extends MLA with a learned top-k token indexer and returns
47 ``(attn_output, attn_weights, topk_indices_or_none)`` to feed shared
48 top-k indices into later layers.
49 """
51 def __init__(
52 self,
53 name: str,
54 config: Any,
55 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
56 **kwargs: Any,
57 ):
58 super().__init__(name, config, submodules=submodules, **kwargs)
59 self.hook_topk_indices = HookPoint()
60 self.hook_dsa_mask = HookPoint()
62 def forward(self, *args: Any, **kwargs: Any) -> Any:
63 if self.original_component is None: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true
64 raise RuntimeError(
65 f"Original component not set for {self.name}. "
66 "Call set_original_component() first."
67 )
69 hf_attn: Any = self.original_component
71 if not self._mla_params_initialized: 71 ↛ 83line 71 didn't jump to line 83 because the condition on line 71 was always true
72 self._q_lora_rank = getattr(hf_attn, "q_lora_rank", None)
73 self._kv_lora_rank = getattr(hf_attn, "kv_lora_rank")
74 self._qk_nope_head_dim = getattr(hf_attn, "qk_nope_head_dim")
75 self._qk_rope_head_dim = getattr(hf_attn, "qk_rope_head_dim")
76 self._v_head_dim = getattr(hf_attn, "v_head_dim")
77 self._qk_head_dim = getattr(
78 hf_attn, "qk_head_dim", self._qk_nope_head_dim + self._qk_rope_head_dim
79 )
80 self._n_heads = getattr(hf_attn, "num_heads")
81 self._mla_params_initialized = True
83 if "hidden_states" in kwargs: 83 ↛ 85line 83 didn't jump to line 85 because the condition on line 83 was always true
84 hidden_states = kwargs.pop("hidden_states")
85 elif len(args) > 0 and isinstance(args[0], torch.Tensor):
86 hidden_states = args[0]
87 args = args[1:]
88 else:
89 raise ValueError("Could not find hidden_states in args or kwargs")
91 position_embeddings = kwargs.pop("position_embeddings", None)
92 attention_mask = kwargs.pop("attention_mask", None)
93 past_key_values = kwargs.pop("past_key_values", None)
94 prev_topk_indices = kwargs.pop("prev_topk_indices", None)
95 position_ids = kwargs.pop("position_ids", None)
97 hidden_states = self.hook_in(hidden_states)
98 batch_size, seq_length = hidden_states.shape[:-1]
100 if self._q_lora_rank is None: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true
101 query_states = hf_attn.q_proj(hidden_states)
102 q_resid = None
103 else:
104 q_resid = hf_attn.q_a_layernorm(hf_attn.q_a_proj(hidden_states))
105 q_resid = self.hook_q_latent(q_resid)
106 query_states = hf_attn.q_b_proj(q_resid)
108 query_states = query_states.view(batch_size, seq_length, -1, self._qk_head_dim).transpose(
109 1, 2
110 )
111 q_nope, q_pe = torch.split(
112 query_states, [self._qk_nope_head_dim, self._qk_rope_head_dim], dim=-1
113 )
115 compressed_kv = hf_attn.kv_a_proj_with_mqa(hidden_states)
116 k_compressed, k_pe = torch.split(
117 compressed_kv, [self._kv_lora_rank, self._qk_rope_head_dim], dim=-1
118 )
119 k_compressed = hf_attn.kv_a_layernorm(k_compressed)
120 k_compressed = self.hook_kv_latent(k_compressed)
122 kv_expanded = hf_attn.kv_b_proj(k_compressed)
123 kv_expanded = kv_expanded.view(
124 batch_size, seq_length, -1, self._qk_nope_head_dim + self._v_head_dim
125 )
126 k_nope, value_states = torch.split(
127 kv_expanded, [self._qk_nope_head_dim, self._v_head_dim], dim=-1
128 )
129 k_nope = k_nope.transpose(1, 2)
130 value_states = value_states.transpose(1, 2)
132 if position_embeddings is not None: 132 ↛ 135line 132 didn't jump to line 135 because the condition on line 132 was always true
133 position_embeddings = self._apply_position_embedding_hooks(position_embeddings)
134 cos, sin = position_embeddings
135 elif self._rotary_emb is not None:
136 position_ids = torch.arange(seq_length, device=hidden_states.device).unsqueeze(0)
137 cos, sin = self._rotary_emb(hidden_states, position_ids)
138 position_embeddings = (cos, sin)
139 else:
140 raise ValueError(
141 "GlmMoeDsaAttentionBridge requires position_embeddings or set_rotary_emb()."
142 )
144 q_pe = _apply_rotary_pos_emb_single(q_pe, cos, sin, unsqueeze_dim=1)
145 k_pe = k_pe.view(batch_size, 1, seq_length, self._qk_rope_head_dim)
146 k_pe = _apply_rotary_pos_emb_single(k_pe, cos, sin, unsqueeze_dim=1)
147 q_pe = self.hook_rot_q(q_pe)
148 k_pe = self.hook_rot_k(k_pe)
149 k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1)
151 query_states = torch.cat([q_nope, q_pe], dim=-1)
152 key_states = torch.cat([k_nope, k_pe], dim=-1)
153 query_states = self.hook_q(query_states)
154 key_states = self.hook_k(key_states)
155 value_states = self.hook_v(value_states)
157 if past_key_values is not None: 157 ↛ 162line 157 didn't jump to line 162 because the condition on line 157 was always true
158 key_states, value_states = past_key_values.update(
159 key_states, value_states, hf_attn.layer_idx
160 )
162 if hf_attn.indexer is not None:
163 if attention_mask is not None and attention_mask.dim() == 4: 163 ↛ 165line 163 didn't jump to line 165 because the condition on line 163 was always true
164 indexer_mask = attention_mask[:, 0, :, :]
165 elif attention_mask is not None:
166 indexer_mask = attention_mask.unsqueeze(1)
167 else:
168 indexer_mask = None
169 if position_ids is None: 169 ↛ 170line 169 didn't jump to line 170 because the condition on line 169 was never true
170 position_ids = torch.arange(seq_length, device=hidden_states.device).unsqueeze(0)
171 topk_indices = hf_attn.indexer(
172 hidden_states,
173 q_resid,
174 position_embeddings,
175 indexer_mask,
176 position_ids,
177 past_key_values=past_key_values,
178 )
179 else:
180 if prev_topk_indices is None: 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 raise ValueError(
182 "Shared DSA layers require top-k indices from a previous "
183 "full indexer layer (prev_topk_indices is None)."
184 )
185 topk_indices = prev_topk_indices
186 topk_indices = self.hook_topk_indices(topk_indices)
188 total_len = key_states.shape[2]
189 index_mask = torch.full(
190 (batch_size, seq_length, total_len),
191 float("-inf"),
192 device=hidden_states.device,
193 dtype=query_states.dtype,
194 )
195 index_mask.scatter_(-1, topk_indices, 0.0)
196 index_mask = self.hook_dsa_mask(index_mask).unsqueeze(1)
197 if attention_mask is not None and attention_mask.dim() == 4: 197 ↛ 199line 197 didn't jump to line 199 because the condition on line 197 was always true
198 attn_scores_mask = index_mask + attention_mask[..., :total_len]
199 elif attention_mask is not None:
200 attn_scores_mask = attention_mask.masked_fill(
201 index_mask == float("-inf"), float("-inf")
202 )
203 else:
204 causal_mask = (
205 torch.arange(total_len, device=hidden_states.device)[None, None, None, :]
206 > torch.arange(q_pe.shape[-2], device=hidden_states.device)[:, None, None]
207 )
208 index_mask = index_mask.masked_fill(causal_mask, float("-inf"))
209 attn_scores_mask = index_mask
211 key_states = _repeat_kv(key_states, hf_attn.num_key_value_groups)
212 value_states = _repeat_kv(value_states, hf_attn.num_key_value_groups)
213 attn_scores = torch.matmul(query_states, key_states.transpose(2, 3)) * hf_attn.scaling
214 attn_scores = attn_scores + attn_scores_mask
215 attn_scores = self.hook_attn_scores(attn_scores)
216 attn_weights = self._softmax_dropout_pattern(
217 attn_scores, upcast_to_fp32=True, target_dtype=query_states.dtype
218 )
219 if self.training and hf_attn.attention_dropout: 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true
220 attn_weights = F.dropout(attn_weights, p=hf_attn.attention_dropout, training=True)
222 attn_output = torch.matmul(attn_weights, value_states)
223 attn_output = attn_output.transpose(1, 2).contiguous()
224 attn_output = attn_output.reshape(batch_size, seq_length, -1)
225 attn_output = hf_attn.o_proj(attn_output)
226 attn_output = self.hook_out(attn_output)
227 return attn_output, attn_weights, topk_indices