add flexattention score_mod support zeros_and

Atest/_inductor/test_flexattnetion.py+55-0

@@ -0,0 +1,55 @@
	1	+import functools
	2	+
	3	+import torch
	4	+import torch_npu
	5	+import torch_npu._inductor # noqa: F401
	6	+from torch._inductor import metrics
	7	+from torch._inductor.utils import run_and_get_code
	8	+from torch.nn.attention.flex_attention import flex_attention
	9	+from torch.testing import FileCheck
	10	+from torch_npu.testing.testcase import TestCase, run_tests
	11	+
	12	+
	13	+class TestFlexAttention(TestCase):
	14	+ def setUp(self):
	15	+ super().setUp()
	16	+ torch._dynamo.reset()
	17	+ metrics.reset()
	18	+
	19	+ def test_epilogue_fused(self):
	20	+ @torch.compile
	21	+ def f(q, k, v):
	22	+ return flex_attention(q, k, v).cos()
	23	+
	24	+ q, k, v = (
	25	+ torch.randn(1, 8, 1024, 64, device="npu") for _ in range(3)
	26	+ )
	27	+ _, code = run_and_get_code(f, q, k, v)
	28	+
	29	+ FileCheck().check("triton_tem_fused").check_not("poi_fused_cos").run(
	30	+ code[0]
	31	+ )
	32	+
	33	+ def test_kernel_options_argument_is_respected(self):
	34	+ make_tensor = functools.partial(
	35	+ torch.randn,
	36	+ (2, 2, 128, 64),
	37	+ device="npu",
	38	+ dtype=torch.float32,
	39	+ requires_grad=True,
	40	+ )
	41	+ q, k, v = make_tensor(), make_tensor(), make_tensor()
	42	+
	43	+ _, code = run_and_get_code(
	44	+ torch.compile(flex_attention),
	45	+ q,
	46	+ k,
	47	+ v,
	48	+ kernel_options={"BLOCK_M": 16},
	49	+ )
	50	+
	51	+ FileCheck().check("BLOCK_M : tl.constexpr = 16").run(code[0])
	52	+
	53	+
	54	+if __name__ == "__main__":
	55	+ run_tests()

Mtorch_npu/_inductor/__init__.py+0-2

	@@ -101,7 +101,6 @@ def _load_triton_backend():
101	)	101	)
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
105	from .fx_passes import patch_pattern_mm_plus_mm, register_fav3_partition_pass	104	from .fx_passes import patch_pattern_mm_plus_mm, register_fav3_partition_pass
106	from .fx_passes.graph_match_pass import (	105	from .fx_passes.graph_match_pass import (
107	post_grad_custom_pass_fuc,	106	post_grad_custom_pass_fuc,
	@@ -226,7 +225,6 @@ def _load_triton_backend():
226	parallel_scheduler()	225	parallel_scheduler()
227		226
228	patch_get_optimization_cflags()	227	patch_get_optimization_cflags()
229	- patch_extract_read_writes()
230	patch_count_bytes()	228	patch_count_bytes()
231	patch_tuning_process()	229	patch_tuning_process()
232		230

Dtorch_npu/_inductor/dependencies.py+0-29

	@@ -1,29 +0,0 @@
1	-from typing import Any, Callable, Sequence
2	-
3	-import sympy
4	-import torch._inductor.dependencies as dependencies
5	-
6	-origin_extract_read_writes = dependencies.extract_read_writes
7	-
8	-def patch_extract_read_writes():
9	- def extract_read_writes(
10	- fn: Callable[..., Any],
11	- *argsizes: Sequence[sympy.Expr],
12	- normalize: bool = False,
13	- prefix: str = "d",
14	- hidden_args: Sequence[list[sympy.Expr]] = (),
15	- ) -> dependencies.ReadWrites:
16	- # NPU does not support normalize load/store
17	- return origin_extract_read_writes(
18	- fn,
19	- *argsizes,
20	- normalize=False,
21	- prefix=prefix,
22	- hidden_args=hidden_args
23	- )
24	-
25	- dependencies.extract_read_writes = extract_read_writes
26	- import torch._inductor.ir as ir
27	- import torch._inductor.scheduler as scheduler
28	- ir.extract_read_writes = extract_read_writes
29	- scheduler.extract_read_writes = extract_read_writes

Mtorch_npu/_inductor/kernel/flex_attention.py+18-2

	@@ -9,7 +9,6 @@ import sympy
9		9
10	import torch	10	import torch
11	from torch._inductor.virtualized import V, ops	11	from torch._inductor.virtualized import V, ops
12	-from torch.utils._ordered_set import OrderedSet
13	from torch.utils._pytree import tree_map	12	from torch.utils._pytree import tree_map
14	from torch.utils._sympy.functions import FloorDiv, Mod	13	from torch.utils._sympy.functions import FloorDiv, Mod
15		14
	@@ -718,6 +717,7 @@ def _get_flex_attention_additional_lowerings():
718	additional_lowerings[aten.bitwise_or.Tensor] = bitwise_or_tensor	717	additional_lowerings[aten.bitwise_or.Tensor] = bitwise_or_tensor
719	additional_lowerings[aten.bitwise_not.default] = bitwise_not_default	718	additional_lowerings[aten.bitwise_not.default] = bitwise_not_default
720	additional_lowerings[aten.remainder.Scalar] = remainder_scalar	719	additional_lowerings[aten.remainder.Scalar] = remainder_scalar
		720	+ additional_lowerings[torch.ops.flex_lib.zeros_and_scatter.default] = zeros_and_scatter_lowering
721		721
722	return additional_lowerings	722	return additional_lowerings
723		723
	@@ -730,7 +730,7 @@ def _build_subgraph_buffer_with_additional_lowerings(args, subgraph):
730	to handle supported fallback operations as pointwise ops.	730	to handle supported fallback operations as pointwise ops.
731	"""	731	"""
732	from torch._inductor.subgraph_lowering import PointwiseSubgraphLowering	732	from torch._inductor.subgraph_lowering import PointwiseSubgraphLowering
733	-	733	+ from torch.utils._ordered_set import OrderedSet
734	additional_lowerings = _get_flex_attention_additional_lowerings()	734	additional_lowerings = _get_flex_attention_additional_lowerings()
735	zeros_and_scatter = torch.ops.flex_lib.zeros_and_scatter.default	735	zeros_and_scatter = torch.ops.flex_lib.zeros_and_scatter.default
736	additional_lowerings[zeros_and_scatter] = zeros_and_scatter_lowering	736	additional_lowerings[zeros_and_scatter] = zeros_and_scatter_lowering
	@@ -742,6 +742,11 @@ def _build_subgraph_buffer_with_additional_lowerings(args, subgraph):
742	)	742	)
743	with V.set_graph_handler(pw_subgraph):	743	with V.set_graph_handler(pw_subgraph):
744	pw_subgraph.run(*args)	744	pw_subgraph.run(*args)
		745	+ # Since we are allowing mutations/buffer creation, we need to register any fresh buffers
		746	+ # creating during the pointwise subgraph lowering
		747	+ if len(pw_subgraph.buffers) > 0:
		748	+ for buffer in pw_subgraph.buffers:
		749	+ V.graph.register_buffer(buffer)
745		750
746	# Older PointwiseSubgraphLowering versions defer approved mutation buffers.	751	# Older PointwiseSubgraphLowering versions defer approved mutation buffers.
747	for buffer in pw_subgraph.buffers:	752	for buffer in pw_subgraph.buffers:
	@@ -1185,6 +1190,7 @@ def _lower_flex_attention_mask_in(
1185	configs = generate_fwd_candidate_configs(	1190	configs = generate_fwd_candidate_configs(
1186	sparse_q_block_size=sparse_q_block_size,	1191	sparse_q_block_size=sparse_q_block_size,
1187	sparse_kv_block_size=sparse_kv_block_size,	1192	sparse_kv_block_size=sparse_kv_block_size,
		1193	+ kernel_options=kernel_options,
1188	)	1194	)
1189	if not configs:	1195	if not configs:
1190	raise RuntimeError(	1196	raise RuntimeError(
	@@ -1428,6 +1434,7 @@ def _lower_flex_attention_backward_mask_in(
1428	sparse_q_block_size=sparse_q_block_size,	1434	sparse_q_block_size=sparse_q_block_size,
1429	sparse_kv_block_size=sparse_kv_block_size,	1435	sparse_kv_block_size=sparse_kv_block_size,
1430	mode=FlexMode.BWD,	1436	mode=FlexMode.BWD,
		1437	+ kernel_options=kernel_options,
1431	)	1438	)
1432	if not configs:	1439	if not configs:
1433	raise RuntimeError(	1440	raise RuntimeError(
	@@ -1662,6 +1669,9 @@ def _register_npu_inductor_flex_attention():
1662	has_explicit_score_mod = bool(	1669	has_explicit_score_mod = bool(
1663	kernel_options.pop(_EXPLICIT_SCORE_MOD_OPTION, False)	1670	kernel_options.pop(_EXPLICIT_SCORE_MOD_OPTION, False)
1664	)	1671	)
		1672	+ # Strip GPU-specific backend selector (e.g. "TRITON"/"FLASH"/"CUDNN") that
		1673	+ # has no meaning on NPU and would leak into Triton constexpr parameters.
		1674	+ kernel_options.pop("BACKEND", None)
1665	# Mark symbols in custom kernel options as static shapes and add guards.	1675	# Mark symbols in custom kernel options as static shapes and add guards.
1666	kernel_options = {	1676	kernel_options = {
1667	k: V.graph.sizevars.guard_int(v)	1677	k: V.graph.sizevars.guard_int(v)
	@@ -1932,6 +1942,7 @@ def _register_npu_inductor_flex_attention():
1932	dict_configs = generate_fwd_candidate_configs(	1942	dict_configs = generate_fwd_candidate_configs(
1933	sparse_q_block_size=SPARSE_Q_BLOCK_SIZE,	1943	sparse_q_block_size=SPARSE_Q_BLOCK_SIZE,
1934	sparse_kv_block_size=SPARSE_KV_BLOCK_SIZE,	1944	sparse_kv_block_size=SPARSE_KV_BLOCK_SIZE,
		1945	+ kernel_options=kernel_options,
1935	)	1946	)
1936		1947
1937	if not dict_configs:	1948	if not dict_configs:
	@@ -2306,6 +2317,9 @@ def _register_npu_inductor_flex_attention():
2306	has_explicit_score_mod = bool(	2317	has_explicit_score_mod = bool(
2307	kernel_options.pop(_EXPLICIT_SCORE_MOD_OPTION, False)	2318	kernel_options.pop(_EXPLICIT_SCORE_MOD_OPTION, False)
2308	)	2319	)
		2320	+ # Strip GPU-specific backend selector (e.g. "TRITON"/"FLASH"/"CUDNN") that
		2321	+ # has no meaning on NPU and would leak into Triton constexpr parameters.
		2322	+ kernel_options.pop("BACKEND", None)
2309	configured_mask_out = bool(	2323	configured_mask_out = bool(
2310	npu_config.flex_attention.flexattention_mask_out	2324	npu_config.flex_attention.flexattention_mask_out
2311	)	2325	)
	@@ -2863,11 +2877,13 @@ def _register_npu_inductor_flex_attention():
2863	sparse_q_block_size=SPARSE_Q_BLOCK_SIZE,	2877	sparse_q_block_size=SPARSE_Q_BLOCK_SIZE,
2864	sparse_kv_block_size=SPARSE_KV_BLOCK_SIZE,	2878	sparse_kv_block_size=SPARSE_KV_BLOCK_SIZE,
2865	mode=FlexMode.BWDDQ,	2879	mode=FlexMode.BWDDQ,
		2880	+ kernel_options=kernel_options,
2866	)	2881	)
2867	bwd_dkdv_dict_configs = generate_bwd_candidate_configs(	2882	bwd_dkdv_dict_configs = generate_bwd_candidate_configs(
2868	sparse_q_block_size=SPARSE_Q_BLOCK_SIZE,	2883	sparse_q_block_size=SPARSE_Q_BLOCK_SIZE,
2869	sparse_kv_block_size=SPARSE_KV_BLOCK_SIZE,	2884	sparse_kv_block_size=SPARSE_KV_BLOCK_SIZE,
2870	mode=FlexMode.BWDDKDV,	2885	mode=FlexMode.BWDDKDV,
		2886	+ kernel_options=kernel_options,
2871	)	2887	)
2872		2888
2873	tasklist_reduce_ub_safe = True	2889	tasklist_reduce_ub_safe = True

Mtorch_npu/_inductor/kernel/flex_attention_config_generator.py+88-0

	@@ -2,6 +2,7 @@
2		2
3	from dataclasses import dataclass	3	from dataclasses import dataclass
4	from enum import Enum	4	from enum import Enum
		5	+from numbers import Integral
5	from typing import Optional	6	from typing import Optional
6		7
7	from torch._inductor import config as inductor_config	8	from torch._inductor import config as inductor_config
	@@ -73,6 +74,7 @@ class FlexAttentionConfigGenerator:
73	sparse_q_block_size: Optional[int] = None,	74	sparse_q_block_size: Optional[int] = None,
74	sparse_kv_block_size: Optional[int] = None,	75	sparse_kv_block_size: Optional[int] = None,
75	mode: FlexMode = FlexMode.FWD,	76	mode: FlexMode = FlexMode.FWD,
		77	+ kernel_options: Optional[dict] = None,
76	):	78	):
77	self.sparse_q_block_size = self._normalize_sparse_block_size(	79	self.sparse_q_block_size = self._normalize_sparse_block_size(
78	sparse_q_block_size	80	sparse_q_block_size
	@@ -81,6 +83,7 @@ class FlexAttentionConfigGenerator:
81	sparse_kv_block_size	83	sparse_kv_block_size
82	)	84	)
83	self.mode = mode	85	self.mode = mode
		86	+ self.kernel_options = dict(kernel_options or {})
84		87
85	self.valid_block_m = self._get_valid_block_sizes(	88	self.valid_block_m = self._get_valid_block_sizes(
86	self.sparse_q_block_size	89	self.sparse_q_block_size
	@@ -97,6 +100,87 @@ class FlexAttentionConfigGenerator:
97	self.valid_block_m = common_blocks	100	self.valid_block_m = common_blocks
98	self.valid_block_n = common_blocks	101	self.valid_block_n = common_blocks
99		102
		103	+ block_m_keys, block_n_keys = self._block_option_keys()
		104	+ user_block_m = self._resolve_user_block(block_m_keys)
		105	+ user_block_n = self._resolve_user_block(block_n_keys)
		106	+ block_m_sparse_sizes = (self.sparse_q_block_size,)
		107	+ block_n_sparse_sizes = (self.sparse_kv_block_size,)
		108	+ if self.mode == FlexMode.BWD:
		109	+ block_m_sparse_sizes += (self.sparse_kv_block_size,)
		110	+ block_n_sparse_sizes += (self.sparse_q_block_size,)
		111	+ self.valid_block_m = self._apply_user_block(
		112	+ self.valid_block_m,
		113	+ user_block_m,
		114	+ block_m_keys,
		115	+ block_m_sparse_sizes,
		116	+ )
		117	+ self.valid_block_n = self._apply_user_block(
		118	+ self.valid_block_n,
		119	+ user_block_n,
		120	+ block_n_keys,
		121	+ block_n_sparse_sizes,
		122	+ )
		123	+
		124	+ def _block_option_keys(self) -> tuple[tuple[str, ...], tuple[str, ...]]:
		125	+ if self.mode == FlexMode.FWD:
		126	+ return (("BLOCK_M",), ("BLOCK_N",))
		127	+ if self.mode == FlexMode.BWD:
		128	+ return (("BLOCK_M1", "BLOCK_N2"), ("BLOCK_N1", "BLOCK_M2"))
		129	+ if self.mode == FlexMode.BWDDQ:
		130	+ return (("BLOCK_M2",), ("BLOCK_N2",))
		131	+ if self.mode == FlexMode.BWDDKDV:
		132	+ return (("BLOCK_M1",), ("BLOCK_N1",))
		133	+ raise ValueError(f"unsupported flex attention mode: {self.mode}")
		134	+
		135	+ def _resolve_user_block(self, option_keys: tuple[str, ...]) -> Optional[int]:
		136	+ user_options = []
		137	+ for key in option_keys:
		138	+ if key not in self.kernel_options:
		139	+ continue
		140	+ value = self.kernel_options[key]
		141	+ if isinstance(value, bool) or not isinstance(value, Integral):
		142	+ raise ValueError(f"{key} must be an integer, got {value!r}")
		143	+ block_size = int(value)
		144	+ if block_size <= 0 or block_size & (block_size - 1):
		145	+ raise ValueError(
		146	+ f"{key} must be a positive power of 2, got {block_size}"
		147	+ )
		148	+ user_options.append((key, block_size))
		149	+ if not user_options:
		150	+ return None
		151	+
		152	+ first_key, block_size = user_options[0]
		153	+ for key, value in user_options[1:]:
		154	+ if value != block_size:
		155	+ raise ValueError(
		156	+ "Conflicting kernel options: "
		157	+ f"{first_key}={block_size} and {key}={value}"
		158	+ )
		159	+ return block_size
		160	+
		161	+ @staticmethod
		162	+ def _apply_user_block(
		163	+ generated_blocks: list[int],
		164	+ user_block: Optional[int],
		165	+ option_keys: tuple[str, ...],
		166	+ sparse_block_sizes: tuple[Optional[int], ...],
		167	+ ) -> list[int]:
		168	+ if user_block is None:
		169	+ return generated_blocks
		170	+ for sparse_block_size in sparse_block_sizes:
		171	+ if sparse_block_size is None:
		172	+ continue
		173	+ if (
		174	+ user_block > sparse_block_size
		175	+ or sparse_block_size % user_block != 0
		176	+ ):
		177	+ option_name = "/".join(option_keys)
		178	+ raise ValueError(
		179	+ f"{option_name}={user_block} is incompatible with "
		180	+ f"sparse block size {sparse_block_size}"
		181	+ )
		182	+ return [user_block]
		183	+
100	def _get_valid_block_sizes(	184	def _get_valid_block_sizes(
101	self, sparse_block_size: Optional[int]	185	self, sparse_block_size: Optional[int]
102	) -> list[int]:	186	) -> list[int]:
	@@ -179,12 +263,14 @@ def prefer_max_tiling_without_benchmark() -> bool:
179	def generate_fwd_candidate_configs(	263	def generate_fwd_candidate_configs(
180	sparse_q_block_size: int,	264	sparse_q_block_size: int,
181	sparse_kv_block_size: int,	265	sparse_kv_block_size: int,
		266	+ kernel_options: Optional[dict] = None,
182	) -> list[dict]:	267	) -> list[dict]:
183	"""Generate valid forward configs."""	268	"""Generate valid forward configs."""
184	return FlexAttentionConfigGenerator(	269	return FlexAttentionConfigGenerator(
185	sparse_q_block_size=sparse_q_block_size,	270	sparse_q_block_size=sparse_q_block_size,
186	sparse_kv_block_size=sparse_kv_block_size,	271	sparse_kv_block_size=sparse_kv_block_size,
187	mode=FlexMode.FWD,	272	mode=FlexMode.FWD,
		273	+ kernel_options=kernel_options,
188	).generate_configs()	274	).generate_configs()
189		275
190		276
	@@ -234,6 +320,7 @@ def generate_bwd_candidate_configs(
234	sparse_q_block_size: int,	320	sparse_q_block_size: int,
235	sparse_kv_block_size: int,	321	sparse_kv_block_size: int,
236	mode: FlexMode,	322	mode: FlexMode,
		323	+ kernel_options: Optional[dict] = None,
237	) -> list[dict]:	324	) -> list[dict]:
238	"""Generate final block configs for a fused or split backward template."""	325	"""Generate final block configs for a fused or split backward template."""
239	if mode not in (FlexMode.BWD, FlexMode.BWDDQ, FlexMode.BWDDKDV):	326	if mode not in (FlexMode.BWD, FlexMode.BWDDQ, FlexMode.BWDDKDV):
	@@ -242,6 +329,7 @@ def generate_bwd_candidate_configs(
242	sparse_q_block_size=sparse_q_block_size,	329	sparse_q_block_size=sparse_q_block_size,
243	sparse_kv_block_size=sparse_kv_block_size,	330	sparse_kv_block_size=sparse_kv_block_size,
244	mode=mode,	331	mode=mode,
		332	+ kernel_options=kernel_options,
245	).generate_configs()	333	).generate_configs()
246		334
247		335

Mtorch_npu/_inductor/lowering_fallback_list.py+0-7

	@@ -58,12 +58,6 @@ NPU_EXTRA_FALLBACK_LIST = [
58	aten.__or__,	58	aten.__or__,
59	aten.__or__.bool,	59	aten.__or__.bool,
60	aten.__or__.int,	60	aten.__or__.int,
61	- aten.__rshift__,
62	- aten.__rshift__.Scalar,
63	- aten.__rshift__.Scalar_out,
64	- aten.__rshift__.Tensor,
65	- aten.__rshift__.Tensor_out,
66	- aten.__rshift__.int,
67	aten.__xor__,	61	aten.__xor__,
68	aten.__xor__.bool,	62	aten.__xor__.bool,
69	aten.__xor__.int,	63	aten.__xor__.int,
	@@ -115,7 +109,6 @@ NPU_EXTRA_FALLBACK_LIST = [
115	aten.add_.Scalar,	109	aten.add_.Scalar,
116	aten.add_.Tensor,	110	aten.add_.Tensor,
117	aten.add_.t,	111	aten.add_.t,
118	- aten.alias.default,
119	aten.any.all_out,	112	aten.any.all_out,
120	aten.any.bool,	113	aten.any.bool,
121	aten.any.dim,	114	aten.any.dim,