Revert "[TSP] last_norm allgather move to model.py (#5924)" (#5961)

This reverts commit 2bb838fed9.
This commit is contained in:
Yuanle Liu
2026-01-09 15:20:40 +08:00
committed by GitHub
parent 1d20957340
commit 8c3513a410
9 changed files with 8 additions and 30 deletions
+2 -2
View File
@@ -1004,7 +1004,7 @@ class MTPProposer(Proposer):
)
# 4. Compute logits, Sample
logits = self.model.compute_logits(hidden_states, forward_meta=self.forward_meta)
logits = self.model.compute_logits(hidden_states)
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"])
@@ -1118,7 +1118,7 @@ class MTPProposer(Proposer):
model_output, self.model_inputs["cum_offsets"], self.forward_meta, self.model_inputs
)
# 4. Compute logits, Sample
logits = self.model.compute_logits(hidden_states, forward_meta=self.forward_meta)
logits = self.model.compute_logits(hidden_states)
sampled_token_ids, sampler_output = self.sampler(
logits,
self.sampling_metadata,