perf: defer dynamo/inductor import for v2.7.1 (#2788)-pytorch-AtomGit

Msetup.py+5-0

	@@ -778,5 +778,10 @@ setup(
778	'torch.backends': [	778	'torch.backends': [
		H H htchu21 天前新增测试主要是 CPU tensor 与backend="eager"，未真正覆盖 NPU 的 inductor、npugraphs、torch.export 初始化链请补充真实 NPU 编译和执行冒烟测试，验证延迟注册后 DeviceInterface、Inductor backend和图模式能力均可用黄桂军 20 天前评论： 20 天前评论：
779	'torch_npu = torch_npu:_autoload',	779	'torch_npu = torch_npu:_autoload',
780	],	780	],
		781	+ 'torch_dynamo_backends': [
		782	+ 'npu = torch_npu.dynamo:_npu_backend_entrypoint',
		783	+ 'npugraph_ex = torch_npu.dynamo:_npugraph_ex_backend_entrypoint',
		784	+ 'npugraphs = torch_npu.dynamo:_npugraphs_backend_entrypoint',
		785	+ ],
781	}	786	}
782	)	787	)

Mtest/_inductor/test_current_device.py+2-1

	@@ -1,4 +1,5 @@
1	import torch	1	import torch
		2	+import torch._dynamo.testing
2	from torch.testing._internal.common_utils import (	3	from torch.testing._internal.common_utils import (
3	run_tests,	4	run_tests,
4	instantiate_parametrized_tests,	5	instantiate_parametrized_tests,
	@@ -30,4 +31,4 @@ instantiate_parametrized_tests(TestCurrentDevice)
30		31
31		32
32	if __name__ == "__main__":	33	if __name__ == "__main__":
33	- run_tests()	34	+ run_tests()

Mtest/_inductor/test_mlir_enable.py+4-1

	@@ -26,6 +26,9 @@ class TestAdd(TestUtils):
26	@parametrize('shape', TestUtils._pointwise_demo_shapes)	26	@parametrize('shape', TestUtils._pointwise_demo_shapes)
27	@parametrize('dtype', ['float32', 'int64'])	27	@parametrize('dtype', ['float32', 'int64'])
28	def test_config_environ_cases(self, shape, dtype):	28	def test_config_environ_cases(self, shape, dtype):
		29	+ # torch._inductor.config is an internal module. Initialize the NPU
		30	+ # config entries explicitly before testing direct config assignment.
		31	+ torch._inductor.config.get_config_copy()
29	torch._inductor.config.npu_backend = "mlir"	32	torch._inductor.config.npu_backend = "mlir"
30	x = self._generate_tensor(shape, dtype)	33	x = self._generate_tensor(shape, dtype)
31	y = self._generate_tensor(shape, dtype)	34	y = self._generate_tensor(shape, dtype)
	@@ -39,4 +42,4 @@ class TestAdd(TestUtils):
39	instantiate_parametrized_tests(TestAdd)	42	instantiate_parametrized_tests(TestAdd)
40		43
41	if __name__ == "__main__":	44	if __name__ == "__main__":
42	- run_tests()	45	+ run_tests()

Atest/dynamo/test_compile_trigger.py+954-0

@@ -0,0 +1,954 @@
	1	+# Owner(s): ["module: dynamo"]
	2	+
	3	+import importlib.util
	4	+import os
	5	+import subprocess
	6	+import sys
	7	+import textwrap
	8	+import unittest
	9	+
	10	+
	11	+class TorchCompileTriggerTests(unittest.TestCase):
	12	+ def run_in_subprocess(self, code):
	13	+ env = os.environ.copy()
	14	+ env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
	15	+ result = subprocess.run(
	16	+ [sys.executable, "-c", textwrap.dedent(code)],
	17	+ capture_output=True,
	18	+ env=env,
	19	+ text=True,
	20	+ timeout=60,
	21	+ )
	22	+ self.assertEqual(
	23	+ result.returncode,
	24	+ 0,
	25	+ f"Subprocess failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}",
	26	+ )
	27	+
	28	+ # Verify import installs only the Dynamo post-import trigger.
	29	+ def test_import_installs_dynamo_post_import_trigger(self):
	30	+ self.run_in_subprocess(
	31	+ """
	32	+ import sys
	33	+ import torch
	34	+
	35	+ def loaded(prefix):
	36	+ return [
	37	+ name for name in sys.modules
	38	+ if name == prefix or name.startswith(prefix + ".")
	39	+ ]
	40	+
	41	+ assert not loaded("torch._dynamo")
	42	+ assert not loaded("torch._inductor")
	43	+ src_compile = torch.compile
	44	+ src_wrapper_init = torch._TorchCompileWrapper.__init__
	45	+
	46	+ import torch_npu
	47	+ from torch_npu.utils import _dynamo
	48	+
	49	+ assert not loaded("torch._dynamo")
	50	+ assert not loaded("torch._inductor")
	51	+ assert torch.compile is src_compile
	52	+ assert torch._TorchCompileWrapper.__init__ is src_wrapper_init
	53	+ assert any(
	54	+ isinstance(finder, _dynamo._DynamoPostImportFinder)
	55	+ for finder in sys.meta_path
	56	+ )
	57	+ assert not _dynamo._lazy_dynamo_setup.has_run
	58	+ """
	59	+ )
	60	+
	61	+ # Verify public compiler APIs work before the first compile.
	62	+ def test_public_compiler_entries_are_available_before_compile(self):
	63	+ self.run_in_subprocess(
	64	+ """
	65	+ import inspect
	66	+ import sys
	67	+ import torch
	68	+ import torch_npu
	69	+ from torch_npu.utils import _dynamo
	70	+
	71	+ marker = torch.compiler.npugraph_mark_step_begin
	72	+ assert marker.__name__ == "npugraph_mark_step_begin"
	73	+ assert str(inspect.signature(marker)) == "()"
	74	+ marker()
	75	+
	76	+ from torch_npu.npu._graph_tree_state import MarkStepBox
	77	+
	78	+ assert MarkStepBox.mark_step_counter == -1
	79	+ assert not _dynamo._lazy_dynamo_setup.has_run
	80	+ assert not _dynamo._lazy_inductor_setup.has_run
	81	+ assert not any(
	82	+ name == "torch._dynamo" or name.startswith("torch._dynamo.")
	83	+ for name in sys.modules
	84	+ )
	85	+ assert not any(
	86	+ name == "torch._inductor" or name.startswith("torch._inductor.")
	87	+ for name in sys.modules
	88	+ )
	89	+
	90	+ backends = torch.compiler.list_backends(exclude_tags=None)
	91	+ assert {"npu", "npugraph_ex", "npugraphs"}.issubset(backends)
	92	+ assert _dynamo._lazy_dynamo_setup.has_run
	93	+ assert not _dynamo._lazy_inductor_setup.has_run
	94	+ assert "torch_npu._inductor" not in sys.modules
	95	+ """
	96	+ )
	97	+
	98	+ # Verify wheel metadata exposes loadable NPU backends without compiler imports.
	99	+ def test_dynamo_backend_entrypoint_metadata_and_cold_load(self):
	100	+ self.run_in_subprocess(
	101	+ """
	102	+ import importlib.metadata
	103	+ import sys
	104	+
	105	+ expected = {
	106	+ "npu": "torch_npu.dynamo:_npu_backend_entrypoint",
	107	+ "npugraph_ex": (
	108	+ "torch_npu.dynamo:_npugraph_ex_backend_entrypoint"
	109	+ ),
	110	+ "npugraphs": "torch_npu.dynamo:_npugraphs_backend_entrypoint",
	111	+ }
	112	+ entry_points = {
	113	+ entry_point.name: entry_point
	114	+ for entry_point in importlib.metadata.entry_points(
	115	+ group="torch_dynamo_backends"
	116	+ )
	117	+ if entry_point.name in expected
	118	+ }
	119	+
	120	+ assert {
	121	+ name: entry_point.value
	122	+ for name, entry_point in entry_points.items()
	123	+ } == expected
	124	+ assert "torch_npu" not in sys.modules
	125	+ assert "torch._dynamo" not in sys.modules
	126	+ assert "torch._inductor" not in sys.modules
	127	+
	128	+ loaded = {
	129	+ name: entry_points[name].load()
	130	+ for name in expected
	131	+ }
	132	+ assert all(callable(backend) for backend in loaded.values())
	133	+ assert "torch._dynamo" not in sys.modules
	134	+ assert "torch._inductor" not in sys.modules
	135	+ """
	136	+ )
	137	+
	138	+ # Verify entry-point loading cannot register the NPU backend twice.
	139	+ def test_backend_entrypoint_can_import_torch_npu_without_duplicate_registration(self):
	140	+ self.run_in_subprocess(
	141	+ """
	142	+ import torch
	143	+ import torch._dynamo
	144	+ from torch._dynamo.backends import registry
	145	+
	146	+ # Model the state produced by setuptools entry-point discovery:
	147	+ # lookup_backend owns the final register_backend call, while
	148	+ # EntryPoint.load imports torch_npu and triggers its lazy setup.
	149	+ class NpuEntryPoint:
	150	+ module = "torch_npu.dynamo"
	151	+
	152	+ def load(self):
	153	+ from torch_npu.dynamo import _npu_backend_entrypoint
	154	+ return _npu_backend_entrypoint
	155	+
	156	+ registry._BACKENDS["npu"] = NpuEntryPoint()
	157	+ backend = registry.lookup_backend("npu")
	158	+
	159	+ assert backend.__name__ == "_npu_backend_entrypoint"
	160	+ assert registry._COMPILER_FNS["npu"] is backend
	161	+
	162	+ # A later setup retry must treat the loaded torch_npu entry point
	163	+ # as an already completed registration.
	164	+ from torch_npu.dynamo import _register_npu_backend
	165	+
	166	+ _register_npu_backend(backend, "npu")
	167	+ assert registry._COMPILER_FNS["npu"] is backend
	168	+ """
	169	+ )
	170	+
	171	+ # Verify backend registration is safe to retry after a partial failure.
	172	+ def test_backend_registration_retry_after_partial_failure(self):
	173	+ self.run_in_subprocess(
	174	+ """
	175	+ import torch
	176	+ import torch_npu
	177	+ import torch._dynamo
	178	+ import torch_npu.dynamo as npu_dynamo
	179	+ from torch._dynamo.backends import registry
	180	+
	181	+ # Start from an undiscovered state so the first registration is
	182	+ # committed before the simulated second-registration failure.
	183	+ for name in ("npu", "npugraph_ex"):
	184	+ registry._BACKENDS.pop(name, None)
	185	+ registry._COMPILER_FNS.pop(name, None)
	186	+
	187	+ original_register = npu_dynamo._register_npu_backend
	188	+ fail_npugraph_ex = True
	189	+
	190	+ def fail_second_registration(backend, name="npu"):
	191	+ if name == "npugraph_ex" and fail_npugraph_ex:
	192	+ raise RuntimeError("simulated npugraph_ex registration failure")
	193	+ return original_register(backend, name)
	194	+
	195	+ npu_dynamo._register_npu_backend = fail_second_registration
	196	+ try:
	197	+ try:
	198	+ npu_dynamo._register_backends()
	199	+ except RuntimeError as error:
	200	+ assert "simulated npugraph_ex" in str(error)
	201	+ else:
	202	+ raise AssertionError("the first registration should fail")
	203	+ finally:
	204	+ npu_dynamo._register_npu_backend = original_register
	205	+
	206	+ assert "npu" in registry._COMPILER_FNS
	207	+ assert "npugraph_ex" not in registry._COMPILER_FNS
	208	+
	209	+ # Retry: the completed npu registration is a no-op, while the
	210	+ # missing npugraph_ex registration is installed normally.
	211	+ npu_dynamo._register_backends()
	212	+ assert "npu" in registry._COMPILER_FNS
	213	+ assert "npugraph_ex" in registry._COMPILER_FNS
	214	+ """
	215	+ )
	216	+
	217	+ # Verify failed lazy setup retries through public backend entries.
	218	+ def test_public_lazy_setup_entries_retry_after_failure(self):
	219	+ self.run_in_subprocess(
	220	+ """
	221	+ import os
	222	+ import sys
	223	+
	224	+ import torch
	225	+ import torch_npu
	226	+ import torch_npu.dynamo as npu_dynamo
	227	+ from torch_npu.utils import _dynamo
	228	+
	229	+ original_patch = _dynamo.patch_dynamo_optimize
	230	+ original_get_backend = npu_dynamo._get_default_backend
	231	+ attempts = []
	232	+
	233	+ def fail_first_dynamo_setup():
	234	+ attempts.append("dynamo")
	235	+ if len(attempts) == 1:
	236	+ raise RuntimeError("simulated Dynamo setup failure")
	237	+ return original_patch()
	238	+
	239	+ _dynamo.patch_dynamo_optimize = fail_first_dynamo_setup
	240	+ npu_dynamo._get_default_backend = (
	241	+ lambda name: npu_dynamo._eager_npu_backend
	242	+ )
	243	+ graph_module = lambda x: x
	244	+ try:
	245	+ try:
	246	+ npu_dynamo._npu_backend_entrypoint(graph_module, [])
	247	+ except RuntimeError as error:
	248	+ assert "simulated Dynamo setup failure" in str(error)
	249	+ else:
	250	+ raise AssertionError("the first backend setup should fail")
	251	+
	252	+ assert not _dynamo._lazy_dynamo_setup.has_run
	253	+ result = npu_dynamo._npu_backend_entrypoint(graph_module, [])
	254	+ assert result is graph_module
	255	+ finally:
	256	+ _dynamo.patch_dynamo_optimize = original_patch
	257	+ npu_dynamo._get_default_backend = original_get_backend
	258	+
	259	+ assert attempts == ["dynamo", "dynamo"]
	260	+ assert _dynamo._lazy_dynamo_setup.has_run
	261	+
	262	+ original_register = _dynamo.register_inductor_npu
	263	+ original_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
	264	+ inductor_attempts = []
	265	+
	266	+ def fail_first_inductor_setup():
	267	+ inductor_attempts.append("inductor")
	268	+ if len(inductor_attempts) == 1:
	269	+ raise RuntimeError("simulated Inductor setup failure")
	270	+ return original_register()
	271	+
	272	+ _dynamo.register_inductor_npu = fail_first_inductor_setup
	273	+ try:
	274	+ try:
	275	+ torch.compile(lambda x: x + 1, backend="inductor")
	276	+ except RuntimeError as error:
	277	+ assert "simulated Inductor setup failure" in str(error)
	278	+ else:
	279	+ raise AssertionError("the first Inductor setup should fail")
	280	+
	281	+ assert not _dynamo._lazy_inductor_setup.has_run
	282	+ assert os.environ.get("TORCHINDUCTOR_NPU_BACKEND") == original_env
	283	+ torch.compile(lambda x: x + 1, backend="inductor")
	284	+ finally:
	285	+ _dynamo.register_inductor_npu = original_register
	286	+
	287	+ assert inductor_attempts == ["inductor", "inductor"]
	288	+ assert _dynamo._lazy_inductor_setup.has_run
	289	+ assert _dynamo.is_inductor_npu_initialized()
	290	+ assert os.environ.get("TORCHINDUCTOR_NPU_BACKEND") == original_env
	291	+ assert "torch_npu._inductor" in sys.modules
	292	+ """
	293	+ )
	294	+
	295	+ # Verify a forked child discards inherited in-progress lazy setup state.
	296	+ def test_public_lazy_setup_recovers_after_fork(self):
	297	+ self.run_in_subprocess(
	298	+ """
	299	+ import os
	300	+ import signal
	301	+ import threading
	302	+
	303	+ import torch_npu
	304	+ import torch_npu.dynamo as npu_dynamo
	305	+ from torch_npu.utils import _dynamo
	306	+
	307	+ parent_pid = os.getpid()
	308	+ original_add = _dynamo.add_dynamo_methods_init
	309	+ original_get_backend = npu_dynamo._get_default_backend
	310	+ entered = threading.Event()
	311	+ release = threading.Event()
	312	+
	313	+ def block_parent_setup():
	314	+ if os.getpid() == parent_pid:
	315	+ entered.set()
	316	+ assert release.wait(timeout=10)
	317	+ return original_add()
	318	+
	319	+ _dynamo.add_dynamo_methods_init = block_parent_setup
	320	+ setup_thread = threading.Thread(target=_dynamo._lazy_dynamo_setup)
	321	+ setup_thread.start()
	322	+ assert entered.wait(timeout=5)
	323	+
	324	+ try:
	325	+ child_pid = os.fork()
	326	+ if child_pid == 0:
	327	+ def timeout(_signal, _frame):
	328	+ raise TimeoutError("lazy setup hung after fork")
	329	+
	330	+ signal.signal(signal.SIGALRM, timeout)
	331	+ signal.alarm(5)
	332	+ try:
	333	+ npu_dynamo._get_default_backend = (
	334	+ lambda name: npu_dynamo._eager_npu_backend
	335	+ )
	336	+ graph_module = lambda x: x
	337	+ result = npu_dynamo._npu_backend_entrypoint(
	338	+ graph_module, []
	339	+ )
	340	+ assert result is graph_module
	341	+ signal.alarm(0)
	342	+ except BaseException as error:
	343	+ os.write(
	344	+ 2,
	345	+ f"child lazy setup failed: {error}\\n".encode(),
	346	+ )
	347	+ os._exit(1)
	348	+ os._exit(0)
	349	+
	350	+ _, status = os.waitpid(child_pid, 0)
	351	+ finally:
	352	+ release.set()
	353	+ setup_thread.join(timeout=10)
	354	+ _dynamo.add_dynamo_methods_init = original_add
	355	+ npu_dynamo._get_default_backend = original_get_backend
	356	+
	357	+ assert os.WIFEXITED(status)
	358	+ assert os.waitstatus_to_exitcode(status) == 0
	359	+ assert not setup_thread.is_alive()
	360	+ """
	361	+ )
	362	+
	363	+ # Verify NPUGraphs rejects unsupported options in every registration order.
	364	+ def test_npugraphs_rejects_options_across_registration_order(self):
	365	+ self.run_in_subprocess(
	366	+ """
	367	+ import contextlib
	368	+ from unittest import mock
	369	+
	370	+ import torch
	371	+ import torch_npu
	372	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
	373	+ from torch_npu.utils import _dynamo, _graph_tree
	374	+
	375	+ gm = object()
	376	+ inputs = [object()]
	377	+ options = {"npu_backend": "mlir"}
	378	+
	379	+ with mock.patch.object(
	380	+ _dynamo, "_lazy_dynamo_setup", lambda: None
	381	+ ), mock.patch.object(
	382	+ _dynamo, "_lazy_inductor_setup", lambda: None
	383	+ ), mock.patch.object(
	384	+ _dynamo,
	385	+ "_NpuBackendScope",
	386	+ lambda backend: contextlib.nullcontext(),
	387	+ ), mock.patch.object(
	388	+ _graph_tree,
	389	+ "npugraphs",
	390	+ lambda model, args, **kwargs: "unexpected",
	391	+ ):
	392	+ for backend in (
	393	+ _npugraphs_backend_entrypoint,
	394	+ _graph_tree.NpugraphsBackend(),
	395	+ ):
	396	+ try:
	397	+ backend(gm, inputs, options=options)
	398	+ except TypeError as error:
	399	+ assert "unexpected keyword argument 'options'" in str(error)
	400	+ else:
	401	+ raise AssertionError("npugraphs must reject options")
	402	+ """
	403	+ )
	404	+
	405	+ # Verify public reset reaches NPUGraphs for every registration order.
	406	+ def test_npugraphs_reset_protocol_across_registration_order(self):
	407	+ for order in ("cold", "hot"):
	408	+ with self.subTest(order=order):
	409	+ initialize_inductor = (
	410	+ "torch.compile(lambda x: x + 1, backend='inductor')"
	411	+ if order == "hot"
	412	+ else ""
	413	+ )
	414	+ self.run_in_subprocess(
	415	+ f"""
	416	+ import sys
	417	+ import types
	418	+ from unittest import mock
	419	+
	420	+ import torch
	421	+ import torch_npu
	422	+
	423	+ {initialize_inductor}
	424	+ torch.compile(lambda x: x + 1, backend="npugraphs")
	425	+
	426	+ from torch._dynamo.backends import registry
	427	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
	428	+
	429	+ backend = registry._COMPILER_FNS["npugraphs"]
	430	+ assert backend is _npugraphs_backend_entrypoint
	431	+ assert hasattr(backend, "reset")
	432	+
	433	+ graph_tree_module = "torch_npu.npu._graph_tree"
	434	+ assert graph_tree_module not in sys.modules
	435	+ backend.reset()
	436	+ assert graph_tree_module not in sys.modules
	437	+
	438	+ reset_calls = []
	439	+ fake_graph_tree = types.ModuleType(graph_tree_module)
	440	+ fake_graph_tree.reset_npugraph_trees = (
	441	+ lambda: reset_calls.append("reset")
	442	+ )
	443	+ with mock.patch.dict(
	444	+ sys.modules,
	445	+ {{graph_tree_module: fake_graph_tree}},
	446	+ ):
	447	+ torch.compiler.reset()
	448	+
	449	+ assert reset_calls == ["reset"]
	450	+ """
	451	+ )
	452	+
	453	+ # Verify every public Export entry initializes only Dynamo on NPU.
	454	+ def test_npu_export_public_entry_and_import_order_matrix(self):
	455	+ cases = {
	456	+ "module_export": (
	457	+ "",
	458	+ "exported = torch.export.export(Model(), (x,))",
	459	+ ),
	460	+ "prebound_export": (
	461	+ "from torch.export import export as export_api",
	462	+ "exported = export_api(Model(), (x,))",
	463	+ ),
	464	+ "prebound_export_for_training": (
	465	+ "from torch.export import export_for_training as export_api",
	466	+ "exported = export_api(Model(), (x,))",
	467	+ ),
	468	+ "prebound_export_for_inference": (
	469	+ "from torch.export import export_for_inference as export_api",
	470	+ "exported = export_api(Model(), (x,))",
	471	+ ),
	472	+ }
	473	+ for name, (pre_import, export_call) in cases.items():
	474	+ with self.subTest(name=name):
	475	+ self.run_in_subprocess(
	476	+ f"""
	477	+ import sys
	478	+ import torch
	479	+ {pre_import}
	480	+ import torch_npu
	481	+
	482	+ stream = torch.npu.Stream()
	483	+
	484	+ class Model(torch.nn.Module):
	485	+ def forward(self, x):
	486	+ x.record_stream(stream)
	487	+ return x + 1
	488	+
	489	+ x = torch.ones(4, device="npu")
	490	+ {export_call}
	491	+ actual = exported.module()(x)
	492	+
	493	+ from torch_npu.utils import _dynamo
	494	+
	495	+ torch.testing.assert_close(actual, x + 1)
	496	+ assert _dynamo._lazy_dynamo_setup.has_run
	497	+ assert not _dynamo._lazy_inductor_setup.has_run
	498	+ assert "torch_npu._inductor" not in sys.modules
	499	+ """
	500	+ )
	501	+
	502	+ # Verify NPU-specific operations retain their Export capture semantics.
	503	+ def test_npu_export_capture_semantics_matrix(self):
	504	+ self.run_in_subprocess(
	505	+ """
	506	+ import sys
	507	+ import torch
	508	+ import torch_npu
	509	+
	510	+ x = torch.ones(4, device="npu")
	511	+ stream = torch.npu.Stream()
	512	+ event = torch.npu.Event()
	513	+
	514	+ class StreamAndEvent(torch.nn.Module):
	515	+ def forward(self, value):
	516	+ event.record()
	517	+ with torch.npu.stream(stream):
	518	+ event.wait(stream)
	519	+ result = value + 1
	520	+ return result
	521	+
	522	+ class Autocast(torch.nn.Module):
	523	+ def forward(self, value):
	524	+ with torch.npu.amp.autocast(dtype=torch.float16):
	525	+ return value * value
	526	+
	527	+ class CurrentDevice(torch.nn.Module):
	528	+ def forward(self, value):
	529	+ return value + torch.npu.current_device()
	530	+
	531	+ class DeviceProperties(torch.nn.Module):
	532	+ def forward(self, value):
	533	+ properties = torch.npu.get_device_properties(
	534	+ torch.npu.current_device()
	535	+ )
	536	+ return value + 1 if properties.total_memory > 0 else value - 1
	537	+
	538	+ class IsAvailable(torch.nn.Module):
	539	+ def forward(self, value):
	540	+ return value + 1 if torch.npu.is_available() else value - 1
	541	+
	542	+ models = (
	543	+ StreamAndEvent,
	544	+ Autocast,
	545	+ CurrentDevice,
	546	+ DeviceProperties,
	547	+ IsAvailable,
	548	+ )
	549	+ for model_type in models:
	550	+ model = model_type()
	551	+ expected = model(x)
	552	+ exported = torch.export.export(model, (x,))
	553	+ actual = exported.module()(x)
	554	+ torch.testing.assert_close(actual, expected)
	555	+
	556	+ from torch_npu.utils import _dynamo
	557	+
	558	+ assert _dynamo._lazy_dynamo_setup.has_run
	559	+ assert not _dynamo._lazy_inductor_setup.has_run
	560	+ assert "torch_npu._inductor" not in sys.modules
	561	+ """
	562	+ )
	563	+
	564	+ # Verify rejected Inductor arguments do not initialize or pollute NPU state.
	565	+ def test_invalid_inductor_arguments_fail_without_npu_initialization(self):
	566	+ cases = {
	567	+ "invalid_mode": (
	568	+ 'torch.compile(lambda x: x + 1, mode="invalid-mode")',
	569	+ "Unrecognized mode=invalid-mode",
	570	+ ),
	571	+ "invalid_option": (
	572	+ "torch.compile(lambda x: x + 1, "
	573	+ 'options={"invalid.option": True})',
	574	+ "Unexpected optimization option invalid.option",
	575	+ ),
	576	+ "invalid_npu_option_type": (
	577	+ "torch.compile(lambda x: x + 1, "
	578	+ 'options={"npu_backend": 1})',
	579	+ "Unexpected type of attr npu_backend",
	580	+ ),
	581	+ }
	582	+ for name, (compile_call, expected_error) in cases.items():
	583	+ with self.subTest(name=name):
	584	+ self.run_in_subprocess(
	585	+ f"""
	586	+ import os
	587	+ import sys
	588	+
	589	+ import torch
	590	+ import torch_npu
	591	+ from torch_npu.utils import _dynamo
	592	+
	593	+ env_name = "TORCHINDUCTOR_NPU_BACKEND"
	594	+ original_env = os.environ.get(env_name)
	595	+ try:
	596	+ {compile_call}
	597	+ except RuntimeError as error:
	598	+ assert {expected_error!r} in str(error), str(error)
	599	+ else:
	600	+ raise AssertionError("invalid compile arguments must fail")
	601	+
	602	+ assert not _dynamo._lazy_inductor_setup.has_run
	603	+ assert not _dynamo.is_inductor_npu_initialized()
	604	+ assert "torch_npu._inductor" not in sys.modules
	605	+ assert os.environ.get(env_name) == original_env
	606	+ """
	607	+ )
	608	+
	609	+ # Verify non-Inductor compile backends do not initialize Inductor.
	610	+ def test_non_inductor_compile_backend_matrix(self):
	611	+ cases = {
	612	+ "eager": (
	613	+ "",
	614	+ 'compiled = torch.compile(Model(), backend="eager", fullgraph=True)',
	615	+ ),
	616	+ "custom": (
	617	+ "custom_backend = lambda graph_module, example_inputs: "
	618	+ "graph_module.forward",
	619	+ "compiled = torch.compile(Model(), backend=custom_backend, fullgraph=True)",
	620	+ ),
	621	+ "npu": (
	622	+ "",
	623	+ 'compiled = torch.compile(Model(), backend="npu", fullgraph=True)',
	624	+ ),
	625	+ }
	626	+ for name, (backend_definition, compile_call) in cases.items():
	627	+ with self.subTest(name=name):
	628	+ allow_missing_torchair = name == "npu"
	629	+ self.run_in_subprocess(
	630	+ f"""
	631	+ import sys
	632	+ import torch
	633	+ import torch_npu
	634	+
	635	+ class Model(torch.nn.Module):
	636	+ def forward(self, x):
	637	+ return torch.sin(x) + 1
	638	+
	639	+ {backend_definition}
	640	+ x = torch.randn(8, device="npu")
	641	+ try:
	642	+ {compile_call}
	643	+ except AssertionError as error:
	644	+ assert {allow_missing_torchair!r}
	645	+ assert "Could not find module torchair" in str(error)
	646	+ else:
	647	+ torch.testing.assert_close(compiled(x), Model()(x))
	648	+
	649	+ from torch_npu.utils import _dynamo
	650	+
	651	+ assert _dynamo._lazy_dynamo_setup.has_run
	652	+ assert not _dynamo._lazy_inductor_setup.has_run
	653	+ assert "torch_npu._inductor" not in sys.modules
	654	+ """
	655	+ )
	656	+
	657	+ # Verify ONNX Dynamo Export initializes only the NPU Dynamo integration.
	658	+ @unittest.skipUnless(importlib.util.find_spec("onnxscript"), "requires onnxscript")
	659	+ def test_npu_onnx_dynamo_export_initialization_chain(self):
	660	+ cases = {
	661	+ "module_export": (
	662	+ "",
	663	+ "result = torch.onnx.export(Model(), (x,), dynamo=True)",
	664	+ ),
	665	+ "prebound_export": (
	666	+ "from torch.onnx import export as onnx_export",
	667	+ "result = onnx_export(Model(), (x,), dynamo=True)",
	668	+ ),
	669	+ "prebound_legacy_dynamo_export": (
	670	+ "from torch.onnx import dynamo_export as onnx_export",
	671	+ "result = onnx_export(Model(), x)",
	672	+ ),
	673	+ }
	674	+ for name, (pre_import, export_call) in cases.items():
	675	+ with self.subTest(name=name):
	676	+ self.run_in_subprocess(
	677	+ f"""
	678	+ import sys
	679	+ import torch
	680	+ {pre_import}
	681	+ import torch_npu
	682	+
	683	+ class Model(torch.nn.Module):
	684	+ def forward(self, x):
	685	+ return torch.sin(x) + 1
	686	+
	687	+ x = torch.randn(8, device="npu")
	688	+ {export_call}
	689	+
	690	+ from torch_npu.utils import _dynamo
	691	+
	692	+ assert result is not None
	693	+ assert _dynamo._lazy_dynamo_setup.has_run
	694	+ assert not _dynamo._lazy_inductor_setup.has_run
	695	+ assert "torch_npu._inductor" not in sys.modules
	696	+ """
	697	+ )
	698	+
	699	+ # Verify pre-imported FSDP receives all NPU patches.
	700	+ def test_fsdp_patch_when_imported_before_torch_npu(self):
	701	+ self.run_in_subprocess(
	702	+ """
	703	+ import torch.distributed.fsdp
	704	+ import torch_npu
	705	+ from torch.distributed.fsdp import sharded_grad_scaler
	706	+ from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
	707	+ FSDPParamGroup,
	708	+ )
	709	+ from torch_npu.distributed.fsdp._add_fsdp_patch import _patched_finalize_backward
	710	+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
	711	+
	712	+ assert FSDPParamGroup.finalize_backward is _patched_finalize_backward
	713	+ assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler
	714	+ """
	715	+ )
	716	+
	717	+ # Verify a real NPU Inductor compile initializes the full stack.
	718	+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
	719	+ def test_npu_inductor_initialization_chain(self):
	720	+ self.run_in_subprocess(
	721	+ """
	722	+ import torch
	723	+ import torch_npu
	724	+
	725	+ def fn(x):
	726	+ return torch.sin(x) + 1
	727	+
	728	+ x = torch.randn(8, device="npu")
	729	+ actual = torch.compile(fn, backend="inductor", fullgraph=True)(x)
	730	+
	731	+ from torch._dynamo.device_interface import get_interface_for_device
	732	+ from torch_npu.utils import _dynamo
	733	+
	734	+ torch.testing.assert_close(actual, fn(x))
	735	+ assert _dynamo._lazy_dynamo_setup.has_run
	736	+ assert _dynamo._lazy_inductor_setup.has_run
	737	+ assert get_interface_for_device("npu").device_count() > 0
	738	+ """
	739	+ )
	740	+
	741	+ # Verify a real NPUGraphs compile initializes the full stack.
	742	+ @unittest.skipUnless(importlib.util.find_spec("triton"), "requires triton-ascend")
	743	+ def test_npu_npugraphs_initialization_chain(self):
	744	+ self.run_in_subprocess(
	745	+ """
	746	+ import torch
	747	+ import torch_npu
	748	+
	749	+ def fn(x):
	750	+ return torch.sin(x) + 1
	751	+
	752	+ x = torch.randn(8, device="npu")
	753	+ actual = torch.compile(fn, backend="npugraphs", fullgraph=True)(x)
	754	+
	755	+ from torch._dynamo.device_interface import get_interface_for_device
	756	+ from torch_npu.utils import _dynamo
	757	+
	758	+ torch.testing.assert_close(actual, fn(x))
	759	+ assert _dynamo._lazy_dynamo_setup.has_run
	760	+ assert _dynamo._lazy_inductor_setup.has_run
	761	+ assert get_interface_for_device("npu").device_count() > 0
	762	+ """
	763	+ )
	764	+
	765	+ # Verify lazy setup completes before compile backend lookup.
	766	+ def test_compile_triggers_setup_before_backend_lookup(self):
	767	+ self.run_in_subprocess(
	768	+ """
	769	+ import torch
	770	+ import torch_npu
	771	+ from torch_npu.utils import _dynamo
	772	+
	773	+ calls = []
	774	+
	775	+ @_dynamo.run_once
	776	+ def fake_setup():
	777	+ calls.append("setup")
	778	+
	779	+ _dynamo._lazy_dynamo_setup = fake_setup
	780	+ from torch._dynamo.backends import registry
	781	+ assert fake_setup.has_run
	782	+
	783	+ compiled = torch.compile(lambda x: x + 1, backend="eager")
	784	+ assert compiled(torch.tensor(1)).item() == 2
	785	+ assert calls == ["setup"]
	786	+ """
	787	+ )
	788	+
	789	+ # Verify the trigger works when Dynamo was imported first.
	790	+ def test_trigger_after_dynamo_was_preimported(self):
	791	+ self.run_in_subprocess(
	792	+ """
	793	+ import sys
	794	+ import torch
	795	+ import torch._dynamo
	796	+ import torch_npu
	797	+ from torch_npu.utils import _dynamo
	798	+
	799	+ assert _dynamo._lazy_dynamo_setup.has_run
	800	+
	801	+ compiled = torch.compile(lambda x: x + 1, backend="eager", fullgraph=True)
	802	+ assert compiled(torch.tensor(1)).item() == 2
	803	+ assert _dynamo._lazy_dynamo_setup.has_run
	804	+ assert not _dynamo._lazy_inductor_setup.has_run
	805	+ assert "torch_npu._inductor" not in sys.modules
	806	+ """
	807	+ )
	808	+
	809	+ # Verify all legacy Dynamo patches remain installed exactly once.
	810	+ def test_dynamo_patch_inventory_is_preserved(self):
	811	+ self.run_in_subprocess(
	812	+ """
	813	+ import torch
	814	+ import torch_npu
	815	+
	816	+ # Importing the Dynamo parent package is the lazy setup boundary.
	817	+ import torch._dynamo
	818	+
	819	+ from torch._dynamo.device_interface import get_interface_for_device
	820	+ from torch._dynamo.variables.builtin import BuiltinVariable
	821	+ from torch._dynamo.variables.builder import VariableBuilder
	822	+ from torch._dynamo.variables.ctx_manager import EventVariable
	823	+ from torch._dynamo.variables.functions import SkipFunctionVariable
	824	+ from torch._dynamo.variables.tensor import TensorVariable
	825	+ from torch._dynamo.variables.torch import constant_fold_functions
	826	+ from torch._dynamo.variables.user_defined import UserDefinedClassVariable
	827	+ from torch._dynamo.utils import common_constant_types
	828	+ from torch_npu.dynamo.trace_rule import (
	829	+ skip_functions_npu,
	830	+ torch_c_binding_in_graph_functions_npu,
	831	+ torch_non_c_binding_in_graph_functions_npu,
	832	+ )
	833	+ from torch_npu.utils import _dynamo
	834	+
	835	+ assert _dynamo._lazy_dynamo_setup.has_run
	836	+ assert get_interface_for_device("npu").device_count() > 0
	837	+
	838	+ # VariableTracker and context-manager patches formerly installed
	839	+ # eagerly by add_dynamo_methods().
	840	+ assert SkipFunctionVariable.__new__.__module__ == "torch_npu.utils._dynamo"
	841	+ assert TensorVariable.call_method.__module__ == "torch_npu.utils._dynamo"
	842	+ assert UserDefinedClassVariable.__new__.__module__ == "torch_npu.utils._dynamo"
	843	+ in_graph_classes = UserDefinedClassVariable._in_graph_classes()
	844	+ assert torch.npu.Event in in_graph_classes
	845	+ assert torch.npu.Stream in in_graph_classes
	846	+ assert torch.npu.fake_record_stream is _dynamo.fake_record_stream
	847	+ assert TensorVariable.method_record_stream.__module__ == "torch_npu.utils._dynamo"
	848	+ assert VariableBuilder._wrap.__module__ == "torch_npu.utils._dynamo"
	849	+ assert BuiltinVariable.call_id.__module__ == "torch_npu.utils._dynamo"
	850	+ assert EventVariable.python_type.__module__ == "torch_npu.utils._dynamo"
	851	+ assert torch._dynamo.optimize.__module__ == "torch_npu.utils._dynamo"
	852	+
	853	+ # Backend and trace-rule registrations formerly performed by
	854	+ # registry_manager._register_dynamo(). Count the maps as well as
	855	+ # checking membership so repeated lazy triggers cannot hide a
	856	+ # duplicate installation.
	857	+ assert {"npu", "npugraph_ex"}.issubset(
	858	+ torch._dynamo.list_backends(exclude_tags=None)
	859	+ )
	860	+ maps = torch._dynamo.trace_rules.torch_name_rule_map
	861	+ assert maps.count(torch_non_c_binding_in_graph_functions_npu) == 1
	862	+ assert maps.count(torch_c_binding_in_graph_functions_npu) == 1
	863	+ assert maps.count(skip_functions_npu) == 1
	864	+ assert constant_fold_functions[torch.npu.current_device]
	865	+ assert constant_fold_functions[torch.npu.get_device_properties]
	866	+ assert constant_fold_functions[torch.npu.is_available]
	867	+ assert torch_npu._C._NPUDeviceProperties in common_constant_types
	868	+ """
	869	+ )
	870	+
	871	+ # Verify all legacy Inductor patches remain installed.
	872	+ def test_inductor_patch_inventory_is_preserved(self):
	873	+ self.run_in_subprocess(
	874	+ """
	875	+ import torch
	876	+ import torch_npu
	877	+
	878	+ # RNG/decomposition patches remain installed at import time.
	879	+ from torch_npu.utils import _inductor
	880	+
	881	+ assert (
	882	+ torch._decomp.decompositions._max_unpoolnd
	883	+ is _inductor._max_unpoolnd_patch
	884	+ )
	885	+ assert torch._prims.rng_prims.philox_rand_offset.__module__ == (
	886	+ "torch_npu.utils._inductor"
	887	+ )
	888	+ assert torch._prims.rng_prims.register_philox_rand.__module__ == (
	889	+ "torch_npu.utils._inductor"
	890	+ )
	891	+ assert torch._prims.rng_prims.get_device.__module__ == (
	892	+ "torch_npu.utils._inductor"
	893	+ )
	894	+
	895	+ # Exercise the new full-Inductor setup boundary without relying on
	896	+ # test ordering or on a prior torch.compile invocation.
	897	+ import torch._dynamo
	898	+ from torch_npu.utils import _dynamo
	899	+
	900	+ _dynamo._lazy_inductor_setup()
	901	+
	902	+ import torch._inductor.compile_fx as compile_fx
	903	+ import torch._inductor.cudagraph_trees as cudagraph_trees
	904	+ import torch._inductor.cudagraph_utils as cudagraph_utils
	905	+ from torch._inductor.codegen.common import get_device_op_overrides
	906	+ from torch._inductor.codecache import CacheBase
	907	+ from torch._inductor.graph import GraphLowering
	908	+ from torch._inductor.utils import GPU_TYPES
	909	+ from torch_npu.utils import _graph_tree
	910	+
	911	+ assert _dynamo._lazy_inductor_setup.has_run
	912	+ assert "npu" in GPU_TYPES
	913	+ assert get_device_op_overrides("npu").__class__.__module__.startswith(
	914	+ "torch_npu._inductor"
	915	+ )
	916	+ assert torch.utils._triton.has_triton.__module__ == (
	917	+ "torch_npu._inductor.utils"
	918	+ )
	919	+ assert torch.utils._triton.has_triton_tma.__module__ == (
	920	+ "torch_npu._inductor.utils"
	921	+ )
	922	+ assert compile_fx.has_triton is torch.utils._triton.has_triton
	923	+ assert GraphLowering.codegen_with_cpp_wrapper.__module__ == (
	924	+ "torch_npu._inductor.graph"
	925	+ )
	926	+ assert CacheBase.get_system.__module__ == (
	927	+ "torch_npu._inductor.codegen.common"
	928	+ )
	929	+
	930	+ # NPUGraph integrations were formerly applied eagerly alongside
	931	+ # the Inductor patches.
	932	+ assert compile_fx.cudagraphify is _graph_tree.npugraphify
	933	+ assert (
	934	+ cudagraph_utils.check_multiple_devices_or_any_cpu_nodes
	935	+ is _graph_tree.check_multiple_devices_or_any_cpu_nodes
	936	+ )
	937	+ assert cudagraph_trees.get_manager.__module__ == (
	938	+ "torch_npu.utils._graph_tree"
	939	+ )
	940	+ assert torch.compiler.npugraph_mark_step_begin is (
	941	+ _graph_tree.npugraph_mark_step_begin
	942	+ )
	943	+
	944	+ config = torch._inductor.config
	945	+ assert config.npu_backend == "default"
	946	+ assert config.enable_shape_handling is False
	947	+ assert config.shape_handling_configs == []
	948	+ assert config.shape_handling_dict is None
	949	+ """
	950	+ )
	951	+
	952	+
	953	+if __name__ == "__main__":
	954	+ unittest.main()

Mtest/npu/test_stream.py+5-0

	@@ -23,6 +23,11 @@ class TestNpuStream(TestCase):
23	def test_get_current_stream_interface(self):	23	def test_get_current_stream_interface(self):
24	from torch_npu._C import _npu_getCurrentRawStream, _npu_getCurrentRawStreamNoWait	24	from torch_npu._C import _npu_getCurrentRawStream, _npu_getCurrentRawStreamNoWait
25	from torch._dynamo.device_interface import get_interface_for_device	25	from torch._dynamo.device_interface import get_interface_for_device
		26	+ from torch_npu.utils._dynamo import _dynamo_register_interface_for_device
		27	+
		28	+ # device_interface is an internal Dynamo module. Initialize its NPU
		29	+ # registration explicitly instead of relying on import torch_npu.
		30	+ _dynamo_register_interface_for_device()
26		31
27	device_number = torch.npu.device_count()	32	device_number = torch.npu.device_count()
28	for i in range(device_number):	33	for i in range(device_number):

Mtest/test_torch_npu_init.py+88-14

	@@ -48,6 +48,7 @@ EXPECTED_LOADED_MODULES = [
48		48
49	EXPECTED_NOT_LOADED_MODULES = [	49	EXPECTED_NOT_LOADED_MODULES = [
50	"torch_npu._C._afd",	50	"torch_npu._C._afd",
		51	+ "torch_npu._inductor",
51	]	52	]
52		53
53	EXPECTED_TOP_LEVEL_ATTRS = [	54	EXPECTED_TOP_LEVEL_ATTRS = [
	@@ -262,20 +263,6 @@ class TestTorchNpuBootstrap(TestCase):
262	import torch.distributed as dist	263	import torch.distributed as dist
263	import torch.distributed.rpc as rpc	264	import torch.distributed.rpc as rpc
264	import torch.distributed.tensor # noqa: F401	265	import torch.distributed.tensor # noqa: F401
265	- from torch._dynamo.device_interface import get_interface_for_device
266	- from torch._dynamo.backends.registry import _BACKENDS
267	- from torch._inductor.codegen.common import device_op_overrides_dict
268	-
269	- iface = get_interface_for_device("npu")
270	- assert iface is not None
271	-
272	- assert "npu" in _BACKENDS, "npu dynamo backend is not registered"
273	- assert "npugraph_ex" in _BACKENDS, (
274	- "npugraph_ex dynamo backend is not registered"
275	- )
276	-
277	- assert "npu" in device_op_overrides_dict
278	- assert device_op_overrides_dict.get("npu") is not None
279		266
280	assert "hccl" in dist.Backend.backend_list	267	assert "hccl" in dist.Backend.backend_list
281	assert "lccl" in dist.Backend.backend_list	268	assert "lccl" in dist.Backend.backend_list
	@@ -616,5 +603,92 @@ class TestTorchNpuBootstrap(TestCase):
616	"""	603	"""
617	)	604	)
618		605
		606	+ def test_14_dtensor_strategies_are_registered_without_dynamo(self):
		607	+ self._run_python(
		608	+ """
		609	+ import sys
		610	+ import torch
		611	+ import torch_npu
		612	+ from torch.distributed.tensor import DTensor
		613	+
		614	+ # Importing torch_npu should register both kinds of NPU DTensor strategy.
		615	+ strategy_funcs = DTensor._op_dispatcher.sharding_propagator.op_strategy_funcs
		616	+ assert torch.ops.npu.npu_rms_norm.default in strategy_funcs
		617	+ assert torch.ops.npu.npu_fusion_attention.default in strategy_funcs
		618	+
		619	+ # Core DTensor registration must not pull in compiler or experimental modules.
		620	+ assert "torch_npu.distributed.tensor" in sys.modules
		621	+ assert "torch_npu.distributed.tensor.experimental" not in sys.modules
		622	+ assert "torch.distributed.tensor.experimental" not in sys.modules
		623	+ assert "torch._dynamo" not in sys.modules
		624	+ assert "torch._inductor" not in sys.modules
		625	+
		626	+ # Execute the compact register_sharding adapter without requiring NPU kernels.
		627	+ import os
		628	+ import tempfile
		629	+ import torch.distributed as dist
		630	+ from torch.distributed.device_mesh import DeviceMesh
		631	+ from torch.distributed.tensor import Replicate
		632	+ from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
		633	+ from torch.distributed.tensor._op_schema import (
		634	+ OpSchema,
		635	+ OpStrategy,
		636	+ PlacementStrategy,
		637	+ )
		638	+
		639	+ fd, path = tempfile.mkstemp()
		640	+ os.close(fd)
		641	+ os.unlink(path)
		642	+ dist.init_process_group(
		643	+ "gloo", init_method=f"file://{path}", rank=0, world_size=1
		644	+ )
		645	+ try:
		646	+ mesh = DeviceMesh("cpu", [0])
		647	+ tensor_meta = TensorMeta(
		648	+ torch.Size([2, 4, 8]), (32, 8, 1), torch.float32
		649	+ )
		650	+ spec = DTensorSpec(mesh, (Replicate(),), tensor_meta=tensor_meta)
		651	+ strategy = OpStrategy([PlacementStrategy(output_specs=spec)])
		652	+ op = torch.ops.npu.npu_rotary_mul.default
		653	+ propagator = DTensor._op_dispatcher.sharding_propagator
		654	+ op_schema = OpSchema(
		655	+ op,
		656	+ (strategy, strategy, strategy, "half"),
		657	+ {},
		658	+ propagator.op_to_schema_info[op],
		659	+ )
		660	+ result = propagator.op_strategy_funcs[op](op_schema)
		661	+ assert len(result.strategies) == 3
		662	+ finally:
		663	+ dist.destroy_process_group()
		664	+ if os.path.exists(path):
		665	+ os.remove(path)
		666	+
		667	+ # The experimental namespace should remain available through lazy access.
		668	+ from torch_npu.distributed.tensor import experimental
		669	+ assert callable(experimental.context_parallel)
		670	+ """
		671	+ )
		672	+
		673	+ def test_15_direct_npu_fsdp_import(self):
		674	+ self._run_python(
		675	+ """
		676	+ import torch_npu.distributed.fsdp as npu_fsdp
		677	+ from torch.distributed.fsdp import sharded_grad_scaler
		678	+ from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
		679	+ FSDPParamGroup,
		680	+ )
		681	+ from torch_npu.distributed.fsdp._add_fsdp_patch import (
		682	+ _patched_finalize_backward,
		683	+ )
		684	+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
		685	+
		686	+ # Direct NPU FSDP import should complete and preserve both public and patch APIs.
		687	+ assert callable(npu_fsdp.fully_shard)
		688	+ assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler
		689	+ assert FSDPParamGroup.finalize_backward is _patched_finalize_backward
		690	+ """
		691	+ )
		692	+
619	if __name__ == "__main__":	693	if __name__ == "__main__":
620	run_tests()	694	run_tests()

Dtest/utils/test_inductor.py+0-33

	@@ -1,33 +0,0 @@
1	-from torch_npu.testing.testcase import TestCase, run_tests
2	-from torch_npu.utils._inductor import NPUDeviceOpOverrides
3	-
4	-
5	-class TestInductor():
6	-
7	- def test_device_guard(self):
8	- overrides = NPUDeviceOpOverrides()
9	- result = overrides.device_guard(1)
10	- expected = "torch_npu.npu._DeviceGuard(1)"
11	- self.assertEqual(result, expected)
12	-
13	- def test_synchronize(self):
14	- overrides = NPUDeviceOpOverrides()
15	- result = overrides.synchronize()
16	- expected = "torch_npu.npu.synchronize()"
17	- self.assertEqual(result, expected)
18	-
19	- def test_set_device(self):
20	- overrides = NPUDeviceOpOverrides()
21	- result = overrides.set_device(0)
22	- expected = "torch_npu.npu.set_device(0)"
23	- self.assertEqual(result, expected)
24	-
25	- def test_import_get_raw_stream_as(self):
26	- overrides = NPUDeviceOpOverrides()
27	- result = overrides.import_get_raw_stream_as("test_name")
28	- expected = "from torch._C import _npu_getCurrentRawStream as test_name"
29	- self.assertEqual(result, expected)
30	-
31	-
32	-if __name__ == "__main__":
33	- run_tests()

Mtorch_npu/_inductor/kernel/bmm.py+5-4

	@@ -25,6 +25,7 @@ from torch._inductor.kernel.mm_common import (
25	mm_configs,	25	mm_configs,
26	mm_options,	26	mm_options,
27	)	27	)
		28	+from torch._inductor.kernel import bmm as inductor_bmm
28		29
29	from .mm import is_contiguous_striding	30	from .mm import is_contiguous_striding
30	from ..utils import use_catlass_template, use_triton_template	31	from ..utils import use_catlass_template, use_triton_template
	@@ -33,10 +34,10 @@ from ..utils import use_catlass_template, use_triton_template
33	log = logging.getLogger("torch._inductor")	34	log = logging.getLogger("torch._inductor")
34	aten = torch.ops.aten	35	aten = torch.ops.aten
35		36
36	-aten_bmm = torch._inductor.kernel.bmm.aten_bmm	37	+aten_bmm = inductor_bmm.aten_bmm
37	-aten_baddbmm = torch._inductor.kernel.bmm.aten_baddbmm	38	+aten_baddbmm = inductor_bmm.aten_baddbmm
38	-bmm_configs = torch._inductor.kernel.bmm.bmm_configs	39	+bmm_configs = inductor_bmm.bmm_configs
39	-bmm_template = torch._inductor.kernel.bmm.bmm_template	40	+bmm_template = inductor_bmm.bmm_template
40		41
41		42
42	def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:	43	def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:

Mtorch_npu/_inductor/utils.py+2-1

	@@ -22,6 +22,7 @@ def patch_is_gpu():
22		22
23		23
24	def patch_has_triton():	24	def patch_has_triton():
		25	+ from torch._inductor import compile_fx
25	from torch.utils._triton import has_triton_package	26	from torch.utils._triton import has_triton_package
26		27
27	@functools.lru_cache(None)	28	@functools.lru_cache(None)
	@@ -60,7 +61,7 @@ def patch_has_triton():
60		61
61	torch.utils._triton.has_triton = has_triton	62	torch.utils._triton.has_triton = has_triton
62	torch._inductor.scheduler.has_triton = has_triton	63	torch._inductor.scheduler.has_triton = has_triton
63	- torch._inductor.compile_fx.has_triton = has_triton	64	+ compile_fx.has_triton = has_triton
64		65
65		66
66	def patch_has_triton_tma():	67	def patch_has_triton_tma():

Mtorch_npu/_init/patches/api_patches.py+0-2

	@@ -3,7 +3,6 @@ from torch_npu._init.patches.patch_manager import PatchManager
3		3
4	@PatchManager.register_patch("api")	4	@PatchManager.register_patch("api")
5	def apply_torch_api_patches():	5	def apply_torch_api_patches():
6	- from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch
7	from torch_npu.multiprocessing.reductions import _add_reductions_methods	6	from torch_npu.multiprocessing.reductions import _add_reductions_methods
8	from torch_npu.utils._module import _apply_module_patch	7	from torch_npu.utils._module import _apply_module_patch
9	from torch_npu.utils._optim import add_optim_method	8	from torch_npu.utils._optim import add_optim_method
	@@ -22,5 +21,4 @@ def apply_torch_api_patches():
22	_add_collect_env_methods()	21	_add_collect_env_methods()
23	add_optim_method()	22	add_optim_method()
24	_add_reductions_methods()	23	_add_reductions_methods()
25	- _apply_fsdp_patch()
26	_add_deterministic_patch()	24	_add_deterministic_patch()

Mtorch_npu/_init/patches/distributed_patches.py+61-10

	@@ -1,5 +1,10 @@
		1	+import importlib.abc
		2	+import importlib.util
		3	+import sys
		4	+
1	import torch	5	import torch
2	import torch.distributed.launcher.api	6	import torch.distributed.launcher.api
		7	+import torch.distributed.nn
3		8
4	import torch_npu	9	import torch_npu
5	from torch_npu._init.patches.patch_manager import PatchManager	10	from torch_npu._init.patches.patch_manager import PatchManager
	@@ -131,21 +136,67 @@ def _apply_wrapped_functions(torch, torch_npu):
131	)	136	)
132		137
133		138
134	-def _apply_sharded_grad_scaler_patch(torch):	139	+def _apply_fsdp_patches():
135	- """
136	- Replace PyTorch FSDP ShardedGradScaler with torch_npu implementation.
137	-
138	- Example:
139	- torch.distributed.fsdp.sharded_grad_scaler.ShardedGradScaler
140	- -> torch_npu.npu.amp.sharded_grad_scaler._ShardedGradScaler
141	- """
142	from torch.distributed.fsdp import sharded_grad_scaler	140	from torch.distributed.fsdp import sharded_grad_scaler
143	-	141	+ from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch
144	from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler	142	from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
145		143
		144	+ _apply_fsdp_patch()
146	sharded_grad_scaler.ShardedGradScaler = _ShardedGradScaler	145	sharded_grad_scaler.ShardedGradScaler = _ShardedGradScaler
147		146
148		147
		148	+class _FSDPPostImportLoader:
		H H htchu21 天前 FSDP 改为 sys.meta_path 延迟 hook，但没有测试首次导入 FSDP 与“FSDP 先于torch_npu 导入”两种顺序下，_apply_fsdp_patch() 和ShardedGradScaler 替换是否仍生效请补充覆盖黄桂军 20 天前评论： 20 天前评论：
		149	+ def __init__(self, loader, finder):
		150	+ self._loader = loader
		151	+ self._finder = finder
		152	+
		153	+ def create_module(self, spec):
		154	+ create_module = getattr(self._loader, "create_module", None)
		155	+ return create_module(spec) if create_module is not None else None
		156	+
		157	+ def exec_module(self, module):
		158	+ self._loader.exec_module(module)
		159	+ setattr(torch.distributed, "fsdp", module)
		160	+ _apply_fsdp_patches()
		R R rmch13 天前 [P2] `_FSDPPostImportLoader.exec_module`：在 `fsdp` 模块初始化过程中重入导入 `def exec_module(self, module): self._loader.exec_module(module) # fsdp 正在初始化 setattr(torch.distributed, "fsdp", module) _apply_fsdp_patches() # ← 此处再次 import torch.distributed.fsdp` `_apply_fsdp_patches()` 内部执行： `from torch.distributed.fsdp import sharded_grad_scaler # re-entrant import` 此时 `torch.distributed.fsdp` 已在 `sys.modules` 中注册但 `exec_module` 尚未返回，模块对象处于部分初始化状态。`from torch.distributed.fsdp import sharded_grad_scaler` 会拿到这个半初始化对象，`sharded_grad_scaler` 属性此时可能还未定义，抛出 `ImportError` 或拿到 `None`，导致 patch 静默失败。建议：不在 `exec_module` 内部通过 `from ... import` 重入，改为直接从已传入的 `module` 对象取属性： `def exec_module(self, module): self._loader.exec_module(module) setattr(torch.distributed, "fsdp", module) # 直接从 module 取，不重入 from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch _apply_fsdp_patch() module.sharded_grad_scaler.ShardedGradScaler = _ShardedGradScaler if self._finder in sys.meta_path: sys.meta_path.remove(self._finder)` 黄桂军 13 天前评论： 13 天前评论：
		161	+ if self._finder in sys.meta_path:
		162	+ sys.meta_path.remove(self._finder)
		163	+
		164	+
		165	+def _find_spec_without_finder(finder, fullname):
		166	+ try:
		167	+ index = sys.meta_path.index(finder)
		168	+ except ValueError:
		169	+ return importlib.util.find_spec(fullname)
		170	+
		171	+ sys.meta_path.pop(index)
		172	+ try:
		173	+ return importlib.util.find_spec(fullname)
		174	+ finally:
		175	+ sys.meta_path.insert(min(index, len(sys.meta_path)), finder)
		176	+
		177	+
		178	+class _FSDPPostImportFinder(importlib.abc.MetaPathFinder):
		179	+ _TARGET = "torch.distributed.fsdp"
		180	+
		181	+ def find_spec(self, fullname, path=None, target=None):
		182	+ if fullname != self._TARGET:
		183	+ return None
		184	+ spec = _find_spec_without_finder(self, fullname)
		185	+ if spec is not None and spec.loader is not None:
		186	+ spec.loader = _FSDPPostImportLoader(spec.loader, self)
		187	+ return spec
		188	+
		189	+
		190	+def _install_fsdp_patch_trigger():
		191	+ if "torch.distributed.fsdp" in sys.modules:
		192	+ _apply_fsdp_patches()
		黄桂军13 天前 renyujin: fsdp patch可以不用打，因为用户会单独调用torch_npu.xxx.fsdp，下来再对一下，看是不是可以不需要提起打patch. 黄桂军 13 天前评论： 13 天前评论：黄桂军13 天前评估：这里不改，把社区代码expriment和dynamo解耦。黄桂军 12 天前评论： 12 天前评论：
		193	+ return
		194	+ if not any(
		195	+ isinstance(finder, _FSDPPostImportFinder) for finder in sys.meta_path
		196	+ ):
		197	+ sys.meta_path.insert(0, _FSDPPostImportFinder())
		198	+
		199	+
149	@PatchManager.register_patch("distributed")	200	@PatchManager.register_patch("distributed")
150	def apply_distributed_methods_patch():	201	def apply_distributed_methods_patch():
151	"""	202	"""
	@@ -158,5 +209,5 @@ def apply_distributed_methods_patch():
158	"""	209	"""
159	_apply_internal_replacements(torch, torch_npu)	210	_apply_internal_replacements(torch, torch_npu)
160	_apply_public_api_aliases(torch, torch_npu)	211	_apply_public_api_aliases(torch, torch_npu)
161	- _apply_sharded_grad_scaler_patch(torch)	212	+ _install_fsdp_patch_trigger()
162	_apply_wrapped_functions(torch, torch_npu)	213	_apply_wrapped_functions(torch, torch_npu)

Mtorch_npu/_init/patches/dynamo_patches.py+0-7

	@@ -6,10 +6,3 @@ def apply_dynamo_methods_patch():
6	from torch_npu.utils._dynamo import add_dynamo_methods	6	from torch_npu.utils._dynamo import add_dynamo_methods
7		7
8	add_dynamo_methods()	8	add_dynamo_methods()
9	-
10	-
11	-@PatchManager.register_patch("dynamo")
12	-def apply_npugraph_tree_patch():
13	- from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
14	-
15	- _apply_npugraph_tree_methods()

Mtorch_npu/_init/registry/registry_manager.py+2-24

	@@ -49,27 +49,6 @@ def _register_distributed():
49	register_distributed_backend_for_npu()	49	register_distributed_backend_for_npu()
50		50
51		51
52	-def _register_dynamo():
53	- """
54	- Register Dynamo integration:
55	- - Dynamo backend
56	- - Dynamo device interface
57	- - NPU trace rules for Dynamo
58	- """
59	- from torch_npu._init.registry.dynamo import (
60	- register_dynamo_backends,
61	- register_dynamo_device_interface,
62	- register_dynamo_trace_rules,
63	- )
64	-
65	- register_dynamo_backends()
66	- register_dynamo_device_interface()
67	-
68	- # Do not repeat this call for register_dynamo_trace_rules appends rules into
69	- # Dynamo's global rules maps.
70	- register_dynamo_trace_rules()
71	-
72	-
73	def _register_rpc():	52	def _register_rpc():
74	"""	53	"""
75	Register and init RPC NPU backend.	54	Register and init RPC NPU backend.
	@@ -92,8 +71,8 @@ def _register_components():
92		71
93	Order matters:	72	Order matters:
94	1. NPU backend is the base capability.	73	1. NPU backend is the base capability.
95	- 2. Distributed and Dynamo depend on NPU backend / _C children.	74	+ 2. Distributed depends on the NPU backend / _C children.
96	- 3. RPC and dtensor are Python-side framework integrations.	75	+ 3. RPC is a Python-side framework integration.
97	4. DefaultDeviceType is set after NPU backend is registered.	76	4. DefaultDeviceType is set after NPU backend is registered.
98	"""	77	"""
99	if not hasattr(torch_npu, "_C"):	78	if not hasattr(torch_npu, "_C"):
	@@ -103,6 +82,5 @@ def _register_components():
103		82
104	_register_npu_backend()	83	_register_npu_backend()
105	_register_distributed()	84	_register_distributed()
106	- _register_dynamo()
107	_register_rpc()	85	_register_rpc()
108	_register_default_gradient_device_type()	86	_register_default_gradient_device_type()

Mtorch_npu/contrib/transfer_to_npu.py+13-1

	@@ -438,6 +438,18 @@ def _compose_wrappers(*wrappers):
438		438
439		439
440	def _init():	440	def _init():
		441	+ # transfer_to_npu patches these modules during its own import. Import them
		442	+ # explicitly instead of relying on torch_npu import side effects.
		443	+ import torch._dynamo.trace_rules # noqa: F401
		444	+ import torch._dynamo.utils # noqa: F401
		445	+ import torch._inductor.runtime.autotune_cache # noqa: F401
		446	+ import torch._inductor.compile_fx # noqa: F401
		447	+ import torch._inductor.utils # noqa: F401
		448	+ import torch._inductor.fx_passes.post_grad # noqa: F401
		449	+ import torch._inductor.fx_passes.joint_graph # noqa: F401
		450	+ import torch._inductor.autotune_process # noqa: F401
		451	+ from torch.distributed.checkpoint import filesystem
		452	+
441	_warning_fn('''	453	_warning_fn('''
442	*************************************************************************************************************	454	*************************************************************************************************************
443	The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..	455	The torch.Tensor.cuda and torch.nn.Module.cuda are replaced with torch.Tensor.npu and torch.nn.Module.npu now..
	@@ -541,7 +553,7 @@ def _init():
541	setattr(torch._inductor.autotune_process, "get_gpu_type", _get_npu_type)	553	setattr(torch._inductor.autotune_process, "get_gpu_type", _get_npu_type)
542		554
543	setattr(torch._utils, '_get_available_device_type', _patch_get_available_device_type)	555	setattr(torch._utils, '_get_available_device_type', _patch_get_available_device_type)
544	- setattr(torch.distributed.checkpoint.filesystem._OverlappingCpuLoader, '__init__',	556	+ setattr(filesystem._OverlappingCpuLoader, '__init__',
545	_patch_OverlappingCpuLoader_init_)	557	_patch_OverlappingCpuLoader_init_)
546		558
547	_replace_to_method_in_allowed_methods()	559	_replace_to_method_in_allowed_methods()

Mtorch_npu/distributed/__init__.py+19-1

	@@ -1,3 +1,5 @@
		1	+import importlib
		2	+
1	__all__ = [	3	__all__ = [
2	"is_hccl_available", "reinit_process_group", "reduce_scatter_tensor_uneven", "all_gather_into_tensor_uneven"	4	"is_hccl_available", "reinit_process_group", "reduce_scatter_tensor_uneven", "all_gather_into_tensor_uneven"
3	]	5	]
	@@ -26,5 +28,21 @@ from torch_npu._C._distributed_c10d import (
26	)	28	)
27		29
28		30
29	-from torch_npu.distributed import fsdp, tensor, nn	31	+from torch_npu.distributed import tensor, nn
		32	+
30	from .distributed_c10d import is_hccl_available, reinit_process_group, _reduce_scatter_tensor_uneven as reduce_scatter_tensor_uneven, _all_gather_into_tensor_uneven as all_gather_into_tensor_uneven	33	from .distributed_c10d import is_hccl_available, reinit_process_group, _reduce_scatter_tensor_uneven as reduce_scatter_tensor_uneven, _all_gather_into_tensor_uneven as all_gather_into_tensor_uneven
		34	+
		35	+
		36	+_LAZY_SUBMODULES = {"fsdp"}
		黄桂军13 天前评估，上面如果将fsdp的expriment和dynamo解耦后，这里是不是也不用修改了。黄桂军 12 天前评论： 12 天前评论：
		37	+
		38	+
		39	+def __getattr__(name):
		40	+ if name not in _LAZY_SUBMODULES:
		41	+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
		42	+ module = importlib.import_module(f"{__name__}.{name}")
		43	+ globals()[name] = module
		44	+ return module
		45	+
		46	+
		47	+def __dir__():
		48	+ return sorted(set(globals()) \| _LAZY_SUBMODULES)

Mtorch_npu/distributed/fsdp/__init__.py+2-0

	@@ -1,4 +1,6 @@
1	import torch_npu.distributed.fsdp._fsdp_collectives	1	import torch_npu.distributed.fsdp._fsdp_collectives
		2	+import torch.distributed.fsdp # noqa: F401
		3	+
2	from ._add_fsdp_patch import fully_shard	4	from ._add_fsdp_patch import fully_shard
3		5
4	fully_shard.__module__ = __name__	6	fully_shard.__module__ = __name__

Mtorch_npu/distributed/tensor/__init__.py+14-1

	@@ -1,3 +1,5 @@
		1	+import importlib
		2	+
1	import torch_npu.distributed.tensor._dtensor_patch # patch before register strategy	3	import torch_npu.distributed.tensor._dtensor_patch # patch before register strategy
2	import torch_npu.distributed.tensor._attention	4	import torch_npu.distributed.tensor._attention
3	import torch_npu.distributed.tensor._math_ops	5	import torch_npu.distributed.tensor._math_ops
	@@ -6,4 +8,15 @@ import torch_npu.distributed.tensor._moe_ops
6	import torch_npu.distributed.tensor._pointwise_ops	8	import torch_npu.distributed.tensor._pointwise_ops
7	import torch_npu.distributed.tensor._sharded_tensor_patch	9	import torch_npu.distributed.tensor._sharded_tensor_patch
8	import torch_npu.distributed.tensor._view_ops	10	import torch_npu.distributed.tensor._view_ops
9	-import torch_npu.distributed.tensor.experimental	11	+
黄桂军13 天前已过期 renyujin；评估这里看是否可以删除黄桂军 13 天前评论： 13 天前评论：
		12	+
		13	+def __getattr__(name):
		14	+ if name != "experimental":
		15	+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
		W W wjlflyer14 天前这里写法如果name不为experimental就会抛异常，就限制了只能加载experimental模块，如果tensor下还有其它模块，这个懒加载方式就不能加载了，再评估一下是否合理。黄桂军 14 天前评论： 14 天前评论：
		16	+ module = importlib.import_module(f"{__name__}.experimental")
		17	+ globals()[name] = module
		18	+ return module
		19	+
		20	+
		21	+def __dir__():
		22	+ return sorted(set(globals()) \| {"experimental"})

Mtorch_npu/distributed/tensor/_attention.py+2-2

	@@ -3,9 +3,8 @@ from typing import cast, Any, Dict, Tuple
3		3
4	import torch	4	import torch
5	from torch.distributed.device_mesh import DeviceMesh	5	from torch.distributed.device_mesh import DeviceMesh
6	-from torch.distributed._tensor.experimental import register_sharding
7	-from torch.distributed._tensor.placement_types import DTensorSpec
8	from torch.distributed.tensor import DTensor, Replicate, Shard	6	from torch.distributed.tensor import DTensor, Replicate, Shard
		7	+from torch.distributed.tensor._dtensor_spec import DTensorSpec
		W W wjlflyer14 天前这里改变import路径原因是什么，是修复问题吗黄桂军 14 天前评论： 14 天前评论：
9	from torch.distributed.tensor._op_schema import (	8	from torch.distributed.tensor._op_schema import (
10	OpInfo,	9	OpInfo,
11	OpSchema,	10	OpSchema,
	@@ -15,6 +14,7 @@ from torch.distributed.tensor._redistribute import redistribute_local_tensor
15		14
16	import torch_npu	15	import torch_npu
17		16
		17	+from ._dtensor_patch import register_sharding
18	from ._common import (	18	from ._common import (
19	get_redistributed_local_args,	19	get_redistributed_local_args,
20	get_redistributed_local_kwargs,	20	get_redistributed_local_kwargs,

Mtorch_npu/distributed/tensor/_dtensor_patch.py+63-0

	@@ -1,17 +1,22 @@
1	# Copyright (c) Meta Platforms, Inc. and affiliates	1	# Copyright (c) Meta Platforms, Inc. and affiliates
2		2
3	import itertools	3	import itertools
		4	+from functools import partial
4	from typing import Callable, Optional	5	from typing import Callable, Optional
5		6
6	import torch	7	import torch
		8	+from torch.distributed.tensor import DTensor
7	from torch.distributed.tensor._dtensor_spec import DTensorSpec	9	from torch.distributed.tensor._dtensor_spec import DTensorSpec
8	from torch.distributed.tensor._op_schema import (	10	from torch.distributed.tensor._op_schema import (
		11	+ _is_inplace_op,
9	OpSchema,	12	OpSchema,
10	OpStrategy,	13	OpStrategy,
11	PlacementStrategy,	14	PlacementStrategy,
12	PlacementList,	15	PlacementList,
		16	+ RuntimeSchemaInfo,
13	TupleStrategy	17	TupleStrategy
14	)	18	)
		19	+from torch.distributed.tensor._ops import utils as dtensor_utils
15	from torch.distributed.tensor._ops.utils import (	20	from torch.distributed.tensor._ops.utils import (
16	generate_redistribute_costs,	21	generate_redistribute_costs,
17	is_tensor_shardable	22	is_tensor_shardable
	@@ -24,6 +29,64 @@ except ImportError:
24	from torch.utils._pytree import register_pytree_node, tree_leaves	29	from torch.utils._pytree import register_pytree_node, tree_leaves
25		30
26		31
		32	+# Adapted from PyTorch v2.7.1's DTensor experimental register_sharding.
		33	+def register_sharding(op):
		W W wjlflyer14 天前这里为什么需要重写register_sharding，如果属于新增patch的话需要评审下。黄桂军 14 天前评论： 14 天前评论：黄桂军13 天前还是看experimental与dynamo的解耦黄桂军 12 天前评论： 12 天前评论：
		34	+ """Register an NPU sharding function without importing DTensor experimental APIs."""
		35	+ def custom_strategy(custom_sharding_fn, op_schema):
		36	+ def strategy_to_spec(strategy):
		37	+ if isinstance(strategy, OpStrategy):
		38	+ return strategy.strategies[0].output_spec
		39	+ if isinstance(strategy, TupleStrategy):
		40	+ return tuple(strategy_to_spec(child) for child in strategy.childs)
		41	+ return strategy
		42	+
		43	+ args_schema = tuple(strategy_to_spec(arg) for arg in op_schema.args_schema)
		44	+ kwargs_schema = {
		45	+ key: strategy_to_spec(value)
		46	+ for key, value in op_schema.kwargs_schema.items()
		47	+ }
		48	+ single_mesh_dim_strategies = [
		49	+ output_specs + input_specs
		50	+ for output_specs, input_specs in custom_sharding_fn(
		51	+ args_schema, *kwargs_schema
		52	+ )
		53	+ ]
		54	+ return dtensor_utils.expand_to_full_mesh_op_strategy(
		55	+ op_schema.get_mesh_from_args(),
		56	+ op_schema,
		57	+ single_mesh_dim_strategies,
		58	+ input_index=len(op_schema.op._schema.returns),
		59	+ inplace_op=_is_inplace_op(op_schema.op),
		60	+ )
		61	+
		62	+ def wrapper(custom_sharding_fn):
		63	+ overloads = op if isinstance(op, list) else [op]
		64	+ for overload in overloads:
		65	+ static_argnum = 100
		66	+ static_kwargkey = []
		67	+ for index, arg in enumerate(overload._schema.arguments):
		68	+ if isinstance(arg.type, torch.IntType) or (
		69	+ isinstance(arg.type, torch.OptionalType)
		70	+ and isinstance(arg.type.getElementType(), torch.IntType)
		71	+ ):
		72	+ static_argnum = min(index, static_argnum)
		73	+ if arg.kwarg_only:
		74	+ static_kwargkey.append(arg.name)
		75	+ schema_info = RuntimeSchemaInfo(
		76	+ static_argnum,
		77	+ static_kwargkey or None,
		78	+ needs_pytree=True,
		79	+ )
		80	+ DTensor._op_dispatcher.sharding_propagator.register_op_strategy(
		81	+ overload,
		82	+ partial(custom_strategy, custom_sharding_fn),
		83	+ schema_info,
		84	+ )
		85	+ return custom_sharding_fn
		86	+
		87	+ return wrapper
		88	+
		89	+
27	def _patched_kwargs_strategy(self) -> tuple[OpStrategy, ...]:	90	def _patched_kwargs_strategy(self) -> tuple[OpStrategy, ...]:
28	kwargs_vals = (	91	kwargs_vals = (
29	tree_leaves(self.kwargs_schema)	92	tree_leaves(self.kwargs_schema)

Mtorch_npu/distributed/tensor/_math_ops.py+2-1

	@@ -26,7 +26,8 @@ from torch.distributed.tensor._ops._math_ops import (
26	)	26	)
27	from torch.distributed.tensor._utils import normalize_to_torch_size	27	from torch.distributed.tensor._utils import normalize_to_torch_size
28	from torch.distributed.tensor import Partial, Replicate, Shard	28	from torch.distributed.tensor import Partial, Replicate, Shard
29	-from torch.distributed.tensor.experimental import register_sharding	29	+
		30	+from ._dtensor_patch import register_sharding
30		31
31	logger = logging.getLogger("torch.distributed.tensor")	32	logger = logging.getLogger("torch.distributed.tensor")
32	npu = torch.ops.npu	33	npu = torch.ops.npu

Mtorch_npu/distributed/tensor/_matrix_ops.py+1-1

	@@ -3,7 +3,6 @@ from typing import cast
3		3
4	import torch	4	import torch
5	import torch_npu	5	import torch_npu
6	-from torch.distributed._tensor.experimental import register_sharding
7	from torch.distributed.tensor import DTensor, Partial, Replicate, Shard	6	from torch.distributed.tensor import DTensor, Partial, Replicate, Shard
8	from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta	7	from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
9	from torch.distributed.tensor._op_schema import (	8	from torch.distributed.tensor._op_schema import (
	@@ -33,6 +32,7 @@ from ._common import (
33	get_redistributed_local_args,	32	get_redistributed_local_args,
34	get_redistributed_local_kwargs,	33	get_redistributed_local_kwargs,
35	)	34	)
		35	+from ._dtensor_patch import register_sharding
36		36
37		37
38	aten = torch.ops.aten	38	aten = torch.ops.aten

Mtorch_npu/distributed/tensor/_moe_ops.py+3-2

	@@ -1,6 +1,7 @@
1	import torch	1	import torch
2	-from torch.distributed._tensor import Partial, Replicate, Shard	2	+from torch.distributed.tensor import Partial, Replicate, Shard
3	-from torch.distributed._tensor.experimental import register_sharding	3	+
		4	+from ._dtensor_patch import register_sharding
		黄桂军13 天前 xiuxiu: 蚂蚁使用v2.9.0，再确认下字节是不是要升2.9 黄桂军 12 天前评论： 12 天前评论：
4		5
5	npu = torch.ops.npu	6	npu = torch.ops.npu
6		7

Mtorch_npu/dynamo/__init__.py+71-5

	@@ -3,9 +3,6 @@ import sys
3	import time	3	import time
4	import warnings	4	import warnings
5		5
6	-from torch._dynamo import register_backend as _register_backend
7	-from torch._dynamo.backends.registry import _BACKENDS
8	-
9	from torch_npu._init.common.warning_utils import _should_print_warning	6	from torch_npu._init.common.warning_utils import _should_print_warning
10	from torch_npu.utils._error_code import ErrCode, pta_error	7	from torch_npu.utils._error_code import ErrCode, pta_error
11		8
	@@ -107,6 +104,15 @@ class _LazyTorchair(_LazyBackend):
107	return getattr(torchair, name)	104	return getattr(torchair, name)
108		105
109		106
		107	+def _install_lazy_torchair():
		108	+ torchair_path = os.path.join(os.path.dirname(__file__), "torchair")
		109	+ if not os.path.exists(torchair_path):
		110	+ return False
		111	+ if "torchair" not in sys.modules:
		112	+ sys.modules["torchair"] = _LazyTorchair("torchair")
		黄桂军13 天前 lijing: setup.py中静态注册torchair是否可以代替这个修改。黄桂军 13 天前评论： 13 天前评论：
		113	+ return True
		114	+
		115	+
110	class _LazyNpuGraphEx(_LazyBackend):	116	class _LazyNpuGraphEx(_LazyBackend):
111	def __init__(self, pkg_name):	117	def __init__(self, pkg_name):
112	self._npugraph_ex = None	118	self._npugraph_ex = None
	@@ -140,7 +146,7 @@ def _lazy_exec(args, *kwargs):
140		146
141		147
142	def _get_default_backend(name):	148	def _get_default_backend(name):
143	- if not os.path.exists(os.path.join(os.path.dirname(__file__), 'torchair')):	149	+ if not _install_lazy_torchair():
144	if _should_print_warning():	150	if _should_print_warning():
145	warnings.warn(	151	warnings.warn(
146	"Register eager implementation for the 'npu' backend of dynamo, "	152	"Register eager implementation for the 'npu' backend of dynamo, "
	@@ -148,7 +154,6 @@ def _get_default_backend(name):
148	return _eager_npu_backend	154	return _eager_npu_backend
149	global _global_backend_name	155	global _global_backend_name
150	_global_backend_name = name	156	_global_backend_name = name
151	- sys.modules['torchair'] = _LazyTorchair('torchair')
152	return _lazy_exec	157	return _lazy_exec
153		158
154		159
	@@ -165,6 +170,27 @@ def _get_npugraph_ex_backend():
165		170
166		171
167	def _register_npu_backend(backend, name="npu"):	172	def _register_npu_backend(backend, name="npu"):
		173	+ from torch._dynamo import register_backend as _register_backend
		174	+ from torch._dynamo.backends.registry import _BACKENDS, _COMPILER_FNS
		175	+
		176	+ registered_backend = _COMPILER_FNS.get(name)
		177	+ if (
		178	+ registered_backend is not None
		179	+ and getattr(registered_backend, "__module__", None) == __name__
		180	+ ):
		181	+ return
		182	+
		183	+ # When a setuptools entry point is currently loading torch_npu, Dynamo has
		184	+ # already put the EntryPoint object in _BACKENDS but has not registered the
		185	+ # loaded callable yet. Leave that registration to lookup_backend(); doing
		186	+ # it here as well makes lookup_backend register the same name twice.
		187	+ pending_backend = _BACKENDS.get(name)
		188	+ if (
		189	+ name not in _COMPILER_FNS
		190	+ and getattr(pending_backend, "module", None) == __name__
		191	+ ):
		192	+ return
		193	+
168	if name in _BACKENDS.keys():	194	if name in _BACKENDS.keys():
169	del _BACKENDS[name]	195	del _BACKENDS[name]
170	_register_backend(backend, name)	196	_register_backend(backend, name)
	@@ -176,3 +202,43 @@ def _register_backends():
176		202
177	_register_npu_backend(global_backend)	203	_register_npu_backend(global_backend)
178	_register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)	204	_register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)
		205	+
		206	+
		207	+def _npu_backend_entrypoint(gm, example_inputs, **kwargs):
		208	+ """Set up the NPU Dynamo backend when an entry point is selected."""
		209	+ from torch_npu.utils._dynamo import _lazy_dynamo_setup
		210	+
		211	+ _lazy_dynamo_setup()
		212	+ return _get_default_backend("npu")(gm, example_inputs, **kwargs)
		213	+
		214	+
		215	+def _npugraph_ex_backend_entrypoint(gm, example_inputs, **kwargs):
		216	+ """Set up the NPUGraph-EX backend when an entry point is selected."""
		217	+ from torch_npu.utils._dynamo import _lazy_dynamo_setup
		218	+
		219	+ _lazy_dynamo_setup()
		220	+ return _exec(gm, example_inputs, **kwargs)
		221	+
		222	+
		223	+class _NpugraphsBackendEntryPoint:
		224	+ """Keep the lazy entry point and its reset protocol in one callable."""
		225	+
		226	+ compiler_name = "npugraphs"
		227	+
		228	+ def __call__(self, gm, example_inputs, **kwargs):
		229	+ from torch_npu.utils._dynamo import _lazy_dynamo_setup, _lazy_inductor_setup
		230	+
		231	+ _lazy_dynamo_setup()
		232	+ _lazy_inductor_setup()
		233	+ from torch_npu.utils._graph_tree import NpugraphsBackend
		234	+
		235	+ return NpugraphsBackend()(gm, example_inputs, **kwargs)
		236	+
		237	+ @staticmethod
		238	+ def reset():
		239	+ graph_tree = sys.modules.get("torch_npu.npu._graph_tree")
		240	+ if graph_tree is not None:
		241	+ graph_tree.reset_npugraph_trees()
		242	+
		243	+
		244	+_npugraphs_backend_entrypoint = _NpugraphsBackendEntryPoint()

Mtorch_npu/npu/__init__.py+13-2

	@@ -151,6 +151,7 @@ import traceback
151	import threading	151	import threading
152	import os	152	import os
153	import re	153	import re
		154	+import importlib
154	import torch	155	import torch
155	from torch.storage import _LegacyStorage, _warn_typed_storage_removal	156	from torch.storage import _LegacyStorage, _warn_typed_storage_removal
156	from torch._utils import classproperty	157	from torch._utils import classproperty
	@@ -172,8 +173,6 @@ from .autocast_utils import * # noqa: F403
172	from .backends import * # noqa: F403	173	from .backends import * # noqa: F403
173	from ._backends import * # noqa: F403	174	from ._backends import * # noqa: F403
174	from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403	175	from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403
175	-from . import npugraph_ex
176	-
177	from .graphs import (	176	from .graphs import (
178	NPUGraph,	177	NPUGraph,
179	graph,	178	graph,
	@@ -194,6 +193,18 @@ from ._npugraph_handlers import (
194	)	193	)
195		194
196		195
		196	+def __getattr__(name):
		197	+ if name == "npugraph_ex":
		198	+ module = importlib.import_module("torch_npu.npu.npugraph_ex")
		199	+ globals()[name] = module
		200	+ return module
		201	+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
		202	+
		203	+
		204	+def __dir__():
		205	+ return sorted(set(globals()) \| {"npugraph_ex"})
		206	+
		207	+
197	config = npu_config._npuConfig()	208	config = npu_config._npuConfig()
198		209
199	matmul = npu_config._allowHF32Matmul()	210	matmul = npu_config._allowHF32Matmul()

Mtorch_npu/npu/_graph_tree.py+1-12

	@@ -100,6 +100,7 @@ from torch.utils import _pytree as pytree
100	from torch.utils.weak import TensorWeakRef	100	from torch.utils.weak import TensorWeakRef
101		101
102	import torch_npu	102	import torch_npu
		103	+from torch_npu.npu._graph_tree_state import MarkStepBox, mark_step_begin
103	from torch_npu.npu import graphs as _npu_graphs	104	from torch_npu.npu import graphs as _npu_graphs
104	from torch_npu._C import (	105	from torch_npu._C import (
105	_npu_NPUAllocator_AllocatorState as AllocatorState,	106	_npu_NPUAllocator_AllocatorState as AllocatorState,
	@@ -267,24 +268,12 @@ local.npu_tree_manager_containers = {}
267	local.npu_tree_manager_locks = defaultdict(threading.Lock)	268	local.npu_tree_manager_locks = defaultdict(threading.Lock)
268		269
269		270
270	-# only incremented by user call of mark_step_begin
271	-class MarkStepBox:
272	- mark_step_counter = 0
273	-
274	-
275	# We need to register this as an object that will be copied over as TLS when new	271	# We need to register this as an object that will be copied over as TLS when new
276	# threads are created in autograd	272	# threads are created in autograd
277	torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)	273	torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers)
278	torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)	274	torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks)
279		275
280		276
281	-def mark_step_begin() -> None:
282	- "Indicates that a new iteration of inference or training is about to begin."
283	-
284	- # iterate down to distinguish from GenerationTracking counter
285	- MarkStepBox.mark_step_counter -= 1
286	-
287	-
288	def reset_npugraph_trees() -> None:	277	def reset_npugraph_trees() -> None:
289	"Clear all npugraph trees"	278	"Clear all npugraph trees"
290	# see shutdown below for why this is necessary	279	# see shutdown below for why this is necessary

Atorch_npu/npu/_graph_tree_state.py+11-0

@@ -0,0 +1,11 @@
	1	+"""Lightweight state shared by the public NPUGraph marker and graph trees."""
	2	+
	3	+
	4	+class MarkStepBox:
	5	+ # Negative values distinguish explicit user steps from Dynamo generations.
	6	+ mark_step_counter = 0
	7	+
	8	+
	9	+def mark_step_begin() -> None:
	10	+ """Indicate that a new inference or training iteration is about to begin."""
	11	+ MarkStepBox.mark_step_counter -= 1

Mtorch_npu/npu/deterministic.py+11-3

	@@ -2,11 +2,19 @@ import torch
2		2
3	from torch import Tensor	3	from torch import Tensor
4	from torch.autograd.function import Function	4	from torch.autograd.function import Function
5	-from torch._dynamo.decorators import forbid_in_graph
6		5
7	__all__ = ["enable_deterministic_with_backward", "disable_deterministic_with_backward"]	6	__all__ = ["enable_deterministic_with_backward", "disable_deterministic_with_backward"]
8		7
9		8
		9	+def _forbid_in_graph(fn):
		10	+ if isinstance(fn, (list, tuple)):
		11	+ return [_forbid_in_graph(x) for x in fn]
		12	+ if not callable(fn):
		13	+ raise AssertionError("forbid_in_graph applies only to callables")
		14	+ fn._dynamo_forbidden = True
		15	+ return fn
		16	+
		17	+
10	class _DeterministicAlgorithmsBeginOp(Function):	18	class _DeterministicAlgorithmsBeginOp(Function):
11		19
12	@staticmethod	20	@staticmethod
	@@ -36,11 +44,11 @@ class _DeterministicAlgorithmsEndOp(Function):
36	return grad_outputs	44	return grad_outputs
37		45
38		46
39	-@forbid_in_graph	47	+@_forbid_in_graph
40	def enable_deterministic_with_backward(tensor: Tensor):	48	def enable_deterministic_with_backward(tensor: Tensor):
41	return _DeterministicAlgorithmsBeginOp.apply(tensor)	49	return _DeterministicAlgorithmsBeginOp.apply(tensor)
42		50
43		51
44	-@forbid_in_graph	52	+@_forbid_in_graph
45	def disable_deterministic_with_backward(tensor: Tensor):	53	def disable_deterministic_with_backward(tensor: Tensor):
46	return _DeterministicAlgorithmsEndOp.apply(tensor)	54	return _DeterministicAlgorithmsEndOp.apply(tensor)

Mtorch_npu/npu/npugraph_ex/__init__.py+2-2

	@@ -30,5 +30,5 @@ def register_replacement(search_fn: SearchFn, replace_fn: ReplaceFn, example_inp
30	return npugraph_ex.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,	30	return npugraph_ex.patterns.pattern_pass_manager.register_replacement(search_fn, replace_fn, example_inputs,
31	trace_fn=trace_fn, extra_check=extra_check,	31	trace_fn=trace_fn, extra_check=extra_check,
32	search_fn_pattern=search_fn_pattern,	32	search_fn_pattern=search_fn_pattern,
33	- scalar_workaround=scalar_workaround,	33	+ scalar_workaround=scalar_workaround,
34	- skip_duplicates=skip_duplicates)	34	+ skip_duplicates=skip_duplicates)

Mtorch_npu/utils/_dynamo.py+351-117

	@@ -1,118 +1,95 @@
1	import importlib	1	import importlib
		2	+import importlib.abc
		3	+import functools
2	import inspect	4	import inspect
3	import logging	5	import logging
4	import os	6	import os
5	import sys	7	import sys
		8	+import threading
6		9
7	import torch	10	import torch
8	import torch_npu	11	import torch_npu
9	from torch import _TorchCompileWrapper	12	from torch import _TorchCompileWrapper
10	-from torch._dynamo import optimize
11	-from torch._dynamo.utils import tensortype_to_dtype
12	-from torch._dynamo.variables.base import VariableTracker
13	-from torch._dynamo.variables.constant import ConstantVariable
14	-from torch._dynamo.variables.ctx_manager import AutocastModeVariable
15	-from torch._dynamo.variables.functions import SkipFunctionVariable
16	-from torch._dynamo.variables.lists import TupleVariable
17	-from torch._dynamo.variables.tensor import TensorVariable
18	-from torch._dynamo.variables.torch import (
19	- TorchCtxManagerClassVariable,
20	- TorchInGraphFunctionVariable,
21	-)
22	-from torch._dynamo.variables.user_defined import UserDefinedClassVariable
23	-from torch_npu.dynamo import _get_global_npu_backend
24		13
25		14
26	use_jit_script = False	15	use_jit_script = False
27	log = logging.getLogger(__name__)	16	log = logging.getLogger(__name__)
28		17
29		18
30	-class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable):	19	+def _create_npu_autocast_mode_variable(func, args, kwargs):
31	- def call_function(self, tx, args, kwargs):	20	+ from torch._dynamo.variables.base import VariableTracker
32	- return NPUAutocastModeVariable.create(self.value, args, kwargs)	21	+ from torch._dynamo.variables.ctx_manager import AutocastModeVariable
		22	+
		23	+ bound_args = inspect.signature(func).bind(args, *kwargs)
		24	+ bound_args.apply_defaults()
		25	+ target_values = []
		26	+ kwargs.clear()
		27	+
		28	+ for key in ["device_type", "dtype", "enabled", "cache_enabled"]:
		29	+ if key == "device_type" and func in [
		30	+ torch_npu.npu.amp.autocast,
		31	+ ]:
		32	+ arg = "npu" if func is torch_npu.npu.amp.autocast else "cpu"
		33	+ else:
		34	+ arg = bound_args.arguments[key]
		35	+ if isinstance(arg, VariableTracker):
		36	+ target_values.append(arg.as_python_constant())
		37	+ else:
		38	+ target_values.append(arg)
		39	+
		40	+ return AutocastModeVariable(target_values, initial_values=None, **kwargs)
33		41
34		42
35	-class NPUAutocastModeVariable(AutocastModeVariable):	43	+def patch_SkipFunctionVariable():
36	- @staticmethod	44	+ from torch._dynamo.variables.functions import SkipFunctionVariable
37	- def create(func, args, kwargs):	45	+ from torch._dynamo.variables.torch import TorchInGraphFunctionVariable
38	- bound_args = inspect.signature(func).bind(args, *kwargs)
39	- bound_args.apply_defaults()
40	- target_values = []
41	- kwargs.clear()
42		46
43	- for key in ["device_type", "dtype", "enabled", "cache_enabled"]:	47	+ def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs):
44	- if key == "device_type" and func in [	48	+ if value in [
45	- torch_npu.npu.amp.autocast,	49	+ torch.npu.stream,
46	- ]:	50	+ torch_npu.npu.stream,
47	- arg = "npu" if func is torch_npu.npu.amp.autocast else "cpu"	51	+ torch_npu.npu.utils.stream,
48	- else:	52	+ ]:
49	- arg = bound_args.arguments[key]	53	+ return TorchInGraphFunctionVariable(value, **kwargs)
50	- if isinstance(arg, VariableTracker):	54	+ return cls.__new__raw(cls)
51	- target_values.append(arg.as_python_constant())
52	- else:
53	- target_values.append(arg)
54		55
55	- var = AutocastModeVariable(target_values, initial_values=None, **kwargs)	56	+ SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__
56	- return var	57	+ SkipFunctionVariable.__new__ = SkipFunctionVariable__new__
57		58
58		59
59	-def UserDefinedClassVariable__new__(cls, value, **kwargs):	60	+def patch_TensorVariable_call_method():
60	- if value in [	61	+ from torch._dynamo.utils import tensortype_to_dtype
61	- torch.npu.amp.autocast,	62	+ from torch._dynamo.variables.constant import ConstantVariable
62	- torch_npu.npu.amp.autocast,	63	+ from torch._dynamo.variables.lists import TupleVariable
63	- torch.npu.amp.autocast_mode.autocast,	64	+ from torch._dynamo.variables.tensor import TensorVariable
64	- torch_npu.npu.amp.autocast_mode.autocast,
65	- ]:
66	- return NPUTorchCtxManagerClassVariable(value, **kwargs)
67	- elif value in [
68	- torch.npu.Stream,
69	- torch_npu.npu.Stream,
70	- torch.npu.streams.Stream,
71	- torch_npu.npu.streams.Stream,
72	- torch_npu.npu.BoolTensor,
73	- torch_npu.npu.ByteTensor,
74	- torch_npu.npu.CharTensor,
75	- torch_npu.npu.DoubleTensor,
76	- torch_npu.npu.FloatTensor,
77	- torch_npu.npu.HalfTensor,
78	- torch_npu.npu.IntTensor,
79	- torch_npu.npu.LongTensor,
80	- torch_npu.npu.ShortTensor,
81	- torch_npu.npu.BFloat16Tensor,
82	- ]:
83	- return TorchInGraphFunctionVariable(value, **kwargs)
84	- return cls.__new__raw(cls)
85		65
		66	+ def TensorVariable_call_method(self, tx, name, args, kwargs):
		67	+ if (
		68	+ name == "type"
		69	+ and self.dtype is not None
		70	+ and len(args) == 0
		71	+ and isinstance(self.device, torch.device)
		72	+ and self.device.type == "npu"
		73	+ ):
		74	+ tensortype = next(
		75	+ k for k, v in tensortype_to_dtype.items() if self.dtype in v
		76	+ )
		77	+ constant_result = ConstantVariable.create(
		78	+ f"torch.npu.{tensortype.__name__}"
		79	+ )
86		80
87	-def SkipFunctionVariable__new__(cls, value, reason=None, **kwargs):	81	+ if len(args) == 1:
88	- if value in [	82	+ return constant_result.getitem_const(args[0])
89	- torch.npu.stream,	83	+ if args:
90	- torch_npu.npu.stream,	84	+ return TupleVariable(
91	- torch_npu.npu.utils.stream,	85	+ [constant_result.getitem_const(a) for a in args]
92	- ]:	86	+ )
93	- return TorchInGraphFunctionVariable(value, **kwargs)	87	+ return constant_result
94	- return cls.__new__raw(cls)
95	-
96	-
97	-def TensorVariable_call_method(self, tx, name, args, kwargs):
98	- if (
99	- name == "type"
100	- and self.dtype is not None
101	- and len(args) == 0
102	- and isinstance(self.device, torch.device)
103	- and self.device.type == "npu"
104	- ):
105	- tensortype = next(k for k, v in tensortype_to_dtype.items() if self.dtype in v)
106	- constant_result = ConstantVariable.create(f"torch.npu.{tensortype.__name__}")
107	-
108	- if len(args) == 1:
109	- return constant_result.getitem_const(args[0])
110	- elif args:
111	- return TupleVariable([constant_result.getitem_const(a) for a in args])
112	- return constant_result
113	- else:
114	return TensorVariable.call_method_raw(self, tx, name, args, kwargs)	88	return TensorVariable.call_method_raw(self, tx, name, args, kwargs)
115		89
		90	+ TensorVariable.call_method_raw = TensorVariable.call_method
		91	+ TensorVariable.call_method = TensorVariable_call_method
		92	+
116		93
117	class _InductorNpuRegistry:	94	class _InductorNpuRegistry:
118	_disabled_register = False	95	_disabled_register = False
	@@ -161,19 +138,23 @@ def register_inductor_npu():
161	_InductorNpuRegistry.register_inductor_npu()	138	_InductorNpuRegistry.register_inductor_npu()
162		139
163		140
164	-def _resolve_npu_backend_from_wrapper(wrapper) -> str:	141	+def _resolve_npu_backend(selected_backend=None) -> str:
165	- """Resolve npu backend with priority: wrapper options > global config > env."""	142	+ """Resolve NPU backend with priority: compile options > config > env."""
166	- wrapper_backend = wrapper.config.get("npu_backend")	143	+ if selected_backend not in (None, "", "default"):
167	- if wrapper_backend not in (None, "", "default"):	144	+ return selected_backend
168	- return wrapper_backend
169		145
170	- global_backend = getattr(torch._inductor.config, "npu_backend", None)	146	+ inductor_config = sys.modules.get("torch._inductor.config")
		147	+ global_backend = getattr(inductor_config, "npu_backend", None)
171	if global_backend not in (None, "", "default"):	148	if global_backend not in (None, "", "default"):
172	return global_backend	149	return global_backend
173		150
174	return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")	151	return os.getenv("TORCHINDUCTOR_NPU_BACKEND", "default")
175		152
176		153
		154	+def _resolve_npu_backend_from_wrapper(wrapper) -> str:
		155	+ return _resolve_npu_backend(wrapper.config.get("npu_backend"))
		156	+
		157	+
177	class _NpuBackendScope:	158	class _NpuBackendScope:
178	"""Apply resolved npu backend for one compile invocation and restore env."""	159	"""Apply resolved npu backend for one compile invocation and restore env."""
179		160
	@@ -182,17 +163,24 @@ class _NpuBackendScope:
182	self._old_env = None	163	self._old_env = None
183		164
184	def __enter__(self):	165	def __enter__(self):
185	- self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")	166	+ self._old_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
186	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend	167	+ try:
187	- register_inductor_npu()	168	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self.backend
		169	+ register_inductor_npu()
		170	+ except BaseException:
		171	+ self._restore_backend_env()
		172	+ raise
188	return self	173	return self
189		174
190	def __exit__(self, exc_type, exc, tb):	175	def __exit__(self, exc_type, exc, tb):
		176	+ self._restore_backend_env()
		177	+ return False
		178	+
		179	+ def _restore_backend_env(self):
191	if self._old_env is None:	180	if self._old_env is None:
192	os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)	181	os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
193	else:	182	else:
194	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env	183	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._old_env
195	- return False
196		184
197		185
198	def patch_inductor_wrapper():	186	def patch_inductor_wrapper():
	@@ -207,15 +195,26 @@ def patch_inductor_wrapper():
207	src_call = _TorchCompileInductorWrapper.__call__	195	src_call = _TorchCompileInductorWrapper.__call__
208		196
209	def new_apply_options(self, options: Optional[dict[str, Any]]):	197	def new_apply_options(self, options: Optional[dict[str, Any]]):
		198	+ if (
		199	+ options is not None
		200	+ and "npu_backend" in options
		201	+ and not isinstance(options["npu_backend"], str)
		202	+ ):
		203	+ backend_type = type(options["npu_backend"]).__name__
		204	+ raise RuntimeError(
		205	+ "Unexpected type of attr npu_backend, "
		206	+ f"got {backend_type} should be str"
		207	+ )
		208	+ src_apply_options(self, options)
210	if options is not None and options.get("enable_shape_handling", False):	209	if options is not None and options.get("enable_shape_handling", False):
211	if not is_inductor_npu_initialized():	210	if not is_inductor_npu_initialized():
212	register_inductor_npu()	211	register_inductor_npu()
213	torch_npu._inductor.patch_shape_handling()	212	torch_npu._inductor.patch_shape_handling()
214	- src_apply_options(self, options)
215		213
216	def new_get_config_copy(self) -> dict[str, Any]:	214	def new_get_config_copy(self) -> dict[str, Any]:
217	ori_dict = src_get_config_copy(self)	215	ori_dict = src_get_config_copy(self)
218	- if self is not torch._inductor.config:	216	+ inductor_config = sys.modules.get("torch._inductor.config")
		217	+ if inductor_config is None or self is not inductor_config:
219	return ori_dict	218	return ori_dict
220	NpuBackendType = Literal["default", "mlir", "dvm"]	219	NpuBackendType = Literal["default", "mlir", "dvm"]
221	if "npu_backend" not in ori_dict:	220	if "npu_backend" not in ori_dict:
	@@ -245,6 +244,7 @@ def patch_inductor_wrapper():
245		244
246	def new_init(self, mode, options, dynamic):	245	def new_init(self, mode, options, dynamic):
247	src_init(self, mode, options, dynamic)	246	src_init(self, mode, options, dynamic)
		247	+ _setup_inductor_for_compile(options)
248	backend = _resolve_npu_backend_from_wrapper(self)	248	backend = _resolve_npu_backend_from_wrapper(self)
249	if backend=="mlir":	249	if backend=="mlir":
250	with _NpuBackendScope(backend):	250	with _NpuBackendScope(backend):
	@@ -265,11 +265,12 @@ def patch_inductor_wrapper():
265	_TorchCompileInductorWrapper.apply_options = new_apply_options	265	_TorchCompileInductorWrapper.apply_options = new_apply_options
266	_TorchCompileInductorWrapper.__init__ = new_init	266	_TorchCompileInductorWrapper.__init__ = new_init
267	ConfigModule.get_config_copy = new_get_config_copy	267	ConfigModule.get_config_copy = new_get_config_copy
268	- torch._inductor.config.get_config_copy()
269		268
270		269
271	def patch_dynamo_optimize():	270	def patch_dynamo_optimize():
272	- src_optimize = optimize	271	+ from torch_npu.dynamo import _get_global_npu_backend
		272	+
		273	+ src_optimize = torch._dynamo.optimize
273		274
274	def npu_optimize(args, *kwargs):	275	def npu_optimize(args, *kwargs):
275	backend = None	276	backend = None
	@@ -293,10 +294,18 @@ def patch_dynamo_optimize():
293		294
294		295
295	def patch_user_defined_class_variable():	296	def patch_user_defined_class_variable():
296	- import functools	297	+ from torch._dynamo.variables.torch import (
		298	+ TorchCtxManagerClassVariable,
		299	+ TorchInGraphFunctionVariable,
		300	+ )
		301	+ from torch._dynamo.variables.user_defined import UserDefinedClassVariable
297		302
298	original_method = UserDefinedClassVariable._in_graph_classes	303	original_method = UserDefinedClassVariable._in_graph_classes
299		304
		305	+ class NPUTorchCtxManagerClassVariable(TorchCtxManagerClassVariable):
		306	+ def call_function(self, tx, args, kwargs):
		307	+ return _create_npu_autocast_mode_variable(self.value, args, kwargs)
		308	+
300	@staticmethod	309	@staticmethod
301	@functools.lru_cache(None)	310	@functools.lru_cache(None)
302	def patched_in_graph_classes():	311	def patched_in_graph_classes():
	@@ -305,7 +314,36 @@ def patch_user_defined_class_variable():
305	result.add(torch.npu.Stream)	314	result.add(torch.npu.Stream)
306	return result	315	return result
307		316
		317	+ def UserDefinedClassVariable__new__(cls, value, **kwargs):
		318	+ if value in [
		319	+ torch.npu.amp.autocast,
		320	+ torch_npu.npu.amp.autocast,
		321	+ torch.npu.amp.autocast_mode.autocast,
		322	+ torch_npu.npu.amp.autocast_mode.autocast,
		323	+ ]:
		324	+ return NPUTorchCtxManagerClassVariable(value, **kwargs)
		325	+ if value in [
		326	+ torch.npu.Stream,
		327	+ torch_npu.npu.Stream,
		328	+ torch.npu.streams.Stream,
		329	+ torch_npu.npu.streams.Stream,
		330	+ torch_npu.npu.BoolTensor,
		331	+ torch_npu.npu.ByteTensor,
		332	+ torch_npu.npu.CharTensor,
		333	+ torch_npu.npu.DoubleTensor,
		334	+ torch_npu.npu.FloatTensor,
		335	+ torch_npu.npu.HalfTensor,
		336	+ torch_npu.npu.IntTensor,
		337	+ torch_npu.npu.LongTensor,
		338	+ torch_npu.npu.ShortTensor,
		339	+ torch_npu.npu.BFloat16Tensor,
		340	+ ]:
		341	+ return TorchInGraphFunctionVariable(value, **kwargs)
		342	+ return cls.__new__raw(cls)
		343	+
308	UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes	344	UserDefinedClassVariable._in_graph_classes = patched_in_graph_classes
		345	+ UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__
		346	+ UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__
309		347
310		348
311	def fake_record_stream(self, s):	349	def fake_record_stream(self, s):
	@@ -386,17 +424,213 @@ def patch_event_variable_python_type():
386	torch._dynamo.variables.ctx_manager.EventVariable.python_type = python_type	424	torch._dynamo.variables.ctx_manager.EventVariable.python_type = python_type
387		425
388		426
		427	+def run_once(f):
		428	+ """Run a function successfully only once, waiting for concurrent callers."""
		429	+ condition = threading.Condition()
		430	+
		431	+ @functools.wraps(f)
		432	+ def wrapper(args, *kwargs):
		433	+ thread_id = threading.get_ident()
		434	+ with condition:
		435	+ while wrapper._is_running:
		436	+ if wrapper._running_thread == thread_id:
		437	+ return None
		438	+ condition.wait()
		439	+ if wrapper.has_run:
		440	+ return None
		441	+ wrapper._is_running = True
		442	+ wrapper._running_thread = thread_id
		443	+
		444	+ try:
		445	+ result = f(args, *kwargs)
		446	+ except BaseException:
		447	+ with condition:
		448	+ wrapper._is_running = False
		449	+ wrapper._running_thread = None
		450	+ condition.notify_all()
		451	+ raise
		452	+
		453	+ with condition:
		454	+ wrapper.has_run = True
		455	+ wrapper._is_running = False
		456	+ wrapper._running_thread = None
		457	+ condition.notify_all()
		458	+ return result
		459	+
		460	+ wrapper.has_run = False
		461	+ wrapper._is_running = False
		462	+ wrapper._running_thread = None
		463	+
		464	+ def reset_after_fork():
		465	+ # The parent thread running f may not exist in the child process.
		466	+ nonlocal condition
		467	+ condition = threading.Condition()
		468	+ wrapper._is_running = False
		469	+ wrapper._running_thread = None
		470	+
		471	+ try:
		472	+ os.register_at_fork(after_in_child=reset_after_fork)
		473	+ except AttributeError:
		474	+ pass
		475	+ return wrapper
		476	+
		477	+
		478	+_COMPLETED_DYNAMO_SETUP_STEPS = set()
		479	+
		480	+
		481	+def _run_dynamo_setup_step(name, setup):
		482	+ """Keep successful setup steps idempotent when a later step fails."""
		483	+ if name in _COMPLETED_DYNAMO_SETUP_STEPS:
		484	+ return
		485	+ setup()
		486	+ _COMPLETED_DYNAMO_SETUP_STEPS.add(name)
		487	+
		488	+
		489	+@run_once
		490	+def _dynamo_register_interface_for_device():
		491	+ from torch._dynamo.device_interface import register_interface_for_device
		492	+ from torch_npu.utils._dynamo_device import NpuInterface
		493	+
		494	+ register_interface_for_device("npu", NpuInterface)
		495	+ for i in range(32):
		496	+ register_interface_for_device(f"npu:{i}", NpuInterface)
		497	+
		498	+
		499	+def _find_spec_without_finder(finder, fullname):
		500	+ """Delegate to the remaining meta-path finders without bypassing them."""
		501	+ try:
		502	+ index = sys.meta_path.index(finder)
		503	+ except ValueError:
		504	+ return importlib.util.find_spec(fullname)
		505	+
		506	+ sys.meta_path.pop(index)
		507	+ try:
		508	+ return importlib.util.find_spec(fullname)
		509	+ finally:
		510	+ sys.meta_path.insert(min(index, len(sys.meta_path)), finder)
		511	+
		512	+
		513	+class _DynamoPostImportLoader(importlib.abc.Loader):
		514	+ def __init__(self, loader, finder):
		515	+ self._loader = loader
		516	+ self._finder = finder
		517	+
		518	+ def create_module(self, spec):
		519	+ create_module = getattr(self._loader, "create_module", None)
		520	+ return create_module(spec) if create_module is not None else None
		521	+
		522	+ def exec_module(self, module):
		523	+ self._loader.exec_module(module)
		524	+ _lazy_dynamo_setup()
		525	+ if self._finder in sys.meta_path:
		526	+ sys.meta_path.remove(self._finder)
		527	+
		528	+
		529	+class _DynamoPostImportFinder(importlib.abc.MetaPathFinder):
		530	+ _target = "torch._dynamo"
		531	+
		532	+ def find_spec(self, fullname, path=None, target=None):
		533	+ if fullname != self._target:
		R R rmch13 天前 [P1] `_DynamoPostImportFinder.find_spec` 仅拦截精确模块名，子模块导入不触发 NPU patch `def find_spec(self, fullname, path=None, target=None): if fullname != self._TARGET: # 只匹配 "torch._dynamo" 精确名 return None` 当用户写 `from torch._dynamo import optimize` 或 `import torch._dynamo.backends` 时，CPython 先以 `"torch._dynamo.optimize"` / `"torch._dynamo.backends"` 调用 `find_spec`，此处直接返回 `None`，`_DynamoPostImportLoader.exec_module` 永远不执行，NPU patch 不安装。只有用户恰好写 `import torch._dynamo`（裸模块名）时才触发，实际使用中命中率很低。建议： `if not (fullname == self._TARGET or fullname.startswith(self._TARGET + ".")): return None # 子模块走到这里时只需委托原 loader，无需替换；可加保护： if fullname != self._TARGET: return _find_spec_without_finder(self, fullname) # 不替换 loader，仅避免递归` 或改为：`_TARGET` 的 loader 替换后，在 `exec_module` 内检查是否已有子模块在 sys.modules，若有则补跑 patch。黄桂军 13 天前评论： 13 天前评论：
		534	+ return None
		535	+ spec = _find_spec_without_finder(self, fullname)
		536	+ if spec is not None and spec.loader is not None:
		537	+ spec.loader = _DynamoPostImportLoader(spec.loader, self)
		538	+ return spec
		539	+
		540	+
		541	+def _install_dynamo_post_import_trigger():
		542	+ """Set up NPU integration whenever Dynamo is first imported."""
		543	+ if "torch._dynamo" in sys.modules:
		544	+ _lazy_dynamo_setup()
		545	+ return
		546	+ if not any(isinstance(finder, _DynamoPostImportFinder) for finder in sys.meta_path):
		547	+ sys.meta_path.insert(0, _DynamoPostImportFinder())
		548	+
		549	+
		550	+@run_once
		551	+def add_dynamo_methods_init():
		552	+ steps = (
		553	+ ("device_interface", _dynamo_register_interface_for_device),
		554	+ ("skip_function_variable", patch_SkipFunctionVariable),
		555	+ ("tensor_variable", patch_TensorVariable_call_method),
		556	+ ("user_defined_class_variable", patch_user_defined_class_variable),
		557	+ ("record_stream", patch_record_stream),
		558	+ ("event_variable", patch_event_variable_python_type),
		559	+ ("variable_builder", patch_variable_builder),
		560	+ ("builtin_variable", patch_builtin_variable),
		561	+ )
		562	+ for name, setup in steps:
		563	+ _run_dynamo_setup_step(name, setup)
		564	+
		565	+
		566	+@run_once
		567	+def _inject_inductor_npu_backend_config():
		568	+ """Inject NPU entries into torch._inductor.config on first use."""
		569	+ torch._inductor.config.get_config_copy()
		570	+
		571	+
		572	+@run_once
		573	+def _lazy_dynamo_setup():
		574	+ """Initialize the Dynamo integration on the first graph-capture operation."""
		575	+ add_dynamo_methods_init()
		576	+
		577	+ from torch_npu.dynamo import _register_backends
		578	+ _run_dynamo_setup_step("backends", _register_backends)
		579	+
		580	+ from torch_npu.dynamo.trace_rule import _patch_npu_trace_rules
		581	+ _run_dynamo_setup_step("trace_rules", _patch_npu_trace_rules)
		582	+
		583	+ _run_dynamo_setup_step("dynamo_optimize", patch_dynamo_optimize)
		584	+
		585	+
		586	+@run_once
		587	+def _lazy_inductor_setup():
		588	+ """Initialize NPU Inductor support only for an Inductor-based backend."""
		589	+ register_inductor_npu()
		590	+
		591	+ from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
		592	+ _apply_npugraph_tree_methods()
		593	+
		594	+ _inject_inductor_npu_backend_config()
		595	+
		596	+
		597	+def _setup_inductor_for_compile(options=None):
		R R rmch20 天前这里对进程级 TORCHINDUCTOR_NPU_BACKEND 的 set/call/restore 没有和 run_once 使用同一把锁。两个线程首次分别以 mlir、dvm 调用 torch.compile 时，后进入者可在前一线程 import/_load_backend 读取环境变量前覆盖它；两个 finally 还可能交错恢复，最终残留错误的环境值，导致 _loaded_backend 与调用方选择不一致。建议将选中的 backend 显式传入 registry，或用共享锁覆盖读取旧值、设置、_lazy_inductor_setup 以及恢复的整个区间，并增加 Barrier 控制的并发回归用例。黄桂军 19 天前评论： 19 天前评论：
		598	+ """Initialize the NPU Inductor backend selected for this compile call."""
		599	+ _lazy_dynamo_setup()
		600	+
		601	+ option_backend = options.get("npu_backend") if isinstance(options, dict) else None
		602	+ selected_backend = _resolve_npu_backend(option_backend)
		603	+
		604	+ old_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
		605	+ if selected_backend not in (None, "", "default"):
		606	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = selected_backend
		607	+ try:
		608	+ _lazy_inductor_setup()
		609	+ finally:
		610	+ if old_backend is None:
		611	+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
		612	+ else:
		613	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = old_backend
		614	+ return selected_backend
		615	+
		616	+
		617	+@run_once
		618	+def install_npugraph_mark_step_trigger():
		619	+ """Expose the public NPUGraph step API without importing compiler internals."""
		620	+ def npugraph_mark_step_begin():
		621	+ from torch_npu.npu._graph_tree_state import mark_step_begin
		622	+ return mark_step_begin()
		623	+
		624	+ torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
		625	+
		626	+
389	def add_dynamo_methods():	627	def add_dynamo_methods():
390	- UserDefinedClassVariable.__new__raw = UserDefinedClassVariable.__new__	628	+ from torch_npu.dynamo import _install_lazy_torchair
391	- UserDefinedClassVariable.__new__ = UserDefinedClassVariable__new__	629	+
392	- SkipFunctionVariable.__new__raw = SkipFunctionVariable.__new__	630	+ _install_lazy_torchair()
393	- SkipFunctionVariable.__new__ = SkipFunctionVariable__new__	631	+ _install_dynamo_post_import_trigger()
394	- TensorVariable.call_method_raw = TensorVariable.call_method	632	+ if "npugraph_ex" not in sys.modules:
395	- TensorVariable.call_method = TensorVariable_call_method	633	+ from torch_npu.dynamo import _LazyNpuGraphEx
396	- patch_dynamo_optimize()	634	+ sys.modules["npugraph_ex"] = _LazyNpuGraphEx("npugraph_ex")
397	patch_inductor_wrapper()	635	patch_inductor_wrapper()
398	- patch_user_defined_class_variable()	636	+ install_npugraph_mark_step_trigger()
399	- patch_record_stream()
400	- patch_event_variable_python_type()
401	- patch_variable_builder()
402	- patch_builtin_variable()

Mtorch_npu/utils/_graph_tree.py+8-5

	@@ -22,7 +22,7 @@ from torch._dynamo.backends.cudagraphs import (
22	get_stack_traces,	22	get_stack_traces,
23	)	23	)
24	from torch._dynamo.backends.debugging import boxed_nop	24	from torch._dynamo.backends.debugging import boxed_nop
25	-from torch._dynamo.backends.registry import register_backend	25	+from torch._dynamo.backends.registry import _COMPILER_FNS, register_backend
26	from torch._inductor import config	26	from torch._inductor import config
27	from torch._inductor.compile_fx import (	27	from torch._inductor.compile_fx import (
28	get_input_idxs_to_check,	28	get_input_idxs_to_check,
	@@ -58,7 +58,7 @@ log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
58		58
59		59
60	def npugraph_mark_step_begin():	60	def npugraph_mark_step_begin():
61	- from torch_npu.npu._graph_tree import mark_step_begin	61	+ from torch_npu.npu._graph_tree_state import mark_step_begin
62	mark_step_begin()	62	mark_step_begin()
63		63
64		64
	@@ -373,9 +373,9 @@ class NpugraphsBackend:
373		373
374	@staticmethod	374	@staticmethod
375	def reset():	375	def reset():
376	- from torch_npu.npu._graph_tree import reset_npugraph_trees	376	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
377		377
378	- reset_npugraph_trees()	378	+ _npugraphs_backend_entrypoint.reset()
379		379
380	@staticmethod	380	@staticmethod
381	def __call__(model, inputs):	381	def __call__(model, inputs):
	@@ -385,7 +385,10 @@ class NpugraphsBackend:
385	def _apply_npugraph_tree_methods():	385	def _apply_npugraph_tree_methods():
386	# aot_npugraphs only applies graphs to the graph. It is also helpful	386	# aot_npugraphs only applies graphs to the graph. It is also helpful
387	# for debugging and can serve as a perf baseline.	387	# for debugging and can serve as a perf baseline.
388	- register_backend(name="npugraphs", compiler_fn=NpugraphsBackend())	388	+ if "npugraphs" not in _COMPILER_FNS:
		389	+ from torch_npu.dynamo import _npugraphs_backend_entrypoint
		390	+
		391	+ register_backend(name="npugraphs", compiler_fn=_npugraphs_backend_entrypoint)
389	torch._inductor.compile_fx.cudagraphify = npugraphify	392	torch._inductor.compile_fx.cudagraphify = npugraphify
390	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes	393	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes
391	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin	394	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin

Mtorch_npu/utils/_inductor.py+8-7

	@@ -1,6 +1,6 @@
1	-from typing import Optional
黄桂军13 天前已过期 lijing:为什么把实现拆出去黄桂军 13 天前评论： 13 天前评论：
2	import operator	1	import operator
3	from functools import reduce	2	from functools import reduce
		3	+from typing import Optional
4		4
5	import torch	5	import torch
6	from torch._prims_common import TensorLike	6	from torch._prims_common import TensorLike
	@@ -26,6 +26,7 @@ def _max_unpoolnd_patch(
26	output.reshape(-1), [indices_flat], self.reshape(-1), accumulate=False	26	output.reshape(-1), [indices_flat], self.reshape(-1), accumulate=False
27	).view(output.shape)	27	).view(output.shape)
28		28
		29	+
29	torch._decomp.decompositions._max_unpoolnd = _max_unpoolnd_patch	30	torch._decomp.decompositions._max_unpoolnd = _max_unpoolnd_patch
30		31
31		32
	@@ -54,8 +55,8 @@ def patch_register_philox_rand():
54	def get_register_philox_rand_patch():	55	def get_register_philox_rand_patch():
55	name = "philox_rand"	56	name = "philox_rand"
56	schema = "(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)" # noqa: B950	57	schema = "(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)" # noqa: B950
57	-	58	+
58	-	59	+
59	def _philox_rand_meta(	60	def _philox_rand_meta(
60	shape: torch.Size,	61	shape: torch.Size,
61	seed: torch.Tensor,	62	seed: torch.Tensor,
	@@ -71,7 +72,7 @@ def patch_register_philox_rand():
71	offset = philox_rand_offset_meta(shape)	72	offset = philox_rand_offset_meta(shape)
72	return (random_values, offset)	73	return (random_values, offset)
73		74
74	-	75	+
75	def _philox_rand(	76	def _philox_rand(
76	shape: torch.Size,	77	shape: torch.Size,
77	seed: torch.Tensor,	78	seed: torch.Tensor,
	@@ -85,13 +86,13 @@ def patch_register_philox_rand():
85	else:	86	else:
86	devices = [device]	87	devices = [device]
87		88
88	- with torch.random.fork_rng(devices, device_type="npu"):	89	+ with torch.random.fork_rng(devices, device_type="npu"):
89	CUDARngStateHelper.set_torch_state_tensor(seed, offset)	90	CUDARngStateHelper.set_torch_state_tensor(seed, offset)
90	random_values = torch.rand(shape, device=device, dtype=dtype)	91	random_values = torch.rand(shape, device=device, dtype=dtype)
91		92
92	return random_values, philox_rand_offset(shape)	93	return random_values, philox_rand_offset(shape)
93		94
94	-	95	+
95	register_rng_prim(	96	register_rng_prim(
96	name=name,	97	name=name,
97	schema=schema,	98	schema=schema,
	@@ -245,4 +246,4 @@ patch_register_run_and_save_rng_state_op()
245	patch_register_run_with_rng_state_op()	246	patch_register_run_with_rng_state_op()
246	patch_philox_rand_offset()	247	patch_philox_rand_offset()
247	patch_register_philox_rand()	248	patch_register_philox_rand()
248	-patch_rng_prims_device()	249	+patch_rng_prims_device()