Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / sink_attention.py: 0%

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"""Parameter-sink attention function.""" 

16 

17from typing import Optional, Sequence 

18 

19import torch # pylint: disable=forbidden-backend-import 

20import omni_training_custom_ops # noqa: F401 # pylint: disable=unused-import 

21 

22 

23def sink_attention( 

24 query: torch.Tensor, 

25 key: torch.Tensor, 

26 value: torch.Tensor, 

27 num_heads: int, 

28 attention_mask: Optional[torch.Tensor] = None, 

29 scale: float = 1.0, 

30 pre_tokens: int = 1048576, 

31 next_tokens: int = 0, 

32 keep_prob: float = 1.0, 

33 sparse_mode: int = 0, 

34 actual_seq_qlen: Optional[Sequence[int]] = None, 

35 actual_seq_kvlen: Optional[Sequence[int]] = None, 

36 sink_num: int = 0, 

37) -> torch.Tensor: 

38 """Compute NPU fused parameter-sink attention in TND layout. 

39 

40 Args: 

41 query: Query tensor in TND layout. 

42 key: Key tensor in TND layout with sink tokens prepended. 

43 value: Value tensor in TND layout with sink tokens prepended. 

44 num_heads: Number of query attention heads. 

45 attention_mask: Optional attention mask. 

46 scale: Scaling factor applied to attention scores. 

47 pre_tokens: Number of preceding tokens considered by sparse attention. 

48 next_tokens: Number of following tokens considered by sparse attention. 

49 keep_prob: Dropout keep probability. 

50 sparse_mode: Sparse attention mode passed to the custom kernel. 

51 actual_seq_qlen: Optional cumulative query sequence lengths. 

52 actual_seq_kvlen: Optional cumulative key/value sequence lengths, 

53 excluding the prepended sink tokens. 

54 sink_num: Number of 64-token sink blocks. 

55 

56 Returns: 

57 Attention output tensor. 

58 """ 

59 return torch.ops.custom.npu_flash_attention_score_enhance( 

60 query, 

61 key, 

62 value, 

63 num_heads, 

64 pse=None, 

65 padding_mask=None, 

66 atten_mask=attention_mask, 

67 scale=scale, 

68 keep_prob=keep_prob, 

69 input_layout="TND", 

70 actual_seq_qlen=actual_seq_qlen, 

71 actual_seq_kvlen=actual_seq_kvlen, 

72 pre_tokens=pre_tokens, 

73 next_tokens=next_tokens, 

74 inner_precise=0, 

75 sparse_mode=sparse_mode, 

76 prefix=[], 

77 sink_num=sink_num, 

78 )[0]