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

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"]))