File size: 640 Bytes
2976eec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | import triton
import triton.language as tl
import torch
def torch_silu_and_mul_ref(input):
"""
Performs the SiLU activation on the first half of the input tensor and
multiplies it element-wise with the second half.
Args:
input (torch.Tensor): Input tensor of shape [..., 2 * d].
param (float): Parameter for the SiLU activation function.
Returns:
torch.Tensor: Output tensor of shape [..., d].
"""
dtype = input.dtype
d = input.size(-1) // 2
A, B = input[:, :d], input[:, d:]
silu_A = A / (1.0 + torch.exp(-A.float()))
output = silu_A * B
return output.to(dtype)
|