变更摘要

本 PR 将 fa_infer_with_arbitrary_func 的 embed（head dim）支持从固定 128 扩展为 32/64/128/256 四档：host 侧 tiling 与 torch extension 校验放开维度限制，并通过 ASCENDC_TPL_SEL_PARAM 按 embed 选择不同的二进制 key，kernel 侧以 binId 模板参数分成三条独立实例化路径（128 档沿用原实现、32/64 档输出半轴运行时化、256 档走新增的 Emb256 n-chunk 实现），同时补充对应 pytest 维度用例。

主要改动

• 对外维度校验放开：TilingFunc 用 EMBED_SUPPORTED 替换原有的 embed != 128 判断，fa_infer_with_arbitrary_func_plugin.cpp 用 HEAD_DIM_SUPPORTED 与 checkHeadDim 允许 q/k/v 最后一维取 32/64/128/256，并新增 q/k/v 三者维度必须一致的检查。
• 按 embed 选择编译 key：新增 fa_infer_with_arbitrary_func_key.h，通过 ASCENDC_TPL_ARGS_DECL / ASCENDC_TPL_SEL 声明 binId（0/1/2），host 侧按 embed == 128 -> 1、embed < 128 -> 2、embed > 128 -> 0 调用 ASCENDC_TPL_SEL_PARAM。
• kernel 入口模板化分派：fa_infer_with_arbitrary_func.cpp 中 kernel 函数改为 template <int binId>，binId == 1 走原 FAInferTla，binId == 2 走 FAInferTla 的 RUNTIME_O=true 版本，其余走新增的 Emb256::FAInferTla256。
• 新增 256 档实现：新文件 fa_infer_embed256_kernel.h 定义 Emb256 命名空间下的 FAInferKernel 与 FAInferTla256，以 CHUNK_WIDTH = 128 将 embed 切成 N_CHUNK 个 chunk，bmm2 双 slot 每 slot 32KB，并用 CHUNK_BARRIER / CHUNK_ACK 跨核 barrier 保证 chunk 间流水排空。
• 输出行宽运行时化与排空事件：FAInferKernel 新增 RUNTIME_O 模板参数与 OWidth()，EpilogueRescaleO 通过 std::conditional_t 在 BlockEpilogueFARescaleORuntime 与原 BlockEpilogue 间选择；新增 FA_DRAIN_RES_EVENT = 15，在尾部循环用 mm1ResReal / notLastTwoLoop 切换排空事件，并在提前返回与结尾补充 PipeBarrier<PIPE_ALL>。
• 测试覆盖与可复现性：pytest 的 D 参数由 [128] 扩展为 [32, 64, 128, 256]，并新增 RANDOM_SEED 与 autouse fixture _reset_random_seed 固定随机种子以保证 fp16 比对结果可重现。

代码审查

审查总结

发现统计：P0 0 / P1 0 / P2 2 / P3 3

逐个文件核查结论：

整体风险判断：中等偏低。 改动主体（embed 32/64 走 runtime rescale、embed 256 走 n-chunk 双 slot + 跨 chunk 屏障、host 侧 TPL key 分派）设计自洽，host/plugin/kernel/key 四处的档位契约一致；主要实质风险集中在 emb256 新文件与主 kernel 的一处守卫不一致（越界 GM 读），以及新档位完全没有 CI 覆盖——即“新代码路径出错不会被流水线发现”。

💬 仅评论

/approve

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.