RecOps Operator List

RecOps is a custom operator set developed by the Rec SDK based on Ascend C for recommendation scenarios. It provides basic operator capabilities for framework components (tf_rec_v1, tf_rec_v2, torch_rec_v1, and torch_rec_v2).

Introduction

Operator Name Description Supported Hardware
concat_jagged_tensor Concatenates two jagged tensors along the dim1 dimension according to the offset and merges them into one tensor. ascend910b, ascend910_93, ascend310p, ascend950
concat_jagged_tensor_grad Splits one tensor into two tensors of potentially different lengths according to the offset. ascend910b, ascend910_93, ascend310p, ascend950
cust_op_by_addr Uses addr as input to query the embedding table, replaces the tf.gather operator, and supports dynamic expansion. ascend910, ascend910b, ascend910_93, ascend950
disentangle_attention Implements the disentangled attention mechanism in the DeBERTa model. ascend910b, ascend910_93, ascend950
fused_lazy_adam Calculates and updates m, v, and variable during the backward update of the LazyAdam optimizer. ascend910b, ascend910_93, ascend950
fused_sgd Calculates and updates parameters during the backward update of the SGD optimizer. ascend910b, ascend910_93, ascend950
gather_for_rank1 Implements index_select for a rank-1 tensor and selects elements from a 1D tensor by index. ascend910b, ascend910_93, ascend310p, ascend950
hstu_dense_backward Implements backpropagation for the HSTU fusion operator and calculates the gradients of Q, K, V, and attn_bias. ascend910b, ascend910_93, ascend950
hstu_dense_backward_fuxi Implements backpropagation for the HSTU-Fuxi fusion operator and calculates gradients in the attention mechanism. ascend910b, ascend910_93, ascend950
hstu_dense_forward Uses the HSTU fusion operator to implement the attention mechanism in recommendation scenarios. It supports GQA and unequal dimensions. ascend910b, ascend910_93, ascend310p, ascend950
hstu_dense_forward_fuxi Based on the HSTU fusion operator, implements the attention mechanism in the Alpha-Fuxi model for recommendation scenarios. It supports timestamp_bias and position_bias, as well as the normal, jagged, and paged layouts. ascend910b, ascend910_93, ascend310p, ascend950
hstu_backward_v2 Implements backpropagation for the HSTU V2 fusion operator and calculates the gradients of Query, Key, Value, and RAB attention scores. ascend950
index_select_for_rank1_backward Implements backpropagation for index_select and calculates gradients. ascend910b, ascend910_93, ascend950
in_linear_silu Before HSTU attention, applies Linear and Silu to the merged and normalized UVQK, then splits it into four tensors: User, Value, Query, and Key. ascend910b, ascend910_93, ascend950
in_linear_silu_backward Implements backpropagation for the in_linear_silu operator and computes the gradients of input x, weight, and bias. ascend910b, ascend910_93, ascend950
lccl Uses the ability of AI Core to directly access the on-chip memory of the peer and uses memory semantics for collective communication (AllToAll, AllUss, and GatherAll). ascend910b, ascend910_93
ln_mul Applies LayerNorm to input X, then performs the gamma and beta computation, and finally multiplies it by input U. ascend910b, ascend910_93, ascend310p, ascend950
multislice_concat Slices the input 2D tensor along the second dimension at the specified positions and lengths, then outputs a tensor composed of several slices. ascend910b, ascend950
norm_multiply_dropout Implements the fusion operator for the layer_norm + multiply + dropout computation logic. ascend910b, ascend910_93, ascend950
norm_multiply_dropout_backward Implements the backpropagation logic for the layer_norm + multiply + dropout computation. ascend910b, ascend910_93, ascend950
pcie_through When large amounts of data are exchanged between the host and device, pcie_through improves swap-in and swap-out performance. ascend910b
relative_attn_bias_backward For the time part of the HSTU model RAB, calculates gradient values during backpropagation for the timestamp parameter. ascend910b, ascend910_93, ascend950
relative_attn_bias_pos Computes the pos part of the HSTU model RAB. ascend910b, ascend910_93, ascend310p, ascend950
relative_attn_bias_time Computes the time part of the HSTU model RAB. ascend910b, ascend910_93, ascend310p, ascend950
reverse_sequence Reverses the second dimension of the input data according to the specified length parameter (seq_lengths). ascend910b, ascend910_93, ascend950
token_mixing Normalizes the sum of x and its transpose x_t. ascend910b, ascend950

Operator Directory Structure

cust_op/ascendc_op/ai_core_op/
├── concat_jagged_tensor/                              # Jagged tensor concatenation
├── concat_jagged_tensor_grad/                         # Jagged tensor concatenation backward
├── cust_op_by_addr/                                   # Address lookup
├── disentangle_attention/                             # Disentangled attention
├── fused_lazy_adam/                                   # LazyAdam optimizer
├── fused_sgd/                                         # SGD optimizer
├── gather_for_rank1/                                  # Rank-1 gather
├── hstu_dense_backward/                               # HSTU backward
├── hstu_dense_backward_fuxi/                          # HSTU-Fuxi backward
├── hstu_dense_forward/                                # HSTU forward
├── hstu_dense_forward_fuxi/                           # HSTU-Fuxi forward
├── hstu_v2/                                           # HSTU V2 backward
├── index_select_for_rank1_backward/                   # index_select backward
├── in_linear_silu/                                    # Linear+Silu fusion
├── in_linear_silu_backward/                           # Linear+Silu backward
├── lccl/                                              # Collective communication
├── ln_mul/                                            # LayerNorm+Multiply fusion
├── multislice_concat/                                 # Multi-slice concatenation
├── norm_multiply_dropout/                             # Norm+Multiply+Dropout fusion
├── norm_multiply_dropout_backward/                    # Norm+Multiply+Dropout backward
├── pcie_through/                                      # PCIe data transfer
├── relative_attn_bias_backward/                       # Relative position bias backward
├── relative_attn_bias_pos/                            # Relative position bias pos
├── relative_attn_bias_time/                           # Relative position bias time
├── reverse_sequence/                                  # Sequence reversal
└── token_mixing/                                      # Token mixing