Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / npu_fusion_attention.py: 25%
126 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"""NPU fusion attention function."""
17from __future__ import annotations
19from collections.abc import Mapping, Sequence
20from dataclasses import dataclass
21from typing import Any, Optional
23import torch # pylint: disable=forbidden-backend-import
24import torch_npu
27@dataclass
28class _FusionAttentionContext:
29 """Normalized arguments required by the NPU fusion-attention operator."""
31 query: torch.Tensor
32 key: torch.Tensor
33 value: torch.Tensor
34 input_layout: str
35 attention_mask: Optional[torch.Tensor]
36 sparse_mode: int
37 query_lengths: Optional[Sequence[int]]
38 key_lengths: Optional[Sequence[int]]
39 batch_size: int
40 query_length: int
41 head_dim: int
42 pre_tokens: int
43 next_tokens: int
44 is_packed: bool
47def _npu_attention_mask(attention_mask: torch.Tensor) -> torch.Tensor:
48 """Convert the Transformers attention mask to the NPU mask convention."""
49 if attention_mask.dtype == torch.bool:
50 return torch.logical_not(attention_mask)
51 return attention_mask != 0
54def _causal_attention_mask(
55 query: torch.Tensor,
56 key: torch.Tensor,
57 sliding_window: Optional[int],
58) -> torch.Tensor:
59 """Build the right-aligned causal mask used by Transformers decoding."""
60 query_length = query.shape[2]
61 key_length = key.shape[2]
62 query_positions = torch.arange(query_length, device=query.device)
63 query_positions = query_positions + max(key_length - query_length, 0)
64 key_positions = torch.arange(key_length, device=query.device)
65 allowed = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
66 if sliding_window is not None:
67 allowed &= key_positions.unsqueeze(0) > (
68 query_positions.unsqueeze(1) - sliding_window
69 )
70 return torch.logical_not(allowed)
73def _length_list(
74 value: torch.Tensor | Sequence[int],
75 *,
76 includes_zero: bool,
77) -> list[int]:
78 """Normalize one cumulative-length representation."""
79 if isinstance(value, torch.Tensor):
80 lengths = [int(item) for item in value.tolist()]
81 else:
82 lengths = [int(item) for item in value]
83 if includes_zero:
84 if not lengths or lengths[0] != 0:
85 raise ValueError("cu_seq_lens must start with zero")
86 lengths = lengths[1:]
87 if not lengths or any(left >= right for left, right in zip(lengths, lengths[1:])):
88 raise ValueError("packed cumulative sequence lengths must be strictly increasing")
89 return lengths
92def _coalesce_lengths(
93 kwargs: dict[str, Any],
94 aliases: tuple[tuple[str, bool], ...],
95 *,
96 name: str,
97) -> list[int] | None:
98 """Read equivalent length arguments and reject conflicting values."""
99 candidates = [
100 (alias, _length_list(kwargs[alias], includes_zero=includes_zero))
101 for alias, includes_zero in aliases
102 if kwargs.get(alias) is not None
103 ]
104 if not candidates:
105 return None
106 first_alias, first = candidates[0]
107 for alias, candidate in candidates[1:]:
108 if candidate != first:
109 raise ValueError(
110 f"conflicting {name} values were provided through "
111 f"{first_alias!r} and {alias!r}"
112 )
113 return first
116def _packed_parameter_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
117 """Expose packed-sequence fields carried by a model-level parameter object."""
118 packed_seq_params = kwargs.get("packed_seq_params")
119 if packed_seq_params is None:
120 return dict(kwargs)
122 resolved_kwargs = dict(kwargs)
123 field_aliases = {
124 "actual_seq_len": "actual_seq_len",
125 "actual_q_len": "actual_q_len",
126 "actual_kv_len": "actual_kv_len",
127 "actual_seq_qlen": "actual_seq_qlen",
128 "actual_seq_kvlen": "actual_seq_kvlen",
129 "cu_seq_lens": "cu_seq_lens_q",
130 "cu_seq_lens_q": "cu_seq_lens_q",
131 "cu_seq_lens_k": "cu_seq_lens_k",
132 "cu_seq_lens_kv": "cu_seq_lens_k",
133 "cu_seqlens_q": "cu_seq_lens_q",
134 "cu_seqlens_k": "cu_seq_lens_k",
135 "cu_seqlens_kv": "cu_seq_lens_k",
136 }
137 found = False
138 for source_name, target_name in field_aliases.items():
139 if isinstance(packed_seq_params, Mapping):
140 value = packed_seq_params.get(source_name)
141 else:
142 value = getattr(packed_seq_params, source_name, None)
143 if value is not None:
144 found = True
145 if resolved_kwargs.get(target_name) is None:
146 resolved_kwargs[target_name] = value
147 if source_name == "cu_seq_lens":
148 if resolved_kwargs.get("cu_seq_lens_k") is None:
149 resolved_kwargs["cu_seq_lens_k"] = value
150 if not found:
151 raise ValueError(
152 "packed_seq_params must provide cumulative query and key/value sequence lengths"
153 )
154 return resolved_kwargs
157def resolve_packed_sequence_lengths(
158 kwargs: dict[str, Any],
159 query_tokens: int,
160 key_tokens: int,
161) -> tuple[list[int] | None, list[int] | None]:
162 """Resolve supported packed-sequence argument aliases.
164 Args:
165 kwargs: Attention keyword arguments or a ``packed_seq_params`` carrier.
166 query_tokens: Number of flattened query tokens.
167 key_tokens: Number of flattened key/value tokens.
169 Returns:
170 Cumulative query and key/value sequence ends without leading zeros.
171 """
172 kwargs = _packed_parameter_kwargs(kwargs)
173 query_lengths = _coalesce_lengths(
174 kwargs,
175 (
176 ("actual_seq_len", False),
177 ("actual_q_len", False),
178 ("actual_seq_qlen", False),
179 ("cu_seq_lens_q", True),
180 ),
181 name="query sequence lengths",
182 )
183 key_lengths = _coalesce_lengths(
184 kwargs,
185 (
186 ("actual_seq_len", False),
187 ("actual_kv_len", False),
188 ("actual_seq_kvlen", False),
189 ("cu_seq_lens_k", True),
190 ),
191 name="key/value sequence lengths",
192 )
193 if (query_lengths is None) != (key_lengths is None):
194 raise ValueError("packed attention requires both query and key/value lengths")
195 if query_lengths is not None:
196 if query_lengths[-1] != query_tokens:
197 raise ValueError("the final query sequence length must equal the query token count")
198 if key_lengths[-1] != key_tokens:
199 raise ValueError(
200 "the final key/value sequence length must equal the key/value token count"
201 )
202 return query_lengths, key_lengths
205def _attention_options(module: torch.nn.Module, kwargs: dict[str, Any]):
206 """Resolve sparse-window and causal options from kwargs and the module."""
207 pre_tokens = kwargs.get("pre_tokens", getattr(module, "pre_tockens", 1048576))
208 next_tokens = kwargs.get("next_tokens", getattr(module, "next_tockens", 0))
209 sparse_mode = kwargs.get("sparse_mode", getattr(module, "sparse_mode", 0))
210 sliding_window = kwargs.get("sliding_window")
211 is_causal = kwargs.get("is_causal", getattr(module, "is_causal", True))
212 if sliding_window is not None:
213 pre_tokens = sliding_window
214 return pre_tokens, next_tokens, sparse_mode, sliding_window, is_causal
217def _prepare_attention_inputs(
218 query: torch.Tensor,
219 key: torch.Tensor,
220 value: torch.Tensor,
221 attention_mask: Optional[torch.Tensor],
222 *,
223 is_packed: bool,
224 is_causal: bool,
225 sliding_window: Optional[int],
226 sparse_mode: int,
227) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str, Optional[torch.Tensor], int]:
228 """Prepare QKV layout, mask, and sparse mode for the NPU operator."""
229 if is_packed:
230 query = query.transpose(1, 2).reshape(-1, query.shape[1], query.shape[-1])
231 key = key.transpose(1, 2).reshape(-1, key.shape[1], key.shape[-1])
232 value = value.transpose(1, 2).reshape(-1, value.shape[1], value.shape[-1])
233 if attention_mask is None and is_causal:
234 npu_mask = torch.ones((2048, 2048), dtype=torch.bool, device=query.device).triu(diagonal=1)
235 sparse_mode = 3
236 else:
237 npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
238 return query, key, value, "TND", npu_mask, sparse_mode
239 if attention_mask is None and is_causal:
240 return query, key, value, "BNSD", _causal_attention_mask(query, key, sliding_window), 0
241 npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
242 return query, key, value, "BNSD", npu_mask, sparse_mode
245def _prepare_fusion_attention_context(
246 module: torch.nn.Module,
247 query: torch.Tensor,
248 key: torch.Tensor,
249 value: torch.Tensor,
250 attention_mask: Optional[torch.Tensor],
251 kwargs: dict[str, Any],
252) -> _FusionAttentionContext:
253 """Normalize model inputs into one NPU fusion-attention invocation."""
254 query_lengths, key_lengths = resolve_packed_sequence_lengths(
255 kwargs,
256 query.shape[0] * query.shape[2],
257 key.shape[0] * key.shape[2],
258 )
259 options = _attention_options(module, kwargs)
260 prepared = _prepare_attention_inputs(
261 query,
262 key,
263 value,
264 attention_mask,
265 is_packed=query_lengths is not None,
266 is_causal=options[4],
267 sliding_window=options[3],
268 sparse_mode=options[2],
269 )
270 return _FusionAttentionContext(
271 *prepared[:5],
272 prepared[5],
273 query_lengths,
274 key_lengths,
275 query.shape[0],
276 query.shape[2],
277 query.shape[-1],
278 options[0],
279 options[1],
280 query_lengths is not None,
281 )
284def npu_fusion_attention_forward(
285 module: torch.nn.Module,
286 query: torch.Tensor,
287 key: torch.Tensor,
288 value: torch.Tensor,
289 attention_mask: Optional[torch.Tensor],
290 dropout: float = 0.0,
291 scaling: Optional[float] = None,
292 **kwargs: Any,
293) -> tuple[torch.Tensor, None]:
294 """Compute NPU fusion attention using the Transformers AttentionInterface contract.
296 Args:
297 module: Owning Transformers attention module.
298 query: Query tensor in BNSD layout.
299 key: Key tensor in BNSD layout.
300 value: Value tensor in BNSD layout.
301 attention_mask: Attention mask using the Transformers mask convention.
302 dropout: Dropout probability.
303 scaling: Attention score scaling factor.
304 **kwargs: Additional attention arguments.
306 Returns:
307 Attention output in BSND layout and ``None`` for attention weights.
309 Raises:
310 ValueError: If sparse attention indices or inconsistent sequence lengths
311 are supplied.
313 Note:
314 QKV inputs use four-dimensional BNSD layout. Packed inputs are flattened
315 to TND when cumulative sequence lengths are supplied through the project
316 ``actual_*`` aliases or Transformers ``cu_seq_lens_*`` names. Float16,
317 bfloat16, and float32 inputs have been verified. A boolean mask uses
318 ``True`` for positions that participate in attention; an additive mask
319 uses zero for those positions.
320 """
321 if kwargs.get("indices") is not None:
322 raise ValueError(
323 "npu_fusion_attention_forward does not consume sparse attention indices; "
324 "select a DSA sparse-attention implementation instead."
325 )
326 context = _prepare_fusion_attention_context(
327 module,
328 query,
329 key,
330 value,
331 attention_mask,
332 kwargs,
333 )
334 output = torch_npu.npu_fusion_attention(
335 context.query,
336 context.key,
337 context.value,
338 context.query.shape[1],
339 context.input_layout,
340 pse=None,
341 padding_mask=None,
342 atten_mask=context.attention_mask,
343 scale=context.head_dim**-0.5 if scaling is None else scaling,
344 pre_tockens=context.pre_tokens,
345 next_tockens=context.next_tokens,
346 keep_prob=1.0 - dropout,
347 inner_precise=0,
348 sparse_mode=context.sparse_mode,
349 actual_seq_qlen=context.query_lengths,
350 actual_seq_kvlen=context.key_lengths,
351 )[0]
352 if context.is_packed:
353 output = output.reshape(
354 context.batch_size,
355 context.query_length,
356 output.shape[1],
357 output.shape[2],
358 )
359 return output, None
360 return output.transpose(1, 2), None