Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / dsa_kl_loss.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 indexer KL loss.""" 

16 

17from typing import Any, Tuple 

18 

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

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

21 

22 

23class _DSAKLLoss(torch.autograd.Function): 

24 """Autograd bridge for the DSA indexer KL-loss custom operator.""" 

25 

26 @staticmethod 

27 def forward( 

28 ctx: Any, 

29 index_query: torch.Tensor, 

30 index_key: torch.Tensor, 

31 merge_weight: torch.Tensor, 

32 query: torch.Tensor, 

33 key: torch.Tensor, 

34 topk_indices: torch.Tensor, 

35 softmax_max: torch.Tensor, 

36 softmax_sum: torch.Tensor, 

37 query_rope: torch.Tensor, 

38 key_rope: torch.Tensor, 

39 actual_seq_qlen: torch.Tensor, 

40 actual_seq_kvlen: torch.Tensor, 

41 scale: float, 

42 loss_coeff: float, 

43 ) -> torch.Tensor: 

44 """Compute the KL loss and save its precomputed indexer gradients.""" 

45 grad_index_query, grad_index_key, grad_merge_weight, inner_loss = ( 

46 torch.ops.custom.npu_sparse_lightning_indexer_grad_kl_loss_enhance( 

47 query=query, 

48 key=key, 

49 query_index=index_query, 

50 key_index=index_key, 

51 weights=merge_weight, 

52 sparse_indices=topk_indices, 

53 softmax_max=softmax_max, 

54 softmax_sum=softmax_sum, 

55 scale_value=scale, 

56 query_rope=query_rope, 

57 key_rope=key_rope, 

58 actual_seq_qlen=actual_seq_qlen, 

59 actual_seq_klen=actual_seq_kvlen, 

60 layout="TND", 

61 sparse_mode=3, 

62 deterministic=torch.are_deterministic_algorithms_enabled(), 

63 sparse_block_size=1, 

64 ) 

65 ) 

66 divisor = query.size(0) 

67 ctx.save_for_backward( 

68 grad_index_query / divisor * loss_coeff, 

69 grad_index_key / divisor * loss_coeff, 

70 grad_merge_weight / divisor * loss_coeff, 

71 ) 

72 return inner_loss.squeeze() / divisor * loss_coeff 

73 

74 @staticmethod 

75 def backward(ctx: Any, grad_loss: torch.Tensor) -> Tuple: 

76 """Scale the precomputed indexer gradients by the upstream gradient.""" 

77 grad_index_query, grad_index_key, grad_merge_weight = ctx.saved_tensors 

78 return ( 

79 grad_index_query * grad_loss, 

80 grad_index_key * grad_loss, 

81 grad_merge_weight * grad_loss, 

82 *((None,) * 11), 

83 ) 

84 

85 

86def dsa_kl_loss( 

87 index_query: torch.Tensor, 

88 index_key: torch.Tensor, 

89 merge_weight: torch.Tensor, 

90 query: torch.Tensor, 

91 key: torch.Tensor, 

92 topk_indices: torch.Tensor, 

93 softmax_max: torch.Tensor, 

94 softmax_sum: torch.Tensor, 

95 query_rope: torch.Tensor, 

96 key_rope: torch.Tensor, 

97 actual_seq_qlen: torch.Tensor, 

98 actual_seq_kvlen: torch.Tensor, 

99 scale: float, 

100 loss_coeff: float, 

101) -> torch.Tensor: 

102 """Compute the DSA indexer KL loss. 

103 

104 ``index_query``, ``index_key``, and ``merge_weight`` are the flattened 

105 tensors returned by :func:`dsa_indexer`. ``topk_indices``, ``softmax_max``, 

106 and ``softmax_sum`` come from :func:`dsa_sparse_attention`. The MLA query, 

107 key, and rotary states use TND layout; sequence lengths are cumulative 

108 int32 tensors, following the custom-operator call contract. 

109 The result is a scalar auxiliary loss. 

110 """ 

111 return _DSAKLLoss.apply( 

112 index_query, 

113 index_key, 

114 merge_weight, 

115 query, 

116 key, 

117 topk_indices, 

118 softmax_max, 

119 softmax_sum, 

120 query_rope, 

121 key_rope, 

122 actual_seq_qlen, 

123 actual_seq_kvlen, 

124 scale, 

125 loss_coeff, 

126 )