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

docs/inference.md

@@ -296,3 +296,17 @@ enable_return_routed_experts = true
 This however is not free, it adds a significant overhead to the HTTP requests as this payload can grow quite large. We reccomend sizing up the env-server pool (`orchestrator.*.source.serve.pool`) to allow for more parallelization on the verifiers side.
 
 Currently this feature is also not supported with CPU KV cache offload, which can have negative impact on the inference throughput.
+
+### Sampling Replay
+
+Truncated sampling (`top_p < 1`, `top_k`) renormalizes the sampling distribution over the surviving "kept set" of tokens. The rollout logprobs reflect that (`logprobs_mode = "processed_logprobs"`), so the trainer must renormalize over the same set — otherwise every importance ratio is biased and training collapses (DeepSeek V3.2's "Keep Sampling Mask", [arXiv:2512.02556](https://arxiv.org/abs/2512.02556) §3.1; Cognition's [SWE-1.7 post](https://cognition.com/blog/swe-1-7)). prime-rl handles this automatically: the kept-set token ids are recorded at sampling time and the trainer renormalizes its logprobs over them.
+
+```toml
+[orchestrator.train.sampling]
+top_p = 0.95
+top_k = 512   # optional, defaults to 512 under truncation (bounds the kept sets)
+```
+
+That's all — there are no replay flags. Truncated train sampling makes the inference server return kept sets (`inference.kept_tokens`, derived from the largest configured top_k) and the trainer replays whatever masks arrive. Configs that would break under renormalized logprobs are rejected: `opd`/`opsd`, the gibberish/repetition filters (removed from the defaults, rejected if explicitly configured), truncation knobs smuggled via `extra_body`, speculative decoding, and Gemma-family (softcapped) lm_heads. Frozen-source envs are exempt.
+
+When launching the inference server standalone, set `inference.kept_tokens` yourself to cover the clients' top_k.

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

@@ -401,6 +401,9 @@ class InferenceConfig(BaseConfig):
     enable_return_routed_experts: bool = False
     """Return routed experts in responses. Forwarded as ``--enable-return-routed-experts``."""
 
+    kept_tokens: int | None = Field(None, ge=1)
+    """Auto-set for sampling replay: capture width for the per-token kept-set sampling masks returned on ``/inference/v1/generate`` responses (``None`` = capture off). Derived by the ``rl`` entrypoint from the largest train-sampling ``top_k`` and persisted into the per-node config. Not meant to be set by hand, except for standalone-launched servers, where it must cover the clients' top_k."""
+
     enable_fp32_lm_head: bool = True
     """Run the lm_head projection in fp32 via a native bf16×bf16 → fp32 GEMM (``torch.mm`` with ``out_dtype=torch.float32``). Stabilizes logprob precision under FP8/bf16 inference, matching SGLang's ``--enable-fp32-lm-head``. Implemented as a monkey-patch over vLLM's LogitsProcessor, activated by setting ``additional_config["fp32_lm_head"] = True`` on the vLLM config."""
 

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

@@ -1,3 +1,4 @@
+import warnings
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
 
@@ -53,6 +54,16 @@ class TrainSamplingConfig(BaseConfig):
     temperature: float = Field(1.0, ge=0, le=2.0)
     """Sampling temperature."""
 
+    top_p: float = Field(1.0, gt=0, le=1.0)
+    """Nucleus (top-p) sampling for train rollouts. Values below 1.0 truncate the sampling
+    distribution; the ``rl`` entrypoint auto-enables sampling replay so trainer and
+    rollout distributions stay consistent — see docs/inference.md (Sampling Replay)."""
+
+    top_k: int | None = Field(None, ge=1)
+    """Top-k sampling for train rollouts. Truncation triggers sampling replay, and
+    a default top-k is injected when only top-p truncates so kept sets stay
+    bounded — see docs/inference.md (Sampling Replay)."""
+
     max_completion_tokens: int | None = None
     """Maximum output tokens per turn. If None, generates until max context length or EOS."""
 
@@ -61,18 +72,46 @@ class TrainSamplingConfig(BaseConfig):
     extra_body: dict[str, Any] = {}
     """Extra body forwarded with each request to the inference server."""
 
+    def truncates_distribution(self) -> bool:
+        return self.top_p < 1.0 or self.top_k is not None
+
+    @model_validator(mode="after")
+    def validate_no_extra_body_truncation(self):
+        """Truncating values must come from the typed fields — the replay policy reads
+        them. Disabled values pass so resolved configs (where ``resolve_env_config``
+        stamped the ``top_k = -1`` / ``min_p = 0.0`` sentinels) re-validate cleanly."""
+        smuggled = [
+            key
+            for key, truncates in (
+                ("top_p", self.extra_body.get("top_p", 1.0) < 1.0),
+                ("top_k", self.extra_body.get("top_k") not in (None, -1, 0)),
+                ("min_p", self.extra_body.get("min_p", 0.0) > 0.0),
+            )
+            if truncates
+        ]
+        if smuggled:
+            raise ValueError(
+                f"extra_body carries truncating {smuggled}; set them as fields on the train "
+                "sampling config instead (they drive sampling replay)."
+            )
+        return self
+
     def to_sampling_args(self) -> dict[str, Any]:
         """Convert to OAI-compatible sampling args dict, omitting None values."""
         args: dict[str, Any] = {
             "temperature": self.temperature,
-            "top_p": 1.0,
+            "top_p": self.top_p,
             "logprobs": True,
         }
         if self.max_completion_tokens is not None:
             args["max_completion_tokens"] = self.max_completion_tokens
 
-        if self.extra_body:
-            args["extra_body"] = dict(self.extra_body)
+        # top_k rides extra_body (like EvalSamplingConfig), overriding the sentinel.
+        extra_body = dict(self.extra_body)
+        if self.top_k is not None:
+            extra_body["top_k"] = self.top_k
+        if extra_body:
+            args["extra_body"] = extra_body
 
         return args
 
@@ -444,6 +483,12 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 ]
 
 
+# Top-k injected on truncated policy sampling that has none: large enough that a
+# 0.95-0.99 nucleus rarely reaches it (the sampling policy is essentially unchanged),
+# small enough to bound the kept-set capture width and trainer mask tensors.
+DEFAULT_TRAIN_TOP_K = 512
+
+
 class OrchestratorConfig(BaseConfig):
     algo: AlgoConfig = GRPOAlgoConfig()
     """Training algorithm: sampling plus the per-token training signal (credit
@@ -614,6 +659,56 @@ def validate_env_algorithms(self):
             env_cfg.algo.validate_env(env_cfg.env)
         return self
 
+    @model_validator(mode="after")
+    def setup_truncated_sampling(self):
+        """Truncated policy sampling trains with sampling replay (rollout
+        logprobs are renormalized — see docs/inference.md, Sampling Replay).
+        Owned here: every truncating config gets a top-k bound (bounds the kept
+        sets); opd/opsd is rejected (full-vocab prefill refs would mix
+        normalizations); the gibberish/repetition filters are pruned or rejected
+        (their full-softmax thresholds misfire on renormalized logprobs).
+        Frozen-source envs sample externally and are exempt."""
+        policy_samplings = [
+            env.sampling for env in self.train.source if env.algo is not None and env.algo.sampling.source == "policy"
+        ] or ([self.train.sampling] if not self.train.source else [])
+        truncating = [sampling for sampling in policy_samplings if sampling.truncates_distribution()]
+        if not truncating:
+            return self
+
+        unbounded = [sampling for sampling in truncating if sampling.top_k is None]
+        if unbounded:
+            warnings.warn(
+                f"Truncated train sampling: defaulting top_k = {DEFAULT_TRAIN_TOP_K} so every kept set is "
+                "bounded and sampling replay stays exact. Set top_k explicitly to override.",
+                stacklevel=2,
+            )
+            for sampling in unbounded:
+                sampling.top_k = DEFAULT_TRAIN_TOP_K
+
+        algos = [env.algo for env in self.train.source if env.algo is not None] or [self.algo]
+        if any(algo.type in ("opd", "opsd") for algo in algos):
+            raise ValueError(
+                "opd/opsd is not supported with truncated train sampling: reference logprobs are full-vocab "
+                "prefill scores while trainer logprobs are renormalized over the kept set, biasing the "
+                "ref_kl term. Remove the truncation (top_p/top_k) or the opd/opsd algo."
+            )
+
+        logprob_filter_types = ("gibberish", "repetition")
+        for slot_name in ("pre_batch_filters", "post_batch_filters"):
+            filters = getattr(self, slot_name)
+            if not any(f.type in logprob_filter_types for f in filters):
+                continue
+            if slot_name in self.model_fields_set:
+                raise ValueError(
+                    f"{slot_name} contains logprob-based filters "
+                    f"({[f.type for f in filters if f.type in logprob_filter_types]}) which misfire under "
+                    "truncated sampling: rollout logprobs are renormalized over the kept set, so "
+                    "full-softmax thresholds over-detect repetition and under-detect gibberish. Remove them "
+                    "from the list (zero_advantage is unaffected)."
+                )
+            setattr(self, slot_name, [f for f in filters if f.type not in logprob_filter_types])
+        return self
+
     @property
     def any_policy_sourced(self) -> bool:
         """True when at least one train env samples rollouts from the live policy."""

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

@@ -570,6 +570,36 @@ def validate_multi_node_requires_router(self):
             raise ValueError("Multi-node deployments require inference.router to front the per-rank engines.")
         return self
 
+    @model_validator(mode="after")
+    def auto_setup_kept_tokens_capture(self):
+        """Size the inference server's kept-set capture to cover the largest train
+        sampling top-k (OrchestratorConfig guarantees truncating configs have one)."""
+        policy_samplings = [
+            env.sampling
+            for env in self.orchestrator.train.source
+            if env.algo is not None and env.algo.sampling.source == "policy"
+        ] or ([self.orchestrator.train.sampling] if not self.orchestrator.train.source else [])
+        top_ks = [sampling.top_k for sampling in policy_samplings if sampling.top_k is not None]
+        if not top_ks:
+            return self
+        if self.inference is None:
+            warnings.warn(
+                "Truncated train sampling with no managed inference server: set "
+                f"`kept_tokens = {max(top_ks)}` on the standalone server's config so it returns "
+                "the sampling masks the trainer replays.",
+                stacklevel=2,
+            )
+            return self
+        derived = max(top_ks)
+        if "kept_tokens" in self.inference.model_fields_set and self.inference.kept_tokens != derived:
+            warnings.warn(
+                f"Overriding inference.kept_tokens = {self.inference.kept_tokens} with {derived}, "
+                "derived from the largest train sampling top_k (keeps sampling replay exact).",
+                stacklevel=2,
+            )
+        self.inference.kept_tokens = derived
+        return self
+
     @model_validator(mode="after")
     def validate_router_replay_without_kv_offload(self):
         if (

src/prime_rl/inference/server.py

@@ -30,6 +30,14 @@ def setup_vllm_env(config: InferenceConfig):
     os.environ["VLLM_USE_DEEP_GEMM"] = deep_gemm_enabled
     os.environ["VLLM_MOE_USE_DEEP_GEMM"] = deep_gemm_enabled
 
+    # Kept-set sampling-mask capture (top-p/top-k replay training). Env vars
+    # rather than additional_config because the sampler patch (engine-core
+    # workers) and the output-capture patch (API server procs, applied at
+    # import time) have no guaranteed vLLM config context; children inherit.
+    if config.kept_tokens is not None:
+        os.environ["PRIME_RETURN_KEPT_TOKENS"] = "1"
+        os.environ["PRIME_KEPT_TOKENS_MAX"] = str(config.kept_tokens)
+
     if config.enable_lora:
         os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "True"