jz-tree offers a framweork for GPU-friendly implementations of tree algorithms in jax (with a CUDA backend). Currently nearest neighbour search and friends-of-friends are supported and deliver top performance.
For installation instructions, please check the documentation!