Implement a matrix-vector product (GEMV) kernel in Triton. Given A of shape (M,N) and x of shape (N,), compute out of shape (M,) as a row-by-vector dot product:
out[i]=j=0∑N−1A[i,j]⋅x[j]This is a row-block reduction: each program owns a contiguous slab of output rows and walks the N dimension in chunks, accumulating the per-row dot products in an fp32 register vector. The pattern is the simplest non-trivial inner reduction in Triton.
Input: A = [[1, 2], [3, 4]], x = [5, 6]
Output: [17, 39]
Input: A = [[1, 0, 1], [0, 1, 0], [1, 1, 1]], x = [2, 3, 4]
Output: [6, 3, 9]
Use a 1D grid over the row blocks. Inside each program the row indices are a vector and the column chunk is another vector.
Initialize an fp32 accumulator of shape (BLOCK_M,) and loop along N, broadcasting the x chunk against the row tile before summing along the column axis.
On the final store, only the M tail can be partial. Reuse the row mask you built once at the top of the kernel.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥128 to absorb reduction drift.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement a matrix-vector product (GEMV) kernel in Triton. Given A of shape (M,N) and x of shape (N,), compute out of shape (M,) as a row-by-vector dot product:
out[i]=j=0∑N−1A[i,j]⋅x[j]This is a row-block reduction: each program owns a contiguous slab of output rows and walks the N dimension in chunks, accumulating the per-row dot products in an fp32 register vector. The pattern is the simplest non-trivial inner reduction in Triton.
Input: A = [[1, 2], [3, 4]], x = [5, 6]
Output: [17, 39]
Input: A = [[1, 0, 1], [0, 1, 0], [1, 1, 1]], x = [2, 3, 4]
Output: [6, 3, 9]
Use a 1D grid over the row blocks. Inside each program the row indices are a vector and the column chunk is another vector.
Initialize an fp32 accumulator of shape (BLOCK_M,) and loop along N, broadcasting the x chunk against the row tile before summing along the column axis.
On the final store, only the M tail can be partial. Reuse the row mask you built once at the top of the kernel.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥128 to absorb reduction drift.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.