Pull Request has been successfully merged by ascend-robot.
（Thank you for cwh19980101's contribution.）
变更摘要
本 PR 为 FA 推理任意函数算子（fa_infer_with_arbitrary_func_forward）新增三个运行时化 kernel 组件，目标是让注意力推理支持 32/64/128/256 等任意 embed 维度：在不修改三方 catlass 库的前提下，将原本依赖编译期常量寻址与 L1/L0 tile 一致的 QK、PV、rescale 组件改写为可按运行时 blockK/blockN 驱动的独立实现。整体思路是按容量档（embed ≤128 与 256 档）放宽 L1/L0 分块约束，并在常见宽度上保留原有常量步长快路径，仅在窄宽/非对齐宽度时回退运行时步长。
主要改动
- 新增运行时 rescale 输出组件
BlockEpilogueFARescaleORuntime：在block_epilogue_fa_rescale_o_runtime.h中实现 epilogue，当运行时行宽n == D_BASE_SIZE时走编译期常量步长快路径，否则回退到新增的FlashUpdateNewRt/FlashUpdateLastNewRt/LastDivNewRt等以rowStride为运行时步长的接口，Cast与 UB→GM 拷贝元素数改按m * n计算，使非 128 宽度下寻址仍正确。 - 新增运行时 PV 组件
BlockMmadFaiPVRuntime：block_mmad_fai_pv_tla_runtime.h中将 L1/L0 校验由 N 轴相等放宽为L1_TILE_N >= L0_TILE_N；256 档（L1_TILE_N > L0_TILE_N）时nLoops恒为 1、tailN取运行时blockN，L0C按容量档L1_TILE_N预留，并将STAGES自适应降为单缓冲以与 QK 组件共享 L0C 硬件池。 - 新增运行时 QK 组件
BlockMmadFaiQKRuntime：block_mmad_fai_qk_tla_runtime.h中将 K 轴约束放宽为L1_TILE_K >= L0_TILE_K，L0A/L0B按L0_TILE_K单块驻留；新增按blockK计算的kLoops循环与initC = (kIdx == 0)累加逻辑，使 256 档下blockK=256分两轮完成S计算，并在 256 档加入首轮PipeBarrier及kIdx == 0后的MTE1_MTE2同步。 - 分页拷贝与事件同步沿用运行时切片语义：PV 侧
CopyInL1B与 QK 侧CopyInL1A/CopyInL1B均按运行时blockK/blockN和kOffset构造GetTile切片，paged cache 分支继续按blockSize逐块搬运。 - 类名与入口签名独立化：三个组件分别使用
BlockMmadFaiQKRuntime、BlockMmadFaiPVRuntime、BlockEpilogueFARescaleORuntime新类名，避免与三方库BlockMmadTla<MmadFAIQK/FAIPV<...>>特化冲突，同时保持原有事件同步与跨核同步逻辑不变。
代码审查
审阅结论
逐文件确认
| 文件 | 结论 |
|---|---|
block_epilogue_fa_rescale_o_runtime.h |
已逐条审阅（含 Fast/Rt 六条 SIMD 路径、nLoops/tailN 边界、MTE3_V/V_MTE3 Set-Wait 配对、VEC2_UB_SIZE 与 m*n 容量关系、Cast 原地收缩）。发现 1 项 P3（行步长假设）。 |
block_mmad_fai_pv_tla_runtime.h |
已逐条审阅（L1B/L0A/L0B/L0C 尺寸核算、STAGES 对事件号与 L0C 池的影响、n-loop 与 M_ALIGN、CrossCoreSetFlag 与 kernel 侧 flag 配对、ctor/dtor 事件平衡）。发现 1 项 P2 + 2 项 P3。 |
block_mmad_fai_qk_tla_runtime.h |
已逐条审阅（k-loop 分块与 initC 累加、L1A 常驻与 isFirstLoop/isLastLoop 门控、CopyInL1A/B 的 kOffset 双向一致性、MTE1/MTE2 排水、事件计数平衡）。功能逻辑未发现缺陷，仅共享 P3（事件号私有副本重复定义，已在 PV 条目中合并指出）。 |
发现统计
- P0：0
- P1：0
- P2：1（PV 256 档
nLoops=1且tileN=blockN越出 L0B tile 的登记尺寸，缺运行时上界保护） - P3：3（新增三头文件未被任何 TU 引用/未参与编译；跨核事件号私有副本与
kernel_utils.h重复维护；rescale 运行时路径把行步长假定为shape[1]）
整体风险判断
本次改动为纯新增的"运行时化"组件副本，未修改任何既有文件，对现网行为零回归风险；同时因三份头文件尚未被任何翻译单元 include，它们当前既不参与编译也不参与运行，PR 声称的 32/64/256 维支持在本提交内实际上尚未生效（这也是本次最需要向作者确认的一点）。
代码本身质量较高：事件 Set/Wait 在 ctor/operator()/dtor 三处按 STAGES 完整配平（含 PV 依赖 QK 构造期 M_MTE1 预置这一隐含耦合，在 128/256 两档下均成立），跨核 flag x 与 16+x 双路等待与 kernel 侧双路置位一致，L1/L0 尺寸经核算在 950 的池容量内。主要遗留风险集中在"运行时基准（blockN/blockK、行 stride）完全依赖调用方自觉"——一旦接线或切块实现与注释假设不符，会表现为静默精度错误而非显式报错，建议在接入 kernel 的同一 PR 中把这几处不变量变成断言。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 0 |
| 🟡 建议 | 4 |
💬 仅评论
- ⚠️ PR 代码变更未包含测试用例（变更文件路径中均不包含 test 或 ut），请确认已经备注原因！
请修正后重新提交，或联系仓库管理员。
/approve
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
合入来源
问题/功能描述
任意mask的FA算子支持32/64/128/256 embedding
母PR：https://gitcode.com/Ascend/ops-rec/pull/244
修改方案描述
本地化catlass组件，支持动态embed32/64/128/256 embed,分为128/256档
128档：处理32/64/128 embed
256档：处理256 embed， n-chunk，256档由于UB不够，使用2 chunk，每次处理128 embed，共同组成256 embed结果。
是否涉及UT/ST
开发自检