已合并
[fix] test_npu_graph_attention_function #28990
zhangqiongwen创建于 1月4日
[fix] test_npu_graph_attention_function #28990
已合并
zhangqiongwen创建于 1月4日
2 个文件变更+8-8
Mtest/_inductor/test_npu_fusion_attention_graph.py+0-2
@@ -1,5 +1,4 @@
1import functools1import functools
2-from unittest import skip
3import sympy2import sympy
4import torch3import torch
5import torch.nn.functional as F4import torch.nn.functional as F
@@ -11,7 +10,6 @@ from torch_npu._inductor.npu_fusion_attention_graph import NpuGraphAttentionFunc
11 10 
12 11 
13class TestNpuFusionAttentionGraph(TestCase):12class TestNpuFusionAttentionGraph(TestCase):
14- @skip("skip for core dump")
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')
Mtorch_npu/_inductor/npu_fusion_attention_graph.py+8-6
@@ -13,9 +13,9 @@ npu_lib = Library("npu_graph", "IMPL", "PrivateUse1")
13meta_lib = Library("npu_graph", "IMPL", "Meta")13meta_lib = Library("npu_graph", "IMPL", "Meta")
14 14 
15npu_def.define(15npu_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)")
17npu_def.define(17npu_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@impl(npu_lib, "npu_fa")21@impl(npu_lib, "npu_fa")
@@ -47,7 +47,7 @@ def npu_fa_backward(*args, **kwargs):
47def npu_fa(query, key, value, head_num, input_layout, pse=None, padding_mask=None,47def 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_num52 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 
92class NpuGraphAttentionFunction(Function):94class 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_tensors129 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(
AtlasAccount
AtlasAccountAtlasAccount1月4日

参数传递错误: 在 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 问题代码:

grad_query, grad_key, grad_value, grad_pse, grad_sink = torch.ops.npu_graph.npu_fa_backward(
    query, key, value, grad_result0, ctx.head_num, ctx.input_layout, pse=pse, padding_mask=padding_mask,
    atten_mask=atten_mask, softmax_max=result1, softmax_sum=result2, softmax_in=result3, attention_in=result0,
    scale_value=ctx.scale, keep_prob=ctx.keep_prob, pre_tockens=ctx.pre_tockens, next_tockens=ctx.next_tockens,
    inner_precise=ctx.inner_precise, seed=result4, offset=result5, numels=result6, prefix=ctx.prefix,
    actual_seq_qlen=ctx.actual_seq_qlen, actual_seq_kvlen=ctx.actual_seq_kvlen, sparse_mode=ctx.sparse_mode,
    gen_mask_parallel=ctx.gen_mask_parallel, sync=ctx.sync
)

修改建议:

根据新的接口定义,seed、offset、numels 应为 int 类型。需要检查 result4、result5、result6 的结构(它们在前向传播中是通过 torch.tensor([seed], requires_grad=False) 等方式创建的),并提取其整数值。例如,使用 result4.item() 或 result4[0].item() 来获取 int 值。同时,需要确保这些值在合理的范围内(例如,seed 应为非负整数)。

此评论由代码审查工具自动生成

likedislike
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,