Triton Fused Multiply-Add Kernel

mediumPROTriton FundamentalsTriton Programming

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

Input

n = 1024, a = 2.0, b = -1.0

Output

True

Reasoning

Inputs: n = 1024 is the length of two random vectors $x$ and $y$ ; a = 2.0 and b = -1.0 are plain scalars that weight them.
Math: every position $i$ is independent, $out_i = a\,x_i + b\,y_i$ . With the sample scalars that is $out_i = 2x_i - y_i$ , so if $x_i = 0.5$ and $y_i = 1.5$ then $out_i = 1.0 - 1.5 = -0.5$ .
Fusion: each $x_i$ and $y_i$ is read once and each $out_i$ written once, with no intermediate vectors for $a\,x$ or $b\,y$ .
Tiling: if each program instance handles a block of $B$ elements, you need $\lceil n/B \rceil$ blocks. For $n = 1500$ and $B = 1024$ that is 2 blocks, and the second has only $1500 - 1024 = 476$ real elements; the other 548 positions must not be read or written.
Output: True means your $out$ matches $a\,x + b\,y$ to within an absolute tolerance of $10^{-5}$ . Notice the scalars can be negative or fractional, so they are runtime values, not constants.

Constraints

Single kernel computing ax + by
Two tensor pointers, two scalar args
Mask the tail block

Learn this in

Memory Access, Autotuning & BenchmarkingGPU Basics with Triton Kernel Fusion: Fused SoftmaxGPU 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 primerPREVIEW

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

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