How to determine the cause of MemcpyD2H transfers in GPU profiler trace

[otoomey]

I am currently trying to diagnose performance issues with our simulator. I generated a profiler trace from the simulator's execution, and I can see that the GPU spends 36% of the time sitting idle:

After checking the trace log, from what I understand the inner loop of our simulator gets stuck waiting for MemcpyD2H transfers to complete:

Do you have any advice on how I can attempt to locate the source of these transfers? Shown in the screenshot is just one iteration out of several hundred that our simulator executes. The whole thing is jit'd. We don't specifically specify which device JAX should use on any of our array instantiations - and it looks the core bits of our simulation are actually running on the GPU as they appear in the trace log under the GPU device. Thank you very much for your help.