From 540d6fee209b4fe8db020d08226d5d67fcfa02b0 Mon Sep 17 00:00:00 2001 From: StonyPort <157573149+zhooooong@users.noreply.github.com> Date: Wed, 26 Nov 2025 15:29:46 +0800 Subject: [PATCH] Support piecewise CUDA graph for embedding models (#13852) Co-authored-by: qiuxuan.lzw --- python/sglang/srt/layers/pooler.py | 2 +- .../sglang/srt/model_executor/model_runner.py | 3 +- .../piecewise_cuda_graph_runner.py | 5 ++- test/srt/test_piecewise_cuda_graph.py | 41 +++++++++++++++++++ 4 files changed, 48 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/pooler.py b/python/sglang/srt/layers/pooler.py index 2c98c856b..f9f8e1a18 100644 --- a/python/sglang/srt/layers/pooler.py +++ b/python/sglang/srt/layers/pooler.py @@ -10,7 +10,7 @@ import torch.nn as nn from transformers import PretrainedConfig from sglang.srt.layers.activation import get_cross_encoder_activation_function -from sglang.srt.model_executor.model_runner import ForwardBatch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch class PoolingType(IntEnum): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index a36311c85..e62f07746 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -88,6 +88,7 @@ from sglang.srt.layers.dp_attention import ( initialize_dp_attention, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.pooler import EmbeddingPoolerOutput from sglang.srt.layers.sampler import Sampler from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model from sglang.srt.lora.lora_manager import LoRAManager @@ -2201,7 +2202,7 @@ class ModelRunner: forward_batch: ForwardBatch, skip_attn_backend_init: bool = False, pp_proxy_tensors=None, - ) -> Union[LogitsProcessorOutput, PPProxyTensors]: + ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]: if self.is_multimodal and should_use_external_mm_preprocess(self.model): data_embedding_funcs = resolve_external_mm_data_embedding_funcs(self.model) diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 3ccf613dc..da57f4273 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -43,6 +43,7 @@ from sglang.srt.layers.dp_attention import ( set_is_extend_in_batch, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.pooler import EmbeddingPoolerOutput from sglang.srt.layers.torchao_utils import save_gemlite_cache from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -584,7 +585,7 @@ class PiecewiseCudaGraphRunner: self, forward_batch: ForwardBatch, **kwargs, - ) -> Union[LogitsProcessorOutput, PPProxyTensors]: + ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]: with enable_piecewise_cuda_graph(), disable_ca_comm(self.model_runner.tp_group): self.model_runner.attn_backend.init_forward_metadata(forward_batch) static_forward_batch = self.replay_prepare(forward_batch, **kwargs) @@ -610,6 +611,8 @@ class PiecewiseCudaGraphRunner: else None ), ) + elif isinstance(output, EmbeddingPoolerOutput): + return output else: assert isinstance(output, PPProxyTensors) # TODO(Yuwei): support PP Support diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index ea2db8fe1..5e2619e32 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/test/srt/test_piecewise_cuda_graph.py @@ -1,9 +1,14 @@ import unittest +import torch + +from sglang import Engine +from sglang.lang.chat_template import get_chat_template_by_model_path from sglang.srt.utils import get_device_sm, kill_process_tree from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( + DEFAULT_IMAGE_URL, DEFAULT_MODEL_NAME_FOR_TEST, DEFAULT_MODEL_NAME_FOR_TEST_MLA, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -330,5 +335,41 @@ class TestPiecewiseCudaGraphInternVL25(CustomTestCase): self.assertGreaterEqual(metrics["score"], 0.70) +class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase): + """Test piecewise CUDA graph with Qwen2.5-VL-3B-Instruct embedding model""" + + def test_embedding(self): + model_path = "Qwen/Qwen2.5-VL-3B-Instruct" + chat_template = get_chat_template_by_model_path(model_path) + text = f"{chat_template.image_token}What is in this picture? Answer: " + + engine = Engine( + model_path=model_path, + enable_multimodal=True, + is_embedding=True, + enable_piecewise_cuda_graph=True, + piecewise_cuda_graph_compiler="eager", + ) + out = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0]["embedding"] + engine.shutdown() + self.assertGreater(len(out), 0) + + engine = Engine( + model_path=model_path, + enable_multimodal=True, + is_embedding=True, + enable_piecewise_cuda_graph=False, + ) + out_without_pcg = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0][ + "embedding" + ] + engine.shutdown() + self.assertGreater(len(out_without_pcg), 0) + + self.assertTrue( + torch.allclose(torch.tensor(out), torch.tensor(out_without_pcg)) + ) + + if __name__ == "__main__": unittest.main()