CANN KV压缩Epilog算子

custom-npu_kv_compress_epilog【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√功能说明将输入x split成nope和ropenope部分进行量化后与rope进行拼接根据slot_mapping中的下标映射-1表示无效数据不进行更新原地更新到kv_compress_cache对应的行上。函数原型custom.kv_compress_epilog(Tensor(a!) kv_compress_cache, Tensor x, Tensor slot_mapping, *, int quant_group_size 64, int quant_mode 2, bool round_scale_flag True) - ()参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、hchead count表示注意力头数、dHead dimension表示注意力头的维度数、T表示bs合轴后的大小。kv_compress_cacheTensor必选参数输入tensor待更新的cache数据。不支持非连续数据格式支持ND数据类型支持float8_e4m3fn和float8_e5m2shape为[ T, d ]d的取值见约束说明。xTensor必选参数输入tensor待量化的输入数据。不支持非连续数据格式支持ND数据类型支持bfloat16shape为[ T, d ]。slot_mappingTensor必选参数输入tensor表示下标映射其值为kv_compress_cache的下标-1代表无效数据不进行更新。不支持非连续数据格式支持ND数据类型支持int32和int64shape为[ T ]。*代表其之前的参数是位置相关的必须按照顺序输入属于必选参数其之后的参数是键值对赋值与位置无关属于可选参数不传入会使用默认值。quant_group_sizeint可选quant_group的大小取值固定为64。quant_modeint可选量化模式取值固定为2表示量化输出并拼接到kv_compress_cache中的scale为float8_e8m0类型。round_scale_flagbool, 可选表示是否计算round_scale默认为True。返回值说明kv_compress_cache 做原地更新操作。约束说明shape 字段取值范围约束 | 字段名 | 取值规则与说明 | |--------------|-------------------------------------------| | d | 计算公式如下说明 | | d | 取值固定为512 | d计算公式如下 $$ ySize d - 64 2 * 64 Ceil((d - 64) / 128) * 1 $$$$ d (128 - ySize % 128) % 128 ySize $$quant_group_size固定为64。该接口支持推理场景下使用。该接口支持aclgraph入图。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。调用示例详见 test_npu_kv_compress_epilog.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考