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

Add overlap scheduling for embeddings code path (#14032)

Browse Source 
Co-authored-by: Satyam Kumar <satyamk@linkedin.com>
This commit is contained in:
satyamk7054
2025-12-24 18:24:18 -08:00
committed by GitHub
parent ecd2d09a6c
commit 38dd4fbb66
6 changed files with 63 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
test/srt/models/test_embedding_models.py
View File

@@ -28,14 +28,16 @@ from sglang.test.test_utils import (
is_in_ci,
)
MODELS = [
("Alibaba-NLP/gte-Qwen2-1.5B-instruct", 1, 1e-5),
("intfloat/e5-mistral-7b-instruct", 1, 1e-5),
("marco/mcdse-2b-v1", 1, 1e-5),
("Qwen/Qwen3-Embedding-8B", 1, 1e-5),
MODEL_TO_CONFIG = {
"Alibaba-NLP/gte-Qwen2-1.5B-instruct": (1, 1e-5),
"intfloat/e5-mistral-7b-instruct": (1, 1e-5),
"marco/mcdse-2b-v1": (1, 1e-5),
"Qwen/Qwen3-Embedding-8B": (1, 1e-5),
# Temporarily disable before this model is fixed
# ("jason9693/Qwen2.5-1.5B-apeach", 1, 1e-5),
]
# "jason9693/Qwen2.5-1.5B-apeach": (1, 1e-5),
}
MODELS = [(key, *MODEL_TO_CONFIG[key]) for key in MODEL_TO_CONFIG]
TORCH_DTYPES = [torch.float16]
@@ -123,11 +125,11 @@ class TestEmbeddingModels(CustomTestCase):
def test_matryoshka_embedding(self):
models_to_test = [
model
for model in MODELS
if "Alibaba-NLP/gte-Qwen2-1.5B-instruct" == model[0]
(
"Alibaba-NLP/gte-Qwen2-1.5B-instruct",
*MODEL_TO_CONFIG["Alibaba-NLP/gte-Qwen2-1.5B-instruct"],
)
]
assert len(models_to_test) == 1
for model, tp_size, prefill_tolerance in models_to_test:
for torch_dtype in TORCH_DTYPES: