变更摘要

该 PR 将 flex attention 的溢出修复移植到最新 master：针对反向传播模板中索引计算可能发生的整数溢出，引入 INDEX_DTYPE（tl.int32/tl.int64）统一索引类型；同时对 dK/dV tasklist reduce 配置增加 UB 预算过滤，并将稀疏块与 tile 块的匹配从「相等」放宽为「整除」以支持子块（sub-block）遍历，从而修复大 shape 或非对齐块场景下的溢出/错误问题。

主要改动

• 索引溢出修复（INDEX_DTYPE）: 在 flexattention_template.py 中将 tile_id、task_id 及各类偏移计算（原 .to(tl.int64)）统一转换为 INDEX_DTYPE，并在 select_algorithm.py 的 NPUTritonTemplate 中新增 _write_index_dtype_define，允许 flex_attention 模板根据 can_use_32bit_indexing 选择 tl.int32 或 tl.int64，避免大 shape 下索引溢出。
• dK/dV tasklist reduce 的 UB 预算过滤: 在 flex_attention.py 新增 _filter_dkdv_tasklist_reduce_configs，按 npu_config.ub_size 的固定比例计算 UB 预算并过滤 BLOCK_N1/head dim 过大的配置；若无配置满足约束，则将 tasklist_reduce_ub_safe 置为 false 并禁用该图的 tasklist 代码生成。
• 子块（sub-block）支持: flex_attention_config_generator.py 将块匹配条件由 block_m == sparse_q_block_size 改为 sparse_q_block_size % block_m == 0，模板中对应将静态断言改为整除判断，并通过 SPARSE_Q_MULTIPLE/SPARSE_KV_MULTIPLE 及 kv_work_pos/kv_sub_block 循环遍历稀疏 KV 块内的子块。
• dK/dV-only 任务数修正: flex_attention_backward_dkdv_only_source 中 NUM_TASKS 由 NUM_KV_BLOCKS * ZKV * HKV 修正为 NUM_KV_BLOCKS * ZQ * HKV，同时相关地址偏移（q_adj1、do_adj1、off_chz1 等）改用 INDEX_DTYPE。

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=45274&platform=gitcode&codeHostingPlatformFlag=gitcode

代码审查

✅ 未发现问题

/approve