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 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_upper_triang_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream)

使用输入上的 2D 上三角掩码计算缩放的 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 NVTETensor softmax_results, NVTETensor output_grads, float scale_factor, cudaStream_t stream)

使用 2D 上三角掩码计算缩放的 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_aligned_causal_masked_softmax_forward(const NVTETensor input, NVTETensor softmax_results, float scale_factor, cudaStream_t stream)

使用与输入矩阵的右下角对齐的隐式 2D 掩码计算缩放的 Softmax 激活。

参数:

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

• softmax_results – [out] 输出张量。

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

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

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

使用与输入矩阵的右下角对齐的隐式 2D 掩码计算缩放的 Softmax 激活的反向传播。

• incoming_grads 是输入张量，包含从后续层接收的梯度。

• softmax_results 是相应的前向 Softmax 操作的输出张量。

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

参数:

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

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

• output_grads – [out] 输出张量。

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

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