chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Satyam Kumar <satyamk@linkedin.com>
This commit is contained in:
satyamk7054
2025-10-28 02:49:36 +08:00
committed by GitHub
co-authored by Satyam Kumar
parent c11b34d599
commit 9fc3e8aac7
15 changed files with 314 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+8
View File
@@ -205,6 +205,14 @@ class ModelConfig:
self.hf_config, "image_token_id", None
) or getattr(self.hf_config, "image_token_index", None)
# matryoshka embeddings
self.matryoshka_dimensions = getattr(
self.hf_config, "matryoshka_dimensions", None
)
self.is_matryoshka = self.matryoshka_dimensions or getattr(
self.hf_config, "is_matryoshka", False
)
@staticmethod
def from_server_args(
server_args: ServerArgs,
python/sglang/srt/entrypoints/engine.py
+4
View File
@@ -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
+1
View File
@@ -126,6 +126,7 @@ class OpenAIServingEmbedding(OpenAIServingBase):
**prompt_kwargs,
rid=request.rid,
priority=request.priority,
dimensions=request.dimensions,
)
return adapted_request, request
python/sglang/srt/layers/pooler.py
+21 -2
View File
@@ -20,7 +20,9 @@ class PoolingType(IntEnum):
@dataclass
class EmbeddingPoolerOutput:
embeddings: torch.Tensor
# Pooler can return list[tensor] instead of tensor if the dimension of each tensor in the batch is different
# due to different per-request matryoshka dim truncation
embeddings: torch.Tensor | list[torch.Tensor]
class Pooler(nn.Module):
@@ -42,6 +44,7 @@ class Pooler(nn.Module):
def forward(
self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
) -> EmbeddingPoolerOutput:
if self.pooling_type == PoolingType.LAST:
last_token_indices = torch.cumsum(forward_batch.extend_seq_lens, dim=0) - 1
pooled_data = hidden_states[last_token_indices]
@@ -53,8 +56,24 @@ class Pooler(nn.Module):
else:
raise ValueError(f"Invalid pooling type: {self.pooling_type}")
if forward_batch.dimensions is not None:
all_same_dimensions = len(set(forward_batch.dimensions)) == 1
if all_same_dimensions:
pooled_data = pooled_data[..., : forward_batch.dimensions[0]]
else:
pooled_data = [
tensor[..., :dim]
for tensor, dim in zip(pooled_data, forward_batch.dimensions)
]
if self.normalize:
pooled_data = nn.functional.normalize(pooled_data, p=2, dim=1)
if isinstance(pooled_data, list):
pooled_data = [
nn.functional.normalize(tensor, p=2, dim=-1)
for tensor in pooled_data
]
else:
pooled_data = nn.functional.normalize(pooled_data, p=2, dim=-1)
return EmbeddingPoolerOutput(embeddings=pooled_data)
python/sglang/srt/managers/io_struct.py
+6
View File
@@ -695,6 +695,9 @@ class EmbeddingReqInput(BaseReq):
# tracing context
trace_context: Optional[Dict] = None
# The number of dimensions the resulting output embeddings should have. It is applicable for Matryoshka Embeddings.
dimensions: Optional[int] = None
def normalize_batch_and_arguments(self):
# at least one of text, input_ids, or image should be provided
if self.text is None and self.input_ids is None and self.image_data is None:
@@ -771,6 +774,7 @@ class EmbeddingReqInput(BaseReq):
video_data=self.video_data[i] if self.video_data is not None else None,
sampling_params=self.sampling_params[i],
rid=self.rid[i],
dimensions=self.dimensions,
http_worker_ipc=self.http_worker_ipc,
)
@@ -791,6 +795,8 @@ class TokenizedEmbeddingReqInput(BaseReq):
data_parallel_rank: Optional[int] = None
# Priority for the request
priority: Optional[int] = None
# The number of dimensions the resulting output embeddings should have. It is applicable for Matryoshka Embeddings.
dimensions: Optional[int] = None
@dataclass
python/sglang/srt/managers/schedule_batch.py
+20
View File
@@ -442,6 +442,7 @@ class Req:
priority: Optional[int] = None,
metrics_collector: Optional[SchedulerMetricsCollector] = None,
extra_key: Optional[str] = None,
dimensions: Optional[int] = None,
http_worker_ipc: Optional[str] = None,
):
# Input and output info
@@ -650,6 +651,9 @@ class Req:
self.tmp_end_idx: int = -1
self.metadata_buffer_index: int = -1
# For Matryoshka embeddings
self.dimensions = dimensions
@property
def seqlen(self):
return len(self.origin_input_ids) + len(self.output_ids)
@@ -1014,6 +1018,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
encoder_lens_cpu: Optional[List[int]] = None
encoder_out_cache_loc: Optional[torch.Tensor] = None
# For matryoshka embeddings
dimensions: Optional[list[int]] = None
# For split prefill
split_index: int = 0
split_prefill_finished: bool = False
@@ -1177,6 +1184,15 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
prefix_lens = [len(r.prefix_indices) for r in reqs]
extend_lens = [r.extend_input_len for r in reqs]
# For matryoshka embeddings
if self.model_config.is_matryoshka and any(
r.dimensions is not None for r in reqs
):
self.dimensions = [
r.dimensions if r.dimensions else self.model_config.hidden_size
for r in reqs
]
token_type_ids = [
r.token_type_ids for r in reqs if r.token_type_ids is not None
]
@@ -1765,6 +1781,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
),
extend_input_logprob_token_ids=self.extend_input_logprob_token_ids,
is_prefill_only=self.is_prefill_only,
dimensions=self.dimensions,
)
def copy(self):
@@ -1873,5 +1890,8 @@ class ModelWorkerBatch:
capture_hidden_mode: CaptureHiddenMode = None
hicache_consumer_index: int = -1
# For matryoshka embeddings
dimensions: Optional[list[int]] = None
# Whether this batch is prefill-only (no token generation needed)
is_prefill_only: bool = False
python/sglang/srt/managers/scheduler.py
+1
View File
@@ -1475,6 +1475,7 @@ class Scheduler(
recv_req.sampling_params,
token_type_ids=recv_req.token_type_ids,
priority=recv_req.priority,
dimensions=recv_req.dimensions,
http_worker_ipc=recv_req.http_worker_ipc,
)
req.tokenizer = self.tokenizer
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+4 -1
View File
@@ -203,7 +203,10 @@ class SchedulerOutputProcessorMixin:
i
].item()
else:
embeddings = embeddings.tolist()
if isinstance(embeddings, torch.Tensor):
embeddings = embeddings.tolist()
else:
embeddings = [tensor.tolist() for tensor in embeddings]
# Check finish conditions
for i, req in enumerate(batch.reqs):
python/sglang/srt/managers/tokenizer_manager.py
+33
View File
@@ -666,6 +666,10 @@ class TokenizerManager(TokenizerCommunicatorMixin):
)
raise ValueError(error_msg)
# Matryoshka embeddings validations
if isinstance(obj, EmbeddingReqInput):
self._validate_for_matryoshka_dim(obj)
if isinstance(obj, GenerateReqInput):
if (
obj.return_hidden_states
@@ -684,6 +688,34 @@ class TokenizerManager(TokenizerCommunicatorMixin):
"Please set `--enable-custom-logit-processor` to enable this feature."
)
def _validate_for_matryoshka_dim(self, obj: EmbeddingReqInput) -> None:
"""Validate the request for Matryoshka dim if it has the field set."""
if obj.dimensions is None:
return
if not self.model_config.is_matryoshka:
raise ValueError(
f"Model '{self.model_config.model_path}' does not support matryoshka representation, "
f"changing output dimensions will lead to poor results."
)
if obj.dimensions < 1:
raise ValueError("Requested dimensions must be greater than 0")
if (
self.model_config.matryoshka_dimensions
and obj.dimensions not in self.model_config.matryoshka_dimensions
):
raise ValueError(
f"Model '{self.model_config.model_path}' only supports {self.model_config.matryoshka_dimensions} matryoshka dimensions, "
f"using other output dimensions will lead to poor results."
)
if obj.dimensions > self.model_config.hidden_size:
raise ValueError(
f"Provided dimensions are greater than max embedding dimension: {self.model_config.hidden_size}"
)
def _validate_input_ids_in_vocab(
self, input_ids: List[int], vocab_size: int
) -> None:
@@ -752,6 +784,7 @@ class TokenizerManager(TokenizerCommunicatorMixin):
sampling_params,
rid=obj.rid,
priority=obj.priority,
dimensions=obj.dimensions,
http_worker_ipc=obj.http_worker_ipc,
)
python/sglang/srt/model_executor/forward_batch_info.py
+4
View File
@@ -320,6 +320,9 @@ class ForwardBatch:
tbo_parent_token_range: Optional[Tuple[int, int]] = None
tbo_children: Optional[List[ForwardBatch]] = None
# For matryoshka embeddings
dimensions: Optional[list[int]] = None
@classmethod
def init_new(
cls,
@@ -361,6 +364,7 @@ class ForwardBatch:
input_embeds=batch.input_embeds,
token_type_ids=batch.token_type_ids,
tbo_split_seq_index=batch.tbo_split_seq_index,
dimensions=batch.dimensions,
)
device = model_runner.device
python/sglang/test/runners.py
+30 -6
View File
@@ -12,10 +12,11 @@
# limitations under the License.
# ==============================================================================
import json
import multiprocessing as mp
import os
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union
from typing import Any, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
@@ -89,7 +90,9 @@ def get_token_ids_logprobs(logits, token_ids):
return logprobs
def _get_sentence_transformer_embedding_model(model_path, torch_dtype):
def _get_sentence_transformer_embedding_model(
model_path, torch_dtype, matryoshka_dim: Optional[int] = None
):
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import is_sentence_transformer_model
@@ -97,6 +100,7 @@ def _get_sentence_transformer_embedding_model(model_path, torch_dtype):
model = SentenceTransformer(
model_path,
model_kwargs={"torch_dtype": torch_dtype},
truncate_dim=matryoshka_dim,
)
else: # if no pre-trained sentence-transformers model
from sentence_transformers import models
@@ -106,7 +110,9 @@ def _get_sentence_transformer_embedding_model(model_path, torch_dtype):
word_embedding_model.get_word_embedding_dimension(),
pooling_mode="lasttoken",
)
model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
model = SentenceTransformer(
modules=[word_embedding_model, pooling_model], truncate_dim=matryoshka_dim
)
return model.cuda()
@@ -135,6 +141,7 @@ class HFRunner:
output_str_only: bool = False,
trust_remote_code: bool = False,
patch_model_do_sample_false: bool = False,
matryoshka_dim: Optional[int] = None,
):
self.model_type = model_type
self.output_str_only = output_str_only
@@ -151,6 +158,7 @@ class HFRunner:
self.out_queue,
model_path,
torch_dtype,
matryoshka_dim,
),
)
self.model_proc.start()
@@ -225,7 +233,14 @@ class HFRunner:
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
return embeddings.contiguous()
def start_model_process(self, in_queue, out_queue, model_path, torch_dtype):
def start_model_process(
self,
in_queue,
out_queue,
model_path,
torch_dtype,
matryoshka_dim: Optional[int] = None,
):
# Apply model-specific patches
monkey_patch_gemma2_sdpa()
@@ -259,7 +274,7 @@ class HFRunner:
self.processor = AutoProcessor.from_pretrained(model_path)
else:
self.model = _get_sentence_transformer_embedding_model(
model_path, torch_dtype
model_path, torch_dtype, matryoshka_dim=matryoshka_dim
)
elif self.model_type == "reward" or self.model_type == "cross_encoder":
from transformers import AutoModelForSequenceClassification
@@ -519,6 +534,7 @@ class SRTRunner:
lora_target_modules: Optional[List[str]] = None,
enable_lora: Optional[bool] = None,
max_loaded_loras: Optional[int] = None,
json_model_override_args: Optional[dict[str, Any]] = None,
lora_eviction_policy: str = "lru",
):
self.model_type = model_type
@@ -566,6 +582,11 @@ class SRTRunner:
lora_target_modules=lora_target_modules,
enable_lora=enable_lora,
max_loaded_loras=max_loaded_loras,
json_model_override_args=(
json.dumps(json_model_override_args)
if json_model_override_args
else "{}"
),
lora_eviction_policy=lora_eviction_policy,
**spec_kwargs,
)
@@ -594,6 +615,7 @@ class SRTRunner:
logprob_start_len: int = 0,
top_k: Optional[int] = None,
token_ids_logprob: Optional[List[int]] = None,
dimensions: Optional[int] = None,
):
if self.is_generation:
return self.forward_generation_raw(
@@ -607,7 +629,9 @@ class SRTRunner:
)
else:
if self.model_type == "embedding":
response = self.engine.encode(prompt=prompts, image_data=image_data)
response = self.engine.encode(
prompt=prompts, image_data=image_data, dimensions=dimensions
)
if isinstance(response, list):
logits = [x["embedding"] for x in response]
else: