custom-npu_hc_post【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 推理系列产品√Ascend 950PR/Ascend 950DT√功能说明hc_post 是 mHC 结构的后处理部分计算逻辑计算过程如下$$ out post.unsqueeze(-1) * x.unsqueeze(-2) torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim2) $$函数原型custom.npu_hc_post(Tensor x, Tensor residual, Tensor post, Tensor comb) - Tensor参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、hchead count表示注意力头数、dhead dimension表示注意力头的维度数。xTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持bfloat16float16float32shape为[b, s, d]或[b * s, d]。residualTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持bfloat16float16float32shape为[b, s, hc, d]或[b * s, hc, d]。postTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持bfloat16float16float32shape为[b, s, hc]或[b * s, hc]。combTensor必选参数输入tensor。不支持非连续数据格式支持ND数据类型支持bfloat16float16float32shape为[b, s, hc, hc]或[b * s, hc, hc]。返回值说明yTensor输出tensor。不支持非连续数据格式支持ND数据类型支持bfloat16float16float32shape与residual的shape一致为[b, s, hc, d]或[b * s, hc, d]。约束说明shape 字段取值范围约束 | 字段名 | 取值规则与说明 | |--------------|-------------------------------------------| | hc | 取值固定为: 4 | | d | 取值固定为4096 |x、residual和y的dtype要保持一致。post和comb的dtype要保持一致。该接口支持推理场景下使用。该接口支持aclgraph入图。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。调用示例详见 test_npu_hc_post.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考