[feat][JAX] Add Apple Silicon MPS support

.claude/docs/backends/jax.md

@@ -19,9 +19,17 @@ Uses Flax/NNX for model definitions.
 - `jax` extra for JAX dependencies.
 - `gpu` extra for JAX on GPU (CUDA).
 - `tpu` extra for JAX on TPU (uses custom index `jax-tpu`).
+- `mps` extra for JAX on Apple Silicon through `jax-mps` (macOS 14+).
 
 ## Tests
 
 ```bash
 uv run --extra dev --extra jax pytest tests/tx/ -v
 ```
+
+On Apple Silicon:
+
+```bash
+JAX_PLATFORMS=mps uv run --isolated --extra dev --extra jax --extra mps --extra tinker \
+  pytest tests/backends/test_jax_backend.py -v
+```

docs/content/docs/getting-started/supported_models.mdx

@@ -29,3 +29,51 @@ The JAX backend supports the following models:
 - Deepseek-V3
 - Llama-3
 - Qwen3 Dense/MoE
+
+### Apple Silicon (experimental)
+
+On macOS 14 or newer, the JAX backend can use Apple Silicon GPUs through the experimental
+[`jax-mps`](https://github.com/tillahoffmann/jax-mps) PJRT plugin. Install the
+JAX and MPS extras together and select the backend explicitly:
+
+```bash
+JAX_PLATFORMS=mps uv run --isolated --extra jax --extra mps --extra tinker \
+  -m skyrl.tinker.api \
+  --base-model Qwen/Qwen3-0.6B \
+  --backend jax \
+  --backend-config '{"max_lora_adapters": 2, "max_lora_rank": 8, "train_micro_batch_size": 1, "sample_max_num_sequences": 4, "gradient_checkpointing": true}'
+```
+
+This path supports the JAX backend's LoRA training and sampling API. It does
+not enable the FSDP or Megatron backends on macOS. `jax-mps` currently supports
+one Apple GPU and does not implement every JAX operation, so start with a small
+dense model and single-device parallelism (`fsdp=ep=tp=1`).
+
+To run one complete RL iteration, keep the server running and execute the
+[Tinker cookbook](https://github.com/thinking-machines-lab/tinker-cookbook) math
+recipe from a separate checkout:
+
+```bash
+TINKER_API_KEY=tml-dummy uv run --isolated --extra math-rl --with tinker==0.22.4 \
+  python -m tinker_cookbook.recipes.math_rl.train \
+  base_url=http://127.0.0.1:8000 \
+  model_name=Qwen/Qwen3-0.6B \
+  renderer_name=qwen3_disable_thinking \
+  lora_rank=8 \
+  env=gsm8k \
+  groups_per_batch=2 \
+  group_size=4 \
+  max_tokens=96 \
+  temperature=1.2 \
+  max_steps=1 \
+  learning_rate=4e-5 \
+  eval_every=0 \
+  save_every=0 \
+  behavior_if_log_dir_exists=delete
+```
+
+This exercises rollout sampling, GSM8K rewards, group-relative advantages,
+the importance-sampling loss, an optimizer update, sampler weight sync, and
+final checkpointing. On an M4 Max with 64 GB of unified memory, the validated
+iteration generated eight rollouts and completed in about 30 seconds after the
+server initialized.

pyproject.toml

@@ -35,6 +35,10 @@ tpu = [
     "jax[tpu]>=0.7.2; sys_platform == 'linux'",
 ]
 
+mps = [
+    "jax-mps>=0.10.10,<0.11; sys_platform == 'darwin' and platform_machine == 'arm64'",
+]
+
 tinker = [
     "tinker>=0.3.0",
     "fastapi[standard]",
@@ -234,9 +238,9 @@ no-build-isolation-package = [
 override-dependencies = [
     "nvidia-resiliency-ext; sys_platform == 'never'",
     "transformer-engine[pytorch]==2.11.0; sys_platform == 'linux'",
-    "transformers>=5.6.1,<=5.8.0; sys_platform == 'linux'",
+    "transformers>=5.6.1,<=5.8.0",
     "megatron-core>=0.16.0; sys_platform == 'linux' and python_version >= '3.12'",
-    "ml_dtypes>=0.5.0; sys_platform == 'linux'",
+    "ml_dtypes>=0.5.0",
     "transformer-engine-cu13; sys_platform == 'never'",
     # `nixl` hard-depends on both nixl-cu12 and nixl-cu13; drop the cu13 variant
     # so it doesn't pull the CUDA-13 stack and bump torch off the cu12 pin.

skyrl/backends/jax.py

@@ -976,17 +976,26 @@ def load_checkpoint(self, checkpoint_path: AnyPath, model_id: str) -> None:
         logger.info(f"Loaded training checkpoint from {checkpoint_path}")
 
     def save_sampler_checkpoint(self, output_path: AnyPath, model_id: str, persist: bool = True) -> None:
-        """Save sampler checkpoint as tar.gz using save_lora_checkpoint."""
+        """Make the current LoRA weights available to the sampler."""
         lora_model = self.models[model_id]
-        save_lora_checkpoint(
-            self.model,
-            self.base_model,
-            lora_model.lora_config,
-            lora_model.adapter_index,
-            output_path,
-            self.process_id,
-        )
-        logger.info(f"Saved LoRA sampler checkpoint to {output_path}")
+        if persist:
+            save_lora_checkpoint(
+                self.model,
+                self.base_model,
+                lora_model.lora_config,
+                lora_model.adapter_index,
+                output_path,
+                self.process_id,
+            )
+            logger.info(f"Saved LoRA sampler checkpoint to {output_path}")
+
+        # Training and sampling share one in-memory model in this backend. Marking
+        # these weights as loaded avoids a redundant archive round trip on the
+        # next sample request, and lets ephemeral RL syncs skip disk entirely.
+        checkpoint_id = output_path.name.removesuffix(".tar.gz")
+        lora_model.loaded_checkpoint_id = checkpoint_id
+        if not persist:
+            logger.info(f"Updated in-memory LoRA sampler weights for model {model_id}")
 
     def load_sampler_checkpoint(self, model_id: str, checkpoint_id: str, checkpoint_path: AnyPath) -> None:
         """Insert sampler weights from checkpoint file."""

skyrl/tinker/engine.py

@@ -650,7 +650,10 @@ def process_save_weights_for_sampler(
 
         with self._checkpoint_status_context(model_id, checkpoint_id, types.CheckpointType.SAMPLER):
             self.backend.save_sampler_checkpoint(output_path, model_id, persist=persist)
-            logger.info(f"Saved sampler checkpoint for model {model_id} to {output_path}")
+            if persist:
+                logger.info(f"Saved sampler checkpoint for model {model_id} to {output_path}")
+            else:
+                logger.info(f"Prepared ephemeral sampler weights for model {model_id}")
 
         # Return path=None when using sampling_session_seq_id and seq_id (SDK expects this)
         if request_data.sampling_session_seq_id is not None and request_data.seq_id is not None:

skyrl/tx/layers/attention.py

@@ -48,7 +48,14 @@ def dot_product_attention(
             implementation="cudnn",
         )
 
-    # CPU/TPU fallback
+    # jax-mps fused SDPA cannot combine causal attention with a padding mask.
+    implementation = "xla" if jax.default_backend() == "mps" and is_causal else None
     return jax.nn.dot_product_attention(
-        q, k, v, scale=scale, mask=attention_mask[:, None, None, :].astype(bool), is_causal=is_causal
+        q,
+        k,
+        v,
+        scale=scale,
+        mask=attention_mask[:, None, None, :].astype(bool),
+        is_causal=is_causal,
+        implementation=implementation,
     )

skyrl/tx/models/deepseekv3.py

@@ -171,13 +171,16 @@ def __call__(
         # Jax attention expects v to have the same shape as k
         v = jnp.pad(v, ((0, 0), (0, 0), (0, 0), (0, self.qk_head_dim - self.v_head_dim)))
 
+        is_causal = kv_cache is None
+        implementation = "xla" if jax.default_backend() == "mps" and is_causal else None
         attn_output = jax.nn.dot_product_attention(
             q,
             k,
             v,
             scale=self.scaling,
             mask=attention_mask[:, None, None, :].astype(bool),
-            is_causal=kv_cache is None,
+            is_causal=is_causal,
+            implementation=implementation,
         )
 
         attn_output = attn_output[:, :, :, : self.v_head_dim].reshape(B, T, self.num_heads * self.v_head_dim)

skyrl/tx/utils/generator.py

@@ -85,6 +85,20 @@ def update_layer(kv_cache, k, v, positions):
         """
         k_cache, v_cache = kv_cache
 
+        if jax.default_backend() == "mps":
+            # jax-mps does not yet support the batched scatter produced by
+            # vmap(dynamic_update_slice).
+            max_start = max(k_cache.shape[1] - k.shape[1], 0)
+            update_start = jnp.clip(positions[:, :1], 0, max_start)
+            update_positions = update_start + jnp.arange(k.shape[1])[None, :]
+            update_mask = jnp.arange(k_cache.shape[1])[None, :, None] == update_positions[:, None, :]
+
+            def update(cache, values):
+                updates = jnp.einsum("bts,bsnh->btnh", update_mask.astype(values.dtype), values)
+                return jnp.where(update_mask.any(axis=-1)[..., None, None], updates, cache)
+
+            return update(k_cache, k), update(v_cache, v)
+
         def update_at_pos(cache_slice, new_val_slice, pos):
             return jax.lax.dynamic_update_slice(cache_slice, new_val_slice, (pos, 0, 0))
 
@@ -274,8 +288,12 @@ def body_fn(s: DecodeState) -> DecodeState:
             stop_pos = jnp.where((s.stop_pos == -1) & is_stop, step + 1, s.stop_pos)
 
             # Update attention mask at per-sequence positions (for left-aligned sequences)
-            batch_idx = jnp.arange(s.attention_mask.shape[0])
-            next_attention_mask = s.attention_mask.at[batch_idx, s.kv_cache.cache_position].set(1)
+            if jax.default_backend() == "mps":
+                position_mask = jnp.arange(s.attention_mask.shape[1])[None, :] == s.kv_cache.cache_position[:, None]
+                next_attention_mask = jnp.where(position_mask, 1, s.attention_mask)
+            else:
+                batch_idx = jnp.arange(s.attention_mask.shape[0])
+                next_attention_mask = s.attention_mask.at[batch_idx, s.kv_cache.cache_position].set(1)
 
             outputs = model(
                 next_token,

tests/backends/test_jax_backend.py

@@ -390,6 +390,23 @@ def test_optim_step_returns_metrics():
     assert no_grad_output.metrics["skyrl.ai/mhc_gradient_norm"] == pytest.approx(0.0)
 
 
+def test_ephemeral_sampler_checkpoint_stays_in_memory(monkeypatch, tmp_path):
+    backend = create_backend(max_lora_adapters=2)
+    model_id = "ephemeral_sampler"
+    create_model(backend, model_id)
+    output_path = tmp_path / "ss0_seq1.tar.gz"
+
+    def fail_if_saved(*args, **kwargs):
+        pytest.fail("ephemeral sampler weights should not be written to disk")
+
+    monkeypatch.setattr("skyrl.backends.jax.save_lora_checkpoint", fail_if_saved)
+
+    backend.save_sampler_checkpoint(output_path, model_id, persist=False)
+
+    assert backend.models[model_id].loaded_checkpoint_id == "ss0_seq1"
+    assert not output_path.exists()
+
+
 def test_gradient_checkpointing():
     """
     Verify gradient checkpointing doesn't affect loss values.

tests/tx/utils/test_generator.py

@@ -1,5 +1,6 @@
 from unittest.mock import MagicMock
 
+import jax
 import jax.numpy as jnp
 from flax import nnx
 
@@ -77,6 +78,22 @@ def make_inputs(batch_size: int, prompt_length: int):
     return input_ids, attention_mask
 
 
+def test_mps_kv_cache_update_layer(monkeypatch):
+    monkeypatch.setattr(jax, "default_backend", lambda: "mps")
+    k_cache = jnp.zeros((2, 5, 1, 2), dtype=jnp.float32)
+    v_cache = jnp.zeros_like(k_cache)
+    k = jnp.array([[[[1.0, 2.0]], [[3.0, 4.0]]], [[[5.0, 6.0]], [[7.0, 8.0]]]])
+    v = k + 10
+    positions = jnp.array([[1, 2], [4, 5]], dtype=jnp.int32)
+
+    updated_k, updated_v = KVCache.update_layer((k_cache, v_cache), k, v, positions)
+
+    expected_k = k_cache.at[0, 1:3].set(k[0]).at[1, 3:5].set(k[1])
+    expected_v = v_cache.at[0, 1:3].set(v[0]).at[1, 3:5].set(v[1])
+    assert jnp.array_equal(updated_k, expected_k)
+    assert jnp.array_equal(updated_v, expected_v)
+
+
 def generator_outputs_equal(output1: GenerateOutput, index1: int, output2: GenerateOutput, index2: int) -> bool:
     """Check if two GenerateOutput objects are equal at the given indices."""
     return (