From e49557b8da66e9c80bdc6659bb869d1403250175 Mon Sep 17 00:00:00 2001
From: Lianmin Zheng <lianminzheng@gmail.com>
Date: Sun, 16 Aug 2026 15:08:19 -0700
Subject: [PATCH] Support model-defined prefill input embedding width (#35002)

Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
---
 .../runner/prefill_cuda_graph_runner.py       | 29 +++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
index 167ea8c6a..a97adb875 100644
--- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
@@ -296,13 +296,19 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self.mamba_track_enabled = self._is_mamba_track_enabled()
 
         # --- buffers ---------------------------------------------------
+        # `hidden_size` here sizes only the multimodal `input_embeds` buffer,
+        # which `general_mm_embed_routine` copies the merged text+media
+        # embeddings into. A model whose merge happens above the embedding width
+        # (e.g. a residual-stream merge) writes a wider tensor than
+        # `config.hidden_size`, so let it declare that width.
+        input_embeds_hidden_size = self._input_embeds_hidden_size()
         self.buffers: PrefillInputBuffers = PrefillInputBuffers.create(
             device=self.device,
             max_bs=self.max_bs,
             max_num_tokens=self.max_num_tokens,
             cache_loc_dtype=self._cache_loc_dtype(),
             is_multimodal=self.is_multimodal,
-            hidden_size=self.model_runner.model_config.hidden_size,
+            hidden_size=input_embeds_hidden_size,
             dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
         )
@@ -316,7 +322,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             max_num_token=self.max_num_tokens,
             cache_loc_dtype=self._cache_loc_dtype(),
             is_multimodal=self.is_multimodal,
-            hidden_size=self.model_runner.model_config.hidden_size,
+            hidden_size=input_embeds_hidden_size,
             embed_dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
             enable_num_token_non_padded=enable_num_token_non_padded(),
@@ -531,6 +537,25 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     def _cache_loc_dtype(self):
         return torch.int64 if not is_npu() else torch.int32
 
+    def _input_embeds_hidden_size(self) -> int:
+        """Width of the `input_embeds` the model writes inside the graph.
+
+        Defaults to `config.hidden_size`; a model that merges multimodal
+        embeddings at a wider width declares it via `input_embeds_hidden_size`.
+
+        `getattr` with a default rather than an always-present field: the
+        property exists only on the handful of model classes whose merge width
+        differs, and adding it to every model in the registry to satisfy a
+        `None` check is not worth it. Same shape as the `hc_hidden_size`
+        opt-in already threaded through `base_runner` and
+        `decode_cuda_graph_runner`.
+        """
+        return getattr(
+            self.model_runner.model,
+            "input_embeds_hidden_size",
+            self.model_runner.model_config.hidden_size,
+        )
+
     def _next_token_logits_buffer(self, rows: int) -> Optional[torch.Tensor]:
         if not self.model_runner.pp_group.is_last_rank:
             return None