Add support for Matryoshka embeddings (#126) (#11142)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

python/sglang/srt/entrypoints/engine.py

@@ -312,6 +312,7 @@ class Engine(EngineBase):
         image_data: Optional[MultimodalDataInputFormat] = None,
         audio_data: Optional[MultimodalDataInputFormat] = None,
         video_data: Optional[MultimodalDataInputFormat] = None,
+        dimensions: Optional[int] = None,
     ) -> Dict:
         """
         The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`.
@@ -322,6 +323,7 @@ class Engine(EngineBase):
             image_data=image_data,
             audio_data=audio_data,
             video_data=video_data,
+            dimensions=dimensions,
         )
         generator = self.tokenizer_manager.generate_request(obj, None)
         ret = self.loop.run_until_complete(generator.__anext__())
@@ -333,6 +335,7 @@ class Engine(EngineBase):
         image_data: Optional[MultimodalDataInputFormat] = None,
         audio_data: Optional[MultimodalDataInputFormat] = None,
         video_data: Optional[MultimodalDataInputFormat] = None,
+        dimensions: Optional[int] = None,
     ) -> Dict:
         """
         Asynchronous version of encode method.
@@ -345,6 +348,7 @@ class Engine(EngineBase):
             image_data=image_data,
             audio_data=audio_data,
             video_data=video_data,
+            dimensions=dimensions,
         )
         generator = self.tokenizer_manager.generate_request(obj, None)
         return await generator.__anext__()

python/sglang/srt/entrypoints/openai/serving_embedding.py

@@ -126,6 +126,7 @@ class OpenAIServingEmbedding(OpenAIServingBase):
             **prompt_kwargs,
             rid=request.rid,
             priority=request.priority,
+            dimensions=request.dimensions,
         )
 
         return adapted_request, request