[Model Runner V2] support piecewise & mixed cudagraph

[izhuhaoran]

Purpose

As titled, this PR supports piecewise & mixed cudagraph for model runner v2

[gemini-code-assist]

Code Review

This pull request introduces support for piecewise and mixed CUDA graphs in the v2 model runner. The changes are well-structured, refactoring the CUDA graph capture logic to handle different modes (FULL, PIECEWISE, FULL_AND_PIECEWISE, etc.) more cleanly. The runtime dispatch logic in the model runner is also updated accordingly. While the implementation for piecewise graphs seems correct, I've found a critical issue in the full graph capture implementation concerning LoRA support.

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

Comment @cursor review or bugbot run to trigger another review on this PR

[LucasWilkinson]

This PR seems to assume that PIECEWISE cudagraphs and FULL cudagprahs will have the same sizes; FULL cudagraphs are upper-bounded by max_num_seqs (or in the case of spec-decode max_num_seqs * (1 + num_speculative_tokens)) while PIECEWISE cudagraphs are upper bounded by max_cudagraph_capture_size. Atleast in V1, for performance I think this should be preserved (doesnt make sense to cut PIECEWISE cudagraphs off at 256 when currently then go up to 512 or 1024.

[izhuhaoran]

This PR seems to assume that PIECEWISE cudagraphs and FULL cudagprahs will have the same sizes; FULL cudagraphs are upper-bounded by max_num_seqs (or in the case of spec-decode max_num_seqs * (1 + num_speculative_tokens)) while PIECEWISE cudagraphs are upper bounded by max_cudagraph_capture_size. Atleast in V1, for performance I think this should be preserved (doesnt make sense to cut PIECEWISE cudagraphs off at 256 when currently then go up to 512 or 1024.

Thanks for this suggestion, already fixed !

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[mergify]

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

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

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[izhuhaoran]

Note: after merging main, runtime errors appear; they’ll be resolved in #33004.

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[njhill]

Thanks a lot for this @izhuhaoran! Great work

I tested it and it gives a huge speedup on blackwell with a small model / decode-heavy workload.

[izhuhaoran]

@njhill Thanks for your review, I've updated the codes. PTAL when you have time.

[njhill]

formatting nits to reduce loc

[izhuhaoran]

formatting nits to reduce loc

Thanks for these suggestions, already reformatted.

[mergify]

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

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

[LucasWilkinson]

Thanks for the contribution! @izhuhaoran

I don't think we need to do it now but we should start thinking about how we can more uniquely identify cudagraphs based on the features/configuration of the captured graph. I dont think num_tokens is sufficient or explicit enough.

Currently this PR just maintains 2 types of batches, uniform and non-uniform and handles that maintaining 2 lists of tokens counts, i.e.

if uniform_decode and self.uniform_decode_cudagraph_sizes:
    return self.uniform_decode_cudagraph_sizes.get(num_tokens)
return self.cudagraph_sizes.get(num_tokens)

If more "batch types" are needed this doesn't feel very scalable, e.g. of features that would have different "batch types" are:

1. lora: Currently (MRV1) we support --specialize-active-lora which captures different cudagraphs for different active LORA counts
2. dynamic spec-decode: we will likely want to capture different uniform graphs with same token counts but a different number of requests, e.g. for we may want 2 graphs for num_tokens = 8, one with 4 requests for num_speculated_tokens = 1, and one with 2 requests for num_speculated_tokens = 3

In MRV1 this is handled by mapping BatchDescriptor -> cudagraph but is not super clean (matching gets a bit cleaner in #34102). I think its worthing thinking a bit into the future about how we might support these cases more naturally.

cc @njhill @WoosukKwon

[WoosukKwon]

@LucasWilkinson Thanks for bringing it up. I do agree that CUDA graph needs more design discussions. I'm accepting this PR to move us forward, but we must revisit this or next week.

[WoosukKwon]

Thanks for the PR! The code looks clean and well-structured, and I think it’s a solid implementation given the current CUDA graph design. Great work! 👍

[izhuhaoran]

@LucasWilkinson @WoosukKwon Thanks for the review and sorry for the late reply due to Chinese New Year. Yes, this PR is only based on the current CUDA graph design, and I agree that it needs further improvement.