已合并
Fixed failed testcases #26361
haiyan8创建于 2025年11月7日
Fixed failed testcases #26361
已合并
共 3 个文件变更+16-7
| @@ -42,7 +42,7 @@ class TestIFAAclgraphUpdate(TestCase): | |||
| 42 | with torch.npu.graph(g): | 42 | with torch.npu.graph(g): |
| 43 | stream = torch.npu.current_stream() | 43 | stream = torch.npu.current_stream() |
| 44 | output = torch.empty(1, 32, 1, 128, dtype=torch.float16, device="npu") | 44 | output = torch.empty(1, 32, 1, 128, dtype=torch.float16, device="npu") |
| 45 | - softmax_lse = torch.empty(1, dtype=torch.float16, device="npu") | 45 | + softmax_lse = torch.empty_like(res_src[1], dtype=torch.float16, device="npu") |
| 46 | event.wait(stream) | 46 | event.wait(stream) |
| 47 | event.reset(stream) | 47 | event.reset(stream) |
| 48 | torch.npu.graph_task_group_begin(stream) | 48 | torch.npu.graph_task_group_begin(stream) |
| @@ -87,7 +87,7 @@ class TestIFAAclgraphUpdate(TestCase): | |||
| 87 | 87 | ||
| 88 | with torch.npu.graph(g, auto_dispatch_capture=True): | 88 | with torch.npu.graph(g, auto_dispatch_capture=True): |
| 89 | output = torch.empty(1, 32, 1, 128, dtype=torch.float16, device="npu") | 89 | output = torch.empty(1, 32, 1, 128, dtype=torch.float16, device="npu") |
| 90 | - softmax_lse = torch.empty(1, dtype=torch.float16, device="npu") | 90 | + softmax_lse = torch.empty_like(res_src[1], dtype=torch.float16, device="npu") |
| 91 | torch_npu.npu_fused_infer_attention_score.out( | 91 | torch_npu.npu_fused_infer_attention_score.out( |
| 92 | query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace, | 92 | query, key, value, num_heads=32, input_layout="BNSD", scale=scale, pre_tokens=65535, workspace=workspace, |
| 93 | next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse]) | 93 | next_tokens=65535, softmax_lse_flag=False, actual_seq_lengths=length, out=[output, softmax_lse]) |
| @@ -128,6 +128,7 @@ class TestIFAAclgraphUpdate(TestCase): | |||
| 128 | self.assertEqual(softmax_lse.cpu(), res_src[1].cpu()) | 128 | self.assertEqual(softmax_lse.cpu(), res_src[1].cpu()) |
| 129 | 129 | ||
| 130 | 130 | ||
| 131 | + | ||
| 131 | def test_npu_fused_infer_attention_score_v2(self): | 132 | def test_npu_fused_infer_attention_score_v2(self): |
| 132 | torch.npu.set_device(0) | 133 | torch.npu.set_device(0) |
| 133 | length = [29] | 134 | length = [29] |
| @@ -175,6 +176,7 @@ class TestIFAAclgraphUpdate(TestCase): | |||
| 175 | self.assertEqual(softmax_lse.cpu(), res_src[1].cpu()) | 176 | self.assertEqual(softmax_lse.cpu(), res_src[1].cpu()) |
| 176 | 177 | ||
| 177 | 178 | ||
| 179 | + | ||
| 178 | def test_npugraph_debug_dump(self): | 180 | def test_npugraph_debug_dump(self): |
| 179 | N, D_in, H, D_out = 640, 4096, 2048, 1024 | 181 | N, D_in, H, D_out = 640, 4096, 2048, 1024 |
| 180 | model = torch.nn.Sequential(torch.nn.Linear(D_in, H), | 182 | model = torch.nn.Sequential(torch.nn.Linear(D_in, H), |
| @@ -340,7 +342,7 @@ class TestPAAclgraphUpdate(TestCase): | |||
| 340 | graph.update(cpu_update_input=[{"context_lens": params.context_lens}]) | 342 | graph.update(cpu_update_input=[{"context_lens": params.context_lens}]) |
| 341 | graph.replay() | 343 | graph.replay() |
| 342 | torch.npu.synchronize() | 344 | torch.npu.synchronize() |
| 343 | - self.assertRtolEqual(output, golden_output) | 345 | + self.assertRtolEqual(output, golden_output, prec16=0.01) |
| 344 | 346 | ||
| 345 | params_new, golden_output = self.preprocess() | 347 | params_new, golden_output = self.preprocess() |
| 346 | params.query.copy_(params_new.query) | 348 | params.query.copy_(params_new.query) |
| @@ -350,7 +352,7 @@ class TestPAAclgraphUpdate(TestCase): | |||
| 350 | graph.update(cpu_update_input=[{"context_lens": params_new.context_lens}]) | 352 | graph.update(cpu_update_input=[{"context_lens": params_new.context_lens}]) |
| 351 | graph.replay() | 353 | graph.replay() |
| 352 | torch.npu.synchronize() | 354 | torch.npu.synchronize() |
| 353 | - self.assertRtolEqual(output, golden_output) | 355 | + self.assertRtolEqual(output, golden_output, prec16=0.01) |
| 354 | 356 | ||
| 355 | if __name__ == "__main__": | 357 | if __name__ == "__main__": |
| 356 | run_tests() | 358 | run_tests() |
| @@ -4,8 +4,11 @@ import pytorch_test_common | |||
| 4 | 4 | ||
| 5 | import torch | 5 | import torch |
| 6 | from pytorch_test_common import skipIfNoNPU | 6 | from pytorch_test_common import skipIfNoNPU |
| 7 | -from torch.onnx import verification | 7 | +from torch.onnx._internal.torchscript_exporter import verification |
| 8 | -from torch.onnx._globals import GLOBALS | 8 | +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS |
| 9 | +from torch.onnx._internal.torchscript_exporter.utils import ( | ||
| 10 | + _trigger_symbolic_function_registration, | ||
| 11 | +) | ||
| 9 | from torch.testing._internal import common_utils | 12 | from torch.testing._internal import common_utils |
| 10 | import torch_npu | 13 | import torch_npu |
| 11 | import torch_npu.testing | 14 | import torch_npu.testing |
| @@ -22,6 +25,7 @@ def _jit_graph_to_onnx_model(graph, operator_export_type, opset_version): | |||
| 22 | """ | 25 | """ |
| 23 | 26 | ||
| 24 | GLOBALS.export_onnx_opset_version = opset_version | 27 | GLOBALS.export_onnx_opset_version = opset_version |
| 28 | + _trigger_symbolic_function_registration() | ||
| 25 | graph = torch.onnx.utils._optimize_graph( | 29 | graph = torch.onnx.utils._optimize_graph( |
| 26 | graph, operator_export_type, params_dict={} | 30 | graph, operator_export_type, params_dict={} |
| 27 | ) | 31 | ) |
| @@ -157,7 +157,10 @@ class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode): | |||
| 157 | # apply tensor | 157 | # apply tensor |
| 158 | workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(*args, **kwargs) | 158 | workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(*args, **kwargs) |
| 159 | output = torch.empty_like(args[0]) | 159 | output = torch.empty_like(args[0]) |
| 160 | - softmax_lse = torch.empty(1, dtype=args[0].dtype, device=args[0].device) | 160 | + if "softmax_lse_flag" in kwargs.keys(): |
| 161 | + softmax_lse = torch.empty(1, dtype=args[0].dtype, device=args[0].device) | ||
| 162 | + else: | ||
| 163 | + softmax_lse = torch.empty(0, dtype=args[0].dtype, device=args[0].device) | ||
| 161 | kwargs["workspace"] = workspace | 164 | kwargs["workspace"] = workspace |
| 162 | kwargs["out"] = [output, softmax_lse] | 165 | kwargs["out"] = [output, softmax_lse] |
| 163 | # begin graph task | 166 | # begin graph task |