[Model Runner V2] Enable piecewise & full CUDA graphs for pipeline parallelism

[ZhanqiuHu]

Summary

Add piecewise CUDA graph capture/replay support for PP in V2 model runner.

Related: #33960

model_runner.py:

• Enable PP cudagraph mode handling
• Add IntermediateTensors buffer during graph replay
• Copy received tensors into the buffer at runtime

cudagraph_utils.py:

• Add intermediate_tensors through capture pipeline
• Handle IntermediateTensors output on non-last PP ranks
• Fix num_reqs divisibility for uniform query length backends

Purpose

V2 model runner did not support CUDA graph capture with PP, falling back to eager mode. This PR adds piecewise CUDA graph capture for PP.

Question

With CUDA graph capture, I ran into AssertionError: TRTLLM decode requires uniform query lengths per request (flashinfer.py:1109) when num_tokens % num_reqs != 0. So I added this workaround to ensure divisibility, but not sure if this is the right approach.

if num_reqs > 0 and num_tokens > num_reqs and num_tokens % num_reqs != 0:
    tokens_per_req = cdiv(num_tokens, num_reqs)
    num_reqs = num_tokens // tokens_per_req
    if num_tokens % num_reqs != 0:
        num_reqs = 1

Test plan

export MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8

# 1. Serve with V2 + PP=2 (piecewise CG, the default)
VLLM_USE_V2_MODEL_RUNNER=1 vllm serve $MODEL -pp 2 --enable-expert-parallel --max-num-seqs 128 --no-enable-prefix-caching

# 2. Throughput benchmark
vllm bench serve --model $MODEL --dataset-name random \
  --random-input-len 2 --random-output-len 512 --num-prompts 128 --num-warmups 16

# 3. Accuracy eval (GSM8K 5-shot)
lm_eval --model local-completions \
  --model_args "model=$MODEL,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=1024" \
  --tasks gsm8k --num_fewshot 5

Test results

Config: Qwen3-30B-A3B-Thinking-2507-FP8, PP=2, 2×B200, --max-num-seqs 128

Performance (128 prompts, input=2, output=512)

Model Runner	CUDA Graph	Req/s	Output tok/s	TTFT (ms)	TPOT (ms)
V2	Eager	13.89	7,111	231	17.5
V2	Piecewise	23.07	11,810	167	10.4
V1 (baseline)	Piecewise	23.23	11,896	165	10.3

Accuracy (GSM8K, 5-shot)

Model Runner	CUDA Graph	strict-match	flexible-extract
V2	Eager	0.7824	0.6732
V2	Piecewise	0.7862	0.6778
V1 (baseline)	Piecewise	0.7824	0.6717

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request enables piecewise CUDA graph capture for pipeline parallelism in the V2 model runner, which was previously unsupported and forced eager mode execution. The changes are well-structured and correctly handle the complexities of this feature. Key changes include downgrading CUDA graph modes for compatibility with pipeline parallelism, introducing a persistent buffer for intermediate tensors to ensure stable memory addresses for graph replay, and updating the graph capture utilities to handle pipeline parallelism constructs. I have one suggestion to improve robustness by adding an assertion to verify the consistency of intermediate tensors between pipeline stages.

[yewentao256]

Thanks for the work!

[yewentao256]

Could you also try #34903 this PR?
Not sure if it is the same issue for enabling full cuda graph

[ZhanqiuHu]

Could you also try #34903 this PR? Not sure if it is the same issue for enabling full cuda graph

Cool! will take a look 👍

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZhanqiuHu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ZhanqiuHu]

Could you also try #34903 this PR? Not sure if it is the same issue for enabling full cuda graph

I think for now the error I hit was only the uniform query lengths assertion with TRTLLM decode, I don't think it has illegal memory access from misaligned attention metadata at the moment.

I will keep that issue in mind, and I think I might run into it later with full cuda graph support.

[ZhanqiuHu]

Friendly ping @WoosukKwon @njhill: This PR has been rebased and is ready for review. Would appreciate any feedback when you get a chance. Thanks!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZhanqiuHu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[WoosukKwon]

LGTM! Thanks for the PR and sorry for the late review.

I've edited the PR for faster merge. Will follow up with the FULL graph support.

[WoosukKwon]

Actually, I ended up merging with #37821 for easier testing.

[ZhanqiuHu]

LGTM! Thanks for the PR and sorry for the late review.

I've edited the PR for faster merge. Will follow up with the FULL graph support.

That sounds great! Thank you!