feat(v1): optional sampler logprobs on SamplingMask

[faresobeid]

What

On top of #2778's packed sampling-mask format, carries the sampler's logprobs at the sampling-mask ids onto the trace graph.

• SamplingMask gets an optional float32 logprobs array, parallel to ids: the sampler's renormalized logprobs of the kept ids.
• SamplingMask.from_sampling_mask(mask, logprobs=None) decodes logprobs from prime-rl's /generate sampling_mask_logprobs: packed float32, base64 {data, shape, dtype}, the same form as perf(v1): accept packed CSR sampling masks from prime-rl's /generate #2778's packed ids, decoded with perf(v1): accept packed CSR sampling masks from prime-rl's /generate #2778's _decode_packed. response_from_generate passes result.get("sampling_mask_logprobs"). Stock vLLM does not send it, so the per-row mask path has no logprobs.
• Attribution: a mask whose logprobs length differs from ids is not attributed (same rule as a mask whose counts do not match the node's sampled positions).
• Aggregation: Branch.sampling_mask concatenates the logprobs only when every mask on the path has them; otherwise logprobs is None and ids/counts still merge. Nodes without a mask (prompt, replayed prefix) do not count.
• The __nd__ serializer/validator handle every field present on the dataclass. Without logprobs, the serialized form and payload are unchanged.

Payload when present: 4 B per kept id.

Consumer

prime-rl score centering, PrimeIntellect-ai/prime-rl#3794 (stacked on prime-rl#3900), pins this head. It sends sampling_mask_logprobs packed. Without that consumer this field is unused.

Stack

Base: #2778 (perf/packed-sampling-mask). Companion: PrimeIntellect-ai/renderers#170 (passes sampling_mask_logprobs and payload through), on top of renderers#174. Order: renderers#174 -> #2778 -> prime-rl#3900; then renderers#170 + this PR -> prime-rl#3794. #2778 is merged into this branch.

Validation

• tests/v1/test_graph.py::test_sampling_mask_logprobs_attributed_aggregated_and_round_tripped: branch concat, wire round trip, mixed turns -> None, mismatched length not attributed.
• tests/v1/test_trace.py::test_sampling_mask_from_packed_matches_lists: packed logprobs decode.
• pytest tests/v1/test_graph.py tests/v1/test_trace.py, ruff check, ruff format --check.

Note

Add optional sampler logprobs to SamplingMask

• Adds an optional float32 logprobs field to SamplingMask in types.py, parallel to the flattened ids, plus support for packed logprobs in the factory
• response_from_generate in train.py now passes renderer-provided mask logprobs into SamplingMask construction
• Serialization and deserialization in graph.py now handle every non-None mask field, so logprobs survive the wire round trip
• Branch.sampling_mask in trace.py concatenates logprobs across turns when every contributing mask has them, and sets them to None otherwise
• Behavioral Change: _attribute_sampling_mask in graph.py rejects masks whose logprobs length does not match the flattened ids; such masks are no longer attached to the assistant node

^{Macroscope summarized cf6ebfe.}

---

Note

Low Risk
Additive optional field with strict length validation; behavior unchanged when logprobs are absent (e.g. stock vLLM).

Overview
Adds optional sampler logprobs on v1 SamplingMask, carried parallel to flattened mask ids from prime-rl /generate through traces and training.

SamplingMask gains an optional float32 logprobs array and from_sampling_mask(..., logprobs) decodes packed sampling_mask_logprobs alongside the existing packed ids/counts. The train client wires that field from generate results into TurnTokens.sampling_mask.

On the graph, _attribute_sampling_mask drops masks when logprobs length does not match ids; Branch.sampling_mask concatenates logprobs only if every attributed node mask includes them (otherwise aggregated logprobs stays None). Serialization/deserialization now round-trips all non-None dataclass fields on the mask, not just ids/counts.

Tests cover multi-turn branch concat, wire round-trip, mixed turns without logprobs, and invalid length rejection.

^{Reviewed by Cursor Bugbot for commit 71ad510. Bugbot is set up for automated code reviews on this repo. Configure here.}

[macroscopeapp]

Approvability

Verdict: Approved at cf6ebfe

Macroscope's review found this PR approvable — This is a focused, backwards-compatible extension that carries optional sampler logprobs through existing SamplingMask ingestion, graph serialization, and branch aggregation. Existing behavior remains unchanged when the optional data is absent, and the new cases have targeted test coverage.

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

Dismissing prior approval to re-evaluate 1e499fb

Dismissing prior approval to re-evaluate cf6ebfe