Slowdown of independent calculations on different devices

[BitPhinix]

I'm trying to run multiple independent calculations on different tpus connected to the same host, for this I have a python script running multiple worker threads, each having one device assigned to it.

I'm being very cautious to ensure all calculation inside the worker is done on the respective device by moving all data to that device before calling any jax functions.

Even with this level of isolation it seems like the independent calculations are heavily impacting calculations on other devices and jax is only running calculation simultaneously (task of worker one typically takes 3 sec, but if another task is running it takes 3 + time of the other task)

This doesn't seem to be happening if I run the same workload batches across the same devices using e.g. pmap.

Is there any inherent reason for it? Is jax using some sort of global lock preventing independent simultaneous calculations on different devices?