Clean up eagle code (#2756)

This commit is contained in:
Lianmin Zheng
2025-01-06 14:54:18 -08:00
committed by GitHub
parent 2855caa481
commit b8574f6953
7 changed files with 138 additions and 128 deletions

View File

@@ -962,10 +962,13 @@ class Scheduler:
self.tp_worker.forward_batch_generation(model_worker_batch)
)
else:
logits_output, next_token_ids, model_worker_batch, spec_info = (
self.draft_worker.forward_batch_speculative_generation(batch)
)
batch.spec_info = spec_info
(
logits_output,
next_token_ids,
model_worker_batch,
num_accepted_tokens,
) = self.draft_worker.forward_batch_speculative_generation(batch)
self.num_generated_tokens += num_accepted_tokens
elif batch.forward_mode.is_idle():
model_worker_batch = batch.get_model_worker_batch()
self.tp_worker.forward_batch_idle(model_worker_batch)