已合并
move_npu_patch_meta_to_module_level #39681
cuiduo创建于 6月30日
move_npu_patch_meta_to_module_level #39681
已合并
共 2 个文件变更+3-4
| @@ -1,7 +1,6 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | 2 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests |
| 3 | from testutils import TestUtils | 3 | from testutils import TestUtils |
| 4 | -import torch_npu | ||
| 5 | 4 | ||
| 6 | 5 | ||
| 7 | class TestLog(TestUtils): | 6 | class TestLog(TestUtils): |
| @@ -16,6 +16,9 @@ from .utils import patch_has_triton, patch_has_triton_tma, patch_is_gpu, get_cur | |||
| 16 | # All backends need npu/cpu/mps device_op_overrides. | 16 | # All backends need npu/cpu/mps device_op_overrides. |
| 17 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system | 17 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system |
| 18 | from .shape_handling import NPUShapeHandling, patch_shape_handling | 18 | from .shape_handling import NPUShapeHandling, patch_shape_handling |
| 19 | +from ._npu_meta_registration import npu_patch_meta | ||
| 20 | + | ||
| 21 | +npu_patch_meta() | ||
| 19 | register_device_op_overrides_npu() | 22 | register_device_op_overrides_npu() |
| 20 | patch_has_triton() | 23 | patch_has_triton() |
| 21 | patch_is_gpu() | 24 | patch_is_gpu() |
| @@ -137,9 +140,6 @@ def _load_triton_backend(): | |||
| 137 | from .graph import patch_count_bytes, patch_run_node | 140 | from .graph import patch_count_bytes, patch_run_node |
| 138 | from .autotune_process import patch_tuning_process, patch_tuning_process_pool | 141 | from .autotune_process import patch_tuning_process, patch_tuning_process_pool |
| 139 | flex_attention._validate_device = _validate_device | 142 | flex_attention._validate_device = _validate_device |
| 140 | - from ._npu_meta_registration import npu_patch_meta | ||
| 141 | - | ||
| 142 | - npu_patch_meta() | ||
| 143 | 143 | ||
| 144 | def _inductor_register_backend_for_device(): | 144 | def _inductor_register_backend_for_device(): |
| 145 | from .codegen.cpp_wrapper import CppWrapperNpu | 145 | from .codegen.cpp_wrapper import CppWrapperNpu |