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

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"""Dense DSA indexer loss using CANN's fused teacher and KL operators.""" 

16 

17from __future__ import annotations 

18 

19from typing import Any, Tuple 

20 

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

22 

23 

24def _dense_indexer_lse( 

25 query_index: torch.Tensor, 

26 key_index: torch.Tensor, 

27 weights: torch.Tensor, 

28 query_lengths: torch.Tensor | None, 

29 key_lengths: torch.Tensor | None, 

30 layout: str, 

31) -> tuple[torch.Tensor, torch.Tensor]: 

32 """Return native indexer maximum and exponential sum, without gradients.""" 

33 # Dense operators remain optional for CPU, sparse and Omni-only execution. 

34 from torch_npu import ( # pylint: disable=import-outside-toplevel,forbidden-backend-import 

35 npu_dense_lightning_indexer_softmax_lse, 

36 ) 

37 

38 with torch.no_grad(): 

39 return npu_dense_lightning_indexer_softmax_lse( 

40 query_index.contiguous(), key_index.contiguous(), weights.contiguous(), 

41 actual_seq_qlen=None if query_lengths is None else query_lengths.tolist(), 

42 actual_seq_klen=None if key_lengths is None else key_lengths.tolist(), 

43 layout=layout, sparse_mode=3, 

44 ) 

45 

46 

47def dsa_dense_indexer_lse( 

48 index_query: torch.Tensor, 

49 index_key: torch.Tensor, 

50 head_weights: torch.Tensor, 

51 *, 

52 query_lengths: torch.Tensor | None = None, 

53 key_lengths: torch.Tensor | None = None, 

54 layout: str = "TND", 

55) -> tuple[torch.Tensor, torch.Tensor]: 

56 """Return detached FP32 indexer max/sum for fused dense KL. 

57 

58 TND lengths are cumulative document ends without a leading zero; BSND 

59 lengths are per-batch valid lengths. TND requires both length tensors. 

60 The two outputs are maximum and exponential sum, rather than logsumexp. 

61 """ 

62 return _dense_indexer_lse( 

63 index_query, index_key, head_weights, query_lengths, key_lengths, layout, 

64 ) 

65 

66 

67class _DenseIndexerKLLoss(torch.autograd.Function): 

68 """Connect CANN's explicit dense KL gradients to indexer projections.""" 

69 

70 @staticmethod 

71 def forward( 

72 ctx: Any, 

73 index_query: torch.Tensor, 

74 index_key: torch.Tensor, 

75 head_weights: torch.Tensor, 

76 query: torch.Tensor, 

77 key: torch.Tensor, 

78 softmax_max: torch.Tensor, 

79 softmax_sum: torch.Tensor, 

80 indexer_lse: tuple[torch.Tensor, torch.Tensor], 

81 scale: float, 

82 query_rope: torch.Tensor | None, 

83 key_rope: torch.Tensor | None, 

84 query_lengths: torch.Tensor | None, 

85 key_lengths: torch.Tensor | None, 

86 layout: str, 

87 loss_coefficient: float, 

88 ) -> torch.Tensor: 

89 """Compute summed KL in CANN and retain token-mean indexer gradients.""" 

90 query_ends = None if query_lengths is None else query_lengths.tolist() 

91 key_ends = None if key_lengths is None else key_lengths.tolist() 

92 token_count = index_query.shape[0] if layout == "TND" else ( 

93 index_query.shape[0] * index_query.shape[1] if query_ends is None else sum(query_ends) 

94 ) 

95 # Load only when the dense native path is selected. 

96 from torch_npu import ( # pylint: disable=import-outside-toplevel,forbidden-backend-import 

97 npu_dense_lightning_indexer_grad_kl_loss, 

98 ) 

99 

100 indexer_max, indexer_sum = indexer_lse 

101 grad_query, grad_key, grad_weights, loss = npu_dense_lightning_indexer_grad_kl_loss( 

102 query.contiguous(), key.contiguous(), index_query.contiguous(), index_key.contiguous(), 

103 head_weights.contiguous(), softmax_max.contiguous(), softmax_sum.contiguous(), 

104 indexer_max.contiguous(), indexer_sum.contiguous(), scale, 

105 query_rope=None if query_rope is None else query_rope.expand( 

106 *query.shape[:-1], query_rope.shape[-1] 

107 ).contiguous(), 

108 key_rope=None if key_rope is None else key_rope.expand( 

109 *key.shape[:-1], key_rope.shape[-1] 

110 ).contiguous(), 

111 actual_seq_qlen=query_ends, actual_seq_klen=key_ends, 

112 layout=layout, sparse_mode=3, 

113 ) 

114 ctx.save_for_backward( 

115 grad_query / token_count * loss_coefficient, 

116 grad_key / token_count * loss_coefficient, 

117 grad_weights / token_count * loss_coefficient, 

118 ) 

119 return loss.squeeze() / token_count * loss_coefficient 

120 

121 @staticmethod 

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

123 """Apply the upstream loss seed only to indexer gradients.""" 

124 grad_query, grad_key, grad_weights = ctx.saved_tensors 

125 return ( 

126 grad_query * grad_loss, 

127 grad_key * grad_loss, 

128 grad_weights * grad_loss, 

129 *((None,) * 12), 

130 ) 

131 

132 

133def dsa_dense_indexer_kl_loss( 

134 index_query: torch.Tensor, 

135 index_key: torch.Tensor, 

136 head_weights: torch.Tensor, 

137 query: torch.Tensor, 

138 key: torch.Tensor, 

139 softmax_max: torch.Tensor, 

140 softmax_sum: torch.Tensor, 

141 indexer_lse: tuple[torch.Tensor, torch.Tensor], 

142 *, 

143 scale: float, 

144 query_rope: torch.Tensor | None = None, 

145 key_rope: torch.Tensor | None = None, 

146 query_lengths: torch.Tensor | None = None, 

147 key_lengths: torch.Tensor | None = None, 

148 layout: str = "TND", 

149 loss_coefficient: float = 1.0, 

150) -> torch.Tensor: 

151 """Return token-mean dense KL; propagate gradients only to indexer inputs. 

152 

153 Main Attention Q/K use unabsorbed D128 states with equal head counts and 

154 separate RoPE. Maximum/sum use [N,T,1] or [B,N,S,1] FP32 layout. 

155 Pass the max/sum pair returned by dsa_dense_indexer_lse as indexer_lse. 

156 Teacher probabilities, explicit gradients and summed loss are computed 

157 inside the CANN fused operator. 

158 """ 

159 return _DenseIndexerKLLoss.apply( 

160 index_query, index_key, head_weights, query, key, softmax_max, softmax_sum, 

161 indexer_lse, scale, query_rope, key_rope, query_lengths, key_lengths, 

162 layout, float(loss_coefficient), 

163 ) 

164 

165 

166__all__ = [ 

167 "dsa_dense_indexer_kl_loss", 

168 "dsa_dense_indexer_lse", 

169]