From c6184b7dc0847cfc622231c6fa3a00483fb1abcd Mon Sep 17 00:00:00 2001 From: siyu Date: Mon, 9 Mar 2026 11:10:40 +0800 Subject: [PATCH] Fix EPD OOM by offloading precomputed_embeddings during chunked prefill (#16503) --- python/sglang/srt/disaggregation/encode_receiver.py | 4 +++- python/sglang/srt/managers/mm_utils.py | 13 +++++++++++++ python/sglang/srt/managers/schedule_batch.py | 8 ++++++++ 3 files changed, 24 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 34cbc1220..5b0574a5a 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -160,7 +160,9 @@ class EmbeddingData: def get_embedding(self, is_concat=False): if is_concat: - return torch.concat([embedding.cuda() for embedding in self.embedding_list]) + return torch.concat( + [embedding.cuda() for embedding in self.embedding_list] + ).to("cpu", non_blocking=True) else: return self.embedding_list diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 48cd1e7fc..be557d454 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1241,6 +1241,19 @@ def general_mm_embed_routine( feature = getattr(mm_item, "feature", None) if isinstance(feature, torch.Tensor) and feature.is_cuda: mm_item.feature = feature.to("cpu", non_blocking=True) + if get_global_server_args().language_only: + precomputed_embeddings = getattr( + mm_item, "precomputed_embeddings", None + ) + if ( + isinstance(precomputed_embeddings, torch.Tensor) + and precomputed_embeddings.is_cuda + ): + mm_item.precomputed_embeddings = ( + precomputed_embeddings.to( + "cpu", non_blocking=True + ) + ) forward_batch.mm_inputs = None forward_batch.mm_input_embeds = input_embeds else: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5597d8bae..98ef554ab 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1639,6 +1639,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # The reference by CudaIpcTensorTransportProxy was cut off, # proactively delete to avoid slow gc. del pixel_values + if get_global_server_args().language_only: + precomputed_embeddings = getattr( + mm_item, "precomputed_embeddings", None + ) + if isinstance(precomputed_embeddings, torch.Tensor): + mm_item.precomputed_embeddings = precomputed_embeddings.to( + self.device, non_blocking=True + ) self.multimodal_inputs = multimodal_inputs self.token_type_ids = token_type_ids_tensor self.seq_lens_sum = sum(seq_lens)