feat(sample-support): capture the sampler's bounded support and trace it across turns

[dyurk-lila]

Note on the diff: This PR is part of a routed-expert-replay / sampler-support series and builds on the PRs below. GitHub cannot show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator
• refactor(generators): assemble routed-expert traces incrementally #1912 — refactor assemble routed-expert traces incrementally
• perf(r3): collate routes through a container-aware pool, and refuse replay under VPP #2078 — perf(r3) collate routes through a container-aware pool
• refactor(wire): generic packed-ndarray codec and an N-field response-body splice #2079 — refactor(wire) generic packed-ndarray codec
• perf(r3): route rollout experts to the trainer packed with cu_seqlens #2080 — perf(r3) packed routed-expert trainer transport
• perf(batch): ship large per-token side channels to the trainer zero-copy #2081 — perf(batch) zero-copy per-token side channels

Reviewing in PR order (lowest number first) shows each incremental change cleanly.

Problem

With top-k, top-p, or min-p sampling, vLLM samples from a filtered set of vocabulary IDs and normalizes probabilities over that set. The trainer normally recomputes log-probabilities over the full vocabulary. Even when model weights and kernels agree, those two values represent different distributions, which biases importance ratios and KL terms.

Replaying the rollout distribution requires recording the support that survived vLLM’s processed sampler for every generated token. This PR establishes that capture, wire, and generator contract. The following PRs carry the support through the training batch and use it in the Megatron and FSDP scorers.

vLLM capture contract

When a training request opts into support capture, the request asks vLLM for flat_logprobs with logprobs=top_k. Each generated token produces one flat row:

[sampled token, top candidate 0, ..., top candidate top_k-1]

The first column supplies the sampled-token log-probability. The remaining columns become one fixed-width int32 support row. Candidates removed by top-p or min-p filtering have processed log-probability -inf and are encoded as trailing -1 values.

vLLM’s approximate top-k/top-p pivot can rarely return a support row that omits the token it sampled. When that happens, the endpoint replaces the weakest valid candidate with the sampled ID and emits a warning. This preserves the invariant required by replay without changing the row width or padding layout.

Capture is rejected unless the configured training sampler can be represented exactly:

• temperature > 0;
• top_k > 1;
• repetition_penalty == 1.0; and
• no arbitrary sampler additional_kwargs.

Evaluation requests opt out per request, so greedy evaluation is unaffected by the training capture settings.

Wire and generator lifecycle

The support array uses PR 27’s packed-array response envelope and is decoded exactly once by RemoteGenerateClient. RemoteGenerateResult.sample_support then carries the validated NumPy array alongside sampled log-probabilities and routed experts.

For multi-turn generation, SampleSupportTrace shares the incremental TokenMetadataTrace lifecycle used by routed-expert metadata:

• generated tokens append their captured rows;
• observation tokens append all--1 rows;
• a temporarily removed EOS retains its row if that EOS is restored at trajectory finalization;
• truncation slices token IDs, log-probabilities, routes, and support consistently; and
• the trace is finalized once per trajectory with an exact row-count check.

Retokenizing a chat history is rejected when per-token side channels are enabled because the new tokenization would no longer align with the captured rows. Step-wise outputs keep one support block aligned with each response step.

Integrated training evidence

The complete sample-support stack was evaluated by training GLM-4.7-Flash, an MoE model, on the easy split of NVIDIA’s public Nemotron-Terminal-Synthetic-Tasks dataset. The comparison used routed-expert replay alone versus routed-expert replay plus support-set replay with top_k=20 and top_p=0.95.

The curves are lightly EMA-smoothed over the raw per-step values and clipped to the shared step range, 11–151:

Mean training reward	Policy entropy

Training reward is nearly indistinguishable between the arms. By the end of the comparison, the top-p arm retains policy entropy around 0.14, while the routed-expert-only arm falls to roughly 0.05.

That retained entropy appears as sampling diversity at evaluation:

Arm	Pass@1	Pass@4	Pass@8
Routed-expert replay	0.514	0.605	0.631
Routed-expert + sample-support replay	0.512	0.620	0.641

Pass@1 is effectively tied, while the support-replay arm improves Pass@4 and Pass@8. These results validate the integrated feature through the trainer-scoring PRs above this capture layer; they are included here because this PR defines the public feature and configuration entry point.

Testing

• vLLM endpoint tests cover flat-row extraction, filtered padding, sampled-token repair, malformed widths, and opt-in behavior.
• Wire tests cover packed support round trips, dtype, shape, byte-count, and padding validation.
• Remote-client tests cover independent and combined route/support capture.
• Generator tests cover batched and agent-loop generation, single- and multi-turn traces, observations, EOS removal/restoration, truncation, train/eval selection, and step-wise output alignment.
• Configuration tests cover every supported sampler constraint and incompatible retokenization/VLM paths.

---

Note

High Risk
Touches the generate data plane, generator token-alignment contracts, and trainer batch layout for MoE router replay. Incorrect packing or support rows would silently bias off-policy correction.

Overview
Adds opt-in sample-support capture so training can later renormalize logprobs over the same filtered vocab vLLM sampled from. Train requests with generator.inference_engine.enable_return_sample_support_set ask vLLM for flat_logprobs/logprobs=top_k; filtered candidates become trailing -1s, and a rare missing sampled ID is repaired in-place. Eval stays greedy and opts out per request. Config refuses temperature ≤ 0, top_k ≤ 1, repetition penalty, extra sampler kwargs, VLMs, and custom chat templates.

SkyRLGymGenerator accumulates support (and routed experts) with SampleSupportTrace / RoutedExpertTrace: generated tokens keep captured rows, observations get all-padding rows, sliced EOS can be restored. Step-wise training still refuses R3 (prompt vs generated-route mismatch) but can emit per-step support blocks.

The generate wire now uses a generic packed-ndarray envelope (load_packed_body splices base64 without materializing huge strings). Trainer-side MoE routes become PackedTensor ([sum(seqlen), layers, topk] + cu_seqlens) instead of a padded rectangle, with pooled collation, CP/TP alignment via align_packed_token_metadata, and VPP+replay rejected. This is capture + transport only; scorers that consume support land in follow-ups.

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

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 9a8327d. Configure here.}

[cursor]

EOS support misaligned after observations

High Severity

When use_conversation_multi_turn is false, restored or synthetic EOS support is appended after final observation padding still held in sample_support_trace. SampleSupportTrace.finalize then keeps the first token_count rows, so observation padding lands on the EOS token and the real EOS support is discarded. Logprobs avoid this because they truncate before appending EOS. The length check still passes, so training silently gets misaligned support.

Additional Locations (1)

• skyrl/backends/skyrl_train/utils/sample_support.py#L51-L59

 

^{Reviewed by Cursor Bugbot for commit 9a8327d. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces a packed tensor representation (PackedTensor) for ragged token-aligned batch fields, such as MoE routing indices and sampler support sets, to optimize memory and serialization overhead. It also adds support for capturing and returning the sampler's bounded top-k support set (rollout_sample_support) from vLLM, refactors the remote inference client with optimized payload splicing, and implements zero-copy out-of-band transfers for large packed tensors. The review feedback highlights several robustness improvements, including catching ValueError in the HTTP retry loop to handle payload corruption, properly handling boolean tensor indexing in PackedTensor, adding defensive checks for None logprobs in the vLLM server actor, and preventing potential IndexError crashes when collating empty rollout expert indices.

[gemini-code-assist]

If load_packed_body raises a ValueError (e.g., due to network corruption or partial responses), it will propagate out of the try block and fail the entire request immediately, bypassing the retry loop. Since orjson.JSONDecodeError inherits from ValueError, catching ValueError (or both) allows transient network corruptions or bad payloads to be retried properly.

Suggested change

	except orjson.JSONDecodeError as exc:
	except (orjson.JSONDecodeError, ValueError) as exc:

[gemini-code-assist]

If index is a boolean mask tensor (e.g., torch.Tensor of type torch.bool), index.tolist() will return a list of booleans (e.g., [True, False, True]). When passed to _gather, torch.as_tensor(list(indices), dtype=torch.long) will convert True to 1 and False to 0, resulting in selecting segments 1, 0, and 1 instead of filtering the segments where the mask is True. Checking for boolean tensors and converting them to integer indices using nonzero() prevents this bug.

Suggested change

	if isinstance(index, torch.Tensor):
	if index.ndim == 0:
	return self.segment(int(index))
	return self._gather(index.tolist())
	if isinstance(index, torch.Tensor):
	if index.dtype == torch.bool:
	index = index.nonzero().flatten()
	if index.ndim == 0:
	return self.segment(int(index))
	return self._gather(index.tolist())

[gemini-code-assist]

If resp.logprobs is None (which can happen in some edge cases or error states in vLLM), _sample_support_from_flat_logprobs will crash with an AttributeError when trying to access logprobs.token_ids. Adding a defensive check for resp.logprobs is None will make this much more robust.

Suggested change

	if capture_sample_support:
	content, support_ids = _sample_support_from_flat_logprobs(
	resp.logprobs,
	sampling_params_dict["top_k"],
	)
	if capture_sample_support:
	if resp.logprobs is None:
	raise ValueError("vLLM did not return logprobs despite sample-support capture being enabled")
	content, support_ids = _sample_support_from_flat_logprobs(
	resp.logprobs,
	sampling_params_dict["top_k"],
	)

[gemini-code-assist]

If rollout_expert_indices is empty (e.g., num_samples == 0), accessing rollout_expert_indices[0] on line 145 will raise an IndexError. Adding a defensive check at the beginning of _collate_rollout_expert_indices prevents unexpected crashes.

    num_samples = len(rollout_expert_indices)
    if num_samples == 0:
        raise ValueError("rollout_expert_indices cannot be empty")