diff --git a/docs_new/docs/basic_usage/native_api.mdx b/docs_new/docs/basic_usage/native_api.mdx index 42c5bd228..eff229316 100644 --- a/docs_new/docs/basic_usage/native_api.mdx +++ b/docs_new/docs/basic_usage/native_api.mdx @@ -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", } ``` diff --git a/docs_new/docs/basic_usage/openai_api_embeddings.mdx b/docs_new/docs/basic_usage/openai_api_embeddings.mdx index a0528a0c4..44e7d8485 100644 --- a/docs_new/docs/basic_usage/openai_api_embeddings.mdx +++ b/docs_new/docs/basic_usage/openai_api_embeddings.mdx @@ -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) diff --git a/docs_new/docs/developer_guide/bench_serving.mdx b/docs_new/docs/developer_guide/bench_serving.mdx index ec9e1db6a..5673ad318 100644 --- a/docs_new/docs/developer_guide/bench_serving.mdx +++ b/docs_new/docs/developer_guide/bench_serving.mdx @@ -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`: diff --git a/docs_new/docs/supported-models/embedding_models.mdx b/docs_new/docs/supported-models/embedding_models.mdx index 86d3ad011..366dcdd9b 100644 --- a/docs_new/docs/supported-models/embedding_models.mdx +++ b/docs_new/docs/supported-models/embedding_models.mdx @@ -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. -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. ## 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"]) + + EmbeddingGemma + `google/embeddinggemma-300m` + N/A + Bidirectional Gemma3 text encoder; auto-detected and served with BCG by default + E5 (Llama/Mistral based) `intfloat/e5-mistral-7b-instruct` diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 0b3936d6d..4bd5e5f01 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -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") diff --git a/python/sglang/srt/configs/embedding_model_spec.py b/python/sglang/srt/configs/embedding_model_spec.py new file mode 100644 index 000000000..c529e1999 --- /dev/null +++ b/python/sglang/srt/configs/embedding_model_spec.py @@ -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, + ) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 6b62e86d7..5194d8dba 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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", {} diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index fa0ab61b8..ac8c25872 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -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: diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index fd216d329..2a843f156 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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 diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 824fb6b22..405918329 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -1183,7 +1183,7 @@ class EmbeddingRequest(BaseModel): class EmbeddingObject(BaseModel): - embedding: List[float] + embedding: Union[List[float], str] index: int object: str = "embedding" diff --git a/python/sglang/srt/entrypoints/openai/serving_embedding.py b/python/sglang/srt/entrypoints/openai/serving_embedding.py index 21ca3aa40..ed52ef832 100644 --- a/python/sglang/srt/entrypoints/openai/serving_embedding.py +++ b/python/sglang/srt/entrypoints/openai/serving_embedding.py @@ -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, ) ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c9bf4cbc1..9275a8526 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index 4c7c3d7f4..09ff64404 100644 --- a/test/registered/bench_fn/test_benchmark_datasets_api.py +++ b/test/registered/bench_fn/test_benchmark_datasets_api.py @@ -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 diff --git a/test/registered/unit/configs/test_embedding_model_spec.py b/test/registered/unit/configs/test_embedding_model_spec.py new file mode 100644 index 000000000..6e1d5018e --- /dev/null +++ b/test/registered/unit/configs/test_embedding_model_spec.py @@ -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() diff --git a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py index fb0554174..0b82c263e 100644 --- a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py +++ b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py @@ -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() diff --git a/test/registered/unit/entrypoints/openai/test_serving_embedding.py b/test/registered/unit/entrypoints/openai/test_serving_embedding.py index 46b5802c9..09a2d9797 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_embedding.py +++ b/test/registered/unit/entrypoints/openai/test_serving_embedding.py @@ -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)