LaelaZorana/spmv-csr-triton-gpu-kernels

spmv-csr-triton

Three Triton kernels for the same sparse matrix vector product, and a selection rule that comes from measurement rather than from the textbook.

The headline is that the textbook answer loses. The load balanced design, which splits the nonzeros into equal chunks and adds into the output with atomics, is what every guide recommends for an uneven matrix. It won one of four measured runs, and on an even matrix it was 7.7 times slower than a plain kernel that gives one program to each row.

The three kernels

All three are correct on the same input and differ only in how the work reaches the hardware.

Kernel	How the work is handed out	Where it hurts
thread_per_row	one lane walks one row	the block waits for its longest row
vector_per_row	one program walks one row in chunks	a short row still costs a whole program
nnz_balanced	equal nonzero chunks per program, atomics into the output	every write contends

What was measured

1,000,000 rows and 32,000,000 nonzeros in every run. Two matrices with the same nonzero count. Uniform is exactly 32 per row. Power law has a median row of 9 or 10 and a handful of rows above 100,000.

Microseconds, median of 20 runs after 5 warmup runs, host timed with a device synchronise before every stop. Every kernel passed allclose against a torch reference at atol 1e-2 and rtol 1e-3 on every run.

Nvidia RTX 4090, torch 2.8.0+cu128, triton 3.4.0

Matrix	Longest row	thread_per_row	vector_per_row	nnz_balanced	Winner
uniform 32	32	1095.7	743.7	1040.7	vector_per_row
power law	1,135,318	98,619.1	1508.0	4288.6	vector_per_row

AMD Instinct MI300X, torch 2.9.1+rocm6.3, triton 3.5.1

Matrix	Longest row	thread_per_row	vector_per_row	nnz_balanced	Winner
uniform 32	32	2166.5	561.3	4339.1	vector_per_row
power law	1,021,722	227,434.7	5068.8	4418.4	nnz_balanced

The exact device strings, the torch and triton versions and the raw timings are in evidence/.

Reading the table

The row length decides the winner, and it decides it by a lot. Give thread_per_row an even matrix and it is within 1.5 times of the best. Give it a power law matrix and it takes 98.6 ms on one card and 227.4 ms on the other, against 1.5 ms and 4.4 ms, so the same kernel on the same nonzero count is 65 and 45 times slower once the row lengths change.

The balanced design is not a safe default. It won once, on the heavy tail on the AMD card, and it won by 13 percent. On the even matrix on that same card it was 4339.1 against 561.3, which is 7.7 times slower, because every one of its 32,000,000 writes is an atomic and the rows were never the problem.

vector_per_row won three of four. That is why pick() returns it unless the matrix is heavy tailed and the card is an AMD one.

Using it

from kernels import get_kernel

spmv = get_kernel("LaelaZorana/spmv-csr-triton")

y = spmv.spmv(values, col_idx, row_ptr, x)            # auto, uses the measured rule
y = spmv.vector_per_row(values, col_idx, row_ptr, x)  # or name one yourself

The auto path reads the longest row against the mean row and the device name, then returns the kernel the measurements chose. Call pick() on its own to see the name without running anything.

For nnz_balanced the kernel walks nonzeros rather than rows, so it needs one row index per nonzero. Build it once with row_ids_from_row_ptr and keep it, because rebuilding it every call costs more than the kernel saves.

What is published here

The runs that chose the default are in evidence/, on both vendors. Both the winning and the losing numbers are there, including the run where the design this repository recommends against comes out fastest.

License

Apache 2.0, see the licence file. The notice file carries the attribution that section 4 asks a derivative work to keep, which is the name and the two cards the numbers were measured on.