sglang 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

sglang/srt/managers/tokenizer_manager.py

@@ -21,7 +21,7 @@ import dataclasses
 import logging
 import multiprocessing as mp
 import os
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import transformers
@@ -38,16 +38,19 @@ from sglang.srt.hf_transformers_utils import (
 )
 from sglang.srt.managers.io_struct import (
     AbortReq,
+    BatchEmbeddingOut,
     BatchStrOut,
     BatchTokenIDOut,
+    EmbeddingReqInput,
     FlushCacheReq,
     GenerateReqInput,
+    TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
 )
 from sglang.srt.mm_utils import expand2square, process_anyres_image
 from sglang.srt.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import is_multimodal_model, load_image
+from sglang.srt.utils import is_generation_model, is_multimodal_model, load_image
 from sglang.utils import get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
@@ -85,31 +88,35 @@ class TokenizerManager:
             trust_remote_code=server_args.trust_remote_code,
             model_overide_args=model_overide_args,
         )
+        self.is_generation = is_generation_model(self.hf_config.architectures)
 
         if server_args.context_length is not None:
             self.context_len = server_args.context_length
         else:
             self.context_len = get_context_length(self.hf_config)
 
-        if is_multimodal_model(self.model_path):
-            self.processor = get_processor(
-                server_args.tokenizer_path,
-                tokenizer_mode=server_args.tokenizer_mode,
-                trust_remote_code=server_args.trust_remote_code,
-            )
-            self.tokenizer = self.processor.tokenizer
-            os.environ["TOKENIZERS_PARALLELISM"] = "false"
-            self.executor = concurrent.futures.ProcessPoolExecutor(
-                initializer=init_global_processor,
-                mp_context=mp.get_context("fork"),
-                initargs=(server_args,),
-            )
+        if server_args.skip_tokenizer_init:
+            self.tokenizer = self.processor = None
         else:
-            self.tokenizer = get_tokenizer(
-                server_args.tokenizer_path,
-                tokenizer_mode=server_args.tokenizer_mode,
-                trust_remote_code=server_args.trust_remote_code,
-            )
+            if is_multimodal_model(self.model_path):
+                self.processor = get_processor(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                )
+                self.tokenizer = self.processor.tokenizer
+                os.environ["TOKENIZERS_PARALLELISM"] = "false"
+                self.executor = concurrent.futures.ProcessPoolExecutor(
+                    initializer=init_global_processor,
+                    mp_context=mp.get_context("fork"),
+                    initargs=(server_args,),
+                )
+            else:
+                self.tokenizer = get_tokenizer(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                )
 
         self.to_create_loop = True
         self.rid_to_state: Dict[str, ReqState] = {}
@@ -133,7 +140,9 @@ class TokenizerManager:
                 image_data, aspect_ratio, grid_pinpoints, self.processor
             )
 
-    async def generate_request(self, obj: GenerateReqInput, request=None):
+    async def generate_request(
+        self, obj: Union[GenerateReqInput, EmbeddingReqInput], request=None
+    ):
         if self.to_create_loop:
             self.create_handle_loop()
 
@@ -144,46 +153,55 @@ class TokenizerManager:
             async for response in self._handle_single_request(obj, request):
                 yield response
         else:
-            if obj.stream:
+            if hasattr(obj, "stream") and obj.stream:
                 raise ValueError("Do not support stream for batch mode.")
 
             async for response in self._handle_batch_request(obj, request):
                 yield response
 
     async def _handle_single_request(
-        self, obj, request, index=None, is_cache_for_prefill=False
+        self,
+        obj: Union[GenerateReqInput, EmbeddingReqInput],
+        request,
+        index=None,
+        is_cache_for_prefill=False,
     ):
         if not is_cache_for_prefill:  # The normal case with a single prompt
             not_use_index = index is None
 
             rid = obj.rid if not_use_index else obj.rid[index]
             input_text = obj.text if not_use_index else obj.text[index]
-            input_ids = (
-                self.tokenizer.encode(input_text)
-                if obj.input_ids is None
-                else obj.input_ids
-            )
-            if not not_use_index and obj.input_ids:
-                input_ids = obj.input_ids[index]
+            if obj.input_ids is None:
+                assert self.tokenizer is not None
+                input_ids = self.tokenizer.encode(input_text)
+            else:
+                input_ids = obj.input_ids if not_use_index else obj.input_ids[index]
 
             self._validate_input_length(input_ids)
 
             sampling_params = self._get_sampling_params(
                 obj.sampling_params if not_use_index else obj.sampling_params[index]
             )
-            pixel_values, image_hash, image_size = await self._get_pixel_values(
-                obj.image_data if not_use_index else obj.image_data[index]
-            )
-            return_logprob = (
-                obj.return_logprob if not_use_index else obj.return_logprob[index]
-            )
-            logprob_start_len = (
-                obj.logprob_start_len if not_use_index else obj.logprob_start_len[index]
-            )
-            top_logprobs_num = (
-                obj.top_logprobs_num if not_use_index else obj.top_logprobs_num[index]
-            )
+
+            if self.is_generation:
+                pixel_values, image_hash, image_size = await self._get_pixel_values(
+                    obj.image_data if not_use_index else obj.image_data[index]
+                )
+                return_logprob = (
+                    obj.return_logprob if not_use_index else obj.return_logprob[index]
+                )
+                logprob_start_len = (
+                    obj.logprob_start_len
+                    if not_use_index
+                    else obj.logprob_start_len[index]
+                )
+                top_logprobs_num = (
+                    obj.top_logprobs_num
+                    if not_use_index
+                    else obj.top_logprobs_num[index]
+                )
         else:  # A prefill request to cache the common prompt for parallel sampling
+            assert self.is_generation
             if obj.text is not None:
                 if isinstance(obj.text, list):
                     input_text = obj.text[index]
@@ -191,7 +209,20 @@ class TokenizerManager:
                 else:
                     input_text = obj.text
                     rid = obj.rid[0]
-                input_ids = self.tokenizer.encode(input_text)
+                if self.tokenizer is not None:
+                    input_ids = self.tokenizer.encode(input_text)
+                else:
+                    assert obj.input_ids is not None
+                    input_ids = obj.input_ids
+                    if isinstance(obj.input_ids, list) and isinstance(
+                        obj.input_ids[0], list
+                    ):
+                        # when obj["input_ids"] is List[List[int]]
+                        input_ids = obj.input_ids[index]
+                        rid = obj.rid[index]
+                    else:
+                        input_ids = obj.input_ids
+                        rid = obj.rid[0]
             else:
                 input_text = None
                 if isinstance(obj.input_ids, list) and isinstance(
@@ -213,19 +244,28 @@ class TokenizerManager:
             logprob_start_len = obj.logprob_start_len[0]
             top_logprobs_num = obj.top_logprobs_num[0]
 
-        tokenized_obj = TokenizedGenerateReqInput(
-            rid,
-            input_text,
-            input_ids,
-            pixel_values,
-            image_hash,
-            image_size,
-            sampling_params,
-            return_logprob,
-            logprob_start_len,
-            top_logprobs_num,
-            obj.stream,
-        )
+        if self.is_generation:
+            tokenized_obj = TokenizedGenerateReqInput(
+                rid,
+                input_text,
+                input_ids,
+                pixel_values,
+                image_hash,
+                image_size,
+                sampling_params,
+                return_logprob,
+                logprob_start_len,
+                top_logprobs_num,
+                obj.stream,
+            )
+        else:  # is embedding
+            tokenized_obj = TokenizedEmbeddingReqInput(
+                rid,
+                input_text,
+                input_ids,
+                sampling_params,
+            )
+
         self.send_to_router.send_pyobj(tokenized_obj)
 
         event = asyncio.Event()
@@ -237,27 +277,33 @@ class TokenizerManager:
             ):
                 yield response
         else:
+            assert self.is_generation
             await self._wait_for_cache_prefill_response(event, state, obj, rid, request)
             yield input_ids
 
-    async def _handle_batch_request(self, obj: GenerateReqInput, request):
+    async def _handle_batch_request(
+        self, obj: Union[GenerateReqInput, EmbeddingReqInput], request
+    ):
         batch_size = obj.batch_size
-        parallel_sample_num = obj.parallel_sample_num
-
-        if parallel_sample_num != 1:
-            # Send prefill requests to cache the common input
-            parallel_sample_num += 1
-            input_id_result = [] if obj.input_ids is None else None
-            for i in range(batch_size):
-                async for input_id in self._handle_single_request(
-                    obj, request, index=i, is_cache_for_prefill=True
-                ):
-                    if input_id_result is not None:
-                        input_id_result.append(input_id)
-            if input_id_result is not None and len(input_id_result) > 1:
-                obj.input_ids = input_id_result
-            elif input_id_result is not None:
-                obj.input_ids = input_id_result[0]
+        if self.is_generation:
+            parallel_sample_num = obj.parallel_sample_num
+
+            if parallel_sample_num != 1:
+                # Send prefill requests to cache the common input
+                parallel_sample_num += 1
+                input_id_result = [] if obj.input_ids is None else None
+                for i in range(batch_size):
+                    async for input_id in self._handle_single_request(
+                        obj, request, index=i, is_cache_for_prefill=True
+                    ):
+                        if input_id_result is not None:
+                            input_id_result.append(input_id)
+                if input_id_result is not None and len(input_id_result) > 1:
+                    obj.input_ids = input_id_result
+                elif input_id_result is not None:
+                    obj.input_ids = input_id_result[0]
+        else:
+            parallel_sample_num = 1
 
         # First send out all requests
         for i in range(batch_size):
@@ -286,23 +332,32 @@ class TokenizerManager:
                         input_text = None
                         input_ids = obj.input_ids[i]
                 sampling_params = self._get_sampling_params(obj.sampling_params[index])
-                pixel_values, image_hash, image_size = await self._get_pixel_values(
-                    obj.image_data[index]
-                )
 
-                tokenized_obj = TokenizedGenerateReqInput(
-                    rid,
-                    input_text,
-                    input_ids,
-                    pixel_values,
-                    image_hash,
-                    image_size,
-                    sampling_params,
-                    obj.return_logprob[index],
-                    obj.logprob_start_len[index],
-                    obj.top_logprobs_num[index],
-                    obj.stream,
-                )
+                if self.is_generation:
+                    pixel_values, image_hash, image_size = await self._get_pixel_values(
+                        obj.image_data[index]
+                    )
+
+                    tokenized_obj = TokenizedGenerateReqInput(
+                        rid,
+                        input_text,
+                        input_ids,
+                        pixel_values,
+                        image_hash,
+                        image_size,
+                        sampling_params,
+                        obj.return_logprob[index],
+                        obj.logprob_start_len[index],
+                        obj.top_logprobs_num[index],
+                        obj.stream,
+                    )
+                else:
+                    tokenized_obj = TokenizedEmbeddingReqInput(
+                        rid,
+                        input_text,
+                        input_ids,
+                        sampling_params,
+                    )
                 self.send_to_router.send_pyobj(tokenized_obj)
 
                 event = asyncio.Event()
@@ -331,17 +386,19 @@ class TokenizerManager:
                                 self.abort_request(rid)
                             raise ValueError(f"Abort request {rid}")
                         continue
-                output_list.append(
-                    self.convert_logprob_style(
-                        state.out_list[-1],
-                        obj.return_logprob[index],
-                        obj.top_logprobs_num[index],
-                        obj.return_text_in_logprobs,
+                if self.is_generation:
+                    output_list.append(
+                        self.convert_logprob_style(
+                            state.out_list[-1],
+                            obj.return_logprob[index],
+                            obj.top_logprobs_num[index],
+                            obj.return_text_in_logprobs,
+                        )
                     )
-                )
+                else:
+                    output_list.append(state.out_list[-1])
                 assert state.finished
                 del self.rid_to_state[rid]
-
         yield output_list
 
     def _validate_input_length(self, input_ids: List[int]):
@@ -370,7 +427,7 @@ class TokenizerManager:
         self,
         event: asyncio.Event,
         state: ReqState,
-        obj: GenerateReqInput,
+        obj: Union[GenerateReqInput, EmbeddingReqInput],
         rid: str,
         request,
     ):
@@ -383,15 +440,23 @@ class TokenizerManager:
                     raise ValueError(f"Abort request {rid}")
                 continue
 
-            out = self.convert_logprob_style(
-                state.out_list[-1],
-                obj.return_logprob,
-                obj.top_logprobs_num,
-                obj.return_text_in_logprobs,
-            )
+            if self.is_generation:
+                out = self.convert_logprob_style(
+                    state.out_list[-1],
+                    obj.return_logprob,
+                    obj.top_logprobs_num,
+                    obj.return_text_in_logprobs,
+                )
+            else:  # isinstance(obj, EmbeddingReqInput)
+                out = state.out_list[-1]
 
+            # Log requests
             if self.server_args.log_requests and state.finished:
-                logger.info(f"in={obj.text}, out={out}")
+                if obj.text is None:
+                    in_obj = {"input_ids": obj.input_ids}
+                else:
+                    in_obj = {"text": obj.text}
+                logger.info(f"in={in_obj}, out={out}")
 
             state.out_list = []
             if state.finished:
@@ -456,19 +521,38 @@ class TokenizerManager:
 
     async def handle_loop(self):
         while True:
-            recv_obj: BatchTokenIDOut = await self.recv_from_detokenizer.recv_pyobj()
-            assert isinstance(recv_obj, BatchStrOut)
-
+            recv_obj: Union[BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut] = (
+                await self.recv_from_detokenizer.recv_pyobj()
+            )
+            assert isinstance(
+                recv_obj, (BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut)
+            ), f"Unexpected obj received: {type(recv_obj)}"
             for i, rid in enumerate(recv_obj.rids):
                 state = self.rid_to_state.get(rid, None)
                 if state is None:
                     continue
 
                 recv_obj.meta_info[i]["id"] = rid
-                out_dict = {
-                    "text": recv_obj.output_strs[i],
-                    "meta_info": recv_obj.meta_info[i],
-                }
+                if isinstance(recv_obj, BatchStrOut):
+                    out_dict = {
+                        "text": recv_obj.output_strs[i],
+                        "meta_info": recv_obj.meta_info[i],
+                    }
+                elif isinstance(recv_obj, BatchTokenIDOut):
+                    read_start = 0 if i == 0 else recv_obj.read_offsets[i - 1]
+                    out_dict = {
+                        "token_ids": recv_obj.decode_ids[
+                            read_start : recv_obj.read_offsets[i]
+                        ],
+                        "meta_info": recv_obj.meta_info[i],
+                    }
+
+                else:
+                    assert isinstance(recv_obj, BatchEmbeddingOut)
+                    out_dict = {
+                        "embedding": recv_obj.embeddings[i],
+                        "meta_info": recv_obj.meta_info[i],
+                    }
                 state.out_list.append(out_dict)
                 state.finished = recv_obj.finished_reason[i] is not None
                 state.event.set()
@@ -508,6 +592,7 @@ class TokenizerManager:
         if not decode_to_text:
             return [(logprob, token_id, None) for logprob, token_id in token_logprobs]
 
+        assert self.tokenizer is not None
         token_ids = [tid for _, tid in token_logprobs]
         token_texts = self.tokenizer.batch_decode(token_ids)
         return [