Support piecewise CUDA graph for embedding models (#13852)

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

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