From 57bfed7515ff1171bc3abe9f6f52c8c99abb8225 Mon Sep 17 00:00:00 2001
From: Dian Ang <23232359+yapdianang@users.noreply.github.com>
Date: Mon, 17 Aug 2026 19:04:36 +0000
Subject: [PATCH 1/5] perf(tinker): prune and order native GSPO updates

---
 .../workers/megatron/megatron_worker.py       |  12 ++-
 .../skyrl_train/workers/worker_utils.py       |  12 +++
 skyrl/backends/skyrl_train_backend.py         |  76 ++++++++++++-
 skyrl/tinker/api.py                           |  25 ++++-
 skyrl/tinker/engine.py                        |  65 ++++++++++-
 .../test_token_based_batching_utils.py        |  13 +++
 .../skyrl_train/test_loss_normalization.py    |  81 ++++++++++++++
 tests/tinker/test_api_validation.py           |  37 +++++++
 tests/tinker/test_engine.py                   | 102 ++++++++++++++++++
 9 files changed, 415 insertions(+), 8 deletions(-)

diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
index d116213abd..f29a8b3907 100644
--- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
+++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
@@ -1254,11 +1254,10 @@ def forward_backward(
             torch.cuda.empty_cache()
 
         # Aggregate metrics across micro-batches
-        all_loss_fn_outputs = []  # Handle separately from scalar metrics
+        loss_fn_output_batches = []
         for m_batch, metrics in zip(micro_buffer, metrics_list):
             # Extract loss_fn_outputs before reduce_metrics (it's not a scalar metric)
-            if "loss_fn_outputs" in metrics:
-                all_loss_fn_outputs.extend(metrics.pop("loss_fn_outputs"))
+            loss_fn_output_batches.append(metrics.pop("loss_fn_outputs", []))
             # Skip fully-padding microbatches: their metrics (clip_ratio=0, policy_entropy=0,
             # ...) are meaningless and would drag down the mean-reduced metrics. Summed
             # metrics (e.g. policy_loss) are unaffected since padding contributes 0, but
@@ -1303,6 +1302,13 @@ def forward_backward(
                 for k, v in moe_metrics.items():
                     status[k] = v
 
+        if not any(loss_fn_output_batches):
+            all_loss_fn_outputs = []
+        elif isinstance(microbatch_iterator, TokenBasedBatchIterator):
+            all_loss_fn_outputs = microbatch_iterator.reorder_and_combine_items(loss_fn_output_batches)
+        else:
+            all_loss_fn_outputs = [item for batch in loss_fn_output_batches for item in batch]
+
         return WorkerOutput(loss_fn_outputs=all_loss_fn_outputs, metrics=status)
 
     def optim_step(self) -> Optional[float]:
diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py
index 2efa37ad6c..71fb9518dc 100644
--- a/skyrl/backends/skyrl_train/workers/worker_utils.py
+++ b/skyrl/backends/skyrl_train/workers/worker_utils.py
@@ -424,6 +424,18 @@ def reorder_and_combine_batches(self, batches: List[TensorBatch]) -> TensorBatch
         reordered_batch.metadata = ref_microbatch.metadata
         return reordered_batch
 
+    def reorder_and_combine_items(self, batches: List[List[dict]]) -> List[dict]:
+        """Restore per-sample microbatch outputs to input order."""
+        ordered = [None] * self.data.batch_size
+        for original_indices, items in zip(self._microbatches, batches):
+            if len(items) < len(original_indices):
+                raise ValueError("Microbatch output has fewer items than input samples")
+            for original_idx, item in zip(original_indices, items):
+                ordered[original_idx] = item
+        if any(item is None for item in ordered):
+            raise ValueError("Microbatch outputs do not cover every input sample")
+        return ordered
+
 
 def get_microbatch_iterator(
     data: TrainingInputBatch, micro_batch_size: int, max_tokens_per_microbatch: int
diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py
index ea08b13d01..e40be0c9ff 100644
--- a/skyrl/backends/skyrl_train_backend.py
+++ b/skyrl/backends/skyrl_train_backend.py
@@ -808,6 +808,21 @@ def _normalize_policy_loss_request(
         if role == "critic":
             return loss_fn, loss_fn_config
 
+        if loss_fn == "gspo":
+            normalized_config = dict(loss_fn_config or {})
+            clip_low_threshold = normalized_config.pop("clip_low_threshold", None)
+            clip_high_threshold = normalized_config.pop("clip_high_threshold", None)
+            if clip_low_threshold is None or clip_high_threshold is None:
+                raise ValueError("loss_fn='gspo' requires clip_low_threshold and clip_high_threshold")
+            normalized_config.update(
+                eps_clip_low=1.0 - clip_low_threshold,
+                eps_clip_high=clip_high_threshold - 1.0,
+                loss_reduction="sequence_mean",
+                use_entropy_loss=False,
+                use_kl_loss=False,
+            )
+            return loss_fn, normalized_config
+
         if loss_fn == "dppo":
             # DPPO thresholds live in the nested `algorithm.dppo` sub-config, but
             # Tinker's loss_fn_config is a flat float dict. Re-nest so the
@@ -832,6 +847,47 @@ def _normalize_policy_loss_request(
             normalized_config["eps_clip_high"] = clip_high_threshold - 1.0
         return ("regular" if loss_fn == "ppo" else "gspo"), normalized_config or None
 
+    @staticmethod
+    def _normalize_gspo_batch(batch: TrainingInputBatch) -> None:
+        """Apply Tinker's sequence-mean GSPO reduction before DP sharding."""
+        advantages = batch["advantages"]
+        loss_mask = advantages.ne(0).to(batch["loss_mask"].dtype)
+        token_counts = loss_mask.sum(dim=-1, keepdim=True)
+        trainable_sequences = token_counts.squeeze(-1).gt(0)
+        num_trainable_sequences = trainable_sequences.sum().clamp(min=1)
+        batch["advantages"] = advantages / token_counts.clamp(min=1) / num_trainable_sequences
+        batch["loss_mask"] = loss_mask
+
+    @staticmethod
+    def _select_trainable_gspo_rows(batch: TrainingInputBatch) -> tuple[TrainingInputBatch, list[int]]:
+        """Drop inactive rows while preserving all-zero-batch optimizer semantics."""
+        trainable_rows = batch["loss_mask"].ne(0).any(dim=-1)
+        trainable_indices = trainable_rows.nonzero(as_tuple=False).flatten()
+        if trainable_indices.numel() == 0 or trainable_indices.numel() == batch.batch_size:
+            return batch, list(range(batch.batch_size))
+
+        selected = TrainingInputBatch(
+            {key: None if value is None else value[trainable_indices] for key, value in batch.items()}
+        )
+        selected.metadata = batch.metadata
+        return selected, trainable_indices.tolist()
+
+    @staticmethod
+    def _get_gspo_pruning_metrics(
+        original_token_counts: list[int], submitted_indices: set[int], start_idx: int, end_idx: int
+    ) -> dict[str, float]:
+        """Report one SDK chunk's pruning counts so combined futures sum to the batch total."""
+        request_indices = range(start_idx, end_idx)
+        request_submitted_indices = [index for index in request_indices if index in submitted_indices]
+        submitted_tokens = sum(original_token_counts[index] for index in request_submitted_indices)
+        original_tokens = sum(original_token_counts[index] for index in request_indices)
+        return {
+            "skyrl.ai/submitted_datums:sum": float(len(request_submitted_indices)),
+            "skyrl.ai/skipped_datums:sum": float(end_idx - start_idx - len(request_submitted_indices)),
+            "skyrl.ai/submitted_tokens:sum": float(submitted_tokens),
+            "skyrl.ai/skipped_tokens:sum": float(original_tokens - submitted_tokens),
+        }
+
     def forward_backward(
         self,
         prepared_batch: types.PreparedModelPassBatch,
@@ -860,6 +916,12 @@ def _forward_backward_single_model_batch(
         ):
             raise ValueError("Critic forward_backward requires values and returns for every response token")
         batch = self._to_training_batch(prepared_batch, role)
+        original_batch_size = batch.batch_size
+        original_token_counts = [int(count) for count in batch["attention_mask"].sum(dim=-1).tolist()]
+        submitted_indices = list(range(original_batch_size))
+        if loss_fn == "gspo":
+            self._normalize_gspo_batch(batch)
+            batch, submitted_indices = self._select_trainable_gspo_rows(batch)
         micro_bs = (
             self._cfg.trainer.micro_train_batch_size_per_gpu if self._cfg.trainer.strategy == "megatron" else None
         )
@@ -897,7 +959,14 @@ def _forward_backward_single_model_batch(
         if pad_size > 0 and per_sample_outputs:
             per_sample_outputs = per_sample_outputs[:-pad_size]
 
+        if len(submitted_indices) != original_batch_size:
+            restored_outputs = [{"logprobs": list(logprobs)} for logprobs in prepared_batch.all_sampling_logprobs]
+            for index, output in zip(submitted_indices, per_sample_outputs, strict=True):
+                restored_outputs[index] = output
+            per_sample_outputs = restored_outputs
+
         metrics = self._extract_metrics(data.metrics)
+        submitted_index_set = set(submitted_indices)
 
         results = {}
         for request_id, _, start_idx, end_idx in prepared_batch.request_batch_slices:
@@ -917,10 +986,15 @@ def _forward_backward_single_model_batch(
                     loss_fn_outputs.append(formatted_output)
             else:
                 loss_fn_outputs = [{} for _ in range(end_idx - start_idx)]
+            request_metrics = dict(metrics)
+            if loss_fn == "gspo":
+                request_metrics.update(
+                    self._get_gspo_pruning_metrics(original_token_counts, submitted_index_set, start_idx, end_idx)
+                )
             results[request_id] = types.ForwardBackwardOutput(
                 loss_fn_output_type=data.loss_fn_output_type,
                 loss_fn_outputs=loss_fn_outputs,
-                metrics=metrics,
+                metrics=request_metrics,
             )
         return results
 
diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py
index e63fcbeaa4..2f88b23560 100644
--- a/skyrl/tinker/api.py
+++ b/skyrl/tinker/api.py
@@ -625,6 +625,8 @@ class ForwardBackwardInput(BaseModel):
     def validate_loss_fn_config_keys(self):
         """Validate loss_fn_config keys based on the selected loss function."""
         if self.loss_fn_config is None:
+            if self.loss_fn == "gspo":
+                raise ValueError("loss_fn='gspo' requires clip_low_threshold and clip_high_threshold.")
             return self
 
         allowed_keys = self._ALLOWED_KEYS_BY_LOSS_FN[self.loss_fn]
@@ -638,6 +640,18 @@ def validate_loss_fn_config_keys(self):
             raise ValueError(
                 f"loss_fn='{self.loss_fn}' does not accept loss_fn_config keys. " f"Received: {invalid_keys}."
             )
+        if self.loss_fn == "gspo":
+            required_keys = {"clip_low_threshold", "clip_high_threshold"}
+            missing_keys = sorted(required_keys - self.loss_fn_config.keys())
+            if missing_keys:
+                raise ValueError(f"loss_fn='gspo' is missing required loss_fn_config keys: {missing_keys}.")
+            clip_low = self.loss_fn_config["clip_low_threshold"]
+            clip_high = self.loss_fn_config["clip_high_threshold"]
+            if not 0.0 <= clip_low <= 1.0 <= clip_high:
+                raise ValueError(
+                    "loss_fn='gspo' requires 0 <= clip_low_threshold <= 1 <= clip_high_threshold; "
+                    f"got {clip_low=} and {clip_high=}."
+                )
         return self
 
     def to_types(self) -> types.ForwardBackwardInput:
@@ -771,12 +785,14 @@ def validate_model_source(self):
 class SaveWeightsRequest(BaseModel):
     model_id: str
     path: str = Field(..., pattern=ID_PATTERN, max_length=ID_MAX_LENGTH)
+    seq_id: int | None = None
     type: Literal["save_weights"] | None = None
 
 
 class LoadWeightsRequest(BaseModel):
     model_id: str
     path: str
+    seq_id: int | None = None
     type: Literal["load_weights"] | None = None
 
 
@@ -859,6 +875,7 @@ class SupportedModel(BaseModel):
 
 class GetServerCapabilitiesResponse(BaseModel):
     supported_models: list[SupportedModel]
+    supported_loss_fns: list[str]
 
 
 class ListCheckpointsResponse(BaseModel):
@@ -1185,6 +1202,7 @@ async def load_weights(request: LoadWeightsRequest, req: Request, session: Async
         request_type=types.RequestType.LOAD_WEIGHTS,
         model_id=request.model_id,
         request_data=types.LoadWeightsInput(source_model_id=source_model_id, checkpoint_id=checkpoint_id),
+        seq_id=request.seq_id,
     )
 
     await session.commit()
@@ -1208,6 +1226,7 @@ async def save_weights(request: SaveWeightsRequest, session: AsyncSession = Depe
         request_type=types.RequestType.SAVE_WEIGHTS,
         model_id=request.model_id,
         request_data=types.SaveWeightsInput(path=request.path),
+        seq_id=request.seq_id,
     )
 
     await session.commit()
@@ -1253,6 +1272,7 @@ async def save_weights_for_sampler(request: SaveWeightsForSamplerRequest, sessio
             seq_id=request.seq_id,
             sampling_session_id=sampling_session_id,
         ),
+        seq_id=request.seq_id,
     )
 
     await session.commit()
@@ -1341,7 +1361,10 @@ async def get_server_capabilities(request: Request):
     supported_models = [
         SupportedModel(model_name=request.app.state.engine_config.base_model),
     ]
-    return GetServerCapabilitiesResponse(supported_models=supported_models)
+    return GetServerCapabilitiesResponse(
+        supported_models=supported_models,
+        supported_loss_fns=sorted(types.SUPPORTED_LOSS_FNS),
+    )
 
 
 class RetrieveFutureRequest(BaseModel):
diff --git a/skyrl/tinker/engine.py b/skyrl/tinker/engine.py
index 25b29e4ee1..bf1f72ed30 100644
--- a/skyrl/tinker/engine.py
+++ b/skyrl/tinker/engine.py
@@ -5,6 +5,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
+from itertools import groupby
 from pathlib import Path
 from typing import Any, Callable
 
@@ -377,6 +378,56 @@ def _find_destructive_barriers(self, session: Session) -> dict[str, int]:
         )
         return dict(session.exec(query).all())
 
+    def _find_next_sequence_ids(self, session: Session) -> dict[str, int]:
+        """Derive each model's next SDK sequence from its persisted futures."""
+        query = (
+            select(FutureDB.model_id, func.max(FutureDB.seq_id))
+            .where(FutureDB.seq_id.is_not(None))
+            .where(FutureDB.status != RequestStatus.PENDING)
+            .group_by(FutureDB.model_id)
+        )
+        return {model_id: last_seq_id + 1 for model_id, last_seq_id in session.exec(query).all()}
+
+    def _find_sequenced_requests(
+        self, session: Session, request_types: set[types.RequestType]
+    ) -> list[tuple[int, str, types.RequestType]]:
+        """Find each model's leading contiguous SDK requests while their types are allowed."""
+        next_sequence_ids = self._find_next_sequence_ids(session)
+        pending = session.exec(
+            select(FutureDB.request_id, FutureDB.model_id, FutureDB.seq_id, FutureDB.request_type)
+            .where(FutureDB.seq_id.is_not(None))
+            .where(FutureDB.status == RequestStatus.PENDING)
+            .order_by(FutureDB.model_id, FutureDB.seq_id)
+        ).all()
+
+        batchable = []
+        for model_id, requests in groupby(pending, key=lambda row: row.model_id):
+            requests = list(requests)
+            expected_seq_id = next_sequence_ids.get(model_id)
+            if expected_seq_id is None:
+                initial_seq_id = requests[0].seq_id
+                if initial_seq_id not in {0, 1}:
+                    continue
+                expected_seq_id = initial_seq_id
+            for request_id, _, seq_id, pending_type in requests:
+                if seq_id != expected_seq_id or pending_type not in request_types:
+                    break
+                batchable.append((request_id, model_id, pending_type))
+                expected_seq_id += 1
+        return batchable
+
+    def _find_sequenced_single_requests(self, session: Session) -> list[tuple[int, str, types.RequestType]]:
+        """Find leading contiguous SDK requests that do not run as model-pass batches."""
+        return self._find_sequenced_requests(
+            session,
+            {
+                types.RequestType.OPTIM_STEP,
+                types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER,
+                types.RequestType.SAVE_WEIGHTS,
+                types.RequestType.LOAD_WEIGHTS,
+            },
+        )
+
     def find_batchable_model_passes(
         self, session: Session, request_type: types.RequestType
     ) -> dict[str, tuple[str, types.ForwardBackwardInput]]:
@@ -399,16 +450,20 @@ def find_batchable_model_passes(
             select(FutureDB.request_id, FutureDB.model_id)
             .where(FutureDB.request_type == request_type)
             .where(FutureDB.status == RequestStatus.PENDING)
+            .where(FutureDB.seq_id.is_(None))
             .order_by(FutureDB.request_id)
         )
         ops = session.exec(query).all()
 
         # Filter: only include ops that come before their model's barrier
         batchable = [
+            (request_id, model_id) for request_id, model_id, _ in self._find_sequenced_requests(session, {request_type})
+        ]
+        batchable.extend(
             (request_id, model_id)
             for request_id, model_id in ops
             if model_id not in barriers or request_id < barriers[model_id]
-        ]
+        )
 
         return {
             str(request_id): (model_id, types.ForwardBackwardInput.model_validate(request_data))
@@ -490,6 +545,7 @@ def find_single_requests(self, session: Session) -> dict[str, tuple[str, types.R
         statement = (
             select(FutureDB.request_id, FutureDB.model_id, FutureDB.request_type)
             .where(FutureDB.status == RequestStatus.PENDING)
+            .where(FutureDB.seq_id.is_(None))
             .where(FutureDB.request_type != types.RequestType.FORWARD_BACKWARD)
             .where(FutureDB.request_type != types.RequestType.FORWARD)
             .where(FutureDB.request_type != types.RequestType.SAMPLE)
@@ -499,15 +555,18 @@ def find_single_requests(self, session: Session) -> dict[str, tuple[str, types.R
         other_futures = session.exec(statement).all()
 
         # Filter: only include ops that come before the first blocked pass for their model
-        other_futures = [
+        legacy_futures = [
             (request_id, model_id, request_type)
             for request_id, model_id, request_type in other_futures
             if model_id not in blocked_pass_barriers or request_id < blocked_pass_barriers[model_id]
         ]
+        sequenced_futures = self._find_sequenced_single_requests(session)
 
         return {
             str(request_id): (model_id, request_type, request_data)
-            for request_id, model_id, request_type, request_data in self._load_requests(session, other_futures)
+            for request_id, model_id, request_type, request_data in self._load_requests(
+                session, sequenced_futures + legacy_futures
+            )
         }
 
     def process_create_model(self, model_id: str, request_data: types.CreateModelInput) -> types.CreateModelOutput:
diff --git a/tests/backends/skyrl_train/test_token_based_batching_utils.py b/tests/backends/skyrl_train/test_token_based_batching_utils.py
index 4fac8a6679..160cbccdc5 100644
--- a/tests/backends/skyrl_train/test_token_based_batching_utils.py
+++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py
@@ -160,6 +160,19 @@ def test_reorder_and_combine(self):
         for i in range(batch.batch_size):
             assert torch.equal(reordered["sequences"][i], batch["sequences"][i])
 
+    def test_reorder_and_combine_items_drops_padding(self):
+        batch = self._make_batch([10, 3, 8, 5])
+        iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=12)
+        output_batches = [
+            [{"sample": index} for index in indices] + [{"sample": "padding"}] for indices in iterator._microbatches
+        ]
+        iterator._num_padding_microbatches = 1
+        output_batches.append([{"sample": "padding-microbatch"}])
+
+        reordered = iterator.reorder_and_combine_items(output_batches)
+
+        assert [output["sample"] for output in reordered] == list(range(batch.batch_size))
+
     def test_get_microbatch_iterator_factory(self):
         batch = self._make_batch([10, 10, 5, 5])
 
diff --git a/tests/tinker/skyrl_train/test_loss_normalization.py b/tests/tinker/skyrl_train/test_loss_normalization.py
index dff334598a..e802a6254e 100644
--- a/tests/tinker/skyrl_train/test_loss_normalization.py
+++ b/tests/tinker/skyrl_train/test_loss_normalization.py
@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import pytest
+import torch
 
 # Skip if skyrl_train_backend.py cannot be imported
 skyrl_train_backend = pytest.importorskip("skyrl.backends.skyrl_train_backend")
@@ -21,6 +22,86 @@ def test_ppo_thresholds_map_to_eps_clip():
     assert config == pytest.approx({"eps_clip_low": 0.2, "eps_clip_high": 0.28})
 
 
+def test_gspo_thresholds_map_to_native_sequence_loss():
+    loss_fn, config = _normalize(
+        None,
+        "policy",
+        "gspo",
+        {"clip_low_threshold": 0.98, "clip_high_threshold": 1.03},
+    )
+    assert loss_fn == "gspo"
+    assert config == {
+        "eps_clip_low": pytest.approx(0.02),
+        "eps_clip_high": pytest.approx(0.03),
+        "loss_reduction": "sequence_mean",
+        "use_entropy_loss": False,
+        "use_kl_loss": False,
+    }
+
+
+def test_gspo_batch_normalization_gives_each_trainable_sequence_equal_weight():
+    batch = skyrl_train_backend.TrainingInputBatch(
+        {
+            "advantages": torch.tensor([[2.0, 2.0, 0.0, 0.0], [-1.0, -1.0, -1.0, 0.0], [0.0] * 4]),
+            "loss_mask": torch.ones(3, 4),
+        }
+    )
+
+    skyrl_train_backend.SkyRLTrainBackend._normalize_gspo_batch(batch)
+
+    assert batch["loss_mask"].tolist() == [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0], [0.0] * 4]
+    assert (batch["advantages"] * batch["loss_mask"]).sum(dim=-1).tolist() == pytest.approx([1.0, -0.5, 0.0])
+
+
+def test_gspo_selects_only_rows_with_gradient_contributions():
+    batch = skyrl_train_backend.TrainingInputBatch(
+        {
+            "sequences": torch.tensor([[1, 2], [3, 0], [4, 5]]),
+            "attention_mask": torch.tensor([[1, 1], [1, 0], [1, 1]]),
+            "loss_mask": torch.tensor([[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]]),
+        }
+    )
+    batch.metadata = {"response_length": 2}
+
+    selected, indices = skyrl_train_backend.SkyRLTrainBackend._select_trainable_gspo_rows(batch)
+
+    assert indices == [0, 2]
+    assert selected["sequences"].tolist() == [[1, 2], [4, 5]]
+    assert selected.metadata == batch.metadata
+
+
+def test_gspo_keeps_entirely_inactive_batch_for_zero_gradient_backward():
+    batch = skyrl_train_backend.TrainingInputBatch(
+        {
+            "sequences": torch.tensor([[1, 2], [3, 4]]),
+            "loss_mask": torch.zeros(2, 2),
+        }
+    )
+
+    selected, indices = skyrl_train_backend.SkyRLTrainBackend._select_trainable_gspo_rows(batch)
+
+    assert selected is batch
+    assert indices == [0, 1]
+
+
+def test_gspo_pruning_metrics_sum_across_sdk_chunks_without_duplication():
+    first = skyrl_train_backend.SkyRLTrainBackend._get_gspo_pruning_metrics([3, 2, 4, 1], {0, 2}, 0, 2)
+    second = skyrl_train_backend.SkyRLTrainBackend._get_gspo_pruning_metrics([3, 2, 4, 1], {0, 2}, 2, 4)
+
+    assert first == {
+        "skyrl.ai/submitted_datums:sum": 1.0,
+        "skyrl.ai/skipped_datums:sum": 1.0,
+        "skyrl.ai/submitted_tokens:sum": 3.0,
+        "skyrl.ai/skipped_tokens:sum": 2.0,
+    }
+    assert second == {
+        "skyrl.ai/submitted_datums:sum": 1.0,
+        "skyrl.ai/skipped_datums:sum": 1.0,
+        "skyrl.ai/submitted_tokens:sum": 4.0,
+        "skyrl.ai/skipped_tokens:sum": 1.0,
+    }
+
+
 def test_dppo_deltas_are_nested_under_dppo():
     loss_fn, config = _normalize(None, "policy", "dppo", {"delta_low": 0.2, "delta_high": 0.3})
     assert loss_fn == "dppo"
diff --git a/tests/tinker/test_api_validation.py b/tests/tinker/test_api_validation.py
index 3207645704..04436d69cf 100644
--- a/tests/tinker/test_api_validation.py
+++ b/tests/tinker/test_api_validation.py
@@ -1,4 +1,5 @@
 import base64
+from types import SimpleNamespace
 
 import pytest
 from pydantic import TypeAdapter, ValidationError
@@ -27,6 +28,42 @@ def test_forward_backward_input_accepts_ppo_threshold_keys():
     assert req.loss_fn_config == {"clip_low_threshold": 0.9, "clip_high_threshold": 1.1}
 
 
+def test_forward_backward_input_accepts_gspo_threshold_keys():
+    req = api.ForwardBackwardInput(
+        data=[_make_datum()],
+        loss_fn="gspo",
+        loss_fn_config={"clip_low_threshold": 0.98, "clip_high_threshold": 1.03},
+    )
+    assert req.to_types().loss_fn == "gspo"
+
+
+@pytest.mark.parametrize(
+    "loss_fn_config",
+    [None, {"clip_low_threshold": 0.98}, {"clip_low_threshold": 1.01, "clip_high_threshold": 1.03}],
+)
+def test_forward_backward_input_rejects_incomplete_or_invalid_gspo_thresholds(loss_fn_config):
+    with pytest.raises(ValidationError, match="loss_fn='gspo'"):
+        api.ForwardBackwardInput(data=[_make_datum()], loss_fn="gspo", loss_fn_config=loss_fn_config)
+
+
+@pytest.mark.asyncio
+async def test_server_capabilities_advertise_native_gspo():
+    request = SimpleNamespace(
+        app=SimpleNamespace(state=SimpleNamespace(engine_config=SimpleNamespace(base_model="test-model")))
+    )
+    capabilities = await api.get_server_capabilities(request)
+
+    assert "gspo" in capabilities.supported_loss_fns
+
+
+def test_weight_requests_preserve_sdk_sequence_ids():
+    save_request = api.SaveWeightsRequest(model_id="model", path="checkpoint", seq_id=7)
+    load_request = api.LoadWeightsRequest(model_id="model", path="tinker://model/weights/checkpoint", seq_id=8)
+
+    assert save_request.seq_id == 7
+    assert load_request.seq_id == 8
+
+
 def test_forward_backward_input_accepts_ppo_value_clip():
     req = api.ForwardBackwardInput(
         data=[_make_datum()],
diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py
index 46af0332b0..c41059c216 100644
--- a/tests/tinker/test_engine.py
+++ b/tests/tinker/test_engine.py
@@ -109,6 +109,15 @@ def test_cleanup_stale_sessions():
             [],
             id="cispo",
         ),
+        pytest.param(
+            "gspo",
+            {"clip_low_threshold": 0.98, "clip_high_threshold": 1.03},
+            [0.1, 0.2, 0.3],
+            [-1.1, -1.0, -0.9],
+            [],
+            [],
+            id="gspo",
+        ),
         pytest.param("ppo_critic", {"value_clip": 0.2}, [], [], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], id="ppo_critic"),
         pytest.param(
             "dppo",
@@ -394,6 +403,99 @@ def test_find_batchable_model_passes_stops_at_barrier(scheduling_engine):
     assert request_data.data[0].loss_fn_inputs.target_tokens.data == [1, 2, 3]
 
 
+def test_find_batchable_model_passes_waits_for_next_sdk_sequence(scheduling_engine):
+    """Later chunks must wait for the first chunk that the SDK deliberately submits last."""
+    engine = scheduling_engine
+    payload = forward_backward_payload()
+    with Session(engine.db_engine) as session:
+        session.add(
+            FutureDB(
+                request_type=types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER,
+                model_id="model_a",
+                seq_id=1,
+                request_data={},
+                status=RequestStatus.COMPLETED,
+            )
+        )
+        later_chunks = [
+            FutureDB(
+                request_type=types.RequestType.FORWARD_BACKWARD,
+                model_id="model_a",
+                seq_id=seq_id,
+                request_data=payload,
+                status=RequestStatus.PENDING,
+            )
+            for seq_id in [3, 4]
+        ]
+        session.add_all(later_chunks)
+        session.commit()
+        later_request_ids = [row.request_id for row in later_chunks]
+
+    with Session(engine.db_engine) as session:
+        assert engine.find_batchable_model_passes(session, types.RequestType.FORWARD_BACKWARD) == {}
+
+    with Session(engine.db_engine) as session:
+        first_chunk = FutureDB(
+            request_type=types.RequestType.FORWARD_BACKWARD,
+            model_id="model_a",
+            seq_id=2,
+            request_data=payload,
+            status=RequestStatus.PENDING,
+        )
+        session.add(first_chunk)
+        session.commit()
+        first_request_id = first_chunk.request_id
+
+    with Session(engine.db_engine) as session:
+        batchable = engine.find_batchable_model_passes(session, types.RequestType.FORWARD_BACKWARD)
+
+    assert list(batchable) == [str(first_request_id), *(str(request_id) for request_id in later_request_ids)]
+
+
+@pytest.mark.parametrize("initial_seq_id", [0, 1])
+def test_find_single_requests_accepts_sdk_sequence_bootstraps(scheduling_engine, initial_seq_id):
+    engine = scheduling_engine
+    with Session(engine.db_engine) as session:
+        initial_request = FutureDB(
+            request_type=types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER,
+            model_id="model_a",
+            seq_id=initial_seq_id,
+            request_data={},
+            status=RequestStatus.PENDING,
+        )
+        session.add(initial_request)
+        session.commit()
+        initial_request_id = initial_request.request_id
+
+    with Session(engine.db_engine) as session:
+        assert engine.find_single_requests(session) == {
+            str(initial_request_id): (
+                "model_a",
+                types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER,
+                {},
+            )
+        }
+
+
+def test_find_single_requests_waits_for_next_sdk_sequence(scheduling_engine):
+    """A destructive request cannot overtake a missing earlier SDK request."""
+    engine = scheduling_engine
+    with Session(engine.db_engine) as session:
+        session.add(
+            FutureDB(
+                request_type=types.RequestType.OPTIM_STEP,
+                model_id="model_a",
+                seq_id=2,
+                request_data={},
+                status=RequestStatus.PENDING,
+            )
+        )
+        session.commit()
+
+    with Session(engine.db_engine) as session:
+        assert engine.find_single_requests(session) == {}
+
+
 def sample_payload(checkpoint_id: str) -> dict:
     return types.SampleInput(
         prompt=types.ModelInput(chunks=[types.EncodedTextChunk(tokens=[1, 2])]),

From 0e3f0755683ff086e73bd32e19f22854eb94487a Mon Sep 17 00:00:00 2001
From: Dian Ang <23232359+yapdianang@users.noreply.github.com>
Date: Tue, 18 Aug 2026 23:22:40 +0000
Subject: [PATCH 2/5] fix(tinker): handle native GSPO review edge cases

---
 .../workers/megatron/megatron_worker.py       |  3 +++
 skyrl/backends/skyrl_train_backend.py         |  2 +-
 skyrl/tinker/engine.py                        |  5 +---
 tests/tinker/test_engine.py                   | 27 +++++++++++--------
 4 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
index f29a8b3907..287ae313b0 100644
--- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
+++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
@@ -1257,6 +1257,9 @@ def forward_backward(
         loss_fn_output_batches = []
         for m_batch, metrics in zip(micro_buffer, metrics_list):
             # Extract loss_fn_outputs before reduce_metrics (it's not a scalar metric)
+            if metrics is None:
+                loss_fn_output_batches.append([])
+                continue
             loss_fn_output_batches.append(metrics.pop("loss_fn_outputs", []))
             # Skip fully-padding microbatches: their metrics (clip_ratio=0, policy_entropy=0,
             # ...) are meaningless and would drag down the mean-reduced metrics. Summed
diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py
index e40be0c9ff..2ad126dec1 100644
--- a/skyrl/backends/skyrl_train_backend.py
+++ b/skyrl/backends/skyrl_train_backend.py
@@ -959,7 +959,7 @@ def _forward_backward_single_model_batch(
         if pad_size > 0 and per_sample_outputs:
             per_sample_outputs = per_sample_outputs[:-pad_size]
 
-        if len(submitted_indices) != original_batch_size:
+        if per_sample_outputs and len(submitted_indices) != original_batch_size:
             restored_outputs = [{"logprobs": list(logprobs)} for logprobs in prepared_batch.all_sampling_logprobs]
             for index, output in zip(submitted_indices, per_sample_outputs, strict=True):
                 restored_outputs[index] = output
diff --git a/skyrl/tinker/engine.py b/skyrl/tinker/engine.py
index bf1f72ed30..537e58e718 100644
--- a/skyrl/tinker/engine.py
+++ b/skyrl/tinker/engine.py
@@ -405,10 +405,7 @@ def _find_sequenced_requests(
             requests = list(requests)
             expected_seq_id = next_sequence_ids.get(model_id)
             if expected_seq_id is None:
-                initial_seq_id = requests[0].seq_id
-                if initial_seq_id not in {0, 1}:
-                    continue
-                expected_seq_id = initial_seq_id
+                expected_seq_id = requests[0].seq_id
             for request_id, _, seq_id, pending_type in requests:
                 if seq_id != expected_seq_id or pending_type not in request_types:
                     break
diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py
index c41059c216..c7604c3a8f 100644
--- a/tests/tinker/test_engine.py
+++ b/tests/tinker/test_engine.py
@@ -477,23 +477,28 @@ def test_find_single_requests_accepts_sdk_sequence_bootstraps(scheduling_engine,
         }
 
 
-def test_find_single_requests_waits_for_next_sdk_sequence(scheduling_engine):
-    """A destructive request cannot overtake a missing earlier SDK request."""
+def test_find_single_requests_bootstraps_resumed_sdk_sequence(scheduling_engine):
     engine = scheduling_engine
     with Session(engine.db_engine) as session:
-        session.add(
-            FutureDB(
-                request_type=types.RequestType.OPTIM_STEP,
-                model_id="model_a",
-                seq_id=2,
-                request_data={},
-                status=RequestStatus.PENDING,
-            )
+        resumed_request = FutureDB(
+            request_type=types.RequestType.OPTIM_STEP,
+            model_id="model_a",
+            seq_id=100,
+            request_data={},
+            status=RequestStatus.PENDING,
         )
+        session.add(resumed_request)
         session.commit()
+        resumed_request_id = resumed_request.request_id
 
     with Session(engine.db_engine) as session:
-        assert engine.find_single_requests(session) == {}
+        assert engine.find_single_requests(session) == {
+            str(resumed_request_id): (
+                "model_a",
+                types.RequestType.OPTIM_STEP,
+                {},
+            )
+        }
 
 
 def sample_payload(checkpoint_id: str) -> dict:

From 860daf06031a329b6d15ceeee5e6a4b7ece71540 Mon Sep 17 00:00:00 2001
From: Dian Ang <23232359+yapdianang@users.noreply.github.com>
Date: Mon, 17 Aug 2026 09:59:28 +0000
Subject: [PATCH 3/5] fix(tinker): keep forwarded sample futures in memory

---
 docs/content/docs/tinker/architecture.mdx     |  9 ++-
 skyrl/tinker/api.py                           | 59 +++++++++-------
 skyrl/tinker/extra/__init__.py                |  7 +-
 skyrl/tinker/extra/external_inference.py      | 19 ++----
 skyrl/tinker/extra/in_memory_future_store.py  | 67 +++++++++++++++++++
 .../extra/skyrl_train_inference_forwarding.py | 22 ++----
 .../skyrl_train/test_async_sample_routing.py  | 32 +++------
 tests/tinker/test_future_waiting.py           | 28 +++++++-
 tests/tinker/test_in_memory_future_store.py   | 61 +++++++++++++++++
 9 files changed, 222 insertions(+), 82 deletions(-)
 create mode 100644 skyrl/tinker/extra/in_memory_future_store.py
 create mode 100644 tests/tinker/test_in_memory_future_store.py

diff --git a/docs/content/docs/tinker/architecture.mdx b/docs/content/docs/tinker/architecture.mdx
index c20947a996..245d4b49f0 100644
--- a/docs/content/docs/tinker/architecture.mdx
+++ b/docs/content/docs/tinker/architecture.mdx
@@ -8,7 +8,7 @@ This page describes how SkyRL implements the Tinker API, including the system ar
 
 The integration is organized in three high-level layers:
 
-1. **API Layer** (`skyrl.tinker.api`) - FastAPI HTTP server that accepts Tinker API requests, stores them in a database, and returns future IDs for async polling
+1. **API Layer** (`skyrl.tinker.api`) - FastAPI HTTP server that accepts Tinker API requests and returns future IDs for async polling. Engine-owned operations use the database; API-forwarded samples use in-memory futures.
 2. **Engine Layer** (`skyrl.tinker.engine`) - Background subprocess that polls the database, batches pending requests, and dispatches them to the backend
 3. **Backend Layer** (`skyrl.backends`) - Translates Tinker operations into training and inference calls, managing Ray workers, FSDP/Megatron training, and vLLM inference
 
@@ -71,6 +71,9 @@ Sampling requests (`sample`) go through the following lifecycle:
 Client calls sample()
     │
     ▼
+API Server (FastAPI)
+    │  Creates an in-memory future
+    ▼
 SkyRL-Train Backend
     │  Converts Tinker SamplingParams → vLLM params
     ▼
@@ -83,10 +86,10 @@ RemoteInferenceClient
 Inference Workers (vLLM)
     │  Generate tokens with logprobs
     ▼
-Results aggregated → GeneratedSequence objects returned
+Results resolve the in-memory future → GeneratedSequence objects returned
 ```
 
-The `sample()` call is a fairly lightweight wrapper around SkyRL-Train's vLLM inference engines. The backend translates Tinker `SamplingParams` to vLLM format (e.g., `stop_strings` → `stop`, `stop_tokens` → `stop_token_ids`) and delegates prompts to the `RemoteInferenceClient`. `RemoteInferenceClient` forwards the requests to an instance of [vLLM Router](https://github.com/vllm-project/router), which handles load balancing and sticky routing across vLLM workers.
+The `sample()` call is a fairly lightweight wrapper around SkyRL-Train's vLLM inference engines. Non-colocated and external inference requests bypass the engine database: the API owns their futures and signals completion with `asyncio.Event`. The backend translates Tinker `SamplingParams` to vLLM format (e.g., `stop_strings` → `stop`, `stop_tokens` → `stop_token_ids`) and delegates prompts to the `RemoteInferenceClient`. `RemoteInferenceClient` forwards the requests to an instance of [vLLM Router](https://github.com/vllm-project/router), which handles load balancing and sticky routing across vLLM workers.
 
 
 ## Weight Sync
diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py
index 2f88b23560..40d66e6b5a 100644
--- a/skyrl/tinker/api.py
+++ b/skyrl/tinker/api.py
@@ -46,6 +46,7 @@
 )
 from skyrl.tinker.extra import (
     ExternalInferenceClient,
+    InMemoryFutureStore,
     SkyRLTrainInferenceForwardingClient,
 )
 from skyrl.tinker.proto_serialization import (
@@ -260,12 +261,17 @@ async def lifespan(app: FastAPI):
     # SkyRL-Train default is colocate_all=True; only opt into forwarding
     # when the operator explicitly sets it to False.
     is_colocated = bool(backend_cfg.get("trainer.placement.colocate_all", True))
+    app.state.external_future_store = InMemoryFutureStore()
     if app.state.engine_config.external_inference_url:
-        app.state.external_inference_client = ExternalInferenceClient(app.state.engine_config, app.state.db_engine)
+        app.state.external_inference_client = ExternalInferenceClient(
+            app.state.engine_config, app.state.external_future_store
+        )
         logger.info(f"External engine configured: {app.state.engine_config.external_inference_url}")
     elif backend_name in ("megatron", "fsdp") and not is_colocated:
         app.state.external_inference_client = SkyRLTrainInferenceForwardingClient(
-            app.state.engine_config, app.state.db_engine
+            app.state.engine_config,
+            app.state.db_engine,
+            app.state.external_future_store,
         )
         logger.info(
             "SkyRL-Train inference forwarding client enabled for non-colocated backend=%s",
@@ -1322,35 +1328,33 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession =
         # Validate that the checkpoint exists and is ready
         await validate_checkpoint(req, model_id, checkpoint_id, types.CheckpointType.SAMPLER, session)
 
-    request_id = await create_future(
-        session=session,
-        request_type=(
-            types.RequestType.EXTERNAL if req.app.state.external_inference_client else types.RequestType.SAMPLE
-        ),
-        model_id=model_id,
-        request_data=types.SampleInput(
-            base_model=base_model,
-            prompt=request.prompt.to_types(),
-            sampling_params=request.sampling_params.to_types(),
-            num_samples=request.num_samples,
-            checkpoint_id=checkpoint_id,
-            # A positive topk implies prompt logprobs: both are read off the same
-            # prompt forward pass, so asking for one asks for the other.
-            prompt_logprobs=bool(request.prompt_logprobs) or request.topk_prompt_logprobs > 0,
-            topk_prompt_logprobs=request.topk_prompt_logprobs,
-            seq_id=request.seq_id,
-            sampling_session_id=request.sampling_session_id,
-        ),
-    )
-
-    await session.commit()
-
     if req.app.state.external_inference_client:
+        request_id = req.app.state.external_future_store.create_future()
         asyncio.create_task(
             req.app.state.external_inference_client.call_and_store_result(
                 request_id, request, model_id, checkpoint_id, base_model=base_model
             )
         )
+    else:
+        request_id = await create_future(
+            session=session,
+            request_type=types.RequestType.SAMPLE,
+            model_id=model_id,
+            request_data=types.SampleInput(
+                base_model=base_model,
+                prompt=request.prompt.to_types(),
+                sampling_params=request.sampling_params.to_types(),
+                num_samples=request.num_samples,
+                checkpoint_id=checkpoint_id,
+                # A positive topk implies prompt logprobs: both are read off the same
+                # prompt forward pass, so asking for one asks for the other.
+                prompt_logprobs=bool(request.prompt_logprobs) or request.topk_prompt_logprobs > 0,
+                topk_prompt_logprobs=request.topk_prompt_logprobs,
+                seq_id=request.seq_id,
+                sampling_session_id=request.sampling_session_id,
+            ),
+        )
+        await session.commit()
 
     return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id))
 
@@ -1377,7 +1381,10 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request):
     request_id = int(request.request_id)
 
     try:
-        row = await wait_for_future(req.app.state.future_waiters, request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
+        if request_id < 0:
+            row = await req.app.state.external_future_store.wait_for_future(request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
+        else:
+            row = await wait_for_future(req.app.state.future_waiters, request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
     except KeyError:
         raise HTTPException(status_code=404, detail="Future not found")
 
diff --git a/skyrl/tinker/extra/__init__.py b/skyrl/tinker/extra/__init__.py
index 727cc03e39..ea194eb88e 100644
--- a/skyrl/tinker/extra/__init__.py
+++ b/skyrl/tinker/extra/__init__.py
@@ -1,6 +1,11 @@
 from skyrl.tinker.extra.external_inference import ExternalInferenceClient
+from skyrl.tinker.extra.in_memory_future_store import InMemoryFutureStore
 from skyrl.tinker.extra.skyrl_train_inference_forwarding import (
     SkyRLTrainInferenceForwardingClient,
 )
 
-__all__ = ["ExternalInferenceClient", "SkyRLTrainInferenceForwardingClient"]
+__all__ = [
+    "ExternalInferenceClient",
+    "InMemoryFutureStore",
+    "SkyRLTrainInferenceForwardingClient",
+]
diff --git a/skyrl/tinker/extra/external_inference.py b/skyrl/tinker/extra/external_inference.py
index c6b644da28..bd4b38275e 100644
--- a/skyrl/tinker/extra/external_inference.py
+++ b/skyrl/tinker/extra/external_inference.py
@@ -1,17 +1,16 @@
 import asyncio
-from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import httpx
 from cloudpathlib import AnyPath
-from sqlmodel.ext.asyncio.session import AsyncSession
 
 from skyrl.backends.renderer import render_model_input
 from skyrl.backends.utils import convert_vllm_prompt_logprobs
 from skyrl.tinker import types
 from skyrl.tinker.config import EngineConfig
-from skyrl.tinker.db_models import FutureDB, RequestStatus
+from skyrl.tinker.db_models import RequestStatus
+from skyrl.tinker.extra.in_memory_future_store import InMemoryFutureStore
 from skyrl.utils.log import logger
 from skyrl.utils.storage import download_and_unpack
 
@@ -41,12 +40,12 @@ def _extract_checkpoint_sync(checkpoint_path: AnyPath, target_dir: Path) -> None
 class ExternalInferenceClient:
     """Client for calling external inference engines (e.g., vLLM)."""
 
-    def __init__(self, engine_config: EngineConfig, db_engine):
+    def __init__(self, engine_config: EngineConfig, future_store: InMemoryFutureStore):
         self.base_url = f"{engine_config.external_inference_url}/v1"
         self.api_key = engine_config.external_inference_api_key
         self.checkpoints_base = engine_config.checkpoints_base
         self.lora_base_dir = engine_config.external_inference_lora_base
-        self.db_engine = db_engine
+        self.future_store = future_store
 
     async def call_and_store_result(
         self,
@@ -57,7 +56,7 @@ async def call_and_store_result(
         *,
         base_model: str | None = None,
     ):
-        """Background task to call external engine and store result in database."""
+        """Call the external engine and resolve its API-process-owned future."""
         try:
             async with httpx.AsyncClient(
                 base_url=self.base_url,
@@ -73,13 +72,7 @@ async def call_and_store_result(
             result = types.ErrorResponse(error=str(e), status="failed")
             status = RequestStatus.FAILED
 
-        async with AsyncSession(self.db_engine) as session:
-            future = await session.get(FutureDB, request_id)
-            # `result_data` is a text column holding pre-serialized JSON.
-            future.result_data = result.model_dump_json()
-            future.status = status
-            future.completed_at = datetime.now(timezone.utc)
-            await session.commit()
+        self.future_store.complete_future(request_id, status, result.model_dump_json())
 
     async def _forward_to_engine(
         self,
diff --git a/skyrl/tinker/extra/in_memory_future_store.py b/skyrl/tinker/extra/in_memory_future_store.py
new file mode 100644
index 0000000000..9a4a82e33f
--- /dev/null
+++ b/skyrl/tinker/extra/in_memory_future_store.py
@@ -0,0 +1,67 @@
+import asyncio
+import itertools
+import time
+from dataclasses import dataclass, field
+
+from skyrl.tinker.db_models import RequestStatus
+
+
+@dataclass
+class _StoredFuture:
+    event: asyncio.Event = field(default_factory=asyncio.Event)
+    status: RequestStatus = RequestStatus.PENDING
+    result_data: str | None = None
+    completed_at: float | None = None
+
+
+class InMemoryFutureStore:
+    """Store API-process-owned futures without routing them through the engine database."""
+
+    def __init__(self, *, terminal_retention_sec: float = 600, max_terminal_futures: int = 10_000):
+        self._terminal_retention_sec = terminal_retention_sec
+        self._max_terminal_futures = max_terminal_futures
+        self._next_id = itertools.count(start=-1, step=-1)
+        self._futures: dict[int, _StoredFuture] = {}
+
+    def create_future(self) -> int:
+        self._cleanup_terminal_futures()
+        request_id = next(self._next_id)
+        self._futures[request_id] = _StoredFuture()
+        return request_id
+
+    def complete_future(self, request_id: int, status: RequestStatus, result_data: str) -> None:
+        future = self._futures[request_id]
+        future.status = status
+        future.result_data = result_data
+        future.completed_at = time.monotonic()
+        future.event.set()
+
+    async def wait_for_future(self, request_id: int, timeout: float) -> tuple[RequestStatus, str | None] | None:
+        future = self._futures.get(request_id)
+        if future is None:
+            raise KeyError(request_id)
+        if future.status == RequestStatus.PENDING:
+            try:
+                await asyncio.wait_for(future.event.wait(), timeout)
+            except asyncio.TimeoutError:
+                return None
+        return future.status, future.result_data
+
+    def _cleanup_terminal_futures(self) -> None:
+        now = time.monotonic()
+        expired = [
+            request_id
+            for request_id, future in self._futures.items()
+            if future.completed_at is not None and now - future.completed_at > self._terminal_retention_sec
+        ]
+        for request_id in expired:
+            del self._futures[request_id]
+
+        terminal = sorted(
+            (future.completed_at, request_id)
+            for request_id, future in self._futures.items()
+            if future.completed_at is not None
+        )
+        overflow = len(terminal) - self._max_terminal_futures
+        for _, request_id in terminal[: max(overflow, 0)]:
+            del self._futures[request_id]
diff --git a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py
index 227db2de92..7cec96d8f7 100644
--- a/skyrl/tinker/extra/skyrl_train_inference_forwarding.py
+++ b/skyrl/tinker/extra/skyrl_train_inference_forwarding.py
@@ -5,7 +5,6 @@
 """
 
 import asyncio
-from datetime import datetime, timezone
 
 import httpx
 from sqlmodel.ext.asyncio.session import AsyncSession
@@ -14,16 +13,18 @@
 from skyrl.backends.utils import convert_vllm_prompt_logprobs
 from skyrl.tinker import types
 from skyrl.tinker.config import EngineConfig
-from skyrl.tinker.db_models import EngineStateDB, FutureDB, RequestStatus
+from skyrl.tinker.db_models import EngineStateDB, RequestStatus
+from skyrl.tinker.extra.in_memory_future_store import InMemoryFutureStore
 from skyrl.utils.log import logger
 
 
 class SkyRLTrainInferenceForwardingClient:
     """Forwards EXTERNAL sample requests to the SkyRL-Train-managed vLLM."""
 
-    def __init__(self, engine_config: EngineConfig, db_engine):
+    def __init__(self, engine_config: EngineConfig, db_engine, future_store: InMemoryFutureStore):
         self.engine_config = engine_config
         self.db_engine = db_engine
+        self.future_store = future_store
         self._cached_proxy_url: str | None = None
         self._cache_lock = asyncio.Lock()
         # Backpressure layered: httpx pool -> vllm-router -> vLLM max_num_seqs.
@@ -71,7 +72,7 @@ async def call_and_store_result(
         *,
         base_model: str | None = None,
     ):
-        """Forward a sample request to vLLM and write the result to FutureDB."""
+        """Forward a sample request to vLLM and resolve its API-process-owned future."""
         try:
             result = await self._forward_with_retry(sample_req, model_id, base_model=base_model)
             status = RequestStatus.COMPLETED
@@ -80,18 +81,7 @@ async def call_and_store_result(
             result = types.ErrorResponse(error=str(e), status="failed")
             status = RequestStatus.FAILED
 
-        async with AsyncSession(self.db_engine) as session:
-            future = await session.get(FutureDB, request_id)
-            if future is None:
-                # Row was deleted between scheduling and completion (cancelled
-                # request, stale-session GC). Nothing to write back.
-                logger.warning("FutureDB row %s missing on completion write — skipping", request_id)
-                return
-            # `result_data` is a text column holding pre-serialized JSON.
-            future.result_data = result.model_dump_json()
-            future.status = status
-            future.completed_at = datetime.now(timezone.utc)
-            await session.commit()
+        self.future_store.complete_future(request_id, status, result.model_dump_json())
 
     async def _forward_with_retry(self, sample_req, model_id: str, *, base_model: str | None) -> types.SampleOutput:
         # httpx.RequestError covers ConnectError, ReadError, TimeoutException, etc.
diff --git a/tests/tinker/skyrl_train/test_async_sample_routing.py b/tests/tinker/skyrl_train/test_async_sample_routing.py
index 72f98751c8..462b5fea69 100644
--- a/tests/tinker/skyrl_train/test_async_sample_routing.py
+++ b/tests/tinker/skyrl_train/test_async_sample_routing.py
@@ -9,8 +9,8 @@
 Coverage:
   - test_engine_state_published: after ``save_weights_for_sampler``, the
     engine's vLLM proxy URL is written to ``EngineStateDB``.
-  - test_sample_uses_external_path: an issued sample creates a future of
-    type ``EXTERNAL`` (not ``SAMPLE``) and resolves successfully.
+  - test_sample_bypasses_future_db: an issued sample resolves without creating
+    a database future.
   - test_sample_concurrent_with_training_is_fast: the central
     parallelism test. While a long-running stream of ``forward_backward``
     + ``optim_step`` calls is in flight, a sample request resolves in
@@ -210,15 +210,14 @@ def test_engine_state_published(server_db_path):
     ), f"expected an http(s) proxy URL, got {row.inference_proxy_url!r}"
 
 
-def test_sample_uses_external_path(server_db_path):
-    """A sample issued through the SDK creates a FutureDB row of type EXTERNAL.
+def test_sample_bypasses_future_db(server_db_path):
+    """A sample issued through the SDK does not use the engine's FutureDB.
 
     This is the "test" half of the design: the API hoists the sample off
-    the engine's serial loop and into the API process's asyncio loop.
+    the engine's serial loop and keeps its future in the API process.
     """
     from sqlmodel import Session, create_engine, func, select
 
-    from skyrl.tinker import types as skyrl_types
     from skyrl.tinker.db_models import FutureDB
 
     proc, db_path, _ = server_db_path
@@ -229,12 +228,11 @@ def test_sample_uses_external_path(server_db_path):
     _train_one_step(tc, tok)
     sampler = tc.save_weights_and_get_sampling_client(name="external_path_a")
 
-    # Snapshot the max future_id before submitting our sample so we can
-    # filter out any EXTERNAL futures from earlier tests.
+    # Snapshot the number of database futures before submitting our sample.
     eng = create_engine(f"sqlite:///{db_path}", echo=False)
     try:
         with Session(eng) as s:
-            max_before = s.exec(select(func.max(FutureDB.request_id))).one() or 0
+            count_before = s.exec(select(func.count()).select_from(FutureDB)).one()
     finally:
         eng.dispose()
 
@@ -245,24 +243,16 @@ def test_sample_uses_external_path(server_db_path):
     ).result()
     assert len(out.sequences) == 1
 
-    # Look for an EXTERNAL future with id > max_before. If async routing
-    # is on, every sample creates exactly one such row.
+    # API-forwarded samples use negative, in-memory future ids and must not
+    # add database work to the engine's scheduling hot path.
     eng = create_engine(f"sqlite:///{db_path}", echo=False)
     try:
         with Session(eng) as s:
-            stmt = (
-                select(FutureDB.request_id, FutureDB.request_type)
-                .where(FutureDB.request_id > max_before)
-                .where(FutureDB.request_type == skyrl_types.RequestType.EXTERNAL)
-            )
-            rows = s.exec(stmt).all()
+            count_after = s.exec(select(func.count()).select_from(FutureDB)).one()
     finally:
         eng.dispose()
 
-    assert len(rows) >= 1, (
-        f"expected at least one EXTERNAL future to be created by the sample call, "
-        f"found {len(rows)}; async sample routing may not be active"
-    )
+    assert count_after == count_before
 
 
 def test_sample_concurrent_with_training_is_fast(server_db_path):
diff --git a/tests/tinker/test_future_waiting.py b/tests/tinker/test_future_waiting.py
index 748f43973f..52ce923495 100644
--- a/tests/tinker/test_future_waiting.py
+++ b/tests/tinker/test_future_waiting.py
@@ -187,11 +187,17 @@ def _count(conn, cursor, statement, parameters, context, executemany):
     assert 0 < len(statements) < 50
 
 
-def _stub_request(async_engine, waiters, headers: dict | None = None):
+def _stub_request(async_engine, waiters, external_future_store=None, headers: dict | None = None):
     from types import SimpleNamespace
 
     return SimpleNamespace(
-        app=SimpleNamespace(state=SimpleNamespace(db_engine=async_engine, future_waiters=waiters)),
+        app=SimpleNamespace(
+            state=SimpleNamespace(
+                db_engine=async_engine,
+                external_future_store=external_future_store,
+                future_waiters=waiters,
+            )
+        ),
         headers=headers or {},
     )
 
@@ -219,6 +225,24 @@ async def test_retrieve_future_returns_completed_result(waiters, async_engine, s
     assert response.body == SAMPLE_RESULT.model_dump_json().encode()
 
 
+@pytest.mark.asyncio
+async def test_retrieve_future_returns_in_memory_external_result(waiters, async_engine):
+    from skyrl.tinker import api
+    from skyrl.tinker.extra import InMemoryFutureStore
+
+    store = InMemoryFutureStore()
+    request_id = store.create_future()
+    store.complete_future(request_id, RequestStatus.COMPLETED, SAMPLE_RESULT.model_dump_json())
+
+    response = await api.retrieve_future(
+        api.RetrieveFutureRequest(request_id=str(request_id)),
+        _stub_request(async_engine, waiters, store),
+    )
+
+    assert response.media_type == "application/json"
+    assert response.body == SAMPLE_RESULT.model_dump_json().encode()
+
+
 @pytest.mark.asyncio
 async def test_retrieve_future_400s_with_the_stored_error(waiters, async_engine, sync_engine):
     """The failure path still decodes the payload, since it inspects the error."""
diff --git a/tests/tinker/test_in_memory_future_store.py b/tests/tinker/test_in_memory_future_store.py
new file mode 100644
index 0000000000..612ecacebf
--- /dev/null
+++ b/tests/tinker/test_in_memory_future_store.py
@@ -0,0 +1,61 @@
+import asyncio
+
+import pytest
+
+from skyrl.tinker.db_models import RequestStatus
+from skyrl.tinker.extra.in_memory_future_store import InMemoryFutureStore
+
+
+@pytest.mark.asyncio
+async def test_concurrent_waiters_receive_completed_result():
+    store = InMemoryFutureStore()
+    request_id = store.create_future()
+    waiters = [asyncio.create_task(store.wait_for_future(request_id, 1)) for _ in range(3)]
+
+    store.complete_future(request_id, RequestStatus.COMPLETED, '{"value":1}')
+
+    assert await asyncio.gather(*waiters) == [
+        (RequestStatus.COMPLETED, '{"value":1}'),
+        (RequestStatus.COMPLETED, '{"value":1}'),
+        (RequestStatus.COMPLETED, '{"value":1}'),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_timed_out_waiter_does_not_cancel_future():
+    store = InMemoryFutureStore()
+    request_id = store.create_future()
+
+    assert await store.wait_for_future(request_id, 0) is None
+    store.complete_future(request_id, RequestStatus.FAILED, '{"error":"boom"}')
+
+    assert await store.wait_for_future(request_id, 1) == (
+        RequestStatus.FAILED,
+        '{"error":"boom"}',
+    )
+
+
+@pytest.mark.asyncio
+async def test_unknown_future_raises_key_error():
+    store = InMemoryFutureStore()
+
+    with pytest.raises(KeyError):
+        await store.wait_for_future(-1, 1)
+
+
+@pytest.mark.asyncio
+async def test_terminal_future_retention_is_bounded():
+    store = InMemoryFutureStore(max_terminal_futures=1)
+    first_id = store.create_future()
+    store.complete_future(first_id, RequestStatus.COMPLETED, "{}")
+    second_id = store.create_future()
+    store.complete_future(second_id, RequestStatus.COMPLETED, "{}")
+
+    store.create_future()
+
+    with pytest.raises(KeyError):
+        await store.wait_for_future(first_id, 1)
+    assert await store.wait_for_future(second_id, 1) == (
+        RequestStatus.COMPLETED,
+        "{}",
+    )

From bfa4f5b936efc064ca8678c67ddd783db731563b Mon Sep 17 00:00:00 2001
From: Dian Ang <23232359+yapdianang@users.noreply.github.com>
Date: Mon, 17 Aug 2026 10:04:39 +0000
Subject: [PATCH 4/5] perf(tinker): evict sample futures in constant time

---
 skyrl/tinker/extra/in_memory_future_store.py | 24 ++++++++------------
 1 file changed, 9 insertions(+), 15 deletions(-)

diff --git a/skyrl/tinker/extra/in_memory_future_store.py b/skyrl/tinker/extra/in_memory_future_store.py
index 9a4a82e33f..a69954085c 100644
--- a/skyrl/tinker/extra/in_memory_future_store.py
+++ b/skyrl/tinker/extra/in_memory_future_store.py
@@ -1,6 +1,7 @@
 import asyncio
 import itertools
 import time
+from collections import deque
 from dataclasses import dataclass, field
 
 from skyrl.tinker.db_models import RequestStatus
@@ -22,6 +23,7 @@ def __init__(self, *, terminal_retention_sec: float = 600, max_terminal_futures:
         self._max_terminal_futures = max_terminal_futures
         self._next_id = itertools.count(start=-1, step=-1)
         self._futures: dict[int, _StoredFuture] = {}
+        self._terminal_futures: deque[tuple[float, int]] = deque()
 
     def create_future(self) -> int:
         self._cleanup_terminal_futures()
@@ -34,7 +36,9 @@ def complete_future(self, request_id: int, status: RequestStatus, result_data: s
         future.status = status
         future.result_data = result_data
         future.completed_at = time.monotonic()
+        self._terminal_futures.append((future.completed_at, request_id))
         future.event.set()
+        self._cleanup_terminal_futures()
 
     async def wait_for_future(self, request_id: int, timeout: float) -> tuple[RequestStatus, str | None] | None:
         future = self._futures.get(request_id)
@@ -49,19 +53,9 @@ async def wait_for_future(self, request_id: int, timeout: float) -> tuple[Reques
 
     def _cleanup_terminal_futures(self) -> None:
         now = time.monotonic()
-        expired = [
-            request_id
-            for request_id, future in self._futures.items()
-            if future.completed_at is not None and now - future.completed_at > self._terminal_retention_sec
-        ]
-        for request_id in expired:
-            del self._futures[request_id]
-
-        terminal = sorted(
-            (future.completed_at, request_id)
-            for request_id, future in self._futures.items()
-            if future.completed_at is not None
-        )
-        overflow = len(terminal) - self._max_terminal_futures
-        for _, request_id in terminal[: max(overflow, 0)]:
+        while self._terminal_futures and (
+            now - self._terminal_futures[0][0] > self._terminal_retention_sec
+            or len(self._terminal_futures) > self._max_terminal_futures
+        ):
+            _, request_id = self._terminal_futures.popleft()
             del self._futures[request_id]

From 26b7ce8e84811159af8cb230d0a688ee12a2038c Mon Sep 17 00:00:00 2001
From: Dian Ang <23232359+yapdianang@users.noreply.github.com>
Date: Tue, 18 Aug 2026 23:15:58 +0000
Subject: [PATCH 5/5] fix(tinker): compose external futures with proto results

---
 skyrl/tinker/api.py                 | 12 +++++++++---
 tests/tinker/test_future_waiting.py |  3 ++-
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py
index 40d66e6b5a..51cbdc27f9 100644
--- a/skyrl/tinker/api.py
+++ b/skyrl/tinker/api.py
@@ -1379,9 +1379,10 @@ class RetrieveFutureRequest(BaseModel):
 async def retrieve_future(request: RetrieveFutureRequest, req: Request):
     """Retrieve the result of an async operation, waiting until it's available."""
     request_id = int(request.request_id)
+    is_external_future = request_id < 0
 
     try:
-        if request_id < 0:
+        if is_external_future:
             row = await req.app.state.external_future_store.wait_for_future(request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
         else:
             row = await wait_for_future(req.app.state.future_waiters, request_id, RETRIEVE_FUTURE_TIMEOUT_SECONDS)
@@ -1391,13 +1392,18 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request):
     if row is None:
         raise HTTPException(status_code=408, detail="Timeout waiting for result")
 
-    status, request_type, result_data = row
+    if is_external_future:
+        status, result_data = row
+        request_type = None
+    else:
+        status, request_type, result_data = row
     if status == RequestStatus.COMPLETED:
         # The SDK retrieves sample/forward/forward_backward results in proto
         # wire format when it advertises support; SDK >= 0.25.0 rejects JSON
         # for these types. Errors and other result types stay JSON.
         if (
-            types.RequestType(request_type) in PROTO_SERIALIZABLE_REQUEST_TYPES
+            request_type is not None
+            and types.RequestType(request_type) in PROTO_SERIALIZABLE_REQUEST_TYPES
             and PROTO_CONTENT_TYPE in req.headers.get("accept", "").lower()
         ):
             return Response(
diff --git a/tests/tinker/test_future_waiting.py b/tests/tinker/test_future_waiting.py
index 52ce923495..bb51fb9318 100644
--- a/tests/tinker/test_future_waiting.py
+++ b/tests/tinker/test_future_waiting.py
@@ -236,9 +236,10 @@ async def test_retrieve_future_returns_in_memory_external_result(waiters, async_
 
     response = await api.retrieve_future(
         api.RetrieveFutureRequest(request_id=str(request_id)),
-        _stub_request(async_engine, waiters, store),
+        _stub_request(async_engine, waiters, store, headers={"accept": api.PROTO_CONTENT_TYPE}),
     )
 
+    # Forwarded samples use the Tinker JSON response even when the client also accepts proto.
     assert response.media_type == "application/json"
     assert response.body == SAMPLE_RESULT.model_dump_json().encode()