support transformers 4.55

modeling_phi3.py

@@ -330,7 +330,14 @@ class Phi3Attention(nn.Module):
                     "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                     "with a layer index."
                 )
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            if not hasattr(past_key_value, 'get_usable_length'):
+                # Transformers >= 4.55
+                past_len = past_key_value.get_seq_length(self.layer_idx) if past_key_value is not None else 0
+                kv_seq_len += past_len
+            else:
+                # Transformers < 4.55
+                kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+
         cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
 
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
@@ -448,7 +455,14 @@ class Phi3FlashAttention2(Phi3Attention):
                     "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                     "with a layer index."
                 )
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            if not hasattr(past_key_value, 'get_usable_length'):
+                # Transformers >= 4.55
+                past_len = past_key_value.get_seq_length(self.layer_idx) if past_key_value is not None else 0
+                kv_seq_len += past_len
+            else:
+                # Transformers < 4.55
+                kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+
 
         # Because the input can be padded, the absolute sequence length depends on the max position id.
         rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)
@@ -738,7 +752,14 @@ class Phi3SdpaAttention(Phi3Attention):
 
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            if not hasattr(past_key_value, 'get_usable_length'):
+                # Transformers >= 4.55
+                past_len = past_key_value.get_seq_length(self.layer_idx) if past_key_value is not None else 0
+                kv_seq_len += past_len
+            else:
+                # Transformers < 4.55
+                kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+
         cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
 
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
@@ -1060,7 +1081,15 @@ class Phi3Model(Phi3PreTrainedModel):
             use_legacy_cache = not isinstance(past_key_values, Cache)
             if use_legacy_cache:
                 past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-            past_key_values_length = past_key_values.get_usable_length(seq_length)
+                if isinstance(past_key_values, Cache) and not hasattr(past_key_values, 'get_usable_length'):
+                    # Transformers >= 4.55
+                    past_key_values_length = past_key_values.get_seq_length()
+                elif isinstance(past_key_values, Cache):
+                    # 4.49 <= Transformers < 4.55
+                    past_key_values_length = past_key_values.get_usable_length(seq_length)
+                else:
+                    # No cache given on first forward, keep length at 0
+                    past_key_values_length = 0
 
         if position_ids is None:
             device = input_ids.device if input_ids is not None else inputs_embeds.device
@@ -1288,19 +1317,35 @@ class Phi3ForCausalLM(Phi3PreTrainedModel):
         # When the first time input length reached long and short factor switching point, enforce re-compute cache
         # It will cause downside of slower at this single token position, however, better than current failure.
         if past_key_values and self.config.rope_scaling and input_ids.shape[1] >= self.config.original_max_position_embeddings + 1:
-            past_length = past_key_values.seen_tokens if isinstance(past_key_values, Cache) else past_key_values[0][0].shape[2]
+            if isinstance(past_key_values, Cache) and not hasattr(past_key_values, 'seen_tokens'):
+                # Transformers > 4.55
+                cache_length = past_key_values.get_seq_length()
+                past_length = cache_length
+            elif isinstance(past_key_values, Cache):
+                # 4.49 <= Transformers < 4.55
+                past_length = past_key_values.seen_tokens
+            else:
+                # Transformers < 4.49
+                past_key_values[0][0].shape[2]
             if past_length <= self.config.original_max_position_embeddings:
                 past_key_values = None
 
+        cache_length = None
+        past_length = None
+        max_cache_length = None
         if past_key_values is not None:
-            if isinstance(past_key_values, Cache):
+            if isinstance(past_key_values, Cache) and not hasattr(past_key_values, 'seen_tokens'):
+                # Transformers > 4.55
+                cache_length = past_key_values.get_seq_length()
+                past_length = cache_length
+                max_cache_length = past_key_values.get_max_cache_shape()
+            elif isinstance(past_key_values, Cache):
+                # 4.49 <= Transformers < 4.55
                 cache_length = past_key_values.get_seq_length()
                 past_length = past_key_values.seen_tokens
-                # Fixing AttributeError: 'DynamicCache' object has no attribute 'get_max_length'
-                # https://github.com/huggingface/transformers/issues/36071
-                # max_cache_length = past_key_values.get_max_length()
                 max_cache_length = past_key_values.get_max_cache_shape()
             else:
+                # Transformers < 4.49
                 cache_length = past_length = past_key_values[0][0].shape[2]
                 max_cache_length = None