[MRV2] Add FULL CUDA graph support with PP

[WoosukKwon]

Should be merged after #35162 and #37818

[gemini-code-assist]

Code Review

This pull request introduces full CUDA graph support with pipeline parallelism (PP) in MRV2. The changes primarily involve modifications to cudagraph_utils.py and model_runner.py to handle intermediate tensors and pipeline parallelism during CUDA graph capture and execution. The code adds logic to manage intermediate tensors for non-first PP ranks, ensuring correct data slicing and copying during graph capture and replay. The changes also remove the previous limitation of skipping CUDA graph capture when pipeline parallelism is enabled.

[gemini-code-assist]

The intermediate_tensors_sliced variable is assigned None when self.is_first_pp_rank is true. However, it is used unconditionally in the model_inputs dictionary on line 345. This could lead to a NameError if the model expects intermediate_tensors to always be present.

                intermediate_tensors_sliced = None
            else:
                intermediate_tensors_sliced = None

            model_inputs = {
                "input_ids": input_buffers.input_ids[:num_tokens],
                "positions": input_buffers.positions[:num_tokens],
                "intermediate_tensors": intermediate_tensors_sliced if intermediate_tensors_sliced is not None else None,
                **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
            }

[gemini-code-assist]

The code copies data from intermediate_tensors to self.intermediate_tensors using copy_. This operation is in-place and modifies self.intermediate_tensors. If the CUDA graph is replayed multiple times with different intermediate_tensors, the initial data in self.intermediate_tensors will be overwritten, potentially leading to incorrect results in subsequent iterations. This is especially problematic if the graph is captured with one set of intermediate tensors and replayed with another.

                    k: v[:n].clone() # Create a copy to avoid modifying the original tensor
                    for k, v in intermediate_tensors.tensors.items()
                },
                intermediate_tensors.kv_connector_output,

[mergify]

Hi @WoosukKwon, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

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

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

[WoosukKwon]

merged with #35162