From d12385ebec472851cf97f52cae6733bc7e41499f Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Sun, 20 Sep 2026 00:51:35 +0000
Subject: [PATCH 1/6] feat(rl): replay sampling support in actor scoring

---
 miles/backends/fsdp_utils/actor.py            | 15 +++++
 miles/backends/megatron_utils/actor.py        |  2 +
 miles/backends/megatron_utils/model.py        | 41 +++++++++----
 .../training_utils/loss_hub/losses.py         |  4 ++
 .../backends/training_utils/sampling_mask.py  | 18 +++++-
 miles/utils/sampling_mask.py                  |  5 ++
 .../backends/megatron_utils/test_model.py     | 55 ++++++++++++++++++
 .../test_shared_ppo_lifecycle.py              | 13 ++++-
 tests/fast/backends/test_fsdp_train.py        | 58 +++++++++++++++++++
 .../training_utils/test_sampling_mask.py      | 45 +++++++++++++-
 10 files changed, 240 insertions(+), 16 deletions(-)

diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py
index 9f98c53b4cc..c77a3fb9ead 100644
--- a/miles/backends/fsdp_utils/actor.py
+++ b/miles/backends/fsdp_utils/actor.py
@@ -21,6 +21,7 @@
 )
 from miles.backends.training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, loss_function
 from miles.backends.training_utils.parallel import get_parallel_state, set_parallel_state
+from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
 from miles.ray.train_actor import TrainRayActor
 from miles.utils import async_utils, train_dump_utils, train_metric_utils
 from miles.utils.context_utils import with_defer
@@ -32,6 +33,7 @@
 from miles.utils.processing_utils import load_processor, load_tokenizer
 from miles.utils.profile_utils import TrainProfiler
 from miles.utils.ray_utils import Box
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils.tracking import init_tracking
 
@@ -369,6 +371,7 @@ def _compute_log_prob(
             active_model.eval()
         else:
             active_model = self.model
+        use_rollout_sampling_mask = model_tag == "actor" and sampling_support_replay_enabled(self.args)
 
         try:
             forward_data_store = []
@@ -392,6 +395,8 @@ def _compute_log_prob(
                             "response_lengths",
                             "max_seq_lens",
                         ]
+                        if use_rollout_sampling_mask:
+                            forward_only_keys.extend(["rollout_sampling_mask_ids", "rollout_sampling_mask_offsets"])
                         batch = get_batch(
                             data_iterator,
                             forward_only_keys,
@@ -399,6 +404,9 @@ def _compute_log_prob(
                             self.args.qkv_format,
                             get_position_ids=True,
                         )
+                        rollout_sampling_mask = None
+                        if use_rollout_sampling_mask:
+                            rollout_sampling_mask = get_rollout_sampling_masks(batch)
 
                         model_args = self._get_model_inputs_args(batch)
                         # keep logits in native bf16 (chunks upcast to fp32 downstream); avoids a full-vocab fp32 tensor (~5GB)
@@ -413,6 +421,7 @@ def _compute_log_prob(
                             response_lengths=batch["response_lengths"],
                             with_entropy=(store_prefix == ""),
                             max_seq_lens=batch.get("max_seq_lens", None),
+                            rollout_sampling_mask=rollout_sampling_mask,
                         )
 
                         batch_result = {
@@ -501,6 +510,11 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
         with routing_replay.stage(routing_replay.REPLAY_BACKWARD), timer("actor_train"):
             data_iterator.reset()
             num_steps_per_rollout = len(num_microbatches)
+            sampling_mask_keys = (
+                ("rollout_sampling_mask_ids", "rollout_sampling_mask_offsets")
+                if sampling_support_replay_enabled(self.args)
+                else ()
+            )
 
             for step_id in range(num_steps_per_rollout):
                 self.optimizer.zero_grad(set_to_none=True)
@@ -523,6 +537,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
                             "returns",
                             "ref_log_probs",
                             "rollout_log_probs",
+                            *sampling_mask_keys,
                         ],
                         self.args.data_pad_size_multiplier,
                         self.args.qkv_format,
diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py
index ee9cd157652..c65d001d8d9 100644
--- a/miles/backends/megatron_utils/actor.py
+++ b/miles/backends/megatron_utils/actor.py
@@ -33,6 +33,7 @@
 from miles.utils.ray_utils import Box
 from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups
 from miles.utils.replay_base import all_replay_managers, routing_replay_manager
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils.structured_log import with_logs
@@ -430,6 +431,7 @@ def compute_log_prob(
                 rollout_id=rollout_id,
                 store_prefix=store_prefix,
                 fp32_output=False,
+                use_rollout_sampling_mask=store_prefix == "" and sampling_support_replay_enabled(self.args),
             )
 
     @with_logs
diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py
index aff3307f551..20182e968b6 100644
--- a/miles/backends/megatron_utils/model.py
+++ b/miles/backends/megatron_utils/model.py
@@ -29,11 +29,13 @@
 from miles.backends.megatron_utils.ft.types import TrainStepOutcome
 from miles.backends.megatron_utils.local_weight_checksum import dump_local_weight_checksums
 from miles.backends.megatron_utils.optimizer_state_reset import reset_optimizer_states
+from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
 from miles.utils.audit_utils.witness.allocator import WitnessInfo
 from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale
 from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase
 from miles.utils.memory_utils import clear_memory
 from miles.utils.multi_lora import is_multi_lora_enabled
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor
 from miles.utils.tracking_utils.structured_log import log_structured
 
@@ -268,6 +270,7 @@ def forward_only(
     rollout_id: int,
     store_prefix: str = "",
     fp32_output: bool = True,
+    use_rollout_sampling_mask: bool = False,
 ) -> dict[str, list[torch.Tensor]]:
     """Run forward passes only and collect non-loss outputs (e.g., logprobs).
 
@@ -283,6 +286,8 @@ def forward_only(
         rollout_id: Rollout identifier (selects the per-rollout dump subdirectory).
         store_prefix: Prefix to prepend to stored output keys.
         fp32_output: Whether Megatron should upcast the complete model output to FP32.
+        use_rollout_sampling_mask: Whether to score over each rollout token's
+            captured sampling support.
 
     Returns:
         Aggregated outputs keyed by ``store_prefix + key``.
@@ -304,17 +309,20 @@ def forward_step(
     ) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
         """Return the model output and its batch-bound loss callback."""
 
+        forward_only_keys = [
+            "tokens",
+            "loss_masks",
+            "multimodal_train_inputs",
+            "total_lengths",
+            "response_lengths",
+            "max_seq_lens",
+            "witness_ids",
+        ]
+        if use_rollout_sampling_mask:
+            forward_only_keys.extend(["rollout_sampling_mask_ids", "rollout_sampling_mask_offsets"])
         batch = get_batch(
             data_iterator,
-            [
-                "tokens",
-                "loss_masks",
-                "multimodal_train_inputs",
-                "total_lengths",
-                "response_lengths",
-                "max_seq_lens",
-                "witness_ids",
-            ],
+            forward_only_keys,
             args.data_pad_size_multiplier,
             args.qkv_format,
             allgather_cp=args.allgather_cp,
@@ -324,6 +332,7 @@ def forward_step(
         packed_seq_params = get_packed_seq_params(batch, args)
         total_lengths = batch["total_lengths"]
         response_lengths = batch["response_lengths"]
+        rollout_sampling_mask = get_rollout_sampling_masks(batch) if use_rollout_sampling_mask else None
 
         if "adapter_token_counts" in batch:
             from megatron.bridge.peft.multi_lora_layers import set_tokens_per_adapter_slot
@@ -342,8 +351,7 @@ def forward_step(
             fp32_output=fp32_output,
         )
 
-        return output_tensor, partial(
-            f,
+        callback_kwargs = dict(
             args=args,
             unconcat_tokens=unconcat_tokens,
             total_lengths=total_lengths,
@@ -351,6 +359,10 @@ def forward_step(
             with_entropy=args.use_rollout_entropy,
             max_seq_lens=batch.get("max_seq_lens", None),
         )
+        if use_rollout_sampling_mask:
+            callback_kwargs["rollout_sampling_mask"] = rollout_sampling_mask
+
+        return output_tensor, partial(f, **callback_kwargs)
 
     # Turn on evaluation mode which disables dropout.
     for model_module in model:
@@ -407,6 +419,12 @@ def run_forward_backward_pass(
 ):
     """One pipeline forward/backward pass over the microbatches; no optimizer interaction."""
 
+    sampling_mask_keys = (
+        ("rollout_sampling_mask_ids", "rollout_sampling_mask_offsets")
+        if sampling_support_replay_enabled(args) and args.loss_type == "policy_loss"
+        else ()
+    )
+
     @dumper_phase_util.wrap_forward_step
     def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[
         torch.Tensor,
@@ -447,6 +465,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
                 "loss_weights",
                 "target_tokens",
                 "sample_indices",
+                *sampling_mask_keys,
             ],
             args.data_pad_size_multiplier,
             args.qkv_format,
diff --git a/miles/backends/training_utils/loss_hub/losses.py b/miles/backends/training_utils/loss_hub/losses.py
index d2d40b06936..60926e17bde 100644
--- a/miles/backends/training_utils/loss_hub/losses.py
+++ b/miles/backends/training_utils/loss_hub/losses.py
@@ -19,7 +19,9 @@
     compute_policy_loss,
 )
 from miles.backends.training_utils.parallel import get_parallel_state
+from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks
 from miles.utils.function_registry import load_function
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.types import RolloutBatch
 
 
@@ -111,6 +113,7 @@ def policy_loss_function(
     total_lengths = batch["total_lengths"]
     max_seq_lens = batch.get("max_seq_lens", None)
     calculate_entropy = args.entropy_coef != 0 or args.observe_training_entropy
+    rollout_sampling_mask = get_rollout_sampling_masks(batch) if sampling_support_replay_enabled(args) else None
 
     log_probs_and_entropy = get_log_probs_and_entropy(
         logits,
@@ -121,6 +124,7 @@ def policy_loss_function(
         with_entropy=calculate_entropy,
         entropy_requires_grad=args.entropy_coef != 0,
         max_seq_lens=max_seq_lens,
+        rollout_sampling_mask=rollout_sampling_mask,
     )
 
     log_probs = log_probs_and_entropy["log_probs"]
diff --git a/miles/backends/training_utils/sampling_mask.py b/miles/backends/training_utils/sampling_mask.py
index a189f13a532..ca285983d76 100644
--- a/miles/backends/training_utils/sampling_mask.py
+++ b/miles/backends/training_utils/sampling_mask.py
@@ -1,10 +1,26 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 
 import torch
 
 from miles.utils.sampling_mask import RolloutSamplingMask
 
 
+def get_rollout_sampling_masks(batch: Mapping[str, object]) -> list[RolloutSamplingMask]:
+    """Reconstruct the per-sample masks carried by the training-data wire."""
+    ids_batch = batch.get("rollout_sampling_mask_ids")
+    offsets_batch = batch.get("rollout_sampling_mask_offsets")
+    if ids_batch is None or offsets_batch is None:
+        raise ValueError("sampling-support actor scoring requires both sampling-mask wire fields")
+    if not isinstance(ids_batch, Sequence) or not isinstance(offsets_batch, Sequence):
+        raise TypeError("rollout sampling-mask ids and offsets must be sequences with one entry per sample")
+    if len(ids_batch) != len(offsets_batch):
+        raise ValueError(f"sampling-mask ids batch size {len(ids_batch)} != offsets batch size {len(offsets_batch)}")
+    return [
+        RolloutSamplingMask(ids=torch.as_tensor(ids), offsets=torch.as_tensor(offsets))
+        for ids, offsets in zip(ids_batch, offsets_batch, strict=True)
+    ]
+
+
 def build_local_sampling_mask(
     logits: torch.Tensor,
     sampling_mask: RolloutSamplingMask,
diff --git a/miles/utils/sampling_mask.py b/miles/utils/sampling_mask.py
index 568d0ebb396..164abf60684 100644
--- a/miles/utils/sampling_mask.py
+++ b/miles/utils/sampling_mask.py
@@ -5,6 +5,11 @@
 import torch
 
 
+def sampling_support_replay_enabled(args) -> bool:
+    """Whether rollout sampling uses a bounded support that must be replayed."""
+    return float(getattr(args, "rollout_top_p", 1.0)) < 1.0 or int(getattr(args, "rollout_top_k", -1)) > 0
+
+
 @dataclass(frozen=True, eq=False)
 class RolloutSamplingMask:
     """One sample's sampling mask: for each response position, the token ids
diff --git a/tests/fast/backends/megatron_utils/test_model.py b/tests/fast/backends/megatron_utils/test_model.py
index cea3a0fa231..c3e69f68c02 100644
--- a/tests/fast/backends/megatron_utils/test_model.py
+++ b/tests/fast/backends/megatron_utils/test_model.py
@@ -2,6 +2,7 @@
 from argparse import Namespace
 from dataclasses import dataclass, field
 from typing import Any
+from unittest.mock import Mock
 
 import pytest
 
@@ -120,3 +121,57 @@ def test_train_one_step_emits_the_train_tag_in_its_structured_event(
             )
 
         assert "train op=train_step rollout=7 step=3 attempt=2 outcome=NORMAL valid_step=true" in caplog.messages
+
+
+def test_forward_only_omits_sampling_mask_for_callbacks_that_do_not_replay_sampling_support(monkeypatch):
+    from miles.backends.megatron_utils import model as model_module
+
+    def forward_backward(**kwargs):
+        output, callback = kwargs["forward_step_func"](kwargs["data_iterator"][0], kwargs["model"][0])
+        return [callback(output)]
+
+    callback_kwargs = {}
+
+    def critic_callback(_output, **kwargs):
+        callback_kwargs.update(kwargs)
+        return {}
+
+    args = Namespace(
+        data_pad_size_multiplier=1,
+        qkv_format="thd",
+        allgather_cp=False,
+        enable_witness=False,
+        use_rollout_entropy=False,
+        custom_megatron_before_log_prob_hook_path=None,
+        seq_length=8,
+        micro_batch_size=1,
+    )
+    batch = {
+        "tokens": None,
+        "input_loss_masks": None,
+        "multimodal_train_inputs": None,
+        "unconcat_tokens": [],
+        "total_lengths": [],
+        "response_lengths": [],
+        "max_seq_lens": None,
+    }
+    dumper = Mock()
+    dumper.wrap_forward_step.side_effect = lambda fn: fn
+
+    monkeypatch.setattr(model_module, "DumperMegatronUtil", lambda *_args, **_kwargs: dumper)
+    monkeypatch.setattr(model_module, "get_model_config", lambda _model: Mock(timers=None))
+    monkeypatch.setattr(model_module, "get_batch", lambda *_args, **_kwargs: batch)
+    monkeypatch.setattr(model_module, "get_packed_seq_params", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(model_module, "get_forward_backward_func", lambda: forward_backward)
+    monkeypatch.setattr(model_module, "mpu", FakeMpu())
+
+    model_module.forward_only(
+        critic_callback,
+        args,
+        [Mock()],
+        [Mock()],
+        [1],
+        rollout_id=0,
+    )
+
+    assert "rollout_sampling_mask" not in callback_kwargs
diff --git a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py
index caa92f15b1a..aa80bb54f06 100644
--- a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py
+++ b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py
@@ -136,18 +136,25 @@ def test_train_keeps_model_resident(actor_module, monkeypatch):
     worker.sleep.assert_not_called()
 
 
-def test_compute_log_prob_keeps_logits_in_model_precision(actor_module, monkeypatch):
+@pytest.mark.parametrize(
+    ("store_prefix", "expected"),
+    [("", True), ("ref_", False), ("teacher_", False)],
+)
+def test_compute_log_prob_replays_sampling_support_only_for_actor_scores(
+    actor_module, monkeypatch, store_prefix, expected
+):
     worker = object.__new__(actor_module.MegatronTrainRayActor)
-    worker.args = Namespace()
+    worker.args = Namespace(rollout_top_p=0.95, rollout_top_k=32)
     worker.model = [object()]
     forward_only = Mock(return_value={"log_probs": []})
     monkeypatch.setattr(actor_module, "forward_only", forward_only)
     monkeypatch.setattr(actor_module, "timer", lambda _name: nullcontext())
 
-    result = worker.compute_log_prob([], [], rollout_id=3)
+    result = worker.compute_log_prob([], [], rollout_id=3, store_prefix=store_prefix)
 
     assert result == {"log_probs": []}
     assert forward_only.call_args.kwargs["fp32_output"] is False
+    assert forward_only.call_args.kwargs["use_rollout_sampling_mask"] is expected
 
 
 def test_save_model_does_not_manage_lifecycle(actor_module, monkeypatch):
diff --git a/tests/fast/backends/test_fsdp_train.py b/tests/fast/backends/test_fsdp_train.py
index 9191bd9a37a..646f8e8354e 100644
--- a/tests/fast/backends/test_fsdp_train.py
+++ b/tests/fast/backends/test_fsdp_train.py
@@ -1,8 +1,12 @@
 from argparse import Namespace
 from collections.abc import Iterator
 from contextlib import contextmanager, nullcontext
+from types import SimpleNamespace
 from unittest.mock import Mock
 
+import pytest
+import torch
+
 from miles.backends.fsdp_utils import actor as actor_module
 from miles.backends.megatron_utils.ft.types import TrainStepOutcome, TrainStepOutput
 
@@ -29,3 +33,57 @@ def test_fsdp_train_debug_rollout_only_returns_a_normal_output(monkeypatch):
 
     assert result == TrainStepOutput(outcome=TrainStepOutcome.NORMAL)
     actor._train_core.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("model_tag", "store_prefix", "top_p", "top_k", "expected"),
+    [
+        ("actor", "", 0.95, 32, True),
+        ("ref", "ref_", 0.95, 32, False),
+        ("actor", "", 1.0, -1, False),
+    ],
+)
+def test_fsdp_only_actor_scores_replay_rollout_sampling_support(
+    monkeypatch, model_tag, store_prefix, top_p, top_k, expected
+):
+    actor = object.__new__(actor_module.FSDPTrainRayActor)
+    actor.args = Namespace(
+        rollout_top_p=top_p,
+        rollout_top_k=top_k,
+        data_pad_size_multiplier=1,
+        qkv_format="thd",
+    )
+    actor.model = Mock(return_value=SimpleNamespace(logits=torch.zeros(1, 1, 4)))
+    actor.ref_model = None
+    actor.precision_policy = object()
+    actor.prof = SimpleNamespace(iterate_train_log_probs=lambda values: values)
+    actor._get_model_inputs_args = lambda _batch: {}
+    data_iterator = Mock()
+    sampling_masks = [object()]
+    captured = {}
+
+    monkeypatch.setattr(actor_module, "timer", _noop_timer)
+    monkeypatch.setattr(actor_module.dist, "get_rank", lambda: 0)
+    monkeypatch.setattr(actor_module, "precision_forward_context", lambda _policy: nullcontext())
+    monkeypatch.setattr(
+        actor_module,
+        "get_batch",
+        lambda *_args, **_kwargs: {
+            "unconcat_tokens": [torch.tensor([0])],
+            "total_lengths": [1],
+            "response_lengths": [1],
+            "max_seq_lens": None,
+        },
+    )
+    monkeypatch.setattr(actor_module, "get_rollout_sampling_masks", lambda _batch: sampling_masks)
+
+    def get_log_probs_and_entropy(**kwargs):
+        captured.update(kwargs)
+        return {"log_probs": [torch.zeros(1)]}
+
+    monkeypatch.setattr(actor_module, "get_log_probs_and_entropy", get_log_probs_and_entropy)
+    monkeypatch.setattr(actor_module, "aggregate_forward_results", lambda *_args, **_kwargs: {})
+
+    actor._compute_log_prob(model_tag, data_iterator, [1], store_prefix=store_prefix)
+
+    assert (captured["rollout_sampling_mask"] is sampling_masks) is expected
diff --git a/tests/fast/backends/training_utils/test_sampling_mask.py b/tests/fast/backends/training_utils/test_sampling_mask.py
index 0a70dce860f..47b90a47b75 100644
--- a/tests/fast/backends/training_utils/test_sampling_mask.py
+++ b/tests/fast/backends/training_utils/test_sampling_mask.py
@@ -6,10 +6,28 @@
 from miles.backends.training_utils import cp_utils
 from miles.backends.training_utils.loss_hub import logit_processors
 from miles.backends.training_utils.loss_hub.math_utils import _calculate_log_probs_and_entropy_true_on_policy
-from miles.backends.training_utils.sampling_mask import build_local_sampling_mask
+from miles.backends.training_utils.sampling_mask import build_local_sampling_mask, get_rollout_sampling_masks
 from miles.utils.sampling_mask import RolloutSamplingMask
 
 
+def test_get_rollout_sampling_masks_reconstructs_transport_batch():
+    masks = get_rollout_sampling_masks(
+        {
+            "rollout_sampling_mask_ids": [torch.tensor([1, 4, 2], dtype=torch.int32)],
+            "rollout_sampling_mask_offsets": [torch.tensor([0, 2, 3], dtype=torch.int64)],
+        }
+    )
+
+    ids, offsets = masks[0]._as_tensors()
+    assert ids.tolist() == [1, 4, 2]
+    assert offsets.tolist() == [0, 2, 3]
+
+
+def test_get_rollout_sampling_masks_requires_complete_transport_batch():
+    with pytest.raises(ValueError, match="requires both sampling-mask wire fields"):
+        get_rollout_sampling_masks({"rollout_sampling_mask_ids": [torch.tensor([1])]})
+
+
 def test_build_local_sampling_mask_selects_original_response_rows_and_tp_shard():
     logits = torch.zeros(2, 4)
     mask = build_local_sampling_mask(
@@ -101,6 +119,31 @@ def test_true_on_policy_masks_logprob_but_keeps_full_vocab_entropy():
     torch.testing.assert_close(entropy, expected_entropy)
 
 
+def test_sampling_support_normalization_has_the_expected_actor_gradient():
+    logits = torch.tensor(
+        [[2.0, 1.0, 0.0, -1.0], [-1.0, 2.0, 1.0, 0.0]],
+        dtype=torch.float64,
+        requires_grad=True,
+    )
+    tokens = torch.tensor([0, 1])
+    sampling_mask = torch.tensor([[True, False, True, False], [False, True, False, False]])
+
+    log_probs, _ = _calculate_log_probs_and_entropy_true_on_policy(
+        logits,
+        tokens,
+        None,
+        sampling_mask=sampling_mask,
+    )
+    log_probs.sum().backward()
+
+    first_support_probs = torch.softmax(logits.detach()[0, [0, 2]], dim=-1)
+    expected_grad = torch.zeros_like(logits)
+    expected_grad[0, 0] = 1 - first_support_probs[0]
+    expected_grad[0, 2] = -first_support_probs[1]
+
+    torch.testing.assert_close(logits.grad, expected_grad)
+
+
 def test_get_log_probs_and_entropy_applies_per_response_sampling_support(monkeypatch):
     parallel_state = SimpleNamespace(
         tp=SimpleNamespace(rank=0, group=None),

From 86562a3e43f14a82c4d5a1e5339ed57ebf0ceefd Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Sun, 20 Sep 2026 00:51:43 +0000
Subject: [PATCH 2/6] feat(rollout): capture native sampling support

---
 miles/rollout/generate_hub/multi_turn.py      |   7 +-
 miles/rollout/generate_hub/single_turn.py     |   6 +-
 .../generate_utils/generate_endpoint_utils.py |  10 +
 miles/rollout/generate_utils/sampling_mask.py | 128 ++++++++++--
 miles/rollout/sglang_rollout.py               |  18 +-
 .../rollout/generate_hub/test_single_turn.py  |   4 +-
 .../generate_utils/test_sampling_mask.py      | 188 +++++++++++++++++-
 7 files changed, 339 insertions(+), 22 deletions(-)

diff --git a/miles/rollout/generate_hub/multi_turn.py b/miles/rollout/generate_hub/multi_turn.py
index 17bea80cea2..0ee40f5cf5f 100644
--- a/miles/rollout/generate_hub/multi_turn.py
+++ b/miles/rollout/generate_hub/multi_turn.py
@@ -51,7 +51,12 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     for _turn in range(args.generate_max_turns):
         # ----------------------- Call inference endpoint -------------------------
 
-        payload, halt_status = compute_request_payload(args, sample.tokens, input.sampling_params)
+        payload, halt_status = compute_request_payload(
+            args,
+            sample.tokens,
+            input.sampling_params,
+            evaluation=input.evaluation,
+        )
         if payload is None:
             sample.status = halt_status
             break
diff --git a/miles/rollout/generate_hub/single_turn.py b/miles/rollout/generate_hub/single_turn.py
index bc744572de3..bab68533115 100644
--- a/miles/rollout/generate_hub/single_turn.py
+++ b/miles/rollout/generate_hub/single_turn.py
@@ -35,7 +35,11 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         input_ids = prompt_ids
 
     payload, halt_status = compute_request_payload(
-        args, input_ids=input_ids, sampling_params=sampling_params, multimodal_inputs=sample.multimodal_inputs
+        args,
+        input_ids=input_ids,
+        sampling_params=sampling_params,
+        multimodal_inputs=sample.multimodal_inputs,
+        evaluation=input.evaluation,
     )
     if payload is None:
         sample.status = halt_status
diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py
index 9279b4f7a5b..fada2bcfc40 100644
--- a/miles/rollout/generate_utils/generate_endpoint_utils.py
+++ b/miles/rollout/generate_utils/generate_endpoint_utils.py
@@ -8,6 +8,7 @@
 import numpy as np
 import pybase64
 
+from miles.rollout.generate_utils.sampling_mask import append_sampling_metadata, should_return_sampling_mask
 from miles.utils.lora import LORA_ADAPTER_NAME, lora_rollout_enabled
 from miles.utils.processing_utils import encode_image_for_rollout_engine, extract_multimodal_train_inputs
 from miles.utils.types import Sample
@@ -54,6 +55,8 @@ def compute_request_payload(
     input_ids: list[int],
     sampling_params: dict,
     multimodal_inputs: dict | None = None,
+    *,
+    evaluation: bool = False,
 ) -> tuple[dict[str, Any] | None, Sample.Status | None]:
     sampling_params = deepcopy(sampling_params)
     max_new_tokens = sampling_params.pop("max_new_tokens", args.rollout_max_response_len)
@@ -62,6 +65,8 @@ def compute_request_payload(
     if max_new_tokens <= 0:
         return None, Sample.Status.TRUNCATED
 
+    return_sampling_mask = should_return_sampling_mask(args, sampling_params, evaluation=evaluation)
+
     payload = {
         "input_ids": input_ids,
         "sampling_params": {**sampling_params, "max_new_tokens": max_new_tokens},
@@ -69,6 +74,8 @@ def compute_request_payload(
         "return_routed_experts": args.use_rollout_routing_replay,
         "return_indexer_topk": args.use_rollout_indexer_replay,
     }
+    if return_sampling_mask:
+        payload["return_sampling_mask"] = True
     if lora_rollout_enabled(args):
         payload["lora_path"] = LORA_ADAPTER_NAME
     if image_data := (multimodal_inputs or {}).get("images"):
@@ -90,6 +97,9 @@ async def update_sample_from_response(
     else:
         new_response_tokens, new_response_log_probs = [], []
 
+    if payload.get("return_sampling_mask", False):
+        new_response_log_probs = append_sampling_metadata(sample, new_response_tokens, output["meta_info"])
+
     # Update sample with tokens directly - avoiding re-tokenization
     sample.tokens = sample.tokens + new_response_tokens
     sample.response_length += len(new_response_tokens)
diff --git a/miles/rollout/generate_utils/sampling_mask.py b/miles/rollout/generate_utils/sampling_mask.py
index e26a6a462ff..4360e286cab 100644
--- a/miles/rollout/generate_utils/sampling_mask.py
+++ b/miles/rollout/generate_utils/sampling_mask.py
@@ -1,24 +1,109 @@
-from collections.abc import Sequence
+from argparse import Namespace
+from collections.abc import Mapping, Sequence
 
-from miles.utils.sampling_mask import RolloutSamplingMask
+from miles.utils.sampling_mask import RolloutSamplingMask, sampling_support_replay_enabled
 from miles.utils.types import Sample
 
 
-def append_forced_sampling_tokens(sample: Sample, token_ids: Sequence[int]) -> None:
-    """Record singleton support for non-sampled tokens inserted by the environment."""
-    sampling_mask = RolloutSamplingMask.from_mask_list([[int(token_id)] for token_id in token_ids])
-    if sample.rollout_sampling_mask is None:
-        if sample.response_length != 0:
-            raise ValueError("cannot initialize a sampling mask after response tokens have already been appended")
-        sample.rollout_sampling_mask = sampling_mask
-        return
+def should_return_sampling_mask(
+    args: Namespace,
+    sampling_params: Mapping[str, object] | None = None,
+    *,
+    evaluation: bool = False,
+) -> bool:
+    """Validate whether a training request must return its realized sampling support."""
+    if evaluation:
+        return False
 
-    if len(sample.rollout_sampling_mask) != sample.response_length:
+    params = sampling_params or {}
+    configured_top_p = float(getattr(args, "rollout_top_p", 1.0))
+    request_top_p = float(configured_top_p if params.get("top_p") is None else params["top_p"])
+    if not 0.0 < request_top_p <= 1.0:
+        raise ValueError(f"training request top_p must be in (0, 1], got {request_top_p}")
+
+    if not sampling_support_replay_enabled(args):
+        raw_top_k = params.get("top_k")
+        request_top_k = -1 if raw_top_k is None else int(raw_top_k)
+        if request_top_p < 1.0 or request_top_k > 0:
+            raise ValueError("bounded training-request sampling requires bounded rollout sampling")
+        return False
+
+    missing_params = [name for name in ("top_p", "top_k", "temperature") if params.get(name) is None]
+    if missing_params:
+        raise ValueError(f"sampling-support replay requires explicit request parameters: {', '.join(missing_params)}")
+
+    request_top_k = int(params["top_k"])
+    if request_top_k <= 0:
+        raise ValueError(f"training request top_k must be positive, got {request_top_k}")
+
+    configured_temperature = float(getattr(args, "rollout_temperature", 1.0))
+    request_temperature = float(params["temperature"])
+    if request_temperature != configured_temperature:
         raise ValueError(
-            f"sampling mask length {len(sample.rollout_sampling_mask)} is not aligned with "
-            f"response_length {sample.response_length} before appending"
+            f"request temperature {request_temperature} does not match --rollout-temperature {configured_temperature}"
         )
-    sample.rollout_sampling_mask = RolloutSamplingMask.concatenate((sample.rollout_sampling_mask, sampling_mask))
+
+    unsupported = {
+        "frequency_penalty": (0, 0.0, None),
+        "presence_penalty": (0, 0.0, None),
+        "repetition_penalty": (1, 1.0, None),
+        "logit_bias": ({}, None),
+    }
+    for name, allowed_values in unsupported.items():
+        if params.get(name) not in allowed_values:
+            raise ValueError(
+                f"{name} is not supported with sampling-support replay because "
+                "the trainer cannot reproduce its logit transformation"
+            )
+    return True
+
+
+def _sampling_mask_from_supports(
+    token_ids: Sequence[int],
+    supports: Sequence[Sequence[int]],
+) -> RolloutSamplingMask:
+    if len(token_ids) != len(supports):
+        raise ValueError(f"sampling support length {len(supports)} != token length {len(token_ids)}")
+
+    for token_id, support in zip(token_ids, supports, strict=True):
+        if not support:
+            raise ValueError("sampling support must contain at least one token")
+        if int(token_id) not in support:
+            raise ValueError(f"sampled token {token_id} is absent from its sampling support")
+    return RolloutSamplingMask.from_mask_list(supports)
+
+
+def append_sampling_metadata(
+    sample: Sample,
+    output_token_ids: Sequence[int],
+    meta_info: dict,
+    *,
+    aborted: bool = False,
+) -> list[float]:
+    """Append SGLang's realized support and return its normalized log-probs."""
+    supports = meta_info.get("output_token_sampling_mask")
+    log_probs = meta_info.get("output_token_sampling_logprobs")
+    if supports is None or log_probs is None:
+        finish_reason = meta_info.get("finish_reason") or {}
+        if (aborted or finish_reason.get("type") == "abort") and not output_token_ids:
+            _append_sampling_mask(sample, RolloutSamplingMask.from_mask_list([]))
+            return []
+        raise ValueError(
+            "SGLang response is missing output_token_sampling_mask or "
+            "output_token_sampling_logprobs; use an SGLang build with the "
+            "native return_sampling_mask primitive"
+        )
+    if len(log_probs) != len(output_token_ids):
+        raise ValueError(f"sampling log-prob length {len(log_probs)} != output token length {len(output_token_ids)}")
+
+    _append_sampling_mask(sample, _sampling_mask_from_supports(output_token_ids, supports))
+    return [float(value) for value in log_probs]
+
+
+def append_forced_sampling_tokens(sample: Sample, token_ids: Sequence[int]) -> None:
+    """Record singleton support for non-sampled tokens inserted by the environment."""
+    sampling_mask = RolloutSamplingMask.from_mask_list([[int(token_id)] for token_id in token_ids])
+    _append_sampling_mask(sample, sampling_mask)
 
 
 def merge_sampling_masks(
@@ -36,3 +121,18 @@ def merge_sampling_masks(
 
     observation_mask = RolloutSamplingMask.from_mask_list([[int(token_id)] for token_id in observation_token_ids])
     return RolloutSamplingMask.concatenate((first_mask, observation_mask, second_mask))
+
+
+def _append_sampling_mask(sample: Sample, sampling_mask: RolloutSamplingMask) -> None:
+    if sample.rollout_sampling_mask is None:
+        if sample.response_length != 0:
+            raise ValueError("cannot initialize a sampling mask after response tokens have already been appended")
+        sample.rollout_sampling_mask = sampling_mask
+        return
+
+    if len(sample.rollout_sampling_mask) != sample.response_length:
+        raise ValueError(
+            f"sampling mask length {len(sample.rollout_sampling_mask)} is not aligned with "
+            f"response_length {sample.response_length} before appending"
+        )
+    sample.rollout_sampling_mask = RolloutSamplingMask.concatenate((sample.rollout_sampling_mask, sampling_mask))
diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py
index 40935a8b632..86c6e2ad6cb 100644
--- a/miles/rollout/sglang_rollout.py
+++ b/miles/rollout/sglang_rollout.py
@@ -42,6 +42,7 @@
 )
 from .generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill
 from .generate_utils.sample_utils import reward_log_summary, sample_text_preview
+from .generate_utils.sampling_mask import append_sampling_metadata, should_return_sampling_mask
 from .rm_hub import async_rm, batched_async_rm
 
 __all__ = ["generate_rollout", "get_model_url"]
@@ -139,7 +140,13 @@ def submit_generate_tasks(self, samples: list[list[Sample]]) -> None:
         self.remaining_batch_size += len(samples)
 
 
-async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample:
+async def generate(
+    args: Namespace,
+    sample: Sample,
+    sampling_params: dict[str, Any],
+    *,
+    evaluation: bool = False,
+) -> Sample:
     """Generate using traditional SGLang router with token-based workflow"""
     if args.ci_test:
         assert isinstance(sample.prompt, str)
@@ -172,11 +179,15 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         sample.status = Sample.Status.TRUNCATED
         return sample
 
+    return_sampling_mask = should_return_sampling_mask(args, sampling_params, evaluation=evaluation)
+
     # Prepare payload for sglang server
     payload = {
         "sampling_params": sampling_params,
         "return_logprob": True,
     }
+    if return_sampling_mask:
+        payload["return_sampling_mask"] = True
     opd_top_k = getattr(args, "opd_log_prob_top_k", 0) or 0
     opd_top_k_strategy = getattr(args, "opd_top_k_strategy", "only-student")
     if getattr(args, "use_opd", False) and opd_top_k > 0 and opd_top_k_strategy != "only-teacher":
@@ -224,6 +235,9 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
     else:
         new_response_tokens, new_response_log_probs = [], []
 
+    if payload.get("return_sampling_mask", False):
+        new_response_log_probs = append_sampling_metadata(sample, new_response_tokens, output["meta_info"])
+
     # Update sample with tokens directly - avoiding re-tokenization
     sample.tokens = sample.tokens + new_response_tokens
     sample.response_length += len(new_response_tokens)
@@ -309,7 +323,7 @@ async def generate_and_rm(
                 )
                 sample = output.samples
             else:
-                sample = await generate(args, sample, sampling_params)
+                sample = await generate(args, sample, sampling_params, evaluation=evaluation)
 
     if sink is not None:
         sink.attempt_end(sample)
diff --git a/tests/fast/rollout/generate_hub/test_single_turn.py b/tests/fast/rollout/generate_hub/test_single_turn.py
index 8c3f221354f..f126d078873 100644
--- a/tests/fast/rollout/generate_hub/test_single_turn.py
+++ b/tests/fast/rollout/generate_hub/test_single_turn.py
@@ -300,10 +300,10 @@ def test_rejected_statuses(self, variant, generation_env, status):
 class TestPayloadStructure:
     def test_sampling_params_passed_through(self, variant, generation_env):
         result = _run_generate(
-            variant, generation_env, sampling_params={"max_new_tokens": 16, "temperature": 0.5, "top_p": 0.9}
+            variant, generation_env, sampling_params={"max_new_tokens": 16, "temperature": 0.5, "top_p": 1.0}
         )
         assert result.requests == [
-            expected_request(variant, sampling_params={"max_new_tokens": 16, "temperature": 0.5, "top_p": 0.9})
+            expected_request(variant, sampling_params={"max_new_tokens": 16, "temperature": 0.5, "top_p": 1.0})
         ]
         assert listify(result.sample) == [expected_sample(variant)]
 
diff --git a/tests/fast/rollout/generate_utils/test_sampling_mask.py b/tests/fast/rollout/generate_utils/test_sampling_mask.py
index a9e8304e0e3..fa9f14a61b0 100644
--- a/tests/fast/rollout/generate_utils/test_sampling_mask.py
+++ b/tests/fast/rollout/generate_utils/test_sampling_mask.py
@@ -1,8 +1,173 @@
-from miles.rollout.generate_utils.sampling_mask import append_forced_sampling_tokens, merge_sampling_masks
-from miles.utils.sampling_mask import RolloutSamplingMask
+from types import SimpleNamespace
+
+import pytest
+
+from miles.rollout.generate_utils.generate_endpoint_utils import compute_request_payload
+from miles.rollout.generate_utils.sampling_mask import (
+    append_forced_sampling_tokens,
+    append_sampling_metadata,
+    merge_sampling_masks,
+    should_return_sampling_mask,
+)
+from miles.utils.sampling_mask import RolloutSamplingMask, sampling_support_replay_enabled
 from miles.utils.types import Sample
 
 
+@pytest.mark.parametrize(
+    ("rollout_top_p", "rollout_top_k", "request_top_p", "request_top_k", "expected"),
+    [
+        (0.95, 32, 0.95, 32, True),
+        (1.0, 32, 1.0, 32, True),
+        (1.0, -1, 1.0, -1, False),
+    ],
+)
+def test_generate_payload_automatically_requests_sampling_mask(
+    rollout_top_p,
+    rollout_top_k,
+    request_top_p,
+    request_top_k,
+    expected,
+):
+    args = SimpleNamespace(
+        rollout_top_p=rollout_top_p,
+        rollout_top_k=rollout_top_k,
+        rollout_temperature=1.0,
+        rollout_max_response_len=16,
+        rollout_max_context_len=None,
+        use_rollout_routing_replay=False,
+        use_rollout_indexer_replay=False,
+    )
+
+    payload, halt_status = compute_request_payload(
+        args,
+        input_ids=[1, 2],
+        sampling_params={
+            "max_new_tokens": 4,
+            "top_p": request_top_p,
+            "top_k": request_top_k,
+            "temperature": 1.0,
+        },
+    )
+
+    assert halt_status is None
+    assert payload.get("return_sampling_mask", False) is expected
+
+
+def test_enabled_replay_accepts_request_top_k_above_default():
+    args = SimpleNamespace(rollout_top_p=1.0, rollout_top_k=32, rollout_temperature=1.0)
+
+    assert should_return_sampling_mask(args, {"top_p": 1.0, "top_k": 64, "temperature": 1.0}) is True
+
+
+def test_unbounded_run_rejects_bounded_training_request():
+    args = SimpleNamespace(rollout_top_p=1.0, rollout_top_k=-1, rollout_temperature=1.0)
+
+    with pytest.raises(ValueError, match="bounded training-request sampling requires bounded rollout sampling"):
+        should_return_sampling_mask(args, {"top_p": 1.0, "top_k": 32, "temperature": 1.0})
+
+
+def test_unbounded_run_treats_an_unset_request_top_k_as_unbounded():
+    args = SimpleNamespace(rollout_top_p=1.0, rollout_top_k=-1, rollout_temperature=1.0)
+
+    assert should_return_sampling_mask(args, {"top_p": 1.0, "top_k": None}) is False
+
+
+def test_evaluation_does_not_request_or_validate_training_sampling_support():
+    args = SimpleNamespace(
+        rollout_top_p=0.95,
+        rollout_top_k=32,
+        rollout_temperature=1.0,
+        rollout_max_response_len=16,
+        rollout_max_context_len=None,
+        use_rollout_routing_replay=False,
+        use_rollout_indexer_replay=False,
+    )
+
+    payload, halt_status = compute_request_payload(
+        args,
+        input_ids=[1, 2],
+        sampling_params={"max_new_tokens": 4, "top_p": 1.0, "top_k": -1, "temperature": 0.5},
+        evaluation=True,
+    )
+
+    assert halt_status is None
+    assert "return_sampling_mask" not in payload
+
+
+def test_training_request_temperature_must_match_actor_scoring_temperature():
+    args = SimpleNamespace(rollout_top_p=0.95, rollout_top_k=32, rollout_temperature=1.0)
+
+    with pytest.raises(ValueError, match="request temperature 0.5"):
+        should_return_sampling_mask(args, {"top_p": 0.95, "top_k": 32, "temperature": 0.5})
+
+
+@pytest.mark.parametrize("missing_param", ["top_p", "top_k", "temperature"])
+def test_top_p_sampling_replay_requires_explicit_request_parameters(missing_param):
+    args = SimpleNamespace(rollout_top_p=0.95, rollout_top_k=32, rollout_temperature=1.0)
+    params = {"top_p": 0.95, "top_k": 32, "temperature": 1.0}
+    del params[missing_param]
+
+    with pytest.raises(ValueError, match=rf"explicit request parameters:.*{missing_param}"):
+        should_return_sampling_mask(args, params)
+
+
+@pytest.mark.parametrize("request_top_k", [-1, 0])
+def test_training_request_top_k_must_be_positive(request_top_k):
+    args = SimpleNamespace(rollout_top_p=0.95, rollout_top_k=32, rollout_temperature=1.0)
+
+    with pytest.raises(ValueError, match="request top_k must be positive"):
+        should_return_sampling_mask(
+            args,
+            {"top_p": 0.95, "top_k": request_top_k, "temperature": 1.0},
+        )
+
+
+@pytest.mark.parametrize(
+    ("name", "value"),
+    [
+        ("frequency_penalty", 0.1),
+        ("presence_penalty", 0.1),
+        ("repetition_penalty", 1.1),
+        ("logit_bias", {"1": 0.1}),
+    ],
+)
+def test_training_request_rejects_unreplayed_logit_transform(name, value):
+    args = SimpleNamespace(rollout_top_p=0.95, rollout_top_k=32, rollout_temperature=1.0)
+
+    with pytest.raises(ValueError, match=rf"{name} is not supported"):
+        should_return_sampling_mask(
+            args,
+            {"top_p": 0.95, "top_k": 32, "temperature": 1.0, name: value},
+        )
+
+
+@pytest.mark.parametrize(
+    ("top_p", "top_k", "expected"),
+    [(0.95, -1, True), (1.0, 32, True), (1.0, -1, False)],
+)
+def test_sampling_support_replay_enabled_by_bounded_filter(top_p, top_k, expected):
+    args = SimpleNamespace(rollout_top_p=top_p, rollout_top_k=top_k)
+    assert sampling_support_replay_enabled(args) is expected
+
+
+def test_append_sampling_metadata_preserves_ragged_support_and_native_logprobs():
+    sample = Sample(tokens=[1, 2])
+    meta_info = {
+        "output_token_sampling_mask": [[10, 4, 7], [11, 3]],
+        "output_token_sampling_logprobs": [-0.25, -0.5],
+    }
+
+    log_probs = append_sampling_metadata(sample, [10, 11], meta_info)
+    sample.tokens.extend([10, 11])
+    sample.response_length = 2
+
+    assert log_probs == [-0.25, -0.5]
+    ids, offsets = sample.rollout_sampling_mask._as_tensors()
+    assert ids.tolist() == [10, 4, 7, 11, 3]
+    assert offsets.tolist() == [0, 3, 5]
+    sample.validate()
+
+
 def test_forced_tokens_append_singleton_support_and_strip_cleanly():
     sample = Sample(
         tokens=[1, 10],
@@ -41,3 +206,22 @@ def test_merge_sampling_masks_inserts_singleton_observation_supports():
 
     assert ids.tolist() == [10, 4, 20, 21, 30, 7, 8]
     assert offsets.tolist() == [0, 2, 3, 4, 7]
+
+
+def test_append_sampling_metadata_rejects_support_without_sampled_token():
+    with pytest.raises(ValueError, match="sampled token 10 is absent"):
+        append_sampling_metadata(
+            Sample(),
+            [10],
+            {
+                "output_token_sampling_mask": [[4, 7]],
+                "output_token_sampling_logprobs": [-0.25],
+            },
+        )
+
+
+def test_abort_before_sampling_does_not_require_sampling_metadata():
+    sample = Sample()
+
+    assert append_sampling_metadata(sample, [], {"finish_reason": {"type": "abort"}}) == []
+    assert len(sample.rollout_sampling_mask) == 0

From 38f60d7f5c29be970dcb2c14d3beeada3dd31d94 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Sun, 20 Sep 2026 00:51:55 +0000
Subject: [PATCH 3/6] feat(session): preserve rollout sampling support

---
 .../rollout/generate_hub/agentic_tool_call.py | 15 ++++-
 .../generate_utils/openai_endpoint_utils.py   |  7 ++-
 miles/rollout/session/config.py               |  6 ++
 miles/rollout/session/core.py                 | 27 ++++++--
 miles/rollout/session/request_args.py         | 24 +++++++
 miles/rollout/session/samples/merge.py        | 11 +++-
 miles/rollout/session/v2/core.py              | 14 +++--
 .../sglang/test_session_server_addition_r3.py |  3 +
 tests/fast/fixtures/session_fixtures.py       |  3 +
 tests/fast/ray/rollout/conftest.py            |  2 +
 .../rollout/generate_hub/test_agentic_v2.py   | 29 +++++++++
 .../test_openai_endpoint_utils.py             | 25 +++++++-
 tests/fast/rollout/session/test_config.py     |  9 +++
 .../fast/rollout/session/test_request_args.py | 20 ++++++
 tests/fast/rollout/session/test_samples.py    | 63 +++++++++++++++++++
 tests/fast/router/test_sessions_v2.py         |  3 +
 tests/fast/utils/workers/test_argv_utils.py   |  3 +
 .../session/bench_session_server_overhead.py  |  3 +
 18 files changed, 250 insertions(+), 17 deletions(-)

diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py
index fdf32a0a815..c4f856ad558 100644
--- a/miles/rollout/generate_hub/agentic_tool_call.py
+++ b/miles/rollout/generate_hub/agentic_tool_call.py
@@ -35,8 +35,10 @@ async def my_agent(
 
 from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
 from miles.rollout.generate_utils.openai_endpoint_utils import OpenAIEndpointTracer
+from miles.rollout.generate_utils.sampling_mask import should_return_sampling_mask
 from miles.rollout.session.v2.metrics import SESSION_ROLLOUT_METRICS_KEY
 from miles.utils.function_registry import load_function
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.types import Sample
 
 logger = logging.getLogger(__name__)
@@ -75,10 +77,21 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     t_start = time.monotonic()
     try:
         logger.debug(f"{log_prefix} Starting agent function call")
+        request_kwargs = build_chat_request_kwargs(input.sampling_params)
+        return_sampling_mask = should_return_sampling_mask(
+            input.args,
+            input.sampling_params,
+            evaluation=input.evaluation,
+        )
+        if return_sampling_mask:
+            request_kwargs["return_sampling_mask"] = True
+        elif input.evaluation and sampling_support_replay_enabled(input.args):
+            # The session server otherwise fills the run's training defaults.
+            request_kwargs["return_sampling_mask"] = False
         agent_metadata = await custom_agent_function(
             base_url=tracer.base_url,
             prompt=input.sample.prompt,
-            request_kwargs=build_chat_request_kwargs(input.sampling_params),
+            request_kwargs=request_kwargs,
             metadata=metadata,
         )
         logger.debug(f"{log_prefix} Agent function returned in {time.monotonic()-t_start:.1f}s")
diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py
index 26c18feed2c..3fea2febac2 100644
--- a/miles/rollout/generate_utils/openai_endpoint_utils.py
+++ b/miles/rollout/generate_utils/openai_endpoint_utils.py
@@ -10,10 +10,12 @@
 from miles.rollout.session.samples.codec import (
     COMPUTED_FIELDS,
     COMPUTED_FIELDS_V2,
+    ROLLOUT_SAMPLING_MASK_FIELDS,
     SamplesReply,
     decode_samples_and_merge_input_sample,
 )
 from miles.utils.http_utils import post, post_bytes_no_retry
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.types import Sample
 
 logger = logging.getLogger(__name__)
@@ -59,11 +61,14 @@ async def create(args: Namespace):
         response = await post(f"{session_url}/sessions", {}, action="post")
         session_id = response["session_id"]
         use_v2 = getattr(args, "use_session_server", None) == "v2"
+        samples_wire_fields = COMPUTED_FIELDS_V2 if use_v2 else COMPUTED_FIELDS
+        if sampling_support_replay_enabled(args):
+            samples_wire_fields += ROLLOUT_SAMPLING_MASK_FIELDS
         return OpenAIEndpointTracer(
             router_url=session_url,
             session_id=session_id,
             session_server_instance_id=session_server_instance_id,
-            samples_wire_fields=COMPUTED_FIELDS_V2 if use_v2 else COMPUTED_FIELDS,
+            samples_wire_fields=samples_wire_fields,
         )
 
     async def collect_samples(
diff --git a/miles/rollout/session/config.py b/miles/rollout/session/config.py
index 58e6a7ec7d0..31ad5afef17 100644
--- a/miles/rollout/session/config.py
+++ b/miles/rollout/session/config.py
@@ -15,6 +15,9 @@ class SessionServerConfig(FrozenStrictBaseModel):
     apply_chat_template_kwargs: dict[str, Any] | None
     use_rollout_routing_replay: bool
     use_rollout_indexer_replay: bool
+    rollout_temperature: float
+    rollout_top_p: float
+    rollout_top_k: int
     sglang_speculative_algorithm: str | None
     num_layers: int | None
     moe_router_topk: int | None
@@ -44,6 +47,9 @@ def compute_session_server_config(
         apply_chat_template_kwargs=args.apply_chat_template_kwargs,
         use_rollout_routing_replay=args.use_rollout_routing_replay,
         use_rollout_indexer_replay=args.use_rollout_indexer_replay,
+        rollout_temperature=args.rollout_temperature,
+        rollout_top_p=args.rollout_top_p,
+        rollout_top_k=args.rollout_top_k,
         sglang_speculative_algorithm=args.sglang_speculative_algorithm,
         num_layers=getattr(args, "num_layers", None),
         moe_router_topk=getattr(args, "moe_router_topk", None),
diff --git a/miles/rollout/session/core.py b/miles/rollout/session/core.py
index 336340caaa0..2f16e3a74e5 100644
--- a/miles/rollout/session/core.py
+++ b/miles/rollout/session/core.py
@@ -2,7 +2,8 @@
 
 HTTP-agnostic: the FastAPI adapter (``sessions.py`` + ``server.py``) turns each request into primitives and calls these methods. Owns one ``SessionRegistry`` (per-session TITO/trajectory state) and one proxy ``backend``.
 
-- ``chat_completions`` strips the R3 replay payloads (``routed_experts`` / ``indexer_topk``) from the client reply copy-on-write; the ``SessionRecord`` keeps the full response for the training path (``GET /sessions/{id}``).
+- ``chat_completions`` strips training-only replay payloads from the client reply copy-on-write; the
+  ``SessionRecord`` keeps the full response for the training path (``GET /sessions/{id}``).
 - ``chat_completions`` holds the per-session lock for prep and state update but not across the proxy call; ``closing`` re-checks and the ``num_assistant`` check gate concurrent DELETE/chat.
 - ``stream: true`` is served as fake streaming: the backend call stays non-streaming (TITO needs the complete message + meta_info) and the full response is re-rendered as a single SSE chunk plus ``data: [DONE]``. Errors all happen before the SSE body is built, so they keep their real status codes as JSON.
 - ``collect_samples`` assembles training Samples from the session's records on the server (compute -> truncate -> merge, synchronously on the loop like the lock-free ``get_session``); deterministic assembly failures return 422 with the assertion text.
@@ -20,13 +21,14 @@
 from miles.rollout.session.errors import SessionNotFoundError, TokenizationError, UpstreamResponseError
 from miles.rollout.session.linear_trajectory import SessionRegistry
 from miles.rollout.session.request_args import filter_turn_args, parse_chat_request
-from miles.rollout.session.samples.codec import encode_samples
+from miles.rollout.session.samples.codec import COMPUTED_FIELDS, ROLLOUT_SAMPLING_MASK_FIELDS, encode_samples
 from miles.rollout.session.samples.merge import (
     compute_samples_from_openai_records,
     merge_samples_with_addition_r3,
     truncate_samples_by_total_tokens,
 )
 from miles.rollout.session.types import GetSessionResponse, SessionRecord
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 
 logger = logging.getLogger(__name__)
 
@@ -67,7 +69,13 @@ def _samples_response(payload: bytes) -> Response:
     return Response(content=payload, status_code=200, media_type="application/octet-stream")
 
 
-_CLIENT_STRIPPED_META_KEYS = ("routed_experts", "indexer_topk")
+_CLIENT_STRIPPED_META_KEYS = (
+    "routed_experts",
+    "indexer_topk",
+    "output_token_sampling_mask",
+    "output_token_sampling_logprobs",
+    "output_token_sampling_mask_length",
+)
 
 
 def _strip_replay_payloads(response: dict) -> dict:
@@ -205,6 +213,9 @@ def __init__(
         # Derived from pause_generation_mode at server bootstrap; session code
         # must depend on this capability, never on the weight-update mode.
         self.use_addition_r3 = use_addition_r3
+        self.samples_wire_fields = COMPUTED_FIELDS
+        if sampling_support_replay_enabled(config):
+            self.samples_wire_fields += ROLLOUT_SAMPLING_MASK_FIELDS
 
     def _maybe_request_addition_r3(
         self, request_body: dict, checkpoint_token_ids: list[int], prompt_token_ids: list[int]
@@ -269,7 +280,9 @@ async def collect_samples(self, session_id: str, *, max_seq_len: int | None) ->
         metadata = self._session_metadata(session_id, session)
         tokenizer = self.registry.tokenizer
         if not session.records:
-            return _samples_response(encode_samples([], metadata, empty_reason="no_records"))
+            return _samples_response(
+                encode_samples([], metadata, empty_reason="no_records", fields=self.samples_wire_fields)
+            )
         try:
             samples = compute_samples_from_openai_records(
                 self.config,
@@ -282,14 +295,16 @@ async def collect_samples(self, session_id: str, *, max_seq_len: int | None) ->
             if max_seq_len is not None:
                 samples = truncate_samples_by_total_tokens(samples, max_seq_len, tokenizer)
             if not samples:
-                return _samples_response(encode_samples([], metadata, empty_reason="all_truncated"))
+                return _samples_response(
+                    encode_samples([], metadata, empty_reason="all_truncated", fields=self.samples_wire_fields)
+                )
             if self.use_addition_r3:
                 samples = [merge_samples_with_addition_r3(self.config, samples, session.records, tokenizer)]
             else:
                 samples = [merge_samples(samples, tokenizer)]
         except (AssertionError, ValueError) as exc:
             return Response(content=str(exc).encode(), status_code=422, media_type="text/plain")
-        return _samples_response(encode_samples(samples, metadata))
+        return _samples_response(encode_samples(samples, metadata, fields=self.samples_wire_fields))
 
     async def delete_session(self, session_id: str) -> Response:
         session = self.registry.get_session(session_id)
diff --git a/miles/rollout/session/request_args.py b/miles/rollout/session/request_args.py
index aabd6479789..3660160e41b 100644
--- a/miles/rollout/session/request_args.py
+++ b/miles/rollout/session/request_args.py
@@ -14,10 +14,12 @@
 from dataclasses import dataclass
 from typing import Any
 
+from miles.rollout.generate_utils.sampling_mask import should_return_sampling_mask
 from miles.rollout.session.config import SessionServerConfig
 from miles.rollout.session.errors import MessageValidationError
 from miles.utils.chat_template_utils.tito_tokenizer import TITOTokenizer, extract_template_args
 from miles.utils.lora import LORA_ADAPTER_NAME, lora_rollout_enabled
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 
 DEFAULT_TURN_ARGS_DROP_KEYS = ("input_ids", "messages")
 
@@ -79,6 +81,28 @@ def resolve_request_args_by_config(
     # TITO needs these on every request: agent-side overrides would break token accumulation.
     request_args["logprobs"] = True
     request_args["return_meta_info"] = True
+
+    sampling_mask_requested = request_args.get("return_sampling_mask")
+    if sampling_mask_requested is False:
+        return_sampling_mask = False
+    else:
+        if sampling_support_replay_enabled(config):
+            for name, value in (
+                ("top_p", config.rollout_top_p),
+                ("top_k", config.rollout_top_k),
+                ("temperature", config.rollout_temperature),
+            ):
+                if request_args.get(name) is None:
+                    request_args[name] = value
+        try:
+            return_sampling_mask = should_return_sampling_mask(config, request_args)
+        except ValueError as exc:
+            raise MessageValidationError(str(exc)) from exc
+    if sampling_mask_requested is True and not return_sampling_mask:
+        raise MessageValidationError("return_sampling_mask requires bounded rollout sampling")
+    if return_sampling_mask:
+        request_args["return_sampling_mask"] = True
+
     # Must be False so stop-token text is trimmed from assistant content;
     # token IDs still come from logprobs below.
     request_args["no_stop_trim"] = False
diff --git a/miles/rollout/session/samples/merge.py b/miles/rollout/session/samples/merge.py
index 3b2456bda19..e993302cf87 100644
--- a/miles/rollout/session/samples/merge.py
+++ b/miles/rollout/session/samples/merge.py
@@ -16,6 +16,7 @@
     get_routed_experts_from_response,
 )
 from miles.rollout.generate_utils.sample_utils import merge_samples
+from miles.rollout.generate_utils.sampling_mask import append_sampling_metadata
 from miles.rollout.session.types import SessionRecord
 from miles.utils.lifecycle import attach_lifecycle_metadata
 from miles.utils.types import Sample, WeightVersionsPerCall
@@ -102,6 +103,7 @@ def _compute_sample_from_openai_record(
     args: Namespace, record: SessionRecord, tokenizer, trim_count: int = 0, *, use_addition_r3: bool = False
 ) -> Sample:
     choice = record.response["choices"][0]
+    finish_reason = choice.get("finish_reason")
 
     prompt_token_ids = record.request.get("input_ids")
     if prompt_token_ids is None:
@@ -111,6 +113,13 @@ def _compute_sample_from_openai_record(
     output_log_probs = [item[0] for item in choice["meta_info"]["output_token_logprobs"]]
 
     sample = Sample()
+    if record.request.get("return_sampling_mask", False):
+        output_log_probs = append_sampling_metadata(
+            sample,
+            output_token_ids,
+            choice["meta_info"],
+            aborted=finish_reason == "abort",
+        )
     sample.tokens = prompt_token_ids + output_token_ids
     sample.rollout_log_probs = output_log_probs
     sample.response = tokenizer.decode(output_token_ids)
@@ -126,7 +135,7 @@ def _compute_sample_from_openai_record(
         sample.strip_last_output_tokens(trim_count, tokenizer)
 
     # TODO unify with Sample.update_from_meta_info
-    match choice["finish_reason"]:
+    match finish_reason:
         case "stop" | "tool_calls":
             sample.status = Sample.Status.COMPLETED
         case "length":
diff --git a/miles/rollout/session/v2/core.py b/miles/rollout/session/v2/core.py
index 2fb3591b376..08f91268123 100644
--- a/miles/rollout/session/v2/core.py
+++ b/miles/rollout/session/v2/core.py
@@ -16,7 +16,7 @@
 )
 from miles.rollout.session.errors import SessionNotFoundError, TokenizationError
 from miles.rollout.session.request_args import filter_turn_args, parse_chat_request
-from miles.rollout.session.samples.codec import COMPUTED_FIELDS_V2, encode_samples
+from miles.rollout.session.samples.codec import COMPUTED_FIELDS_V2, ROLLOUT_SAMPLING_MASK_FIELDS, encode_samples
 from miles.rollout.session.types import GetSessionResponse, SessionRecord
 from miles.rollout.session.v2.metrics import SESSION_ROLLOUT_METRICS_KEY, build_session_rollout_metrics
 from miles.rollout.session.v2.session_state import (
@@ -26,6 +26,7 @@
 )
 from miles.rollout.session.v2.utils import build_leaf_material, tree_metadata
 from miles.utils.function_registry import load_function
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 
 logger = logging.getLogger(__name__)
 
@@ -39,6 +40,9 @@ def __init__(
         self, backend, registry: SessionRegistryV2, config, session_server_instance_id=None, *, use_addition_r3=False
     ):
         super().__init__(backend, registry, config, session_server_instance_id, use_addition_r3=use_addition_r3)
+        self.samples_wire_fields = COMPUTED_FIELDS_V2
+        if sampling_support_replay_enabled(config):
+            self.samples_wire_fields += ROLLOUT_SAMPLING_MASK_FIELDS
         # Import-path only in production: function_registry is process-local.
         self.sample_picker = load_function(config.session_sample_picker_path, sync_required=True)
         self.sample_postprocessor = load_function(config.session_sample_postprocessor_path, sync_required=True)
@@ -89,7 +93,7 @@ async def collect_samples(
             metadata["agent"] = agent_metadata
         if not session.tree.nodes:
             return _samples_response(
-                encode_samples([], metadata, empty_reason="no_records", fields=COMPUTED_FIELDS_V2)
+                encode_samples([], metadata, empty_reason="no_records", fields=self.samples_wire_fields)
             )
 
         try:
@@ -105,7 +109,7 @@ async def collect_samples(
             return Response(content=str(exc).encode(), status_code=422, media_type="text/plain")
         if not material:
             return _samples_response(
-                encode_samples([], metadata, empty_reason="all_truncated", fields=COMPUTED_FIELDS_V2)
+                encode_samples([], metadata, empty_reason="all_truncated", fields=self.samples_wire_fields)
             )
 
         # Hook lane: a policy bug is a deterministic 422 carrying the hook's
@@ -127,13 +131,13 @@ async def collect_samples(
             return Response(content=body.encode(), status_code=422, media_type="text/plain")
         if not samples:
             return _samples_response(
-                encode_samples([], metadata, empty_reason="all_truncated", fields=COMPUTED_FIELDS_V2)
+                encode_samples([], metadata, empty_reason="all_truncated", fields=self.samples_wire_fields)
             )
         # Hooks may inspect or mutate session metadata, so publish the
         # authoritative server-owned value only at the wire boundary.
         if self.config.sglang_speculative_algorithm is not None:
             metadata[SESSION_ROLLOUT_METRICS_KEY] = build_session_rollout_metrics(session_id, session.tree.nodes)
-        return _samples_response(encode_samples(samples, metadata, fields=COMPUTED_FIELDS_V2))
+        return _samples_response(encode_samples(samples, metadata, fields=self.samples_wire_fields))
 
     async def chat_completions(
         self, session_id: str, *, method: str, query: str, headers: dict, body: bytes
diff --git a/tests/e2e/sglang/test_session_server_addition_r3.py b/tests/e2e/sglang/test_session_server_addition_r3.py
index 01dd44ff911..a77efc2222f 100644
--- a/tests/e2e/sglang/test_session_server_addition_r3.py
+++ b/tests/e2e/sglang/test_session_server_addition_r3.py
@@ -79,6 +79,9 @@ def _serve_session(backend_url: str) -> Iterator[str]:
         use_session_server="v1",
         use_rollout_routing_replay=True,
         use_rollout_indexer_replay=False,
+        rollout_temperature=1.0,
+        rollout_top_p=1.0,
+        rollout_top_k=-1,
         save_debug_trajectory_data=None,
         pause_generation_mode="in_place",
         num_layers=_NUM_LAYERS,
diff --git a/tests/fast/fixtures/session_fixtures.py b/tests/fast/fixtures/session_fixtures.py
index 26ae2f4ce33..d185d8a6d29 100644
--- a/tests/fast/fixtures/session_fixtures.py
+++ b/tests/fast/fixtures/session_fixtures.py
@@ -18,6 +18,9 @@ def make_session_server_config(**overrides: Any) -> SessionServerConfig:
         apply_chat_template_kwargs=None,
         use_rollout_routing_replay=False,
         use_rollout_indexer_replay=False,
+        rollout_temperature=1.0,
+        rollout_top_p=1.0,
+        rollout_top_k=-1,
         sglang_speculative_algorithm=None,
         num_layers=None,
         moe_router_topk=None,
diff --git a/tests/fast/ray/rollout/conftest.py b/tests/fast/ray/rollout/conftest.py
index 13a17f9233a..7b311d5a9e9 100644
--- a/tests/fast/ray/rollout/conftest.py
+++ b/tests/fast/ray/rollout/conftest.py
@@ -33,6 +33,8 @@ def make_args(**overrides: Any) -> Namespace:
         n_samples_per_eval_prompt=4,
         rollout_max_response_len=512,
         rollout_temperature=1.0,
+        rollout_top_p=1.0,
+        rollout_top_k=-1,
         over_sampling_batch_size=None,
         rollout_global_dataset=False,
         num_rollout=1,
diff --git a/tests/fast/rollout/generate_hub/test_agentic_v2.py b/tests/fast/rollout/generate_hub/test_agentic_v2.py
index ae3357b9f2a..e6d138e160a 100644
--- a/tests/fast/rollout/generate_hub/test_agentic_v2.py
+++ b/tests/fast/rollout/generate_hub/test_agentic_v2.py
@@ -85,6 +85,35 @@ async def test_success_returns_list_and_forwards_agent_metadata(monkeypatch):
     assert tracer.agent_metadata == {"agent_result": "done"}
 
 
+@pytest.mark.asyncio
+async def test_evaluation_explicitly_disables_session_sampling_replay(monkeypatch):
+    sample = Sample(status=Sample.Status.COMPLETED, response="done", response_length=1, tokens=[1])
+    tracer = _Tracer(SamplesReply(samples=[sample], session_metadata={}, empty_reason=None))
+    seen_request_kwargs = None
+
+    async def fake_create(args):
+        return tracer
+
+    async def fake_agent(**kwargs):
+        nonlocal seen_request_kwargs
+        seen_request_kwargs = kwargs["request_kwargs"]
+        return None
+
+    monkeypatch.setattr(agentic_tool_call.OpenAIEndpointTracer, "create", fake_create)
+    monkeypatch.setattr(agentic_tool_call, "load_function", lambda path: fake_agent)
+    training_input = _generate_input(rollout_top_p=0.95, rollout_top_k=32, rollout_temperature=1.0)
+    generate_input = GenerateFnInput(
+        state=training_input.state,
+        sample=training_input.sample,
+        sampling_params=training_input.sampling_params,
+        evaluation=True,
+    )
+
+    await agentic_tool_call.generate(generate_input)
+
+    assert seen_request_kwargs["return_sampling_mask"] is False
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(("input_rollout_id", "expected_rollout_id"), [(None, 7), (11, 11)])
 async def test_success_assigns_shared_rollout_id_to_v2_leaves(monkeypatch, input_rollout_id, expected_rollout_id):
diff --git a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py
index 4f78cfaaefc..4db262db14f 100644
--- a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py
+++ b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py
@@ -16,7 +16,12 @@
 
 import miles.utils.http_utils as http_utils
 from miles.rollout.generate_utils.openai_endpoint_utils import OpenAIEndpointTracer
-from miles.rollout.session.samples.codec import COMPUTED_FIELDS, COMPUTED_FIELDS_V2, encode_samples
+from miles.rollout.session.samples.codec import (
+    COMPUTED_FIELDS,
+    COMPUTED_FIELDS_V2,
+    ROLLOUT_SAMPLING_MASK_FIELDS,
+    encode_samples,
+)
 from miles.utils.http_utils import post_bytes_no_retry
 from miles.utils.types import Sample
 
@@ -330,8 +335,22 @@ async def fake_post(url: str, payload: dict, action: str = "post"):
 
     monkeypatch.setattr("miles.rollout.generate_utils.openai_endpoint_utils.post", fake_post)
 
-    def args(version):
-        return SimpleNamespace(session_server_addrs=["127.0.0.1:7000"], use_session_server=version)
+    def args(version, top_p=1.0, top_k=-1):
+        return SimpleNamespace(
+            session_server_addrs=["127.0.0.1:7000"],
+            use_session_server=version,
+            rollout_top_p=top_p,
+            rollout_top_k=top_k,
+        )
 
     assert (await OpenAIEndpointTracer.create(args(True))).samples_wire_fields == COMPUTED_FIELDS
     assert (await OpenAIEndpointTracer.create(args("v2"))).samples_wire_fields == COMPUTED_FIELDS_V2
+    assert (await OpenAIEndpointTracer.create(args(True, 0.95))).samples_wire_fields == (
+        COMPUTED_FIELDS + ROLLOUT_SAMPLING_MASK_FIELDS
+    )
+    assert (await OpenAIEndpointTracer.create(args("v2", 0.95))).samples_wire_fields == (
+        COMPUTED_FIELDS_V2 + ROLLOUT_SAMPLING_MASK_FIELDS
+    )
+    assert (await OpenAIEndpointTracer.create(args(True, top_k=32))).samples_wire_fields == (
+        COMPUTED_FIELDS + ROLLOUT_SAMPLING_MASK_FIELDS
+    )
diff --git a/tests/fast/rollout/session/test_config.py b/tests/fast/rollout/session/test_config.py
index 030f4de5118..693994b3ec3 100644
--- a/tests/fast/rollout/session/test_config.py
+++ b/tests/fast/rollout/session/test_config.py
@@ -16,6 +16,9 @@
     "apply_chat_template_kwargs": "apply_chat_template_kwargs",
     "use_rollout_routing_replay": "use_rollout_routing_replay",
     "use_rollout_indexer_replay": "use_rollout_indexer_replay",
+    "rollout_temperature": "rollout_temperature",
+    "rollout_top_p": "rollout_top_p",
+    "rollout_top_k": "rollout_top_k",
     "sglang_speculative_algorithm": "sglang_speculative_algorithm",
     "num_layers": "num_layers",
     "moe_router_topk": "moe_router_topk",
@@ -49,6 +52,9 @@
     apply_chat_template_kwargs={"enable_thinking": True},
     use_rollout_routing_replay=True,
     use_rollout_indexer_replay=True,
+    rollout_temperature=0.7,
+    rollout_top_p=0.95,
+    rollout_top_k=32,
     sglang_speculative_algorithm="EAGLE",
     num_layers=61,
     moe_router_topk=8,
@@ -127,6 +133,9 @@ def test_missing_optional_args_fall_back_to_none(self):
     apply_chat_template_kwargs=None,
     use_rollout_routing_replay=False,
     use_rollout_indexer_replay=False,
+    rollout_temperature=1.0,
+    rollout_top_p=1.0,
+    rollout_top_k=-1,
     sglang_speculative_algorithm=None,
     num_layers=None,
     moe_router_topk=None,
diff --git a/tests/fast/rollout/session/test_request_args.py b/tests/fast/rollout/session/test_request_args.py
index 46f4d092659..0f386a576fa 100644
--- a/tests/fast/rollout/session/test_request_args.py
+++ b/tests/fast/rollout/session/test_request_args.py
@@ -38,6 +38,26 @@ def test_replay_flags_follow_the_launch_flags(self):
         assert wire["return_routed_experts"] is True
         assert wire["return_indexer_topk"] is True
 
+    def test_sampling_replay_uses_launch_defaults_without_overwriting_request_filters(self):
+        config = make_session_server_config(rollout_top_p=0.95, rollout_top_k=32, rollout_temperature=0.7)
+
+        wire, _ = resolve_request_args_by_config({"top_p": 0.8, "custom_params": {"caller_option": "kept"}}, config)
+
+        assert wire["return_sampling_mask"] is True
+        assert wire["temperature"] == 0.7
+        assert wire["top_p"] == 0.8
+        assert wire["top_k"] == 32
+        assert wire["custom_params"] == {"caller_option": "kept"}
+
+    def test_sampling_replay_can_be_disabled_per_request(self):
+        config = make_session_server_config(rollout_top_p=0.95, rollout_top_k=32)
+
+        wire, _ = resolve_request_args_by_config({"return_sampling_mask": False}, config)
+
+        assert wire["return_sampling_mask"] is False
+        assert "top_p" not in wire
+        assert "top_k" not in wire
+
     @pytest.mark.parametrize("field", ["input_ids", "routed_experts_start_len", "logprob_start_len", "lora_path"])
     def test_client_tito_control_fields_are_rejected(self, field):
         with pytest.raises(MessageValidationError, match=f"{field}="):
diff --git a/tests/fast/rollout/session/test_samples.py b/tests/fast/rollout/session/test_samples.py
index a172f583f05..9e85b41a900 100644
--- a/tests/fast/rollout/session/test_samples.py
+++ b/tests/fast/rollout/session/test_samples.py
@@ -41,6 +41,8 @@ def _make_record(
     routed_experts: str | None = None,
     routed_experts_start_len: int | None = None,
     weight_versions: list[dict[str, str | int]] | None = None,
+    sampling_masks: list[list[int]] | None = None,
+    sampling_log_probs: list[float] | None = None,
 ) -> SessionRecord:
     """Build a minimal session record mimicking SGLang's response format.
 
@@ -74,6 +76,10 @@ def _make_record(
         request["routed_experts_start_len"] = routed_experts_start_len
     if weight_versions is not None:
         meta_info["weight_versions"] = weight_versions
+    if sampling_masks is not None:
+        request["return_sampling_mask"] = True
+        meta_info["output_token_sampling_mask"] = sampling_masks
+        meta_info["output_token_sampling_logprobs"] = sampling_log_probs
     return SessionRecord(
         timestamp=0.0,
         method="POST",
@@ -180,6 +186,63 @@ def test_record_without_any_weight_version_still_contributes_a_call(self):
 
         assert samples[0].weight_versions == [WeightVersionsPerCall(spans=[])]
 
+    def test_single_record_uses_native_sampling_support_and_log_probs(self):
+        tok = _mock_tokenizer()
+        record = _make_record(
+            prompt_token_ids=[1, 2, 3],
+            output_token_ids=[10, 11],
+            output_log_probs=[-2.5, -2.6],
+            sampling_masks=[[10, 4, 7], [11, 3]],
+            sampling_log_probs=[-0.5, -0.6],
+        )
+
+        (sample,) = compute_samples_from_openai_records(_ARGS, [record], tok)
+
+        assert sample.rollout_log_probs == [-0.5, -0.6]
+        ids, offsets = sample.rollout_sampling_mask._as_tensors()
+        assert ids.tolist() == [10, 4, 7, 11, 3]
+        assert offsets.tolist() == [0, 3, 5]
+        sample.validate()
+
+    def test_abort_without_sampling_metadata_is_non_trainable(self):
+        tok = _mock_tokenizer()
+        record = _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[], finish_reason="abort")
+        record.request["return_sampling_mask"] = True
+
+        (sample,) = compute_samples_from_openai_records(_ARGS, [record], tok)
+
+        assert sample.status == Sample.Status.ABORTED
+        assert sample.reward is None
+        assert len(sample.rollout_sampling_mask) == 0
+
+    def test_successful_turn_still_requires_sampling_metadata(self):
+        tok = _mock_tokenizer()
+        record = _make_record(prompt_token_ids=[1, 2, 3], output_token_ids=[10])
+        record.request["return_sampling_mask"] = True
+
+        with pytest.raises(ValueError, match="missing output_token_sampling_mask"):
+            compute_samples_from_openai_records(_ARGS, [record], tok)
+
+    def test_later_abort_preserves_terminal_status(self):
+        tok = _mock_tokenizer()
+        records = [
+            _make_record(
+                prompt_token_ids=[1, 2],
+                output_token_ids=[10],
+                sampling_masks=[[10, 11]],
+                sampling_log_probs=[-0.2],
+            ),
+            _make_record(prompt_token_ids=[1, 2, 10, 20], output_token_ids=[], finish_reason="abort"),
+        ]
+        records[1].request["return_sampling_mask"] = True
+
+        samples = compute_samples_from_openai_records(_ARGS, records, tok)
+        merged = merge_samples(samples, tok)
+
+        assert merged.status == Sample.Status.ABORTED
+        assert merged.tokens == [1, 2, 10, 20]
+        assert len(merged.rollout_sampling_mask) == merged.response_length == 2
+
     def test_multiple_records_produce_multiple_samples(self):
         tok = _mock_tokenizer()
         records = [
diff --git a/tests/fast/router/test_sessions_v2.py b/tests/fast/router/test_sessions_v2.py
index 342930fd229..58b91770adc 100644
--- a/tests/fast/router/test_sessions_v2.py
+++ b/tests/fast/router/test_sessions_v2.py
@@ -47,6 +47,9 @@ def process_fn(prompt: str) -> ProcessResult:
             "tito_model": "default",
             "use_rollout_routing_replay": False,
             "use_rollout_indexer_replay": False,
+            "rollout_temperature": 1.0,
+            "rollout_top_p": 1.0,
+            "rollout_top_k": -1,
             "sglang_speculative_algorithm": None,
             "num_layers": None,
             "moe_router_topk": None,
diff --git a/tests/fast/utils/workers/test_argv_utils.py b/tests/fast/utils/workers/test_argv_utils.py
index f90000de69e..2877b82fca7 100644
--- a/tests/fast/utils/workers/test_argv_utils.py
+++ b/tests/fast/utils/workers/test_argv_utils.py
@@ -104,6 +104,9 @@ def test_real_worker_configs_roundtrip(self):
             apply_chat_template_kwargs={"enable_thinking": False},
             use_rollout_routing_replay=True,
             use_rollout_indexer_replay=False,
+            rollout_temperature=1.0,
+            rollout_top_p=1.0,
+            rollout_top_k=-1,
             sglang_speculative_algorithm=None,
             num_layers=None,
             moe_router_topk=None,
diff --git a/tests/manual/session/bench_session_server_overhead.py b/tests/manual/session/bench_session_server_overhead.py
index e7e5ba80f30..42864afcc6c 100644
--- a/tests/manual/session/bench_session_server_overhead.py
+++ b/tests/manual/session/bench_session_server_overhead.py
@@ -354,6 +354,9 @@ def _build_server_config(
         tito_model=bench_args.tito_model,
         use_rollout_routing_replay=True,
         use_rollout_indexer_replay=False,
+        rollout_temperature=1.0,
+        rollout_top_p=1.0,
+        rollout_top_k=-1,
         sglang_speculative_algorithm=None,
         num_layers=None,
         moe_router_topk=None,

From ee9670f143a96169e8f56dd7056a602240be2973 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Sun, 20 Sep 2026 00:52:00 +0000
Subject: [PATCH 4/6] fix(rl): validate bounded sampling replay contracts

---
 miles/utils/arguments.py           | 49 ++++++++++++++++++-
 tests/fast/utils/test_arguments.py | 78 ++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 2 deletions(-)

diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index cff760f2eac..2d2f24fcc15 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -23,6 +23,7 @@
 from miles.utils.megatron_args_utils import compute_megatron_world_size_except_dp
 from miles.utils.object_store import ObjectStoreBackend
 from miles.utils.run_uuid import RUN_UUID_LENGTH, generate_run_uuid, validate_run_uuid
+from miles.utils.sampling_mask import sampling_support_replay_enabled
 from miles.utils.tracking_utils.ci_history import RECORD_DIR_ENV
 
 logger = logging.getLogger(__name__)
@@ -555,10 +556,24 @@ def add_rollout_arguments(parser):
                 help="the temperature for the inference engine during rollout.",
             )
             parser.add_argument(
-                "--rollout-top-p", type=float, default=1.0, help="the top-p for the inference engine during rollout."
+                "--rollout-top-p",
+                type=float,
+                default=1.0,
+                help=(
+                    "the top-p for the inference engine during rollout. Values below 1 enable "
+                    "sampling-support replay and require --use-miles-router plus a positive "
+                    "--rollout-top-k."
+                ),
             )
             parser.add_argument(
-                "--rollout-top-k", type=int, default=-1, help="the top-k for the inference engine during rollout."
+                "--rollout-top-k",
+                type=int,
+                default=-1,
+                help=(
+                    "the top-k for the inference engine during rollout. Positive values enable "
+                    "sampling-support replay. SGLang's --sampling-mask-max-tokens is the physical "
+                    "returned-support limit because cutoff ties can retain more than top-k tokens."
+                ),
             )
             parser.add_argument(
                 "--rollout-max-context-len",
@@ -2932,6 +2947,36 @@ def miles_validate_args(args):
             "have known issues in SGLang and need to be fixed."
         )
 
+    if not 0.0 < args.rollout_top_p <= 1.0:
+        raise ValueError(f"--rollout-top-p must be in (0, 1], got {args.rollout_top_p}")
+    if args.rollout_top_k != -1 and args.rollout_top_k < 1:
+        raise ValueError(f"--rollout-top-k must be -1 or at least 1, got {args.rollout_top_k}")
+    if sampling_support_replay_enabled(args):
+        if args.rollout_top_k == -1:
+            raise ValueError(
+                "--rollout-top-p below 1 requires a positive --rollout-top-k; "
+                "top-p alone does not bound the returned support size"
+            )
+        if args.recompute_logprobs_via_prefill:
+            raise ValueError(
+                "sampling-support replay cannot be combined with --recompute-logprobs-via-prefill; "
+                "prefill scoring does not preserve the rollout sampling support"
+            )
+        if not args.use_miles_router:
+            raise ValueError(
+                "sampling-support replay currently requires --use-miles-router; "
+                "the SGLang v0.5.20 gateway's typed request schema drops "
+                "return_sampling_mask before forwarding the request"
+            )
+        if args.kl_coef != 0 or args.use_kl_loss or args.use_opd:
+            # The actor still produces full-vocabulary logits, but replay currently exposes only the
+            # support-normalized actor score to the loss. These objectives can be enabled once the loss
+            # path also preserves an unmasked actor score from the same forward pass.
+            raise ValueError(
+                "sampling-support replay cannot currently be combined with reference KL or teacher distillation; "
+                "those objectives require a separate full-policy actor score"
+            )
+
     if not args.use_session_server and args.tito_model != TITOTokenizerType.DEFAULT.value:
         raise ValueError(
             f"--tito-model={args.tito_model} requires --use-session-server; "
diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py
index 4f8873e8092..bc0a3f40381 100644
--- a/tests/fast/utils/test_arguments.py
+++ b/tests/fast/utils/test_arguments.py
@@ -230,6 +230,84 @@ def test_recompute_logprobs_via_prefill_flag_is_parsed():
     assert args.recompute_logprobs_via_prefill is True
 
 
+@pytest.mark.parametrize(
+    ("extra", "message"),
+    [
+        (["--rollout-top-p", "0"], "--rollout-top-p must be in"),
+        (["--rollout-top-k", "0"], "--rollout-top-k must be -1 or at least 1"),
+        (
+            ["--rollout-top-p", "0.95"],
+            "--rollout-top-p below 1 requires a positive --rollout-top-k",
+        ),
+        (
+            [
+                "--rollout-top-p",
+                "0.95",
+                "--rollout-top-k",
+                "32",
+                "--use-miles-router",
+                "--true-on-policy-mode",
+                "--recompute-logprobs-via-prefill",
+            ],
+            "sampling-support replay cannot be combined with --recompute-logprobs-via-prefill",
+        ),
+        (
+            ["--rollout-top-p", "0.95", "--rollout-top-k", "32"],
+            "currently requires --use-miles-router",
+        ),
+        (
+            [
+                "--rollout-top-p",
+                "0.95",
+                "--rollout-top-k",
+                "32",
+                "--use-miles-router",
+                "--kl-coef",
+                "0.1",
+            ],
+            "cannot currently be combined with reference KL or teacher distillation",
+        ),
+        (
+            [
+                "--rollout-top-p",
+                "0.95",
+                "--rollout-top-k",
+                "32",
+                "--use-miles-router",
+                "--use-kl-loss",
+            ],
+            "cannot currently be combined with reference KL or teacher distillation",
+        ),
+        (
+            [
+                "--rollout-top-p",
+                "0.95",
+                "--rollout-top-k",
+                "32",
+                "--use-miles-router",
+                "--use-opd",
+            ],
+            "cannot currently be combined with reference KL or teacher distillation",
+        ),
+    ],
+)
+def test_sampling_support_arguments_fail_closed(extra, message):
+    parser = argparse.ArgumentParser()
+    get_miles_extra_args_provider()(parser)
+    args = parser.parse_args(extra + ["--num-rollout", "1"] + REQUIRED_ARGS)
+
+    with pytest.raises(ValueError, match=message):
+        miles_validate_args(args)
+
+
+def test_finite_top_k_enables_sampling_support_replay():
+    parser = argparse.ArgumentParser()
+    get_miles_extra_args_provider()(parser)
+    args = parser.parse_args(["--rollout-top-k", "32", "--use-miles-router", "--num-rollout", "1"] + REQUIRED_ARGS)
+
+    miles_validate_args(args)
+
+
 def test_sglang_parallel_sizes_keep_server_args_destinations():
     parser = add_sglang_arguments(argparse.ArgumentParser())
     args = parser.parse_args(

From eb5391180d9512beafddf45811e1116b5517ec03 Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Sun, 20 Sep 2026 00:52:06 +0000
Subject: [PATCH 5/6] test(e2e): validate sampling-support replay

---
 .../test_qwen3_0_6b_top_p_sampling_replay.py  | 106 ++++++++++++++++++
 1 file changed, 106 insertions(+)
 create mode 100644 tests/e2e/megatron/test_qwen3_0_6b_top_p_sampling_replay.py

diff --git a/tests/e2e/megatron/test_qwen3_0_6b_top_p_sampling_replay.py b/tests/e2e/megatron/test_qwen3_0_6b_top_p_sampling_replay.py
new file mode 100644
index 00000000000..c994f274885
--- /dev/null
+++ b/tests/e2e/megatron/test_qwen3_0_6b_top_p_sampling_replay.py
@@ -0,0 +1,106 @@
+import os
+
+from tests.ci.ci_register import register_cuda_ci
+from tests.ci.metric_history import register_ci_gate
+
+import miles.utils.external_utils.command_utils as U
+
+register_cuda_ci(
+    est_time=900,
+    suite="stage-c-2-gpu-h200",
+    labels=["megatron", "sglang", "replay"],
+    hardware=["hopper", "blackwell"],
+)
+register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")
+
+MODEL_NAME = "Qwen3-0.6B"
+MODEL_TYPE = "qwen3-0.6B"
+NUM_GPUS = 2
+
+
+def prepare() -> None:
+    U.exec_command_cpu("mkdir -p /root/models /root/datasets")
+    U.exec_command_cpu(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.hf_download_dataset("zhuzilin/dapo-math-17k")
+    U.convert_checkpoint(
+        model_name=MODEL_NAME,
+        megatron_model_type=MODEL_TYPE,
+        num_gpus_per_node=NUM_GPUS,
+        dir_dst="/root/models",
+    )
+
+
+def execute() -> None:
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ --ref-load /root/models/{MODEL_NAME}_torch_dist "
+
+    rollout_args = (
+        "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl "
+        "--input-key prompt "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rm-type deepscaler "
+        "--num-rollout 1 "
+        "--rollout-batch-size 2 "
+        "--n-samples-per-prompt 2 "
+        "--rollout-max-response-len 128 "
+        "--rollout-temperature 1 "
+        "--rollout-top-p 0.8 "
+        "--rollout-top-k 32 "
+        "--global-batch-size 4 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    algorithm_args = "--advantage-estimator grpo "
+
+    megatron_args = (
+        "--tensor-model-parallel-size 2 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 1 "
+        "--expert-tensor-parallel-size 1 "
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+    )
+
+    sglang_args = (
+        "--rollout-num-gpus-per-engine 2 "
+        "--sglang-mem-fraction-static 0.7 "
+        "--sglang-decode-log-interval 1000 "
+        "--use-miles-router "
+    )
+
+    train_args = (
+        ckpt_args
+        + rollout_args
+        + optimizer_args
+        + algorithm_args
+        + megatron_args
+        + sglang_args
+        + U.get_default_wandb_args(__file__)
+        + "--ci-test --actor-num-nodes 1 --actor-num-gpus-per-node 2 --colocate "
+    )
+
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=NUM_GPUS,
+        megatron_model_type=MODEL_TYPE,
+    )
+
+
+if __name__ == "__main__":
+    prepare()
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute()

From 9f68c6135215a7a4dbf099b0b6f45588a8b76b4c Mon Sep 17 00:00:00 2001
From: Nan <nan@modal.com>
Date: Sun, 20 Sep 2026 00:52:12 +0000
Subject: [PATCH 6/6] docs(rl): document sampling-support replay

---
 docs/advanced/index.md                   |  9 ++-
 docs/advanced/sampling-support-replay.md | 94 ++++++++++++++++++++++++
 docs/docs.json                           |  1 +
 docs/user-guide/cli-reference.md         |  7 +-
 4 files changed, 107 insertions(+), 4 deletions(-)
 create mode 100644 docs/advanced/sampling-support-replay.md

diff --git a/docs/advanced/index.md b/docs/advanced/index.md
index 3b16906451f..46b7cf8ee9a 100644
--- a/docs/advanced/index.md
+++ b/docs/advanced/index.md
@@ -7,7 +7,7 @@ This section covers the Miles features that the Core-features section of the
 homepage points at: low-precision training (FP8 / MXFP8 / NVFP4 / INT4 QAT),
 Rollout Routing Replay for MoE, fast weight updates over P2P RDMA,
 disaggregated RL rollout through an external service, fault tolerance,
-speculative decoding, and LoRA training and serving.
+speculative decoding, sampling-support replay, and LoRA training and serving.
 
 <CardGroup cols={2}>
 
@@ -38,6 +38,13 @@ speculative decoding, and LoRA training and serving.
 
   </Card>
 
+  <Card title="Sampling-Support Replay" icon="filter" href="/advanced/sampling-support-replay">
+
+    Preserve top-p and top-k rollout probabilities when the actor recomputes
+    token log probabilities for training.
+
+  </Card>
+
   <Card title="On-Policy Distillation" icon="graduation-cap" href="/advanced/on-policy-distillation">
 
     Train a student on its own rollouts while matching teacher token
diff --git a/docs/advanced/sampling-support-replay.md b/docs/advanced/sampling-support-replay.md
new file mode 100644
index 00000000000..bfce83ac39f
--- /dev/null
+++ b/docs/advanced/sampling-support-replay.md
@@ -0,0 +1,94 @@
+---
+title: Sampling-Support Replay
+description: Train against the same bounded sampling distribution used during rollout.
+---
+
+Top-p and top-k change the distribution that generates a rollout token. They
+remove tokens from the vocabulary and renormalize the remaining probability
+mass. Recomputing a token's log probability over the full vocabulary therefore
+does not reproduce its behavior-policy probability, even when the rollout and
+trainer weights are identical.
+
+Sampling-support replay preserves that distribution. For every generated token,
+SGLang returns the realized set of token IDs that remained after sampling
+filters. Miles transports this ragged support with the sample and applies it to
+the actor logits before the softmax:
+
+$$
+q_\theta(a_t \mid S_t) =
+\frac{\exp(z_\theta(a_t) / T)}
+{\sum_{j \in S_t} \exp(z_\theta(j) / T)}.
+$$
+
+The support `S_t` is fixed rollout data; `z_θ` comes from the actor being
+trained. This gives PPO/GRPO a denominator and numerator in the same sampled
+probability space while retaining gradients through the current actor logits.
+
+## Enable replay
+
+```bash
+ROLLOUT_ARGS+=(
+  --rollout-temperature 1.0
+  --rollout-top-p 0.95
+  --rollout-top-k 64
+  --use-miles-router
+)
+```
+
+Replay is enabled whenever `--rollout-top-p` is below `1` or
+`--rollout-top-k` is positive. A finite top-p run must also set a positive
+top-k so that the captured support is bounded. Miles asks SGLang to capture its
+native sampling support with `return_sampling_mask` and uses the returned
+support-normalized log probability as the rollout log probability.
+
+`--rollout-top-k` is the default for rollout requests, not a global upper bound
+on request-specific top-k values. SGLang owns support capacity through
+`--sglang-sampling-mask-max-tokens` (or the matching server-group override) and
+rejects requests whose realized support cannot be represented.
+
+## Correctness requirements
+
+Miles rejects configurations that it cannot replay faithfully:
+
+- Every training request must state its top-p, top-k, and temperature.
+- Request temperature must match `--rollout-temperature`.
+- Frequency, presence, and repetition penalties and `logit_bias` are not
+  supported because the trainer does not replay those logit transformations.
+- `--recompute-logprobs-via-prefill` is incompatible because that path does not
+  preserve the per-token support.
+- The Miles router is currently required as a transport compatibility measure.
+  The SGLang v0.5.20 gateway deserializes chat requests through a typed schema
+  that does not retain `return_sampling_mask`, so the engine never receives the
+  capture request. The gateway's non-streaming response path does preserve raw
+  response bytes; response stripping is not the issue.
+
+SGLang owns compatibility with custom logit processors and the physical support
+limit. Speculative decoding is not restricted by Miles; it works when the
+SGLang backend returns the same native support and normalized log probability
+for every accepted output token.
+
+Tool and environment tokens are recorded with singleton support. Multi-turn
+sessions concatenate supports across assistant turns and singleton supports
+across observations, so every response token remains aligned with exactly one
+support. Evaluation requests do not capture this training-only metadata.
+
+## Current objective limitation
+
+The actor produces full-vocabulary logits, but the current loss interface
+returns one actor score per token. With replay enabled, that score is normalized
+over the captured support for the policy ratio. Reference KL and on-policy
+distillation also need the actor score normalized over the full vocabulary, so
+Miles currently rejects `--use-kl-loss`, nonzero `--kl-coef`, and `--use-opd`
+with replay.
+
+This is an interface limitation, not a mathematical conflict. Both scores can
+be derived from the same actor forward pass once the loss path carries them as
+separate values.
+
+## Monitor replay
+
+Track `train/train_rollout_logprob_abs_diff`. On the first update from identical
+rollout and actor weights, the value should be close to the numerical tolerance
+of the two inference paths. Later, it also includes legitimate policy staleness
+and weight updates, so interpret it together with version-lag and clipping
+metrics rather than as a standalone correctness test.
diff --git a/docs/docs.json b/docs/docs.json
index 7f57d0d115b..e11cba21af2 100644
--- a/docs/docs.json
+++ b/docs/docs.json
@@ -220,6 +220,7 @@
        "advanced/low-precision",
        "advanced/int4-qat",
        "advanced/speculative-decoding",
+       "advanced/sampling-support-replay",
        "advanced/lora",
        "advanced/on-policy-distillation"
       ]
diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md
index 7d2a17affa6..6cd51bc84fe 100644
--- a/docs/user-guide/cli-reference.md
+++ b/docs/user-guide/cli-reference.md
@@ -80,7 +80,8 @@ then push up until you OOM.
 | Flag | Default | What |
 |---|---|---|
 | `--rollout-temperature` | `1.0` | Sampling temperature. |
-| `--rollout-top-p` | `1.0` | Top-p truncation. |
+| `--rollout-top-p` | `1.0` | Top-p truncation. Values below `1` enable [sampling-support replay](/advanced/sampling-support-replay) and require a positive top-k. |
+| `--rollout-top-k` | `-1` | Top-k truncation. Positive values enable [sampling-support replay](/advanced/sampling-support-replay). |
 | `--rollout-max-response-len` | `–` | Max tokens per response. |
 | `--rollout-stop-token-ids` | model default | Stop token IDs. Override when generations don't stop. |
 | `--apply-chat-template` | off | Apply the tokenizer's chat template. |
@@ -186,8 +187,8 @@ Sections mirror the launch-script argument groups.
 |---|---|---|---|
 | `--rollout-max-response-len` | int | – | Max tokens per response. |
 | `--rollout-temperature` | float | `1.0` | Sampling temperature. |
-| `--rollout-top-p` | float | `1.0` | Top-p truncation. |
-| `--rollout-top-k` | int | `-1` | Top-k truncation (-1 disables). |
+| `--rollout-top-p` | float | `1.0` | Top-p truncation. Values below `1` require bounded [sampling-support replay](/advanced/sampling-support-replay). |
+| `--rollout-top-k` | int | `-1` | Top-k truncation (`-1` disables). Positive values enable [sampling-support replay](/advanced/sampling-support-replay). |
 | `--rollout-stop` | str+ | – | Stop strings. |
 | `--rollout-stop-token-ids` | int+ | – | Stop token IDs. |