From d87f63a2eb51fefe537534333eb9e6ec31c884f5 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/2] perf(v1): accept packed CSR sampling masks from
 prime-rl's /generate

---
 tests/unit/test_trace.py | 25 +++++++++++++++++++++++++
 verifiers/v1/types.py    | 16 +++++++++++++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_trace.py b/tests/unit/test_trace.py
index b4fd9ed279..cc094d3a93 100644
--- a/tests/unit/test_trace.py
+++ b/tests/unit/test_trace.py
@@ -241,3 +241,28 @@ def test_transcript(make_full_trace):
     assert "[tool search]\nTOOL RESULT: Paris is the capital." in transcript
     assert "It is Paris." in transcript
     assert "SECRET REASONING" not in transcript  # reasoning is excluded
+
+
+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 8f1a6ceeeb06e23d1421e478929cff9dfb2ca722 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 2/2] 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."""