perf(sample-support): reuse replay scores for entropy

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)

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)

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):

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