CANN闪电索引器算子文档

custom.lightning_indexer【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 推理系列产品√功能说明算子功能用于高效处理索引数据。计算公式$$ Indices(query,key,weights)Topk(broadcast_vmul(relu(query · key)), weights) $$函数原型custom.lightning_indexer(query, key, weights) - Tensor参数说明说明query、key、weights参数维度含义BBatch Size表示输入样本批量大小、SSequence Length表示输入样本序列长度、HHead Size表示hidden层的大小、NHead Num表示多头数、DHead Dim表示hidden层最小的单元尺寸且满足DH/N。queryTensor必选参数不支持非连续的Tensor数据格式支持ND数据类型支持float16。keyTensor必选参数不支持非连续的Tensor数据格式支持ND数据类型支持float16。weightsTensor必选参数不支持非连续的Tensor数据格式支持ND数据类型支持float16。返回值说明outTensor公式中的输出数据类型支持int32。数据格式支持ND。约束说明该接口支持推理场景下使用。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。参数key、value的N仅支持1。参数query中的D和key的D值相等为128。调用示例详见test_lightning_indexer.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考