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

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