Coverage for transformer_lens/model_bridge/generalized_components/mla_attention.py: 68%
190 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Multi-Head Latent Attention (MLA) bridge component for DeepSeek models.
3MLA compresses Q and KV into lower-dimensional latent spaces via LoRA-style
4projections before standard attention. This component reimplements the MLA
5forward path step-by-step with hooks at each meaningful stage, exposing:
7- hook_q_latent / hook_kv_latent: compressed representations (the information bottleneck)
8- hook_q / hook_k / hook_v: final Q/K/V entering attention (post-decompression, post-RoPE)
9- hook_rot_q / hook_rot_k: after RoPE on the rope portion splits
10- hook_attn_scores / hook_pattern: pre/post-softmax attention weights
11- hook_z: pre-output-projection (alias for o.hook_in)
12- hook_result: per-head output (z_h @ W_O_h, pre-sum), gated by use_attn_result
13"""
15from __future__ import annotations
17from typing import Any, Dict, Optional
19import torch
21from transformer_lens.hook_points import HookPoint
22from transformer_lens.model_bridge.generalized_components.attention import (
23 AttentionBridge,
24)
25from transformer_lens.model_bridge.generalized_components.base import (
26 GeneralizedComponent,
27)
28from transformer_lens.model_bridge.generalized_components.position_embedding_hooks_mixin import (
29 PositionEmbeddingHooksMixin,
30)
33def _rotate_half(x: torch.Tensor) -> torch.Tensor:
34 """Rotate half of the hidden dims of the input (standard RoPE helper)."""
35 x1 = x[..., : x.shape[-1] // 2]
36 x2 = x[..., x.shape[-1] // 2 :]
37 return torch.cat((-x2, x1), dim=-1)
40def _apply_rotary_pos_emb(
41 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
42) -> tuple[torch.Tensor, torch.Tensor]:
43 """Apply rotary position embedding to q and k tensors."""
44 cos = cos.unsqueeze(1) # [batch, 1, seq, dim]
45 sin = sin.unsqueeze(1)
46 q_embed = (q * cos) + (_rotate_half(q) * sin)
47 k_embed = (k * cos) + (_rotate_half(k) * sin)
48 return q_embed, k_embed
51def _apply_rotary_pos_emb_interleave(
52 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
53) -> tuple[torch.Tensor, torch.Tensor]:
54 """Apply interleaved-pair rotary position embedding (transformers >= 5.13).
56 Pairs even-dimension elements with the following odd dimension, matching
57 HF's ``apply_rotary_pos_emb_interleave`` (used by GLM-MoE-DSA and DeepSeek-V3
58 when ``rope_interleave=True``).
59 """
60 cos = cos[..., : cos.shape[-1] // 2].unsqueeze(1)
61 sin = sin[..., : sin.shape[-1] // 2].unsqueeze(1)
62 q1, q2 = q[..., 0::2], q[..., 1::2]
63 k1, k2 = k[..., 0::2], k[..., 1::2]
64 q_embed = torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
65 k_embed = torch.cat([k1 * cos - k2 * sin, k2 * cos + k1 * sin], dim=-1)
66 return q_embed, k_embed
69def _apply_rotary_complex(
70 q: torch.Tensor, k: torch.Tensor, freqs_cis: torch.Tensor
71) -> tuple[torch.Tensor, torch.Tensor]:
72 """Apply rotary position embedding via complex multiplication (DeepSeek-V2 style).
74 DeepSeek-V2 uses ``freqs_cis = torch.polar(ones, freqs)`` (complex exponentials)
75 instead of the standard (cos, sin) pair. This matches the V2 HF implementation of
76 ``apply_rotary_emb``.
78 Args:
79 q: Query rope portion [batch, heads, seq, rope_dim].
80 k: Key rope portion [batch, 1, seq, rope_dim].
81 freqs_cis: Complex rotary frequencies [batch, seq, rope_dim // 2].
83 Returns:
84 Tuple of rotated (q, k) tensors with same dtype and shape as inputs.
85 """
86 freqs = freqs_cis.unsqueeze(1) # [batch, 1, seq, rope_dim // 2]
87 # torch.complex copies, dodging view_as_complex's stride requirements: the rope
88 # slices start at an odd offset whenever qk_nope_head_dim or kv_lora_rank is odd.
89 q_pairs = q.float().reshape(*q.shape[:-1], -1, 2)
90 k_pairs = k.float().reshape(*k.shape[:-1], -1, 2)
91 q_c = torch.complex(q_pairs[..., 0], q_pairs[..., 1])
92 k_c = torch.complex(k_pairs[..., 0], k_pairs[..., 1])
93 q_rot = torch.view_as_real(q_c * freqs.to(q_c.device)).flatten(3).type_as(q)
94 k_rot = torch.view_as_real(k_c * freqs.to(k_c.device)).flatten(3).type_as(k)
95 return q_rot, k_rot
98class MLAAttentionBridge(PositionEmbeddingHooksMixin, AttentionBridge):
99 """Bridge for DeepSeek's Multi-Head Latent Attention (MLA).
101 Reimplements the MLA forward path with hooks at each computation stage.
102 Standard W_Q/W_K/W_V properties are not available on MLA models — use
103 the submodule weight access (q_a_proj, q_b_proj, etc.) instead.
104 """
106 # MLA has no standard q/k/v submodules — override to empty
107 property_aliases: Dict[str, str] = {}
109 hook_aliases = {
110 "hook_z": "o.hook_in",
111 }
113 # MLA's forward never forks the residual pre-LN; suppress dead HookPoints.
114 supports_split_qkv_fork: bool = False
115 # o_proj is a plain concat-heads projection, so z_h @ W_O_h is well-defined.
116 supports_attn_result: bool = True
118 def __init__(
119 self,
120 name: str,
121 config: Any,
122 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
123 **kwargs: Any,
124 ):
125 super().__init__(name, config, submodules=submodules, **kwargs)
126 self._init_position_embedding_hooks()
127 # The per-head hook_result computation lives in this class's forward; a
128 # subclass that replaces forward (GLM DSA) does not inherit it, so it
129 # must not advertise support it cannot honor.
130 if type(self).forward is not MLAAttentionBridge.forward:
131 self.supports_attn_result = False
133 self.hook_q_latent = HookPoint() # Compressed Q (post q_a_layernorm)
134 self.hook_kv_latent = HookPoint() # Compressed KV (post kv_a_layernorm)
135 self.hook_q = HookPoint() # Final Q entering attention (post-RoPE concat)
136 self.hook_k = HookPoint() # Final K entering attention (post-RoPE concat)
137 self.hook_v = HookPoint() # V from kv_b_proj split
138 self.hook_rot_q = HookPoint() # Q rope portion after RoPE
139 self.hook_rot_k = HookPoint() # K rope portion after RoPE
140 # hook_q/k/v fire after the head transpose, on [batch, heads, pos, d]; the
141 # rotary hooks are transposed back to [batch, pos, heads, d] for the user
142 # by the TransposeRotaryHeads conversion, so they keep the default axis.
143 for hook in (self.hook_q, self.hook_k, self.hook_v):
144 hook.pos_dim = 2
146 # MLA params lazy-initialized from HF module (bridge config lacks these fields)
147 self._mla_params_initialized = False
149 def forward(self, *args: Any, **kwargs: Any) -> Any:
150 """Reimplemented MLA forward with hooks at each computation stage.
152 Follows the DeepseekV3Attention forward path, calling into HF submodules
153 individually and firing hooks at each meaningful stage.
154 """
155 if self.original_component is None: 155 ↛ 156line 155 didn't jump to line 156 because the condition on line 155 was never true
156 raise RuntimeError(
157 f"Original component not set for {self.name}. "
158 "Call set_original_component() first."
159 )
161 hf_attn: Any = self.original_component
163 if not self._mla_params_initialized:
164 self._q_lora_rank = getattr(hf_attn, "q_lora_rank", None)
165 self._kv_lora_rank = getattr(hf_attn, "kv_lora_rank", 512)
166 self._qk_nope_head_dim = getattr(hf_attn, "qk_nope_head_dim", 128)
167 self._qk_rope_head_dim = getattr(hf_attn, "qk_rope_head_dim", 64)
168 self._v_head_dim = getattr(hf_attn, "v_head_dim", 128)
169 self._qk_head_dim = self._qk_nope_head_dim + self._qk_rope_head_dim
170 self._n_heads = getattr(hf_attn, "num_heads", 32)
171 hf_config = getattr(hf_attn, "config", None)
172 self._rope_interleave = (
173 getattr(hf_config, "rope_interleave", False) if hf_config else False
174 )
175 self._mla_params_initialized = True
177 # --- Extract inputs ---
178 if "hidden_states" in kwargs:
179 hidden_states = kwargs.pop("hidden_states")
180 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 180 ↛ 184line 180 didn't jump to line 184 because the condition on line 180 was always true
181 hidden_states = args[0]
182 args = args[1:]
183 else:
184 raise ValueError("Could not find hidden_states in args or kwargs")
186 position_embeddings = kwargs.pop("position_embeddings", None)
187 attention_mask = kwargs.pop("attention_mask", None)
189 hidden_states = self.hook_in(hidden_states)
191 batch_size, seq_length = hidden_states.shape[:2]
193 # --- Query path ---
194 if self._q_lora_rank is None:
195 # Direct projection (no compression)
196 q_states = hf_attn.q_proj(hidden_states)
197 else:
198 # Two-stage compression: q_a_proj → q_a_layernorm → q_b_proj
199 q_compressed = hf_attn.q_a_proj(hidden_states)
200 q_compressed = hf_attn.q_a_layernorm(q_compressed)
201 q_compressed = self.hook_q_latent(q_compressed)
202 q_states = hf_attn.q_b_proj(q_compressed)
204 # Reshape to [batch, n_heads, seq, qk_head_dim]
205 q_states = q_states.view(batch_size, seq_length, -1, self._qk_head_dim).transpose(1, 2)
206 # Split into nope (non-RoPE) and pe (RoPE) portions
207 q_pass, q_rot = torch.split(
208 q_states, [self._qk_nope_head_dim, self._qk_rope_head_dim], dim=-1
209 )
211 # --- KV path ---
212 # kv_a_proj_with_mqa outputs [compressed_kv || k_pe]
213 compressed_kv_full = hf_attn.kv_a_proj_with_mqa(hidden_states)
214 # Split: compressed KV latent (for kv_b_proj) and k rope portion (for direct RoPE)
215 # Note: k_pe is split off here and goes directly to RoPE — hook_kv_latent
216 # captures only the compressed_kv portion that enters the decompression path.
217 k_pass, k_rot = torch.split(
218 compressed_kv_full, [self._kv_lora_rank, self._qk_rope_head_dim], dim=-1
219 )
221 # Compress → normalize → decompress the KV latent
222 k_pass = hf_attn.kv_a_layernorm(k_pass)
223 k_pass = self.hook_kv_latent(k_pass)
224 k_pass = hf_attn.kv_b_proj(k_pass)
226 # Reshape to [batch, n_heads, seq, nope+v_head]
227 key_shape = (batch_size, seq_length, -1, self._qk_nope_head_dim + self._v_head_dim)
228 k_pass = k_pass.view(key_shape).transpose(1, 2)
229 # Split K nope portion and V
230 k_pass, value_states = torch.split(
231 k_pass, [self._qk_nope_head_dim, self._v_head_dim], dim=-1
232 )
234 # k_rot is [batch, seq, rope_dim] → [batch, 1, seq, rope_dim] for broadcasting
235 k_rot = k_rot.view(batch_size, 1, seq_length, self._qk_rope_head_dim)
237 # --- RoPE ---
238 # DeepSeek-V2 passes a complex freqs_cis tensor; V3 passes a (cos, sin) tuple.
239 # Detect the format and apply the appropriate rotation.
240 cos = sin = None
241 if position_embeddings is not None: 241 ↛ 252line 241 didn't jump to line 252 because the condition on line 241 was always true
242 position_embeddings = self._apply_position_embedding_hooks(position_embeddings)
243 if isinstance(position_embeddings, torch.Tensor) and position_embeddings.is_complex():
244 # V2-style: complex exponential freqs_cis
245 q_rot, k_rot = _apply_rotary_complex(q_rot, k_rot, position_embeddings)
246 elif self._rope_interleave: 246 ↛ 250line 246 didn't jump to line 250 because the condition on line 246 was always true
247 cos, sin = position_embeddings
248 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
249 else:
250 cos, sin = position_embeddings
251 q_rot, k_rot = _apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
252 elif self._rotary_emb is not None:
253 # Fallback: compute from rotary_emb if position_embeddings not passed
254 position_ids = torch.arange(seq_length, device=hidden_states.device).unsqueeze(0)
255 emb = self._rotary_emb(hidden_states, position_ids)
256 if isinstance(emb, torch.Tensor) and emb.is_complex():
257 q_rot, k_rot = _apply_rotary_complex(q_rot, k_rot, emb)
258 elif self._rope_interleave:
259 cos, sin = emb
260 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
261 else:
262 cos, sin = emb
263 q_rot, k_rot = _apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
264 else:
265 raise ValueError(
266 "MLAAttentionBridge requires position_embeddings or set_rotary_emb() "
267 "to be called before forward."
268 )
269 q_rot = self.hook_rot_q(q_rot)
270 k_rot = self.hook_rot_k(k_rot)
272 # Expand k_rot to match the number of heads
273 k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
275 # Concatenate nope + rope portions to form final Q and K
276 query_states = torch.cat((q_pass, q_rot), dim=-1)
277 key_states = torch.cat((k_pass, k_rot), dim=-1)
279 # Fire final Q/K/V hooks — these are the tensors entering attention
280 query_states = self.hook_q(query_states)
281 key_states = self.hook_k(key_states)
282 value_states = self.hook_v(value_states)
284 # --- KV Cache ---
285 past_key_values = kwargs.pop("past_key_values", None)
286 cache_position = kwargs.pop("cache_position", None)
287 if past_key_values is not None:
288 cache_kwargs: dict = {"cache_position": cache_position}
289 if cos is not None:
290 cache_kwargs["cos"] = cos
291 if sin is not None:
292 cache_kwargs["sin"] = sin
293 key_states, value_states = past_key_values.update(
294 key_states, value_states, hf_attn.layer_idx, cache_kwargs
295 )
297 # --- Attention computation (no V padding — only needed for flash attention) ---
298 # Read the HF module's scaling: DeepSeek-V2/V3 multiply the base
299 # qk_head_dim^-0.5 by yarn mscale^2 under long-context rope configs.
300 # trust_remote_code DeepSeek-V2 names the same quantity softmax_scale.
301 scaling = getattr(hf_attn, "scaling", None)
302 if scaling is None:
303 scaling = getattr(hf_attn, "softmax_scale", None)
304 if scaling is None:
305 scaling = self._qk_head_dim ** (-0.5)
306 attn_scores = torch.matmul(query_states, key_states.transpose(-2, -1)) * scaling
308 # HF's sdpa path passes attention_mask=None and relies on is_causal, so the
309 # causal mask has to be synthesized here when no 4D mask is given.
310 attn_scores = self._apply_reconstruct_attention_mask(
311 attn_scores, attention_mask, seq_len=key_states.shape[-2], q_seq_len=seq_length
312 )
314 attn_scores = self.hook_attn_scores(attn_scores)
315 attn_weights = self._softmax_dropout_pattern(
316 attn_scores, upcast_to_fp32=True, target_dtype=query_states.dtype
317 )
319 # Weighted sum of values
320 attn_output = torch.matmul(attn_weights, value_states)
322 # --- Output projection ---
323 attn_output = attn_output.transpose(1, 2).contiguous()
324 attn_output = attn_output.reshape(batch_size, seq_length, -1)
325 o_bridge = getattr(self, "o", None)
326 if (
327 bool(getattr(self.config, "use_attn_result", False))
328 and isinstance(o_bridge, GeneralizedComponent)
329 and o_bridge.original_component is not None
330 ):
331 # Per-head output pre-sum across heads. Fire hook_z (o.hook_in) on
332 # the pre-projection tensor first so a patch there flows into the
333 # per-head computation — matches the default path, where o_proj's
334 # bridge fires it inside the projection call.
335 attn_output = o_bridge.hook_in(attn_output)
336 # Head count from the tensor, not cfg: remote-code V2 modules may
337 # not expose num_heads, and the flat width is n_heads * v_head_dim
338 # by construction.
339 n_heads = attn_output.shape[-1] // self._v_head_dim
340 z_4d = attn_output.reshape(batch_size, seq_length, n_heads, self._v_head_dim)
341 attn_output = self._compute_per_head_result(z_4d, n_heads, self._v_head_dim)
342 else:
343 attn_output = hf_attn.o_proj(attn_output)
345 attn_output = self.hook_out(attn_output)
346 return attn_output, attn_weights
348 def get_random_inputs(
349 self,
350 batch_size: int = 2,
351 seq_len: int = 8,
352 device: Optional[torch.device] = None,
353 dtype: Optional[torch.dtype] = None,
354 ) -> Dict[str, Any]:
355 """Generate test inputs with hidden_states, position_embeddings, and attention_mask."""
356 if device is None:
357 device = torch.device("cpu")
358 if dtype is None:
359 dtype = torch.float32
361 # Try bridge config (d_model), then HF attention's config (hidden_size), then fallback
362 d_model = None
363 if self.config and hasattr(self.config, "d_model"):
364 d_model = self.config.d_model
365 if d_model is None and self.original_component is not None:
366 hf_cfg = getattr(self.original_component, "config", None)
367 if hf_cfg is not None:
368 d_model = getattr(hf_cfg, "hidden_size", None)
369 if d_model is None:
370 d_model = 256
371 inputs: Dict[str, Any] = {
372 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype)
373 }
375 # Generate position_embeddings from rotary_emb if available,
376 # otherwise create dummy (cos=1, sin=0) embeddings
377 rope_head_dim = self._qk_rope_head_dim if self._mla_params_initialized else 64
378 if self._rotary_emb is not None:
379 try:
380 dummy_input = inputs["hidden_states"]
381 position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
382 position_embeddings = self._rotary_emb(dummy_input, position_ids)
383 inputs["position_embeddings"] = position_embeddings
384 except Exception:
385 cos = torch.ones(1, seq_len, rope_head_dim, device=device, dtype=dtype)
386 sin = torch.zeros(1, seq_len, rope_head_dim, device=device, dtype=dtype)
387 inputs["position_embeddings"] = (cos, sin)
388 else:
389 cos = torch.ones(1, seq_len, rope_head_dim, device=device, dtype=dtype)
390 sin = torch.zeros(1, seq_len, rope_head_dim, device=device, dtype=dtype)
391 inputs["position_embeddings"] = (cos, sin)
393 inputs["attention_mask"] = None
394 return inputs
396 def __getattr__(self, name: str) -> Any:
397 """Raise clear error for standard weight properties that don't apply to MLA."""
398 if name in ("W_Q", "W_K", "W_V", "W_O", "b_Q", "b_K", "b_V", "b_O"):
399 raise NotImplementedError(
400 f"{name} is not available on MLA (Multi-Head Latent Attention) models. "
401 f"MLA uses compressed projections instead of standard Q/K/V. "
402 f"Access weights via submodules: q_a_proj, q_b_proj, kv_a_proj_with_mqa, "
403 f"kv_b_proj, o (o_proj)."
404 )
405 return super().__getattr__(name)