torch_npu接口列表

本章节包含常用自定义接口，包括创建tensor及计算类操作。

表1 torch_npu API

API名称	说明
（beta）torch_npu._npu_dropout	不使用种子（seed）进行dropout结果计数。与torch.dropout相似，优化NPU设备实现。
（beta）torch_npu.copy_memory_	从src拷贝元素到self张量，并返回self。
（beta）torch_npu.empty_with_format	返回一个填充未初始化数据的张量。
（beta）torch_npu.fast_gelu	快速高斯误差线性单元激活函数（Fast Gaussian Error Linear Units activation function）对输入的每个元素计算FastGelu的前向结果。
（beta）torch_npu.npu_alloc_float_status	为溢出检测模式申请tensor作为入参。
（beta）torch_npu.npu_anchor_response_flags	在单个特征图中生成锚点的责任标志。
（beta）torch_npu.npu_apply_adam	获取adam优化器的计算结果。
（beta）torch_npu.npu_batch_nms	根据batch分类计算输入框评分，通过评分排序，删除评分高于阈值（iou_threshold）的框，支持多批多类处理。通过NonMaxSuppression（nms）操作可有效删除冗余的输入框，提高检测精度。NonMaxSuppression：抑制不是极大值的元素，搜索局部的极大值，常用于计算机视觉任务中的检测类模型。
（beta）torch_npu.npu_bert_apply_adam	针对bert模型，获取adam优化器的计算结果。
（beta）torch_npu.npu_bmmV2	将矩阵“a”乘以矩阵“b”，生成“a*b”。
（beta）torch_npu.npu_bounding_box_decode	根据rois和deltas生成标注框。自定义Faster R-CNN算子。
（beta）torch_npu.npu_bounding_box_encode	计算标注框和ground truth真值框之间的坐标变化。自定义Faster R-CNN算子。
（beta）torch_npu.npu_broadcast	返回self张量的新视图，其单维度扩展，结果连续。张量也可以扩展更多维度，新的维度添加在最前面。
（beta）torch_npu.npu_ciou	应用基于NPU的CIoU操作。在DIoU的基础上增加了penalty term，并propose CIoU。
（beta）torch_npu.npu_clear_float_status	清除溢出检测相关标志位。
（beta）torch_npu.npu_confusion_transpose	混淆reshape和transpose运算。
（beta）torch_npu.npu_conv_transpose2d	在由多个输入平面组成的输入图像上应用一个2D转置卷积算子，有时这个过程也被称为“反卷积”。
（beta）torch_npu.npu_conv2d	在由多个输入平面组成的输入图像上应用一个2D卷积。
（beta）torch_npu.npu_conv3d	在由多个输入平面组成的输入图像上应用一个3D卷积。
（beta）torch_npu.npu_convolution	在由多个输入平面组成的输入图像上应用一个2D或3D卷积。
（beta）torch_npu.npu_convolution_transpose	在由多个输入平面组成的输入图像上应用一个2D或3D转置卷积算子，有时这个过程也被称为“反卷积”。
（beta）torch_npu.npu_deformable_conv2d	使用预期输入计算变形卷积输出（deformed convolution output）。
（beta）torch_npu.npu_diou	应用基于NPU的DIoU操作。考虑到目标之间距离，以及距离和范围的重叠率，不同目标或边界需趋于稳定。
（beta）torch_npu.npu_dtype_cast	执行张量数据类型（dtype）转换。
（beta）torch_npu.npu_format_cast	修改NPU张量的格式。
（beta）torch_npu.npu_format_cast_	原地修改self张量格式，与src格式保持一致。
（beta）torch_npu.npu_get_float_status	获取溢出检测结果。
（beta）torch_npu.npu_giou	首先计算两个框的最小封闭面积和IoU，然后计算封闭区域中不属于两个框的封闭面积的比例，最后从IoU中减去这个比例，得到GIoU。
（beta）torch_npu.npu_grid_assign_positive	执行position-sensitive的候选区域池化梯度计算。
（beta）torch_npu.npu_gru	计算DynamicGRUV2。
（beta）torch_npu.npu_indexing	以begin为起始索引，end为结束索引，strides为步长，对输入张量进行切片。
（beta）torch_npu.npu_iou	根据ground-truth和预测区域计算交并比（IoU）或前景交叉比（IoF）。
（beta）torch_npu.npu_layer_norm_eval	对层归一化结果进行计算。与torch.nn.functional.layer_norm相同，优化NPU设备实现。
（beta）torch_npu.npu_linear	将矩阵“a”乘以矩阵“b”，生成“a*b”。
（beta）torch_npu.npu_lstm	计算DynamicRNN。
（beta）torch_npu.npu_max	使用dim对最大结果进行计算。类似于torch.max，优化NPU设备实现。
（beta）torch_npu.npu_min	使用dim对最小结果进行计算。类似于torch.min，优化NPU设备实现。
（beta）torch_npu.npu_mish	按元素计算self的双曲正切。
（beta）torch_npu.npu_nms_rotated	按分数降序选择旋转标注框的子集。
（beta）torch_npu.npu_nms_v4	按分数降序选择标注框的子集。
（beta）torch_npu.npu_nms_with_mask	生成值0或1，用于nms算子确定有效位。
（beta）torch_npu.npu_one_hot	返回一个one-hot张量。input中index表示的位置采用on_value值，而其他所有位置采用off_value的值。
（beta）torch_npu.npu_pad	填充张量。
（beta）torch_npu.npu_ps_roi_pooling	执行Position Sensitive ROI Pooling。
（beta）torch_npu.npu_ptiou	根据ground-truth和预测区域计算交并比（IoU）或前景交叉比（IoF）。
（beta）torch_npu.npu_random_choice_with_mask	混洗非零元素的index。
（beta）torch_npu.npu_reshape	reshape张量。仅更改张量shape，其数据不变。
（beta）torch_npu.npu_roi_align	从特征图中获取ROI特征矩阵。自定义Faster R-CNN算子。
（beta）torch_npu.npu_rotated_iou	计算旋转框的IoU。
（beta）torch_npu.npu_rotated_overlaps	计算旋转框的重叠面积。
（beta）torch_npu.npu_sign_bits_pack	将float类型1位Adam打包为uint8。
（beta）torch_npu.npu_sign_bits_unpack	将uint8类型1位Adam拆包为float。
（beta）torch_npu.npu_silu	计算self的Swish。Swish是一种激活函数，计算公式为' x * sigmoid(x) '。
（beta）torch_npu.npu_slice	从张量中提取切片。
（beta）torch_npu.npu_softmax_cross_entropy_with_logits	计算softmax的交叉熵cost。
（beta）torch_npu.npu_sort_v2	沿给定维度，对输入张量元素进行升序排序（不返回索引）。若dim未设置，则选择输入的最后一个维度。如果descending为True，则元素将按值降序排序。
（beta）torch_npu.npu_transpose	返回原始张量视图，其维度已permute，结果连续。
（beta）torch_npu.npu_yolo_boxes_encode	根据YOLO的锚点框（anchor box）和真值框（ground-truth box）生成标注框。自定义mmdetection算子。
（beta）torch_npu.npu_fused_attention_score	实现“Transformer attention score”的融合计算逻辑，主要将matmul、transpose、add、softmax、dropout、batchmatmul、permute等计算进行了融合。
（beta）torch_npu.npu_multi_head_attention	实现Transformer模块中的MultiHeadAttention计算逻辑。
（beta）torch_npu.npu_rms_norm	RmsNorm算子是大模型常用的归一化操作，相比LayerNorm算子，其去掉了减去均值的部分。
（beta）torch_npu.npu_dropout_with_add_softmax	实现axpy_v2、softmax_v2、drop_out_domask_v3功能。
torch_npu.npu_rotary_mul	实现RotaryEmbedding旋转位置编码。
torch_npu.npu_scaled_masked_softmax	计算输入张量x缩放并按照mask遮蔽后的Softmax结果。
（beta）torch_npu.npu_swiglu	提供swiglu的激活函数。
（beta）torch_npu.one_	用1填充self张量。
torch_npu.npu_group_norm_swish	计算输入`input`的组归一化结果`y`，均值`mean`，标准差的倒数`rstd`，以及swish的输出。
torch_npu.npu_cross_entropy_loss	将原生CrossEntropyLoss中的log_softmax和nll_loss融合，降低计算时使用的内存。接口允许计算zloss。
torch_npu.npu_advance_step_flashattn	在NPU上实现vLLM库中advance_step_flashattn的功能，在每个生成步骤中原地更新input_tokens，input_positions，seq_lens和slot_mapping。
torch_npu.npu_all_gather_base_mm	TP切分场景下，实现allgather和matmul的融合，融合算子内部实现通信和计算流水并行。
torch_npu.npu_anti_quant	将INT4或者INT8数据反量化为FP16或者BF16，其中输入是INT4类型时，将每8个数据看作是一个INT32数据。
torch_npu.npu_convert_weight_to_int4pack	将数据类型为int32的输入tensor打包为int4存放，每8个int4数据通过一个int32数据承载，并进行交叠排放。
torch_npu.npu_dynamic_quant	为输入的张量进行pre-token对称动态量化。
torch_npu.npu_dynamic_quant_asymmetric	对输入的张量进行pertoken非对称动态量化。其中输入的最后一个维度对应一个token，每个token作为一组进行量化。
torch_npu.npu_fast_gelu	快速高斯误差线性单元激活函数（Fast Gaussian Error Linear Units activation function），对输入的每个元素计算FastGelu；输入是具有任何有效形状的张量。
torch_npu.npu_ffn	该FFN算子提供MoeFFN和FFN的计算功能。在没有专家分组（expert_tokens为空）时是FFN，有专家分组时是MoeFFN。
torch_npu.npu_fused_infer_attention_score	适配增量&全量推理场景的FlashAttention算子，既可以支持全量计算场景（PromptFlashAttention），也可支持增量计算场景（IncreFlashAttention）。
torch_npu.npu_fused_infer_attention_score_v2	适配增量&全量推理场景的FlashAttention算子，既可以支持全量计算场景（PromptFlashAttention），也可支持增量计算场景（IncreFlashAttention）。新增支持MLA（MultiHead Latent Attention）全量化功能。
torch_npu.npu_fusion_attention	实现“Transformer Attention Score”的融合计算。
torch_npu.npu_gelu	计算高斯误差线性单元的激活函数。
torch_npu.npu_gelu_mul	对输入Tensor进行高斯误差线性单元（GELU）与MUL乘法结合的复合计算操作。
torch_npu.npu_group_norm_silu	计算输入self的组归一化结果out、均值meanOut、标准差的倒数rstdOut、以及silu的输出。
torch_npu.npu_group_quant	对输入的张量进行分组量化操作。
torch_npu.npu_grouped_matmul	npu_grouped_matmul是一种对多个矩阵乘法（matmul）操作进行分组计算的高效方法。
torch_npu.npu_grouped_matmul_swiglu_quant_v2	npu_grouped_matmul_swiglu_quant_v2是一种融合分组矩阵乘法（GroupedMatmul）、SwiGLu混合激活函数、量化（quant）的计算方法，用于模型性能优化。
torch_npu.npu_incre_flash_attention	增量FA实现。
torch_npu.npu_kv_quant_sparse_flash_attention	Sparse Flash Attention伪量化实现。
torch_npu.npu_lightning_indexer	获取每一个token对应的Top-k个位置（全量化）。
torch_npu.npu_mla_prolog	推理场景下，Multi-Head Latent Attention前处理计算接口。
torch_npu.npu_mla_prolog_v2	推理场景下，Multi-Head Latent Attention前处理计算的增强接口。
（beta）torch_npu.npu_mla_prolog_v3	推理场景下，Multi-Head Latent Attention前处理计算的增强接口。
torch_npu.npu_mm_all_reduce_base	TP切分场景下，实现mm和all_reduce的融合，融合算子内部实现计算和通信流水并行。
torch_npu.npu_mm_reduce_scatter_base	TP切分场景下，实现matmul和reduce_scatter的融合，融合算子内部实现计算和通信流水并行。
torch_npu.npu_moe_compute_expert_tokens	MoE计算中，通过二分查找的方式查找每个专家处理的最后一行的位置。
torch_npu.npu_moe_finalize_routing	MoE计算中，最后处理合并MoE FFN的输出结果。
torch_npu.npu_moe_gating_top_k_softmax	MoE计算中，对gating的输出做Softmax计算，取topk操作。
torch_npu.npu_moe_init_routing	MoE的routing计算，根据torch_npu.npu_moe_gating_top_k_softmax的计算结果做routing处理。
torch_npu.npu_prefetch	提供网络weight预取功能，将需要预取的权重搬到L2 Cache中（当前仅支持权重的预取，暂不支持KV cache的预取）。尤其在做较大Tensor的MatMul计算且需要搬移到L2 Cache的操作时，可通过该接口提前预取权重，适当提高模型性能，具体效果取决与用户采用的并行方式和配置。
torch_npu.npu_prompt_flash_attention	全量FA实现。
torch_npu.npu_quant_lightning_indexer	获取每一个token对应的Top-k个位置。
torch_npu.npu_quant_matmul	完成量化的矩阵乘计算，最小支持输入维度为2维，最大支持输入维度为6维。
torch_npu.npu_quant_matmul_reduce_sum	完成量化的分组矩阵计算，然后所有组的矩阵计算结果相加后输出。
torch_npu.npu_quant_scatter	先将updates进行量化，然后将updates中的值按指定的轴axis和索引indices更新self中的值，并将结果保存到输出tensor，self本身的数据不变。
torch_npu.npu_recurrent_gated_delta_rule	完成变步长的Recurrent Gated Delta Rule计算。
torch_npu.npu_quant_scatter_	先将updates进行量化，然后将updates中的值按指定的轴axis和索引indices更新self中的值，self中的数据被改变。
torch_npu.npu_quantize	对输入的张量进行量化处理。
torch_npu.npu_scatter_nd_update	将updates中的值按指定的索引indices更新self中的值，并将结果保存到输出tensor，self本身的数据不变。
torch_npu.npu_sparse_flash_attention	Sparse Flash Attention实现。
torch_npu.npu_sparse_lightning_indexer_grad_kl_loss	该接口实现了npu_lightning_indexer的反向功能，并融合了Loss的计算。
torch_npu.npu_top_k_top_p	对原始输入`logits`进行`top-k`和`top-p`采样过滤。
torch_npu.npu_top_k_top_p_sample	根据输入词频`logits`、`top_k`/`top_p`采样参数、随机采样权重分布`q`，进行topK-topP-Sample采样计算，输出每个batch的最大词频`logits_select_idx`，以及topK-topP采样后的词频分布`logits_top_kp_select`。
torch_npu.npu_scatter_pa_kv_cache	更新KvCache中指定位置的`key`和`value`。
torch_npu.npu_scatter_nd_update_	将updates中的值按指定的索引indices更新self中的值，并将结果保存到输出tensor，self中的数据被改变。
torch_npu.npu_trans_quant_param	完成量化计算参数scale数据类型的转换。
torch_npu.npu_weight_quant_batchmatmul	该接口用于实现矩阵乘计算中的weight输入和输出的量化操作，支持pertensor，perchannel，pergroup多场景量化。
torch_npu.scatter_update	将tensor updates中的值按指定的轴axis和索引indices更新tensor data中的值，并将结果保存到输出tensor，data本身的数据不变。
torch_npu.scatter_update_	将tensor updates中的值按指定的轴axis和索引indices更新tensor data中的值，并将结果保存到输出tensor，data本身的数据被改变。
torch_npu.empty_with_swapped_memory	申请一个device信息为NPU且实际内存在host侧的特殊Tensor。
torch_npu.erase_stream	Tensor通过`record_stream`在内存池上添加的已被stream使用的标记后，可以通过该接口移除该标记。
torch_npu.npu_gather_sparse_index	从输入Tensor的指定维度，按照`index`中的下标序号提取元素，保存到输出Tensor中。
torch_npu.npu_moe_distribute_combine	先进行reduce_scatterv通信，再进行alltoallv通信，最后将接收的数据整合（乘权重再相加）。
torch_npu.npu_moe_distribute_dispatch	对Token数据先进行量化（可选），再进行EP（Expert Parallelism）域的alltoallv通信，再进行TP（Tensor Parallelism）域的allgatherv通信（可选）。
torch_npu.npu_moe_distribute_combine_v2	先进行reduce_scatterv通信，再进行alltoallv通信，最后将接收的数据整合（乘权重再相加）。
torch_npu.npu_moe_distribute_dispatch_v2	对Token数据先进行量化（可选），再进行EP（Expert Parallelism）域的alltoallv通信，再进行TP（Tensor Parallelism）域的allgatherv通信（可选）。
torch_npu.npu_moe_gating_top_k	MoE计算中，对输入x做Sigmoid计算，对计算结果分组进行排序，最后根据分组排序的结果选取前k个专家。
torch_npu.npu_moe_init_routing_v2	MoE（Mixture of Expert）的routing计算，根据torch_npu.npu_moe_gating_top_k_softmax的计算结果做routing处理，支持不量化和动态量化模式。
torch_npu.npu_swiglu_quant	在swiglu激活函数后添加quant操作，实现输入x的SwiGluQuant计算，支持int8或int4量化输出，支持MoE场景和非MoE场景（group_index为空），支持分组量化，支持动态/静态量化。
torch_npu.npu_dequant_swiglu_quant	对张量x做dequant反量化+swiglu+quant量化操作，同时支持分组。
torch_npu.npu_kv_rmsnorm_rope_cache	融合了MLA（Multi-head Latent Attention）结构中RMSNorm归一化计算与RoPE（Rotary Position Embedding）位置编码以及更新KVCache的ScatterUpdate操作。
torch_npu.npu_interleave_rope	针对单输入x进行旋转位置编码。
torch_npu.npu_moe_re_routing	MoE网络中，进行AlltoAll操作从其他卡上拿到需要算的token后，将token按照专家序重新排列。
torch_npu.matmul_checksum	提供基于原生torch.matmul和Tensor.matmul接口的aicore错误硬件故障接口，内部执行矩阵计算结果校验过程，并对校验误差和实时计算的校验门限进行对比，判断校验误差是否超越门限，若超越则认为发生了aicore错误。
torch_npu.npu_alltoallv_gmm	MoE网络中，完成路由专家AlltoAllv、Permute、GroupedMatMul融合并实现与共享专家MatMul并行融合，先通信后计算。
torch_npu.npu_gmm_alltoallv	MoE网络中，完成路由专家GroupedMatMul、AlltoAllv融合并实现与共享专家MatMul并行融合，先计算后通信。
torch_npu.npu_moe_distribute_combine_add_rms_norm	完成moe_distribute_combine+add+rms_norm融合。需与torch_npu.npu_moe_distribute_dispatch配套使用，相当于按npu_moe_distribute_dispatch算子收集数据的路径原路返回后对数据进行add_rms_norm操作。
torch_npu.npu_transpose_batchmatmul	完成张量input与张量weight的矩阵乘计算。
torch_npu.npu_moe_update_expert	提供MoE网络中常用的EPLB（Expert Parallelism Load Balancer）算法进行冗余专家部署，解决负载不均衡问题。MoeUpdateExpert算子可以完成每个token的topK个专家逻辑专家号到物理专家实例号的映射，支持根据阈值对token发送的topK个专家进行剪枝。
torch_npu.npu_dynamic_block_quant	对输入张量，通过给定的`row_block_size`和`col_block_size`将输入划分成多个数据块，以数据块为基本粒度进行量化。
torch_npu.set_device_limit	设置指定device上算子在执行时所使用的cube和vector核数。
torch_npu.get_device_limit	获取指定device上算子执行的cube和vector核数。
torch_npu.set_stream_limit	设置指定流上算子在执行时所使用的cube和vector核数。
torch_npu.reset_stream_limit	恢复指定流上算子在执行时所使用的cube和vector核数为默认配置。
torch_npu.get_stream_limit	获取指定流上算子执行的cube和vector核数。
torch_npu.npu_sim_exponential_	原地获取输入张量中元素的指数分布采样。