2.4_fix_mtp_forward_meta (#5977)

2026-04-22 16:07:51 +08:00 · 2026-01-10 00:41:36 +08:00
parent 3e69022698
commit bbe9731f46
1 changed files with 3 additions and 1 deletions
@@ -1014,7 +1014,9 @@ class MTPProposer(Proposer):
                # 4. Compute logits, Sample
                logits = self.model.compute_logits(hidden_states, forward_meta=self.forward_meta)
                if self.enable_logprob and self.enable_draft_logprob and substep == 0:
-                    first_token_logits = self.model.compute_logits(self.model_inputs["first_token_hidden_states"])
+                    first_token_logits = self.model.compute_logits(
+                        self.model_inputs["first_token_hidden_states"], forward_meta=self.forward_meta
+                    )

                    speculate_get_logits(
                        self.model_inputs["draft_logits"],