EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Let ModelRunner take InputMetadata as input, instead of ScheduleBatch (#1541)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-09-29 20:28:45 -07:00
committed by GitHub
parent 55b974f96f
commit 3f0fe08d37
12 changed files with 142 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/bench_latency.py
View File

@@ -225,14 +225,16 @@ def extend(reqs, model_runner):
tree_cache=None,
)
batch.prepare_for_extend(model_runner.model_config.vocab_size)
logits_output = model_runner.forward(batch)
input_metadata = batch.get_input_metadata()
logits_output = model_runner.forward(input_metadata)
next_token_ids = model_runner.sample(logits_output, batch).tolist()
return next_token_ids, logits_output.next_token_logits, batch
def decode(input_token_ids, batch, model_runner):
batch.prepare_for_decode(input_token_ids)
logits_output = model_runner.forward(batch)
input_metadata = batch.get_input_metadata()
logits_output = model_runner.forward(input_metadata)
next_token_ids = model_runner.sample(logits_output, batch).tolist()
return next_token_ids, logits_output.next_token_logits