fix: fix vlm cuda graph shape stability (#30868)

python/sglang/srt/compilation/compile.py

@@ -86,6 +86,22 @@ def _mark_dynamic_on_value(val, dims):
         # else: ignore (None or non-tensor)
 
 
+_MROPE_TOKEN_AXIS_ARGUMENTS = frozenset(
+    {"positions", "position_ids", "mrope_positions"}
+)
+
+
+def _runtime_dynamic_dim_for_argument(name: str) -> int:
+    """Return the runtime-sized tensor dimension for a compiled argument.
+
+    Most compiler-facing tensors are batch/token-major and vary on dim 0.
+    mRoPE positions instead have shape ``[rope_axis, token]``; their runtime
+    token length is therefore the final dimension. ``maybe_mark_dynamic``
+    accepts ``-1`` for that final dimension irrespective of tensor rank.
+    """
+    return -1 if name in _MROPE_TOKEN_AXIS_ARGUMENTS else 0
+
+
 def _infer_dynamic_arg_dims_from_annotations(forward_fn):
     sig = inspect.signature(forward_fn)
     dyn = {}
@@ -96,7 +112,7 @@ def _infer_dynamic_arg_dims_from_annotations(forward_fn):
             ann is torch.Tensor
             or getattr(getattr(ann, "__args__", [None])[0], "__name__", "") == "Tensor"
         ):
-            dyn[name] = 0
+            dyn[name] = _runtime_dynamic_dim_for_argument(name)
         elif getattr(ann, "__name__", "") in ("IntermediateTensors",) or any(
             getattr(a, "__name__", "") == "IntermediateTensors"
             for a in getattr(ann, "__args__", [])
@@ -104,12 +120,33 @@ def _infer_dynamic_arg_dims_from_annotations(forward_fn):
             dyn[name] = 0
         elif ann == "torch.Tensor" or ann == "Optional[torch.Tensor]":
             # For future import annotations (e.g. from __future__ import annotations), the annotation is a string
-            dyn[name] = 0
+            dyn[name] = _runtime_dynamic_dim_for_argument(name)
     if not dyn:
         raise ValueError("No dynamic dims inferred; pass dynamic_arg_dims explicitly.")
     return dyn
 
 
+def _mark_dynamic_forward_batch(forward_batch) -> None:
+    """Mark runtime-sized ForwardBatch tensors before the first PCG trace.
+
+    ``ForwardBatch`` is deliberately not annotated as a Tensor, so the
+    annotation-based argument inference above cannot see its tensor fields.
+    PCG nevertheless reads sequence- and token-sized metadata from it in
+    attention layers. Leaving those dimensions static gives Dynamo a guard on
+    the dummy capture shape and silently creates a new backend for a real VLM
+    request. All direct tensor fields are batch-major except mRoPE positions,
+    whose token axis is the final dimension.
+    """
+    if forward_batch is None:
+        return
+
+    for name, value in vars(forward_batch).items():
+        if not isinstance(value, torch.Tensor) or value.ndim == 0:
+            continue
+        dims = _runtime_dynamic_dim_for_argument(name)
+        _mark_dynamic_on_value(value, dims)
+
+
 def install_torch_compiled(
     module: torch.nn.Module,
     *,
@@ -129,9 +166,8 @@ def install_torch_compiled(
     if backend_factory is None:
         from sglang.srt.compilation.backend import SGLangBackend
 
-        backend_factory = lambda gm, ex: SGLangBackend(compile_config, graph_pool)(
-            gm, ex
-        )
+        def backend_factory(gm, ex):
+            return SGLangBackend(compile_config, graph_pool)(gm, ex)
 
     compiled_codes: list[type(original_code)] = []
     state = {"compiled": False, "compiled_callable": None}
@@ -172,13 +208,16 @@ def install_torch_compiled(
                 val = ba.arguments[name]
                 if val is not None:
                     _mark_dynamic_on_value(val, dims)
+        _mark_dynamic_forward_batch(ba.arguments.get("forward_batch"))
 
         # Avoid cross-instance cache reuse
         torch._dynamo.eval_frame.remove_from_cache(unbound_fwd.__code__)
 
         bound = types.MethodType(unbound_fwd, self)
         compiled_callable = torch.compile(
-            bound, fullgraph=fullgraph, backend=backend_factory
+            bound,
+            fullgraph=fullgraph,
+            backend=backend_factory,
         )
 
         # Trigger Dynamo so bytecode hook can capture

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -18,11 +18,9 @@ from sglang.srt.compilation.compile_phase import (
     is_in_torch_compile_warmup,
 )
 from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
-from sglang.srt.utils import is_hip
 from sglang.srt.utils.common import print_warning_once
 
 logger = logging.getLogger(__name__)
-_is_hip = is_hip()
 
 
 @dataclasses.dataclass
@@ -155,22 +153,21 @@ class CUDAPiecewiseBackend:
 
             # During normal capture (PiecewiseCudaGraphRunner.capture()),
             # set_pcg_capture_stream() guarantees a valid stream. However,
-            # Dynamo may silently recompile on HIP/MLA serving batches whose
-            # token count exceeds the captured range. The replacement backend
-            # has no capture stream; fall back there instead of crashing while
-            # preserving the original assertion on other platforms.
+            # Dynamo may silently recompile serving batches when a dynamic
+            # multimodal input introduces a previously unseen guard. The
+            # replacement backend has no capture stream, so it cannot safely
+            # create a CUDA graph. Execute that subgraph normally instead of
+            # crashing the scheduler; subsequent matching shapes still use
+            # their captured graphs.
             stream = get_pcg_capture_stream()
-            if _is_hip and stream is None:
+            if stream is None:
                 print_warning_once(
-                    "PCG capture stream is not set; likely a Dynamo runtime "
-                    "recompilation. Falling back to eager execution for this "
+                    "PCG capture stream is not set. This can be a Dynamo runtime "
+                    "recompilation or an optional VLM branch pre-warmed outside "
+                    "CUDA graph capture; falling back to eager execution for this "
                     "subgraph."
                 )
                 return entry.runnable(*args)
-            assert (
-                stream is not None
-            ), "PCG capture stream is not set, please check if runtime recompilation happened"
-
             if self.compile_config.get_enable_debug_mode():
                 input_addresses = [
                     x.data_ptr() for x in args if isinstance(x, torch.Tensor)

python/sglang/srt/compilation/xpu_piecewise_backend.py

@@ -65,7 +65,7 @@ class XPUPiecewiseBackend(CUDAPiecewiseBackend):
             # at 8192 tokens when the capture grid tops out at 512). The
             # recompiled backend instance has no capture stream; fall back to
             # eager for that sub-graph instead of crashing the scheduler.
-            # Mirrors the HIP fallback in CUDAPiecewiseBackend.__call__.
+            # Mirrors the CUDA fallback in CUDAPiecewiseBackend.__call__.
             stream = get_pcg_capture_stream()
             if stream is None:
                 print_warning_once(

python/sglang/srt/layers/attention/vision.py

@@ -85,6 +85,7 @@ from sglang.srt.layers.linear import (
 )
 from sglang.srt.layers.quantization import QuantizationConfig
 from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb
+from sglang.srt.layers.rotary_embedding.utils import apply_rotary_pos_emb_native_eager
 from sglang.srt.runtime_context import get_server_args
 from sglang.srt.utils import add_prefix
 
@@ -203,11 +204,15 @@ def resolve_seqlens(
     return resolved_seqlens
 
 
-def resolve_max_seqlen(source, cu_seqlens: torch.Tensor) -> int:
-    """Return max segment length, caching it on a stable carrier so the
-    device->host sync (.item()) happens once per forward instead of once per layer.
+def resolve_max_seqlen(
+    source: torch.Tensor | SingletonCache | None, cu_seqlens: torch.Tensor
+) -> int:
+    """Return the maximum segment length, caching only on ``SingletonCache``.
+
+    Raw tensors have no mutable instance dictionary, so caching on them would
+    raise ``AttributeError``. They use the same calculation without a cache.
     """
-    if isinstance(source, SingletonCache) or isinstance(source, torch.Tensor):
+    if isinstance(source, SingletonCache):
         cached = getattr(source, "_max_seqlen", None)
         if cached is None:
             seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
@@ -1331,7 +1336,18 @@ class VisionAttention(nn.Module):
                 cos = torch.cat([cos, cos], dim=-1)
                 sin = torch.cat([sin, sin], dim=-1)
 
-            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+            # `apply_rotary_pos_emb` is torch.compile-decorated. Its first
+            # specialization may otherwise be compiled while a ViT CUDA graph
+            # is being captured, which makes Inductor attempt an illegal
+            # CPU-to-CUDA copy. The eager version is captured as part of the
+            # graph, so its pointwise work is still replayed without launch
+            # overhead.
+            rotary_fn = (
+                apply_rotary_pos_emb_native_eager
+                if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get()
+                else apply_rotary_pos_emb
+            )
+            q, k = rotary_fn(q, k, cos, sin)
             q = q.view(original_q_shape)
             k = k.view(original_k_shape)
 

python/sglang/srt/layers/rotary_embedding/utils.py

@@ -70,8 +70,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@torch.compile(dynamic=True, backend=get_compiler_backend())
-def apply_rotary_pos_emb_native(
+def apply_rotary_pos_emb_native_eager(
     q: torch.Tensor,
     k: torch.Tensor,
     cos: torch.Tensor,
@@ -94,6 +93,17 @@ def apply_rotary_pos_emb_native(
     return q_embed, k_embed
 
 
+@torch.compile(dynamic=True, backend=get_compiler_backend())
+def apply_rotary_pos_emb_native(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim=1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return apply_rotary_pos_emb_native_eager(q, k, cos, sin, unsqueeze_dim)
+
+
 def apply_rotary_pos_emb_npu(
     q: torch.Tensor,
     k: torch.Tensor,

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

@@ -544,6 +544,71 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         attn_backend.init_forward_metadata(fb)
         self._run_forward(fb, num_tokens)
 
+    def run_dummy_multimodal_deepstack_forward(
+        self, language_model: torch.nn.Module, num_tokens: int
+    ) -> bool:
+        """Warm the tensor-valued deepstack branch before serving requests.
+
+        The regular PCG dummy is text-only. Qwen3-VL only provides
+        ``input_deepstack_embeds`` after visual encoding, so leaving this
+        branch cold makes the first image request synchronously recompile the
+        language model. The model/signature checks keep this a no-op for
+        non-deepstack architectures.
+        """
+        if (
+            "input_deepstack_embeds"
+            not in inspect.signature(language_model.forward).parameters
+        ):
+            return False
+
+        num_deepstack = getattr(self.model_runner.model, "num_deepstack_embeddings", 0)
+        if num_deepstack <= 0:
+            return False
+
+        hidden_size = (
+            getattr(getattr(language_model, "config", None), "hidden_size", None)
+            or self.model_runner.model_config.hidden_size
+        )
+        fb, attn_backend = self.capture_prepare(num_tokens)
+        attn_backend.init_forward_metadata(fb)
+        deepstack_embeds = torch.zeros(
+            (num_tokens, hidden_size * num_deepstack),
+            dtype=self.model_runner.dtype,
+            device=self.device,
+        )
+        torch._dynamo.maybe_mark_dynamic(deepstack_embeds, 0)
+
+        fb.dp_local_start_pos = fb.dp_local_num_tokens = None
+        set_dp_buffer_len(
+            fb.global_dp_buffer_len,
+            num_tokens,
+            fb.dp_padding_mode.is_max_len(),
+            fb.global_num_tokens_cpu,
+        )
+        set_is_extend_in_batch(False)
+
+        with (
+            forward_context(
+                ForwardContext(attn_backend=self.model_runner.attn_backend)
+            ),
+            set_tc_piecewise_forward_context(
+                fb,
+                self.attention_layers,
+                self.quant_config,
+                self.moe_layers,
+                self.moe_fusions,
+                dsa_indexers=self.dsa_indexers,
+            ),
+        ):
+            language_model.forward(
+                fb.input_ids,
+                self._get_layer_model_positions(fb),
+                fb,
+                input_embeds=fb.input_embeds,
+                input_deepstack_embeds=deepstack_embeds,
+            )
+        return True
+
     def _has_inactive_dp_rank(self, forward_batch: ForwardBatch) -> bool:
         # DSV4 DP attention / DeepEP collectives need every DP rank to enter
         # the same replay path. Sparse-DP batches (one or more ranks with

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -200,6 +200,14 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
                                     f"Compiling num tokens ({num_tokens=})"
                                 )
                             cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
+
+                # Qwen3-VL deepstack embeddings are produced only after
+                # visual encoding. First trace the tensor branch above, then
+                # execute it once outside the compile-warmup marker so its
+                # regular kernel/JIT warmup also happens during startup.
+                cuda_graph_runner.run_dummy_multimodal_deepstack_forward(
+                    inner_model, cuda_graph_runner.capture_num_tokens[-1]
+                )
             finally:
                 _toggle_multi_platform_ops(inner_model, reverse=True, num_tokens=16)
 

test/registered/cpu/test_rope.py

@@ -11,6 +11,7 @@ from sglang.srt.layers.rotary_embedding.rope_variant import (
     DeepseekScalingRotaryEmbedding,
     apply_rotary_pos_emb_native,
 )
+from sglang.srt.layers.rotary_embedding.utils import apply_rotary_pos_emb_native_eager
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -275,9 +276,14 @@ class TestROPE(CustomTestCase):
             cos = torch.rand(num_tokens, head_size).to(sincos_dtype)
             sin = torch.rand(num_tokens, head_size).to(sincos_dtype)
             q_out_ref, k_out_ref = apply_rotary_pos_emb_native(query, key, cos, sin)
+            q_out_eager, k_out_eager = apply_rotary_pos_emb_native_eager(
+                query, key, cos, sin
+            )
             q_out_sgl, k_out_sgl = torch.ops.sgl_kernel.apply_rotary_pos_emb_cpu(
                 query, key, cos, sin
             )
+            torch.testing.assert_close(q_out_ref, q_out_eager)
+            torch.testing.assert_close(k_out_ref, k_out_eager)
             torch.testing.assert_close(q_out_ref, q_out_sgl, atol=1e-2, rtol=1e-2)
             torch.testing.assert_close(k_out_ref, k_out_sgl, atol=1e-2, rtol=1e-2)
 

test/registered/cuda_graph/test_cuda_piecewise_backend.py

@@ -0,0 +1,56 @@
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA piecewise backend requires CUDA", allow_module_level=True)
+
+import sglang.srt.compilation.cuda_piecewise_backend as cuda_backend
+from sglang.srt.compilation.cuda_piecewise_backend import (
+    ConcreteSizeEntry,
+    CUDAPiecewiseBackend,
+)
+
+
+def test_runtime_recompile_without_capture_stream_falls_back(monkeypatch):
+    """A Dynamo replacement backend cannot capture outside a PCG session."""
+    compile_config = MagicMock()
+    compile_config.get_capture_sizes.return_value = []
+    backend = CUDAPiecewiseBackend(
+        graph=MagicMock(),
+        compile_config=compile_config,
+        inductor_config={},
+        graph_pool=None,
+        piecewise_compile_index=0,
+        total_piecewise_compiles=1,
+        sym_shape_indices=[0],
+        compiled_graph_for_general_shape=MagicMock(),
+        sglang_backend=MagicMock(),
+    )
+    backend.first_run_finished = True
+    fallback = MagicMock(return_value="fallback-result")
+    backend.concrete_size_entries = {
+        4: ConcreteSizeEntry(
+            runtime_shape=4,
+            need_to_compile=False,
+            use_cudagraph=True,
+            runnable=fallback,
+            num_finished_warmup=1,
+        )
+    }
+
+    monkeypatch.setattr(cuda_backend, "get_pcg_capture_stream", lambda: None)
+    monkeypatch.setattr(cuda_backend, "is_in_torch_compile_warmup", lambda: False)
+    monkeypatch.setattr(cuda_backend, "print_warning_once", lambda _message: None)
+
+    assert backend(4) == "fallback-result"
+    fallback.assert_called_once_with(4)
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/compilation/test_torch_compile_decoration.py

@@ -0,0 +1,78 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+from sglang.srt.compilation.compile import (
+    _infer_dynamic_arg_dims_from_annotations,
+    _mark_dynamic_forward_batch,
+    _runtime_dynamic_dim_for_argument,
+)
+
+
+class _MropeModel:
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch,
+    ):
+        return input_ids, positions, forward_batch
+
+
+class _StringAnnotatedMropeModel:
+    def forward(
+        self,
+        input_ids: "torch.Tensor",
+        positions: "torch.Tensor",
+        forward_batch,
+    ):
+        return input_ids, positions, forward_batch
+
+
+def test_positions_marks_the_token_axis_dynamic_for_mrope_and_1d_rope():
+    dynamic_dims = _infer_dynamic_arg_dims_from_annotations(_MropeModel.forward)
+    string_dynamic_dims = _infer_dynamic_arg_dims_from_annotations(
+        _StringAnnotatedMropeModel.forward
+    )
+
+    assert dynamic_dims["input_ids"] == 0
+    assert dynamic_dims["positions"] == -1
+    assert string_dynamic_dims["positions"] == -1
+
+
+def test_runtime_dynamic_dim_uses_the_token_axis_for_mrope_metadata():
+    assert _runtime_dynamic_dim_for_argument("positions") == -1
+    assert _runtime_dynamic_dim_for_argument("position_ids") == -1
+    assert _runtime_dynamic_dim_for_argument("mrope_positions") == -1
+    assert _runtime_dynamic_dim_for_argument("input_ids") == 0
+
+
+def test_forward_batch_marks_token_and_batch_metadata_dynamic():
+    batch = SimpleNamespace(
+        input_embeds=torch.empty(8, 16),
+        seq_lens=torch.empty(2, dtype=torch.int64),
+        mrope_positions=torch.empty(3, 8, dtype=torch.int64),
+        scalar=torch.tensor(1),
+    )
+    marked = []
+
+    def record_mark_dynamic(value, dims):
+        marked.append((id(value), tuple(dims)))
+
+    with patch("torch._dynamo.maybe_mark_dynamic", side_effect=record_mark_dynamic):
+        _mark_dynamic_forward_batch(batch)
+
+    assert (id(batch.input_embeds), (0,)) in marked
+    assert (id(batch.seq_lens), (0,)) in marked
+    assert (id(batch.mrope_positions), (1,)) in marked
+    assert all(value_id != id(batch.scalar) for value_id, _ in marked)
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/layers/attention/test_vision_seqlens.py

@@ -0,0 +1,25 @@
+import pytest
+import torch
+
+from sglang.srt.layers.attention.vision import SingletonCache, resolve_max_seqlen
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def test_resolve_max_seqlen_accepts_raw_tensor_without_attribute_cache():
+    cu_seqlens = torch.tensor([0, 2, 7], dtype=torch.int32)
+
+    assert resolve_max_seqlen(cu_seqlens, cu_seqlens) == 5
+
+
+def test_resolve_max_seqlen_caches_on_singleton_carrier():
+    source = SingletonCache()
+    cu_seqlens = torch.tensor([0, 2, 7], dtype=torch.int32)
+
+    assert resolve_max_seqlen(source, cu_seqlens) == 5
+    assert source._max_seqlen == 5
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))