flexattention: port flex attention overflow fixes to master-pytorch-AtomGit

Mtorch_npu/_inductor/kernel/flex_attention.py+61-7

	@@ -171,6 +171,45 @@ def _is_named_ir_node(value: Any) -> bool:
171	_EXPLICIT_SCORE_MOD_OPTION = "_NPU_EXPLICIT_SCORE_MOD"	171	_EXPLICIT_SCORE_MOD_OPTION = "_NPU_EXPLICIT_SCORE_MOD"
172	_STREAMING_BLOCK_MASK_TARGET_BYTES = 256 * 1024 * 1024	172	_STREAMING_BLOCK_MASK_TARGET_BYTES = 256 * 1024 * 1024
173	_STREAMING_BLOCK_MASK_BYTES_PER_ELEMENT = 8	173	_STREAMING_BLOCK_MASK_BYTES_PER_ELEMENT = 8
		174	+_TASKLIST_REDUCE_UB_BUDGET_NUMERATOR = 4
		175	+_TASKLIST_REDUCE_UB_BUDGET_DENOMINATOR = 5
		176	+_TASKLIST_REDUCE_FP32_TILE_COUNT = 3
		177	+_TASKLIST_REDUCE_FP32_BYTES = 4
		178	+
		179	+
		180	+def _filter_dkdv_tasklist_reduce_configs(
		181	+ configs: list[dict[str, Any]],
		182	+ qk_head_dim: int,
		183	+ v_head_dim: int,
		184	+) -> list[dict[str, Any]]:
		185	+ """Keep dK/dV tiles whose tasklist reduction fits in UB."""
		186	+ max_head_dim = max(int(qk_head_dim), int(v_head_dim))
		187	+ ub_budget = (
		188	+ int(npu_config.ub_size) * _TASKLIST_REDUCE_UB_BUDGET_NUMERATOR
		189	+ // _TASKLIST_REDUCE_UB_BUDGET_DENOMINATOR
		190	+ )
		191	+
		192	+ safe_configs = []
		193	+ for cfg in configs:
		194	+ reduce_ub_bytes = (
		195	+ _TASKLIST_REDUCE_FP32_TILE_COUNT
		196	+ * int(cfg["BLOCK_N1"])
		197	+ * max_head_dim
		198	+ * _TASKLIST_REDUCE_FP32_BYTES
		199	+ )
		200	+ if reduce_ub_bytes <= ub_budget:
		201	+ safe_configs.append(cfg)
		202	+
		203	+ log.info(
		204	+ "tasklist-reduce UB filter: kept=%d/%d qk_head_dim=%d "
		205	+ "v_head_dim=%d budget_bytes=%d",
		206	+ len(safe_configs),
		207	+ len(configs),
		208	+ int(qk_head_dim),
		209	+ int(v_head_dim),
		210	+ ub_budget,
		211	+ )
		212	+ return safe_configs
174		213
175		214
176	def create_zero_int_tensor_fake(x) -> torch.Tensor:	215	def create_zero_int_tensor_fake(x) -> torch.Tensor:
	@@ -652,16 +691,10 @@ def _get_flex_attention_additional_lowerings():
652	These lowerings allow supported fallback operations to be lowered as pointwise	691	These lowerings allow supported fallback operations to be lowered as pointwise
653	ops in the score_mod and mask_mod subgraphs.	692	ops in the score_mod and mask_mod subgraphs.
654	"""	693	"""
655	- from torch._inductor.lowering import make_pointwise, index_impl	694	+ from torch._inductor.lowering import make_pointwise
656		695
657	additional_lowerings = {}	696	additional_lowerings = {}
658		697
659	- def index_pointwise(x, indices):
660	- return index_impl(x, indices, check=True)
661	-
662	- additional_lowerings[aten.index] = index_pointwise
663	- additional_lowerings[aten.index.Tensor] = index_pointwise
664	-
665	bitwise_and_fn = make_pointwise(ops.bitwise_and)	698	bitwise_and_fn = make_pointwise(ops.bitwise_and)
666		699
667	def bitwise_and_tensor(a, b):	700	def bitwise_and_tensor(a, b):
	@@ -2688,6 +2721,26 @@ def _register_npu_inductor_flex_attention():
2688	mode=FlexMode.BWDDKDV,	2721	mode=FlexMode.BWDDKDV,
2689	)	2722	)
2690		2723
		2724	+ tasklist_reduce_ub_safe = True
		2725	+ if (
		2726	+ flexattention_mask_out
		2727	+ and npu_config.flex_attention.bwd_dkdv_tasklist
		2728	+ and not bwd_has_dynamic_shape
		2729	+ ):
		2730	+ tasklist_safe_dkdv_configs = _filter_dkdv_tasklist_reduce_configs(
		2731	+ bwd_dkdv_dict_configs,
		2732	+ qk_head_dim=kernel_options["QK_HEAD_DIM"],
		2733	+ v_head_dim=kernel_options["V_HEAD_DIM"],
		2734	+ )
		2735	+ if tasklist_safe_dkdv_configs:
		2736	+ bwd_dkdv_dict_configs = tasklist_safe_dkdv_configs
		2737	+ else:
		2738	+ tasklist_reduce_ub_safe = False
		2739	+ log.warning(
		2740	+ "No dK/dV config satisfies the tasklist-reduce UB "
		2741	+ "constraint; disabling tasklist codegen for this graph."
		2742	+ )
		2743	+
2691	log.debug(	2744	log.debug(
2692	"bwd dict_configs count: dq=%d dkdv=%d",	2745	"bwd dict_configs count: dq=%d dkdv=%d",
2693	len(bwd_dq_dict_configs),	2746	len(bwd_dq_dict_configs),
	@@ -2848,6 +2901,7 @@ def _register_npu_inductor_flex_attention():
2848	not flexattention_mask_out	2901	not flexattention_mask_out
2849	or not npu_config.flex_attention.bwd_dkdv_tasklist	2902	or not npu_config.flex_attention.bwd_dkdv_tasklist
2850	or bwd_has_dynamic_shape	2903	or bwd_has_dynamic_shape
		2904	+ or not tasklist_reduce_ub_safe
2851	):	2905	):
2852	return {}	2906	return {}
2853		2907

Mtorch_npu/_inductor/kernel/flex_attention_config_generator.py+8-2

	@@ -162,8 +162,8 @@ class FlexAttentionConfigGenerator:
162	and self.sparse_kv_block_size is not None	162	and self.sparse_kv_block_size is not None
163	):	163	):
164	return (	164	return (
165	- block_m == self.sparse_q_block_size	165	+ self.sparse_q_block_size % block_m == 0
166	- and block_n == self.sparse_kv_block_size	166	+ and self.sparse_kv_block_size % block_n == 0
167	)	167	)
168	return True	168	return True
169		169
	@@ -210,6 +210,12 @@ def build_sparse_mask_candidate_configs(
210	for block in (128, 64)	210	for block in (128, 64)
211	if block <= sparse_kv_block_size and sparse_kv_block_size % block == 0	211	if block <= sparse_kv_block_size and sparse_kv_block_size % block == 0
212	]	212	]
		213	+ if not mask_block_m_candidates or not mask_block_n_candidates:
		214	+ raise ValueError(
		215	+ "Q and KV block size must be divisible by BLOCK_M and BLOCK_N. We "
		216	+ f"got Q_BLOCK_SIZE={sparse_q_block_size} and "
		217	+ f"KV_BLOCK_SIZE={sparse_kv_block_size}."
		218	+ )
213	return [	219	return [
214	{	220	{
215	"MASK_BLOCK_M": block_m,	221	"MASK_BLOCK_M": block_m,

Mtorch_npu/_inductor/kernel/flexattention_template.py+54-36

	@@ -551,8 +551,9 @@ compute_flex_attention_sparse_mask_in_loop_no_load_balance = r"""
551	NUM_Q_TILES = tl.cdiv(Q_LEN, BLOCK_M)	551	NUM_Q_TILES = tl.cdiv(Q_LEN, BLOCK_M)
552		552
553	for tile_id in range(tl.program_id(0), NUM_Q_TILES * ZQ * HQ, tl.num_programs(0)):	553	for tile_id in range(tl.program_id(0), NUM_Q_TILES * ZQ * HQ, tl.num_programs(0)):
554	- q_start = tile_id % NUM_Q_TILES	554	+ tile_id_index = tile_id.to(INDEX_DTYPE)
555	- off_zh = tile_id // NUM_Q_TILES	555	+ q_start = tile_id_index % NUM_Q_TILES
		556	+ off_zh = tile_id_index // NUM_Q_TILES
556	off_zq = off_zh // HQ	557	off_zq = off_zh // HQ
557	off_hq = off_zh % HQ	558	off_hq = off_zh % HQ
558	off_zkv = off_zq % ZKV	559	off_zkv = off_zq % ZKV
	@@ -894,8 +895,10 @@ flex_attention_backward_qmajor_dq_source = r"""
894	KV_LEN = {{size("K", 2)}}	895	KV_LEN = {{size("K", 2)}}
895	MATMUL_PRECISION = Q.dtype.element_ty	896	MATMUL_PRECISION = Q.dtype.element_ty
896		897
897	- tl.static_assert(BLOCK_M2 == SPARSE_Q_BLOCK_SIZE)	898	+ tl.static_assert(SPARSE_Q_BLOCK_SIZE % BLOCK_M2 == 0)
898	- tl.static_assert(BLOCK_N2 == SPARSE_KV_BLOCK_SIZE)	899	+ tl.static_assert(SPARSE_KV_BLOCK_SIZE % BLOCK_N2 == 0)
		900	+ SPARSE_Q_MULTIPLE: tl.constexpr = SPARSE_Q_BLOCK_SIZE // BLOCK_M2
		901	+ SPARSE_KV_MULTIPLE: tl.constexpr = SPARSE_KV_BLOCK_SIZE // BLOCK_N2
899		902
900	pid = tl.program_id(0).to(tl.int32)	903	pid = tl.program_id(0).to(tl.int32)
901	num_core = tl.num_programs(0).to(tl.int32)	904	num_core = tl.num_programs(0).to(tl.int32)
	@@ -924,15 +927,17 @@ flex_attention_backward_qmajor_dq_source = r"""
924	stride_full_kv_idx_blk = {{stride("FULL_KV_IDX", 3)}}	927	stride_full_kv_idx_blk = {{stride("FULL_KV_IDX", 3)}}
925		928
926	for task_id in range(pid, DQ_NUM_TASKS, num_core):	929	for task_id in range(pid, DQ_NUM_TASKS, num_core):
927	- q_block = task_id % DQ_NUM_Q_BLOCKS	930	+ task_id_index = task_id.to(INDEX_DTYPE)
928	- off_zq = (task_id // DQ_NUM_Q_BLOCKS) // HQ	931	+ q_block = task_id_index % DQ_NUM_Q_BLOCKS
929	- off_hq = (task_id // DQ_NUM_Q_BLOCKS) % HQ	932	+ off_zq = (task_id_index // DQ_NUM_Q_BLOCKS) // HQ
		933	+ off_hq = (task_id_index // DQ_NUM_Q_BLOCKS) % HQ
930	off_hkv = off_hq // GQA_SHARED_HEADS	934	off_hkv = off_hq // GQA_SHARED_HEADS
931	off_zkv = off_zq % ZKV	935	off_zkv = off_zq % ZKV
932	sparse_idx_z = off_zq % SPARSE_Z	936	sparse_idx_z = off_zq % SPARSE_Z
933	sparse_h = off_hq % SPARSE_HQ	937	sparse_h = off_hq % SPARSE_HQ
934		938
935	q_start = q_block * BLOCK_M2	939	q_start = q_block * BLOCK_M2
		940	+ sparse_q_block = q_block // SPARSE_Q_MULTIPLE
936	offs_m = q_start + tl.arange(0, BLOCK_M2)	941	offs_m = q_start + tl.arange(0, BLOCK_M2)
937		942
938	q_base = Q + stride_qz * off_zq + stride_qh * off_hq	943	q_base = Q + stride_qz * off_zq + stride_qh * off_hq
	@@ -940,7 +945,7 @@ flex_attention_backward_qmajor_dq_source = r"""
940	v_base = V + stride_vz * off_zkv + stride_vh * off_hkv	945	v_base = V + stride_vz * off_zkv + stride_vh * off_hkv
941	do_base = DO + stride_doz * off_zq + stride_doh * off_hq	946	do_base = DO + stride_doz * off_zq + stride_doh * off_hq
942	dq_base = DQ + stride_dqz * off_zq + stride_dqh * off_hq	947	dq_base = DQ + stride_dqz * off_zq + stride_dqh * off_hq
943	- off_chz = ((off_zq * HQ + off_hq) * Q_LEN).to(tl.int64)	948	+ off_chz = ((off_zq * HQ + off_hq) * Q_LEN).to(INDEX_DTYPE)
944	lse_base = LSE + off_chz	949	lse_base = LSE + off_chz
945	delta_base = DELTA + off_chz	950	delta_base = DELTA + off_chz
946		951
	@@ -964,12 +969,12 @@ flex_attention_backward_qmajor_dq_source = r"""
964	kv_num_offset = (	969	kv_num_offset = (
965	sparse_idx_z * stride_kv_num_blks_z	970	sparse_idx_z * stride_kv_num_blks_z
966	+ sparse_h * stride_kv_num_blks_h	971	+ sparse_h * stride_kv_num_blks_h
967	- + q_block * stride_kv_num_blks_m	972	+ + sparse_q_block * stride_kv_num_blks_m
968	)	973	)
969	kv_idx_offset = (	974	kv_idx_offset = (
970	sparse_idx_z * stride_kv_idx_z	975	sparse_idx_z * stride_kv_idx_z
971	+ sparse_h * stride_kv_idx_h	976	+ sparse_h * stride_kv_idx_h
972	- + q_block * stride_kv_idx_m	977	+ + sparse_q_block * stride_kv_idx_m
973	)	978	)
974	{% if TORCHINDUCTOR_FLEXATTENTION_MASKOUT %}	979	{% if TORCHINDUCTOR_FLEXATTENTION_MASKOUT %}
975	stride_q_offsets_z = {{stride("Q_OFFSETS", 0)}}	980	stride_q_offsets_z = {{stride("Q_OFFSETS", 0)}}
	@@ -978,15 +983,21 @@ flex_attention_backward_qmajor_dq_source = r"""
978	q_offsets_idx = (	983	q_offsets_idx = (
979	sparse_idx_z * stride_q_offsets_z	984	sparse_idx_z * stride_q_offsets_z
980	+ sparse_h * stride_q_offsets_h	985	+ sparse_h * stride_q_offsets_h
981	- + q_block * stride_q_offsets_q	986	+ + sparse_q_block * stride_q_offsets_q
982	)	987	)
983	q_offset_base = tl.load(arg_Q_OFFSETS + q_offsets_idx)	988	q_offset_base = tl.load(arg_Q_OFFSETS + q_offsets_idx)
984	{% endif %}	989	{% endif %}
985		990
986	kv_num_blocks = tl.load(arg_KV_NUM_BLKS + kv_num_offset)	991	kv_num_blocks = tl.load(arg_KV_NUM_BLKS + kv_num_offset)
987	- for blk_pos in range(0, kv_num_blocks):	992	+ for kv_work_pos in range(0, kv_num_blocks * SPARSE_KV_MULTIPLE):
		993	+ blk_pos = kv_work_pos // SPARSE_KV_MULTIPLE
		994	+ kv_sub_block = kv_work_pos % SPARSE_KV_MULTIPLE
988	kv_sparse_idx = tl.load(arg_KV_IDX + kv_idx_offset + blk_pos * stride_kv_idx_blk)	995	kv_sparse_idx = tl.load(arg_KV_IDX + kv_idx_offset + blk_pos * stride_kv_idx_blk)
989	- offs_n = kv_sparse_idx * SPARSE_KV_BLOCK_SIZE + tl.arange(0, BLOCK_N2)	996	+ offs_n = (
		997	+ kv_sparse_idx * SPARSE_KV_BLOCK_SIZE
		998	+ + kv_sub_block * BLOCK_N2
		999	+ + tl.arange(0, BLOCK_N2)
		1000	+ )
990	k = tl.load(	1001	k = tl.load(
991	k_base + offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kd,	1002	k_base + offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kd,
992	mask=offs_n[:, None] < KV_LEN,	1003	mask=offs_n[:, None] < KV_LEN,
	@@ -1005,7 +1016,7 @@ flex_attention_backward_qmajor_dq_source = r"""
1005	n = get_bounded_indices(offs_n[None, :], KV_LEN if (not IS_DIVISIBLE or not SAFE_HEAD_DIM) else None)	1016	n = get_bounded_indices(offs_n[None, :], KV_LEN if (not IS_DIVISIBLE or not SAFE_HEAD_DIM) else None)
1006	{% if TORCHINDUCTOR_FLEXATTENTION_MASKOUT %}	1017	{% if TORCHINDUCTOR_FLEXATTENTION_MASKOUT %}
1007	flat_blk = q_offset_base + blk_pos	1018	flat_blk = q_offset_base + blk_pos
1008	- offs_m_local = offs_m[:, None] - q_block * SPARSE_Q_BLOCK_SIZE	1019	+ offs_m_local = offs_m[:, None] - sparse_q_block * SPARSE_Q_BLOCK_SIZE
1009	offs_n_local = offs_n[None, :] - kv_sparse_idx * SPARSE_KV_BLOCK_SIZE	1020	offs_n_local = offs_n[None, :] - kv_sparse_idx * SPARSE_KV_BLOCK_SIZE
1010	mask_offsets = offs_m_local * SPARSE_MASK_STRIDE_M + offs_n_local	1021	mask_offsets = offs_m_local * SPARSE_MASK_STRIDE_M + offs_n_local
1011	mask_mod_output = tl.load(	1022	mask_mod_output = tl.load(
	@@ -1064,17 +1075,23 @@ flex_attention_backward_qmajor_dq_source = r"""
1064	full_kv_num_offset = (	1075	full_kv_num_offset = (
1065	sparse_idx_z * stride_full_kv_num_blks_z	1076	sparse_idx_z * stride_full_kv_num_blks_z
1066	+ sparse_h * stride_full_kv_num_blks_h	1077	+ sparse_h * stride_full_kv_num_blks_h
1067	- + q_block * stride_full_kv_num_blks_m	1078	+ + sparse_q_block * stride_full_kv_num_blks_m
1068	)	1079	)
1069	full_kv_idx_offset = (	1080	full_kv_idx_offset = (
1070	sparse_idx_z * stride_full_kv_idx_z	1081	sparse_idx_z * stride_full_kv_idx_z
1071	+ sparse_h * stride_full_kv_idx_h	1082	+ sparse_h * stride_full_kv_idx_h
1072	- + q_block * stride_full_kv_idx_m	1083	+ + sparse_q_block * stride_full_kv_idx_m
1073	)	1084	)
1074	full_kv_num_blocks = tl.load(arg_FULL_KV_NUM_BLKS + full_kv_num_offset)	1085	full_kv_num_blocks = tl.load(arg_FULL_KV_NUM_BLKS + full_kv_num_offset)
1075	- for blk_pos in range(0, full_kv_num_blocks):	1086	+ for kv_work_pos in range(0, full_kv_num_blocks * SPARSE_KV_MULTIPLE):
		1087	+ blk_pos = kv_work_pos // SPARSE_KV_MULTIPLE
		1088	+ kv_sub_block = kv_work_pos % SPARSE_KV_MULTIPLE
1076	kv_sparse_idx = tl.load(arg_FULL_KV_IDX + full_kv_idx_offset + blk_pos * stride_full_kv_idx_blk)	1089	kv_sparse_idx = tl.load(arg_FULL_KV_IDX + full_kv_idx_offset + blk_pos * stride_full_kv_idx_blk)
1077	- offs_n = kv_sparse_idx * SPARSE_KV_BLOCK_SIZE + tl.arange(0, BLOCK_N2)	1090	+ offs_n = (
		1091	+ kv_sparse_idx * SPARSE_KV_BLOCK_SIZE
		1092	+ + kv_sub_block * BLOCK_N2
		1093	+ + tl.arange(0, BLOCK_N2)
		1094	+ )
1078	k = tl.load(	1095	k = tl.load(
1079	k_base + offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kd,	1096	k_base + offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kd,
1080	mask=offs_n[:, None] < KV_LEN,	1097	mask=offs_n[:, None] < KV_LEN,
	@@ -1210,7 +1227,7 @@ flex_attention_backward_dkdv_only_source = r"""
1210	MATMUL_PRECISION = Q.dtype.element_ty	1227	MATMUL_PRECISION = Q.dtype.element_ty
1211		1228
1212	NUM_KV_BLOCKS = tl.cdiv(KV_LEN, BLOCK_N1)	1229	NUM_KV_BLOCKS = tl.cdiv(KV_LEN, BLOCK_N1)
1213	- NUM_TASKS = NUM_KV_BLOCKS * ZKV * HKV	1230	+ NUM_TASKS = NUM_KV_BLOCKS * ZQ * HKV
1214		1231
1215	pid = tl.program_id(0).to(tl.int32)	1232	pid = tl.program_id(0).to(tl.int32)
1216	num_core = tl.num_programs(0).to(tl.int32)	1233	num_core = tl.num_programs(0).to(tl.int32)
	@@ -1229,9 +1246,10 @@ flex_attention_backward_dkdv_only_source = r"""
1229	stride_q_idx_n = {{stride("Q_IDX", 2)}}	1246	stride_q_idx_n = {{stride("Q_IDX", 2)}}
1230		1247
1231	for task_id in range(pid, NUM_TASKS, num_core):	1248	for task_id in range(pid, NUM_TASKS, num_core):
1232	- kv_start_block = task_id % NUM_KV_BLOCKS	1249	+ task_id_index = task_id.to(INDEX_DTYPE)
1233	- off_zq = (task_id // NUM_KV_BLOCKS) // HKV	1250	+ kv_start_block = task_id_index % NUM_KV_BLOCKS
1234	- off_hkv = (task_id // NUM_KV_BLOCKS) % HKV	1251	+ off_zq = (task_id_index // NUM_KV_BLOCKS) // HKV
		1252	+ off_hkv = (task_id_index // NUM_KV_BLOCKS) % HKV
1235	off_zkv = off_zq % ZKV	1253	off_zkv = off_zq % ZKV
1236	sparse_idx_z = off_zq % SPARSE_Z	1254	sparse_idx_z = off_zq % SPARSE_Z
1237		1255
	@@ -1239,9 +1257,9 @@ flex_attention_backward_dkdv_only_source = r"""
1239	offs_n1 = start_n1 + tl.arange(0, BLOCK_N1)	1257	offs_n1 = start_n1 + tl.arange(0, BLOCK_N1)
1240	pid_mask = kv_start_block // SPARSE_KV_MULTIPLE	1258	pid_mask = kv_start_block // SPARSE_KV_MULTIPLE
1241		1259
1242	- k_adj = (stride_kh * off_hkv + stride_kz * off_zkv).to(tl.int64)	1260	+ k_adj = (stride_kh * off_hkv + stride_kz * off_zkv).to(INDEX_DTYPE)
1243	- v_adj = (stride_vh * off_hkv + stride_vz * off_zkv).to(tl.int64)	1261	+ v_adj = (stride_vh * off_hkv + stride_vz * off_zkv).to(INDEX_DTYPE)
1244	- dv_adj = (stride_dvh * off_hkv + stride_dvz * off_zq).to(tl.int64)	1262	+ dv_adj = (stride_dvh * off_hkv + stride_dvz * off_zq).to(INDEX_DTYPE)
1245		1263
1246	K1 = K + k_adj	1264	K1 = K + k_adj
1247	V1 = V + v_adj	1265	V1 = V + v_adj
	@@ -1263,9 +1281,9 @@ flex_attention_backward_dkdv_only_source = r"""
1263	for off_g in range(0, GQA_SHARED_HEADS):	1281	for off_g in range(0, GQA_SHARED_HEADS):
1264	off_hq1 = off_hkv * GQA_SHARED_HEADS + off_g	1282	off_hq1 = off_hkv * GQA_SHARED_HEADS + off_g
1265		1283
1266	- q_adj1 = (stride_qh * off_hq1 + stride_qz * off_zq).to(tl.int64)	1284	+ q_adj1 = (stride_qh * off_hq1 + stride_qz * off_zq).to(INDEX_DTYPE)
1267	- do_adj1 = (stride_doh * off_hq1 + stride_doz * off_zq).to(tl.int64)	1285	+ do_adj1 = (stride_doh * off_hq1 + stride_doz * off_zq).to(INDEX_DTYPE)
1268	- off_chz1 = ((off_zq * HQ + off_hq1) * Q_LEN).to(tl.int64)	1286	+ off_chz1 = ((off_zq * HQ + off_hq1) * Q_LEN).to(INDEX_DTYPE)
1269		1287
1270	Q1 = Q + q_adj1	1288	Q1 = Q + q_adj1
1271	DO1 = DO + do_adj1	1289	DO1 = DO + do_adj1
	@@ -1707,7 +1725,7 @@ flex_attention_backward_dkdv_tasklist_source = (
1707	work_start = tl.load(TASK_OFFSETS + meta_id)	1725	work_start = tl.load(TASK_OFFSETS + meta_id)
1708	work_end = tl.load(TASK_OFFSETS + meta_id + 1)	1726	work_end = tl.load(TASK_OFFSETS + meta_id + 1)
1709	for work_idx in range(work_start, work_end):	1727	for work_idx in range(work_start, work_end):
1710	- off_hkv = tl.load(WORK_ITEMS + work_idx * 5 + 0).to(tl.int64)	1728	+ off_hkv = tl.load(WORK_ITEMS + work_idx * 5 + 0).to(INDEX_DTYPE)
1711	kv_start_block = tl.load(WORK_ITEMS + work_idx * 5 + 1)	1729	kv_start_block = tl.load(WORK_ITEMS + work_idx * 5 + 1)
1712	{% if not TASKLIST_NO_SPLIT %}	1730	{% if not TASKLIST_NO_SPLIT %}
1713	sub_id = tl.load(WORK_ITEMS + work_idx * 5 + 2)	1731	sub_id = tl.load(WORK_ITEMS + work_idx * 5 + 2)
	@@ -1722,9 +1740,9 @@ flex_attention_backward_dkdv_tasklist_source = (
1722	start_n1 = kv_start_block * BLOCK_N1	1740	start_n1 = kv_start_block * BLOCK_N1
1723	offs_n1 = start_n1 + tl.arange(0, BLOCK_N1)	1741	offs_n1 = start_n1 + tl.arange(0, BLOCK_N1)
1724		1742
1725	- k_adj = (stride_kh * off_hkv + stride_kz * off_zkv).to(tl.int64)	1743	+ k_adj = (stride_kh * off_hkv + stride_kz * off_zkv).to(INDEX_DTYPE)
1726	- v_adj = (stride_vh * off_hkv + stride_vz * off_zkv).to(tl.int64)	1744	+ v_adj = (stride_vh * off_hkv + stride_vz * off_zkv).to(INDEX_DTYPE)
1727	- dv_adj = (stride_dvh * off_hkv + stride_dvz * off_zq).to(tl.int64)	1745	+ dv_adj = (stride_dvh * off_hkv + stride_dvz * off_zq).to(INDEX_DTYPE)
1728	K1 = K + k_adj	1746	K1 = K + k_adj
1729	V1 = V + v_adj	1747	V1 = V + v_adj
1730	{% if TASKLIST_NO_SPLIT %}	1748	{% if TASKLIST_NO_SPLIT %}
	@@ -1751,9 +1769,9 @@ flex_attention_backward_dkdv_tasklist_source = (
1751	for off_g in range(0, GQA_SHARED_HEADS):	1769	for off_g in range(0, GQA_SHARED_HEADS):
1752	off_hq1 = off_hkv * GQA_SHARED_HEADS + off_g	1770	off_hq1 = off_hkv * GQA_SHARED_HEADS + off_g
1753	sparse_idx_hq1 = off_hq1 % SPARSE_HQ	1771	sparse_idx_hq1 = off_hq1 % SPARSE_HQ
1754	- q_adj1 = (stride_qh * off_hq1 + stride_qz * off_zq).to(tl.int64)	1772	+ q_adj1 = (stride_qh * off_hq1 + stride_qz * off_zq).to(INDEX_DTYPE)
1755	- do_adj1 = (stride_doh * off_hq1 + stride_doz * off_zq).to(tl.int64)	1773	+ do_adj1 = (stride_doh * off_hq1 + stride_doz * off_zq).to(INDEX_DTYPE)
1756	- off_chz1 = ((off_zq * HQ + off_hq1) * Q_LEN).to(tl.int64)	1774	+ off_chz1 = ((off_zq * HQ + off_hq1) * Q_LEN).to(INDEX_DTYPE)
1757	Q1 = Q + q_adj1	1775	Q1 = Q + q_adj1
1758	DO1 = DO + do_adj1	1776	DO1 = DO + do_adj1
1759	LSE1 = LSE + off_chz1	1777	LSE1 = LSE + off_chz1
	@@ -1959,7 +1977,7 @@ flex_attention_backward_dkdv_reduce_source = r"""
1959	KV_LEN = {{size("DK", 2)}}	1977	KV_LEN = {{size("DK", 2)}}
1960		1978
1961	split_base_id = tl.program_id(0).to(tl.int32)	1979	split_base_id = tl.program_id(0).to(tl.int32)
1962	- off_hkv = tl.load(SPLIT_BASES + split_base_id * 3 + 0).to(tl.int64)	1980	+ off_hkv = tl.load(SPLIT_BASES + split_base_id * 3 + 0).to(INDEX_DTYPE)
1963	kv_block = tl.load(SPLIT_BASES + split_base_id * 3 + 1)	1981	kv_block = tl.load(SPLIT_BASES + split_base_id * 3 + 1)
1964	split_count = tl.load(SPLIT_BASES + split_base_id * 3 + 2)	1982	split_count = tl.load(SPLIT_BASES + split_base_id * 3 + 2)
1965	off_z = tl.zeros_like(off_hkv)	1983	off_z = tl.zeros_like(off_hkv)

Mtorch_npu/_inductor/select_algorithm.py+29-5

	@@ -395,6 +395,24 @@ class NPUTritonTemplate(TritonTemplate):
395	self.codegen_kernel_name = codegen_kernel_name or f"triton_{name}"	395	self.codegen_kernel_name = codegen_kernel_name or f"triton_{name}"
396	self.compile_options = compile_options or NPUTemplateCompileOption()	396	self.compile_options = compile_options or NPUTemplateCompileOption()
397		397
		398	+ def _write_index_dtype_define(
		399	+ self,
		400	+ defines: StringIO,
		401	+ numel: sympy.Expr,
		402	+ buffers: Any,
		403	+ ) -> None:
		404	+ can_use_32bit_indexing = TritonScheduling.can_use_32bit_indexing(
		405	+ numel, buffers
		406	+ )
		407	+ is_flex_attention = self.name.startswith("flex_attention")
		408	+ if not can_use_32bit_indexing and not is_flex_attention:
		409	+ raise NotImplementedError(
		410	+ "64-bit indexing is not yet implemented for triton templates"
		411	+ )
		412	+ if is_flex_attention:
		413	+ index_dtype = "tl.int32" if can_use_32bit_indexing else "tl.int64"
		414	+ defines.write(f"INDEX_DTYPE : tl.constexpr = {index_dtype}\n")
		415	+
398	def make_runtime_renderer_factory(	416	def make_runtime_renderer_factory(
399	self,	417	self,
400	*,	418	*,
	@@ -421,6 +439,14 @@ class NPUTritonTemplate(TritonTemplate):
421	if self.name.startswith("flex_attention") and name == "generate_with_caching":	439	if self.name.startswith("flex_attention") and name == "generate_with_caching":
422	continue	440	continue
423	defines.write(f"{name} : tl.constexpr = {value}\n")	441	defines.write(f"{name} : tl.constexpr = {value}\n")
		442	+ fake_out = ir.Buffer(name="buf_out", layout=layout)
		443	+ numel = sympy_product(layout.size)
		444	+ if self.manual_output_buffer is None:
		445	+ buffers = itertools.chain(input_nodes, (fake_out,))
		446	+ else:
		447	+ buffers = input_nodes
		448	+ numel = sympy_product(call_sizes)
		449	+ self._write_index_dtype_define(defines, numel, buffers)
424	kernel_options = {	450	kernel_options = {
425	"defines": defines.getvalue(),	451	"defines": defines.getvalue(),
426	"num_stages": num_stages,	452	"num_stages": num_stages,
	@@ -547,12 +573,10 @@ class NPUTritonTemplate(TritonTemplate):
547	else:	573	else:
548	buffers = checked_input_nodes	574	buffers = checked_input_nodes
549	numel = sympy_product(call_sizes or layout.size)	575	numel = sympy_product(call_sizes or layout.size)
550	- if not TritonScheduling.can_use_32bit_indexing(numel, buffers):	576	+ self._write_index_dtype_define(defines, numel, buffers)
551	- raise NotImplementedError(
552	- "64-bit indexing is not yet implemented for triton templates"
553	- )
554		577
555	- defines.write("INDEX_DTYPE : tl.constexpr = tl.int32\n")	578	+ if not self.name.startswith("flex_attention"):
		579	+ defines.write("INDEX_DTYPE : tl.constexpr = tl.int32\n")
556	defines = defines.getvalue()	580	defines = defines.getvalue()
557		581
558	if call_sizes is None:	582	if call_sizes is None: