feat: sampling mask pass-through

[mikasenghaas]

Summary

• Build SamplingMask from a sampling_mask payload.
• Keep sampling-mask trace types independent of the inference producer.
• Require the published renderers==0.1.12.dev2 feature floor.
• Store masks as flat int32 ids and per-token counts arrays.
• Use sampling_mask across response tokens, message nodes, branch views, wire serialization, and public exports.
• Keep Branch.sampling_mask aligned to the branch token ids for prime-rl.

Pairs with PrimeIntellect-ai/renderers#144 and prime-rl's native sampling-replay PR.

Breaking

• Rename vf.KeptTokens to vf.SamplingMask.
• Rename TurnTokens.kept_tokens, MessageNode.kept_tokens, and Branch.kept_tokens to sampling_mask.
• Update serialized trace consumers to read the sampling_mask field.

Verification

• Merged current main without conflicts.
• uv sync --upgrade-package renderers: lockfile regenerated with renderers==0.1.12.dev2.
• uv run pytest tests/v1/test_graph.py tests/v1/test_trace.py -q: 21 passed.
• Ruff, format, type, lock, and diff checks pass.

🤖 Generated with Claude Code

Note

Replace KeptTokens with SamplingMask across v1 verifier pipeline

• Adds the SamplingMask dataclass in types.py with ids and counts fields plus from_sampling_mask, and swaps kept_tokens for sampling_mask on TurnTokens and MessageNode
• train.py now reads sampling_mask from the generate result and builds it via SamplingMask.from_sampling_mask; response_from_generate populates TurnTokens.sampling_mask
• graph.py renames _attribute_kept_tokens to _attribute_sampling_mask, which validates counts/ids lengths and assigns the SamplingMask directly (no base64 decode or buffer copy)
• trace.py excludes sampling_mask from disk serialization and exposes Branch.sampling_mask, which scatters per-node counts onto sampled positions and concatenates ids
• Bumps renderers[multimodal] dependency from >=0.1.11 to >=0.1.12.dev2 in pyproject.toml
• Behavioral Change: KeptTokens is removed from the verifiers.v1 public API (__all__) and SamplingMask replaces it; any downstream code importing KeptTokens or reading kept_tokens on TurnTokens, MessageNode, or Branch must switch to sampling_mask. _attribute_sampling_mask no longer base64-decodes or copies buffers — callers must provide a valid SamplingMask payload directly

Changes since #2460 opened

• Removed producer-specific references from sampling mask documentation [3ff2553]

^{Macroscope summarized 67af753.}

---

Note

Medium Risk
Breaking rename across the v1 training/trace wire contract; incorrect mask alignment would silently drop masks at attribution, and consumers must upgrade renderers in lockstep.

Overview
Replaces the KeptTokens / kept_tokens surface with SamplingMask / sampling_mask so training traces align with the new renderers generate payload (sampling_mask as nested token-id rows). Public exports, TurnTokens, MessageNode, wire serialization, and Branch.sampling_mask all use the new names; disk dumps still exclude the tensor field via EXCLUDE_FIELDS.

SamplingMask holds flat int32 ids and per-completion-token counts, with from_sampling_mask to build from the engine’s list[list[int]]. The train client maps result["sampling_mask"] into TurnTokens; graph attribution validates row counts against sampled positions and stores the mask on the assistant node (no base64 decode/copy path like the old kept-set flow).

Dependency: renderers[multimodal] is pinned to >=0.1.12.dev2 (lockfile updated). Downstream code that imported KeptTokens or read kept_tokens on traces must switch to sampling_mask.

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

[macroscopeapp]

Approvability

Verdict: Not approved

Macroscope's review found this PR not approvable — This is a cross-project training-trace integration that adopts a new renderer payload and replaces the public and wire contract from KeptTokens/kept_tokens to SamplingMask/sampling_mask. Although ordinary generation is largely unchanged without a mask, the breaking API/schema migration and prerelease dependency require human review.

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