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
docs_new/docs/basic_usage/native_api.mdx
+2
View File
@@ -64,6 +64,7 @@ Get the information of the model.
- `has_audio_understanding`: Whether the model has audio-understanding capability.
- `model_type`: The model type from the HuggingFace config (e.g., "qwen2", "llama").
- `architectures`: The model architectures from the HuggingFace config (e.g., ["Qwen2ForCausalLM"]).
- `embedding`: The resolved embedding-serving plan. It includes pooling, normalization, execution and attention style, Matryoshka dimensions, cache policy, and effective BCG prefill settings. This field is available when the model configuration exposes an embedding capability contract.
```python Example
url = f"http://localhost:{port}/get_model_info"
@@ -85,6 +86,7 @@ assert response_json.keys() == {
"has_audio_understanding",
"model_type",
"architectures",
"embedding",
}
```
docs_new/docs/basic_usage/openai_api_embeddings.mdx
+21 -3
View File
@@ -12,7 +12,7 @@ This tutorial covers the embedding APIs for embedding models. For a list of the
## Launch A Server
Launch the server in your terminal and wait for it to initialize. Remember to add `--is-embedding` to the command.
Launch the server in your terminal and wait for it to initialize. Native encoder embedding architectures and `google/embeddinggemma-300m` are detected automatically. Decoder-style embedding models still require `--is-embedding`.
@@ -22,8 +22,8 @@ from sglang.utils import wait_for_server, print_highlight, terminate_process
embedding_process, port = launch_server_cmd(
"""
python3 -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct \
--host 0.0.0.0 --is-embedding --log-level warning
sglang serve --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct \
--is-embedding --log-level warning
"""
)
@@ -117,6 +117,24 @@ input_ids_embedding = json.loads(subprocess.check_output(curl_ids, shell=True))[
print_highlight(f"Input IDs embedding (first 10): {input_ids_embedding[:10]}")
```
## Compact Base64 Responses
Set `encoding_format` to `base64` when JSON arrays would dominate response size. The encoded value contains little-endian FP32 values and can be decoded by OpenAI-compatible clients.
```python Example
response = requests.post(
f"http://localhost:{port}/v1/embeddings",
json={
"model": "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
"input": text,
"encoding_format": "base64",
},
)
base64_embedding = response.json()["data"][0]["embedding"]
print_highlight(f"Base64 embedding: {base64_embedding[:20]}...")
```
```python Example
terminate_process(embedding_process)
docs_new/docs/developer_guide/bench_serving.mdx
+33
View File
@@ -16,6 +16,7 @@ This guide explains how to benchmark online serving throughput and latency using
- `sglang` / `sglang-native`: `POST /generate`
- `sglang-oai`, `vllm`, `lmdeploy`: `POST /v1/completions`
- `sglang-oai-chat`, `vllm-chat`, `lmdeploy-chat`: `POST /v1/chat/completions`
- `sglang-embedding`, `vllm-embedding`: `POST /v1/embeddings`
- `trt` (TensorRT-LLM): `POST /v2/models/ensemble/generate_stream`
- `gserver`: Custom server (Not Implemented yet in this script)
- `truss`: `POST /v1/models/model:predict`
@@ -55,6 +56,38 @@ python3 -m sglang.bench_serving \
--model meta-llama/Llama-3.1-8B-Instruct
```
### Fair embedding comparison
Use the two embedding backends with the same model, tokenizer, input length, prompt count, and concurrency. The benchmark reports input-token throughput and end-to-end latency; embeddings have no decode-side TTFT or TPOT.
```bash Command
# Start either server on the same hardware and precision, then run one at a time.
python3 -m sglang.bench_serving \
--backend sglang-embedding \
--model google/embeddinggemma-300m \
--dataset-name random \
--random-input-len 2048 \
--num-prompts 300 \
--max-concurrency 64 \
--warmup-requests 3 \
--flush-cache
```
```bash Command
# vLLM's cache reset endpoint requires VLLM_SERVER_DEV_MODE=1 at server startup.
python3 -m sglang.bench_serving \
--backend vllm-embedding \
--model google/embeddinggemma-300m \
--dataset-name random \
--random-input-len 2048 \
--num-prompts 300 \
--max-concurrency 64 \
--warmup-requests 3 \
--flush-cache
```
`--flush-cache` calls `/flush_cache` for SGLang and `/reset_prefix_cache` for vLLM after warmup. For vLLM, start the server with `VLLM_SERVER_DEV_MODE=1`; without it the benchmark fails loudly rather than accidentally measuring warm-cache performance.
### Datasets
Select with `--dataset-name`:
docs_new/docs/supported-models/embedding_models.mdx
+23 -13
View File
@@ -5,7 +5,7 @@ description: Dense and sparse embedding models with FlashInfer acceleration and
SGLang provides robust support for embedding models by integrating efficient serving mechanisms with its flexible programming interface. This integration allows for streamlined handling of embedding tasks, facilitating faster and more accurate retrieval and semantic search operations. SGLang's architecture enables better resource utilization and reduced latency in embedding model deployment.
<Warning>
Embedding models are executed with `--is-embedding` flag and some may require `--trust-remote-code`
Native encoder embedding architectures and `google/embeddinggemma-300m` are detected automatically. Decoder-style embedding models require `--is-embedding`; add `--trust-remote-code` when the model requires it.
</Warning>
## Quick Start
@@ -13,13 +13,21 @@ Embedding models are executed with `--is-embedding` flag and some may require `-
### Launch Server
```bash
python3 -m sglang.launch_server \
sglang serve \
--model-path Qwen/Qwen3-Embedding-4B \
--is-embedding \
--host 0.0.0.0 \
--port 30000
--is-embedding
```
### EmbeddingGemma
EmbeddingGemma uses bidirectional attention and is auto-detected, so its best default command is simply:
```bash
sglang serve --model-path google/embeddinggemma-300m
```
On CUDA, SGLang automatically uses breakable CUDA graph (BCG) for its full encoder prefill and disables incompatible radix-cache and chunked-prefill behavior. Do not add the deprecated piecewise CUDA graph knobs.
### Client Request
```python
@@ -30,7 +38,7 @@ url = "http://127.0.0.1:30000"
payload = {
"model": "Qwen/Qwen3-Embedding-4B",
"input": "What is the capital of France?",
"encoding_format": "float"
"encoding_format": "float" # or "base64" for compact FP32 responses
}
response = requests.post(url + "/v1/embeddings", json=payload).json()
@@ -43,12 +51,10 @@ print("Embedding:", response["data"][0]["embedding"])
For multimodal models like GME that support both text and images:
```bash
python3 -m sglang.launch_server \
sglang serve \
--model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct \
--is-embedding \
--chat-template gme-qwen2-vl \
--host 0.0.0.0 \
--port 30000
--chat-template gme-qwen2-vl
```
```python Example
@@ -85,11 +91,9 @@ print("Embeddings:", [x.get("embedding") for x in response.get("data", [])])
If the model config already includes `matryoshka_dimensions` or `is_matryoshka` then no override is needed. Otherwise, you can use `--json-model-override-args` as below:
```bash Command
python3 -m sglang.launch_server \
sglang serve \
--model-path Qwen/Qwen3-Embedding-0.6B \
--is-embedding \
--host 0.0.0.0 \
--port 30000 \
--json-model-override-args '{"matryoshka_dimensions": [128, 256, 512, 1024, 1536]}'
```
@@ -133,6 +137,12 @@ print("Embedding:", response["data"][0]["embedding"])
</tr>
</thead>
<tbody>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>EmbeddingGemma</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`google/embeddinggemma-300m`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>N/A</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Bidirectional Gemma3 text encoder; auto-detected and served with BCG by default</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>E5 (Llama/Mistral based)</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`intfloat/e5-mistral-7b-instruct`</td>
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
test/registered/bench_fn/test_benchmark_datasets_api.py
+35 -1
View File
@@ -11,7 +11,7 @@ import unittest
from collections import Counter
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import numpy as np
from PIL import Image
@@ -42,6 +42,13 @@ from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time
from sglang.benchmark.datasets.openai_dataset import sample_openai_requests
from sglang.benchmark.datasets.random import sample_random_requests
from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests
from sglang.benchmark.serving import (
_BACKEND_API_PATHS,
_EMBEDDING_BACKENDS,
ASYNC_REQUEST_FUNCS,
async_request_openai_embeddings,
flush_server_cache,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=40, suite="base-a-test-cpu")
@@ -87,6 +94,33 @@ def create_lightweight_tokenizer() -> PreTrainedTokenizerFast:
return hf_tokenizer
class TestEmbeddingBenchmarkBackends(unittest.TestCase):
def test_vllm_embedding_reuses_the_openai_embedding_request_path(self):
self.assertIn("vllm-embedding", _EMBEDDING_BACKENDS)
self.assertIs(
ASYNC_REQUEST_FUNCS["vllm-embedding"], async_request_openai_embeddings
)
self.assertEqual(_BACKEND_API_PATHS["vllm-embedding"], "/v1/embeddings")
def test_embedding_cache_flush_uses_the_engine_specific_endpoint(self):
with (
patch("sglang.benchmark.serving.get_auth_headers", return_value={}),
patch("sglang.benchmark.serving.requests.post") as post,
):
post.return_value = MagicMock()
flush_server_cache("http://127.0.0.1:8000", "vllm-embedding")
post.assert_called_once_with(
"http://127.0.0.1:8000/reset_prefix_cache", headers={}
)
post.reset_mock()
flush_server_cache("http://127.0.0.1:30000", "sglang-embedding")
post.assert_called_once_with(
"http://127.0.0.1:30000/flush_cache", headers={}
)
class DummyProcessor:
def __init__(self, tokenizer: PreTrainedTokenizerFast):
self.tokenizer = tokenizer
test/registered/unit/configs/test_embedding_model_spec.py
+122
View File
@@ -0,0 +1,122 @@
import unittest
from types import SimpleNamespace
from sglang.srt.configs.embedding_model_spec import (
AttentionPattern,
BCGEligibility,
BCGPrefillPolicy,
EmbeddingExecution,
EmbeddingTask,
PoolingStrategy,
embedding_support_matrix,
resolve_embedding_model_spec,
resolved_embedding_plan,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestEmbeddingModelSpec(unittest.TestCase):
def test_embedding_gemma_declares_full_encoder_bcg_contract(self):
spec = resolve_embedding_model_spec(
["Gemma3TextModel"],
is_embedding_requested=False,
is_embedding_gemma=True,
)
self.assertEqual(spec.family, "embeddinggemma")
self.assertEqual(spec.task, EmbeddingTask.EMBED)
self.assertEqual(spec.pooling, PoolingStrategy.MEAN)
self.assertTrue(spec.auto_enable_embedding)
self.assertTrue(spec.safe_disable_kv_cache)
self.assertEqual(spec.bcg_prefill_policy, BCGPrefillPolicy.FULL_ENCODER)
self.assertEqual(spec.execution, EmbeddingExecution.ENCODER_ONLY)
self.assertEqual(spec.attention, AttentionPattern.BIDIRECTIONAL)
self.assertEqual(spec.bcg_eligibility, BCGEligibility.FULL_ENCODER)
def test_encoder_embedding_models_enable_embedding_mode_automatically(self):
spec = resolve_embedding_model_spec(
["BertModel"],
is_embedding_requested=False,
is_embedding_gemma=False,
)
self.assertEqual(spec.family, "bert")
self.assertEqual(spec.task, EmbeddingTask.EMBED)
self.assertEqual(spec.pooling, PoolingStrategy.CLS)
self.assertFalse(spec.requires_embedding_flag)
self.assertTrue(spec.auto_enable_embedding)
def test_support_matrix_is_derived_from_the_same_registry(self):
matrix = embedding_support_matrix()
by_architecture = {row["architecture"]: row for row in matrix}
self.assertEqual(len(matrix), 7)
self.assertEqual(by_architecture["BertModel"]["family"], "bert")
self.assertEqual(by_architecture["BertModel"]["attention"], "bidirectional")
self.assertTrue(by_architecture["CLIPModel"]["supports_multimodal"])
self.assertEqual(
by_architecture["Gemma3TextModel (use_bidirectional_attention=true)"][
"bcg_eligibility"
],
"full_encoder",
)
def test_resolved_plan_reports_effective_runtime_knobs(self):
spec = resolve_embedding_model_spec(
["Gemma3TextModel"],
is_embedding_requested=False,
is_embedding_gemma=True,
)
plan = resolved_embedding_plan(
spec,
server_args=SimpleNamespace(
is_embedding=True,
cuda_graph_config=SimpleNamespace(
prefill=SimpleNamespace(
backend="breakable", max_bs=16384, bs=[1024, 16384]
)
),
prefill_only_disable_kv_cache=True,
disable_radix_cache=True,
chunked_prefill_size=-1,
),
model_config=SimpleNamespace(
is_matryoshka=False, matryoshka_dimensions=None
),
)
self.assertTrue(plan["enabled"])
self.assertTrue(plan["bcg"]["enabled"])
self.assertEqual(plan["bcg"]["capture_token_budget"], 16384)
self.assertEqual(plan["bcg"]["capture_batch_sizes"], [1024, 16384])
self.assertTrue(plan["cache"]["kv_cache_disabled"])
self.assertTrue(plan["cache"]["radix_cache_disabled"])
self.assertTrue(plan["cache"]["chunked_prefill_disabled"])
def test_decoder_embedding_intent_does_not_assume_encoder_fast_path(self):
spec = resolve_embedding_model_spec(
["Qwen3ForCausalLM"],
is_embedding_requested=True,
is_embedding_gemma=False,
)
self.assertEqual(spec.family, "explicit_decoder_embedding")
self.assertEqual(spec.task, EmbeddingTask.EMBED)
self.assertFalse(spec.safe_disable_kv_cache)
self.assertEqual(spec.bcg_prefill_policy, BCGPrefillPolicy.DEFAULT)
def test_unknown_generation_model_has_no_embedding_contract_without_intent(self):
spec = resolve_embedding_model_spec(
["Qwen3ForCausalLM"],
is_embedding_requested=False,
is_embedding_gemma=False,
)
self.assertEqual(spec.task, EmbeddingTask.NONE)
self.assertEqual(spec.family, "none")
if __name__ == "__main__":
unittest.main()
test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py
+19
View File
@@ -4,6 +4,7 @@ import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.srt.configs.embedding_model_spec import resolve_embedding_model_spec
from sglang.srt.configs.model_config import (
is_multimodal_piecewise_cuda_graph_supported,
)
@@ -128,6 +129,24 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
self.assertEqual(args.cuda_graph_config.decode.backend, Backend.DISABLED)
self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE)
def test_encoder_embedding_model_enables_embedding_mode_without_flag(self):
args = ServerArgs(model_path="dummy")
args.is_embedding = False
args.model_config = SimpleNamespace(
embedding_model_spec=resolve_embedding_model_spec(
["BertModel"],
is_embedding_requested=False,
is_embedding_gemma=False,
),
is_multimodal=False,
hf_config=SimpleNamespace(architectures=["BertModel"]),
)
with patch.object(args, "get_model_config", return_value=args.model_config):
args._handle_model_capability_adjustments()
self.assertTrue(args.is_embedding)
if __name__ == "__main__":
unittest.main()
test/registered/unit/entrypoints/openai/test_serving_embedding.py
+25
View File
@@ -2,9 +2,11 @@
Unit tests for the OpenAIServingEmbedding class from serving_embedding.py.
"""
import base64
import importlib
import importlib.abc
import importlib.machinery
import struct
import sys
import types
import unittest
@@ -339,6 +341,29 @@ class ServingEmbeddingTestCase(unittest.TestCase):
self.image_only_multimodal_req
)
def test_base64_embedding_response_uses_little_endian_float32(self):
response = self.serving_embedding._build_embedding_response(
[{"embedding": [0.25, -1.5], "meta_info": {"prompt_tokens": 2}}],
encoding_format="base64",
)
encoded_embedding = response.data[0].embedding
self.assertIsInstance(encoded_embedding, str)
self.assertEqual(
struct.unpack("<2f", base64.b64decode(encoded_embedding)), (0.25, -1.5)
)
self.assertEqual(response.usage.prompt_tokens, 2)
def test_rejects_unknown_embedding_encoding_format(self):
invalid_request = EmbeddingRequest(
model="test-model", input="hello", encoding_format="binary"
)
self.assertIn(
"encoding_format must be either",
self.serving_embedding._validate_request(invalid_request),
)
if __name__ == "__main__":
unittest.main(verbosity=2)