QuantLightningIndexerMetadata【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 推理系列产品√功能说明API功能QuantLightningIndexerMetadata是QuantLightningIndexer的前置算子通过AICPU为QuantLightningIndexer算子生成分核结果包括每个核需要处理的数据的起始点、结束点等内容随后QuantLightningIndexer根据该分核结果进行实际计算。主要计算过程为获取每个batch的基本块大小并计算负载。计算所有batch的总负载和总的基本块个数。为每个核分配负载并记录分核结果分核结果包括每个核需要处理的数据的起始点、结束点等内容。函数原型custom.npu_quant_lightning_indexer_metadata(num_heads_q, num_heads_k, head_dim, query_quant_mode, key_quant_mode, *, actual_seq_lengths_queryNone, actual_seq_lengths_keyNone, batch_size0, max_seqlen_q0, max_seqlen_k0, layout_queryBSND, layout_keyBSND, sparse_count2048, sparse_mode3, pre_tokens9223372036854775807, next_tokens9223372036854775807, cmp_ratio1, devicenpu:0) - TensorTransformer QuantLightningIndexerMetadata算子实现参考: QuantLightningIndexerMetadata参数说明num_heads_qint必选参数表示公式中的$Q$的多头数目前仅支持64。num_heads_kint必选参数表示公式中的$\tilde{K}$的多头数目前仅支持1。head_dimint必选参数表示注意力头的维度。query_quant_modeTensor必选参数用于标识query的量化模式当前支持Per-Token-Head量化模式当前仅支持传入0。key_quant_modeTensor必选参数用于标识输入key的量化模式当前支持Per-Token-Head量化模式当前仅支持传入0。*代表其之前的参数是位置相关的必须按照顺序输入之后的参数是可选参数位置无关不赋值会使用默认值。actual_seq_lengths_queryTensor可选参数表示不同Batch中query的有效token数数据类型支持int32。如果不指定seqlen可传入None表示和query的shape的S长度相同。该入参中每个Batch的有效token数不超过query中的维度S大小且不小于0。支持长度为B的一维tensor。当layout_query为TND时该入参必须传入且以该入参元素的数量作为B值该入参中每个元素的值表示当前batch与之前所有batch的token数总和即前缀和因此后一个元素的值必须大于等于前一个元素的值。不能出现负值。actual_seq_lengths_keyTensor可选参数表示不同Batch中压缩前原始key的有效token数数据类型支持int32。如果不指定seqlen可传入None表示和key的shape的S长度相同。该参数中每个Batch的原始有效token数除以压缩率后不超过key中的维度S大小且不小于0支持长度为B的一维tensor。当layout_key为TND或PA_BSND时该入参必须传入layout_key为TND该参数中每个元素的值表示当前batch与之前所有batch的token数总和即前缀和因此后一个元素的值必须大于等于前一个元素的值。batch_sizeint可选参数表示输入样本批量大小默认值为None。max_seqlen_qint可选参数表示每个Batch中的q的有效token数。max_seqlen_kint可选参数表示每个Batch中的k的有效token数。cmp_ratioint可选参数用于稀疏计算表示key的压缩倍数。数据类型支持int32。支持1/2/4/8/16/32/64/128默认值为1。layout_querystr可选参数用于标识输入query的数据排布格式当前支持BSND、TND默认值BSND。layout_keystr可选参数用于标识输入key的数据排布格式当前仅支持默认值PA_BSND。sparse_countint可选参数代表topK阶段需要保留的block数量支持[1, 2048]默认值为2048数据类型支持int32。sparse_modeint可选参数表示sparse的模式支持0/3默认值为3数据类型支持int32。 sparse_mode为0时代表defaultMask模式。sparse_mode为3时代表rightDownCausal模式的mask对应以右顶点为划分的下三角场景。pre_tokensint可选参数用于稀疏计算表示attention需要和前几个Token计算关联。数据类型支持int64仅支持默认值2^63-1。next_tokensint可选参数用于稀疏计算表示attention需要和后几个Token计算关联。数据类型支持int64仅支持默认值2^63-1。devicestr可选参数用于获取设备信息默认值为None。返回值说明metadataTensor每个cube核上LightningIndexer计算任务的Batch、Head、以及 Q 和 K 的分块的索引以及每个vector核上LightningDecode的规约任务索引。约束说明该接口支持推理场景下使用。该接口支持aclgraph模式。Tensor不能全传None。Atlas A3 推理系列产品 调用示例支持单算子模式调用和aclgraph模式调用作为QuantLightningIndexer算子的前序算子调用示例见QuantLightningIndexer调用示例。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考