Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / gqa_attention.py: 49%
225 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"""Grouped-query attention module using the Hyper high-performance layout."""
17from __future__ import annotations
19from collections.abc import Callable, Mapping, Sequence
20from typing import Any
22# This package provides PyTorch-specific high-performance modules.
23# pylint: disable=forbidden-backend-import
24import torch # pylint: disable=forbidden-backend-import
25from torch import nn
26from hyper_parallel.components.checkpoint.weight_conversion import WeightConverter
28from hyper_parallel.components.checkpoint import InterleaveGateQKV, InterleaveQKV
29from hyper_parallel.models.replacement import module_replacement
30from hyper_parallel.components.functional import (
31 apply_rotary_pos_emb,
32 apply_rotary_pos_emb_interleave,
33 npu_fusion_attention_forward,
34)
37def rotate_half(x: torch.Tensor) -> torch.Tensor:
38 """Rotate the two halves of the last dimension as Transformers does."""
39 first = x[..., : x.shape[-1] // 2]
40 second = x[..., x.shape[-1] // 2 :]
41 return torch.cat((-second, first), dim=-1)
44def _projections_can_fuse(projections: tuple[nn.Linear, ...]) -> bool:
45 """Return whether projections share the layout needed for fusion."""
46 biases = tuple(projection.bias for projection in projections)
47 return (
48 len({projection.in_features for projection in projections}) == 1
49 and len({projection.weight.requires_grad for projection in projections}) == 1
50 and (all(bias is None for bias in biases) or all(bias is not None for bias in biases))
51 and (biases[0] is None or len({bias.requires_grad for bias in biases}) == 1)
52 )
55def _apply_gqa_rope(
56 query: torch.Tensor,
57 key: torch.Tensor,
58 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
59 rotary_interleaved: bool,
60 head_dim: int,
61) -> tuple[torch.Tensor, torch.Tensor]:
62 """Apply the configured rotary convention to grouped-query attention."""
63 if position_embeddings is None:
64 return query, key
65 cos, sin = position_embeddings
66 if rotary_interleaved:
67 return apply_rotary_pos_emb_interleave(query, key, cos, sin, unsqueeze_dim=2)
68 if cos.shape[-1] >= head_dim:
69 return apply_rotary_pos_emb(query, key, cos, sin, unsqueeze_dim=2)
70 rotary_dim = cos.shape[-1]
71 query_rot, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
72 key_rot, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
73 cos = cos.unsqueeze(2)
74 sin = sin.unsqueeze(2)
75 query = torch.cat((query_rot * cos + rotate_half(query_rot) * sin, query_pass), dim=-1)
76 key = torch.cat((key_rot * cos + rotate_half(key_rot) * sin, key_pass), dim=-1)
77 return query, key
80@module_replacement
81class GQAAttention(nn.Module):
82 """Transformers-compatible GQA using a grouped ``linear_qkv`` layout.
84 The source module may expose a fused ``qkv_proj`` or separate
85 ``q_proj``/``k_proj``/``v_proj`` layers. Its checkpoint layout is converted
86 to the per-KV-head grouping consumed by the high-performance kernel.
87 Construction only creates the target structure; ``make_transforms``
88 declares the checkpoint conversion.
89 """
91 def _initialize_dimensions(self, module: nn.Module) -> None:
92 """Initialize head dimensions and grouped-QKV layout metadata."""
93 config = self.config
94 self.num_heads = getattr(module, "num_attention_heads", config.num_attention_heads)
95 self.num_key_value_heads = getattr(module, "num_key_value_heads", config.num_key_value_heads)
96 if self.num_heads % self.num_key_value_heads:
97 raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
98 self.num_key_value_groups = self.num_heads // self.num_key_value_heads
99 self.qk_head_dim = getattr(
100 module, "head_dim", getattr(config, "head_dim", config.hidden_size // self.num_heads)
101 )
102 self.v_head_dim = (
103 getattr(module, "v_head_dim", None) or getattr(config, "v_head_dim", None) or self.qk_head_dim
104 )
105 self.qkv_split_sizes = (
106 self.num_key_value_groups * self.qk_head_dim,
107 self.qk_head_dim,
108 self.v_head_dim,
109 )
110 self.qkv_group_width = sum(self.qkv_split_sizes)
112 def _initialize_attention_options(self, module: nn.Module) -> None:
113 """Copy supported attention options and reject unsupported variants."""
114 config = self.config
115 self.scaling = getattr(module, "scaling", self.qk_head_dim**-0.5)
116 self.attention_dropout = getattr(module, "attention_dropout", 0.0)
117 if isinstance(self.attention_dropout, nn.Dropout):
118 self.attention_dropout = self.attention_dropout.p
119 self.is_causal = getattr(module, "is_causal", True)
120 self.sliding_window = getattr(module, "sliding_window", getattr(config, "sliding_window", None))
121 self.rotary_interleaved = getattr(
122 module,
123 "rotary_interleaved",
124 getattr(config, "rope_interleave", getattr(config, "rope_interleaved", False)),
125 )
126 if getattr(module, "param_sink_number", 0) > 0:
127 raise ValueError(
128 "GQAAttention only implements standard GQA; parameter-sink attention "
129 "must use the sink_attention high-performance interface."
130 )
131 self.attn_groupnorm = bool(getattr(module, "attn_groupnorm", False))
132 self.attn_elementwise_gate = bool(getattr(module, "attn_elementwise_gate", False))
133 if self.attn_groupnorm or self.attn_elementwise_gate:
134 raise ValueError(
135 "GQAAttention only implements standard GQA; group normalization and "
136 "elementwise attention gates require a dedicated replacement."
137 )
139 def _source_qkv(self, module: nn.Module) -> tuple[nn.Linear, tuple[str, ...]]:
140 """Validate and return the source QKV projection layout."""
141 self._source_is_fused = hasattr(module, "qkv_proj")
142 if self._source_is_fused:
143 source_qkv = module.qkv_proj
144 if not isinstance(source_qkv, nn.Linear):
145 raise TypeError("qkv_proj must be torch.nn.Linear")
146 expected_size = (
147 self.num_heads * self.qk_head_dim
148 + self.num_key_value_heads * self.qk_head_dim
149 + self.num_key_value_heads * self.v_head_dim
150 )
151 if source_qkv.out_features != expected_size:
152 raise ValueError("qkv_proj output size is incompatible with the configured GQA dimensions")
153 return source_qkv, ("qkv_proj",)
154 required = ("q_proj", "k_proj", "v_proj")
155 if any(not hasattr(module, name) for name in required):
156 raise TypeError("GQAAttention requires qkv_proj or q_proj/k_proj/v_proj on the source module")
157 projections = tuple(getattr(module, name) for name in required)
158 if any(not isinstance(projection, nn.Linear) for projection in projections):
159 raise TypeError("q_proj, k_proj, and v_proj must be torch.nn.Linear")
160 expected_sizes = (
161 self.num_heads * self.qk_head_dim,
162 self.num_key_value_heads * self.qk_head_dim,
163 self.num_key_value_heads * self.v_head_dim,
164 )
165 if tuple(projection.out_features for projection in projections) != expected_sizes:
166 raise ValueError("Q/K/V projection sizes are incompatible with the configured GQA dimensions")
167 if not _projections_can_fuse(projections):
168 raise ValueError("Q, K, and V projections cannot be represented by one fused projection")
169 return module.q_proj, required
171 def _initialize_projections(self, module: nn.Module, source_qkv: nn.Linear) -> None:
172 """Create the fused QKV projection and attach the source output projection."""
173 projection_size = self.num_key_value_heads * self.qkv_group_width
174 self.linear_qkv = nn.Linear(
175 source_qkv.in_features,
176 projection_size,
177 bias=source_qkv.bias is not None,
178 device=source_qkv.weight.device,
179 dtype=source_qkv.weight.dtype,
180 )
181 self.linear_qkv.weight.requires_grad_(source_qkv.weight.requires_grad)
182 if self.linear_qkv.bias is not None:
183 self.linear_qkv.bias.requires_grad_(source_qkv.bias.requires_grad)
184 if not hasattr(module, "o_proj") or not isinstance(module.o_proj, nn.Linear):
185 raise TypeError("GQAAttention requires o_proj on the source module")
186 if module.o_proj.in_features != self.num_heads * self.v_head_dim:
187 raise ValueError("o_proj input size is incompatible with the configured GQA dimensions")
188 self.o_proj = module.o_proj
189 self.q_norm = getattr(module, "q_norm", None)
190 self.k_norm = getattr(module, "k_norm", None)
192 def __init__(
193 self,
194 *,
195 module: nn.Module,
196 module_fqn: str = "",
197 context: Mapping[str, Any] | None = None,
198 attention_interface: Callable[..., tuple[torch.Tensor, torch.Tensor | None]] = (
199 npu_fusion_attention_forward
200 ),
201 ) -> None:
202 """Build the high-performance module from a Transformers attention module."""
203 super().__init__()
204 del module_fqn, context
205 self.attention_interface = attention_interface
206 config = module.config
207 self.config = config
208 self.layer_idx = getattr(module, "layer_idx", getattr(module, "layer_number", None))
209 self._initialize_dimensions(module)
210 self._initialize_attention_options(module)
211 source_qkv, source_projection_names = self._source_qkv(module)
212 self._source_projection_names = source_projection_names
213 self._initialize_projections(module, source_qkv)
214 self.train(module.training)
216 def make_transforms(self) -> list[WeightConverter]:
217 """Describe reversible source-checkpoint to high-performance conversion."""
218 if self._source_is_fused:
219 qkv_sources = "qkv_proj.weight"
220 bias_sources: str | list[str] = "qkv_proj.bias"
221 else:
222 qkv_sources = [f"{name}.weight" for name in self._source_projection_names]
223 bias_sources = [f"{name}.bias" for name in self._source_projection_names]
225 transforms: list[WeightConverter] = [
226 WeightConverter(
227 source_patterns=qkv_sources,
228 target_patterns="linear_qkv.weight",
229 operations=[
230 InterleaveQKV(
231 self.num_key_value_heads,
232 self.num_key_value_groups,
233 self.qk_head_dim,
234 self.v_head_dim,
235 self._source_is_fused,
236 )
237 ],
238 )
239 ]
240 if self.linear_qkv.bias is not None:
241 transforms.append(
242 WeightConverter(
243 source_patterns=bias_sources,
244 target_patterns="linear_qkv.bias",
245 operations=[
246 InterleaveQKV(
247 self.num_key_value_heads,
248 self.num_key_value_groups,
249 self.qk_head_dim,
250 self.v_head_dim,
251 self._source_is_fused,
252 )
253 ],
254 )
255 )
256 return transforms
258 def forward(
259 self,
260 hidden_states: torch.Tensor,
261 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
262 attention_mask: torch.Tensor | None = None,
263 past_key_values: Any | None = None,
264 actual_seq_len: torch.Tensor | Sequence[int] | None = None,
265 **kwargs: Any,
266 ) -> tuple[torch.Tensor, torch.Tensor | None]:
267 """Run GQA with the same external contract as Transformers attention."""
268 input_shape = hidden_states.shape[:-1]
269 qkv_states = self.linear_qkv(hidden_states).view(
270 *input_shape,
271 self.num_key_value_heads,
272 self.qkv_group_width,
273 )
274 query_states, key_states, value_states = torch.split(
275 qkv_states, self.qkv_split_sizes, dim=-1
276 )
277 query_states = query_states.reshape(
278 *input_shape, self.num_heads, self.qk_head_dim
279 )
281 if self.q_norm is not None:
282 query_states = self.q_norm(query_states)
283 if self.k_norm is not None:
284 key_states = self.k_norm(key_states)
286 query_states, key_states = _apply_gqa_rope(
287 query_states,
288 key_states,
289 position_embeddings,
290 self.rotary_interleaved,
291 self.qk_head_dim,
292 )
293 query_states = query_states.transpose(1, 2)
294 key_states = key_states.transpose(1, 2)
295 value_states = value_states.transpose(1, 2)
297 if past_key_values is not None:
298 key_states, value_states = past_key_values.update(
299 key_states, value_states, self.layer_idx
300 )
302 attn_output, attn_weights = self.attention_interface(
303 self,
304 query_states,
305 key_states,
306 value_states,
307 attention_mask,
308 dropout=self.attention_dropout if self.training else 0.0,
309 scaling=self.scaling,
310 sliding_window=self.sliding_window,
311 actual_seq_len=actual_seq_len,
312 **kwargs,
313 )
314 attn_output = attn_output.reshape(*input_shape, -1).contiguous()
315 attn_output = self.o_proj(attn_output)
316 return attn_output, attn_weights
319@module_replacement
320class GatedGQAAttention(nn.Module):
321 """GQA with a sigmoid output gate derived from each query head.
323 The source ``q_proj`` is expected to store ``[query, gate]`` rows for each
324 query head. This is the gated-query layout used by Qwen3.5 full attention.
325 """
327 def _initialize_dimensions(self, module: nn.Module) -> None:
328 """Initialize gated grouped-QKV dimensions and attention options."""
329 config = self.config
330 self.num_heads = getattr(module, "num_attention_heads", config.num_attention_heads)
331 self.num_key_value_heads = getattr(module, "num_key_value_heads", config.num_key_value_heads)
332 if self.num_heads % self.num_key_value_heads:
333 raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
334 self.num_key_value_groups = self.num_heads // self.num_key_value_heads
335 self.qk_head_dim = getattr(
336 module, "head_dim", getattr(config, "head_dim", config.hidden_size // self.num_heads)
337 )
338 self.v_head_dim = (
339 getattr(module, "v_head_dim", None) or getattr(config, "v_head_dim", None) or self.qk_head_dim
340 )
341 self.qkv_split_sizes = (
342 self.num_key_value_groups * self.qk_head_dim,
343 self.qk_head_dim,
344 self.v_head_dim,
345 )
346 self.qkv_group_width = sum(self.qkv_split_sizes)
347 self.gated_qkv_group_width = self.qkv_group_width + self.num_key_value_groups * self.qk_head_dim
348 self.scaling = getattr(module, "scaling", self.qk_head_dim**-0.5)
349 self.attention_dropout = getattr(module, "attention_dropout", 0.0)
350 if isinstance(self.attention_dropout, nn.Dropout):
351 self.attention_dropout = self.attention_dropout.p
352 self.is_causal = getattr(module, "is_causal", True)
353 self.sliding_window = getattr(module, "sliding_window", getattr(config, "sliding_window", None))
354 self.rotary_interleaved = getattr(
355 module,
356 "rotary_interleaved",
357 getattr(config, "rope_interleave", getattr(config, "rope_interleaved", False)),
358 )
360 def _validate_source_projections(self, module: nn.Module) -> None:
361 """Validate separate gated-query, key, value, and output projections."""
362 projections = (module.q_proj, module.k_proj, module.v_proj)
363 if any(not isinstance(projection, nn.Linear) for projection in projections):
364 raise TypeError("q_proj, k_proj, and v_proj must be torch.nn.Linear")
365 expected_sizes = (
366 2 * self.num_heads * self.qk_head_dim,
367 self.num_key_value_heads * self.qk_head_dim,
368 self.num_key_value_heads * self.v_head_dim,
369 )
370 if tuple(projection.out_features for projection in projections) != expected_sizes:
371 raise ValueError(
372 "Q/gate, K, and V projection sizes are incompatible with the configured gated GQA dimensions"
373 )
374 if not _projections_can_fuse(projections):
375 raise ValueError("Q/gate, K, and V projections cannot be represented by one fused projection")
376 attention_output_size = self.num_heads * self.v_head_dim
377 if self.num_heads * self.qk_head_dim != attention_output_size:
378 raise ValueError("GatedGQAAttention requires the gate width to match the attention output width")
379 if not isinstance(module.o_proj, nn.Linear):
380 raise TypeError("o_proj must be torch.nn.Linear")
381 if module.o_proj.in_features != attention_output_size:
382 raise ValueError("o_proj input size is incompatible with the configured gated GQA dimensions")
384 def _initialize_projection(self, module: nn.Module) -> None:
385 """Create the fused gated-QKV projection and attach source output layers."""
386 self.linear_qkv = nn.Linear(
387 module.q_proj.in_features,
388 self.num_key_value_heads * self.gated_qkv_group_width,
389 bias=module.q_proj.bias is not None,
390 device=module.q_proj.weight.device,
391 dtype=module.q_proj.weight.dtype,
392 )
393 self.linear_qkv.weight.requires_grad_(module.q_proj.weight.requires_grad)
394 if self.linear_qkv.bias is not None:
395 self.linear_qkv.bias.requires_grad_(module.q_proj.bias.requires_grad)
396 self.q_norm = module.q_norm
397 self.k_norm = module.k_norm
398 self.o_proj = module.o_proj
400 def __init__(
401 self,
402 *,
403 module: nn.Module,
404 module_fqn: str = "",
405 context: Mapping[str, Any] | None = None,
406 attention_interface: Callable[..., tuple[torch.Tensor, torch.Tensor | None]] = (
407 npu_fusion_attention_forward
408 ),
409 ) -> None:
410 """Build gated GQA from a source with separate Q/gate, K, and V projections."""
411 super().__init__()
412 del module_fqn, context
413 if hasattr(module, "qkv_proj"):
414 raise ValueError(
415 "GatedGQAAttention requires separate q_proj/k_proj/v_proj projections"
416 )
417 required = ("q_proj", "k_proj", "v_proj", "q_norm", "k_norm", "o_proj")
418 if any(not hasattr(module, name) for name in required):
419 raise TypeError(
420 "GatedGQAAttention requires q_proj, k_proj, v_proj, q_norm, "
421 "k_norm, and o_proj on the source module"
422 )
424 self.attention_interface = attention_interface
425 self.config = module.config
426 self.layer_idx = getattr(module, "layer_idx", None)
427 self._initialize_dimensions(module)
428 self._validate_source_projections(module)
429 self._initialize_projection(module)
430 self.train(module.training)
432 def make_transforms(self) -> list[WeightConverter]:
433 """Describe reversible gated-query and grouped-QKV conversion."""
434 source_weights = ["q_proj.weight", "k_proj.weight", "v_proj.weight"]
435 operation = InterleaveGateQKV(
436 self.num_key_value_heads,
437 self.num_key_value_groups,
438 self.qk_head_dim,
439 self.v_head_dim,
440 )
441 transforms = [
442 WeightConverter(
443 source_patterns=source_weights,
444 target_patterns="linear_qkv.weight",
445 operations=[operation],
446 )
447 ]
448 if self.linear_qkv.bias is not None:
449 transforms.append(
450 WeightConverter(
451 source_patterns=["q_proj.bias", "k_proj.bias", "v_proj.bias"],
452 target_patterns="linear_qkv.bias",
453 operations=[
454 InterleaveGateQKV(
455 self.num_key_value_heads,
456 self.num_key_value_groups,
457 self.qk_head_dim,
458 self.v_head_dim,
459 )
460 ],
461 )
462 )
463 return transforms
465 def forward(
466 self,
467 hidden_states: torch.Tensor,
468 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
469 attention_mask: torch.Tensor | None = None,
470 past_key_values: Any | None = None,
471 actual_seq_len: torch.Tensor | Sequence[int] | None = None,
472 **kwargs: Any,
473 ) -> tuple[torch.Tensor, torch.Tensor | None]:
474 """Run gated GQA with the same external contract as Transformers attention."""
475 input_shape = hidden_states.shape[:-1]
476 qkv_states = self.linear_qkv(hidden_states).view(
477 *input_shape,
478 self.num_key_value_heads,
479 self.gated_qkv_group_width,
480 )
481 query_states, key_states, value_states, gate = torch.split(
482 qkv_states,
483 (*self.qkv_split_sizes, self.num_key_value_groups * self.qk_head_dim),
484 dim=-1,
485 )
486 query_states = query_states.reshape(
487 *input_shape, self.num_heads, self.qk_head_dim
488 )
489 gate = gate.reshape(*input_shape, -1)
491 if self.q_norm is not None:
492 query_states = self.q_norm(query_states)
493 if self.k_norm is not None:
494 key_states = self.k_norm(key_states)
496 query_states, key_states = _apply_gqa_rope(
497 query_states,
498 key_states,
499 position_embeddings,
500 self.rotary_interleaved,
501 self.qk_head_dim,
502 )
503 query_states = query_states.transpose(1, 2)
504 key_states = key_states.transpose(1, 2)
505 value_states = value_states.transpose(1, 2)
507 if past_key_values is not None:
508 key_states, value_states = past_key_values.update(
509 key_states, value_states, self.layer_idx
510 )
512 attn_output, attn_weights = self.attention_interface(
513 self,
514 query_states,
515 key_states,
516 value_states,
517 attention_mask,
518 dropout=self.attention_dropout if self.training else 0.0,
519 scaling=self.scaling,
520 sliding_window=self.sliding_window,
521 actual_seq_len=actual_seq_len,
522 **kwargs,
523 )
524 attn_output = attn_output.reshape(*input_shape, -1).contiguous()
525 attn_output = attn_output * torch.sigmoid(gate)
526 attn_output = self.o_proj(attn_output)
527 return attn_output, attn_weights