activation.h

Activation functions.

Functions

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

Compute GELU activation of the input.

Parameters:

  • input – [in] Input tensor for GELU activation.

  • output – [inout] Output tensor.

  • stream – [in] CUDA stream used for the operation.

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

Compute GeGLU of the input.

Parameters:

  • input – [in] Input tensor of shape [N, H * 2]. It computes GELU([N, :H]) x [N, H:]

  • output – [inout] Output tensor of shape [N, H].

  • stream – [in] CUDA stream used for the operation.

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

Compute GeGLU gradient.

Parameters:

  • grad – [in] Input tensor of shape [N, H].

  • input – [in] Input tensor of shape [N, H * 2].

  • output – [inout] Output tensor of shape [N, H * 2].

  • stream – [in] CUDA stream used for the operation.