Supported sharded GPU simulation with mjwarp

### The feature, motivation and pitch

It would be very useful to support multi-GPU mjwarp simulations.

### Problem

Running **mjwarp simulations** across multiple GPUs is currently difficult due to limitations in sharding and device handling.

### Challenges
1. **Device-side model/data creation**: The current API is not compatible with replicated or data-parallel sharding.
2. **Sharding awareness in JAX FFI integration**: The `mjx/third_party/warp/jax_experimental/ffi.py` file (and other code locations) assumes the first visible device instead of handling different devices properly.
3. **Collision data layout**: mjwarp stores collision data in a flattened layout, making partitioning and sharding more complex.

### Alternatives

_No response_

### Additional context

_No response_