chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages 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

9
python/sglang/srt/managers/tp_worker.py
View File

@@ -21,6 +21,7 @@ import logging
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
from sglang.srt.managers.io_struct import UpdateWeightReqInput
from sglang.srt.model_executor.forward_batch_info import InputMetadata
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import broadcast_pyobj, is_multimodal_model, set_random_seed
@@ -105,13 +106,13 @@ class ModelTpWorker:
self.random_seed,
)
def forward_batch_generation(self, batch):
logits_output = self.model_runner.forward(batch)
def forward_batch_generation(self, input_metadata: InputMetadata, batch):
logits_output = self.model_runner.forward(input_metadata)
next_token_ids = self.model_runner.sample(logits_output, batch)
return logits_output, next_token_ids
def forward_batch_embedding(self, batch):
logits_output = self.model_runner.forward(batch)
def forward_batch_embedding(self, input_metadata: InputMetadata):
logits_output = self.model_runner.forward(input_metadata)
embeddings = logits_output.embeddings.tolist()
return embeddings