Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / dsa_sparse_attention_rescale.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 rescaling for parameter-sink attention.""" 

16 

17from typing import Any, Tuple 

18 

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

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

21import torch_npu 

22from einops import rearrange 

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

24 

25 

26def _sparse_attention_forward( 

27 query_nope, compressed_kv, query_rope, key_rope, topk_indices, 

28 batch_size, sequence_length, scale, actual_seq_qlen, actual_seq_kvlen, 

29): 

30 """Run sparse attention and restore the padded output's BSND layout.""" 

31 query_nope, compressed_kv, query_rope, key_rope = [ 

32 rearrange(tensor, "b s n d -> (b s) n d") 

33 for tensor in (query_nope, compressed_kv, query_rope, key_rope) 

34 ] 

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

36 query_nope, compressed_kv, compressed_kv, topk_indices, scale, 

37 block_table=None, 

38 actual_seq_lengths_query=actual_seq_qlen, 

39 actual_seq_lengths_kv=actual_seq_kvlen, 

40 query_rope=query_rope, 

41 key_rope=key_rope, 

42 sparse_block_size=1, 

43 layout_query="TND", 

44 layout_kv="TND", 

45 sparse_mode=3, 

46 attention_mode=2, 

47 return_softmax_lse=True, 

48 ) 

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

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

51 output = rearrange(output, "(b s) n d -> b s n d", b=batch_size, s=sequence_length) 

52 return output, softmax_max, softmax_sum 

53 

54 

55def _sink_attention_forward(query_nope, query_rope, sink_key, sink_value, num_heads, scale, keep_prob): 

56 """Run the sink branch in SBH layout.""" 

57 query = torch.cat([query_nope, query_rope], dim=-1) 

58 return torch_npu.npu_fusion_attention( 

59 rearrange(query, "b s n d -> s b (n d)"), 

60 rearrange(sink_key, "b s n d -> s b (n d)"), 

61 rearrange(sink_value, "b s n d -> s b (n d)"), 

62 num_heads, "SBH", 

63 pse=None, 

64 padding_mask=None, 

65 atten_mask=None, 

66 scale=scale, 

67 keep_prob=keep_prob, 

68 inner_precise=0, 

69 sparse_mode=0, 

70 actual_seq_qlen=None, 

71 actual_seq_kvlen=None, 

72 )[:3] 

73 

74 

75def _rescale_sparse_outputs( 

76 output, softmax_max, softmax_sum, sink_output, sink_softmax_max, sink_softmax_sum, 

77 batch_size, sequence_length, num_heads, 

78): 

79 """Combine the two attention outputs using stable softmax weights.""" 

80 sink_output = rearrange(sink_output, "s b (n d) -> b s n d", n=num_heads) 

81 softmax_max = softmax_max.squeeze(0).view(batch_size, sequence_length, num_heads) 

82 softmax_sum = softmax_sum.squeeze(0).view(batch_size, sequence_length, num_heads) 

83 sink_softmax_max = sink_softmax_max[:, :, :, 0].transpose(1, 2) 

84 sink_softmax_sum = sink_softmax_sum[:, :, :, 0].transpose(1, 2) 

85 combined_max = torch.maximum(softmax_max, sink_softmax_max) 

86 output_sum = softmax_sum * torch.exp(softmax_max - combined_max) 

87 sink_sum = sink_softmax_sum * torch.exp(sink_softmax_max - combined_max) 

88 combined_sum = output_sum + sink_sum 

89 output_scale = (output_sum / combined_sum).unsqueeze(-1) 

90 sink_scale = (sink_sum / combined_sum).unsqueeze(-1) 

91 return (output * output_scale + sink_output * sink_scale).to(dtype=output.dtype), output_scale, sink_scale 

92 

93 

94def _sparse_attention_backward(ctx, grad_rescaled_output): 

95 """Compute sparse gradients and restore their BSND layout.""" 

96 ( 

97 query_nope, compressed_kv, query_rope, key_rope, sink_key, _, topk_indices, 

98 softmax_max, softmax_sum, _, _, rescaled_output, output_scale, _, 

99 ) = ctx.saved_tensors 

100 query_nope, compressed_kv, query_rope, key_rope = [ 

101 rearrange(tensor, "b s n d -> (b s) n d") 

102 for tensor in (query_nope, compressed_kv, query_rope, key_rope) 

103 ] 

104 grad_output = rearrange(output_scale * grad_rescaled_output, "b s n d -> (b s) n d") 

105 rescaled_output = rearrange(rescaled_output, "b s n d -> (b s) n d") 

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

107 grad_output = grad_output[:, :, :-query_rope.size(-1)] 

108 rescaled_output = rescaled_output[:, :, :-query_rope.size(-1)] 

109 gradients = torch.ops.custom.npu_sparse_flash_attention_grad_enhance( 

110 query_nope, compressed_kv, compressed_kv, topk_indices, 

111 grad_output.to(sink_key.dtype), rescaled_output, softmax_max, softmax_sum, ctx.scale, 

112 sparse_block_size=1, 

113 actual_seq_qlen=ctx.actual_seq_qlen, 

114 actual_seq_kvlen=ctx.actual_seq_kvlen, 

115 query_rope=query_rope, 

116 key_rope=key_rope, 

117 layout="TND", 

118 sparse_mode=3, 

119 attention_mode=2, 

120 deterministic=torch.are_deterministic_algorithms_enabled(), 

121 ) 

122 return tuple( 

123 rearrange(tensor, "(b s) n d -> b s n d", b=ctx.batch_size, s=ctx.sequence_length) 

124 for tensor in gradients 

125 ) 

126 

127 

128def _sink_attention_backward(ctx, grad_rescaled_output): 

129 """Compute sink gradients and restore their BSND layout.""" 

130 ( 

131 query_nope, _, query_rope, _, sink_key, sink_value, _, _, _, 

132 sink_softmax_max, sink_softmax_sum, rescaled_output, _, sink_scale, 

133 ) = ctx.saved_tensors 

134 query = rearrange(torch.cat([query_nope, query_rope], dim=-1), "b s n d -> s b (n d)") 

135 grad_output = rearrange(sink_scale * grad_rescaled_output, "b s n d -> s b (n d)") 

136 grad_query, grad_key, grad_value = torch_npu.npu_fusion_attention_grad( 

137 query, 

138 rearrange(sink_key, "b s n d -> s b (n d)"), 

139 rearrange(sink_value, "b s n d -> s b (n d)"), 

140 grad_output.to(sink_key.dtype), ctx.num_heads, "SBH", 

141 pse=None, 

142 padding_mask=None, 

143 atten_mask=None, 

144 softmax_max=sink_softmax_max, 

145 softmax_sum=sink_softmax_sum, 

146 attention_in=rearrange(rescaled_output, "b s n d -> s b (n d)"), 

147 scale_value=ctx.scale, 

148 inner_precise=0, 

149 keep_prob=ctx.keep_prob, 

150 actual_seq_qlen=None, 

151 actual_seq_kvlen=None, 

152 sparse_mode=0, 

153 )[:3] 

154 return ( 

155 rearrange(grad_query, "s b (n d) -> b s n d", n=ctx.num_heads), 

156 rearrange(grad_key, "s b (n d) -> b s n d", n=sink_key.size(2)), 

157 rearrange(grad_value, "s b (n d) -> b s n d", n=sink_value.size(2)), 

158 ) 

159 

160 

161class _SparseAttentionRescale(torch.autograd.Function): 

162 """Autograd bridge for sparse attention with separate sink parameters.""" 

163 

164 @staticmethod 

165 def forward( 

166 ctx: Any, 

167 query_nope: torch.Tensor, 

168 compressed_kv: torch.Tensor, 

169 query_rope: torch.Tensor, 

170 key_rope: torch.Tensor, 

171 sink_key: torch.Tensor, 

172 sink_value: torch.Tensor, 

173 topk_indices: torch.Tensor, 

174 batch_size: int, 

175 sequence_length: int, 

176 num_heads: int, 

177 scale: float, 

178 keep_prob: float, 

179 actual_seq_qlen: torch.Tensor, 

180 actual_seq_kvlen: torch.Tensor, 

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

182 """Run sparse and sink attention and rescale their outputs.""" 

183 output, softmax_max, softmax_sum = _sparse_attention_forward( 

184 query_nope, compressed_kv, query_rope, key_rope, topk_indices, 

185 batch_size, sequence_length, scale, actual_seq_qlen, actual_seq_kvlen, 

186 ) 

187 sink_output, sink_softmax_max, sink_softmax_sum = _sink_attention_forward( 

188 query_nope, query_rope, sink_key, sink_value, num_heads, scale, keep_prob, 

189 ) 

190 rescaled_output, output_scale, sink_scale = _rescale_sparse_outputs( 

191 output, softmax_max, softmax_sum, sink_output, sink_softmax_max, sink_softmax_sum, 

192 batch_size, sequence_length, num_heads, 

193 ) 

194 

195 ctx.save_for_backward( 

196 query_nope, 

197 compressed_kv, 

198 query_rope, 

199 key_rope, 

200 sink_key, 

201 sink_value, 

202 topk_indices, 

203 softmax_max, 

204 softmax_sum, 

205 sink_softmax_max, 

206 sink_softmax_sum, 

207 rescaled_output, 

208 output_scale, 

209 sink_scale, 

210 ) 

211 ctx.batch_size = batch_size 

212 ctx.sequence_length = sequence_length 

213 ctx.num_heads = num_heads 

214 ctx.scale = scale 

215 ctx.keep_prob = keep_prob 

216 ctx.actual_seq_qlen = actual_seq_qlen 

217 ctx.actual_seq_kvlen = actual_seq_kvlen 

218 return rescaled_output, softmax_max, softmax_sum 

219 

220 @staticmethod 

221 def backward( 

222 ctx: Any, 

223 grad_rescaled_output: torch.Tensor, 

224 grad_softmax_max: torch.Tensor, 

225 grad_softmax_sum: torch.Tensor, 

226 ) -> tuple: 

227 """Run the explicit sparse- and fusion-attention backward operators.""" 

228 del grad_softmax_max, grad_softmax_sum 

229 grad_query_nope, grad_key, grad_value, grad_query_rope, grad_key_rope = _sparse_attention_backward( 

230 ctx, grad_rescaled_output, 

231 ) 

232 sink_grad_query, sink_grad_key, sink_grad_value = _sink_attention_backward(ctx, grad_rescaled_output) 

233 sink_grad_query_nope, sink_grad_query_rope = torch.split( 

234 sink_grad_query, 

235 [ctx.saved_tensors[0].size(-1), ctx.saved_tensors[2].size(-1)], 

236 dim=-1, 

237 ) 

238 return ( 

239 grad_query_nope + sink_grad_query_nope, 

240 grad_key + grad_value, 

241 grad_query_rope + sink_grad_query_rope, 

242 grad_key_rope, 

243 sink_grad_key, 

244 sink_grad_value, 

245 *((None,) * 8), 

246 ) 

247 

248 

249def dsa_sparse_attention_rescale( 

250 query_nope: torch.Tensor, 

251 compressed_kv: torch.Tensor, 

252 query_rope: torch.Tensor, 

253 key_rope: torch.Tensor, 

254 sink_key: torch.Tensor, 

255 sink_value: torch.Tensor, 

256 topk_indices: torch.Tensor, 

257 batch_size: int, 

258 sequence_length: int, 

259 num_heads: int, 

260 scale: float, 

261 keep_prob: float, 

262 actual_seq_qlen: torch.Tensor, 

263 actual_seq_kvlen: torch.Tensor, 

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

265 """Combine DSA sparse attention and parameter-sink attention. 

266 

267 The four DSA states and both sink states use BSND layout. ``topk_indices`` 

268 is the int32 output from :func:`dsa_indexer`; sequence lengths are 

269 cumulative int32 NPU tensors. The returned BSND output is followed by the 

270 sparse-attention softmax maximum and sum used by :func:`dsa_kl_loss`. 

271 """ 

272 return _SparseAttentionRescale.apply( 

273 query_nope, 

274 compressed_kv, 

275 query_rope, 

276 key_rope, 

277 sink_key, 

278 sink_value, 

279 topk_indices, 

280 batch_size, 

281 sequence_length, 

282 num_heads, 

283 scale, 

284 keep_prob, 

285 actual_seq_qlen, 

286 actual_seq_kvlen, 

287 )