[RL] Enable bounded top-p sampling-support replay

[nanjiangwill]

Summary

• Enable sampling-support replay only when --rollout-top-p is below 1; require a positive --rollout-top-k as the hard support bound.
• Request SGLang realized support and support-normalized rollout log-probability across standard, legacy, multi-turn, tool-call, and session rollout paths.
• Apply the captured support to actor forward scoring and policy loss in FSDP and Megatron.
• Keep top_p=1 on the existing path: no capture request, sampling-mask transport, reconstruction, or actor masking, including when top_k is finite.
• Reject training requests whose sampling distribution cannot be replayed faithfully, including unbounded top-p, mismatched temperature, unsupported logit transformations, incomplete support, or a router that does not preserve the native sampling-mask request.

Why

SGLang samples each token from the distribution remaining after its sampling filters. Scoring that token with a full-vocabulary actor softmax is a different policy. Correct policy ratios require both sides to use the same realized action space: SGLang returns the support and its normalized rollout log-probability, and Miles normalizes actor log-probability over that support.

SGLang remains the source of truth; Miles does not reconstruct top-p or top-k filtering. Run-level top_k only bounds the returned support and transport size. It does not enable replay by itself. Per-request top-p may vary below 1, while request top-k must stay within the configured bound and temperature must match actor scoring.

Reference, teacher, value, evaluation, and entropy paths remain unchanged. Forced tool or environment tokens use singleton support because they were not sampled by the policy.

API changes

• --rollout-top-p < 1 currently requires --use-miles-router; argument validation fails before startup otherwise because the model gateway does not forward return_sampling_mask for chat-completions requests.
• PD disaggregation is not enabled for this mode because Miles router currently rejects PD. server_group.py carries TODO(#2596) beside the engine environment setup with the bounded SGLang metadata-buffer configuration required when Miles-router PD support is added.
• compute_request_payload and the legacy rollout generate accept a keyword-only evaluation flag so evaluation can skip training-only capture.
• Megatron forward_only accepts use_rollout_sampling_mask because only actor scoring consumes captured support; reference, teacher, and value forwards do not.

Series

1. [RL] Add sampling-support log-prob primitives #2200 adds optional sampling-support log-probability primitives.
2. [RL] Represent and transport rollout sampling support #2595 adds the compact sample representation and trainer transport.
3. [RL] Enable bounded top-p sampling-support replay #2596 (this PR) enables bounded top-p capture and actor replay.

Both prerequisite PRs are merged; this PR is one commit on current main.

Validation

• Repository pre-commit hooks pass on the complete diff.
• Focused CPU tests pass for sampling-support capture, CSR transport, session/tool merging, request validation, and actor scoring.
• The registered 2×H200 Qwen3-0.6B E2E passed against current main: live SGLang TP=2 rollout, sampling-mask transport, Megatron TP=2 actor scoring, training, and weight update. The actor/rollout log-probability absolute difference was 0.00756 against the 0.03 CI threshold.
• The E2E uses the existing Miles router because it preserves SGLang-native request extensions. The SGLang model-gateway chat-completions schema currently drops return_sampling_mask; Miles rejects that configuration during argument validation and also fails closed if required response metadata is absent.

Acknowledgements

Thank you to Zilin Zhu for adding top-p masking support to slime in THUDM/slime#2102.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.