Auto-research with codex: How I achieved a 232x Faster Kernel

Auto-research with Codex: Achieving a 232x Faster Kernel

The recent auto-research themed contest hosted by GPU Mode in collaboration with Core Automation challenged participants to implement batched square compact-Householder QR factorization, also known as QR decomposition. Placing 12th out of 183 competitors, the author achieved a 232x speedup over the baseline solution. This post aims to detail the approach, learnings, and obstacles encountered during the contest, which marked the author's first serious attempt at auto-research.

During the 14-day contest, the author made over 1500 submissions, showcasing their dedication to improving the kernel's performance. While the author had some background in GPU kernel optimization, they were an underdog compared to more experienced participants like the CUDA Colonel, a principal engineer at NVIDIA. However, armed with the basics and recent knowledge of GatedDeltaNet, the author was well-equipped to tackle the challenge.

The problem statement required the implementation of batched square FP32 CUDA matrices A with shape batch x n x n, returning the same compact Householder QR representation as torch.geqrf(A). This involved returning an H matrix whose upper triangle is R and whose lower triangle stores Householder vectors, along with a tau vector of reflector coefficients.

Correct submissions were evaluated based on runtime performance across various matrix sizes and conditioning cases, with the leaderboard ranking them using the geometric mean.

The key requirements for the contest included producing an orthogonal matrix Q, with columns that are unit-length and perpendicular to each other, and an upper triangular matrix R. The checker rebuilt Q using torch.linalg.householder_product(H, tau), took R = triu(H), and verified A ≈ QR, Q⊤Q ≈ I, and Q⊤A ≈ R. While the contest did not explicitly ask for the dense Q and R, the compact Householder version was necessary to reconstruct Q and read R from the upper triangle.

The compact Householder version proved to be crucial for the contest, as it allowed the checker to reconstruct Q and read R from the upper triangle. The author's journey through the contest involved learning about QR decomposition, its various methods, and implementing a blocked Householder algorithm with the trailing WY-update as the main architecture. This approach proved to be effective in achieving the desired speedup.