Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / mla_attention.py: 0%
106 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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.models.replacement import module_replacement
31from hyper_parallel.components.functional import apply_rotary_pos_emb, apply_rotary_pos_emb_interleave
32from hyper_parallel.components.functional import npu_fusion_attention_forward
35@dataclass
36class _MLALatents:
37 batch_size: int
38 sequence_length: int
39 query_pass: torch.Tensor
40 query_rope: torch.Tensor
41 kv_nope: torch.Tensor
42 key_rope: torch.Tensor
45@module_replacement
46class MLAAttention(nn.Module):
47 """Transformers-compatible MLA using a fused latent projection.
49 ``q_a_proj`` and ``kv_a_proj_with_mqa`` are fused into ``linear_qkv``.
50 Unchanged child modules retain their source names and checkpoint layout.
51 Construction only creates the target structure; ``make_transforms``
52 declares the checkpoint conversion.
53 """
55 def __init__(
56 self,
57 *,
58 module: nn.Module,
59 module_fqn: str = "",
60 context: Mapping[str, Any] | None = None,
61 attention_interface: Callable[..., tuple[torch.Tensor, torch.Tensor | None]] = (
62 npu_fusion_attention_forward
63 ),
64 ) -> None:
65 """Build the high-performance module from a Transformers MLA module."""
66 super().__init__()
67 del module_fqn, context
68 self.attention_interface = attention_interface
69 required = (
70 "q_a_proj",
71 "q_a_layernorm",
72 "q_b_proj",
73 "kv_a_proj_with_mqa",
74 "kv_a_layernorm",
75 "kv_b_proj",
76 "o_proj",
77 )
78 missing = [name for name in required if not hasattr(module, name)]
79 if missing:
80 raise TypeError(f"MLAAttention source module is missing: {missing}")
82 config = module.config
83 self.config = config
84 self.layer_idx = getattr(module, "layer_idx", getattr(module, "layer_number", None))
85 self.num_heads = getattr(module, "num_heads", config.num_attention_heads)
86 self.num_key_value_heads = getattr(
87 module,
88 "num_key_value_heads",
89 getattr(config, "num_key_value_heads", self.num_heads),
90 )
91 if self.num_key_value_heads is None:
92 self.num_key_value_heads = self.num_heads
93 self.num_key_value_groups = self.num_heads // self.num_key_value_heads
94 self.q_lora_rank = getattr(module, "q_lora_rank", config.q_lora_rank)
95 self.kv_lora_rank = getattr(module, "kv_lora_rank", config.kv_lora_rank)
96 if self.q_lora_rank is None:
97 raise ValueError("MLAAttention requires the Q-LoRA projection path")
98 self.qk_rope_head_dim = getattr(module, "qk_rope_head_dim", config.qk_rope_head_dim)
99 self.qk_nope_head_dim = getattr(module, "qk_nope_head_dim", config.qk_nope_head_dim)
100 self.v_head_dim = getattr(module, "v_head_dim", config.v_head_dim)
101 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
102 self.scaling = getattr(module, "scaling", self.qk_head_dim**-0.5)
103 self.attention_dropout = getattr(module, "attention_dropout", 0.0)
104 if isinstance(self.attention_dropout, nn.Dropout):
105 self.attention_dropout = self.attention_dropout.p
106 self.is_causal = getattr(module, "is_causal", True)
107 self.sliding_window = getattr(
108 module, "sliding_window", getattr(config, "sliding_window", None)
109 )
110 self.rotary_interleaved = bool(
111 getattr(module, "rotary_interleaved", getattr(config, "rope_interleave", False))
112 )
113 if getattr(module, "param_sink_number", 0) > 0:
114 raise ValueError(
115 "MLAAttention only implements standard MLA; parameter-sink attention "
116 "must use the sink high-performance interfaces."
117 )
118 self.use_mome = bool(getattr(module, "use_mome", False))
119 if self.use_mome:
120 raise NotImplementedError(
121 "MLAAttention only implements standard MLA; MOME must use the "
122 "aggregate_hidden high-performance path."
123 )
125 q_a = module.q_a_proj
126 kv_a = module.kv_a_proj_with_mqa
127 self._q_latent_output_size = q_a.out_features
128 self._kv_latent_output_size = kv_a.out_features
129 can_fuse = (
130 q_a.in_features == kv_a.in_features
131 and q_a.weight.requires_grad == kv_a.weight.requires_grad
132 and (q_a.bias is None) == (kv_a.bias is None)
133 and (q_a.bias is None or q_a.bias.requires_grad == kv_a.bias.requires_grad)
134 )
135 if not can_fuse:
136 raise ValueError(
137 "Q and KV latent projections cannot be represented by one fused projection"
138 )
139 self.linear_qkv = nn.Linear(
140 q_a.in_features,
141 q_a.out_features + kv_a.out_features,
142 bias=q_a.bias is not None,
143 device=q_a.weight.device,
144 dtype=q_a.weight.dtype,
145 )
146 self.linear_qkv.weight.requires_grad_(q_a.weight.requires_grad)
147 if q_a.bias is not None:
148 self.linear_qkv.bias.requires_grad_(q_a.bias.requires_grad)
150 self.q_a_layernorm = module.q_a_layernorm
151 self.kv_a_layernorm = module.kv_a_layernorm
152 self.q_b_proj = module.q_b_proj
153 self.kv_b_proj = module.kv_b_proj
154 self.o_proj = module.o_proj
156 self.train(module.training)
158 def make_transforms(self) -> list[WeightConverter]:
159 """Describe reversible source-checkpoint to high-performance conversion."""
160 transforms: list[WeightConverter] = [
161 WeightConverter(
162 source_patterns=[
163 "q_a_proj.weight",
164 "kv_a_proj_with_mqa.weight",
165 ],
166 target_patterns="linear_qkv.weight",
167 operations=[
168 ConcatenateWithSections(
169 sections=(
170 self._q_latent_output_size,
171 self._kv_latent_output_size,
172 ),
173 dim=0,
174 )
175 ],
176 ),
177 ]
178 if self.linear_qkv.bias is not None:
179 transforms.insert(
180 1,
181 WeightConverter(
182 source_patterns=[
183 "q_a_proj.bias",
184 "kv_a_proj_with_mqa.bias",
185 ],
186 target_patterns="linear_qkv.bias",
187 operations=[
188 ConcatenateWithSections(
189 sections=(
190 self._q_latent_output_size,
191 self._kv_latent_output_size,
192 ),
193 dim=0,
194 )
195 ],
196 ),
197 )
198 return transforms
200 def _project_latents(self, hidden_states: torch.Tensor) -> _MLALatents:
201 """Project hidden states into query and compressed KV latent states."""
202 batch_size, sequence_length = hidden_states.shape[:-1]
203 latent_states = self.linear_qkv(hidden_states)
204 query_latent, kv_nope, key_rope = torch.split(
205 latent_states,
206 (self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim),
207 dim=-1,
208 )
209 query_states = self.q_b_proj(self.q_a_layernorm(query_latent)).view(
210 batch_size,
211 sequence_length,
212 self.num_heads,
213 self.qk_head_dim,
214 )
215 query_pass, query_rope = torch.split(
216 query_states,
217 (self.qk_nope_head_dim, self.qk_rope_head_dim),
218 dim=-1,
219 )
220 return _MLALatents(
221 batch_size,
222 sequence_length,
223 query_pass,
224 query_rope.transpose(1, 2),
225 self.kv_a_layernorm(kv_nope).view(
226 batch_size, 1, sequence_length, self.kv_lora_rank
227 ),
228 key_rope.view(batch_size, 1, sequence_length, self.qk_rope_head_dim),
229 )
231 def _project_attention_inputs(
232 self,
233 hidden_states: torch.Tensor,
234 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
235 past_key_values: Any | None,
236 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
237 """Build query, key, and value tensors from compressed latent states."""
238 latents = self._project_latents(hidden_states)
239 query_rope, key_rope = latents.query_rope, latents.key_rope
240 if position_embeddings is not None:
241 cos, sin = position_embeddings
242 rope_fn = apply_rotary_pos_emb_interleave if self.rotary_interleaved else apply_rotary_pos_emb
243 query_rope, key_rope = rope_fn(query_rope, key_rope, cos, sin)
244 kv_nope = latents.kv_nope
245 if past_key_values is not None:
246 kv_nope, key_rope = past_key_values.update(kv_nope, key_rope, self.layer_idx)
247 kv_states = self.kv_b_proj(kv_nope).view(
248 latents.batch_size,
249 kv_nope.shape[2],
250 self.num_heads,
251 self.qk_nope_head_dim + self.v_head_dim,
252 ).transpose(1, 2)
253 key_nope, value_states = torch.split(
254 kv_states,
255 (self.qk_nope_head_dim, self.v_head_dim),
256 dim=-1,
257 )
258 return (
259 torch.cat((latents.query_pass.transpose(1, 2), query_rope), dim=-1),
260 torch.cat((key_nope, key_rope.expand(-1, self.num_heads, -1, -1)), dim=-1),
261 value_states,
262 )
264 def forward(
265 self,
266 hidden_states: torch.Tensor,
267 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
268 attention_mask: torch.Tensor | None = None,
269 past_key_values: Any | None = None,
270 actual_seq_len: torch.Tensor | Sequence[int] | None = None,
271 **kwargs: Any,
272 ) -> tuple[torch.Tensor, torch.Tensor | None]:
273 """Run MLA with the same external contract as Transformers attention."""
274 batch_size, seq_length = hidden_states.shape[:-1]
275 query_states, key_states, value_states = self._project_attention_inputs(
276 hidden_states,
277 position_embeddings,
278 past_key_values,
279 )
281 attn_output, attn_weights = self.attention_interface(
282 self,
283 query_states,
284 key_states,
285 value_states,
286 attention_mask,
287 dropout=self.attention_dropout if self.training else 0.0,
288 scaling=self.scaling,
289 sliding_window=self.sliding_window,
290 actual_seq_len=actual_seq_len,
291 **kwargs,
292 )
293 attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
294 attn_output = self.o_proj(attn_output)
295 return attn_output, attn_weights