Closed
[Requirement|需求建议]: quant_lightning_indexer_v2 支持 arch22 fast 模板 #7022
doubinbinAdd on  13 days agoClosed at  2 days ago
doubinbin
doubinbin
13 days ago Create

Backgroud(背景信息)

本需求将源算子的 arch22 fast 模板迁移到主仓,并把模板判定收敛到主算子 host tiling,覆盖 decode 等小 gSize 场景;A3/arch22 fast 模板的分核数据迁入 metadata 尾部统一扩展区,与 PR #14840 引入的 metadata 动态布局基线衔接。

Benefit / Necessity (价值/作用)

  • arch22 上支持 gSize=24/32 的 fast 模板,覆盖 num_heads_q=24/32、layout_q=TND、layout_k=TND/PA_BBND 等 decode 场景;
  • fast 判定移入主算子 host tiling;slow 分核回归标准记录区,fast 分核数据独立存放于统一扩展区,消除 fast 槽位复用歧义;
  • 与上游 metadata 动态布局契约(8 int32 头部 + 动态核记录 + 尾部 256 int32 扩展区)保持一致;
  • TilingKey 保持不变,不引入额外的模板/二进制兼容负担。

Design(设计方案)

  • 新增 arch22 fast 模板(kernel / cube / vector),运行时由 TilingData 中的 fastUsedCoreNum 选择 fast/slow,TilingKey 保持不变;
  • host tiling 新增 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;
  • fast 分核由 kernel 读取 cu_seqlens_q 计算,host 仅按上界预留 blockDim 与 workspace;
  • metadata AICPU 在 fast 模板时按 3 槽连续布局(enable/tile 起点/tile 终点,≤24 group,3×24=72 int32 ≤ 256)写入尾部统一扩展区,标准记录区仅承载 slow 分核;fast 场景 metadata 允许为空;
  • fast kernel 从 metadata 头部读实际核数计算扩展区基址,按 3 槽步长读取,并做越界保护;
  • arch22 放开 layout_k=TND(仅 fast);PA_BBND key 0 轴非连续校验补齐 A3(DAV_2201);
  • 同步更新算子中英文文档与 torch 扩展;
  • 验证:910B3 上 fast gSize=32/24、slow gSize=64 精度通过(来自分支提交记录)。

相关实现 PR:https://gitcode.com/cann/ops-transformer/merge_requests/15377

likedislike
doubinbindoubinbin
13 days ago Added label: requirement
周云飞
周云飞Collaborators
13 days ago Comments:

/assign @doubinbin

likedislike
CANN-robotCANN-robotCollaborators
13 days ago Assign doubinbin as responsible person
周云飞
周云飞Collaborators
13 days ago Comments:

感谢提供详细的设计方案和关联 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 的适用情况或性能收益。我们会继续评估该方案及其兼容性影响。

likedislike
doubinbindoubinbin
13 days ago Associated pull request
doubinbindoubinbin
13 days ago Modify description of issue
doubinbindoubinbin
13 days ago Modify description of issue
doubinbindoubinbin
11 days ago Change title to "[Requirement|需求建议]: quant_lightning_indexer_v2 支持 arch22 fast 模板", originally titled "[Requirement|需求建议]: quant_lightning_indexer_v2 支持 arch22 fast 模板并解除 metadata 分核依赖".
doubinbindoubinbin
10 days ago Associated pull request
doubinbindoubinbin
10 days ago Modify description of issue
周云飞周云飞Collaborators
2 days ago changed custom state from 进行中 to 已完成
周云飞周云飞Collaborators
2 days ago Close issue
周云飞
周云飞Collaborators
2 days ago Comments:

已合入 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

likedislike