Custom Operator

During recommendation training, some operators are not implemented on NPUs, or some operators are implemented on NPUs but the performance is poor, which cannot meet the recommendation training requirements. Rec SDK Torch provides the following Table 1 Custom operator list to support or accelerate NPU training of recommendation models.

Some custom operators have been bound to open-source APIs (the backend implementation of open-source APIs is forwarded to the NPU). After the Rec SDK Torch software package is imported, the custom operators can be directly called through open-source APIs. Other operators need to be called through the mxrec module registered at the PTA layer. For details, see the README file in the operator directory under RecSDK/cust_op/ascendc_op/ai_core_op/ in the Rec SDK Torch code library.

Custom operators are used for high-performance computing. When calling a custom operator, ensure that the input parameters meet the operator constraints, parameter types, and parameter shapes. Otherwise, the operator may fail to be executed due to out-of-bounds array or insufficient video memory.

**Table 1** Custom operator list
Operator Name	Open-Source API	Support Open-Source API Calling (Yes/No)	Function
asynchronous_complete_cumsum	torch.ops.fbgemm.asynchronous_complete_cumsum	Yes	Accumulates the input offset.
dense_to_jagged	torch.ops.fbgemm.dense_to_jagged	Yes	Converts a padded dense tensor into a jagged tensor.
gather_for_rank1	torch.index_select	No	Implements the index_select function when the shape of the input parameter is 1. This operator is optimized for special shapes.
index_select_for_rank1_backward	torch.index_select	No	Reverse operator of gather_for_rank1.
hstu_dense_forward	None	-	The hstu operator is used to implement the attention mechanism in recommendation scenarios. Forward operator of the hstu operator. The hstu operator contains the forward and backward operators, which have been bound to Autograd.
hstu_dense_backward	None	-	Reverse operator of the hstu operator.
hstu_dense_forward_fuxi	None	-	Adapts to the fuxi-alpha model based on the hstu forward operator.
hstu_dense_backward_fuxi	None	-	Adapts to the fuxi-alpha model based on the hstu reverse operator.
jagged_to_padded_dense	torch.ops.fbgemm.jagged_to_padded_dense	Yes	Converts a jagged tensor into a padded dense tensor. This operator and the dense_to_jagged operator are forward and reverse operators of each other and have been bound to Autograd.
permute_1D_sparse_data	torch.ops.fbgemm.permute_1D_sparse_data	Yes	Transposes 1D sparse data. This operator does not have an independent directory and is implemented by calling the 2D operator. The implementation is the same as that of FBGEMM. For details about how to call this operator, see the fbgemm API document.
permute2d_sparse_data	torch.ops.fbgemm.permute_2D_sparse_data	Yes	Transposes 2D sparse data.
relative_attn_bias_pos	None	-	Fused operator relative_attn_bias with the pos parameter.
relative_attn_bias_time	None	-	Fused operator relative_attn_bias with the time parameter.
relative_attn_bias_backward	None	-	Reverse operator of relative_attn_bias.
split_embedding_codegen_lookup_adagrad_function	fbgemm_gpu.split_table_batched_embeddings_ops_training.SplitTableBatchedEmbeddingBagsCodegen	Yes	Implements the forward function of SplitTableBatchedEmbeddingBagsCodegen table lookup in the FBGEMM library.
backward_codegen_adagrad_unweighted_exact	fbgemm_gpu.split_table_batched_embeddings_ops_training.SplitTableBatchedEmbeddingBagsCodegen	Yes	Implements the reverse update function of SplitTableBatchedEmbeddingBagsCodegen table lookup in the FBGEMM library. Different optimizers use different update algorithms.
disetangle_attention	None	-	Implements disetangle_attention in the DeBERTa model to decouple attention.
dense_embedding_codegen_lookup_function	torch.ops.fbgemm. dense_embedding_nobag_codegen_forward_unweighted_cuda	Yes	Implements the forward function of the dense_embedding_nobag_codegen_forward_unweighted_cuda table lookup in the FBGEMM GPU library.
dense_embedding_codegen_lookup_function_grad	None	-	Implements the backward gradient accumulation function of the dense_embedding_nobag_codegen_forward_unweighted_cuda table lookup in the FBGEMM GPU library.

The table lists only the custom operators in the PyTorch scenario in the ai_core_op path.

Parent topic: API Reference