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

embedding: centralize capabilities and complete OpenAI compatibility (#32481)

Browse Source
This commit is contained in:
Mick
2026-07-30 10:28:52 +08:00
committed by GitHub
parent 313a518bee
commit 22faf9fef8
16 changed files with 728 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/serving.py
+25 -6
View File
@@ -937,6 +937,7 @@ ASYNC_REQUEST_FUNCS = {
"sglang-embedding": async_request_openai_embeddings,
"vllm": async_request_openai_completions,
"vllm-chat": async_request_openai_chat_completions,
"vllm-embedding": async_request_openai_embeddings,
"lmdeploy": async_request_openai_completions,
"lmdeploy-chat": async_request_openai_chat_completions,
"trt": async_request_trt_llm,
@@ -954,12 +955,24 @@ _BACKEND_API_PATHS = {
"sglang-embedding": "/v1/embeddings",
"vllm": "/v1/completions",
"vllm-chat": "/v1/chat/completions",
"vllm-embedding": "/v1/embeddings",
"lmdeploy": "/v1/completions",
"lmdeploy-chat": "/v1/chat/completions",
"trt": "/v2/models/ensemble/generate_stream",
"truss": "/v1/models/model:predict",
}
_EMBEDDING_BACKENDS = frozenset(("sglang-embedding", "vllm-embedding"))
def flush_server_cache(base_url: str, backend: str) -> None:
"""Flush an engine's prefix cache after benchmark warmup."""
cache_endpoint = (
"/reset_prefix_cache" if backend.startswith("vllm") else "/flush_cache"
)
response = requests.post(base_url + cache_endpoint, headers=get_auth_headers())
response.raise_for_status()
@dataclass
class BenchmarkMetrics:
@@ -1419,9 +1432,14 @@ async def benchmark(
f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..."
)
# Flush cache
if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache:
requests.post(base_url + "/flush_cache", headers=get_auth_headers())
# Flush cache after warmup so the measured run does not benefit from
# request-local prefix reuse. vLLM exposes a different, development-mode
# endpoint for the same purpose.
should_flush_cache = (
"sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")
) or flush_cache
if should_flush_cache:
flush_server_cache(base_url, backend)
time.sleep(1.0)
@@ -1595,7 +1613,7 @@ async def benchmark(
"Total input vision tokens:", metrics.total_input_vision
)
)
is_embedding = backend == "sglang-embedding"
is_embedding = backend in _EMBEDDING_BACKENDS
if not is_embedding:
print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
print(
@@ -1963,6 +1981,7 @@ def run_benchmark(args_: argparse.Namespace):
"sglang-oai": 30000,
"lmdeploy": 23333,
"vllm": 8000,
"vllm-embedding": 8000,
"trt": 8000,
"gserver": 9988,
"truss": 8080,
@@ -2014,14 +2033,14 @@ def run_benchmark(args_: argparse.Namespace):
print("No model specified or found. Please provide a model using `--model`.")
sys.exit(1)
if args.backend != "sglang-embedding" and not check_chat_template(args.model):
if args.backend not in _EMBEDDING_BACKENDS and not check_chat_template(args.model):
print(
"\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n"
"Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n"
)
if (
args.backend == "sglang-embedding"
args.backend in _EMBEDDING_BACKENDS
and args.dataset_name in _EMBEDDING_UNSUPPORTED_DATASETS
):
print(f"{args.dataset_name} dataset is unsupported for embeddings benchmark")
python/sglang/srt/configs/embedding_model_spec.py
+344
View File
@@ -0,0 +1,344 @@
"""Embedding-model capabilities resolved from model architecture and server intent.
This module is deliberately declarative. Server-argument resolution, model
implementations, documentation, and benchmarks can consume the same contract
without each reimplementing a partial list of embedding architectures.
"""
from dataclasses import dataclass
from enum import Enum
from typing import Any, Sequence
class EmbeddingTask(str, Enum):
NONE = "none"
EMBED = "embed"
CLASSIFY = "classify"
class PoolingStrategy(str, Enum):
MODEL_DEFINED = "model_defined"
CLS = "cls"
LAST = "last"
MEAN = "mean"
class BCGPrefillPolicy(str, Enum):
"""Whether a model has a validated breakable-CUDA-graph prefill policy."""
DEFAULT = "default"
FULL_ENCODER = "full_encoder"
class EmbeddingExecution(str, Enum):
"""The model path that produces an embedding."""
NONE = "none"
ENCODER_ONLY = "encoder_only"
DECODER_POOLING = "decoder_pooling"
MULTIMODAL = "multimodal"
CLASSIFICATION = "classification"
class AttentionPattern(str, Enum):
NONE = "none"
BIDIRECTIONAL = "bidirectional"
CAUSAL = "causal"
class BCGEligibility(str, Enum):
"""Whether a model has a validated BCG prefill strategy."""
DISABLED = "disabled"
STANDARD = "standard"
FULL_ENCODER = "full_encoder"
@dataclass(frozen=True)
class EmbeddingModelSpec:
"""Resolved capabilities for an embedding or pooling model.
``auto_enable_embedding`` is intentionally narrow: it is true only when a
checkpoint unambiguously declares a pooling/embedding-only architecture.
Decoder checkpoints trained for embeddings still require explicit user
intent.
"""
family: str
task: EmbeddingTask
execution: EmbeddingExecution
attention: AttentionPattern
pooling: PoolingStrategy
normalize: bool
postprocessor: str
tokenizer_special_tokens: str
supports_dimensions: bool
supports_token_embeddings: bool
supports_multimodal: bool
requires_embedding_flag: bool
auto_enable_embedding: bool
bidirectional_attention: bool
bcg_prefill_policy: BCGPrefillPolicy
safe_disable_radix_cache: bool = False
safe_disable_chunked_prefill: bool = False
safe_disable_kv_cache: bool = False
@property
def bcg_eligibility(self) -> BCGEligibility:
if self.bcg_prefill_policy == BCGPrefillPolicy.FULL_ENCODER:
return BCGEligibility.FULL_ENCODER
if self.task == EmbeddingTask.EMBED:
return BCGEligibility.STANDARD
return BCGEligibility.DISABLED
def as_dict(self) -> dict[str, Any]:
"""Return a stable, JSON-serializable description of this contract."""
return {
"family": self.family,
"task": self.task.value,
"execution": self.execution.value,
"attention": self.attention.value,
"pooling": self.pooling.value,
"normalize": self.normalize,
"postprocessor": self.postprocessor,
"tokenizer_special_tokens": self.tokenizer_special_tokens,
"supports_dimensions": self.supports_dimensions,
"supports_token_embeddings": self.supports_token_embeddings,
"supports_multimodal": self.supports_multimodal,
"requires_embedding_flag": self.requires_embedding_flag,
"auto_enable_embedding": self.auto_enable_embedding,
"bcg_eligibility": self.bcg_eligibility.value,
"safe_disable_kv_cache": self.safe_disable_kv_cache,
"safe_disable_radix_cache": self.safe_disable_radix_cache,
"safe_disable_chunked_prefill": self.safe_disable_chunked_prefill,
}
_EMBEDDING_ARCHITECTURES = {
"BertModel": ("bert", EmbeddingExecution.ENCODER_ONLY, PoolingStrategy.CLS),
"CLIPModel": ("clip", EmbeddingExecution.MULTIMODAL, PoolingStrategy.LAST),
"Contriever": (
"contriever",
EmbeddingExecution.ENCODER_ONLY,
PoolingStrategy.MODEL_DEFINED,
),
"LlamaEmbeddingModel": (
"llama_embedding",
EmbeddingExecution.DECODER_POOLING,
PoolingStrategy.LAST,
),
"MistralModel": (
"mistral_embedding",
EmbeddingExecution.DECODER_POOLING,
PoolingStrategy.LAST,
),
"XLMRobertaModel": (
"xlm_roberta",
EmbeddingExecution.ENCODER_ONLY,
PoolingStrategy.CLS,
),
}
_CLASSIFICATION_ARCHITECTURES = {
"BertForSequenceClassification",
"LlamaForSequenceClassification",
"LlamaForSequenceClassificationWithNormal_Weights",
"Qwen2ForSequenceClassification",
"Qwen3ForSequenceClassification",
"XLMRobertaForSequenceClassification",
}
def embedding_support_matrix() -> list[dict[str, Any]]:
"""Return registry-backed rows for docs, tooling, and CI coverage checks.
Entries describe architecture-level guarantees. Runtime-only properties,
such as the available Matryoshka dimensions and captured BCG sizes, belong
to :func:`resolved_embedding_plan`.
"""
rows = []
for architecture, (family, execution, pooling) in _EMBEDDING_ARCHITECTURES.items():
spec = _native_embedding_spec(family, execution, pooling)
rows.append({"architecture": architecture, **spec.as_dict()})
rows.append(
{
"architecture": "Gemma3TextModel (use_bidirectional_attention=true)",
**_embedding_gemma_spec().as_dict(),
}
)
return rows
def _embedding_gemma_spec() -> EmbeddingModelSpec:
return EmbeddingModelSpec(
family="embeddinggemma",
task=EmbeddingTask.EMBED,
execution=EmbeddingExecution.ENCODER_ONLY,
attention=AttentionPattern.BIDIRECTIONAL,
pooling=PoolingStrategy.MEAN,
normalize=True,
postprocessor="sentence_transformers",
tokenizer_special_tokens="model_default",
supports_dimensions=False,
supports_token_embeddings=False,
supports_multimodal=False,
requires_embedding_flag=False,
auto_enable_embedding=True,
bidirectional_attention=True,
bcg_prefill_policy=BCGPrefillPolicy.FULL_ENCODER,
safe_disable_radix_cache=True,
safe_disable_chunked_prefill=True,
safe_disable_kv_cache=True,
)
def _native_embedding_spec(
family: str, execution: EmbeddingExecution, pooling: PoolingStrategy
) -> EmbeddingModelSpec:
return EmbeddingModelSpec(
family=family,
task=EmbeddingTask.EMBED,
execution=execution,
attention=(
AttentionPattern.CAUSAL
if execution == EmbeddingExecution.DECODER_POOLING
else AttentionPattern.BIDIRECTIONAL
),
pooling=pooling,
normalize=True,
postprocessor="model_defined",
tokenizer_special_tokens="model_default",
supports_dimensions=False,
supports_token_embeddings=False,
supports_multimodal=execution == EmbeddingExecution.MULTIMODAL,
requires_embedding_flag=False,
auto_enable_embedding=True,
bidirectional_attention=execution != EmbeddingExecution.DECODER_POOLING,
bcg_prefill_policy=BCGPrefillPolicy.DEFAULT,
)
def resolved_embedding_plan(
spec: EmbeddingModelSpec, *, server_args: Any, model_config: Any
) -> dict[str, Any]:
"""Combine static capabilities with the effective server configuration.
This boundary deliberately accepts duck-typed arguments so the declarative
registry remains independent of ServerArgs and ModelConfig import cycles.
"""
prefill_graph = getattr(
getattr(server_args, "cuda_graph_config", None), "prefill", None
)
backend = getattr(prefill_graph, "backend", None)
backend_value = getattr(backend, "value", backend)
capture_sizes = getattr(prefill_graph, "bs", None) or []
max_capture_tokens = getattr(prefill_graph, "max_bs", None)
return {
**spec.as_dict(),
"enabled": bool(getattr(server_args, "is_embedding", False)),
"supports_dimensions": bool(getattr(model_config, "is_matryoshka", False)),
"matryoshka_dimensions": list(
getattr(model_config, "matryoshka_dimensions", None) or []
),
"bcg": {
"prefill_backend": backend_value,
"enabled": backend_value == "breakable",
"capture_token_budget": max_capture_tokens,
"capture_batch_sizes": list(capture_sizes),
},
"cache": {
"kv_cache_disabled": bool(
getattr(server_args, "prefill_only_disable_kv_cache", False)
),
"radix_cache_disabled": bool(
getattr(server_args, "disable_radix_cache", False)
),
"chunked_prefill_disabled": getattr(
server_args, "chunked_prefill_size", None
)
== -1,
},
}
def resolve_embedding_model_spec(
architectures: Sequence[str] | None,
*,
is_embedding_requested: bool,
is_embedding_gemma: bool,
) -> EmbeddingModelSpec:
"""Resolve a conservative embedding capability description.
Unknown architectures remain ``NONE`` unless the caller explicitly asks
for embedding mode. This preserves today's CausalLM behavior while
allowing model-specific support to be added in one location.
"""
architecture_set = set(architectures or ())
if is_embedding_gemma:
return _embedding_gemma_spec()
if architecture_set & _CLASSIFICATION_ARCHITECTURES:
return EmbeddingModelSpec(
family="sequence_classification",
task=EmbeddingTask.CLASSIFY,
execution=EmbeddingExecution.CLASSIFICATION,
attention=AttentionPattern.NONE,
pooling=PoolingStrategy.MODEL_DEFINED,
normalize=False,
postprocessor="model_defined",
tokenizer_special_tokens="model_default",
supports_dimensions=False,
supports_token_embeddings=False,
supports_multimodal=False,
requires_embedding_flag=False,
auto_enable_embedding=False,
bidirectional_attention=False,
bcg_prefill_policy=BCGPrefillPolicy.DEFAULT,
)
for architecture, (family, execution, pooling) in _EMBEDDING_ARCHITECTURES.items():
if architecture in architecture_set:
return _native_embedding_spec(family, execution, pooling)
if is_embedding_requested:
return EmbeddingModelSpec(
family="explicit_decoder_embedding",
task=EmbeddingTask.EMBED,
execution=EmbeddingExecution.DECODER_POOLING,
attention=AttentionPattern.CAUSAL,
pooling=PoolingStrategy.MODEL_DEFINED,
normalize=True,
postprocessor="model_defined",
tokenizer_special_tokens="model_default",
supports_dimensions=False,
supports_token_embeddings=False,
supports_multimodal=False,
requires_embedding_flag=True,
auto_enable_embedding=False,
bidirectional_attention=False,
bcg_prefill_policy=BCGPrefillPolicy.DEFAULT,
)
return EmbeddingModelSpec(
family="none",
task=EmbeddingTask.NONE,
execution=EmbeddingExecution.NONE,
attention=AttentionPattern.NONE,
pooling=PoolingStrategy.MODEL_DEFINED,
normalize=False,
postprocessor="none",
tokenizer_special_tokens="model_default",
supports_dimensions=False,
supports_token_embeddings=False,
supports_multimodal=False,
requires_embedding_flag=False,
auto_enable_embedding=False,
bidirectional_attention=False,
bcg_prefill_policy=BCGPrefillPolicy.DEFAULT,
)
python/sglang/srt/configs/model_config.py
+6
View File
@@ -25,6 +25,7 @@ from typing import Any, List, Optional, Set, Union
import torch
from transformers import PretrainedConfig
from sglang.srt.configs.embedding_model_spec import resolve_embedding_model_spec
from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config
from sglang.srt.environ import envs
from sglang.srt.layers.quantization import QUANTIZATION_METHODS
@@ -293,6 +294,11 @@ class ModelConfig:
)
self.hf_text_config = get_hf_text_config(self.hf_config)
self.is_embedding_gemma = is_embedding_gemma(self.hf_text_config)
self.embedding_model_spec = resolve_embedding_model_spec(
self.hf_config.architectures,
is_embedding_requested=bool(is_embedding),
is_embedding_gemma=self.is_embedding_gemma,
)
rope_scaling = getattr(self.hf_text_config, "rope_parameters", None) or getattr(
self.hf_text_config, "rope_scaling", {}
python/sglang/srt/entrypoints/grpc_bridge.py
+8
View File
@@ -15,6 +15,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Optional
from pydantic import ValidationError
from sglang.srt.configs.embedding_model_spec import resolved_embedding_plan
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
logger = logging.getLogger(__name__)
@@ -382,6 +383,13 @@ class RuntimeHandle:
"model_type": getattr(model_config.hf_config, "model_type", None),
"architectures": getattr(model_config.hf_config, "architectures", None),
}
embedding_model_spec = getattr(model_config, "embedding_model_spec", None)
if embedding_model_spec is not None:
result["embedding"] = resolved_embedding_plan(
embedding_model_spec,
server_args=self.server_args,
model_config=model_config,
)
return json.dumps(result, default=str)
def get_server_info(self) -> str:
python/sglang/srt/entrypoints/http_server.py
+8
View File
@@ -61,6 +61,7 @@ from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
from sglang.srt.configs.embedding_model_spec import resolved_embedding_plan
from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode
from sglang.srt.entrypoints.anthropic.protocol import (
@@ -717,6 +718,13 @@ async def model_info():
"weight_version": _global_state.tokenizer_manager.server_args.weight_version,
# "hf_config": model_config.hf_config.to_dict(),
}
embedding_model_spec = getattr(model_config, "embedding_model_spec", None)
if embedding_model_spec is not None:
result["embedding"] = resolved_embedding_plan(
embedding_model_spec,
server_args=_global_state.tokenizer_manager.server_args,
model_config=model_config,
)
return result
python/sglang/srt/entrypoints/openai/protocol.py
+1 -1
View File
@@ -1183,7 +1183,7 @@ class EmbeddingRequest(BaseModel):
class EmbeddingObject(BaseModel):
embedding: List[float]
embedding: Union[List[float], str]
index: int
object: str = "embedding"
python/sglang/srt/entrypoints/openai/serving_embedding.py
+28 -3
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
import base64
import struct
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import jinja2
@@ -41,6 +43,12 @@ class OpenAIServingEmbedding(OpenAIServingBase):
def _validate_request(self, request: EmbeddingRequest) -> Optional[str]:
"""Validate that the input is not empty or whitespace only."""
if request.encoding_format not in ("float", "base64"):
return (
"encoding_format must be either 'float' or 'base64', "
f"got {request.encoding_format!r}"
)
if not (input := request.input):
return "Input cannot be empty"
@@ -254,10 +262,25 @@ class OpenAIServingEmbedding(OpenAIServingBase):
if not isinstance(ret, list):
ret = [ret]
response = self._build_embedding_response(ret)
response = self._build_embedding_response(ret, request.encoding_format)
return response
def _build_embedding_response(self, ret: List[Dict[str, Any]]) -> EmbeddingResponse:
@staticmethod
def _serialize_embedding(
embedding: List[float], encoding_format: str
) -> Union[List[float], str]:
if encoding_format == "float":
return embedding
# OpenAI-compatible base64 embeddings contain contiguous little-endian
# float32 values. Explicit packing prevents the wire format from
# depending on the host's native byte order or Python float width.
encoded = struct.pack(f"<{len(embedding)}f", *embedding)
return base64.b64encode(encoded).decode("ascii")
def _build_embedding_response(
self, ret: List[Dict[str, Any]], encoding_format: str = "float"
) -> EmbeddingResponse:
"""Build the embedding response"""
embedding_objects = []
prompt_tokens = 0
@@ -265,7 +288,9 @@ class OpenAIServingEmbedding(OpenAIServingBase):
for idx, ret_item in enumerate(ret):
embedding_objects.append(
EmbeddingObject(
embedding=ret_item["embedding"],
embedding=self._serialize_embedding(
ret_item["embedding"], encoding_format
),
index=idx,
)
)
python/sglang/srt/server_args.py
+28 -1
View File
@@ -41,6 +41,7 @@ from sglang.srt.arg_groups.argparse_actions import (
DeprecatedStoreTrueAction,
LoRAPathAction,
)
from sglang.srt.configs.embedding_model_spec import BCGPrefillPolicy
from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_spec_by_arch
from sglang.srt.connector import ConnectorType
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
@@ -3626,7 +3627,33 @@ class ServerArgs:
# whose values depend on later prompt tokens, so both are invalid.
# Breakable CUDA Graph captures one complete prefill and is the graph
# mode validated for this encoder-style attention.
if getattr(model_config, "is_embedding_gemma", False):
# Native encoder architectures declare a pooling-only task and do not
# need the legacy --is-embedding intent flag. Decoder checkpoints still
# require that explicit opt-in because their architecture alone does
# not distinguish embedding from generation serving.
#
# ``_handle_model_capability_adjustments`` is also exercised directly
# by a few focused tests that use a small ModelConfig stand-in. Keep
# the old predicate as a compatibility fallback while production
# ModelConfig instances use the central capability contract.
embedding_model_spec = getattr(model_config, "embedding_model_spec", None)
if (
embedding_model_spec is not None
and embedding_model_spec.auto_enable_embedding
and not self.is_embedding
):
self.is_embedding = True
logger.info(
"Embedding architecture detected: enabling embedding mode automatically."
)
is_embedding_gemma = (
embedding_model_spec is not None
and embedding_model_spec.bcg_prefill_policy == BCGPrefillPolicy.FULL_ENCODER
)
if embedding_model_spec is None:
is_embedding_gemma = getattr(model_config, "is_embedding_gemma", False)
if is_embedding_gemma:
# This is an encoder-only model even though its HF architecture is
# named Gemma3TextModel. Marking it as embedding mode enables the
# FlashAttention raw-K/V fast path, which does not write or read