Triton SiLU (Swish) Kernel

mediumPROElementwise & Fused OpsTriton Programming

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

Input

n = 1024

Output

True

Reasoning

The input value n = 1024 determines the size of the input array for the SiLU kernel implementation.
The silu_kernel function calculates the SiLU activation using the formula $out = x \cdot \sigma(x)$ , where $\sigma(x)$ is the sigmoid function, for each element in the input array.
The result from the silu_kernel function is compared to the result from torch.nn.functional.silu using torch.allclose to check for numerical equivalence within a certain tolerance.
The comparison yields True if the results are close enough, indicating that the custom silu_kernel implementation matches the PyTorch reference implementation.

Constraints

out = x * tl.sigmoid(x)
Mask the tail block

More from Triton Programming

Triton Single-Query Attention Kernelhard Triton Row Cumulative Sumhard Triton Stable LogSumExp Kernelhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium

Concept primerPREVIEW

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 * \text{sigmoid}(x)$ , where $\text{sigmoid}(x) = \frac{1}{1 + e^{-x}}$ . 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.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

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.

solution.py

Python 3.13

GPU · T4

Test results