Fix MTP cache overflow during generation

modeling_bailing_moe_v3.py

@@ -1273,6 +1273,7 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         output_router_logits: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        run_mtp: bool = False,
         **kwargs,
     ) -> Union[Tuple, MoeV3ModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -1359,7 +1360,11 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
         all_router_logits = () if output_router_logits else None
         next_decoder_cache = None
         layers = self.layers[: -self.num_nextn_predict_layers] if self.num_nextn_predict_layers > 0 else self.layers
-        mtp_layers = self.layers[-self.num_nextn_predict_layers :] if self.num_nextn_predict_layers > 0 else None
+        mtp_layers = (
+            self.layers[-self.num_nextn_predict_layers :]
+            if run_mtp and self.num_nextn_predict_layers > 0
+            else None
+        )
 
         # tptest miss causal_mask = create_causal_mask(
 
@@ -1424,10 +1429,10 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
                         hidden_states,
                         attention_mask,
                         position_ids,
-                        past_key_values,
+                        None,
                         output_attentions,
                         output_router_logits,
-                        use_cache,
+                        False,
                         position_embeddings,
                     )
                 else:
@@ -1436,10 +1441,10 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
                         hidden_states,
                         attention_mask=attention_mask,
                         position_ids=position_ids,
-                        past_key_value=past_key_values,
+                        past_key_value=None,
                         output_attentions=output_attentions,
                         output_router_logits=output_router_logits,
-                        use_cache=use_cache,
+                        use_cache=False,
                         position_embeddings=position_embeddings,
                     )
                 if mtp_hidden_states is None:
@@ -1450,9 +1455,6 @@ class BailingMoeV3Model(BailingMoeV3PreTrainedModel):
                 if output_hidden_states:
                     all_hidden_states += (hidden_states,)
 
-                if use_cache:
-                    next_decoder_cache = layer_outputs[2 if output_attentions else 1]
-
                 if output_attentions:
                     all_self_attns += (layer_outputs[1],)
 
@@ -1554,6 +1556,7 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
             output_router_logits if output_router_logits is not None else self.config.output_router_logits
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        run_mtp = labels is not None and self.num_nextn_predict_layers > 0
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
             input_ids=input_ids,
@@ -1566,6 +1569,7 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
             output_hidden_states=output_hidden_states,
             output_router_logits=output_router_logits,
             return_dict=return_dict,
+            run_mtp=run_mtp,
             **kwargs,
         )
 
@@ -1580,11 +1584,9 @@ class BailingMoeV3ForCausalLM(BailingMoeV3PreTrainedModel, GenerationMixin):
             loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
 
         all_mtp_logits = None
-        if self.num_nextn_predict_layers > 0:
-            mtp_hidden_states = outputs.mtp_hidden_states
+        if run_mtp:
             shift_labels_mtp = None
-            for i in range(self.num_nextn_predict_layers):
-                mtp_hidden_states = mtp_hidden_states[i]
+            for mtp_hidden_states in outputs.mtp_hidden_states:
                 mtp_logits = self.lm_head(mtp_hidden_states).float()
                 if all_mtp_logits is None:
                     all_mtp_logits = []