It is an area where I am still learning, but as far as I know, the badness of things is relative. CPU-GPU communication is the major thing to avoid. In my CUDA

It is an area where I am still learning, but as far as I know, the badness of things is relative. CPU-GPU communication is the major thing to avoid. In my CUDA llama3.c fork, making changes to reduce CPU-GPU communication generally increases performance. That means combining kernels and reducing cudaMemcpy and cudaMalloc calls. CUDA graphs likely could give further reductions, but I have yet to run out of ideas for more normal methods for getting them. To give a concrete example of CPU-GPU communication in a tight loop being bad (much like what you were planning to do), when I initially moved softmax to the GPU, I had been invoking the kernel in a host side loop (since attention calls softmax multiple times per pass), and it slowed down inference by a few tokens per second versus the CPU softmax. I eliminated the host side loop by using the grid dimension so that 1 CUDA call did all of the loop iterations implicitly and inference became a few tokens per second faster than the CPU softmax. As for branches, avoiding them on the GPU where possible can be beneficial, as long as you are not doing something more expensive like CPU-GPU communication in their place. For example, in my softmax kernel, rather than do a branch to calculate the maximum value inside my loop, I use fmaxf(). This is better since there is no branch, and fmaxf is incredibly cheap since it is a single PTX instruction (which presumably maps to a single SASS intruction). GPGPU grew out of work for having programmable shaders on GPUs and a big part of that was branching. As long as you are minimizing branch divergence and only using necessary branches, you are fine as far as branching on a GPU are concerned.