softmax.h

函数

void nvte_scaled_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream)

计算输入的缩放softmax激活。

参数

• input – [in] Softmax的输入张量。

• softmax_results – [out] 输出张量。

• scale_factor – [in] 输入张量的标量。

• stream – [in] 用于操作的CUDA流。

void nvte_scaled_softmax_backward(const NVTETensor incoming_grads, const NVTETensor softmax_results, NVTETensor output_grads, float scale_factor, cudaStream_t stream)

计算缩放softmax激活的反向传播。

• incoming_grads 是包含从下一层接收到的梯度的输入张量。

• softmax_results 是对应前向softmax操作的输出张量。

• output_grads 是包含计算出的梯度的输出张量。

参数

• incoming_grads – [in] 反向传播的输入梯度张量。

• softmax_results – [in] Softmax前向传播的输出张量。

• output_grads – [out] 输出张量。

• scale_factor – [in] 输出张量的标量。

• stream – [in] 用于操作的CUDA流。

void nvte_scaled_masked_softmax_forward(const NVTETensor input, const NVTETensor mask, NVTETensor softmax_results, float scale_factor, cudaStream_t stream)

计算输入的缩放掩码softmax激活。

参数

• input – [in] Softmax的输入张量。

• mask – [in] 输入张量的掩码。

• softmax_results – [out] 输出张量。

• scale_factor – [in] 输入张量的标量。

• stream – [in] 用于操作的CUDA流。

void nvte_scaled_masked_softmax_backward(const NVTETensor incoming_grads, const softmax_results, NVTETensor output_grads, float scale_factor, cudaStream_t stream)

计算缩放掩码softmax激活的反向传播。

• incoming_grads 是包含从下一层接收到的梯度的输入张量。

• softmax_results 是对应前向softmax操作的输出张量。

• output_grads 是包含计算出的梯度的输出张量。

参数

• incoming_grads – [in] 反向传播的输入梯度张量。

• softmax_results – [in] Softmax前向传播的输出张量。

• output_grads – [out] 输出张量。

• scale_factor – [in] 输出张量的标量。

• stream – [in] 用于操作的CUDA流。

void nvte_scaled_upper_triang_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream)

使用二维上三角掩码计算输入的缩放softmax激活。

参数

• input – [in] Softmax的输入张量。

• softmax_results – [out] 输出张量。

• scale_factor – [in] 输入张量的标量。

• stream – [in] 用于操作的CUDA流。

void nvte_scaled_upper_triang_masked_softmax_backward(const NVTETensor incoming_grads, const softmax_results, NVTETensor output_grads, float scale_factor, cudaStream_t stream)

使用二维上三角掩码计算缩放softmax激活的反向传播。

• incoming_grads 是包含从下一层接收到的梯度的输入张量。

• softmax_results 是对应前向softmax操作的输出张量。

• output_grads 是包含计算出的梯度的输出张量。

参数

• incoming_grads – [in] 反向传播的输入梯度张量。

• softmax_results – [in] Softmax前向传播的输出张量。

• output_grads – [out] 输出张量。

• scale_factor – [in] 输出张量的标量。

• stream – [in] 用于操作的CUDA流。