Set history_cross_kv_seqlens to 0 by default (InternLM#2666)

irexyc · Oct 28, 2024 · a8b2765 · a8b2765
1 parent 538f618
commit a8b2765
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py
@@ -385,7 +385,7 @@ class SchedulerSequence:
     mrope_position_ids: Optional[Tensor] = None
     mrope_position_delta: Optional[int] = None
     cross_attention_states: Optional[Tensor] = None
-    history_cross_kv_seqlens: Optional[int] = None
+    history_cross_kv_seqlens: int = 0
 
     def __post_init__(self):
         """post init."""
@@ -489,12 +489,11 @@ def num_all_tokens(self):
 
     def num_all_cross_tokens(self):
         """num of all cross tokens."""
-        if self.history_cross_kv_seqlens is None:
-            if self.cross_attention_states is None:
-                self.history_cross_kv_seqlens = 0
-            else:
-                self.history_cross_kv_seqlens = self.cross_attention_states.shape[  # noqa
-                    -2]
+        if self.cross_attention_states is None:
+            self.history_cross_kv_seqlens = 0
+        else:
+            self.history_cross_kv_seqlens = self.cross_attention_states.shape[
+                -2]
         return self.history_cross_kv_seqlens
 
     def update_token_ids(self,