chore: import upstream snapshot with attribution

python/sglang/srt/compilation/compile_phase.py

@@ -0,0 +1,57 @@
+"""torch.compile-internal phase markers used by the tc_piecewise backend.
+
+Two pieces of state, both private to the torch.compile path (the
+``cuda_piecewise_backend`` FX backend and the runner that drives it):
+
+* ``_in_torch_compile_warmup`` — true during the warmup-compile loop
+  where we run the compiled callable to trigger inductor compilation
+  but explicitly do **not** capture into a CUDA graph yet.
+  ``cuda_piecewise_backend`` reads this to short-circuit the capture
+  branch.
+* ``_pcg_capture_stream`` — the CUDA stream on which the runner is
+  performing capture, surfaced so the FX backend can use the same
+  stream for its own ``torch.cuda.graph(...)`` calls.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+import torch
+
+_in_torch_compile_warmup = False
+_pcg_capture_stream: torch.cuda.Stream | None = None
+
+
+def is_in_torch_compile_warmup() -> bool:
+    """True while inside the tc_piecewise warmup-compile pass. Strict subset of
+    ``torch.compiler.is_compiling()``.
+    """
+    return _in_torch_compile_warmup
+
+
+@contextmanager
+def enable_torch_compile_warmup():
+    """Mark the enclosed scope as the tc_piecewise warmup-compile pass. The FX
+    piecewise backend uses this to skip CUDA graph capture during warmup.
+    """
+    global _in_torch_compile_warmup
+    _in_torch_compile_warmup = True
+    try:
+        yield
+    finally:
+        _in_torch_compile_warmup = False
+
+
+def get_pcg_capture_stream() -> torch.cuda.Stream | None:
+    return _pcg_capture_stream
+
+
+@contextmanager
+def set_pcg_capture_stream(stream: torch.cuda.Stream):
+    global _pcg_capture_stream
+    _pcg_capture_stream = stream
+    try:
+        yield
+    finally:
+        _pcg_capture_stream = None