Support embedding input as a list (#1014)

python/sglang/srt/managers/tokenizer_manager.py

@@ -153,9 +153,7 @@ class TokenizerManager:
             async for response in self._handle_single_request(obj, request):
                 yield response
         else:
-            if isinstance(obj, EmbeddingReqInput):
+            if hasattr(obj, "stream") and obj.stream:
-                raise NotImplementedError("Please send only one prompt in each request")
-            if obj.stream:
                 raise ValueError("Do not support stream for batch mode.")
 
             async for response in self._handle_batch_request(obj, request):
@@ -283,8 +281,11 @@ class TokenizerManager:
             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
+        if self.is_generation:
             parallel_sample_num = obj.parallel_sample_num
 
             if parallel_sample_num != 1:
@@ -301,6 +302,8 @@ class TokenizerManager:
                     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):
@@ -329,6 +332,8 @@ class TokenizerManager:
                         input_text = None
                         input_ids = obj.input_ids[i]
                 sampling_params = self._get_sampling_params(obj.sampling_params[index])
+
+                if self.is_generation:
                     pixel_values, image_hash, image_size = await self._get_pixel_values(
                         obj.image_data[index]
                     )
@@ -346,11 +351,19 @@ class TokenizerManager:
                         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()
                 state = ReqState([], False, event)
                 self.rid_to_state[rid] = state
+
         # Then wait for all responses
         output_list = []
         for i in range(batch_size):
@@ -373,6 +386,7 @@ class TokenizerManager:
                                 self.abort_request(rid)
                             raise ValueError(f"Abort request {rid}")
                         continue
+                if self.is_generation:
                     output_list.append(
                         self.convert_logprob_style(
                             state.out_list[-1],
@@ -381,6 +395,8 @@ class TokenizerManager:
                             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

python/sglang/test/runners.py

@@ -219,11 +219,9 @@ class SRTRunner:
                 output_strs=output_strs, top_input_logprobs=top_input_logprobs
             )
         else:
-            logits = []
+            response = self.runtime.encode(prompts)
-            for prompt in prompts:
-                response = self.runtime.encode(prompt)
             response = json.loads(response)
-                logits.append(response["embedding"])
+            logits = [x["embedding"] for x in response]
             return ModelOutput(embed_logits=logits)
 
     def __enter__(self):

test/srt/test_embedding_openai_server.py

@@ -38,8 +38,9 @@ class TestOpenAIServer(unittest.TestCase):
             num_prompt_tokens = len(self.tokenizer.encode(prompt))
 
         if use_list_input:
-            prompt_arg = [prompt_input, prompt_input]
+            prompt_arg = [prompt_input] * 2
             num_prompts = len(prompt_arg)
+            num_prompt_tokens *= num_prompts
         else:
             prompt_arg = prompt_input
             num_prompts = 1
@@ -70,7 +71,7 @@ class TestOpenAIServer(unittest.TestCase):
     def test_embedding(self):
         # TODO the fields of encoding_format, dimensions, user are skipped
         # TODO support use_list_input
-        for use_list_input in [False]:
+        for use_list_input in [False, True]:
             for token_input in [False, True]:
                 self.run_embedding(use_list_input, token_input)