feat: build KeptTokens from native sampling_mask

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,
         ),
     )

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:

verifiers/v1/types.py

@@ -1,10 +1,12 @@
 from collections.abc import Iterable
 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
 
 
 class TextContentPart(BaseModel):
@@ -187,13 +189,23 @@
 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 @@
     # 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)