[feat][JAX] Add Apple Silicon MPS support

[bvolpato]

Summary

Add opt-in Apple Silicon support for SkyRL's JAX backend using jax-mps.

Changes

• Add an mps dependency extra for Apple Silicon macOS.
• Restore dependency resolution for the current JAX and Transformers stack.
• Route causal attention through XLA when the jax-mps SDPA patch is active.
• Add MPS-compatible fallbacks for attention-mask and KV-cache updates.
• Preserve dynamic_update_slice boundary behavior in the KV-cache fallback.
• Document setup, supported workflows, and current limitations.
• Add focused backend, generator, attention, and cache-boundary coverage.

Validation

Current Linux validation after rebasing onto main:

• JAX backend (19 passed)
• Generator and attention regressions (8 passed)
• DeepSeek focused suite
• Ruff, Black, Gitleaks, and focused type checks
• Documentation production build (66 pages)
• uv lock --check and git diff --check

Earlier branch validation on an M4 Max covered the MPS JAX backend, DeepSeek-V3 against its Hugging Face reference, and Qwen3-0.6B forward, backward, optimizer update, and sampling. Those hardware results predate the latest rebase and fallback fixes; current fixes have CPU regression coverage but have not been rerun on Apple hardware.

Limitations

• jax-mps remains experimental and supports one Apple GPU.
• Some JAX operations remain unsupported by the plugin.
• Causal attention uses the XLA implementation when the jax-mps patch is active.
• KV-cache updates use a dense einsum fallback until batched scatter support lands upstream.
• Buffer-donation warnings may still appear.

[gemini-code-assist]

Code Review

This pull request introduces experimental support for Apple Silicon (macOS 14+) in the JAX backend using the jax-mps PJRT plugin. It includes documentation updates, dependency adjustments in pyproject.toml, and custom workarounds in the attention layers and KV cache update logic to bypass current jax-mps limitations. Additionally, it implements ephemeral in-memory sampler checkpoints to avoid redundant disk writes. Feedback on the changes highlights a behavioral discrepancy in the custom MPS KV cache update logic, where negative start indices are incorrectly wrapped instead of clamped to zero as standard JAX dynamic_update_slice does. Suggestions have been provided to simplify this logic and update the corresponding unit tests.