[Fix] Avoid applying cuda graph input-buffer registry on non-cuda devices - #27549
Conversation
|
/tag-run-ci-label |
![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 restricts the eager input-buffer view creation in forward_decode, forward_extend, and forward_idle to when the device is 'cuda'. The reviewer noted that using self.device == 'cuda' is too restrictive as it disables the eager input-buffer registry on other platforms that support CUDA graphs, such as 'musa' or out-of-tree platforms. They suggested broadening the check to support these platforms consistently with the rest of the codebase.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| pp_proxy_tensors=None, | ||
| ) -> Union[LogitsProcessorOutput, PPProxyTensors]: | ||
| if not self.server_args.enable_pdmux: | ||
| if not self.server_args.enable_pdmux and self.device == 'cuda': |
There was a problem hiding this comment.
Using self.device == 'cuda' is too restrictive and will disable the eager input-buffer registry on other platforms that support CUDA graphs, such as musa or out-of-tree platforms. To maintain consistency with the rest of the codebase (e.g., the CUDA graph recapture check at line 1719), we should allow all platforms that support CUDA graphs.
| if not self.server_args.enable_pdmux and self.device == 'cuda': | |
| if not self.server_args.enable_pdmux and (self.device in ["cuda", "musa"] or (current_platform.is_out_of_tree() and current_platform.support_cuda_graph())): |
| return (ret, can_run_graph) | ||
|
|
||
| if not self.server_args.enable_pdmux: | ||
| if not self.server_args.enable_pdmux and self.device == 'cuda': |
There was a problem hiding this comment.
Using self.device == 'cuda' is too restrictive and will disable the eager input-buffer registry on other platforms that support CUDA graphs, such as musa or out-of-tree platforms. To maintain consistency with the rest of the codebase (e.g., the CUDA graph recapture check at line 1719), we should allow all platforms that support CUDA graphs.
| if not self.server_args.enable_pdmux and self.device == 'cuda': | |
| if not self.server_args.enable_pdmux and (self.device in ["cuda", "musa"] or (current_platform.is_out_of_tree() and current_platform.support_cuda_graph())): |
| # indices and trigger SWA mapping use-after-free. | ||
| if forward_batch.batch_size > 0: | ||
| if not self.server_args.enable_pdmux: | ||
| if not self.server_args.enable_pdmux and self.device == 'cuda': |
There was a problem hiding this comment.
Using self.device == 'cuda' is too restrictive and will disable the eager input-buffer registry on other platforms that support CUDA graphs, such as musa or out-of-tree platforms. To maintain consistency with the rest of the codebase (e.g., the CUDA graph recapture check at line 1719), we should allow all platforms that support CUDA graphs.
| if not self.server_args.enable_pdmux and self.device == 'cuda': | |
| if not self.server_args.enable_pdmux and (self.device in ["cuda", "musa"] or (current_platform.is_out_of_tree() and current_platform.support_cuda_graph())): |
|
Hi @ch-wan would you help review this PR? Does the approach that Gemini suggests make more sense? Thanks. |
|
Based on codex, the error is caused by |
2 participants
Motivation
The feature introduced in PR #27407 is not compatible with non-CUDA devices (at least not for Xeon CPU). Xeon CI failure is observed after the merging of this PR.
Modifications
Add gating logics so that the feature can be applied only on cuda device.
Accuracy Tests
N/A
Speed Tests and Profiling
N/A
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): 🚫 Run #27127564313
Latest PR Test (Extra): ❌ Run #27127564042