| 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) | |