已开启
[Requirement|需求建议]: QuantBatchMatmulV4/WeightQuantBatchMatmulV2新增inferDtype(rt2.0迁移配套) #6096
zhangquanxin创建于  2 天前
zhangquanxin
zhangquanxin成员
2 天前 创建

Background(背景信息)

rt2.0 架构下,QuantBatchMatmulV4、WeightQuantBatchMatmulV2、TransData 三个算子的 infershape/inferDtype 正在从 canndev 仓(rt1.0)迁移至开源仓(rt2.0)。其中两算子的 infershape 已在 op_host 目录存在,但 inferDtype 缺失

  • QuantBatchMatmulV4:输出 y dtype 由 REQUIRED_ATTR(dtype, Int) 声明(支持 DT_FLOAT16/DT_BF16),与输入 dtype(FP8/INT8 等)不同,图模式下若不注册 InferDataType,默认推导会得到错误的输出 dtype。canndev rt1.0 侧已有 InferDataTypeQuantBatchMatmulV4(matmul_infer_fns.cc),迁移后开源仓需具备对等能力;
  • WeightQuantBatchMatmulV2:输出 dtype 由 dtype attr 与 quant_scale 可选输入(C8 场景)共同决定(INT8 或同 x dtype)。canndev rt1.0 侧已有 InferDataTypeForWeightQuantBatchMatmulV2,迁移后开源仓需具备对等能力。

Benefit / Necessity(价值/作用)

  1. Ascend950 场景下 canndev 仓 rt1.0 注册将被 TILING_REG_FLAG 宏隔离,开源仓 rt2.0 的 infershape/inferDtype 成为唯一来源,图模式 dtype 推导能力必须完整;
  2. 图模式(GE 图编译)下输出 dtype 可正确推导(输出 desc 未显式设置时由 InferDataType 补全),保障迁移后行为与 rt1.0 对等;
  3. 满足开源仓目录规范:inferDtype 放 op_graph 目录,与 infershape(op_host)解耦。

Design(设计方案)

按开源仓 op_graph 目录规范新增 <op>_graph_infer.cpp,注册方式 IMPL_OP_INFERSHAPE(X).InferDataType(...),推导逻辑与 canndev 原版保持一致(仅适配头文件引用),并保留进入/退出日志便于定位:

  1. matmul/quant_batch_matmul_v4/op_graph/quant_batch_matmul_v4_graph_infer.cppy dtype = static_cast<ge::DataType>(dtype attr[0])
  2. matmul/weight_quant_batch_matmul_v2/op_graph/weight_quant_batch_matmul_v2_graph_infer.cpp:迁移 canndev InferDataTypeForWeightQuantBatchMatmulV2 全部逻辑(attr 索引 3、quant_scale 输入索引 4 的 C8 判断、x dtype 兜底)。

cmake 自动收集 op_graph/*.cpp,无需修改构建脚本。

likedislike
zhangquanxin
zhangquanxin成员
2 天前 评论:

/assign @zhangquanxin

likedislike
CANN-robotCANN-robot成员
2 天前 将 zhangquanxin 设为负责人