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)