Conversation
Summary of ChangesHello, 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 enhances the Qwen3-VL-MoE model's inference performance by integrating CUDA graph capabilities into its prefill operations. The changes focus on reducing CPU overhead and improving execution efficiency for repetitive computational steps, such as attention and deepstack feature processing, by allowing these sequences of GPU operations to be captured and replayed as single units. 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. 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 refactors the Qwen3VLMOETransformerLayerInfer class by introducing wrapper functions for _context_attention_kernel and apply_deepstack_features. The _apply_deepstack_features_wrapper_run method specifically adds logic to handle CUDA graph capturing for prefill operations, ensuring proper tensor handling and graph management. It also introduces new imports for InferStateInfo and tensor_to_no_ref_tensor. Feedback suggests that the type hints for infer_state and new_infer_state in the new wrapper method should be more specific (Qwen3VLInferStateInfo) for better type safety, and redundant return statements can be removed for cleaner code.
| def _apply_deepstack_features_wrapper_run( | ||
| self, | ||
| input_embeddings: torch.Tensor, | ||
| infer_state: InferStateInfo, | ||
| layer_num: int, | ||
| ): | ||
| if torch.cuda.is_current_stream_capturing(): | ||
| input_embeddings = input_embeddings.contiguous() | ||
| _input_embeddings = tensor_to_no_ref_tensor(input_embeddings) | ||
| pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() | ||
| pre_capture_graph.__exit__(None, None, None) | ||
|
|
||
| infer_state.prefill_cuda_graph_create_graph_obj() | ||
| infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() | ||
|
|
||
| def apply_func(new_infer_state: InferStateInfo): | ||
| apply_deepstack_features( | ||
| input_embeddings=_input_embeddings, | ||
| infer_state=new_infer_state, | ||
| layer_num=layer_num, | ||
| ) | ||
| return | ||
|
|
||
| infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=apply_func, after_graph=pre_capture_graph) | ||
| else: | ||
| apply_deepstack_features( | ||
| input_embeddings=input_embeddings, | ||
| infer_state=infer_state, | ||
| layer_num=layer_num, | ||
| ) | ||
|
|
||
| return |
There was a problem hiding this comment.
The type hints for infer_state and new_infer_state should be the more specific Qwen3VLInferStateInfo instead of InferStateInfo. The apply_deepstack_features function requires attributes that are specific to Qwen3VLInferStateInfo, so using the correct type hint improves type safety and maintainability. Additionally, the redundant return statements can be removed for cleaner code.
def _apply_deepstack_features_wrapper_run(
self,
input_embeddings: torch.Tensor,
infer_state: Qwen3VLInferStateInfo,
layer_num: int,
):
if torch.cuda.is_current_stream_capturing():
input_embeddings = input_embeddings.contiguous()
_input_embeddings = tensor_to_no_ref_tensor(input_embeddings)
pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph()
pre_capture_graph.__exit__(None, None, None)
infer_state.prefill_cuda_graph_create_graph_obj()
infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__()
def apply_func(new_infer_state: Qwen3VLInferStateInfo):
apply_deepstack_features(
input_embeddings=_input_embeddings,
infer_state=new_infer_state,
layer_num=layer_num,
)
infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=apply_func, after_graph=pre_capture_graph)
else:
apply_deepstack_features(
input_embeddings=input_embeddings,
infer_state=infer_state,
layer_num=layer_num,
)1 participant
No description provided.