文件最后提交记录最后更新时间
12 天前
1 天前
13 天前
19 天前
19 天前
8 个月前
13 天前
README

RopeQuantKvcache

产品支持情况

产品 是否支持
Ascend 950PR/Ascend 950DT
Atlas A3 训练系列产品/Atlas A3 推理系列产品 ×
Atlas A2 训练系列产品/Atlas A2 推理系列产品
Atlas 200I/500 A2 推理产品 ×
Atlas 推理系列产品 ×
Atlas 训练系列产品 ×
Kirin X90 处理器系列产品
Kirin 9030 处理器系列产品

功能说明

  • 算子功能:对输入张量的尾轴进行切分,划分为q、k、v,对q、k进行旋转位置编码;当前Kernel写出q,并将量化后的k、v按照indice更新到kCacheRef和vCacheRef,保留的k、v输出当前未写入。

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
qkv 输入 Device侧的aclTensor,需要切分的张量。 FLOAT16 ND
cos 输入 Device侧的aclTensor,用于旋转位置编码的张量。 FLOAT16 ND
sin 输入 Device侧的aclTensor,用于旋转位置编码的张量。 FLOAT16 ND
quant_scale 输入 Device侧的aclTensor,表示量化缩放参数的张量。 FLOAT32 ND
quant_offset 输入 Device侧的aclTensor,表示量化偏移量的张量。 INT32 ND
k_cache 输入 用于原地更新的输入。 INT8 ND
v_cache 输入 用于原地更新的输入。 INT8 ND
indice 输入 用于更新量化结果的下标 INT32 ND
size_splits 属性 用于对qkv进行切分。 INT64 -
layout 属性 表示qkv的数据排布方式。 String -
kv_output 属性 控制是否输出原本的k、v。 BOOL -
q 输出 切分出的q执行旋转位置编码后的结果。 FLOAT16 ND
k 输出 切分出的k执行旋转位置编码后的结果。 FLOAT16 ND
v 输出 切分出的v。 FLOAT16 ND
k_cache 输出 切分出的k执行旋转位置编码并量化后的结果。 INT8 ND
v_cache 输出 切分出的v量化后的结果。 INT8 ND

约束说明

  • qkv为3维[B, S, H]H=(Nq+Nkv+Nkv)*Dcossin为4维[B, S, 1, D],其B、S、D必须与切分后的k一致,不支持Broadcast。
  • quant_scalequant_offset均为1维[D];kCache和vCache均为4维[B, S_cache, Nkv, D]且shape相同,S_cache必须能容纳更新位置;indice的元素数为B。
  • size_splits为3项[Nq*D, Nkv*D, Nkv*D],三项之和必须等于qkv尾轴H;layoutkv_output当前均未被实现消费。
  • 不支持动态Shape(-1)和动态Rank(-2)。

调用说明

调用方式 调用样例 说明
图模式调用 test_geir_rope_quant_kvcache 通过图模式方式调用算子。