/assign @doubinbin


感谢提供详细的设计方案和关联 PR。这项需求主要涉及 arch22 的 fast 模板扩展,以及主算子与 metadata 算子之间的职责调整,适合按功能扩展和兼容性优化来评估。
需求分析
按 Issue 描述,将 fast 模板判定放到主算子的 host tiling,并由 kernel 读取 cu_seqlens_q 计算分核,可以减少主算子对 metadata fast 槽位的依赖,也有利于降低 fast 场景的调用约束。相应代价是主算子需要承担分核计算,并保证 host 预留的 blockDim、workspace 能覆盖 kernel 实际使用范围;性能收益仍需结合目标场景确认。
处理方向
所提方案可作为评估方向:在 arch22 上针对 query 为 int8、gSize 为 24/32、layout_q=TND、layout_k=TND/PA_BBND,且满足 sparse_mode=3、cmp_ratio=1、return_value=false 的组合启用 fast 模板,仅在该路径允许 metadata 为空;slow 路径继续保留 metadata 必传的约束。
其中,layout_k=TND 的放开需要与 key 维度解析保持一致,避免仅放宽参数检查、后续仍按原布局取维度。保持 TilingKey 不变有助于减少模板选择规则的变动,但还需确认 TilingData 字段及 metadata fast 槽位调整对现有调用的兼容性,不能仅凭 TilingKey 不变就认定二进制兼容。
Issue 已关联 PR !15122,但目前材料不足以确认其实际实现范围和合入状态。文中提到的 910B3 精度通过属于分支提交记录中的验证信息,尚不能据此确认 Ascend910_93 的适用情况或性能收益。我们会继续评估该方案及其兼容性影响。


已合入 PR #15377,为 quant_lightning_indexer_v2 的 arch22 fast 模板增加了适配:
- 增加 fast 模板参数组合校验,支持 gSize 为 24 或 32;要求 query 为 INT8、布局为 TND,key 布局为 TND 或 PA_BBND,并满足相应稀疏模式等条件。
- 增加 fast 路径的分核、workspace 计算及相关 tiling 数据,并引入 arch22 fast kernel 头文件。
- 调整 arch22 的 key 布局和连续性检查,同时保留 gSize 为 64 的 slow 路径约束。
关联 PR:PR #15377


Backgroud(背景信息)
本需求将源算子的 arch22 fast 模板迁移到主仓,并把模板判定收敛到主算子 host tiling,覆盖 decode 等小 gSize 场景;A3/arch22 fast 模板的分核数据迁入 metadata 尾部统一扩展区,与 PR #14840 引入的 metadata 动态布局基线衔接。
Benefit / Necessity (价值/作用)
num_heads_q=24/32、layout_q=TND、layout_k=TND/PA_BBND等 decode 场景;Design(设计方案)
fastUsedCoreNum选择 fast/slow,TilingKey 保持不变;IsFastTemplateEligible/CheckTemplateCombinationSupport,fast 模板要求:arch22、query dtype int8、gSize 24/32、layout_q=TND、layout_k=TND/PA_BBND、sparse_mode=3、cmp_ratio=1、return_value=false;cu_seqlens_q计算,host 仅按上界预留 blockDim 与 workspace;layout_k=TND(仅 fast);PA_BBND key 0 轴非连续校验补齐 A3(DAV_2201);相关实现 PR:https://gitcode.com/cann/ops-transformer/merge_requests/15377