From 7bf1a0bc22a2341333316778b0d14b477fa4339b Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 27 Jul 2026 13:44:42 +0800 Subject: [PATCH 1/6] refactor: centralize embedding model capabilities --- .../srt/configs/embedding_model_spec.py | 150 ++++++++++++++++++ python/sglang/srt/configs/model_config.py | 6 + python/sglang/srt/server_args.py | 24 ++- .../unit/configs/test_embedding_model_spec.py | 66 ++++++++ 4 files changed, 245 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/configs/embedding_model_spec.py create mode 100644 test/registered/unit/configs/test_embedding_model_spec.py 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 000000000000..1548bf4eaa9d --- /dev/null +++ b/python/sglang/srt/configs/embedding_model_spec.py @@ -0,0 +1,150 @@ +"""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 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" + + +@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 an embedding-only architecture. Decoder + checkpoints trained for embeddings still require explicit user intent. + """ + + family: str + task: EmbeddingTask + pooling: PoolingStrategy + normalize: 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 + + +_EMBEDDING_ARCHITECTURES = { + "BertModel": ("bert", PoolingStrategy.CLS), + "CLIPModel": ("clip", PoolingStrategy.LAST), + "Contriever": ("contriever", PoolingStrategy.MODEL_DEFINED), + "LlamaEmbeddingModel": ("llama_embedding", PoolingStrategy.LAST), + "MistralModel": ("mistral_embedding", PoolingStrategy.LAST), + "XLMRobertaModel": ("xlm_roberta", PoolingStrategy.CLS), +} + +_CLASSIFICATION_ARCHITECTURES = { + "BertForSequenceClassification", + "LlamaForSequenceClassification", + "LlamaForSequenceClassificationWithNormal_Weights", + "Qwen2ForSequenceClassification", + "Qwen3ForSequenceClassification", + "XLMRobertaForSequenceClassification", +} + + +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 EmbeddingModelSpec( + family="embeddinggemma", + task=EmbeddingTask.EMBED, + pooling=PoolingStrategy.MEAN, + normalize=True, + 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, + ) + + if architecture_set & _CLASSIFICATION_ARCHITECTURES: + return EmbeddingModelSpec( + family="sequence_classification", + task=EmbeddingTask.CLASSIFY, + pooling=PoolingStrategy.MODEL_DEFINED, + normalize=False, + requires_embedding_flag=False, + auto_enable_embedding=False, + bidirectional_attention=False, + bcg_prefill_policy=BCGPrefillPolicy.DEFAULT, + ) + + for architecture, (family, pooling) in _EMBEDDING_ARCHITECTURES.items(): + if architecture in architecture_set: + return EmbeddingModelSpec( + family=family, + task=EmbeddingTask.EMBED, + pooling=pooling, + normalize=True, + requires_embedding_flag=True, + auto_enable_embedding=False, + bidirectional_attention=False, + bcg_prefill_policy=BCGPrefillPolicy.DEFAULT, + ) + + if is_embedding_requested: + return EmbeddingModelSpec( + family="explicit_decoder_embedding", + task=EmbeddingTask.EMBED, + pooling=PoolingStrategy.MODEL_DEFINED, + normalize=True, + requires_embedding_flag=True, + auto_enable_embedding=False, + bidirectional_attention=False, + bcg_prefill_policy=BCGPrefillPolicy.DEFAULT, + ) + + return EmbeddingModelSpec( + family="none", + task=EmbeddingTask.NONE, + pooling=PoolingStrategy.MODEL_DEFINED, + normalize=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 a2839a89de2b..a7f836283962 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -25,6 +25,7 @@ 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 @@ def __init__( ) 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/server_args.py b/python/sglang/srt/server_args.py index 9bf05be7ed6b..152977b90e57 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -41,6 +41,7 @@ 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 ( @@ -3630,7 +3631,28 @@ def _handle_model_capability_adjustments(self): # 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): + # ``_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/unit/configs/test_embedding_model_spec.py b/test/registered/unit/configs/test_embedding_model_spec.py new file mode 100644 index 000000000000..01300ff72518 --- /dev/null +++ b/test/registered/unit/configs/test_embedding_model_spec.py @@ -0,0 +1,66 @@ +import unittest + +from sglang.srt.configs.embedding_model_spec import ( + BCGPrefillPolicy, + EmbeddingTask, + PoolingStrategy, + resolve_embedding_model_spec, +) +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) + + def test_encoder_embedding_models_still_require_explicit_intent(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.assertTrue(spec.requires_embedding_flag) + self.assertFalse(spec.auto_enable_embedding) + + 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() From 9cd56a04dd071af90fbc182b96b3ef4360ce45df Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 27 Jul 2026 13:46:52 +0800 Subject: [PATCH 2/6] feat: support base64 embedding responses --- .../sglang/srt/entrypoints/openai/protocol.py | 2 +- .../entrypoints/openai/serving_embedding.py | 31 +++++++++++++++++-- .../openai/test_serving_embedding.py | 25 +++++++++++++++ 3 files changed, 54 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 824fb6b220ee..405918329dd1 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 21ca3aa40b66..ed52ef8323b1 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 @@ def _request_id_prefix(self) -> str: 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 @@ async def _handle_non_streaming_request( 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 @@ def _build_embedding_response(self, ret: List[Dict[str, Any]]) -> EmbeddingRespo 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/test/registered/unit/entrypoints/openai/test_serving_embedding.py b/test/registered/unit/entrypoints/openai/test_serving_embedding.py index 46b5802c9324..09a2d9797151 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 @@ def test_non_jinja_template_errors_are_raised_as_value_error(self): 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) From 7a22d48e60abebff9339948bc03545d6d0c1c12d Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 27 Jul 2026 13:48:29 +0800 Subject: [PATCH 3/6] docs: clarify embedding defaults and BCG --- .../basic_usage/openai_api_embeddings.mdx | 24 +++++++++++-- .../supported-models/embedding_models.mdx | 36 ++++++++++++------- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/docs_new/docs/basic_usage/openai_api_embeddings.mdx b/docs_new/docs/basic_usage/openai_api_embeddings.mdx index a0528a0c407e..483893426341 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. Most decoder-style embedding models require `--is-embedding`. `google/embeddinggemma-300m` is detected automatically and does not need that flag. @@ -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/supported-models/embedding_models.mdx b/docs_new/docs/supported-models/embedding_models.mdx index 86d3ad01161f..ea637a4f3894 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` +Most decoder-style embedding models require `--is-embedding`, and some require `--trust-remote-code`. Encoder-only and sequence-classification architectures are recognized automatically. `google/embeddinggemma-300m` is also detected automatically. ## 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` From 1610c58a927747c781389b797dc9b38a9f220694 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 27 Jul 2026 13:51:52 +0800 Subject: [PATCH 4/6] bench: add vllm embedding backend --- .../docs/developer_guide/bench_serving.mdx | 33 +++++++++++++++++ python/sglang/benchmark/serving.py | 31 ++++++++++++---- .../bench_fn/test_benchmark_datasets_api.py | 36 ++++++++++++++++++- 3 files changed, 93 insertions(+), 7 deletions(-) diff --git a/docs_new/docs/developer_guide/bench_serving.mdx b/docs_new/docs/developer_guide/bench_serving.mdx index ec9e1db6a85a..5673ad318988 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/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 0b3936d6dfa7..4bd5e5f0105a 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -937,6 +937,7 @@ async def _call_profile_pd(profile_urls: List[Tuple[str, str]], mode: str) -> No "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 @@ async def _call_profile_pd(profile_urls: List[Tuple[str, str]], mode: str) -> No "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 limited_request_func(request_func_input, pbar): 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 limited_request_func(request_func_input, pbar): "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/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index 4c7c3d7f4f8f..09ff644040d6 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 @@ 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.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 From d1a64f589bdae41e435f53e3074fea13a2062b2c Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 27 Jul 2026 13:59:41 +0800 Subject: [PATCH 5/6] feat: auto-enable native encoder embeddings --- .../basic_usage/openai_api_embeddings.mdx | 2 +- .../supported-models/embedding_models.mdx | 2 +- .../srt/configs/embedding_model_spec.py | 9 +++++---- python/sglang/srt/server_args.py | 5 +++++ .../unit/configs/test_embedding_model_spec.py | 6 +++--- .../test_multimodal_piecewise_cuda_graph.py | 19 +++++++++++++++++++ 6 files changed, 34 insertions(+), 9 deletions(-) diff --git a/docs_new/docs/basic_usage/openai_api_embeddings.mdx b/docs_new/docs/basic_usage/openai_api_embeddings.mdx index 483893426341..44e7d8485fce 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. Most decoder-style embedding models require `--is-embedding`. `google/embeddinggemma-300m` is detected automatically and does not need that flag. +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`. diff --git a/docs_new/docs/supported-models/embedding_models.mdx b/docs_new/docs/supported-models/embedding_models.mdx index ea637a4f3894..366dcdd9b3e8 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. -Most decoder-style embedding models require `--is-embedding`, and some require `--trust-remote-code`. Encoder-only and sequence-classification architectures are recognized automatically. `google/embeddinggemma-300m` is also detected automatically. +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 diff --git a/python/sglang/srt/configs/embedding_model_spec.py b/python/sglang/srt/configs/embedding_model_spec.py index 1548bf4eaa9d..9e1860874051 100644 --- a/python/sglang/srt/configs/embedding_model_spec.py +++ b/python/sglang/srt/configs/embedding_model_spec.py @@ -35,8 +35,9 @@ 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 an embedding-only architecture. Decoder - checkpoints trained for embeddings still require explicit user intent. + checkpoint unambiguously declares a pooling/embedding-only architecture. + Decoder checkpoints trained for embeddings still require explicit user + intent. """ family: str @@ -120,8 +121,8 @@ def resolve_embedding_model_spec( task=EmbeddingTask.EMBED, pooling=pooling, normalize=True, - requires_embedding_flag=True, - auto_enable_embedding=False, + requires_embedding_flag=False, + auto_enable_embedding=True, bidirectional_attention=False, bcg_prefill_policy=BCGPrefillPolicy.DEFAULT, ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 152977b90e57..a8a2b3b73822 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3631,6 +3631,11 @@ def _handle_model_capability_adjustments(self): # 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. + # 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 diff --git a/test/registered/unit/configs/test_embedding_model_spec.py b/test/registered/unit/configs/test_embedding_model_spec.py index 01300ff72518..39eeb683dfa6 100644 --- a/test/registered/unit/configs/test_embedding_model_spec.py +++ b/test/registered/unit/configs/test_embedding_model_spec.py @@ -26,7 +26,7 @@ def test_embedding_gemma_declares_full_encoder_bcg_contract(self): self.assertTrue(spec.safe_disable_kv_cache) self.assertEqual(spec.bcg_prefill_policy, BCGPrefillPolicy.FULL_ENCODER) - def test_encoder_embedding_models_still_require_explicit_intent(self): + def test_encoder_embedding_models_enable_embedding_mode_automatically(self): spec = resolve_embedding_model_spec( ["BertModel"], is_embedding_requested=False, @@ -36,8 +36,8 @@ def test_encoder_embedding_models_still_require_explicit_intent(self): self.assertEqual(spec.family, "bert") self.assertEqual(spec.task, EmbeddingTask.EMBED) self.assertEqual(spec.pooling, PoolingStrategy.CLS) - self.assertTrue(spec.requires_embedding_flag) - self.assertFalse(spec.auto_enable_embedding) + self.assertFalse(spec.requires_embedding_flag) + self.assertTrue(spec.auto_enable_embedding) def test_decoder_embedding_intent_does_not_assume_encoder_fast_path(self): spec = resolve_embedding_model_spec( 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 fb0554174afa..0b82c263e0f9 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 @@ 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 @@ def test_embedding_gemma_forces_breakable_prefill(self): 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() From 32895357257ba513b8e097741809847e773b0752 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 27 Jul 2026 15:21:40 +0800 Subject: [PATCH 6/6] feat: expose resolved embedding capabilities --- docs_new/docs/basic_usage/native_api.mdx | 2 + .../srt/configs/embedding_model_spec.py | 255 +++++++++++++++--- python/sglang/srt/entrypoints/grpc_bridge.py | 8 + python/sglang/srt/entrypoints/http_server.py | 8 + .../unit/configs/test_embedding_model_spec.py | 56 ++++ 5 files changed, 298 insertions(+), 31 deletions(-) diff --git a/docs_new/docs/basic_usage/native_api.mdx b/docs_new/docs/basic_usage/native_api.mdx index 42c5bd228319..eff2293166d9 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/python/sglang/srt/configs/embedding_model_spec.py b/python/sglang/srt/configs/embedding_model_spec.py index 9e1860874051..c529e19993de 100644 --- a/python/sglang/srt/configs/embedding_model_spec.py +++ b/python/sglang/srt/configs/embedding_model_spec.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from enum import Enum -from typing import Sequence +from typing import Any, Sequence class EmbeddingTask(str, Enum): @@ -30,6 +30,30 @@ class BCGPrefillPolicy(str, Enum): 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. @@ -42,8 +66,15 @@ class EmbeddingModelSpec: 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 @@ -52,14 +83,61 @@ class EmbeddingModelSpec: 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", PoolingStrategy.CLS), - "CLIPModel": ("clip", PoolingStrategy.LAST), - "Contriever": ("contriever", PoolingStrategy.MODEL_DEFINED), - "LlamaEmbeddingModel": ("llama_embedding", PoolingStrategy.LAST), - "MistralModel": ("mistral_embedding", PoolingStrategy.LAST), - "XLMRobertaModel": ("xlm_roberta", PoolingStrategy.CLS), + "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 = { @@ -72,6 +150,121 @@ class EmbeddingModelSpec: } +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, *, @@ -88,51 +281,44 @@ def resolve_embedding_model_spec( architecture_set = set(architectures or ()) if is_embedding_gemma: - return EmbeddingModelSpec( - family="embeddinggemma", - task=EmbeddingTask.EMBED, - pooling=PoolingStrategy.MEAN, - normalize=True, - 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, - ) + 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, pooling) in _EMBEDDING_ARCHITECTURES.items(): + for architecture, (family, execution, pooling) in _EMBEDDING_ARCHITECTURES.items(): if architecture in architecture_set: - return EmbeddingModelSpec( - family=family, - task=EmbeddingTask.EMBED, - pooling=pooling, - normalize=True, - requires_embedding_flag=False, - auto_enable_embedding=True, - bidirectional_attention=False, - bcg_prefill_policy=BCGPrefillPolicy.DEFAULT, - ) + 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, @@ -142,8 +328,15 @@ def resolve_embedding_model_spec( 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, diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index fa0ab61b859f..ac8c25872d45 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -15,6 +15,7 @@ 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 @@ def get_model_info(self) -> str: "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 fd216d329ebb..2a843f156079 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -61,6 +61,7 @@ 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/test/registered/unit/configs/test_embedding_model_spec.py b/test/registered/unit/configs/test_embedding_model_spec.py index 39eeb683dfa6..6e1d5018e39c 100644 --- a/test/registered/unit/configs/test_embedding_model_spec.py +++ b/test/registered/unit/configs/test_embedding_model_spec.py @@ -1,10 +1,16 @@ 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 @@ -25,6 +31,9 @@ def test_embedding_gemma_declares_full_encoder_bcg_contract(self): 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( @@ -39,6 +48,53 @@ def test_encoder_embedding_models_enable_embedding_mode_automatically(self): 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"],