qwen3_vl support prefill cuda graph feature

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -515,9 +515,27 @@ def _context_forward(self, input_ids, infer_state: InferStateInfo):
         input_embs = pre_method(cuda_input_ids, infer_state, self.pre_post_weight)
         input_tensors = [input_embs]
 
-        def prefill_func(input_tensors, infer_state):
+        # prefill cuda graph 在 qwen3 vl 上的前几层由于特殊的处理，导致目前无法支持cuda graph
+        from lightllm.utils.config_utils import is_qwen3_vl
+
+        if is_qwen3_vl():
+            no_graph_layer_num = 3
+        else:
+            no_graph_layer_num = 0
+
+        def no_graph_prefill_func(input_tensors, infer_state):
+            _input_embs = input_tensors[0]
+            for i in range(no_graph_layer_num):
+                layer = self.layers_infer[i]
+                layer_method = (layer.context_forward, layer.tpsp_context_forward)[run_mode_index]
+                _input_embs = layer_method(_input_embs, infer_state, self.trans_layers_weight[i])
+            return [_input_embs]
+
+        input_tensors = no_graph_prefill_func(input_tensors=input_tensors, infer_state=infer_state)
+
+        def graph_prefill_func(input_tensors, infer_state):
             _input_embs = input_tensors[0]
-            for i in range(self.layers_num):
+            for i in range(no_graph_layer_num, self.layers_num):
                 layer = self.layers_infer[i]
                 layer_method = (layer.context_forward, layer.tpsp_context_forward)[run_mode_index]
                 _input_embs = layer_method(_input_embs, infer_state, self.trans_layers_weight[i])
@@ -531,7 +549,7 @@ def prefill_func(input_tensors, infer_state):
             )
             if self.prefill_graph.need_capture(handle_token_num=finded_handle_token_num):
                 output_tensors: List[torch.Tensor] = self.prefill_graph.capture_prefill(
-                    prefill_func=prefill_func,
+                    prefill_func=graph_prefill_func,
                     input_tensors=input_tensors,
                     infer_state=infer_state,
                 )
@@ -542,7 +560,8 @@ def prefill_func(input_tensors, infer_state):
 
         else:
             g_cache_manager.cache_env_in()
-            output_tensors: List[torch.Tensor] = prefill_func(input_tensors, infer_state)
+            input_tensors = no_graph_prefill_func(input_tensors=input_tensors, infer_state=infer_state)
+            output_tensors: List[torch.Tensor] = graph_prefill_func(input_tensors, infer_state)
             g_cache_manager.cache_env_out()
 
         input_embs = output_tensors[0]

lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py

@@ -70,7 +70,7 @@ def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_wei
         input1 = None
         self._post_cache_kv(cache_kv, infer_state, layer_weight)
 
-        o = self.__context_attention_wrapper_run(
+        o = self._context_attention_wrapper_run(
             q=q, cache_kv=cache_kv, infer_state=infer_state, layer_weight=layer_weight
         )
 
@@ -116,7 +116,7 @@ def tpsp_context_forward(self, input_embdings: torch.Tensor, infer_state: InferS
         input1 = None
         self._post_cache_kv(cache_kv, infer_state, layer_weight)
 
-        o = self.__context_attention_wrapper_run(
+        o = self._context_attention_wrapper_run(
             q=q, cache_kv=cache_kv, infer_state=infer_state, layer_weight=layer_weight
         )
 
@@ -148,7 +148,7 @@ def tpsp_token_forward(self, input_embdings: torch.Tensor, infer_state: InferSta
         input_embdings.add_(ffn_out.view(-1, self.embed_dim_))
         return input_embdings
 
-    def __context_attention_wrapper_run(
+    def _context_attention_wrapper_run(
         self, q: torch.Tensor, cache_kv: torch.Tensor, infer_state: InferStateInfo, layer_weight
     ) -> torch.Tensor:
         if torch.cuda.is_current_stream_capturing():

lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py

@@ -63,7 +63,7 @@ def context_forward(self, input_embdings, infer_state: Qwen3VLInferStateInfo, la
         q, cache_kv = self._get_qkv(input1, infer_state, layer_weight)
         input1 = None
         self._post_cache_kv(cache_kv, infer_state, layer_weight)
-        o = self._context_attention_kernel(q, cache_kv, infer_state, layer_weight)
+        o = self._context_attention_wrapper_run(q, cache_kv, infer_state, layer_weight)
         q = None
         o = self._get_o(o, infer_state, layer_weight)
         if self.tp_world_size_ > 1:

lightllm/utils/config_utils.py

@@ -132,3 +132,16 @@ def get_fixed_kv_len():
         return len(model_cfg["prompt_cache_token_ids"])
     else:
         return 0
+
+
+@lru_cache(maxsize=None)
+def is_qwen3_vl():
+    from lightllm.utils.llm_utils import get_llm_model_class
+    from lightllm.models import Qwen3VLTpPartModel, Qwen3VLMOETpPartModel
+
+    model_class = get_llm_model_class()
+
+    if model_class in [Qwen3VLTpPartModel, Qwen3VLMOETpPartModel]:
+        return True
+    else:
+        return False