import triton backend for mm and triton cv fusion to master branch-pytorch-AtomGit

已合并

import triton backend for mm and triton cv fusion to master branch #44574

shi-yufeng99创建于 8月14日从shi-yufeng99/pytorch： 0814_master_triton_cv_branch1 合入到 Ascend/pytorch： master

import triton backend for mm and triton cv fusion to master branch #44574

已合并

shi-yufeng99创建于 8月14日

从shi-yufeng99/pytorch 0814_master_triton_cv_branch1合入到Ascend/pytorch master

共 7 个文件变更+1072-15

Mtorch_npu/_inductor/__init__.py+4-1

	@@ -71,6 +71,7 @@ def _load_dvm_backend():
71	def _load_triton_backend():	71	def _load_triton_backend():
72	_apply_common_patches()	72	_apply_common_patches()
73	import torch	73	import torch
		74	+ torch._inductor.runtime.benchmarking.GPU_BENCHMARK_DEVICE_TYPES = ("cuda", "xpu", "mtia", "npu")
74	has_triton = torch.utils._triton.has_triton()	75	has_triton = torch.utils._triton.has_triton()
75	if not has_triton:	76	if not has_triton:
76	import warnings	77	import warnings
	@@ -89,7 +90,7 @@ def _load_triton_backend():
89	from .codecache import patch_get_cpp_wrapper_header	90	from .codecache import patch_get_cpp_wrapper_header
90	from .export import patch_aot_load	91	from .export import patch_aot_load
91	from .codegen._sizevars import patch_simplify	92	from .codegen._sizevars import patch_simplify
92	- from .codegen.ir import patch_indexing, patch_loop_body	93	+ from .codegen.ir import patch_fixed_indexer, patch_indexing, patch_loop_body
93	from .cpp_builder import (	94	from .cpp_builder import (
94	patch_get_cpp_torch_device_options,	95	patch_get_cpp_torch_device_options,
95	patch_get_optimization_cflags,	96	patch_get_optimization_cflags,
	@@ -202,6 +203,8 @@ def _load_triton_backend():
202	patch_num_splits()	203	patch_num_splits()
203	patch_loop_body()	204	patch_loop_body()
204	patch_indexing()	205	patch_indexing()
		206	+ patch_fixed_indexer()
		207	+
205	patch_create_device_properties()	208	patch_create_device_properties()
206	patch_load_cached_autotuning()	209	patch_load_cached_autotuning()
207	patch_triton_heuristics_cached_autotune()	210	patch_triton_heuristics_cached_autotune()

Mtorch_npu/_inductor/codegen/ir.py+68-0

	@@ -1972,6 +1972,74 @@ def patch_loop_body():
1972	LoopBody.substitube_indirect_index = substitube_indirect_index	1972	LoopBody.substitube_indirect_index = substitube_indirect_index
1973		1973
1974		1974
		1975	+def patch_fixed_indexer():
		1976	+ """Patch ``torch._inductor.ir._fixed_indexer`` to NOT skip size-1 dimensions.
		1977	+
		1978	+ Root cause background
		1979	+ ---------------------
		1980	+ The upstream ``_fixed_indexer`` contains an optimisation that skips the
		1981	+ ``idx * stride`` term for every dimension whose size equals 1::
		1982	+
		1983	+ for idx, st, sz in zip(index, stride, size):
		1984	+ if sz != 1: # <-- size-1 dims are dropped
		1985	+ result = result + idx * st
		1986	+
		1987	+ This is valid when idx is a real element index that is always 0 for a
		1988	+ size-1 dimension. However, inside a triton template (e.g. the mm /
		1989	+ bmm epilogue) the "index" handed to ``store_output`` is a block tensor
		1990	+ such as ``idx_n = rn[None, :]`` where
		1991	+ ``rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)``. ``idx_n`` ranges over
		1992	+ ``[0, BLOCK_N)`` — it is not bounded by the dimension size; the
		1993	+ out-of-bounds elements are simply masked away by ``mask = idx_n < N``.
		1994	+
		1995	+ When the output has a size-1 dimension (e.g. mm with N=1, output shape
		1996	+ ``[M, 1]``) the optimisation drops ``idx_n`` entirely, so
		1997	+ ``FixedLayout.make_indexer()`` returns ``idx_m`` instead of the full
		1998	+ contiguous expression ``idx_m + idx_n``.
		1999	+
		2000	+ Inside ``TritonTemplateKernel.store_output`` this makes
		2001	+ ``output_index (= idx_m) != contiguous_index (= idx_m + idx_n)``, so the
		2002	+ store falls back to the non-contiguous path and emits::
		2003	+
		2004	+ tl.store(out_ptr0 + tl.broadcast_to(idx_m, [BLOCK_M, BLOCK_N]), acc, mask)
		2005	+
		2006	+ ``idx_m`` has shape ``[BLOCK_M, 1]``; ``broadcast_to`` expands it to
		2007	+ ``[BLOCK_M, BLOCK_N]``. triton-ascend's MLIR backend cannot lower this
		2008	+ broadcast inside ``tt.store`` and raises::
		2009	+
		2010	+ MLIRCompilationError: 'tt.store' op failed to verify that
		2011	+ value type matches ptr type
		2012	+
		2013	+ The fix
		2014	+ -------
		2015	+ Always include every ``idx * stride`` term (remove the ``if sz != 1``
		2016	+ guard). For non-template callers the index of a size-1 dimension is a
		2017	+ concrete 0, so ``0 * stride == 0`` and sympy simplifies it away — the
		2018	+ generated code is unchanged. For template callers the full expression
		2019	+ ``idx_m + idx_n`` is preserved, which equals ``contiguous_index`` and
		2020	+ makes ``store_output`` use the correctly-shaped ``xindex`` variable::
		2021	+
		2022	+ tl.store(out_ptr0 + tl.broadcast_to(xindex, [BLOCK_M, BLOCK_N]), acc, mask)
		2023	+ """
		2024	+ from torch._inductor import ir as torch_ir
		2025	+
		2026	+ def _fixed_indexer_no_skip(size, stride=None, offset=torch_ir.Integer(0)):
		2027	+ """A closure containing math to read a given element."""
		2028	+
		2029	+ def indexer(index):
		2030	+ assert stride is not None and len(index) == len(stride)
		2031	+ assert len(index) == len(size)
		2032	+ result = offset
		2033	+ for idx, st, sz in zip(index, stride, size):
		2034	+ # NPU patch: keep the term even when sz == 1 (see docstring).
		2035	+ result = result + idx * st
		2036	+ return result
		2037	+
		2038	+ return indexer
		2039	+
		2040	+ torch_ir._fixed_indexer = _fixed_indexer_no_skip
		2041	+
		2042	+
1975	def patch_indexing():	2043	def patch_indexing():
1976	# todo: move patch function to loop_body.py and _sizevars.py	2044	# todo: move patch function to loop_body.py and _sizevars.py
1977	CaptureIndexing.index_select = loop_body_block_index_select	2045	CaptureIndexing.index_select = loop_body_block_index_select

Mtorch_npu/_inductor/codegen/scheduling.py+1-1

	@@ -300,7 +300,7 @@ class NPUTritonScheduling(TritonScheduling):
300	src_code = (	300	src_code = (
301	f"{kernel.imports_for_benchmark_kernel()}\n"	301	f"{kernel.imports_for_benchmark_kernel()}\n"
302	f"{src_code}\n"	302	f"{src_code}\n"
303	- f"{kernel.codegen_kernel_benchmark(num_gb, grid).getvalue()}"	303	+ f"{kernel.codegen_kernel_benchmark(num_gb).getvalue()}"
304	)	304	)
305		305
306	if only_gen_src_code:	306	if only_gen_src_code:

Mtorch_npu/_inductor/kernel/bmm.py+227-8

	@@ -1,11 +1,14 @@
1	-import logging	1	+import logging
		2	+from typing import Any, Dict, List
2		3
3	import torch	4	import torch
4	from torch._inductor.codegen.rocm.ck_universal_gemm_template import CKGemmTemplate	5	from torch._inductor.codegen.rocm.ck_universal_gemm_template import CKGemmTemplate
5		6
6	from torch._inductor import ir, lowering as L	7	from torch._inductor import ir, lowering as L
7	-from torch._inductor.select_algorithm import autotune_select_algorithm	8	+from torch._inductor.select_algorithm import (
8	-	9	+ autotune_select_algorithm,
		10	+ SymbolicGridFn,
		11	+)
9	from torch._inductor.utils import (	12	from torch._inductor.utils import (
10	use_aten_gemm_kernels,	13	use_aten_gemm_kernels,
11	use_ck_template,	14	use_ck_template,
	@@ -17,16 +20,214 @@ from torch._inductor.kernel.mm_common import (
17	_is_static_problem,	20	_is_static_problem,
18	mm_args,	21	mm_args,
19	)	22	)
20	-from torch._inductor.kernel import bmm as inductor_bmm
21		23
22	from .mm import is_contiguous_striding	24	from .mm import is_contiguous_striding
23	-from ..utils import use_catlass_template	25	+from ..select_algorithm import NPUTritonTemplate
		26	+from ..utils import use_catlass_template, use_triton_template
		27	+
24		28
25	log = logging.getLogger("torch._inductor")	29	log = logging.getLogger("torch._inductor")
26	aten = torch.ops.aten	30	aten = torch.ops.aten
27		31
28	-aten_bmm = inductor_bmm.aten_bmm	32	+aten_bmm = torch._inductor.kernel.bmm.aten_bmm
29	-aten_baddbmm = inductor_bmm.aten_baddbmm	33	+aten_baddbmm = torch._inductor.kernel.bmm.aten_baddbmm
		34	+
		35	+
		36	+# ---------------------------------------------------------------------------
		37	+# NPU Triton BMM Template (for CV / epilogue fusion with batch dimension)
		38	+# ---------------------------------------------------------------------------
		39	+# Uses triton_bmm.py.jinja which extends the mm template with a batch
		40	+# dimension (idx_q = tl.program_id(1)). Grid is (MN_tiles, batch, 1).
		41	+# The {{store_output}} placeholder supports epilogue fusion (e.g. relu).
		42	+
		43	+@SymbolicGridFn
		44	+def npu_bmm_grid(b, m, n, meta, *, cdiv):
		45	+ """Grid function for NPU bmm triton template.
		46	+
		47	+ Returns (num_mn_tiles, batch, 1) where num_mn_tiles covers all M*N blocks.
		48	+ """
		49	+ return (cdiv(m, meta["BLOCK_M"]) * cdiv(n, meta["BLOCK_N"]), b, 1)
		50	+
		51	+
		52	+# Inline template source (previously loaded from templates/triton_bmm.py.jinja).
		53	+# Kept as a string constant so the kernel no longer depends on the external
		54	+# .jinja file at runtime.
		55	+_BMM_TEMPLATE = """{{def_kernel("A", "B")}}
		56	+ M = {{size("A", -2)}}
		57	+ N = {{size("B", -1)}}
		58	+ K = {{size("A", -1)}}
		59	+
		60	+ stride_aq = {{stride("A", 0)}}
		61	+ stride_am = {{stride("A", 1)}}
		62	+ stride_ak = {{stride("A", 2)}}
		63	+
		64	+ stride_bq = {{stride("B", 0)}}
		65	+ stride_bk = {{stride("B", 1)}}
		66	+ stride_bn = {{stride("B", 2)}}
		67	+
		68	+ # based on triton.ops.matmul
		69	+ pid = tl.program_id(0).to(INDEX_DTYPE)
		70	+ grid_m = (M + BLOCK_M - 1) // BLOCK_M
		71	+ grid_n = (N + BLOCK_N - 1) // BLOCK_N
		72	+
		73	+ # re-order program ID for better L2 performance
		74	+ width = GROUP_M * grid_n
		75	+ group_id = pid // width
		76	+ group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
		77	+ pid_m = group_id * GROUP_M + (pid % group_size)
		78	+ pid_n = (pid % width) // group_size
		79	+
		80	+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
		81	+ rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
		82	+
		83	+ # batch dimension index — precompute batch offsets before K-loop
		84	+ # to avoid redundant multiply inside the hot loop
		85	+ idx_q = tl.program_id(1).to(INDEX_DTYPE)
		86	+ a_batch_off = idx_q * stride_aq
		87	+ b_batch_off = idx_q * stride_bq
		88	+
		89	+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
		90	+
		91	+ for k_start in range(0, K, BLOCK_K):
		92	+ offs_k = k_start + tl.arange(0, BLOCK_K)
		93	+ {% if EVEN_K %}
		94	+ a = tl.load(A + (rm[:, None] * stride_am + offs_k[None, :] * stride_ak + a_batch_off))
		95	+ b = tl.load(B + (offs_k[:, None] * stride_bk + rn[None, :] * stride_bn + b_batch_off))
		96	+ {% else %}
		97	+ # K is not a multiple of BLOCK_K: mask out-of-bounds elements
		98	+ k_mask = offs_k < K
		99	+ a = tl.load(A + (rm[:, None] * stride_am + offs_k[None, :] * stride_ak + a_batch_off), mask=k_mask[None, :], other=0.0)
		100	+ b = tl.load(B + (offs_k[:, None] * stride_bk + rn[None, :] * stride_bn + b_batch_off), mask=k_mask[:, None], other=0.0)
		101	+ {% endif %}
		102	+ acc = tl.dot(a, b, acc=acc, allow_tf32=ALLOW_TF32, out_dtype=ACC_TYPE)
		103	+
		104	+ # rematerialize rm, rn and idx_q to save registers
		105	+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
		106	+ rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
		107	+ idx_q = tl.program_id(1).to(INDEX_DTYPE)
		108	+ idx_m = rm[:, None]
		109	+ idx_n = rn[None, :]
		110	+ mask = (idx_m < M) & (idx_n < N)
		111	+
		112	+ # inductor generates a suffix
		113	+ {{store_output(("idx_q", "idx_m", "idx_n"), "acc", "mask", val_shape=("BLOCK_M", "BLOCK_N"))}}
		114	+"""
		115	+
		116	+npu_triton_bmm_template = NPUTritonTemplate(
		117	+ name="npu_triton_bmm",
		118	+ grid=npu_bmm_grid,
		119	+ source=_BMM_TEMPLATE,
		120	+ debug=False,
		121	+)
		122	+
		123	+
		124	+def _get_npu_bmm_configs(
		125	+ m: int,
		126	+ n: int,
		127	+ k: int,
		128	+) -> List[Dict[str, Any]]:
		129	+ """Generate tiling configs for NPU triton bmm template.
		130	+
		131	+ Same tiling shapes as mm, adapted for batched matmul.
		132	+ """
		133	+ configs: List[Dict[str, Any]] = []
		134	+
		135	+ tile_shapes = [
		136	+ (64, 64, 32),
		137	+ (64, 128, 32),
		138	+ (128, 64, 32),
		139	+ (128, 128, 32),
		140	+ (64, 64, 64),
		141	+ (128, 64, 64),
		142	+ (64, 128, 64),
		143	+ (128, 128, 64), # large tile for big BMM shapes
		144	+ # --- added by performance optimization (autotune-discovered) ---
		145	+ # These configs were found to be significantly faster on ascend950PR
		146	+ # for large BMM shapes (e.g. B=80, M=200, N=1280, K=640).
		147	+ # BLOCK_N=256 improves N-dimension tiling for wide output matrices.
		148	+ (128, 256, 64),
		149	+ (64, 256, 64),
		150	+ # BLOCK_K=128 reduces K-loop iterations for large K dimensions.
		151	+ (128, 128, 128),
		152	+ (64, 128, 128),
		153	+ # BLOCK_K=256 further reduces K-loop iterations (e.g. K=1280 → 5 iters
		154	+ # instead of 10 with BLOCK_K=128). Best config for large K on ascend950PR.
		155	+ (128, 256, 256),
		156	+ (128, 128, 256),
		157	+ # --- end of performance optimization additions ---
		158	+ (32, 64, 32),
		159	+ (64, 32, 32),
		160	+ (32, 32, 32),
		161	+ ]
		162	+
		163	+ for block_m, block_n, block_k in tile_shapes:
		164	+ # Dynamically compute EVEN_K: True only when K is an exact multiple
		165	+ # of BLOCK_K, so the template can skip the K-boundary mask
		166	+ # for performance while staying correct when K is not aligned.
		167	+ even_k = (k % block_k == 0)
		168	+ # Use GROUP_M=[1, 8] for large tiles (BLOCK_M>=128 and BLOCK_N>=128),
		169	+ # GROUP_M=[8] for small tiles. GROUP_M=1 (row-major traversal) is
		170	+ # better for shapes with small grid_m (e.g. M=200, BLOCK_M=128 → grid_m=2).
		171	+ if block_m >= 128 and block_n >= 128:
		172	+ group_m_values = [1, 8]
		173	+ else:
		174	+ group_m_values = [8]
		175	+ for group_m in group_m_values:
		176	+ for num_stages in [2, 3]:
		177	+ for num_warps in [4, 8]:
		178	+ configs.append({
		179	+ "BLOCK_M": block_m,
		180	+ "BLOCK_N": block_n,
		181	+ "BLOCK_K": block_k,
		182	+ "GROUP_M": group_m,
		183	+ "num_stages": num_stages,
		184	+ "num_warps": num_warps,
		185	+ "ALLOW_TF32": "False",
		186	+ "ACC_TYPE": "tl.float32",
		187	+ "EVEN_K": even_k,
		188	+ })
		189	+
		190	+ return configs
		191	+
		192	+
		193	+def add_npu_triton_bmm_choices(
		194	+ choices: List[ir.ChoiceCaller],
		195	+ layout: "ir.Layout",
		196	+ mat1: "ir.IRNode",
		197	+ mat2: "ir.IRNode",
		198	+ m: int,
		199	+ n: int,
		200	+ k: int,
		201	+) -> None:
		202	+ """Add NPU Triton bmm template choices to the choices list.
		203	+
		204	+ The bmm template handles the batch dimension via tl.program_id(1) and
		205	+ supports epilogue fusion via {{store_output}}.
		206	+ """
		207	+ input_nodes = [mat1, mat2]
		208	+ configs = _get_npu_bmm_configs(m, n, k)
		209	+
		210	+ for cfg in configs:
		211	+ num_stages = cfg.pop("num_stages")
		212	+ num_warps = cfg.pop("num_warps")
		213	+
		214	+ try:
		215	+ choice = npu_triton_bmm_template.generate(
		216	+ input_nodes=input_nodes,
		217	+ layout=layout,
		218	+ num_stages=num_stages,
		219	+ num_warps=num_warps,
		220	+ **cfg,
		221	+ )
		222	+ if choice is not None:
		223	+ choices.append(choice)
		224	+ except Exception:
		225	+ log.debug(
		226	+ "Failed to generate NPU triton bmm choice with config %s",
		227	+ cfg,
		228	+ exc_info=True,
		229	+ )
		230	+
30		231
31	def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:	232	def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:
32	"""	233	"""
	@@ -44,7 +245,6 @@ def is_batch_stride_largest_or_zero(mat1, mat2, layout) -> bool:
44	def _register_npu_inductor_bmm():	245	def _register_npu_inductor_bmm():
45	@L.register_lowering(aten.bmm)	246	@L.register_lowering(aten.bmm)
46	def tuned_bmm(mat1, mat2, *, layout=None):	247	def tuned_bmm(mat1, mat2, *, layout=None):
47	-
48	if all(x.get_device().type == "cpu" for x in [mat1, mat2]):	248	if all(x.get_device().type == "cpu" for x in [mat1, mat2]):
49	# decompose to small ops when memory bound	249	# decompose to small ops when memory bound
50	if mat1.get_size()[1] == 1 or mat2.get_size()[2] == 1:	250	if mat1.get_size()[1] == 1 or mat2.get_size()[2] == 1:
	@@ -119,6 +319,25 @@ def _register_npu_inductor_bmm():
119	if use_ck_template(layout):	319	if use_ck_template(layout):
120	CKGemmTemplate.add_ck_gemm_choices(choices, layout, [mat1, mat2])	320	CKGemmTemplate.add_ck_gemm_choices(choices, layout, [mat1, mat2])
121		321
		322	+ # Add NPU Triton bmm template choices for CV (Compute/Vector) fusion.
		323	+ # The bmm template handles the batch dimension via tl.program_id(1)
		324	+ # and supports epilogue fusion via {{store_output}}.
		325	+ if is_nonzero and use_triton_template(layout):
		326	+ try:
		327	+ add_npu_triton_bmm_choices(
		328	+ choices, layout, mat1, mat2, m, n, k
		329	+ )
		330	+ log.debug(
		331	+ "NPU Triton CV fusion: added triton bmm template choices "
		332	+ "for bmm(%d, %d, %d), total choices now %d",
		333	+ m,
		334	+ n,
		335	+ k,
		336	+ len(choices),
		337	+ )
		338	+ except Exception:
		339	+ log.warning("Failed to add NPU triton bmm template choices", exc_info=True)
		340	+
122	if len(choices) == 0:	341	if len(choices) == 0:
123	log.warning("No choices for GEMM, using ATen backend as fallback")	342	log.warning("No choices for GEMM, using ATen backend as fallback")
124	choices.append(aten_bmm.bind((mat1, mat2), layout))	343	choices.append(aten_bmm.bind((mat1, mat2), layout))

Mtorch_npu/_inductor/kernel/mm.py+741-4

	@@ -1,4 +1,5 @@
1	import logging	1	import logging
		2	+from typing import Any, Dict, List
2		3
3	import torch	4	import torch
4		5
	@@ -7,7 +8,7 @@ from torch._inductor.codegen.cpp_gemm_template import CppGemmTemplate
7	import torch._inductor.kernel	8	import torch._inductor.kernel
8	from torch._inductor.virtualized import V	9	from torch._inductor.virtualized import V
9		10
10	-from torch._inductor import config as inductor_config	11	+from torch._inductor import config as inductor_config, ir
11	from torch._inductor.codegen.rocm.ck_universal_gemm_template import CKGemmTemplate	12	from torch._inductor.codegen.rocm.ck_universal_gemm_template import CKGemmTemplate
12	from torch._inductor.codegen.wrapper import PythonWrapperCodegen	13	from torch._inductor.codegen.wrapper import PythonWrapperCodegen
13	from torch._inductor.ir import FixedLayout, FlexibleLayout	14	from torch._inductor.ir import FixedLayout, FlexibleLayout
	@@ -23,11 +24,18 @@ from torch._inductor.utils import (
23	)	24	)
24	from torch._inductor.kernel.mm_common import (	25	from torch._inductor.kernel.mm_common import (
25	_is_static_problem,	26	_is_static_problem,
		27	+ addmm_epilogue,
26	mm_args,	28	mm_args,
		29	+ mm_grid,
27	)	30	)
		31	+# Note: addmm_epilogue is imported above and used to build the epilogue_fn
		32	+# that fuses `beta * bias + alpha * acc` into the template's store_output.
		33	+from torch._inductor.select_algorithm import SymbolicGridFn
28		34
29	from ..codegen.catlass.gemm_template import CATLASS1xGemmTemplate	35	from ..codegen.catlass.gemm_template import CATLASS1xGemmTemplate
30	-from ..utils import use_catlass_template	36	+from ..select_algorithm import NPUTritonTemplate
		37	+from ..utils import use_catlass_template, use_triton_template
		38	+
31		39
32	log = logging.getLogger("torch._inductor")	40	log = logging.getLogger("torch._inductor")
33	aten = torch.ops.aten	41	aten = torch.ops.aten
	@@ -37,6 +45,651 @@ aten_mm = torch._inductor.kernel.mm.aten_mm
37	aten_addmm = torch._inductor.kernel.mm.aten_addmm	45	aten_addmm = torch._inductor.kernel.mm.aten_addmm
38	mm_template = torch._inductor.kernel.mm.mm_template	46	mm_template = torch._inductor.kernel.mm.mm_template
39		47
		48	+
		49	+# ---------------------------------------------------------------------------
		50	+# NPU Triton MM Template (for CV / epilogue fusion)
		51	+# ---------------------------------------------------------------------------
		52	+# Following the community torch design:
		53	+# - The jinja template lives in templates/triton_mm.py.jinja
		54	+# - The NPUTritonTemplate instance is created here, loading that jinja file
		55	+# - Tiling configs are generated by _get_npu_mm_configs() below
		56	+# - Choices are added in tuned_mm() via add_npu_triton_mm_choices()
		57	+#
		58	+# When the scheduler detects that the matmul output is consumed by pointwise
		59	+# ops (add, relu, etc.), it fuses them into the template's {{store_output}}
		60	+# epilogue — this is the "CV fusion" mechanism.
		61	+# ---------------------------------------------------------------------------
		62	+
		63	+_MM_TEMPLATE = """{{def_kernel("A", "B")}}
		64	+ M = {{size("A", 0)}}
		65	+ N = {{size("B", 1)}}
		66	+ K = {{size("A", 1)}}
		67	+ if M * N == 0:
		68	+ # early exit due to zero-size input(s)
		69	+ return
		70	+ stride_am = {{stride("A", 0)}}
		71	+ stride_ak = {{stride("A", 1)}}
		72	+ stride_bk = {{stride("B", 0)}}
		73	+ stride_bn = {{stride("B", 1)}}
		74	+
		75	+ # based on triton.ops.matmul
		76	+ pid = tl.program_id(0).to(INDEX_DTYPE)
		77	+ grid_m = (M + BLOCK_M - 1) // BLOCK_M
		78	+ grid_n = (N + BLOCK_N - 1) // BLOCK_N
		79	+
		80	+ # re-order program ID for better L2 performance
		81	+ # Use super-grouping (GROUP_M) for improved L2 cache hit rate
		82	+ width = GROUP_M * grid_n
		83	+ group_id = pid // width
		84	+ group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
		85	+ pid_m = group_id * GROUP_M + (pid % group_size)
		86	+ pid_n = (pid % width) // group_size
		87	+
		88	+ tl.assume(pid_m >= 0)
		89	+ tl.assume(pid_n >= 0)
		90	+
		91	+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
		92	+ rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
		93	+
		94	+ offs_k = tl.arange(0, BLOCK_K)
		95	+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
		96	+
		97	+ for k_idx in range(0, tl.cdiv(K, BLOCK_K)):
		98	+ offs_a_k = k_idx * BLOCK_K + tl.arange(0, BLOCK_K)
		99	+ {% if EVEN_K %}
		100	+ a = tl.load(A + (rm[:, None] * stride_am + offs_a_k[None, :] * stride_ak))
		101	+ b = tl.load(B + (offs_a_k[:, None] * stride_bk + rn[None, :] * stride_bn))
		102	+ {% else %}
		103	+ # K is not a multiple of BLOCK_K: mask out-of-bounds elements
		104	+ k_mask = offs_a_k < K
		105	+ a = tl.load(A + (rm[:, None] * stride_am + offs_a_k[None, :] * stride_ak), mask=k_mask[None, :], other=0.0)
		106	+ b = tl.load(B + (offs_a_k[:, None] * stride_bk + rn[None, :] * stride_bn), mask=k_mask[:, None], other=0.0)
		107	+ {% endif %}
		108	+ acc += tl.dot(a, b, allow_tf32=ALLOW_TF32, out_dtype=ACC_TYPE)
		109	+
		110	+ # rematerialize rm and rn to save registers
		111	+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
		112	+ rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
		113	+ idx_m = rm[:, None]
		114	+ idx_n = rn[None, :]
		115	+ mask = (idx_m < M) & (idx_n < N)
		116	+
		117	+ # inductor generates a suffix
		118	+ {{store_output(("idx_m", "idx_n"), "acc", "mask", val_shape=("BLOCK_M", "BLOCK_N"))}}
		119	+"""
		120	+
		121	+
		122	+npu_triton_mm_template = NPUTritonTemplate(
		123	+ name="npu_triton_mm",
		124	+ grid=mm_grid,
		125	+ source=_MM_TEMPLATE,
		126	+ debug=False,
		127	+)
		128	+
		129	+
		130	+# ---------------------------------------------------------------------------
		131	+# NPU Persistent Triton MM Template (diagonal core division)
		132	+# ---------------------------------------------------------------------------
		133	+# Uses new_triton_mm.py.jinja which implements a persistent kernel strategy:
		134	+# each program handles multiple tiles via `for block_idx in range(pid, NUM_BLOCKS, num_cores)`.
		135	+# For large matrices (NUM_BLOCKS_M >= 8 and NUM_BLOCKS_N >= 8), a diagonal
		136	+# super-grouping strategy improves L2 cache hit rate.
		137	+# The {{store_output}} placeholder is inside the for loop, so indent_width=8.
		138	+
		139	+# ===================================================================
		140	+# Persistent kernel with super-grouping (diagonal tile traversal)
		141	+# ===================================================================
		142	+# Inspired by the community persistent_tma_mm template but adapted
		143	+# for Ascend NPU (no TMA — uses plain tl.load).
		144	+#
		145	+# Core idea:
		146	+# * Launch NUM_SMS programs (one per AI Core).
		147	+# * Each program serially processes multiple output tiles:
		148	+# tile_id = start_pid, start_pid + NUM_SMS, start_pid + 2*NUM_SMS, ...
		149	+# * Super-grouping (GROUP_M) reorders tiles so that programs running
		150	+# concurrently access overlapping rows of A, improving L2 cache
		151	+# hit rate.
		152	+#
		153	+# Key NPU constraint: GROUP_M is guaranteed to be a factor of
		154	+# NUM_BLOCKS_M (enforced in config generation), so every modulo /
		155	+# division below is (runtime) op (compile-time constant) — safe for
		156	+# the NPU Triton backend.
		157	+# ===================================================================
		158	+
		159	+_PERSISTENT_MM_TEMPLATE = """
		160	+{{def_kernel("A", "B")}}
		161	+ M = {{size("A", 0)}}
		162	+ N = {{size("B", 1)}}
		163	+ K = {{size("A", 1)}}
		164	+ if M * N == 0:
		165	+ # early exit due to zero-size input(s)
		166	+ return
		167	+ stride_am = {{stride("A", 0)}}
		168	+ stride_ak = {{stride("A", 1)}}
		169	+ stride_bk = {{stride("B", 0)}}
		170	+ stride_bn = {{stride("B", 1)}}
		171	+
		172	+ start_pid = tl.program_id(0).to(INDEX_DTYPE)
		173	+
		174	+ # Offsets shared across all tiles handled by this program
		175	+ rk_init = tl.arange(0, BLOCK_K)
		176	+
		177	+ # Iterate over the tiles assigned to this program.
		178	+ # NUM_TILES_PER_PROGRAM is a compile-time constant (= ceil(NUM_BLOCKS / NUM_SMS)).
		179	+ for tile_iter in range(NUM_TILES_PER_PROGRAM):
		180	+ tile_id = start_pid + tile_iter * NUM_SMS
		181	+ if tile_id < NUM_BLOCKS:
		182	+ # ---- Super-grouping tile reordering ----
		183	+ # Map linear tile_id → (pid_m, pid_n) with diagonal traversal.
		184	+ # WIDTH = GROUP_M * NUM_BLOCKS_N (compile-time constant).
		185	+ group_id = tile_id // WIDTH
		186	+ pid_m = group_id * GROUP_M + (tile_id % GROUP_M)
		187	+ pid_n = (tile_id % WIDTH) // GROUP_M
		188	+
		189	+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
		190	+ rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
		191	+
		192	+ # ---- K-loop: accumulate matmul ----
		193	+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
		194	+ for k_start in range(0, K, BLOCK_K):
		195	+ offs_k = k_start + rk_init
		196	+ {% if EVEN_K %}
		197	+ a = tl.load(A + (rm[:, None] * stride_am + offs_k[None, :] * stride_ak))
		198	+ b = tl.load(B + (offs_k[:, None] * stride_bk + rn[None, :] * stride_bn))
		199	+ {% else %}
		200	+ k_mask = offs_k < K
		201	+ a = tl.load(A + (rm[:, None] * stride_am + offs_k[None, :] * stride_ak), mask=k_mask[None, :], other=0.0)
		202	+ b = tl.load(B + (offs_k[:, None] * stride_bk + rn[None, :] * stride_bn), mask=k_mask[:, None], other=0.0)
		203	+ {% endif %}
		204	+ acc += tl.dot(a, b, allow_tf32=ALLOW_TF32, out_dtype=ACC_TYPE)
		205	+
		206	+ # ---- Store output (inductor generates epilogue suffix) ----
		207	+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
		208	+ rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
		209	+ idx_m = rm[:, None]
		210	+ idx_n = rn[None, :]
		211	+ mask = (idx_m < M) & (idx_n < N)
		212	+ {{store_output(("idx_m", "idx_n"), "acc", "mask", val_shape=("BLOCK_M", "BLOCK_N"), indent_width=12)}}
		213	+"""
		214	+
		215	+
		216	+@SymbolicGridFn
		217	+def npu_persistent_mm_grid(M: int, N: int, meta: Dict[str, Any], *, cdiv, min):
		218	+ """Grid function for persistent mm template.
		219	+
		220	+ Launches min(NUM_BLOCKS, NUM_SMS) programs. Each program iterates
		221	+ over multiple tiles via `tile_id = start_pid + tile_iter * NUM_SMS`.
		222	+ """
		223	+ num_blocks = cdiv(M, meta["BLOCK_M"]) * cdiv(N, meta["BLOCK_N"])
		224	+ return (min(num_blocks, meta["NUM_SMS"]), 1, 1)
		225	+
		226	+
		227	+npu_persistent_mm_template = NPUTritonTemplate(
		228	+ name="npu_persistent_mm",
		229	+ grid=npu_persistent_mm_grid,
		230	+ source=_PERSISTENT_MM_TEMPLATE,
		231	+ debug=False,
		232	+)
		233	+
		234	+
		235	+# ---------------------------------------------------------------------------
		236	+# Dynamic-shape helpers
		237	+# ---------------------------------------------------------------------------
		238	+# When torch.compile(dynamic=True) is used, M / N / K arrive as symbolic
		239	+# sympy expressions (e.g. ``s97``) rather than concrete ints. The config
		240	+# generation functions below need concrete values for two purposes:
		241	+#
		242	+# 1. GROUP_M selection (non-persistent template) — a hint is safe
		243	+# because GROUP_M only influences L2 cache super-grouping; the actual
		244	+# grid is recomputed at runtime inside the kernel.
		245	+#
		246	+# 2. NUM_BLOCKS / NUM_TILES_PER_PROGRAM / WIDTH (persistent template) —
		247	+# these are baked into the kernel as ``tl.constexpr`` and must
		248	+# match the real runtime shape, so a hint is not safe; the
		249	+# persistent template is skipped entirely for dynamic shapes.
		250	+# ---------------------------------------------------------------------------
		251	+
		252	+def _is_symbolic_dim(val) -> bool:
		253	+ """Return True if val is a symbolic (non-statically-known) dimension.
		254	+
		255	+ Uses the same utility that ``_is_static_problem`` relies on
		256	+ (``PythonWrapperCodegen.statically_known_int_or_none``), so plain
		257	+ ints and sympy integers are considered static, while expressions
		258	+ with free symbols (e.g. ``s97``) are considered symbolic.
		259	+ """
		260	+ return PythonWrapperCodegen.statically_known_int_or_none(val) is None
		261	+
		262	+
		263	+def _hint_int(val, fallback: int = 1) -> int:
		264	+ """Resolve a (possibly symbolic) dimension to a concrete int hint.
		265	+
		266	+ For statically-known dimensions the actual value is returned.
		267	+ For symbolic dimensions ``V.graph.sizevars.size_hint`` is used to
		268	+ obtain a concrete estimate from the tracing context's example
		269	+ inputs (the same mechanism that ``use_catlass_template`` uses for
		270	+ ``size_hint(m * n * k)``).
		271	+
		272	+ This is only safe for config selection (choosing GROUP_M, tile
		273	+ sizes, etc.) — never for values that must match the runtime shape
		274	+ exactly (such as the persistent template's NUM_BLOCKS).
		275	+ """
		276	+ static = PythonWrapperCodegen.statically_known_int_or_none(val)
		277	+ if static is not None:
		278	+ return static
		279	+ try:
		280	+ return V.graph.sizevars.size_hint(val)
		281	+ except Exception:
		282	+ return fallback
		283	+
		284	+
		285	+def _get_npu_mm_configs(
		286	+ m: int,
		287	+ n: int,
		288	+ k: int,
		289	+ *,
		290	+ max_block_dim: int = 256,
		291	+) -> List[Dict[str, Any]]:
		292	+ """Generate tiling configs for NPU triton matmul template.
		293	+
		294	+ Generates a set of (BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M, num_stages,
		295	+ num_warps) configurations suitable for NPU hardware.
		296	+
		297	+ The config space is designed around the Ascend Cube unit:
		298	+ - Tile dimensions are multiples of 16 (Cube unit granularity).
		299	+ - Larger tiles (256×128, 128×256) increase the
		300	+ compute-to-memory-access ratio, keeping the Cube unit busy.
		301	+ - Larger BLOCK_K (128) reduces loop overhead for big-K problems.
		302	+
		303	+ Args:
		304	+ max_block_dim: Cap BLOCK_M and BLOCK_N at this value. Set to 128
		305	+ for addmm (the bias epilogue fusion adds register pressure).
		306	+ """
		307	+ # Resolve m to a concrete hint for GROUP_M selection.
		308	+ # GROUP_M only affects L2 cache super-grouping behaviour, not kernel
		309	+ # correctness — the actual grid (grid_m, grid_n) is computed at runtime
		310	+ # inside the kernel from the real M, N values. Using a hint here simply
		311	+ # picks a reasonable GROUP_M for the expected shape.
		312	+ m_hint = _hint_int(m)
		313	+
		314	+ configs: List[Dict[str, Any]] = []
		315	+
		316	+ # Core tiling shapes: (BLOCK_M, BLOCK_N, BLOCK_K)
		317	+ # Organised from large (high arithmetic intensity) to small.
		318	+ tile_shapes = [
		319	+ # --- Large tiles: high compute/memory ratio for big matrices ---
		320	+ (256, 128, 64), # big-M: e.g. M=12800
		321	+ (128, 256, 64), # big-N
		322	+ (128, 128, 128), # big-K: fewer K-loop iters
		323	+ (128, 64, 128),
		324	+ (64, 128, 128),
		325	+ # --- Medium tiles: balanced ---
		326	+ (128, 128, 64),
		327	+ (128, 64, 64),
		328	+ (64, 128, 64),
		329	+ (64, 64, 64),
		330	+ # --- Small tiles: for small matrices / tail effects ---
		331	+ (64, 64, 32),
		332	+ (32, 64, 32),
		333	+ (64, 32, 32),
		334	+ # --- Small-N tiles: for skinny output matrices (e.g. N=32) ---
		335	+ # Larger BLOCK_M/BLOCK_K compensate for the tiny N dimension.
		336	+ (128, 32, 64),
		337	+ (128, 32, 128),
		338	+ (64, 32, 64),
		339	+ (64, 32, 128),
		340	+ # --- Small-M tiles: for skinny A matrices (e.g. M=32) ---
		341	+ (32, 128, 64),
		342	+ (32, 128, 128),
		343	+ (32, 64, 64),
		344	+ (32, 64, 128),
		345	+ # --- Large-K tiles: reduce K-loop iterations for big-K problems.
		346	+ # Kept small in M*N to stay within the 256 KB register file.
		347	+ (64, 64, 256),
		348	+ (32, 128, 256),
		349	+ (128, 32, 256),
		350	+ (32, 64, 256),
		351	+ ]
		352	+
		353	+ for block_m, block_n, block_k in tile_shapes:
		354	+ # Skip tiles that exceed the max block dimension (register safety).
		355	+ if block_m > max_block_dim or block_n > max_block_dim:
		356	+ continue
		357	+ # EVEN_K: keep using the original (possibly symbolic) k.
		358	+ # For dynamic k, (k % block_k == 0) evaluates to False via sympy
		359	+ # structural equality, which correctly selects the masked K-loop
		360	+ # path — safe for any runtime K value.
		361	+ even_k = (k % block_k == 0)
		362	+ # Choose GROUP_M adaptively: cap at 8, but use smaller values
		363	+ # when grid_m is small to avoid degenerate super-grouping.
		364	+ # Use m_hint (concrete) so the comparison works for dynamic shapes.
		365	+ grid_m = (m_hint + block_m - 1) // block_m
		366	+ if grid_m >= 8:
		367	+ group_m = 8
		368	+ elif grid_m >= 4:
		369	+ group_m = 4
		370	+ else:
		371	+ group_m = 1
		372	+ # Use num_stages=4 for BLOCK_K=128 to improve pipelining.
		373	+ # For BLOCK_K=256, limit to [2] to avoid register overflow.
		374	+ if block_k >= 256:
		375	+ stages_list = [2]
		376	+ elif block_k >= 128:
		377	+ stages_list = [2, 3, 4]
		378	+ else:
		379	+ stages_list = [2, 3]
		380	+ for num_stages in stages_list:
		381	+ for num_warps in [4, 8]:
		382	+ configs.append({
		383	+ "BLOCK_M": block_m,
		384	+ "BLOCK_N": block_n,
		385	+ "BLOCK_K": block_k,
		386	+ "GROUP_M": group_m,
		387	+ "num_stages": num_stages,
		388	+ "num_warps": num_warps,
		389	+ "ALLOW_TF32": "False",
		390	+ "ACC_TYPE": "tl.float32",
		391	+ "EVEN_K": even_k,
		392	+ })
		393	+
		394	+ return configs
		395	+
		396	+
		397	+def add_npu_triton_mm_choices(
		398	+ choices: List[ir.ChoiceCaller],
		399	+ layout: "ir.Layout",
		400	+ mat1: "ir.IRNode",
		401	+ mat2: "ir.IRNode",
		402	+ m: int,
		403	+ n: int,
		404	+ k: int,
		405	+) -> None:
		406	+ """Add NPU Triton matmul template choices to the choices list.
		407	+
		408	+ Generates multiple TritonTemplateCaller choices from the NPU triton mm
		409	+ template with different tiling configs, enabling the autotune mechanism
		410	+ to select the best config.
		411	+ """
		412	+ input_nodes = [mat1, mat2]
		413	+ configs = _get_npu_mm_configs(m, n, k)
		414	+
		415	+ for cfg in configs:
		416	+ # Extract kernel launch params (not passed as constexpr defines)
		417	+ num_stages = cfg.pop("num_stages")
		418	+ num_warps = cfg.pop("num_warps")
		419	+
		420	+ try:
		421	+ choice = npu_triton_mm_template.generate(
		422	+ input_nodes=input_nodes,
		423	+ layout=layout,
		424	+ num_stages=num_stages,
		425	+ num_warps=num_warps,
		426	+ **cfg,
		427	+ )
		428	+ if choice is not None:
		429	+ choices.append(choice)
		430	+ except Exception:
		431	+ log.debug(
		432	+ "Failed to generate NPU triton mm choice with config %s",
		433	+ cfg,
		434	+ exc_info=True,
		435	+ )
		436	+
		437	+
		438	+def _choose_group_m(num_blocks_m: int) -> int:
		439	+ """Choose GROUP_M as the largest factor of num_blocks_m not exceeding 8.
		440	+
		441	+ This guarantees num_blocks_m % GROUP_M == 0, so the super-grouping
		442	+ tile assignment never needs runtime % runtime (which hangs the NPU
		443	+ Triton backend). Falls back to 1 if no larger factor exists.
		444	+ """
		445	+ for g in [8, 4, 2, 1]:
		446	+ if num_blocks_m % g == 0:
		447	+ return g
		448	+ return 1
		449	+
		450	+
		451	+def _get_npu_persistent_mm_configs(
		452	+ m: int,
		453	+ n: int,
		454	+ k: int,
		455	+) -> List[Dict[str, Any]]:
		456	+ """Generate tiling configs for the NPU persistent triton matmul template.
		457	+
		458	+ Each config includes NUM_BLOCKS_M, NUM_BLOCKS_N, NUM_BLOCKS, GROUP_M,
		459	+ WIDTH, NUM_SMS and NUM_TILES_PER_PROGRAM which are required by the
		460	+ persistent + super-grouping kernel strategy.
		461	+
		462	+ Key invariants:
		463	+ * GROUP_M is always a factor of NUM_BLOCKS_M, so all modulo /
		464	+ integer-division operations in the template are
		465	+ (runtime value) % (compile-time constant) — safe for the NPU backend.
		466	+ * NUM_SMS and NUM_TILES_PER_PROGRAM are compile-time constants
		467	+ (required for the outer loop bound).
		468	+
		469	+ .. note::
		470	+ The persistent kernel bakes ``NUM_BLOCKS``, ``NUM_TILES_PER_PROGRAM``,
		471	+ ``WIDTH``, ``NUM_BLOCKS_M`` and ``NUM_BLOCKS_N`` into the kernel as
		472	+ ``tl.constexpr``. These must match the actual runtime shape:
		473	+
		474	+ * ``NUM_BLOCKS`` too small → output tiles missed → wrong results
		475	+ * ``NUM_TILES_PER_PROGRAM`` too small → tiles not covered →
		476	+ wrong results
		477	+ * ``WIDTH`` / ``GROUP_M`` derived from wrong ``NUM_BLOCKS_N`` →
		478	+ tiles mapped to wrong ``(pid_m, pid_n)`` → wrong results
		479	+
		480	+ For dynamic (symbolic) shapes we cannot guarantee these compile-time
		481	+ constants match the runtime shape, so the persistent template is
		482	+ skipped and the non-persistent template (which computes the grid
		483	+ at runtime) is used instead.
		484	+ """
		485	+ # --- Dynamic-shape guard ---
		486	+ # If any of M, N, K is symbolic, we cannot safely determine the
		487	+ # compile-time tile constants. Return an empty config list so that
		488	+ # no persistent-mm choices are added; the non-persistent template
		489	+ # (npu_triton_mm_template) will still be available.
		490	+ if _is_symbolic_dim(m) or _is_symbolic_dim(n) or _is_symbolic_dim(k):
		491	+ log.debug(
		492	+ "NPU persistent mm: skipping for dynamic shape "
		493	+ "(m=%s, n=%s, k=%s) — compile-time tile constants "
		494	+ "cannot be safely determined for symbolic dimensions; "
		495	+ "falling back to non-persistent triton mm template.",
		496	+ m, n, k,
		497	+ )
		498	+ return []
		499	+
		500	+ configs: List[Dict[str, Any]] = []
		501	+
		502	+ # Core tiling shapes — aligned with the Cube-unit-optimised set.
		503	+ tile_shapes = [
		504	+ (256, 128, 64), # big-M
		505	+ (128, 256, 64), # big-N
		506	+ (128, 128, 128), # big-K
		507	+ (128, 64, 128),
		508	+ (64, 128, 128),
		509	+ (128, 128, 64),
		510	+ (128, 64, 64),
		511	+ (64, 128, 64),
		512	+ (64, 64, 64),
		513	+ # --- Small-N tiles: for skinny output matrices (e.g. N=32) ---
		514	+ (128, 32, 64),
		515	+ (128, 32, 128),
		516	+ (64, 32, 64),
		517	+ (64, 32, 128),
		518	+ # --- Small-M tiles: for skinny A matrices (e.g. M=32) ---
		519	+ (32, 128, 64),
		520	+ (32, 128, 128),
		521	+ (32, 64, 64),
		522	+ (32, 64, 128),
		523	+ # --- Large-K tiles: reduce K-loop iterations for big-K problems.
		524	+ (64, 64, 256),
		525	+ (32, 128, 256),
		526	+ (128, 32, 256),
		527	+ (32, 64, 256),
		528	+ ]
		529	+
		530	+ # Number of AI Cores to use for persistent execution.
		531	+ # Ascend 950PR has 56 AI Cores; we try several core counts so the
		532	+ # autotuner can pick the best parallelism vs. per-program work.
		533	+ num_cores_options = [8, 16, 32, 56]
		534	+
		535	+ for block_m, block_n, block_k in tile_shapes:
		536	+ num_blocks_m = (m + block_m - 1) // block_m
		537	+ num_blocks_n = (n + block_n - 1) // block_n
		538	+ num_blocks = num_blocks_m * num_blocks_n
		539	+
		540	+ # Skip if the tile is larger than the matrix (no benefit).
		541	+ if num_blocks == 0:
		542	+ continue
		543	+
		544	+ # Choose GROUP_M as a factor of num_blocks_m (≤ 8) so that
		545	+ # group_size == GROUP_M always holds (no partial last group).
		546	+ group_m = _choose_group_m(num_blocks_m)
		547	+ width = group_m * num_blocks_n
		548	+ even_k = (k % block_k == 0)
		549	+
		550	+ for num_cores in num_cores_options:
		551	+ # Don't launch more programs than there are blocks
		552	+ actual_num_cores = min(num_cores, num_blocks)
		553	+ if actual_num_cores == 0:
		554	+ continue
		555	+ # Compute tiles per program (compile-time constant for loop bound)
		556	+ num_tiles_per_program = (num_blocks + actual_num_cores - 1) // actual_num_cores
		557	+ # Use num_stages=4 for BLOCK_K=128 to improve pipelining.
		558	+ # For BLOCK_K=256, limit to [2] to avoid register overflow.
		559	+ if block_k >= 256:
		560	+ stages_list = [2]
		561	+ elif block_k >= 128:
		562	+ stages_list = [2, 3, 4]
		563	+ else:
		564	+ stages_list = [2, 3]
		565	+ for num_stages in stages_list:
		566	+ for num_warps in [4, 8]:
		567	+ configs.append({
		568	+ "BLOCK_M": block_m,
		569	+ "BLOCK_N": block_n,
		570	+ "BLOCK_K": block_k,
		571	+ "GROUP_M": group_m,
		572	+ "NUM_BLOCKS_M": num_blocks_m,
		573	+ "NUM_BLOCKS_N": num_blocks_n,
		574	+ "NUM_BLOCKS": num_blocks,
		575	+ "WIDTH": width,
		576	+ "NUM_SMS": actual_num_cores,
		577	+ "NUM_TILES_PER_PROGRAM": num_tiles_per_program,
		578	+ "num_stages": num_stages,
		579	+ "num_warps": num_warps,
		580	+ "ALLOW_TF32": "False",
		581	+ "ACC_TYPE": "tl.float32",
		582	+ "EVEN_K": even_k,
		583	+ })
		584	+
		585	+ return configs
		586	+
		587	+
		588	+def add_npu_persistent_mm_choices(
		589	+ choices: List[ir.ChoiceCaller],
		590	+ layout: "ir.Layout",
		591	+ mat1: "ir.IRNode",
		592	+ mat2: "ir.IRNode",
		593	+ m: int,
		594	+ n: int,
		595	+ k: int,
		596	+) -> None:
		597	+ """Add NPU persistent Triton matmul template choices to the choices list.
		598	+
		599	+ Uses the diagonal core division template (new_triton_mm.py.jinja) which
		600	+ supports epilogue fusion via {{store_output}}.
		601	+ """
		602	+ input_nodes = [mat1, mat2]
		603	+ configs = _get_npu_persistent_mm_configs(m, n, k)
		604	+
		605	+ for cfg in configs:
		606	+ num_stages = cfg.pop("num_stages")
		607	+ num_warps = cfg.pop("num_warps")
		608	+
		609	+ try:
		610	+ choice = npu_persistent_mm_template.generate(
		611	+ input_nodes=input_nodes,
		612	+ layout=layout,
		613	+ num_stages=num_stages,
		614	+ num_warps=num_warps,
		615	+ **cfg,
		616	+ )
		617	+ if choice is not None:
		618	+ choices.append(choice)
		619	+ except Exception:
		620	+ log.debug(
		621	+ "Failed to generate NPU persistent mm choice with config %s",
		622	+ cfg,
		623	+ exc_info=True,
		624	+ )
		625	+
		626	+
		627	+def add_npu_triton_addmm_choices(
		628	+ choices: List[ir.ChoiceCaller],
		629	+ layout: "ir.Layout",
		630	+ inp: "ir.IRNode",
		631	+ mat1: "ir.IRNode",
		632	+ mat2: "ir.IRNode",
		633	+ m: int,
		634	+ n: int,
		635	+ k: int,
		636	+ alpha: float = 1,
		637	+ beta: float = 1,
		638	+) -> None:
		639	+ """Add NPU Triton addmm template choices to the choices list.
		640	+
		641	+ Reuses the same ``npu_triton_mm_template`` (triton_mm.py.jinja) as plain
		642	+ ``mm``. The bias term is fused into the kernel epilogue via the
		643	+ ``prefix_args`` + ``epilogue_fn`` mechanism:
		644	+
		645	+ * ``input_nodes = [inp, mat1, mat2]`` — bias (``inp``) is placed first so
		646	+ it becomes a prefix arg: it is NOT named in ``{{def_kernel("A", "B")}}``
		647	+ but is automatically loaded inside ``{{store_output}}`` and passed to
		648	+ ``epilogue_fn``.
		649	+ * ``prefix_args=1`` tells the codegen that the first input node is the
		650	+ bias/epilogue input (not a matmul operand).
		651	+ * ``epilogue_fn=addmm_epilogue(dtype, alpha, beta)`` generates
		652	+ ``beta * bias + alpha * acc`` at store time, fusing the addmm bias
		653	+ addition into the matmul kernel's epilogue.
		654	+
		655	+ This mirrors how the community ``tuned_addmm`` reuses ``mm_template`` with
		656	+ ``prefix_args=1`` and ``addmm_epilogue``.
		657	+ """
		658	+ # bias first (prefix arg), then the two matmul operands named A, B in jinja
		659	+ input_nodes = [inp, mat1, mat2]
		660	+ # Use max_block_dim=128 for addmm: the bias epilogue fusion (loading
		661	+ # and broadcasting the bias vector) adds register pressure, so 256-dim
		662	+ # tiles can cause "cc overflow" on the NPU.
		663	+ configs = _get_npu_mm_configs(m, n, k, max_block_dim=128)
		664	+
		665	+ # Build the epilogue function that fuses betabias + alphaacc.
		666	+ epilogue_fn = addmm_epilogue(layout.dtype, alpha, beta)
		atomgit-bot8月14日 🟠 High Priority 第 574 行调用 `addmm_epilogue(layout.dtype, alpha, beta)`，但本模块（mm.py）从未导入 `addmm_epilogue`。文件头（第 27-31 行）只从 `torch._inductor.kernel.mm_common` 导入了 `_is_static_problem`、`mm_args`、`mm_grid`，第 32-33 行的注释声称"addmm_epilogue is imported above"与事实不符；grep 确认 `addmm_epilogue` 在模块内无定义、无导入。\n\n变更行为链：`tuned_addmm` → `add_npu_triton_addmm_choices`（第 873 行）→ 第 574 行 `addmm_epilogue(...)` 抛 `NameError: name 'addmm_epilogue' is not defined` → 被第 893-894 行的 `except Exception` 捕获并仅打 warning → 所有 addmm 的 Triton 模板 choice 一个都加不进去（第 574 行在 for 循环之前执行）。\n\n失败模式：新增的 addmm CV/epilogue 融合功能 100% 静默失效，`choices` 只含 catlass/aten 等其它后端，功能形同虚设且无任何报错提示（只有一条 warning）。\n\n修复方向：在 `from torch._inductor.kernel.mm_common import (...)` 中加入 `addmm_epilogue`（该函数是 torch 上游定义在 mm_common 中的 epilogue 构造器）。建议：在文件头 `from torch._inductor.kernel.mm_common import (...)` 块中补充导入 `addmm_epilogue`，使第 574 行的调用有定义。不准确？ S shi-yufeng99 8月17日评论： 8月17日评论：
		667	+
		668	+ for cfg in configs:
		669	+ # Extract kernel launch params (not passed as constexpr defines)
		670	+ num_stages = cfg.pop("num_stages")
		671	+ num_warps = cfg.pop("num_warps")
		672	+
		673	+ try:
		674	+ choice = npu_triton_mm_template.generate(
		675	+ input_nodes=input_nodes,
		676	+ layout=layout,
		677	+ num_stages=num_stages,
		678	+ num_warps=num_warps,
		679	+ prefix_args=1,
		680	+ epilogue_fn=epilogue_fn,
		681	+ **cfg,
		682	+ )
		683	+ if choice is not None:
		684	+ choices.append(choice)
		685	+ except Exception:
		686	+ log.debug(
		687	+ "Failed to generate NPU triton addmm choice with config %s",
		688	+ cfg,
		689	+ exc_info=True,
		690	+ )
		691	+
		692	+
40	def is_contiguous_striding(size, stride) -> bool:	693	def is_contiguous_striding(size, stride) -> bool:
41	def is_contiguous_row_major(stride, size) -> bool:	694	def is_contiguous_row_major(stride, size) -> bool:
42	# to support non-contiguous row-major input	695	# to support non-contiguous row-major input
	@@ -97,6 +750,48 @@ def _register_npu_inductor_mm():
97	[mat1, mat2],	750	[mat1, mat2],
98	)	751	)
99		752
		753	+ # Add NPU Triton matmul template choices for CV (Compute/Vector) fusion.
		754	+ # The triton template supports epilogue fusion: when the matmul output
		755	+ # is consumed by pointwise ops (e.g. add, relu), the scheduler fuses
		756	+ # them into the kernel's epilogue via {{store_output}}.
		757	+ if is_nonzero and use_triton_template(layout):
		758	+ try:
		759	+ add_npu_triton_mm_choices(
		760	+ choices, layout, mat1, mat2, m, n, k
		761	+ )
		762	+ log.debug(
		763	+ "NPU Triton CV fusion: added triton mm template choices "
		764	+ "for mm(%d, %d, %d), total choices now %d",
		765	+ m,
		766	+ n,
		767	+ k,
		768	+ len(choices),
		769	+ )
		770	+ except Exception:
		771	+ log.warning("Failed to add NPU triton mm template choices", exc_info=True)
		772	+
		773	+ # Add NPU persistent Triton matmul template choices.
		774	+ # Uses new_triton_mm.py.jinja with a persistent kernel strategy:
		775	+ # each program serially processes multiple output tiles, improving
		776	+ # L2 cache reuse via super-grouping and reducing launch overhead.
		777	+ # Best for large matrices where num_tiles >> num_cores.
		778	+ # Also supports epilogue fusion via {{store_output}}.
		779	+ if is_nonzero and use_triton_template(layout):
		780	+ try:
		781	+ add_npu_persistent_mm_choices(
		782	+ choices, layout, mat1, mat2, m, n, k
		783	+ )
		784	+ log.debug(
		785	+ "NPU persistent mm: added persistent mm template choices "
		786	+ "for mm(%d, %d, %d), total choices now %d",
		787	+ m,
		788	+ n,
		789	+ k,
		790	+ len(choices),
		791	+ )
		792	+ except Exception:
		793	+ log.warning("Failed to add NPU persistent mm template choices", exc_info=True)
		794	+
100	input_nodes = [mat1, mat2]	795	input_nodes = [mat1, mat2]
101		796
102	if (	797	if (
	@@ -157,10 +852,21 @@ def _register_npu_inductor_addmm():
157	except NotImplementedError:	852	except NotImplementedError:
158	is_contiguous_input_tmp = False	853	is_contiguous_input_tmp = False
159		854
160	- if not (	855	+ # Determine whether the catlass template is available for this problem.
		856	+ # We no longer fall back immediately when catlass is unavailable: the
		857	+ # NPU Triton addmm template (reusing npu_triton_mm_template with
		858	+ # prefix_args=1 + addmm_epilogue) can still be tried below.
		859	+ catlass_available = (
161	is_contiguous_input_tmp	860	is_contiguous_input_tmp
162	and use_catlass_template("addmm", layout_tmp if layout is None else layout, m0, n0, k0)	861	and use_catlass_template("addmm", layout_tmp if layout is None else layout, m0, n0, k0)
163	- ):	862	+ )
		863	+ # Determine whether the NPU Triton template is available.
		864	+ triton_available = is_contiguous_input_tmp and use_triton_template(
		865	+ layout_tmp if layout is None else layout
		866	+ )
		867	+
		868	+ # If neither catlass nor triton templates are available, fall back.
		869	+ if not (catlass_available or triton_available):
164	return fallback_handler(aten.addmm.default)(inp, mat1, mat2, alpha=alpha, beta=beta)	870	return fallback_handler(aten.addmm.default)(inp, mat1, mat2, alpha=alpha, beta=beta)
165		871
166	ordered_kwargs_for_cpp_kernel = ("beta", "alpha")	872	ordered_kwargs_for_cpp_kernel = ("beta", "alpha")
	@@ -248,6 +954,37 @@ def _register_npu_inductor_addmm():
248	has_bias=True,	954	has_bias=True,
249	)	955	)
250		956
		957	+ # Add NPU Triton addmm template choices.
		958	+ # Reuses npu_triton_mm_template (triton_mm.py.jinja) with the bias
		959	+ # fused into the epilogue via prefix_args=1 + addmm_epilogue.
		960	+ # This enables the autotuner to try a Triton kernel for addmm and
		961	+ # also supports CV epilogue fusion (e.g. downstream pointwise ops
		962	+ # fused into {{store_output}}).
		963	+ if is_nonzero and use_triton_template(layout):
		964	+ try:
		965	+ add_npu_triton_addmm_choices(
		966	+ choices,
		967	+ layout,
		968	+ inp_expanded,
		969	+ mat1,
		970	+ mat2,
		971	+ m,
		972	+ n,
		973	+ k,
		974	+ alpha=alpha,
		975	+ beta=beta,
		976	+ )
		977	+ log.debug(
		978	+ "NPU Triton CV fusion: added triton addmm template choices "
		979	+ "for addmm(%d, %d, %d), total choices now %d",
		980	+ m,
		981	+ n,
		982	+ k,
		983	+ len(choices),
		984	+ )
		985	+ except Exception:
		986	+ log.warning("Failed to add NPU triton addmm template choices", exc_info=True)
		987	+
251	add_aten_fallback = False	988	add_aten_fallback = False
252	if len(choices) == 0:	989	if len(choices) == 0:
253	log.warning("No choices for GEMM, using ATen backend as fallback")	990	log.warning("No choices for GEMM, using ATen backend as fallback")

Mtorch_npu/_inductor/runtime/triton_heuristics.py+30-1

	@@ -60,7 +60,6 @@ from .symbolic_grouping import (
60	from torch._inductor.runtime.triton_heuristics import ( # noqa: F401	60	from torch._inductor.runtime.triton_heuristics import ( # noqa: F401
61	fixed_config,	61	fixed_config,
62	user_autotune,	62	user_autotune,
63	- template,
64	FixedGrid,	63	FixedGrid,
65	SequentialComboKernelGrid,	64	SequentialComboKernelGrid,
66	PrecomputedGrid,	65	PrecomputedGrid,
	@@ -3285,6 +3284,36 @@ def persistent_reduction(
3285	heuristic_type=HeuristicType.PERSISTENT_REDUCTION,	3284	heuristic_type=HeuristicType.PERSISTENT_REDUCTION,
3286	)	3285	)
3287		3286
		3287	+
		3288	+def template(
		3289	+ num_stages,
		3290	+ num_warps,
		3291	+ triton_meta,
		3292	+ num_consumer_groups=0,
		3293	+ num_buffers_warp_spec=0,
		3294	+ filename=None,
		3295	+ inductor_meta=None,
		3296	+ **kwargs,
		3297	+):
		3298	+ """
		3299	+ Compile a triton template
		3300	+ """
		3301	+ # Prepare the base configuration
		3302	+ config_args = {
		3303	+ "num_stages": num_stages,
		3304	+ "num_warps": num_warps,
		3305	+ }
		3306	+
		3307	+ return cached_autotune(
		3308	+ None,
		3309	+ [triton.Config({}, **config_args)],
		3310	+ triton_meta=triton_meta,
		3311	+ inductor_meta=inductor_meta,
		3312	+ heuristic_type=HeuristicType.TEMPLATE,
		3313	+ filename=filename,
		3314	+ )
		3315	+
		3316	+
3288	def foreach(size_hints, triton_meta, num_warps, filename=None, inductor_meta=None):	3317	def foreach(size_hints, triton_meta, num_warps, filename=None, inductor_meta=None):
3289	"""	3318	"""
3290	Compile a triton foreach kernel	3319	Compile a triton foreach kernel

Mtorch_npu/_inductor/scheduler.py+1-0

	@@ -298,6 +298,7 @@ def patch_scheduler():
298		298
299	if not async_compile.use_process_pool():	299	if not async_compile.use_process_pool():
300	fut = None	300	fut = None
		301	+ mod = PyCodeCache.load(src_code_or_mod)
301	else:	302	else:
302	mod = PyCodeCache.load(src_code_or_mod)	303	mod = PyCodeCache.load(src_code_or_mod)
303	fut = async_compile.triton(	304	fut = async_compile.triton(