refactor: remove ModelWorkerBatch indirection (#25516)

This commit is contained in:
Liangsheng Yin
2026-05-17 18:05:39 -07:00
committed by GitHub
parent b3803164cb
commit 58ece60703
21 changed files with 430 additions and 559 deletions
+2 -4
View File
@@ -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