已合并
[fix] test_npu_graph_attention_function #28990
zhangqiongwen创建于 1月4日
[fix] test_npu_graph_attention_function #28990
已合并
共 2 个文件变更+8-8
| @@ -1,5 +1,4 @@ | |||
| 1 | import functools | 1 | import functools |
| 2 | -from unittest import skip | ||
| 3 | import sympy | 2 | import sympy |
| 4 | import torch | 3 | import torch |
| 5 | import torch.nn.functional as F | 4 | import torch.nn.functional as F |
| @@ -11,7 +10,6 @@ from torch_npu._inductor.npu_fusion_attention_graph import NpuGraphAttentionFunc | |||
| 11 | 10 | ||
| 12 | 11 | ||
| 13 | class TestNpuFusionAttentionGraph(TestCase): | 12 | class TestNpuFusionAttentionGraph(TestCase): |
| 14 | - | ||
| 15 | def test_npu_graph_attention_function(self): | 13 | def test_npu_graph_attention_function(self): |
| 16 | query = torch.randn(2, 4, 8, 16, device='npu', requires_grad=True) | 14 | query = torch.randn(2, 4, 8, 16, device='npu', requires_grad=True) |
| 17 | key = torch.randn(2, 4, 8, 16, device='npu') | 15 | key = torch.randn(2, 4, 8, 16, device='npu') |
| @@ -13,9 +13,9 @@ npu_lib = Library("npu_graph", "IMPL", "PrivateUse1") | |||
| 13 | meta_lib = Library("npu_graph", "IMPL", "Meta") | 13 | meta_lib = Library("npu_graph", "IMPL", "Meta") |
| 14 | 14 | ||
| 15 | npu_def.define( | 15 | npu_def.define( |
| 16 | - "npu_fa(Tensor query, Tensor key, Tensor value, int head_num, str input_layout, Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, float scale=1., float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, int[]? prefix=None, int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, bool sync=False) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)") | 16 | + "npu_fa(Tensor query, Tensor key, Tensor value, int head_num, str input_layout, Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, float scale=1., float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, int[]? prefix=None, int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, bool sync=False, str softmax_layout=\"\", Tensor? sink=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)") |
| 17 | npu_def.define( | 17 | npu_def.define( |
| 18 | - "npu_fa_backward(Tensor query, Tensor key, Tensor value, Tensor dy, int head_num, str input_layout, *, Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, Tensor? softmax_max=None, Tensor? softmax_sum=None, Tensor? softmax_in=None, Tensor? attention_in=None, float scale_value=1., float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, Tensor? seed=None, Tensor? offset=None, Tensor? numels=None, int[]? prefix=None, int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, bool sync=False) -> (Tensor, Tensor, Tensor, Tensor)") | 18 | + "npu_fa_backward(Tensor query, Tensor key, Tensor value, Tensor dy, int head_num, str input_layout, *, Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, Tensor? softmax_max=None, Tensor? softmax_sum=None, Tensor? softmax_in=None, Tensor? attention_in=None, float scale_value=1., float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, Tensor? seed=None, Tensor? offset=None, Tensor? numels=None, int[]? prefix=None, int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, bool sync=False, str softmax_layout=\"\", Tensor? sink=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor)") |
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | 21 | ||
| @@ -47,7 +47,7 @@ def npu_fa_backward(*args, **kwargs): | |||
| 47 | def npu_fa(query, key, value, head_num, input_layout, pse=None, padding_mask=None, | 47 | def npu_fa(query, key, value, head_num, input_layout, pse=None, padding_mask=None, |
| 48 | atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, next_tockens=2147483647, | 48 | atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, next_tockens=2147483647, |
| 49 | inner_precise=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, | 49 | inner_precise=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, |
| 50 | - gen_mask_parallel=True, sync=False): | 50 | + gen_mask_parallel=True, sync=False, softmax_layout="", sink=None): |
| 51 | B = query.size(0) | 51 | B = query.size(0) |
| 52 | N = head_num | 52 | N = head_num |
| 53 | S1 = query.size(2) | 53 | S1 = query.size(2) |
| @@ -81,12 +81,14 @@ def npu_fa_backward(query, key, value, dy, head_num, input_layout, *, pse=None, | |||
| 81 | softmax_max=None, softmax_sum=None, softmax_in=None, attention_in=None, scale_value=1.0, | 81 | softmax_max=None, softmax_sum=None, softmax_in=None, attention_in=None, scale_value=1.0, |
| 82 | keep_prob=1.0, pre_tockens=2147483647, next_tockens=2147483647, inner_precise=0, seed=0, offset=0, | 82 | keep_prob=1.0, pre_tockens=2147483647, next_tockens=2147483647, inner_precise=0, seed=0, offset=0, |
| 83 | numels=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, | 83 | numels=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, |
| 84 | - gen_mask_parallel=True, sync=False): | 84 | + gen_mask_parallel=True, sync=False, softmax_layout="", sink=None): |
| 85 | dq = torch.empty_like(query, dtype=query.dtype, device='meta').contiguous() | 85 | dq = torch.empty_like(query, dtype=query.dtype, device='meta').contiguous() |
| 86 | dk = torch.empty_like(key, dtype=query.dtype, device='meta').contiguous() | 86 | dk = torch.empty_like(key, dtype=query.dtype, device='meta').contiguous() |
| 87 | dv = torch.empty_like(value, dtype=query.dtype, device='meta').contiguous() | 87 | dv = torch.empty_like(value, dtype=query.dtype, device='meta').contiguous() |
| 88 | dpse = torch.empty([0], dtype=query.dtype, device='meta').contiguous() | 88 | dpse = torch.empty([0], dtype=query.dtype, device='meta').contiguous() |
| 89 | - return (torch.empty_like(dq), torch.empty_like(dk), torch.empty_like(dv), torch.empty_like(dpse) if pse else None) | 89 | + dsink = torch.empty([], device='meta') if sink is None else torch.empty_like(sink, dtype=sink.dtype, device='meta').contiguous() |
| 90 | + return (torch.empty_like(dq), torch.empty_like(dk), torch.empty_like(dv), | ||
| 91 | + torch.empty_like(dpse) if pse else None, dsink) | ||
| 90 | 92 | ||
| 91 | 93 | ||
| 92 | class NpuGraphAttentionFunction(Function): | 94 | class NpuGraphAttentionFunction(Function): |
| @@ -127,7 +129,7 @@ class NpuGraphAttentionFunction(Function): | |||
| 127 | query, key, value, pse, padding_mask, atten_mask, result1, result2, result3, result0, result4, result5, result6 = ctx.saved_tensors | 129 | query, key, value, pse, padding_mask, atten_mask, result1, result2, result3, result0, result4, result5, result6 = ctx.saved_tensors |
| 128 | # 反向传播逻辑 | 130 | # 反向传播逻辑 |
| 129 | # 这里假设有一个实现反向传播的函数 `npu_fusion_attention_backward` | 131 | # 这里假设有一个实现反向传播的函数 `npu_fusion_attention_backward` |
| 130 | - grad_query, grad_key, grad_value, grad_pse = torch.ops.npu_graph.npu_fa_backward( | 132 | + grad_query, grad_key, grad_value, grad_pse, grad_sink = torch.ops.npu_graph.npu_fa_backward( |
| 131 | query, key, value, grad_result0, ctx.head_num, ctx.input_layout, pse=pse, padding_mask=padding_mask, | 133 | query, key, value, grad_result0, ctx.head_num, ctx.input_layout, pse=pse, padding_mask=padding_mask, |
| 132 | atten_mask=atten_mask, softmax_max=result1, softmax_sum=result2, softmax_in=result3, attention_in=result0, | 134 | atten_mask=atten_mask, softmax_max=result1, softmax_sum=result2, softmax_in=result3, attention_in=result0, |
| 133 | scale_value=ctx.scale, keep_prob=ctx.keep_prob, pre_tockens=ctx.pre_tockens, next_tockens=ctx.next_tockens, | 135 | scale_value=ctx.scale, keep_prob=ctx.keep_prob, pre_tockens=ctx.pre_tockens, next_tockens=ctx.next_tockens, |
参数传递错误: 在 NpuGraphAttentionFunction.backward 中,调用 torch.ops.npu_graph.npu_fa_backward 时,seed、offset、numels 参数传递的是 result4、result5、result6(Tensor 对象),但根据新的接口定义,这些参数应为 int 类型。这会导致类型不匹配,可能引发运行时错误或未定义行为。
问题类型: 参数传递错误 文件路径:
torch_npu/_inductor/npu_fusion_attention_graph.py行号: 132 问题代码:修改建议:
此评论由代码审查工具自动生成