Coverage for transformer_lens/model_bridge/generalized_components/mla_attention.py: 67%
194 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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 q_c = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
88 k_c = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
89 q_rot = torch.view_as_real(q_c * freqs.to(q_c.device)).flatten(3).type_as(q)
90 k_rot = torch.view_as_real(k_c * freqs.to(k_c.device)).flatten(3).type_as(k)
91 return q_rot, k_rot
94class MLAAttentionBridge(PositionEmbeddingHooksMixin, AttentionBridge):
95 """Bridge for DeepSeek's Multi-Head Latent Attention (MLA).
97 Reimplements the MLA forward path with hooks at each computation stage.
98 Standard W_Q/W_K/W_V properties are not available on MLA models — use
99 the submodule weight access (q_a_proj, q_b_proj, etc.) instead.
100 """
102 # MLA has no standard q/k/v submodules — override to empty
103 property_aliases: Dict[str, str] = {}
105 hook_aliases = {
106 "hook_z": "o.hook_in",
107 }
109 # MLA's forward never forks the residual pre-LN; suppress dead HookPoints.
110 supports_split_qkv_fork: bool = False
111 # o_proj is a plain concat-heads projection, so z_h @ W_O_h is well-defined.
112 supports_attn_result: bool = True
114 def __init__(
115 self,
116 name: str,
117 config: Any,
118 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
119 **kwargs: Any,
120 ):
121 super().__init__(name, config, submodules=submodules, **kwargs)
122 self._init_position_embedding_hooks()
123 # The per-head hook_result computation lives in this class's forward; a
124 # subclass that replaces forward (GLM DSA) does not inherit it, so it
125 # must not advertise support it cannot honor.
126 if type(self).forward is not MLAAttentionBridge.forward:
127 self.supports_attn_result = False
129 self.hook_q_latent = HookPoint() # Compressed Q (post q_a_layernorm)
130 self.hook_kv_latent = HookPoint() # Compressed KV (post kv_a_layernorm)
131 self.hook_q = HookPoint() # Final Q entering attention (post-RoPE concat)
132 self.hook_k = HookPoint() # Final K entering attention (post-RoPE concat)
133 self.hook_v = HookPoint() # V from kv_b_proj split
134 self.hook_rot_q = HookPoint() # Q rope portion after RoPE
135 self.hook_rot_k = HookPoint() # K rope portion after RoPE
136 # hook_q/k/v fire after the head transpose, on [batch, heads, pos, d]; the
137 # rotary hooks are transposed back to [batch, pos, heads, d] for the user
138 # by the TransposeRotaryHeads conversion, so they keep the default axis.
139 for hook in (self.hook_q, self.hook_k, self.hook_v):
140 hook.pos_dim = 2
142 # MLA params lazy-initialized from HF module (bridge config lacks these fields)
143 self._mla_params_initialized = False
145 def forward(self, *args: Any, **kwargs: Any) -> Any:
146 """Reimplemented MLA forward with hooks at each computation stage.
148 Follows the DeepseekV3Attention forward path, calling into HF submodules
149 individually and firing hooks at each meaningful stage.
150 """
151 if self.original_component is None: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 raise RuntimeError(
153 f"Original component not set for {self.name}. "
154 "Call set_original_component() first."
155 )
157 hf_attn: Any = self.original_component
159 if not self._mla_params_initialized:
160 self._q_lora_rank = getattr(hf_attn, "q_lora_rank", None)
161 self._kv_lora_rank = getattr(hf_attn, "kv_lora_rank", 512)
162 self._qk_nope_head_dim = getattr(hf_attn, "qk_nope_head_dim", 128)
163 self._qk_rope_head_dim = getattr(hf_attn, "qk_rope_head_dim", 64)
164 self._v_head_dim = getattr(hf_attn, "v_head_dim", 128)
165 self._qk_head_dim = self._qk_nope_head_dim + self._qk_rope_head_dim
166 self._n_heads = getattr(hf_attn, "num_heads", 32)
167 hf_config = getattr(hf_attn, "config", None)
168 self._rope_interleave = (
169 getattr(hf_config, "rope_interleave", False) if hf_config else False
170 )
171 self._mla_params_initialized = True
173 # --- Extract inputs ---
174 if "hidden_states" in kwargs:
175 hidden_states = kwargs.pop("hidden_states")
176 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 176 ↛ 180line 176 didn't jump to line 180 because the condition on line 176 was always true
177 hidden_states = args[0]
178 args = args[1:]
179 else:
180 raise ValueError("Could not find hidden_states in args or kwargs")
182 position_embeddings = kwargs.pop("position_embeddings", None)
183 attention_mask = kwargs.pop("attention_mask", None)
185 hidden_states = self.hook_in(hidden_states)
187 batch_size, seq_length = hidden_states.shape[:2]
189 # --- Query path ---
190 if self._q_lora_rank is None:
191 # Direct projection (no compression)
192 q_states = hf_attn.q_proj(hidden_states)
193 else:
194 # Two-stage compression: q_a_proj → q_a_layernorm → q_b_proj
195 q_compressed = hf_attn.q_a_proj(hidden_states)
196 q_compressed = hf_attn.q_a_layernorm(q_compressed)
197 q_compressed = self.hook_q_latent(q_compressed)
198 q_states = hf_attn.q_b_proj(q_compressed)
200 # Reshape to [batch, n_heads, seq, qk_head_dim]
201 q_states = q_states.view(batch_size, seq_length, -1, self._qk_head_dim).transpose(1, 2)
202 # Split into nope (non-RoPE) and pe (RoPE) portions
203 q_pass, q_rot = torch.split(
204 q_states, [self._qk_nope_head_dim, self._qk_rope_head_dim], dim=-1
205 )
207 # --- KV path ---
208 # kv_a_proj_with_mqa outputs [compressed_kv || k_pe]
209 compressed_kv_full = hf_attn.kv_a_proj_with_mqa(hidden_states)
210 # Split: compressed KV latent (for kv_b_proj) and k rope portion (for direct RoPE)
211 # Note: k_pe is split off here and goes directly to RoPE — hook_kv_latent
212 # captures only the compressed_kv portion that enters the decompression path.
213 k_pass, k_rot = torch.split(
214 compressed_kv_full, [self._kv_lora_rank, self._qk_rope_head_dim], dim=-1
215 )
217 # Compress → normalize → decompress the KV latent
218 k_pass = hf_attn.kv_a_layernorm(k_pass)
219 k_pass = self.hook_kv_latent(k_pass)
220 k_pass = hf_attn.kv_b_proj(k_pass)
222 # Reshape to [batch, n_heads, seq, nope+v_head]
223 key_shape = (batch_size, seq_length, -1, self._qk_nope_head_dim + self._v_head_dim)
224 k_pass = k_pass.view(key_shape).transpose(1, 2)
225 # Split K nope portion and V
226 k_pass, value_states = torch.split(
227 k_pass, [self._qk_nope_head_dim, self._v_head_dim], dim=-1
228 )
230 # k_rot is [batch, seq, rope_dim] → [batch, 1, seq, rope_dim] for broadcasting
231 k_rot = k_rot.view(batch_size, 1, seq_length, self._qk_rope_head_dim)
233 # --- RoPE ---
234 # DeepSeek-V2 passes a complex freqs_cis tensor; V3 passes a (cos, sin) tuple.
235 # Detect the format and apply the appropriate rotation.
236 cos = sin = None
237 if position_embeddings is not None: 237 ↛ 251line 237 didn't jump to line 251 because the condition on line 237 was always true
238 position_embeddings = self._apply_position_embedding_hooks(position_embeddings)
239 if isinstance(position_embeddings, torch.Tensor) and position_embeddings.is_complex():
240 # V2-style: complex exponential freqs_cis
241 q_rot, k_rot = _apply_rotary_complex(q_rot, k_rot, position_embeddings)
242 elif self._rope_interleave: 242 ↛ 246line 242 didn't jump to line 246 because the condition on line 242 was always true
243 cos, sin = position_embeddings
244 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
245 else:
246 cos, sin = position_embeddings
247 if self._rope_interleave:
248 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
249 else:
250 q_rot, k_rot = _apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
251 elif self._rotary_emb is not None:
252 # Fallback: compute from rotary_emb if position_embeddings not passed
253 position_ids = torch.arange(seq_length, device=hidden_states.device).unsqueeze(0)
254 emb = self._rotary_emb(hidden_states, position_ids)
255 if isinstance(emb, torch.Tensor) and emb.is_complex():
256 q_rot, k_rot = _apply_rotary_complex(q_rot, k_rot, emb)
257 elif self._rope_interleave:
258 cos, sin = emb
259 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
260 else:
261 cos, sin = emb
262 if self._rope_interleave:
263 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
264 else:
265 q_rot, k_rot = _apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
266 else:
267 raise ValueError(
268 "MLAAttentionBridge requires position_embeddings or set_rotary_emb() "
269 "to be called before forward."
270 )
271 q_rot = self.hook_rot_q(q_rot)
272 k_rot = self.hook_rot_k(k_rot)
274 # Expand k_rot to match the number of heads
275 k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
277 # Concatenate nope + rope portions to form final Q and K
278 query_states = torch.cat((q_pass, q_rot), dim=-1)
279 key_states = torch.cat((k_pass, k_rot), dim=-1)
281 # Fire final Q/K/V hooks — these are the tensors entering attention
282 query_states = self.hook_q(query_states)
283 key_states = self.hook_k(key_states)
284 value_states = self.hook_v(value_states)
286 # --- KV Cache ---
287 past_key_values = kwargs.pop("past_key_values", None)
288 cache_position = kwargs.pop("cache_position", None)
289 if past_key_values is not None:
290 cache_kwargs: dict = {"cache_position": cache_position}
291 if cos is not None:
292 cache_kwargs["cos"] = cos
293 if sin is not None:
294 cache_kwargs["sin"] = sin
295 key_states, value_states = past_key_values.update(
296 key_states, value_states, hf_attn.layer_idx, cache_kwargs
297 )
299 # --- Attention computation (no V padding — only needed for flash attention) ---
300 # Read the HF module's scaling: DeepSeek-V2/V3 multiply the base
301 # qk_head_dim^-0.5 by yarn mscale^2 under long-context rope configs.
302 # trust_remote_code DeepSeek-V2 names the same quantity softmax_scale.
303 scaling = getattr(hf_attn, "scaling", None)
304 if scaling is None:
305 scaling = getattr(hf_attn, "softmax_scale", None)
306 if scaling is None:
307 scaling = self._qk_head_dim ** (-0.5)
308 attn_scores = torch.matmul(query_states, key_states.transpose(-2, -1)) * scaling
310 if attention_mask is not None:
311 attention_mask = self._normalize_compatibility_mask_sentinel(attention_mask)
312 attn_scores = attn_scores + attention_mask
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)