flashinfer.activation.gelu_tanh_and_mul

flashinfer.activation.gelu_tanh_and_mul(input: torch.Tensor, out: torch.Tensor = None) torch.Tensor

Fused GeLU Tanh and Mul operation.

gelu(tanh(input[..., :hidden_size])) * input[..., hidden_size:]

Parameters:
  • input (torch.Tensor) – Input tensor, shape (…, 2 * hidden_size).

  • out (Optional[torch.Tensor]) – The the output tensor, if specified, the kernel will update this tensor inplace.

Returns:

output – Output tensor, shape (…, hidden_size).

Return type:

torch.Tensor