feat: Add ACLGraph update plans-pytorch-AtomGit

Atest/_inductor/test_aclgraph_update_plan_compile.py+737-0

@@ -0,0 +1,737 @@
	1	+import json
	2	+import os
	3	+import subprocess
	4	+import sys
	5	+from unittest import mock
	6	+
	7	+import torch
	8	+from torch._inductor import config
	9	+from torch._inductor.codegen.common import IndentedBuffer
	10	+from torch._inductor.virtualized import V
	11	+from torch_npu._inductor._aclgraph_update_plan import (
	12	+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
	13	+ append_inductor_aclgraph_update_plan_for_codegen_node,
	14	+)
	15	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.wrapper import (
	16	+ NpuMlirSubgraphPythonWrapperCodegen,
	17	+ NpuMlirWrapperCodeGen,
	18	+)
	19	+from torch_npu._inductor.codegen.wrapper import (
	20	+ _NPUKernelCodegenMixin,
	21	+ NPUSubgraphPythonWrapperCodegen,
	22	+)
	23	+from torch.testing._internal.common_utils import run_tests
	24	+
	25	+import torch_npu
	26	+import torch_npu._inductor
	27	+from torch_npu.testing.common_utils import SupportedDevices
	28	+
	29	+from testutils import TestUtils
	30	+
	31	+
	32	+def _make_ifa_inputs():
	33	+ q = torch.randn(1, 32, 1, 128, dtype=torch.float16, device="npu")
	34	+ k = torch.randn(1, 32, 1, 128, dtype=torch.float16, device="npu")
	35	+ v = torch.randn(1, 32, 1, 128, dtype=torch.float16, device="npu")
	36	+ return q, k, v
	37	+
	38	+
	39	+def _ifa_with_const_actual_seq_lengths(query, key, value):
	40	+ out, _ = torch_npu.npu_fused_infer_attention_score(
	41	+ query,
	42	+ key,
	43	+ value,
	44	+ num_heads=32,
	45	+ input_layout="BNSD",
	46	+ scale=128.0,
	47	+ pre_tokens=65535,
	48	+ next_tokens=65535,
	49	+ softmax_lse_flag=False,
	50	+ actual_seq_lengths=[37],
	51	+ )
	52	+ return out
	53	+
	54	+
	55	+def _ifa_with_actual_seq_lengths_kv(query, key, value):
	56	+ out, _ = torch_npu.npu_fused_infer_attention_score(
	57	+ query,
	58	+ key,
	59	+ value,
	60	+ num_heads=32,
	61	+ input_layout="BNSD",
	62	+ scale=128.0,
	63	+ pre_tokens=65535,
	64	+ next_tokens=65535,
	65	+ softmax_lse_flag=False,
	66	+ actual_seq_lengths=[37],
	67	+ actual_seq_lengths_kv=[1],
	68	+ )
	69	+ return out
	70	+
	71	+
	72	+def _ifa_with_runtime_actual_seq_lengths(query, key, value, actual_seq_lengths):
	73	+ out, _ = torch_npu.npu_fused_infer_attention_score(
	74	+ query,
	75	+ key,
	76	+ value,
	77	+ num_heads=32,
	78	+ input_layout="BNSD",
	79	+ scale=128.0,
	80	+ pre_tokens=65535,
	81	+ next_tokens=65535,
	82	+ softmax_lse_flag=False,
	83	+ actual_seq_lengths=actual_seq_lengths,
	84	+ )
	85	+ return out
	86	+
	87	+
	88	+def _ifa_v2_with_const_actual_seq_qlen(query, key, value):
	89	+ out, _ = torch_npu.npu_fused_infer_attention_score_v2(
	90	+ query,
	91	+ key,
	92	+ value,
	93	+ num_query_heads=32,
	94	+ input_layout="BNSD",
	95	+ softmax_scale=128.0,
	96	+ pre_tokens=65535,
	97	+ next_tokens=65535,
	98	+ return_softmax_lse=False,
	99	+ actual_seq_qlen=[1],
	100	+ )
	101	+ return out
	102	+
	103	+
	104	+def _two_ifa_with_const_actual_seq_lengths(query, key, value):
	105	+ out1, _ = torch_npu.npu_fused_infer_attention_score(
	106	+ query,
	107	+ key,
	108	+ value,
	109	+ num_heads=32,
	110	+ input_layout="BNSD",
	111	+ scale=128.0,
	112	+ pre_tokens=65535,
	113	+ next_tokens=65535,
	114	+ softmax_lse_flag=False,
	115	+ actual_seq_lengths=[37],
	116	+ )
	117	+ out2, _ = torch_npu.npu_fused_infer_attention_score(
	118	+ query,
	119	+ key,
	120	+ out1,
	121	+ num_heads=32,
	122	+ input_layout="BNSD",
	123	+ scale=128.0,
	124	+ pre_tokens=65535,
	125	+ next_tokens=65535,
	126	+ softmax_lse_flag=False,
	127	+ actual_seq_lengths=[41],
	128	+ )
	129	+ return out2
	130	+
	131	+
	132	+def _run_and_get_code_without_reset(fn, *args):
	133	+ from torch._inductor.graph import GraphLowering
	134	+
	135	+ source_codes = []
	136	+
	137	+ def save_output_code(code):
	138	+ source_codes.append(code)
	139	+
	140	+ with mock.patch.object(GraphLowering, "save_output_code", save_output_code):
	141	+ result = fn(*args)
	142	+ return result, source_codes
	143	+
	144	+
	145	+def _compiled_code(
	146	+ fn,
	147	+ *args,
	148	+ cudagraphs=True,
	149	+ cudagraph_trees=True,
	150	+ graph_partition=False,
	151	+ npu_backend=None,
	152	+):
	153	+ torch._dynamo.reset()
	154	+ old_cudagraphs = config.triton.cudagraphs
	155	+ old_cudagraph_trees = config.triton.cudagraph_trees
	156	+ old_force_disable_caches = config.force_disable_caches
	157	+ old_graph_partition = config.graph_partition
	158	+ try:
	159	+ config.triton.cudagraphs = cudagraphs
	160	+ config.triton.cudagraph_trees = cudagraph_trees
	161	+ config.force_disable_caches = True
	162	+ config.graph_partition = graph_partition
	163	+ options = {"npu_backend": npu_backend} if npu_backend is not None else None
	164	+ compiled = torch.compile(
	165	+ fn,
	166	+ backend="inductor",
	167	+ fullgraph=True,
	168	+ options=options,
	169	+ )
	170	+ _, codes = _run_and_get_code_without_reset(compiled, *args)
	171	+ finally:
	172	+ config.triton.cudagraphs = old_cudagraphs
	173	+ config.triton.cudagraph_trees = old_cudagraph_trees
	174	+ config.force_disable_caches = old_force_disable_caches
	175	+ config.graph_partition = old_graph_partition
	176	+ torch._dynamo.reset()
	177	+ return "\n".join(codes)
	178	+
	179	+
	180	+def _compiled_code_with_backend_in_subprocess(npu_backend):
	181	+ script = f"""
	182	+import json
	183	+import torch
	184	+
	185	+from test_aclgraph_update_plan_compile import (
	186	+ _compiled_code,
	187	+ _ifa_with_const_actual_seq_lengths,
	188	+ _make_ifa_inputs,
	189	+)
	190	+
	191	+torch.npu.set_device(0)
	192	+code = _compiled_code(
	193	+ _ifa_with_const_actual_seq_lengths,
	194	+ *_make_ifa_inputs(),
	195	+ npu_backend={npu_backend!r},
	196	+)
	197	+print("ACLGRAPH_CODE_BEGIN")
	198	+print(json.dumps(code))
	199	+print("ACLGRAPH_CODE_END")
	200	+"""
	201	+ try:
	202	+ result = subprocess.run(
	203	+ [sys.executable, "-c", script],
	204	+ env=os.environ.copy(),
	205	+ stdout=subprocess.PIPE,
	206	+ stderr=subprocess.PIPE,
	207	+ text=True,
	208	+ check=True,
	209	+ )
	210	+ except subprocess.CalledProcessError as exc:
	211	+ raise AssertionError(
	212	+ f"{npu_backend} subprocess compile failed.\n"
	213	+ f"stdout:\n{exc.stdout}\n"
	214	+ f"stderr:\n{exc.stderr}"
	215	+ ) from exc
	216	+ begin = "ACLGRAPH_CODE_BEGIN"
	217	+ end = "ACLGRAPH_CODE_END"
	218	+ if begin not in result.stdout or end not in result.stdout:
	219	+ raise AssertionError(
	220	+ f"Failed to collect generated code from {npu_backend} subprocess.\n"
	221	+ f"stdout:\n{result.stdout}\n"
	222	+ f"stderr:\n{result.stderr}"
	223	+ )
	224	+ payload = result.stdout.split(begin, 1)[1].split(end, 1)[0].strip()
	225	+ return json.loads(payload)
	226	+
	227	+
	228	+class TestACLGraphUpdatePlanCompile(TestUtils):
	229	+
	230	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	231	+ def test_ifa_no_graph_partition_codegen_attaches_plan_to_call_function(self):
	232	+ torch.npu.set_device(0)
	233	+ torch._dynamo.reset()
	234	+
	235	+ code = _compiled_code(
	236	+ _ifa_with_const_actual_seq_lengths,
	237	+ *_make_ifa_inputs(),
	238	+ graph_partition=False,
	239	+ )
	240	+
	241	+ self.assertIn("_torch_npu_aclgraph_update_plan", code)
	242	+ self.assertIn("def call(args):", code)
	243	+ self.assertIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	244	+ self.assertEqual(code.count(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}"), 1)
	245	+ self.assertNotIn("def partition_0(args):", code)
	246	+ self.assertNotIn(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	247	+ self.assertNotIn(f"\n{ACLGRAPH_UPDATE_PLAN_GLOBAL} = ", code)
	248	+ self.assertIn("npu_fused_infer_attention_score.default", code)
	249	+ self.assertIn("actual_seq_lengths", code)
	250	+ self.assertIn("'value': 37", code)
	251	+
	252	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	253	+ def test_ifa_codegen_plan_survives_generated_code_reload(self):
	254	+ torch.npu.set_device(0)
	255	+ torch._dynamo.reset()
	256	+
	257	+ code = _compiled_code(_ifa_with_const_actual_seq_lengths, *_make_ifa_inputs())
	258	+ namespace = {}
	259	+ exec(compile(code, "<aclgraph_update_plan_test>", "exec"), namespace)
	260	+
	261	+ plan = getattr(namespace["call"], ACLGRAPH_UPDATE_PLAN_GLOBAL)
	262	+ self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score.default")
	263	+ self.assertEqual(
	264	+ plan[0]["updates"]["actual_seq_lengths"],
	265	+ {
	266	+ "kind": "list",
	267	+ "items": [{"kind": "constant", "value": 37}],
	268	+ },
	269	+ )
	270	+ if "actual_seq_lengths_kv" in plan[0]["updates"]:
	271	+ self.assertEqual(plan[0]["updates"]["actual_seq_lengths_kv"], {"kind": "none"})
	272	+
	273	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	274	+ def test_ifa_codegen_emits_multiple_actual_seq_keys(self):
	275	+ torch.npu.set_device(0)
	276	+ torch._dynamo.reset()
	277	+
	278	+ code = _compiled_code(_ifa_with_actual_seq_lengths_kv, *_make_ifa_inputs())
	279	+
	280	+ self.assertIn("actual_seq_lengths", code)
	281	+ self.assertIn("actual_seq_lengths_kv", code)
	282	+ self.assertIn("'value': 37", code)
	283	+ self.assertIn("'value': 1", code)
	284	+
	285	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	286	+ def test_ifa_codegen_recompiles_when_guarded_actual_seq_list_changes(self):
	287	+ torch.npu.set_device(0)
	288	+ torch._dynamo.reset()
	289	+
	290	+ old_cudagraphs = config.triton.cudagraphs
	291	+ old_cudagraph_trees = config.triton.cudagraph_trees
	292	+ old_force_disable_caches = config.force_disable_caches
	293	+ try:
	294	+ config.triton.cudagraphs = True
	295	+ config.triton.cudagraph_trees = True
	296	+ config.force_disable_caches = True
	297	+ compiled = torch.compile(
	298	+ _ifa_with_runtime_actual_seq_lengths,
	299	+ backend="inductor",
	300	+ fullgraph=True,
	301	+ )
	302	+ code_37 = "\n".join(
	303	+ _run_and_get_code_without_reset(compiled, *_make_ifa_inputs(), [37])[1]
	304	+ )
	305	+ code_41 = "\n".join(
	306	+ _run_and_get_code_without_reset(compiled, *_make_ifa_inputs(), [41])[1]
	307	+ )
	308	+ finally:
	309	+ config.triton.cudagraphs = old_cudagraphs
	310	+ config.triton.cudagraph_trees = old_cudagraph_trees
	311	+ config.force_disable_caches = old_force_disable_caches
	312	+ torch._dynamo.reset()
	313	+
	314	+ self.assertIn("'value': 37", code_37)
	315	+ self.assertNotIn("'value': 37", code_41)
	316	+ self.assertTrue(
	317	+ "'value': 41" in code_41
	318	+ or ("'kind': 'input'" in code_41 and "'index': 3" in code_41)
	319	+ )
	320	+
	321	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	322	+ def test_ifa_v2_codegen_emits_aclgraph_update_plan(self):
	323	+ torch.npu.set_device(0)
	324	+ torch._dynamo.reset()
	325	+
	326	+ code = _compiled_code(_ifa_v2_with_const_actual_seq_qlen, *_make_ifa_inputs())
	327	+
	328	+ self.assertIn("_torch_npu_aclgraph_update_plan", code)
	329	+ self.assertIn("npu_fused_infer_attention_score_v2.default", code)
	330	+ self.assertIn("actual_seq_qlen", code)
	331	+ self.assertIn("'value': 1", code)
	332	+
	333	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	334	+ def test_ifa_graph_partition_codegen_attaches_plan_to_partition_function(self):
	335	+ class Graph:
	336	+ cpp_wrapper = False
	337	+ disable_cudagraphs_reason = None
	338	+
	339	+ plan = [
	340	+ {
	341	+ "op": "npu_fused_infer_attention_score.default",
	342	+ "updates": {
	343	+ "actual_seq_lengths": {
	344	+ "kind": "list",
	345	+ "items": [{"kind": "constant", "value": 37}],
	346	+ }
	347	+ },
	348	+ }
	349	+ ]
	350	+
	351	+ old_cudagraphs = config.triton.cudagraphs
	352	+ old_cudagraph_trees = config.triton.cudagraph_trees
	353	+ old_graph_partition = config.graph_partition
	354	+ try:
	355	+ config.triton.cudagraphs = True
	356	+ config.triton.cudagraph_trees = True
	357	+ config.graph_partition = True
	358	+
	359	+ wrapper = object.__new__(NPUSubgraphPythonWrapperCodegen)
	360	+ wrapper.launcher_fn_name = "partition_0"
	361	+ with V.set_graph_handler(Graph()):
	362	+ wrapper.torch_npu_aclgraph_update_plan = plan
	363	+ result = IndentedBuffer()
	364	+ wrapper.generate_after_suffix(result)
	365	+ finally:
	366	+ config.triton.cudagraphs = old_cudagraphs
	367	+ config.triton.cudagraph_trees = old_cudagraph_trees
	368	+ config.graph_partition = old_graph_partition
	369	+
	370	+ code = result.getvalue()
	371	+ self.assertIn(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	372	+ self.assertEqual(code.count(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}"), 1)
	373	+ self.assertNotIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	374	+ self.assertNotIn(f"\n{ACLGRAPH_UPDATE_PLAN_GLOBAL} = ", code)
	375	+ self.assertIn(repr(plan), code)
	376	+
	377	+ def test_mlir_dvm_wrapper_appends_aclgraph_update_plan_for_extern_kernel(self):
	378	+ class Graph:
	379	+ cpp_wrapper = False
	380	+ disable_cudagraphs_reason = None
	381	+
	382	+ class Arg:
	383	+ def __init__(self, name):
	384	+ self.name = name
	385	+
	386	+ class Schema:
	387	+ arguments = [
	388	+ Arg("query"),
	389	+ Arg("key"),
	390	+ Arg("value"),
	391	+ Arg("num_heads"),
	392	+ Arg("input_layout"),
	393	+ Arg("actual_seq_lengths"),
	394	+ ]
	395	+
	396	+ class Target:
	397	+ __name__ = "npu_fused_infer_attention_score.default"
	398	+ _schema = Schema()
	399	+
	400	+ class Value:
	401	+ def __init__(self, name):
	402	+ self.name = name
	403	+
	404	+ def get_name(self):
	405	+ return self.name
	406	+
	407	+ class Kernel:
	408	+ op_overload = Target()
	409	+ inputs = [Value("arg0_1"), Value("arg1_1"), Value("arg2_1")]
	410	+ constant_args = [32, "BNSD", [37]]
	411	+ kwargs = {}
	412	+ layout = object()
	413	+
	414	+ def get_name(self):
	415	+ return "buf0"
	416	+
	417	+ def get_origin_node(self):
	418	+ return None
	419	+
	420	+ def get_kernel_name(self):
	421	+ return "torch.ops.npu.npu_fused_infer_attention_score.default"
	422	+
	423	+ old_cudagraphs = config.triton.cudagraphs
	424	+ old_cudagraph_trees = config.triton.cudagraph_trees
	425	+ old_graph_partition = config.graph_partition
	426	+ try:
	427	+ config.triton.cudagraphs = True
	428	+ config.triton.cudagraph_trees = True
	429	+ config.graph_partition = False
	430	+
	431	+ wrapper = object.__new__(NpuMlirWrapperCodeGen)
	432	+ wrapper.launcher_fn_name = "call"
	433	+ wrapper.declare = ""
	434	+ wrapper.ending = ""
	435	+ wrapper.supports_intermediate_hooks = False
	436	+ wrapper.get_graph_input_names = lambda: ["arg0_1", "arg1_1", "arg2_1"]
	437	+ wrapper.get_graph_inputs = lambda: {}
	438	+ wrapper.writeline = lambda line: None
	439	+
	440	+ with V.set_graph_handler(Graph()):
	441	+ append_inductor_aclgraph_update_plan_for_codegen_node(wrapper, Kernel())
	442	+ result = IndentedBuffer()
	443	+ wrapper.generate_after_suffix(result)
	444	+ finally:
	445	+ config.triton.cudagraphs = old_cudagraphs
	446	+ config.triton.cudagraph_trees = old_cudagraph_trees
	447	+ config.graph_partition = old_graph_partition
	448	+
	449	+ self.assertEqual(
	450	+ wrapper.torch_npu_aclgraph_update_plan,
	451	+ [
	452	+ {
	453	+ "op": "npu_fused_infer_attention_score.default",
	454	+ "updates": {
	455	+ "actual_seq_lengths": {
	456	+ "kind": "list",
	457	+ "items": [{"kind": "constant", "value": 37}],
	458	+ }
	459	+ },
	460	+ }
	461	+ ],
	462	+ )
	463	+ code = result.getvalue()
	464	+ self.assertIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	465	+ self.assertIn(repr(wrapper.torch_npu_aclgraph_update_plan), code)
	466	+
	467	+ def test_mlir_dvm_subgraph_wrapper_emits_aclgraph_update_plan(self):
	468	+ class Graph:
	469	+ cpp_wrapper = False
	470	+ disable_cudagraphs_reason = None
	471	+
	472	+ plan = [
	473	+ {
	474	+ "op": "npu_fused_infer_attention_score.default",
	475	+ "updates": {
	476	+ "actual_seq_lengths": {
	477	+ "kind": "list",
	478	+ "items": [{"kind": "constant", "value": 37}],
	479	+ }
	480	+ },
	481	+ }
	482	+ ]
	483	+
	484	+ old_cudagraphs = config.triton.cudagraphs
	485	+ old_cudagraph_trees = config.triton.cudagraph_trees
	486	+ old_graph_partition = config.graph_partition
	487	+ try:
	488	+ config.triton.cudagraphs = True
	489	+ config.triton.cudagraph_trees = True
	490	+ config.graph_partition = True
	491	+
	492	+ with V.set_graph_handler(Graph()):
	493	+ wrapper = object.__new__(NpuMlirSubgraphPythonWrapperCodegen)
	494	+ wrapper.launcher_fn_name = "partition_0"
	495	+ wrapper.subgraph_name = "partition_0"
	496	+ wrapper.torch_npu_aclgraph_update_plan = plan
	497	+ result = IndentedBuffer()
	498	+ wrapper.generate_after_suffix(result)
	499	+ finally:
	500	+ config.triton.cudagraphs = old_cudagraphs
	501	+ config.triton.cudagraph_trees = old_cudagraph_trees
	502	+ config.graph_partition = old_graph_partition
	503	+
	504	+ code = result.getvalue()
	505	+ self.assertIn(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	506	+ self.assertNotIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
	507	+ self.assertIn(repr(plan), code)
	508	+
	509	+ def test_mlir_dvm_wrapper_does_not_inherit_default_npu_codegen_mixin(self):
	510	+ self.assertFalse(issubclass(NpuMlirWrapperCodeGen, _NPUKernelCodegenMixin))
	511	+ self.assertFalse(issubclass(NpuMlirSubgraphPythonWrapperCodegen, _NPUKernelCodegenMixin))
	512	+
	513	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	514	+ def test_ifa_codegen_preserves_multiple_plan_entry_order(self):
	515	+ torch.npu.set_device(0)
	516	+ torch._dynamo.reset()
	517	+
	518	+ code = _compiled_code(
	519	+ _two_ifa_with_const_actual_seq_lengths,
	520	+ *_make_ifa_inputs(),
	521	+ )
	522	+
	523	+ self.assertGreaterEqual(
	524	+ code.count("'op': 'npu_fused_infer_attention_score.default'"),
	525	+ 2,
	526	+ )
	527	+ first_update = code.find("'value': 37")
	528	+ second_update = code.find("'value': 41")
	529	+ self.assertGreaterEqual(first_update, 0)
	530	+ self.assertGreater(second_update, first_update)
	531	+
	532	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	533	+ def test_ifa_codegen_skips_aclgraph_update_plan_without_cudagraphs(self):
	534	+ torch.npu.set_device(0)
	535	+ torch._dynamo.reset()
	536	+
	537	+ code = _compiled_code(
	538	+ _ifa_with_const_actual_seq_lengths,
	539	+ *_make_ifa_inputs(),
	540	+ cudagraphs=False,
	541	+ )
	542	+
	543	+ self.assertNotIn("_torch_npu_aclgraph_update_plan", code)
	544	+
	545	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	546	+ def test_ifa_codegen_skips_aclgraph_update_plan_without_cudagraph_trees(self):
	547	+ torch.npu.set_device(0)
	548	+ torch._dynamo.reset()
	549	+
	550	+ code = _compiled_code(
	551	+ _ifa_with_const_actual_seq_lengths,
	552	+ *_make_ifa_inputs(),
	553	+ cudagraph_trees=False,
	554	+ )
	555	+
	556	+ self.assertNotIn("_torch_npu_aclgraph_update_plan", code)
	557	+
	558	+ def test_wrapper_plan_gate_respects_cudagraph_disable_reason(self):
	559	+ from torch._inductor.virtualized import V
	560	+ from torch_npu._inductor._aclgraph_update_plan.codegen import (
	561	+ should_generate_inductor_aclgraph_update_plan,
	562	+ )
	563	+
	564	+ class Graph:
	565	+ cpp_wrapper = False
	566	+ disable_cudagraphs_reason = "unsupported"
	567	+
	568	+ old_cudagraphs = config.triton.cudagraphs
	569	+ old_cudagraph_trees = config.triton.cudagraph_trees
	570	+ try:
	571	+ config.triton.cudagraphs = True
	572	+ config.triton.cudagraph_trees = True
	573	+ with V.set_graph_handler(Graph()):
	574	+ self.assertFalse(should_generate_inductor_aclgraph_update_plan())
	575	+ finally:
	576	+ config.triton.cudagraphs = old_cudagraphs
	577	+ config.triton.cudagraph_trees = old_cudagraph_trees
	578	+
	579	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	580	+ def test_ifa_cudagraph_tree_receives_aclgraph_update_plan(self):
	581	+ torch.npu.set_device(0)
	582	+ torch._dynamo.reset()
	583	+
	584	+ import torch_npu.npu._graph_tree as graph_tree
	585	+
	586	+ old_cudagraphs = config.triton.cudagraphs
	587	+ old_cudagraph_trees = config.triton.cudagraph_trees
	588	+ old_force_disable_caches = config.force_disable_caches
	589	+ old_slow_path_asserts = config.triton.slow_path_cudagraph_asserts
	590	+ original_update = graph_tree.update_aclgraph_records_for_graph
	591	+ seen_plans = []
	592	+
	593	+ def collect_plan(plan, graph, inputs):
	594	+ seen_plans.append(plan)
	595	+ return original_update(plan, graph, inputs)
	596	+
	597	+ try:
	598	+ config.triton.cudagraphs = True
	599	+ config.triton.cudagraph_trees = True
	600	+ config.force_disable_caches = True
	601	+ config.triton.slow_path_cudagraph_asserts = False
	602	+ graph_tree.update_aclgraph_records_for_graph = collect_plan
	603	+
	604	+ compiled = torch.compile(
	605	+ _ifa_with_const_actual_seq_lengths,
	606	+ backend="inductor",
	607	+ fullgraph=True,
	608	+ )
	609	+ inputs = _make_ifa_inputs()
	610	+ expected = _ifa_with_const_actual_seq_lengths(*inputs)
	611	+ actual = compiled(*inputs)
	612	+ torch.testing.assert_close(
	613	+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
	614	+ )
	615	+
	616	+ inputs = _make_ifa_inputs()
	617	+ expected = _ifa_with_const_actual_seq_lengths(*inputs)
	618	+ actual = compiled(*inputs)
	619	+ torch.testing.assert_close(
	620	+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
	621	+ )
	622	+ finally:
	623	+ graph_tree.update_aclgraph_records_for_graph = original_update
	624	+ config.triton.cudagraphs = old_cudagraphs
	625	+ config.triton.cudagraph_trees = old_cudagraph_trees
	626	+ config.force_disable_caches = old_force_disable_caches
	627	+ config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts
	628	+ torch._dynamo.reset()
	629	+
	630	+ self.assertTrue(seen_plans)
	631	+ self.assertTrue(any(plan for plan in seen_plans))
	632	+ plan = next(plan for plan in seen_plans if plan)
	633	+ self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score.default")
	634	+ self.assertEqual(
	635	+ plan[0]["updates"]["actual_seq_lengths"],
	636	+ {"kind": "list", "items": [{"kind": "constant", "value": 37}]},
	637	+ )
	638	+
	639	+ def test_npugraphify_keeps_aclgraph_update_plan_on_callable_attribute(self):
	640	+ import torch_npu.npu._graph_tree as graph_tree
	641	+
	642	+ expected_plan = [{"op": "test.op", "updates": {}}]
	643	+
	644	+ def model(args):
	645	+ return args
	646	+
	647	+ setattr(model, ACLGRAPH_UPDATE_PLAN_GLOBAL, expected_plan)
	648	+
	649	+ captured = {}
	650	+
	651	+ def fake_add_function(args, *kwargs):
	652	+ captured["arg_count"] = len(args)
	653	+ captured["model"] = args[0]
	654	+ return lambda inputs: inputs, []
	655	+
	656	+ manager = mock.Mock()
	657	+ manager.add_function.side_effect = fake_add_function
	658	+ with mock.patch(
	659	+ "torch_npu.npu._graph_tree.get_container",
	660	+ return_value=mock.Mock(get_tree_manager=mock.Mock(return_value=manager)),
	661	+ ):
	662	+ graph_tree.npugraphify(
	663	+ model,
	664	+ [],
	665	+ device_index=0,
	666	+ is_backward=False,
	667	+ is_inference=True,
	668	+ )
	669	+
	670	+ self.assertEqual(captured["arg_count"], 8)
	671	+ self.assertIs(
	672	+ getattr(captured["model"], ACLGRAPH_UPDATE_PLAN_GLOBAL),
	673	+ expected_plan,
	674	+ )
	675	+
	676	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	677	+ def test_ifa_v2_cudagraph_tree_receives_aclgraph_update_plan(self):
	678	+ torch.npu.set_device(0)
	679	+ torch._dynamo.reset()
	680	+
	681	+ import torch_npu.npu._graph_tree as graph_tree
	682	+
	683	+ old_cudagraphs = config.triton.cudagraphs
	684	+ old_cudagraph_trees = config.triton.cudagraph_trees
	685	+ old_force_disable_caches = config.force_disable_caches
	686	+ old_slow_path_asserts = config.triton.slow_path_cudagraph_asserts
	687	+ original_update = graph_tree.update_aclgraph_records_for_graph
	688	+ seen_plans = []
	689	+
	690	+ def collect_plan(plan, graph, inputs):
	691	+ seen_plans.append(plan)
	692	+ return original_update(plan, graph, inputs)
	693	+
	694	+ try:
	695	+ config.triton.cudagraphs = True
	696	+ config.triton.cudagraph_trees = True
	697	+ config.force_disable_caches = True
	698	+ config.triton.slow_path_cudagraph_asserts = False
	699	+ graph_tree.update_aclgraph_records_for_graph = collect_plan
	700	+
	701	+ compiled = torch.compile(
	702	+ _ifa_v2_with_const_actual_seq_qlen,
	703	+ backend="inductor",
	704	+ fullgraph=True,
	705	+ )
	706	+ inputs = _make_ifa_inputs()
	707	+ expected = _ifa_v2_with_const_actual_seq_qlen(*inputs)
	708	+ actual = compiled(*inputs)
	709	+ torch.testing.assert_close(
	710	+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
	711	+ )
	712	+
	713	+ inputs = _make_ifa_inputs()
	714	+ expected = _ifa_v2_with_const_actual_seq_qlen(*inputs)
	715	+ actual = compiled(*inputs)
	716	+ torch.testing.assert_close(
	717	+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
	718	+ )
	719	+ finally:
	720	+ graph_tree.update_aclgraph_records_for_graph = original_update
	721	+ config.triton.cudagraphs = old_cudagraphs
	722	+ config.triton.cudagraph_trees = old_cudagraph_trees
	723	+ config.force_disable_caches = old_force_disable_caches
	724	+ config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts
	725	+ torch._dynamo.reset()
	726	+
	727	+ self.assertTrue(any(plan for plan in seen_plans))
	728	+ plan = next(plan for plan in seen_plans if plan)
	729	+ self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score_v2.default")
	730	+ self.assertEqual(
	731	+ plan[0]["updates"]["actual_seq_qlen"],
	732	+ {"kind": "list", "items": [{"kind": "constant", "value": 1}]},
	733	+ )
	734	+
	735	+
	736	+if __name__ == "__main__":
	737	+ run_tests()

Atest/npu/test_aclgraph_update_plan.py+609-0

@@ -0,0 +1,609 @@
	1	+import unittest
	2	+
	3	+from torch_npu.npu._aclgraph_update_plan import (
	4	+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
	5	+ resolve_aclgraph_update_plan,
	6	+ validate_aclgraph_update_plan,
	7	+)
	8	+from torch_npu._inductor._aclgraph_update_plan.codegen import (
	9	+ build_aclgraph_update_plan_entry_for_inductor,
	10	+)
	11	+from torch_npu.npu._aclgraph_update_plan.resolver import (
	12	+ build_cpu_update_input_for_graph,
	13	+)
	14	+
	15	+
	16	+class TestACLGraphUpdatePlan(unittest.TestCase):
	17	+ def setUp(self):
	18	+ from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	19	+
	20	+ class Handler:
	21	+ UPDATE_SPECS = {
	22	+ "npu_fusion_attention_v3.default": [
	23	+ ("arg", 14, "actual_seq_qlen"),
	24	+ ("arg", 15, "actual_seq_kvlen"),
	25	+ ],
	26	+ "npu_fusion_attention_v3.out": [
	27	+ ("arg", 14, "actual_seq_qlen"),
	28	+ ("arg", 15, "actual_seq_kvlen"),
	29	+ ],
	30	+ "npu_fused_infer_attention_score.default": [
	31	+ ("arg", 5, "actual_seq_lengths"),
	32	+ ("arg", 6, "actual_seq_lengths_kv"),
	33	+ ],
	34	+ "npu_fused_infer_attention_score_v2.default": [
	35	+ ("arg", 7, "actual_seq_qlen"),
	36	+ ("arg", 8, "actual_seq_kvlen"),
	37	+ ],
	38	+ }
	39	+
	40	+ @classmethod
	41	+ def get_update_specs(cls, op_name):
	42	+ return cls.UPDATE_SPECS.get(op_name, [])
	43	+
	44	+ self._old_handlers = dict(_NPU_GRAPH_OP_HANDLERS)
	45	+ _NPU_GRAPH_OP_HANDLERS.update({
	46	+ "npu_fusion_attention_v3.default": Handler,
	47	+ "npu_fusion_attention_v3.out": Handler,
	48	+ "npu_fused_infer_attention_score.default": Handler,
	49	+ "npu_fused_infer_attention_score_v2.default": Handler,
	50	+ })
	51	+
	52	+ def tearDown(self):
	53	+ from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	54	+
	55	+ _NPU_GRAPH_OP_HANDLERS.clear()
	56	+ _NPU_GRAPH_OP_HANDLERS.update(self._old_handlers)
	57	+
	58	+ def test_build_inductor_plan_maps_graph_input_sources(self):
	59	+ class Arg:
	60	+ def __init__(self, name):
	61	+ self.name = name
	62	+
	63	+ class Schema:
	64	+ arguments = [
	65	+ Arg("query"),
	66	+ Arg("key"),
	67	+ Arg("value"),
	68	+ Arg("head_num"),
	69	+ Arg("input_layout"),
	70	+ Arg("pse"),
	71	+ Arg("padding_mask"),
	72	+ Arg("atten_mask"),
	73	+ Arg("scale"),
	74	+ Arg("keep_prob"),
	75	+ Arg("pre_tockens"),
	76	+ Arg("next_tockens"),
	77	+ Arg("inner_precise"),
	78	+ Arg("prefix"),
	79	+ Arg("actual_seq_qlen"),
	80	+ Arg("actual_seq_kvlen"),
	81	+ ]
	82	+
	83	+ class Target:
	84	+ __name__ = "npu_fusion_attention_v3.default"
	85	+ _schema = Schema()
	86	+
	87	+ class Value:
	88	+ def __init__(self, name):
	89	+ self.name = name
	90	+
	91	+ def get_name(self):
	92	+ return self.name
	93	+
	94	+ actual = Value("arg0_1")
	95	+ q = Value("arg1_1")
	96	+ k = Value("arg2_1")
	97	+ v = Value("arg3_1")
	98	+
	99	+ self.assertEqual(
	100	+ build_aclgraph_update_plan_entry_for_inductor(
	101	+ Target(),
	102	+ (
	103	+ q, k, v, 1, "TND", None, None, None, 1.0, 1.0,
	104	+ 2147483647, 2147483647, 0, None, actual, actual,
	105	+ ),
	106	+ {},
	107	+ ["arg0_1", "arg1_1", "arg2_1", "arg3_1"],
	108	+ {},
	109	+ ),
	110	+ {
	111	+ "op": "npu_fusion_attention_v3.default",
	112	+ "updates": {
	113	+ "actual_seq_qlen": {"kind": "input", "index": 0},
	114	+ "actual_seq_kvlen": {"kind": "input", "index": 0},
	115	+ },
	116	+ },
	117	+ )
	118	+
	119	+ def test_build_inductor_plan_skips_fa3_bnsd_like_runtime_handler(self):
	120	+ class Arg:
	121	+ def __init__(self, name):
	122	+ self.name = name
	123	+
	124	+ class Schema:
	125	+ arguments = [
	126	+ Arg("query"),
	127	+ Arg("key"),
	128	+ Arg("value"),
	129	+ Arg("head_num"),
	130	+ Arg("input_layout"),
	131	+ Arg("pse"),
	132	+ Arg("padding_mask"),
	133	+ Arg("atten_mask"),
	134	+ Arg("scale"),
	135	+ Arg("keep_prob"),
	136	+ Arg("pre_tockens"),
	137	+ Arg("next_tockens"),
	138	+ Arg("inner_precise"),
	139	+ Arg("prefix"),
	140	+ Arg("actual_seq_qlen"),
	141	+ Arg("actual_seq_kvlen"),
	142	+ ]
	143	+
	144	+ class Target:
	145	+ __name__ = "npu_fusion_attention_v3.default"
	146	+ _schema = Schema()
	147	+
	148	+ class Value:
	149	+ def __init__(self, name):
	150	+ self.name = name
	151	+
	152	+ def get_name(self):
	153	+ return self.name
	154	+
	155	+ actual = Value("arg0_1")
	156	+ q = Value("arg1_1")
	157	+ k = Value("arg2_1")
	158	+ v = Value("arg3_1")
	159	+ args_prefix = (q, k, v, 1)
	160	+ args_suffix = (
	161	+ None, None, None, 1.0, 1.0,
	162	+ 2147483647, 2147483647, 0, None, actual, actual,
	163	+ )
	164	+
	165	+ self.assertIsNone(
	166	+ build_aclgraph_update_plan_entry_for_inductor(
	167	+ Target(),
	168	+ args_prefix + ("BNSD",) + args_suffix,
	169	+ {},
	170	+ ["arg0_1", "arg1_1", "arg2_1", "arg3_1"],
	171	+ {},
	172	+ )
	173	+ )
	174	+ self.assertEqual(
	175	+ build_aclgraph_update_plan_entry_for_inductor(
	176	+ Target(),
	177	+ args_prefix + ("TND",) + args_suffix,
	178	+ {},
	179	+ ["arg0_1", "arg1_1", "arg2_1", "arg3_1"],
	180	+ {},
	181	+ ),
	182	+ {
	183	+ "op": "npu_fusion_attention_v3.default",
	184	+ "updates": {
	185	+ "actual_seq_qlen": {"kind": "input", "index": 0},
	186	+ "actual_seq_kvlen": {"kind": "input", "index": 0},
	187	+ },
	188	+ },
	189	+ )
	190	+
	191	+ def test_build_inductor_plan_ignores_unhandled_ops(self):
	192	+ class Arg:
	193	+ def __init__(self, name):
	194	+ self.name = name
	195	+
	196	+ class Schema:
	197	+ arguments = [Arg("actual_seq_qlen")]
	198	+
	199	+ class Target:
	200	+ __name__ = "unhandled_attention.default"
	201	+ _schema = Schema()
	202	+
	203	+ self.assertIsNone(
	204	+ build_aclgraph_update_plan_entry_for_inductor(
	205	+ Target(),
	206	+ [4],
	207	+ {},
	208	+ [],
	209	+ {},
	210	+ )
	211	+ )
	212	+
	213	+ def test_build_inductor_plan_filters_to_handler_update_specs(self):
	214	+ from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	215	+
	216	+ class Handler:
	217	+ @classmethod
	218	+ def get_update_specs(cls, op_name):
	219	+ return [("arg", 6, "actual_seq_lengths_kv")]
	220	+
	221	+ class Arg:
	222	+ def __init__(self, name):
	223	+ self.name = name
	224	+
	225	+ class Schema:
	226	+ arguments = [
	227	+ Arg("query"),
	228	+ Arg("key"),
	229	+ Arg("value"),
	230	+ Arg("pse_shift"),
	231	+ Arg("atten_mask"),
	232	+ Arg("actual_seq_lengths"),
	233	+ Arg("actual_seq_lengths_kv"),
	234	+ ]
	235	+
	236	+ class Target:
	237	+ __name__ = "npu_fused_infer_attention_score.default"
	238	+ _schema = Schema()
	239	+
	240	+ _NPU_GRAPH_OP_HANDLERS["npu_fused_infer_attention_score.default"] = Handler
	241	+ self.assertEqual(
	242	+ build_aclgraph_update_plan_entry_for_inductor(
	243	+ Target(),
	244	+ ["q", "k", "v", None, None, [15], [100]],
	245	+ {},
	246	+ [],
	247	+ {},
	248	+ ),
	249	+ {
	250	+ "op": "npu_fused_infer_attention_score.default",
	251	+ "updates": {
	252	+ "actual_seq_lengths_kv": {"kind": "list", "items": [
	253	+ {"kind": "constant", "value": 100},
	254	+ ]},
	255	+ },
	256	+ },
	257	+ )
	258	+
	259	+ def test_build_inductor_plan_for_ifa_v1_positional_actual_seq_lengths(self):
	260	+ class Arg:
	261	+ def __init__(self, name):
	262	+ self.name = name
	263	+
	264	+ class Schema:
	265	+ arguments = [
	266	+ Arg("query"),
	267	+ Arg("key"),
	268	+ Arg("value"),
	269	+ Arg("pse_shift"),
	270	+ Arg("atten_mask"),
	271	+ Arg("actual_seq_lengths"),
	272	+ Arg("actual_seq_lengths_kv"),
	273	+ ]
	274	+
	275	+ class Target:
	276	+ __name__ = "npu_fused_infer_attention_score.default"
	277	+ _schema = Schema()
	278	+
	279	+ self.assertEqual(
	280	+ build_aclgraph_update_plan_entry_for_inductor(
	281	+ Target(),
	282	+ ["q", "k", "v", None, None, [15], [100]],
	283	+ {},
	284	+ [],
	285	+ {},
	286	+ ),
	287	+ {
	288	+ "op": "npu_fused_infer_attention_score.default",
	289	+ "updates": {
	290	+ "actual_seq_lengths": {"kind": "list", "items": [
	291	+ {"kind": "constant", "value": 15},
	292	+ ]},
	293	+ "actual_seq_lengths_kv": {"kind": "list", "items": [
	294	+ {"kind": "constant", "value": 100},
	295	+ ]},
	296	+ },
	297	+ },
	298	+ )
	299	+
	300	+ def test_build_inductor_plan_for_ifa_v2_positional_actual_seq_qlen(self):
	301	+ class Arg:
	302	+ def __init__(self, name):
	303	+ self.name = name
	304	+
	305	+ class Schema:
	306	+ arguments = [
	307	+ Arg("query"),
	308	+ Arg("key"),
	309	+ Arg("value"),
	310	+ Arg("query_rope"),
	311	+ Arg("key_rope"),
	312	+ Arg("pse_shift"),
	313	+ Arg("atten_mask"),
	314	+ Arg("actual_seq_qlen"),
	315	+ Arg("actual_seq_kvlen"),
	316	+ ]
	317	+
	318	+ class Target:
	319	+ __name__ = "npu_fused_infer_attention_score_v2.default"
	320	+ _schema = Schema()
	321	+
	322	+ self.assertEqual(
	323	+ build_aclgraph_update_plan_entry_for_inductor(
	324	+ Target(),
	325	+ ["q", "k", "v", None, None, None, None, [16], [128]],
	326	+ {},
	327	+ [],
	328	+ {},
	329	+ ),
	330	+ {
	331	+ "op": "npu_fused_infer_attention_score_v2.default",
	332	+ "updates": {
	333	+ "actual_seq_qlen": {"kind": "list", "items": [
	334	+ {"kind": "constant", "value": 16},
	335	+ ]},
	336	+ "actual_seq_kvlen": {"kind": "list", "items": [
	337	+ {"kind": "constant", "value": 128},
	338	+ ]},
	339	+ },
	340	+ },
	341	+ )
	342	+
	343	+ def test_build_inductor_plan_rejects_unlifted_tensor_actual_seq_constant(self):
	344	+ import torch
	345	+
	346	+ class Arg:
	347	+ def __init__(self, name):
	348	+ self.name = name
	349	+
	350	+ class Schema:
	351	+ arguments = [
	352	+ Arg("query"),
	353	+ Arg("key"),
	354	+ Arg("value"),
	355	+ Arg("pse_shift"),
	356	+ Arg("atten_mask"),
	357	+ Arg("actual_seq_lengths"),
	358	+ ]
	359	+
	360	+ class Target:
	361	+ __name__ = "npu_fused_infer_attention_score.default"
	362	+ _schema = Schema()
	363	+
	364	+ with self.assertRaisesRegex(RuntimeError, "Tensor constant"):
	365	+ build_aclgraph_update_plan_entry_for_inductor(
	366	+ Target(),
	367	+ ["q", "k", "v", None, None, torch.tensor([15])],
	368	+ {},
	369	+ [],
	370	+ {},
	371	+ )
	372	+
	373	+ def test_resolve_input_and_constant_sources(self):
	374	+ new_inputs = ["q", "qlen", "kvlen"]
	375	+ plan = [
	376	+ {
	377	+ "op": "npu_fusion_attention_v3.out",
	378	+ "updates": {
	379	+ "actual_seq_qlen": {"kind": "input", "index": 1},
	380	+ "actual_seq_kvlen": {"kind": "list", "items": [
	381	+ {"kind": "constant", "value": 4},
	382	+ ]},
	383	+ },
	384	+ }
	385	+ ]
	386	+
	387	+ self.assertEqual(ACLGRAPH_UPDATE_PLAN_GLOBAL, "_torch_npu_aclgraph_update_plan")
	388	+ self.assertEqual(
	389	+ resolve_aclgraph_update_plan(plan, new_inputs),
	390	+ [{"actual_seq_qlen": "qlen", "actual_seq_kvlen": [4]}],
	391	+ )
	392	+
	393	+ def test_resolve_list_source_with_input_and_constant_items(self):
	394	+ new_inputs = ["q", 10000]
	395	+ plan = [
	396	+ {
	397	+ "op": "npu_fused_infer_attention_score.default",
	398	+ "updates": {
	399	+ "actual_seq_lengths": {"kind": "list", "items": [
	400	+ {"kind": "constant", "value": 15},
	401	+ {"kind": "input", "index": 1},
	402	+ ]},
	403	+ "actual_seq_lengths_kv": {"kind": "list", "items": [
	404	+ {"kind": "input", "index": 1},
	405	+ ]},
	406	+ },
	407	+ }
	408	+ ]
	409	+
	410	+ self.assertEqual(
	411	+ resolve_aclgraph_update_plan(plan, new_inputs),
	412	+ [{"actual_seq_lengths": [15, 10000], "actual_seq_lengths_kv": [10000]}],
	413	+ )
	414	+
	415	+ def test_resolve_rejects_out_of_range_input_index(self):
	416	+ plan = [
	417	+ {
	418	+ "op": "npu_fusion_attention_v3.out",
	419	+ "updates": {
	420	+ "actual_seq_qlen": {"kind": "input", "index": 3},
	421	+ },
	422	+ }
	423	+ ]
	424	+
	425	+ with self.assertRaisesRegex(RuntimeError, "out of range"):
	426	+ resolve_aclgraph_update_plan(plan, ["only_one_input"])
	427	+
	428	+ def test_validate_plan_rejects_length_mismatch(self):
	429	+ with self.assertRaisesRegex(RuntimeError, "length mismatch"):
	430	+ validate_aclgraph_update_plan(
	431	+ [{"op": "npu_fusion_attention_v3.out", "updates": {}}],
	432	+ [],
	433	+ )
	434	+
	435	+ def test_validate_plan_rejects_missing_plan_with_cache_hint(self):
	436	+ class Record:
	437	+ class Op:
	438	+ __name__ = "npu_fused_infer_attention_score.default"
	439	+
	440	+ op_cache_entry = Op()
	441	+ kwargs = {"actual_seq_lengths": None}
	442	+
	443	+ with self.assertRaisesRegex(RuntimeError, "cached compiled code"):
	444	+ validate_aclgraph_update_plan([], [Record()])
	445	+
	446	+ def test_validate_plan_rejects_op_mismatch(self):
	447	+ class Record:
	448	+ class Op:
	449	+ __name__ = "npu_fused_infer_attention_score.out"
	450	+
	451	+ op_cache_entry = Op()
	452	+ kwargs = {}
	453	+
	454	+ with self.assertRaisesRegex(RuntimeError, "op mismatch"):
	455	+ validate_aclgraph_update_plan(
	456	+ [{"op": "npu_fusion_attention_v3.out", "updates": {}}],
	457	+ [Record()],
	458	+ )
	459	+
	460	+ def test_validate_plan_rejects_invalid_entry_shape(self):
	461	+ class Record:
	462	+ class Op:
	463	+ __name__ = "npu_fusion_attention_v3.out"
	464	+
	465	+ op_cache_entry = Op()
	466	+ kwargs = {"actual_seq_qlen": None}
	467	+
	468	+ with self.assertRaisesRegex(RuntimeError, "invalid plan entry"):
	469	+ validate_aclgraph_update_plan(
	470	+ [{"updates": {
	471	+ "actual_seq_qlen": {"kind": "constant", "value": 4},
	472	+ }}],
	473	+ [Record()],
	474	+ )
	475	+
	476	+ def test_validate_plan_rejects_invalid_source_shape(self):
	477	+ class Record:
	478	+ class Op:
	479	+ __name__ = "npu_fusion_attention_v3.out"
	480	+
	481	+ op_cache_entry = Op()
	482	+ kwargs = {"actual_seq_qlen": None}
	483	+
	484	+ with self.assertRaisesRegex(RuntimeError, "invalid source"):
	485	+ validate_aclgraph_update_plan(
	486	+ [{"op": "npu_fusion_attention_v3.out", "updates": {
	487	+ "actual_seq_qlen": "not_a_source",
	488	+ }}],
	489	+ [Record()],
	490	+ )
	491	+
	492	+ def test_validate_plan_allows_default_out_compatibility(self):
	493	+ class Record:
	494	+ class Op:
	495	+ __name__ = "npu_fusion_attention_v3.out"
	496	+
	497	+ op_cache_entry = Op()
	498	+ kwargs = {"actual_seq_qlen": None}
	499	+
	500	+ validate_aclgraph_update_plan(
	501	+ [{"op": "npu_fusion_attention_v3.default", "updates": {
	502	+ "actual_seq_qlen": {"kind": "list", "items": [
	503	+ {"kind": "constant", "value": 4},
	504	+ ]},
	505	+ }}],
	506	+ [Record()],
	507	+ )
	508	+
	509	+ def test_validate_plan_reads_legacy_handler_update_specs(self):
	510	+ from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	511	+
	512	+ class Handler:
	513	+ UPDATE_SPECS = {
	514	+ "npu_fusion_attention_v3.out": [
	515	+ ("kwarg", "actual_seq_qlen", "actual_seq_qlen"),
	516	+ ],
	517	+ }
	518	+
	519	+ class Record:
	520	+ class Op:
	521	+ __name__ = "npu_fusion_attention_v3.out"
	522	+
	523	+ op_cache_entry = Op()
	524	+ kwargs = {}
	525	+
	526	+ _NPU_GRAPH_OP_HANDLERS["npu_fusion_attention_v3.out"] = Handler
	527	+ validate_aclgraph_update_plan(
	528	+ [{"op": "npu_fusion_attention_v3.default", "updates": {
	529	+ "actual_seq_qlen": {"kind": "constant", "value": 4},
	530	+ }}],
	531	+ [Record()],
	532	+ )
	533	+
	534	+ def test_validate_plan_rejects_empty_updates(self):
	535	+ class Record:
	536	+ class Op:
	537	+ __name__ = "npu_fusion_attention_v3.out"
	538	+
	539	+ op_cache_entry = Op()
	540	+ kwargs = {}
	541	+
	542	+ with self.assertRaisesRegex(RuntimeError, "no updates"):
	543	+ validate_aclgraph_update_plan(
	544	+ [{"op": "npu_fusion_attention_v3.out", "updates": {}}],
	545	+ [Record()],
	546	+ )
	547	+
	548	+ def test_validate_plan_rejects_unsupported_constant(self):
	549	+ class Record:
	550	+ class Op:
	551	+ __name__ = "npu_fusion_attention_v3.out"
	552	+
	553	+ op_cache_entry = Op()
	554	+ kwargs = {"actual_seq_qlen": None}
	555	+
	556	+ with self.assertRaisesRegex(RuntimeError, "unsupported constant"):
	557	+ validate_aclgraph_update_plan(
	558	+ [{"op": "npu_fusion_attention_v3.out", "updates": {
	559	+ "actual_seq_qlen": {"kind": "constant", "value": object()},
	560	+ }}],
	561	+ [Record()],
	562	+ )
	563	+
	564	+ def test_validate_plan_rejects_unsupported_nested_constant(self):
	565	+ class Record:
	566	+ class Op:
	567	+ __name__ = "npu_fused_infer_attention_score.default"
	568	+
	569	+ op_cache_entry = Op()
	570	+ kwargs = {"actual_seq_lengths": None}
	571	+
	572	+ with self.assertRaisesRegex(RuntimeError, "unsupported constant"):
	573	+ validate_aclgraph_update_plan(
	574	+ [{"op": "npu_fused_infer_attention_score.default", "updates": {
	575	+ "actual_seq_lengths": {"kind": "list", "items": [
	576	+ {"kind": "constant", "value": object()},
	577	+ ]},
	578	+ }}],
	579	+ [Record()],
	580	+ )
	581	+
	582	+ def test_build_cpu_update_input_for_graph_tree(self):
	583	+ class Record:
	584	+ class Op:
	585	+ __name__ = "npu_fusion_attention_v3.out"
	586	+
	587	+ op_cache_entry = Op()
	588	+ kwargs = {
	589	+ "actual_seq_qlen": None,
	590	+ "actual_seq_kvlen": None,
	591	+ }
	592	+
	593	+ plan = [
	594	+ {
	595	+ "op": "npu_fusion_attention_v3.out",
	596	+ "updates": {
	597	+ "actual_seq_qlen": {"kind": "input", "index": 0},
	598	+ "actual_seq_kvlen": {"kind": "input", "index": 1},
	599	+ },
	600	+ }
	601	+ ]
	602	+
	603	+ self.assertEqual(
	604	+ build_cpu_update_input_for_graph(plan, ["qlen", "kvlen"], [Record()]),
	605	+ [{"actual_seq_qlen": "qlen", "actual_seq_kvlen": "kvlen"}],
	606	+ )
	607	+
	608	+if __name__ == "__main__":
	609	+ unittest.main()

Mtest/npu/test_npugraph_handler.py+22-0

	@@ -55,6 +55,28 @@ class TestNpuGraphHandlerBuiltinRegistration(TestCase):
55	f"Expected handler for '{op_name}' not found in registry",	55	f"Expected handler for '{op_name}' not found in registry",
56	)	56	)
57		57
		58	+ def test_ifa_update_specs_cover_actual_seq_args(self):
		59	+ for op_name in (
		60	+ "npu_fused_infer_attention_score",
		61	+ "npu_fused_infer_attention_score.default",
		62	+ "npu_fused_infer_attention_score.out",
		63	+ ):
		64	+ with self.subTest(op_name=op_name):
		65	+ specs = _NPU_GRAPH_OP_HANDLERS[op_name].get_update_specs(op_name)
		66	+ self.assertIn(("arg", 5, "actual_seq_lengths"), specs)
		67	+ self.assertIn(("arg", 6, "actual_seq_lengths_kv"), specs)
		68	+
		69	+ def test_ifa_v2_update_specs_cover_actual_seq_args(self):
		70	+ for op_name in (
		71	+ "npu_fused_infer_attention_score_v2",
		72	+ "npu_fused_infer_attention_score_v2.default",
		73	+ "npu_fused_infer_attention_score_v2.out",
		74	+ ):
		75	+ with self.subTest(op_name=op_name):
		76	+ specs = _NPU_GRAPH_OP_HANDLERS[op_name].get_update_specs(op_name)
		77	+ self.assertIn(("arg", 7, "actual_seq_qlen"), specs)
		78	+ self.assertIn(("arg", 8, "actual_seq_kvlen"), specs)
		79	+
58		80
59	if __name__ == "__main__":	81	if __name__ == "__main__":
60	run_tests()	82	run_tests()

Mtest/torch_npu_schema.json+3-0

	@@ -2330,6 +2330,9 @@
2330	"torch_npu.npu.NpuGraphOpHandler.update_args": {	2330	"torch_npu.npu.NpuGraphOpHandler.update_args": {
2331	"signature": "(dispatch_record, update_input)"	2331	"signature": "(dispatch_record, update_input)"
2332	},	2332	},
		2333	+ "torch_npu.npu.NpuGraphOpHandler.get_update_specs": {
		2334	+ "signature": "(op_name)"
		2335	+ },
2333	"torch_npu.npu.NpuGraphOpHandler.record_wrap_kwarg": {	2336	"torch_npu.npu.NpuGraphOpHandler.record_wrap_kwarg": {
2334	"signature": "(key, value, tensor_param_names)"	2337	"signature": "(key, value, tensor_param_names)"
2335	},	2338	},

Atorch_npu/_inductor/_aclgraph_update_plan/__init__.py+12-0

@@ -0,0 +1,12 @@
	1	+from torch_npu._inductor._aclgraph_update_plan.codegen import (
	2	+ append_inductor_aclgraph_update_plan_for_codegen_node,
	3	+ emit_inductor_aclgraph_update_plan_for_wrapper,
	4	+)
	5	+from torch_npu.npu._aclgraph_update_plan import ACLGRAPH_UPDATE_PLAN_GLOBAL
	6	+
	7	+
	8	+__all__ = [
	9	+ "ACLGRAPH_UPDATE_PLAN_GLOBAL",
	10	+ "append_inductor_aclgraph_update_plan_for_codegen_node",
	11	+ "emit_inductor_aclgraph_update_plan_for_wrapper",
	12	+]

Atorch_npu/_inductor/_aclgraph_update_plan/codegen.py+245-0

@@ -0,0 +1,245 @@
	1	+from typing import Any, Callable, Dict, Optional, Sequence
	2	+
	3	+import sympy
	4	+
	5	+from torch_npu.npu._aclgraph_update_plan.resolver import (
	6	+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
	7	+ _get_update_specs,
	8	+ _normalize_op_name,
	9	+)
	10	+from torch_npu.utils._error_code import ErrCode, pta_error
	11	+
	12	+
	13	+def _handler_for_op(op_name: str) -> Optional[Any]:
	14	+ import torch_npu.npu._npugraph_handlers # noqa: F401
	15	+ from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	16	+
	OO O openLiBingCI6月10日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月10日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月11日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月11日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月11日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月11日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月11日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月11日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月12日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月12日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月12日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月12日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月12日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月12日此条代码评论区间+12至+16 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	17	+ return _NPU_GRAPH_OP_HANDLERS.get(op_name)
	OO O openLiBingCI6月16日此条代码评论区间+10至+17 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 O openLiBingCI6月16日此条代码评论区间+10至+17 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。
	18	+
	19	+
	20	+def _matches_aclgraph_update_exclusion(op_name: str, bound: Dict[str, Any]) -> bool:
	21	+ normalized_op = _normalize_op_name(op_name)
	22	+ if normalized_op in {
	23	+ "npu_fusion_attention_v3",
	24	+ "npu_fusion_attention_grad_v3",
	25	+ }:
	26	+ return bound.get("input_layout") == "BNSD"
	27	+ return False
	28	+
	29	+
	30	+def _literal_to_source(value: Any) -> Dict[str, Any]:
	31	+ import torch
	32	+
	33	+ if isinstance(value, torch.Tensor):
	34	+ raise RuntimeError(
	35	+ "Unsupported ACLGraph update Tensor constant; Tensor update values must be graph inputs",
	36	+ pta_error(ErrCode.PARAM),
	37	+ )
	38	+ if value is None:
	39	+ return {"kind": "none"}
	40	+ if isinstance(value, (int, float, bool, str)):
	41	+ return {"kind": "constant", "value": value}
	42	+ raise RuntimeError(
	43	+ f"Unsupported ACLGraph update source: {value!r}",
	44	+ pta_error(ErrCode.PARAM),
	45	+ )
	46	+
	47	+
	48	+def _maybe_sympy_expr(value: Any) -> Optional[sympy.Expr]:
	49	+ if isinstance(value, sympy.Expr):
	50	+ return value
	51	+ import torch
	52	+
	53	+ if isinstance(value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
	54	+ return value.node.expr
	55	+ return None
	56	+
	57	+
	58	+def _try_get_name(value: Any) -> Optional[str]:
	59	+ import torch
	60	+
	61	+ if isinstance(value, torch.Tensor):
	62	+ return None
	63	+
	64	+ if hasattr(value, "get_name"):
	65	+ try:
	66	+ return value.get_name()
	67	+ except (AttributeError, NotImplementedError):
	68	+ pass
	69	+
	70	+ data = getattr(value, "data", None)
	71	+ if data is not None and data is not value:
	72	+ name = _try_get_name(data)
	73	+ if name is not None:
	74	+ return name
	75	+
	76	+ if hasattr(value, "unwrap_view"):
	77	+ try:
	78	+ return _try_get_name(value.unwrap_view())
	79	+ except (AttributeError, NotImplementedError):
	80	+ pass
	81	+
	82	+ return None
	83	+
	84	+
	85	+def _exprs_equal(lhs: sympy.Expr, rhs: sympy.Expr) -> bool:
	86	+ try:
	87	+ return bool(sympy.simplify(lhs - rhs) == 0)
	88	+ except Exception:
	89	+ return lhs == rhs
	90	+
	91	+
	92	+def _lookup_graph_input_index(
	93	+ value: Any,
	94	+ graph_input_names: Sequence[str],
	95	+ graph_inputs: Dict[str, Any],
	96	+) -> Optional[int]:
	97	+ name = _try_get_name(value)
	98	+ if name in graph_input_names:
	99	+ return list(graph_input_names).index(name)
	100	+
	101	+ expr = _maybe_sympy_expr(value)
	102	+ for idx, input_name in enumerate(graph_input_names):
	103	+ input_value = graph_inputs.get(input_name)
	104	+ if value is input_value:
	105	+ return idx
	106	+ input_expr = _maybe_sympy_expr(input_value)
	107	+ if expr is not None and input_expr is not None and _exprs_equal(expr, input_expr):
	108	+ return idx
	109	+
	110	+ return None
	111	+
	112	+
	113	+def _inductor_value_to_source(
	114	+ value: Any,
	115	+ graph_input_names: Sequence[str],
	116	+ graph_inputs: Dict[str, Any],
	117	+) -> Dict[str, Any]:
	118	+ index = _lookup_graph_input_index(value, graph_input_names, graph_inputs)
	119	+ if index is not None:
	120	+ return {"kind": "input", "index": index}
	121	+ if isinstance(value, (list, tuple)):
	122	+ return {
	123	+ "kind": "list",
	124	+ "items": [
	125	+ _inductor_value_to_source(item, graph_input_names, graph_inputs)
	126	+ for item in value
	127	+ ],
	128	+ }
	129	+ return _literal_to_source(value)
	130	+
	131	+
	132	+def _bind_by_schema(target: Any, args: Sequence[Any], kwargs: Dict[str, Any]) -> Dict[str, Any]:
	133	+ schema = getattr(target, "_schema", None)
	134	+ if schema is None:
	135	+ return {}
	136	+
	137	+ bound = {}
	138	+ args = list(args)
	139	+ for idx, arg in enumerate(schema.arguments):
	140	+ name = arg.name
	141	+ if idx < len(args):
	142	+ bound[name] = args[idx]
	143	+ elif name in kwargs:
	144	+ bound[name] = kwargs[name]
	145	+ return bound
	146	+
	147	+
	148	+def _build_update_plan_entry_from_bound_args(
	149	+ op_name: str,
	150	+ bound: Dict[str, Any],
	151	+ source_builder: Callable[[Any], Dict[str, Any]],
	152	+ supported_keys: set,
	153	+) -> Optional[Dict[str, Any]]:
	154	+ updates = {}
	155	+ for key, value in bound.items():
	156	+ if key in supported_keys:
	157	+ updates[key] = source_builder(value)
	158	+ if not updates:
	159	+ return None
	160	+ return {"op": op_name, "updates": updates}
	161	+
	162	+
	163	+def build_aclgraph_update_plan_entry_for_inductor(
	164	+ op_overload: Any,
	165	+ args: Sequence[Any],
	166	+ kwargs: Dict[str, Any],
	167	+ graph_input_names: Sequence[str],
	168	+ graph_inputs: Dict[str, Any],
	169	+) -> Optional[Dict[str, Any]]:
	170	+ op_name = getattr(op_overload, "__name__", str(op_overload))
	171	+ handler_cls = _handler_for_op(op_name)
	172	+ if handler_cls is None:
	173	+ return None
	174	+
	175	+ bound = _bind_by_schema(op_overload, args, kwargs)
	176	+ if _matches_aclgraph_update_exclusion(op_name, bound):
	177	+ return None
	178	+
	179	+ supported_keys = {key for _, _, key in _get_update_specs(handler_cls, op_name)}
	180	+ return _build_update_plan_entry_from_bound_args(
	181	+ op_name,
	182	+ bound,
	183	+ lambda value: _inductor_value_to_source(value, graph_input_names, graph_inputs),
	184	+ supported_keys,
	185	+ )
	186	+
	187	+
	188	+def should_generate_inductor_aclgraph_update_plan() -> bool:
	189	+ from torch._inductor import config
	190	+ from torch._inductor.virtualized import V
	191	+
	192	+ return (
	193	+ config.triton.cudagraphs
	194	+ and config.triton.cudagraph_trees
	195	+ and getattr(V.graph, "disable_cudagraphs_reason", None) is None
	196	+ )
	197	+
	198	+
	199	+def append_inductor_aclgraph_update_plan_for_codegen_node(wrapper: Any, node: Any) -> None:
	200	+ if not should_generate_inductor_aclgraph_update_plan():
	201	+ return
	202	+
	203	+ op_overload = getattr(node, "op_overload", None)
	204	+ if op_overload is None:
	205	+ return
	206	+
	207	+ if hasattr(node, "unflatten_args"):
	208	+ args, kwargs = node.unflatten_args(node.inputs, node.constant_args)
	209	+ else:
	210	+ args = [node.inputs, node.constant_args]
	211	+ kwargs = getattr(node, "kwargs", {})
	212	+
	213	+ entry = build_aclgraph_update_plan_entry_for_inductor(
	214	+ op_overload,
	215	+ args,
	216	+ kwargs,
	217	+ wrapper.get_graph_input_names(),
	218	+ wrapper.get_graph_inputs(),
	219	+ )
	220	+ if entry is None:
	221	+ return
	222	+
	223	+ if not hasattr(wrapper, "torch_npu_aclgraph_update_plan"):
	224	+ wrapper.torch_npu_aclgraph_update_plan = []
	225	+ wrapper.torch_npu_aclgraph_update_plan.append(entry)
	226	+
	227	+
	228	+def emit_inductor_aclgraph_update_plan_for_wrapper(
	229	+ wrapper: Any,
	230	+ result: Any,
	231	+ is_graph_partition_subgraph: bool,
	232	+) -> None:
	233	+ from torch._inductor import config
	234	+
	235	+ if not should_generate_inductor_aclgraph_update_plan():
	236	+ return
	237	+
	238	+ plan = getattr(wrapper, "torch_npu_aclgraph_update_plan", [])
	239	+ if not plan:
	240	+ return
	241	+
	242	+ if config.graph_partition and not is_graph_partition_subgraph:
	243	+ return
	244	+
	245	+ result.writeline(f"{wrapper.launcher_fn_name}.{ACLGRAPH_UPDATE_PLAN_GLOBAL} = {plan!r}")

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/codegen/wrapper.py+39-1

	@@ -16,6 +16,10 @@ from ... import codecache
16	from torch._inductor.codegen.common import (	16	from torch._inductor.codegen.common import (
17	IndentedBuffer,	17	IndentedBuffer,
18	)	18	)
		19	+from torch_npu._inductor._aclgraph_update_plan import (
		20	+ append_inductor_aclgraph_update_plan_for_codegen_node,
		21	+ emit_inductor_aclgraph_update_plan_for_wrapper,
		22	+)
19		23
20		24
21	class NpuMlirWrapperCodeGen(PythonWrapperCodegen):	25	class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
	@@ -37,7 +41,7 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
37	raise ValueError("subgraph_name must be provided for python wrapper")	41	raise ValueError("subgraph_name must be provided for python wrapper")
38	if parent_wrapper is None:	42	if parent_wrapper is None:
39	raise ValueError("parent_wrapper must be provided for python wrapper")	43	raise ValueError("parent_wrapper must be provided for python wrapper")
40	- return SubgraphPythonWrapperCodegen(	44	+ return NpuMlirSubgraphPythonWrapperCodegen(
41	subgraph_name, parent_wrapper, partition_signatures	45	subgraph_name, parent_wrapper, partition_signatures
42	)	46	)
43	return NpuMlirWrapperCodeGen()	47	return NpuMlirWrapperCodeGen()
	@@ -126,6 +130,22 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
126	f"run_intermediate_hooks({origin_node.name!r}, {output_name})"	130	f"run_intermediate_hooks({origin_node.name!r}, {output_name})"
127	)	131	)
128		132
		133	+ def generate_extern_kernel_alloc(self, extern_kernel):
		134	+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
		135	+ super().generate_extern_kernel_alloc(extern_kernel)
		136	+
		137	+ def generate_fallback_kernel(self, node) -> None:
		138	+ append_inductor_aclgraph_update_plan_for_codegen_node(self, node)
		139	+ super().generate_fallback_kernel(node)
		140	+
		141	+ def generate_after_suffix(self, result: IndentedBuffer) -> None:
		142	+ super().generate_after_suffix(result)
		143	+ emit_inductor_aclgraph_update_plan_for_wrapper(
		144	+ self,
		145	+ result,
		146	+ is_graph_partition_subgraph=False,
		147	+ )
		148	+
129	def write_get_raw_stream(self, device_idx: int, graph=None) -> str:	149	def write_get_raw_stream(self, device_idx: int, graph=None) -> str:
130	self.write_triton_header_once()	150	self.write_triton_header_once()
131	name = f"stream{device_idx}"	151	name = f"stream{device_idx}"
	@@ -186,3 +206,21 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
186		206
187	def generate_return(self, output_refs: list[str]) -> None:	207	def generate_return(self, output_refs: list[str]) -> None:
188	super().generate_return(output_refs)	208	super().generate_return(output_refs)
		209	+
		210	+
		211	+class NpuMlirSubgraphPythonWrapperCodegen(SubgraphPythonWrapperCodegen):
		212	+ def generate_extern_kernel_alloc(self, extern_kernel):
		213	+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
		214	+ super().generate_extern_kernel_alloc(extern_kernel)
		215	+
		216	+ def generate_fallback_kernel(self, node) -> None:
		217	+ append_inductor_aclgraph_update_plan_for_codegen_node(self, node)
		218	+ super().generate_fallback_kernel(node)
		219	+
		220	+ def generate_after_suffix(self, result: IndentedBuffer) -> None:
		221	+ super().generate_after_suffix(result)
		222	+ emit_inductor_aclgraph_update_plan_for_wrapper(
		223	+ self,
		224	+ result,
		225	+ is_graph_partition_subgraph=True,
		226	+ )

Mtorch_npu/_inductor/codegen/wrapper.py+60-29

	@@ -23,6 +23,63 @@ from torch._inductor.codegen.wrapper import BufferLike, WrapperLine
23	from torch._inductor import ir	23	from torch._inductor import ir
24	import torch_npu.npu.aclnn	24	import torch_npu.npu.aclnn
25	from ..fx_passes.utils.schedule_node_utils import is_multi_stream	25	from ..fx_passes.utils.schedule_node_utils import is_multi_stream
		26	+from torch_npu._inductor._aclgraph_update_plan import (
		27	+ append_inductor_aclgraph_update_plan_for_codegen_node,
		28	+ emit_inductor_aclgraph_update_plan_for_wrapper,
		29	+)
		30	+
		31	+
		32	+def _is_codegen_graph_partition_subgraph(wrapper) -> bool:
		33	+ try:
		34	+ from torch._inductor.utils import is_codegen_graph_partition_subgraph
		35	+ except ImportError:
		36	+ return isinstance(wrapper, SubgraphPythonWrapperCodegen)
		37	+ try:
		38	+ return is_codegen_graph_partition_subgraph(wrapper)
		39	+ except AttributeError:
		40	+ return isinstance(wrapper, SubgraphPythonWrapperCodegen)
		41	+
		42	+
		43	+class _NPUKernelCodegenMixin:
		44	+ # This mixin must appear before the PyTorch wrapper base in the MRO so NPU
		45	+ # hooks run first, then cooperative super() continues into the base wrapper.
		46	+ # generate numel expr for range_tree_node
		47	+ def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
		48	+ expr = f"{kernel_name}_{node.name}_numel"
		49	+ simplified = V.graph.sizevars.simplify(numel_expr)
		50	+ # Ensure all PRECOMPUTED_SIZE symbols in the simplified expression
		51	+ # are defined before we emit the line that uses them.
		52	+ for sym in simplified.free_symbols:
		53	+ self.ensure_size_computed(sym)
		54	+ self.writeline(f"{expr} = {pexpr(simplified)}")
		55	+ # We can get symbolic expressions here, like s0*64
		56	+ # It is fine to have them here, but we need to handle them correctly as their own type
		57	+ # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
		58	+ # scalars as well.
		59	+ # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
		60	+ # constant now, need type info. I agree, this needs type info, and while this is not true type info
		61	+ # it suffices as a type hint for the purposes of producing the correct code for this type.
		62	+ return SymbolicCallArg(expr, numel_expr)
		63	+
		64	+ # don't assert
		65	+ def codegen_input_size_asserts(self) -> None:
		66	+ pass
		67	+
		68	+ def generate_fallback_kernel(self, node: ir.FallbackKernel) -> None:
		69	+ append_inductor_aclgraph_update_plan_for_codegen_node(self, node)
		70	+ super().generate_fallback_kernel(node)
		71	+
		72	+ def generate_extern_kernel_alloc(self, extern_kernel):
		73	+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
		74	+ super().generate_extern_kernel_alloc(extern_kernel)
		75	+
		76	+ def generate_after_suffix(self, result: IndentedBuffer) -> None:
		77	+ super().generate_after_suffix(result)
		78	+ emit_inductor_aclgraph_update_plan_for_wrapper(
		79	+ self,
		80	+ result,
		81	+ _is_codegen_graph_partition_subgraph(self),
		82	+ )
26		83
27		84
28	@dataclasses.dataclass	85	@dataclasses.dataclass
	@@ -54,7 +111,7 @@ class NPUMultiOutputLine(MultiOutputLine):
54	f"{self.multi_stream_intent}{self.wrapper.declare}{self.result_name} = {value}{self.wrapper.ending}"	111	f"{self.multi_stream_intent}{self.wrapper.declare}{self.result_name} = {value}{self.wrapper.ending}"
55	)	112	)
56		113
57	-class NPUPythonWrapperCodeGen(PythonWrapperCodegen):	114	+class NPUPythonWrapperCodeGen(_NPUKernelCodegenMixin, PythonWrapperCodegen):
58	def __init__(self):	115	def __init__(self):
59	super().__init__()	116	super().__init__()
60	self.buffer_args_multi_stream_intent = {}	117	self.buffer_args_multi_stream_intent = {}
	@@ -92,24 +149,6 @@ class NPUPythonWrapperCodeGen(PythonWrapperCodegen):
92	"import torch_npu._inductor.runtime.triton_heuristics as triton_heuristics"	149	"import torch_npu._inductor.runtime.triton_heuristics as triton_heuristics"
93	)	150	)
94		151
95	- # generate numel expr for range_tree_node
96	- def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
97	- expr = f"{kernel_name}_{node.name}_numel"
98	- simplified = V.graph.sizevars.simplify(numel_expr)
99	- # Ensure all PRECOMPUTED_SIZE symbols in the simplified expression
100	- # are defined before we emit the line that uses them.
101	- for sym in simplified.free_symbols:
102	- self.ensure_size_computed(sym)
103	- self.writeline(f"{expr} = {pexpr(simplified)}")
104	- # We can get symbolic expressions here, like s0*64
105	- # It is fine to have them here, but we need to handle them correctly as their own type
106	- # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
107	- # scalars as well.
108	- # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
109	- # constant now, need type info. I agree, this needs type info, and while this is not true type info
110	- # it suffices as a type hint for the purposes of producing the correct code for this type.
111	- return SymbolicCallArg(expr, numel_expr)
112	-
113	def generate_save_uncompiled_kernels(self):	152	def generate_save_uncompiled_kernels(self):
114	# remove incorrect grid=(0,0,0) param	153	# remove incorrect grid=(0,0,0) param
115	self.wrapper_call.splice(	154	self.wrapper_call.splice(
	@@ -126,10 +165,6 @@ class NPUPythonWrapperCodeGen(PythonWrapperCodegen):
126	"""	165	"""
127	)	166	)
128		167
129	- # don't assert
130	- def codegen_input_size_asserts(self) -> None:
131	- pass
132	-
133	def write_prefix(self) -> None:	168	def write_prefix(self) -> None:
134	super().write_prefix()	169	super().write_prefix()
135	if torch_npu.npu.aclnn._use_static_aclnn_kernel:	170	if torch_npu.npu.aclnn._use_static_aclnn_kernel:
	@@ -156,7 +191,6 @@ class NPUPythonWrapperCodeGen(PythonWrapperCodegen):
156	self.wrapper_call.writeline('static_kernel_complier.__exit__(*exc_info)')	191	self.wrapper_call.writeline('static_kernel_complier.__exit__(*exc_info)')
157	super().generate_return(output_refs)	192	super().generate_return(output_refs)
158		193
159	-
160	def get_buffer_define_multi_stream_by_name(self, name):	194	def get_buffer_define_multi_stream_by_name(self, name):
161	multi_stream_intent_str = ""	195	multi_stream_intent_str = ""
162	if name and self.buffer_define_multi_stream is not None and name in self.buffer_define_multi_stream.keys():	196	if name and self.buffer_define_multi_stream is not None and name in self.buffer_define_multi_stream.keys():
	@@ -594,8 +628,5 @@ class NPUPythonWrapperCodeGen(PythonWrapperCodegen):
594	i += 1	628	i += 1
595	return sub_streams_line_no	629	return sub_streams_line_no
596		630
597	-class NPUSubgraphPythonWrapperCodegen(SubgraphPythonWrapperCodegen):	631	+class NPUSubgraphPythonWrapperCodegen(_NPUKernelCodegenMixin, SubgraphPythonWrapperCodegen):
598	- def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):	632	+ pass
599	- expr = f"{kernel_name}_{node.name}_numel"
600	- self.writeline(f"{expr} = {pexpr(numel_expr)}")
601	- return SymbolicCallArg(expr, numel_expr)

Atorch_npu/npu/_aclgraph_update_plan/__init__.py+18-0

@@ -0,0 +1,18 @@
	1	+from torch_npu.npu._aclgraph_update_plan.resolver import (
	2	+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
	3	+ build_cpu_update_input_for_graph,
	4	+ resolve_aclgraph_update_plan,
	5	+ update_aclgraph_records_for_graph,
	6	+ validate_aclgraph_update_plan,
	7	+ validate_aclgraph_update_plan_for_graph,
	8	+)
	9	+
	10	+
	11	+__all__ = [
	12	+ "ACLGRAPH_UPDATE_PLAN_GLOBAL",
	13	+ "build_cpu_update_input_for_graph",
	14	+ "resolve_aclgraph_update_plan",
	15	+ "update_aclgraph_records_for_graph",
	16	+ "validate_aclgraph_update_plan",
	17	+ "validate_aclgraph_update_plan_for_graph",
	18	+]

Atorch_npu/npu/_aclgraph_update_plan/resolver.py+262-0

@@ -0,0 +1,262 @@
	1	+from typing import Any, Dict, List, Sequence
	2	+
	3	+from torch_npu.utils._error_code import ErrCode, pta_error
	4	+
	5	+
	6	+ACLGRAPH_UPDATE_PLAN_GLOBAL = "_torch_npu_aclgraph_update_plan"
	7	+
	8	+
	9	+def _normalize_op_name(op_name: str) -> str:
	10	+ for suffix in (".default", ".out"):
	11	+ if op_name.endswith(suffix):
	12	+ return op_name[: -len(suffix)]
	13	+ return op_name
	14	+
	15	+
	16	+def _op_names_compatible(expected_op: str, actual_op: str) -> bool:
	17	+ return expected_op == actual_op or _normalize_op_name(expected_op) == _normalize_op_name(actual_op)
	18	+
	19	+
	20	+def _get_update_specs(handler_cls: Any, op_name: str) -> List[Any]:
	21	+ get_update_specs = getattr(handler_cls, "get_update_specs", None)
	22	+ if get_update_specs is not None:
	23	+ return get_update_specs(op_name)
	24	+ return getattr(handler_cls, "UPDATE_SPECS", {}).get(op_name, [])
	25	+
	26	+
	27	+def _consumable_keys(record: Any) -> set:
	28	+ from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	29	+
	30	+ op_name = record.op_cache_entry.__name__
	31	+ keys = set(getattr(record, "kwargs", {}).keys())
	32	+ handler_cls = _NPU_GRAPH_OP_HANDLERS.get(op_name)
	33	+ if handler_cls is not None:
	34	+ keys.update(key for _, _, key in _get_update_specs(handler_cls, op_name))
	35	+ return keys
	36	+
	37	+
	38	+def resolve_aclgraph_update_plan(
	39	+ plan: Sequence[Dict[str, Any]],
	40	+ new_inputs: Sequence[Any],
	41	+) -> List[Dict[str, Any]]:
	42	+ cpu_update_input: List[Dict[str, Any]] = []
	43	+ for entry_idx, entry in enumerate(plan or []):
	44	+ updates = entry.get("updates", {})
	45	+ resolved = {}
	46	+ for key, source in updates.items():
	47	+ resolved[key] = _resolve_source(entry_idx, key, source, new_inputs)
	48	+ cpu_update_input.append(resolved)
	49	+ return cpu_update_input
	50	+
	51	+
	52	+def validate_aclgraph_update_plan(
	53	+ plan: Sequence[Dict[str, Any]],
	54	+ graph_dispatch_records: Sequence[Any],
	55	+) -> None:
	56	+ plan = plan or []
	57	+ if not plan and graph_dispatch_records:
	58	+ ops = [record.op_cache_entry.__name__ for record in graph_dispatch_records]
	59	+ raise RuntimeError(
	60	+ "Captured updatable ACLGraph operators but missing ACLGraph update plan: "
	61	+ f"{ops}. This may be caused by reusing cached compiled code generated "
	62	+ "with ACLGraph disabled or by using the npugraphs backend, which does "
	63	+ "not support ACLGraph update plans yet.",
	64	+ pta_error(ErrCode.PARAM),
	65	+ )
	66	+
	67	+ if len(plan) != len(graph_dispatch_records):
	68	+ raise RuntimeError(
	69	+ "ACLGraph update plan length mismatch: "
	70	+ f"plan has {len(plan)} entries but graph captured "
	71	+ f"{len(graph_dispatch_records)} updatable records",
	72	+ pta_error(ErrCode.PARAM),
	73	+ )
	74	+
	75	+ for idx, (entry, record) in enumerate(zip(plan, graph_dispatch_records)):
	76	+ if not isinstance(entry, dict):
	77	+ raise RuntimeError(
	78	+ f"ACLGraph update plan has invalid plan entry at index {idx}: {entry!r}",
	79	+ pta_error(ErrCode.PARAM),
	80	+ )
	81	+ expected_op = entry.get("op")
	82	+ actual_op = record.op_cache_entry.__name__
	83	+ if not isinstance(expected_op, str):
	84	+ raise RuntimeError(
	85	+ "ACLGraph update plan has invalid plan entry: "
	86	+ f"entry {idx} op must be a string, got {expected_op!r}",
	87	+ pta_error(ErrCode.PARAM),
	88	+ )
	89	+ if not _op_names_compatible(expected_op, actual_op):
	90	+ raise RuntimeError(
	91	+ "ACLGraph update plan op mismatch: "
	92	+ f"entry {idx} expects {expected_op!r}, captured {actual_op!r}",
	93	+ pta_error(ErrCode.PARAM),
	94	+ )
	95	+
	96	+ updates = entry.get("updates", {})
	97	+ if not isinstance(updates, dict):
	98	+ raise RuntimeError(
	99	+ "ACLGraph update plan has invalid plan entry: "
	100	+ f"entry {idx} updates must be a dict, got {updates!r}",
	101	+ pta_error(ErrCode.PARAM),
	102	+ )
	103	+ if not updates:
	104	+ raise RuntimeError(
	105	+ "ACLGraph update plan entry has no updates: "
	106	+ f"entry {idx}, op {actual_op!r}",
	107	+ pta_error(ErrCode.PARAM),
	108	+ )
	109	+
	110	+ consumable = _consumable_keys(record)
	111	+ unknown = sorted(set(updates) - consumable)
	112	+ if unknown:
	113	+ raise RuntimeError(
	114	+ "ACLGraph update plan has key(s) that captured op cannot consume: "
	115	+ f"entry {idx}, op {actual_op!r}, keys {unknown}",
	116	+ pta_error(ErrCode.PARAM),
	117	+ )
	118	+
	119	+ for key, source in updates.items():
	120	+ _validate_source(idx, key, source)
	121	+
	122	+
	123	+def build_cpu_update_input_for_graph(
	124	+ plan: Sequence[Dict[str, Any]],
	125	+ new_inputs: Sequence[Any],
	126	+ graph_dispatch_records: Sequence[Any],
	127	+) -> List[Dict[str, Any]]:
	128	+ validate_aclgraph_update_plan(plan, graph_dispatch_records)
	129	+ return resolve_aclgraph_update_plan(plan, new_inputs)
	130	+
	131	+
	132	+def validate_aclgraph_update_plan_for_graph(
	133	+ plan: Sequence[Dict[str, Any]],
	134	+ graph: Any,
	135	+) -> None:
	136	+ if graph is None or not graph.auto_dispatch_capture:
	137	+ return
	138	+ validate_aclgraph_update_plan(
	139	+ plan,
	140	+ graph.graph_dispatch_mode.graph_dispatch_records,
	141	+ )
	142	+
	143	+
	144	+def update_aclgraph_records_for_graph(
	145	+ plan: Sequence[Dict[str, Any]],
	146	+ graph: Any,
	147	+ new_inputs: Sequence[Any],
	148	+) -> bool:
	149	+ if graph is None or not graph.auto_dispatch_capture:
	150	+ return False
	151	+ if not plan:
	152	+ return False
	153	+
	154	+ graph.update(resolve_aclgraph_update_plan(plan, new_inputs))
	155	+ return True
	156	+
	157	+
	158	+def _resolve_source(
	159	+ entry_idx: int,
	160	+ key: str,
	161	+ source: Dict[str, Any],
	162	+ new_inputs: Sequence[Any],
	163	+) -> Any:
	164	+ if not isinstance(source, dict):
	165	+ raise RuntimeError(
	166	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	167	+ f"has invalid source {source!r}",
	168	+ pta_error(ErrCode.PARAM),
	169	+ )
	170	+ kind = source.get("kind")
	171	+ if kind == "input":
	172	+ index = source.get("index")
	173	+ if not isinstance(index, int) or index < 0 or index >= len(new_inputs):
	174	+ raise RuntimeError(
	175	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	176	+ f"input index {index} is out of range for {len(new_inputs)} inputs",
	177	+ pta_error(ErrCode.PARAM),
	178	+ )
	179	+ return new_inputs[index]
	180	+ if kind == "constant":
	181	+ return source.get("value")
	182	+ if kind == "none":
	183	+ return None
	184	+ if kind == "list":
	185	+ items = source.get("items")
	186	+ if not isinstance(items, list):
	187	+ raise RuntimeError(
	188	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	189	+ f"has invalid list source items {items!r}",
	190	+ pta_error(ErrCode.PARAM),
	191	+ )
	192	+ return [_resolve_source(entry_idx, key, item, new_inputs) for item in items]
	193	+ raise RuntimeError(
	194	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	195	+ f"has unsupported source kind {kind!r}",
	196	+ pta_error(ErrCode.PARAM),
	197	+ )
	198	+
	199	+
	200	+def _validate_literal_constant(value: Any) -> None:
	201	+ import torch
	202	+
	203	+ if isinstance(value, torch.Tensor):
	204	+ raise RuntimeError(
	205	+ "Unsupported ACLGraph update Tensor constant; Tensor update values must be graph inputs",
	206	+ pta_error(ErrCode.PARAM),
	207	+ )
	208	+ if value is None or isinstance(value, (int, float, bool, str)):
	209	+ return
	210	+ if isinstance(value, (list, tuple)):
	211	+ for item in value:
	212	+ _validate_literal_constant(item)
	213	+ return
	214	+ raise RuntimeError(
	215	+ f"Unsupported ACLGraph update constant: {value!r}",
	216	+ pta_error(ErrCode.PARAM),
	217	+ )
	218	+
	219	+
	220	+def _validate_source(entry_idx: int, key: str, source: Dict[str, Any]) -> None:
	221	+ if not isinstance(source, dict):
	222	+ raise RuntimeError(
	223	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	224	+ f"has invalid source {source!r}",
	225	+ pta_error(ErrCode.PARAM),
	226	+ )
	227	+ kind = source.get("kind")
	228	+ if kind == "input":
	229	+ index = source.get("index")
	230	+ if not isinstance(index, int) or index < 0:
	231	+ raise RuntimeError(
	232	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	233	+ f"has invalid input index {index}",
	234	+ pta_error(ErrCode.PARAM),
	235	+ )
	236	+ elif kind == "constant":
	237	+ try:
	238	+ _validate_literal_constant(source.get("value"))
	239	+ except RuntimeError as exc:
	240	+ raise RuntimeError(
	241	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	242	+ f"has unsupported constant value {source.get('value')!r}",
	243	+ pta_error(ErrCode.PARAM),
	244	+ ) from exc
	245	+ elif kind == "none":
	246	+ return
	247	+ elif kind == "list":
	248	+ items = source.get("items")
	249	+ if not isinstance(items, list):
	250	+ raise RuntimeError(
	251	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	252	+ f"has invalid list source items {items!r}",
	253	+ pta_error(ErrCode.PARAM),
	254	+ )
	255	+ for item in items:
	256	+ _validate_source(entry_idx, key, item)
	257	+ else:
	258	+ raise RuntimeError(
	259	+ f"ACLGraph update plan entry {entry_idx} key {key} "
	260	+ f"has unsupported source kind {kind!r}",
	261	+ pta_error(ErrCode.PARAM),
	262	+ )

Mtorch_npu/npu/_graph_tree.py+29-1

	@@ -101,9 +101,15 @@ from torch.utils._ordered_set import OrderedSet
101	from torch.utils.weak import TensorWeakRef	101	from torch.utils.weak import TensorWeakRef
102		102
103	import torch_npu	103	import torch_npu
		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,
106	_set_cached_tensors_enabled as _set_cached_tensors_enabled)	107	_set_cached_tensors_enabled as _set_cached_tensors_enabled)
		108	+from torch_npu.npu._aclgraph_update_plan.resolver import (
		109	+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
		110	+ update_aclgraph_records_for_graph,
		111	+ validate_aclgraph_update_plan_for_graph,
		112	+)
107	import torch_npu.npu.aclnn	113	import torch_npu.npu.aclnn
108		114
109	if TYPE_CHECKING:	115	if TYPE_CHECKING:
	@@ -771,6 +777,9 @@ class NPUGraphNode:
771	if not isinstance(inputs, (list, tuple)):	777	if not isinstance(inputs, (list, tuple)):
772	raise RuntimeError("check isinstance(inputs, (list, tuple))")	778	raise RuntimeError("check isinstance(inputs, (list, tuple))")
773	self.wrapped_function = wrapped_function	779	self.wrapped_function = wrapped_function
		780	+ self.aclgraph_update_plan = getattr(
		781	+ wrapped_function.model, ACLGRAPH_UPDATE_PLAN_GLOBAL, None
		782	+ ) or []
774	self.id = graph_id	783	self.id = graph_id
775	self.device = device_index	784	self.device = device_index
776	self.stack_traces = stack_traces	785	self.stack_traces = stack_traces
	@@ -1065,10 +1074,19 @@ class NPUGraphNode:
1065	def run(self, new_inputs: List[InputType]) -> OutputType:	1074	def run(self, new_inputs: List[InputType]) -> OutputType:
1066	log.debug("NPUGRAPH-TREE Node Run node=%s", self.id)	1075	log.debug("NPUGRAPH-TREE Node Run node=%s", self.id)
1067	self.check_static_inputs_are_stable(new_inputs)	1076	self.check_static_inputs_are_stable(new_inputs)
1068	-	1077	+ aclgraph_update_submitted = update_aclgraph_records_for_graph(
		1078	+ self.aclgraph_update_plan,
		1079	+ self.graph,
		1080	+ new_inputs,
		1081	+ )
1069	self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs)	1082	self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs)
1070		1083
1071	self.run_graph()	1084	self.run_graph()
		1085	+ if aclgraph_update_submitted:
		1086	+ # Ensure the next ACLGraph update does not record reusable external events before this replay resets them.
		1087	+ self.graph.graph_dispatch_mode.update_stream.wait_stream(
		1088	+ torch.npu.current_stream()
		1089	+ )
1072		1090
1073	outputs = self.reconstruct_outputs()	1091	outputs = self.reconstruct_outputs()
1074	new_inputs.clear()	1092	new_inputs.clear()
	@@ -1232,6 +1250,8 @@ class NPUGraphNode:
1232		1250
1233	check_memory_pool(self.device, self.npu_graphs_pool, memory)	1251	check_memory_pool(self.device, self.npu_graphs_pool, memory)
1234		1252
		1253	+ aclgraph_update_inputs = list(inputs)
		1254	+
1235	with preserve_rng_state(), torch.npu.device(	1255	with preserve_rng_state(), torch.npu.device(
1236	self.device	1256	self.device
1237	), clear_cublas_manager(), torch.npu.graph(	1257	), clear_cublas_manager(), torch.npu.graph(
	@@ -1242,6 +1262,14 @@ class NPUGraphNode:
1242	), get_history_recording():	1262	), get_history_recording():
1243	static_outputs = model(inputs)	1263	static_outputs = model(inputs)
1244		1264
		1265	+ validate_aclgraph_update_plan_for_graph(self.aclgraph_update_plan, self.graph)
		1266	+ update_aclgraph_records_for_graph(
		1267	+ self.aclgraph_update_plan,
		1268	+ self.graph,
		1269	+ aclgraph_update_inputs,
		1270	+ )
		1271	+ aclgraph_update_inputs.clear()
		1272	+
1245	# running model should reclaim memory	1273	# running model should reclaim memory
1246	if not len(inputs) == 0:	1274	if not len(inputs) == 0:
1247	raise RuntimeError("check len(inputs) == 0 fail")	1275	raise RuntimeError("check len(inputs) == 0 fail")

Mtorch_npu/npu/_npugraph_handlers/_fa3_graph_handler.py+30-14

	@@ -43,6 +43,21 @@ class _FA3TensorListOutHandler(NpuGraphOpHandler):
43	class FA3ForwardHandler(_FA3TensorListOutHandler):	43	class FA3ForwardHandler(_FA3TensorListOutHandler):
44	"""FA v3 forward: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""	44	"""FA v3 forward: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""
45		45
		46	+ UPDATE_SPECS = {
		47	+ "npu_fusion_attention_v3": [
		48	+ ("arg", 14, "actual_seq_qlen"),
		49	+ ("arg", 15, "actual_seq_kvlen"),
		50	+ ],
		51	+ "npu_fusion_attention_v3.default": [
		52	+ ("arg", 14, "actual_seq_qlen"),
		53	+ ("arg", 15, "actual_seq_kvlen"),
		54	+ ],
		55	+ "npu_fusion_attention_v3.out": [
		56	+ ("arg", 14, "actual_seq_qlen"),
		57	+ ("arg", 15, "actual_seq_kvlen"),
		58	+ ],
		59	+ }
		60	+
46	@classmethod	61	@classmethod
47	def should_handle(cls, func, args, kwargs):	62	def should_handle(cls, func, args, kwargs):
48	"""BNSD layout bypasses handler entirely — no dispatch record, no update."""	63	"""BNSD layout bypasses handler entirely — no dispatch record, no update."""
	@@ -53,13 +68,6 @@ class FA3ForwardHandler(_FA3TensorListOutHandler):
53	return False	68	return False
54	return True	69	return True
55		70
56	- @classmethod
57	- def update_args(cls, record, update_input):
58	- if "actual_seq_qlen" in update_input and len(record.args) > 14:
59	- record.args[14] = update_input["actual_seq_qlen"]
60	- if "actual_seq_kvlen" in update_input and len(record.args) > 15:
61	- record.args[15] = update_input["actual_seq_kvlen"]
62	-
63	@classmethod	71	@classmethod
64	def prepare_capture(cls, func, args, kwargs):	72	def prepare_capture(cls, func, args, kwargs):
65	func_out = torch_npu.npu_fusion_attention_v3.out	73	func_out = torch_npu.npu_fusion_attention_v3.out
	@@ -126,6 +134,21 @@ class FA3ForwardHandler(_FA3TensorListOutHandler):
126	class FA3BackwardHandler(_FA3TensorListOutHandler):	134	class FA3BackwardHandler(_FA3TensorListOutHandler):
127	"""FA v3 backward: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""	135	"""FA v3 backward: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""
128		136
		137	+ UPDATE_SPECS = {
		138	+ "npu_fusion_attention_grad_v3": [
		139	+ ("arg", 21, "actual_seq_qlen"),
		140	+ ("arg", 22, "actual_seq_kvlen"),
		141	+ ],
		142	+ "npu_fusion_attention_grad_v3.default": [
		143	+ ("arg", 21, "actual_seq_qlen"),
		144	+ ("arg", 22, "actual_seq_kvlen"),
		145	+ ],
		146	+ "npu_fusion_attention_grad_v3.out": [
		147	+ ("arg", 21, "actual_seq_qlen"),
		148	+ ("arg", 22, "actual_seq_kvlen"),
		149	+ ],
		150	+ }
		151	+
129	@classmethod	152	@classmethod
130	def should_handle(cls, func, args, kwargs):	153	def should_handle(cls, func, args, kwargs):
131	"""BNSD layout bypasses handler entirely — no dispatch record, no update."""	154	"""BNSD layout bypasses handler entirely — no dispatch record, no update."""
	@@ -136,13 +159,6 @@ class FA3BackwardHandler(_FA3TensorListOutHandler):
136	return False	159	return False
137	return True	160	return True
138		161
139	- @classmethod
140	- def update_args(cls, record, update_input):
141	- if "actual_seq_qlen" in update_input and len(record.args) > 21:
142	- record.args[21] = update_input["actual_seq_qlen"]
143	- if "actual_seq_kvlen" in update_input and len(record.args) > 22:
144	- record.args[22] = update_input["actual_seq_kvlen"]
145	-
146	@classmethod	162	@classmethod
147	def prepare_capture(cls, func, args, kwargs):	163	def prepare_capture(cls, func, args, kwargs):
148	func_out = torch_npu.npu_fusion_attention_grad_v3.out	164	func_out = torch_npu.npu_fusion_attention_grad_v3.out

Mtorch_npu/npu/_npugraph_handlers/ifa_handler.py+30-9

	@@ -6,7 +6,8 @@ This module defines the NPU Graph operator handlers for the
6		6
7	Structure: ``_TensorListOutHandler`` provides ``postprocess_result`` (return	7	Structure: ``_TensorListOutHandler`` provides ``postprocess_result`` (return
8	kwargs["out"]). ``IFAv1DefaultHandler`` and ``IFAv2DefaultHandler`` inherit	8	kwargs["out"]). ``IFAv1DefaultHandler`` and ``IFAv2DefaultHandler`` inherit
9	-it and each implement ``update_args`` and ``prepare_capture``; both	9	+it and declare ``UPDATE_SPECS`` + ``prepare_capture``; the spec-driven
		10	+``update_args`` from the base class handles all updates uniformly. Both
10	``.default`` and ``.out`` are registered on the same handler class.	11	``.default`` and ``.out`` are registered on the same handler class.
11	"""	12	"""
12	__all__ = []	13	__all__ = []
	@@ -37,10 +38,20 @@ class _TensorListOutHandler(NpuGraphOpHandler):
37	class _IFAv1DefaultHandler(_TensorListOutHandler):	38	class _IFAv1DefaultHandler(_TensorListOutHandler):
38	"""IFA v1: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""	39	"""IFA v1: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""
39		40
40	- @classmethod	41	+ UPDATE_SPECS = {
41	- def update_args(cls, record, update_input):	42	+ "npu_fused_infer_attention_score": [
42	- if "actual_seq_lengths_kv" in update_input and len(record.args) >= 7:	43	+ ("arg", 5, "actual_seq_lengths"),
43	- record.args[6] = update_input["actual_seq_lengths_kv"]	44	+ ("arg", 6, "actual_seq_lengths_kv"),
		45	+ ],
		46	+ "npu_fused_infer_attention_score.default": [
		47	+ ("arg", 5, "actual_seq_lengths"),
		48	+ ("arg", 6, "actual_seq_lengths_kv"),
		49	+ ],
		50	+ "npu_fused_infer_attention_score.out": [
		51	+ ("arg", 5, "actual_seq_lengths"),
		52	+ ("arg", 6, "actual_seq_lengths_kv"),
		53	+ ],
		54	+ }
44		55
45	@classmethod	56	@classmethod
46	def prepare_capture(cls, func, args, kwargs):	57	def prepare_capture(cls, func, args, kwargs):
	@@ -82,10 +93,20 @@ class _IFAv1DefaultHandler(_TensorListOutHandler):
82	class _IFAv2DefaultHandler(_TensorListOutHandler):	93	class _IFAv2DefaultHandler(_TensorListOutHandler):
83	"""IFA v2: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""	94	"""IFA v2: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""
84		95
85	- @classmethod	96	+ UPDATE_SPECS = {
86	- def update_args(cls, record, update_input):	97	+ "npu_fused_infer_attention_score_v2": [
87	- if "actual_seq_kvlen" in update_input and len(record.args) >= 9:	98	+ ("arg", 7, "actual_seq_qlen"),
88	- record.args[8] = update_input["actual_seq_kvlen"]	99	+ ("arg", 8, "actual_seq_kvlen"),
		100	+ ],
		101	+ "npu_fused_infer_attention_score_v2.default": [
		102	+ ("arg", 7, "actual_seq_qlen"),
		103	+ ("arg", 8, "actual_seq_kvlen"),
		104	+ ],
		105	+ "npu_fused_infer_attention_score_v2.out": [
		106	+ ("arg", 7, "actual_seq_qlen"),
		107	+ ("arg", 8, "actual_seq_kvlen"),
		108	+ ],
		109	+ }
89		110
90	@classmethod	111	@classmethod
91	def prepare_capture(cls, func, args, kwargs):	112	def prepare_capture(cls, func, args, kwargs):

Mtorch_npu/npu/_npugraph_handlers/npugraph_handler.py+50-13

	@@ -42,20 +42,34 @@ class NpuGraphOpHandler:
42	parameter is ``cls``, not ``self``). There is no instance; the global	42	parameter is ``cls``, not ``self``). There is no instance; the global
43	registry stores class objects directly. This structurally prevents	43	registry stores class objects directly. This structurally prevents
44	storing mutable per-invocation state. Class-level constants (e.g.	44	storing mutable per-invocation state. Class-level constants (e.g.
45	- ``_OP_ARG_SPECS``) are accessible via ``cls``.	45	+ ``UPDATE_SPECS``) are accessible via ``cls``.
46		46
47	- Users should inherit this class and override the needed hooks:	47	+ Declarative update_args via UPDATE_SPECS: Subclasses declare which
		48	+ update keys they consume and where those keys live in args/kwargs via the
		49	+ ``UPDATE_SPECS`` class attribute. The base ``update_args`` implementation
		50	+ walks the spec and applies updates uniformly. Subclasses normally do not
		51	+ need to override ``update_args``.
		52	+
		53	+ Schema:
		54	+
		55	+ .. code-block:: python
		56	+
		57	+ UPDATE_SPECS: Dict[op_name, List[Tuple[Literal["arg", "kwarg"], int \| str, key_name]]]
		58	+
		59	+ Example:
48		60
49	.. code-block:: python	61	.. code-block:: python
50		62
51	@register_npu_graph_handler(["my_op", "my_op.default"])	63	@register_npu_graph_handler(["my_op", "my_op.default"])
52	class MyHandler(NpuGraphOpHandler):	64	class MyHandler(NpuGraphOpHandler):
53	- @classmethod	65	+ UPDATE_SPECS = {
54	- def update_args(cls, record, update_input):	66	+ "my_op": [("arg", 2, "batch")],
55	- if "batch" in update_input and len(record.args) >= 3:	67	+ "my_op.default": [("arg", 2, "batch")],
56	- record.args[2] = update_input["batch"]	68	+ }
57	"""	69	"""
58		70
		71	+ UPDATE_SPECS = {}
		72	+
59	@classmethod	73	@classmethod
60	def prepare_capture(cls, func, args, kwargs):	74	def prepare_capture(cls, func, args, kwargs):
61	r"""Prepare operator call before graph-task recording.	75	r"""Prepare operator call before graph-task recording.
	@@ -100,19 +114,42 @@ class NpuGraphOpHandler:
100	return result	114	return result
101		115
102	@classmethod	116	@classmethod
103	- def update_args(cls, dispatch_record, update_input):	117	+ def get_update_specs(cls, op_name):
104	- r"""Apply operator-specific indexed-arg updates.	118	+ r"""Return per-op update specs.
105		119
106	- Framework-level kwargs updates are handled by the dispatch skeleton.	120	+ Args:
107	- Override this hook only when update values must be applied to	121	+ op_name (str): Operator dispatch name (e.g. ``"_npu_paged_attention.default"``).
108	- arguments by index.	122	+
		123	+ Returns:
		124	+ List of ``(loc, idx_or_name, key)`` tuples. ``loc`` is ``"arg"`` or
		125	+ ``"kwarg"``; ``idx_or_name`` is an int index for ``"arg"`` or a
		126	+ string name for ``"kwarg"``; ``key`` is the user-facing update key.
		127	+ Empty list if this op is not in ``UPDATE_SPECS``.
		128	+ """
		129	+ return cls.UPDATE_SPECS.get(op_name, [])
		130	+
		131	+ @classmethod
		132	+ def update_args(cls, dispatch_record, update_input):
		133	+ r"""Apply operator-specific updates by walking ``UPDATE_SPECS``.
		134	+
		135	+ Default implementation reads the spec for this op and assigns the
		136	+ matching key's value from ``update_input`` to the recorded args/kwargs
		137	+ slot. Subclasses normally do not need to override this; declare
		138	+ ``UPDATE_SPECS`` instead.
109		139
110	Args:	140	Args:
111	dispatch_record (_GraphDispatchRecord): Recorded operator call.	141	dispatch_record (_GraphDispatchRecord): Recorded operator call.
112	- Args can be modified via ``dispatch_record.args[i]``.
113	update_input (dict): User-provided update payload.	142	update_input (dict): User-provided update payload.
114	"""	143	"""
115	- pass	144	+ specs = cls.get_update_specs(dispatch_record.op_cache_entry.__name__)
		145	+ for loc, idx_or_name, key in specs:
		146	+ if key not in update_input:
		147	+ continue
		148	+ if loc == "arg":
		149	+ if len(dispatch_record.args) > idx_or_name:
		150	+ dispatch_record.args[idx_or_name] = update_input[key]
		151	+ elif loc == "kwarg":
		152	+ dispatch_record.kwargs[idx_or_name] = update_input[key]
116		153
117	@classmethod	154	@classmethod
118	def record_wrap_kwarg(cls, key, value, tensor_param_names):	155	def record_wrap_kwarg(cls, key, value, tensor_param_names):

Mtorch_npu/npu/_npugraph_handlers/simple_handler.py+6-15

	@@ -12,21 +12,12 @@ from .npugraph_handler import NpuGraphOpHandler, register_npu_graph_handler
12	class _SimpleGraphHandler(NpuGraphOpHandler):	12	class _SimpleGraphHandler(NpuGraphOpHandler):
13	"""Handler for PA (Paged Attention) and MLA operators.	13	"""Handler for PA (Paged Attention) and MLA operators.
14		14
15	- Attributes:	15	+ Update behavior is declarative via :attr:`UPDATE_SPECS`. The base class's
16	- _OP_ARG_SPECS (dict[str, tuple[int, str]]): Specifies	16	+ spec-driven ``update_args`` walks this map and assigns the matching
17	- ``op_name -> (arg_index, update_key)`` for each supported	17	+ update_input key's value to the recorded arg slot.
18	- operator.
19	"""	18	"""
20		19
21	- _OP_ARG_SPECS = {	20	+ UPDATE_SPECS = {
22	- "_npu_paged_attention.default": (7, "context_lens"),	21	+ "_npu_paged_attention.default": [("arg", 7, "context_lens")],
23	- "npu_multi_head_latent_attention.out": (5, "context_lens"),	22	+ "npu_multi_head_latent_attention.out": [("arg", 5, "context_lens")],
24	}	23	}
25	-
26	- @classmethod
27	- def update_args(cls, record, update_input):
28	- spec = cls._OP_ARG_SPECS.get(record.op_cache_entry.__name__)
29	- if spec:
30	- arg_index, key = spec
31	- if key in update_input and len(record.args) >= (arg_index + 1):
32	- record.args[arg_index] = update_input[key]

Mtorch_npu/npu/graphs.py+0-1

	@@ -587,7 +587,6 @@ class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode):
587	graph_task_update_end(self.update_stream)	587	graph_task_update_end(self.update_stream)
588	record.event.record(self.update_stream)	588	record.event.record(self.update_stream)
589		589
590	-
591	# Python shim helps Sphinx process docstrings more reliably.	590	# Python shim helps Sphinx process docstrings more reliably.
592	class NPUGraph(torch_npu._C._NPUGraph):	591	class NPUGraph(torch_npu._C._NPUGraph):
593	r"""Wrapper around a NPU graph.	592	r"""Wrapper around a NPU graph.

Mtorch_npu/utils/_graph_tree.py+7-8

	@@ -110,6 +110,7 @@ def npugraphify(
110	mutated_input_idxs: Tuple[int, ...] = (),	110	mutated_input_idxs: Tuple[int, ...] = (),
111	) -> Callable[..., Any]:	111	) -> Callable[..., Any]:
112	from torch_npu.npu._graph_tree import npugraphify_impl as new_npugraphify_impl	112	from torch_npu.npu._graph_tree import npugraphify_impl as new_npugraphify_impl
		113	+
113	npugraphify_fn: Callable[..., Any]	114	npugraphify_fn: Callable[..., Any]
114	if config.triton.cudagraph_trees:	115	if config.triton.cudagraph_trees:
115	npugraphify_fn = functools.partial(	116	npugraphify_fn = functools.partial(
	@@ -389,13 +390,11 @@ def _apply_npugraph_tree_methods():
389	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes	390	torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes
390	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin	391	torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
391		392
392	- # Bridge upstream callers of `torch._inductor.cudagraph_trees.get_manager`
393	- # to the NPU manager registry. The only upstream call sites are
394	- # `_inductor/output_code.py:maybe_handle_backward_generation` (used when
395	- # forward was cudagraph'd but backward is not, to drive the cudagraph
396	- # generation state machine) and `_dynamo/backends/cudagraphs.py`. NPU
397	- # registers its manager under `torch_npu.npu._graph_tree`, so without
398	- # this forward those upstream paths raise AttributeError or return None.
399	import torch._inductor.cudagraph_trees as _upstream_cgt # noqa: F401	393	import torch._inductor.cudagraph_trees as _upstream_cgt # noqa: F401
400	- from torch_npu.npu._graph_tree import get_manager as _npu_get_manager	394	+
		395	+ def _npu_get_manager(args, *kwargs):
		396	+ from torch_npu.npu._graph_tree import get_manager
		397	+
		398	+ return get_manager(args, *kwargs)
		399	+
401	_upstream_cgt.get_manager = _npu_get_manager	400	_upstream_cgt.get_manager = _npu_get_manager