From ba32dc2173d385b3ac9ff648e398b2f642509fd6 Mon Sep 17 00:00:00 2001
From: Mika Senghaas <mail@mikasenghaas.de>
Date: Tue, 11 Aug 2026 17:12:45 +0000
Subject: [PATCH 1/3] feat: top-p/top-k train sampling with sampling replay

Squash of feat/top-p-mask-replay (PR #2979) onto current main, adapting to
the vllm pass-through inference config, multi-tenant removal, and the v0
env compat drop.

Co-authored-by: fares <fares@primeintellect.ai>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/inference.md                             |  14 +
 .../src/prime_rl/configs/inference.py         |   3 +
 .../src/prime_rl/configs/orchestrator.py      | 101 +++++++-
 .../src/prime_rl/configs/rl.py                |  30 +++
 src/prime_rl/inference/server.py              |   8 +
 src/prime_rl/inference/vllm/kept_tokens.py    | 243 ++++++++++++++++++
 src/prime_rl/inference/vllm/server.py         |  24 ++
 src/prime_rl/inference/vllm/serving_tokens.py |  14 +
 .../inference/vllm/worker/__init__.py         |   5 +
 src/prime_rl/orchestrator/envs.py             |   6 +
 src/prime_rl/orchestrator/train_sink.py       |   9 +
 src/prime_rl/orchestrator/trajectories.py     |  23 +-
 src/prime_rl/trainer/batch.py                 |  59 ++++-
 src/prime_rl/trainer/model.py                 |   6 +
 src/prime_rl/trainer/models/layers/lm_head.py |  82 +++++-
 .../trainer/models/layers/lm_head_gemma.py    |   8 +-
 src/prime_rl/trainer/rl/data.py               |  18 ++
 src/prime_rl/trainer/rl/loss.py               |  35 ++-
 src/prime_rl/trainer/rl/train.py              |  15 +-
 src/prime_rl/transport/__init__.py            |   2 +
 src/prime_rl/transport/types.py               |  14 +
 21 files changed, 698 insertions(+), 21 deletions(-)
 create mode 100644 src/prime_rl/inference/vllm/kept_tokens.py

diff --git a/docs/inference.md b/docs/inference.md
index 92bd96c377..6748c0d759 100644
--- a/docs/inference.md
+++ b/docs/inference.md
@@ -293,3 +293,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 d551ac7243..ce3b864a45 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -463,6 +463,9 @@ class InferenceConfig(BaseConfig):
     use_pd_kv_transfer: bool = False
     """Auto-set for disaggregated P/D: emit the NIXL transfer connector. Persisted into the per-node config (which drops ``deployment``) so the connector is still built per worker. Not meant to be set by hand."""
 
+    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 74af6c0de8..f4608ed05d 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
 
@@ -48,6 +49,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."""
 
@@ -56,18 +67,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
 
@@ -390,6 +429,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
@@ -557,6 +602,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 6d70c3fae1..fce89be4f0 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -574,6 +574,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 7e0cbc737b..832d6ba0b6 100644
--- a/src/prime_rl/inference/server.py
+++ b/src/prime_rl/inference/server.py
@@ -26,6 +26,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.vllm.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 728574aba9..585fe1a423 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -25,6 +25,7 @@
     monkey_patch_strip_routed_experts_from_chat,
     monkey_patch_tokenize_params_validation,
 )
+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
@@ -41,6 +42,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).
@@ -229,6 +236,23 @@ def server(config: InferenceConfig):
     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 21c4445086..d080dfa911 100644
--- a/src/prime_rl/orchestrator/envs.py
+++ b/src/prime_rl/orchestrator/envs.py
@@ -136,6 +136,12 @@ def __init__(self, config: TrainSourceConfig, address: str, sampler: Sampler, al
         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 5f052c14ed..dc9725f062 100644
--- a/src/prime_rl/orchestrator/train_sink.py
+++ b/src/prime_rl/orchestrator/train_sink.py
@@ -204,6 +204,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 a1282ad0e9..3075afe84d 100644
--- a/src/prime_rl/trainer/batch.py
+++ b/src/prime_rl/trainer/batch.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 
-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
@@ -278,6 +278,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
@@ -383,6 +405,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
@@ -406,6 +431,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]
@@ -436,6 +463,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(
@@ -448,6 +482,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,
@@ -541,6 +576,9 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> 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] = []
@@ -555,6 +593,7 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> 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
 
     for sample in bin_content.samples:
         sample_len = len(sample.input_ids)
@@ -592,6 +631,10 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> 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
 
     sequence_lengths = [len(sample.input_ids) for sample in bin_content.samples]
     assert sum(sequence_lengths) == len(input_ids), (sequence_lengths, len(input_ids))
@@ -607,6 +650,7 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> 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,
@@ -733,6 +777,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
@@ -769,6 +815,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:
@@ -781,6 +836,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 aab4e35290..f2e1959379 100644
--- a/src/prime_rl/trainer/rl/data.py
+++ b/src/prime_rl/trainer/rl/data.py
@@ -1,6 +1,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
@@ -36,6 +37,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
@@ -122,6 +127,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc
             "lora_num_tokens": torch.tensor([input_ids.shape[0]], dtype=torch.int32),
             "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,
@@ -151,6 +157,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch:
             "lora_num_tokens": torch.tensor([self.seq_len], dtype=torch.int32),
             "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,
@@ -206,6 +213,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),
@@ -226,6 +243,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 5946fcb3b7..8e049688b9 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -30,6 +30,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,
@@ -338,6 +339,8 @@ def train(config: TrainerConfig):
                 # 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
@@ -358,6 +361,10 @@ def train(config: TrainerConfig):
             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
 
@@ -378,6 +385,8 @@ def train(config: TrainerConfig):
                 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")
@@ -410,6 +419,7 @@ def train(config: TrainerConfig):
                     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:
@@ -418,7 +428,10 @@ def train(config: TrainerConfig):
                 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 06b9719e5e..4abdb7312c 100644
--- a/src/prime_rl/transport/__init__.py
+++ b/src/prime_rl/transport/__init__.py
@@ -7,6 +7,7 @@
     FileSystemMicroBatchSender,
 )
 from prime_rl.transport.types import (
+    KeptTokens,
     MicroBatch,
     RoutedExperts,
     TrainingSample,
@@ -48,6 +49,7 @@ def setup_micro_batch_receiver(
     "MicroBatchSender",
     "TrainingSample",
     "MicroBatch",
+    "KeptTokens",
     "RoutedExperts",
     "setup_micro_batch_sender",
     "setup_micro_batch_receiver",
diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py
index 21fe63f810..155727246a 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
+
+
 # Produced by the orchestrator's train sink; consumed in-process by
 # ``prepare_batch``, which packs samples into per-rank ``MicroBatch``es.
 class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
@@ -69,6 +76,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
+
 
 # Orchestrator -> Trainer
 class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
@@ -97,3 +108,6 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
     rl_weights: list[float] | None = None
     ce_weights: list[float] | None = None
     ref_kl_weights: list[float] | None = None
+
+    # See TrainingSample.kept_tokens; appended last for wire-layout stability.
+    kept_tokens: KeptTokens | None = None

From 4b894afad517fa58641a9ea1652c13fa7540b7c7 Mon Sep 17 00:00:00 2001
From: Mika Senghaas <mail@mikasenghaas.de>
Date: Tue, 11 Aug 2026 17:21:00 +0000
Subject: [PATCH 2/3] feat: ship the kept-set extension on the id tensor only

Nothing between sampler and API process pairs logprob ids and values
column-wise, so the -inf float filler rows were pure IPC overhead.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/prime_rl/inference/vllm/kept_tokens.py | 16 ++++++++++++----
 src/prime_rl/inference/vllm/server.py      |  2 +-
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/prime_rl/inference/vllm/kept_tokens.py b/src/prime_rl/inference/vllm/kept_tokens.py
index 4da7932294..d5785eaca2 100644
--- a/src/prime_rl/inference/vllm/kept_tokens.py
+++ b/src/prime_rl/inference/vllm/kept_tokens.py
@@ -8,10 +8,17 @@
 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.
+   to each ``LogprobsTensors`` id row (ids only — nothing between sampler and
+   API process pairs ids and logprobs column-wise); 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``.
+
+vLLM is growing native support — ``enable_return_sampling_mask``
+(vllm-project/vllm#49577, unreleased) — with the same semantics and
+constraints; once it ships in a released version this module reduces to the
+``routed_experts``-style API-layer glue (``KeptTokensCapture`` + serializer).
 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.
@@ -165,12 +172,13 @@ def capturing_sample(*sample_args, **sample_kwargs):
         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.
+        # Only the id tensor grows: the splitter reads ids alone, nothing between
+        # sampler and API process pairs the two tensors column-wise (LogprobsLists
+        # slices rows), and skipping a float extension halves the IPC overhead.
         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),
+            logprobs=stock.logprobs,
             selected_token_ranks=stock.selected_token_ranks,
             cu_num_generated_tokens=stock.cu_num_generated_tokens,
         )
diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py
index 585fe1a423..942bac0d6f 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -250,7 +250,7 @@ def server(config: InferenceConfig):
             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?)."
+                f"Got logprobs_mode={getattr(args, 'logprobs_mode', None)!r} (inference.vllm override?)."
             )
 
     # Set the worker extension class based on the broadcast backend

From d09488d57ae4fa32fc31a31b3d69eb66130d7e76 Mon Sep 17 00:00:00 2001
From: Mika Senghaas <mail@mikasenghaas.de>
Date: Tue, 11 Aug 2026 17:57:16 +0000
Subject: [PATCH 3/3] feat: gate sampling-mask capture on additional_config,
 not env vars

The enable flag rides vLLM's additional_config as
enable_return_sampling_mask (named after the in-flight native vLLM flag,
vllm-project/vllm#49577), snapshotted at Sampler.__init__ like the fp32
patches. The API-process patches are data-driven off the separator id and
install unconditionally. The capture width is a fixed constant; the
orchestrator rejects train-sampling top_k above it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 docs/inference.md                             |   4 +-
 .../src/prime_rl/configs/inference.py         |   6 +-
 .../src/prime_rl/configs/orchestrator.py      |  24 +++-
 .../src/prime_rl/configs/rl.py                |  23 ++--
 src/prime_rl/inference/server.py              |   8 --
 src/prime_rl/inference/vllm/kept_tokens.py    | 105 ++++++++++--------
 src/prime_rl/inference/vllm/server.py         |  26 +++--
 src/prime_rl/inference/vllm/serving_tokens.py |  15 ++-
 .../inference/vllm/worker/__init__.py         |   4 +-
 src/prime_rl/orchestrator/train_sink.py       |   6 +-
 10 files changed, 120 insertions(+), 101 deletions(-)

diff --git a/docs/inference.md b/docs/inference.md
index 6748c0d759..24c032895a 100644
--- a/docs/inference.md
+++ b/docs/inference.md
@@ -304,6 +304,6 @@ 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.
+That's all — there are no replay flags. Truncated train sampling makes the inference server return kept sets (`inference.enable_return_sampling_mask`, auto-enabled; the capture width is fixed at 512, so train-sampling `top_k` above 512 is rejected) 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.
+When launching the inference server standalone, set `inference.enable_return_sampling_mask = true` yourself; clients must sample with `top_k <= 512`.
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 ce3b864a45..75ca2c89cb 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py
@@ -463,8 +463,8 @@ class InferenceConfig(BaseConfig):
     use_pd_kv_transfer: bool = False
     """Auto-set for disaggregated P/D: emit the NIXL transfer connector. Persisted into the per-node config (which drops ``deployment``) so the connector is still built per worker. Not meant to be set by hand."""
 
-    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_return_sampling_mask: bool = False
+    """Auto-set for sampling replay: return per-token kept-set sampling masks on ``/inference/v1/generate`` responses, at a fixed capture width of 512 (the orchestrator bounds train-sampling ``top_k`` to match). Named after vLLM's in-flight native flag (vllm-project/vllm#49577); until that releases it activates prime-rl's capture patches via ``additional_config``. Auto-enabled by the ``rl`` entrypoint under truncated train sampling and persisted into the per-node config; set by hand only for standalone-launched servers."""
 
     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."""
@@ -626,6 +626,8 @@ def to_namespace(self) -> Namespace:
             additional_config["fp32_lm_head"] = True
         if self.enable_fp32_router_logits:
             additional_config["fp32_router_logits"] = True
+        if self.enable_return_sampling_mask:
+            additional_config["enable_return_sampling_mask"] = True
         if additional_config:
             namespace.additional_config = additional_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 f4608ed05d..06d936d663 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
@@ -429,10 +429,13 @@ 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
+# Top-k injected on truncated policy sampling that has none, and the hard upper
+# bound for explicit top-k: it equals the inference server's fixed kept-set
+# capture width (SAMPLING_MASK_MAX in prime_rl/inference/vllm/kept_tokens.py),
+# so kept sets never overflow and replay stays exact. Large enough that a
+# 0.95-0.99 nucleus rarely reaches it (the sampling policy is essentially
+# unchanged), small enough to bound the trainer mask tensors.
+TRAIN_TOP_K_BOUND = 512
 
 
 class OrchestratorConfig(BaseConfig):
@@ -618,15 +621,24 @@ def setup_truncated_sampling(self):
         if not truncating:
             return self
 
+        oversized = [sampling.top_k for sampling in truncating if (sampling.top_k or 0) > TRAIN_TOP_K_BOUND]
+        if oversized:
+            raise ValueError(
+                f"Truncated train sampling with top_k = {max(oversized)} exceeds the inference server's "
+                f"fixed kept-set capture width ({TRAIN_TOP_K_BOUND}): overflowing kept sets would be "
+                "dropped and silently bias the replayed importance ratios. Use top_k <= "
+                f"{TRAIN_TOP_K_BOUND}."
+            )
+
         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 "
+                f"Truncated train sampling: defaulting top_k = {TRAIN_TOP_K_BOUND} 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
+                sampling.top_k = TRAIN_TOP_K_BOUND
 
         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):
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 fce89be4f0..af6b529680 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -575,33 +575,26 @@ def validate_multi_node_requires_router(self):
         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)."""
+    def auto_setup_sampling_mask_capture(self):
+        """Truncated train sampling needs the inference server to return the kept-set
+        sampling masks the trainer replays (OrchestratorConfig guarantees truncating
+        configs are bounded by the fixed capture width)."""
         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:
+        if not any(sampling.truncates_distribution() for sampling in policy_samplings):
             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.",
+                "`enable_return_sampling_mask = true` 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
+        self.inference.enable_return_sampling_mask = True
         return self
 
     @model_validator(mode="after")
diff --git a/src/prime_rl/inference/server.py b/src/prime_rl/inference/server.py
index 832d6ba0b6..7e0cbc737b 100644
--- a/src/prime_rl/inference/server.py
+++ b/src/prime_rl/inference/server.py
@@ -26,14 +26,6 @@ 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.vllm.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
index d5785eaca2..3a0b2347f0 100644
--- a/src/prime_rl/inference/vllm/kept_tokens.py
+++ b/src/prime_rl/inference/vllm/kept_tokens.py
@@ -10,26 +10,31 @@
 1. Engine-core worker: append ``[-1 separator | kept ids, -1 padded]`` columns
    to each ``LogprobsTensors`` id row (ids only — nothing between sampler and
    API process pairs ids and logprobs column-wise); everything downstream is
-   width-agnostic.
+   width-agnostic. Gated per engine on ``additional_config
+   ["enable_return_sampling_mask"]``, snapshotted at ``Sampler.__init__``
+   (fp32_lm_head-style) — no env vars.
 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``.
-
-vLLM is growing native support — ``enable_return_sampling_mask``
-(vllm-project/vllm#49577, unreleased) — with the same semantics and
-constraints; once it ships in a released version this module reduces to the
-``routed_experts``-style API-layer glue (``KeptTokensCapture`` + serializer).
+   attach to the finished ``CompletionOutput``. Purely data-driven off the
+   separator id, so it installs unconditionally.
 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.
+The orchestrator therefore bounds train-sampling ``top_k`` to
+``SAMPLING_MASK_MAX`` so kept sets never overflow the capture width.
+
+vLLM is growing native support under the same flag name —
+``enable_return_sampling_mask`` (vllm-project/vllm#49577, unreleased) — with
+the same semantics and constraints; once it ships in a released version this
+module reduces to the ``routed_experts``-style API-layer glue
+(``KeptTokensCapture`` + serializer).
 """
 
 from __future__ import annotations
 
-import os
 from collections.abc import AsyncIterator
 from typing import Any
 
@@ -37,21 +42,21 @@
 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
+# Enable flag in vLLM's additional_config, named after the in-flight native
+# vLLM flag (vllm-project/vllm#49577); set from inference.enable_return_sampling_mask.
+SAMPLING_MASK_KEY = "enable_return_sampling_mask"
+
+# Fixed kept-set capture width. Not configurable: the orchestrator rejects
+# train-sampling top_k above this, so capture never overflows and replay is
+# exact at every position.
+SAMPLING_MASK_MAX = 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)
-
+SEPARATOR = -1
 
-def kept_tokens_enabled() -> bool:
-    return os.environ.get(KEPT_TOKENS_ENV) == "1"
+EMPTY_KEPT_ROW = np.empty(0, dtype=np.int32)
 
 
 def serialize_kept_tokens(kept_token_ids: list[np.ndarray] | None, num_tokens: int) -> dict[str, Any] | None:
@@ -99,32 +104,43 @@ async def __aiter__(self):
 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.
+    Self-gates on ``additional_config["enable_return_sampling_mask"]``,
+    snapshotted at ``Sampler.__init__`` where vLLM guarantees a
+    ``set_current_vllm_config()`` context (same mechanism as fp32_lm_head).
+    When enabled, 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, and the V2 model
+    runner's separate Sampler class would leave it inert — the server launcher
+    rejects both combinations.
     """
     import torch
-    from vllm import envs
+    from vllm.config import get_current_vllm_config
     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)))
+    cap = SAMPLING_MASK_MAX
+    original_init = Sampler.__init__
     original_forward = Sampler.forward
 
-    def _forward(self, logits, sampling_metadata, predict_bonus_token=False, logprobs_mode_override=None):
+    def patched_init(self, *args, **kwargs):
+        original_init(self, *args, **kwargs)
+        additional_config = get_current_vllm_config().additional_config or {}
+        self._prime_return_sampling_mask = bool(additional_config.get(SAMPLING_MASK_KEY, False))
+        if self._prime_return_sampling_mask:
+            logger.warning("Kept-set sampling-mask capture ENABLED for this Sampler instance (cap=%d).", cap)
+
+    def patched_forward(self, logits, sampling_metadata, predict_bonus_token=False, logprobs_mode_override=None):
+        if not getattr(self, "_prime_return_sampling_mask", False):
+            return original_forward(self, logits, sampling_metadata, predict_bonus_token, logprobs_mode_override)
+
         captured: dict[str, torch.Tensor | None] = {}
         original_sample = self.sample
 
@@ -170,12 +186,12 @@ def capturing_sample(*sample_args, **sample_kwargs):
         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)
+        ext_ids = ext_ids.to(ids_dtype).masked_fill_(~valid, SEPARATOR)
 
         # Only the id tensor grows: the splitter reads ids alone, nothing between
         # sampler and API process pairs the two tensors column-wise (LogprobsLists
         # slices rows), and skipping a float extension halves the IPC overhead.
-        separator_ids = torch.full((num_rows, 1), _SEPARATOR, dtype=ids_dtype, device=device)
+        separator_ids = torch.full((num_rows, 1), SEPARATOR, dtype=ids_dtype, device=device)
         output.logprobs_tensors = LogprobsTensors(
             logprob_token_ids=torch.cat([stock.logprob_token_ids, separator_ids, ext_ids], dim=1),
             logprobs=stock.logprobs,
@@ -184,9 +200,10 @@ def capturing_sample(*sample_args, **sample_kwargs):
         )
         return output
 
-    _forward._prime_rl_kept_tokens = True
-    Sampler.forward = _forward
-    logger.warning("Installed kept-tokens sampler patch (cap=%d).", cap)
+    patched_forward._prime_rl_kept_tokens = True
+    Sampler.__init__ = patched_init
+    Sampler.forward = patched_forward
+    logger.info("Installed kept-tokens sampler patch (self-gates on additional_config[%r]).", SAMPLING_MASK_KEY)
 
 
 def monkey_patch_kept_tokens_output_capture():
@@ -209,7 +226,7 @@ def monkey_patch_kept_tokens_output_capture():
     original_update = LogprobsProcessor._update_sample_logprobs
     original_new_completion_output = RequestState._new_completion_output
 
-    def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> None:
+    def patched_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.
@@ -218,9 +235,9 @@ def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> 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)
+        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))
+            kept_rows.extend([EMPTY_KEPT_ROW] * len(token_ids))
             return original_update(self, logprobs_lists)
 
         split = int(separators[0])
@@ -237,7 +254,7 @@ def _update_sample_logprobs(self, logprobs_lists: LogprobsLists) -> None:
             ),
         )
 
-    def _new_completion_output(self, *args, **kwargs):
+    def patched_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)
@@ -245,7 +262,7 @@ def _new_completion_output(self, *args, **kwargs):
                 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
+    patched_update_sample_logprobs._prime_rl_kept_tokens = True
+    LogprobsProcessor._update_sample_logprobs = patched_update_sample_logprobs
+    RequestState._new_completion_output = patched_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 942bac0d6f..7c830bcbe8 100644
--- a/src/prime_rl/inference/vllm/server.py
+++ b/src/prime_rl/inference/vllm/server.py
@@ -25,7 +25,7 @@
     monkey_patch_strip_routed_experts_from_chat,
     monkey_patch_tokenize_params_validation,
 )
-from prime_rl.inference.vllm.kept_tokens import kept_tokens_enabled, monkey_patch_kept_tokens_output_capture
+from prime_rl.inference.vllm.kept_tokens import 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
@@ -44,10 +44,9 @@
 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()
+# dicts and attach it to the finished CompletionOutput. Data-driven off the
+# separator id, so rows without extensions pass through untouched.
+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,22 +235,27 @@ def server(config: InferenceConfig):
     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
+    if config.enable_return_sampling_mask:
+        # All three 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 "
+                "enable_return_sampling_mask 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."
+                "Sampler.forward. Disable speculative_config or the sampling-mask capture."
             )
         if getattr(args, "logprobs_mode", None) != "processed_logprobs":
             raise ValueError(
-                "kept_tokens capture requires logprobs_mode='processed_logprobs' (the "
+                "enable_return_sampling_mask 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} (inference.vllm override?)."
             )
+        if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") == "1":
+            raise ValueError(
+                "enable_return_sampling_mask does not support VLLM_USE_V2_MODEL_RUNNER=1: the V2 "
+                "runner samples through a separate Sampler class the capture patch never sees."
+            )
 
     # 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 a60efec161..dbade1ae24 100644
--- a/src/prime_rl/inference/vllm/serving_tokens.py
+++ b/src/prime_rl/inference/vllm/serving_tokens.py
@@ -47,7 +47,7 @@
 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.kept_tokens import KeptTokensCapture
 from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture
 
 
@@ -329,10 +329,10 @@ 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
+        # Unconditional: outputs only carry kept rows when the engine's
+        # sampling-mask capture is enabled, so this is a no-op otherwise.
+        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
@@ -360,9 +360,8 @@ 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)
+        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 a271fbf517..abc3bc21e4 100644
--- a/src/prime_rl/inference/vllm/worker/__init__.py
+++ b/src/prime_rl/inference/vllm/worker/__init__.py
@@ -26,6 +26,6 @@
 # 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
+# Install kept-tokens sampler patch (sampling replay); self-gates on
+# additional_config["enable_return_sampling_mask"]
 monkey_patch_kept_tokens_sampler()
diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py
index dc9725f062..4d3662e429 100644
--- a/src/prime_rl/orchestrator/train_sink.py
+++ b/src/prime_rl/orchestrator/train_sink.py
@@ -209,9 +209,9 @@ async def process_group(self, group_id: uuid.UUID) -> None:
                     # 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."
+                        "carry no kept-set sampling masks. Set `enable_return_sampling_mask = true` on "
+                        "the inference server config (the rl entrypoint does this automatically) and "
+                        "make sure the server runs prime-rl's vLLM patches."
                     )
 
         if self.pre_filters: