Star857
1352
代码
代码
Issues1676
Pull Requests2249
项目讨论
Wiki
分析
Star857
1352
已合并
feat: Add ACLGraph update plans #36821
luochao60创建于 5月26日
feat: Add ACLGraph update plans #36821
已合并
luochao60创建于 5月26日
从luochao60/pytorch_9490Pynative_refactor_aclgraph_update_20260519_2.7.1合入到Ascend/pytorchv2.7.1
共 18 个文件变更+2112-124
Atest/_inductor/test_aclgraph_update_plan_compile.py+687-0
@@ -0,0 +1,687 @@
1+from unittest import mock
2+ 
3+import torch
4+from torch._inductor import config
5+from torch._inductor.codegen.common import IndentedBuffer
6+from torch._inductor.virtualized import V
7+from torch_npu._inductor._aclgraph_update_plan import (
8+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
9+)
10+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.wrapper import (
11+ NpuMlirSubgraphPythonWrapperCodegen,
12+ NpuMlirWrapperCodeGen,
13+)
14+from torch_npu._inductor.codegen.wrapper import (
15+ _NPUKernelCodegenMixin,
16+ NPUSubgraphPythonWrapperCodegen,
17+)
18+from torch.testing._internal.common_utils import run_tests
19+ 
20+import torch_npu
21+import torch_npu._inductor
22+from torch_npu.testing.common_utils import SupportedDevices
23+ 
24+from testutils import TestUtils
25+ 
26+ 
27+def _make_ifa_inputs():
28+ q = torch.randn(1, 32, 1, 128, dtype=torch.float16, device="npu")
29+ k = torch.randn(1, 32, 1, 128, dtype=torch.float16, device="npu")
30+ v = torch.randn(1, 32, 1, 128, dtype=torch.float16, device="npu")
31+ return q, k, v
32+ 
33+ 
34+def _ifa_with_const_actual_seq_lengths(query, key, value):
35+ out, _ = torch_npu.npu_fused_infer_attention_score(
36+ query,
37+ key,
38+ value,
39+ num_heads=32,
40+ input_layout="BNSD",
41+ scale=128.0,
42+ pre_tokens=65535,
43+ next_tokens=65535,
44+ softmax_lse_flag=False,
45+ actual_seq_lengths=[37],
46+ )
47+ return out
48+ 
49+ 
50+def _ifa_with_actual_seq_lengths_kv(query, key, value):
51+ out, _ = torch_npu.npu_fused_infer_attention_score(
52+ query,
53+ key,
54+ value,
55+ num_heads=32,
56+ input_layout="BNSD",
57+ scale=128.0,
58+ pre_tokens=65535,
59+ next_tokens=65535,
60+ softmax_lse_flag=False,
61+ actual_seq_lengths=[37],
62+ actual_seq_lengths_kv=[1],
63+ )
64+ return out
65+ 
66+ 
67+def _ifa_with_runtime_actual_seq_lengths(query, key, value, actual_seq_lengths):
68+ out, _ = torch_npu.npu_fused_infer_attention_score(
69+ query,
70+ key,
71+ value,
72+ num_heads=32,
73+ input_layout="BNSD",
74+ scale=128.0,
75+ pre_tokens=65535,
76+ next_tokens=65535,
77+ softmax_lse_flag=False,
78+ actual_seq_lengths=actual_seq_lengths,
79+ )
80+ return out
81+ 
82+ 
83+def _ifa_v2_with_const_actual_seq_qlen(query, key, value):
84+ out, _ = torch_npu.npu_fused_infer_attention_score_v2(
85+ query,
86+ key,
87+ value,
88+ num_query_heads=32,
89+ input_layout="BNSD",
90+ softmax_scale=128.0,
91+ pre_tokens=65535,
92+ next_tokens=65535,
93+ return_softmax_lse=False,
94+ actual_seq_qlen=[1],
95+ )
96+ return out
97+ 
98+ 
99+def _two_ifa_with_const_actual_seq_lengths(query, key, value):
100+ out1, _ = torch_npu.npu_fused_infer_attention_score(
101+ query,
102+ key,
103+ value,
104+ num_heads=32,
105+ input_layout="BNSD",
106+ scale=128.0,
107+ pre_tokens=65535,
108+ next_tokens=65535,
109+ softmax_lse_flag=False,
110+ actual_seq_lengths=[37],
111+ )
112+ out2, _ = torch_npu.npu_fused_infer_attention_score(
113+ query,
114+ key,
115+ out1,
116+ num_heads=32,
117+ input_layout="BNSD",
118+ scale=128.0,
119+ pre_tokens=65535,
120+ next_tokens=65535,
121+ softmax_lse_flag=False,
122+ actual_seq_lengths=[41],
123+ )
124+ return out2
125+ 
126+ 
127+def _run_and_get_code_without_reset(fn, *args):
128+ from torch._inductor.graph import GraphLowering
129+ 
130+ source_codes = []
131+ 
132+ def save_output_code(code):
133+ source_codes.append(code)
134+ 
135+ with mock.patch.object(GraphLowering, "save_output_code", save_output_code):
136+ result = fn(*args)
137+ return result, source_codes
138+ 
139+ 
140+def _compiled_code(
141+ fn,
142+ *args,
143+ cudagraphs=True,
144+ cudagraph_trees=True,
145+ graph_partition=False,
146+ npu_backend=None,
147+):
148+ torch._dynamo.reset()
149+ old_cudagraphs = config.triton.cudagraphs
150+ old_cudagraph_trees = config.triton.cudagraph_trees
151+ old_force_disable_caches = config.force_disable_caches
152+ old_graph_partition = config.graph_partition
153+ try:
154+ config.triton.cudagraphs = cudagraphs
155+ config.triton.cudagraph_trees = cudagraph_trees
156+ config.force_disable_caches = True
157+ config.graph_partition = graph_partition
158+ options = {"npu_backend": npu_backend} if npu_backend is not None else None
159+ compiled = torch.compile(
160+ fn,
161+ backend="inductor",
162+ fullgraph=True,
163+ options=options,
164+ )
165+ _, codes = _run_and_get_code_without_reset(compiled, *args)
166+ finally:
167+ config.triton.cudagraphs = old_cudagraphs
168+ config.triton.cudagraph_trees = old_cudagraph_trees
169+ config.force_disable_caches = old_force_disable_caches
170+ config.graph_partition = old_graph_partition
171+ torch._dynamo.reset()
172+ return "\n".join(codes)
173+ 
174+ 
175+class TestACLGraphUpdatePlanCompile(TestUtils):
176+ 
177+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
178+ def test_ifa_no_graph_partition_codegen_attaches_plan_to_call_function(self):
179+ torch.npu.set_device(0)
180+ torch._dynamo.reset()
181+ 
182+ code = _compiled_code(
183+ _ifa_with_const_actual_seq_lengths,
184+ *_make_ifa_inputs(),
185+ graph_partition=False,
186+ )
187+ 
188+ self.assertIn("_torch_npu_aclgraph_update_plan", code)
189+ self.assertIn("def call(args):", code)
190+ self.assertIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
191+ self.assertEqual(code.count(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}"), 1)
192+ self.assertNotIn("def partition_0(args):", code)
193+ self.assertNotIn(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
194+ self.assertNotIn(f"\n{ACLGRAPH_UPDATE_PLAN_GLOBAL} = ", code)
195+ self.assertIn("npu_fused_infer_attention_score.default", code)
196+ self.assertIn("actual_seq_lengths", code)
197+ self.assertIn("'value': 37", code)
198+ 
199+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
200+ def test_ifa_codegen_plan_survives_generated_code_reload(self):
201+ torch.npu.set_device(0)
202+ torch._dynamo.reset()
203+ 
204+ code = _compiled_code(_ifa_with_const_actual_seq_lengths, *_make_ifa_inputs())
205+ namespace = {}
206+ exec(compile(code, "<aclgraph_update_plan_test>", "exec"), namespace)
207+ 
208+ self.assertEqual(
209+ getattr(namespace["call"], ACLGRAPH_UPDATE_PLAN_GLOBAL),
210+ [
211+ {
212+ "op": "npu_fused_infer_attention_score.default",
213+ "updates": {
214+ "actual_seq_lengths": {
215+ "kind": "list",
216+ "items": [{"kind": "constant", "value": 37}],
217+ }
218+ },
219+ },
220+ ],
221+ )
222+ 
223+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
224+ def test_ifa_codegen_emits_multiple_actual_seq_keys(self):
225+ torch.npu.set_device(0)
226+ torch._dynamo.reset()
227+ 
228+ code = _compiled_code(_ifa_with_actual_seq_lengths_kv, *_make_ifa_inputs())
229+ 
230+ self.assertIn("actual_seq_lengths", code)
231+ self.assertIn("actual_seq_lengths_kv", code)
232+ self.assertIn("'value': 37", code)
233+ self.assertIn("'value': 1", code)
234+ 
235+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
236+ def test_ifa_codegen_recompiles_when_guarded_actual_seq_list_changes(self):
237+ torch.npu.set_device(0)
238+ torch._dynamo.reset()
239+ 
240+ old_cudagraphs = config.triton.cudagraphs
241+ old_cudagraph_trees = config.triton.cudagraph_trees
242+ old_force_disable_caches = config.force_disable_caches
243+ try:
244+ config.triton.cudagraphs = True
245+ config.triton.cudagraph_trees = True
246+ config.force_disable_caches = True
247+ compiled = torch.compile(
248+ _ifa_with_runtime_actual_seq_lengths,
249+ backend="inductor",
250+ fullgraph=True,
251+ )
252+ code_37 = "\n".join(
253+ _run_and_get_code_without_reset(compiled, *_make_ifa_inputs(), [37])[1]
254+ )
255+ code_41 = "\n".join(
256+ _run_and_get_code_without_reset(compiled, *_make_ifa_inputs(), [41])[1]
257+ )
258+ finally:
259+ config.triton.cudagraphs = old_cudagraphs
260+ config.triton.cudagraph_trees = old_cudagraph_trees
261+ config.force_disable_caches = old_force_disable_caches
262+ torch._dynamo.reset()
263+ 
264+ self.assertIn("'value': 37", code_37)
265+ self.assertNotIn("'value': 37", code_41)
266+ self.assertTrue(
267+ "'value': 41" in code_41
268+ or ("'kind': 'input'" in code_41 and "'index': 3" in code_41)
269+ )
270+ 
271+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
272+ def test_ifa_v2_codegen_emits_aclgraph_update_plan(self):
273+ torch.npu.set_device(0)
274+ torch._dynamo.reset()
275+ 
276+ code = _compiled_code(_ifa_v2_with_const_actual_seq_qlen, *_make_ifa_inputs())
277+ 
278+ self.assertIn("_torch_npu_aclgraph_update_plan", code)
279+ self.assertIn("npu_fused_infer_attention_score_v2.default", code)
280+ self.assertIn("actual_seq_qlen", code)
281+ self.assertIn("'value': 1", code)
282+ 
283+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
284+ def test_ifa_graph_partition_codegen_attaches_plan_to_partition_function(self):
285+ class Graph:
286+ cpp_wrapper = False
287+ disable_cudagraphs_reason = None
288+ 
289+ plan = [
290+ {
291+ "op": "npu_fused_infer_attention_score.default",
292+ "updates": {
293+ "actual_seq_lengths": {
294+ "kind": "list",
295+ "items": [{"kind": "constant", "value": 37}],
296+ }
297+ },
298+ }
299+ ]
300+ 
301+ old_cudagraphs = config.triton.cudagraphs
302+ old_cudagraph_trees = config.triton.cudagraph_trees
303+ old_graph_partition = config.graph_partition
304+ try:
305+ config.triton.cudagraphs = True
306+ config.triton.cudagraph_trees = True
307+ config.graph_partition = True
308+ 
309+ wrapper = object.__new__(NPUSubgraphPythonWrapperCodegen)
310+ wrapper.launcher_fn_name = "partition_0"
311+ with V.set_graph_handler(Graph()):
312+ wrapper.torch_npu_aclgraph_update_plan = plan
313+ result = IndentedBuffer()
314+ wrapper.generate_after_suffix(result)
315+ finally:
316+ config.triton.cudagraphs = old_cudagraphs
317+ config.triton.cudagraph_trees = old_cudagraph_trees
318+ config.graph_partition = old_graph_partition
319+ 
320+ code = result.getvalue()
321+ self.assertIn(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
322+ self.assertEqual(code.count(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}"), 1)
323+ self.assertNotIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
324+ self.assertNotIn(f"\n{ACLGRAPH_UPDATE_PLAN_GLOBAL} = ", code)
325+ self.assertIn(repr(plan), code)
326+ 
327+ def test_mlir_dvm_wrapper_appends_aclgraph_update_plan_for_extern_kernel(self):
328+ class Graph:
329+ cpp_wrapper = False
330+ disable_cudagraphs_reason = None
331+ 
332+ class Arg:
333+ def __init__(self, name):
334+ self.name = name
335+ 
336+ class Schema:
337+ arguments = [
338+ Arg("query"),
339+ Arg("key"),
340+ Arg("value"),
341+ Arg("num_heads"),
342+ Arg("input_layout"),
343+ Arg("actual_seq_lengths"),
344+ ]
345+ 
346+ class Target:
347+ __name__ = "npu_fused_infer_attention_score.default"
348+ _schema = Schema()
349+ 
350+ class Value:
351+ def __init__(self, name):
352+ self.name = name
353+ 
354+ def get_name(self):
355+ return self.name
356+ 
357+ class Kernel:
358+ op_overload = Target()
359+ inputs = [Value("arg0_1"), Value("arg1_1"), Value("arg2_1")]
360+ constant_args = [32, "BNSD", [37]]
361+ kwargs = {}
362+ layout = object()
363+ 
364+ def get_name(self):
365+ return "buf0"
366+ 
367+ def get_origin_node(self):
368+ return None
369+ 
370+ def get_kernel_name(self):
371+ return "torch.ops.npu.npu_fused_infer_attention_score.default"
372+ 
373+ old_cudagraphs = config.triton.cudagraphs
374+ old_cudagraph_trees = config.triton.cudagraph_trees
375+ old_graph_partition = config.graph_partition
376+ try:
377+ config.triton.cudagraphs = True
378+ config.triton.cudagraph_trees = True
379+ config.graph_partition = False
380+ 
381+ wrapper = object.__new__(NpuMlirWrapperCodeGen)
382+ wrapper.launcher_fn_name = "call"
383+ wrapper.declare = ""
384+ wrapper.ending = ""
385+ wrapper.supports_intermediate_hooks = False
386+ wrapper.get_graph_input_names = lambda: ["arg0_1", "arg1_1", "arg2_1"]
387+ wrapper.get_graph_inputs = lambda: {}
388+ wrapper.writeline = lambda line: None
389+ 
390+ with V.set_graph_handler(Graph()):
391+ wrapper.generate_extern_kernel_alloc(Kernel(), [])
392+ result = IndentedBuffer()
393+ wrapper.generate_after_suffix(result)
394+ finally:
395+ config.triton.cudagraphs = old_cudagraphs
396+ config.triton.cudagraph_trees = old_cudagraph_trees
397+ config.graph_partition = old_graph_partition
398+ 
399+ self.assertEqual(
400+ wrapper.torch_npu_aclgraph_update_plan,
401+ [
402+ {
403+ "op": "npu_fused_infer_attention_score.default",
404+ "updates": {
405+ "actual_seq_lengths": {
406+ "kind": "list",
407+ "items": [{"kind": "constant", "value": 37}],
408+ }
409+ },
410+ }
411+ ],
412+ )
413+ code = result.getvalue()
414+ self.assertIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
415+ self.assertIn(repr(wrapper.torch_npu_aclgraph_update_plan), code)
416+ 
417+ def test_mlir_dvm_subgraph_wrapper_emits_aclgraph_update_plan(self):
418+ class Graph:
419+ cpp_wrapper = False
420+ disable_cudagraphs_reason = None
421+ 
422+ plan = [
423+ {
424+ "op": "npu_fused_infer_attention_score.default",
425+ "updates": {
426+ "actual_seq_lengths": {
427+ "kind": "list",
428+ "items": [{"kind": "constant", "value": 37}],
429+ }
430+ },
431+ }
432+ ]
433+ 
434+ old_cudagraphs = config.triton.cudagraphs
435+ old_cudagraph_trees = config.triton.cudagraph_trees
436+ old_graph_partition = config.graph_partition
437+ try:
438+ config.triton.cudagraphs = True
439+ config.triton.cudagraph_trees = True
440+ config.graph_partition = True
441+ 
442+ with V.set_graph_handler(Graph()):
443+ wrapper = object.__new__(NpuMlirSubgraphPythonWrapperCodegen)
444+ wrapper.launcher_fn_name = "partition_0"
445+ wrapper.subgraph_name = "partition_0"
446+ wrapper.torch_npu_aclgraph_update_plan = plan
447+ result = IndentedBuffer()
448+ wrapper.generate_after_suffix(result)
449+ finally:
450+ config.triton.cudagraphs = old_cudagraphs
451+ config.triton.cudagraph_trees = old_cudagraph_trees
452+ config.graph_partition = old_graph_partition
453+ 
454+ code = result.getvalue()
455+ self.assertIn(f"partition_0.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
456+ self.assertNotIn(f"call.{ACLGRAPH_UPDATE_PLAN_GLOBAL}", code)
457+ self.assertIn(repr(plan), code)
458+ 
459+ def test_mlir_dvm_wrapper_does_not_inherit_default_npu_codegen_mixin(self):
460+ self.assertFalse(issubclass(NpuMlirWrapperCodeGen, _NPUKernelCodegenMixin))
461+ self.assertFalse(issubclass(NpuMlirSubgraphPythonWrapperCodegen, _NPUKernelCodegenMixin))
462+ 
463+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
464+ def test_ifa_codegen_preserves_multiple_plan_entry_order(self):
465+ torch.npu.set_device(0)
466+ torch._dynamo.reset()
467+ 
468+ code = _compiled_code(
469+ _two_ifa_with_const_actual_seq_lengths,
470+ *_make_ifa_inputs(),
471+ )
472+ 
473+ self.assertGreaterEqual(
474+ code.count("'op': 'npu_fused_infer_attention_score.default'"),
475+ 2,
476+ )
477+ first_update = code.find("'value': 37")
478+ second_update = code.find("'value': 41")
479+ self.assertGreaterEqual(first_update, 0)
480+ self.assertGreater(second_update, first_update)
481+ 
482+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
483+ def test_ifa_codegen_skips_aclgraph_update_plan_without_cudagraphs(self):
484+ torch.npu.set_device(0)
485+ torch._dynamo.reset()
486+ 
487+ code = _compiled_code(
488+ _ifa_with_const_actual_seq_lengths,
489+ *_make_ifa_inputs(),
490+ cudagraphs=False,
491+ )
492+ 
493+ self.assertNotIn("_torch_npu_aclgraph_update_plan", code)
494+ 
495+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
496+ def test_ifa_codegen_skips_aclgraph_update_plan_without_cudagraph_trees(self):
497+ torch.npu.set_device(0)
498+ torch._dynamo.reset()
499+ 
500+ code = _compiled_code(
501+ _ifa_with_const_actual_seq_lengths,
502+ *_make_ifa_inputs(),
503+ cudagraph_trees=False,
504+ )
505+ 
506+ self.assertNotIn("_torch_npu_aclgraph_update_plan", code)
507+ 
508+ def test_wrapper_plan_gate_respects_cudagraph_disable_reason(self):
509+ from torch._inductor.virtualized import V
510+ from torch_npu._inductor._aclgraph_update_plan.codegen import (
511+ should_generate_inductor_aclgraph_update_plan,
512+ )
513+ 
514+ class Graph:
515+ cpp_wrapper = False
516+ disable_cudagraphs_reason = "unsupported"
517+ 
518+ old_cudagraphs = config.triton.cudagraphs
519+ old_cudagraph_trees = config.triton.cudagraph_trees
520+ try:
521+ config.triton.cudagraphs = True
522+ config.triton.cudagraph_trees = True
523+ with V.set_graph_handler(Graph()):
524+ self.assertFalse(should_generate_inductor_aclgraph_update_plan())
525+ finally:
526+ config.triton.cudagraphs = old_cudagraphs
527+ config.triton.cudagraph_trees = old_cudagraph_trees
528+ 
529+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
530+ def test_ifa_cudagraph_tree_receives_aclgraph_update_plan(self):
531+ torch.npu.set_device(0)
532+ torch._dynamo.reset()
533+ 
534+ import torch_npu.npu._graph_tree as graph_tree
535+ 
536+ old_cudagraphs = config.triton.cudagraphs
537+ old_cudagraph_trees = config.triton.cudagraph_trees
538+ old_force_disable_caches = config.force_disable_caches
539+ old_slow_path_asserts = config.triton.slow_path_cudagraph_asserts
540+ original_update = graph_tree.update_aclgraph_records_for_graph
541+ seen_plans = []
542+ 
543+ def collect_plan(plan, graph, inputs):
544+ seen_plans.append(plan)
545+ return original_update(plan, graph, inputs)
546+ 
547+ try:
548+ config.triton.cudagraphs = True
549+ config.triton.cudagraph_trees = True
550+ config.force_disable_caches = True
551+ config.triton.slow_path_cudagraph_asserts = False
552+ graph_tree.update_aclgraph_records_for_graph = collect_plan
553+ 
554+ compiled = torch.compile(
555+ _ifa_with_const_actual_seq_lengths,
556+ backend="inductor",
557+ fullgraph=True,
558+ )
559+ inputs = _make_ifa_inputs()
560+ expected = _ifa_with_const_actual_seq_lengths(*inputs)
561+ actual = compiled(*inputs)
562+ torch.testing.assert_close(
563+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
564+ )
565+ 
566+ inputs = _make_ifa_inputs()
567+ expected = _ifa_with_const_actual_seq_lengths(*inputs)
568+ actual = compiled(*inputs)
569+ torch.testing.assert_close(
570+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
571+ )
572+ finally:
573+ graph_tree.update_aclgraph_records_for_graph = original_update
574+ config.triton.cudagraphs = old_cudagraphs
575+ config.triton.cudagraph_trees = old_cudagraph_trees
576+ config.force_disable_caches = old_force_disable_caches
577+ config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts
578+ torch._dynamo.reset()
579+ 
580+ self.assertTrue(seen_plans)
581+ self.assertTrue(any(plan for plan in seen_plans))
582+ plan = next(plan for plan in seen_plans if plan)
583+ self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score.default")
584+ self.assertEqual(
585+ plan[0]["updates"]["actual_seq_lengths"],
586+ {"kind": "list", "items": [{"kind": "constant", "value": 37}]},
587+ )
588+ 
589+ def test_npugraphify_keeps_aclgraph_update_plan_on_callable_attribute(self):
590+ import torch_npu.npu._graph_tree as graph_tree
591+ 
592+ expected_plan = [{"op": "test.op", "updates": {}}]
593+ 
594+ def model(args):
595+ return args
596+ 
597+ setattr(model, ACLGRAPH_UPDATE_PLAN_GLOBAL, expected_plan)
598+ 
599+ captured = {}
600+ 
601+ def fake_add_function(*args, **kwargs):
602+ captured["arg_count"] = len(args)
603+ captured["model"] = args[0]
604+ return lambda inputs: inputs, []
605+ 
606+ manager = mock.Mock()
607+ manager.add_function.side_effect = fake_add_function
608+ with mock.patch(
609+ "torch_npu.npu._graph_tree.get_container",
610+ return_value=mock.Mock(get_tree_manager=mock.Mock(return_value=manager)),
611+ ):
612+ graph_tree.npugraphify(
613+ model,
614+ [],
615+ device_index=0,
616+ is_backward=False,
617+ is_inference=True,
618+ )
619+ 
620+ self.assertEqual(captured["arg_count"], 8)
621+ self.assertIs(
622+ getattr(captured["model"], ACLGRAPH_UPDATE_PLAN_GLOBAL),
623+ expected_plan,
624+ )
625+ 
626+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
627+ def test_ifa_v2_cudagraph_tree_receives_aclgraph_update_plan(self):
628+ torch.npu.set_device(0)
629+ torch._dynamo.reset()
630+ 
631+ import torch_npu.npu._graph_tree as graph_tree
632+ 
633+ old_cudagraphs = config.triton.cudagraphs
634+ old_cudagraph_trees = config.triton.cudagraph_trees
635+ old_force_disable_caches = config.force_disable_caches
636+ old_slow_path_asserts = config.triton.slow_path_cudagraph_asserts
637+ original_update = graph_tree.update_aclgraph_records_for_graph
638+ seen_plans = []
639+ 
640+ def collect_plan(plan, graph, inputs):
641+ seen_plans.append(plan)
642+ return original_update(plan, graph, inputs)
643+ 
644+ try:
645+ config.triton.cudagraphs = True
646+ config.triton.cudagraph_trees = True
647+ config.force_disable_caches = True
648+ config.triton.slow_path_cudagraph_asserts = False
649+ graph_tree.update_aclgraph_records_for_graph = collect_plan
650+ 
651+ compiled = torch.compile(
652+ _ifa_v2_with_const_actual_seq_qlen,
653+ backend="inductor",
654+ fullgraph=True,
655+ )
656+ inputs = _make_ifa_inputs()
657+ expected = _ifa_v2_with_const_actual_seq_qlen(*inputs)
658+ actual = compiled(*inputs)
659+ torch.testing.assert_close(
660+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
661+ )
662+ 
663+ inputs = _make_ifa_inputs()
664+ expected = _ifa_v2_with_const_actual_seq_qlen(*inputs)
665+ actual = compiled(*inputs)
666+ torch.testing.assert_close(
667+ actual.cpu(), expected.cpu(), rtol=1e-3, atol=1e-3
668+ )
669+ finally:
670+ graph_tree.update_aclgraph_records_for_graph = original_update
671+ config.triton.cudagraphs = old_cudagraphs
672+ config.triton.cudagraph_trees = old_cudagraph_trees
673+ config.force_disable_caches = old_force_disable_caches
674+ config.triton.slow_path_cudagraph_asserts = old_slow_path_asserts
675+ torch._dynamo.reset()
676+ 
677+ self.assertTrue(any(plan for plan in seen_plans))
678+ plan = next(plan for plan in seen_plans if plan)
679+ self.assertEqual(plan[0]["op"], "npu_fused_infer_attention_score_v2.default")
680+ self.assertEqual(
681+ plan[0]["updates"]["actual_seq_qlen"],
682+ {"kind": "list", "items": [{"kind": "constant", "value": 1}]},
683+ )
684+ 
685+ 
686+if __name__ == "__main__":
687+ 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 
59if __name__ == "__main__":81if __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+ 
17+ return _NPU_GRAPH_OP_HANDLERS.get(op_name)
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+34-7
@@ -6,13 +6,17 @@ import sympy
6from torch._inductor.virtualized import V6from torch._inductor.virtualized import V
7from torch._inductor import config, ir7from torch._inductor import config, ir
8from torch._inductor.codegen.wrapper import (8from torch._inductor.codegen.wrapper import (
9- PythonWrapperCodegen, 9+ PythonWrapperCodegen,
10- pexpr, 10+ pexpr,
11 cache_on_self,11 cache_on_self,
12 SubgraphPythonWrapperCodegen,12 SubgraphPythonWrapperCodegen,
13 counters,13 counters,
14)14)
15import torch_npu15import torch_npu
16+from torch_npu._inductor._aclgraph_update_plan import (
17+ append_inductor_aclgraph_update_plan_for_codegen_node,
18+ emit_inductor_aclgraph_update_plan_for_wrapper,
19+)
16 20 
17from ... import codecache21from ... import codecache
18 22 
@@ -25,9 +29,9 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
25 29 
26 @staticmethod30 @staticmethod
27 def create(31 def create(
28- is_subgraph: bool, 32+ is_subgraph: bool,
29- subgraph_name: str, 33+ subgraph_name: str,
30- parent_wrapper: PythonWrapperCodegen, 34+ parent_wrapper: PythonWrapperCodegen,
31 partition_signatures: Optional[ir.GraphPartitionSignature] = None35 partition_signatures: Optional[ir.GraphPartitionSignature] = None
32 ):36 ):
33 if is_subgraph:37 if is_subgraph:
@@ -39,7 +43,7 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
39 raise ValueError(43 raise ValueError(
40 "parent_wrapper must be provided for python wrapper"44 "parent_wrapper must be provided for python wrapper"
41 )45 )
42- return SubgraphPythonWrapperCodegen(46+ return NpuMlirSubgraphPythonWrapperCodegen(
43 subgraph_name, parent_wrapper, partition_signatures47 subgraph_name, parent_wrapper, partition_signatures
44 )48 )
45 return NpuMlirWrapperCodeGen()49 return NpuMlirWrapperCodeGen()
@@ -93,6 +97,7 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
93 )97 )
94 98 
95 def generate_extern_kernel_alloc(self, extern_kernel, args):99 def generate_extern_kernel_alloc(self, extern_kernel, args):
100+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
96 # If it's a NoneLayout then the extern_kernel should essentially be101 # If it's a NoneLayout then the extern_kernel should essentially be
97 # treated as if it doesn't return anything102 # treated as if it doesn't return anything
98 no_return = isinstance(extern_kernel.layout, ir.NoneLayout)103 no_return = isinstance(extern_kernel.layout, ir.NoneLayout)
@@ -125,6 +130,14 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
125 f"run_intermediate_hooks({origin_node.name!r}, {output_name})"130 f"run_intermediate_hooks({origin_node.name!r}, {output_name})"
126 )131 )
127 132 
133+ def generate_after_suffix(self, result) -> None:
134+ super().generate_after_suffix(result)
135+ emit_inductor_aclgraph_update_plan_for_wrapper(
136+ self,
137+ result,
138+ is_graph_partition_subgraph=False,
139+ )
140+ 
128 def write_get_raw_stream(self, device_idx: int, graph=None) -> str:141 def write_get_raw_stream(self, device_idx: int, graph=None) -> str:
129 self.write_triton_header_once()142 self.write_triton_header_once()
130 name = f"stream{device_idx}"143 name = f"stream{device_idx}"
@@ -133,7 +146,7 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
133 f"torch_npu.npu.set_device({device_idx})"146 f"torch_npu.npu.set_device({device_idx})"
134 )147 )
135 return name148 return name
136- 149+ 
137 def generate_kernel_call(150 def generate_kernel_call(
138 self,151 self,
139 kernel_name,152 kernel_name,
@@ -206,3 +219,17 @@ class NpuMlirWrapperCodeGen(PythonWrapperCodegen):
206 self.wrapper_call.writeline('exc_info=(None, None, None)')219 self.wrapper_call.writeline('exc_info=(None, None, None)')
207 self.wrapper_call.writeline('static_kernel_complier.__exit__(*exc_info)')220 self.wrapper_call.writeline('static_kernel_complier.__exit__(*exc_info)')
208 super().generate_return(output_refs)221 super().generate_return(output_refs)
222+ 
223+ 
224+class NpuMlirSubgraphPythonWrapperCodegen(SubgraphPythonWrapperCodegen):
225+ def generate_extern_kernel_alloc(self, extern_kernel, args):
226+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
227+ super().generate_extern_kernel_alloc(extern_kernel, args)
228+ 
229+ def generate_after_suffix(self, result) -> None:
230+ super().generate_after_suffix(result)
231+ emit_inductor_aclgraph_update_plan_for_wrapper(
232+ self,
233+ result,
234+ is_graph_partition_subgraph=True,
235+ )
Mtorch_npu/_inductor/codegen/wrapper.py+48-34
@@ -20,17 +20,45 @@ from torch._inductor.codegen.common import DeferredLine, WorkspaceArg, IndentedB
20from torch._inductor.codegen.wrapper import BufferLike, WrapperLine20from torch._inductor.codegen.wrapper import BufferLike, WrapperLine
21from torch._inductor import ir21from torch._inductor import ir
22import torch_npu.npu.aclnn22import torch_npu.npu.aclnn
23+from torch_npu._inductor._aclgraph_update_plan import (
24+ append_inductor_aclgraph_update_plan_for_codegen_node,
25+ emit_inductor_aclgraph_update_plan_for_wrapper,
26+)
23from ..fx_passes.utils.schedule_node_utils import is_multi_stream27from ..fx_passes.utils.schedule_node_utils import is_multi_stream
24 28 
25 29 
26-class NPUSubgraphPythonWrapperCodegen(SubgraphPythonWrapperCodegen):30+class _NPUKernelCodegenMixin:
31+ # This mixin must appear before the PyTorch wrapper base in the MRO so NPU
32+ # hooks run first, then cooperative super() continues into the base wrapper.
33+ # generate numel expr for range_tree_node
27 def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):34 def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
28 expr = f"{kernel_name}_{node.name}_numel"35 expr = f"{kernel_name}_{node.name}_numel"
29- self.writeline(f"{expr} = {pexpr(numel_expr)}")36+ simplified = V.graph.sizevars.simplify(numel_expr)
37+ # Ensure all PRECOMPUTED_SIZE symbols in the *simplified* expression
38+ # are defined before we emit the line that uses them.
39+ for sym in simplified.free_symbols:
40+ self.ensure_size_computed(sym)
41+ self.writeline(f"{expr} = {pexpr(simplified)}")
30 return SymbolicCallArg(expr, numel_expr)42 return SymbolicCallArg(expr, numel_expr)
31 43 
44+ # don't assert
45+ def codegen_input_size_asserts(self) -> None:
46+ pass
32 47 
33-class NPUWrapperCodeGen(PythonWrapperCodegen):48+ def generate_extern_kernel_alloc(self, extern_kernel, args):
49+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
50+ super().generate_extern_kernel_alloc(extern_kernel, args)
51+ 
52+ def generate_after_suffix(self, result: IndentedBuffer) -> None:
53+ super().generate_after_suffix(result)
54+ emit_inductor_aclgraph_update_plan_for_wrapper(
55+ self,
56+ result,
57+ _is_codegen_graph_partition_subgraph(self),
58+ )
59+ 
60+ 
61+class NPUWrapperCodeGen(_NPUKernelCodegenMixin, PythonWrapperCodegen):
34 def __init__(self):62 def __init__(self):
35 super().__init__()63 super().__init__()
36 self.buffer_args_multi_stream_intent = {}64 self.buffer_args_multi_stream_intent = {}
@@ -68,24 +96,6 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
68 "import torch_npu._inductor.runtime.triton_heuristics as triton_heuristics"96 "import torch_npu._inductor.runtime.triton_heuristics as triton_heuristics"
69 )97 )
70 98 
71- # generate numel expr for range_tree_node
72- def generate_node_numel_expr(self, kernel_name: str, node, numel_expr):
73- expr = f"{kernel_name}_{node.name}_numel"
74- simplified = V.graph.sizevars.simplify(numel_expr)
75- # Ensure all PRECOMPUTED_SIZE symbols in the *simplified* expression
76- # are defined before we emit the line that uses them.
77- for sym in simplified.free_symbols:
78- self.ensure_size_computed(sym)
79- self.writeline(f"{expr} = {pexpr(simplified)}")
80- # We can get symbolic expressions here, like s0*64
81- # It is fine to have them here, but we need to handle them correctly as their own type
82- # This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
83- # scalars as well.
84- # This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
85- # constant now, need type info. I agree, this needs type info, and while this is not true type info
86- # it suffices as a type hint for the purposes of producing the correct code for this type.
87- return SymbolicCallArg(expr, numel_expr)
88- 
89 def generate_save_uncompiled_kernels(self):99 def generate_save_uncompiled_kernels(self):
90 # remove incorrect grid=(0,0,0) param100 # remove incorrect grid=(0,0,0) param
91 self.wrapper_call.splice(101 self.wrapper_call.splice(
@@ -102,10 +112,6 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
102 """112 """
103 )113 )
104 114 
105- # don't assert
106- def codegen_input_size_asserts(self) -> None:
107- pass
108- 
109 def get_next_kernel_suffix(self) -> str:115 def get_next_kernel_suffix(self) -> str:
110 iter_val = copy.copy(self._names_iter)116 iter_val = copy.copy(self._names_iter)
111 return f"{next(iter_val)}"117 return f"{next(iter_val)}"
@@ -198,6 +204,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
198 204 
199 def generate_extern_kernel_alloc(self, extern_kernel, args):205 def generate_extern_kernel_alloc(self, extern_kernel, args):
200 if is_multi_stream():206 if is_multi_stream():
207+ append_inductor_aclgraph_update_plan_for_codegen_node(self, extern_kernel)
201 no_return = isinstance(extern_kernel.layout, NoneLayout)208 no_return = isinstance(extern_kernel.layout, NoneLayout)
202 output_name = extern_kernel.get_name()209 output_name = extern_kernel.get_name()
203 origin_node = extern_kernel.get_origin_node()210 origin_node = extern_kernel.get_origin_node()
@@ -364,8 +371,8 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
364 multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(new_name)371 multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(new_name)
365 return f"{multi_stream_intent_str}{self.declare_maybe_reference}{new_name} = {old_name}{del_line}{self.ending} {self.comment} reuse"372 return f"{multi_stream_intent_str}{self.declare_maybe_reference}{new_name} = {old_name}{del_line}{self.ending} {self.comment} reuse"
366 return super().codegen_exact_buffer_reuse(old_name, new_name, del_line)373 return super().codegen_exact_buffer_reuse(old_name, new_name, del_line)
367- 374+ 
368- 375+ 
369 def codegen_deferred_allocation(self, name: str, view: ir.ReinterpretView) -> None:376 def codegen_deferred_allocation(self, name: str, view: ir.ReinterpretView) -> None:
370 if is_multi_stream():377 if is_multi_stream():
371 multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(name)378 multi_stream_intent_str = self.get_buffer_define_multi_stream_by_name(name)
@@ -446,7 +453,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
446 if isinstance(line, str) and "main_stream = " in line:453 if isinstance(line, str) and "main_stream = " in line:
447 stream_line = idx454 stream_line = idx
448 break455 break
449- 456+ 
450 if stream_line == -1:457 if stream_line == -1:
451 for line in self.lines:458 for line in self.lines:
452 if isinstance(line, WrapperLine):459 if isinstance(line, WrapperLine):
@@ -508,8 +515,7 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
508 self.kernel_declarations.getvaluewithlinemap(),515 self.kernel_declarations.getvaluewithlinemap(),
509 )516 )
510 return super()._generate(is_inference)517 return super()._generate(is_inference)
511- 518+ 
512-
513 def handle_cross_stream_del_buf(self):519 def handle_cross_stream_del_buf(self):
514 total_lines = len(self.lines)520 total_lines = len(self.lines)
515 sub_streams_line_no = self.get_sub_streams_line_no()521 sub_streams_line_no = self.get_sub_streams_line_no()
@@ -521,12 +527,12 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
521 tab_value = self.buffer_args_multi_stream_intent[keys[0]]527 tab_value = self.buffer_args_multi_stream_intent[keys[0]]
522 if idx > sub_stream_line[0] and idx < sub_stream_line[1] and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() not in self.buffer_args_multi_stream_intent.keys():528 if idx > sub_stream_line[0] and idx < sub_stream_line[1] and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() not in self.buffer_args_multi_stream_intent.keys():
523 self.buffer_args_multi_stream_intent[line.node.get_name()] = tab_value529 self.buffer_args_multi_stream_intent[line.node.get_name()] = tab_value
524- 530+ 
525 n = len(sub_streams_line_no)531 n = len(sub_streams_line_no)
526 for i in range(1, n):532 for i in range(1, n):
527 prev_end = sub_streams_line_no[i-1][1]533 prev_end = sub_streams_line_no[i-1][1]
528 curr_start = sub_streams_line_no[i][0]534 curr_start = sub_streams_line_no[i][0]
529- 535+ 
530 if prev_end < idx < curr_start and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():536 if prev_end < idx < curr_start and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():
531 self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)537 self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)
532 538 
@@ -534,8 +540,8 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
534 last_end = sub_streams_line_no[-1][1]540 last_end = sub_streams_line_no[-1][1]
535 if last_end < idx < total_lines and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():541 if last_end < idx < total_lines and isinstance(line, WrapperLine) and hasattr(line, "node") and line.node.get_name() in self.buffer_args_multi_stream_intent.keys():
536 self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)542 self.buffer_args_multi_stream_intent.pop(line.node.get_name(), None)
537- 543+ 
538- 544+ 
539 def get_sub_streams_line_no(self):545 def get_sub_streams_line_no(self):
540 sub_streams_line_no = []546 sub_streams_line_no = []
541 i = 0547 i = 0
@@ -565,3 +571,11 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
565 continue571 continue
566 i += 1572 i += 1
567 return sub_streams_line_no573 return sub_streams_line_no
574+ 
575+ 
576+def _is_codegen_graph_partition_subgraph(wrapper: PythonWrapperCodegen) -> bool:
577+ return isinstance(wrapper, SubgraphPythonWrapperCodegen)
578+ 
579+ 
580+class NPUSubgraphPythonWrapperCodegen(_NPUKernelCodegenMixin, SubgraphPythonWrapperCodegen):
581+ pass
Mtorch_npu/_logging/_internal.py+1-2
@@ -41,9 +41,8 @@ def _add_logging_module():
41 torch._logging._internal.register_log("shmem", "torch_npu.symmetric_memory")41 torch._logging._internal.register_log("shmem", "torch_npu.symmetric_memory")
42 torch._logging._internal.register_log("env", "torch_npu.env")42 torch._logging._internal.register_log("env", "torch_npu.env")
43 torch._logging._internal.register_log("acl", "torch_npu.acl")43 torch._logging._internal.register_log("acl", "torch_npu.acl")
44- torch._logging._internal.register_log("aclgraph", "torch_npu.aclgraph")44+ torch._logging._internal.register_log("aclgraph", "torch_npu.npugraph")
45 torch._logging._internal.register_log("npugraph", "torch_npu.npugraph")45 torch._logging._internal.register_log("npugraph", "torch_npu.npugraph")
46- torch._logging._internal.register_log("cudagraphs", "torch_npu.npugraph")
47 46 
48 47 
49def _update_log_state_from_env():48def _update_log_state_from_env():
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+30-14
@@ -100,9 +100,15 @@ from torch.utils import _pytree as pytree
100from torch.utils.weak import TensorWeakRef100from torch.utils.weak import TensorWeakRef
101 101 
102import torch_npu102import torch_npu
103+from torch_npu.npu import graphs as _npu_graphs
103from torch_npu._C import (104from torch_npu._C import (
104 _npu_NPUAllocator_AllocatorState as AllocatorState,105 _npu_NPUAllocator_AllocatorState as AllocatorState,
105 _set_cached_tensors_enabled as _set_cached_tensors_enabled)106 _set_cached_tensors_enabled as _set_cached_tensors_enabled)
107+from torch_npu.npu._aclgraph_update_plan.resolver import (
108+ ACLGRAPH_UPDATE_PLAN_GLOBAL,
109+ update_aclgraph_records_for_graph,
110+ validate_aclgraph_update_plan_for_graph,
111+)
106import torch_npu.npu.aclnn112import torch_npu.npu.aclnn
107 113 
108if TYPE_CHECKING:114if TYPE_CHECKING:
@@ -114,7 +120,7 @@ StorageWeakRefPointer = int
114StorageDataPtr = int120StorageDataPtr = int
115NBytes = int121NBytes = int
116S = TypeVar("S", bound="StorageWeakRefWrapper")122S = TypeVar("S", bound="StorageWeakRefWrapper")
117-log = logging.getLogger("torch_npu.npugraph")123+log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
118 124 
119 125 
120@dataclasses.dataclass(frozen=True)126@dataclasses.dataclass(frozen=True)
@@ -634,7 +640,7 @@ class NPUWarmupNode:
634 s = storage()640 s = storage()
635 if s is not None:641 if s is not None:
636 non_npugraph_inps_storage_ptrs.add(s._cdata)642 non_npugraph_inps_storage_ptrs.add(s._cdata)
637- 643+ 
638 if not len(new_inputs) == 0:644 if not len(new_inputs) == 0:
639 raise RuntimeError("check len(new_inputs) == 0 fail")645 raise RuntimeError("check len(new_inputs) == 0 fail")
640 646 
@@ -765,6 +771,9 @@ class NPUGraphNode:
765 if not isinstance(inputs, (list, tuple)):771 if not isinstance(inputs, (list, tuple)):
766 raise RuntimeError("check isinstance(inputs, (list, tuple))")772 raise RuntimeError("check isinstance(inputs, (list, tuple))")
767 self.wrapped_function = wrapped_function773 self.wrapped_function = wrapped_function
774+ self.aclgraph_update_plan = getattr(
775+ wrapped_function.model, ACLGRAPH_UPDATE_PLAN_GLOBAL, None
776+ ) or []
768 self.id = graph_id777 self.id = graph_id
769 self.device = device_index778 self.device = device_index
770 self.stack_traces = stack_traces779 self.stack_traces = stack_traces
@@ -836,7 +845,7 @@ class NPUGraphNode:
836 )845 )
837 846 
838 self.non_static_input_idx: LevelList[int] = [847 self.non_static_input_idx: LevelList[int] = [
839- i 848+ i
840 for i in range(len(inputs))849 for i in range(len(inputs))
841 if i not in self.static_input_idxs850 if i not in self.static_input_idxs
842 ]851 ]
@@ -1059,12 +1068,19 @@ class NPUGraphNode:
1059 def run(self, new_inputs: List[InputType]) -> OutputType:1068 def run(self, new_inputs: List[InputType]) -> OutputType:
1060 log.debug("NPUGRAPH-TREE Node Run node=%s", self.id)1069 log.debug("NPUGRAPH-TREE Node Run node=%s", self.id)
1061 self.check_static_inputs_are_stable(new_inputs)1070 self.check_static_inputs_are_stable(new_inputs)
1062- for item in new_inputs:1071+ aclgraph_update_submitted = update_aclgraph_records_for_graph(
1063- if isinstance(item, torch.Tensor) and item.dtype == torch.int32 and item.device.type == "cpu":1072+ self.aclgraph_update_plan,
1064- self.graph.update(cpu_update_input=[{"context_lens": item}, {"actual_seq_lengths_kv": item}])1073+ self.graph,
1074+ new_inputs,
1075+ )
1065 self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs)1076 self._copy_inputs_and_remove_from_src(self.reconstructed_inputs, new_inputs)
1066 1077 
1067 self.run_graph()1078 self.run_graph()
1079+ if aclgraph_update_submitted:
1080+ # Ensure the next ACLGraph update does not record reusable external events before this replay resets them.
1081+ self.graph.graph_dispatch_mode.update_stream.wait_stream(
1082+ torch.npu.current_stream()
1083+ )
1068 1084 
1069 outputs = self.reconstruct_outputs()1085 outputs = self.reconstruct_outputs()
1070 new_inputs.clear()1086 new_inputs.clear()
@@ -1228,11 +1244,7 @@ class NPUGraphNode:
1228 1244 
1229 check_memory_pool(self.device, self.npu_graphs_pool, memory)1245 check_memory_pool(self.device, self.npu_graphs_pool, memory)
1230 1246 
1231- cpu_tensor = None1247+ aclgraph_update_inputs = list(inputs)
1232- for item in inputs:
1233- if isinstance(item, torch.Tensor) and item.dtype == torch.int32 and item.device.type == "cpu":
1234- cpu_tensor = item.clone()
1235- del item
1236 1248 
1237 with preserve_rng_state(), torch.npu.device(1249 with preserve_rng_state(), torch.npu.device(
1238 self.device1250 self.device
@@ -1245,9 +1257,13 @@ class NPUGraphNode:
1245 ), get_history_recording():1257 ), get_history_recording():
1246 static_outputs = model(inputs)1258 static_outputs = model(inputs)
1247 1259 
1248- if cpu_tensor is not None:1260+ validate_aclgraph_update_plan_for_graph(self.aclgraph_update_plan, self.graph)
1249- self.graph.update(cpu_update_input=[{"context_lens": cpu_tensor}, 1261+ update_aclgraph_records_for_graph(
1250- {"actual_seq_lengths_kv": cpu_tensor}])1262+ self.aclgraph_update_plan,
1263+ self.graph,
1264+ aclgraph_update_inputs,
1265+ )
1266+ aclgraph_update_inputs.clear()
1251 1267 
1252 # running model should reclaim memory1268 # running model should reclaim memory
1253 if not len(inputs) == 0:1269 if not len(inputs) == 0:
Mtorch_npu/npu/_npugraph_handlers/_fa3_graph_handler.py+30-14
@@ -43,6 +43,21 @@ class _FA3TensorListOutHandler(NpuGraphOpHandler):
43class FA3ForwardHandler(_FA3TensorListOutHandler):43class 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 @classmethod61 @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 False68 return False
54 return True69 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 @classmethod71 @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.out73 func_out = torch_npu.npu_fusion_attention_v3.out
@@ -126,6 +134,21 @@ class FA3ForwardHandler(_FA3TensorListOutHandler):
126class FA3BackwardHandler(_FA3TensorListOutHandler):134class 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 @classmethod152 @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 False159 return False
137 return True160 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 @classmethod162 @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.out164 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 
7Structure: ``_TensorListOutHandler`` provides ``postprocess_result`` (return7Structure: ``_TensorListOutHandler`` provides ``postprocess_result`` (return
8kwargs["out"]). ``IFAv1DefaultHandler`` and ``IFAv2DefaultHandler`` inherit8kwargs["out"]). ``IFAv1DefaultHandler`` and ``IFAv2DefaultHandler`` inherit
9-it and each implement ``update_args`` and ``prepare_capture``; both9+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):
37class _IFAv1DefaultHandler(_TensorListOutHandler):38class _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- @classmethod41+ 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 @classmethod56 @classmethod
46 def prepare_capture(cls, func, args, kwargs):57 def prepare_capture(cls, func, args, kwargs):
@@ -82,10 +93,20 @@ class _IFAv1DefaultHandler(_TensorListOutHandler):
82class _IFAv2DefaultHandler(_TensorListOutHandler):93class _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- @classmethod96+ 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 @classmethod111 @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 global42 parameter is ``cls``, not ``self``). There is no instance; the global
43 registry stores **class objects** directly. This structurally prevents43 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:: python61 .. 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- @classmethod65+ 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 @classmethod73 @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 result114 return result
101 115 
102 @classmethod116 @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 to121+ 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- pass144+ 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 @classmethod154 @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
12class _SimpleGraphHandler(NpuGraphOpHandler):12class _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]]): Specifies16+ spec-driven ``update_args`` walks this map and assigns the matching
17- ``op_name -> (arg_index, update_key)`` for each supported17+ 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+13-14
@@ -61,7 +61,7 @@ from torch_npu._C import ( # noqa: F401
61)61)
62 62 
63 63 
64-log = logging.getLogger("torch_npu.npugraph")64+log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
65 65 
66 66 
67def is_current_stream_capturing():67def is_current_stream_capturing():
@@ -229,18 +229,18 @@ def _print_npugraph_tensor_impl(input, tensor_name=None):
229 if device.type == "cpu":229 if device.type == "cpu":
230 _print_callback_pending(tensor_name, input)230 _print_callback_pending(tensor_name, input)
231 return231 return
232- 232+ 
233 if device.type != "npu":233 if device.type != "npu":
234 return234 return
235 235 
236 device_index = device.index236 device_index = device.index
237 save_stream = _get_save_tensor_stream(device_index)237 save_stream = _get_save_tensor_stream(device_index)
238- 238+ 
239 # Record event on the original compute stream before switching239 # Record event on the original compute stream before switching
240 event1 = torch.npu.Event()240 event1 = torch.npu.Event()
241 event2 = torch.npu.Event()241 event2 = torch.npu.Event()
242 event1.record()242 event1.record()
243- 243+ 
244 with torch.npu.stream(save_stream):244 with torch.npu.stream(save_stream):
245 # Wait for the original stream to complete before D2H245 # Wait for the original stream to complete before D2H
246 event1.wait()246 event1.wait()
@@ -255,7 +255,7 @@ def _print_npugraph_tensor_impl(input, tensor_name=None):
255 )255 )
256 # Mark save_stream completion256 # Mark save_stream completion
257 event2.record()257 event2.record()
258- 258+ 
259 # Wait for save_stream to complete (back to original stream now)259 # Wait for save_stream to complete (back to original stream now)
260 event2.wait()260 event2.wait()
261 261 
@@ -268,19 +268,19 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
268 if device.type == "cpu":268 if device.type == "cpu":
269 torch.save(input, _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))269 torch.save(input, _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))
270 return270 return
271- 271+ 
272 if device.type != "npu":272 if device.type != "npu":
273 return273 return
274 274 
275 device_index = device.index275 device_index = device.index
276 save_stream = _get_save_tensor_stream(device_index)276 save_stream = _get_save_tensor_stream(device_index)
277 final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)277 final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)
278- 278+ 
279 # Record event on the original compute stream before switching279 # Record event on the original compute stream before switching
280 event1 = torch.npu.Event()280 event1 = torch.npu.Event()
281 event2 = torch.npu.Event()281 event2 = torch.npu.Event()
282 event1.record()282 event1.record()
283- 283+ 
284 with torch.npu.stream(save_stream):284 with torch.npu.stream(save_stream):
285 # Wait for the original stream to complete before D2H285 # Wait for the original stream to complete before D2H
286 event1.wait()286 event1.wait()
@@ -295,7 +295,7 @@ def _save_npugraph_tensor_impl(input, save_path=None, overwrite=False):
295 )295 )
296 # Mark save_stream completion296 # Mark save_stream completion
297 event2.record()297 event2.record()
298- 298+ 
299 # Wait for save_stream to complete (back to original stream now)299 # Wait for save_stream to complete (back to original stream now)
300 event2.wait()300 event2.wait()
301 301 
@@ -305,19 +305,19 @@ def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=Fals
305 if device.type == "cpu":305 if device.type == "cpu":
306 torch.save(list(input), _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))306 torch.save(list(input), _build_save_npugraph_tensor_path(save_path, overwrite=overwrite))
307 return307 return
308- 308+ 
309 if device.type != "npu":309 if device.type != "npu":
310 return310 return
311 311 
312 device_index = device.index312 device_index = device.index
313 save_stream = _get_save_tensor_stream(device_index)313 save_stream = _get_save_tensor_stream(device_index)
314 final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)314 final_path = _build_save_npugraph_tensor_path(save_path, device_index, overwrite)
315- 315+ 
316 # Record event on the original compute stream before switching316 # Record event on the original compute stream before switching
317 event1 = torch.npu.Event()317 event1 = torch.npu.Event()
318 event2 = torch.npu.Event()318 event2 = torch.npu.Event()
319 event1.record()319 event1.record()
320- 320+ 
321 with torch.npu.stream(save_stream):321 with torch.npu.stream(save_stream):
322 # Wait for the original stream to complete before D2H322 # Wait for the original stream to complete before D2H
323 event1.wait()323 event1.wait()
@@ -332,7 +332,7 @@ def _save_npugraph_tensor_tensor_list_impl(input, save_path=None, overwrite=Fals
332 )332 )
333 # Mark save_stream completion333 # Mark save_stream completion
334 event2.record()334 event2.record()
335- 335+ 
336 # Wait for save_stream to complete (back to original stream now)336 # Wait for save_stream to complete (back to original stream now)
337 event2.wait()337 event2.wait()
338 338 
@@ -586,7 +586,6 @@ class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode):
586 graph_task_update_end(self.update_stream)586 graph_task_update_end(self.update_stream)
587 record.event.record(self.update_stream)587 record.event.record(self.update_stream)
588 588 
589- 
590# Python shim helps Sphinx process docstrings more reliably.589# Python shim helps Sphinx process docstrings more reliably.
591class NPUGraph(torch_npu._C._NPUGraph):590class NPUGraph(torch_npu._C._NPUGraph):
592 r"""Wrapper around a NPU graph.591 r"""Wrapper around a NPU graph.
Mtorch_npu/utils/_graph_tree.py+12-2
@@ -54,7 +54,7 @@ from torch.multiprocessing.reductions import StorageWeakRef
54import torch_npu.npu.aclnn54import torch_npu.npu.aclnn
55 55 
56 56 
57-log = logging.getLogger("torch_npu.aclgraph")57+log = torch._logging.getArtifactLogger("torch_npu.npugraph", "cudagraphs")
58 58 
59 59 
60def npugraph_mark_step_begin():60def npugraph_mark_step_begin():
@@ -104,6 +104,7 @@ def npugraphify(
104 mutated_input_idxs: Tuple[int, ...] = (),104 mutated_input_idxs: Tuple[int, ...] = (),
105) -> Callable[..., Any]:105) -> Callable[..., Any]:
106 from torch_npu.npu._graph_tree import npugraphify_impl as new_npugraphify_impl106 from torch_npu.npu._graph_tree import npugraphify_impl as new_npugraphify_impl
107+ 
107 npugraphify_fn: Callable[..., Any]108 npugraphify_fn: Callable[..., Any]
108 if config.triton.cudagraph_trees:109 if config.triton.cudagraph_trees:
109 npugraphify_fn = functools.partial(110 npugraphify_fn = functools.partial(
@@ -222,7 +223,7 @@ def npugraphify_impl(
222 223 
223 else:224 else:
224 copy_indices = [225 copy_indices = [
225- idx 226+ idx
226 for idx in range(len(static_inputs))227 for idx in range(len(static_inputs))
227 if idx not in static_input_idxs228 if idx not in static_input_idxs
228 ]229 ]
@@ -378,3 +379,12 @@ def _apply_npugraph_tree_methods():
378 torch._inductor.compile_fx.cudagraphify = npugraphify379 torch._inductor.compile_fx.cudagraphify = npugraphify
379 torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes380 torch._inductor.cudagraph_utils.check_multiple_devices_or_any_cpu_nodes = check_multiple_devices_or_any_cpu_nodes
380 torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin381 torch.compiler.npugraph_mark_step_begin = npugraph_mark_step_begin
382+ 
383+ import torch._inductor.cudagraph_trees as _upstream_cgt # noqa: F401
384+ 
385+ def _npu_get_manager(*args, **kwargs):
386+ from torch_npu.npu._graph_tree import get_manager
387+ 
388+ return get_manager(*args, **kwargs)
389+ 
390+ _upstream_cgt.get_manager = _npu_get_manager