feat(v1): optional sampler logprobs on SamplingMask

tests/unit/test_graph.py

@@ -190,6 +190,82 @@ def test_routed_experts_none_when_absent():
     assert trace.branches[-1].routed_experts is None
 
 
+def _two_turn_sampling_mask_trace(masks: list[vf.SamplingMask]) -> vf.Trace:
+    trace = vf.Trace(
+        agent=vf.AgentInfo(config=vf.AgentConfig()),
+        task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="x")),
+    )
+    user = vf.UserMessage(content="u1")
+    turns = [
+        ([user], [10, 11, 12], [20, 21], [(0, 2)]),
+        (
+            [user, vf.AssistantMessage(content="a1"), vf.UserMessage(content="u2")],
+            [10, 11, 12, 20, 21, 30, 31],
+            [40, 41],
+            [(0, 2), None, (5, 7)],
+        ),
+    ]
+    for i, ((messages, prompt_ids, completion_ids, spans), mask) in enumerate(
+        zip(turns, masks)
+    ):
+        graph.prepare_turn(trace, messages).commit(
+            vf.Response(
+                id=str(i),
+                created=0,
+                model="t",
+                message=vf.AssistantMessage(content=f"a{i + 1}"),
+                finish_reason="stop",
+                tokens=TurnTokens(
+                    prompt_ids=prompt_ids,
+                    completion_ids=completion_ids,
+                    message_spans=spans,
+                    sampling_mask=mask,
+                ),
+            )
+        )
+    return trace
+
+
+def _mask(
+    rows: list[list[int]], logprobs: list[float] | None = None
+) -> vf.SamplingMask:
+    mask = vf.SamplingMask.from_sampling_mask(rows)
+    if logprobs is not None:
+        mask.logprobs = np.array(logprobs, np.float32)
+    return mask
+
+
+def test_sampling_mask_logprobs_attributed_aggregated_and_round_tripped():
+    """Optional sampler logprobs ride along with the mask ids: concatenated on the branch,
+    kept by the wire round trip, and dropped (None) when any turn lacks them. A mask whose
+    logprobs do not match its ids is not attributed."""
+    first = _mask([[20, 5], [21]], [-0.1, -2.4, 0.0])
+    second = _mask([[40], [41, 7, 8]], [0.0, -0.5, -1.0, -3.0])
+    trace = _two_turn_sampling_mask_trace([first, second])
+
+    branch_mask = trace.branches[-1].sampling_mask
+    assert branch_mask is not None
+    assert len(branch_mask.counts) == len(trace.branches[-1].token_ids)
+    assert branch_mask.ids.tolist() == [20, 5, 21, 40, 41, 7, 8]
+    assert branch_mask.logprobs.tolist() == pytest.approx(
+        [-0.1, -2.4, 0.0, 0.0, -0.5, -1.0, -3.0]
+    )
+    restored = type(trace).model_validate(trace.model_dump()).branches[-1].sampling_mask
+    assert restored is not None
+    for name in ("ids", "counts", "logprobs"):
+        np.testing.assert_array_equal(
+            getattr(restored, name), getattr(branch_mask, name)
+        )
+
+    mixed = _two_turn_sampling_mask_trace([first, _mask([[40], [41, 7, 8]])])
+    assert mixed.branches[-1].sampling_mask.logprobs is None
+    assert mixed.branches[-1].sampling_mask.ids.tolist() == branch_mask.ids.tolist()
+
+    mismatched = _mask([[40], [41, 7, 8]], [0.0, -0.5])
+    trace = _two_turn_sampling_mask_trace([first, mismatched])
+    assert trace.nodes[-1].sampling_mask is None
+
+
 def test_tool_call_hash_matches_v0_content_and_arguments_normalization():
     left = vf.AssistantMessage(
         content=None,

tests/unit/test_trace.py

@@ -250,19 +250,21 @@ def test_sampling_mask_from_packed_matches_lists():
 
     from verifiers.v1.types import SamplingMask
 
-    rows = [[4, 7], [], [9]]
-    ids, counts = np.array([4, 7, 9], np.int32), np.array([2, 0, 1], np.int32)
-    packed = {
-        name: {
+    def pack(a):
+        return {
             "data": base64.b64encode(a.tobytes()).decode(),
             "shape": list(a.shape),
-            "dtype": "int32",
+            "dtype": a.dtype.name,
         }
-        for name, a in (("ids", ids), ("counts", counts))
-    }
+
+    rows = [[4, 7], [], [9]]
+    ids, counts = np.array([4, 7, 9], np.int32), np.array([2, 0, 1], np.int32)
+    logprobs = np.array([-0.5, -1.0, 0.0], np.float32)
+    packed = {"ids": pack(ids), "counts": pack(counts)}
     for mask in (
         SamplingMask.from_sampling_mask(rows),
-        SamplingMask.from_sampling_mask(packed),
+        SamplingMask.from_sampling_mask(packed, pack(logprobs)),
     ):
         np.testing.assert_array_equal(mask.ids, ids)
         np.testing.assert_array_equal(mask.counts, counts)
+    np.testing.assert_array_equal(mask.logprobs, logprobs)

verifiers/v1/clients/train.py

@@ -178,7 +178,9 @@ def response_from_generate(
             mm_placeholders=mm_placeholders,
             mm_token_type_id_map=mm_token_type_id_map,
             routed_experts=result.get("routed_experts"),
-            sampling_mask=SamplingMask.from_sampling_mask(mask)
+            sampling_mask=SamplingMask.from_sampling_mask(
+                mask, result.get("sampling_mask_logprobs")
+            )
             if (mask := result.get("sampling_mask"))
             else None,
         ),

verifiers/v1/graph.py

@@ -209,8 +209,9 @@ def serialize_sampling_mask(self, mask: SamplingMask | None) -> dict | None:
         if mask is None:
             return None
         return {
-            "ids": _encode_ndarray(mask.ids),
-            "counts": _encode_ndarray(mask.counts),
+            name: _encode_ndarray(array)
+            for name, array in vars(mask).items()
+            if array is not None
         }
 
     @field_validator("sampling_mask", mode="before")
@@ -219,10 +220,7 @@ def deserialize_sampling_mask(cls, value: Any) -> SamplingMask | None:
         if value is None or isinstance(value, SamplingMask):
             return value
         if isinstance(value, dict):
-            return SamplingMask(
-                ids=_decode_ndarray(value["ids"]),
-                counts=_decode_ndarray(value["counts"]),
-            )
+            return SamplingMask(**{k: _decode_ndarray(v) for k, v in value.items()})
         raise TypeError(f"cannot build SamplingMask from {type(value).__name__}")
 
 
@@ -699,8 +697,10 @@ def _attribute_sampling_mask(
     if payload is None:
         return
     node = trace.nodes[assistant_id]
-    if len(payload.counts) != sum(node.mask) or int(payload.counts.sum()) != len(
-        payload.ids
+    if (
+        len(payload.counts) != sum(node.mask)
+        or int(payload.counts.sum()) != len(payload.ids)
+        or (payload.logprobs is not None and len(payload.logprobs) != len(payload.ids))
     ):
         return
     node.sampling_mask = payload

verifiers/v1/trace.py

@@ -333,19 +333,29 @@ def sampling_mask(self) -> SamplingMask | None:
             return None
         # Attribution validates each mask against the node's sampled positions.
         ids_parts: list[np.ndarray] = []
+        logprobs_parts: list[np.ndarray | None] = []
         counts_parts: list[np.ndarray] = []
         for node in self.nodes:
             counts = np.zeros(len(node.mask), dtype=np.int32)
             if node.sampling_mask is not None and len(node.sampling_mask.counts):
                 counts[np.nonzero(node.mask)[0]] = node.sampling_mask.counts
                 ids_parts.append(node.sampling_mask.ids)
+                logprobs_parts.append(node.sampling_mask.logprobs)
             counts_parts.append(counts)
         ids = (
             np.concatenate(ids_parts).astype(np.int32, copy=False)
             if ids_parts
             else np.zeros(0, dtype=np.int32)
         )
-        return SamplingMask(ids=ids, counts=np.concatenate(counts_parts))
+        # Logprobs only when every mask on the path carries them.
+        logprobs = (
+            np.concatenate([np.zeros(0, np.float32), *logprobs_parts])
+            if all(part is not None for part in logprobs_parts)
+            else None
+        )
+        return SamplingMask(
+            ids=ids, counts=np.concatenate(counts_parts), logprobs=logprobs
+        )
 
     @property
     def usage(self) -> Usage | None:

verifiers/v1/types.py

@@ -202,22 +202,30 @@ class SamplingMask:
     """Sampling masks stored as flat int32 `ids` and `counts` arrays.
 
     Each row contains the token ids that survived sampling filters for one completion
-    token. Row boundaries are recovered from `counts`.
+    token. Row boundaries are recovered from `counts`. The optional float32 `logprobs`,
+    parallel to `ids`, are the sampler's renormalized logprobs of those ids.
     """
 
     ids: Any
     counts: Any
+    logprobs: Any = None
 
     @classmethod
     def from_sampling_mask(
-        cls, sampling_mask: list[list[int]] | dict[str, Any]
+        cls,
+        sampling_mask: list[list[int]] | dict[str, Any],
+        logprobs: dict[str, Any] | None = None,
     ) -> "SamplingMask":
         """Build from vLLM's one-list-per-token masks, or from prime-rl's packed CSR
-        form: `{"ids", "counts"}`, each a base64 `{data, shape, dtype}` int32 array."""
+        form: `{"ids", "counts"}`, each a base64 `{data, shape, dtype}` int32 array.
+        `logprobs` is packed float32 parallel to the ids (only prime-rl sends it)."""
+        if logprobs is not None:
+            logprobs = _decode_packed(logprobs, np.float32)
         if isinstance(sampling_mask, dict):
             return cls(
                 ids=_decode_packed(sampling_mask["ids"], np.int32),
                 counts=_decode_packed(sampling_mask["counts"], np.int32),
+                logprobs=logprobs,
             )
         counts = np.fromiter(
             (len(row) for row in sampling_mask),
@@ -229,7 +237,7 @@ def from_sampling_mask(
             if int(counts.sum())
             else np.empty(0, dtype=np.int32)
         )
-        return cls(ids=ids, counts=counts)
+        return cls(ids=ids, counts=counts, logprobs=logprobs)
 
 
 def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: