diff --git a/docs/inference.md b/docs/inference.md
index 8a76ab477a..f50196e2cb 100644
--- a/docs/inference.md
+++ b/docs/inference.md
@@ -296,3 +296,17 @@ enable_return_routed_experts = true
 This however is not free, it adds a significant overhead to the HTTP requests as this payload can grow quite large. We reccomend sizing up the env-server pool (`orchestrator.*.source.serve.pool`) to allow for more parallelization on the verifiers side.
 
 Currently this feature is also not supported with CPU KV cache offload, which can have negative impact on the inference throughput.
+
+### Sampling Replay
+
+Truncated sampling (`top_p < 1`, `top_k`) renormalizes the sampling distribution over the surviving "kept set" of tokens. The rollout logprobs reflect that (`logprobs_mode = "processed_logprobs"`), so the trainer must renormalize over the same set — otherwise every importance ratio is biased and training collapses (DeepSeek V3.2's "Keep Sampling Mask", [arXiv:2512.02556](https://arxiv.org/abs/2512.02556) §3.1; Cognition's [SWE-1.7 post](https://cognition.com/blog/swe-1-7)). prime-rl handles this automatically: the kept-set token ids are recorded at sampling time and the trainer renormalizes its logprobs over them.
+
+```toml
+[orchestrator.train.sampling]
+top_p = 0.95
+top_k = 512   # optional, defaults to 512 under truncation (bounds the kept sets)
+```
+
+That's all — there are no replay flags. Truncated train sampling makes the inference server return kept sets (`inference.kept_tokens`, derived from the largest configured top_k) and the trainer replays whatever masks arrive. Configs that would break under renormalized logprobs are rejected: `opd`/`opsd`, the gibberish/repetition filters (removed from the defaults, rejected if explicitly configured), truncation knobs smuggled via `extra_body`, speculative decoding, and Gemma-family (softcapped) lm_heads. Frozen-source envs are exempt.
+
+When launching the inference server standalone, set `inference.kept_tokens` yourself to cover the clients' top_k.
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
index baab060f95..6d42f303ff 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -401,6 +401,9 @@ class InferenceConfig(BaseConfig):
     enable_return_routed_experts: bool = False
     """Return routed experts in responses. Forwarded as ``--enable-return-routed-experts``."""
 
+    kept_tokens: int | None = Field(None, ge=1)
+    """Auto-set for sampling replay: capture width for the per-token kept-set sampling masks returned on ``/inference/v1/generate`` responses (``None`` = capture off). Derived by the ``rl`` entrypoint from the largest train-sampling ``top_k`` and persisted into the per-node config. Not meant to be set by hand, except for standalone-launched servers, where it must cover the clients' top_k."""
+
     enable_fp32_lm_head: bool = True
     """Run the lm_head projection in fp32 via a native bf16×bf16 → fp32 GEMM (``torch.mm`` with ``out_dtype=torch.float32``). Stabilizes logprob precision under FP8/bf16 inference, matching SGLang's ``--enable-fp32-lm-head``. Implemented as a monkey-patch over vLLM's LogitsProcessor, activated by setting ``additional_config["fp32_lm_head"] = True`` on the vLLM config."""
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
index 434a361f28..a69f9b61e2 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
@@ -1,3 +1,4 @@
+import warnings
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
 
@@ -53,6 +54,16 @@ class TrainSamplingConfig(BaseConfig):
     temperature: float = Field(1.0, ge=0, le=2.0)
     """Sampling temperature."""
 
+    top_p: float = Field(1.0, gt=0, le=1.0)
+    """Nucleus (top-p) sampling for train rollouts. Values below 1.0 truncate the sampling
+    distribution; the ``rl`` entrypoint auto-enables sampling replay so trainer and
+    rollout distributions stay consistent — see docs/inference.md (Sampling Replay)."""
+
+    top_k: int | None = Field(None, ge=1)
+    """Top-k sampling for train rollouts. Truncation triggers sampling replay, and
+    a default top-k is injected when only top-p truncates so kept sets stay
+    bounded — see docs/inference.md (Sampling Replay)."""
+
     max_completion_tokens: int | None = None
     """Maximum output tokens per turn. If None, generates until max context length or EOS."""
 
@@ -61,18 +72,46 @@ class TrainSamplingConfig(BaseConfig):
     extra_body: dict[str, Any] = {}
     """Extra body forwarded with each request to the inference server."""
 
+    def truncates_distribution(self) -> bool:
+        return self.top_p < 1.0 or self.top_k is not None
+
+    @model_validator(mode="after")
+    def validate_no_extra_body_truncation(self):
+        """Truncating values must come from the typed fields — the replay policy reads
+        them. Disabled values pass so resolved configs (where ``resolve_env_config``
+        stamped the ``top_k = -1`` / ``min_p = 0.0`` sentinels) re-validate cleanly."""
+        smuggled = [
+            key
+            for key, truncates in (
+                ("top_p", self.extra_body.get("top_p", 1.0) < 1.0),
+                ("top_k", self.extra_body.get("top_k") not in (None, -1, 0)),
+                ("min_p", self.extra_body.get("min_p", 0.0) > 0.0),
+            )
+            if truncates
+        ]
+        if smuggled:
+            raise ValueError(
+                f"extra_body carries truncating {smuggled}; set them as fields on the train "
+                "sampling config instead (they drive sampling replay)."
+            )
+        return self
+
     def to_sampling_args(self) -> dict[str, Any]:
         """Convert to OAI-compatible sampling args dict, omitting None values."""
         args: dict[str, Any] = {
             "temperature": self.temperature,
-            "top_p": 1.0,
+            "top_p": self.top_p,
             "logprobs": True,
         }
         if self.max_completion_tokens is not None:
             args["max_completion_tokens"] = self.max_completion_tokens
 
-        if self.extra_body:
-            args["extra_body"] = dict(self.extra_body)
+        # top_k rides extra_body (like EvalSamplingConfig), overriding the sentinel.
+        extra_body = dict(self.extra_body)
+        if self.top_k is not None:
+            extra_body["top_k"] = self.top_k
+        if extra_body:
+            args["extra_body"] = extra_body
 
         return args
 
@@ -444,6 +483,12 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 ]
 
 
+# Top-k injected on truncated policy sampling that has none: large enough that a
+# 0.95-0.99 nucleus rarely reaches it (the sampling policy is essentially unchanged),
+# small enough to bound the kept-set capture width and trainer mask tensors.
+DEFAULT_TRAIN_TOP_K = 512
+
+
 class OrchestratorConfig(BaseConfig):
     algo: AlgoConfig = GRPOAlgoConfig()
     """Training algorithm: sampling plus the per-token training signal (credit
@@ -614,6 +659,56 @@ def validate_env_algorithms(self):
             env_cfg.algo.validate_env(env_cfg.env)
         return self
 
+    @model_validator(mode="after")
+    def setup_truncated_sampling(self):
+        """Truncated policy sampling trains with sampling replay (rollout
+        logprobs are renormalized — see docs/inference.md, Sampling Replay).
+        Owned here: every truncating config gets a top-k bound (bounds the kept
+        sets); opd/opsd is rejected (full-vocab prefill refs would mix
+        normalizations); the gibberish/repetition filters are pruned or rejected
+        (their full-softmax thresholds misfire on renormalized logprobs).
+        Frozen-source envs sample externally and are exempt."""
+        policy_samplings = [
+            env.sampling for env in self.train.source if env.algo is not None and env.algo.sampling.source == "policy"
+        ] or ([self.train.sampling] if not self.train.source else [])
+        truncating = [sampling for sampling in policy_samplings if sampling.truncates_distribution()]
+        if not truncating:
+            return self
+
+        unbounded = [sampling for sampling in truncating if sampling.top_k is None]
+        if unbounded:
+            warnings.warn(
+                f"Truncated train sampling: defaulting top_k = {DEFAULT_TRAIN_TOP_K} so every kept set is "
+                "bounded and sampling replay stays exact. Set top_k explicitly to override.",
+                stacklevel=2,
+            )
+            for sampling in unbounded:
+                sampling.top_k = DEFAULT_TRAIN_TOP_K
+
+        algos = [env.algo for env in self.train.source if env.algo is not None] or [self.algo]
+        if any(algo.type in ("opd", "opsd") for algo in algos):
+            raise ValueError(
+                "opd/opsd is not supported with truncated train sampling: reference logprobs are full-vocab "
+                "prefill scores while trainer logprobs are renormalized over the kept set, biasing the "
+                "ref_kl term. Remove the truncation (top_p/top_k) or the opd/opsd algo."
+            )
+
+        logprob_filter_types = ("gibberish", "repetition")
+        for slot_name in ("pre_batch_filters", "post_batch_filters"):
+            filters = getattr(self, slot_name)
+            if not any(f.type in logprob_filter_types for f in filters):
+                continue
+            if slot_name in self.model_fields_set:
+                raise ValueError(
+                    f"{slot_name} contains logprob-based filters "
+                    f"({[f.type for f in filters if f.type in logprob_filter_types]}) which misfire under "
+                    "truncated sampling: rollout logprobs are renormalized over the kept set, so "
+                    "full-softmax thresholds over-detect repetition and under-detect gibberish. Remove them "
+                    "from the list (zero_advantage is unaffected)."
+                )
+            setattr(self, slot_name, [f for f in filters if f.type not in logprob_filter_types])
+        return self
+
     @property
     def any_policy_sourced(self) -> bool:
         """True when at least one train env samples rollouts from the live policy."""
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
index 8a4a02ae21..226e2c2b28 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -570,6 +570,36 @@ def validate_multi_node_requires_router(self):
             raise ValueError("Multi-node deployments require inference.router to front the per-rank engines.")
         return self
 
+    @model_validator(mode="after")
+    def auto_setup_kept_tokens_capture(self):
+        """Size the inference server's kept-set capture to cover the largest train
+        sampling top-k (OrchestratorConfig guarantees truncating configs have one)."""
+        policy_samplings = [
+            env.sampling
+            for env in self.orchestrator.train.source
+            if env.algo is not None and env.algo.sampling.source == "policy"
+        ] or ([self.orchestrator.train.sampling] if not self.orchestrator.train.source else [])
+        top_ks = [sampling.top_k for sampling in policy_samplings if sampling.top_k is not None]
+        if not top_ks:
+            return self
+        if self.inference is None:
+            warnings.warn(
+                "Truncated train sampling with no managed inference server: set "
+                f"`kept_tokens = {max(top_ks)}` on the standalone server's config so it returns "
+                "the sampling masks the trainer replays.",
+                stacklevel=2,
+            )
+            return self
+        derived = max(top_ks)
+        if "kept_tokens" in self.inference.model_fields_set and self.inference.kept_tokens != derived:
+            warnings.warn(
+                f"Overriding inference.kept_tokens = {self.inference.kept_tokens} with {derived}, "
+                "derived from the largest train sampling top_k (keeps sampling replay exact).",
+                stacklevel=2,
+            )
+        self.inference.kept_tokens = derived
+        return self
+
     @model_validator(mode="after")
     def validate_router_replay_without_kv_offload(self):
         if (
diff --git a/src/prime_rl/inference/server.py b/src/prime_rl/inference/server.py
index e44d9d9e4e..bb811729ff 100644
--- a/src/prime_rl/inference/server.py
+++ b/src/prime_rl/inference/server.py
@@ -30,6 +30,14 @@ def setup_vllm_env(config: InferenceConfig):
     os.environ["VLLM_USE_DEEP_GEMM"] = deep_gemm_enabled
     os.environ["VLLM_MOE_USE_DEEP_GEMM"] = deep_gemm_enabled
 
+    # Kept-set sampling-mask capture (top-p/top-k replay training). Env vars
+    # rather than additional_config because the sampler patch (engine-core
+    # workers) and the output-capture patch (API server procs, applied at
+    # import time) have no guaranteed vLLM config context; children inherit.
+    if config.kept_tokens is not None:
+        os.environ["PRIME_RETURN_KEPT_TOKENS"] = "1"
+        os.environ["PRIME_KEPT_TOKENS_MAX"] = str(config.kept_tokens)
+
     if config.enable_lora:
         os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "True"
 
diff --git a/src/prime_rl/inference/vllm/kept_tokens.py b/src/prime_rl/inference/vllm/kept_tokens.py
new file mode 100644
index 0000000000..4da7932294
--- /dev/null
+++ b/src/prime_rl/inference/vllm/kept_tokens.py
@@ -0,0 +1,243 @@
+"""Kept-token (sampling mask) capture for top-p/top-k replay training.
+
+Truncated sampling (top-p/top-k) renormalizes the sampling distribution
+over a per-token "kept set"; the trainer replays these sets to renormalize its
+own logprobs identically (DeepSeek V3.2 "Keep Sampling Mask", arXiv:2512.02556
+§3.1). vLLM materializes the mask (it's the finite entries of the processed
+logprobs) but never returns it, and its inter-process output structs are fixed
+msgspec/dataclass schemas — so the kept ids ride the existing logprobs channel:
+
+1. Engine-core worker: append ``[-1 separator | kept ids, -1 padded]`` columns
+   to each ``LogprobsTensors`` row; everything downstream is width-agnostic.
+2. API process: split the extension back off before vLLM builds logprob dicts
+   (stock consumers see stock columns), accumulate the ragged rows per request,
+   attach to the finished ``CompletionOutput``.
+3. ``/inference/v1/generate``: serialize as base64
+   ``{"ids": int32 concat, "counts": int32 per completion token}``. Kept sets
+   are decode-only, so PD-disaggregated serving needs no router changes.
+
+A count of 0 means no usable kept set (above the capture width, or the
+position wasn't truncated); the trainer falls back to full-vocab logprobs.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import AsyncIterator
+from typing import Any
+
+import numpy as np
+import pybase64
+from vllm.outputs import RequestOutput
+
+KEPT_TOKENS_ENV = "PRIME_RETURN_KEPT_TOKENS"
+KEPT_TOKENS_MAX_ENV = "PRIME_KEPT_TOKENS_MAX"
+# Fallback only — setup_vllm_env always stamps the env var from inference.kept_tokens.
+KEPT_TOKENS_MAX_DEFAULT = 512
+
+# Separator/padding token id in the widened logprobs rows. Never a valid
+# vocab id, and stock vLLM never emits it (top-k indices and requested
+# logprob_token_ids are always >= 0).
+_SEPARATOR = -1
+
+_EMPTY_KEPT_ROW = np.empty(0, dtype=np.int32)
+
+
+def kept_tokens_enabled() -> bool:
+    return os.environ.get(KEPT_TOKENS_ENV) == "1"
+
+
+def serialize_kept_tokens(kept_token_ids: list[np.ndarray] | None, num_tokens: int) -> dict[str, Any] | None:
+    """Encode per-position kept-set rows as compact base64 raw bytes.
+
+    Returns ``{"ids": b64(int32 concat), "counts": b64(int32[num_tokens])}``
+    or None when nothing was captured. ``counts[i]`` is the kept-set size
+    for completion token i (0 = absent); ``ids`` is the concatenation of
+    all rows in position order.
+    """
+    if not kept_token_ids:
+        return None
+
+    # Stop-token trimming can leave fewer response tokens than sampling steps.
+    rows = kept_token_ids[:num_tokens]
+    if len(rows) < num_tokens:
+        rows = rows + [np.empty(0, dtype=np.int32)] * (num_tokens - len(rows))
+
+    counts = np.fromiter((len(row) for row in rows), dtype=np.int32, count=num_tokens)
+    if not int(counts.sum()):
+        return None
+    ids = np.ascontiguousarray(np.concatenate(rows).astype(np.int32, copy=False))
+    return {
+        "ids": pybase64.b64encode(memoryview(ids)).decode("ascii"),
+        "counts": pybase64.b64encode(memoryview(np.ascontiguousarray(counts))).decode("ascii"),
+    }
+
+
+class KeptTokensCapture:
+    """Records ``kept_token_ids`` off streamed ``RequestOutput``s per choice index."""
+
+    def __init__(self, generator: AsyncIterator[RequestOutput]):
+        self._generator = generator
+        self.kept_tokens: dict[int, dict[str, Any]] = {}
+
+    async def __aiter__(self):
+        async for request_output in self._generator:
+            for output in request_output.outputs:
+                encoded = serialize_kept_tokens(getattr(output, "kept_token_ids", None), len(output.token_ids))
+                if encoded is not None:
+                    self.kept_tokens[output.index] = encoded
+            yield request_output
+
+
+def monkey_patch_kept_tokens_sampler():
+    """Widen sampler logprobs rows with the kept-set extension (engine-core process).
+
+    Intercepts ``self.sample`` for the duration of ``Sampler.forward`` to grab
+    the full processed logprobs the stock forward discards; the kept set per
+    row is their finite entries. Requires ``logprobs_mode="processed_logprobs"``,
+    which also forces the sampling path that materializes the mask (FlashInfer's
+    fused sampler doesn't). Speculative decoding bypasses this patch entirely —
+    the server launcher rejects that combination.
+    """
+    import torch
+    from vllm import envs
+    from vllm.logger import init_logger
+    from vllm.v1.outputs import LogprobsTensors
+    from vllm.v1.sample.sampler import Sampler
+
+    if not kept_tokens_enabled():
+        return
+    if envs.VLLM_USE_V2_MODEL_RUNNER:
+        # The V2 runner samples through a separate Sampler class; capture would be inert.
+        raise ValueError("VLLM_USE_V2_MODEL_RUNNER does not yet support: kept-tokens capture")
+    if getattr(Sampler.forward, "_prime_rl_kept_tokens", False):
+        return
+
+    logger = init_logger(__name__)
+    cap = int(os.environ.get(KEPT_TOKENS_MAX_ENV, str(KEPT_TOKENS_MAX_DEFAULT)))
+    original_forward = Sampler.forward
+
+    def _forward(self, logits, sampling_metadata, predict_bonus_token=False, logprobs_mode_override=None):
+        captured: dict[str, torch.Tensor | None] = {}
+        original_sample = self.sample
+
+        def capturing_sample(*sample_args, **sample_kwargs):
+            sampled, processed_logprobs = original_sample(*sample_args, **sample_kwargs)
+            captured["processed_logprobs"] = processed_logprobs
+            return sampled, processed_logprobs
+
+        # Instance attribute shadows the bound method for this call only;
+        # the model runner drives the sampler single-threaded.
+        self.sample = capturing_sample
+        try:
+            output = original_forward(self, logits, sampling_metadata, predict_bonus_token, logprobs_mode_override)
+        finally:
+            del self.sample
+
+        processed_logprobs = captured.get("processed_logprobs")
+        logprobs_mode = logprobs_mode_override or self.logprobs_mode
+        num_logprobs = sampling_metadata.max_num_logprobs
+        if (
+            processed_logprobs is None
+            or logprobs_mode != "processed_logprobs"
+            or output.logprobs_tensors is None
+            # logprobs=-1 (full vocab) and scoring requests need no extension
+            or num_logprobs is None
+            or num_logprobs < 0
+            or sampling_metadata.logprob_token_ids
+        ):
+            return output
+
+        stock = output.logprobs_tensors
+        num_rows = stock.logprob_token_ids.shape[0]
+        if processed_logprobs.shape[0] != num_rows:
+            return output
+
+        # Fixed width `cap + 1` keeps this device-side (no host sync to stall the
+        # engine loop): a finite entry in the extra column means the kept set
+        # exceeds the cap, and such rows — like untruncated/greedy ones — ship an
+        # empty extension with only the separator marking alignment.
+        ids_dtype = stock.logprob_token_ids.dtype
+        device = processed_logprobs.device
+        width = min(cap + 1, processed_logprobs.shape[-1])
+        ext_logprobs, ext_ids = processed_logprobs.topk(width, dim=-1)
+        finite = ext_logprobs > float("-inf")
+        valid = finite & ~finite[:, -1:]
+        ext_ids = ext_ids.to(ids_dtype).masked_fill_(~valid, _SEPARATOR)
+
+        # The splitter reads only id columns; the logprob extension is -inf filler.
+        separator_ids = torch.full((num_rows, 1), _SEPARATOR, dtype=ids_dtype, device=device)
+        extension_logprobs = torch.full((num_rows, width + 1), float("-inf"), device=device)
+        output.logprobs_tensors = LogprobsTensors(
+            logprob_token_ids=torch.cat([stock.logprob_token_ids, separator_ids, ext_ids], dim=1),
+            logprobs=torch.cat([stock.logprobs, extension_logprobs], dim=1),
+            selected_token_ranks=stock.selected_token_ranks,
+            cu_num_generated_tokens=stock.cu_num_generated_tokens,
+        )
+        return output
+
+    _forward._prime_rl_kept_tokens = True
+    Sampler.forward = _forward
+    logger.warning("Installed kept-tokens sampler patch (cap=%d).", cap)
+
+
+def monkey_patch_kept_tokens_output_capture():
+    """Split kept-set extensions off logprobs rows in the API process.
+
+    Strips the extension before vLLM builds per-position logprob dicts and
+    attaches the accumulated rows to the finished ``CompletionOutput``.
+    Detection is data-driven (the separator id), so rows without extensions
+    pass through untouched.
+    """
+    from vllm.logger import init_logger
+    from vllm.v1.engine.logprobs import LogprobsProcessor
+    from vllm.v1.engine.output_processor import RequestState
+    from vllm.v1.outputs import LogprobsLists
+
+    if getattr(LogprobsProcessor._update_sample_logprobs, "_prime_rl_kept_tokens", False):
+        return
+
+    logger = init_logger(__name__)
+    original_update = LogprobsProcessor._update_sample_logprobs
+    original_new_completion_output = RequestState._new_completion_output
+
+    def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> None:
+        token_ids, logprobs, ranks, cu_num_generated_tokens = logprobs_lists
+        # Append one kept row per position even on extension-less steps, so rows
+        # stay position-aligned if steps start (or stop) carrying separators.
+        kept_rows: list[np.ndarray] | None = getattr(self, "_prime_kept_token_ids", None)
+        if kept_rows is None:
+            kept_rows = self._prime_kept_token_ids = []
+
+        # Rows in one update come from one step's batch tensor: same separator column.
+        separators = np.nonzero(token_ids[0] == _SEPARATOR)[0] if token_ids.size else np.empty(0, dtype=np.int64)
+        if not separators.size:
+            kept_rows.extend([_EMPTY_KEPT_ROW] * len(token_ids))
+            return original_update(self, logprobs_lists)
+
+        split = int(separators[0])
+        for extension in token_ids[:, split + 1 :]:
+            kept_rows.append(np.ascontiguousarray(extension[extension >= 0], dtype=np.int32))
+
+        return original_update(
+            self,
+            LogprobsLists(
+                token_ids[:, :split],
+                logprobs[:, :split],
+                ranks,
+                cu_num_generated_tokens,
+            ),
+        )
+
+    def _new_completion_output(self, *args, **kwargs):
+        output = original_new_completion_output(self, *args, **kwargs)
+        if output.finish_reason is not None and self.logprobs_processor is not None:
+            kept_rows = getattr(self.logprobs_processor, "_prime_kept_token_ids", None)
+            if kept_rows is not None:
+                output.kept_token_ids = kept_rows
+        return output
+
+    _update_sample_logprobs._prime_rl_kept_tokens = True
+    LogprobsProcessor._update_sample_logprobs = _update_sample_logprobs
+    RequestState._new_completion_output = _new_completion_output
+    logger.info("Installed kept-tokens output capture patch (splits -1-separated logprobs extensions).")
diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index 28e2f572ff..0db20e00a2 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -27,6 +27,7 @@
     monkey_patch_tokenize_params_validation,
     monkey_patch_vllm_padded_input_scrub,
 )
+from prime_rl.inference.vllm.kept_tokens import kept_tokens_enabled, monkey_patch_kept_tokens_output_capture
 
 # NOTE: Fix harmony stop token propagation for GPT-OSS models
 # Upstream issue still open: https://github.com/vllm-project/vllm/issues/22519
@@ -46,6 +47,12 @@
 # routed_experts from chat responses since the server-wide enable flag has no
 # per-request toggle.
 monkey_patch_strip_routed_experts_from_chat()
+# NOTE: Kept-set sampling masks (top-p/top-k replay) ride the logprobs rows as a
+# -1-separated extension; split it off before vLLM builds per-position logprob
+# dicts and attach it to the finished CompletionOutput. No-op without
+# PRIME_RETURN_KEPT_TOKENS=1 (set by setup_vllm_env before this module loads).
+if kept_tokens_enabled():
+    monkey_patch_kept_tokens_output_capture()
 # NOTE: vLLM hard-codes a 120s DP coordinator startup timeout, which the rank-0
 # API server blows through when all engine-core ranks on the node are loading
 # weights concurrently (multi-node disaggregated deployments).
@@ -236,6 +243,23 @@ def server(config: InferenceConfig, vllm_extra: dict[str, Any] | None = None):
     assert args is not None
     validate_parsed_serve_args(args)
 
+    if config.kept_tokens is not None:
+        # Both would leave the sampler patch silently inert (no kept sets ever
+        # emitted) while the trainer keeps replaying nothing — exactly the
+        # top-p bias the feature exists to fix. Fail fast instead.
+        if getattr(args, "speculative_config", None):
+            raise ValueError(
+                "kept_tokens capture is incompatible with speculative decoding: vLLM's "
+                "RejectionSampler builds logprobs via gather_logprobs, bypassing the patched "
+                "Sampler.forward. Disable speculative_config or the kept-set capture."
+            )
+        if getattr(args, "logprobs_mode", None) != "processed_logprobs":
+            raise ValueError(
+                "kept_tokens capture requires logprobs_mode='processed_logprobs' (the "
+                "default): the kept set is recovered from the truncation-masked logprobs. "
+                f"Got logprobs_mode={getattr(args, 'logprobs_mode', None)!r} (vllm_extra override?)."
+            )
+
     # Set the worker extension class based on the broadcast backend
     args.worker_extension_cls = WORKER_EXTENSION_CLS[config.weight_broadcast.type]
 
diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py
index e14a5ac83e..a60efec161 100644
--- a/src/prime_rl/inference/vllm/serving_tokens.py
+++ b/src/prime_rl/inference/vllm/serving_tokens.py
@@ -47,11 +47,15 @@
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
+from prime_rl.inference.vllm.kept_tokens import KeptTokensCapture, kept_tokens_enabled
 from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture
 
 
 class PrimeRlGenerateResponseChoice(GenerateResponseChoice):
     routed_experts: dict[str, Any] | None = None
+    # Kept-set sampling masks for top-p/top-k replay training: base64 raw
+    # bytes {"ids": int32 concat, "counts": int32 per completion token}.
+    kept_tokens: dict[str, Any] | None = None
 
 
 class PrimeRlGenerateResponse(GenerateResponse):
@@ -324,6 +328,12 @@ async def serve_tokens_full_generator(  # type: ignore[override]
             )
             result_generator = capture
 
+        # Capture kept-set sampling masks (top-p/top-k replay) the same way.
+        kept_capture: KeptTokensCapture | None = None
+        if kept_tokens_enabled():
+            kept_capture = KeptTokensCapture(result_generator)
+            result_generator = kept_capture
+
         # Always capture the final ``RequestOutput`` so we can attach a
         # ``usage`` block to the response. The router parses ``usage`` for
         # per-run billing metrics; without it the cache-discount counter
@@ -350,6 +360,10 @@ async def serve_tokens_full_generator(  # type: ignore[override]
                 kv_transfer_params=response.kv_transfer_params,
             )
 
+        if kept_capture is not None:
+            for choice in response.choices:
+                choice.kept_tokens = kept_capture.kept_tokens.get(choice.index)
+
         if final_capture.final_res is not None:
             response.usage = _build_usage(final_capture.final_res)
 
diff --git a/src/prime_rl/inference/vllm/worker/__init__.py b/src/prime_rl/inference/vllm/worker/__init__.py
index 1369f4ab0d..a271fbf517 100644
--- a/src/prime_rl/inference/vllm/worker/__init__.py
+++ b/src/prime_rl/inference/vllm/worker/__init__.py
@@ -7,6 +7,7 @@
     monkey_patch_minimax_m2_for_lora,
     monkey_patch_no_moe_lora,
 )
+from prime_rl.inference.vllm.kept_tokens import monkey_patch_kept_tokens_sampler
 
 logger = logging.getLogger(__name__)
 
@@ -24,3 +25,7 @@
 
 # Install fp32 router logits patch; self-gates on additional_config["fp32_router_logits"]
 monkey_patch_fp32_router_logits()
+
+# Install kept-tokens sampler patch (sampling replay); no-op unless
+# PRIME_RETURN_KEPT_TOKENS=1
+monkey_patch_kept_tokens_sampler()
diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py
index 3bffb3504d..19425415f9 100644
--- a/src/prime_rl/orchestrator/envs.py
+++ b/src/prime_rl/orchestrator/envs.py
@@ -259,6 +259,12 @@ def __init__(self, config: TrainSourceConfig, sampler: Sampler, algorithm: Algor
         self.sampler = sampler
         self.algorithm = algorithm
         self.sampling_args = sampler.sampling_args(config.sampling.to_sampling_args())
+        # Truncated policy sampling must ship the kept-set masks the trainer replays.
+        self.requires_kept_masks = (
+            config.sampling.truncates_distribution()
+            and config.algo is not None
+            and config.algo.sampling.source == "policy"
+        )
 
 
 class EvalEnv(Env):
diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py
index fb0dc6abc1..ef1aee928d 100644
--- a/src/prime_rl/orchestrator/train_sink.py
+++ b/src/prime_rl/orchestrator/train_sink.py
@@ -217,6 +217,15 @@ async def process_group(self, group_id: uuid.UUID) -> None:
         for r in survivors:
             for sample in r.samples:
                 sample.temperatures = [temperature] * len(sample.token_ids)
+                if env.requires_kept_masks and sample.kept_tokens is None:
+                    # Rollout logprobs are kept-renormalized; training without the masks
+                    # silently biases every importance ratio.
+                    raise RuntimeError(
+                        f"env '{env_name}' samples with truncation (top_p/top_k) but its rollouts "
+                        "carry no kept-set sampling masks. Set `kept_tokens` on the inference server "
+                        "config (the rl entrypoint derives it automatically) and make sure the server "
+                        "runs prime-rl's vLLM patches."
+                    )
 
         if self.pre_filters:
             apply_filters(self.pre_filters, survivors)
diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py
index 8053453ac8..aa67425117 100644
--- a/src/prime_rl/orchestrator/trajectories.py
+++ b/src/prime_rl/orchestrator/trajectories.py
@@ -21,7 +21,7 @@
 import verifiers.v1 as vf
 
 from prime_rl.transport import TrainingSample
-from prime_rl.transport.types import EncodedTensor, RoutedExperts
+from prime_rl.transport.types import EncodedTensor, KeptTokens, RoutedExperts
 from prime_rl.utils.logger import get_logger
 
 
@@ -66,6 +66,26 @@ def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExp
     return RoutedExperts(data=arr.tobytes(), shape=list(arr.shape), dtype=str(arr.dtype))
 
 
+def _encode_kept_tokens(kept: vf.KeptTokens | None, num_tokens: int) -> KeptTokens | None:
+    """The branch's kept-set sampling masks (`Branch.kept_tokens`) -> the transport
+    `KeptTokens` the trainer replays. Realigns `counts` to `num_tokens`
+    (truncating drops the tail's ids too) as a backstop — `Branch.kept_tokens` already
+    guarantees alignment, mirroring `_encode_routed_experts`. A 0 count just means no
+    replay for that position, so partial coverage stays safe."""
+    if kept is None:
+        return None
+    ids, counts = kept.ids, kept.counts
+    if len(counts) > num_tokens:
+        counts = counts[:num_tokens]
+        ids = ids[: int(counts.sum())]
+    elif len(counts) < num_tokens:
+        counts = np.concatenate([counts, np.zeros(num_tokens - len(counts), dtype=np.int32)])
+    return KeptTokens(
+        ids=np.ascontiguousarray(ids, dtype=np.int32).tobytes(),
+        counts=np.ascontiguousarray(counts, dtype=np.int32).tobytes(),
+    )
+
+
 def iter_trainable_branches(trace: vf.Trace) -> Iterator[tuple[vf.Branch, list[bool]]]:
     """Yield each branch that yields a training sample, with its trainable-token mask.
 
@@ -126,6 +146,7 @@ def trace_to_samples(
                 mm_kwargs=mm_kwargs,
                 mm_token_type_ids=mm_token_type_ids,
                 routed_experts=_encode_routed_experts(branch.routed_experts, len(token_ids)),
+                kept_tokens=_encode_kept_tokens(branch.kept_tokens, len(token_ids)),
             )
         )
     if not samples:
diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py
index 4eb8ff4baa..a56c9bb1aa 100644
--- a/src/prime_rl/trainer/batch.py
+++ b/src/prime_rl/trainer/batch.py
@@ -5,7 +5,7 @@
 import numpy as np
 
 from prime_rl.trainer.utils import balanced_partition
-from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample
+from prime_rl.transport.types import EncodedTensor, KeptTokens, MicroBatch, RoutedExperts, TrainingSample
 
 # Backfill value per component weight stream when a packed sample doesn't
 # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl
@@ -42,6 +42,28 @@ def _pad_routed_experts(micro_batch: MicroBatch, padding_size: int) -> None:
     routed_experts.shape[0] += padding_size
 
 
+_KEPT_ITEMSIZE = np.dtype(np.int32).itemsize
+
+
+def _empty_kept_tokens(num_tokens: int) -> KeptTokens:
+    return KeptTokens(ids=b"", counts=b"\0" * (num_tokens * _KEPT_ITEMSIZE))
+
+
+def _slice_kept_tokens(kept_tokens: KeptTokens, seq_len: int) -> KeptTokens:
+    counts = np.frombuffer(kept_tokens.counts, dtype=np.int32)[:seq_len]
+    return KeptTokens(
+        ids=kept_tokens.ids[: int(counts.sum()) * _KEPT_ITEMSIZE],
+        counts=counts.tobytes(),
+    )
+
+
+def _pad_kept_tokens(micro_batch: MicroBatch, padding_size: int) -> None:
+    kept_tokens = micro_batch.kept_tokens
+    assert kept_tokens is not None
+    # Padding tokens carry no sampling mask (count 0), so only counts grow.
+    kept_tokens.counts += b"\0" * (padding_size * _KEPT_ITEMSIZE)
+
+
 def _slice_encoded(tensor: EncodedTensor, n_rows: int) -> EncodedTensor:
     """First `n_rows` rows of a dim-0-stacked encoded tensor (e.g. pixel_values, image_grid_thw)."""
     row = int(np.prod(tensor.shape[1:])) if len(tensor.shape) > 1 else 1
@@ -147,6 +169,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
     routed_experts = (
         _copy_routed_experts(training_example.routed_experts) if training_example.routed_experts is not None else None
     )
+    # No copy needed: KeptTokens holds immutable bytes, and _pad_kept_tokens only
+    # ever mutates _materialize_bin's own accumulator.
+    kept_tokens = training_example.kept_tokens
 
     if len(input_ids) > seq_len:
         # Multimodal: never split an image's placeholder block — cut to a whole-image boundary
@@ -170,6 +195,8 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
             ref_kl_weights = ref_kl_weights[:cut]
         if routed_experts is not None:
             routed_experts = _slice_routed_experts(routed_experts, cut)
+        if kept_tokens is not None:
+            kept_tokens = _slice_kept_tokens(kept_tokens, cut)
         if mm_token_type_ids is not None:
             mm_token_type_ids = mm_token_type_ids[:cut]
         env_names = env_names[:cut]
@@ -200,6 +227,13 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
         )
         assert len(routed_experts.data) == len(input_ids) * _routed_experts_row_size(routed_experts)
 
+    if kept_tokens is not None:
+        kept_counts = np.frombuffer(kept_tokens.counts, dtype=np.int32)
+        assert len(kept_counts) == len(input_ids), (
+            f"kept_tokens counts: {len(kept_counts)}, input_ids: {len(input_ids)}"
+        )
+        assert len(kept_tokens.ids) == int(kept_counts.sum()) * _KEPT_ITEMSIZE
+
     assert len(env_names) == len(input_ids), f"env_names: {len(env_names)}, input_ids: {len(input_ids)}"
 
     return MicroBatch(
@@ -212,6 +246,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
         ref_logprobs=ref_logprobs,
         temperatures=temperatures,
         routed_experts=routed_experts,
+        kept_tokens=kept_tokens,
         mm_token_type_ids=mm_token_type_ids,
         env_names=env_names,
         mm_kwargs=mm_kwargs,
@@ -313,6 +348,9 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
     # A weight stream materializes as soon as one packed sample carries it; the
     # samples that lack it get the stream's identity fill (STREAM_FILL).
     has_stream = {name: any(getattr(s, name) is not None for _, s in bin_content.samples) for name in STREAM_FILL}
+    # Kept sets are per-token optional (unlike routed_experts): samples without
+    # them get zero-count backfill instead of constraining packing.
+    has_kept_tokens = any(sample.kept_tokens is not None for _, sample in bin_content.samples)
 
     input_ids: list[int] = []
     loss_mask: list[bool] = []
@@ -327,6 +365,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
     streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL}
     seq_lens: list[int] = []
     routed_experts: RoutedExperts | None = None
+    kept_tokens: KeptTokens | None = KeptTokens(ids=b"", counts=b"") if has_kept_tokens else None
     lora_num_tokens = [0] * num_loras
 
     for lora_idx, sample in bin_content.samples:
@@ -365,6 +404,10 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
                     mm_kwargs[key].data += sample.mm_kwargs[key].data
                     mm_kwargs[key].shape[0] += sample.mm_kwargs[key].shape[0]
         seq_lens.extend(sample.seq_lens)
+        if kept_tokens is not None:
+            sample_kept = sample.kept_tokens if sample.kept_tokens is not None else _empty_kept_tokens(sample_len)
+            kept_tokens.ids += sample_kept.ids
+            kept_tokens.counts += sample_kept.counts
         lora_num_tokens[lora_idx] += sample_len
 
     sequence_lengths = [len(sample.input_ids) for _, sample in bin_content.samples]
@@ -382,6 +425,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
         temperatures=temperatures,
         lora_num_tokens=lora_num_tokens,
         routed_experts=routed_experts,
+        kept_tokens=kept_tokens,
         mm_token_type_ids=mm_token_type_ids,
         env_names=env_names,
         mm_kwargs=mm_kwargs,
@@ -513,6 +557,8 @@ def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBa
         micro_batch.mm_token_type_ids.extend([0] * padding_size)
     if micro_batch.routed_experts is not None:
         _pad_routed_experts(micro_batch, padding_size)
+    if micro_batch.kept_tokens is not None:
+        _pad_kept_tokens(micro_batch, padding_size)
     micro_batch.env_names.extend([""] * padding_size)
 
     return micro_batch
@@ -549,6 +595,15 @@ def _assert_token_arrays_aligned(micro_batch: MicroBatch) -> None:
         assert micro_batch.routed_experts.shape[0] == num_tokens, (
             f"routed_experts misaligned after packing: {micro_batch.routed_experts.shape[0]} != {num_tokens} tokens"
         )
+    if micro_batch.kept_tokens is not None:
+        kept_counts = np.frombuffer(micro_batch.kept_tokens.counts, dtype=np.int32)
+        assert len(kept_counts) == num_tokens, (
+            f"kept_tokens misaligned after packing: {len(kept_counts)} != {num_tokens} tokens"
+        )
+        assert len(micro_batch.kept_tokens.ids) == int(kept_counts.sum()) * _KEPT_ITEMSIZE, (
+            f"kept_tokens ids/counts inconsistent after packing: "
+            f"{len(micro_batch.kept_tokens.ids)} bytes != {int(kept_counts.sum())} ids"
+        )
 
 
 def _make_dummy_batch(source: MicroBatch) -> MicroBatch:
@@ -561,6 +616,8 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch:
     dummy.rl_weights = None
     dummy.ce_weights = None
     dummy.ref_kl_weights = None
+    # Fully loss-masked, so replaying sampling masks would be pure wasted work.
+    dummy.kept_tokens = None
     return dummy
 
 
diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py
index d2d6703819..3b3ae8ac3d 100644
--- a/src/prime_rl/trainer/model.py
+++ b/src/prime_rl/trainer/model.py
@@ -1373,6 +1373,7 @@ def forward(
     labels: Int[Tensor, "batch seq"] | None = None,
     temperature: Tensor | None = None,
     routed_experts: Int[Tensor, "batch seq layers topk"] | None = None,
+    kept_tokens: Int[Tensor, "batch seq kept"] | None = None,
     # Generic multimodal kwargs (e.g. {"pixel_values": ...,
     # "image_grid_thw": ...} for Qwen3-VL; just {"pixel_values": ...}
     # for Gemma3). Passed straight through to ``model(**kwargs)`` so
@@ -1390,6 +1391,11 @@ def forward(
         "temperature": temperature,
     }
 
+    # Kept-set sampling masks are consumed by the injected prime lm_head; HF
+    # forwards don't know the kwarg, so only pass it when present.
+    if kept_tokens is not None:
+        kwargs["kept_tokens"] = kept_tokens
+
     if mm_kwargs:
         # Forward the per-model multimodal tensors verbatim, plus the
         # renderer-supplied ``mm_token_type_ids`` (renderer owns the
diff --git a/src/prime_rl/trainer/models/layers/lm_head.py b/src/prime_rl/trainer/models/layers/lm_head.py
index 6bff021d20..55ad77ace6 100644
--- a/src/prime_rl/trainer/models/layers/lm_head.py
+++ b/src/prime_rl/trainer/models/layers/lm_head.py
@@ -44,6 +44,7 @@ def forward(
         hidden_states: torch.Tensor,
         labels: torch.Tensor | None = None,
         temperature: Tensor | None = None,
+        kept_tokens: Tensor | None = None,
     ) -> PrimeLmOutput:
         assert labels is not None, "FusedOutputLinear requires labels for chunked logprob computation"
         assert temperature is not None, "FusedOutputLinear requires per-token temperatures"
@@ -52,9 +53,11 @@ def forward(
         hidden_states = hidden_states.reshape(b * s, h).contiguous()
         labels = labels.reshape(b * s).contiguous()
         inv_t = 1.0 / temperature.reshape(b * s).contiguous()  # [N]
+        if kept_tokens is not None:
+            kept_tokens = kept_tokens.reshape(b * s, kept_tokens.shape[-1]).contiguous()
 
         logprobs, entropy = _SequenceChunkedLogProbEntropyFn.apply(
-            hidden_states, self.weight, labels, inv_t, self.chunk_size
+            hidden_states, self.weight, labels, inv_t, self.chunk_size, kept_tokens
         )
 
         logprobs = logprobs.reshape(b, s)
@@ -67,9 +70,14 @@ def __init__(self, in_features: int, out_features: int):
         super().__init__(in_features, out_features, bias=False)
 
     def forward(
-        self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: Tensor | None = None
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        temperature: Tensor | None = None,
+        kept_tokens: Tensor | None = None,
     ) -> PrimeLmOutput:
-        # VanillaOutputLinear just returns logits - temperature scaling is done externally in train.py
+        # VanillaOutputLinear just returns logits - temperature scaling and
+        # kept-set replay are done externally in train.py
         return PrimeLmOutput(logits=super().forward(hidden_states))
 
 
@@ -86,6 +94,23 @@ def _online_logsumexp_and_weighted_update(
     return m_new, s_new, t_new
 
 
+def kept_replay_mask(kept_tokens: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+    """Positions replayable against their sampling mask: non-empty (-1 is padding)
+    and containing the label (a label outside its mask means misaligned data —
+    fall back to full-vocab rather than emit a corrupt logprob)."""
+    return (kept_tokens >= 0).any(dim=-1) & (kept_tokens == labels.unsqueeze(-1)).any(dim=-1)
+
+
+def _kept_local_indices(
+    kept_chunk: torch.Tensor, vocab_start: int, vocab_end: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Kept-set ids (int64) mapped into a vocab chunk: clamped local indices
+    plus the in-chunk validity mask (-1 entries are padding)."""
+    in_range = (kept_chunk >= vocab_start) & (kept_chunk < vocab_end)
+    local = (kept_chunk - vocab_start).clamp(0, vocab_end - vocab_start - 1)
+    return local, in_range
+
+
 class _SequenceChunkedLogProbEntropyFn(torch.autograd.Function):
     @staticmethod
     def forward(  # type: ignore[override]
@@ -95,9 +120,15 @@ def forward(  # type: ignore[override]
         labels: torch.Tensor,  # [N]
         inv_temperature: torch.Tensor,  # [N]
         chunk_size: int,
+        kept_tokens: torch.Tensor | None = None,  # [N, K] int32, -1-padded
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Returns per-token logprobs and entropy by chunking over flattened sequence tokens.
+
+        Positions with a usable ``kept_tokens`` mask (see ``kept_replay_mask``) are
+        renormalized over the kept set: ``logprob = scaled_logits[label] -
+        logsumexp(scaled_logits[kept])``. Other positions — and entropy, a
+        full-distribution diagnostic — keep full-vocab normalization.
         """
         assert hidden.dim() == 2, f"expected hidden [N,H], got {tuple(hidden.shape)}"
         assert weight.dim() == 2, f"expected weight [V,H], got {tuple(weight.shape)}"
@@ -107,6 +138,10 @@ def forward(  # type: ignore[override]
         assert hidden.shape[1] == weight.shape[1], "hidden/weight H mismatch"
         assert hidden.shape[0] == inv_temperature.shape[0], "hidden/inv_temperature N mismatch"
         assert chunk_size > 0
+        if kept_tokens is not None:
+            assert kept_tokens.dim() == 2 and kept_tokens.shape[0] == hidden.shape[0], (
+                f"expected kept_tokens [N,K], got {tuple(kept_tokens.shape)}"
+            )
 
         device = hidden.device
         n = hidden.shape[0]
@@ -115,6 +150,7 @@ def forward(  # type: ignore[override]
         logprobs = torch.empty((n,), device=device, dtype=torch.float32)
         entropy = torch.empty((n,), device=device, dtype=torch.float32)
         logz = torch.empty((n,), device=device, dtype=torch.float32)
+        replay = torch.zeros((n,), device=device, dtype=torch.bool) if kept_tokens is not None else None
 
         for start in range(0, n, chunk_size):
             end = min(start + chunk_size, n)
@@ -128,6 +164,13 @@ def forward(  # type: ignore[override]
             t = torch.zeros((token_count,), device=device, dtype=torch.float32)
             target_logits = torch.zeros((token_count,), device=device, dtype=torch.float32)
 
+            kept_chunk = kept_tokens[start:end].to(torch.long) if kept_tokens is not None else None
+            if kept_chunk is not None:
+                replay_chunk = kept_replay_mask(kept_chunk, labels_chunk)
+                # Each kept id lives in exactly one vocab chunk; collect its logit
+                # into a [tokens, K] buffer and logsumexp once after the loop.
+                kept_logits = torch.full_like(kept_chunk, float("-inf"), dtype=torch.float32)
+
             for vocab_start in range(0, vocab, vocab_chunk_size):
                 vocab_end = min(vocab_start + vocab_chunk_size, vocab)
                 weight_chunk = weight[vocab_start:vocab_end]
@@ -136,17 +179,26 @@ def forward(  # type: ignore[override]
 
                 m, s, t = _online_logsumexp_and_weighted_update(m, s, t, scaled_logits)
 
+                if kept_chunk is not None:
+                    local, in_range = _kept_local_indices(kept_chunk, vocab_start, vocab_end)
+                    kept_logits = torch.where(in_range, scaled_logits.gather(1, local), kept_logits)
+
                 mask = (labels_chunk >= vocab_start) & (labels_chunk < vocab_end)
                 if torch.any(mask):
                     idx = (labels_chunk[mask] - vocab_start).to(torch.long)
                     target_logits[mask] = scaled_logits[mask, idx]
 
-            logz_chunk = m + torch.log(s)
+            logz_full = m + torch.log(s)
+            if kept_chunk is not None:
+                logz_chunk = torch.where(replay_chunk, torch.logsumexp(kept_logits, dim=-1), logz_full)
+                replay[start:end] = replay_chunk
+            else:
+                logz_chunk = logz_full
             logz[start:end] = logz_chunk
             logprobs[start:end] = target_logits - logz_chunk
-            entropy[start:end] = logz_chunk - (t / s)
+            entropy[start:end] = logz_full - (t / s)
 
-        ctx.save_for_backward(hidden, weight, labels, inv_temperature, logz)
+        ctx.save_for_backward(hidden, weight, labels, inv_temperature, logz, kept_tokens, replay)
         ctx.chunk_size = chunk_size
 
         return logprobs, entropy
@@ -157,7 +209,7 @@ def backward(ctx, grad_logprobs: torch.Tensor, grad_entropy: torch.Tensor | None
             "Backward through entropy is not implemented in FusedOutputLinear"
         )
 
-        hidden, weight, labels, inv_temperature, logz = ctx.saved_tensors
+        hidden, weight, labels, inv_temperature, logz, kept_tokens, replay = ctx.saved_tensors
         chunk_size: int = ctx.chunk_size
 
         n, _ = hidden.shape
@@ -175,12 +227,24 @@ def backward(ctx, grad_logprobs: torch.Tensor, grad_entropy: torch.Tensor | None
             grad_chunk = grad_logprobs[start:end].to(torch.float32)
             inv_t_chunk = inv_temperature[start:end].unsqueeze(-1)
             logz_chunk = logz[start:end]
+            kept_chunk = kept_tokens[start:end].to(torch.long) if kept_tokens is not None else None
+            replay_chunk = replay[start:end] if replay is not None else None
 
             for vocab_start in range(0, vocab, vocab_chunk_size):
                 vocab_end = min(vocab_start + vocab_chunk_size, vocab)
                 weight_chunk = weight[vocab_start:vocab_end]
                 logits_chunk = hidden_chunk @ weight_chunk.t()
                 scaled_logits = logits_chunk.to(torch.float32) * inv_t_chunk
+
+                if kept_chunk is not None:
+                    # Replayed rows get softmax gradient only on kept ids. Mask non-kept
+                    # logits to -inf BEFORE the exp: a non-kept logit above the kept-set
+                    # logZ would overflow exp() to inf (and inf * 0 = NaN in the grads).
+                    local, in_range = _kept_local_indices(kept_chunk, vocab_start, vocab_end)
+                    kept_indicator = torch.zeros(scaled_logits.shape, dtype=torch.int8, device=scaled_logits.device)
+                    kept_indicator.scatter_add_(1, local, in_range.to(torch.int8))
+                    non_kept = replay_chunk.unsqueeze(-1) & (kept_indicator == 0)
+                    scaled_logits.masked_fill_(non_kept, float("-inf"))
                 probs = torch.exp(scaled_logits - logz_chunk.unsqueeze(-1))
 
                 grad_logits = (-grad_chunk).unsqueeze(-1) * probs
@@ -195,7 +259,7 @@ def backward(ctx, grad_logprobs: torch.Tensor, grad_entropy: torch.Tensor | None
                 if needs_weight:
                     grad_weight[vocab_start:vocab_end].add_(grad_logits.to(weight.dtype).t() @ hidden_chunk)
 
-        return grad_hidden, grad_weight, None, None, None
+        return grad_hidden, grad_weight, None, None, None, None
 
 
 def inject_prime_lm_head(
@@ -258,6 +322,7 @@ def new_forward(
         labels: torch.Tensor | None = None,
         logits_to_keep: int = 0,
         temperature: torch.Tensor | None = None,
+        kept_tokens: torch.Tensor | None = None,
         **kwargs: object,
     ) -> PrimeLmOutput:
         # For VLM with images, don't create position_ids - let model compute MRoPE internally
@@ -281,6 +346,7 @@ def new_forward(
             hidden_states[:, slice_indices, :],
             labels[:, slice_indices] if labels is not None else None,
             temperature=temperature[:, slice_indices] if temperature is not None else None,
+            kept_tokens=kept_tokens[:, slice_indices] if kept_tokens is not None else None,
         )
 
     # Bind the new forward to the model
diff --git a/src/prime_rl/trainer/models/layers/lm_head_gemma.py b/src/prime_rl/trainer/models/layers/lm_head_gemma.py
index 7fee1a8d87..610dd88dae 100644
--- a/src/prime_rl/trainer/models/layers/lm_head_gemma.py
+++ b/src/prime_rl/trainer/models/layers/lm_head_gemma.py
@@ -23,9 +23,11 @@ def forward(
         hidden_states: torch.Tensor,
         labels: torch.Tensor | None = None,
         temperature: Tensor | None = None,
+        kept_tokens: Tensor | None = None,
     ) -> PrimeLmOutput:
         assert labels is not None, "GemmaFusedOutputLinear requires labels for chunked logprob computation"
         assert temperature is not None, "GemmaFusedOutputLinear requires per-token temperatures"
+        assert kept_tokens is None, "kept-set replay is not supported with Gemma softcapped lm_heads"
 
         b, s, h = hidden_states.shape
         hidden_states = hidden_states.reshape(b * s, h).contiguous()
@@ -47,7 +49,11 @@ def __init__(self, in_features: int, out_features: int, softcap: float):
         self.softcap = softcap
 
     def forward(
-        self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: Tensor | None = None
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        temperature: Tensor | None = None,
+        kept_tokens: Tensor | None = None,
     ) -> PrimeLmOutput:
         logits = super().forward(hidden_states)
         logits = self.softcap * torch.tanh(logits / self.softcap)
diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py
index 1ad250cf72..0a52118f22 100644
--- a/src/prime_rl/trainer/rl/data.py
+++ b/src/prime_rl/trainer/rl/data.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import TypedDict
 
+import numpy as np
 import torch
 from jaxtyping import Bool, Float, Int
 from torch import Tensor
@@ -39,6 +40,10 @@ class TensorMicroBatch(TypedDict):
     # MoE router replay
     routed_experts: Int[Tensor, "batch seq layers topk"] | None
 
+    # Kept-set sampling masks: kept token ids per position, -1-padded to the
+    # micro batch's max kept-set size (an all--1 row means no mask).
+    kept_tokens: Int[Tensor, "batch seq kept"] | None
+
     # Generic multimodal kwargs — flat dict matching the model's forward
     # signature (e.g. ``{"pixel_values": ..., "image_grid_thw": ...}`` for
     # Qwen3-VL; ``{"pixel_values": ...}`` for Gemma3-VL). The trainer
@@ -132,6 +137,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc
             "lora_num_tokens": lora_num_tokens,
             "seq_lens": torch.tensor(sequence_lengths, dtype=torch.long),
             "routed_experts": None,
+            "kept_tokens": None,
             "mm_kwargs": None,
             "mm_token_type_ids": None,
             "rl_weights": None,
@@ -165,6 +171,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch:
             "lora_num_tokens": lora_num_tokens,
             "seq_lens": torch.tensor([self.seq_len], dtype=torch.long),
             "routed_experts": None,
+            "kept_tokens": None,
             "mm_kwargs": None,
             "mm_token_type_ids": None,
             "rl_weights": None,
@@ -246,6 +253,16 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch:
                 .to(torch.int32)
                 .unsqueeze(0)
             )
+        kept_tokens = None
+        packed_kept_tokens = micro_batch.kept_tokens
+        if packed_kept_tokens is not None:
+            counts = np.frombuffer(packed_kept_tokens.counts, dtype=np.int32)
+            ids = np.frombuffer(packed_kept_tokens.ids, dtype=np.int32)
+            # Boolean assignment fills row-major, matching the flat concat order.
+            max_kept = max(int(counts.max()), 1) if counts.size else 1
+            padded = np.full((len(counts), max_kept), -1, dtype=np.int32)
+            padded[np.arange(max_kept)[None, :] < counts[:, None]] = ids
+            kept_tokens = torch.from_numpy(padded).unsqueeze(0)
         return TensorMicroBatch(
             input_ids=torch.tensor(micro_batch.input_ids, dtype=torch.long).unsqueeze(0),
             position_ids=torch.tensor(micro_batch.position_ids, dtype=torch.long).unsqueeze(0),
@@ -265,6 +282,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch:
             if micro_batch.mm_token_type_ids is not None
             else None,
             routed_experts=routed_experts,
+            kept_tokens=kept_tokens,
             rl_weights=torch.tensor(micro_batch.rl_weights, dtype=torch.float).unsqueeze(0)
             if micro_batch.rl_weights is not None
             else None,
diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py
index 7d9dcce88b..29ebabaee8 100644
--- a/src/prime_rl/trainer/rl/loss.py
+++ b/src/prime_rl/trainer/rl/loss.py
@@ -7,6 +7,7 @@
 from torch import Tensor
 
 from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, IPOLossConfig, LossConfig
+from prime_rl.trainer.models.layers.lm_head import kept_replay_mask
 from prime_rl.utils.utils import import_object
 
 
@@ -54,6 +55,27 @@ def selective_log_softmax(
     return torch.gather(logprobs, dim=-1, index=index.unsqueeze(-1)).squeeze(-1)
 
 
+@jaxtyped(typechecker=typechecker)
+def selective_log_softmax_with_kept(
+    logits: Float[Tensor, "batch seq vocab"],
+    index: Int[Tensor, "batch seq"],
+    kept_tokens: Int[Tensor, "batch seq kept"],
+) -> Float[Tensor, "batch seq"]:
+    """Per-token logprobs with kept-set (sampling-mask) replay: positions with a
+    usable mask (see ``kept_replay_mask``) get ``logits[index] -
+    logsumexp(logits[kept])``, others full-vocab. Non-replayed rows are zeroed
+    before the logsumexp so the unselected ``where`` branch can't emit NaN grads.
+    """
+    full_logprobs = selective_log_softmax(logits, index)
+    replay = kept_replay_mask(kept_tokens, index)
+    kept_logits = torch.gather(logits, -1, kept_tokens.clamp_min(0).long())
+    kept_logits = torch.where(kept_tokens >= 0, kept_logits, float("-inf"))
+    kept_logits = torch.where(replay.unsqueeze(-1), kept_logits, 0.0)
+    logz_kept = torch.logsumexp(kept_logits, dim=-1)
+    target_logits = torch.gather(logits, -1, index.unsqueeze(-1)).squeeze(-1)
+    return torch.where(replay, target_logits - logz_kept, full_logprobs)
+
+
 @jaxtyped(typechecker=typechecker)
 @torch.compile(dynamic=True)
 def compute_entropy(shifted_logits: Float[Tensor, "batch seq vocab"]) -> Float[Tensor, "batch seq"]:
@@ -63,14 +85,15 @@ def compute_entropy(shifted_logits: Float[Tensor, "batch seq vocab"]) -> Float[T
     return entropy
 
 
-def shift_tensor_left(t: Float[Tensor, "batch seq"]) -> Float[Tensor, "batch seq"]:
-    """Shifts the tensor one token to the left.
+def shift_tensor_left(t: Tensor, pad_value: float = 0.0) -> Tensor:
+    """Shifts the tensor one position to the left along dim 1.
 
-    Used to create labels from input_ids: labels[i] = input_ids[i+1].
-    The last position is padded with 0 (a valid token index) since this value
-    will be shifted off by shift_tensor_right and never used.
+    Used to create labels from input_ids: labels[i] = input_ids[i+1]. The last
+    position is padded with ``pad_value`` (0 is a valid token index but gets
+    shifted off by shift_tensor_right and never used). Works for [batch, seq]
+    labels and label-aligned [batch, seq, ...] fields like kept_tokens.
     """
-    return torch.cat([t[:, 1:], torch.full((t.shape[0], 1), 0, device=t.device, dtype=t.dtype)], dim=1)
+    return torch.cat([t[:, 1:], torch.full_like(t[:, :1], pad_value)], dim=1)
 
 
 def shift_tensor_right(t: Float[Tensor, "batch seq"], pad_value: float | None = None) -> Float[Tensor, "batch seq"]:
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 295d3ad15e..9ffb95af2d 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -31,6 +31,7 @@
     compute_loss,
     compute_importance_ratio_and_mismatch_kl,
     selective_log_softmax,
+    selective_log_softmax_with_kept,
     setup_rl_loss_fn,
     shift_tensor_left,
     shift_tensor_right,
@@ -365,6 +366,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                 # we could've gotten routed experts from the inference server, but we didn't enable router replay
                 routed_experts = None
 
+            kept_tokens = micro_batch["kept_tokens"].to("cuda") if micro_batch["kept_tokens"] is not None else None
+
             # Multimodal kwargs are an opaque per-model dict (e.g.
             # {"pixel_values": ..., "image_grid_thw": ...} for Qwen3-VL,
             # just {"pixel_values": ...} for Gemma3-VL) — we move every
@@ -385,6 +388,10 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
             seq_lens = micro_batch["seq_lens"].to("cuda")
 
             labels = shift_tensor_left(input_ids)
+            if kept_tokens is not None:
+                # Kept sets ride at the sampled token's own position (like inference
+                # logprobs); shift to align with the label each position predicts.
+                kept_tokens = shift_tensor_left(kept_tokens, pad_value=-1)
 
             seq_lens_are_pre_shard = False
 
@@ -405,6 +412,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                 labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size)
                 if routed_experts is not None and not defer_vlm_cp_to_model:
                     routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_size)
+                if kept_tokens is not None and not defer_vlm_cp_to_model:
+                    kept_tokens = shard_for_cp(kept_tokens, cp_rank=cp_rank, cp_world_size=cp_size)
 
             if config.model.lora:
                 lora_num_tokens = micro_batch["lora_num_tokens"].to("cuda")
@@ -437,6 +446,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                     seq_lens=seq_lens,
                     seq_lens_are_pre_shard=seq_lens_are_pre_shard,
                     routed_experts=routed_experts,
+                    kept_tokens=kept_tokens,
                 )
 
             if out.get("logprobs") is None:
@@ -445,7 +455,10 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                 logits = out["logits"]
                 # Per-token temperature scaling: temperatures is [batch, seq], logits is [batch, seq, vocab]
                 scaled_logits = logits / temperatures.unsqueeze(-1)
-                out["logprobs"] = selective_log_softmax(scaled_logits, labels)
+                if kept_tokens is not None:
+                    out["logprobs"] = selective_log_softmax_with_kept(scaled_logits, labels, kept_tokens)
+                else:
+                    out["logprobs"] = selective_log_softmax(scaled_logits, labels)
                 out["entropy"] = compute_entropy(scaled_logits)
             # else: FusedOutputLinear was used - logprobs already computed with per-token temperatures
 
diff --git a/src/prime_rl/transport/__init__.py b/src/prime_rl/transport/__init__.py
index 7e49ecbc81..52e4acf4aa 100644
--- a/src/prime_rl/transport/__init__.py
+++ b/src/prime_rl/transport/__init__.py
@@ -9,6 +9,7 @@
     FileSystemTrainingBatchSender,
 )
 from prime_rl.transport.types import (
+    KeptTokens,
     MicroBatch,
     RoutedExperts,
     TrainingBatch,
@@ -72,6 +73,7 @@ def setup_micro_batch_receiver(
     "TrainingSample",
     "TrainingBatch",
     "MicroBatch",
+    "KeptTokens",
     "RoutedExperts",
     "setup_training_batch_sender",
     "setup_training_batch_receiver",
diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py
index 121da68d92..845eade7c6 100644
--- a/src/prime_rl/transport/types.py
+++ b/src/prime_rl/transport/types.py
@@ -18,6 +18,13 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru
     dtype: str
 
 
+# Kept-set sampling masks for top-p/top-k replay: flat int32 id bytes plus an
+# int32 count per token position (0 = no mask); len(ids) == 4 * counts.sum().
+class KeptTokens(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
+    ids: bytes
+    counts: bytes
+
+
 # Orchestrator -> Packer
 class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
     """A single training example — one branch of a rollout as a flat token sequence.
@@ -68,6 +75,10 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr
     # samples without live rl member tokens (the trainer raises otherwise).
     advantages: list[float] | None = None
 
+    # Last field on purpose: array_like structs encode positionally, so appending
+    # keeps the wire layout of earlier fields stable across versions.
+    kept_tokens: KeptTokens | None = None
+
 
 class TrainingBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
     """A batch of training examples with metadata for transport."""
@@ -109,3 +120,6 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
     # Packer-derived metadata used for run-local token exports.
     run_id: str | None = None
     run_step: int | None = None
+
+    # See TrainingSample.kept_tokens; appended last for wire-layout stability.
+    kept_tokens: KeptTokens | None = None