Triton Scalar Multiply Kernel

easyFREETriton FundamentalsTriton Programming

Problem Statement

Write a Triton kernel that multiplies every element of a 1D tensor by a runtime scalar: out = x * scale.

Background

Scalars can be passed as ordinary kernel arguments (no pointer needed) and used directly inside the kernel.

Your Task

Implement mul_kernel and run(n=1024, scale=3.0) that returns whether the output equals x * scale.

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, scale = 3.0

Output

True

Reasoning

The input values are n = 1024 and scale = 3.0, which are used to allocate a 1D tensor of length n and a scalar value scale.
The mul_kernel function is launched, which multiplies every element of the input tensor x by the scale value, resulting in an output tensor out where each element is calculated as $out_i = x_i \cdot scale$ .
The output tensor out is then compared to a reference tensor calculated using PyTorch, where each element is also $x_i \cdot scale$ , using the torch.allclose function with a small tolerance.
The comparison returns True if all elements of the output tensor are close to the corresponding elements of the reference tensor, indicating that the Triton kernel produced the correct result.

Constraints

Pass scale as a plain kernel argument
Mask out-of-bounds lanes
out = x * scale

Learn this in

Introduction to TritonGPU Basics with Triton Your First Kernel: Vector AdditionGPU Basics with Triton Memory Access, Autotuning & BenchmarkingGPU Basics with Triton

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 primer

Background Knowledge

The problem involves implementing a Triton kernel, which is a fundamental concept in the Triton Programming collection. Triton is a programming language and framework for writing high-performance, GPU-accelerated code. In this context, a kernel refers to a small program that runs on the GPU, performing a specific computation. The problem requires writing a kernel that performs a simple scalar multiplication operation, where every element of a 1D tensor is multiplied by a runtime scalar.

To understand this problem, it's essential to have a basic understanding of GPU programming and tensor operations. In GPU programming, data is typically divided into smaller chunks and processed in parallel by multiple threads. Tensors are multi-dimensional arrays used to represent data in various deep learning and scientific computing applications. In this case, we're working with a 1D tensor, which can be thought of as a vector. The scalar multiplication operation is a fundamental operation in linear algebra, where every element of a vector is multiplied by a scalar value.

The problem also involves using Triton and PyTorch, two popular frameworks for deep learning and GPU-accelerated computing. Triton provides a Python-like syntax for writing high-performance kernels, while PyTorch provides a dynamic computation graph and automatic differentiation for building and training neural networks. The problem requires using these frameworks to implement the kernel and test its correctness.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Define a Triton kernel that takes a 1D tensor and a scalar value as input
Use Triton's programming model to perform the scalar multiplication operation in parallel across the tensor
Launch the kernel on the GPU and store the result in a new tensor
Compare the result with a reference implementation using PyTorch to verify correctness

The algorithm pattern involved is a simple element-wise operation, where every element of the input tensor is multiplied by the scalar value. This operation can be parallelized easily, making it a good fit for GPU acceleration.

Step-by-Step Strategy

To implement the solution, follow these steps:

Define a Triton kernel function mul_kernel that takes a 1D tensor x and a scalar value scale as input.
Use Triton's programming model to perform the scalar multiplication operation in parallel across the tensor. This will involve using Triton's tl module to launch a kernel with multiple threads.
Store the result of the scalar multiplication operation in a new tensor out.
Define a run function that allocates input data on the GPU, launches the mul_kernel kernel, and stores the result in a new tensor.
Use PyTorch to create a reference implementation of the scalar multiplication operation and compare the result with the output of the mul_kernel kernel.

Common Pitfalls

Some common pitfalls to watch out for when implementing this solution include:

Forgetting to allocate input data on the GPU before launching the kernel
Incorrectly launching the kernel with the wrong number of threads or blocks
Failing to synchronize the kernel launch with the host thread
Incorrectly comparing the result of the kernel with the reference implementation

Time & Space Complexity

The time complexity of the solution will depend on the size of the input tensor and the number of threads used to launch the kernel. In general, the time complexity will be O(n), where n is the size of the input tensor. The space complexity will also be O(n), as we need to store the input tensor, the output tensor, and any temporary data used by the kernel.

solution.py

Python 3.13

GPU · T4

Test results