Triton SiLU (Swish) Kernel
Problem Statement
Implement the SiLU / Swish activation out = x * sigmoid(x).
Background
Triton provides tl.sigmoid. SiLU is used in modern transformer MLP blocks.
Your Task
Implement silu_kernel and run(n=1024) comparing to torch.nn.functional.silu.
How it is tested
Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.
Example:
n = 1024
True
- The input value
n = 1024determines the size of the input array for the SiLU kernel implementation. - The
silu_kernelfunction calculates the SiLU activation using the formula out=xโ ฯ(x), where ฯ(x) is the sigmoid function, for each element in the input array. - The result from the
silu_kernelfunction is compared to the result fromtorch.nn.functional.siluusingtorch.allcloseto check for numerical equivalence within a certain tolerance. - The comparison yields
Trueif the results are close enough, indicating that the customsilu_kernelimplementation matches the PyTorch reference implementation.
Constraints:
- out = x * tl.sigmoid(x)
- Mask the tail block
Background Knowledge
The SiLU (Swish) activation function is a type of activation function used in deep learning models, particularly in transformer-based architectures. It is defined as out=xโsigmoid(x), where sigmoid(x)=1+eโx1โ. This function is also known as the Swish activation function. The sigmoid function is a common activation function used in neural networks, and it maps any real-valued number to a value between 0 and 1.
In the context of the Triton Programming collection, the problem requires implementing the SiLU activation function using Triton's tl.sigmoid function. Triton is a programming language and framework for writing high-performance, GPU-accelerated code. The goal is to implement a custom kernel, silu_kernel, that computes the SiLU activation function and compare its output to the reference implementation provided by torch.nn.functional.silu.
The element-wise operations are a key concept in this problem, as the SiLU activation function is applied element-wise to the input tensor. This means that each element of the input tensor is processed independently, and the output tensor has the same shape as the input tensor. Understanding how to perform element-wise operations efficiently on the GPU is crucial for achieving high performance in this problem.
Algorithm/Approach
The general approach to solving this problem involves the following steps:
- Define a custom kernel function, silu_kernel, that takes an input tensor and computes the SiLU activation function element-wise.
- Use Triton's tl.sigmoid function to compute the sigmoid of the input tensor.
- Multiply the input tensor with the sigmoid of the input tensor to compute the final output.
- Compare the output of the custom kernel with the reference implementation provided by torch.nn.functional.silu using torch.allclose.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.