Support piecewise CUDA graph for embedding models (#13852)

Co-authored-by: qiuxuan.lzw <qiuxuan.lzw@alibaba-inc.com>

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

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)

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

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()