There is a bug in this repository for the tolerance since it uses a tolerance of 1e-3. Not the claimed tolerance of 1e-4 (which is already lower than the standard 1e-5 in Sakana AI's CUDA agentic engineering work)
The tolerance was set to 1e-4 in the search portion of the code, but the main.py overrides it to 1e-3.
This means that the LLM often aggressively tries to use FP16 or fast math operations while sacrificing full precision, but reverts back to FP32 operations for some portions of the kernel like accumelation.
It makes sense that you could easily get 2x or even more speedup with such a low tolerance. Since it is effectively only half the precision that FP32 allows you. One could likely get 2x speedup just by converting all operations to FP16 with such a loose tolerance.
There is a bug in this repository for the tolerance since it uses a tolerance of
1e-3. Not the claimed tolerance of1e-4(which is already lower than the standard1e-5in Sakana AI's CUDA agentic engineering work)The tolerance was set to
1e-4in the search portion of the code, but themain.pyoverrides it to1e-3.This means that the LLM often aggressively tries to use FP16 or fast math operations while sacrificing full precision, but reverts back to FP32 operations for some portions of the kernel like accumelation.
It makes sense that you could easily get 2x or even more speedup with such a low tolerance. Since it is effectively only half the precision that FP32 allows you. One could likely get 2x speedup just by converting all operations to FP16 with such a loose tolerance.