[REQUEST] Temperature / top-k / top-p sampling on HybridEngineRollout's CUDA graph and continuous batching paths #8934

Is your feature request related to a problem? Please describe.

HybridEngineRolloutConfig(use_graph_capture=True) only takes effect for greedy decoding. In HybridEngineRollout.generate() the graph path is chosen with use_graph_capture and is_greedy. With any temperature > 0, the rollout silently runs Hugging Face generate() instead.

Continuous batching (SamplingConfig.continuous_batch_size, #8368), with or without its CUDA graph (#8732), is greedy only as well.

RL rollouts (GRPO, PPO, on-policy distillation with sampled responses) sample at temperature 1. For them, graph capture and continuous batching are unavailable. Enabling non-greedy sampling would allow RL rollout generation to be greatly accelerated by these features.

Describe the solution you'd like

Only the decode forward is captured in the graph. The next token is chosen on the host from the graph's output logits after each replay (static_logits[:, -1, :].argmax(...)). So sampling needs no change to the captured graph:

• generate() sends every non-continuous request to _generate_graph when use_graph_capture=True.
• One token selector, _token_selector(sampling), picks the first token (from the prefill logits) and each decoded token on all three paths: graph decode, continuous batching, and continuous batching with graphs.
  • argmax when temperature <= 0, as today;
  • otherwise, Hugging Face's TemperatureLogitsWarper, TopKLogitsWarper and TopPLogitsWarper, applied in generate()'s order on float32 logits, then torch.multinomial. This is the same distribution generate() samples from.
• Continuous batching keeps applying the model's repetition penalty, now before the warpers, which is generate()'s order. The temperature check in _validate_continuous_inputs is removed. Models without cache-class support, which prefill through module.generate(), get the same sampling arguments as the main generate() call, through one shared helper.
• Documentation: a short "CUDA Graph Decode" section in inference-engine.rst. It also notes that other settings from the model's generation_config (for example a repetition penalty) are not applied on the graph path, which is already true for greedy graph decode.

Planned tests (dispatch and sampling correctness)

• Dispatch: with use_graph_capture=True, sampled requests go to the graph path instead of generate(), and continuous batching accepts temperature > 0 and samples (CPU).
• The token selector draws from the distribution SamplingConfig describes: 20,000 draws from fixed logits against an independently computed reference, for temperature, top-k, top-p and their combination. No token outside the top-k / top-p support may appear, and every token's frequency must lie within 4 standard errors (CPU).
• Continuous batching applies the repetition penalty before the warpers (CPU, deterministic).
• End to end on GPU, on a tiny random Llama: sampled tokens on each of the three paths match the model's own distribution, checked against a plain forward pass.

Known issues outside this change

On master, two cache and mask gaps stop some paths from running on a Hugging Face model with SDPA attention, greedy included. This change doesn't touch either; #8602 fixes both.

• Plain graph path (_generate_graph), Transformers 4.57.x: the decode step's 2-D attention mask makes Transformers call padding_mask.all() on the host during capture (reported by 0z5a on [REQUEST] OPSD Profile and improve HybridEngine rollout performance #8197). feat(rollout): segment-based kernel injection for hybrid-engine greedy decode #8602 passes a pre-built 4-D decode mask.
• Plain graph path and continuous batching without graphs, Transformers 5.x: DeepSpeedStaticCache has no get_query_offset (also Support Transformers 5.20 cache query offsets in DeepSpeedStaticCache #8852). feat(rollout): segment-based kernel injection for hybrid-engine greedy decode #8602 adds it.

Continuous batching with graphs builds its own 4-D mask and runs on both versions. Until #8602 lands, the end-to-end GPU tests of the affected paths fail with these errors.

Possible future improvement

Sample inside the graph. Here the token is still picked on the host after each replay: a handful of small kernels per step, outside the captured forward. Capturing the selection too would:

• remove that per-step launch overhead;
• open the way to capturing several decode steps in one graph.

It needs graph-safe RNG state for torch.multinomial and graph-safe top-k / top-p filtering.

Additional context

• Related: feat(rollout): add continuous batching generation prototype #8368 and Enable CUDA Graph for continuous rollouts #8732. [REQUEST] OPSD Profile and improve HybridEngine rollout performance #8197 lists "improving sampled-decode performance" as a follow-up direction.