feat: top-p/top-k train sampling with sampling replay

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.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.enable_return_sampling_mask = true` yourself; clients must sample with `top_k <= 512`.

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."""
 
+    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."""
 
@@ -623,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
 

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,15 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 ]
 
 
+# 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):
     algo: AlgoConfig = GRPOAlgoConfig()
     """Training algorithm: sampling plus the per-token training signal (credit
@@ -557,6 +605,65 @@ 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
+
+        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 = {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 = 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):
+            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."""

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -574,6 +574,29 @@ 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_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 [])
+        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 "
+                "`enable_return_sampling_mask = true` on the standalone server's config so it "
+                "returns the sampling masks the trainer replays.",
+                stacklevel=2,
+            )
+            return self
+        self.inference.enable_return_sampling_mask = True
+        return self
+
     @model_validator(mode="after")
     def validate_router_replay_without_kv_offload(self):
         if (