Comment by Alexander-Barth

Julia is actually quite nice for this. If you prefer a python-like approach consider Triton from openai, numba (https://numba.pydata.org/) or CuTe DSL from Nvidia.