chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Tiny extract ModelRunnerOutput (#15400)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-18 22:18:45 +08:00
committed by GitHub
parent c5f4e20f2f
commit ad9616f13a
7 changed files with 34 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/bench_one_batch.py
View File

@@ -373,7 +373,7 @@ def extend(reqs, model_runner):
_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)
logits_output, _ = model_runner.forward(forward_batch)
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
@@ -385,7 +385,7 @@ def decode(input_token_ids, batch, model_runner):
_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)
logits_output, _ = model_runner.forward(forward_batch)
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