Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / npu_fusion_attention.py: 91%
190 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"""NPU fusion attention function."""
17from __future__ import annotations
19from collections.abc import Mapping, Sequence
20from dataclasses import dataclass
21from numbers import Integral
22from typing import Any, Optional
24import torch # pylint: disable=forbidden-backend-import
27_QUERY_LENGTH_ALIASES = (
28 ("actual_seq_len", False), ("actual_q_len", False), ("actual_seq_qlen", False),
29 ("cu_seq_lens", True), ("cu_seq_lens_q", True), ("cu_seqlens_q", True),
30)
31_KEY_LENGTH_ALIASES = (
32 ("actual_seq_len", False), ("actual_kv_len", False), ("actual_seq_kvlen", False),
33 ("cu_seq_lens", True), ("cu_seq_lens_k", True), ("cu_seq_lens_kv", True),
34 ("cu_seqlens_k", True), ("cu_seqlens_kv", True),
35)
36PACKED_SEQUENCE_ARGUMENTS = frozenset(
37 name for name, _ in _QUERY_LENGTH_ALIASES + _KEY_LENGTH_ALIASES
38) | {"packed_seq_params"}
40_COMPRESSED_CAUSAL_MASK_SIZE = 2048
43@dataclass
44class _FusionAttentionContext:
45 """Normalized arguments required by the NPU fusion-attention operator."""
47 query: torch.Tensor
48 key: torch.Tensor
49 value: torch.Tensor
50 input_layout: str
51 attention_mask: Optional[torch.Tensor]
52 sparse_mode: int
53 query_lengths: Optional[Sequence[int]]
54 key_lengths: Optional[Sequence[int]]
55 batch_size: int
56 query_length: int
57 head_dim: int
58 pre_tokens: int
59 next_tokens: int
60 is_packed: bool
61 valid_rows: Optional[torch.Tensor]
62 indices: Optional[torch.Tensor] = None
65def _npu_attention_mask(attention_mask: torch.Tensor) -> torch.Tensor:
66 """Convert the Transformers attention mask to the NPU mask convention."""
67 if attention_mask.dtype == torch.bool:
68 return torch.logical_not(attention_mask)
69 return attention_mask != 0
72def _causal_attention_mask(
73 query: torch.Tensor,
74 key: torch.Tensor,
75 sliding_window: Optional[int],
76) -> torch.Tensor:
77 """Build the right-aligned causal mask used by Transformers decoding."""
78 query_length = query.shape[2]
79 key_length = key.shape[2]
80 query_positions = torch.arange(query_length, device=query.device)
81 query_positions = query_positions + max(key_length - query_length, 0)
82 key_positions = torch.arange(key_length, device=query.device)
83 allowed = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
84 if sliding_window is not None:
85 allowed &= key_positions.unsqueeze(0) > (
86 query_positions.unsqueeze(1) - sliding_window
87 )
88 return torch.logical_not(allowed)
91def _get_compressed_causal_mask(device: torch.device) -> torch.Tensor:
92 """Return the mask required by Ascend compressed causal sparse modes."""
93 return torch.ones(
94 (_COMPRESSED_CAUSAL_MASK_SIZE, _COMPRESSED_CAUSAL_MASK_SIZE),
95 dtype=torch.bool,
96 device=device,
97 ).triu(diagonal=1)
100def _length_list(
101 value: torch.Tensor | Sequence[int],
102 *,
103 includes_zero: bool,
104) -> list[int]:
105 """Normalize one cumulative-length representation."""
106 if isinstance(value, torch.Tensor):
107 if value.ndim != 1:
108 raise ValueError("packed cumulative sequence lengths must be one-dimensional")
109 value = value.tolist()
110 if any(isinstance(item, bool) or not isinstance(item, Integral) for item in value):
111 raise ValueError("packed cumulative sequence lengths must be integers")
112 lengths = [int(item) for item in value]
113 if includes_zero:
114 if not lengths or lengths[0] != 0:
115 raise ValueError("cu_seq_lens must start with zero")
116 lengths = lengths[1:]
117 if not lengths or any(left >= right for left, right in zip([0] + lengths[:-1], lengths)):
118 raise ValueError("packed cumulative sequence lengths must be strictly increasing")
119 return lengths
122def _coalesce_lengths(
123 kwargs: dict[str, Any],
124 aliases: tuple[tuple[str, bool], ...],
125 *,
126 name: str,
127) -> list[int] | None:
128 """Read equivalent length arguments and reject conflicting values."""
129 candidates = []
130 sources = (("", kwargs), ("packed_seq_params.", kwargs.get("packed_seq_params")))
131 for prefix, source in sources:
132 if source is None:
133 continue
134 for alias, includes_zero in aliases:
135 value = source.get(alias) if isinstance(source, Mapping) else getattr(source, alias, None)
136 if value is not None:
137 candidates.append((prefix + alias, _length_list(value, includes_zero=includes_zero)))
138 if not candidates:
139 return None
140 first_alias, first = candidates[0]
141 for alias, candidate in candidates[1:]:
142 if candidate != first:
143 raise ValueError(
144 f"conflicting {name} values were provided through "
145 f"{first_alias!r} and {alias!r}"
146 )
147 return first
150def resolve_packed_sequence_lengths(
151 kwargs: dict[str, Any],
152 query_tokens: int,
153 key_tokens: int,
154) -> tuple[list[int] | None, list[int] | None]:
155 """Resolve supported packed-sequence argument aliases.
157 Args:
158 kwargs: Attention keyword arguments or a ``packed_seq_params`` carrier.
159 query_tokens: Number of flattened query tokens.
160 key_tokens: Number of flattened key/value tokens.
162 Returns:
163 Cumulative query and key/value sequence ends without leading zeros.
164 """
165 query_lengths = _coalesce_lengths(
166 kwargs, _QUERY_LENGTH_ALIASES, name="query sequence lengths",
167 )
168 key_lengths = _coalesce_lengths(
169 kwargs, _KEY_LENGTH_ALIASES, name="key/value sequence lengths",
170 )
171 if kwargs.get("packed_seq_params") is not None and query_lengths is None and key_lengths is None:
172 raise ValueError("packed_seq_params must provide cumulative query and key/value sequence lengths")
173 if (query_lengths is None) != (key_lengths is None):
174 raise ValueError("packed attention requires both query and key/value lengths")
175 if query_lengths is not None:
176 if query_lengths[-1] != query_tokens:
177 raise ValueError("the final query sequence length must equal the query token count")
178 if key_lengths[-1] != key_tokens:
179 raise ValueError(
180 "the final key/value sequence length must equal the key/value token count"
181 )
182 return query_lengths, key_lengths
185def _attention_options(
186 module: torch.nn.Module,
187 kwargs: dict[str, Any],
188 query_length: int,
189 key_length: int,
190 *,
191 is_packed: bool,
192):
193 """Resolve sparse-window and causal options from kwargs and the module."""
194 pre_tokens = kwargs.get("pre_tokens", getattr(module, "pre_tockens", 1048576))
195 sparse_mode = kwargs.get("sparse_mode", getattr(module, "sparse_mode", 0))
196 sliding_window = kwargs.get("sliding_window")
197 is_causal = kwargs.get("is_causal", getattr(module, "is_causal", True))
198 next_tokens = kwargs.get("next_tokens", getattr(module, "next_tockens", 0 if is_causal else 2147483647))
199 if sliding_window is not None:
200 pre_tokens = sliding_window
201 elif not is_packed and is_causal and query_length == key_length and pre_tokens < query_length:
202 sliding_window = pre_tokens + 1
203 return pre_tokens, next_tokens, sparse_mode, sliding_window, is_causal
206def _prepare_attention_inputs(
207 query: torch.Tensor,
208 key: torch.Tensor,
209 value: torch.Tensor,
210 attention_mask: Optional[torch.Tensor],
211 *,
212 is_packed: bool,
213 is_causal: bool,
214 sliding_window: Optional[int],
215 sparse_mode: int,
216) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str, Optional[torch.Tensor], int]:
217 """Prepare QKV layout, mask, and sparse mode for the NPU operator."""
218 if is_packed:
219 query = query.transpose(1, 2).reshape(-1, query.shape[1], query.shape[-1])
220 key = key.transpose(1, 2).reshape(-1, key.shape[1], key.shape[-1])
221 value = value.transpose(1, 2).reshape(-1, value.shape[1], value.shape[-1])
222 if attention_mask is None and is_causal:
223 npu_mask = _get_compressed_causal_mask(query.device)
224 sparse_mode = 3
225 else:
226 npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
227 return query, key, value, "TND", npu_mask, sparse_mode
228 if attention_mask is None and is_causal:
229 if sparse_mode in (2, 4):
230 return query, key, value, "BNSD", _get_compressed_causal_mask(query.device), sparse_mode
231 return query, key, value, "BNSD", _causal_attention_mask(query, key, sliding_window), 0
232 if attention_mask is not None and attention_mask.ndim == 2:
233 attention_mask = attention_mask.to(device=query.device, dtype=torch.bool)[:, None, None, :]
234 attention_mask = attention_mask.expand(-1, 1, query.shape[2], -1)
235 npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
236 if npu_mask is not None and is_causal and sparse_mode == 0:
237 npu_mask = npu_mask | _causal_attention_mask(query, key, sliding_window)
238 return query, key, value, "BNSD", npu_mask, sparse_mode
241def _prepare_padding_attention(
242 context: _FusionAttentionContext,
243 sliding_window: Optional[int],
244) -> bool:
245 """Gather valid QKV tokens directly into TND, retaining output indices."""
246 attention_mask = context.attention_mask
247 if attention_mask.shape[-1] != context.query_length:
248 raise ValueError("the padding mask length must equal the query/key sequence length")
249 attention_mask = attention_mask.to(device=context.query.device, dtype=torch.bool).expand(context.batch_size, -1)
250 context.attention_mask = attention_mask
251 if bool(attention_mask.all()):
252 context.attention_mask = None
253 return False
254 if sliding_window is not None:
255 # Removing interior holes would change the original window distances.
256 starts = attention_mask[:, 0].int() + (attention_mask[:, 1:] & ~attention_mask[:, :-1]).sum(dim=-1)
257 if bool((starts > 1).any()):
258 return False
259 context.indices = attention_mask.flatten().nonzero(as_tuple=False).flatten()
260 lengths = attention_mask.sum(dim=-1, dtype=torch.int32)
261 # Empty samples have no selected tokens and need no segment in the NPU call.
262 context.query_lengths = tuple(lengths[lengths > 0].cumsum(0).cpu().tolist())
263 context.key_lengths = context.query_lengths
264 context.query, context.key, context.value = (
265 tensor.transpose(1, 2).flatten(0, 1).index_select(0, context.indices)
266 for tensor in (context.query, context.key, context.value)
267 )
268 context.input_layout = "TND"
269 context.is_packed = True
270 context.attention_mask = _get_compressed_causal_mask(context.query.device)
271 context.sparse_mode = 3 if sliding_window is None else 4
272 if sliding_window is not None:
273 context.pre_tokens = sliding_window - 1
274 return True
277def _prepare_attention_mask(
278 context: _FusionAttentionContext,
279 is_causal: bool,
280 sliding_window: Optional[int],
281) -> None:
282 """Select padding, compressed causal/window, or explicit-mask preparation."""
283 equal_lengths = context.query.shape[2] == context.key.shape[2]
284 compressible = not context.is_packed and equal_lengths and is_causal and context.next_tokens == 0
285 compressible = compressible and (sliding_window is None or sliding_window > 0)
286 if compressible and context.attention_mask is not None and context.attention_mask.ndim == 2:
287 if _prepare_padding_attention(context, sliding_window):
288 return
289 if not context.is_packed and context.attention_mask is None and is_causal:
290 context.sparse_mode = 0
291 if compressible:
292 if sliding_window is not None:
293 context.sparse_mode = 4
294 context.pre_tokens = sliding_window - 1
295 elif context.query_length > _COMPRESSED_CAUSAL_MASK_SIZE:
296 context.sparse_mode = 2
297 prepared = _prepare_attention_inputs(
298 context.query, context.key, context.value, context.attention_mask,
299 is_packed=context.is_packed,
300 is_causal=is_causal,
301 sliding_window=sliding_window,
302 sparse_mode=context.sparse_mode,
303 )
304 (
305 context.query, context.key, context.value,
306 context.input_layout, context.attention_mask, context.sparse_mode,
307 ) = prepared
310def _prepare_fusion_attention_context(
311 module: torch.nn.Module,
312 query: torch.Tensor,
313 key: torch.Tensor,
314 value: torch.Tensor,
315 attention_mask: Optional[torch.Tensor],
316 kwargs: dict[str, Any],
317) -> _FusionAttentionContext:
318 """Normalize model inputs into one NPU fusion-attention invocation."""
319 query_lengths, key_lengths = resolve_packed_sequence_lengths(
320 kwargs,
321 query.shape[0] * query.shape[2],
322 key.shape[0] * key.shape[2],
323 )
324 options = _attention_options(
325 module, kwargs, query.shape[2], key.shape[2], is_packed=query_lengths is not None,
326 )
327 context = _FusionAttentionContext(
328 query, key, value, "BNSD", attention_mask, options[2],
329 query_lengths,
330 key_lengths,
331 query.shape[0],
332 query.shape[2],
333 query.shape[-1],
334 options[0],
335 options[1],
336 query_lengths is not None,
337 None,
338 )
339 _prepare_attention_mask(context, options[4], options[3])
340 if not context.is_packed and context.attention_mask is not None and context.sparse_mode == 0:
341 context.valid_rows = (~context.attention_mask).any(dim=-1)
342 # Give fully masked rows one finite softmax entry, then suppress their
343 # output and upstream gradient. Some NPU kernels otherwise return nonzero rows.
344 context.attention_mask[..., 0] &= context.valid_rows
345 return context
348def npu_fusion_attention_forward(
349 module: torch.nn.Module,
350 query: torch.Tensor,
351 key: torch.Tensor,
352 value: torch.Tensor,
353 attention_mask: Optional[torch.Tensor],
354 dropout: float = 0.0,
355 scaling: Optional[float] = None,
356 **kwargs: Any,
357) -> tuple[torch.Tensor, None]:
358 """Compute NPU fusion attention using the Transformers AttentionInterface contract.
360 Args:
361 module: Owning Transformers attention module.
362 query: Query tensor in BNSD layout.
363 key: Key tensor in BNSD layout.
364 value: Value tensor in BNSD layout.
365 attention_mask: Attention mask using the Transformers mask convention.
366 dropout: Dropout probability.
367 scaling: Attention score scaling factor.
368 **kwargs: Additional attention arguments.
370 Returns:
371 Attention output in BSND layout and ``None`` for attention weights.
373 Raises:
374 ValueError: If sparse attention indices or inconsistent sequence lengths
375 are supplied.
377 Note:
378 QKV inputs use four-dimensional BNSD layout. Packed inputs are flattened
379 to TND when cumulative sequence lengths are supplied through the project
380 ``actual_*`` aliases or Transformers ``cu_seq_lens_*`` names. Float16,
381 bfloat16, and float32 inputs have been verified. A boolean mask uses
382 ``True`` for positions that participate in attention; an additive mask
383 uses zero for those positions. A two-dimensional padding mask uses
384 one/True for valid tokens. Equal-length causal padding attention uses
385 variable-length attention and restores padded queries as zeros. Windowed
386 attention uses this path for contiguous padding; masks with interior
387 holes retain their original positional window through an explicit mask.
388 """
389 # Length parsing is also used by CPU CP validation; load the optional kernel only here.
390 import torch_npu # pylint: disable=C0415
392 if kwargs.get("indices") is not None:
393 raise ValueError(
394 "npu_fusion_attention_forward does not consume sparse attention indices; "
395 "select a DSA sparse-attention implementation instead."
396 )
397 context = _prepare_fusion_attention_context(
398 module,
399 query,
400 key,
401 value,
402 attention_mask,
403 kwargs,
404 )
405 if context.indices is not None and context.indices.numel() == 0:
406 output = query.new_zeros((query.shape[0], query.shape[2], query.shape[1], query.shape[3]))
407 return output + context.query.sum() + context.key.sum() + context.value.sum(), None
408 output = torch_npu.npu_fusion_attention(
409 context.query.contiguous(),
410 context.key.contiguous(),
411 context.value.contiguous(),
412 context.query.shape[1],
413 context.input_layout,
414 pse=None,
415 padding_mask=None,
416 atten_mask=context.attention_mask,
417 scale=context.head_dim**-0.5 if scaling is None else scaling,
418 pre_tockens=context.pre_tokens,
419 next_tockens=context.next_tokens,
420 keep_prob=1.0 - dropout,
421 inner_precise=0,
422 sparse_mode=context.sparse_mode,
423 actual_seq_qlen=context.query_lengths,
424 actual_seq_kvlen=context.key_lengths,
425 )[0]
426 if context.valid_rows is not None:
427 output = output.masked_fill(~context.valid_rows.unsqueeze(-1), 0.0)
428 if context.indices is not None:
429 restored = output.new_zeros((context.batch_size * context.query_length, *output.shape[1:]))
430 output = restored.index_copy(0, context.indices, output)
431 if context.is_packed:
432 output = output.reshape(
433 context.batch_size,
434 context.query_length,
435 output.shape[1],
436 output.shape[2],
437 )
438 return output, None
439 return output.transpose(1, 2), None