Clean up eagle code (#2756)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user