feat(trace): carry sampler probability heads to training branches

verifiers/v1/__init__.py

@@ -138,6 +138,7 @@
     Tool,
     ToolCall,
     ToolMessage,
+    TopLogprobs,
     TurnTokens,
     Usage,
     UserMessage,
@@ -260,6 +261,7 @@
     "Branch",
     "TurnTokens",
     "SamplingMask",
+    "TopLogprobs",
     "Timing",
     "TimeSpan",
     "TimeSplit",

verifiers/v1/clients/train.py

@@ -28,6 +28,7 @@
     SamplingMask,
     Tool,
     ToolCall,
+    TopLogprobs,
     TurnTokens,
     Usage,
     WeightVersionSpan,
@@ -185,6 +186,13 @@ def response_from_generate(
             weight_versions=[WeightVersionSpan(**span) for span in spans]
             if (spans := result.get("weight_versions"))
             else None,
+            top_logprobs=(
+                TopLogprobs.from_top_logprobs(
+                    result["completion_top_ids"], result["completion_top_logprobs"]
+                )
+                if result.get("completion_top_ids") is not None
+                else None
+            ),
         ),
     )
 

verifiers/v1/graph.py

@@ -24,7 +24,7 @@
 import hashlib
 import json
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -33,6 +33,7 @@
     ConfigDict,
     Field,
     FieldSerializationInfo,
+    ValidationInfo,
     field_serializer,
     field_validator,
 )
@@ -49,6 +50,7 @@
     TextContentPart,
     Tool,
     ToolMessage,
+    TopLogprobs,
     WeightVersionSpan,
 )
 
@@ -170,6 +172,9 @@ class MessageNode(BaseModel):
     completion, as vLLM returns them). Assistant nodes only; None when the engine labels
     none."""
 
+    top_logprobs: SkipJsonSchema[TopLogprobs | None] = None
+    """Sampler top-k probabilities for sampled tokens; unnormalized over the head."""
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @property
@@ -210,26 +215,28 @@ def deserialize_ndarray_field(cls, value: Any) -> np.ndarray | None:
             return _decode_ndarray(value)
         raise TypeError(f"cannot build ndarray field from {type(value).__name__}")
 
-    @field_serializer("sampling_mask")
-    def serialize_sampling_mask(self, mask: SamplingMask | None) -> dict | None:
-        if mask is None:
+    @field_serializer("sampling_mask", "top_logprobs")
+    def serialize_sampling_arrays(
+        self, payload: SamplingMask | TopLogprobs | None
+    ) -> dict | None:
+        if payload is None:
             return None
         return {
-            "ids": _encode_ndarray(mask.ids),
-            "counts": _encode_ndarray(mask.counts),
+            field.name: _encode_ndarray(getattr(payload, field.name))
+            for field in fields(payload)
         }
 
-    @field_validator("sampling_mask", mode="before")
+    @field_validator("sampling_mask", "top_logprobs", mode="before")
     @classmethod
-    def deserialize_sampling_mask(cls, value: Any) -> SamplingMask | None:
-        if value is None or isinstance(value, SamplingMask):
+    def deserialize_sampling_arrays(
+        cls, value: Any, info: ValidationInfo
+    ) -> SamplingMask | TopLogprobs | None:
+        kind = SamplingMask if info.field_name == "sampling_mask" else TopLogprobs
+        if value is None or isinstance(value, kind):
             return value
         if isinstance(value, dict):
-            return SamplingMask(
-                ids=_decode_ndarray(value["ids"]),
-                counts=_decode_ndarray(value["counts"]),
-            )
-        raise TypeError(f"cannot build SamplingMask from {type(value).__name__}")
+            return kind(**{key: _decode_ndarray(array) for key, array in value.items()})
+        raise TypeError(f"cannot build {kind.__name__} from {type(value).__name__}")
 
 
 def _canonical_tool_arguments(arguments: str) -> str:
@@ -823,6 +830,23 @@ def _project_prompt_attribution(
     return projected_spans, projected_is_content
 
 
+def _attribute_top_logprobs(
+    trace: Trace, assistant_id: int, payload: TopLogprobs | None
+) -> None:
+    """Attach a completion-aligned top-k sampling head to the assistant node."""
+    if payload is None:
+        return
+    node = trace.nodes[assistant_id]
+    if (
+        len(payload.counts) != sum(node.mask)
+        or (payload.counts < 0).any()
+        or int(payload.counts.sum()) != len(payload.ids)
+        or len(payload.ids) != len(payload.logprobs)
+    ):
+        raise ValueError("Top logprobs must align with every sampled node token")
+    node.top_logprobs = payload
+
+
 def _commit_turn(turn: PendingTurn, response: Response) -> int:
     trace = turn.trace
     prompt = turn.prompt
@@ -1020,6 +1044,9 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int:
     _attribute_sampling_mask(
         trace, assistant_id, tokens.sampling_mask if tokens else None
     )
+    _attribute_top_logprobs(
+        trace, assistant_id, tokens.top_logprobs if tokens else None
+    )
 
     return assistant_id
 

verifiers/v1/trace.py

@@ -31,6 +31,7 @@
     SamplingMask,
     Tool,
     ToolMessage,
+    TopLogprobs,
     Usage,
     content_text,
 )
@@ -44,6 +45,7 @@
         "__all__": {
             "routed_experts",
             "sampling_mask",
+            "top_logprobs",
         }
     }
 }
@@ -347,6 +349,24 @@ def sampling_mask(self) -> SamplingMask | None:
         )
         return SamplingMask(ids=ids, counts=np.concatenate(counts_parts))
 
+    @property
+    def top_logprobs(self) -> TopLogprobs | None:
+        """Top-k sampling heads aligned to this branch's token ids."""
+        heads = [n.top_logprobs for n in self.nodes if n.top_logprobs is not None]
+        if not heads:
+            return None
+        counts = []
+        for node in self.nodes:
+            rows = np.zeros(len(node.mask), dtype=np.int32)
+            if node.top_logprobs is not None:
+                rows[np.nonzero(node.mask)[0]] = node.top_logprobs.counts
+            counts.append(rows)
+        return TopLogprobs(
+            ids=np.concatenate([head.ids for head in heads]),
+            logprobs=np.concatenate([head.logprobs for head in heads]),
+            counts=np.concatenate(counts),
+        )
+
     @property
     def usage(self) -> Usage | None:
         return Usage.aggregate(c.usage for c in self.calls if c.usage is not None)

verifiers/v1/types.py

@@ -249,6 +249,41 @@ def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray:
     return array.astype(dtype, copy=False)
 
 
+@dataclass
+class TopLogprobs:
+    """Top-k sampling head stored as flat `ids`, `logprobs` and `counts` arrays.
+
+    Each row holds one completion token's top-k sampler candidates in descending
+    probability order, independent of the sampled action. Row boundaries come from
+    `counts`; `ids` and `logprobs` share the same flat layout.
+    """
+
+    ids: Any
+    logprobs: Any
+    counts: Any
+
+    @classmethod
+    def from_top_logprobs(
+        cls, ids: list[list[int]], logprobs: list[list[float]]
+    ) -> "TopLogprobs":
+        if len(ids) != len(logprobs):
+            raise ValueError("top ids and logprobs must have one row per token")
+        counts = np.fromiter(
+            (len(row) for row in ids),
+            dtype=np.int32,
+            count=len(ids),
+        )
+        if any(len(ids[i]) != len(logprobs[i]) for i in range(len(ids))):
+            raise ValueError("top ids and logprobs rows must be the same length")
+        return cls(
+            ids=np.fromiter((token for row in ids for token in row), dtype=np.int32),
+            logprobs=np.fromiter(
+                (p for row in logprobs for p in row), dtype=np.float32
+            ),
+            counts=counts,
+        )
+
+
 class TurnTokens(BaseModel):
     """Training tokens from renderer tokenization or provider-returned token IDs."""
 
@@ -288,6 +323,10 @@ class TurnTokens(BaseModel):
     # attributed to the assistant node by the turn's `commit`, then dropped.
     weight_versions: list[WeightVersionSpan] | None = Field(default=None, exclude=True)
 
+    # Transient carrier (excluded): per-completion-token top-k sampling heads,
+    # attributed to the assistant node by the turn's `commit`, then dropped.
+    top_logprobs: TopLogprobs | None = Field(default=None, exclude=True)
+
 
 class Response(BaseModel):
     id: str