Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / attention_rescale.py: 32%
57 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
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"""Fusion-attention rescaling for parameter-sink attention."""
17from typing import Any, Optional, Sequence, Tuple
19import torch # pylint: disable=forbidden-backend-import
20import torch_npu
21from einops import rearrange
24def _rescale_attention_outputs(
25 output: torch.Tensor,
26 softmax_max: torch.Tensor,
27 softmax_sum: torch.Tensor,
28 sink_output: torch.Tensor,
29 sink_softmax_max: torch.Tensor,
30 sink_softmax_sum: torch.Tensor,
31 sequence_length: int,
32 batch_size: int,
33 num_heads: int,
34) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
35 """Merge normal and sink attention outputs using their softmax statistics."""
36 output = rearrange(output, "(b s) n d -> s b n d", b=batch_size, s=sequence_length)
37 sink_output = rearrange(sink_output, "s b (n d) -> s b n d", n=num_heads)
39 softmax_sum = softmax_sum[:, :, 0].view(batch_size, sequence_length, num_heads).transpose(0, 1)
40 softmax_max = softmax_max[:, :, 0].view(batch_size, sequence_length, num_heads).transpose(0, 1)
41 sink_softmax_sum = rearrange(sink_softmax_sum[:, :, :, 0], "b n s -> s b n")
42 sink_softmax_max = rearrange(sink_softmax_max[:, :, :, 0], "b n s -> s b n")
44 combined_max = torch.maximum(softmax_max, sink_softmax_max)
45 output_sum = softmax_sum * torch.exp(softmax_max - combined_max)
46 sink_sum = sink_softmax_sum * torch.exp(sink_softmax_max - combined_max)
47 combined_sum = output_sum + sink_sum
48 output_scale = (output_sum / combined_sum).unsqueeze(3)
49 sink_scale = (sink_sum / combined_sum).unsqueeze(3)
50 return (
51 (output * output_scale + sink_output * sink_scale).to(dtype=output.dtype),
52 output_scale,
53 sink_scale,
54 )
57def _normal_attention_backward(ctx, grad_rescaled_output, saved_tensors):
58 """Compute the normal branch's Q/K/V gradients."""
59 (
60 query, key, value, sink_key, _, attention_mask, softmax_max, softmax_sum,
61 _, _, rescaled_output, output_scale, _,
62 ) = saved_tensors
63 grad_output = rearrange(output_scale * grad_rescaled_output, "s b n d -> (b s) n d")
64 return torch_npu.npu_fusion_attention_grad(
65 query, key, value, grad_output.to(sink_key.dtype), ctx.num_heads, "TND",
66 pse=None,
67 padding_mask=None,
68 atten_mask=attention_mask,
69 softmax_max=softmax_max,
70 softmax_sum=softmax_sum,
71 attention_in=rearrange(rescaled_output, "s b n d -> (b s) n d"),
72 scale_value=ctx.scale,
73 pre_tockens=ctx.pre_tokens,
74 next_tockens=ctx.next_tokens,
75 inner_precise=0,
76 keep_prob=ctx.keep_prob,
77 actual_seq_qlen=ctx.actual_seq_qlen,
78 actual_seq_kvlen=ctx.actual_seq_kvlen,
79 sparse_mode=ctx.sparse_mode,
80 softmax_layout="TND",
81 )[:3]
84def _sink_attention_backward(ctx, grad_rescaled_output, saved_tensors):
85 """Compute sink gradients and restore the query gradient's TND layout."""
86 (
87 query, _, _, sink_key, sink_value, _, _, _,
88 sink_softmax_max, sink_softmax_sum, rescaled_output, _, sink_scale,
89 ) = saved_tensors
90 sink_query = rearrange(query, "(b s) n d -> s b (n d)", b=ctx.batch_size, s=ctx.sequence_length)
91 grad_output = rearrange(sink_scale * grad_rescaled_output, "s b n d -> s b (n d)")
92 grad_query, grad_key, grad_value = torch_npu.npu_fusion_attention_grad(
93 sink_query, sink_key, sink_value, grad_output.to(sink_key.dtype), ctx.num_heads, "SBH",
94 pse=None,
95 padding_mask=None,
96 atten_mask=None,
97 softmax_max=sink_softmax_max,
98 softmax_sum=sink_softmax_sum,
99 attention_in=rearrange(rescaled_output, "s b n d -> s b (n d)"),
100 scale_value=ctx.scale,
101 inner_precise=0,
102 keep_prob=ctx.keep_prob,
103 actual_seq_qlen=None,
104 actual_seq_kvlen=None,
105 sparse_mode=0,
106 )[:3]
107 grad_query = rearrange(
108 grad_query, "s b (n d) -> (b s) n d", b=ctx.batch_size, s=ctx.sequence_length, n=ctx.num_heads,
109 )
110 return grad_query, grad_key, grad_value
113class _AttentionRescale(torch.autograd.Function):
114 """Autograd bridge for fusion attention with separate sink parameters."""
116 @staticmethod
117 def forward(
118 ctx: Any,
119 query: torch.Tensor,
120 key: torch.Tensor,
121 value: torch.Tensor,
122 sink_key: torch.Tensor,
123 sink_value: torch.Tensor,
124 attention_mask: torch.Tensor,
125 batch_size: int,
126 sequence_length: int,
127 num_heads: int,
128 scale: float,
129 pre_tokens: int,
130 next_tokens: int,
131 keep_prob: float,
132 sparse_mode: int,
133 actual_seq_qlen: Optional[Sequence[int]],
134 actual_seq_kvlen: Optional[Sequence[int]],
135 ) -> Tuple[torch.Tensor, torch.Tensor]:
136 """Run normal and sink fusion attention and rescale their outputs."""
137 output, softmax_max, softmax_sum = torch_npu.npu_fusion_attention(
138 query,
139 key,
140 value,
141 num_heads,
142 "TND",
143 pse=None,
144 padding_mask=None,
145 atten_mask=attention_mask,
146 scale=scale,
147 pre_tockens=pre_tokens,
148 next_tockens=next_tokens,
149 keep_prob=keep_prob,
150 inner_precise=0,
151 sparse_mode=sparse_mode,
152 actual_seq_qlen=actual_seq_qlen,
153 actual_seq_kvlen=actual_seq_kvlen,
154 softmax_layout="TND",
155 )[:3]
156 sink_query = rearrange(query, "(b s) n d -> s b (n d)", b=batch_size, s=sequence_length)
157 sink_output, sink_softmax_max, sink_softmax_sum = torch_npu.npu_fusion_attention(
158 sink_query,
159 sink_key,
160 sink_value,
161 num_heads,
162 "SBH",
163 pse=None,
164 padding_mask=None,
165 atten_mask=None,
166 scale=scale,
167 keep_prob=keep_prob,
168 inner_precise=0,
169 sparse_mode=0,
170 actual_seq_qlen=None,
171 actual_seq_kvlen=None,
172 )[:3]
173 rescaled_output, output_scale, sink_scale = _rescale_attention_outputs(
174 output,
175 softmax_max,
176 softmax_sum,
177 sink_output,
178 sink_softmax_max,
179 sink_softmax_sum,
180 sequence_length,
181 batch_size,
182 num_heads,
183 )
184 ctx.save_for_backward(
185 query,
186 key,
187 value,
188 sink_key,
189 sink_value,
190 attention_mask,
191 softmax_max,
192 softmax_sum,
193 sink_softmax_max,
194 sink_softmax_sum,
195 rescaled_output,
196 output_scale,
197 sink_scale,
198 )
199 ctx.batch_size = batch_size
200 ctx.sequence_length = sequence_length
201 ctx.num_heads = num_heads
202 ctx.scale = scale
203 ctx.pre_tokens = pre_tokens
204 ctx.next_tokens = next_tokens
205 ctx.keep_prob = keep_prob
206 ctx.sparse_mode = sparse_mode
207 ctx.actual_seq_qlen = actual_seq_qlen
208 ctx.actual_seq_kvlen = actual_seq_kvlen
209 return rescaled_output, softmax_max
211 @staticmethod
212 def backward(ctx: Any, grad_rescaled_output: torch.Tensor, grad_softmax_max: torch.Tensor) -> tuple:
213 """Run the explicit fusion-attention backward operators.
215 Args:
216 ctx: Autograd context containing the saved forward tensors.
217 grad_rescaled_output: Gradient of the rescaled attention output.
218 grad_softmax_max: Unused gradient of the auxiliary softmax maximum.
219 """
220 del grad_softmax_max
221 saved_tensors = ctx.saved_tensors
222 grad_query, grad_key, grad_value = _normal_attention_backward(ctx, grad_rescaled_output, saved_tensors)
223 sink_grad_query, sink_grad_key, sink_grad_value = _sink_attention_backward(
224 ctx, grad_rescaled_output, saved_tensors,
225 )
226 return (
227 grad_query + sink_grad_query,
228 grad_key,
229 grad_value,
230 sink_grad_key,
231 sink_grad_value,
232 *((None,) * 11),
233 )
236def attention_rescale(
237 query: torch.Tensor,
238 key: torch.Tensor,
239 value: torch.Tensor,
240 sink_key: torch.Tensor,
241 sink_value: torch.Tensor,
242 attention_mask: torch.Tensor,
243 batch_size: int,
244 sequence_length: int,
245 num_heads: int,
246 scale: float,
247 pre_tokens: int,
248 next_tokens: int,
249 keep_prob: float,
250 sparse_mode: int,
251 actual_seq_qlen: Optional[Sequence[int]],
252 actual_seq_kvlen: Optional[Sequence[int]],
253) -> Tuple[torch.Tensor, torch.Tensor]:
254 """Combine normal and parameter-sink fusion-attention results.
256 Query, key, and value use TND layout. Sink key and value use SBH layout.
257 Sequence lengths are cumulative Python sequences, as required by the
258 NPU TND fusion-attention interface.
260 Args:
261 query, key, value: Inputs for the normal attention branch.
262 sink_key, sink_value: Key and value inputs for the sink branch.
263 attention_mask: Mask applied to the normal attention branch.
264 batch_size, sequence_length: Input batch and sequence dimensions.
265 num_heads: Number of attention heads.
266 scale: Attention score scale.
267 pre_tokens, next_tokens: Sparse-attention token limits.
268 keep_prob: Attention dropout keep probability.
269 sparse_mode: Sparse-attention mode.
270 actual_seq_qlen, actual_seq_kvlen: Cumulative query and key/value lengths.
272 Returns:
273 Rescaled attention output in SBND layout and the normal attention
274 softmax maximum.
275 """
276 return _AttentionRescale.apply(
277 query,
278 key,
279 value,
280 sink_key,
281 sink_value,
282 attention_mask,
283 batch_size,
284 sequence_length,
285 num_heads,
286 scale,
287 pre_tokens,
288 next_tokens,
289 keep_prob,
290 sparse_mode,
291 actual_seq_qlen,
292 actual_seq_kvlen,
293 )