[Feature][aclgraph] replace hard-coded dispatch with Registry + Template Method for NPU Graph op handlers-pytorch-AtomGit

已合并

[Feature][aclgraph] replace hard-coded dispatch with Registry + Template Method for NPU Graph op handlers #30977

suhaibo创建于 2月25日从suhaibo/pytorch： v2.7.1 合入到 Ascend/pytorch： v2.7.1

[Feature][aclgraph] replace hard-coded dispatch with Registry + Template Method for NPU Graph op handlers #30977

已合并

suhaibo创建于 2月25日

从suhaibo/pytorch v2.7.1合入到Ascend/pytorch v2.7.1

共 9 个文件变更+638-146

Mtest/allowlist_for_publicAPI.json+3-1

	@@ -2856,7 +2856,9 @@
2856	"BFloat16Tensor",	2856	"BFloat16Tensor",
2857	"restart_device",	2857	"restart_device",
2858	"stop_device",	2858	"stop_device",
2859	- "set_op_timeout_ms"	2859	+ "set_op_timeout_ms",
		2860	+ "NpuGraphOpHandler",
		2861	+ "register_npu_graph_handler"
2860	],	2862	],
2861	"torch_npu.profiler": [	2863	"torch_npu.profiler": [
2862	"ProfilerActivity",	2864	"ProfilerActivity",

Atest/npu/test_npugraph_handler.py+60-0

@@ -0,0 +1,60 @@
	1	+from torch_npu.testing.testcase import TestCase, run_tests
	2	+from torch_npu.npu import (
	3	+ NpuGraphOpHandler,
	4	+ register_npu_graph_handler,
	5	+)
	6	+from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
	7	+
	8	+
	9	+class TestNpuGraphHandlerRegistry(TestCase):
	10	+
	11	+ def setUp(self):
	12	+ self._snapshot = dict(_NPU_GRAPH_OP_HANDLERS)
	13	+
	14	+ def tearDown(self):
	15	+ _NPU_GRAPH_OP_HANDLERS.clear()
	16	+ _NPU_GRAPH_OP_HANDLERS.update(self._snapshot)
	17	+
	18	+ def test_register_single_name(self):
	19	+ @register_npu_graph_handler("test_op_single")
	20	+ class _H(NpuGraphOpHandler):
	21	+ pass
	22	+
	23	+ self.assertIn("test_op_single", _NPU_GRAPH_OP_HANDLERS)
	24	+ self.assertIs(_NPU_GRAPH_OP_HANDLERS["test_op_single"], _H)
	25	+
	26	+ def test_register_multiple_names(self):
	27	+ @register_npu_graph_handler(["test_op_a", "test_op_a.default"])
	28	+ class _H(NpuGraphOpHandler):
	29	+ pass
	30	+
	31	+ self.assertIn("test_op_a", _NPU_GRAPH_OP_HANDLERS)
	32	+ self.assertIn("test_op_a.default", _NPU_GRAPH_OP_HANDLERS)
	33	+ self.assertIs(_NPU_GRAPH_OP_HANDLERS["test_op_a"], _H)
	34	+
	35	+
	36	+class TestNpuGraphHandlerBuiltinRegistration(TestCase):
	37	+
	38	+ EXPECTED_OPS = [
	39	+ "npu_fused_infer_attention_score",
	40	+ "npu_fused_infer_attention_score.default",
	41	+ "npu_fused_infer_attention_score.out",
	42	+ "npu_fused_infer_attention_score_v2",
	43	+ "npu_fused_infer_attention_score_v2.default",
	44	+ "npu_fused_infer_attention_score_v2.out",
	45	+ "_npu_paged_attention.default",
	46	+ "npu_multi_head_latent_attention.out",
	47	+ ]
	48	+
	49	+ def test_all_expected_ops_registered(self):
	50	+ for op_name in self.EXPECTED_OPS:
	51	+ with self.subTest(op_name=op_name):
	52	+ self.assertIn(
	53	+ op_name,
	54	+ _NPU_GRAPH_OP_HANDLERS,
	55	+ f"Expected handler for '{op_name}' not found in registry",
	56	+ )
	57	+
	58	+
	59	+if __name__ == "__main__":
	60	+ run_tests()

Mtest/torch_npu_schema.json+18-0

	@@ -2282,6 +2282,24 @@
2282	"torch_npu.npu.npugraph_ex.register_replacement": {	2282	"torch_npu.npu.npugraph_ex.register_replacement": {
2283	"signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False)"	2283	"signature": "(search_fn, replace_fn, example_inputs, trace_fn=<function fwd_only>, extra_check=<function _return_true>, search_fn_pattern=None, scalar_workaround=None, skip_duplicates=False)"
2284	},	2284	},
		2285	+ "torch_npu.npu.register_npu_graph_handler": {
		2286	+ "signature": "(op_names)"
		2287	+ },
		2288	+ "torch_npu.npu.NpuGraphOpHandler": {
		2289	+ "signature": "()"
		2290	+ },
		2291	+ "torch_npu.npu.NpuGraphOpHandler.prepare_capture": {
		2292	+ "signature": "(func, args, kwargs)"
		2293	+ },
		2294	+ "torch_npu.npu.NpuGraphOpHandler.postprocess_result": {
		2295	+ "signature": "(result, kwargs)"
		2296	+ },
		2297	+ "torch_npu.npu.NpuGraphOpHandler.update_args": {
		2298	+ "signature": "(dispatch_record, update_input)"
		2299	+ },
		2300	+ "torch_npu.npu.NpuGraphOpHandler.record_wrap_kwarg": {
		2301	+ "signature": "(key, value, tensor_param_names)"
		2302	+ },
2285	"torch_npu.distributed.run.parse_args": {	2303	"torch_npu.distributed.run.parse_args": {
2286	"signature": "(args)"	2304	"signature": "(args)"
2287	},	2305	},

Mtorch_npu/npu/__init__.py+8-1

	@@ -134,7 +134,9 @@ __all__ = [
134	"reset_peak_host_memory_stats",	134	"reset_peak_host_memory_stats",
135	"set_deterministic_level",	135	"set_deterministic_level",
136	"use_compatible_impl",	136	"use_compatible_impl",
137	- "are_compatible_impl_enabled"	137	+ "are_compatible_impl_enabled",
		138	+ "NpuGraphOpHandler",
		139	+ "register_npu_graph_handler",
138	]	140	]
139		141
140	from typing import Tuple, Union, List, cast, Optional	142	from typing import Tuple, Union, List, cast, Optional
	@@ -178,6 +180,11 @@ from .graphs import (
178	graph_task_update_end,	180	graph_task_update_end,
179	)	181	)
180		182
		183	+from ._npugraph_handlers import (
		184	+ NpuGraphOpHandler,
		185	+ register_npu_graph_handler,
		186	+)
		187	+
181	# init profiler	188	# init profiler
182	if not torch_npu._C._profiler_init():	189	if not torch_npu._C._profiler_init():
183	raise RuntimeError("proflier initialization failed" + prof_error(ErrCode.UNAVAIL))	190	raise RuntimeError("proflier initialization failed" + prof_error(ErrCode.UNAVAIL))

Atorch_npu/npu/_npugraph_handlers/__init__.py+25-0

@@ -0,0 +1,25 @@
	1	+"""NPU Graph Operator Handler Framework (Registry + Template Method).
	2	+
	3	+This package exposes the public API for the NPU Graph operator handler
	4	+mechanism and ensures that all built-in handlers are registered at
	5	+import time.
	6	+
	7	+Public API (re-exported from :mod:`npugraph_handler`)
	8	+-----------------------------------------------------
	9	+- :class:`NpuGraphOpHandler` -- base class for operator handlers.
	10	+- :func:`register_npu_graph_handler` -- class decorator to register a handler.
	11	+
	12	+"""
	13	+
	14	+# Re-export public API from the core module
	15	+from .npugraph_handler import (
	16	+ NpuGraphOpHandler,
	17	+ register_npu_graph_handler,
	18	+)
	19	+
	20	+# Auto-register built-in handlers, Import built-in handler classes to trigger their ``@register_npu_graph_handler`` decorators.
	21	+from .ifa_handler import ( # noqa: F401
	22	+ _IFAv1DefaultHandler,
	23	+ _IFAv2DefaultHandler,
	24	+)
	25	+from .simple_handler import _SimpleGraphHandler # noqa: F401

Atorch_npu/npu/_npugraph_handlers/ifa_handler.py+122-0

@@ -0,0 +1,122 @@
	1	+"""IFA (Infer Fused Attention) v1 / v2 Graph Handlers.
	2	+
	3	+This module defines the NPU Graph operator handlers for the
	4	+``npu_fused_infer_attention_score`` (v1) and
	5	+``npu_fused_infer_attention_score_v2`` (v2) operator families.
	6	+
	7	+Structure: ``_TensorListOutHandler`` provides ``postprocess_result`` (return
	8	+kwargs["out"]). ``IFAv1DefaultHandler`` and ``IFAv2DefaultHandler`` inherit
	9	+it and each implement ``update_args`` and ``prepare_capture``; both
	10	+``.default`` and ``.out`` are registered on the same handler class.
	11	+"""
	12	+__all__ = []
	13	+
	14	+import torch_npu
	15	+from .npugraph_handler import NpuGraphOpHandler, register_npu_graph_handler
	16	+
	17	+
	18	+class _TensorListOutHandler(NpuGraphOpHandler):
	19	+ """Base for operators whose ``out`` kwarg is a ``TensorList``.
	20	+
	21	+ Returns ``kwargs["out"]`` from ``postprocess_result`` so callers get a
	22	+ Python list instead of the raw C++ return.
	23	+ """
	24	+
	25	+ @classmethod
	26	+ def postprocess_result(cls, result, kwargs):
	27	+ return kwargs["out"]
	28	+
	29	+
	30	+# ========================= IFA v1 ================================
	31	+
	32	+@register_npu_graph_handler([
	33	+ "npu_fused_infer_attention_score",
	34	+ "npu_fused_infer_attention_score.default",
	35	+ "npu_fused_infer_attention_score.out",
	36	+])
	37	+class _IFAv1DefaultHandler(_TensorListOutHandler):
	38	+ """IFA v1: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""
	39	+
	40	+ @classmethod
	41	+ def update_args(cls, record, update_input):
	42	+ if "actual_seq_lengths_kv" in update_input and len(record.args) >= 7:
	43	+ record.args[6] = update_input["actual_seq_lengths_kv"]
	44	+
	45	+ @classmethod
	46	+ def prepare_capture(cls, func, args, kwargs):
	47	+ func_out = torch_npu.npu_fused_infer_attention_score.out
	48	+ if func is func_out:
	49	+ return func, args, kwargs
	50	+
	51	+ workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
	52	+ args, *kwargs
	53	+ )
	54	+ out_args = [args[0], args[2]]
	55	+ out_kwargs_keys = [
	56	+ "input_layout",
	57	+ "quant_scale2",
	58	+ "block_table",
	59	+ "num_heads",
	60	+ "num_key_value_heads",
	61	+ "softmax_lse_flag",
	62	+ "query_rope",
	63	+ ]
	64	+ out_kwargs = {k: kwargs[k] for k in out_kwargs_keys if k in kwargs}
	65	+ output, softmax_lse = (
	66	+ torch_npu._npu_fused_infer_attention_score_infer_output(
	67	+ out_args, *out_kwargs
	68	+ )
	69	+ )
	70	+ kwargs["workspace"] = workspace
	71	+ kwargs["out"] = [output, softmax_lse]
	72	+ return func_out, args, kwargs
	73	+
	74	+
	75	+# ========================= IFA v2 ================================
	76	+
	77	+@register_npu_graph_handler([
	78	+ "npu_fused_infer_attention_score_v2",
	79	+ "npu_fused_infer_attention_score_v2.default",
	80	+ "npu_fused_infer_attention_score_v2.out",
	81	+])
	82	+class _IFAv2DefaultHandler(_TensorListOutHandler):
	83	+ """IFA v2: ``.default`` pre-allocates and swaps to ``.out``; ``.out`` passthrough."""
	84	+
	85	+ @classmethod
	86	+ def update_args(cls, record, update_input):
	87	+ if "actual_seq_kvlen" in update_input and len(record.args) >= 9:
	88	+ record.args[8] = update_input["actual_seq_kvlen"]
	89	+
	90	+ @classmethod
	91	+ def prepare_capture(cls, func, args, kwargs):
	92	+ func_out = torch_npu.npu_fused_infer_attention_score_v2.out
	93	+ if func is func_out:
	94	+ return func, args, kwargs
	95	+
	96	+ workspace = (
	97	+ torch_npu._npu_fused_infer_attention_score_v2_get_max_workspace(
	98	+ args, *kwargs
	99	+ )
	100	+ )
	101	+ out_args = [args[0], args[2]]
	102	+ out_kwargs_keys = [
	103	+ "query_dtype",
	104	+ "value_dtype",
	105	+ "input_layout",
	106	+ "quant_scale_out",
	107	+ "block_table",
	108	+ "num_query_heads",
	109	+ "num_key_value_heads",
	110	+ "return_softmax_lse",
	111	+ "query_rope",
	112	+ "out_dtype",
	113	+ ]
	114	+ out_kwargs = {k: kwargs[k] for k in out_kwargs_keys if k in kwargs}
	115	+ output, softmax_lse = (
	116	+ torch_npu._npu_fused_infer_attention_score_v2_infer_output(
	117	+ out_args, *out_kwargs
	118	+ )
	119	+ )
	120	+ kwargs["workspace"] = workspace
	121	+ kwargs["out"] = [output, softmax_lse]
	122	+ return func_out, args, kwargs

Atorch_npu/npu/_npugraph_handlers/npugraph_handler.py+212-0

@@ -0,0 +1,212 @@
	1	+"""NPU Graph Operator Handler -- base class, global registry, and utilities.
	2	+
	3	+This module contains the core building blocks of the NPU Graph operator
	4	+handler framework:
	5	+
	6	+- :class:`NpuGraphOpHandler` -- abstract base class for handlers.
	7	+- :data:`_NPU_GRAPH_OP_HANDLERS` -- global registry (dict).
	8	+- :func:`register_npu_graph_handler` -- class-decorator for registration.
	9	+
	10	+Design Contracts
	11	+----------------
	12	+1. Stateless Handler -- All hook methods are ``@classmethod``; the
	13	+ registry stores class objects, never instances.
	14	+2. Function Replacement Consistency -- When ``prepare_capture``
	15	+ substitutes ``func`` with ``actual_func``, the handler registered for
	16	+ ``actual_func.__name__`` must implement a compatible
	17	+ ``update_args`` (typically guaranteed by a shared intermediate base).
	18	+
	19	+See Also
	20	+--------
	21	+``torch_npu/npu/graphs.py`` for the template-method skeleton that
	22	+consumes this registry.
	23	+"""
	24	+
	25	+import logging
	26	+from copy import deepcopy
	27	+
	28	+import torch
	29	+from torch_npu._C import _weak_ref_tensor as TensorWeakRef
	30	+
	31	+logger = logging.getLogger(__name__)
	32	+
	33	+
	34	+class NpuGraphOpHandler:
	35	+ r"""Base class for NPU Graph operator handlers.
	36	+
	37	+ Subclasses override ``@classmethod`` hooks to customize capture / update
	38	+ behavior for specific operators, while the framework keeps
	39	+ stream / event / task-group orchestration in a common template.
	40	+
	41	+ Stateless by design -- All hook methods are ``@classmethod`` (first
	42	+ parameter is ``cls``, not ``self``). There is no instance; the global
	43	+ registry stores class objects directly. This structurally prevents
	44	+ storing mutable per-invocation state. Class-level constants (e.g.
	45	+ ``_OP_ARG_SPECS``) are accessible via ``cls``.
	46	+
	47	+ Users should inherit this class and override the needed hooks:
	48	+
	49	+ .. code-block:: python
	50	+
	51	+ @register_npu_graph_handler(["my_op", "my_op.default"])
	52	+ class MyHandler(NpuGraphOpHandler):
	53	+ @classmethod
	54	+ def update_args(cls, record, update_input):
	55	+ if "batch" in update_input and len(record.args) >= 3:
	56	+ record.args[2] = update_input["batch"]
	57	+ """
	58	+
	59	+ @classmethod
	60	+ def prepare_capture(cls, func, args, kwargs):
	61	+ r"""Prepare operator call before graph-task recording.
	62	+
	63	+ This hook runs before ``graph_task_group_begin`` and can be used
	64	+ for operator-specific preprocessing such as workspace allocation,
	65	+ output pre-allocation, or switching from ``.default`` to ``.out``
	66	+ overloads.
	67	+
	68	+ .. note:: Function Replacement Contract
	69	+
	70	+ If ``actual_func`` differs from ``func``, ensure the handler
	71	+ registered for ``actual_func.__name__`` implements a compatible
	72	+ ``update_args``. The recommended approach is to share a common
	73	+ base class that defines ``update_args`` (see ``_IFAv1Base`` /
	74	+ ``_IFAv2Base``).
	75	+
	76	+ Args:
	77	+ func (OpOverload): Original operator callable.
	78	+ args (tuple): Original arguments.
	79	+ kwargs (dict): Original keyword arguments.
	80	+
	81	+ Returns:
	82	+ tuple[Callable, tuple, dict]: ``(actual_func, args, kwargs)`` to
	83	+ execute during recording.
	84	+ """
	85	+ return func, args, kwargs
	86	+
	87	+ @classmethod
	88	+ def postprocess_result(cls, result, kwargs):
	89	+ r"""Post-process operator return value after recording.
	90	+
	91	+ Called after ``graph_task_group_end`` and dispatch-record creation.
	92	+
	93	+ Args:
	94	+ result: Raw return value from ``actual_func(args, *kwargs)``.
	95	+ kwargs (dict): Current keyword arguments (may contain ``"out"``).
	96	+
	97	+ Returns:
	98	+ Final value returned to the Python caller.
	99	+ """
	100	+ return result
	101	+
	102	+ @classmethod
	103	+ def update_args(cls, dispatch_record, update_input):
	104	+ r"""Apply operator-specific indexed-arg updates.
	105	+
	106	+ Framework-level kwargs updates are handled by the dispatch skeleton.
	107	+ Override this hook only when update values must be applied to
	108	+ arguments by index.
	109	+
	110	+ Args:
	111	+ dispatch_record (_GraphDispatchRecord): Recorded operator call.
	112	+ Args can be modified via ``dispatch_record.args[i]``.
	113	+ update_input (dict): User-provided update payload.
	114	+ """
	115	+ pass
	116	+
	117	+ @classmethod
	118	+ def record_wrap_kwarg(cls, key, value, tensor_param_names):
	119	+ r"""Convert a kwarg value into record-time storage representation.
	120	+
	121	+ Called only during the capture phase (creating the dispatch
	122	+ record). The purpose of ``TensorWeakRef`` conversion is to avoid the
	123	+ Python-side record holding strong references to NPU tensors, letting
	124	+ the C++ graph runtime manage tensor memory lifetimes.
	125	+
	126	+ .. note::
	127	+
	128	+ The update phase uses direct assignment for kwargs
	129	+ (``record.kwargs[key] = update_input[key]``), consistent with the
	130	+ original implementation and with how ``update_args`` handles
	131	+ arguments. Update is a short-lived "assign -> replay"
	132	+ flow where weak-ref conversion is unnecessary.
	133	+
	134	+ Logic (consistent with original, with list/tuple generalisation):
	135	+
	136	+ - ``None`` -> ``None`` (fast path)
	137	+ - ``list`` / ``tuple`` -> element-wise: NPU Tensor -> ``TensorWeakRef``,
	138	+ else -> ``deepcopy`` (replaces old hardcoded ``if key == "out"``
	139	+ that assumed exactly 2 Tensors). Only NPU tensors are wrapped;
	140	+ CPU tensors use ``deepcopy`` because ``TensorWeakRef`` from
	141	+ torch_npu._C is only valid for NPU tensors.
	142	+ - Single value where ``key`` in ``tensor_param_names`` and value is an
	143	+ NPU Tensor -> ``TensorWeakRef``
	144	+ - Everything else -> ``deepcopy``
	145	+
	146	+ Args:
	147	+ key (str): Kwarg name.
	148	+ value: Kwarg value.
	149	+ tensor_param_names (list[str]): Tensor-typed kwarg names parsed
	150	+ from operator schema.
	151	+
	152	+ Returns:
	153	+ Stored value for the dispatch record.
	154	+ """
	155	+ if value is None:
	156	+ return None
	157	+
	158	+ def _is_npu_tensor(t):
	159	+ return torch.is_tensor(t) and "npu" in str(t.device)
	160	+
	161	+ if isinstance(value, (list, tuple)):
	162	+ wrapped = [
	163	+ TensorWeakRef(t) if _is_npu_tensor(t) else deepcopy(t)
	164	+ for t in value
	165	+ ]
	166	+ return type(value)(wrapped)
	167	+
	168	+ if key in tensor_param_names and _is_npu_tensor(value):
	169	+ return TensorWeakRef(value)
	170	+
	171	+ return deepcopy(value)
	172	+
	173	+
	174	+# ---------------------------------------------------------------------------
	175	+# Global Registry
	176	+# ---------------------------------------------------------------------------
	177	+
	178	+_NPU_GRAPH_OP_HANDLERS = {}
	179	+
	180	+
	181	+def register_npu_graph_handler(op_names):
	182	+ r"""Register an operator handler via class decorator.
	183	+
	184	+ The decorated class itself (not an instance) is stored in the global
	185	+ registry. All hook methods must be ``@classmethod``.
	186	+
	187	+ Args:
	188	+ op_names (str or list[str]): Operator names resolved from
	189	+ ``func.__name__`` in ``__torch_dispatch__`` (for example,
	190	+ ``"my_op"``, ``"my_op.default"``, ``"my_op.out"``).
	191	+
	192	+ Returns:
	193	+ A class decorator.
	194	+
	195	+ Example::
	196	+
	197	+ @register_npu_graph_handler(["my_op", "my_op.default"])
	198	+ class MyHandler(NpuGraphOpHandler):
	199	+ ...
	200	+ """
	201	+ def decorator(cls):
	202	+ names = op_names if isinstance(op_names, (list, tuple)) else [op_names]
	203	+ for name in names:
	204	+ if name in _NPU_GRAPH_OP_HANDLERS:
	205	+ existing = _NPU_GRAPH_OP_HANDLERS[name].__name__
	206	+ logger.warning(
	207	+ f"NpuGraphOpHandler for '{name}' is being overridden: "
	208	+ f"{existing} -> {cls.__name__}"
	209	+ )
	210	+ _NPU_GRAPH_OP_HANDLERS[name] = cls # store class, not instance
	211	+ return cls
	212	+ return decorator

Atorch_npu/npu/_npugraph_handlers/simple_handler.py+32-0

@@ -0,0 +1,32 @@
	1	+"""Simple Graph Handlers (Paged Attention / MLA)."""
	2	+
	3	+__all__ = []
	4	+
	5	+from .npugraph_handler import NpuGraphOpHandler, register_npu_graph_handler
	6	+
	7	+
	8	+@register_npu_graph_handler([
	9	+ "_npu_paged_attention.default",
	10	+ "npu_multi_head_latent_attention.out",
	11	+])
	12	+class _SimpleGraphHandler(NpuGraphOpHandler):
	13	+ """Handler for PA (Paged Attention) and MLA operators.
	14	+
	15	+ Attributes:
	16	+ _OP_ARG_SPECS (dict[str, tuple[int, str]]): Specifies
	17	+ ``op_name -> (arg_index, update_key)`` for each supported
	18	+ operator.
	19	+ """
	20	+
	21	+ _OP_ARG_SPECS = {
	22	+ "_npu_paged_attention.default": (7, "context_lens"),
	23	+ "npu_multi_head_latent_attention.out": (5, "context_lens"),
	24	+ }
	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+158-144

	@@ -1,8 +1,17 @@
1	-__all__ = ["is_current_stream_capturing", "graph_pool_handle", "graph_task_group_begin",	1	+__all__ = [
2	- "graph_task_group_end", "graph_task_update_begin", "graph_task_update_end",	2	+ "is_current_stream_capturing",
3	- "NPUGraph", "graph", "make_graphed_callables"]	3	+ "graph_pool_handle",
		4	+ "graph_task_group_begin",
		5	+ "graph_task_group_end",
		6	+ "graph_task_update_begin",
		7	+ "graph_task_update_end",
		8	+ "NPUGraph",
		9	+ "graph",
		10	+ "make_graphed_callables",
		11	+]
4		12
5	import gc	13	import gc
		14	+import logging
6	import re	15	import re
7	import typing	16	import typing
8	from copy import deepcopy	17	from copy import deepcopy
	@@ -12,10 +21,12 @@ from typing import List, Dict, Any, Optional, Tuple
12	import torch	21	import torch
13	import torch_npu._C	22	import torch_npu._C
14	from torch_npu._C import _weak_ref_tensor as TensorWeakRef	23	from torch_npu._C import _weak_ref_tensor as TensorWeakRef
		24	+from torch_npu.npu._npugraph_handlers.npugraph_handler import _NPU_GRAPH_OP_HANDLERS
15	from torch_npu.utils._error_code import ErrCode, pta_error	25	from torch_npu.utils._error_code import ErrCode, pta_error
16	from torch_npu._compiler._config import force_npugraph_gc	26	from torch_npu._compiler._config import force_npugraph_gc
17	from .utils import _dummy_type	27	from .utils import _dummy_type
18		28
		29	+
19	if not hasattr(torch_npu._C, "_NPUStreamBase"):	30	if not hasattr(torch_npu._C, "_NPUStreamBase"):
20	# Define dummy base classes	31	# Define dummy base classes
21	torch_npu._C.__dict__["_NPUGraph"] = _dummy_type("_NPUGraph")	32	torch_npu._C.__dict__["_NPUGraph"] = _dummy_type("_NPUGraph")
	@@ -77,7 +88,7 @@ def graph_task_update_end(stream):
77		88
78	@dataclass	89	@dataclass
79	class _GraphDispatchRecord:	90	class _GraphDispatchRecord:
80	- """存储单次操作的完整记录"""	91	+ """Record of a single dispatched operator call during graph capture."""
81	event: Any = None	92	event: Any = None
82	handle: Any = None	93	handle: Any = None
83	kwargs: Dict[str, Any] = field(default_factory=dict)	94	kwargs: Dict[str, Any] = field(default_factory=dict)
	@@ -86,6 +97,14 @@ class _GraphDispatchRecord:
86		97
87		98
88	class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode):	99	class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode):
		100	+ """Template-method skeleton for NPU Graph capture and update.
		101	+
		102	+ The skeleton keeps the common stream / event / task-group orchestration
		103	+ and delegates operator-specific logic to
		104	+ :class:`~torch_npu.npu.NpuGraphOpHandler` classes
		105	+ looked up in ``_NPU_GRAPH_OP_HANDLERS``.
		106	+ """
		107	+
89	tensor_schema_name = {}	108	tensor_schema_name = {}
90	update_stream = None	109	update_stream = None
91		110
	@@ -102,153 +121,145 @@ class _GraphDispatchMode(torch.utils._python_dispatch.TorchDispatchMode):
102	return True	121	return True
103		122
104	@classmethod	123	@classmethod
105	- def update_schema(cls, name, schame):	124	+ def update_schema(cls, name, schema):
106	if name in cls.tensor_schema_name:	125	if name in cls.tensor_schema_name:
107	return	126	return
108	# match: Tensor q_nope, Tensor? mask=None, Tensor(a!) output	127	# match: Tensor q_nope, Tensor? mask=None, Tensor(a!) output
109	pattern = r'Tensor(?:\(a!\)\|\?)?\s+(\w+)'	128	pattern = r'Tensor(?:\(a!\)\|\?)?\s+(\w+)'
110	- cls.tensor_schema_name[name] = re.findall(pattern, schame)	129	+ cls.tensor_schema_name[name] = re.findall(pattern, schema)
111	-
112	- def update_capture_record(self, cpu_update_input):
113	- if len(cpu_update_input) == 1:
114	- new_list = [cpu_update_input[0].copy() for _ in range(len(self.graph_dispatch_records))]
115	- cpu_update_input = new_list
116	- if len(self.graph_dispatch_records) != len(self.graph_dispatch_records):
117	- raise RuntimeError(f"Currently, there are {len(self.graph_dispatch_records)} operators that need to be updated by capture, "
118	- f"and there are only {len(self.graph_dispatch_records)} elements in the incoming cpu_update_input list", pta_error(ErrCode.PARAM))
119	- with torch.npu.stream(self.update_stream):
120	- for graph_dispatch_record, update_input in zip(self.graph_dispatch_records, cpu_update_input):
121	- graph_task_update_begin(self.update_stream, graph_dispatch_record.handle)
122	- for key in update_input:
123	- if key in graph_dispatch_record.kwargs:
124	- graph_dispatch_record.kwargs[key] = update_input[key]
125		130
126	- # When parameters are passed through args, position is position of the updated parameter.	131	+ # -----------------------------------------------------------------
127	- if graph_dispatch_record.op_cache_entry.__name__ in ["npu_fused_infer_attention_score", "npu_fused_infer_attention_score.out"]:	132	+ # Capture skeleton
128	- position, key = 6, "actual_seq_lengths_kv"	133	+ # -----------------------------------------------------------------
129	- elif graph_dispatch_record.op_cache_entry.__name__ in ["npu_fused_infer_attention_score_v2", "npu_fused_infer_attention_score_v2.out"]:
130	- position, key = 8, "actual_seq_kvlen"
131	- elif graph_dispatch_record.op_cache_entry.__name__ == "_npu_paged_attention.default":
132	- position, key = 7, "context_lens"
133	- elif graph_dispatch_record.op_cache_entry.__name__ == "npu_multi_head_latent_attention.out":
134	- position, key = 5, "context_lens"
135	- if len(graph_dispatch_record.args) >= (position + 1):
136	- graph_dispatch_record.args[position] = update_input[key]
137	- graph_dispatch_record.op_cache_entry(graph_dispatch_record.args, *graph_dispatch_record.kwargs)
138	- graph_task_update_end(self.update_stream)
139	- graph_dispatch_record.event.record(self.update_stream)
140		134
141	- def _append_dispatch_record(self, event, handle, args, kwargs, func):	135	+ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
		136	+ # Registry stores class objects; call via cls.method()
		137	+ handler_cls = _NPU_GRAPH_OP_HANDLERS.get(func.__name__)
		138	+
		139	+ if handler_cls:
		140	+ # 1) Common: obtain stream and event
		141	+ stream = torch_npu.npu.current_stream()
		142	+ event = torch.npu.ExternalEvent()
		143	+ event.wait(stream)
		144	+ event.reset(stream)
		145	+
		146	+ # 2) Operator-specific: preprocessing (workspace, output pre-alloc, func swap)
		147	+ actual_func, args, kwargs = handler_cls.prepare_capture(
		148	+ func, args, kwargs
		149	+ )
		150	+
		151	+ # 3) Common: parse operator schema
		152	+ self.update_schema(
		153	+ str(actual_func.__name__), str(actual_func._schema)
		154	+ )
		155	+
		156	+ # 4) Common: record graph task group
		157	+ graph_task_group_begin(stream)
		158	+ result = actual_func(args, *kwargs)
		159	+ handle = graph_task_group_end(stream)
		160	+
		161	+ # 5) Common: create dispatch record (delegate kwarg conversion to handler)
		162	+ self.graph_dispatch_records.append(
		163	+ self._append_dispatch_record(
		164	+ event, handle, args, kwargs, actual_func, handler_cls
		165	+ )
		166	+ )
		167	+
		168	+ # 6) Operator-specific: post-process return value
		169	+ return handler_cls.postprocess_result(result, kwargs)
		170	+
		171	+ return func(args, *kwargs)
		172	+
		173	+ def _append_dispatch_record(
		174	+ self, event, handle, args, kwargs, func, handler_cls
		175	+ ):
		176	+ """Create a dispatch record, converting args / kwargs to weak-refs or deep copies.
		177	+
		178	+ ``handler_cls`` is a required parameter (class object) guaranteed by
		179	+ the capture skeleton.
		180	+ """
142	args_ref = []	181	args_ref = []
143	for element in args:	182	for element in args:
144	if torch.is_tensor(element) and "npu" in str(element.device):	183	if torch.is_tensor(element) and "npu" in str(element.device):
145	args_ref.append(TensorWeakRef(element))	184	args_ref.append(TensorWeakRef(element))
146	else:	185	else:
147	args_ref.append(deepcopy(element))	186	args_ref.append(deepcopy(element))
148	- kwargs_ref = {}
149	- for key, vaule in kwargs.items():
150	- if key == "out":
151	- kwargs_ref[key] = [TensorWeakRef(vaule[0]), TensorWeakRef(vaule[1])]
152	- elif key in self.tensor_schema_name[str(func.__name__)]:
153	- kwargs_ref[key] = TensorWeakRef(vaule)
154	- else:
155	- kwargs_ref[key] = deepcopy(vaule)
156	- return _GraphDispatchRecord(event=event, handle=handle, kwargs=kwargs_ref, args=list(args_ref), op_cache_entry=func)
157		187
158	- def __torch_dispatch__(self, func, types, args=(), kwargs=None):	188	+ tensor_param_names = self.tensor_schema_name.get(
159	- if func.__name__ in ["npu_fused_infer_attention_score", "npu_fused_infer_attention_score.default"]:	189	+ str(func.__name__), []
160	- func_out = torch_npu.npu_fused_infer_attention_score.out	190	+ )
161	- self.update_schema(str(func_out.__name__), str(func_out._schema))	191	+ kwargs_ref = {}
162	- stream = torch_npu.npu.current_stream()	192	+ for key, value in kwargs.items():
163	- event = torch.npu.ExternalEvent()	193	+ kwargs_ref[key] = handler_cls.record_wrap_kwarg(
164	- event.wait(stream)	194	+ key, value, tensor_param_names
165	- event.reset(stream)	195	+ )
166	- # apply tensor	196	+
167	- workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(args, *kwargs)	197	+ return _GraphDispatchRecord(
168	- out_args = [args[0], args[2]]	198	+ event=event,
169	- out_kwargs_keys = [	199	+ handle=handle,
170	- 'input_layout',	200	+ kwargs=kwargs_ref,
171	- 'quant_scale2',	201	+ args=list(args_ref),
172	- 'block_table',	202	+ op_cache_entry=func,
173	- 'num_heads',	203	+ )
174	- 'num_key_value_heads',	204	+
175	- 'softmax_lse_flag',	205	+ # -----------------------------------------------------------------
176	- 'query_rope']	206	+ # Update skeleton
177	- out_kwargs = {key: kwargs[key] for key in out_kwargs_keys if key in kwargs}	207	+ # -----------------------------------------------------------------
178	- output, softmax_lse = torch_npu._npu_fused_infer_attention_score_infer_output(out_args, *out_kwargs)	208	+
179	- kwargs["workspace"] = workspace	209	+ def update_capture_record(self, cpu_update_input):
180	- kwargs["out"] = [output, softmax_lse]	210	+ if len(cpu_update_input) == 1:
181	- # begin graph task	211	+ new_list = [
182	- graph_task_group_begin(stream)	212	+ cpu_update_input[0].copy()
183	- func_out(args, *kwargs)	213	+ for _ in range(len(self.graph_dispatch_records))
184	- handle = graph_task_group_end(stream)	214	+ ]
185	- # save state for update	215	+ cpu_update_input = new_list
186	- self.graph_dispatch_records.append(	216	+
187	- self._append_dispatch_record(event, handle, args, kwargs, func_out))	217	+ # BUG FIX: original code compared len(self.graph_dispatch_records)
188	- return kwargs["out"]	218	+ # with itself -- always True. Corrected to compare with
189	- elif func.__name__ in ["npu_fused_infer_attention_score_v2", "npu_fused_infer_attention_score_v2.default"]:	219	+ # cpu_update_input.
190	- func_out = torch_npu.npu_fused_infer_attention_score_v2.out	220	+ if len(cpu_update_input) != len(self.graph_dispatch_records):
191	- self.update_schema(str(func_out.__name__), str(func_out._schema))	221	+ raise RuntimeError(
192	- stream = torch_npu.npu.current_stream()	222	+ f"Currently, there are {len(self.graph_dispatch_records)} "
193	- event = torch.npu.ExternalEvent()	223	+ f"operators that need to be updated by capture, and there "
194	- event.wait(stream)	224	+ f"are only {len(cpu_update_input)} elements in the incoming "
195	- event.reset(stream)	225	+ f"cpu_update_input list",
196	- # apply tensor	226	+ pta_error(ErrCode.PARAM),
197	- workspace = torch_npu._npu_fused_infer_attention_score_v2_get_max_workspace(args, *kwargs)	227	+ )
198	- out_args = [args[0], args[2]]	228	+
199	- out_kwargs_keys = [	229	+ with torch.npu.stream(self.update_stream):
200	- 'query_dtype',	230	+ for record, update_input in zip(
201	- 'value_dtype',	231	+ self.graph_dispatch_records, cpu_update_input
202	- 'input_layout',	232	+ ):
203	- 'quant_scale_out',	233	+ graph_task_update_begin(self.update_stream, record.handle)
204	- 'block_table',	234	+
205	- 'num_query_heads',	235	+ # 7) Common: update matching kwargs (direct assignment,
206	- 'num_key_value_heads',	236	+ # consistent with original implementation).
207	- 'return_softmax_lse',	237	+ # Capture-phase uses TensorWeakRef/deepcopy to avoid
208	- 'query_rope',	238	+ # strong references; update is a short "assign -> replay"
209	- 'out_dtype']	239	+ # flow where direct assignment suffices.
210	- out_kwargs = {key: kwargs[key] for key in out_kwargs_keys if key in kwargs}	240	+ for key in update_input:
211	- output, softmax_lse = torch_npu._npu_fused_infer_attention_score_v2_infer_output(out_args, *out_kwargs)	241	+ if key in record.kwargs:
212	- kwargs["workspace"] = workspace	242	+ record.kwargs[key] = update_input[key]
213	- kwargs["out"] = [output, softmax_lse]	243	+
214	- # begin graph task	244	+ # 8) Operator-specific: update args by index.
215	- graph_task_group_begin(stream)	245	+ # Defensive assert -- unregistered ops go through
216	- func_out(args, *kwargs)	246	+ # passthrough in capture and never produce a dispatch
217	- handle = graph_task_group_end(stream)	247	+ # record, so handler_cls must not be None here.
218	- # save state for update	248	+ handler_cls = _NPU_GRAPH_OP_HANDLERS.get(
219	- self.graph_dispatch_records.append(	249	+ record.op_cache_entry.__name__
220	- self._append_dispatch_record(event, handle, args, kwargs, func_out))	250	+ )
221	- return kwargs["out"]	251	+ if handler_cls is None:
222	- elif func.__name__ in ["npu_fused_infer_attention_score.out", "npu_fused_infer_attention_score_v2.out"]:	252	+ raise RuntimeError(
223	- self.update_schema(str(func.__name__), str(func._schema))	253	+ f"No handler for recorded op: {record.op_cache_entry.__name__}. "
224	- stream = torch_npu.npu.current_stream()	254	+ f"This indicates the handler was unregistered between capture and update.",
225	- event = torch.npu.ExternalEvent()	255	+ pta_error(ErrCode.PARAM),
226	- event.wait(stream)	256	+ )
227	- event.reset(stream)	257	+ handler_cls.update_args(record, update_input)
228	- # begin graph task	258	+
229	- graph_task_group_begin(stream)	259	+ # 9) Common: replay the operator
230	- func(args, *kwargs)	260	+ record.op_cache_entry(record.args, *record.kwargs)
231	- handle = graph_task_group_end(stream)	261	+ graph_task_update_end(self.update_stream)
232	- # save state for update	262	+ record.event.record(self.update_stream)
233	- self.graph_dispatch_records.append(
234	- self._append_dispatch_record(event, handle, args, kwargs, func))
235	- return kwargs["out"]
236	- elif func.__name__ in ["_npu_paged_attention.default", "npu_multi_head_latent_attention.out"]:
237	- self.update_schema(str(func.__name__), str(func._schema))
238	- stream = torch_npu.npu.current_stream()
239	- event = torch.npu.ExternalEvent()
240	- event.wait(stream)
241	- event.reset(stream)
242	- # begin graph task
243	- graph_task_group_begin(stream)
244	- result = func(args, *kwargs)
245	- handle = graph_task_group_end(stream)
246	- # save state for update
247	- self.graph_dispatch_records.append(
248	- self._append_dispatch_record(event, handle, args, kwargs, func))
249	- return result
250	- else:
251	- return func(args, *kwargs)
252		263
253		264
254	# Python shim helps Sphinx process docstrings more reliably.	265	# Python shim helps Sphinx process docstrings more reliably.
	@@ -261,7 +272,7 @@ class NPUGraph(torch_npu._C._NPUGraph):
261		272
262	def __new__(cls):	273	def __new__(cls):
263	return super().__new__(cls)	274	return super().__new__(cls)
264	-	275	+
265	def __init__(self):	276	def __init__(self):
266	self.graph_dispatch_mode = _GraphDispatchMode()	277	self.graph_dispatch_mode = _GraphDispatchMode()
267	self.auto_dispatch_capture = False	278	self.auto_dispatch_capture = False
	@@ -315,8 +326,11 @@ class NPUGraph(torch_npu._C._NPUGraph):
315		326
316	def update(self, cpu_update_input):	327	def update(self, cpu_update_input):
317	if not self.auto_dispatch_capture:	328	if not self.auto_dispatch_capture:
318	- raise RuntimeError("The current graph configuration does not support update,"	329	+ raise RuntimeError(
319	- "Try to capture by setting auto_dispatch_capture=True during capture", pta_error(ErrCode.PARAM))	330	+ "The current graph configuration does not support update,"
		331	+ "Try to capture by setting auto_dispatch_capture=True during capture",
		332	+ pta_error(ErrCode.PARAM),
		333	+ )
320	self.graph_dispatch_mode.update_capture_record(cpu_update_input)	334	self.graph_dispatch_mode.update_capture_record(cpu_update_input)
321		335
322	def debug_dump(self, debug_path):	336	def debug_dump(self, debug_path):
	@@ -324,8 +338,8 @@ class NPUGraph(torch_npu._C._NPUGraph):
324		338
325	Arguments:	339	Arguments:
326	debug_path (required): Path to dump the graph to.	340	debug_path (required): Path to dump the graph to.
327	- """	341	+ """
328	- return super().debug_dump(debug_path)	342	+ return super().debug_dump(debug_path)
329		343
330		344
331	class graph:	345	class graph:
	@@ -699,4 +713,4 @@ def make_graphed_callables(
699	if just_one_callable:	713	if just_one_callable:
700	return ret[0]	714	return ret[0]
701		715
702	- return tuple(ret)	716	+ return tuple(ret)