[Fix] bcg: preserve Qwen3-VL DeepStack inputs during replay

python/sglang/srt/configs/model_config.py

@@ -2401,6 +2401,8 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
     "Qwen3_5MoeForConditionalGeneration",
     "Qwen4ExpForConditionalGeneration",
     "MuseGlimmerForConditionalGeneration",
+    "Qwen3VLForConditionalGeneration",
+    "Qwen3VLMoeForConditionalGeneration",
     "MuseSparkForConditionalGeneration",
     "KimiK3ForConditionalGeneration",
     "KimiK25ForConditionalGeneration",

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -831,6 +831,7 @@ def build_prefill_registry(
     is_multimodal: bool = False,
     hidden_size: int = 0,
     embed_dtype: Optional[torch.dtype] = None,
+    deepstack_replay_width: int = 0,
     enable_mamba_track: bool = False,
     enable_num_token_non_padded: bool = False,
     require_gathered_buffer: bool = False,
@@ -925,6 +926,17 @@ def _bs(bs: int, _mt: int) -> Tuple[int, ...]:
                 copy_from_fb=False,
             )
         )
+        if is_multimodal and deepstack_replay_width > 0:
+            slots.append(
+                GraphSlot(
+                    "input_deepstack_embeds",
+                    lambda _bs2, mt: (mt, deepstack_replay_width),
+                    embed_dtype,
+                    axis="tokens",
+                    padding_policy=PaddingPolicy.ZERO,
+                    copy_from_fb=False,
+                )
+            )
     if enable_mamba_track:
         slots.append(GraphSlot("mamba_track_indices", _bs, torch.int64, axis="bs"))
         slots.append(GraphSlot("mamba_track_mask", _bs, torch.bool, axis="bs"))

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -44,7 +44,15 @@
 from collections.abc import Sequence
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+    Protocol,
+    Union,
+    runtime_checkable,
+)
 
 import torch
 import tqdm
@@ -241,6 +249,57 @@ def _slice_output_rows(output: Any, num_tokens: int) -> Any:
     raise TypeError(f"Unsupported full prefill CUDA graph output: {type(output)}")
 
 
+@runtime_checkable
+class _SupportsDeepStackReplay(Protocol):
+    def get_num_deepstack_embeddings(self) -> int: ...
+
+
+def _resolve_deepstack_replay_width(
+    *, model: Any, hidden_size: int, uses_eager_prefill_tail: bool
+) -> int:
+    # Only the eager-tail backends (BREAKABLE/FULL) bind and refresh this slot;
+    # tc_piecewise captures the outer forward and never reads it.
+    if not uses_eager_prefill_tail:
+        return 0
+    if isinstance(model, _SupportsDeepStackReplay):
+        return hidden_size * model.get_num_deepstack_embeddings()
+    return 0
+
+
+def _refresh_deepstack_replay_slot(
+    slot: torch.Tensor, deepstack_embeds: Optional[torch.Tensor]
+) -> None:
+    """Sync the captured DeepStack slot for this replay.
+
+    Absence zeroes the slot so previous rows cannot bleed through the LM's
+    ``add_``; a non-empty tensor that does not fit raises instead — replaying
+    with zeros would silently serve the request without its visual part.
+    """
+    if deepstack_embeds is None or deepstack_embeds.numel() == 0:
+        slot.zero_()
+        return
+
+    num_tokens = slot.shape[0]
+    if (
+        deepstack_embeds.ndim != slot.ndim
+        or deepstack_embeds.shape[0] > num_tokens
+        or deepstack_embeds.shape[1:] != slot.shape[1:]
+        or deepstack_embeds.dtype != slot.dtype
+        or deepstack_embeds.device != slot.device
+    ):
+        raise RuntimeError(
+            "input_deepstack_embeds does not fit the captured BCG replay slot: "
+            f"got shape={tuple(deepstack_embeds.shape)}, "
+            f"dtype={deepstack_embeds.dtype}, device={deepstack_embeds.device}; "
+            f"slot holds up to {num_tokens} rows of {tuple(slot.shape[1:])} "
+            f"{slot.dtype} on {slot.device}."
+        )
+
+    num_rows = deepstack_embeds.shape[0]
+    slot[:num_rows].copy_(deepstack_embeds)
+    slot[num_rows:].zero_()
+
+
 @dataclass(frozen=True)
 class _ChunkedPrefixCaptureBuffers:
     """Runner-owned tensors shared by every prefix and token-bucket variant.
@@ -430,6 +489,12 @@ def __init__(self, model_runner: ModelRunner):
 
         self.mamba_track_enabled = self._is_mamba_track_enabled()
 
+        deepstack_replay_width = _resolve_deepstack_replay_width(
+            model=self.model_runner.model,
+            hidden_size=self.model_runner.model_config.hidden_size,
+            uses_eager_prefill_tail=self._uses_eager_prefill_tail(),
+        )
+
         # --- buffers ---------------------------------------------------
         # `hidden_size` here sizes only the multimodal `input_embeds` buffer,
         # which `general_mm_embed_routine` copies the merged text+media
@@ -446,6 +511,7 @@ def __init__(self, model_runner: ModelRunner):
             hidden_size=input_embeds_hidden_size,
             dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
+            deepstack_replay_width=deepstack_replay_width,
             enable_input_embeds=(self.is_multimodal or self._use_draft_input_embeds),
             pp_size=self.pp_size,
             is_first_pp_rank=self.model_runner.pp_group.is_first_rank,
@@ -471,6 +537,7 @@ def __init__(self, model_runner: ModelRunner):
             is_multimodal=self.is_multimodal,
             hidden_size=input_embeds_hidden_size,
             embed_dtype=self.model_runner.dtype,
+            deepstack_replay_width=deepstack_replay_width,
             enable_mamba_track=self.mamba_track_enabled,
             # FullCG always pads to a capture bucket. Models that mask padded
             # hidden rows need the live boundary even without expert parallelism.
@@ -874,6 +941,13 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int):
                 kwargs = _build_layer_model_forward_kwargs(
                     self.layer_model, forward_batch, pp_proxy_tensors
                 )
+                # DeepStack replay slot: bound by name like the other optional
+                # transformer inputs above, but sourced from the runner's
+                # capture buffer rather than the forward batch.
+                if self.buffer_registry.has_slot("input_deepstack_embeds"):
+                    kwargs["input_deepstack_embeds"] = self.buffer_registry.get_slot(
+                        "input_deepstack_embeds"
+                    ).slice_for(1, num_tokens)
                 if pp_proxy_tensors is not None:
                     input_ids = None
                     for embeds_name in ("input_embeds", "inputs_embeds"):
@@ -2183,6 +2257,13 @@ def replay_layer_forward(*args, **layer_kwargs):
                 and self.buffer_registry.has_slot("input_embeds")
             ):
                 self._fill_input_embeds_slot(args, layer_kwargs, static_num_tokens)
+            if self.buffer_registry.has_slot("input_deepstack_embeds"):
+                _refresh_deepstack_replay_slot(
+                    slot=self.buffer_registry.get_slot(
+                        "input_deepstack_embeds"
+                    ).slice_for(1, static_num_tokens),
+                    deepstack_embeds=layer_kwargs.get("input_deepstack_embeds"),
+                )
             hs = self.backend.replay(shape_key, static_forward_batch, **kwargs)
             hs = self._restore_qwen_bcg_hc_output(hs)
             return _slice_output_rows(hs, raw_num_tokens) if full_path else hs

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -280,6 +280,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
     positions: torch.Tensor
     input_embeds: Optional[torch.Tensor]
     mrope_positions: Optional[torch.Tensor]
+    input_deepstack_embeds: Optional[torch.Tensor]
     pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
 
     @classmethod
@@ -294,6 +295,7 @@ def create(
         hidden_size: int,
         dtype: torch.dtype,
         enable_mamba_track: bool,
+        deepstack_replay_width: int = 0,
         enable_input_embeds: Optional[bool] = None,
         pp_size: int = 1,
         is_first_pp_rank: bool = False,
@@ -331,8 +333,15 @@ def create(
 
             if is_multimodal:
                 mrope_positions = torch.zeros((3, max_num_tokens), dtype=torch.int64)
+                if deepstack_replay_width > 0:
+                    input_deepstack_embeds = torch.zeros(
+                        (max_num_tokens, deepstack_replay_width), dtype=dtype
+                    )
+                else:
+                    input_deepstack_embeds = None
             else:
                 mrope_positions = None
+                input_deepstack_embeds = None
 
             pp_proxy_tensors = (
                 allocate_pp_proxy_tensors(
@@ -359,5 +368,6 @@ def create(
             positions=positions,
             input_embeds=input_embeds,
             mrope_positions=mrope_positions,
+            input_deepstack_embeds=input_deepstack_embeds,
             pp_proxy_tensors=pp_proxy_tensors,
         )

python/sglang/srt/models/qwen3_vl.py

@@ -1411,6 +1411,10 @@ def __init__(
         # For EAGLE3 support
         self.capture_aux_hidden_states = False
 
+    # Opt-in for the prefill graph runner's DeepStack replay slot; 0 means no slot.
+    def get_num_deepstack_embeddings(self) -> int:
+        return self.num_deepstack_embeddings
+
     def separate_deepstack_embeds(self, embedding):
         assert embedding.shape[-1] % (1 + self.num_deepstack_embeddings) == 0, (
             f"hidden_state of {embedding.shape} should be divisible by ({1 + self.num_deepstack_embeddings})"