[Model Runner V2] Refactor CudaGraphManager #29583
Conversation
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the CudaGraphManager to improve its structure and generalize CUDA graph capturing. The logic for determining CUDA graph sizes, preparing inputs for capture, and orchestrating the capture process is extracted into separate utility functions. This improves modularity and readability. The key change is the generalization from capturing graphs based on batch_size to num_tokens, which allows for more flexible support for different batch types (prefill, decode, mixed). I found one critical issue in the implementation that needs to be addressed.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ 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".
Uh oh!
There was an error while loading. Please reload this page.