activation.h

激活函数。

函数

void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算输入的 GELU 激活函数。

参数
  • input[in] GELU 激活函数的输入张量。

  • output[inout] 输出张量。

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

void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算 GELU 激活函数的梯度。

参数
  • grad[in] 输入梯度。

  • input[in] GELU 激活函数的输入张量。

  • output[inout] 输出张量。

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

void nvte_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算输入的 GeGLU。

参数
  • input[in] 形状为 [N, H * 2] 的输入张量。

  • output[inout] 形状为 [N, H] 的输出张量。它计算 GELU(input[N, :H]) x input[N, H:]

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

void nvte_dgeglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算 GeGLU 梯度。

参数
  • grad[in] 形状为 [N, H] 的输入梯度。

  • input[in] 形状为 [N, H * 2] 的前向输入张量。

  • output[inout] 形状为 [N, H * 2] 的输出梯度。

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

void nvte_relu(const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算输入的 RELU 激活函数。

参数
  • input[in] RELU 激活函数的输入张量。

  • output[inout] 输出张量。

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

void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算 RELU 激活函数的梯度。

参数
  • grad[in] 输入梯度。

  • input[in] RELU 激活函数的输入张量。

  • output[inout] 输出张量。

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

void nvte_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算输入的 SwiGLU 激活函数。

参数
  • input[in] 形状为 [N, H * 2] 的输入张量。

  • output[inout] 形状为 [N, H] 的输出张量。它计算 Swish(input[N, :H]) x input[N, H:]

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

void nvte_dswiglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算 SwiGLU 梯度。

参数
  • grad[in] 形状为 [N, H] 的输入梯度。

  • input[in] 形状为 [N, H * 2] 的前向输入张量。

  • output[inout] 形状为 [N, H * 2] 的输出梯度。

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

void nvte_reglu(const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算输入的 ReGLU 激活函数。

参数
  • input[in] 形状为 [N, H * 2] 的输入张量。

  • output[inout] 形状为 [N, H] 的输出张量。它计算 ReLU(input[N, :H]) x input[N, H:]

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

void nvte_dreglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream)

计算 ReGLU 梯度。

参数
  • grad[in] 形状为 [N, H] 的输入梯度。

  • input[in] 形状为 [N, H * 2] 的前向输入张量。

  • output[inout] 形状为 [N, H * 2] 的输出梯度。

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