已合并
fix_nfa_grad_meta #38633
Ambi创建于 6月16日
fix_nfa_grad_meta #38633
已合并
共 3 个文件变更+74-93
| @@ -30,7 +30,7 @@ class TestRunWithRngState(TestUtils): | |||
| 30 | 30 | ||
| 31 | return res1, res2 | 31 | return res1, res2 |
| 32 | 32 | ||
| 33 | - @parametrize("shape", [(10,)]) | 33 | + @parametrize("shape", [(20,)]) |
| 34 | 34 | ||
| 35 | def test_rng_state_with_compile(self, shape, dtype): | 35 | def test_rng_state_with_compile(self, shape, dtype): |
| 36 | device = "npu" | 36 | device = "npu" |
| @@ -127,56 +127,7 @@ def disable_implicit_decomposition(): | |||
| 127 | op_override.py_kernels.pop(DispatchKey.CompositeImplicitAutograd) | 127 | op_override.py_kernels.pop(DispatchKey.CompositeImplicitAutograd) |
| 128 | 128 | ||
| 129 | 129 | ||
| 130 | -def _patch_run_node(tracer, node, args, kwargs, nnmodule): | ||
| 131 | - op = node.op | ||
| 132 | - | ||
| 133 | - with set_current_node(node): | ||
| 134 | - | ||
| 135 | - def make_error_message(e): | ||
| 136 | - return f"Failed running {op} {node.target}(*{args}, **{kwargs}):\n" + str(e) | ||
| 137 | - | ||
| 138 | - try: | ||
| 139 | - if op == "call_function": | ||
| 140 | - # patch start | ||
| 141 | - if 'npu.npu_fusion_attention' in str(node.target): | ||
| 142 | - if 'actual_seq_qlen' in kwargs: | ||
| 143 | - kwargs['actual_seq_qlen'] = list(kwargs['actual_seq_qlen']) | ||
| 144 | - if 'actual_seq_kvlen' in kwargs: | ||
| 145 | - kwargs['actual_seq_kvlen'] = list(kwargs['actual_seq_kvlen']) | ||
| 146 | - # patch end | ||
| 147 | - return node.target(*args, **kwargs) | ||
| 148 | - elif op == "call_method": | ||
| 149 | - return getattr(args[0], node.target)(*args[1:], **kwargs) | ||
| 150 | - elif op == "call_module": | ||
| 151 | - if nnmodule is None: | ||
| 152 | - raise RuntimeError( | ||
| 153 | - f"Module {node.target} not found in the current module" | ||
| 154 | - ) | ||
| 155 | - return nnmodule(*args, **kwargs) | ||
| 156 | - elif op == "get_attr": | ||
| 157 | - return tracer.output_graph.get_submodule(node.target) | ||
| 158 | - elif op == "placeholder": | ||
| 159 | - if "example_value" not in node.meta: | ||
| 160 | - raise RuntimeError( | ||
| 161 | - f"placeholder {node.target} has no example value" | ||
| 162 | - ) | ||
| 163 | - return node.meta["example_value"] | ||
| 164 | - | ||
| 165 | - except (NotImplementedError, UnsupportedFakeTensorException) as e: | ||
| 166 | - # NB: mimic how wrap_fake_exception does it | ||
| 167 | - from torch._dynamo.exc import unimplemented | ||
| 168 | - | ||
| 169 | - unimplemented(make_error_message(e), from_exc=e) | ||
| 170 | - except Exception as e: | ||
| 171 | - raise RuntimeError(make_error_message(e)).with_traceback( | ||
| 172 | - e.__traceback__ | ||
| 173 | - ) from e | ||
| 174 | - | ||
| 175 | - raise AssertionError(op) | ||
| 176 | - | ||
| 177 | - | ||
| 178 | disable_implicit_decomposition() | 130 | disable_implicit_decomposition() |
| 179 | -torch._dynamo.utils.run_node = _patch_run_node | ||
| 180 | 131 | ||
| 181 | 132 | ||
| 182 | from torch._dynamo.backends import common | 133 | from torch._dynamo.backends import common |
| @@ -1,15 +1,28 @@ | |||
| 1 | +from typing import ( | ||
| 2 | + Callable, | ||
| 3 | + Optional, | ||
| 4 | + Sequence, | ||
| 5 | + List, | ||
| 6 | + Tuple | ||
| 7 | +) | ||
| 8 | + | ||
| 1 | import torch | 9 | import torch |
| 2 | import torch_npu | 10 | import torch_npu |
| 3 | import torch.library | 11 | import torch.library |
| 4 | from torch.library import Library | 12 | from torch.library import Library |
| 5 | 13 | ||
| 6 | -from typing import Callable, Optional, Sequence, List, Tuple | 14 | + |
| 15 | +# Not good implementation, but no other way | ||
| 16 | +def get_current_raw_stream(device): | ||
| 17 | + return torch.npu.current_stream(device).npu_stream | ||
| 18 | + | ||
| 7 | 19 | ||
| 8 | NPU_STREAMS = {} | 20 | NPU_STREAMS = {} |
| 9 | NPU_EVENTS = {} | 21 | NPU_EVENTS = {} |
| 10 | 22 | ||
| 11 | # create a library to hold the custom op | 23 | # create a library to hold the custom op |
| 12 | npu_stream_lib = Library("npu_stream", "FRAGMENT") # noqa | 24 | npu_stream_lib = Library("npu_stream", "FRAGMENT") # noqa |
| 25 | +inductor_npu_lib = Library("inductor_npu", "FRAGMENT") # noqa | ||
| 13 | 26 | ||
| 14 | def direct_register_custom_op( | 27 | def direct_register_custom_op( |
OO 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。    Ambi 6月18日 评论： 6月18日 评论： | |||
| 15 | op_name: str, | 28 | op_name: str, |
| @@ -183,7 +196,6 @@ def graph_break( | |||
| 183 | outputs.append(torch.ops.npu_utils.graph_break(inp)) | 196 | outputs.append(torch.ops.npu_utils.graph_break(inp)) |
| 184 | return outputs | 197 | return outputs |
| 185 | 198 | ||
| 186 | -inductor_npu_lib = Library("inductor_npu", "FRAGMENT") # noqa | ||
| 187 | 199 | ||
| 188 | def npu_fusion_attention( | 200 | def npu_fusion_attention( |
| 189 | query: torch.Tensor, | 201 | query: torch.Tensor, |
| @@ -204,7 +216,9 @@ def npu_fusion_attention( | |||
| 204 | actual_seq_kvlen: Optional[torch.Tensor] = None, | 216 | actual_seq_kvlen: Optional[torch.Tensor] = None, |
| 205 | sparse_mode: int = 0, | 217 | sparse_mode: int = 0, |
| 206 | gen_mask_parallel: bool = True, | 218 | gen_mask_parallel: bool = True, |
| 207 | - sync: bool = False | 219 | + sync: bool = False, |
| 220 | + softmax_layout: str = "", | ||
| 221 | + sink: Optional[torch.Tensor] = None | ||
| 208 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | 222 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 209 | prefix = prefix.tolist() if prefix is not None else prefix | 223 | prefix = prefix.tolist() if prefix is not None else prefix |
| 210 | actual_seq_qlen = actual_seq_qlen.tolist() if actual_seq_qlen is not None else actual_seq_qlen | 224 | actual_seq_qlen = actual_seq_qlen.tolist() if actual_seq_qlen is not None else actual_seq_qlen |
| @@ -228,14 +242,20 @@ def npu_fusion_attention( | |||
| 228 | actual_seq_kvlen=actual_seq_kvlen, | 242 | actual_seq_kvlen=actual_seq_kvlen, |
| 229 | sparse_mode=sparse_mode, | 243 | sparse_mode=sparse_mode, |
| 230 | gen_mask_parallel=gen_mask_parallel, | 244 | gen_mask_parallel=gen_mask_parallel, |
| 231 | - sync=sync | 245 | + sync=sync, |
| 246 | + softmax_layout=softmax_layout, | ||
| 247 | + sink=sink | ||
| 232 | ) | 248 | ) |
| 233 | 249 | ||
| 234 | - seed = torch.tensor([seed], device='npu', dtype=torch.int64) | 250 | + seed = torch.tensor([seed], dtype=torch.int64) |
| 235 | - offset = torch.tensor([offset], device='npu', dtype=torch.int64) | 251 | + offset = torch.tensor([offset], dtype=torch.int64) |
| 236 | - numels = torch.tensor([numels], device='npu', dtype=torch.int64) | 252 | + numels = torch.tensor([numels], dtype=torch.int64) |
| 253 | + | ||
| 254 | + attn_ret = (attention_score, softmax_max, softmax_sum, softmax_out, ) | ||
| 255 | + dropout_ret = (seed, offset, numels,) | ||
| 256 | + | ||
| 257 | + return *attn_ret, *dropout_ret | ||
| 237 | 258 | ||
| 238 | - return attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels | ||
| 239 | 259 | ||
| 240 | def npu_fusion_attention_fake( | 260 | def npu_fusion_attention_fake( |
| 241 | query: torch.Tensor, | 261 | query: torch.Tensor, |
| @@ -256,10 +276,11 @@ def npu_fusion_attention_fake( | |||
| 256 | actual_seq_kvlen: Optional[torch.Tensor] = None, | 276 | actual_seq_kvlen: Optional[torch.Tensor] = None, |
| 257 | sparse_mode: int = 0, | 277 | sparse_mode: int = 0, |
| 258 | gen_mask_parallel: bool = True, | 278 | gen_mask_parallel: bool = True, |
| 259 | - sync: bool = False | 279 | + sync: bool = False, |
| 280 | + softmax_layout: str = "", | ||
| 281 | + sink: Optional[torch.Tensor] = None | ||
| 260 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | 282 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 261 | B = query.size(0) | 283 | B = query.size(0) |
| 262 | - N = head_num | ||
| 263 | S1 = query.size(2) | 284 | S1 = query.size(2) |
| 264 | S2 = key.size(2) | 285 | S2 = key.size(2) |
| 265 | 286 | ||
| @@ -277,17 +298,15 @@ def npu_fusion_attention_fake( | |||
| 277 | softmax_max = torch.empty([B, head_num, S1, 8], dtype=torch.float32, device=query.device) | 298 | softmax_max = torch.empty([B, head_num, S1, 8], dtype=torch.float32, device=query.device) |
| 278 | softmax_sum = torch.empty([B, head_num, S1, 8], dtype=torch.float32, device=query.device) | 299 | softmax_sum = torch.empty([B, head_num, S1, 8], dtype=torch.float32, device=query.device) |
| 279 | softmax_out = torch.empty([0], dtype=query.dtype, device=query.device) | 300 | softmax_out = torch.empty([0], dtype=query.dtype, device=query.device) |
| 280 | - seed = torch.empty([1], dtype=torch.int64, device=query.device) | 301 | + seed = torch.empty([1], dtype=torch.int64, device='cpu') |
| 281 | - offset = torch.empty([1], dtype=torch.int64, device=query.device) | 302 | + offset = torch.empty([1], dtype=torch.int64, device='cpu') |
| 282 | - numels = torch.empty([1], dtype=torch.int64, device=query.device) | 303 | + numels = torch.empty([1], dtype=torch.int64, device='cpu') |
| 304 | + | ||
| 305 | + attn_ret = (attention_score, softmax_max, softmax_sum, softmax_out, ) | ||
| 306 | + dropout_ret = (seed, offset, numels, ) | ||
| 307 | + | ||
| 308 | + return *attn_ret, *dropout_ret | ||
| 283 | 309 | ||
| 284 | - return (attention_score, | ||
| 285 | - softmax_max, | ||
| 286 | - softmax_sum, | ||
| 287 | - softmax_out, | ||
| 288 | - seed, | ||
| 289 | - offset, | ||
| 290 | - numels) | ||
| 291 | 310 | ||
| 292 | direct_register_custom_op( | 311 | direct_register_custom_op( |
| 293 | op_name="npu_fusion_attention", | 312 | op_name="npu_fusion_attention", |
| @@ -327,8 +346,10 @@ def npu_fusion_attention_grad( | |||
| 327 | actual_seq_kvlen: Optional[torch.Tensor] = None, | 346 | actual_seq_kvlen: Optional[torch.Tensor] = None, |
| 328 | sparse_mode: int = 0, | 347 | sparse_mode: int = 0, |
| 329 | gen_mask_parallel: bool = True, | 348 | gen_mask_parallel: bool = True, |
| 330 | - sync: bool = False | 349 | + sync: bool = False, |
| 331 | - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | 350 | + softmax_layout: str = "", |
| 351 | + sink: Optional[torch.Tensor] = None | ||
| 352 | + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| 332 | prefix = prefix.tolist() if prefix is not None else prefix | 353 | prefix = prefix.tolist() if prefix is not None else prefix |
| 333 | actual_seq_qlen = actual_seq_qlen.tolist() if actual_seq_qlen is not None else actual_seq_qlen | 354 | actual_seq_qlen = actual_seq_qlen.tolist() if actual_seq_qlen is not None else actual_seq_qlen |
| 334 | actual_seq_kvlen = actual_seq_kvlen.tolist() if actual_seq_kvlen is not None else actual_seq_kvlen | 355 | actual_seq_kvlen = actual_seq_kvlen.tolist() if actual_seq_kvlen is not None else actual_seq_kvlen |
| @@ -337,15 +358,16 @@ def npu_fusion_attention_grad( | |||
| 337 | offset = offset.item() | 358 | offset = offset.item() |
| 338 | numels = numels.item() | 359 | numels = numels.item() |
| 339 | 360 | ||
| 340 | - dq, dk, dv, dpse = torch.ops.npu.npu_fusion_attention_grad( | 361 | + dq, dk, dv, dpse, dsink = torch.ops.npu.npu_fusion_attention_grad( |
| 341 | query, key, value, dy, head_num, input_layout, pse=pse, padding_mask=padding_mask, atten_mask=atten_mask, | 362 | query, key, value, dy, head_num, input_layout, pse=pse, padding_mask=padding_mask, atten_mask=atten_mask, |
| 342 | softmax_max=softmax_max, softmax_sum=softmax_sum, softmax_in=softmax_in, attention_in=attention_in, scale_value=scale_value, | 363 | softmax_max=softmax_max, softmax_sum=softmax_sum, softmax_in=softmax_in, attention_in=attention_in, scale_value=scale_value, |
| 343 | keep_prob=keep_prob, pre_tockens=pre_tockens, next_tockens=next_tockens, inner_precise=inner_precise, seed=seed, offset=offset, | 364 | keep_prob=keep_prob, pre_tockens=pre_tockens, next_tockens=next_tockens, inner_precise=inner_precise, seed=seed, offset=offset, |
| 344 | numels=numels, prefix=prefix, actual_seq_qlen=actual_seq_qlen, actual_seq_kvlen=actual_seq_kvlen, sparse_mode=sparse_mode, | 365 | numels=numels, prefix=prefix, actual_seq_qlen=actual_seq_qlen, actual_seq_kvlen=actual_seq_kvlen, sparse_mode=sparse_mode, |
| 345 | - gen_mask_parallel=gen_mask_parallel, sync=sync | 366 | + gen_mask_parallel=gen_mask_parallel, sync=sync, softmax_layout=softmax_layout |
| 346 | ) | 367 | ) |
| 347 | 368 | ||
| 348 | - return dq, dk, dv, dpse | 369 | + return dq, dk, dv, dpse if pse else None, dsink if sink else None |
| 370 | + | ||
| 349 | 371 | ||
| 350 | def npu_fusion_attention_grad_fake( | 372 | def npu_fusion_attention_grad_fake( |
| 351 | query: torch.Tensor, | 373 | query: torch.Tensor, |
| @@ -375,13 +397,16 @@ def npu_fusion_attention_grad_fake( | |||
| 375 | actual_seq_kvlen: Optional[torch.Tensor] = None, | 397 | actual_seq_kvlen: Optional[torch.Tensor] = None, |
| 376 | sparse_mode: int = 0, | 398 | sparse_mode: int = 0, |
| 377 | gen_mask_parallel: bool = True, | 399 | gen_mask_parallel: bool = True, |
| 378 | - sync: bool = False | 400 | + sync: bool = False, |
| 379 | - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | 401 | + softmax_layout: str = "", |
| 402 | + sink: Optional[torch.Tensor] = None | ||
| 403 | + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| 380 | dq = torch.empty_like(query, dtype=query.dtype, device=query.device).contiguous() | 404 | dq = torch.empty_like(query, dtype=query.dtype, device=query.device).contiguous() |
| 381 | dk = torch.empty_like(key, dtype=query.dtype, device=query.device).contiguous() | 405 | dk = torch.empty_like(key, dtype=query.dtype, device=query.device).contiguous() |
| 382 | dv = torch.empty_like(value, dtype=query.dtype, device=query.device).contiguous() | 406 | dv = torch.empty_like(value, dtype=query.dtype, device=query.device).contiguous() |
| 383 | dpse = torch.empty([0], dtype=query.dtype, device=query.device).contiguous() | 407 | dpse = torch.empty([0], dtype=query.dtype, device=query.device).contiguous() |
| 384 | - return dq, dk, dv, dpse if pse else None | 408 | + dsink = torch.empty_like(sink) if sink is not None else None |
| 409 | + return dq, dk, dv, dpse if pse else None, dsink | ||
| 385 | 410 | ||
| 386 | direct_register_custom_op( | 411 | direct_register_custom_op( |
| 387 | op_name="npu_fusion_attention_grad", | 412 | op_name="npu_fusion_attention_grad", |
| @@ -392,19 +417,22 @@ direct_register_custom_op( | |||
| 392 | dispatch_key='PrivateUse1' | 417 | dispatch_key='PrivateUse1' |
| 393 | ) | 418 | ) |
| 394 | 419 | ||
| 420 | + | ||
| 395 | class InductorNpuAttentionFunction(torch.autograd.Function): | 421 | class InductorNpuAttentionFunction(torch.autograd.Function): |
| 396 | 422 | ||
| 397 | def forward(ctx, query, key, value, head_num, input_layout, pse=None, padding_mask=None, atten_mask=None, scale=1.0, | 423 | def forward(ctx, query, key, value, head_num, input_layout, pse=None, padding_mask=None, atten_mask=None, scale=1.0, |
| 398 | keep_prob=1.0, pre_tockens=2147483647, next_tockens=2147483647, inner_precise=0, prefix=None, | 424 | keep_prob=1.0, pre_tockens=2147483647, next_tockens=2147483647, inner_precise=0, prefix=None, |
| 399 | - actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, gen_mask_parallel=True, sync=False): | 425 | + actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, gen_mask_parallel=True, sync=False, |
| 426 | + softmax_layout="", sink=None): | ||
| 400 | attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels = torch.ops.inductor_npu.npu_fusion_attention( | 427 | attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels = torch.ops.inductor_npu.npu_fusion_attention( |
| 401 | query, key, value, head_num, input_layout, pse=pse, padding_mask=padding_mask, atten_mask=atten_mask, | 428 | query, key, value, head_num, input_layout, pse=pse, padding_mask=padding_mask, atten_mask=atten_mask, |
| 402 | scale=scale, keep_prob=keep_prob, pre_tockens=pre_tockens, next_tockens=next_tockens, | 429 | scale=scale, keep_prob=keep_prob, pre_tockens=pre_tockens, next_tockens=next_tockens, |
| 403 | inner_precise=inner_precise, prefix=prefix, actual_seq_qlen=actual_seq_qlen, | 430 | inner_precise=inner_precise, prefix=prefix, actual_seq_qlen=actual_seq_qlen, |
| 404 | - actual_seq_kvlen=actual_seq_kvlen, sparse_mode=sparse_mode, gen_mask_parallel=gen_mask_parallel, sync=sync | 431 | + actual_seq_kvlen=actual_seq_kvlen, sparse_mode=sparse_mode, gen_mask_parallel=gen_mask_parallel, sync=sync, |
| 432 | + softmax_layout=softmax_layout, sink=sink | ||
| 405 | ) | 433 | ) |
| 406 | ctx.save_for_backward(query, key, value, pse, padding_mask, atten_mask, actual_seq_qlen, actual_seq_kvlen,\ | 434 | ctx.save_for_backward(query, key, value, pse, padding_mask, atten_mask, actual_seq_qlen, actual_seq_kvlen,\ |
| 407 | - softmax_max, softmax_sum, softmax_out, attention_score, seed, offset, numels) | 435 | + softmax_max, softmax_sum, softmax_out, attention_score, seed, offset, numels, sink) |
| 408 | ctx.head_num = head_num | 436 | ctx.head_num = head_num |
| 409 | ctx.input_layout = input_layout | 437 | ctx.input_layout = input_layout |
| 410 | ctx.scale = scale | 438 | ctx.scale = scale |
| @@ -413,39 +441,41 @@ class InductorNpuAttentionFunction(torch.autograd.Function): | |||
| 413 | ctx.next_tockens = next_tockens | 441 | ctx.next_tockens = next_tockens |
| 414 | ctx.inner_precise = inner_precise | 442 | ctx.inner_precise = inner_precise |
| 415 | ctx.prefix = prefix | 443 | ctx.prefix = prefix |
| 416 | - # ctx.actual_seq_qlen = actual_seq_qlen | ||
| 417 | - # ctx.actual_seq_kvlen = actual_seq_kvlen | ||
| 418 | ctx.sparse_mode = sparse_mode | 444 | ctx.sparse_mode = sparse_mode |
| 419 | ctx.gen_mask_parallel = gen_mask_parallel | 445 | ctx.gen_mask_parallel = gen_mask_parallel |
| 420 | ctx.sync = sync | 446 | ctx.sync = sync |
| 447 | + ctx.softmax_layout = softmax_layout | ||
| 421 | 448 | ||
| 422 | - return attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels | 449 | + attn_ret = (attention_score, softmax_max, softmax_sum, softmax_out, ) |
| 450 | + dropout_ret = (seed, offset, numels) | ||
| 451 | + | ||
| 452 | + return *attn_ret, *dropout_ret | ||
| 423 | 453 | ||
| 424 | 454 | ||
| 425 | def backward(ctx, grad_attention_score, grad_softmax_max, grad_softmax_sum, grad_softmax_out, grad_seed, grad_offset, grad_numels): | 455 | def backward(ctx, grad_attention_score, grad_softmax_max, grad_softmax_sum, grad_softmax_out, grad_seed, grad_offset, grad_numels): |
| 426 | query, key, value, pse, padding_mask, atten_mask, actual_seq_qlen, actual_seq_kvlen, \ | 456 | query, key, value, pse, padding_mask, atten_mask, actual_seq_qlen, actual_seq_kvlen, \ |
| 427 | - softmax_max, softmax_sum, softmax_out, attention_score, seed, offset, numels = ctx.saved_tensors | 457 | + softmax_max, softmax_sum, softmax_out, attention_score, seed, offset, numels, sink = ctx.saved_tensors |
| 428 | - grad_query, grad_key, grad_value, grad_pse = torch.ops.inductor_npu.npu_fusion_attention_grad( | 458 | + grad_query, grad_key, grad_value, grad_pse, grad_sink = torch.ops.inductor_npu.npu_fusion_attention_grad( |
| 429 | query, key, value, grad_attention_score, ctx.head_num, ctx.input_layout, pse=pse, padding_mask=padding_mask, | 459 | query, key, value, grad_attention_score, ctx.head_num, ctx.input_layout, pse=pse, padding_mask=padding_mask, |
| 430 | atten_mask=atten_mask, softmax_max=softmax_max, softmax_sum=softmax_sum, softmax_in=softmax_out, attention_in=attention_score, | 460 | atten_mask=atten_mask, softmax_max=softmax_max, softmax_sum=softmax_sum, softmax_in=softmax_out, attention_in=attention_score, |
| 431 | scale_value=ctx.scale, keep_prob=ctx.keep_prob, pre_tockens=ctx.pre_tockens, next_tockens=ctx.next_tockens, | 461 | scale_value=ctx.scale, keep_prob=ctx.keep_prob, pre_tockens=ctx.pre_tockens, next_tockens=ctx.next_tockens, |
| 432 | inner_precise=ctx.inner_precise, seed=seed, offset=offset, numels=numels, prefix=None, | 462 | inner_precise=ctx.inner_precise, seed=seed, offset=offset, numels=numels, prefix=None, |
| 433 | actual_seq_qlen=actual_seq_qlen, actual_seq_kvlen=actual_seq_kvlen, sparse_mode=ctx.sparse_mode, | 463 | actual_seq_qlen=actual_seq_qlen, actual_seq_kvlen=actual_seq_kvlen, sparse_mode=ctx.sparse_mode, |
| 434 | - gen_mask_parallel=ctx.gen_mask_parallel, sync=ctx.sync | 464 | + gen_mask_parallel=ctx.gen_mask_parallel, sync=ctx.sync, softmax_layout=ctx.softmax_layout, sink=sink |
| 435 | ) | 465 | ) |
| 436 | - return ( | 466 | + return (grad_query, grad_key, grad_value, None, None, grad_pse, ) + (None, ) * 14 + (grad_sink, ) |
| 437 | - grad_query, grad_key, grad_value, None, None, grad_pse, None, None, None, None, None, None, None, None, None, | 467 | + |
| 438 | - None, None, None, None, None, None, None, None, None, None, None) | ||
| 439 | 468 | ||
| 440 | def inductor_npu_fusion_attention(query, key, value, head_num, input_layout, pse=None, padding_mask=None, | 469 | def inductor_npu_fusion_attention(query, key, value, head_num, input_layout, pse=None, padding_mask=None, |
| 441 | atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, | 470 | atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, |
| 442 | next_tockens=2147483647, | 471 | next_tockens=2147483647, |
| 443 | inner_precise=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, | 472 | inner_precise=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, |
| 444 | - gen_mask_parallel=True, sync=False): | 473 | + gen_mask_parallel=True, sync=False, softmax_layout="", sink=None): |
| 445 | return InductorNpuAttentionFunction.apply(query, key, value, head_num, input_layout, pse, padding_mask, | 474 | return InductorNpuAttentionFunction.apply(query, key, value, head_num, input_layout, pse, padding_mask, |
| 446 | atten_mask, scale, keep_prob, pre_tockens, next_tockens, | 475 | atten_mask, scale, keep_prob, pre_tockens, next_tockens, |
| 447 | inner_precise, prefix, actual_seq_qlen, actual_seq_kvlen, sparse_mode, | 476 | inner_precise, prefix, actual_seq_qlen, actual_seq_kvlen, sparse_mode, |
| 448 | - gen_mask_parallel, sync) | 477 | + gen_mask_parallel, sync, softmax_layout, sink) |
| 478 | + | ||
| 449 | 479 | ||
| 450 | def apply_inductor_npu_attention_patch(): | 480 | def apply_inductor_npu_attention_patch(): |
| 451 | - torch.ops.npu.npu_fusion_attention = inductor_npu_fusion_attention | 481 | + torch.ops.npu.npu_fusion_attention = inductor_npu_fusion_attention |
此条代码评论区间+23至+27
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。