feat(v1): by-handle payload segments on traces

tests/unit/test_graph.py

@@ -740,3 +740,60 @@ def test_prompt_supplied_assistant_messages_are_not_sampled_turns():
     assert [n.sampled for n in trace.nodes] == [False, False, False, True]
     assert trace.num_turns == 1
     assert trace.assistant_messages == [response]
+
+
+def test_payload_segments_attributed_per_node_and_rebased_per_branch():
+    """By-handle rows tile the nodes a turn creates. A bridged turn's first row belongs to the
+    prior sampled turn's final token; a reused non-sampled prefix row is dropped. Each branch
+    rebases its nodes' segments onto its own token positions."""
+    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")
+
+    def commit(messages, content, prompt_ids, spans, pos, rows):
+        segment = vf.PayloadSegment(
+            field="routed_experts",
+            file="f",
+            offset=pos * 2,
+            pos=pos,
+            rows=rows,
+            dtype="uint8",
+            shape=[2],
+        )
+        graph.prepare_turn(trace, messages).commit(
+            vf.Response(
+                id=content,
+                created=0,
+                model="t",
+                message=vf.AssistantMessage(content=content),
+                finish_reason="stop",
+                tokens=TurnTokens(
+                    prompt_ids=prompt_ids,
+                    completion_ids=[40, 41],
+                    message_spans=spans,
+                    payload=[segment],
+                ),
+            )
+        )
+
+    def rows(branch):
+        return [(s.pos, s.rows, s.offset) for s in branch.payload]
+
+    commit([user], "a1", [10, 11, 12], [(0, 3)], 0, 4)
+    # Bridged: request row 4 is a1's unforwarded final token.
+    commit(
+        [user, vf.AssistantMessage(content="a1"), vf.UserMessage(content="u2")],
+        "a2",
+        [10, 11, 12, 40, 41, 30, 31],
+        [(0, 3), None, (5, 7)],
+        4,
+        4,
+    )
+    # Fork off the shared user prompt: re-forwarded rows 0-2 are dropped.
+    commit([user], "b1", [10, 11, 12], [(0, 3)], 0, 4)
+
+    a2, b1 = (b for b in trace.branches if len(b.token_ids) in (9, 5))
+    assert rows(a2) == [(0, 3, 0), (3, 1, 6), (4, 1, 8), (5, 2, 10), (7, 1, 14)]
+    assert rows(b1) == [(0, 3, 0), (3, 1, 6)]

verifiers/v1/__init__.py

@@ -153,6 +153,8 @@
         Message,
         MessageContent,
         Messages,
+        NodeSegment,
+        PayloadSegment,
         Request,
         Response,
         SamplingMask,
@@ -260,8 +262,10 @@
     "ModelCall": "verifiers.v1.trace",
     "ModelContext": "verifiers.v1.clients",
     "NetworkPolicyConfig": "verifiers.v1.configs.runtime",
+    "NodeSegment": "verifiers.v1.types",
     "PATCH_CAP_BYTES": "verifiers.v1.utils.git",
     "ParentLink": "verifiers.v1.semantic",
+    "PayloadSegment": "verifiers.v1.types",
     "PolicyEvent": "verifiers.v1.trace",
     "PolicySpan": "verifiers.v1.episode",
     "PrimeConfig": "verifiers.v1.configs.runtime",
@@ -440,6 +444,8 @@
     "Branch",
     "TurnTokens",
     "SamplingMask",
+    "PayloadSegment",
+    "NodeSegment",
     "TopLogprobs",
     "Timing",
     "TimeSpan",

verifiers/v1/clients/train.py

@@ -25,6 +25,7 @@
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
+    PayloadSegment,
     Response,
     SamplingMask,
     Tool,
@@ -193,6 +194,9 @@ def response_from_generate(
             sampling_mask=SamplingMask.from_sampling_mask(mask)
             if (mask := result.get("sampling_mask"))
             else None,
+            payload=[PayloadSegment(**segment) for segment in segments]
+            if (segments := result.get("payload"))
+            else None,
             weight_versions=[WeightVersionSpan(**span) for span in spans]
             if (spans := result.get("weight_versions"))
             else None,

verifiers/v1/graph.py

@@ -48,6 +48,8 @@
     FloatStream,
     IntStream,
     Message,
+    NodeSegment,
+    PayloadSegment,
     Response,
     SamplingMask,
     TextContentPart,
@@ -778,6 +780,43 @@ def _attribute_routed_experts(
         off = end
 
 
+def _attribute_payload(
+    trace: Trace,
+    prefix_node_ids: list[int],
+    new_node_ids: list[int],
+    path_len: int,
+    segments: list[PayloadSegment] | None,
+) -> None:
+    """Append this turn's by-handle rows to `Trace.payload`, node by node. The request's
+    positions tile `[path_len:]` over the nodes created this turn. The one position before,
+    the prefix's final token, belongs to the prefix's last node when that node was sampled:
+    the previous turn never forwarded its final token, so this prefill is the first to
+    produce its row. Rows of other reused prefix positions (a turn that re-forwarded its
+    history) are dropped. Rows are never rewritten, only appended."""
+    if not segments:
+        return
+    # (node, request lo, request hi, node start)
+    spans: list[tuple[int, int, int, int]] = []
+    last = next(
+        (nid for nid in reversed(prefix_node_ids) if len(trace.nodes[nid].token_ids)),
+        None,
+    )
+    if last is not None and trace.nodes[last].sampled:
+        spans.append(
+            (last, path_len - 1, path_len, path_len - len(trace.nodes[last].token_ids))
+        )
+    start = path_len
+    for nid in new_node_ids:
+        end = start + len(trace.nodes[nid].token_ids)
+        spans.append((nid, start, end, start))
+        start = end
+    for segment in segments:
+        for nid, lo, hi, node_start in spans:
+            part = segment.clip(lo, hi, shift=-node_start)
+            if part is not None:
+                trace.payload.append(NodeSegment(node=nid, segment=part))
+
+
 def _attribute_sampling_mask(
     trace: Trace, assistant_id: int, payload: SamplingMask | None
 ) -> None:
@@ -1092,6 +1131,10 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int:
         trace, assistant_id, tokens.top_logprobs if tokens else None
     )
 
+    _attribute_payload(
+        trace, prefix, new_node_ids, path_len, tokens.payload if tokens else None
+    )
+
     return assistant_id
 
 

verifiers/v1/serve/delta.py

@@ -46,6 +46,7 @@
 LIST_FIELDS = (
     "nodes",
     "calls",
+    "payload",
     "errors",
     "extra_usage",
     "request_rewrites",

verifiers/v1/trace.py

@@ -29,6 +29,8 @@
     FinishReason,
     Message,
     Messages,
+    NodeSegment,
+    PayloadSegment,
     SamplingMask,
     Tool,
     ToolMessage,
@@ -48,9 +50,11 @@
             "sampling_mask",
             "top_logprobs",
         }
-    }
+    },
+    "payload": True,
 }
-"""Raw tensor fields kept on the msgpack wire but excluded from disk serialization."""
+"""Raw tensor fields and short-lived payload handles, kept on the msgpack wire but excluded
+from disk serialization."""
 
 
 class TimeSpan(BaseModel):
@@ -209,6 +213,9 @@ class Branch(BaseModel):
     """Whether this physical path contributes a training sample."""
     mm_token_type_id_map: dict[int, int] = Field(default_factory=dict)
     """The trace's `mm_token_type_id_map`, carried so `mm_token_type_ids` is self-contained."""
+    payload: list[PayloadSegment] = Field(default_factory=list)
+    """The trace's payload segments on this path, positions relative to `token_ids`, in
+    the order they were appended (a later segment wins where two overlap)."""
 
     @property
     def messages(self) -> Messages:
@@ -413,6 +420,10 @@ class Trace(BaseModel, Generic[DataT, StateT, AgentConfigT]):
     """The message graph, including physical and semantic parent links."""
     calls: list[ModelCall] = Field(default_factory=list)
     """Every model call; automatically recorded at intercept time + linked into `nodes`."""
+    payload: list[NodeSegment] = Field(default_factory=list)
+    """By-handle per-token side arrays (router-replay ids, sampling masks), append-only:
+    each turn's `commit` attributes the inference server's segments to the nodes they cover.
+    Handles point into a short-lived shared directory; kept off disk records."""
     mm_token_type_id_map: dict[int, int] = Field(default_factory=dict)
     """Special-token id -> modality marker (1 = image placeholder, 2 = video placeholder)
     from the renderer that tokenized this trace, stamped at turn commit. Applied to a
@@ -562,6 +573,11 @@ def branches(self) -> list[Branch]:
                 path.append(nid)
                 nid = self.nodes[nid].parent
             path.reverse()
+            node_starts: dict[int, int] = {}
+            offset = 0
+            for n in path:
+                node_starts[n] = offset
+                offset += len(self.nodes[n].token_ids)
             is_compaction_attempt = any(
                 link.type == "compaction_attempt"
                 for link in self.nodes[leaf].semantic_parents
@@ -576,6 +592,13 @@ def branches(self) -> list[Branch]:
                         or leaf in accepted_compaction_attempts
                     ),
                     mm_token_type_id_map=self.mm_token_type_id_map,
+                    payload=[
+                        entry.segment.model_copy(
+                            update={"pos": entry.segment.pos + node_starts[entry.node]}
+                        )
+                        for entry in self.payload
+                        if entry.node in node_starts
+                    ],
                 )
             )
         return branches

verifiers/v1/types.py

@@ -299,6 +299,48 @@ def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray:
     return array.astype(dtype, copy=False)
 
 
+class PayloadSegment(BaseModel):
+    """`rows` consecutive rows of one per-token side array (`field`, e.g. `routed_experts`,
+    `sampling_mask`) that the inference server wrote to `file` at byte `offset`, as a
+    C-contiguous `[rows, *shape]` `dtype` array. `pos` is the token position of the first
+    row in whatever sequence holds the segment: the request on `TurnTokens`, the node on
+    `Trace.payload`, the branch on `Branch.payload`. Only the handle travels; the trainer
+    reads the rows."""
+
+    field: str
+    file: str
+    offset: int
+    pos: int
+    rows: int
+    dtype: str
+    shape: list[int]
+
+    @property
+    def end(self) -> int:
+        return self.pos + self.rows
+
+    def clip(self, lo: int, hi: int, shift: int = 0) -> "PayloadSegment | None":
+        """The rows inside positions `[lo, hi)`, moved by `shift`; None when empty."""
+        start, end = max(self.pos, lo), min(self.end, hi)
+        if start >= end:
+            return None
+        row_bytes = np.dtype(self.dtype).itemsize * int(np.prod(self.shape))
+        return self.model_copy(
+            update={
+                "offset": self.offset + (start - self.pos) * row_bytes,
+                "pos": start + shift,
+                "rows": end - start,
+            }
+        )
+
+
+class NodeSegment(BaseModel):
+    """A payload segment attributed to `Trace.nodes[node]`, positions node-local."""
+
+    node: int
+    segment: PayloadSegment
+
+
 @dataclass
 class TopLogprobs:
     """Top-k sampling head stored as flat `ids`, `logprobs` and `counts` arrays.
@@ -369,6 +411,9 @@ class TurnTokens(BaseModel):
     # Transient carrier (excluded): per-completion-token sampling masks,
     # attributed to the assistant node by the turn's `commit`, then dropped.
     sampling_mask: SamplingMask | None = Field(default=None, exclude=True)
+    # Transient carrier (excluded): by-handle side arrays from `generate`, request-absolute
+    # positions, attributed per node onto `Trace.payload` by the turn's `commit`.
+    payload: list[PayloadSegment] | None = Field(default=None, exclude=True)
     # Transient carrier (excluded): the engine's weight-version spans over the completion,
     # attributed to the assistant node by the turn's `commit`, then dropped.
     weight_versions: list[WeightVersionSpan] | None = Field(default=None, exclude=True)