From 38dd4fbb669ad4fb22b49e3b33cb3b1cd003abb1 Mon Sep 17 00:00:00 2001 From: satyamk7054 <43010011+satyamk7054@users.noreply.github.com> Date: Wed, 24 Dec 2025 18:24:18 -0800 Subject: [PATCH] Add overlap scheduling for embeddings code path (#14032) Co-authored-by: Satyam Kumar --- benchmark/prefill_only/util.py | 1 + python/sglang/srt/managers/schedule_batch.py | 5 ++- python/sglang/srt/managers/scheduler.py | 43 ++++++++++++++++--- .../scheduler_output_processor_mixin.py | 3 ++ python/sglang/srt/server_args.py | 6 ++- test/srt/models/test_embedding_models.py | 24 ++++++----- 6 files changed, 63 insertions(+), 19 deletions(-) diff --git a/benchmark/prefill_only/util.py b/benchmark/prefill_only/util.py index 3b3855916..2451239d6 100644 --- a/benchmark/prefill_only/util.py +++ b/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}") diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 72a0f711f..dbf761807 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a814ab1cd..d1bb2c2ed 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 83050f76b..7b68cae46 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ac6cb90e2..60157c888 100644 --- a/python/sglang/srt/server_args.py +++ b/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) diff --git a/test/srt/models/test_embedding_models.py b/test/srt/models/test_embedding_models.py index a93e762cf..d046166ae 100644 --- a/test/srt/models/test_embedding_models.py +++ b/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: