Enable CUDA Graph for continuous rollouts

[nathon-lee]

Enable CUDA Graph for continuous rollouts

Summary

Enable experimental CUDA Graph decode replay for greedy Hybrid Engine rollouts that use
SamplingConfig.continuous_batch_size.

• Keep a fixed-capacity physical slot layout for each captured graph.
• Preallocate static decode inputs, additive causal masks, position IDs, write positions, and
  DeepSpeedStaticCache rows.
• Refill retired physical slots with prefetched requests while masking inactive rows, so replay always uses stable
  tensor addresses and a fixed batch shape.
• Reuse a captured graph across matching rollout calls. Rebuild it when the model parameter storage/version, device,
  dtype, batch capacity, prompt width, or cache shape changes.
• Use a preallocated four-dimensional additive mask during capture, avoiding Transformers mask-construction
  operations that are not CUDA-graph-safe while preserving the model's configured attention implementation.
• Preserve generated-token attention masks across decode steps and clear
  retired rows before reuse.
• Preserve the existing upstream continuous-batching behavior for aligned decode fronts, cache trimming, legacy prefill, repetition penalty, and statistics.

Requirements and limitations

CUDA Graph capture currently requires CUDA inputs, greedy decoding, one sample per prompt, equal padded prompt
widths, align_decode_fronts=False, and enable_cache_trimming=False. The graph path rejects aligned decode fronts
and cache trimming explicitly because those modes move cache rows or columns, whereas a captured graph requires fixed
physical addresses.

Models that explicitly report no cache-class support remain rejected. Models with unknown support use the legacy
prefill path and should be validated against eager generation before production use.

Validation

pytest -q \
  tests/unit/runtime/rollout/test_continuous_batching.py \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py
# 49 passed

pre-commit run --files \
  deepspeed/runtime/rollout/hybrid_engine_rollout.py \
  docs/code-docs/source/inference-engine.rst \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py
# passed

On an NVIDIA RTX A4500 with PyTorch 2.9.1+cu128 and Transformers 4.57.1, locally initialized small GPT-2 and Llama
models produced token-identical eager and CUDA Graph outputs. Two different GPT-2 prompts reused one capture
(capture_count=1).

The same validation ran with public openai-community/gpt2 pretrained weights in FP16. Two fixed-length,
left-padded prompts with seeds 1234 and 1235 produced token-identical eager and CUDA Graph outputs, including
attention masks, while reusing one capture (capture_count=1).

Performance

The former synthetic continuous-rollout sample included graph capture in every rollout and is not representative of
the reused steady-state path. Production benchmarking should report both first-call latency, which includes capture,
and steady-state reuse separately.

For a local sanity check on the RTX A4500, a randomly initialized medium GPT-2 model with BF16 disabled, batch size
1, prompt length 32, response length 32, two warmup calls, and five steady-state measurements produced:

Mode	Median latency	Throughput
Eager	250.88 ms	127.55 tokens/s
CUDA Graph reuse	31.45 ms	1,017.62 tokens/s

This is a local synthetic sanity check, not a production-model result. Production-model measurements should use fixed
prompts and report first-call and steady-state results independently.

CC @delock @minjiaz @PKUWZP

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-10T10:02:19.995038Z	c280ea3	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c280ea3b13

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".