fix tp capture in vit cuda graph#17255
Conversation
narutolhy
requested review from
JustinTong0323,
mickqian and
yhyang201
as code owners
Summary of ChangesHello @narutolhy, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves an issue in the ViT CUDA Graph runner where Tensor Parallelism communication was not correctly integrated into the CUDA Graph capture process, leading to potential instability. By conditionally incorporating capture-aware communicators during graph creation, it ensures proper synchronization of TP communication within the graph, thereby improving the reliability and correctness of ViT model execution under Tensor Parallelism without affecting non-TP or non-graph execution. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 effectively addresses a critical issue concerning CUDA Graph capture in ViT models when Tensor Parallelism is enabled. The introduction of nullcontext and the conditional use of tp_group.ca_comm.capture() ensure that TP communication is properly integrated into the CUDA Graph capture phase, enhancing the correctness and stability of ViT graph execution. The changes are well-scoped and directly resolve the problem described in the motivation.
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
1 similar comment
|
/rerun-failed-ci |
narutolhy
force-pushed
the
fix_vit_cuda_graph_all_reduce
branch
from
7c3fa93 to
577ea21
Compare
|
/rerun-failed-ci |
1 similar comment
|
/rerun-failed-ci |
narutolhy
force-pushed
the
fix_vit_cuda_graph_all_reduce
branch
2 times, most recently
from
114e8a8 to
11bbeed
Compare
|
/tag-and-rerun-ci |
narutolhy
force-pushed
the
fix_vit_cuda_graph_all_reduce
branch
from
11bbeed to
ef6ed25
Compare
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
10 similar comments
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
4 similar comments
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
3 participants
Motivation
When using the ViT CUDA Graph runner with Tensor Parallelism enabled, CUDA Graph capture may implicitly involve TP communication paths (e.g., all-reduce or capture-aware communicators). However, the existing implementation does not align the communicator’s capture lifecycle with torch.cuda.graph(...), which can lead to incorrect capture behavior, runtime errors, or hangs during replay in TP setups.
This PR ensures that TP communication (when capture-aware) is properly included in the CUDA Graph capture phase, improving correctness and stability for ViT graph execution under Tensor Parallelism.
Modifications
Updated python/sglang/srt/multimodal/vit_cuda_graph_runner.py to correctly handle TP communication during CUDA Graph capture.
Introduced nullcontext to provide a zero-overhead fallback when no capture-aware communicator is present.
Retrieved the current TP group via get_tp_group() and conditionally wrapped CUDA Graph capture with tp_group.ca_comm.capture() when available.
The final capture logic becomes:
Use ca_comm.capture() only if the TP group provides a capture-aware communicator.
Preserve the original behavior when no such communicator exists.
This change is strictly scoped to the CUDA Graph creation path and does not affect non-TP or non-graph execution.
Accuracy Tests
covered in test_vlm_vit_cuda_graph.py
Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci