diff --git a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py
index cb3e526c95..98059023f1 100644
--- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py
+++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py
@@ -2,6 +2,7 @@
 
 from dataclasses import dataclass
 
+import numpy as np
 import torch
 
 from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import (
@@ -205,3 +206,51 @@ def scatter_packed_token_values_to_batch(
     )
     batch_values[output_mask] = values[packed_mask]
     return batch_values
+
+
+class TokenMetadataTrace:
+    """Accumulate arrays whose first dimension is aligned to tokens."""
+
+    def __init__(self) -> None:
+        self._chunks: list[np.ndarray] = []
+        self._schema: tuple[tuple[int, ...], np.dtype] | None = None
+        self._num_rows = 0
+        self._finalized = False
+
+    @property
+    def num_rows(self) -> int:
+        return self._num_rows
+
+    def append(self, rows: np.ndarray, *, expected_rows: int) -> None:
+        if self._finalized:
+            raise RuntimeError("token metadata trace is already finalized")
+        if isinstance(expected_rows, bool) or not isinstance(expected_rows, int) or expected_rows < 0:
+            raise ValueError(f"expected_rows must be a non-negative integer, got {expected_rows!r}")
+        if not isinstance(rows, np.ndarray):
+            raise TypeError("token metadata rows must be a NumPy array")
+        if rows.ndim < 1:
+            raise ValueError("token metadata must have a token-row dimension")
+        if rows.shape[0] != expected_rows:
+            raise ValueError(f"token metadata has {rows.shape[0]} rows, expected {expected_rows}")
+        if not rows.flags.c_contiguous:
+            raise ValueError("token metadata rows must be contiguous")
+
+        schema = (rows.shape[1:], rows.dtype)
+        if self._schema is None:
+            self._schema = schema
+        elif schema != self._schema:
+            raise ValueError(f"token metadata schema changed from {self._schema} to {schema}")
+
+        self._chunks.append(rows)
+        self._num_rows += expected_rows
+
+    def finalize(self, *, expected_rows: int) -> np.ndarray:
+        if self._finalized:
+            raise RuntimeError("token metadata trace is already finalized")
+        if self._num_rows != expected_rows:
+            raise ValueError(f"token metadata trace has {self._num_rows} rows, expected {expected_rows}")
+        if not self._chunks:
+            raise ValueError("token metadata trace has no chunks")
+
+        self._finalized = True
+        return self._chunks[0] if len(self._chunks) == 1 else np.concatenate(self._chunks, axis=0)
diff --git a/skyrl/backends/skyrl_train/distributed/ulysses/utils.py b/skyrl/backends/skyrl_train/distributed/ulysses/utils.py
index f51105b053..0ab9ede056 100644
--- a/skyrl/backends/skyrl_train/distributed/ulysses/utils.py
+++ b/skyrl/backends/skyrl_train/distributed/ulysses/utils.py
@@ -106,10 +106,10 @@ def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: Proc
     return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)
 
 
-def _pad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor:
+def _pad_tensor(x: Tensor, dim: int, padding_size: int, padding_value: int = 0) -> Tensor:
     shape = list(x.shape)
     shape[dim] = padding_size
-    pad = torch.zeros(shape, dtype=x.dtype, device=x.device)
+    pad = torch.full(shape, padding_value, dtype=x.dtype, device=x.device)
     return torch.cat([x, pad], dim=dim)
 
 
@@ -267,6 +267,7 @@ def ulysses_pad_and_slice_inputs(
     position_ids_rmpad: Optional[torch.Tensor] = None,
     attention_mask_rmpad: Optional[torch.Tensor] = None,
     sp_size: int = 1,
+    input_padding_value: int = 0,
 ):
     """
     Pad and slice input_ids to be divisible by sp_size
@@ -277,9 +278,11 @@ def ulysses_pad_and_slice_inputs(
     The is the utility of pre-forward for ulysses sequence parallelism
 
     Args:
-        input_ids_rmpad: shape of [bsz, seqlen]
+        input_ids_rmpad: shape of [bsz, seqlen, ...]. Trailing dimensions are
+            preserved so token-aligned metadata can use the same partition.
         position_ids_rmpad: shape of [bsz, seqlen]
         sp_size (int): ulysses sequence parallelism size
+        input_padding_value: Value for padded entries in ``input_ids_rmpad``.
 
     Returns:
         torch.Tensor: padded and sliced input_ids
@@ -294,10 +297,15 @@ def ulysses_pad_and_slice_inputs(
     group = get_ulysses_sequence_parallel_group()
     if group is None:
         raise ValueError("`sp_size` > 1 but no ulysses sequence parallel group set.")
-    _, total_seq_len = input_ids_rmpad.shape
+    total_seq_len = input_ids_rmpad.size(1)
     pad_size = (sp_size - total_seq_len % sp_size) % sp_size
     if pad_size > 0:
-        input_ids_rmpad = torch.nn.functional.pad(input_ids_rmpad, (0, pad_size), value=0)
+        input_ids_rmpad = _pad_tensor(
+            input_ids_rmpad,
+            dim=1,
+            padding_size=pad_size,
+            padding_value=input_padding_value,
+        )
         if position_ids_rmpad is not None:
             pad_pos_ids = (
                 torch.arange(pad_size, device=position_ids_rmpad.device)
diff --git a/skyrl/backends/skyrl_train/inference_servers/base.py b/skyrl/backends/skyrl_train/inference_servers/base.py
index aac9ee6f6d..73d2fc9ae5 100644
--- a/skyrl/backends/skyrl_train/inference_servers/base.py
+++ b/skyrl/backends/skyrl_train/inference_servers/base.py
@@ -34,6 +34,7 @@ class InferenceEngineInput(TypedDict):
     # Optional prefix-cache salt forwarded to vLLM as the request ``cache_salt`` so cache blocks are
     # only shared between requests carrying the same salt. See ``GeneratorConfig.use_cache_salt``.
     cache_salt: Optional[str]
+    routed_experts_prompt_starts: Optional[List[int]]
 
 
 class InferenceEngineOutput(TypedDict):
@@ -50,6 +51,7 @@ class InferenceEngineOutput(TypedDict):
     response_logprobs: Optional[List[List[float]]]
     prompt_logprobs: Optional[List[List[float]]]  # per-prompt-token logprobs under the current model
     rollout_expert_indices: Optional[List[RoutedExpertIndices]]
+    rollout_sample_support: Optional[List[List[List[int]]]]
 
 
 class InferenceEngineInterface(ABC):
diff --git a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py
index a9b8f43a2c..978b087789 100644
--- a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py
+++ b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py
@@ -104,3 +104,20 @@ def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices
     if compact.dtype != dtype:
         raise ValueError(f"packed routed_experts uses non-canonical dtype {dtype.name}; expected {compact.dtype.name}")
     return compact
+
+
+def clamp_sampled_logprobs(sampled: np.ndarray) -> Tuple[list[dict[str, float]], int]:
+    """Build ``logprobs.content`` from a flat array of sampled-token logprobs.
+
+    The array form of :func:`build_logprobs_content`, for vLLM's flat-logprobs
+    rows. A non-finite entry here reaches the client as JSON ``null`` (orjson
+    emits ``null`` rather than raising), which then fails tensor conversion in
+    preprocessing, so it has to be floored before serialization.
+
+    NaN needs the same treatment as ``-inf`` and is easier to miss: callers that
+    screen support columns with ``isneginf`` do not catch it.
+    """
+    finite = np.isfinite(sampled)
+    num_clamped = int((~finite).sum())
+    values = np.where(finite, sampled, CLAMPED_LOGPROB).tolist()
+    return [{"logprob": value} for value in values], num_clamped
diff --git a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py
index 7bfbd42a64..ba18a9a451 100644
--- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py
+++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py
@@ -76,6 +76,10 @@
 from skyrl.backends.skyrl_train.inference_servers.generate_wire import (
     decode_packed_routed_experts,
 )
+from skyrl.backends.skyrl_train.inference_servers.sample_support_set_wire import (
+    decode_sample_support_set,
+)
+from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices
 from skyrl.backends.utils import convert_vllm_prompt_logprobs
 from skyrl.env_vars import (
     SKYRL_GENERATE_CONCURRENCY_PER_ENGINE,
@@ -167,6 +171,165 @@ class SampleResponse(TypedDict):
     topk_prompt_logprobs: Optional[List[Optional[List[Tuple[int, float]]]]]
 
 
+@dataclass(frozen=True)
+class RemoteGenerateResult:
+    """Raw token generation result returned by ``RemoteGenerateClient``."""
+
+    raw_response: Dict[str, Any]
+    response_ids: List[int]
+    response_logprobs: Optional[List[float]]
+    stop_reason: str
+    routed_experts: Optional[RoutedExpertIndices]
+    sample_support: Optional[List[List[int]]]
+
+
+@dataclass
+class RemoteGenerateClient:
+    """Reusable HTTP client for one raw-token generation request."""
+
+    proxy_url: str
+    _session: Optional[aiohttp.ClientSession] = field(default=None, init=False, repr=False)
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        current_loop = asyncio.get_running_loop()
+        if self._session is not None and not self._session.closed and self._session.loop != current_loop:
+            self._session = None
+        if self._session is None or self._session.closed:
+            connector = aiohttp.TCPConnector(
+                limit=SKYRL_HTTP_CONNECTION_LIMIT,
+                keepalive_timeout=2,
+            )
+            self._session = aiohttp.ClientSession(
+                connector=connector,
+                timeout=aiohttp.ClientTimeout(total=None),
+            )
+        return self._session
+
+    async def _post(self, url: str, json: Dict[str, Any], headers: Optional[Dict[str, str]] = None) -> Any:
+        """POST JSON with retry on transient connection and response-decoding failures."""
+        session = await self._get_session()
+        last_exc: Optional[Exception] = None
+        for attempt in range(_DATA_PLANE_RETRIES):
+            try:
+                async with session.post(url, json=json, headers=headers) as resp:
+                    try:
+                        body = orjson.loads(await resp.read())
+                    except orjson.JSONDecodeError as exc:
+                        if 400 <= resp.status < 500:
+                            text = await resp.text()
+                            raise aiohttp.ClientResponseError(
+                                resp.request_info,
+                                resp.history,
+                                status=resp.status,
+                                message=text or resp.reason,
+                                headers=resp.headers,
+                            ) from exc
+                        last_exc = exc
+                        logger.debug(f"retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}")
+                        await asyncio.sleep(1)
+                        continue
+                    raise_for_status(resp, body)
+                    return body
+            except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as exc:
+                last_exc = exc
+                logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}")
+                await asyncio.sleep(1)
+        if last_exc is None:
+            raise RuntimeError(f"POST failed without an exception for {url=}")
+        raise last_exc
+
+    async def generate(
+        self,
+        *,
+        prompt_token_ids: List[int],
+        sampling_params: Dict[str, Any],
+        session_id: Optional[Any],
+        model: str,
+        return_routed_experts: bool = False,
+        routed_experts_prompt_start: Optional[int] = None,
+        return_sample_support: bool = False,
+        mm_features: Optional[MultiModalFeatures] = None,
+        cache_salt: Optional[str] = None,
+    ) -> RemoteGenerateResult:
+        """Generate one raw-token completion with optional replay metadata."""
+        if routed_experts_prompt_start is not None:
+            if not return_routed_experts:
+                raise ValueError("routed_experts_prompt_start requires return_routed_experts=True")
+            if (
+                isinstance(routed_experts_prompt_start, bool)
+                or not isinstance(routed_experts_prompt_start, int)
+                or not 0 <= routed_experts_prompt_start <= len(prompt_token_ids)
+            ):
+                raise ValueError("routed_experts_prompt_start must be an integer within the prompt")
+
+        use_skyrl_endpoint = return_routed_experts or return_sample_support
+        path = "/skyrl/v1/generate" if use_skyrl_endpoint else "/inference/v1/generate"
+        request_sampling_params = dict(sampling_params)
+        if routed_experts_prompt_start is not None:
+            request_sampling_params["routed_experts_prompt_start"] = routed_experts_prompt_start
+        payload: Dict[str, Any] = {
+            "sampling_params": request_sampling_params,
+            "model": model,
+            "token_ids": prompt_token_ids,
+        }
+        if return_sample_support:
+            payload["return_sample_support"] = True
+        if mm_features:
+            payload["features"] = mm_features
+        # `cache_salt` is a top-level request field (forwarded to vLLM's TokensPrompt), not a sampling
+        # param.
+        if cache_salt is not None:
+            payload["cache_salt"] = cache_salt
+
+        headers = {"Content-Type": "application/json"}
+        if session_id:
+            headers["X-Session-ID"] = str(session_id)
+
+        response = await self._post(f"{self.proxy_url}{path}", json=payload, headers=headers)
+        choice = response["choices"][0]
+        token_ids = choice["token_ids"]
+        logprobs = choice.get("logprobs")
+        response_logprobs = None
+        if logprobs is not None:
+            logprobs_content = logprobs.get("content", [])
+            if logprobs_content:
+                response_logprobs = [logprob_info["logprob"] for logprob_info in logprobs_content]
+
+        routed_experts = None
+        if return_routed_experts:
+            packed_routed_experts = choice.get("routed_experts")
+            if not isinstance(packed_routed_experts, dict):
+                raise ValueError("/skyrl/v1/generate must return packed routed_experts")
+            routed_experts = decode_packed_routed_experts(packed_routed_experts)
+
+        sample_support = (
+            decode_sample_support_set(choice["rollout_sample_support"]).tolist() if return_sample_support else None
+        )
+
+        return RemoteGenerateResult(
+            raw_response=response,
+            response_ids=token_ids,
+            response_logprobs=response_logprobs,
+            stop_reason=choice["finish_reason"],
+            routed_experts=routed_experts,
+            sample_support=sample_support,
+        )
+
+    async def aclose(self) -> None:
+        if self._session is not None and not self._session.closed:
+            await self._session.close()
+        self._session = None
+
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+        state["_session"] = None
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._session = None
+
+
 @dataclass
 class RemoteInferenceClient(InferenceEngineInterface):
     """
@@ -215,6 +378,9 @@ class RemoteInferenceClient(InferenceEngineInterface):
     enable_return_routed_experts: bool = False
     """Whether to return routed expert indices (R3 / rollout router replay)."""
 
+    enable_return_sample_support_set: bool = False
+    """Whether to return sampled-token support sets for replay."""
+
     uses_lora_weight_sync: bool = False
     """True when the trainer syncs LoRA adapters (rather than full/merged weights). When True,
     `sleep()` is forced to level=1: level=2 discards the base model from VRAM with no CPU backup,
@@ -225,7 +391,7 @@ class RemoteInferenceClient(InferenceEngineInterface):
     """Optional HF tokenizer for local tokenize/detokenize (avoids HTTP round-trips)."""
 
     # Private fields excluded from repr for cleaner output
-    _session: Optional[aiohttp.ClientSession] = field(default=None, repr=False)
+    _generate_client: Optional[RemoteGenerateClient] = field(default=None, repr=False)
     _world_size: Optional[Tuple[int, int]] = field(default=None, repr=False)
     _gen_sem: Optional[asyncio.Semaphore] = field(default=None, repr=False)
     _detok_sem: Optional[asyncio.Semaphore] = field(default=None, repr=False)
@@ -282,66 +448,16 @@ def _get_semaphores(self) -> Tuple[Optional[asyncio.Semaphore], Optional[asyncio
             self._sem_loop = current_loop
         return self._gen_sem, self._detok_sem
 
+    def _get_generate_client(self) -> RemoteGenerateClient:
+        if self._generate_client is None:
+            self._generate_client = RemoteGenerateClient(proxy_url=self.proxy_url)
+        return self._generate_client
+
     async def _get_session(self) -> aiohttp.ClientSession:
-        """Get or create the aiohttp session."""
-        # Re-use the existing session object if it is not closed.
-        # Note that we also create a new session object if the event loop has changed, since
-        # aiohttp.ClientSession is tied to the event loop.
-        current_loop = asyncio.get_running_loop()
-        if self._session is not None and not self._session.closed and self._session.loop != current_loop:
-            # Event loop changed - the old session is unusable (bound to a dead loop).
-            self._session = None
-        if self._session is None or self._session.closed:
-            # keepalive_timeout must be shorter than the server's timeout_keep_alive
-            # (uvicorn default: 5s). Otherwise aiohttp reuses connections the server
-            # has already closed, causing ECONNRESET under high concurrency.
-            connector = aiohttp.TCPConnector(
-                limit=SKYRL_HTTP_CONNECTION_LIMIT,
-                keepalive_timeout=2,
-            )
-            self._session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=None))
-        return self._session
+        return await self._get_generate_client()._get_session()
 
     async def _post(self, url: str, json: Dict[str, Any], headers: Optional[Dict[str, str]] = None) -> Any:
-        """POST with retry + backoff on transient connection errors.
-
-        Between generate bursts the pool's keep-alive connections go stale
-        (server closes them after ``timeout_keep_alive``).  An immediate
-        retry would grab another stale connection from the same pool, so we
-        sleep briefly to let the connector detect and purge dead sockets
-        before the next attempt.
-        """
-        session = await self._get_session()
-        last_exc: Optional[Exception] = None
-        for attempt in range(_DATA_PLANE_RETRIES):
-            try:
-                async with session.post(url, json=json, headers=headers) as resp:
-                    try:
-                        body = orjson.loads(await resp.read())
-                    except orjson.JSONDecodeError as e:
-                        if 400 <= resp.status < 500:
-                            # Non-JSON client error (e.g. plain text 422 from vllm-router).
-                            # Raise immediately — client errors won't succeed on retry.
-                            text = await resp.text()
-                            raise aiohttp.ClientResponseError(
-                                resp.request_info,
-                                resp.history,
-                                status=resp.status,
-                                message=text or resp.reason,
-                                headers=resp.headers,
-                            )
-                        last_exc = e
-                        logger.debug(f"retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {e}")
-                        await asyncio.sleep(1)
-                        continue
-                    raise_for_status(resp, body)
-                    return body
-            except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as e:
-                last_exc = e
-                logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {e}")
-                await asyncio.sleep(1)
-                continue
-        raise last_exc  # type: ignore[misc]
+        return await self._get_generate_client()._post(url, json=json, headers=headers)
 
     # ---------------------------
     # Data Plane
@@ -406,6 +522,12 @@ async def generate(
         session_ids = input_batch.get("session_ids")
         mm_features = input_batch.get("mm_features")
         cache_salt = input_batch.get("cache_salt")
+        routed_experts_prompt_starts = input_batch.get("routed_experts_prompt_starts")
+        if routed_experts_prompt_starts is not None:
+            if not self.enable_return_routed_experts:
+                raise ValueError("routed_experts_prompt_starts requires enable_return_routed_experts=True")
+            if len(routed_experts_prompt_starts) != len(prompt_token_ids):
+                raise ValueError("routed_experts_prompt_starts must have one entry per prompt")
         get_logprobs = sampling_params.get("logprobs") is not None
 
         # Two semaphores decouple the generate and detokenize stages:
@@ -429,6 +551,9 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]:
                     sampling_params=sampling_params,
                     session_id=session_ids[idx] if session_ids and idx < len(session_ids) else None,
                     mm_features=mm_features[idx] if mm_features and idx < len(mm_features) else None,
+                    routed_experts_prompt_start=(
+                        routed_experts_prompt_starts[idx] if routed_experts_prompt_starts is not None else None
+                    ),
                     model=model,
                     cache_salt=cache_salt,
                 )
@@ -438,6 +563,9 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]:
                     sampling_params=sampling_params,
                     session_id=session_ids[idx] if session_ids and idx < len(session_ids) else None,
                     mm_features=mm_features[idx] if mm_features and idx < len(mm_features) else None,
+                    routed_experts_prompt_start=(
+                        routed_experts_prompt_starts[idx] if routed_experts_prompt_starts is not None else None
+                    ),
                     model=model,
                     cache_salt=cache_salt,
                 )
@@ -454,6 +582,11 @@ async def _throttled_detokenize(token_ids: List[int]) -> str:
         rollout_expert_indices = (
             [result["routed_experts"] for result in raw_results] if self.enable_return_routed_experts else None
         )
+        rollout_sample_support = (
+            [result["rollout_sample_support"] for result in raw_results]
+            if self.enable_return_sample_support_set
+            else None
+        )
 
         return InferenceEngineOutput(
             responses=responses,
@@ -461,6 +594,7 @@ async def _throttled_detokenize(token_ids: List[int]) -> str:
             response_ids=[r["response_ids"] for r in raw_results],
             response_logprobs=[r["response_logprobs"] for r in raw_results] if get_logprobs else None,
             rollout_expert_indices=rollout_expert_indices,
+            rollout_sample_support=rollout_sample_support,
         )
 
     async def _generate_single(
@@ -471,64 +605,25 @@ async def _generate_single(
         model: str,
         mm_features: Optional[MultiModalFeatures] = None,
         cache_salt: Optional[str] = None,
+        routed_experts_prompt_start: Optional[int] = None,
     ) -> Dict[str, Any]:
-        """
-        Generate completion for a single prompt.
-
-        With keep-mode pause, in-flight requests are frozen by the vLLM
-        scheduler and resume where they left off after /resume. No retry
-        logic is needed.
-
-        Returns:
-            Dict with keys: stop_reason, response_ids, response_logprobs
-        """
-        url = (
-            f"{self.proxy_url}/skyrl/v1/generate"
-            if self.enable_return_routed_experts
-            else f"{self.proxy_url}/inference/v1/generate"
+        result = await self._get_generate_client().generate(
+            prompt_token_ids=prompt_token_ids,
+            sampling_params=sampling_params,
+            session_id=session_id,
+            model=model,
+            return_routed_experts=self.enable_return_routed_experts,
+            routed_experts_prompt_start=routed_experts_prompt_start,
+            return_sample_support=self.enable_return_sample_support_set,
+            mm_features=mm_features,
+            cache_salt=cache_salt,
         )
-
-        payload: dict[str, Any] = {
-            "sampling_params": sampling_params,
-            "model": model,
-            "token_ids": prompt_token_ids,
-        }
-        if mm_features:
-            payload["features"] = mm_features
-        # `cache_salt` is a top-level request field (forwarded to vLLM's TokensPrompt), not a sampling
-        # param.
-        if cache_salt is not None:
-            payload["cache_salt"] = cache_salt
-
-        headers = {"Content-Type": "application/json"}
-        if session_id:
-            headers["X-Session-ID"] = str(session_id)
-
-        response = await self._post(url, json=payload, headers=headers)
-
-        choice = response["choices"][0]
-        token_ids = choice["token_ids"]
-        stop_reason = choice["finish_reason"]
-
-        response_logprobs: Optional[List[float]] = None
-        logprobs = choice.get("logprobs")
-        if logprobs is not None:
-            logprobs_content = logprobs.get("content", [])
-            if logprobs_content:
-                response_logprobs = [logprob_info["logprob"] for logprob_info in logprobs_content]
-
-        routed_experts = None
-        if self.enable_return_routed_experts:
-            packed_routed_experts = choice.get("routed_experts")
-            if not isinstance(packed_routed_experts, dict):
-                raise ValueError("/skyrl/v1/generate must return packed routed_experts")
-            routed_experts = decode_packed_routed_experts(packed_routed_experts)
-
         return {
-            "stop_reason": stop_reason,
-            "response_ids": token_ids,
-            "response_logprobs": response_logprobs,
-            "routed_experts": routed_experts,
+            "stop_reason": result.stop_reason,
+            "response_ids": result.response_ids,
+            "response_logprobs": result.response_logprobs,
+            "routed_experts": result.routed_experts,
+            "rollout_sample_support": result.sample_support,
         }
 
     async def _render_for_sample(
@@ -1405,9 +1500,8 @@ async def get_world_size(self) -> Tuple[int, int]:
 
     async def teardown(self) -> None:
         """Close HTTP session."""
-        if self._session and not self._session.closed:
-            await self._session.close()
-            self._session = None
+        if self._generate_client is not None:
+            await self._generate_client.aclose()
 
     async def __aenter__(self) -> "RemoteInferenceClient":
         """Async context manager entry."""
@@ -1424,7 +1518,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
     def __getstate__(self) -> dict:
         """Exclude non-serializable fields from pickle."""
         state = self.__dict__.copy()
-        state["_session"] = None
         state["_gen_sem"] = None
         state["_detok_sem"] = None
         state["_sem_loop"] = None
@@ -1433,19 +1526,12 @@ def __getstate__(self) -> dict:
     def __setstate__(self, state: dict) -> None:
         """Restore state after unpickling."""
         self.__dict__.update(state)
-        self._session = None
         self._gen_sem = None
         self._detok_sem = None
         self._sem_loop = None
 
-    async def aclose(self):
-        if self._session is not None:
-            try:
-                await self._session.close()
-            except Exception as e:
-                logger.warning(f"Encountered exception {e} while closing client session")
-                pass
-            self._session = None
+    async def aclose(self) -> None:
+        await self.teardown()
 
 
 def raise_for_status(resp: aiohttp.ClientResponse, body: Optional[Any] = None) -> None:
diff --git a/skyrl/backends/skyrl_train/inference_servers/sample_support_set_wire.py b/skyrl/backends/skyrl_train/inference_servers/sample_support_set_wire.py
new file mode 100644
index 0000000000..4288e5042d
--- /dev/null
+++ b/skyrl/backends/skyrl_train/inference_servers/sample_support_set_wire.py
@@ -0,0 +1,73 @@
+"""Packed wire format for dense sampler-support vocab IDs."""
+
+from typing import Literal, TypedDict
+
+import numpy as np
+import pybase64
+
+
+class PackedSampleSupportSet(TypedDict):
+    data: str
+    shape: list[int]
+    dtype: Literal["int32"]
+
+
+_DTYPE = np.dtype("int32")
+
+
+def _validated_dense_support(values: np.ndarray) -> np.ndarray:
+    if not isinstance(values, np.ndarray):
+        raise TypeError("sample support must be a NumPy array")
+    if values.ndim != 2:
+        raise ValueError(f"sample support must have shape [tokens, top_k], got {values.shape!r}")
+    if not np.issubdtype(values.dtype, np.integer):
+        raise ValueError("sample support must contain integer vocab IDs")
+    if values.size:
+        minimum = int(values.min())
+        maximum = int(values.max())
+        if minimum < -1 or maximum > np.iinfo(np.int32).max:
+            raise ValueError("sample support IDs must be -1 padding or non-negative int32 vocab IDs")
+        if np.any((values[:, :-1] == -1) & (values[:, 1:] >= 0)):
+            raise ValueError("sample support padding must be trailing -1 values")
+    return np.ascontiguousarray(values, dtype=_DTYPE)
+
+
+def encode_sample_support_set(values: np.ndarray) -> PackedSampleSupportSet:
+    dense = _validated_dense_support(values)
+    return {
+        "data": pybase64.b64encode(memoryview(dense)).decode("ascii"),
+        "shape": list(dense.shape),
+        "dtype": "int32",
+    }
+
+
+def decode_sample_support_set(payload: PackedSampleSupportSet) -> np.ndarray:
+    if not isinstance(payload, dict):
+        raise TypeError("packed sample support must be a dictionary")
+    try:
+        data_b64 = payload["data"]
+        raw_shape = payload["shape"]
+        dtype_name = payload["dtype"]
+    except KeyError as exc:
+        raise ValueError("packed sample support is missing data/shape/dtype") from exc
+    if dtype_name != "int32":
+        raise ValueError("packed sample support dtype must be int32")
+    if not isinstance(data_b64, str):
+        raise ValueError("packed sample support data must be a base64 string")
+    if not isinstance(raw_shape, list) or len(raw_shape) != 2:
+        raise ValueError("packed sample support shape must be [tokens, top_k]")
+    if any(not isinstance(dim, int) or isinstance(dim, bool) for dim in raw_shape):
+        raise ValueError("packed sample support shape must contain integers")
+    if any(dim < 0 for dim in raw_shape):
+        raise ValueError("packed sample support shape must be non-negative")
+    try:
+        data = pybase64.b64decode_as_bytearray(data_b64, validate=True) if data_b64 else bytearray()
+    except ValueError as exc:
+        raise ValueError("packed sample support data must be valid base64") from exc
+    expected_nbytes = raw_shape[0] * raw_shape[1] * _DTYPE.itemsize
+    if len(data) != expected_nbytes:
+        raise ValueError(
+            f"packed sample support byte-size mismatch: data={len(data)} bytes, expected={expected_nbytes}"
+        )
+    dense = np.frombuffer(data, dtype=_DTYPE).reshape(raw_shape)
+    return _validated_dense_support(dense)
diff --git a/skyrl/backends/skyrl_train/inference_servers/setup.py b/skyrl/backends/skyrl_train/inference_servers/setup.py
index 437d4804f3..84d481d596 100644
--- a/skyrl/backends/skyrl_train/inference_servers/setup.py
+++ b/skyrl/backends/skyrl_train/inference_servers/setup.py
@@ -293,6 +293,7 @@ def build_new_inference_client(
         server_urls=server_setup.server_urls,
         model_name=ie_cfg.served_model_name or cfg.trainer.policy.model.path,
         enable_return_routed_experts=ie_cfg.enable_return_routed_experts,
+        enable_return_sample_support_set=ie_cfg.enable_return_sample_support_set,
         uses_lora_weight_sync=_uses_lora_weight_sync(cfg),
         data_parallel_size=ie_cfg.data_parallel_size,
         tokenizer=tokenizer,
diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py
index 4772096066..0e9e4a1040 100644
--- a/skyrl/backends/skyrl_train/inference_servers/utils.py
+++ b/skyrl/backends/skyrl_train/inference_servers/utils.py
@@ -84,6 +84,7 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
     args: Namespace = parser.parse_args(args=[])
 
     ie_cfg = cfg.generator.inference_engine
+    sample_support_top_k = cfg.generator.sampling_params.top_k
     overrides = dict(
         model=cfg.trainer.policy.model.path,
         tensor_parallel_size=ie_cfg.tensor_parallel_size,
@@ -120,6 +121,9 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
         # Overridable via generator.inference_engine.engine_init_kwargs.trust_remote_code below.
         trust_remote_code=True,
     )
+    if ie_cfg.enable_return_sample_support_set:
+        overrides["max_logprobs"] = sample_support_top_k
+        overrides["logprobs_mode"] = "processed_logprobs"
     for key, value in overrides.items():
         setattr(args, key, value)
 
diff --git a/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py b/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py
index 7c35bd0efd..35c7499dfc 100644
--- a/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py
+++ b/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py
@@ -10,6 +10,7 @@
 from typing import List, Optional, Tuple
 
 import httpx
+import numpy as np
 import orjson
 import uvicorn
 import vllm.envs as envs
@@ -23,6 +24,7 @@
     init_app_state,
 )
 from vllm.inputs import TokensPrompt
+from vllm.logprobs import FlatLogprobs
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import SamplingParams as VLLMSamplingParams
 from vllm.usage.usage_lib import UsageContext
@@ -37,9 +39,13 @@
 from skyrl.backends.skyrl_train.inference_servers.generate_wire import (
     CLAMPED_LOGPROB,
     build_logprobs_content,
+    clamp_sampled_logprobs,
     pack_routed_experts,
 )
 from skyrl.backends.skyrl_train.inference_servers.protocols import ServerActorProtocol
+from skyrl.backends.skyrl_train.inference_servers.sample_support_set_wire import (
+    encode_sample_support_set,
+)
 from skyrl.env_vars import (
     SKYRL_HTTP_CONNECTION_LIMIT,
     SKYRL_VLLM_DP_PORT_OFFSET,
@@ -49,6 +55,60 @@
 logger = logging.getLogger(__name__)
 
 
+def _sample_support_from_flat_logprobs(
+    logprobs: FlatLogprobs,
+    top_k: int,
+) -> tuple[list[dict[str, float]], np.ndarray]:
+    """Extract sampled scores and post-filter support from vLLM's flat rows."""
+    # vLLM emits [sampled token, top-1, ..., top-k] for every generated token.
+    row_width = top_k + 1
+    token_ids = np.asarray(logprobs.token_ids, dtype=np.int32).reshape(-1, row_width)
+    processed_logprobs = np.asarray(logprobs.logprobs).reshape(-1, row_width)
+    support_ids = np.where(np.isneginf(processed_logprobs[:, 1:]), np.int32(-1), token_ids[:, 1:])
+    sampled_logprobs, num_clamped = clamp_sampled_logprobs(processed_logprobs[:, 0])
+    if num_clamped:
+        logger.warning(
+            "sample-support: clamped %d/%d non-finite sampled logprob(s) to %s",
+            num_clamped,
+            len(sampled_logprobs),
+            CLAMPED_LOGPROB,
+        )
+
+    # Repair rows whose sampled token is absent from the captured support.
+    #
+    # The support row is columns 1: (the top_k neighbors); column 0 is the token
+    # vLLM actually sampled. vLLM's approximate Triton top-k/top-p pivot (in
+    # processed_logprobs mode) can leave slightly more than top_k survivors, so the
+    # sampled token can rank just beyond top_k and be absent from cols 1:. When that
+    # happens the downstream sample-support invariant
+    # (assert_sampled_tokens_in_sample_support_sets / build_sample_support_replay)
+    # hard-crashes the whole run. To preserve it we overwrite the row's WEAKEST valid
+    # member (vLLM returns top-k descending, so the trailing valid slot is weakest)
+    # with the sampled id. Overwriting keeps width == top_k, inserts the sampled id
+    # exactly once (so the trainer's renorm denominator is not double-counted), and
+    # leaves any trailing -1 padding intact (the weakest valid col precedes the pad).
+    sampled = token_ids[:, 0]
+    valid = support_ids >= 0
+    present = np.any(support_ids == sampled[:, None], axis=1)
+    # Only repair rows that already have >=1 valid member (matches the downstream
+    # has_support condition); a fully -1 row is left untouched.
+    missing = (~present) & valid.any(axis=1)
+    if np.any(missing):
+        rows = np.flatnonzero(missing)
+        weakest_col = valid.sum(axis=1) - 1  # last valid (weakest) column per row
+        support_ids[rows, weakest_col[rows]] = sampled[rows]
+        logger.warning(
+            "sample-support repair: %d token(s) had the sampled id absent from top-%d "
+            "support; overwrote the weakest member to preserve the invariant (vLLM "
+            "approx top-k/top-p pivot artifact); example: sampled token %d at row %d",
+            rows.size,
+            top_k,
+            int(sampled[rows[0]]),
+            int(rows[0]),
+        )
+    return sampled_logprobs, support_ids
+
+
 class VLLMServerActor(ServerActorProtocol):
     """
     Ray actor that runs a vLLM OpenAI-compatible API server.
@@ -421,7 +481,10 @@ async def _skyrl_generate(request: Request):
             token_ids = body["token_ids"]
             sampling_params_dict = body.get("sampling_params", {})
             cache_salt = body.get("cache_salt")
-
+            capture_sample_support = body.get("return_sample_support", False)
+            if capture_sample_support:
+                sampling_params_dict["flat_logprobs"] = True
+                sampling_params_dict["logprobs"] = sampling_params_dict["top_k"]
             sampling_params = VLLMSamplingParams(**sampling_params_dict)
             # `cache_salt` salts vLLM's prefix cache; vLLM rejects an empty salt, so attach only when set.
             if cache_salt is not None:
@@ -442,7 +505,15 @@ async def _skyrl_generate(request: Request):
             finish_reason = resp.finish_reason
 
             logprobs = None
-            if resp.logprobs is not None:
+            sample_support = None
+            if capture_sample_support:
+                content, sample_support_ids = _sample_support_from_flat_logprobs(
+                    resp.logprobs,
+                    sampling_params_dict["top_k"],
+                )
+                logprobs = {"content": content}
+                sample_support = encode_sample_support_set(sample_support_ids)
+            elif resp.logprobs is not None:
                 content, num_clamped = build_logprobs_content(token_ids_out, resp.logprobs)
                 if num_clamped:
                     logger.warning(
@@ -462,6 +533,7 @@ async def _skyrl_generate(request: Request):
                         "finish_reason": finish_reason,
                         "logprobs": logprobs,
                         "routed_experts": routed_experts,
+                        "rollout_sample_support": sample_support,
                     }
                 ]
             }
diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py
index 1542adeff3..da17296ab3 100644
--- a/skyrl/backends/skyrl_train/training_batch.py
+++ b/skyrl/backends/skyrl_train/training_batch.py
@@ -485,6 +485,7 @@ class TrainingInput(TypedDict, total=False):
     rollout_logprobs: Optional[Float[torch.Tensor, "batch_size response_len"]]  # sampling policy; off-policy corr.
     rollout_expert_indices: Optional[Integer[torch.Tensor, "batch_size seq_len layer_num topk"]]  # MoE router replay
     router_padding_mask: Optional[Bool[torch.Tensor, "batch_size seq_len"]]  # True = no captured route (skip in replay)
+    sample_support_ids: Optional[Integer[torch.Tensor, "batch_size seq_len topk"]]  # sampler support, -1 = padding
     pixel_values: Optional[TensorList]  # list of `batch_size` [num_patches_i, dim] tensors
     image_grid_thw: Optional[TensorList]  # list of `batch_size` [num_images_i, 3] tensors
 
@@ -545,6 +546,15 @@ def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int)
             additional_dims = tensor.shape[1:]
             padding_tensor = torch.ones(pad_size, *additional_dims, dtype=torch.bool, device=tensor.device)
             new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
+        elif key == "sample_support_ids":
+            additional_dims = tensor.shape[1:]
+            padding_tensor = torch.full(
+                (pad_size, *additional_dims),
+                -1,
+                dtype=tensor.dtype,
+                device=tensor.device,
+            )
+            new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
         else:
             # Copy row 0 `pad_size` times. Loss masked so values don't affect the loss. Just need valid shape/dtype.
             assert tensor.shape[0] > 0, f"Cannot pad empty tensor field {key!r}"
diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py
index b3caec7c8b..a35bcc543a 100644
--- a/skyrl/backends/skyrl_train/utils/routed_experts.py
+++ b/skyrl/backends/skyrl_train/utils/routed_experts.py
@@ -1,11 +1,68 @@
+from collections.abc import Sequence
 from typing import TypeAlias
 
 import numpy as np
 
+from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import (
+    TokenMetadataTrace,
+)
+
 RoutedExpertIndices: TypeAlias = np.ndarray
 ROUTED_EXPERT_DTYPES = frozenset({np.dtype(np.uint8), np.dtype(np.int16), np.dtype(np.int32)})
 
 
+class RoutedExpertTrace:
+    """Accumulate routed experts across incremental generation calls."""
+
+    def __init__(self) -> None:
+        self._metadata = TokenMetadataTrace()
+        self._schema: tuple[int, int, np.dtype] | None = None
+
+    @property
+    def prompt_start(self) -> int:
+        return self._metadata.num_rows
+
+    def record_generation(
+        self,
+        *,
+        prompt_token_count: int,
+        generated_token_count: int,
+        routed_experts: RoutedExpertIndices,
+    ) -> None:
+        if prompt_token_count < self.prompt_start:
+            raise ValueError("routed-expert prompt start exceeds prompt length")
+        if generated_token_count < 1:
+            raise ValueError("routed-expert generation must produce at least one token")
+
+        expected_rows = prompt_token_count - self.prompt_start + generated_token_count - 1
+        compact = compact_routed_expert_indices(routed_experts)
+        if self._schema is None:
+            self._schema = (*compact.shape[1:], compact.dtype)
+        self._metadata.append(compact, expected_rows=expected_rows)
+
+    def finalize(self, *, token_count: int, loss_mask: Sequence[int]) -> RoutedExpertIndices:
+        if len(loss_mask) != token_count:
+            raise ValueError(f"loss mask has {len(loss_mask)} entries, expected {token_count}")
+        if self.prompt_start > token_count:
+            raise ValueError(f"routed-expert trace has {self.prompt_start} rows for {token_count} tokens")
+
+        if any(loss_mask[self.prompt_start + 1 : token_count]):
+            for source_index in range(self.prompt_start, token_count - 1):
+                if loss_mask[source_index + 1] != 0:
+                    raise ValueError(f"missing routed-expert row for loss-active target at token {source_index + 1}")
+
+        padding_count = token_count - self.prompt_start
+        if padding_count:
+            if self._schema is None:
+                raise ValueError("cannot pad routed-expert trace before any routes are captured")
+            num_layers, topk, dtype = self._schema
+            padding_row = np.arange(topk, dtype=dtype)
+            padding = np.broadcast_to(padding_row, (padding_count, num_layers, topk)).copy()
+            self._metadata.append(padding, expected_rows=padding_count)
+
+        return self._metadata.finalize(expected_rows=token_count)
+
+
 def compact_routed_expert_indices(routed_experts: RoutedExpertIndices) -> RoutedExpertIndices:
     """Validate and compact a routed-expert array to the canonical integer dtype."""
     if not isinstance(routed_experts, np.ndarray):
diff --git a/skyrl/backends/skyrl_train/utils/sample_support_replay.py b/skyrl/backends/skyrl_train/utils/sample_support_replay.py
new file mode 100644
index 0000000000..22c211213e
--- /dev/null
+++ b/skyrl/backends/skyrl_train/utils/sample_support_replay.py
@@ -0,0 +1,422 @@
+"""Support-conditioned logprobs for bounded sampler replay."""
+
+from dataclasses import dataclass
+
+import torch
+
+from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import (
+    TokenMetadataLayout,
+    align_token_metadata,
+    scatter_packed_token_values_to_batch,
+)
+from skyrl.backends.skyrl_train.utils.torch_utils import logprobs_from_logits
+
+
+@dataclass(frozen=True)
+class SampleSupportScores:
+    """Support-conditioned scores and the rows backed by recorded support."""
+
+    logprobs: torch.Tensor
+    entropy: torch.Tensor | None
+    valid_mask: torch.Tensor
+
+
+def _selected_hidden_projection(
+    hidden: torch.Tensor,
+    token_ids: torch.Tensor,
+    local_mask: torch.Tensor,
+    lm_head_weight: torch.Tensor,
+    temperature: float,
+    chunk_size: int | None,
+    invalid_value: float,
+) -> torch.Tensor:
+    """Project selected candidate pairs without materializing vocabulary logits."""
+    num_rows, width = token_ids.shape
+    row_ids = torch.arange(num_rows, device=hidden.device).unsqueeze(1).expand(-1, width).reshape(-1)
+    flat_token_ids = token_ids.reshape(-1)
+    flat_mask = local_mask.reshape(-1)
+    output = torch.empty(flat_token_ids.shape, dtype=torch.float32, device=hidden.device)
+    # Bound the temporary [candidate pairs, hidden] projection for wide supports.
+    pair_chunk_size = flat_token_ids.numel() if chunk_size is None else chunk_size
+    for start in range(0, flat_token_ids.numel(), pair_chunk_size):
+        end = min(start + pair_chunk_size, flat_token_ids.numel())
+        selected_hidden = hidden.index_select(0, row_ids[start:end]).to(lm_head_weight.dtype)
+        selected_weight = lm_head_weight.index_select(0, flat_token_ids[start:end])
+        projected = (selected_hidden * selected_weight).sum(dim=-1) / temperature
+        output[start:end] = torch.where(
+            flat_mask[start:end],
+            projected.to(torch.float32),
+            invalid_value,
+        )
+    return output.reshape(num_rows, width)
+
+
+def sample_support_scores(
+    logits_or_hidden: torch.Tensor,
+    sampled_ids: torch.Tensor,
+    support_ids: torch.Tensor,
+    *,
+    vocab_start_index: int,
+    vocab_end_index: int,
+    tp_group: torch.distributed.ProcessGroup | None,
+    lm_head_weight: torch.Tensor | None = None,
+    temperature: float = 1.0,
+    chunk_size: int | None = None,
+    compute_entropy: bool,
+    entropy_requires_grad: bool,
+) -> SampleSupportScores:
+    """Compute sampled-token logprobs and optional entropy over recorded support."""
+    if logits_or_hidden.shape[:-1] != sampled_ids.shape or support_ids.shape[:-1] != sampled_ids.shape:
+        raise ValueError(
+            "logits, sampled_ids, and support_ids must have matching prefix shapes, got "
+            f"{logits_or_hidden.shape[:-1]}, {sampled_ids.shape}, and {support_ids.shape[:-1]}"
+        )
+    if support_ids.dtype != torch.int32:
+        raise ValueError(f"sample support must use int32 vocab ids, got {support_ids.dtype}")
+    if temperature <= 0:
+        raise ValueError("temperature must be positive")
+    if entropy_requires_grad and not compute_entropy:
+        raise ValueError("entropy gradients require compute_entropy=True")
+
+    flat_source = logits_or_hidden.reshape(-1, logits_or_hidden.shape[-1])
+    flat_sampled = sampled_ids.reshape(-1).long()
+    flat_support = support_ids.reshape(-1, support_ids.shape[-1]).long()
+    valid_members = flat_support >= 0
+    valid_rows = valid_members.any(dim=-1)
+    local_members = valid_members & (flat_support >= vocab_start_index) & (flat_support < vocab_end_index)
+    local_support_ids = (flat_support - vocab_start_index).clamp(0, vocab_end_index - vocab_start_index - 1)
+    local_sample_mask = (flat_sampled >= vocab_start_index) & (flat_sampled < vocab_end_index)
+    local_sample_ids = (flat_sampled - vocab_start_index).clamp(0, vocab_end_index - vocab_start_index - 1)
+
+    compute_dtype = (
+        torch.float32 if logits_or_hidden.dtype in (torch.float16, torch.bfloat16) else logits_or_hidden.dtype
+    )
+    if lm_head_weight is None:
+        local_values = flat_source.gather(1, local_support_ids).to(compute_dtype)
+        local_values = torch.where(local_members, local_values, float("-inf"))
+        local_sampled = flat_source.gather(1, local_sample_ids.unsqueeze(1)).squeeze(1).to(compute_dtype)
+        local_sampled = torch.where(local_sample_mask, local_sampled, 0.0)
+    else:
+        if lm_head_weight.shape[0] != vocab_end_index - vocab_start_index:
+            raise ValueError("lm_head_weight rows do not match the configured vocabulary shard")
+        local_values = _selected_hidden_projection(
+            flat_source,
+            local_support_ids,
+            local_members,
+            lm_head_weight,
+            temperature,
+            chunk_size,
+            float("-inf"),
+        )
+        local_sampled = _selected_hidden_projection(
+            flat_source,
+            local_sample_ids.unsqueeze(1),
+            local_sample_mask.unsqueeze(1),
+            lm_head_weight,
+            temperature,
+            chunk_size,
+            0.0,
+        ).squeeze(1)
+
+    local_max = local_values.detach().amax(dim=-1)
+    global_max = local_max.clone()
+    if tp_group is not None and torch.distributed.get_world_size(tp_group) > 1:
+        torch.distributed.all_reduce(global_max, op=torch.distributed.ReduceOp.MAX, group=tp_group)
+    safe_max = torch.where(valid_rows, global_max, 0.0)
+
+    local_exp = torch.where(local_members, (local_values - safe_max.unsqueeze(1)).exp(), 0.0)
+    local_sum = local_exp.sum(dim=-1)
+    local_stats = [local_sum, local_sampled]
+    if compute_entropy:
+        entropy_values = local_values if entropy_requires_grad else local_values.detach()
+        entropy_exp = local_exp if entropy_requires_grad else local_exp.detach()
+        shifted_values = torch.where(local_members, entropy_values - safe_max.unsqueeze(1), 0.0)
+        local_stats.append((entropy_exp * shifted_values).sum(dim=-1))
+    # Numerator, denominator, and optional entropy statistic share one TP SUM collective.
+    local_stats = torch.stack(local_stats)
+    global_stats = local_stats.detach().clone()
+    if tp_group is not None and torch.distributed.get_world_size(tp_group) > 1:
+        torch.distributed.all_reduce(global_stats, op=torch.distributed.ReduceOp.SUM, group=tp_group)
+    global_stats = global_stats + local_stats - local_stats.detach()
+    denominator, sampled_score = global_stats[:2]
+    logprobs = sampled_score - safe_max - torch.where(valid_rows, denominator, 1.0).log()
+    logprobs = torch.where(valid_rows, logprobs, 0.0)
+    entropy = None
+    if compute_entropy:
+        entropy_denominator = denominator if entropy_requires_grad else denominator.detach()
+        shifted_score_sum = global_stats[2] if entropy_requires_grad else global_stats[2].detach()
+        safe_denominator = torch.where(valid_rows, entropy_denominator, 1.0)
+        entropy = safe_denominator.log() - shifted_score_sum / safe_denominator
+        entropy = torch.where(valid_rows, entropy, 0.0).reshape(sampled_ids.shape)
+    return SampleSupportScores(
+        logprobs=logprobs.reshape(sampled_ids.shape),
+        entropy=entropy,
+        valid_mask=valid_rows.reshape(sampled_ids.shape),
+    )
+
+
+def synthetic_eos_logprobs(
+    logits_or_hidden: torch.Tensor,
+    sampled_ids: torch.Tensor,
+    synthetic_eos_mask: torch.Tensor,
+    *,
+    vocab_start_index: int,
+    vocab_end_index: int,
+    tp_group: torch.distributed.ProcessGroup | None,
+    inference_only: bool,
+    lm_head_weight: torch.Tensor | None = None,
+    temperature: float = 1.0,
+    chunk_size: int | None = None,
+    fused_backend: str = "torch",
+    metadata_layout: TokenMetadataLayout | None = None,
+    trajectory_ids: torch.Tensor | None = None,
+    num_trajectories: int | None = None,
+) -> torch.Tensor:
+    """Compute ordinary logprobs for EOS tokens appended after vLLM generation."""
+    if synthetic_eos_mask.shape != sampled_ids.shape:
+        raise ValueError("synthetic_eos_mask and sampled_ids must have matching shapes")
+
+    if trajectory_ids is not None:
+        if trajectory_ids.shape != synthetic_eos_mask.shape:
+            raise ValueError("trajectory_ids and synthetic_eos_mask must have matching shapes")
+        if num_trajectories is None or num_trajectories <= 0:
+            raise ValueError("num_trajectories must be positive when trajectory_ids are provided")
+        flat_trajectory_ids = trajectory_ids.reshape(-1).to(torch.long)
+        capacity = num_trajectories
+    elif metadata_layout is not None and metadata_layout.padded_sequence_lengths is not None:
+        if synthetic_eos_mask.shape[0] != 1:
+            raise ValueError("Packed synthetic EOS metadata must have a singleton batch dimension")
+        if metadata_layout.cu_seqlens_padded is None:
+            raise ValueError("Packed synthetic EOS fallback requires padded sequence boundaries")
+        if any(length <= 0 for length in metadata_layout.padded_sequence_lengths):
+            raise ValueError("Synthetic EOS fallback requires non-empty trajectory segments")
+        expected_tokens = metadata_layout.aligned_sequence_length // metadata_layout.context_parallel_size
+        if expected_tokens != synthetic_eos_mask.numel():
+            raise ValueError("Synthetic EOS layout does not match the model token layout")
+        lengths = (
+            metadata_layout.cu_seqlens_padded.to(
+                device=synthetic_eos_mask.device,
+                dtype=torch.long,
+            ).diff()
+            // metadata_layout.context_parallel_size
+        )
+        capacity = lengths.shape[0]
+        flat_trajectory_ids = torch.repeat_interleave(
+            torch.arange(capacity, device=lengths.device),
+            lengths,
+            output_size=synthetic_eos_mask.numel(),
+        )
+    else:
+        if synthetic_eos_mask.shape[0] == 0 or synthetic_eos_mask.shape[1] == 0:
+            raise ValueError("Synthetic EOS fallback requires non-empty trajectory segments")
+        lengths = torch.full(
+            (synthetic_eos_mask.shape[0],),
+            synthetic_eos_mask.shape[1],
+            dtype=torch.long,
+            device=synthetic_eos_mask.device,
+        )
+        capacity = lengths.shape[0]
+        flat_trajectory_ids = torch.repeat_interleave(
+            torch.arange(capacity, device=lengths.device),
+            lengths,
+            output_size=synthetic_eos_mask.numel(),
+        )
+
+    # Preprocessing permits at most one unsupported loss-bearing EOS per
+    # trajectory. Select one fixed slot for every trajectory so TP collectives
+    # never depend on the number of EOS fallbacks in this microbatch.
+    token_indices = torch.arange(synthetic_eos_mask.numel(), device=synthetic_eos_mask.device)
+    sentinel = synthetic_eos_mask.numel()
+    valid_trajectory = (flat_trajectory_ids >= 0) & (flat_trajectory_ids < capacity)
+    candidate_indices = torch.where(synthetic_eos_mask.reshape(-1) & valid_trajectory, token_indices, sentinel)
+    selected_indices = torch.full(
+        (capacity,),
+        sentinel,
+        dtype=torch.long,
+        device=synthetic_eos_mask.device,
+    ).scatter_reduce(
+        0,
+        flat_trajectory_ids.clamp(0, capacity - 1),
+        candidate_indices,
+        reduce="amin",
+        include_self=True,
+    )
+    has_selection = selected_indices != sentinel
+    selected_indices = torch.where(has_selection, selected_indices, 0)
+
+    flat_source = logits_or_hidden.reshape(-1, logits_or_hidden.shape[-1])
+    flat_targets = sampled_ids.reshape(-1)
+    selected_source = flat_source.index_select(0, selected_indices)
+    selected_targets = flat_targets.index_select(0, selected_indices)
+    source_is_full_vocab_logits = lm_head_weight is None and tp_group is None
+    source_is_tp_sharded_logits = lm_head_weight is None and tp_group is not None
+    if source_is_full_vocab_logits:
+        selected = logprobs_from_logits(selected_source, selected_targets, inplace_backward=False)
+    elif source_is_tp_sharded_logits:
+        from skyrl.backends.skyrl_train.distributed.megatron.model_utils import (
+            DistributedLogprob,
+        )
+
+        selected = DistributedLogprob.apply(
+            selected_source.unsqueeze(0),
+            selected_targets.unsqueeze(0),
+            vocab_start_index,
+            vocab_end_index,
+            tp_group,
+            inference_only,
+        ).squeeze(0)
+    else:
+        from skyrl.backends.skyrl_train.distributed.megatron.model_utils import (
+            _fused_lm_head_logprob_apply,
+        )
+
+        if temperature != 1.0:
+            lm_head_weight = lm_head_weight / temperature
+        selected_chunk_size = (
+            selected_source.shape[0] if chunk_size is None else min(chunk_size, selected_source.shape[0])
+        )
+        selected = _fused_lm_head_logprob_apply(
+            fused_backend,
+            selected_source.unsqueeze(0),
+            lm_head_weight,
+            selected_targets.unsqueeze(0),
+            vocab_start_index,
+            vocab_end_index,
+            selected_chunk_size,
+            tp_group,
+            inference_only,
+        ).squeeze(0)
+    selected = torch.where(has_selection, selected, 0.0).to(torch.float32)
+    output = torch.zeros(sampled_ids.numel(), dtype=torch.float32, device=logits_or_hidden.device)
+    output = output.scatter_add(0, selected_indices, selected)
+    return output.reshape(sampled_ids.shape)
+
+
+def aligned_sample_support_scores(
+    logits_or_hidden: torch.Tensor,
+    sampled_ids: torch.Tensor,
+    support_ids: torch.Tensor,
+    loss_mask: torch.Tensor,
+    *,
+    vocab_start_index: int,
+    vocab_end_index: int,
+    tp_group: torch.distributed.ProcessGroup | None,
+    inference_only: bool,
+    compute_entropy: bool,
+    entropy_requires_grad: bool,
+    lm_head_weight: torch.Tensor | None = None,
+    temperature: float = 1.0,
+    chunk_size: int | None = None,
+    fused_backend: str = "torch",
+    metadata_layout: TokenMetadataLayout | None = None,
+    trajectory_ids: torch.Tensor | None = None,
+    num_trajectories: int | None = None,
+) -> SampleSupportScores:
+    """Apply bounded replay and the synthetic-EOS exception to aligned tokens."""
+    scores = sample_support_scores(
+        logits_or_hidden,
+        sampled_ids,
+        support_ids,
+        vocab_start_index=vocab_start_index,
+        vocab_end_index=vocab_end_index,
+        tp_group=tp_group,
+        compute_entropy=compute_entropy,
+        entropy_requires_grad=entropy_requires_grad,
+        lm_head_weight=lm_head_weight,
+        temperature=temperature if lm_head_weight is not None else 1.0,
+        chunk_size=chunk_size,
+    )
+    synthetic_eos_mask = loss_mask & ~scores.valid_mask
+    eos_logprobs = synthetic_eos_logprobs(
+        logits_or_hidden,
+        sampled_ids,
+        synthetic_eos_mask,
+        vocab_start_index=vocab_start_index,
+        vocab_end_index=vocab_end_index,
+        tp_group=tp_group,
+        inference_only=inference_only,
+        lm_head_weight=lm_head_weight,
+        temperature=temperature if lm_head_weight is not None else 1.0,
+        chunk_size=chunk_size,
+        fused_backend=fused_backend,
+        metadata_layout=metadata_layout,
+        trajectory_ids=trajectory_ids,
+        num_trajectories=num_trajectories,
+    )
+    return SampleSupportScores(
+        logprobs=torch.where(synthetic_eos_mask, eos_logprobs, scores.logprobs),
+        entropy=scores.entropy,
+        valid_mask=scores.valid_mask,
+    )
+
+
+def compute_sample_support_scores(
+    logits_or_hidden: torch.Tensor,
+    sequences: torch.Tensor,
+    loss_mask: torch.Tensor | None,
+    sample_support_ids: torch.Tensor | None,
+    num_actions: int,
+    *,
+    packed: bool,
+    metadata_layout: TokenMetadataLayout | None,
+    vocab_start_index: int,
+    vocab_end_index: int,
+    tp_group: torch.distributed.ProcessGroup | None,
+    inference_only: bool,
+    lm_head_weight: torch.Tensor | None,
+    temperature: float,
+    chunk_size: int | None,
+    fused_backend: str,
+    compute_entropy: bool,
+    entropy_requires_grad: bool,
+) -> SampleSupportScores:
+    """Compute dense support-conditioned scores in canonical trainer layout."""
+    if sample_support_ids is None:
+        raise ValueError("sample-support replay is enabled but the microbatch has no recorded support")
+    if loss_mask is None:
+        raise ValueError("sample-support replay requires the response loss mask")
+
+    target_loss_mask = torch.zeros_like(sequences, dtype=torch.bool)
+    target_loss_mask[:, -num_actions:] = loss_mask.to(torch.bool)
+    if packed:
+        if metadata_layout is None:
+            raise ValueError("Packed sample-support replay requires the shared token metadata layout")
+        aligned_sampled_ids = align_token_metadata(sequences, metadata_layout, 0, next_token=True)
+        aligned_support_ids = align_token_metadata(sample_support_ids, metadata_layout, -1, next_token=True)
+        aligned_loss_mask = align_token_metadata(target_loss_mask, metadata_layout, False, next_token=True)
+    else:
+        aligned_sampled_ids = sequences[:, 1:]
+        aligned_support_ids = sample_support_ids[:, 1:]
+        aligned_loss_mask = target_loss_mask[:, 1:]
+
+    aligned_source = logits_or_hidden if packed else logits_or_hidden[:, :-1]
+    scores = aligned_sample_support_scores(
+        aligned_source,
+        aligned_sampled_ids,
+        aligned_support_ids,
+        aligned_loss_mask,
+        vocab_start_index=vocab_start_index,
+        vocab_end_index=vocab_end_index,
+        tp_group=tp_group,
+        inference_only=inference_only,
+        compute_entropy=compute_entropy,
+        entropy_requires_grad=entropy_requires_grad,
+        lm_head_weight=lm_head_weight,
+        temperature=temperature,
+        chunk_size=chunk_size,
+        fused_backend=fused_backend,
+        metadata_layout=metadata_layout if packed else None,
+    )
+
+    if packed:
+        assert metadata_layout is not None
+        return SampleSupportScores(
+            logprobs=scatter_packed_token_values_to_batch(scores.logprobs, metadata_layout, 0),
+            entropy=(
+                scatter_packed_token_values_to_batch(scores.entropy, metadata_layout, 0)
+                if scores.entropy is not None
+                else None
+            ),
+            valid_mask=scatter_packed_token_values_to_batch(scores.valid_mask, metadata_layout, False),
+        )
+    return scores
diff --git a/skyrl/backends/skyrl_train/utils/torch_utils.py b/skyrl/backends/skyrl_train/utils/torch_utils.py
index 3da9dc4077..d47af2e782 100644
--- a/skyrl/backends/skyrl_train/utils/torch_utils.py
+++ b/skyrl/backends/skyrl_train/utils/torch_utils.py
@@ -133,7 +133,7 @@ def logprobs_from_logits(
     Returns:
         Tensor: Log-probabilities of the target labels, shape logits.shape[:-1].
     """
-    if FLASH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE:
+    if FLASH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE and logits.is_cuda:
         batch_dim = logits.shape[:-1]
         last_dim = logits.shape[-1]
         logits = logits.reshape(-1, last_dim)
diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py
index 91a4a25047..6c5a379877 100644
--- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py
+++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py
@@ -30,6 +30,7 @@
 )
 from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import is_fp8_enabled
 from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import (
+    TokenMetadataLayout,
     build_token_metadata_layout,
 )
 from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits
@@ -51,6 +52,9 @@
     setup_per_microbatch_replay_backward,
     setup_per_microbatch_replay_forward,
 )
+from skyrl.backends.skyrl_train.utils.sample_support_replay import (
+    compute_sample_support_scores,
+)
 from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean
 from skyrl.backends.skyrl_train.workers.worker_utils import (
     compute_minibatch_rollout_logprob_diff_metrics,
@@ -264,7 +268,7 @@ def forward(
             self._assert_vlm_supported()
         forward_backward_func = get_forward_backward_func()
 
-        def collection_func(logits, data):
+        def collection_func(logits, *, data, metadata_layout: TokenMetadataLayout | None):
             sequences = data["sequences"]
             packed_seq_params = data.get("packed_seq_params")
             packed_targets = data.get("packed_targets")
@@ -285,7 +289,28 @@ def collection_func(logits, data):
             if temperature != 1.0 and not fused_lm_head:
                 logits.div_(temperature)
 
-            if fused_lm_head and packed_seq_params is not None and packed_targets is not None:
+            shard_vocab_size = lm_head_weight.shape[0] if fused_lm_head else logits.shape[-1]
+            if self.cfg.algorithm.enable_sample_support_replay:
+                token_logprobs = compute_sample_support_scores(
+                    logits,
+                    sequences,
+                    data.get("loss_mask"),
+                    data.get("sample_support_ids"),
+                    data["num_actions"],
+                    packed=packed_seq_params is not None,
+                    metadata_layout=metadata_layout,
+                    vocab_start_index=tp_rank * shard_vocab_size,
+                    vocab_end_index=(tp_rank + 1) * shard_vocab_size,
+                    tp_group=tp_grp,
+                    lm_head_weight=lm_head_weight if fused_lm_head else None,
+                    temperature=temperature,
+                    inference_only=True,
+                    chunk_size=self.cfg.logprobs_chunk_size,
+                    fused_backend=self._fused_lm_head_backend,
+                    compute_entropy=False,
+                    entropy_requires_grad=False,
+                ).logprobs
+            elif fused_lm_head and packed_seq_params is not None and packed_targets is not None:
                 token_logprobs = from_parallel_hidden_to_logprobs_packed_sequences(
                     logits,  # decoder hidden states [1, T, H]
                     lm_head_weight,
@@ -355,6 +380,7 @@ def forward_step(batch_iter, model):
             fp8_enabled = is_fp8_enabled(getattr(model_config, "fp8", None))
             rollout_expert_indices = batch.pop("rollout_expert_indices", None)
             router_padding_mask = batch.pop("router_padding_mask", None)
+            sample_support_ids = batch.get("sample_support_ids")
 
             sequences = batch["sequences"]
             attention_mask = batch["attention_mask"].to(bool)
@@ -362,6 +388,12 @@ def forward_step(batch_iter, model):
             sub_seq_lengths_field = batch.get("sub_seq_lengths")
             sub_seq_lengths = [t.tolist() for t in sub_seq_lengths_field] if sub_seq_lengths_field is not None else None
             batch["sub_seq_lengths_list"] = sub_seq_lengths
+            if (
+                sample_support_ids is not None
+                and sub_seq_lengths is not None
+                and any(len(row_lengths) > 1 for row_lengths in sub_seq_lengths)
+            ):
+                raise ValueError("sample-support replay does not support controller-packed multi-subsequence rows")
 
             vlm_inputs = {}
             if batch.get("pixel_values") is not None and mpu.get_pipeline_model_parallel_rank() == 0:
@@ -398,7 +430,7 @@ def forward_step(batch_iter, model):
                     new_position_ids = None
 
             metadata_layout = None
-            if rollout_expert_indices is not None:
+            if rollout_expert_indices is not None or (sample_support_ids is not None and packed_seq_params is not None):
                 metadata_layout = build_token_metadata_layout(
                     attention_mask,
                     attention_mask.device,
@@ -456,7 +488,7 @@ def forward_step(batch_iter, model):
                     post_process=mpu.is_pipeline_last_stage(ignore_virtual=True),
                 )
 
-            return outputs, partial(collection_func, data=batch)
+            return outputs, partial(collection_func, data=batch, metadata_layout=metadata_layout)
 
         batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.actor_module))
 
@@ -562,7 +594,7 @@ def forward_backward_mini_batch(
             # NOTE: users can provide a custom loss config class, so we need to use the same class after applying overrides
             loss_config = type(loss_config).from_dict_config(new_loss_config)
 
-        def loss_func(logits, data):
+        def loss_func(logits, *, data, metadata_layout: TokenMetadataLayout | None):
             sequences = data["sequences"]
             packed_seq_params = data.get("packed_seq_params")
             packed_targets = data.get("packed_targets")
@@ -590,11 +622,6 @@ def loss_func(logits, data):
             # grad are never materialized.
             fused_lm_head = self._fused_lm_head and data.get("lm_head_weight") is not None
             lm_head_weight = data.get("lm_head_weight")
-            if fused_lm_head and loss_config.use_entropy_loss:
-                raise NotImplementedError(
-                    "fused_lm_head_logprob does not support use_entropy_loss=True "
-                    "(the fused entropy is a no-grad metric)."
-                )
             if fused_lm_head:
                 _v_local = int(lm_head_weight.shape[0])
                 fused_vocab_start, fused_vocab_end = tp_rank * _v_local, (tp_rank + 1) * _v_local
@@ -603,7 +630,34 @@ def loss_func(logits, data):
             if temperature != 1.0 and not fused_lm_head:
                 logits.div_(temperature)
 
-            if fused_lm_head and packed_seq_params is not None and packed_targets is not None:
+            shard_vocab_size = lm_head_weight.shape[0] if fused_lm_head else logits.shape[-1]
+            support_entropy = None
+            support_entropy_mask = None
+            if self.cfg.algorithm.enable_sample_support_replay:
+                compute_support_entropy = resolved_loss_name != "cross_entropy"
+                support_scores = compute_sample_support_scores(
+                    logits,
+                    sequences,
+                    loss_mask,
+                    data.get("sample_support_ids"),
+                    num_actions,
+                    packed=packed_seq_params is not None,
+                    metadata_layout=metadata_layout,
+                    vocab_start_index=tp_rank * shard_vocab_size,
+                    vocab_end_index=(tp_rank + 1) * shard_vocab_size,
+                    tp_group=tp_grp,
+                    lm_head_weight=lm_head_weight if fused_lm_head else None,
+                    temperature=temperature,
+                    inference_only=forward_only,
+                    chunk_size=self.cfg.logprobs_chunk_size,
+                    fused_backend=self._fused_lm_head_backend,
+                    compute_entropy=compute_support_entropy,
+                    entropy_requires_grad=compute_support_entropy and loss_config.use_entropy_loss,
+                )
+                token_logprobs = support_scores.logprobs
+                support_entropy = support_scores.entropy
+                support_entropy_mask = support_scores.valid_mask
+            elif fused_lm_head and packed_seq_params is not None and packed_targets is not None:
                 token_logprobs = from_parallel_hidden_to_logprobs_packed_sequences(
                     logits,  # decoder hidden states [1, T, H]
                     lm_head_weight,
@@ -808,7 +862,16 @@ def loss_func(logits, data):
 
             # RL path: add optional KL/entropy terms
             with torch.set_grad_enabled(loss_config.use_entropy_loss):
-                if fused_lm_head and packed_seq_params is not None and packed_targets is not None:
+                if support_entropy is not None and support_entropy_mask is not None:
+                    action_entropy = support_entropy[:, -num_actions:]
+                    action_entropy_mask = support_entropy_mask[:, -num_actions:] & loss_mask.to(torch.bool)
+                    entropy = masked_mean(action_entropy, action_entropy_mask)
+                    entropy_for_loss = entropy
+                elif fused_lm_head and loss_config.use_entropy_loss:
+                    raise NotImplementedError(
+                        "Differentiable full-vocabulary entropy is not supported with the fused LM head"
+                    )
+                elif fused_lm_head and packed_seq_params is not None and packed_targets is not None:
                     entropy, entropy_for_loss = from_parallel_hidden_to_entropy_packed_sequences(
                         logits,  # decoder hidden states [1, T, H]
                         lm_head_weight,
@@ -956,6 +1019,7 @@ def forward_step(batch_iter, model):
             fp8_enabled = is_fp8_enabled(getattr(model_config, "fp8", None))
             rollout_expert_indices = batch.pop("rollout_expert_indices", None)
             router_padding_mask = batch.pop("router_padding_mask", None)
+            sample_support_ids = batch.get("sample_support_ids")
 
             sequences = batch["sequences"]
             attention_mask = batch["attention_mask"].to(bool)
@@ -971,6 +1035,12 @@ def forward_step(batch_iter, model):
             sub_seq_lengths_field = batch.get("sub_seq_lengths")
             sub_seq_lengths = [t.tolist() for t in sub_seq_lengths_field] if sub_seq_lengths_field is not None else None
             batch["sub_seq_lengths_list"] = sub_seq_lengths
+            if (
+                sample_support_ids is not None
+                and sub_seq_lengths is not None
+                and any(len(row_lengths) > 1 for row_lengths in sub_seq_lengths)
+            ):
+                raise ValueError("sample-support replay does not support controller-packed multi-subsequence rows")
 
             vlm_inputs = {}
             if batch.get("pixel_values") is not None and mpu.get_pipeline_model_parallel_rank() == 0:
@@ -1020,7 +1090,7 @@ def forward_step(batch_iter, model):
             is_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=True)
 
             metadata_layout = None
-            if rollout_expert_indices is not None:
+            if rollout_expert_indices is not None or (sample_support_ids is not None and packed_seq_params is not None):
                 metadata_layout = build_token_metadata_layout(
                     attention_mask,
                     attention_mask.device,
@@ -1131,7 +1201,7 @@ def depad(tensor):
             if rollout_expert_indices is not None:
                 setup_per_microbatch_replay_backward()
 
-            return outputs, partial(loss_func, data=batch)
+            return outputs, partial(loss_func, data=batch, metadata_layout=metadata_layout)
 
         # batch should be a list of micro-batches
         batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.actor_module))
diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
index 65dc8ef328..7909384022 100644
--- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
+++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py
@@ -668,6 +668,10 @@ def _forward_logprobs(self, data: TrainingInputBatch) -> torch.Tensor:
                     "num_actions": micro.metadata["response_length"],
                     "rollout_expert_indices": (rollout_expert_indices if self.enable_router_replay else None),
                     "router_padding_mask": micro.get("router_padding_mask") if self.enable_router_replay else None,
+                    "sample_support_ids": (
+                        micro.get("sample_support_ids") if self.cfg.algorithm.enable_sample_support_replay else None
+                    ),
+                    "loss_mask": micro.get("loss_mask"),
                     "sub_seq_lengths": micro.get("sub_seq_lengths"),
                     **vlm_inputs,
                 }
@@ -790,6 +794,13 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d
                         dtype=value.dtype,
                         device=device,
                     )
+                elif key == "sample_support_ids":
+                    pad_tensor = torch.full(
+                        (pad_count, *value.shape[1:]),
+                        -1,
+                        dtype=value.dtype,
+                        device=device,
+                    )
                 elif key == "response_mask":
                     # response_mask should be zeros for padded samples
                     pad_tensor = torch.zeros((pad_count, *value.shape[1:]), dtype=value.dtype, device=device)
@@ -1057,6 +1068,9 @@ def forward(
                     "response_mask": experience.response_mask,
                     "rollout_expert_indices": rollout_expert_indices if self.enable_router_replay else None,
                     "router_padding_mask": experience.router_padding_mask if self.enable_router_replay else None,
+                    "sample_support_ids": (
+                        experience.sample_support_ids if self.cfg.algorithm.enable_sample_support_replay else None
+                    ),
                     "sub_seq_lengths": experience.sub_seq_lengths,
                     **vlm_inputs,
                 }
@@ -1183,6 +1197,9 @@ def forward_backward(
                     "response_mask": experience.response_mask,
                     "rollout_expert_indices": rollout_expert_indices if self.enable_router_replay else None,
                     "router_padding_mask": experience.router_padding_mask if self.enable_router_replay else None,
+                    "sample_support_ids": (
+                        experience.sample_support_ids if self.cfg.algorithm.enable_sample_support_replay else None
+                    ),
                     # used with global sequence packing (None when token-based batching is active)
                     "sub_seq_lengths": experience.sub_seq_lengths,
                     "is_padding_batch": (
diff --git a/skyrl/backends/skyrl_train/workers/model_wrapper.py b/skyrl/backends/skyrl_train/workers/model_wrapper.py
index 8235cddb73..810683a528 100644
--- a/skyrl/backends/skyrl_train/workers/model_wrapper.py
+++ b/skyrl/backends/skyrl_train/workers/model_wrapper.py
@@ -27,6 +27,9 @@
     ulysses_pad_and_slice_inputs,
 )
 from skyrl.backends.skyrl_train.training_batch import TensorList
+from skyrl.backends.skyrl_train.utils.sample_support_replay import (
+    aligned_sample_support_scores,
+)
 from skyrl.backends.skyrl_train.utils.torch_utils import (
     chunked_entropy_from_logits,
     logprobs_from_logits,
@@ -250,9 +253,28 @@ def forward(
         pixel_values: Optional[TensorList] = None,
         image_grid_thw: Optional[TensorList] = None,
         mm_token_type_ids: Optional[torch.Tensor] = None,
+        sample_support_ids: Optional[torch.Tensor] = None,
+        loss_mask: Optional[torch.Tensor] = None,
+        enable_sample_support_replay: bool = False,
     ) -> torch.Tensor:
         """Returns action log probs"""
         has_image_inputs = pixel_values is not None or image_grid_thw is not None
+        if enable_sample_support_replay and sample_support_ids is None:
+            raise ValueError("sample-support replay is enabled but the microbatch has no recorded support")
+        if enable_sample_support_replay and loss_mask is None:
+            raise ValueError("sample-support replay is enabled but the microbatch has no loss mask")
+        support_ids_fwd = sample_support_ids if enable_sample_support_replay else None
+        target_loss_mask_fwd = None
+        support_trajectory_ids_fwd = None
+        support_entropy = None
+        support_entropy_mask = None
+        if enable_sample_support_replay:
+            # The loss mask is response-only. Place it in the full token layout so it
+            # follows support IDs through the same unpadding, target shift, and SP slice.
+            target_loss_mask_fwd = torch.zeros_like(sequences, dtype=torch.bool)
+            target_loss_mask_fwd[:, sequences.shape[1] - loss_mask.shape[1] :] = loss_mask.to(torch.bool)
+            support_trajectory_ids_fwd = torch.arange(sequences.shape[0], device=sequences.device).unsqueeze(1)
+            support_trajectory_ids_fwd = support_trajectory_ids_fwd.expand_as(sequences)
         if self.is_vlm:
             # VLMs use model specific 3D positional IDs, meaning sequence packing can not be supported.
             # Sequence packing requires computing position IDs, but position IDs for VLMs are 3D and require
@@ -286,9 +308,24 @@ def forward(
                 position_ids_fwd, _, _, _, _ = unpad_input(position_ids.unsqueeze(-1), attention_mask)
                 # (nnz, 1) -> (1, nnz)
                 position_ids_fwd = position_ids_fwd.transpose(0, 1)
+                if support_ids_fwd is not None:
+                    support_ids_fwd = support_ids_fwd.flatten(0, 1).index_select(0, nnz_indices).unsqueeze(0)
+                    target_loss_mask_fwd = target_loss_mask_fwd.flatten().index_select(0, nnz_indices).unsqueeze(0)
+                    support_trajectory_ids_fwd = (
+                        support_trajectory_ids_fwd.flatten().index_select(0, nnz_indices).unsqueeze(0)
+                    )
                 attention_mask_fwd = None  # no attention mask with FA 2
 
         sequences_rolled = torch.roll(sequences_fwd, shifts=-1, dims=1)
+        support_ids_rolled = torch.roll(support_ids_fwd, shifts=-1, dims=1) if support_ids_fwd is not None else None
+        target_loss_mask_rolled = (
+            torch.roll(target_loss_mask_fwd, shifts=-1, dims=1) if target_loss_mask_fwd is not None else None
+        )
+        support_trajectory_ids_rolled = (
+            torch.roll(support_trajectory_ids_fwd, shifts=-1, dims=1)
+            if support_trajectory_ids_fwd is not None
+            else None
+        )
         if self.sequence_parallel_size > 1:
             # NOTE: don't pass any attn mask with sample packing
             attention_mask_fwd = None if self.remove_microbatch_padding else attention_mask_fwd
@@ -301,6 +338,27 @@ def forward(
             sequences_rolled, _, _, _ = ulysses_pad_and_slice_inputs(
                 sequences_rolled, None, None, self.sequence_parallel_size
             )
+            if support_ids_rolled is not None:
+                support_ids_rolled, _, _, _ = ulysses_pad_and_slice_inputs(
+                    support_ids_rolled,
+                    None,
+                    None,
+                    self.sequence_parallel_size,
+                    input_padding_value=-1,
+                )
+                target_loss_mask_rolled, _, _, _ = ulysses_pad_and_slice_inputs(
+                    target_loss_mask_rolled,
+                    None,
+                    None,
+                    self.sequence_parallel_size,
+                )
+                support_trajectory_ids_rolled, _, _, _ = ulysses_pad_and_slice_inputs(
+                    support_trajectory_ids_rolled,
+                    None,
+                    None,
+                    self.sequence_parallel_size,
+                    input_padding_value=-1,
+                )
 
         if self.is_vlm:
             # NOTE: transformers v5 introduced `mm_token_type_ids` to distinguish text
@@ -333,12 +391,32 @@ def forward(
         logits_BSV = output["logits"]
         logits_BSV.div_(temperature)
 
-        # NOTE: this is slightly inaccurate with sample packing because last token from nth seq -> first token of n+1th seq loss is added.
-        log_probs = logprobs_from_logits(
-            logits_BSV,
-            sequences_rolled,
-            inplace_backward=True,
-        )
+        if enable_sample_support_replay:
+            assert (
+                support_ids_rolled is not None
+                and target_loss_mask_rolled is not None
+                and support_trajectory_ids_rolled is not None
+            )
+            support_scores = aligned_sample_support_scores(
+                logits_BSV,
+                sequences_rolled,
+                support_ids_rolled,
+                target_loss_mask_rolled,
+                vocab_start_index=0,
+                vocab_end_index=logits_BSV.shape[-1],
+                tp_group=None,
+                inference_only=not torch.is_grad_enabled(),
+                compute_entropy=compute_entropy,
+                entropy_requires_grad=compute_entropy and entropy_requires_grad,
+                trajectory_ids=support_trajectory_ids_rolled,
+                num_trajectories=sequences.shape[0],
+            )
+            log_probs = support_scores.logprobs
+            support_entropy = support_scores.entropy
+            support_entropy_mask = support_scores.valid_mask
+        else:
+            # NOTE: this is slightly inaccurate with sample packing because last token from nth seq -> first token of n+1th seq loss is added.
+            log_probs = logprobs_from_logits(logits_BSV, sequences_rolled, inplace_backward=True)
 
         # gather output if sp > 1
         if self.sequence_parallel_size > 1:
@@ -346,6 +424,13 @@ def forward(
             log_probs = gather_outputs_and_unpad(
                 log_probs, gather_dim=dim, unpad_dim=dim, padding_size=pad_size
             )  # shape can be (1, nnz) - with packing or (B, S) - without packing
+            if support_entropy is not None and support_entropy_mask is not None:
+                support_entropy = gather_outputs_and_unpad(
+                    support_entropy, gather_dim=dim, unpad_dim=dim, padding_size=pad_size
+                )
+                support_entropy_mask = gather_outputs_and_unpad(
+                    support_entropy_mask, gather_dim=dim, unpad_dim=dim, padding_size=pad_size
+                )
 
         if self.remove_microbatch_padding:
             # add padding back - postprocess logprobs to be compatible with original tensor
@@ -354,36 +439,39 @@ def forward(
             log_probs = pad_input(
                 log_probs.transpose(0, 1), indices=nnz_indices, batch=batch_size, seqlen=seqlen
             ).squeeze(-1)
+            if support_entropy is not None and support_entropy_mask is not None:
+                support_entropy = pad_input(
+                    support_entropy.transpose(0, 1), indices=nnz_indices, batch=batch_size, seqlen=seqlen
+                ).squeeze(-1)
+                support_entropy_mask = pad_input(
+                    support_entropy_mask.transpose(0, 1), indices=nnz_indices, batch=batch_size, seqlen=seqlen
+                ).squeeze(-1)
 
         if compute_entropy:
-            # For sample packing: entropy is calculated on unpacked data, so no attention mask needed
-            # For non-sample packing: pass the attention mask to exclude padding tokens
-            entropy_mask = None
-            if not self.remove_microbatch_padding:
-                # Non-sample packing: pass attention mask to handle padding
-                # Use attention_mask_fwd which may be sliced (if sequence_parallel_size > 1) or full
-                entropy_mask = attention_mask_fwd
-
-            entropy_BS = self.chunked_entropy_from_logits_fn(
-                logits_BSV,
-                requires_grad=entropy_requires_grad,
-                attention_mask=entropy_mask,
-                chunk_size=self.logprobs_chunk_size,
-            )
-
-            if self.sequence_parallel_size > 1:
-                dim = entropy_BS.ndim - 1
-                entropy_BS = gather_outputs_and_unpad(
-                    entropy_BS, gather_dim=dim, unpad_dim=dim, padding_size=pad_size
-                )  # shape can be (1, nnz) - with packing or (B,S) - without packing
-            if self.remove_microbatch_padding:
-                entropy_BS = pad_input(
-                    entropy_BS.transpose(0, 1), indices=nnz_indices, batch=batch_size, seqlen=seqlen
-                ).squeeze(
-                    -1
-                )  # (1, nnz) -> (B, S)
+            if support_entropy is not None and support_entropy_mask is not None:
+                output["entropy"] = support_entropy
+                output["entropy_mask"] = support_entropy_mask
+            else:
+                # For sample packing: entropy is calculated on unpacked data, so no attention mask needed
+                # For non-sample packing: pass the attention mask to exclude padding tokens
+                entropy_mask = None if self.remove_microbatch_padding else attention_mask_fwd
+                entropy_BS = self.chunked_entropy_from_logits_fn(
+                    logits_BSV,
+                    requires_grad=entropy_requires_grad,
+                    attention_mask=entropy_mask,
+                    chunk_size=self.logprobs_chunk_size,
+                )
 
-            output["entropy"] = entropy_BS
+                if self.sequence_parallel_size > 1:
+                    dim = entropy_BS.ndim - 1
+                    entropy_BS = gather_outputs_and_unpad(
+                        entropy_BS, gather_dim=dim, unpad_dim=dim, padding_size=pad_size
+                    )
+                if self.remove_microbatch_padding:
+                    entropy_BS = pad_input(
+                        entropy_BS.transpose(0, 1), indices=nnz_indices, batch=batch_size, seqlen=seqlen
+                    ).squeeze(-1)
+                output["entropy"] = entropy_BS
 
         if isinstance(num_actions, list):
             if len(num_actions) == 1:
diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py
index 3e4ad83a0b..47c84b7f4d 100644
--- a/skyrl/backends/skyrl_train/workers/worker.py
+++ b/skyrl/backends/skyrl_train/workers/worker.py
@@ -890,6 +890,7 @@ def _forward_backward_micro(
         loss_mask = experience.loss_mask
         response_mask = experience.response_mask
         rollout_action_logprobs = experience.rollout_logprobs
+        sample_support_replay = self.cfg.algorithm.enable_sample_support_replay
 
         # Determine which loss function to use
         resolved_loss_name = loss_fn if loss_fn is not None else self.cfg.algorithm.policy_loss_type
@@ -924,6 +925,9 @@ def _forward_backward_micro(
                 entropy_requires_grad=self.cfg.algorithm.use_entropy_loss,
                 pixel_values=experience.pixel_values,
                 image_grid_thw=experience.image_grid_thw,
+                sample_support_ids=experience.sample_support_ids if sample_support_replay else None,
+                loss_mask=loss_mask if sample_support_replay else None,
+                enable_sample_support_replay=sample_support_replay,
             )
             # loss function
             # TODO: recompute advantages
@@ -1000,7 +1004,12 @@ def _forward_backward_micro(
                 # batch_size, seqlen
                 entropy_BS = output["entropy"]
                 entropy_BS = entropy_BS[:, -num_actions - 1 : -1]
-                entropy = masked_mean(entropy_BS, loss_mask)
+                entropy_loss_mask = loss_mask
+                if "entropy_mask" in output:
+                    entropy_loss_mask = entropy_loss_mask.to(torch.bool) & output["entropy_mask"][
+                        :, -num_actions - 1 : -1
+                    ].to(torch.bool)
+                entropy = masked_mean(entropy_BS, entropy_loss_mask)
 
             if self.cfg.algorithm.use_entropy_loss:
                 entropy_loss_term = entropy * self.cfg.algorithm.entropy_loss_coef
@@ -1177,6 +1186,7 @@ def _forward_micro_with_loss(
         loss_mask = experience.loss_mask
         response_mask = experience.response_mask
         rollout_action_logprobs = experience.rollout_logprobs
+        sample_support_replay = self.cfg.algorithm.enable_sample_support_replay
 
         current_loss_fn = PolicyLossRegistry.get(loss_fn)
 
@@ -1199,6 +1209,9 @@ def _forward_micro_with_loss(
                 entropy_requires_grad=False,
                 pixel_values=experience.pixel_values,
                 image_grid_thw=experience.image_grid_thw,
+                sample_support_ids=experience.sample_support_ids if sample_support_replay else None,
+                loss_mask=loss_mask if sample_support_replay else None,
+                enable_sample_support_replay=sample_support_replay,
             )
             policy_loss, _ = current_loss_fn(
                 action_log_probs,
@@ -1260,6 +1273,7 @@ def _forward_micro_batch(self, micro_batch: TrainingInputBatch) -> TrainingOutpu
         attention_mask = micro_batch["attention_mask"]
         pixel_values = micro_batch.get("pixel_values", None)
         image_grid_thw = micro_batch.get("image_grid_thw", None)
+        sample_support_replay = self.cfg.algorithm.enable_sample_support_replay
 
         with torch.no_grad(), torch.autocast(dtype=torch.bfloat16, device_type="cuda"):
             policy_logprob = self.model(
@@ -1270,6 +1284,10 @@ def _forward_micro_batch(self, micro_batch: TrainingInputBatch) -> TrainingOutpu
                 temperature=self.cfg.algorithm.temperature,
                 pixel_values=pixel_values,
                 image_grid_thw=image_grid_thw,
+                # Policy ratios require recomputed logprobs to use the same support normalization.
+                sample_support_ids=micro_batch["sample_support_ids"] if sample_support_replay else None,
+                loss_mask=micro_batch["loss_mask"] if sample_support_replay else None,
+                enable_sample_support_replay=sample_support_replay,
             )
         policy_logprob = policy_logprob.to("cpu")
         output = TrainingOutputBatch(
@@ -1563,14 +1581,19 @@ def _forward_micro_batch(self, micro_batch: TrainingInputBatch) -> TrainingOutpu
         attention_mask = micro_batch["attention_mask"]
         pixel_values = micro_batch.get("pixel_values", None)
         image_grid_thw = micro_batch.get("image_grid_thw", None)
+        sample_support_replay = self.cfg.algorithm.enable_sample_support_replay
         with torch.no_grad(), torch.autocast(dtype=torch.bfloat16, device_type="cuda"):
             log_probs = self.model(
                 sequences,
                 response_length,
                 attention_mask,
                 return_output=False,
+                temperature=self.cfg.algorithm.temperature if sample_support_replay else 1.0,
                 pixel_values=pixel_values,
                 image_grid_thw=image_grid_thw,
+                sample_support_ids=micro_batch["sample_support_ids"] if sample_support_replay else None,
+                loss_mask=micro_batch["loss_mask"] if sample_support_replay else None,
+                enable_sample_support_replay=sample_support_replay,
             )
         log_probs = log_probs.to("cpu")
         output = TrainingOutputBatch(
diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py
index 2efa37ad6c..838ba14d8f 100644
--- a/skyrl/backends/skyrl_train/workers/worker_utils.py
+++ b/skyrl/backends/skyrl_train/workers/worker_utils.py
@@ -163,6 +163,7 @@ def batch_to_experience(batch: TrainingInputBatch):
             num_actions=batch.metadata["response_length"],  # int
             rollout_logprobs=batch.get("rollout_logprobs"),
             rollout_expert_indices=batch.get("rollout_expert_indices"),
+            sample_support_ids=batch.get("sample_support_ids"),
             router_padding_mask=batch.get("router_padding_mask"),
             # additional info
             # can be used to log metrics etc for micro-batches in the worker
@@ -334,6 +335,14 @@ def _create_padding_microbatch(self) -> TrainingInputBatch:
             )
         if self.data.get("router_padding_mask") is not None:
             data["router_padding_mask"] = torch.ones((batch_size, seq_len), dtype=torch.bool, device=device)
+        if self.data.get("sample_support_ids") is not None:
+            ref_tensor = self.data["sample_support_ids"]
+            data["sample_support_ids"] = torch.full(
+                (batch_size, *ref_tensor.shape[1:]),
+                -1,
+                dtype=ref_tensor.dtype,
+                device=device,
+            )
         data.metadata = {}
         if self.data.metadata:
             data.metadata.update(self.data.metadata)
diff --git a/skyrl/benchmarks/bench_dense_sample_support_wire.py b/skyrl/benchmarks/bench_dense_sample_support_wire.py
new file mode 100644
index 0000000000..955c66e0f4
--- /dev/null
+++ b/skyrl/benchmarks/bench_dense_sample_support_wire.py
@@ -0,0 +1,61 @@
+"""Measure dense sampler-support list versus packed JSON transport."""
+
+import statistics
+import time
+
+import numpy as np
+import orjson
+
+from skyrl.backends.skyrl_train.inference_servers.sample_support_set_wire import (
+    decode_sample_support_set,
+    encode_sample_support_set,
+)
+
+
+def _median_ms(fn, iterations: int) -> float:
+    durations = []
+    for _ in range(iterations):
+        start = time.perf_counter()
+        fn()
+        durations.append((time.perf_counter() - start) * 1000)
+    return statistics.median(durations)
+
+
+def _report(label: str, support: np.ndarray, iterations: int) -> None:
+    list_json = orjson.dumps(support.tolist())
+    packed_json = orjson.dumps(encode_sample_support_set(support))
+
+    print(f"scenario={label} shape={support.shape} iterations={iterations}")
+    print(f"list_json_bytes={len(list_json)} packed_json_bytes={len(packed_json)}")
+    print(f"size_ratio={len(packed_json) / len(list_json):.4f}")
+    print(
+        "list_encode_ms=%.3f packed_encode_ms=%.3f"
+        % (
+            _median_ms(lambda: orjson.dumps(support.tolist()), iterations),
+            _median_ms(lambda: orjson.dumps(encode_sample_support_set(support)), iterations),
+        )
+    )
+    print(
+        "list_decode_ms=%.3f packed_decode_ms=%.3f"
+        % (
+            _median_ms(lambda: orjson.loads(list_json), iterations),
+            _median_ms(lambda: decode_sample_support_set(orjson.loads(packed_json)).tolist(), iterations),
+        )
+    )
+
+
+def main() -> None:
+    rng = np.random.default_rng(17)
+    tokens, iterations = 4096, 100
+    for top_k in (8, 64):
+        full = rng.integers(0, 152_064, size=(tokens, top_k), dtype=np.int32)
+        variable = full.copy()
+        widths = rng.integers(1, top_k + 1, size=tokens)
+        variable[np.arange(top_k)[None, :] >= widths[:, None]] = -1
+
+        _report(f"top_k_{top_k}_full", full, iterations)
+        _report(f"top_k_{top_k}_top_p_variable", variable, iterations)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py
index 57c7521910..18f30e84c6 100644
--- a/skyrl/train/config/config.py
+++ b/skyrl/train/config/config.py
@@ -916,6 +916,8 @@ class AlgorithmConfig(BaseConfig):
     Enabled Truncated Importance Sampling (TIS) as proposed in https://fengyao.notion.site/off-policy-rl."""
     off_policy_correction: OffPolicyCorrectionConfig = field(default_factory=OffPolicyCorrectionConfig)
     """See https://docs.skyrl.ai/docs/algorithms/off_policy_correction for a full guide."""
+    enable_sample_support_replay: bool = False
+    """Renormalize policy logprobs over the sampler's recorded support."""
     sapo: SAPOConfig = field(default_factory=SAPOConfig)
     """Only used when ``policy_loss_type="sapo"``."""
     value_clip: float = 0.2
@@ -1142,6 +1144,8 @@ class InferenceEngineConfig(BaseConfig):
     enable_return_routed_experts: bool = False
     """Return per-layer expert routing indices, for rollout router replay (R3) when training an MoE model.
     Used together with ``trainer.policy.megatron_config.moe_enable_routing_replay``."""
+    enable_return_sample_support_set: bool = False
+    """Return the bounded post-filter sampler support for each generated token."""
     max_num_batched_tokens: int = 8192
     """vLLM continuous-batching parameter: maximum number of tokens to pack into a batch."""
     enforce_eager: bool = False
@@ -1737,6 +1741,28 @@ def __post_init__(self):
         if self.trainer.algorithm.temperature is None:
             self.trainer.algorithm.temperature = self.generator.sampling_params.temperature
 
+        capture_sample_support = self.generator.inference_engine.enable_return_sample_support_set
+        replay_sample_support = self.trainer.algorithm.enable_sample_support_replay
+        sampling_params = self.generator.sampling_params
+        if replay_sample_support and not capture_sample_support:
+            raise ValueError(
+                "trainer.algorithm.enable_sample_support_replay requires "
+                "generator.inference_engine.enable_return_sample_support_set"
+            )
+        if replay_sample_support and self.trainer.strategy not in {"megatron", "fsdp"}:
+            raise ValueError("sample-support replay requires trainer.strategy=megatron or fsdp")
+        if capture_sample_support:
+            if sampling_params.temperature <= 0:
+                raise ValueError("sample-support capture requires generator.sampling_params.temperature > 0")
+            if sampling_params.top_k <= 1:
+                raise ValueError("sample-support capture requires generator.sampling_params.top_k > 1")
+            if sampling_params.repetition_penalty != 1.0:
+                raise ValueError("sample-support capture requires repetition_penalty=1.0")
+            if sampling_params.additional_kwargs:
+                raise ValueError("sample-support capture does not support sampling_params.additional_kwargs")
+            if self.generator.vision_language_generator:
+                raise ValueError("sample-support capture does not support vision_language_generator")
+
         if self.data.dataloader.num_workers is None:
             self.data.dataloader.num_workers = 8
         if self.data.dataloader.persistent_workers and self.data.dataloader.num_workers == 0:
diff --git a/skyrl/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py
index fb7ac6e2f7..9e6734d9b9 100644
--- a/skyrl/train/dataset/preprocess.py
+++ b/skyrl/train/dataset/preprocess.py
@@ -283,6 +283,75 @@ def convert_prompts_responses_to_batch_tensors(
     )
 
 
+def build_dense_sample_support(
+    rollout_sample_support: Optional[List[List[List[int]]]],
+    response_ids: List[List[int]],
+    loss_masks: List[List[int]],
+    sequence_length: int,
+    top_k: int,
+    eos_token_id: int,
+) -> Optional[Integer[torch.Tensor, "batch seq_len topk"]]:
+    """Validate and left-pad per-token sampler support for replay."""
+    if rollout_sample_support is None:
+        return None
+    if len(rollout_sample_support) != len(response_ids):
+        raise ValueError("rollout_sample_support must have one entry per trajectory")
+    if len(loss_masks) != len(response_ids):
+        raise ValueError("loss_masks must have one entry per trajectory")
+
+    support = torch.full((len(response_ids), sequence_length, top_k), -1, dtype=torch.int32)
+    int32_max = int(np.iinfo(np.int32).max)
+    for sample_index, (sample_rows, sampled_tokens, sample_loss_mask) in enumerate(
+        zip(rollout_sample_support, response_ids, loss_masks, strict=True)
+    ):
+        if len(sample_rows) != len(sampled_tokens):
+            raise ValueError(
+                f"rollout_sample_support[{sample_index}] has {len(sample_rows)} rows for "
+                f"{len(sampled_tokens)} response tokens"
+            )
+        if len(sample_loss_mask) != len(sampled_tokens):
+            raise ValueError(
+                f"loss_masks[{sample_index}] has {len(sample_loss_mask)} entries for "
+                f"{len(sampled_tokens)} response tokens"
+            )
+
+        sample_support = torch.full((len(sample_rows), top_k), -1, dtype=torch.int64)
+        for token_index, row in enumerate(sample_rows):
+            if row:
+                if len(row) != top_k:
+                    raise ValueError("rollout_sample_support rows must match generator.sampling_params.top_k")
+                sample_support[token_index] = torch.as_tensor(row, dtype=torch.int64)
+
+        valid = sample_support >= 0
+        if torch.any((sample_support < -1) | (sample_support > int32_max)):
+            raise ValueError("rollout_sample_support vocab ids must fit non-negative int32")
+        if torch.any(valid & ((~valid).cumsum(dim=1) > 0)):
+            raise ValueError("rollout_sample_support padding must use trailing -1 values")
+
+        sampled = torch.as_tensor(sampled_tokens, dtype=torch.int64).unsqueeze(1)
+        loss_bearing = torch.as_tensor(sample_loss_mask, dtype=torch.bool)
+        has_support = valid.any(dim=1)
+        unsupported_loss = loss_bearing & ~has_support
+        if torch.count_nonzero(unsupported_loss) > 1:
+            raise ValueError(f"rollout_sample_support[{sample_index}] has more than one loss-bearing unsupported token")
+        unsupported_non_eos = unsupported_loss & (sampled.squeeze(1) != eos_token_id)
+        if torch.any(unsupported_non_eos):
+            token_index = int(torch.where(unsupported_non_eos)[0][0])
+            raise ValueError(
+                f"rollout_sample_support[{sample_index}][{token_index}] is empty for a loss-bearing non-EOS token"
+            )
+        missing = loss_bearing & has_support & ~torch.any(sample_support == sampled, dim=1)
+        if torch.any(missing):
+            missing_token = sampled_tokens[int(torch.where(missing)[0][0])]
+            raise ValueError(f"sampled token {missing_token} is missing from rollout_sample_support")
+
+        start = sequence_length - len(sampled_tokens)
+        if start < 0:
+            raise ValueError("response tokens exceed the sample-support sequence width")
+        support[sample_index, start:] = sample_support.to(torch.int32)
+    return support
+
+
 def compute_prompt_boundaries(uids: List[str]) -> List[Tuple[int, int]]:
     """Compute per-prompt ``(start, end)`` slices from a flat ``uids`` list.
 
diff --git a/skyrl/train/dataset/replay_buffer.py b/skyrl/train/dataset/replay_buffer.py
index 6f627e31ed..20f4626a44 100644
--- a/skyrl/train/dataset/replay_buffer.py
+++ b/skyrl/train/dataset/replay_buffer.py
@@ -78,6 +78,7 @@ class Experience:
     # Per-row sub-sequence lengths for sequence packing (one 1-D int tensor per
     # packed row); ``None`` when packing is off.
     sub_seq_lengths: Optional[TensorList] = None
+    sample_support_ids: Optional[Integer[torch.Tensor, "batch seq_len topk"]] = None
 
     @torch.no_grad()
     def to_device(self, device: torch.device) -> None:
@@ -102,6 +103,8 @@ def to_device(self, device: torch.device) -> None:
             self.rollout_logprobs = to(self.rollout_logprobs, device)
         if self.rollout_expert_indices is not None:
             self.rollout_expert_indices = to(self.rollout_expert_indices, device)
+        if self.sample_support_ids is not None:
+            self.sample_support_ids = to(self.sample_support_ids, device)
         if self.router_padding_mask is not None:
             self.router_padding_mask = to(self.router_padding_mask, device)
         if self.pixel_values is not None:
@@ -133,6 +136,8 @@ def pin_memory(self):
             self.rollout_logprobs = self.rollout_logprobs.pin_memory()
         if self.rollout_expert_indices is not None:
             self.rollout_expert_indices = self.rollout_expert_indices.pin_memory()
+        if self.sample_support_ids is not None:
+            self.sample_support_ids = self.sample_support_ids.pin_memory()
         if self.router_padding_mask is not None:
             self.router_padding_mask = self.router_padding_mask.pin_memory()
         return self
diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index 81792e0a2a..0b0838e1b0 100644
--- a/skyrl/train/generators/base.py
+++ b/skyrl/train/generators/base.py
@@ -52,6 +52,7 @@ class GeneratorOutput(TypedDict):
     # record its split.
     trajectory_time_splits: Optional[Dict[str, List[float]]]
     rollout_expert_indices: Optional[List[RoutedExpertIndices]]
+    rollout_sample_support: Optional[List[List[List[int]]]]
     # Applicable only for step-wise training
     is_last_step: Optional[List[bool]]
     # Per-row env metrics (one dict per row in the flattened batch). Used by
diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index 26b1102e6c..e9e6bb6a9d 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -13,17 +13,24 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import uuid4
 
+import numpy as np
 import torch
 from loguru import logger
 from tqdm.asyncio import tqdm
 
 import skyrl_gym
+from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import (
+    TokenMetadataTrace,
+)
 from skyrl.backends.skyrl_train.inference_servers.base import (
     ConversationType,
     InferenceEngineInput,
     InferenceEngineInterface,
 )
-from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices
+from skyrl.backends.skyrl_train.utils.routed_experts import (
+    RoutedExpertIndices,
+    RoutedExpertTrace,
+)
 from skyrl.train.config import GeneratorConfig, SkyRLGymConfig
 from skyrl.train.generators.base import (
     GeneratorInput,
@@ -52,6 +59,7 @@ class TrajectoryOutput:
     rollout_logprobs: Optional[List[float]]
     env_metrics: Dict[str, Any]
     rollout_expert_indices: Optional[RoutedExpertIndices] = None
+    rollout_sample_support: Optional[List[List[int]]] = None
     pixel_values: Optional[torch.Tensor] = None
     image_grid_thw: Optional[torch.Tensor] = None
     # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: agent loops may
@@ -83,7 +91,8 @@ class AgentLoopState:
     rollout_logprobs: Optional[List[float]]
     response_end_idx: Optional[int]
     done: bool
-    rollout_expert_indices: Optional[RoutedExpertIndices] = None
+    routed_expert_trace: Optional[RoutedExpertTrace] = None
+    sample_support_trace: Optional[TokenMetadataTrace] = None
 
 
 @dataclass
@@ -93,13 +102,22 @@ class TurnOutput:
     output_logprobs: Optional[List[float]]
     new_obs: ConversationType
     obs_ids: List[int]
-    rollout_expert_indices: Optional[RoutedExpertIndices]
     reward: Optional[float]
+    rollout_sample_support: Optional[np.ndarray] = None
     added_eos: bool = False
 
-    def get_turn_rollout_expert_indices(self) -> Optional[RoutedExpertIndices]:
-        """Return only routes that the inference model actually executed."""
-        return self.rollout_expert_indices
+    def get_turn_rollout_sample_support(self) -> Optional[np.ndarray]:
+        if self.rollout_sample_support is None:
+            return None
+        padding_count = int(self.added_eos) + len(self.obs_ids)
+        if not padding_count:
+            return self.rollout_sample_support
+        padding = np.full(
+            (padding_count, self.rollout_sample_support.shape[1]),
+            -1,
+            dtype=self.rollout_sample_support.dtype,
+        )
+        return np.concatenate((self.rollout_sample_support, padding), axis=0)
 
     def get_turn_loss_mask(self) -> List[int]:
         """
@@ -363,6 +381,8 @@ async def agent_loop(
             current_sampling_params: dict = (
                 sampling_params if sampling_params is not None else asdict(self.generator_cfg.sampling_params)
             )
+            capture_sample_support = self.generator_cfg.inference_engine.enable_return_sample_support_set
+            sample_support_width = current_sampling_params["top_k"] if capture_sample_support else 0
 
             # Accumulate per-step rewards. Format: (reward, response_end_token_idx)
             per_step_rewards: List[Tuple[float, Optional[int]]] = []
@@ -379,6 +399,10 @@ async def agent_loop(
                 rollout_logprobs=[] if get_logprobs else None,
                 response_end_idx=None,
                 done=False,
+                routed_expert_trace=(
+                    RoutedExpertTrace() if self.generator_cfg.inference_engine.enable_return_routed_experts else None
+                ),
+                sample_support_trace=TokenMetadataTrace() if capture_sample_support and not is_step_wise else None,
             )
 
             while not agent_loop_state.done:
@@ -401,11 +425,13 @@ async def agent_loop(
                     agent_loop_state.loss_mask = []
                     agent_loop_state.rollout_logprobs = None
 
+                routed_expert_trace = agent_loop_state.routed_expert_trace
                 engine_input = InferenceEngineInput(
                     prompt_token_ids=[agent_loop_state.input_ids],
                     session_ids=[session_id],
                     sampling_params=sampling_params,
                     cache_salt=cache_salt,
+                    routed_experts_prompt_starts=[routed_expert_trace.prompt_start] if routed_expert_trace else None,
                 )
                 llm_call_start_time = time.monotonic()
                 engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name)
@@ -426,6 +452,27 @@ async def agent_loop(
                         raise ValueError(
                             "Rollout expert indices bookkeeping is not supported with custom chat template"
                         )
+                if routed_expert_trace is not None:
+                    if rollout_expert_indices is None:
+                        raise ValueError("R3 generation did not return routed expert indices")
+                    routed_expert_trace.record_generation(
+                        prompt_token_count=len(agent_loop_state.input_ids),
+                        generated_token_count=len(output_ids),
+                        routed_experts=rollout_expert_indices,
+                    )
+                sample_support_rows = None
+                if capture_sample_support:
+                    sample_support_rows = np.asarray(
+                        engine_output["rollout_sample_support"][0],
+                        dtype=np.int32,
+                        order="C",
+                    ).reshape(-1, sample_support_width)
+                    if self.custom_chat_template is not None:
+                        raise ValueError("Sample-support bookkeeping is not supported with custom chat template")
+                    if sample_support_rows.shape[0] != len(output_ids):
+                        raise ValueError(
+                            f"Sample support has {sample_support_rows.shape[0]} rows for {len(output_ids)} tokens"
+                        )
                 # Append eos when sampling_params.stop is not None. Does not affect 3.a as chat templates add eos_token.
                 # sampling_params is not None for eval, but None for training (which uses engine.sampling_params which are from cfg)
                 stop_strs = current_sampling_params.get("stop", None)
@@ -459,6 +506,10 @@ async def agent_loop(
                     )
                     output = env_step_output["postprocessed_action"]
                     output_ids = self.tokenizer.encode(output, add_special_tokens=False)
+                    if routed_expert_trace is not None:
+                        raise ValueError("R3 bookkeeping is incompatible with postprocessed_action")
+                    if sample_support_rows is not None:
+                        raise ValueError("Sample-support bookkeeping is incompatible with postprocessed_action")
 
                 obs_ids = self.get_obs_ids_from_obs(new_obs, agent_loop_state.done)
 
@@ -470,8 +521,8 @@ async def agent_loop(
                     new_obs=new_obs,
                     reward=step_reward,
                     obs_ids=obs_ids,
+                    rollout_sample_support=sample_support_rows,
                     added_eos=added_eos,
-                    rollout_expert_indices=rollout_expert_indices,
                 )
 
                 if is_step_wise:
@@ -482,6 +533,7 @@ async def agent_loop(
                     # agent loop only tracks loss mask and rollout logprobs for this turn with step_wise training
                     turn_loss_mask = turn_output.get_turn_loss_mask()
                     turn_response_logprobs: Optional[List[float]] = turn_output.get_turn_rollout_logprobs()
+                    turn_sample_support = turn_output.get_turn_rollout_sample_support()
 
                     per_step_output = TrajectoryOutput(
                         response_ids=turn_response_ids,
@@ -491,7 +543,9 @@ async def agent_loop(
                         rollout_logprobs=turn_response_logprobs,
                         stop_reason=stop_reason,
                         env_metrics=env.get_metrics() if agent_loop_state.done else {},
-                        rollout_expert_indices=turn_output.get_turn_rollout_expert_indices(),
+                        rollout_sample_support=(
+                            turn_sample_support.tolist() if turn_sample_support is not None else None
+                        ),
                     )
                     agent_loop_output.step_outputs.append(per_step_output)
 
@@ -523,6 +577,7 @@ async def agent_loop(
             prompt_ids = agent_loop_state.input_ids[:initial_prompt_length]
             rollout_logprobs = None
             rollout_expert_indices_out = None
+            rollout_sample_support_out = None
             response_ids = None
 
             # Prepare the final loss_mask, response_ids and rollout_logprobs .
@@ -553,10 +608,6 @@ async def agent_loop(
                     rollout_logprobs = agent_loop_state.rollout_logprobs[
                         : agent_loop_state.response_end_idx - initial_prompt_length + 1
                     ]
-                if agent_loop_state.rollout_expert_indices is not None:
-                    rollout_expert_indices_out = agent_loop_state.rollout_expert_indices[
-                        : agent_loop_state.response_end_idx + 1
-                    ]
                 # fix index for per_step_rewards
                 per_step_rewards = [(reward, idx - initial_prompt_length) for reward, idx in per_step_rewards]
                 assert len(loss_mask) == len(
@@ -571,8 +622,26 @@ async def agent_loop(
                     loss_mask.append(1)
                     if rollout_logprobs is not None:
                         rollout_logprobs.append(0.0)
+                    if agent_loop_state.sample_support_trace is not None:
+                        padding = np.full((1, sample_support_width), -1, dtype=np.int32)
+                        agent_loop_state.sample_support_trace.append(padding, expected_rows=1)
                     appended_eos_token = True
 
+            if agent_loop_state.routed_expert_trace is not None and agent_loop_state.routed_expert_trace.prompt_start:
+                rollout_expert_indices_out = agent_loop_state.routed_expert_trace.finalize(
+                    token_count=len(prompt_ids) + len(response_ids),
+                    loss_mask=[0] * len(prompt_ids) + loss_mask,
+                )
+            if agent_loop_state.sample_support_trace is not None and agent_loop_state.sample_support_trace.num_rows:
+                sample_support_rows = agent_loop_state.sample_support_trace.finalize(
+                    expected_rows=agent_loop_state.sample_support_trace.num_rows
+                )
+                if sample_support_rows.shape[0] < len(response_ids):
+                    raise ValueError(
+                        f"Sample-support trace has {sample_support_rows.shape[0]} rows for {len(response_ids)} tokens"
+                    )
+                rollout_sample_support_out = sample_support_rows[: len(response_ids)].tolist()
+
             if self.generator_cfg.step_wise_trajectories:
                 for per_step_output, (reward, resp_end_idx) in zip(agent_loop_output.step_outputs, per_step_rewards):
                     per_token_reward = [0.0] * len(per_step_output.response_ids)
@@ -591,6 +660,7 @@ async def agent_loop(
                     rollout_logprobs=rollout_logprobs,
                     env_metrics=env_metrics,
                     rollout_expert_indices=rollout_expert_indices_out,
+                    rollout_sample_support=rollout_sample_support_out,
                 )
 
             agent_loop_output = self._post_process_agent_loop_output(
@@ -758,6 +828,7 @@ async def generate_batched(
         stop_reasons = engine_output["stop_reasons"]
         logprobs = engine_output.get("response_logprobs", None)
         raw_rollout_expert_indices = engine_output.get("rollout_expert_indices", None)
+        raw_rollout_sample_support = engine_output.get("rollout_sample_support", None)
 
         truncated_responses = []
         rewards = []
@@ -765,6 +836,9 @@ async def generate_batched(
         env_metrics = []
         truncated_logprobs: Optional[List[List[float]]] = [] if logprobs is not None else None
         truncated_indices: Optional[List[RoutedExpertIndices]] = [] if raw_rollout_expert_indices is not None else None
+        truncated_sample_support: Optional[List[List[List[int]]]] = (
+            [] if raw_rollout_sample_support is not None else None
+        )
 
         for i, (output, response, env, env_class) in enumerate(zip(outputs, responses, envs, env_classes)):
             # step on environment and compute reward
@@ -783,6 +857,8 @@ async def generate_batched(
                 sample_indices = raw_rollout_expert_indices[i]
                 prompt_len = len(prompt_token_ids[i])
                 truncated_indices.append(sample_indices[: prompt_len + len(response)])
+            if raw_rollout_sample_support is not None:
+                truncated_sample_support.append(raw_rollout_sample_support[i][: len(response)])
 
             # Get environment-specific metrics
             env_metrics.append(env.get_metrics())
@@ -804,6 +880,7 @@ async def generate_batched(
             "rollout_metrics": rollout_metrics,
             "rollout_logprobs": truncated_logprobs,
             "rollout_expert_indices": truncated_indices,
+            "rollout_sample_support": truncated_sample_support,
         }
 
         return generator_output
@@ -943,6 +1020,16 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
         else:
             rollout_expert_indices = None
 
+        if self.generator_cfg.step_wise_trajectories:
+            sample_support_values = [
+                step_output.rollout_sample_support for output in all_outputs for step_output in output.step_outputs
+            ]
+        else:
+            sample_support_values = [output.rollout_sample_support for output in all_outputs]
+        rollout_sample_support = (
+            sample_support_values if any(value is not None for value in sample_support_values) else None
+        )
+
         rollout_metrics = get_rollout_metrics(
             responses,
             rewards,
@@ -976,6 +1063,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             "trajectory_generation_times": out_trajectory_generation_times,
             "trajectory_time_splits": out_trajectory_time_splits,
             "rollout_expert_indices": rollout_expert_indices,
+            "rollout_sample_support": rollout_sample_support,
             "is_last_step": is_last_step,
             "env_metrics": env_metrics,
         }
@@ -1047,8 +1135,6 @@ def _update_agent_state_by_retokenizing_chat_history(
         agent_loop_state.response_end_idx = None
         # `logprobs` are not computed because retokenizing breaks token-in-token-out
         agent_loop_state.rollout_logprobs = None
-        # indices are not meaningful when retokenizing
-        agent_loop_state.rollout_expert_indices = None
         return agent_loop_state
 
     def _update_agent_loop_state_with_multiturn_chat_template(
@@ -1100,17 +1186,12 @@ def _update_agent_loop_state_with_multiturn_chat_template(
         loss_mask_for_turn = turn_output.get_turn_loss_mask()
         rollout_logprobs_for_turn = turn_output.get_turn_rollout_logprobs()
 
-        # use the raw rollout expert indices without any appending of observation tokens
-        # this will be overwritten each turn, so we don't need to append observation tokens to it
-        rollout_expert_indices_for_turn = turn_output.rollout_expert_indices
-
         if self.generator_cfg.step_wise_trajectories:
             # cumulative input_ids is not tracked for step wise training
             agent_loop_state.response_end_idx = len(turn_output.output_ids) - 1
-            # no running loss_mask, `rollout_logprobs`, or `rollout_expert_indices` are tracked for step-wise training
+            # no running loss_mask or rollout logprobs are tracked for step-wise training
             agent_loop_state.loss_mask = None
             agent_loop_state.rollout_logprobs = None
-            agent_loop_state.rollout_expert_indices = None
         else:
             # Directly append turn output
             turn_ids = turn_output.output_ids + turn_output.obs_ids
@@ -1119,11 +1200,9 @@ def _update_agent_loop_state_with_multiturn_chat_template(
             agent_loop_state.loss_mask += loss_mask_for_turn
             if agent_loop_state.rollout_logprobs is not None and rollout_logprobs_for_turn is not None:
                 agent_loop_state.rollout_logprobs += rollout_logprobs_for_turn
-            if rollout_expert_indices_for_turn is not None:
-                # overwrite the existing rollout inference indices, since the inference engine should
-                # return the expert indices for the entire sequence including each turn's input
-                # and the final response should not have an observation appended to it
-                agent_loop_state.rollout_expert_indices = rollout_expert_indices_for_turn
+            turn_sample_support = turn_output.get_turn_rollout_sample_support()
+            if agent_loop_state.sample_support_trace is not None and turn_sample_support is not None:
+                agent_loop_state.sample_support_trace.append(turn_sample_support, expected_rows=len(turn_ids))
 
         return agent_loop_state
 
@@ -1187,6 +1266,15 @@ def _update_agent_loop_state_with_singleturn_chat_template(
             rollout_logprobs_for_turn = turn_output.output_logprobs[: len(new_resp_tokens)] + [0.0] * len(
                 obs_ids_to_add
             )
+        turn_sample_support = None
+        if turn_output.rollout_sample_support is not None:
+            generated_support = turn_output.rollout_sample_support[: len(new_resp_tokens)]
+            observation_support = np.full(
+                (len(obs_ids_to_add), generated_support.shape[1]),
+                -1,
+                dtype=generated_support.dtype,
+            )
+            turn_sample_support = np.concatenate((generated_support, observation_support), axis=0)
 
         # Directly append turn output
         agent_loop_state.response_end_idx = len(agent_loop_state.input_ids) + len(new_resp_tokens) - 1
@@ -1194,13 +1282,6 @@ def _update_agent_loop_state_with_singleturn_chat_template(
         agent_loop_state.loss_mask += loss_mask_for_turn
         if agent_loop_state.rollout_logprobs is not None and rollout_logprobs_for_turn is not None:
             agent_loop_state.rollout_logprobs += rollout_logprobs_for_turn
-        if (
-            self.generator_cfg.inference_engine.enable_return_routed_experts
-            and turn_output.rollout_expert_indices is not None
-        ):
-            # overwrite the existing rollout inference indices, since the inference engine should
-            # return the expert indices for the entire sequence including each turn's input and observation tokens
-            # and the final response should not have an observation appended to it
-            agent_loop_state.rollout_expert_indices = turn_output.rollout_expert_indices
-
+        if agent_loop_state.sample_support_trace is not None and turn_sample_support is not None:
+            agent_loop_state.sample_support_trace.append(turn_sample_support, expected_rows=len(turn_ids))
         return agent_loop_state
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index c84ea8a5aa..ef38fcc052 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -822,6 +822,7 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput:
     is_token_level_rewards = isinstance(gen_out["rewards"][0], list)
     has_logprobs = gen_out.get("rollout_logprobs") is not None
     has_stop_reasons = gen_out.get("stop_reasons") is not None
+    has_sample_support = gen_out.get("rollout_sample_support") is not None
 
     # Per-field output accumulators.
     # Fields that we take from all the entries in the merge group
@@ -829,6 +830,7 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput:
     out_response_ids: List[List[int]] = []
     out_loss_masks: List[List[int]] = []
     out_logprobs: Optional[List[List[float]]] = [] if has_logprobs else None
+    out_sample_support: Optional[List[List[List[int]]]] = [] if has_sample_support else None
     # If per-token rewards, we keep appending. If per-turn rewards, we only take from the last turn.
     out_rewards: list = []
 
@@ -842,16 +844,21 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput:
     acc_response: List[int] = list(gen_out["response_ids"][0])
     acc_loss_mask: List[int] = list(gen_out["loss_masks"][0])
     acc_logprobs: Optional[List[float]] = list(gen_out["rollout_logprobs"][0]) if has_logprobs else None
+    acc_sample_support: Optional[List[List[int]]] = (
+        [list(row) for row in gen_out["rollout_sample_support"][0]] if has_sample_support else None
+    )
     acc_rewards_tokens: Optional[List[float]] = list(gen_out["rewards"][0]) if is_token_level_rewards else None
     last = 0
 
     def flush():
-        nonlocal acc_prompt, acc_response, acc_loss_mask, acc_logprobs, acc_rewards_tokens, last
+        nonlocal acc_prompt, acc_response, acc_loss_mask, acc_logprobs, acc_sample_support, acc_rewards_tokens, last
         out_prompt_ids.append(acc_prompt)
         out_response_ids.append(acc_response)
         out_loss_masks.append(acc_loss_mask)
         if has_logprobs:
             out_logprobs.append(acc_logprobs)
+        if has_sample_support:
+            out_sample_support.append(acc_sample_support)
         out_rewards.append(acc_rewards_tokens if is_token_level_rewards else gen_out["rewards"][last])
         if has_stop_reasons:
             out_stop_reasons.append(gen_out["stop_reasons"][last])
@@ -869,6 +876,9 @@ def flush():
             acc_response = list(gen_out["response_ids"][i])
             acc_loss_mask = list(gen_out["loss_masks"][i])
             acc_logprobs = list(gen_out["rollout_logprobs"][i]) if has_logprobs else None
+            acc_sample_support = (
+                [list(row) for row in gen_out["rollout_sample_support"][i]] if has_sample_support else None
+            )
             acc_rewards_tokens = list(gen_out["rewards"][i]) if is_token_level_rewards else None
             last = i
             continue
@@ -883,6 +893,8 @@ def flush():
         acc_loss_mask.extend([0] * len(obs_delta))
         if acc_logprobs is not None:
             acc_logprobs.extend([0.0] * len(obs_delta))
+        if acc_sample_support is not None:
+            acc_sample_support.extend([] for _ in obs_delta)
         if acc_rewards_tokens is not None:
             acc_rewards_tokens.extend([0.0] * len(obs_delta))
 
@@ -891,6 +903,8 @@ def flush():
         acc_loss_mask.extend(gen_out["loss_masks"][i])
         if acc_logprobs is not None:
             acc_logprobs.extend(gen_out["rollout_logprobs"][i])
+        if acc_sample_support is not None:
+            acc_sample_support.extend(gen_out["rollout_sample_support"][i])
         if acc_rewards_tokens is not None:
             acc_rewards_tokens.extend(gen_out["rewards"][i])
 
@@ -905,6 +919,7 @@ def flush():
         "loss_masks": out_loss_masks,
         "stop_reasons": out_stop_reasons,
         "rollout_logprobs": out_logprobs,
+        "rollout_sample_support": out_sample_support,
         "trajectory_ids": out_trajectory_ids,
         "rollout_expert_indices": None,
         "is_last_step": out_is_last_step,
diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py
index ae5eec4dfc..2ba72c8393 100644
--- a/skyrl/train/trainer.py
+++ b/skyrl/train/trainer.py
@@ -53,6 +53,7 @@
 from skyrl.train.config import SkyRLTrainConfig
 from skyrl.train.dataset import PromptDataset
 from skyrl.train.dataset.preprocess import (
+    build_dense_sample_support,
     compute_prompt_boundaries,
     compute_prompt_mini_batch_boundaries,
     convert_prompts_responses_to_batch_tensors,
@@ -873,6 +874,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis
 
         logprobs: Optional[List[List[float]]] = generator_output.get("rollout_logprobs", None)
         rollout_expert_indices = generator_output.get("rollout_expert_indices", None)
+        rollout_sample_support = generator_output.get("rollout_sample_support", None)
 
         pixel_values = generator_output.get("pixel_values", None)
         image_grid_thw = generator_output.get("image_grid_thw", None)
@@ -911,6 +913,14 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis
                 attention_masks_tensor,
                 [len(indices) for indices in rollout_expert_indices],
             )
+        sample_support_ids = build_dense_sample_support(
+            rollout_sample_support,
+            response_ids,
+            loss_masks,
+            sequences_tensor.shape[1],
+            self.cfg.generator.sampling_params.top_k,
+            self.tokenizer.eos_token_id,
+        )
 
         # sanity check for off_policy_correction
         off_policy_correction = self.cfg.trainer.algorithm.off_policy_correction
@@ -933,6 +943,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis
                 "rollout_logprobs": rollout_logprobs_tensor,
                 "rollout_expert_indices": rollout_expert_indices_tensor,
                 "router_padding_mask": router_padding_mask,
+                "sample_support_ids": sample_support_ids,
                 "pixel_values": pixel_values,
                 "image_grid_thw": image_grid_thw,
             },
@@ -1328,6 +1339,8 @@ def fwd_logprobs_values_reward(
             fwd_keys.append("rollout_expert_indices")
         if training_input.get("router_padding_mask") is not None:
             fwd_keys.append("router_padding_mask")
+        if training_input.get("sample_support_ids") is not None:
+            fwd_keys.extend(["sample_support_ids", "loss_mask"])
         if training_input.get("pixel_values") is not None:
             fwd_keys.append("pixel_values")
         if training_input.get("image_grid_thw") is not None:
diff --git a/skyrl/train/utils/trainer_utils.py b/skyrl/train/utils/trainer_utils.py
index 0d91977ecb..f28c9ba8c9 100644
--- a/skyrl/train/utils/trainer_utils.py
+++ b/skyrl/train/utils/trainer_utils.py
@@ -702,6 +702,7 @@ def validate_generator_output(num_prompts: int, generator_output: GeneratorOutpu
             "stop_reasons",
             "trajectory_ids",
             "rollout_expert_indices",
+            "rollout_sample_support",
             "is_last_step",
             "pixel_values",
             "image_grid_thw",
@@ -731,6 +732,11 @@ def validate_generator_output(num_prompts: int, generator_output: GeneratorOutpu
                 f"Response ids and rollout logprobs must have the same length, "
                 f"for sample {i} got {len(response_ids)} and {len(generator_output['rollout_logprobs'][i])}"
             )
+        if generator_output.get("rollout_sample_support") is not None:
+            assert len(response_ids) == len(generator_output["rollout_sample_support"][i]), (
+                "Response ids and rollout sample support must have the same length, "
+                f"for sample {i} got {len(response_ids)} and {len(generator_output['rollout_sample_support'][i])}"
+            )
 
     # loss masks should be non-zero for at least one element for trainer
     if np.concatenate(generator_output["loss_masks"]).sum() == 0:
diff --git a/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py
index ff8f3dc7de..73ba160579 100644
--- a/tests/backends/skyrl_train/distributed/test_token_metadata.py
+++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py
@@ -1,10 +1,15 @@
 import sys
 import types
 
+import numpy as np
 import pytest
 import torch
 
 from skyrl.backends.skyrl_train.distributed.megatron import token_metadata
+from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import (
+    TokenMetadataTrace,
+)
+from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertTrace
 
 
 @pytest.fixture
@@ -86,3 +91,63 @@ def test_packed_layout_aligns_next_token_metadata_and_scatters_rows(monkeypatch,
 
     assert aligned.tolist() == [[11, 12, -1, -1, 21, -1, -1, -1]]
     assert batch_values.tolist() == [[0.0, 1.0, 2.0], [0.0, 0.0, 5.0]]
+
+
+def test_token_metadata_trace_chunks_and_independent_schema() -> None:
+    trace, other = TokenMetadataTrace(), TokenMetadataTrace()
+    trace.append(np.ones((2, 3), dtype=np.int32), expected_rows=2)
+    trace.append(np.zeros((1, 3), dtype=np.int32), expected_rows=1)
+    other.append(np.empty((0, 4), dtype=np.float32), expected_rows=0)
+
+    with pytest.raises(ValueError, match="expected 4"):
+        trace.finalize(expected_rows=4)
+    result = trace.finalize(expected_rows=3)
+    assert result.shape == (3, 3)
+    assert other.finalize(expected_rows=0).shape == (0, 4)
+    with pytest.raises(RuntimeError, match="already finalized"):
+        trace.finalize(expected_rows=3)
+
+
+@pytest.mark.parametrize(
+    ("rows", "expected", "match"),
+    [
+        (np.ones((2, 2), dtype=np.int32), 1, "has 2 rows"),
+        (np.ones((2, 2), dtype=np.int32)[:, ::2], 2, "contiguous"),
+        (np.ones((1, 3), dtype=np.int32), 1, "schema changed"),
+        (np.ones((1, 2), dtype=np.int16), 1, "schema changed"),
+    ],
+)
+def test_token_metadata_trace_rejects_invalid_chunks(rows, expected, match) -> None:
+    trace = TokenMetadataTrace()
+    if rows.shape[0] == 1:
+        trace.append(np.ones((1, 2), dtype=np.int32), expected_rows=1)
+    with pytest.raises(ValueError, match=match):
+        trace.append(rows, expected_rows=expected)
+
+
+def routes(rows: int) -> np.ndarray:
+    return np.arange(rows * 4, dtype=np.int32).reshape(rows, 2, 2) % 8
+
+
+def test_routed_expert_trace_tracks_multiturn_suffix_and_terminal_gap() -> None:
+    trace = RoutedExpertTrace()
+    trace.record_generation(prompt_token_count=3, generated_token_count=2, routed_experts=routes(4))
+    assert trace.prompt_start == 4
+    trace.record_generation(prompt_token_count=7, generated_token_count=2, routed_experts=routes(4))
+
+    result = trace.finalize(token_count=9, loss_mask=[0, 0, 0, 1, 1, 0, 0, 1, 1])
+    assert result.shape == (9, 2, 2) and result.dtype == np.uint8
+    assert np.array_equal(result[-1, 0], [0, 1])
+
+
+@pytest.mark.parametrize("active", [False, True])
+def test_routed_expert_trace_only_pads_masked_suffix(active: bool) -> None:
+    trace = RoutedExpertTrace()
+    trace.record_generation(prompt_token_count=3, generated_token_count=1, routed_experts=routes(3))
+    mask = [0, 0, 0, 0, int(active)]
+    if active:
+        with pytest.raises(ValueError, match="loss-active target"):
+            trace.finalize(token_count=5, loss_mask=mask)
+    else:
+        result = trace.finalize(token_count=5, loss_mask=mask)
+        assert np.array_equal(result[-2:, 0], [[0, 1], [0, 1]])
diff --git a/tests/backends/skyrl_train/distributed/test_ulysses_token_metadata.py b/tests/backends/skyrl_train/distributed/test_ulysses_token_metadata.py
new file mode 100644
index 0000000000..7b71f77b4b
--- /dev/null
+++ b/tests/backends/skyrl_train/distributed/test_ulysses_token_metadata.py
@@ -0,0 +1,156 @@
+import torch
+
+from skyrl.backends.skyrl_train.distributed.ulysses import utils
+from skyrl.backends.skyrl_train.utils.sample_support_replay import (
+    aligned_sample_support_scores,
+)
+
+
+def test_ulysses_padding_preserves_trailing_metadata_dimensions(monkeypatch):
+    metadata = torch.arange(2 * 5 * 3).reshape(2, 5, 3)
+    rank = 0
+
+    def slice_for_rank(tensor, dim, padding):
+        return tensor.chunk(2, dim=dim)[rank]
+
+    monkeypatch.setattr(utils, "get_ulysses_sequence_parallel_group", lambda: object())
+    monkeypatch.setattr(utils, "slice_input_tensor", slice_for_rank)
+
+    slices = []
+    for rank in range(2):
+        sliced, positions, attention_mask, pad_size = utils.ulysses_pad_and_slice_inputs(
+            metadata,
+            sp_size=2,
+            input_padding_value=-1,
+        )
+        slices.append(sliced)
+
+    assert sliced.shape == (2, 3, 3)
+    expected = torch.cat((metadata, torch.full((2, 1, 3), -1)), dim=1)
+    assert torch.equal(torch.cat(slices, dim=1), expected)
+    assert positions is None
+    assert attention_mask is None
+    assert pad_size == 1
+
+
+def test_ulysses_sample_support_matches_unsharded_values_and_gradients(monkeypatch):
+    rank = 0
+
+    def slice_for_rank(tensor, dim, padding):
+        return tensor.chunk(2, dim=dim)[rank]
+
+    monkeypatch.setattr(utils, "get_ulysses_sequence_parallel_group", lambda: object())
+    monkeypatch.setattr(utils, "slice_input_tensor", slice_for_rank)
+
+    logits = torch.randn(1, 5, 7, dtype=torch.float64, requires_grad=True)
+    sampled_ids = torch.tensor([[2, 3, 4, 5, 6]])
+    support_ids = torch.tensor([[[2, 0], [3, 1], [4, 2], [5, 3], [-1, -1]]], dtype=torch.int32)
+    loss_mask = torch.ones((1, 5), dtype=torch.bool)
+    sharded_logprobs = []
+    for rank in range(2):
+        local_logits, _, _, _ = utils.ulysses_pad_and_slice_inputs(logits, sp_size=2)
+        local_sampled_ids, _, _, _ = utils.ulysses_pad_and_slice_inputs(sampled_ids, sp_size=2)
+        local_support_ids, _, _, _ = utils.ulysses_pad_and_slice_inputs(
+            support_ids,
+            sp_size=2,
+            input_padding_value=-1,
+        )
+        local_loss_mask, _, _, _ = utils.ulysses_pad_and_slice_inputs(loss_mask, sp_size=2)
+        sharded_logprobs.append(
+            aligned_sample_support_scores(
+                local_logits,
+                local_sampled_ids,
+                local_support_ids,
+                local_loss_mask,
+                vocab_start_index=0,
+                vocab_end_index=logits.shape[-1],
+                tp_group=None,
+                inference_only=False,
+                compute_entropy=False,
+                entropy_requires_grad=False,
+            ).logprobs
+        )
+
+    actual = torch.cat(sharded_logprobs, dim=1)[:, : logits.shape[1]]
+    actual.sum().backward()
+    actual_grad = logits.grad.clone()
+
+    reference_logits = logits.detach().clone().requires_grad_(True)
+    expected = aligned_sample_support_scores(
+        reference_logits,
+        sampled_ids,
+        support_ids,
+        loss_mask,
+        vocab_start_index=0,
+        vocab_end_index=logits.shape[-1],
+        tp_group=None,
+        inference_only=False,
+        compute_entropy=False,
+        entropy_requires_grad=False,
+    ).logprobs
+    expected.sum().backward()
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_ulysses_synthetic_eos_uses_fixed_trajectory_capacity(monkeypatch):
+    rank = 0
+
+    def slice_for_rank(tensor, dim, padding):
+        return tensor.chunk(2, dim=dim)[rank]
+
+    monkeypatch.setattr(utils, "get_ulysses_sequence_parallel_group", lambda: object())
+    monkeypatch.setattr(utils, "slice_input_tensor", slice_for_rank)
+
+    logits = torch.randn(1, 6, 9, dtype=torch.float64, requires_grad=True)
+    sampled_ids = torch.tensor([[0, 1, 2, 3, 4, 5]])
+    support_ids = torch.tensor(
+        [[[0, 6], [1, 7], [-1, -1], [3, 8], [4, 0], [-1, -1]]],
+        dtype=torch.int32,
+    )
+    loss_mask = torch.ones((1, 6), dtype=torch.bool)
+    trajectory_ids = torch.tensor([[0, 0, 0, 1, 1, 1]])
+    sharded_logprobs = []
+    for rank in range(2):
+        local = [
+            utils.ulysses_pad_and_slice_inputs(tensor, sp_size=2)[0]
+            for tensor in (logits, sampled_ids, support_ids, loss_mask, trajectory_ids)
+        ]
+        sharded_logprobs.append(
+            aligned_sample_support_scores(
+                *local[:4],
+                vocab_start_index=0,
+                vocab_end_index=logits.shape[-1],
+                tp_group=None,
+                inference_only=False,
+                trajectory_ids=local[4],
+                num_trajectories=2,
+                compute_entropy=False,
+                entropy_requires_grad=False,
+            ).logprobs
+        )
+
+    actual = torch.cat(sharded_logprobs, dim=1)
+    actual.sum().backward()
+    actual_grad = logits.grad.clone()
+
+    reference_logits = logits.detach().clone().requires_grad_(True)
+    expected = aligned_sample_support_scores(
+        reference_logits,
+        sampled_ids,
+        support_ids,
+        loss_mask,
+        vocab_start_index=0,
+        vocab_end_index=logits.shape[-1],
+        tp_group=None,
+        inference_only=False,
+        trajectory_ids=trajectory_ids,
+        num_trajectories=2,
+        compute_entropy=False,
+        entropy_requires_grad=False,
+    ).logprobs
+    expected.sum().backward()
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
diff --git a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py
index ad3372e06f..b76166787e 100644
--- a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py
+++ b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py
@@ -40,6 +40,21 @@ def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch):
     # tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py
 
 
+@pytest.mark.vllm
+def test_sample_support_uses_processed_top_k_logprobs():
+    cfg = SkyRLTrainConfig.from_cli_overrides(
+        [
+            "generator.inference_engine.enable_return_sample_support_set=true",
+            "generator.sampling_params.top_k=8",
+        ]
+    )
+
+    args = build_vllm_cli_args(cfg)
+
+    assert args.max_logprobs == 8
+    assert args.logprobs_mode == "processed_logprobs"
+
+
 def test_resolve_policy_model_name_uses_served_model_name():
     cfg = SkyRLTrainConfig()
     cfg.trainer.policy.model.path = "base-model"
diff --git a/tests/backends/skyrl_train/inference_servers/test_dense_sample_support_wire.py b/tests/backends/skyrl_train/inference_servers/test_dense_sample_support_wire.py
new file mode 100644
index 0000000000..bd127370e2
--- /dev/null
+++ b/tests/backends/skyrl_train/inference_servers/test_dense_sample_support_wire.py
@@ -0,0 +1,60 @@
+import numpy as np
+import orjson
+import pytest
+
+from skyrl.backends.skyrl_train.inference_servers.sample_support_set_wire import (
+    decode_sample_support_set,
+    encode_sample_support_set,
+)
+
+
+def test_packed_round_trip_preserves_dense_support():
+    support = np.array([[7, 152064, -1, -1], [9, 10, 11, -1], [12, -1, -1, -1]], dtype=np.int32)
+
+    packed = encode_sample_support_set(support)
+    restored = decode_sample_support_set(orjson.loads(orjson.dumps(packed)))
+
+    np.testing.assert_array_equal(restored, support)
+    assert restored.dtype == np.int32
+    assert restored.flags.c_contiguous
+
+
+def test_packed_round_trip_preserves_empty_token_dimension():
+    support = np.empty((0, 8), dtype=np.int32)
+
+    packed = encode_sample_support_set(support)
+
+    assert packed["shape"] == [0, 8]
+    np.testing.assert_array_equal(decode_sample_support_set(packed), support)
+
+
+def test_packed_wire_rejects_legacy_list_payload():
+    with pytest.raises(TypeError, match="dictionary"):
+        decode_sample_support_set([[7, 8]])
+
+
+@pytest.mark.parametrize(
+    ("support", "message"),
+    [
+        (np.array([[1.5, 2.5]]), "integer"),
+        (np.array([[1, -1, 2]]), "trailing"),
+        (np.array([[-2, 1]]), "-1 padding"),
+        (np.array([[2**31, -1]], dtype=np.int64), "int32"),
+    ],
+)
+def test_wire_rejects_invalid_dense_support(support, message):
+    with pytest.raises(ValueError, match=message):
+        encode_sample_support_set(support)
+
+
+@pytest.mark.parametrize("shape", [[1], [1, 2, 3], [1.5, 2], ["1", 2], [True, 2], [-1, 2]])
+def test_packed_wire_rejects_invalid_shape(shape):
+    with pytest.raises(ValueError, match="shape"):
+        decode_sample_support_set({"data": "", "shape": shape, "dtype": "int32"})
+
+
+def test_packed_wire_rejects_bad_base64_and_size():
+    with pytest.raises(ValueError, match="valid base64"):
+        decode_sample_support_set({"data": "!", "shape": [1, 1], "dtype": "int32"})
+    with pytest.raises(ValueError, match="byte-size mismatch"):
+        decode_sample_support_set({"data": "", "shape": [1, 1], "dtype": "int32"})
diff --git a/tests/backends/skyrl_train/inference_servers/test_generate_wire.py b/tests/backends/skyrl_train/inference_servers/test_generate_wire.py
index ac2a6219ae..013aadee72 100644
--- a/tests/backends/skyrl_train/inference_servers/test_generate_wire.py
+++ b/tests/backends/skyrl_train/inference_servers/test_generate_wire.py
@@ -12,6 +12,7 @@
 from skyrl.backends.skyrl_train.inference_servers.generate_wire import (
     CLAMPED_LOGPROB,
     build_logprobs_content,
+    clamp_sampled_logprobs,
     decode_packed_routed_experts,
     pack_routed_experts,
 )
@@ -183,3 +184,27 @@ def test_decode_rejects_noncanonical_dtype():
 
     with pytest.raises(ValueError, match="non-canonical dtype"):
         decode_packed_routed_experts(payload)
+
+
+@pytest.mark.parametrize("bad", [float("-inf"), float("inf"), float("nan")])
+def test_clamp_sampled_logprobs_floors_non_finite(bad):
+    content, num_clamped = clamp_sampled_logprobs(np.array([-0.5, bad, -1.25], dtype=np.float32))
+    assert [e["logprob"] for e in content] == [-0.5, CLAMPED_LOGPROB, -1.25]
+    assert num_clamped == 1
+
+
+def test_clamp_sampled_logprobs_leaves_finite_rows_alone():
+    sampled = np.array([-0.5, -1.25, -3.0], dtype=np.float64)
+    content, num_clamped = clamp_sampled_logprobs(sampled)
+    assert [e["logprob"] for e in content] == [-0.5, -1.25, -3.0]
+    assert num_clamped == 0
+
+
+def test_clamp_sampled_logprobs_output_is_json_serializable():
+    # NaN is the case an isneginf screen misses, and orjson would emit it as null.
+    content, _ = clamp_sampled_logprobs(np.array([float("nan"), float("-inf")], dtype=np.float32))
+    assert all(math.isfinite(e["logprob"]) for e in orjson.loads(orjson.dumps(content)))
+
+
+def test_clamp_sampled_logprobs_empty():
+    assert clamp_sampled_logprobs(np.array([], dtype=np.float32)) == ([], 0)
diff --git a/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py b/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py
index 75267428b6..f064a4fcff 100644
--- a/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py
+++ b/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py
@@ -22,8 +22,12 @@
 from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import (
     SKYRL_LORA_ADAPTER_NAME,
     PauseMode,
+    RemoteGenerateClient,
     RemoteInferenceClient,
 )
+from skyrl.backends.skyrl_train.inference_servers.sample_support_set_wire import (
+    encode_sample_support_set,
+)
 from skyrl.backends.skyrl_train.inference_servers.setup import (
     build_new_inference_client,
 )
@@ -35,6 +39,7 @@ def create_mock_vllm_server(server_id: int) -> FastAPI:
     app = FastAPI()
     app.state.last_generate_features = None
     app.state.last_generate_model = None
+    app.state.last_generate_sampling_params = None
     app.state.last_chat_model = None
     app.state.last_completion_model = None
     app.state.last_render_model = None
@@ -63,6 +68,10 @@ async def get_finished():
     async def get_last_generate_features():
         return {"features": app.state.last_generate_features}
 
+    @app.get("/test/last_generate_sampling_params")
+    async def get_last_generate_sampling_params():
+        return app.state.last_generate_sampling_params
+
     @app.get("/test/last_models")
     async def get_last_models():
         return {
@@ -104,6 +113,7 @@ async def completions(request: Request):
     async def generate(request: Request):
         body = await request.json()  # Consume body
         sp = body.get("sampling_params", {})
+        app.state.last_generate_sampling_params = sp
         input_token_ids = body.get("token_ids", [])
         app.state.last_generate_model = body.get("model")
         n = sp.get("n", 1)
@@ -438,8 +448,7 @@ def test_serialization(self, mock_servers):
         assert restored.proxy_url == client.proxy_url
         assert restored.server_urls == client.server_urls
         assert restored.model_name == client.model_name
-        # Session should be None after unpickling
-        assert restored._session is None
+        assert restored._generate_client is None
 
 
 class TestDataPlane:
@@ -480,13 +489,47 @@ async def test_generate_decodes_packed_routed_experts(self, mock_servers):
             enable_return_routed_experts=True,
         )
         try:
-            result = await client.generate({"prompt_token_ids": [[1, 2, 3]]})
+            result = await client.generate({"prompt_token_ids": [[1, 2, 3]], "routed_experts_prompt_starts": [1]})
+            async with httpx.AsyncClient() as http:
+                captured = (await http.get(f"{mock_servers['proxy_url']}/test/last_generate_sampling_params")).json()
         finally:
             await client.teardown()
 
         assert len(result["rollout_expert_indices"]) == 1
         assert result["rollout_expert_indices"][0].dtype == np.uint8
         assert np.array_equal(result["rollout_expert_indices"][0], np.arange(12).reshape(3, 2, 2))
+        assert captured["routed_experts_prompt_start"] == 1
+
+    @pytest.mark.asyncio
+    async def test_external_generate_client_requests_sample_support(self, monkeypatch):
+        generate_client = RemoteGenerateClient(proxy_url="http://unused")
+        captured = {}
+
+        async def fake_post(url, json, headers):
+            captured.update(url=url, json=json, headers=headers)
+            return {
+                "choices": [
+                    {
+                        "token_ids": [7],
+                        "finish_reason": "stop",
+                        "logprobs": {"content": [{"logprob": -0.1}]},
+                        "rollout_sample_support": encode_sample_support_set(np.asarray([[7, 8]], dtype=np.int32)),
+                    }
+                ]
+            }
+
+        monkeypatch.setattr(generate_client, "_post", fake_post)
+        result = await generate_client.generate(
+            prompt_token_ids=[1, 2],
+            sampling_params={},
+            session_id=None,
+            model="default",
+            return_sample_support=True,
+        )
+
+        assert captured["url"].endswith("/skyrl/v1/generate")
+        assert captured["json"]["return_sample_support"] is True
+        assert result.sample_support == [[7, 8]]
 
     @pytest.mark.asyncio
     async def test_generate_rejects_list_routed_experts(self, monkeypatch):
@@ -508,7 +551,7 @@ async def return_list_routes(*args, **kwargs):
                 ]
             }
 
-        monkeypatch.setattr(client, "_post", return_list_routes)
+        monkeypatch.setattr(client._get_generate_client(), "_post", return_list_routes)
         with pytest.raises(ValueError, match="must return packed"):
             await client._generate_single([1], {}, None, "model")
 
@@ -1025,8 +1068,7 @@ async def test_async_context_manager(self, mock_servers):
             result = await client.resume()
             assert len(result) == 2
 
-        # Session should be closed after exiting context
-        assert client._session is None or client._session.closed
+        assert client._generate_client is None or client._generate_client._session is None
 
 
 async def _get_lora_registries(server_urls: List[str]) -> List[Dict[str, str]]:
diff --git a/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py b/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py
new file mode 100644
index 0000000000..feba2d6959
--- /dev/null
+++ b/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py
@@ -0,0 +1,143 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+pytest.importorskip("vllm")
+
+from skyrl.backends.skyrl_train.inference_servers.sample_support_set_wire import (
+    decode_sample_support_set,
+)
+from skyrl.backends.skyrl_train.inference_servers.vllm_server_actor import (
+    VLLMServerActor,
+    _sample_support_from_flat_logprobs,
+)
+
+pytestmark = pytest.mark.vllm
+
+
+def test_flat_logprobs_extracts_sampled_scores_and_support_rows():
+    flat_logprobs = SimpleNamespace(
+        token_ids=[7, 7, 8, 9, 4, 3, 4, 5],
+        logprobs=[-0.1, -0.1, -0.2, -0.3, -0.4, -0.2, -0.4, -0.6],
+    )
+
+    sampled, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=3)
+
+    assert sampled == [{"logprob": -0.1}, {"logprob": -0.4}]
+    np.testing.assert_array_equal(support, [[7, 8, 9], [3, 4, 5]])
+
+
+def test_flat_logprobs_replaces_top_p_masked_candidates():
+    flat_logprobs = SimpleNamespace(
+        token_ids=[7, 7, 8, 9],
+        logprobs=[-0.1, -0.1, -0.2, float("-inf")],
+    )
+
+    _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=3)
+
+    np.testing.assert_array_equal(support, [[7, 8, -1]])
+
+
+def test_flat_logprobs_repairs_sampled_token_absent_from_support():
+    # Three rows, top_k=3 (row_width=4):
+    #   Row A: sampled id (100) absent from a fully-valid support row -> repair.
+    #   Row B: sampled id (7) already present -> unchanged.
+    #   Row C: sampled id (5) absent from a support row that has trailing -1 padding.
+    top_k = 3
+    flat_logprobs = SimpleNamespace(
+        token_ids=[100, 8, 9, 10, 7, 7, 8, 9, 5, 6, 7, 8],
+        logprobs=[
+            -0.1,
+            -0.2,
+            -0.3,
+            -0.4,  # row A: all valid
+            -0.1,
+            -0.1,
+            -0.2,
+            -0.3,  # row B: all valid
+            -0.4,
+            -0.5,
+            -0.6,
+            float("-inf"),  # row C: last col filtered -> padding
+        ],
+    )
+
+    _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=top_k)
+    sampled_ids = [100, 7, 5]
+
+    # (b) each row keeps width == top_k
+    assert all(row.size == top_k for row in support)
+
+    # (a) every row's support now contains its sampled id
+    for sampled_id, row in zip(sampled_ids, support):
+        assert sampled_id in row
+
+    # (c) the sampled id appears exactly once per repaired row (no duplicate)
+    assert np.count_nonzero(support[0] == 100) == 1
+    assert np.count_nonzero(support[2] == 5) == 1
+
+    # (d) trailing -1 padding preserved on the padded row
+    assert support[2][-1] == -1
+
+    # (e) the unaffected row (sampled already present) is unchanged
+    np.testing.assert_array_equal(support[1], [7, 8, 9])
+
+    # Concrete expected repair: weakest (trailing) valid member overwritten.
+    np.testing.assert_array_equal(support[0], [8, 9, 100])
+    np.testing.assert_array_equal(support[2], [6, 5, -1])
+
+
+def test_flat_logprobs_top_k_one_repairs_single_support_column():
+    # top_k == 1 (row_width == 2): a single support column that must hold the sampled id.
+    flat_logprobs = SimpleNamespace(
+        token_ids=[42, 9],
+        logprobs=[-0.1, -0.2],
+    )
+
+    _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=1)
+
+    np.testing.assert_array_equal(support, [[42]])
+
+
+def test_skyrl_generate_returns_packed_sample_support():
+    class FakeEngine:
+        sampling_params = None
+
+        async def generate(self, prompt, sampling_params, request_id):
+            self.sampling_params = sampling_params
+            yield SimpleNamespace(
+                outputs=[
+                    SimpleNamespace(
+                        token_ids=[7],
+                        finish_reason="stop",
+                        logprobs=SimpleNamespace(
+                            token_ids=[7, 7, 8],
+                            logprobs=[-0.1, -0.1, -0.2],
+                        ),
+                        routed_experts=None,
+                    )
+                ]
+            )
+
+    app = FastAPI()
+    engine = FakeEngine()
+    VLLMServerActor._add_custom_endpoints(app, engine, SimpleNamespace(enable_lora=False))
+
+    with TestClient(app) as client:
+        response = client.post(
+            "/skyrl/v1/generate",
+            json={
+                "token_ids": [1, 2],
+                "sampling_params": {"temperature": 1.0, "top_k": 2},
+                "return_sample_support": True,
+            },
+        )
+
+    assert response.status_code == 200
+    assert engine.sampling_params.flat_logprobs is True
+    assert engine.sampling_params.logprobs == 2
+    packed = response.json()["choices"][0]["rollout_sample_support"]
+    np.testing.assert_array_equal(decode_sample_support_set(packed), [[7, 8]])
diff --git a/tests/backends/skyrl_train/test_token_based_batching_utils.py b/tests/backends/skyrl_train/test_token_based_batching_utils.py
index 4fac8a6679..dd9d4710bc 100644
--- a/tests/backends/skyrl_train/test_token_based_batching_utils.py
+++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py
@@ -203,6 +203,7 @@ def test_padding_microbatch_uses_unique_dummy_routes(self):
         batch = self._make_batch([4, 4], num_actions=2)
         batch["rollout_expert_indices"] = torch.full((2, 4, 2, 3), 7, dtype=torch.int16)
         batch["router_padding_mask"] = torch.zeros((2, 4), dtype=torch.bool)
+        batch["sample_support_ids"] = torch.full((2, 4, 8), 7, dtype=torch.int32)
         iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=8)
 
         padding = iterator._create_padding_microbatch()
@@ -210,6 +211,7 @@ def test_padding_microbatch_uses_unique_dummy_routes(self):
         expected = torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(padding["rollout_expert_indices"])
         assert torch.equal(padding["rollout_expert_indices"], expected)
         assert torch.all(padding["router_padding_mask"])
+        assert torch.all(padding["sample_support_ids"] == -1)
 
     def test_multimodal_tensorlist_microbatching(self):
         """Token-based microbatching must gather TensorList fields (multi-modal pixel_values /
diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py
index e562f5ed48..dd6546040d 100644
--- a/tests/backends/skyrl_train/test_train_batch.py
+++ b/tests/backends/skyrl_train/test_train_batch.py
@@ -552,6 +552,7 @@ def test_tensor_batch_none_tensor_list():
     "rollout_logprobs",
     "rollout_expert_indices",
     "router_padding_mask",
+    "sample_support_ids",
     "pixel_values",
     "image_grid_thw",
 }
@@ -578,6 +579,7 @@ def _make_full_training_batch(batch_size: int = 4, seq_len: int = 5) -> Training
         "rollout_logprobs": torch.randn(batch_size, seq_len),
         "rollout_expert_indices": torch.randint(0, 8, (batch_size, seq_len, 2, 3), dtype=torch.long),
         "router_padding_mask": torch.zeros((batch_size, seq_len), dtype=torch.bool),
+        "sample_support_ids": torch.randint(0, 100, (batch_size, seq_len, 4), dtype=torch.int32),
         "pixel_values": TensorList([torch.randn(i + 1, 3) for i in range(batch_size)]),  # batch_size * (i + 1) * 3
         "image_grid_thw": TensorList([torch.tensor([[1, 2, 3]]) for _ in range(batch_size)]),  # batch_size * 1 * 3
     }
@@ -643,11 +645,14 @@ def test_pad_batch_all_fields():
     assert torch.equal(padded["rollout_expert_indices"][:batch_size], batch["rollout_expert_indices"])
     expected_routes = torch.tensor([0, 1, 2]).expand_as(padded["rollout_expert_indices"][batch_size:])
     assert torch.equal(padded["rollout_expert_indices"][batch_size:], expected_routes)
+    assert torch.equal(padded["sample_support_ids"][:batch_size], batch["sample_support_ids"])
+    assert torch.all(padded["sample_support_ids"][batch_size:] == -1)
 
     regular_tensor_keys = EXPECTED_TRAINING_INPUT_FIELDS - {
         "loss_mask",
         "rollout_expert_indices",
         "router_padding_mask",
+        "sample_support_ids",
         "pixel_values",
         "image_grid_thw",
     }
diff --git a/tests/backends/skyrl_train/utils/test_sample_support_replay.py b/tests/backends/skyrl_train/utils/test_sample_support_replay.py
new file mode 100644
index 0000000000..09afb00bac
--- /dev/null
+++ b/tests/backends/skyrl_train/utils/test_sample_support_replay.py
@@ -0,0 +1,353 @@
+import sys
+import types
+
+import pytest
+import torch
+
+from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import (
+    TokenMetadataLayout,
+)
+from skyrl.backends.skyrl_train.utils.sample_support_replay import (
+    aligned_sample_support_scores,
+    sample_support_scores,
+    synthetic_eos_logprobs,
+)
+
+
+def _reference(logits, sampled_ids, support_ids):
+    outputs = []
+    for row_logits, sampled_id, support in zip(
+        logits.reshape(-1, logits.shape[-1]),
+        sampled_ids.reshape(-1),
+        support_ids.reshape(-1, support_ids.shape[-1]),
+        strict=True,
+    ):
+        members = support[support >= 0].long()
+        outputs.append(
+            row_logits.new_zeros(())
+            if members.numel() == 0
+            else row_logits[sampled_id] - torch.logsumexp(row_logits[members], dim=0)
+        )
+    return torch.stack(outputs).reshape(sampled_ids.shape)
+
+
+def _reference_entropy(logits, support_ids):
+    outputs = []
+    for row_logits, support in zip(
+        logits.reshape(-1, logits.shape[-1]),
+        support_ids.reshape(-1, support_ids.shape[-1]),
+        strict=True,
+    ):
+        members = support[support >= 0].long()
+        if members.numel() == 0:
+            outputs.append(row_logits.new_zeros(()))
+        else:
+            member_logprobs = torch.log_softmax(row_logits[members], dim=0)
+            outputs.append(-(member_logprobs.exp() * member_logprobs).sum())
+    return torch.stack(outputs).reshape(support_ids.shape[:-1])
+
+
+def test_support_logprobs_match_reference_values_and_gradients():
+    logits = torch.randn(2, 3, 11, dtype=torch.float64, requires_grad=True)
+    sampled_ids = torch.tensor([[2, 5, 1], [8, 3, 7]])
+    support_ids = torch.tensor(
+        [
+            [[2, 4, 6, -1], [5, -1, -1, -1], [-1, -1, -1, -1]],
+            [[8, 0, 9, 4], [3, 2, -1, -1], [7, 1, 5, -1]],
+        ],
+        dtype=torch.int32,
+    )
+
+    scores = sample_support_scores(
+        logits,
+        sampled_ids,
+        support_ids,
+        vocab_start_index=0,
+        vocab_end_index=logits.shape[-1],
+        tp_group=None,
+        compute_entropy=False,
+        entropy_requires_grad=False,
+    )
+    expected = _reference(logits, sampled_ids, support_ids)
+
+    assert scores.valid_mask.tolist() == [[True, True, False], [True, True, True]]
+    torch.testing.assert_close(scores.logprobs, expected)
+    scores.logprobs.sum().backward()
+    actual_grad = logits.grad.clone()
+
+    reference_logits = logits.detach().clone().requires_grad_(True)
+    _reference(reference_logits, sampled_ids, support_ids).sum().backward()
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_support_entropy_matches_reference_values_and_gradients():
+    logits = torch.randn(2, 3, 11, dtype=torch.float64, requires_grad=True)
+    sampled_ids = torch.tensor([[2, 5, 1], [8, 3, 7]])
+    support_ids = torch.tensor(
+        [
+            [[2, 4, 6, -1], [5, -1, -1, -1], [-1, -1, -1, -1]],
+            [[8, 0, 9, 4], [3, 2, -1, -1], [7, 1, 5, -1]],
+        ],
+        dtype=torch.int32,
+    )
+
+    scores = sample_support_scores(
+        logits,
+        sampled_ids,
+        support_ids,
+        vocab_start_index=0,
+        vocab_end_index=logits.shape[-1],
+        tp_group=None,
+        compute_entropy=True,
+        entropy_requires_grad=True,
+    )
+    expected_logprobs = _reference(logits, sampled_ids, support_ids)
+    expected_entropy = _reference_entropy(logits, support_ids)
+    assert scores.entropy is not None
+    torch.testing.assert_close(scores.logprobs, expected_logprobs)
+    torch.testing.assert_close(scores.entropy, expected_entropy)
+    assert scores.valid_mask.tolist() == [[True, True, False], [True, True, True]]
+
+    (scores.logprobs + scores.entropy).sum().backward()
+    actual_grad = logits.grad.clone()
+    reference_logits = logits.detach().clone().requires_grad_(True)
+    (
+        _reference(reference_logits, sampled_ids, support_ids) + _reference_entropy(reference_logits, support_ids)
+    ).sum().backward()
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_support_entropy_metric_is_detached():
+    logits = torch.randn(1, 2, 7, dtype=torch.float64, requires_grad=True)
+    scores = sample_support_scores(
+        logits,
+        torch.tensor([[2, 5]]),
+        torch.tensor([[[2, 4], [5, 1]]], dtype=torch.int32),
+        vocab_start_index=0,
+        vocab_end_index=logits.shape[-1],
+        tp_group=None,
+        compute_entropy=True,
+        entropy_requires_grad=False,
+    )
+
+    assert scores.entropy is not None
+    assert not scores.entropy.requires_grad
+
+
+def test_entropy_gradients_require_entropy_computation():
+    with pytest.raises(ValueError, match="compute_entropy=True"):
+        sample_support_scores(
+            torch.randn(1, 5),
+            torch.tensor([1]),
+            torch.tensor([[1, 2]], dtype=torch.int32),
+            vocab_start_index=0,
+            vocab_end_index=5,
+            tp_group=None,
+            compute_entropy=False,
+            entropy_requires_grad=True,
+        )
+
+
+def test_support_entropy_excludes_synthetic_eos():
+    logits = torch.tensor([[[1.0, 0.0, 2.0, -1.0, 0.5], [0.0, 1.0, -1.0, 2.0, 0.5]]])
+    sampled_ids = torch.tensor([[2, 4]])
+    support = torch.tensor([[[2, 1], [-1, -1]]], dtype=torch.int32)
+
+    scores = aligned_sample_support_scores(
+        logits,
+        sampled_ids,
+        support,
+        loss_mask=torch.ones((1, 2), dtype=torch.bool),
+        vocab_start_index=0,
+        vocab_end_index=logits.shape[-1],
+        tp_group=None,
+        inference_only=False,
+        compute_entropy=True,
+        entropy_requires_grad=False,
+    )
+
+    assert scores.entropy is not None
+    assert torch.isfinite(scores.logprobs).all()
+    torch.testing.assert_close(scores.entropy, _reference_entropy(logits, support))
+    assert scores.valid_mask.tolist() == [[True, False]]
+
+
+def test_fused_selected_projection_matches_explicit_logits_with_pair_chunking():
+    temperature = 0.7
+    hidden = torch.randn(2, 3, 5, dtype=torch.float64, requires_grad=True)
+    weight = torch.randn(9, 5, dtype=torch.float64, requires_grad=True)
+    sampled_ids = torch.tensor([[1, 4, 7], [2, 5, 8]])
+    support_ids = torch.tensor(
+        [
+            [[1, 0, 3], [4, 6, -1], [7, -1, -1]],
+            [[2, 1, 8], [5, 4, -1], [8, 0, 6]],
+        ],
+        dtype=torch.int32,
+    )
+
+    fused = sample_support_scores(
+        hidden,
+        sampled_ids,
+        support_ids,
+        vocab_start_index=0,
+        vocab_end_index=weight.shape[0],
+        tp_group=None,
+        lm_head_weight=weight,
+        temperature=temperature,
+        chunk_size=4,
+        compute_entropy=False,
+        entropy_requires_grad=False,
+    )
+    fused.logprobs.sum().backward()
+    fused_hidden_grad = hidden.grad.clone()
+    fused_weight_grad = weight.grad.clone()
+
+    explicit_hidden = hidden.detach().clone().requires_grad_(True)
+    explicit_weight = weight.detach().clone().requires_grad_(True)
+    explicit_logits = (explicit_hidden @ explicit_weight.T) / temperature
+    explicit = sample_support_scores(
+        explicit_logits,
+        sampled_ids,
+        support_ids,
+        vocab_start_index=0,
+        vocab_end_index=weight.shape[0],
+        tp_group=None,
+        compute_entropy=False,
+        entropy_requires_grad=False,
+    )
+    explicit.logprobs.sum().backward()
+
+    torch.testing.assert_close(fused.logprobs, explicit.logprobs, check_dtype=False)
+    torch.testing.assert_close(fused_hidden_grad, explicit_hidden.grad, rtol=1e-5, atol=1e-6)
+    torch.testing.assert_close(fused_weight_grad, explicit_weight.grad, rtol=1e-5, atol=1e-6)
+
+
+def test_support_ids_must_be_int32():
+    with pytest.raises(ValueError, match="int32"):
+        sample_support_scores(
+            torch.randn(1, 5),
+            torch.tensor([1]),
+            torch.tensor([[1, 2]], dtype=torch.int64),
+            vocab_start_index=0,
+            vocab_end_index=5,
+            tp_group=None,
+            compute_entropy=False,
+            entropy_requires_grad=False,
+        )
+
+
+def _install_fake_distributed_logprob(monkeypatch, calls):
+    model_utils = types.ModuleType("skyrl.backends.skyrl_train.distributed.megatron.model_utils")
+
+    class DistributedLogprob:
+        @staticmethod
+        def apply(source, targets, *args):
+            calls.append(source.shape)
+            return source.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
+
+    model_utils.DistributedLogprob = DistributedLogprob
+    monkeypatch.setitem(sys.modules, model_utils.__name__, model_utils)
+    return model_utils
+
+
+def test_synthetic_eos_uses_one_fixed_slot_per_unpacked_trajectory(monkeypatch):
+    calls = []
+    _install_fake_distributed_logprob(monkeypatch, calls)
+    logits = torch.arange(3 * 4 * 5, dtype=torch.float64).reshape(3, 4, 5).requires_grad_(True)
+    sampled_ids = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4], [2, 3, 4, 0]])
+    synthetic_eos_mask = torch.tensor(
+        [[False, False, True, False], [False, False, False, False], [False, True, False, False]]
+    )
+
+    actual = synthetic_eos_logprobs(
+        logits,
+        sampled_ids,
+        synthetic_eos_mask,
+        vocab_start_index=0,
+        vocab_end_index=5,
+        tp_group=object(),
+        inference_only=False,
+    )
+
+    expected = torch.zeros_like(actual)
+    expected[0, 2] = logits.detach()[0, 2, 2]
+    expected[2, 1] = logits.detach()[2, 1, 3]
+    torch.testing.assert_close(actual, expected)
+    assert calls == [torch.Size([1, 3, 5])]
+
+    actual.sum().backward()
+    expected_grad = torch.zeros_like(logits)
+    expected_grad[0, 2, 2] = 1
+    expected_grad[2, 1, 3] = 1
+    torch.testing.assert_close(logits.grad, expected_grad)
+
+
+def test_synthetic_eos_uses_packed_cp_trajectory_segments(monkeypatch):
+    calls = []
+    _install_fake_distributed_logprob(monkeypatch, calls)
+    logits = torch.arange(4 * 5, dtype=torch.float64).reshape(1, 4, 5).requires_grad_(True)
+    sampled_ids = torch.tensor([[0, 1, 2, 3]])
+    synthetic_eos_mask = torch.tensor([[False, True, False, True]])
+    layout = TokenMetadataLayout(
+        attention_mask=torch.ones((2, 3), dtype=torch.bool),
+        sequence_lengths=[3, 3],
+        aligned_sequence_length=8,
+        padded_sequence_lengths=[4, 4],
+        cu_seqlens_padded=torch.tensor([0, 4, 8], dtype=torch.int32),
+        context_parallel_size=2,
+        context_parallel_rank=0,
+    )
+
+    actual = synthetic_eos_logprobs(
+        logits,
+        sampled_ids,
+        synthetic_eos_mask,
+        vocab_start_index=0,
+        vocab_end_index=5,
+        tp_group=object(),
+        inference_only=False,
+        metadata_layout=layout,
+    )
+
+    expected = torch.zeros_like(actual)
+    expected[0, 1] = logits.detach()[0, 1, 1]
+    expected[0, 3] = logits.detach()[0, 3, 3]
+    torch.testing.assert_close(actual, expected)
+    assert calls == [torch.Size([1, 2, 5])]
+
+
+def test_synthetic_eos_fused_projection_keeps_capacity_and_chunk_bound(monkeypatch):
+    calls = []
+    model_utils = _install_fake_distributed_logprob(monkeypatch, calls)
+
+    def fused_apply(backend, hidden, weight, targets, start, end, chunk_size, group, inference_only):
+        calls.append((hidden.shape, chunk_size))
+        return hidden[..., 0]
+
+    model_utils._fused_lm_head_logprob_apply = fused_apply
+    hidden = torch.arange(3 * 4 * 2, dtype=torch.float64).reshape(3, 4, 2).requires_grad_(True)
+    sampled_ids = torch.zeros((3, 4), dtype=torch.long)
+    synthetic_eos_mask = torch.tensor(
+        [[False, False, False, False], [False, True, False, False], [False, False, False, False]]
+    )
+
+    actual = synthetic_eos_logprobs(
+        hidden,
+        sampled_ids,
+        synthetic_eos_mask,
+        vocab_start_index=0,
+        vocab_end_index=5,
+        tp_group=object(),
+        inference_only=False,
+        lm_head_weight=torch.ones((5, 2), dtype=torch.float64),
+        chunk_size=2,
+    )
+
+    expected = torch.zeros_like(actual)
+    expected[1, 1] = hidden.detach()[1, 1, 0]
+    torch.testing.assert_close(actual, expected)
+    assert calls == [(torch.Size([1, 3, 2]), 2)]
+    actual.sum().backward()
+    expected_grad = torch.zeros_like(hidden)
+    expected_grad[1, 1, 0] = 1
+    torch.testing.assert_close(hidden.grad, expected_grad)
diff --git a/tests/backends/skyrl_train/utils/test_torch_utils.py b/tests/backends/skyrl_train/utils/test_torch_utils.py
index 13ea7ded6d..b18b1fdc24 100644
--- a/tests/backends/skyrl_train/utils/test_torch_utils.py
+++ b/tests/backends/skyrl_train/utils/test_torch_utils.py
@@ -6,6 +6,7 @@
 from skyrl.backends.skyrl_train.utils.torch_utils import (
     chunked_cross_entropy_from_log_probs,
     chunked_entropy_from_logits,
+    logprobs_from_logits,
 )
 
 
@@ -35,6 +36,27 @@ def test_chunked_cross_entropy_from_logprobs():
     assert torch.allclose(result_BS[1, 0], torch.tensor(expected_uniform_entropy), atol=1e-4)
 
 
+def test_logprobs_from_logits_uses_torch_path_for_cpu(monkeypatch):
+    def fail_if_called(*args, **kwargs):
+        raise AssertionError("flash-attn cross entropy requires CUDA tensors")
+
+    monkeypatch.setattr(
+        "skyrl.backends.skyrl_train.utils.torch_utils.FLASH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE",
+        True,
+    )
+    monkeypatch.setattr(
+        "skyrl.backends.skyrl_train.utils.torch_utils.logprobs_from_logits_flash_attn",
+        fail_if_called,
+    )
+    logits = torch.tensor([[[1.0, 2.0, 3.0]]])
+    labels = torch.tensor([[2]])
+
+    actual = logprobs_from_logits(logits, labels)
+
+    expected = torch.log_softmax(logits, dim=-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1)
+    torch.testing.assert_close(actual, expected)
+
+
 def test_chunked_entropy_from_logits():
     # Define a small log-probability tensor (batch_size=2, seqlen=3, vocab_size=4)
     logits = [
diff --git a/tests/backends/skyrl_train/workers/test_fsdp_dense_sample_support.py b/tests/backends/skyrl_train/workers/test_fsdp_dense_sample_support.py
new file mode 100644
index 0000000000..b995b623df
--- /dev/null
+++ b/tests/backends/skyrl_train/workers/test_fsdp_dense_sample_support.py
@@ -0,0 +1,354 @@
+import importlib
+import importlib.machinery
+import sys
+import types
+
+import pytest
+import torch
+from torch import nn
+
+
+@pytest.fixture
+def model_wrapper(monkeypatch):
+    def unpad_input(tensor, attention_mask):
+        indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+        unpadded = tensor.flatten(0, 1).index_select(0, indices)
+        return unpadded, indices, None, None, None
+
+    def pad_input(tensor, indices, batch, seqlen):
+        padded = tensor.new_zeros((batch * seqlen, *tensor.shape[1:]))
+        return padded.index_copy(0, indices, tensor).reshape(batch, seqlen, *tensor.shape[1:])
+
+    flash_attn = types.ModuleType("flash_attn")
+    flash_attn.__spec__ = importlib.machinery.ModuleSpec("flash_attn", loader=None)
+    bert_padding = types.ModuleType("flash_attn.bert_padding")
+    bert_padding.__spec__ = importlib.machinery.ModuleSpec("flash_attn.bert_padding", loader=None)
+    bert_padding.pad_input = pad_input
+    bert_padding.unpad_input = unpad_input
+    flash_attn.bert_padding = bert_padding
+
+    peft = types.ModuleType("peft")
+    peft.LoraConfig = object
+    peft.TaskType = types.SimpleNamespace(CAUSAL_LM="CAUSAL_LM")
+    peft.get_peft_model = lambda model, config: model
+    peft_tuners = types.ModuleType("peft.tuners")
+    peft_lora = types.ModuleType("peft.tuners.lora")
+    peft_lora.LoraLayer = nn.Module
+    peft_tuners.lora = peft_lora
+    peft.tuners = peft_tuners
+
+    ulysses = types.ModuleType("skyrl.backends.skyrl_train.distributed.ulysses")
+    ulysses_utils = types.ModuleType("skyrl.backends.skyrl_train.distributed.ulysses.utils")
+    ulysses_utils.gather_outputs_and_unpad = None
+    ulysses_utils.ulysses_pad_and_slice_inputs = None
+    ulysses.utils = ulysses_utils
+
+    stubs = {
+        "flash_attn": flash_attn,
+        "flash_attn.bert_padding": bert_padding,
+        "peft": peft,
+        "peft.tuners": peft_tuners,
+        "peft.tuners.lora": peft_lora,
+        "skyrl.backends.skyrl_train.distributed.ulysses": ulysses,
+        "skyrl.backends.skyrl_train.distributed.ulysses.utils": ulysses_utils,
+    }
+    for name, module in stubs.items():
+        monkeypatch.setitem(sys.modules, name, module)
+
+    module_name = "skyrl.backends.skyrl_train.workers.model_wrapper"
+    monkeypatch.delitem(sys.modules, module_name, raising=False)
+    module = importlib.import_module(module_name)
+    yield module
+    sys.modules.pop(module_name, None)
+
+
+class _FakeCausalLM(nn.Module):
+    def __init__(self, sequence_length: int, vocab_size: int):
+        super().__init__()
+        self.logits = nn.Parameter(torch.randn(sequence_length, vocab_size, dtype=torch.float64))
+
+    def forward(self, input_ids, **kwargs):
+        return {"logits": (self.logits * 1.0).unsqueeze(0).expand(input_ids.shape[0], -1, -1)}
+
+
+def _reference(logits, sequences, support):
+    values = []
+    for position in range(sequences.shape[1] - 1):
+        sampled = sequences[0, position + 1]
+        members = support[0, position + 1]
+        members = members[members >= 0].long()
+        values.append(logits[position, sampled] - torch.logsumexp(logits[position, members], dim=0))
+    return torch.stack(values).unsqueeze(0)
+
+
+def test_fsdp_forward_matches_dense_reference_values_and_gradients(model_wrapper):
+    sequences = torch.tensor([[1, 2, 3, 4]])
+    attention_mask = torch.ones_like(sequences)
+    support = torch.tensor(
+        [[[-1, -1, -1], [2, 5, -1], [3, -1, -1], [4, 0, 6]]],
+        dtype=torch.int32,
+    )
+    model = _FakeCausalLM(sequence_length=4, vocab_size=7)
+    wrapper = model_wrapper.HFModelWrapper(model, bf16=False)
+
+    actual = wrapper(
+        sequences,
+        num_actions=3,
+        attention_mask=attention_mask,
+        sample_support_ids=support,
+        loss_mask=torch.ones((1, 3), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    )
+    actual.sum().backward()
+    actual_grad = model.logits.grad.clone()
+
+    reference_logits = model.logits.detach().clone().requires_grad_(True)
+    expected = _reference(reference_logits, sequences, support)
+    expected.sum().backward()
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_fsdp_dense_support_runs_optimizer_step(model_wrapper):
+    sequences = torch.tensor([[1, 2, 3]])
+    support = torch.tensor([[[-1, -1], [2, 4], [3, 1]]], dtype=torch.int32)
+    model = _FakeCausalLM(sequence_length=3, vocab_size=5)
+    wrapper = model_wrapper.HFModelWrapper(model, bf16=False)
+    optimizer = torch.optim.SGD(wrapper.parameters(), lr=0.1)
+    before = model.logits.detach().clone()
+
+    loss = -wrapper(
+        sequences,
+        num_actions=2,
+        attention_mask=torch.ones_like(sequences),
+        sample_support_ids=support,
+        loss_mask=torch.ones((1, 2), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    ).mean()
+    loss.backward()
+    optimizer.step()
+
+    assert torch.isfinite(loss)
+    assert not torch.equal(model.logits, before)
+
+
+def test_fsdp_dense_support_skips_full_sequence_vocabulary_logprobs(monkeypatch, model_wrapper):
+    # The fixed-capacity EOS fallback may score B rows; the ordinary B*S path must stay disabled.
+    def fail_if_called(*args, **kwargs):
+        raise AssertionError("full-sequence vocabulary logprobs should not run during support replay")
+
+    monkeypatch.setattr(model_wrapper, "logprobs_from_logits", fail_if_called)
+    sequences = torch.tensor([[1, 2, 3]])
+    support = torch.tensor([[[-1, -1], [2, 4], [3, 1]]], dtype=torch.int32)
+    wrapper = model_wrapper.HFModelWrapper(_FakeCausalLM(sequence_length=3, vocab_size=5), bf16=False)
+
+    actual = wrapper(
+        sequences,
+        num_actions=2,
+        attention_mask=torch.ones_like(sequences),
+        sample_support_ids=support,
+        loss_mask=torch.ones((1, 2), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    )
+
+    assert torch.isfinite(actual).all()
+
+
+def test_fsdp_replay_entropy_skips_full_vocabulary_entropy(monkeypatch, model_wrapper):
+    sequences = torch.tensor([[1, 2, 3, 4]])
+    support = torch.tensor(
+        [[[-1, -1, -1], [2, 5, -1], [3, -1, -1], [4, 0, 6]]],
+        dtype=torch.int32,
+    )
+    model = _FakeCausalLM(sequence_length=4, vocab_size=7)
+    wrapper = model_wrapper.HFModelWrapper(model, bf16=False)
+
+    def fail_if_called(*args, **kwargs):
+        raise AssertionError("full-vocabulary entropy should not run during support replay")
+
+    monkeypatch.setattr(wrapper, "chunked_entropy_from_logits_fn", fail_if_called)
+    _, output = wrapper(
+        sequences,
+        num_actions=3,
+        attention_mask=torch.ones_like(sequences),
+        return_output=True,
+        compute_entropy=True,
+        entropy_requires_grad=True,
+        sample_support_ids=support,
+        loss_mask=torch.ones((1, 3), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    )
+
+    expected = []
+    for position in range(3):
+        members = support[0, position + 1]
+        members = members[members >= 0].long()
+        member_logprobs = torch.log_softmax(model.logits[position, members], dim=0)
+        expected.append(-(member_logprobs.exp() * member_logprobs).sum())
+    expected = torch.stack(expected).unsqueeze(0)
+
+    torch.testing.assert_close(output["entropy"][:, :-1], expected)
+    assert output["entropy_mask"][:, :-1].all()
+    assert output["entropy"].requires_grad
+
+
+def test_fsdp_synthetic_eos_uses_full_vocabulary_logprob(model_wrapper):
+    sequences = torch.tensor([[1, 2, 3]])
+    support = torch.tensor([[[-1, -1], [2, 4], [-1, -1]]], dtype=torch.int32)
+    model = _FakeCausalLM(sequence_length=3, vocab_size=5)
+    wrapper = model_wrapper.HFModelWrapper(model, bf16=False)
+
+    actual = wrapper(
+        sequences,
+        num_actions=2,
+        attention_mask=torch.ones_like(sequences),
+        sample_support_ids=support,
+        loss_mask=torch.ones((1, 2), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    )
+    actual.sum().backward()
+    actual_grad = model.logits.grad.clone()
+
+    reference_logits = model.logits.detach().clone().requires_grad_(True)
+    supported = reference_logits[0, 2] - torch.logsumexp(reference_logits[0, [2, 4]], dim=0)
+    eos = reference_logits[1, 3] - torch.logsumexp(reference_logits[1], dim=0)
+    expected = torch.stack((supported, eos)).unsqueeze(0)
+    expected.sum().backward()
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_fsdp_packed_microbatch_matches_dense_reference(model_wrapper):
+    sequences = torch.tensor([[1, 2, 3, 4], [0, 5, 6, 7]])
+    attention_mask = torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]])
+    support = torch.full((2, 4, 2), -1, dtype=torch.int32)
+    support[0, 2:] = torch.tensor([[3, 8], [4, 0]], dtype=torch.int32)
+    support[1, 2:] = torch.tensor([[6, 1], [7, 2]], dtype=torch.int32)
+    model = _FakeCausalLM(sequence_length=7, vocab_size=9)
+    wrapper = model_wrapper.HFModelWrapper(
+        model,
+        use_flash_attention_2=True,
+        bf16=False,
+        remove_microbatch_padding=True,
+    )
+
+    actual = wrapper(
+        sequences,
+        num_actions=2,
+        attention_mask=attention_mask,
+        sample_support_ids=support,
+        loss_mask=torch.ones((2, 2), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    )
+    actual.sum().backward()
+    actual_grad = model.logits.grad.clone()
+
+    reference_logits = model.logits.detach().clone().requires_grad_(True)
+    expected = torch.stack(
+        (
+            torch.stack(
+                (
+                    reference_logits[1, 3] - torch.logsumexp(reference_logits[1, [3, 8]], dim=0),
+                    reference_logits[2, 4] - torch.logsumexp(reference_logits[2, [4, 0]], dim=0),
+                )
+            ),
+            torch.stack(
+                (
+                    reference_logits[4, 6] - torch.logsumexp(reference_logits[4, [6, 1]], dim=0),
+                    reference_logits[5, 7] - torch.logsumexp(reference_logits[5, [7, 2]], dim=0),
+                )
+            ),
+        )
+    )
+    expected.sum().backward()
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_fsdp_packed_microbatch_supports_one_synthetic_eos_per_trajectory(model_wrapper):
+    sequences = torch.tensor([[1, 2, 3, 4], [0, 5, 6, 7]])
+    attention_mask = torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]])
+    support = torch.full((2, 4, 2), -1, dtype=torch.int32)
+    support[0, 2] = torch.tensor([3, 8], dtype=torch.int32)
+    support[1, 2] = torch.tensor([6, 1], dtype=torch.int32)
+    model = _FakeCausalLM(sequence_length=7, vocab_size=9)
+    wrapper = model_wrapper.HFModelWrapper(
+        model,
+        use_flash_attention_2=True,
+        bf16=False,
+        remove_microbatch_padding=True,
+    )
+
+    actual = wrapper(
+        sequences,
+        num_actions=2,
+        attention_mask=attention_mask,
+        sample_support_ids=support,
+        loss_mask=torch.ones((2, 2), dtype=torch.bool),
+        enable_sample_support_replay=True,
+    )
+    actual.sum().backward()
+    actual_grad = model.logits.grad.clone()
+
+    reference_logits = model.logits.detach().clone().requires_grad_(True)
+    expected = torch.stack(
+        (
+            torch.stack(
+                (
+                    reference_logits[1, 3] - torch.logsumexp(reference_logits[1, [3, 8]], dim=0),
+                    reference_logits[2, 4] - torch.logsumexp(reference_logits[2], dim=0),
+                )
+            ),
+            torch.stack(
+                (
+                    reference_logits[4, 6] - torch.logsumexp(reference_logits[4, [6, 1]], dim=0),
+                    reference_logits[5, 7] - torch.logsumexp(reference_logits[5], dim=0),
+                )
+            ),
+        )
+    )
+    expected.sum().backward()
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad, reference_logits.grad)
+
+
+def test_fsdp_capture_only_matches_feature_disabled(model_wrapper):
+    sequences = torch.tensor([[1, 2, 3]])
+    attention_mask = torch.ones_like(sequences)
+    support = torch.tensor([[[-1, -1], [2, 4], [3, 1]]], dtype=torch.int32)
+    wrapper = model_wrapper.HFModelWrapper(_FakeCausalLM(sequence_length=3, vocab_size=5), bf16=False)
+
+    expected = wrapper(sequences, num_actions=2, attention_mask=attention_mask)
+    actual = wrapper(
+        sequences,
+        num_actions=2,
+        attention_mask=attention_mask,
+        sample_support_ids=support,
+        enable_sample_support_replay=False,
+    )
+
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"loss_mask": torch.ones((1, 2), dtype=torch.bool)}, "no recorded support"),
+        ({"sample_support_ids": torch.full((1, 3, 2), -1, dtype=torch.int32)}, "no loss mask"),
+    ],
+)
+def test_fsdp_replay_requires_support_and_loss_mask(model_wrapper, kwargs, message):
+    wrapper = model_wrapper.HFModelWrapper(_FakeCausalLM(sequence_length=3, vocab_size=5), bf16=False)
+
+    with pytest.raises(ValueError, match=message):
+        wrapper(
+            torch.tensor([[1, 2, 3]]),
+            num_actions=2,
+            attention_mask=torch.ones((1, 3), dtype=torch.long),
+            enable_sample_support_replay=True,
+            **kwargs,
+        )
diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py
index f1e33fa17d..46af0332b0 100644
--- a/tests/tinker/test_engine.py
+++ b/tests/tinker/test_engine.py
@@ -340,6 +340,8 @@ def processor(group):
     assert processed_models == ["model_a", "model_b"]
     with Session(engine.db_engine) as session:
         assert all(f.status == RequestStatus.COMPLETED for f in session.exec(select(FutureDB)).all())
+
+
 def forward_backward_payload() -> dict:
     return types.ForwardBackwardInput(
         data=[
diff --git a/tests/train/dataset/test_sample_support_preprocess.py b/tests/train/dataset/test_sample_support_preprocess.py
new file mode 100644
index 0000000000..7a97a9a0a2
--- /dev/null
+++ b/tests/train/dataset/test_sample_support_preprocess.py
@@ -0,0 +1,69 @@
+import pytest
+import torch
+
+from skyrl.train.dataset.preprocess import build_dense_sample_support
+
+
+def test_sample_support_is_response_aligned_int32():
+    support = build_dense_sample_support(
+        [[[10, 11, -1], [12, -1, -1]], [[20, 21, 22]]],
+        [[10, 12], [20]],
+        [[1, 1], [1]],
+        sequence_length=5,
+        top_k=3,
+        eos_token_id=2,
+    )
+
+    assert support is not None
+    assert support.dtype == torch.int32
+    assert support[0].tolist() == [[-1, -1, -1]] * 3 + [[10, 11, -1], [12, -1, -1]]
+    assert support[1].tolist() == [[-1, -1, -1]] * 4 + [[20, 21, 22]]
+
+
+def test_empty_loss_masked_rows_and_loss_bearing_synthetic_eos_are_preserved():
+    support = build_dense_sample_support(
+        [[[7, 8], [], []]],
+        [[7, 9, 2]],
+        [[1, 0, 1]],
+        sequence_length=4,
+        top_k=2,
+        eos_token_id=2,
+    )
+
+    assert support.tolist() == [[[-1, -1], [7, 8], [-1, -1], [-1, -1]]]
+
+
+def test_empty_loss_bearing_non_eos_is_rejected():
+    with pytest.raises(ValueError, match="loss-bearing non-EOS"):
+        build_dense_sample_support(
+            [[[]]],
+            [[3]],
+            [[1]],
+            sequence_length=2,
+            top_k=2,
+            eos_token_id=2,
+        )
+
+
+def test_multiple_loss_bearing_unsupported_eos_are_rejected():
+    with pytest.raises(ValueError, match="more than one loss-bearing unsupported token"):
+        build_dense_sample_support(
+            [[[], []]],
+            [[2, 2]],
+            [[1, 1]],
+            sequence_length=2,
+            top_k=2,
+            eos_token_id=2,
+        )
+
+
+def test_loss_bearing_sampled_token_must_be_in_support():
+    with pytest.raises(ValueError, match="sampled token 3 is missing"):
+        build_dense_sample_support(
+            [[[2, 4]]],
+            [[3]],
+            [[1]],
+            sequence_length=2,
+            top_k=2,
+            eos_token_id=2,
+        )
diff --git a/tests/train/generators/test_datatypes.py b/tests/train/generators/test_datatypes.py
index 12c9efdba2..580cbd97f5 100644
--- a/tests/train/generators/test_datatypes.py
+++ b/tests/train/generators/test_datatypes.py
@@ -31,7 +31,6 @@ def test_turn_output(output_ids, observation_ids, output_logprobs, added_eos, ex
         output_logprobs=output_logprobs,
         new_obs=[],
         obs_ids=observation_ids,
-        rollout_expert_indices=None,
         added_eos=added_eos,
         reward=1.0,
     )
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index e3ee7a1f5c..0ad7f2a695 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -30,6 +30,7 @@ def test_generator_output_concatenation():
         "rollout_metrics",
         "rollout_logprobs",
         "rollout_expert_indices",
+        "rollout_sample_support",
         # optional but present in the signature
         "trajectory_ids",
         "trajectory_generation_times",
@@ -238,6 +239,7 @@ def test_case1_response_only_assistant(self):
             "rollout_logprobs": [[-0.5], [-0.3, -0.4]],
             "trajectory_ids": [tid, tid],
             "rollout_expert_indices": None,
+            "rollout_sample_support": [[[20, 21]], [[40, 42], [41, 43]]],
             "is_last_step": [False, True],
         }
 
@@ -251,6 +253,7 @@ def test_case1_response_only_assistant(self):
         assert merged["loss_masks"] == [[1, 0, 1, 1]]
         # logprobs: A1=-0.5, O2=0.0, A2_tok1=-0.3, A2_tok2=-0.4
         assert merged["rollout_logprobs"] == [[-0.5, 0.0, -0.3, -0.4]]
+        assert merged["rollout_sample_support"] == [[[20, 21], [], [40, 42], [41, 43]]]
         # rewards: A1=1.0, O2=0.0, A2_tok1=0.0, A2_tok2=5.0
         assert merged["rewards"] == [[1.0, 0.0, 0.0, 5.0]]
         assert merged["stop_reasons"] == ["eos"]
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 6d03d937b6..fb618d28a2 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -23,20 +23,22 @@
 MOCK_TOKENIZER_ENCODED_IDS = [1, 2, 3, 4]
 
 
-def test_turn_output_keeps_uncaptured_suffix_out_of_routes():
-    routes = np.asarray([[[2, 3]], [[4, 5]]], dtype=np.uint8)
+def test_turn_output_masks_uncaptured_suffix():
     output = TurnOutput(
         output="answer",
         output_ids=[10, 11, 4],
         output_logprobs=None,
         new_obs=[],
         obs_ids=[20, 21],
-        rollout_expert_indices=routes,
         reward=1.0,
+        rollout_sample_support=np.array([[10, 100], [11, 101]], dtype=np.int32),
         added_eos=True,
     )
 
-    assert output.get_turn_rollout_expert_indices() is routes
+    np.testing.assert_array_equal(
+        output.get_turn_rollout_sample_support(),
+        np.array([[10, 100], [11, 101], [-1, -1], [-1, -1], [-1, -1]], dtype=np.int32),
+    )
     assert output.get_turn_loss_mask() == [1, 1, 0, 0, 0]
 
 
@@ -413,6 +415,73 @@ def mock_generate(_, model=None):
     assert output.stop_reason == "stop"
 
 
+@pytest.mark.asyncio
+@patch("skyrl_gym.make")
+async def test_agent_loop_uses_incremental_replay_metadata_traces(
+    mock_make,
+    mock_tokenizer,
+    mock_llm,
+    mock_env,
+    generator_cfg,
+    mock_env_cfg,
+):
+    generator_cfg.batched = False
+    generator_cfg.max_turns = 2
+    generator_cfg.use_conversation_multi_turn = True
+    generator_cfg.inference_engine.enable_return_routed_experts = True
+    generator_cfg.inference_engine.enable_return_sample_support_set = True
+    generator_cfg.sampling_params.top_k = 2
+    mock_make.return_value = mock_env
+    mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {})
+
+    mock_env.step.side_effect = [
+        BaseTextEnvStepOutput(observations=[{"role": "user", "content": "next"}], reward=1.0, done=done, metadata={})
+        for done in (False, True)
+    ]
+    prompt_starts = []
+    generation_index = 0
+
+    def generate(input_batch, model=None):
+        nonlocal generation_index
+        prompt_tokens = input_batch["prompt_token_ids"][0]
+        prompt_start = input_batch["routed_experts_prompt_starts"][0]
+        prompt_starts.append(prompt_start)
+        output_ids = [10, 11]
+        num_route_rows = len(prompt_tokens) - prompt_start + len(output_ids) - 1
+        routes = np.arange(num_route_rows * 4, dtype=np.int32).reshape(num_route_rows, 2, 2) % 8
+        sample_support = [[10, 100 + generation_index], [11, 110 + generation_index]]
+        generation_index += 1
+        return {
+            "responses": ["mocked output"],
+            "response_ids": [output_ids],
+            "stop_reasons": ["stop"],
+            "rollout_expert_indices": [routes],
+            "rollout_sample_support": [sample_support],
+        }
+
+    mock_llm.generate = AsyncMock(side_effect=generate)
+    generator = SkyRLGymGenerator(
+        generator_cfg=generator_cfg,
+        skyrl_gym_cfg=mock_env_cfg,
+        inference_engine_client=mock_llm,
+        tokenizer=mock_tokenizer,
+    )
+    generator.base_conversation_token_ids = []
+
+    output = await generator.agent_loop(
+        [{"role": "user", "content": "Start"}],
+        mock_env_cfg.env_class,
+        {},
+        max_tokens=32,
+        max_input_length=64,
+    )
+
+    assert prompt_starts == [0, 5]
+    assert output.rollout_sample_support[:2] == [[10, 100], [11, 110]]
+    assert output.rollout_sample_support[-2:] == [[10, 101], [11, 111]]
+    assert all(row == [-1, -1] for row in output.rollout_sample_support[2:-2])
+
+
 @pytest.mark.asyncio
 @patch("skyrl_gym.make")
 async def test_generate_batched(mock_make, mock_tokenizer, mock_llm, mock_env, generator_cfg, mock_env_cfg):
diff --git a/tests/train/test_config.py b/tests/train/test_config.py
index 5afc726c43..f2eadcb510 100644
--- a/tests/train/test_config.py
+++ b/tests/train/test_config.py
@@ -154,6 +154,76 @@ def test_trainer_config_rejects_invalid_vocab_entropy_chunking(field_name, value
         TrainerConfig(**{field_name: value})
 
 
+@pytest.mark.parametrize("strategy", ["megatron", "fsdp"])
+def test_sample_support_replay_requires_capture(strategy):
+    with pytest.raises(ValueError, match="enable_sample_support_replay requires"):
+        SkyRLTrainConfig.from_cli_overrides(
+            [
+                f"trainer.strategy={strategy}",
+                "trainer.algorithm.enable_sample_support_replay=true",
+            ]
+        )
+
+
+@pytest.mark.parametrize(
+    ("override", "message"),
+    [
+        ("generator.sampling_params.temperature=0", "temperature > 0"),
+        ("generator.sampling_params.top_k=1", "top_k > 1"),
+        ("generator.sampling_params.repetition_penalty=1.1", "repetition_penalty=1.0"),
+        ("generator.sampling_params.additional_kwargs.foo=bar", "additional_kwargs"),
+    ],
+)
+def test_sample_support_capture_rejects_unsupported_sampling_modifiers(override, message):
+    with pytest.raises(ValueError, match=message):
+        SkyRLTrainConfig.from_cli_overrides(
+            [
+                "generator.inference_engine.enable_return_sample_support_set=true",
+                "generator.sampling_params.top_k=8",
+                override,
+            ]
+        )
+
+
+@pytest.mark.parametrize("strategy", ["megatron", "fsdp"])
+def test_sample_support_replay_accepts_top_k_top_p_and_min_p(strategy):
+    cfg = SkyRLTrainConfig.from_cli_overrides(
+        [
+            "trainer.algorithm.enable_sample_support_replay=true",
+            f"trainer.strategy={strategy}",
+            "generator.inference_engine.enable_return_sample_support_set=true",
+            "generator.sampling_params.top_k=8",
+            "generator.sampling_params.top_p=0.9",
+            "generator.sampling_params.min_p=0.05",
+        ]
+    )
+
+    assert cfg.trainer.algorithm.enable_sample_support_replay
+
+
+def test_sample_support_replay_rejects_unimplemented_strategy():
+    with pytest.raises(ValueError, match="requires trainer.strategy=megatron or fsdp"):
+        SkyRLTrainConfig.from_cli_overrides(
+            [
+                "trainer.algorithm.enable_sample_support_replay=true",
+                "trainer.strategy=jax",
+                "generator.inference_engine.enable_return_sample_support_set=true",
+                "generator.sampling_params.top_k=17",
+            ]
+        )
+
+
+def test_sample_support_capture_rejects_vlm():
+    with pytest.raises(ValueError, match="vision_language_generator"):
+        SkyRLTrainConfig.from_cli_overrides(
+            [
+                "generator.inference_engine.enable_return_sample_support_set=true",
+                "generator.sampling_params.top_k=8",
+                "generator.vision_language_generator=true",
+            ]
+        )
+
+
 def test_cli_overrides_plus_prefix_rejected():
     with pytest.raises(ValueError, match="The '\\+' prefix"):
         SkyRLTrainConfig.from_cli_overrides(["+new_field=value"])
diff --git a/tests/train/test_trainer.py b/tests/train/test_trainer.py
index eb70dcc1f3..6f8d575b10 100644
--- a/tests/train/test_trainer.py
+++ b/tests/train/test_trainer.py
@@ -114,6 +114,37 @@ def _get_test_data(trainer: RayPPOTrainer):
     return data
 
 
+def test_fwd_logprobs_preserves_sample_support_and_loss_mask(dummy_config):
+    dummy_config.trainer.critic.model.path = None
+    trainer = object.__new__(RayPPOTrainer)
+    trainer.cfg = dummy_config
+    trainer.ref_model = None
+    trainer.dispatch = MagicMock()
+    trainer.all_metrics = {}
+    trainer._skip_policy_forward = MagicMock(return_value=False)
+    seen = {}
+
+    def execute_forward_pass(model, batch, **kwargs):
+        seen.update(batch)
+        return torch.zeros((2, 3))
+
+    trainer._execute_forward_pass = execute_forward_pass
+    batch = TrainingInputBatch(
+        {
+            "sequences": torch.ones((2, 4), dtype=torch.long),
+            "attention_mask": torch.ones((2, 4), dtype=torch.long),
+            "loss_mask": torch.ones((2, 3)),
+            "sample_support_ids": torch.tensor([[[1, -1]] * 4, [[2, -1]] * 4], dtype=torch.int32),
+        }
+    )
+    batch.metadata = {"response_length": 3}
+
+    trainer.fwd_logprobs_values_reward(batch)
+
+    assert torch.equal(seen["sample_support_ids"], batch["sample_support_ids"])
+    assert torch.equal(seen["loss_mask"], batch["loss_mask"])
+
+
 def test_calculate_kl_create_experience_batched(dummy_config):
     trainer = RayPPOTrainer(
         cfg=dummy_config,
diff --git a/tests/train/test_trainer_utils.py b/tests/train/test_trainer_utils.py
index 4fa401f1cf..2c80650574 100644
--- a/tests/train/test_trainer_utils.py
+++ b/tests/train/test_trainer_utils.py
@@ -395,6 +395,7 @@ def test_handle_replace_sampling_sufficient_good_samples():
         "rollout_metrics": None,
         "rollout_logprobs": [[0.1, 0.2], [0.3, 0.4], [0.5, 0.25], [0.15, 0.25], [0.1, 0.2], [0.3, 0.4]],
         "rollout_expert_indices": [np.asarray([[[i, i + 1]]], dtype=np.uint8) for i in range(6)],
+        "rollout_sample_support": [[[i, i + 10], [i + 1, i + 11]] for i in range(6)],
     }
     uids = ["uid1", "uid1", "uid2", "uid2", "uid3", "uid3"]  # 2 samples per prompt
     sampling_config = {"n_samples_per_prompt": 2, "min_replace_ratio": 0.3}
@@ -416,6 +417,11 @@ def test_handle_replace_sampling_sufficient_good_samples():
     }
     for response, routes in zip(result_output["response_ids"], result_output["rollout_expert_indices"]):
         assert np.array_equal(routes, route_by_response[tuple(response)])
+    support_by_response = dict(
+        zip(map(tuple, generator_output["response_ids"]), generator_output["rollout_sample_support"])
+    )
+    for response, support in zip(result_output["response_ids"], result_output["rollout_sample_support"]):
+        assert support == support_by_response[tuple(response)]
 
     # Check that bad uid2 samples were replaced with good samples
     uid2_indices = [i for i, uid in enumerate(result_uids) if uid == "uid2"]
@@ -654,6 +660,7 @@ def test_filter_generator_output():
         "rollout_metrics": {"metric": "value"},
         "rollout_logprobs": [[0.16, 0.4], [0.1, 0.2], [0.3, 0.4]],
         "rollout_expert_indices": routes,
+        "rollout_sample_support": [[[7, 70], [8, 80]], [[9, 90], [10, 100]], [[11, 110], [12, 120]]],
     }
     kept_indices = [0, 2]  # Keep first and third samples
 
@@ -668,6 +675,7 @@ def test_filter_generator_output():
     assert filtered["rollout_logprobs"] == [[0.16, 0.4], [0.3, 0.4]]
     assert filtered["rollout_expert_indices"][0] is routes[0]
     assert filtered["rollout_expert_indices"][1] is routes[2]
+    assert filtered["rollout_sample_support"] == [[[7, 70], [8, 80]], [[11, 110], [12, 120]]]
 
 
 def test_zero_variance_filter_mixed_groups():