From 6092721594034f17f50d7063f42cbfd57898171e Mon Sep 17 00:00:00 2001 From: Binyao Jiang Date: Tue, 20 Jan 2026 14:06:06 -0800 Subject: [PATCH] [Piecewise] Fix PCG issue for multimodal and embedding model that wraps language_model (#17290) --- python/sglang/srt/model_executor/model_runner.py | 3 ++- .../sglang/srt/model_executor/piecewise_cuda_graph_runner.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d15a2f611..af67d52e1 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2027,9 +2027,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Collect attention layers and moe layers from the model self.model.model = resolve_language_model(self.model) + language_model = getattr(self.model, "language_model", self.model) self.attention_layers = [] self.moe_layers = [] - for layer in self.model.model.layers: + for layer in language_model.model.layers: if hasattr(layer, "self_attn"): if hasattr(layer.self_attn, "attn"): self.attention_layers.append(layer.self_attn.attn) 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 a6bd7ab56..51e72aebc 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -240,8 +240,11 @@ class PiecewiseCudaGraphRunner: set_graph_pool_id(get_global_graph_memory_pool()) with enable_piecewise_cuda_graph(): + language_model = getattr( + self.model_runner.model, "language_model", self.model_runner.model + ) with patch_model( - self.model_runner.model.model, self.compile_config.compiler + language_model.model, self.compile_config.compiler ) as patched_model: install_torch_compiled( patched_model,