Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / dsa_sparse_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"""DSA sparse-attention function.""" 

16 

17from typing import Tuple 

18 

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

20import torch.nn.functional as F # pylint: disable=forbidden-backend-import 

21from einops import rearrange 

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

23 

24 

25def dsa_sparse_attention( 

26 query_nope: torch.Tensor, 

27 compressed_kv: torch.Tensor, 

28 query_rope: torch.Tensor, 

29 key_rope: torch.Tensor, 

30 topk_indices: torch.Tensor, 

31 scale: float, 

32 actual_seq_qlen: torch.Tensor, 

33 actual_seq_kvlen: torch.Tensor, 

34) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 

35 """Compute NPU DSA sparse flash attention without parameter-sink rescaling. 

36 

37 Args: 

38 query_nope: Absorbed MLA query in BSND layout. 

39 compressed_kv: Compressed MLA key/value in BSND layout. 

40 query_rope: Query rotary states in BSND layout. 

41 key_rope: Key rotary states in BSND layout. 

42 topk_indices: Int32 ``[T, N_kv, K]`` output from :func:`dsa_indexer`. 

43 scale: Attention score scaling factor. 

44 actual_seq_qlen: Cumulative query lengths as an int32 NPU tensor. 

45 actual_seq_kvlen: Cumulative key lengths as an int32 NPU tensor. 

46 

47 Returns: 

48 Attention output in BSND layout, followed by the softmax maximum and 

49 sum required by :func:`dsa_kl_loss`. The output head dimension contains 

50 the compressed-KV dimension followed by zero padding for the RoPE 

51 dimension required by the sparse-attention backward operator. 

52 """ 

53 query_tnd = rearrange(query_nope, "b s n d -> (b s) n d") 

54 key_tnd = rearrange(compressed_kv, "b s n d -> (b s) n d") 

55 query_rope_tnd = rearrange(query_rope, "b s n d -> (b s) n d") 

56 key_rope_tnd = rearrange(key_rope, "b s n d -> (b s) n d") 

57 output, softmax_max, softmax_sum = torch.ops.custom.npu_sparse_flash_attention_enhance( 

58 query_tnd, 

59 key_tnd, 

60 key_tnd, 

61 topk_indices, 

62 scale, 

63 block_table=None, 

64 actual_seq_lengths_query=actual_seq_qlen, 

65 actual_seq_lengths_kv=actual_seq_kvlen, 

66 query_rope=query_rope_tnd, 

67 key_rope=key_rope_tnd, 

68 sparse_block_size=1, 

69 layout_query="TND", 

70 layout_kv="TND", 

71 sparse_mode=3, 

72 attention_mode=2, 

73 return_softmax_lse=True, 

74 ) 

75 if query_rope.size(-1) > 0: 

76 output = F.pad(output, [0, query_rope.size(-1)]) 

77 output = rearrange(output, "(b s) n d -> b s n d", b=query_nope.shape[0]) 

78 return output, softmax_max, softmax_sum