From dc361d862ef902b8f1a322405feee43a334dff97 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:51:01 +0000 Subject: [PATCH 1/4] perf(v1): accept packed CSR sampling masks from prime-rl's /generate --- tests/v1/test_trace.py | 25 +++++++++++++++++++++++++ verifiers/v1/types.py | 16 +++++++++++++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/v1/test_trace.py b/tests/v1/test_trace.py index fe31f422fd..fa429b16e7 100644 --- a/tests/v1/test_trace.py +++ b/tests/v1/test_trace.py @@ -767,3 +767,28 @@ def test_semantic_edge_set_accepts_deep_acyclic_chain(): ) assert len(edge_set.edges) == 2_000 + + +def test_sampling_mask_from_packed_matches_lists(): + import base64 + + import numpy as np + + 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: { + "data": base64.b64encode(a.tobytes()).decode(), + "shape": list(a.shape), + "dtype": "int32", + } + for name, a in (("ids", ids), ("counts", counts)) + } + for mask in ( + SamplingMask.from_sampling_mask(rows), + SamplingMask.from_sampling_mask(packed), + ): + np.testing.assert_array_equal(mask.ids, ids) + np.testing.assert_array_equal(mask.counts, counts) diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 79a01de60a..926592e53c 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -1,3 +1,4 @@ +import base64 from collections.abc import Iterable from dataclasses import dataclass from typing import Annotated, Any, Literal @@ -208,7 +209,20 @@ class SamplingMask: counts: Any @classmethod - def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "SamplingMask": + def from_sampling_mask( + cls, sampling_mask: list[list[int]] | dict[str, Any] + ) -> "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.""" + if isinstance(sampling_mask, dict): + ids, counts = ( + np.frombuffer( + base64.b64decode(sampling_mask[key]["data"]), + dtype=sampling_mask[key]["dtype"], + ).astype(np.int32, copy=False) + for key in ("ids", "counts") + ) + return cls(ids=ids, counts=counts) counts = np.fromiter( (len(row) for row in sampling_mask), dtype=np.int32, From 71ad510b5fe15580c46c50345eb92dd258361058 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:59:18 +0000 Subject: [PATCH 2/4] feat(v1): optional sampler logprobs on SamplingMask --- tests/v1/test_graph.py | 76 +++++++++++++++++++++++++++++++++++ tests/v1/test_trace.py | 18 +++++---- verifiers/v1/clients/train.py | 4 +- verifiers/v1/graph.py | 16 ++++---- verifiers/v1/trace.py | 12 +++++- verifiers/v1/types.py | 31 +++++++++----- 6 files changed, 128 insertions(+), 29 deletions(-) diff --git a/tests/v1/test_graph.py b/tests/v1/test_graph.py index a4d8a25779..0ca30da55e 100644 --- a/tests/v1/test_graph.py +++ b/tests/v1/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, diff --git a/tests/v1/test_trace.py b/tests/v1/test_trace.py index fa429b16e7..69099990f7 100644 --- a/tests/v1/test_trace.py +++ b/tests/v1/test_trace.py @@ -776,19 +776,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) diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 60275cb16c..c58e8f561a 100644 --- a/verifiers/v1/clients/train.py +++ b/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, ), diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 270f087d3e..4302fc03f2 100644 --- a/verifiers/v1/graph.py +++ b/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 diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index c3b4f6c58d..280ad2b0a6 100644 --- a/verifiers/v1/trace.py +++ b/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: diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 926592e53c..20a441f8f6 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -202,27 +202,31 @@ 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): - ids, counts = ( - np.frombuffer( - base64.b64decode(sampling_mask[key]["data"]), - dtype=sampling_mask[key]["dtype"], - ).astype(np.int32, copy=False) - for key in ("ids", "counts") + return cls( + ids=_decode_packed(sampling_mask["ids"], np.int32), + counts=_decode_packed(sampling_mask["counts"], np.int32), + logprobs=logprobs, ) - return cls(ids=ids, counts=counts) counts = np.fromiter( (len(row) for row in sampling_mask), dtype=np.int32, @@ -233,7 +237,12 @@ 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: + array = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) + return array.astype(dtype, copy=False) class TurnTokens(BaseModel): From ec0bbed9e609a7c63581d0694920d71045a30c42 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:48:00 +0000 Subject: [PATCH 3/4] refactor(v1): one _decode_packed helper for packed sampling-mask arrays --- verifiers/v1/types.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 926592e53c..0e15ad9f8b 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -215,14 +215,10 @@ def from_sampling_mask( """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.""" if isinstance(sampling_mask, dict): - ids, counts = ( - np.frombuffer( - base64.b64decode(sampling_mask[key]["data"]), - dtype=sampling_mask[key]["dtype"], - ).astype(np.int32, copy=False) - for key in ("ids", "counts") + return cls( + ids=_decode_packed(sampling_mask["ids"], np.int32), + counts=_decode_packed(sampling_mask["counts"], np.int32), ) - return cls(ids=ids, counts=counts) counts = np.fromiter( (len(row) for row in sampling_mask), dtype=np.int32, @@ -236,6 +232,11 @@ def from_sampling_mask( return cls(ids=ids, counts=counts) +def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: + array = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) + return array.astype(dtype, copy=False) + + class TurnTokens(BaseModel): """Training tokens from renderer tokenization or provider-returned token IDs.""" From 8a90f9587d37ce18dc66ecba16425d1ecc01589f Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:39:23 +0000 Subject: [PATCH 4/4] fix(v1): omit absent sampling-mask logprobs from the wire --- tests/unit/test_graph.py | 2 ++ verifiers/v1/graph.py | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py index 46447bf399..a80bd495fc 100644 --- a/tests/unit/test_graph.py +++ b/tests/unit/test_graph.py @@ -260,6 +260,8 @@ def test_sampling_mask_logprobs_attributed_aggregated_and_round_tripped(): mixed = _two_turn_sampling_mask_trace([first, _mask([[40], [41, 7, 8]])]) assert mixed.branches[-1].sampling_mask.logprobs is None + restored = type(mixed).model_validate(mixed.model_dump()).branches[-1].sampling_mask + assert restored.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]) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 707036ac3c..3efd5c4c29 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -222,8 +222,9 @@ def serialize_sampling_arrays( if payload is None: return None return { - field.name: _encode_ndarray(getattr(payload, field.name)) + field.name: _encode_ndarray(value) for field in fields(payload) + if (value := getattr(payload, field.name)) is not None } @field_validator("sampling_mask", "top_logprobs", mode="before")