[feat][JAX] Add Apple Silicon MPS support

[bvolpato]

Summary

Add opt-in Apple Silicon GPU support for the JAX backend through the open-source
jax-mps PJRT plugin.

Apple's closed-source jax-metal plugin cannot consume StableHLO from SkyRL's
current JAX 0.10.x stack. jax-mps targets JAX/JAXLIB 0.10.x and executes JAX
programs through MLX on Metal. This refreshes the approach from #1332 against
current main and jax-mps 0.10.10.

Changes

• add an explicit mps extra for Apple Silicon on macOS 14+
• restore transformers and ml_dtypes resolution on macOS
• use XLA attention when causal MPS attention also has a padding mask
• avoid unsupported batched scatter operations during MPS KV-cache and mask updates
• cover the MPS KV-cache fallback, including boundary clamping
• document installation, limitations, and a local Tinker server command

Testing

Validated on an M4 Max MacBook Pro with 64 GB unified memory.

• MPS JAX backend suite: 18 passed
• CPU JAX backend suite: 18 passed
• Generator suite: 7 passed
• DeepSeek-V3 tiny model on MPS: passed against the Hugging Face reference
• Qwen3-0.6B on MPS: forward/backward, optimizer step, and post-update sampling passed
  • gradient norm: 76.5
  • parameter delta norm: 0.0587403
• pre-commit: ruff, black, and gitleaks passed
• documentation production build passed
• uv lock --check passed

JAX_PLATFORMS=mps uv run --isolated \
  --extra dev --extra jax --extra mps --extra tinker \
  pytest tests/backends/test_jax_backend.py -v

Repro Command

JAX_PLATFORMS=mps uv run --isolated --extra jax --extra mps --extra tinker \
  -m skyrl.tinker.api \
  --base-model trl-internal-testing/tiny-Qwen3ForCausalLM \
  --backend jax \
  --backend-config '{"max_lora_adapters": 2, "max_lora_rank": 8, "train_micro_batch_size": 1, "sample_max_num_sequences": 1}'

Downsides

jax-mps is experimental, supports one Apple GPU, does not support every JAX
operation, and currently emits warnings because buffer donation is unavailable.
Causal attention uses the XLA implementation instead of fused SDPA on MPS.

Until tillahoffmann/jax-mps#223 is released, MPS decoding uses a dense mask/einsum KV-cache
update instead of batched dynamic_update_slice; this can increase decode memory
traffic for long contexts. CPU, CUDA, and TPU paths retain their existing
implementations.