已关闭
[Requirement|需求建议]: torch_extension新增 block_sparse_attention适配接口 #4397
liwenguihw创建于  8月18日关闭于  8月27日
liwenguihw
liwenguihw成员
8月18日 创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud(背景信息)

为支持BlockSparseAttention算子Mxfp4、MXFP8等量化功能,因此需要新增加aclnnBlockSparseAttentionV3接口,新增参数pQuantScaleOptional、quantMode和dstTypeMax,需要针对该接口适配torch接口block_sparse_attention。

Benefit / Necessity (价值/作用)

价值:通过torch接口方便用户调用aclnnBlockSparseAttentionV3, 使能BlockSparseAttention的量化功能。
应用场景:多模态场景。

Design(设计方案)

新增aclnn接口名:aclnnBlockSparseAttentionV3
block_sparse_attention算子def文件新增pQuantScaleOptional、quantMode和dstTypeMax参数;
torch_extension: block_sparse_attention

实现方案:
新增aclnn接口aclnnBlockSparseAttentionV3,较v2版本上新增pQuantScaleOptional、quantMode和dstTypeMax参数;其中pQuantScaleOptional是p的量化压缩因子,为FP8,HiFloat8预留参数,使能用户自定义静态量化scale;quantMode=0表示非量化,1表示FP8静态量化,2表示MXFP4 OCP量化,3表示MXFP4 CX量化(后续新增量化可根据quantMode进行区分);dstTypeMax目前是在MXFP4量化的CX量化时,让用户自定义量化量程,支持设置该值0.0,或者[6.0, 12.0],0.0:代表Amax(DType)为量化结果数据类型的最大值,[6.0, 12.0]:取值为6.0-12.0代表Amax(DType)为传入值。
torch 接口block_sparse_attention封装对aclnnBlockSparseAttentionV3的调用。

likedislike
liwenguihwliwenguihw成员
8月18日 添加了label:requirement
liwenguihwliwenguihw成员
8月18日 将 liwenguihw 设为负责人
CANN-robotCANN-robot成员
8月27日 关闭了 issue
CANN-robotCANN-robot成员
8月27日 添加了label:resolved
YYuZhengzhong成员
9月1日 删除了关联的pull request:BSA 算子新增torch 接口