[nxphi47]

Carry sampler probability heads from renderer responses onto assistant nodes and training branches for prime-rl #4012, paired with renderers #180.

TopLogprobs is a flat IDs/logprobs/counts carrier. It shares the existing raw-array serializer with sampling masks, checks sampled-token alignment when attributed, and gives context or uncaptured positions zero counts when flattening a branch. Public trace records exclude the arrays. Empty completions preserve an empty carrier.

Validation: 179 unit tests passed; a temporary probe covered multi-turn alignment, mixed capture, empty completions, malformed rows/counts, MessagePack round trips, and record exclusion. All contributor hooks and the configured type gate passed using uv 0.13.1. The full suite was attempted: four Docker Compose integration cases failed because the local Docker daemon is unavailable; provider/runtime cases were skipped where prerequisites were absent.

Updated onto current main without rewriting the existing PR history.

Note

Carry top-k sampler logprob heads onto graph nodes and branches

• Adds a TopLogprobs dataclass in types.py that flattens top-k token IDs, log probabilities, and per-token counts, with a validating from_top_logprobs converter.
• response_from_generate in train.py parses completion top IDs/logprobs into a transient TurnTokens.top_logprobs field.
• At turn commit, _attribute_top_logprobs in graph.py validates payload alignment against the assistant node's sampling mask and attaches it to the node; misaligned payloads raise ValueError.
• Branch.top_logprobs in trace.py concatenates node-level heads in node order for branch-aligned training data. Serialization is shared with SamplingMask; top_logprobs is excluded from disk serialization via EXCLUDE_FIELDS but kept in the msgpack wire format.
• Risk: MessageNode.serialize_sampling_arrays and deserialize_sampling_arrays now handle both SamplingMask and TopLogprobs; misaligned top-logprob payloads at commit raise ValueError.

📊 Macroscope summarized 2b8718e. 5 files reviewed, 2 issues evaluated, 2 issues filtered, 0 comments posted

🗂️ Filtered Issues

verifiers/v1/graph.py — 0 comments posted, 2 evaluated, 2 filtered

• line 843: counts is only checked for negativity and its truncated total. A malformed numeric array such as np.array([1.5, 0.5]) for a two-token completion with two flat ids passes all four conditions (int(sum)==len(ids)), then is stored. Branch.top_logprobs later scatters it into an int32 array, silently converting it to [1, 0] while retaining both ids/logprobs; the branch therefore has counts that no longer describe the flat head and training consumes misaligned sampler probabilities. Require an integral count dtype/values (and validate the exact sum without int()) before attaching it. [ Out of scope (post-validation triage) ]
• line 844: The alignment test uses len(payload.ids)/len(payload.logprobs), which measures only axis 0 and never enforces the documented flat layout. For example, a one-token completion with ids and logprobs shaped (1, 2) and counts=[1] passes every check, is attached, and Branch.top_logprobs returns 2-D candidate arrays despite its counts declaring one flat candidate. Downstream training then receives an invalid head layout (or fails when it assumes a 1-D CSR array). Reject non-1-D ids/logprobs/counts before assigning the payload. [ Out of scope (post-validation triage) ]

[macroscopeapp]

Approvability

Verdict: Needs Human Approval

Macroscope's review found this PR requires human approval — This PR adds a new sampler-probability data path from generation responses through serialized trace nodes into training branches. The cross-component runtime and wire-format contract is broader than a self-contained additive change and merits human review.

^{You can add or adjust custom eligibility rules. Learn more.}