feat(inductor): partition FA v3 TND+dropout via inductor pattern matcher-pytorch-AtomGit

Atest/_inductor/test_fav3_partition_pass.py+225-0

@@ -0,0 +1,225 @@
	1	+# Owner(s): ["module: inductor"]
	2	+"""End-to-end tests for FA v3 graph partition pass."""
	3	+
	4	+import functools
	5	+
	6	+import numpy as np
	7	+import torch_npu
	8	+from testutils import TestUtils
	9	+from torch_npu.testing.common_utils import SupportedDevices
	10	+
	11	+import torch
	12	+from torch.testing._internal.common_utils import run_tests
	13	+
	14	+
	15	+# Loose tolerance for dropout paths (RNG ordering between eager and compiled
	16	+# wrapper differs slightly even when the underlying FA v3 op runs in eager).
	17	+_HEAD_NUM = 8
	18	+_HEAD_DIM = 64
	19	+_DTYPE = torch.float16
	20	+
	21	+
	22	+def _make_bnsd_inputs(batch=2, seq=128):
	23	+ shape = (batch, _HEAD_NUM, seq, _HEAD_DIM)
	24	+ q = torch.randn(shape, dtype=_DTYPE, device="npu", requires_grad=False)
	25	+ k = torch.randn(shape, dtype=_DTYPE, device="npu", requires_grad=False)
	26	+ v = torch.randn(shape, dtype=_DTYPE, device="npu", requires_grad=False)
	27	+ return q, k, v
	28	+
	29	+
	30	+def _make_tnd_inputs(batch=2, seq=128):
	31	+ total_tokens = batch * seq
	32	+ shape = (total_tokens, _HEAD_NUM, _HEAD_DIM)
	33	+ q = torch.randn(shape, dtype=_DTYPE, device="npu", requires_grad=False)
	34	+ k = torch.randn(shape, dtype=_DTYPE, device="npu", requires_grad=False)
	35	+ v = torch.randn(shape, dtype=_DTYPE, device="npu", requires_grad=False)
	36	+ # FA v3 schema requires actual_seq_qlen / actual_seq_kvlen to be Tensor (CPU int64).
	37	+ cu_seqlens = torch.tensor(
	38	+ np.arange(1, batch + 1) * seq,
	39	+ dtype=torch.int64,
	40	+ )
	41	+ return q, k, v, cu_seqlens
	42	+
	43	+
	44	+def _fa3_bnsd(q, k, v, keep_prob):
	45	+ out = torch.ops.npu.npu_fusion_attention_v3(
	46	+ q,
	47	+ k,
	48	+ v,
	49	+ _HEAD_NUM,
	50	+ "BNSD",
	51	+ keep_prob=keep_prob,
	52	+ )
	53	+ return out[0]
	54	+
	55	+
	56	+def _fa3_tnd(q, k, v, actual_seq_qlen, actual_seq_kvlen, keep_prob):
	57	+ out = torch.ops.npu.npu_fusion_attention_v3(
	58	+ q,
	59	+ k,
	60	+ v,
	61	+ _HEAD_NUM,
	62	+ "TND",
	63	+ actual_seq_qlen=actual_seq_qlen,
	64	+ actual_seq_kvlen=actual_seq_kvlen,
	65	+ keep_prob=keep_prob,
	66	+ )
	67	+ return out[0]
	68	+
	69	+
	70	+class TestFAv3PartitionPass(TestUtils):
	71	+ """End-to-end tests for `register_fav3_partition_pass`."""
	72	+
	73	+ @classmethod
	74	+ def setUpClass(cls):
	75	+ super().setUpClass()
	76	+ # Importing the inductor subpackage triggers register_fav3_partition_pass()
	77	+ # via the wiring in torch_npu/_inductor/__init__.py.
	78	+ from torch_npu._inductor.fx_passes import _proxy_ops, fav3_partition_pass
	79	+
	80	+ cls._pmp = fav3_partition_pass._FAV3_PMP
	81	+ cls._proxy_targets = set(_proxy_ops.PROXY_TARGETS.values())
	82	+
	83	+ # Disable inductor's persistent FX graph cache so each test triggers a
	84	+ # fresh post_grad_passes run -- otherwise the second compile of the
	85	+ # same callable hits cached output, our pass is skipped, and rewrite
	86	+ # counters falsely report 0.
	87	+ cls._fx_cache_orig = torch._inductor.config.fx_graph_cache
	88	+ torch._inductor.config.fx_graph_cache = False
	89	+ from torch_npu._inductor import config as npu_config
	90	+
	91	+ npu_config.npugraph_trees.disable_cpu_input_check = True
	92	+
	93	+ @classmethod
	94	+ def tearDownClass(cls):
	95	+ torch._inductor.config.fx_graph_cache = cls._fx_cache_orig
	96	+ super().tearDownClass()
	97	+ from torch_npu._inductor import config as npu_config
	98	+
	99	+ npu_config.npugraph_trees.disable_cpu_input_check = False
	100	+
	101	+ def setUp(self):
	102	+ super().setUp()
	103	+ # Wrap _FAV3_PMP.apply to count rewrites this test causes.
	104	+ self._rewrite_count = 0
	105	+ original_apply = type(self)._pmp.apply
	106	+
	107	+ @functools.wraps(original_apply)
	108	+ def counting_apply(gm):
	109	+ graph = gm.graph if hasattr(gm, "graph") else gm
	110	+ before = sum(
	111	+ 1
	112	+ for n in graph.nodes
	113	+ if n.op == "call_function" and n.target in type(self)._proxy_targets
	114	+ )
	115	+ result = original_apply(gm)
	116	+ after = sum(
	117	+ 1
	118	+ for n in graph.nodes
	119	+ if n.op == "call_function" and n.target in type(self)._proxy_targets
	120	+ )
	121	+ self._rewrite_count += max(0, after - before)
	122	+ return result
	123	+
	124	+ self._original_apply = original_apply
	125	+ type(self)._pmp.apply = counting_apply
	126	+ torch._dynamo.reset()
	127	+
	128	+ def tearDown(self):
	129	+ type(self)._pmp.apply = self._original_apply
	130	+ torch._dynamo.reset()
	131	+ super().tearDown()
	132	+
	133	+ # ---- helpers --------------------------------------------------------
	134	+
	135	+ def _run_eager_then_compiled(self, fn, args, compile_mode):
	136	+ torch.manual_seed(42)
	137	+ torch_npu.npu.manual_seed(42)
	138	+ eager_out = fn(*args)
	139	+
	140	+ torch.manual_seed(42)
	141	+ torch_npu.npu.manual_seed(42)
	142	+ compiled_fn = torch.compile(fn, mode=compile_mode, dynamic=False)
	143	+ compiled_out = compiled_fn(*args)
	144	+ return eager_out, compiled_out
	145	+
	146	+ # ---- scenarios ------------------------------------------------------
	147	+
	148	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	149	+ def test_tnd_dropout_reduce_overhead_partitions(self):
	150	+ """Regression target: TND + dropout must be partitioned out of ACLgraph."""
	151	+ torch.manual_seed(0)
	152	+ q, k, v, cu = _make_tnd_inputs()
	153	+ eager_out, compiled_out = self._run_eager_then_compiled(
	154	+ _fa3_tnd,
	155	+ (q, k, v, cu, cu, 0.5),
	156	+ compile_mode="reduce-overhead",
	157	+ )
	158	+
	159	+ # If the partition pass didn't rewrite, ACLgraph capture would crash
	160	+ # before reaching here -- so reaching the assertion is itself half the
	161	+ # signal. We additionally require a measurable rewrite.
	162	+ self.assertGreaterEqual(
	163	+ self._rewrite_count,
	164	+ 1,
	165	+ f"expected >=1 FA v3 node to be rewritten, got {self._rewrite_count}",
	166	+ )
	167	+ self.assertEqual(eager_out, compiled_out)
	168	+
	169	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	170	+ def test_tnd_no_dropout_reduce_overhead_no_rewrite(self):
	171	+ """TND with keep_prob=1.0 is safe in ACLgraph; no rewrite expected."""
	172	+ torch.manual_seed(0)
	173	+ q, k, v, cu = _make_tnd_inputs()
	174	+ eager_out, compiled_out = self._run_eager_then_compiled(
	175	+ _fa3_tnd,
	176	+ (q, k, v, cu, cu, 1.0),
	177	+ compile_mode="reduce-overhead",
	178	+ )
	179	+
	180	+ self.assertEqual(
	181	+ self._rewrite_count,
	182	+ 0,
	183	+ f"expected no FA v3 rewrite, got {self._rewrite_count}",
	184	+ )
	185	+ self.assertEqual(eager_out, compiled_out)
	186	+
	187	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	188	+ def test_bnsd_dropout_reduce_overhead_no_rewrite(self):
	189	+ """BNSD always safe regardless of dropout; no rewrite expected."""
	190	+ torch.manual_seed(0)
	191	+ q, k, v = _make_bnsd_inputs()
	192	+ eager_out, compiled_out = self._run_eager_then_compiled(
	193	+ _fa3_bnsd,
	194	+ (q, k, v, 0.5),
	195	+ compile_mode="reduce-overhead",
	196	+ )
	197	+
	198	+ self.assertEqual(
	199	+ self._rewrite_count,
	200	+ 0,
	201	+ f"expected no FA v3 rewrite for BNSD, got {self._rewrite_count}",
	202	+ )
	203	+ self.assertEqual(eager_out, compiled_out)
	204	+
	205	+ @SupportedDevices(["Ascend910B", "Ascend910_93"])
	206	+ def test_tnd_dropout_default_mode_early_out(self):
	207	+ """cudagraphs=OFF early-out: pass must be a no-op even with TND+dropout."""
	208	+ torch.manual_seed(0)
	209	+ q, k, v, cu = _make_tnd_inputs()
	210	+ eager_out, compiled_out = self._run_eager_then_compiled(
	211	+ _fa3_tnd,
	212	+ (q, k, v, cu, cu, 0.5),
	213	+ compile_mode="default",
	214	+ )
	215	+
	216	+ self.assertEqual(
	217	+ self._rewrite_count,
	218	+ 0,
	219	+ f"cudagraphs=OFF early-out: expected no rewrite, got {self._rewrite_count}",
	220	+ )
	221	+ self.assertEqual(eager_out, compiled_out)
	222	+
	223	+
	224	+if __name__ == "__main__":
	225	+ run_tests()

Mtorch_npu/_inductor/__init__.py+2-1

	@@ -102,7 +102,7 @@ def _load_triton_backend():
102	from .codegen.cpp_utils import patch_device_to_aten	102	from .codegen.cpp_utils import patch_device_to_aten
103	from .decomposition import _register_triton_decompositions	103	from .decomposition import _register_triton_decompositions
104	from .dependencies import patch_extract_read_writes	104	from .dependencies import patch_extract_read_writes
105	- from .fx_passes import patch_pattern_mm_plus_mm	105	+ from .fx_passes import patch_pattern_mm_plus_mm, register_fav3_partition_pass
106	from .fx_passes.graph_match_pass import (	106	from .fx_passes.graph_match_pass import (
107	post_grad_custom_pass_fuc,	107	post_grad_custom_pass_fuc,
108	pre_grad_custom_pass_fuc,	108	pre_grad_custom_pass_fuc,
	@@ -219,6 +219,7 @@ def _load_triton_backend():
219		219
220	pre_grad_custom_pass_fuc()	220	pre_grad_custom_pass_fuc()
221	post_grad_custom_pass_fuc()	221	post_grad_custom_pass_fuc()
		222	+ register_fav3_partition_pass()
222	if os.environ.get("ENABLE_PARALLEL_SCHEDULER", "false").lower() == "true":	223	if os.environ.get("ENABLE_PARALLEL_SCHEDULER", "false").lower() == "true":
223	from .fx_passes.parallel_scheduler_pass import parallel_scheduler	224	from .fx_passes.parallel_scheduler_pass import parallel_scheduler
224		225

Mtorch_npu/_inductor/fx_passes/__init__.py+5-0

	@@ -1 +1,6 @@
		1	+"""torch_npu inductor FX pass extensions."""
		2	+
		3	+__all__ = ["register_fav3_partition_pass"]
		4	+
		5	+from .fav3_partition_pass import register_fav3_partition_pass
1	from .post_grad import patch_pattern_mm_plus_mm	6	from .post_grad import patch_pattern_mm_plus_mm

Atorch_npu/_inductor/fx_passes/_proxy_ops.py+65-0

@@ -0,0 +1,65 @@
	1	+"""Proxy ops for FA v3 graph partition.
	2	+
	3	+Each proxy mirrors the schema of the underlying ``npu_fusion_attention_v3`` /
	4	+``npu_fusion_attention_grad_v3`` op, but is registered with
	5	+``Tag.cudagraph_unsafe`` so that inductor's scheduler partitions any FX node
	6	+whose ``target`` was rewritten to it. The kernel transparently forwards to the
	7	+original op, so the eager fallback path produced by graph partition runs the
	8	+unmodified FA v3 implementation.
	9	+"""
	10	+
	11	+__all__ = ["PROXY_TARGETS"]
	12	+
	13	+import torch
	14	+
	15	+
	16	+_LIB = torch.library.Library("npu", "FRAGMENT")
	17	+
	18	+
	19	+def _clone_schema_under_new_name(original_overload, new_unqualified_name: str) -> str:
	20	+ schema_str = str(original_overload._schema)
	21	+ head = original_overload._schema.name
	22	+ overload = original_overload._schema.overload_name
	23	+ if overload:
	24	+ head = f"{head}.{overload}"
	25	+ new_head = f"npu::{new_unqualified_name}"
	26	+ if overload:
	27	+ new_head = f"{new_head}.{overload}"
	28	+ if not schema_str.startswith(head):
	29	+ raise RuntimeError(
	30	+ f"Unexpected schema prefix for {original_overload}: {schema_str!r}"
	31	+ )
	32	+ return new_head + schema_str[len(head) :]
	33	+
	34	+
	35	+def _register_proxy(original_overload, new_unqualified_name: str):
	36	+ _LIB.define(
	37	+ _clone_schema_under_new_name(original_overload, new_unqualified_name),
	38	+ tags=[torch._C.Tag.cudagraph_unsafe],
	39	+ )
	40	+ proxy_overload = getattr(
	41	+ getattr(torch.ops.npu, new_unqualified_name),
	42	+ original_overload._overloadname,
	43	+ )
	44	+
	45	+ def _kernel(args, *kwargs):
	46	+ return original_overload(args, *kwargs)
	47	+
	48	+ _LIB.impl(proxy_overload.name(), _kernel, "CompositeExplicitAutograd")
	49	+ # Fake/meta impl: delegate to the original op so AOT tracing reuses its
	50	+ # device-propagation rules (FA v3 inputs mix npu q/k/v with cpu
	51	+ # actual_seq_qlen/kvlen tensors, which the default FakeTensor logic rejects).
	52	+ torch.library.register_fake(proxy_overload.name(), _kernel, lib=_LIB)
	53	+ return proxy_overload
	54	+
	55	+
	56	+PROXY_TARGETS = {
	57	+ torch.ops.npu.npu_fusion_attention_v3.default: _register_proxy(
	58	+ torch.ops.npu.npu_fusion_attention_v3.default,
	59	+ "npu_fusion_attention_v3_unsafe",
	60	+ ),
	61	+ torch.ops.npu.npu_fusion_attention_grad_v3.default: _register_proxy(
	62	+ torch.ops.npu.npu_fusion_attention_grad_v3.default,
	63	+ "npu_fusion_attention_grad_v3_unsafe",
	64	+ ),
	65	+}

Atorch_npu/_inductor/fx_passes/fav3_partition_pass.py+102-0

@@ -0,0 +1,102 @@
	1	+"""FA v3 graph partition pass via Inductor's PatternMatcherPass + POST_GRAD_PATTERNS.
	2	+
	3	+Identifies FA v3 forward/backward FX nodes that hit the dropout-on-TND path
	4	+(incompatible with ACLgraph capture) and rewrites their ``target`` to the
	5	+proxy ops registered in :mod:`._proxy_ops`. Those proxies carry
	6	+``Tag.cudagraph_unsafe``, so inductor's scheduler partitions them out of the
	7	+captured graph and falls back to eager.
	8	+
	9	+Wired up via :func:`register_fav3_partition_pass`, called from
	10	+``torch_npu/_inductor/__init__.py``.
	11	+"""
	12	+
	13	+__all__ = ["register_fav3_partition_pass"]
	14	+
	15	+import torch
	16	+from torch._inductor import config as inductor_config
	17	+from torch._inductor.fx_passes.post_grad import POST_GRAD_PATTERNS
	18	+from torch._inductor.pattern_matcher import (
	19	+ CallFunctionVarArgs,
	20	+ PatternMatcherPass,
	21	+ register_graph_pattern,
	22	+)
	23	+from torch.fx.operator_schemas import normalize_function
	24	+
	25	+# Returns True if installed CANN version >= given version.
	26	+from torch_npu.npu.utils import _is_gte_cann_version
	27	+
	28	+_CANN_VERSION = _is_gte_cann_version("9.1.0")
	29	+from . import _proxy_ops
	30	+
	31	+
	32	+
	33	+_PASS_KEY = "fav3_partition"
	34	+
	35	+# Module-level singleton -- Python module cache guarantees one-time registration
	36	+# of the rules below. Do NOT importlib.reload this module (would re-append
	37	+# handlers to _FAV3_PMP.patterns and raise on re-define in _proxy_ops).
	38	+_FAV3_PMP = PatternMatcherPass(pass_name=_PASS_KEY)
	39	+
	40	+
	41	+def _make_check_and_handler(target):
	42	+ proxy = _proxy_ops.PROXY_TARGETS[target]
	43	+
	44	+ def _check(match):
	45	+ # Early-out: when ACLgraph isn't going to capture, scheduler skips
	46	+ # partitioning entirely (scheduler.py: "partition includes all ops
	47	+ # when cudagraphs is disabled"); rewriting only adds a dispatcher hop.
	48	+ if not inductor_config.triton.cudagraphs:
	49	+ return False
	50	+
	51	+ node = match.nodes[0]
	52	+ normalized = normalize_function(
	53	+ node.target, node.args, node.kwargs, normalize_to_only_use_kwargs=True
	54	+ )
	55	+ if normalized is None:
	56	+ return False
	57	+ _, kwargs = normalized
	58	+
	59	+ keep_prob = kwargs.get("keep_prob", 1.0)
	60	+ input_layout = kwargs.get("input_layout")
	61	+ if not isinstance(input_layout, str):
	62	+ return False
	63	+ if not isinstance(keep_prob, (int, float)):
	64	+ return False
	65	+ if input_layout.upper() == "TND":
	66	+ if float(keep_prob) < 1.0:
	67	+ return True
	68	+ if float(keep_prob) == 1.0 and not _CANN_VERSION:
	69	+ return True
	70	+ return False
	71	+
	72	+ def _handler(match, args, *kwargs):
	73	+ node = match.nodes[0]
	74	+ node.target = proxy
	75	+
	76	+ return _check, _handler
	77	+
	78	+
	79	+def _register_rule_for(target):
	80	+ check, handler = _make_check_and_handler(target)
	81	+ register_graph_pattern(
	82	+ CallFunctionVarArgs(target),
	83	+ extra_check=check,
	84	+ pass_dict=_FAV3_PMP,
	85	+ )(handler)
	86	+
	87	+
	88	+# Module-level rule registration. Both fwd and bwd targets feed into the same
	89	+# _FAV3_PMP instance; inductor will dispatch on whichever matches per-graph.
	90	+_register_rule_for(torch.ops.npu.npu_fusion_attention_v3.default)
	91	+_register_rule_for(torch.ops.npu.npu_fusion_attention_grad_v3.default)
	92	+
	93	+
	94	+def register_fav3_partition_pass() -> None:
	95	+ """Mount the PMP into inductor's post-grad fusion dispatch.
	96	+
	97	+ Called once from ``torch_npu/_inductor/__init__.py`` alongside the other
	98	+ inductor extension hooks. Idempotent: dict assignment is overwrite-safe;
	99	+ setdefault preserves any user-supplied value.
	100	+ """
	101	+ POST_GRAD_PATTERNS[_PASS_KEY] = _FAV3_PMP
	102	+ inductor_config.post_grad_fusion_options.setdefault(_PASS_KEY, {})