flashinfer.activation.gelu_and_mul

flashinfer.activation.gelu_and_mul(input: Tensor, out: Tensor = None, enable_pdl: bool | None = None) → Tensor

融合的 GeLU 和 Mul 操作。

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

参数:

• input (torch.Tensor) – 输入张量，形状为 (…, 2 * hidden_size)。

• out (Optional[torch.Tensor]) – 输出张量，如果指定，则内核将就地更新此张量。

• enable_pdl (bool) – 是否启用 程序依赖启动

返回值:

output – 输出张量，形状为 (…, hidden_size)。

返回值类型:

torch.Tensor