Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / attention_rescale.py: 0%
56 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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):
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 ) = ctx.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):
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 ) = ctx.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."""
214 del grad_softmax_max
215 grad_query, grad_key, grad_value = _normal_attention_backward(ctx, grad_rescaled_output)
216 sink_grad_query, sink_grad_key, sink_grad_value = _sink_attention_backward(ctx, grad_rescaled_output)
217 return (
218 grad_query + sink_grad_query,
219 grad_key,
220 grad_value,
221 sink_grad_key,
222 sink_grad_value,
223 *((None,) * 11),
224 )
227def attention_rescale(
228 query: torch.Tensor,
229 key: torch.Tensor,
230 value: torch.Tensor,
231 sink_key: torch.Tensor,
232 sink_value: torch.Tensor,
233 attention_mask: torch.Tensor,
234 batch_size: int,
235 sequence_length: int,
236 num_heads: int,
237 scale: float,
238 pre_tokens: int,
239 next_tokens: int,
240 keep_prob: float,
241 sparse_mode: int,
242 actual_seq_qlen: Optional[Sequence[int]],
243 actual_seq_kvlen: Optional[Sequence[int]],
244) -> Tuple[torch.Tensor, torch.Tensor]:
245 """Combine normal and parameter-sink fusion-attention results.
247 Query, key, and value use TND layout. Sink key and value use SBH layout.
248 Sequence lengths are cumulative Python sequences, as required by the
249 torch-npu TND fusion-attention interface.
251 Returns:
252 Rescaled attention output in SBND layout and the normal attention
253 softmax maximum.
254 """
255 return _AttentionRescale.apply(
256 query,
257 key,
258 value,
259 sink_key,
260 sink_value,
261 attention_mask,
262 batch_size,
263 sequence_length,
264 num_heads,
265 scale,
266 pre_tokens,
267 next_tokens,
268 keep_prob,
269 sparse_mode,
270 actual_seq_qlen,
271 actual_seq_kvlen,
272 )