Question Regarding the Changing Duration of CPU Calls to a Compiled TPU Kernel

[yutation]

Hi,

I wrote code where the CPU repeatedly invokes the same compiled TPU kernel at a high frequency. Initially, I observed that the CPU invocation duration is significantly shorter than the TPU kernel execution time, indicating asynchronous behavior. However, after approximately 30 iterations, the CPU invocation duration increases and roughly equals the TPU kernel execution time, suggesting a shift to synchronous behavior.

The relevant code and corresponding trace are shown below:

def function_to_test_2( iterations: int,  state_matrix: jnp.ndarray) -> jnp.ndarray:
    for i in range(iterations):
        state_matrix = tpu_kernel_jit(state_matrix, state_matrix)
    return state_matrix

This behavior leads me to suspect that there may be a command buffer limit or queue depth restriction for pending TPU operations, which eventually causes the CPU to block and wait for TPU execution to catch up.
Could Someone Explain This Behavior?

Thanks,
Jingtian