fix

Author: hiworldwzj

lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py

@@ -19,6 +19,7 @@
 from lightllm.models.qwen3_vl.triton_kernel.deepstack_multimodal_emb import apply_deepstack_features
 from lightllm.models.qwen2_vl.layer_infer.transformer_layer_infer import Qwen2VLTransformerLayerInfer
 from lightllm.models.qwen3.triton_kernel.qk_norm import qk_rmsnorm_forward
+from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
 
 
 class Qwen3VLTransformerLayerInfer(Qwen2VLTransformerLayerInfer):
@@ -77,9 +78,42 @@ def context_forward(self, input_embdings, infer_state: Qwen3VLInferStateInfo, la
         if self.tp_world_size_ > 1:
             all_reduce(ffn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False)
         input_embdings.add_(ffn_out.view(-1, self.embed_dim_))
-        apply_deepstack_features(
+        self._apply_deepstack_features_wrapper_run(
             input_embeddings=input_embdings,
             infer_state=infer_state,
             layer_num=self.layer_num_,
         )
         return input_embdings
+
+    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