Add overlap scheduling for embeddings code path (#14032)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

benchmark/prefill_only/util.py

@@ -553,6 +553,7 @@ async def process_results(
     print(f"  Server type:           {config.server_type}")
     print(f"  HTTP mode:             {http_mode}")
     print(f"  Target RPS:            {rps}")
+    print(f"  Achieved RPS:          {len(all_results) / total_duration:.2f}")
     print(f"  Item count:            {item_count}")
     print(f"  Distribution:          {config.distribution}")
     print(f"  Unique requests generated: {config.num_unique_requests}")

python/sglang/srt/managers/schedule_batch.py

@@ -1996,7 +1996,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         self.orig_seq_lens = self.orig_seq_lens[keep_indices_device]
         self.out_cache_loc = None
         self.seq_lens_sum = self.seq_lens.sum().item()
-        self.output_ids = self.output_ids[keep_indices_device]
+
+        if self.output_ids is not None:
+            self.output_ids = self.output_ids[keep_indices_device]
+
         self.mamba_track_indices = None
         self.mamba_track_mask = None
         self.mamba_track_seqlens = None

python/sglang/srt/managers/scheduler.py

@@ -210,6 +210,25 @@ GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300))
 @dataclass
 class EmbeddingBatchResult:
     embeddings: torch.Tensor
+    copy_done: Optional[torch.cuda.Event] = None
+
+    def copy_to_cpu(self):
+        """Copy embeddings tensor to CPU in overlap scheduling."""
+
+        if isinstance(self.embeddings, torch.Tensor):
+            self.copy_done = torch.get_device_module(self.embeddings.device).Event()
+            self.embeddings = self.embeddings.to("cpu", non_blocking=True)
+        else:
+            assert isinstance(self.embeddings, list)
+            if len(self.embeddings) == 0:
+                return
+
+            self.copy_done = torch.get_device_module(self.embeddings[0].device).Event()
+            self.embeddings = [
+                emb.to("cpu", non_blocking=True) for emb in self.embeddings
+            ]
+
+        self.copy_done.record()
 
 
 class Scheduler(
@@ -1083,7 +1102,9 @@ class Scheduler(
     @DynamicGradMode()
     def event_loop_overlap(self):
         """A scheduler loop that overlaps the CPU processing and GPU computation."""
-        self.result_queue: Deque[Tuple[ScheduleBatch, GenerationBatchResult]] = deque()
+        self.result_queue: Deque[
+            Tuple[ScheduleBatch, Union[GenerationBatchResult, EmbeddingBatchResult]]
+        ] = deque()
 
         def pop_and_process():
             # Process the results of the last batch
@@ -1124,7 +1145,8 @@ class Scheduler(
 
             # Run sample of the current batch
             # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed.
-            self.launch_batch_sample_if_needed(batch_result)
+            if self.is_generation:
+                self.launch_batch_sample_if_needed(batch_result)
 
             # Update last_batch
             self.last_batch = batch
@@ -2249,8 +2271,19 @@ class Scheduler(
             ret = batch_result
         else:  # embedding or reward model
             model_worker_batch = batch.get_model_worker_batch()
-            embeddings = self.tp_worker.forward_batch_embedding(model_worker_batch)
-            ret = EmbeddingBatchResult(embeddings=embeddings)
+
+            if self.enable_overlap:
+                self.record_batch_in_overlap(model_worker_batch)
+                with self.forward_stream_ctx:
+                    self.forward_stream.wait_stream(self.default_stream)
+                    embeddings = self.tp_worker.forward_batch_embedding(
+                        model_worker_batch
+                    )
+                    ret = EmbeddingBatchResult(embeddings=embeddings)
+                    ret.copy_to_cpu()
+            else:
+                embeddings = self.tp_worker.forward_batch_embedding(model_worker_batch)
+                ret = EmbeddingBatchResult(embeddings=embeddings)
 
         # Capture prefill end time for EXTEND mode
         if batch.forward_mode == ForwardMode.EXTEND:
@@ -2262,7 +2295,7 @@ class Scheduler(
 
     def launch_batch_sample_if_needed(
         self, batch_result: GenerationBatchResult
-    ) -> Union[GenerationBatchResult, EmbeddingBatchResult]:
+    ) -> Union[GenerationBatchResult]:
         # TODO(lsyin): make the delayed sample a default behavior after
         # unifying the forward_batch_generation interface (related to spec V2).
         if batch_result is None or batch_result.delay_sample_func is None:

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -224,6 +224,9 @@ class SchedulerOutputProcessorMixin:
                     )
 
         else:  # embedding or reward model
+            if result.copy_done is not None:
+                result.copy_done.synchronize()
+
             is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
 
             embeddings = result.embeddings

python/sglang/srt/server_args.py

@@ -1431,9 +1431,11 @@ class ServerArgs:
                         self.disable_radix_cache = True
                         self.disable_overlap_schedule = False
 
-        if not self.get_model_config().is_generation:
+        if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
             self.disable_overlap_schedule = True
-            logger.warning("Overlap scheduler is disabled for embedding models.")
+            logger.warning(
+                f"Overlap scheduler is disabled when using sparse head for embedding model."
+            )
 
         # TRTLLM AllReduce Fusion supports SM90/100/120, enable it by default
         # for models with explicit support (DeepseekV3, GptOss, Glm4Moe, Qwen3Moe)

test/srt/models/test_embedding_models.py

@@ -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: