Coverage for transformer_lens/model_bridge/generalized_components/mla_attention.py: 65%
183 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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"""
14from __future__ import annotations
16from typing import Any, Dict, Optional
18import torch
20from transformer_lens.hook_points import HookPoint
21from transformer_lens.model_bridge.generalized_components.attention import (
22 AttentionBridge,
23)
24from transformer_lens.model_bridge.generalized_components.base import (
25 GeneralizedComponent,
26)
27from transformer_lens.model_bridge.generalized_components.position_embedding_hooks_mixin import (
28 PositionEmbeddingHooksMixin,
29)
32def _rotate_half(x: torch.Tensor) -> torch.Tensor:
33 """Rotate half of the hidden dims of the input (standard RoPE helper)."""
34 x1 = x[..., : x.shape[-1] // 2]
35 x2 = x[..., x.shape[-1] // 2 :]
36 return torch.cat((-x2, x1), dim=-1)
39def _apply_rotary_pos_emb(
40 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
41) -> tuple[torch.Tensor, torch.Tensor]:
42 """Apply rotary position embedding to q and k tensors."""
43 cos = cos.unsqueeze(1) # [batch, 1, seq, dim]
44 sin = sin.unsqueeze(1)
45 q_embed = (q * cos) + (_rotate_half(q) * sin)
46 k_embed = (k * cos) + (_rotate_half(k) * sin)
47 return q_embed, k_embed
50def _apply_rotary_pos_emb_interleave(
51 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
52) -> tuple[torch.Tensor, torch.Tensor]:
53 """Apply interleaved-pair rotary position embedding (transformers >= 5.13).
55 Pairs even-dimension elements with the following odd dimension, matching
56 HF's ``apply_rotary_pos_emb_interleave`` (used by GLM-MoE-DSA and DeepSeek-V3
57 when ``rope_interleave=True``).
58 """
59 cos = cos[..., : cos.shape[-1] // 2].unsqueeze(1)
60 sin = sin[..., : sin.shape[-1] // 2].unsqueeze(1)
61 q1, q2 = q[..., 0::2], q[..., 1::2]
62 k1, k2 = k[..., 0::2], k[..., 1::2]
63 q_embed = torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
64 k_embed = torch.cat([k1 * cos - k2 * sin, k2 * cos + k1 * sin], dim=-1)
65 return q_embed, k_embed
68def _apply_rotary_complex(
69 q: torch.Tensor, k: torch.Tensor, freqs_cis: torch.Tensor
70) -> tuple[torch.Tensor, torch.Tensor]:
71 """Apply rotary position embedding via complex multiplication (DeepSeek-V2 style).
73 DeepSeek-V2 uses ``freqs_cis = torch.polar(ones, freqs)`` (complex exponentials)
74 instead of the standard (cos, sin) pair. This matches the V2 HF implementation of
75 ``apply_rotary_emb``.
77 Args:
78 q: Query rope portion [batch, heads, seq, rope_dim].
79 k: Key rope portion [batch, 1, seq, rope_dim].
80 freqs_cis: Complex rotary frequencies [batch, seq, rope_dim // 2].
82 Returns:
83 Tuple of rotated (q, k) tensors with same dtype and shape as inputs.
84 """
85 freqs = freqs_cis.unsqueeze(1) # [batch, 1, seq, rope_dim // 2]
86 q_c = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
87 k_c = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
88 q_rot = torch.view_as_real(q_c * freqs.to(q_c.device)).flatten(3).type_as(q)
89 k_rot = torch.view_as_real(k_c * freqs.to(k_c.device)).flatten(3).type_as(k)
90 return q_rot, k_rot
93class MLAAttentionBridge(PositionEmbeddingHooksMixin, AttentionBridge):
94 """Bridge for DeepSeek's Multi-Head Latent Attention (MLA).
96 Reimplements the MLA forward path with hooks at each computation stage.
97 Standard W_Q/W_K/W_V properties are not available on MLA models — use
98 the submodule weight access (q_a_proj, q_b_proj, etc.) instead.
99 """
101 # MLA has no standard q/k/v submodules — override to empty
102 property_aliases: Dict[str, str] = {}
104 hook_aliases = {
105 "hook_result": "hook_out",
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
112 def __init__(
113 self,
114 name: str,
115 config: Any,
116 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
117 **kwargs: Any,
118 ):
119 super().__init__(name, config, submodules=submodules, **kwargs)
120 self._init_position_embedding_hooks()
122 self.hook_q_latent = HookPoint() # Compressed Q (post q_a_layernorm)
123 self.hook_kv_latent = HookPoint() # Compressed KV (post kv_a_layernorm)
124 self.hook_q = HookPoint() # Final Q entering attention (post-RoPE concat)
125 self.hook_k = HookPoint() # Final K entering attention (post-RoPE concat)
126 self.hook_v = HookPoint() # V from kv_b_proj split
127 self.hook_rot_q = HookPoint() # Q rope portion after RoPE
128 self.hook_rot_k = HookPoint() # K rope portion after RoPE
130 # MLA params lazy-initialized from HF module (bridge config lacks these fields)
131 self._mla_params_initialized = False
133 def forward(self, *args: Any, **kwargs: Any) -> Any:
134 """Reimplemented MLA forward with hooks at each computation stage.
136 Follows the DeepseekV3Attention forward path, calling into HF submodules
137 individually and firing hooks at each meaningful stage.
138 """
139 if self.original_component is None: 139 ↛ 140line 139 didn't jump to line 140 because the condition on line 139 was never true
140 raise RuntimeError(
141 f"Original component not set for {self.name}. "
142 "Call set_original_component() first."
143 )
145 hf_attn: Any = self.original_component
147 if not self._mla_params_initialized:
148 self._q_lora_rank = getattr(hf_attn, "q_lora_rank", None)
149 self._kv_lora_rank = getattr(hf_attn, "kv_lora_rank", 512)
150 self._qk_nope_head_dim = getattr(hf_attn, "qk_nope_head_dim", 128)
151 self._qk_rope_head_dim = getattr(hf_attn, "qk_rope_head_dim", 64)
152 self._v_head_dim = getattr(hf_attn, "v_head_dim", 128)
153 self._qk_head_dim = self._qk_nope_head_dim + self._qk_rope_head_dim
154 self._n_heads = getattr(hf_attn, "num_heads", 32)
155 hf_config = getattr(hf_attn, "config", None)
156 self._rope_interleave = (
157 getattr(hf_config, "rope_interleave", False) if hf_config else False
158 )
159 self._mla_params_initialized = True
161 # --- Extract inputs ---
162 if "hidden_states" in kwargs:
163 hidden_states = kwargs.pop("hidden_states")
164 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 164 ↛ 168line 164 didn't jump to line 168 because the condition on line 164 was always true
165 hidden_states = args[0]
166 args = args[1:]
167 else:
168 raise ValueError("Could not find hidden_states in args or kwargs")
170 position_embeddings = kwargs.pop("position_embeddings", None)
171 attention_mask = kwargs.pop("attention_mask", None)
173 hidden_states = self.hook_in(hidden_states)
175 batch_size, seq_length = hidden_states.shape[:2]
177 # --- Query path ---
178 if self._q_lora_rank is None:
179 # Direct projection (no compression)
180 q_states = hf_attn.q_proj(hidden_states)
181 else:
182 # Two-stage compression: q_a_proj → q_a_layernorm → q_b_proj
183 q_compressed = hf_attn.q_a_proj(hidden_states)
184 q_compressed = hf_attn.q_a_layernorm(q_compressed)
185 q_compressed = self.hook_q_latent(q_compressed)
186 q_states = hf_attn.q_b_proj(q_compressed)
188 # Reshape to [batch, n_heads, seq, qk_head_dim]
189 q_states = q_states.view(batch_size, seq_length, -1, self._qk_head_dim).transpose(1, 2)
190 # Split into nope (non-RoPE) and pe (RoPE) portions
191 q_pass, q_rot = torch.split(
192 q_states, [self._qk_nope_head_dim, self._qk_rope_head_dim], dim=-1
193 )
195 # --- KV path ---
196 # kv_a_proj_with_mqa outputs [compressed_kv || k_pe]
197 compressed_kv_full = hf_attn.kv_a_proj_with_mqa(hidden_states)
198 # Split: compressed KV latent (for kv_b_proj) and k rope portion (for direct RoPE)
199 # Note: k_pe is split off here and goes directly to RoPE — hook_kv_latent
200 # captures only the compressed_kv portion that enters the decompression path.
201 k_pass, k_rot = torch.split(
202 compressed_kv_full, [self._kv_lora_rank, self._qk_rope_head_dim], dim=-1
203 )
205 # Compress → normalize → decompress the KV latent
206 k_pass = hf_attn.kv_a_layernorm(k_pass)
207 k_pass = self.hook_kv_latent(k_pass)
208 k_pass = hf_attn.kv_b_proj(k_pass)
210 # Reshape to [batch, n_heads, seq, nope+v_head]
211 key_shape = (batch_size, seq_length, -1, self._qk_nope_head_dim + self._v_head_dim)
212 k_pass = k_pass.view(key_shape).transpose(1, 2)
213 # Split K nope portion and V
214 k_pass, value_states = torch.split(
215 k_pass, [self._qk_nope_head_dim, self._v_head_dim], dim=-1
216 )
218 # k_rot is [batch, seq, rope_dim] → [batch, 1, seq, rope_dim] for broadcasting
219 k_rot = k_rot.view(batch_size, 1, seq_length, self._qk_rope_head_dim)
221 # --- RoPE ---
222 # DeepSeek-V2 passes a complex freqs_cis tensor; V3 passes a (cos, sin) tuple.
223 # Detect the format and apply the appropriate rotation.
224 cos = sin = None
225 if position_embeddings is not None: 225 ↛ 239line 225 didn't jump to line 239 because the condition on line 225 was always true
226 position_embeddings = self._apply_position_embedding_hooks(position_embeddings)
227 if isinstance(position_embeddings, torch.Tensor) and position_embeddings.is_complex():
228 # V2-style: complex exponential freqs_cis
229 q_rot, k_rot = _apply_rotary_complex(q_rot, k_rot, position_embeddings)
230 elif self._rope_interleave: 230 ↛ 234line 230 didn't jump to line 234 because the condition on line 230 was always true
231 cos, sin = position_embeddings
232 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
233 else:
234 cos, sin = position_embeddings
235 if self._rope_interleave:
236 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
237 else:
238 q_rot, k_rot = _apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
239 elif self._rotary_emb is not None:
240 # Fallback: compute from rotary_emb if position_embeddings not passed
241 position_ids = torch.arange(seq_length, device=hidden_states.device).unsqueeze(0)
242 emb = self._rotary_emb(hidden_states, position_ids)
243 if isinstance(emb, torch.Tensor) and emb.is_complex():
244 q_rot, k_rot = _apply_rotary_complex(q_rot, k_rot, emb)
245 elif self._rope_interleave:
246 cos, sin = emb
247 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
248 else:
249 cos, sin = emb
250 if self._rope_interleave:
251 q_rot, k_rot = _apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
252 else:
253 q_rot, k_rot = _apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
254 else:
255 raise ValueError(
256 "MLAAttentionBridge requires position_embeddings or set_rotary_emb() "
257 "to be called before forward."
258 )
259 q_rot = self.hook_rot_q(q_rot)
260 k_rot = self.hook_rot_k(k_rot)
262 # Expand k_rot to match the number of heads
263 k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
265 # Concatenate nope + rope portions to form final Q and K
266 query_states = torch.cat((q_pass, q_rot), dim=-1)
267 key_states = torch.cat((k_pass, k_rot), dim=-1)
269 # Fire final Q/K/V hooks — these are the tensors entering attention
270 query_states = self.hook_q(query_states)
271 key_states = self.hook_k(key_states)
272 value_states = self.hook_v(value_states)
274 # --- KV Cache ---
275 past_key_values = kwargs.pop("past_key_values", None)
276 cache_position = kwargs.pop("cache_position", None)
277 if past_key_values is not None:
278 cache_kwargs: dict = {"cache_position": cache_position}
279 if cos is not None:
280 cache_kwargs["cos"] = cos
281 if sin is not None:
282 cache_kwargs["sin"] = sin
283 key_states, value_states = past_key_values.update(
284 key_states, value_states, hf_attn.layer_idx, cache_kwargs
285 )
287 # --- Attention computation (no V padding — only needed for flash attention) ---
288 # Read the HF module's scaling: DeepSeek-V2/V3 multiply the base
289 # qk_head_dim^-0.5 by yarn mscale^2 under long-context rope configs.
290 # trust_remote_code DeepSeek-V2 names the same quantity softmax_scale.
291 scaling = getattr(hf_attn, "scaling", None)
292 if scaling is None:
293 scaling = getattr(hf_attn, "softmax_scale", None)
294 if scaling is None:
295 scaling = self._qk_head_dim ** (-0.5)
296 attn_scores = torch.matmul(query_states, key_states.transpose(-2, -1)) * scaling
298 if attention_mask is not None:
299 attention_mask = self._normalize_compatibility_mask_sentinel(attention_mask)
300 attn_scores = attn_scores + attention_mask
302 attn_scores = self.hook_attn_scores(attn_scores)
303 attn_weights = self._softmax_dropout_pattern(
304 attn_scores, upcast_to_fp32=True, target_dtype=query_states.dtype
305 )
307 # Weighted sum of values
308 attn_output = torch.matmul(attn_weights, value_states)
310 # --- Output projection ---
311 attn_output = attn_output.transpose(1, 2).contiguous()
312 attn_output = attn_output.reshape(batch_size, seq_length, -1)
313 attn_output = hf_attn.o_proj(attn_output)
315 attn_output = self.hook_out(attn_output)
316 return attn_output, attn_weights
318 def get_random_inputs(
319 self,
320 batch_size: int = 2,
321 seq_len: int = 8,
322 device: Optional[torch.device] = None,
323 dtype: Optional[torch.dtype] = None,
324 ) -> Dict[str, Any]:
325 """Generate test inputs with hidden_states, position_embeddings, and attention_mask."""
326 if device is None:
327 device = torch.device("cpu")
328 if dtype is None:
329 dtype = torch.float32
331 # Try bridge config (d_model), then HF attention's config (hidden_size), then fallback
332 d_model = None
333 if self.config and hasattr(self.config, "d_model"):
334 d_model = self.config.d_model
335 if d_model is None and self.original_component is not None:
336 hf_cfg = getattr(self.original_component, "config", None)
337 if hf_cfg is not None:
338 d_model = getattr(hf_cfg, "hidden_size", None)
339 if d_model is None:
340 d_model = 256
341 inputs: Dict[str, Any] = {
342 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype)
343 }
345 # Generate position_embeddings from rotary_emb if available,
346 # otherwise create dummy (cos=1, sin=0) embeddings
347 rope_head_dim = self._qk_rope_head_dim if self._mla_params_initialized else 64
348 if self._rotary_emb is not None:
349 try:
350 dummy_input = inputs["hidden_states"]
351 position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
352 position_embeddings = self._rotary_emb(dummy_input, position_ids)
353 inputs["position_embeddings"] = position_embeddings
354 except Exception:
355 cos = torch.ones(1, seq_len, rope_head_dim, device=device, dtype=dtype)
356 sin = torch.zeros(1, seq_len, rope_head_dim, device=device, dtype=dtype)
357 inputs["position_embeddings"] = (cos, sin)
358 else:
359 cos = torch.ones(1, seq_len, rope_head_dim, device=device, dtype=dtype)
360 sin = torch.zeros(1, seq_len, rope_head_dim, device=device, dtype=dtype)
361 inputs["position_embeddings"] = (cos, sin)
363 inputs["attention_mask"] = None
364 return inputs
366 def __getattr__(self, name: str) -> Any:
367 """Raise clear error for standard weight properties that don't apply to MLA."""
368 if name in ("W_Q", "W_K", "W_V", "W_O", "b_Q", "b_K", "b_V", "b_O"):
369 raise NotImplementedError(
370 f"{name} is not available on MLA (Multi-Head Latent Attention) models. "
371 f"MLA uses compressed projections instead of standard Q/K/V. "
372 f"Access weights via submodules: q_a_proj, q_b_proj, kv_a_proj_with_mqa, "
373 f"kv_b_proj, o (o_proj)."
374 )
375 return super().__getattr__(name)