dev_fix_mtp_forward_meta (#5976)

2026-04-23 00:17:25 +08:00 · 2026-01-10 00:40:56 +08:00
parent 9db48ecb34
commit 62bd92f9ba
1 changed files with 3 additions and 1 deletions
@@ -1006,7 +1006,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"],