From 5436f8f3c14a2f76f23dc3265816a95376f8d5d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=97=AB=E9=B9=8F=E5=85=A8?= Date: Thu, 11 Sep 2025 19:38:23 +0800 Subject: [PATCH] support npu_fused_infer_attention_score.default for auto_dispatch --- torch_npu/npu/graphs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_npu/npu/graphs.py b/torch_npu/npu/graphs.py index b7b312db64..5496af8151 100644 --- a/torch_npu/npu/graphs.py +++ b/torch_npu/npu/graphs.py @@ -146,7 +146,7 @@ class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode): return _GraphDispatchRecord(event=event, handle=handle, kwargs=kwargs_ref, args=tuple(args_ref), op_cache_entry=func) def __torch_dispatch__(self, func, types, args=(), kwargs=None): - if func.__name__ == "npu_fused_infer_attention_score": + if func.__name__ in ["npu_fused_infer_attention_score", "npu_fused_infer_attention_score.default"]: func_out = torch_npu.npu_fused_infer_attention_score.out self.update_schema(str(func_out.__name__), str(func_out._schema)) stream = torch_npu.npu.current_stream() -- Gitee