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/2] 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/2] 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