Token capture finalizer: batch the group staging fetch, and give the pool a placement policy + a splittable timer

[ZhiyuLi-Nvidia]

Follow-ups from review of #3837 (gate-authoritative token capture). Neither blocks that PR — both were raised there as non-blocking comments and are filed here so they are not lost on merge.

Both concern the same component, the CPU finalizer pool, and they interact: you cannot pick between them without the timer split in (2).

  SingleController ── group ──► FinalizerActor pool (default 2, unplaced)
                                     │
                                     ├─ 1 fetch per rollout  ──► TQ    (1)
                                     └─ per-token verify in Python     (2)

  row_assembly/rollouts_ms measures BOTH ──► cannot tell which dominates

---

1. finalize_group issues N staging fetches where 1 would do

Ref: #3837 (comment)

finalize_group calls finalize_rollout per rollout, and each does its own
fetch_for_finalization → kv_batch_get. All N rollouts' staging keys are known before
the first fetch (the receipts are already in hand, and receipt → staging_keys is pure
local work), so one batched read over the union would serve the group.

  now        for rollout in group:  fetch(that rollout's keys)      2N messages
  proposed   phase 1  parse N receipts -> all keys (no I/O)
             phase 2  ONE fetch over the union                       2 messages
             phase 3  per-rollout verify from a key -> row map

Each kv_batch_get is constant in key count, but its kv_retrieve_meta hop lands on the
TQ controller's single request thread, and the pool is bounded, so the N calls serialize
inside one actor. finalize/calls_per_rollout == 1.0 in observed runs, meaning each
current call fetches a single key — per-call overhead dominates payload.

Not measured. Byte volume is unchanged; the claim is (N-1) fewer serialized round trips
per group.

Two guards have to become per-key first — this is the actual work, not the batching:

• the row-count check and _row_to_base_snapshot are both request-scoped today, so one
  missing or malformed row would reject all N rollouts instead of one. Rows carry their own
  identity (staging_key derived from rollout_id_utf8 + model_call_id_utf8, already
  asserted against the requested key), so both can be made per-key.
• direct mode's extras probe falls back to the base schema on any exception, and
  kv_batch_get readiness is all-or-nothing across a batch — so one row that staged no
  routes degrades every row in the request. Reading STAGING_FIELDS first, then
  [ROUTED_EXPERTS_FIELD] for only the keys with encoding != 0, makes readiness hold by
  construction. This one has value independent of the batching.

---

2. Finalizer pool has no placement policy and no way to size itself

Ref: #3837 (comment)

2a. Placement

  actor                      scheduling_strategy
  ─────────────────────────  ────────────────────────────────
  GPU worker groups          PlacementGroupSchedulingStrategy
  GenerationRouterActor      NodeAffinity(driver, soft=False)
  NemoGym actor              NodeAffinity(driver, soft=True)
  FinalizerActor             — none —

FinalizerActor.remote(...) takes no .options(), so num_cpus=1 is the whole policy and
Ray places the pool on whatever node has a free logical CPU. Correctness is unaffected (no
GPU state; the actor connects rather than bootstraps), but network locality to the storage
units is left to luck, and under Mooncake each TQ client mounts its own segment — so N
finalizers reserve N × global_segment_size of DRAM on unpredictable nodes, which
num_cpus=1 does not express.

Fix: match GenerationRouterActor, with soft=True since the pool scales with
num_finalizer_workers and must be allowed to spill.

2b. Sizing — and why this is CPU-bound, not I/O-bound

This is the part that decides whether more actors help at all, so stating the evidence
rather than asserting it.

The architecture already concedes it. The pool is N separate Ray processes, each
reserving a whole CPU. If the work were I/O-bound, a thread pool or an inline
asyncio.to_thread on the controller would have been the cheaper answer. Separate
processes buy exactly one thing — separate GILs. num_cpus=1 is then the correct
reservation, since a GIL-bound actor cannot use more than about one core anyway.

The work is per-token pure Python, repeated per staged call:

  per staged call, over delta_len tokens
    tq_token_sink.py:527   token_ids_delta      [int(t) for t in ....tolist()]
                  :529-530 token_mask_delta     [float(m) for m in ....tolist()]
                  :531-533 logprobs_delta       [float(p) for p in ....tolist()]
                  :479     id/text fields       bytes(int(v) for v in ....tolist())
    Gym verify_and_linearize                    digest recompute + validity scan
  × calls/rollout × rollouts/group × groups/step

Every one of those is .tolist() (allocates a Python list) followed by an interpreter loop
over it. Only the digest's hashing leaves Python.

Two honest caveats:

• The digest and validity passes now live in Gym's verify_and_linearize
  (blackbox_finalizer.py:222-223 — "All base token/digest/lineage/terminal semantics
  belong to Gym; the finalizer never re-verifies them"). They still run on the finalizer
  actor's CPU, but I could not read that implementation from my workspace, so its cost is
  inferred from its inputs. The three list conversions above are directly verifiable in RL.
• The digest is recomputed, not trusted — that is the ledger's security property, so it
  cannot be skipped. It does mean the token data is hashed a second time after the writer
  already hashed it once.

Why the default is untested rather than wrong. num_finalizer_workers: 2 is set by no
recipe, and the one functional test runs at num_generations_per_prompt=2 /
max_buffered_rollouts=4 — far too small to saturate. Single-call math workloads
(finalize/calls_per_rollout == 1.0, short deltas) will be fine at 2. The multi-turn
agentic chains this feature exists for multiply every term above, and that is the regime
nothing has exercised.

The obstacle to picking a number is one merged timer: row_assembly/rollouts_ms spans the blocking fetch and
the Python passes.

  fetch_ms dominant   ->  do item 1, don't add actors
  verify_ms dominant  ->  add actors, ceiling = free CPUs

Splitting it into row_assembly/fetch_ms and a derived verify_ms is ~3 lines and makes the
choice decidable. finalize/queue_depth and finalize/queue_wait_ms already exist, so rising
queue depth is the undersize signal once the split lands.

---

Suggested order: the timer split first (cheapest, and it decides the rest), then the
placement fix, then the group-fetch batching if fetch_ms turns out to dominate.