Implement a fused multiply-add (FMA) kernel in Triton. Given a scalar a and two 1D float tensors x and y:
out[i]=a⋅x[i]+y[i]Like vector addition this is a pure map, but it introduces two new beats: a scalar parameter passed by value alongside the tensor pointers, and a fused two-input arithmetic expression that the Triton backend lowers to a single hardware FMA instruction per lane.
Input: a = 2.0, x = [1.0, 2.0, 3.0, 4.0], y = [1.0, 1.0, 1.0, 1.0]
Output: [3.0, 5.0, 7.0, 9.0]
Input: a = -1.0, x = [1.0, 2.0, 3.0], y = [10.0, 10.0, 10.0]
Output: [9.0, 8.0, 7.0]
Forward the scalar a from the launcher straight into the kernel call. Triton broadcasts a Python float to every lane automatically.
Write the body as a single expression a⋅x+y. Splitting it into a separate multiply and add gives the compiler more freedom but loses the explicit FMA opportunity.
Tail mask discipline is identical to vector add: build one mask at the top, reuse it on both loads and the store.
torch.float32 on cuda.atol=1e-3.Accepts: number
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement a fused multiply-add (FMA) kernel in Triton. Given a scalar a and two 1D float tensors x and y:
out[i]=a⋅x[i]+y[i]Like vector addition this is a pure map, but it introduces two new beats: a scalar parameter passed by value alongside the tensor pointers, and a fused two-input arithmetic expression that the Triton backend lowers to a single hardware FMA instruction per lane.
Input: a = 2.0, x = [1.0, 2.0, 3.0, 4.0], y = [1.0, 1.0, 1.0, 1.0]
Output: [3.0, 5.0, 7.0, 9.0]
Input: a = -1.0, x = [1.0, 2.0, 3.0], y = [10.0, 10.0, 10.0]
Output: [9.0, 8.0, 7.0]
Forward the scalar a from the launcher straight into the kernel call. Triton broadcasts a Python float to every lane automatically.
Write the body as a single expression a⋅x+y. Splitting it into a separate multiply and add gives the compiler more freedom but loses the explicit FMA opportunity.
Tail mask discipline is identical to vector add: build one mask at the top, reuse it on both loads and the store.
torch.float32 on cuda.atol=1e-3.Accepts: number
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.