My question is if the while-loop logic can be run on the GPU to save time. Let’s say the latency of communication between CPU and GPU is 1ms and we output 100 tokens. The GPU runs the model, and output logits. The CPU samples from the logits with some strategy, e.g. top-k and maybe also some kind of beam search etc, and decides on the output token for this step. Then it takes the token and send it back to the GPU for the next timestamp prediction. If we output 100 tokens, only the communication of this whole process takes 100ms. If this part of the logic was on the GPU, we would have saved this communication time. I hope I explained myself correctly.