perf(sample-support): carry the captured support to the trainer as one packed field

[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
• feat(sample-support): capture the sampler's bounded support and trace it across turns #2082 — feat(sample-support) capture and trace sampler support

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

Problem

PR 30 produces one [response_tokens, top_k] support array per trajectory. Expanding those arrays to [batch, sequence_length, top_k] would allocate a large prompt region containing only -1 padding and repeat the same copies that packed routed-expert transport removed.

This PR carries support from generator output to trainer workers as one packed response-token field, without yet changing model scores.

Training representation

rollout_sample_support is a PackedTensor with:

• values: [sum(response_lengths), top_k] int32 support IDs; and
• cu_seqlens: one response-segment boundary per trajectory.

Unlike routed experts, which cover real prompt and response tokens, sample support covers response tokens only. Preprocessing requires exactly one support row per response token and a common support width across the batch. It performs one pooled fill into the final packed allocation and retains PR 29’s zero-copy object-store transport.

The field is threaded through generator-output concatenation and step-wise merging, training-input construction, slicing and chunking, replay-buffer storage, device transfer, pinning, and both sample- and token-based batch iterators.

Model-position alignment

The packed support values are not padded or shifted through every intermediate layout. Instead, each microbatch derives a cheap int64 row-ID channel:

1. assign every packed support row its index in values;
2. place response row r at the model position whose logits predict that response token;
3. use -1 where no support row exists; and
4. gather support values by row ID only at the scorer boundary introduced in the next PR.

Row IDs are derived after slicing because PackedTensor.chunk, padding, and microbatch selection rebase the packed row space. Persisting IDs from the global batch would point at the wrong support rows.

Padding rules

Packed fields have explicit, field-specific padding rules:

• routed experts need one valid dummy route for each attended padding token; and
• response-only support needs an empty segment for a synthetic padding batch.

The rules are registered centrally so adding another packed field cannot silently inherit routed-expert semantics. Ordinary batch padding copies the source segment length where appropriate, while data-parallel equalization uses the field’s synthetic-row rule.

Testing

• Support collation tests cover dtype and width validation, exact response lengths, read-only/non-contiguous NumPy inputs, and pooled filling.
• Row-ID tests cover left-padded prompts, variable response lengths, next-token prediction offsets, packed microbatches, slicing, and invalid bounds.
• Training-batch tests cover selection, chunking, concatenation, repetition, serialization, zero-copy transport, batch padding, and synthetic microbatches.
• Generator-output tests cover concatenation and prefix-aware step merging with observation padding rows.
• Trainer validation tests enforce the distinct token domains for routed experts and sample support.

---

Note

Medium Risk
Changes how per-token training side channels are collated, padded, sliced, and aligned for Megatron layouts. Incorrect packing or padding would silently misalign support rows with response tokens.

Overview
Carries captured sampler support into training as one packed field instead of a dense [batch, seq, top_k] rectangle. rollout_sample_support is a PackedTensor of [sum(response_lengths), top_k] IDs plus cu_seqlens, covering response tokens only (unlike packed MoE routes, which cover real prompt+response tokens).

Preprocessing validates a shared top-k width and exact response lengths, then fills the packed buffer with the same pooled first-touch path used for routes. Zero-copy object-store transport is reused. Packed fields now have explicit padding rules: dummy distinct-expert rows for R3, empty segments for synthetic support rows, so a new packed field cannot inherit router-replay padding.

Support is not expanded through every layout. Each microbatch derives int64 row IDs after slice/pad (because packing rebases the row space), placing response row r at the logit that predicts that token and -1 elsewhere. Scoring is not wired up yet.

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

[gemini-code-assist]

Code Review

This pull request introduces a packed side-channel mechanism to efficiently transmit and handle per-token metadata (such as MoE routing indices and sampler support sets) during training and generation, replacing padded representations with a PackedTensor layout to optimize memory usage. Feedback on these changes highlights three key improvement opportunities: catching ValueError alongside JSONDecodeError in the remote inference client to allow retries on transient errors while immediately raising permanent layout drift errors; returning an empty array instead of None in skyrl_gym_generator.py when the sample support trace has zero rows to prevent batch concatenation failures; and optimizing PackedTensor.cat to return the single batch directly when only one element is provided.

[gemini-code-assist]

If load_packed_body raises a ValueError (e.g., due to a truncated response or unterminated base64), it will not be caught by except orjson.JSONDecodeError, causing the request to fail immediately without retrying. We should catch ValueError as well, but re-raise it if it is a permanent layout or drift error to preserve the intended behavior for permanent failures.

                    except (orjson.JSONDecodeError, ValueError) as exc:
                        if isinstance(exc, ValueError) and not isinstance(exc, orjson.JSONDecodeError):
                            if any(msg in str(exc) for msg in ("drifted", "unspliced", "found no envelope")):
                                raise

[gemini-code-assist]

If agent_loop_state.sample_support_trace is not None but num_rows is 0 (e.g., if the environment starts in a terminal state and the loop is skipped), rollout_sample_support_out remains None. This will cause concatenate_generator_outputs to fail with a ValueError due to a mix of None and np.ndarray values in the batch. We should return an empty array of shape (0, sample_support_width) when num_rows is 0 to ensure consistency.

Suggested change

	if agent_loop_state.sample_support_trace is not None and agent_loop_state.sample_support_trace.num_rows:
	rollout_sample_support_out = agent_loop_state.sample_support_trace.finalize(
	token_count=len(response_ids),
	extra_rows=final_observation_token_count,
	)
	if agent_loop_state.sample_support_trace is not None:
	if agent_loop_state.sample_support_trace.num_rows:
	rollout_sample_support_out = agent_loop_state.sample_support_trace.finalize(
	token_count=len(response_ids),
	extra_rows=final_observation_token_count,
	)
	else:
	rollout_sample_support_out = np.empty((0, sample_support_width), dtype=SAMPLE_SUPPORT_DTYPE)

[gemini-code-assist]

When batches contains only a single PackedTensor, we can avoid the overhead of torch.cat and cu_seqlens_from_lengths by returning the single batch directly. This is a simple and effective performance optimization.

    @staticmethod
    def cat(batches: Sequence["PackedTensor"]) -> "PackedTensor":
        if not batches:
            raise ValueError("cannot cat an empty list of packed batches")
        if len(batches) == 1:
            return batches[0]
        lengths = torch.cat([batch.sequence_lengths for batch in batches])
        return PackedTensor(
            torch.cat([batch.values for batch in batches], dim=0),
            cu_seqlens_from_lengths(lengths, device=batches[0].device),
        )

[cursor]

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

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

[cursor]

Unclamped logprobs break sample-support wire

High Severity

When return_sample_support is on, sampled logprobs are taken straight from flat processed_logprobs and never passed through build_logprobs_content. Non-finite sampled scores that vLLM emits intermittently therefore reach orjson.dumps, which rejects them and fails the /skyrl/v1/generate response. That breaks the capture path that training relies on for support-aware logprob renormalization.

Additional Locations (1)

• skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py#L516-L548

 

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