recover need_capture

Author: STwangyingrui

lightllm/common/basemodel/basemodel.py

@@ -381,9 +381,7 @@ def _decode(
             )
             infer_state.init_some_extra_state(self, padded_model_input.input_ids)
 
-            # Check if a graph needs to be captured.
-            # get_graph returns None if a graph for the batch_size doesn't exist.
-            if self.graph.get_graph(find_graph_batch_size) is None:
+            if self.graph.need_capture(find_graph_batch_size):
                 infer_state.is_cuda_graph = True
                 model_output: ModelOutput = self.graph.capture_decode(
                     self._token_forward, padded_model_input.input_ids, infer_state
@@ -574,9 +572,7 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
             )
             infer_state1.init_some_extra_state(self, padded_model_input1.input_ids)
 
-            # Check if a graph needs to be captured.
-            # get_graph returns None if a graph for the batch_size doesn't exist.
-            if self.graph.get_graph(find_graph_batch_size) is None:
+            if self.graph.need_capture(find_graph_batch_size):
                 infer_state0.is_cuda_graph = True
                 infer_state1.is_cuda_graph = True
 

lightllm/common/basemodel/cuda_graph.py

@@ -49,15 +49,15 @@ def __init__(self, max_batch_size=8, max_len_in_batch=8192):
     def can_run(self, batch_size, max_len_in_batch):
         return batch_size <= self.max_batch_size and max_len_in_batch <= self.graph_max_len_in_batch
 
-    def get_graph(self, batch_size):
+    def need_capture(self, batch_size):
         # We assume batch_size has already been adjusted to the closest supported graph batch size
         # If the graph already exists, get it and move it to the most recently used position.
         if batch_size in self.graph:
             find_graph = self.graph.pop(batch_size)  # Dequeue the graph
             self.graph[batch_size] = find_graph  # Enqueue the graph for LRU
-            return find_graph
+            return False
         else:
-            return None
+            return True
 
     def evict_oldest_graph(self):
         if self.graph: