From d4801be447738522d2c62b49857c83787293be6f Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 19 Jul 2026 22:35:51 +0800 Subject: [PATCH] fix: fix vlm cuda graph shape stability (#30868) --- python/sglang/srt/compilation/compile.py | 51 ++++++++++-- .../srt/compilation/cuda_piecewise_backend.py | 23 +++--- .../srt/compilation/xpu_piecewise_backend.py | 2 +- python/sglang/srt/layers/attention/vision.py | 26 +++++-- .../srt/layers/rotary_embedding/utils.py | 14 +++- .../runner/prefill_cuda_graph_runner.py | 65 ++++++++++++++++ .../tc_piecewise_cuda_graph_backend.py | 8 ++ test/registered/cpu/test_rope.py | 6 ++ .../cuda_graph/test_cuda_piecewise_backend.py | 56 +++++++++++++ .../test_torch_compile_decoration.py | 78 +++++++++++++++++++ .../layers/attention/test_vision_seqlens.py | 25 ++++++ 11 files changed, 327 insertions(+), 27 deletions(-) create mode 100644 test/registered/cuda_graph/test_cuda_piecewise_backend.py create mode 100644 test/registered/unit/compilation/test_torch_compile_decoration.py create mode 100644 test/registered/unit/layers/attention/test_vision_seqlens.py diff --git a/python/sglang/srt/compilation/compile.py b/python/sglang/srt/compilation/compile.py index 1e1526fa5..c39175a5c 100644 --- a/python/sglang/srt/compilation/compile.py +++ b/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 diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index a750eb967..a6ba34e74 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/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) diff --git a/python/sglang/srt/compilation/xpu_piecewise_backend.py b/python/sglang/srt/compilation/xpu_piecewise_backend.py index ca269d135..50118700a 100644 --- a/python/sglang/srt/compilation/xpu_piecewise_backend.py +++ b/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( diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 27d7bae1c..f11ed2722 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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) diff --git a/python/sglang/srt/layers/rotary_embedding/utils.py b/python/sglang/srt/layers/rotary_embedding/utils.py index 882c68241..bfc43ee66 100644 --- a/python/sglang/srt/layers/rotary_embedding/utils.py +++ b/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, 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 34cc4043a..4e77a442f 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 @@ -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 diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py index bf6039f66..94d5f2681 100644 --- a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py +++ b/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) diff --git a/test/registered/cpu/test_rope.py b/test/registered/cpu/test_rope.py index d829a0622..4413519c5 100644 --- a/test/registered/cpu/test_rope.py +++ b/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) diff --git a/test/registered/cuda_graph/test_cuda_piecewise_backend.py b/test/registered/cuda_graph/test_cuda_piecewise_backend.py new file mode 100644 index 000000000..a7a80e3d8 --- /dev/null +++ b/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"])) diff --git a/test/registered/unit/compilation/test_torch_compile_decoration.py b/test/registered/unit/compilation/test_torch_compile_decoration.py new file mode 100644 index 000000000..82d7b2464 --- /dev/null +++ b/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"])) diff --git a/test/registered/unit/layers/attention/test_vision_seqlens.py b/test/registered/unit/layers/attention/test_vision_seqlens.py new file mode 100644 index 000000000..6e2ddf605 --- /dev/null +++ b/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"]))