agentic RL top-p replay

[willisguo14]

summary

replay SGLang's recorded per-token sampling support when rollout-top-p < 1. (SGLang PR)

scope

agentic RL

--rollout-top-p 0.9
--train-backend megatron
--use-session-server
--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate

out of scope: other generation paths, PP, speculative decoding, PD disaggregation, true-on-policy mode

tests

PYTHONPATH="$REPO/src:$REPO/submodules/miles:$REPO/submodules/sglang/python:$REPO/submodules/harbor/src:$REPO/submodules/megatron-
lm${PYTHONPATH:+:$PYTHONPATH}"

cd "$REPO/submodules/miles"

python -m pytest -q \
  tests/fast/router/test_sessions.py \
  tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py

CUDA_VISIBLE_DEVICES='' TORCH_COMPILE_DISABLE=1 \
python -m pytest -q \
  tests/fast/backends/megatron_utils/test_top_p_replay.py

python -m torch.distributed.run \
  --standalone --nnodes=1 --nproc_per_node=4 \
  tests/e2e/precision/test_top_p_sampling_replay.py

[DavidBellamy]

LGTM