Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / mla_attention.py: 93%
117 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Multi-head latent attention module using the Hyper projection layout."""
17from __future__ import annotations
19from collections.abc import Callable, Mapping, Sequence
20from dataclasses import dataclass
21from typing import Any
23# This package provides PyTorch-specific high-performance modules.
24# pylint: disable=forbidden-backend-import
25import torch # pylint: disable=forbidden-backend-import
26from torch import nn
27from hyper_parallel.components.checkpoint.weight_conversion import WeightConverter
29from hyper_parallel.components.checkpoint import ConcatenateWithSections
30from hyper_parallel.components.functional.npu_fusion_attention import npu_fusion_attention_forward
31from hyper_parallel.models.replacement import module_replacement
34def prepare_mla_attention_mask(
35 attention_mask: torch.Tensor | None,
36 query_length: int,
37 key_length: int,
38) -> torch.Tensor | None:
39 """Disambiguate MLA Q/K masks before calling a shared attention interface.
41 Args:
42 attention_mask: Optional mask. A two-dimensional Q/K matrix takes
43 precedence over padding in MLA, including when batch equals query length.
44 query_length: Number of query tokens visible to the attention interface.
45 key_length: Number of key tokens visible to the attention interface.
47 Returns:
48 A broadcastable four-dimensional view for a Q/K matrix. Other inputs
49 pass through unchanged, preserving padding and compressed causal paths.
50 """
51 if (attention_mask is not None and attention_mask.ndim == 2
52 and attention_mask.shape == (query_length, key_length)):
53 # Shared FA interprets 2D masks as batch padding; MLA also accepts Q/K matrices.
54 return attention_mask[None, None, :, :]
55 return attention_mask
58@dataclass
59class _MLALatents:
60 batch_size: int
61 sequence_length: int
62 query_pass: torch.Tensor
63 query_rope: torch.Tensor
64 kv_nope: torch.Tensor
65 key_rope: torch.Tensor
68@module_replacement
69class MLAAttention(nn.Module):
70 """Transformers-compatible MLA using a fused latent projection.
72 ``q_a_proj`` and ``kv_a_proj_with_mqa`` are fused into ``linear_qkv``.
73 Unchanged child modules retain their source names and checkpoint layout.
74 Construction only creates the target structure; ``make_transforms``
75 declares the checkpoint conversion.
76 """
78 def __init__(
79 self,
80 *,
81 module: nn.Module,
82 module_fqn: str = "",
83 context: Mapping[str, Any] | None = None,
84 attention_interface: Callable[..., tuple[torch.Tensor, torch.Tensor | None]] = (
85 npu_fusion_attention_forward
86 ),
87 ) -> None:
88 """Build the high-performance module from a Transformers MLA module."""
89 super().__init__()
90 del module_fqn, context
91 self.attention_interface = attention_interface
92 required = (
93 "q_a_proj",
94 "q_a_layernorm",
95 "q_b_proj",
96 "kv_a_proj_with_mqa",
97 "kv_a_layernorm",
98 "kv_b_proj",
99 "o_proj",
100 )
101 missing = [name for name in required if not hasattr(module, name)]
102 if missing:
103 raise TypeError(f"MLAAttention source module is missing: {missing}")
105 config = module.config
106 self.config = config
107 self.layer_idx = getattr(module, "layer_idx", getattr(module, "layer_number", None))
108 self.num_heads = getattr(module, "num_heads", config.num_attention_heads)
109 self.num_key_value_heads = getattr(
110 module,
111 "num_key_value_heads",
112 getattr(config, "num_key_value_heads", self.num_heads),
113 )
114 if self.num_key_value_heads is None:
115 self.num_key_value_heads = self.num_heads
116 self.num_key_value_groups = self.num_heads // self.num_key_value_heads
117 self.q_lora_rank = getattr(module, "q_lora_rank", config.q_lora_rank)
118 self.kv_lora_rank = getattr(module, "kv_lora_rank", config.kv_lora_rank)
119 if self.q_lora_rank is None:
120 raise ValueError("MLAAttention requires the Q-LoRA projection path")
121 self.qk_rope_head_dim = getattr(module, "qk_rope_head_dim", config.qk_rope_head_dim)
122 self.qk_nope_head_dim = getattr(module, "qk_nope_head_dim", config.qk_nope_head_dim)
123 self.v_head_dim = getattr(module, "v_head_dim", config.v_head_dim)
124 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
125 self.scaling = getattr(module, "scaling", self.qk_head_dim**-0.5)
126 self.attention_dropout = getattr(module, "attention_dropout", 0.0)
127 if isinstance(self.attention_dropout, nn.Dropout):
128 self.attention_dropout = self.attention_dropout.p
129 self.is_causal = getattr(module, "is_causal", True)
130 self.sliding_window = getattr(
131 module, "sliding_window", getattr(config, "sliding_window", None)
132 )
133 self.rotary_interleaved = bool(
134 getattr(module, "rotary_interleaved", getattr(config, "rope_interleave", False))
135 )
136 if getattr(module, "param_sink_number", 0) > 0:
137 raise ValueError(
138 "MLAAttention only implements standard MLA; parameter-sink attention "
139 "must use the sink high-performance interfaces."
140 )
141 self.use_mome = bool(getattr(module, "use_mome", False))
142 if self.use_mome:
143 raise NotImplementedError(
144 "MLAAttention only implements standard MLA; MOME must use the "
145 "aggregate_hidden high-performance path."
146 )
148 q_a = module.q_a_proj
149 kv_a = module.kv_a_proj_with_mqa
150 self._q_latent_output_size = q_a.out_features
151 self._kv_latent_output_size = kv_a.out_features
152 can_fuse = (
153 q_a.in_features == kv_a.in_features
154 and q_a.weight.requires_grad == kv_a.weight.requires_grad
155 and (q_a.bias is None) == (kv_a.bias is None)
156 and (q_a.bias is None or q_a.bias.requires_grad == kv_a.bias.requires_grad)
157 )
158 if not can_fuse:
159 raise ValueError(
160 "Q and KV latent projections cannot be represented by one fused projection"
161 )
162 self.linear_qkv = nn.Linear(
163 q_a.in_features,
164 q_a.out_features + kv_a.out_features,
165 bias=q_a.bias is not None,
166 device=q_a.weight.device,
167 dtype=q_a.weight.dtype,
168 )
169 self.linear_qkv.weight.requires_grad_(q_a.weight.requires_grad)
170 if q_a.bias is not None:
171 self.linear_qkv.bias.requires_grad_(q_a.bias.requires_grad)
173 self.q_a_layernorm = module.q_a_layernorm
174 self.kv_a_layernorm = module.kv_a_layernorm
175 self.q_b_proj = module.q_b_proj
176 self.kv_b_proj = module.kv_b_proj
177 self.o_proj = module.o_proj
179 self.train(module.training)
181 def make_transforms(self) -> list[WeightConverter]:
182 """Describe reversible source-checkpoint to high-performance conversion."""
183 transforms: list[WeightConverter] = [
184 WeightConverter(
185 source_patterns=[
186 "q_a_proj.weight",
187 "kv_a_proj_with_mqa.weight",
188 ],
189 target_patterns="linear_qkv.weight",
190 operations=[
191 ConcatenateWithSections(
192 sections=(
193 self._q_latent_output_size,
194 self._kv_latent_output_size,
195 ),
196 dim=0,
197 )
198 ],
199 ),
200 ]
201 if self.linear_qkv.bias is not None:
202 transforms.insert(
203 1,
204 WeightConverter(
205 source_patterns=[
206 "q_a_proj.bias",
207 "kv_a_proj_with_mqa.bias",
208 ],
209 target_patterns="linear_qkv.bias",
210 operations=[
211 ConcatenateWithSections(
212 sections=(
213 self._q_latent_output_size,
214 self._kv_latent_output_size,
215 ),
216 dim=0,
217 )
218 ],
219 ),
220 )
221 return transforms
223 def project_latent_inputs(
224 self, hidden_states: torch.Tensor,
225 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
226 """Compute shared latent inputs before head expansion.
228 Args:
229 hidden_states: Input tokens with shape [batch, sequence, hidden].
231 Returns:
232 Normalized query and KV latents, and the unrotated shared RoPE key.
233 """
234 latent_states = self.linear_qkv(hidden_states)
235 query_latent, kv_nope, key_rope = torch.split(
236 latent_states,
237 (self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim),
238 dim=-1,
239 )
240 return self.q_a_layernorm(query_latent), self.kv_a_layernorm(kv_nope), key_rope
242 def _project_latents(self, hidden_states: torch.Tensor) -> _MLALatents:
243 """Project hidden states into query and compressed KV latent states."""
244 batch_size, sequence_length = hidden_states.shape[:-1]
245 query_latent, kv_nope, key_rope = self.project_latent_inputs(hidden_states)
246 query_states = self.q_b_proj(query_latent).view(
247 batch_size,
248 sequence_length,
249 self.num_heads,
250 self.qk_head_dim,
251 )
252 query_pass, query_rope = torch.split(
253 query_states,
254 (self.qk_nope_head_dim, self.qk_rope_head_dim),
255 dim=-1,
256 )
257 return _MLALatents(
258 batch_size,
259 sequence_length,
260 query_pass,
261 query_rope.transpose(1, 2),
262 kv_nope.view(
263 batch_size, 1, sequence_length, self.kv_lora_rank
264 ),
265 key_rope.view(batch_size, 1, sequence_length, self.qk_rope_head_dim),
266 )
268 def _project_attention_inputs(
269 self,
270 hidden_states: torch.Tensor,
271 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
272 past_key_values: Any | None,
273 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
274 """Build query, key, and value tensors from compressed latent states."""
275 latents = self._project_latents(hidden_states)
276 query_rope, key_rope = latents.query_rope, latents.key_rope
277 if position_embeddings is not None:
278 # Defer the optional NPU dependency until rotary embedding is used.
279 from hyper_parallel.components.functional import ( # pylint: disable=C0415
280 apply_rotary_pos_emb, apply_rotary_pos_emb_interleave,
281 )
283 cos, sin = position_embeddings
284 rope_fn = apply_rotary_pos_emb_interleave if self.rotary_interleaved else apply_rotary_pos_emb
285 query_rope, key_rope = rope_fn(query_rope, key_rope, cos, sin)
286 kv_nope = latents.kv_nope
287 if past_key_values is not None:
288 kv_nope, key_rope = past_key_values.update(kv_nope, key_rope, self.layer_idx)
289 kv_states = self.kv_b_proj(kv_nope).view(
290 latents.batch_size,
291 kv_nope.shape[2],
292 self.num_heads,
293 self.qk_nope_head_dim + self.v_head_dim,
294 ).transpose(1, 2)
295 key_nope, value_states = torch.split(
296 kv_states,
297 (self.qk_nope_head_dim, self.v_head_dim),
298 dim=-1,
299 )
300 return (
301 torch.cat((latents.query_pass.transpose(1, 2), query_rope), dim=-1),
302 torch.cat((key_nope, key_rope.expand(-1, self.num_heads, -1, -1)), dim=-1),
303 value_states,
304 )
306 def project_output(self, attention_output: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
307 """Project [B, local_tokens, heads * value_dim] to [B, local_tokens, hidden].
309 This model-math extension point lets gated MLA multiply its gate before
310 Wo; the default MLA keeps the original output projection unchanged.
311 This method owns model math only; CP communication belongs to its adapter.
313 Args:
314 attention_output: Flattened local-token attention output.
315 hidden_states: Original local-token input, available to gated variants.
317 Returns:
318 Projected local hidden states.
319 """
320 del hidden_states
321 return self.o_proj(attention_output)
323 def forward(
324 self,
325 hidden_states: torch.Tensor,
326 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
327 attention_mask: torch.Tensor | None = None,
328 past_key_values: Any | None = None,
329 actual_seq_len: torch.Tensor | Sequence[int] | None = None,
330 **kwargs: Any,
331 ) -> tuple[torch.Tensor, torch.Tensor | None]:
332 """Run MLA with the same external contract as Transformers attention.
334 Args:
335 hidden_states: Input tokens in batch, sequence, hidden order.
336 position_embeddings: Optional precomputed RoPE cosine and sine.
337 attention_mask: Optional mask in the attention interface convention.
338 past_key_values: Optional cache for incremental decoding.
339 actual_seq_len: Optional cumulative packed document ends.
340 **kwargs: Additional arguments for the attention interface.
342 Returns:
343 Projected attention output and optional attention weights.
344 """
345 batch_size, seq_length = hidden_states.shape[:-1]
346 query_states, key_states, value_states = self._project_attention_inputs(
347 hidden_states,
348 position_embeddings,
349 past_key_values,
350 )
352 attention_mask = prepare_mla_attention_mask(
353 attention_mask, query_states.shape[2], key_states.shape[2],
354 )
355 attn_output, attn_weights = self.attention_interface(
356 self,
357 query_states,
358 key_states,
359 value_states,
360 attention_mask,
361 dropout=self.attention_dropout if self.training else 0.0,
362 scaling=self.scaling,
363 sliding_window=self.sliding_window,
364 actual_seq_len=actual_seq_len,
365 **kwargs,
366 )
367 attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
368 # Gated MLA derives its gate from the original input before Wo; ordinary
369 # MLA only applies Wo. Keep this math extension shared by CP1 and CP.
370 attn_output = self.project_output(attn_output, hidden_states)
371 return attn_output, attn_weights