refactor: remove ModelWorkerBatch indirection (#25516)
This commit is contained in:
@@ -457,8 +457,7 @@ def extend(reqs, model_runner):
|
||||
)
|
||||
batch.prepare_for_extend()
|
||||
_maybe_prepare_mlp_sync_batch(batch, model_runner)
|
||||
model_worker_batch = batch.get_model_worker_batch()
|
||||
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
|
||||
forward_batch = ForwardBatch.init_new(batch, model_runner)
|
||||
logits_output = model_runner.forward(forward_batch).logits_output
|
||||
next_token_ids = model_runner.sample(logits_output, forward_batch)
|
||||
return next_token_ids, logits_output.next_token_logits, batch
|
||||
@@ -469,8 +468,7 @@ def decode(input_token_ids, batch, model_runner):
|
||||
batch.output_ids = input_token_ids
|
||||
batch.prepare_for_decode()
|
||||
_maybe_prepare_mlp_sync_batch(batch, model_runner)
|
||||
model_worker_batch = batch.get_model_worker_batch()
|
||||
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
|
||||
forward_batch = ForwardBatch.init_new(batch, model_runner)
|
||||
logits_output = model_runner.forward(forward_batch).logits_output
|
||||
next_token_ids = model_runner.sample(logits_output, forward_batch)
|
||||
return next_token_ids, logits_output.next_token_logits
|
||||
|
||||
Reference in New Issue
Block a user