Fix: Incorrect parameters passed to forward_batch_generation (#5506) (#5511)

This commit is contained in:
u4lr451
2025-04-18 09:49:59 +08:00
committed by GitHub
parent c08a717c77
commit 211c7b31b8

View File

@@ -271,14 +271,11 @@ class EAGLEWorker(TpModelWorker):
)
elif batch.forward_mode.is_idle():
model_worker_batch = batch.get_model_worker_batch()
logits_output, next_token_ids, _ = (
self.target_worker.forward_batch_generation(
ForwardBatch.init_new(
model_worker_batch, self.target_worker.model_runner
)
)
logits_output, next_token_ids = self.target_worker.forward_batch_generation(
model_worker_batch
)
return logits_output, next_token_ids, model_worker_batch.bid, 0, False
return logits_output, next_token_ids, model_worker_batch.bid, 0
else:
logits_output, next_token_ids, bid = self.forward_target_extend(batch)
with self.draft_tp_context(self.draft_model_runner.tp_group):