fix: b_ready_cache_len default value (#366)

Co-authored-by: baishihao <baishihao@sensetime.com>
ModelTC · Mar 20, 2024 · 9fed5e9 · 9fed5e9
1 parent 859a48b
commit 9fed5e9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
@@ -201,7 +201,10 @@ def _prefill(
         infer_state.b_req_idx = b_req_idx
         infer_state.b_start_loc = b_start_loc
         infer_state.b_seq_len = b_seq_len
-        infer_state.b_ready_cache_len = b_ready_cache_len
+        if b_ready_cache_len is not None:
+            infer_state.b_ready_cache_len = b_ready_cache_len
+        else:
+            infer_state.b_ready_cache_len = torch.zeros_like(b_seq_len, dtype=b_seq_len.dtype, device=b_seq_len.device)
         infer_state.multimodal_params = multimodal_params
 
         infer_state.mem_manager = self.mem_manager