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

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 Any, Tuple 

18 

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

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

21from einops import rearrange 

22 

23from hyper_parallel.tools.logging import get_logger 

24 

25logger = get_logger() 

26 

27 

28def _load_sparse_attention_grad(): 

29 """Defer the optional NPU dependency until sparse-attention backward.""" 

30 from cann_ops_transformer.ops import sparse_flash_attention_grad # pylint: disable=import-outside-toplevel 

31 

32 return sparse_flash_attention_grad 

33 

34 

35def _sparse_attention_forward( 

36 query: torch.Tensor, 

37 key: torch.Tensor, 

38 indices: torch.Tensor, 

39 scale: float, 

40 query_rope: torch.Tensor | None, 

41 key_rope: torch.Tensor | None, 

42 query_lengths: torch.Tensor | None, 

43 key_lengths: torch.Tensor | None, 

44 layout: str, 

45 attention_op: Any, 

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

47 """Call torch_npu SparseFlashAttention with the V3.2 DSA contract.""" 

48 kwargs = { 

49 "query_rope": None if query_rope is None else query_rope.contiguous(), 

50 "key_rope": None if key_rope is None else key_rope.contiguous(), 

51 "sparse_block_size": 1, 

52 "layout_query": layout, 

53 "layout_kv": layout, 

54 "sparse_mode": 3, 

55 "attention_mode": 2, 

56 "return_softmax_lse": True, 

57 } 

58 if layout == "TND": 

59 kwargs.update({"actual_seq_lengths_query": query_lengths, "actual_seq_lengths_kv": key_lengths}) 

60 key = key.contiguous() 

61 return attention_op(query.contiguous(), key, key, indices.contiguous(), scale, **kwargs) 

62 

63 

64def _sparse_attention_backward( 

65 query: torch.Tensor, 

66 key: torch.Tensor, 

67 indices: torch.Tensor, 

68 grad_output: torch.Tensor, 

69 output: torch.Tensor, 

70 maximum: torch.Tensor, 

71 total: torch.Tensor, 

72 scale: float, 

73 query_rope: torch.Tensor | None, 

74 key_rope: torch.Tensor | None, 

75 query_lengths: torch.Tensor | None, 

76 key_lengths: torch.Tensor | None, 

77 layout: str, 

78 grad_op: Any, 

79) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: 

80 """Let CANN GradV2 merge gradients for the shared K/V tensor.""" 

81 key = key.contiguous() 

82 result = grad_op( 

83 query.contiguous(), key, None, indices.contiguous(), grad_output.contiguous(), 

84 output.contiguous(), maximum.contiguous(), total.contiguous(), None, scale, 1, 

85 query_rope=None if query_rope is None else query_rope.contiguous(), 

86 key_rope=None if key_rope is None else key_rope.contiguous(), 

87 actual_seq_qlen=query_lengths, 

88 actual_seq_kvlen=key_lengths, 

89 layout=layout, 

90 sparse_mode=3, 

91 attention_mode=2, 

92 ) 

93 grad_query, grad_key, _, grad_query_rope, grad_key_rope, _ = result 

94 return ( 

95 grad_query, 

96 grad_key, 

97 None if query_rope is None else grad_query_rope, 

98 None if key_rope is None else grad_key_rope, 

99 ) 

100 

101 

102class _SparseAttention(torch.autograd.Function): 

103 """Pair torch_npu sparse-attention forward with the available CANN backward API.""" 

104 

105 @staticmethod 

106 def forward( 

107 ctx: Any, 

108 query: torch.Tensor, 

109 key: torch.Tensor, 

110 indices: torch.Tensor, 

111 query_rope: torch.Tensor | None, 

112 key_rope: torch.Tensor | None, 

113 qlen: torch.Tensor | None, 

114 klen: torch.Tensor | None, 

115 scale: float, 

116 layout: str, 

117 attention_op: Any, 

118 grad_op: Any, 

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

120 """Run sparse attention and save inputs needed by its backward pass.""" 

121 outputs = _sparse_attention_forward( 

122 query, key, indices, scale, query_rope, key_rope, qlen, klen, layout, attention_op, 

123 ) 

124 optional_inputs = (query_rope, key_rope, qlen, klen) 

125 ctx.optional_input_mask = tuple(tensor is not None for tensor in optional_inputs) 

126 saved_inputs = [query, key, indices] 

127 saved_inputs.extend(tensor for tensor in optional_inputs if tensor is not None) 

128 ctx.save_for_backward(*saved_inputs, *outputs) 

129 ctx.scale, ctx.layout, ctx.grad_op = scale, layout, grad_op 

130 ctx.mark_non_differentiable(*outputs[1:]) 

131 return outputs 

132 

133 @staticmethod 

134 def backward( 

135 ctx: Any, 

136 grad_output: torch.Tensor, 

137 gradmax: torch.Tensor | None, 

138 gradsum: torch.Tensor | None, 

139 ) -> tuple[torch.Tensor | None, ...]: 

140 """Compute gradients for differentiable sparse-attention inputs.""" 

141 del gradmax, gradsum 

142 saved = iter(ctx.saved_tensors) 

143 query, key, indices = (next(saved) for _ in range(3)) 

144 optional_inputs = tuple(next(saved) if present else None for present in ctx.optional_input_mask) 

145 query_rope, key_rope, qlen, klen = optional_inputs 

146 output, maximum, total = (next(saved) for _ in range(3)) 

147 gradients = _sparse_attention_backward( 

148 query, key, indices, grad_output, output, maximum, total, ctx.scale, 

149 query_rope, key_rope, qlen, klen, ctx.layout, ctx.grad_op(), 

150 ) 

151 return gradients[0], gradients[1], None, *gradients[2:], None, None, None, None, None, None 

152 

153 

154def _cann_sparse_attention( 

155 query: torch.Tensor, 

156 key: torch.Tensor, 

157 sparse_indices: torch.Tensor, 

158 scale: float, 

159 *, 

160 query_rope: torch.Tensor | None, 

161 key_rope: torch.Tensor | None, 

162 query_lengths: torch.Tensor | None, 

163 key_lengths: torch.Tensor | None, 

164 layout: str, 

165 attention_op: Any, 

166 grad_op: Any, 

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

168 """Run CANN sparse attention with the matching backward API and return statistics.""" 

169 return _SparseAttention.apply( 

170 query, key, sparse_indices, query_rope, key_rope, 

171 query_lengths, key_lengths, scale, layout, attention_op, grad_op, 

172 ) 

173 

174 

175def dsa_sparse_attention( 

176 query_nope: torch.Tensor, 

177 compressed_kv: torch.Tensor, 

178 query_rope: torch.Tensor, 

179 key_rope: torch.Tensor, 

180 topk_indices: torch.Tensor, 

181 scale: float, 

182 actual_seq_qlen: torch.Tensor, 

183 actual_seq_kvlen: torch.Tensor, 

184 *, 

185 backend: str = "cann", 

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

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

188 

189 Args: 

190 query_nope: Absorbed MLA query in BSND layout. 

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

192 query_rope: Query rotary states in BSND layout. 

193 key_rope: Key rotary states in BSND layout. 

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

195 scale: Attention score scaling factor. 

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

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

198 backend: ``"cann"`` (default) or the legacy ``"omni"`` backend. 

199 

200 Returns: 

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

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

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

204 dimension required by the sparse-attention backward operator. 

205 """ 

206 if backend not in ("cann", "omni"): 

207 raise ValueError(f"backend must be 'cann' or 'omni', got {backend!r}") 

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

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

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

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

212 # Load only the selected optional backend so either can be used independently. 

213 # pylint: disable=import-outside-toplevel 

214 if backend == "cann": 

215 import torch_npu # pylint: disable=forbidden-backend-import 

216 

217 output, softmax_max, softmax_sum = _cann_sparse_attention( 

218 query_tnd, 

219 key_tnd, 

220 topk_indices, 

221 scale, 

222 query_rope=query_rope_tnd, 

223 key_rope=key_rope_tnd, 

224 query_lengths=actual_seq_qlen, 

225 key_lengths=actual_seq_kvlen, 

226 layout="TND", 

227 attention_op=torch_npu.npu_sparse_flash_attention, 

228 grad_op=_load_sparse_attention_grad, 

229 ) 

230 else: 

231 logger.warning( 

232 "dsa_sparse_attention backend='omni' is deprecated and will be removed in a future release; " 

233 "use backend='cann'." 

234 ) 

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

236 

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

238 query_tnd, 

239 key_tnd, 

240 key_tnd, 

241 topk_indices, 

242 scale, 

243 block_table=None, 

244 actual_seq_lengths_query=actual_seq_qlen, 

245 actual_seq_lengths_kv=actual_seq_kvlen, 

246 query_rope=query_rope_tnd, 

247 key_rope=key_rope_tnd, 

248 sparse_block_size=1, 

249 layout_query="TND", 

250 layout_kv="TND", 

251 sparse_mode=3, 

252 attention_mode=2, 

253 return_softmax_lse=True, 

254 ) 

255 # pylint: enable=import-outside-toplevel 

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

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

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

259 return output, softmax_max, softmax_sum