custom-npu_hc_pre_sinkhorn【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR/Ascend 950DT√功能说明hc_pre_sinkhorn 负责 hc_pre 的 sinkhorn 部分的计算处理详细计算过程参考test_npu_hc_pre_sinkhorn.py函数原型custom.npu_hc_pre_sinkhorn(Tensor mixes, Tensor rsqrt, Tensor hc_scale, Tensor hc_base, Tensor x, int hc_mult4, int hc_sinkhorn_iters20, float hc_eps1e-5) - (Tensor, Tensor, Tensor)参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、hchead count表示注意力头数、dhead dimension表示注意力头的维度数、T表示bs合轴后的大小。mixesTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持floatshape为[T, hc_mix]或[b, s, hc_mix]。rsqrtTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持floatshape为[T, 1]或[b, s, 1]。hc_scaleTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持floatshape为[3]。hc_baseTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持floatshape为[hc_mix]。xTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持bfloat16shape为[T, hc_mult, d]或[b, s, hc_mult, d]。hc_multint固定为4。hc_sinkhorn_itersint, 可选取值固定为20。hc_epsfloat, 可选计算过程中的$\epsilon$参数Host侧参数。仅支持double类型默认值为1e-05。返回值说明yTensor输出tensor。数据格式支持ND数据类型支持bfloat16shape为[T, d]或[b, s, d]。postTensor输出tensor。数据格式支持ND数据类型支持floatshape为[T, hc_mult]或[b, s, hc_mult]。comb_fragTensor输出tensor。数据格式支持ND数据类型支持floatshape为[T, hc_mult, hc_mult]或[b, s, hc_mult, hc_mult]。约束说明shape 字段取值范围约束 | 字段名 | 取值规则与说明 | |--------------|-------------------------------------------| | hc_mult | 取值固定为: 4 | | d | 取值固定为4096 | | hc_mix | 取值固定为: 24 |该接口支持推理场景下使用。该接口支持aclgraph入图。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。调用示例详见 test_npu_hc_pre_sinkhorn.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考