From f0d4154412c86edaeb7afbadd1b04665b5170d6b Mon Sep 17 00:00:00 2001
From: Mika Senghaas <mail@mikasenghaas.de>
Date: Sat, 29 Aug 2026 00:16:45 +0000
Subject: [PATCH] feat: build KeptTokens from native sampling_mask

vLLM 0.28 returns kept-set sampling masks natively (sampling_mask,
one list of surviving vocab ids per completion token) instead of the
custom base64 kept_tokens payload. Convert to the flat int32 ids/counts
arrays on the train client; graph attribution validates alignment and
attaches the arrays as-is.
---
 verifiers/v1/clients/train.py |  4 ++--
 verifiers/v1/graph.py         |  7 ++-----
 verifiers/v1/types.py         | 20 ++++++++++++++++----
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py
index 528b8e12a..d9f9b5b63 100644
--- a/verifiers/v1/clients/train.py
+++ b/verifiers/v1/clients/train.py
@@ -162,8 +162,8 @@ def response_from_generate(
             multi_modal_data=result.get("multi_modal_data"),
             mm_token_type_id_map=mm_token_type_id_map,
             routed_experts=result.get("routed_experts"),
-            kept_tokens=KeptTokens(**kept)
-            if (kept := result.get("kept_tokens"))
+            kept_tokens=KeptTokens.from_sampling_mask(mask)
+            if (mask := result.get("sampling_mask"))
             else None,
         ),
     )
diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py
index 165504bcd..b6f1455b6 100644
--- a/verifiers/v1/graph.py
+++ b/verifiers/v1/graph.py
@@ -526,13 +526,10 @@ def _attribute_kept_tokens(
     that doesn't line up with the node's sampled tokens is dropped, not misaligned."""
     if payload is None:
         return
-    counts = np.frombuffer(binascii.a2b_base64(payload.counts), dtype=np.int32)
-    ids = np.frombuffer(binascii.a2b_base64(payload.ids), dtype=np.int32)
     node = trace.nodes[assistant_id]
-    if len(counts) != sum(node.mask) or int(counts.sum()) != len(ids):
+    if len(payload.counts) != sum(node.mask) or int(payload.counts.sum()) != len(payload.ids):
         return
-    # Own the buffers — the payload views reference the turn's response bytes.
-    node.kept_tokens = KeptTokens(ids=ids.copy(), counts=counts.copy())
+    node.kept_tokens = payload
 
 
 def _commit_turn(turn: PendingTurn, response: Response) -> int:
diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py
index 03cf38b05..7f643972e 100644
--- a/verifiers/v1/types.py
+++ b/verifiers/v1/types.py
@@ -2,6 +2,8 @@
 from dataclasses import dataclass
 from typing import Annotated, Any, Literal
 
+import numpy as np
+
 from pydantic import AliasChoices, BaseModel, ConfigDict, Field
 from renderers.base import MultiModalData
 from typing_extensions import TypedDict
@@ -187,13 +189,23 @@ class RoutedExperts(TypedDict):
 class KeptTokens:
     """Kept-set sampling masks for sampling replay: `ids` (every kept set concatenated
     in position order) and `counts` (kept-set size per completion token; 0 = no usable
-    mask). Base64 blobs straight off the `generate` response on the `TurnTokens`
-    carrier; decoded to flat int32 arrays on `MessageNode` (`len(ids) == sum(counts)`,
-    row boundaries recovered from `counts`)."""
+    mask) as flat int32 arrays (`len(ids) == sum(counts)`, row boundaries recovered
+    from `counts`). Built from vLLM's native `sampling_mask` payload on the `generate`
+    response (one list of surviving vocab ids per completion token)."""
 
     ids: Any
     counts: Any
 
+    @classmethod
+    def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "KeptTokens":
+        counts = np.fromiter((len(row) for row in sampling_mask), dtype=np.int32, count=len(sampling_mask))
+        ids = (
+            np.concatenate([np.asarray(row, dtype=np.int32) for row in sampling_mask])
+            if int(counts.sum())
+            else np.empty(0, dtype=np.int32)
+        )
+        return cls(ids=ids, counts=counts)
+
 
 class TurnTokens(BaseModel):
     """Training tokens from renderer tokenization or provider-returned token IDs."""
@@ -224,7 +236,7 @@ class TurnTokens(BaseModel):
     # Transient carrier (excluded): the kept-set sampling masks from `generate` (token ids
     # surviving top-p/top-k truncation, per completion token), attributed to the assistant
     # node by the turn's `commit`, then dropped. None unless the engine ran with
-    # `enable_return_kept_tokens`.
+    # `--return-sampling-mask`.
     kept_tokens: KeptTokens | None = Field(default=None, exclude=True)