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