Triton Fused Multiply-Add Kernel
Problem Statement
Implement a fused multiply-add kernel: out = a * x + b * y for two 1D tensors and two runtime scalars a, b.
Background
Fusing multiple elementwise operations into one kernel means each element is read and written exactly once, saving memory bandwidth versus separate multiply and add passes.
Your Task
Implement fma_kernel and run(n=1024, a=2.0, b=-1.0).
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, a = 2.0, b = -1.0
True
- The
runfunction allocates two 1D tensorsxandyof sizen=1024on the GPU. - It then launches the
fma_kernelwith inputsx,y, and scalarsa=2.0,b=-1.0, computing the output tensoroutas out=a⋅x+b⋅y=2.0⋅x−1.0⋅y. - The result is compared to a reference output computed using PyTorch's built-in functions,
torch_reference = 2.0 * x - 1.0 * y. - The function returns
Trueif the two outputs are close, as determined bytorch.allclose(triton_out, torch_reference), indicating that the Triton kernel produced the correct result.
Constraints:
- Single kernel computing ax + by
- Two tensor pointers, two scalar args
- Mask the tail block
Background Knowledge
The problem involves implementing a fused multiply-add kernel, which is a fundamental operation in linear algebra and deep learning. This operation combines two element-wise operations, multiplication and addition, into a single step. The goal is to compute out = a * x + b * y for two 1D tensors x and y, and two runtime scalars a and b. Understanding the concept of element-wise operations is crucial, as it involves performing operations on corresponding elements of two or more tensors.
In the context of GPU programming, fusing multiple operations into one kernel can significantly improve performance by reducing memory bandwidth usage. This is because each element is read and written exactly once, minimizing the number of memory accesses. The Triton programming language is designed to efficiently execute such operations on GPUs, providing a high-level interface for writing custom kernels. Familiarity with PyTorch and its tensor operations is also essential, as the problem involves comparing the output of the custom kernel with a PyTorch reference implementation.
To tackle this problem, it's essential to understand the basics of GPU architecture and how to optimize kernel launches for maximum performance. This includes understanding thread blocks, warps, and memory coalescing. Additionally, knowledge of numerical stability and floating-point arithmetic is necessary to ensure accurate results, especially when dealing with large tensors and scalars.
Algorithm/Approach
The general approach to solving this problem involves:
- Defining a custom Triton kernel that performs the fused multiply-add operation
- Launching the kernel on the GPU using Triton's kernel launch API
- Allocating input tensors x and y on the GPU and passing them to the kernel
- Comparing the output of the custom kernel with a PyTorch reference implementation 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.