Fixed `past_key_values` evaluation in `MiniCPMModel.forward`

modeling_minicpm.py

@@ -1954,18 +1954,21 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
         past_key_values_length = 0
 
         if use_cache:
-            use_legacy_cache = not isinstance(past_key_values, Cache)
-            if use_legacy_cache:
+            # Reject old tuple-style cache, but allow None (first forward pass)
+            if past_key_values is not None and not isinstance(past_key_values, Cache):
                 raise ValueError(
                     'You must use the new past_key_values format, such as the Cache class, instead of the old tuple format.'
                 )
-
+ 
+            # Initialize cache if None (first forward pass)
+            if past_key_values is None:
+                if self.config.sparse_config is not None and torch.cuda.is_available():
+                    past_key_values = InfLLMv2Cache(config=self.config, num_hidden_layers=self.config.num_hidden_layers)
+                else:
+                    past_key_values = DynamicCache()
+ 
             # Calculate the usable length of past key values
-            past_key_values_length = past_key_values.get_seq_length() if isinstance(past_key_values, InfLLMv2Cache) else 0
-
-            # Initialize InfLLMv2Cache if needed
-            if self.config.sparse_config is not None and torch.cuda.is_available() and past_key_values_length == 0:
-                past_key_values = InfLLMv2Cache(config = self.config, num_hidden_layers=self.config.num_hidden_layers)
+            past_key_values_length = past_key_values.get_seq_length()
 
         if position_ids is None:
             device = input_ids.device if input_ids is not None else inputs_embeds.device
@@ -2047,7 +2050,7 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
 
         next_cache = None
         if use_cache:
-            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+            next_cache = next_decoder_cache
         if not return_dict:
             return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
         return BaseModelOutputWithPast(