[AMD] Make breakable CUDA graph run on ROCm/HIP (#28173)

docs_new/docs/advanced_features/breakable_cuda_graph.mdx

@@ -139,8 +139,8 @@ Some models fork work onto secondary CUDA streams (e.g., for overlapped computat
 
 ## Compatibility
 
-- **NVIDIA CUDA only.** Breakable CUDA graph is not supported on ROCm/HIP or other non-CUDA platforms. On unsupported platforms, `--debug-cuda-graph` is automatically disabled with a warning.
+- **CUDA and ROCm/HIP.** Breakable CUDA graph runs on both NVIDIA and AMD GPUs. Other platforms (NPU, CPU, MPS, XPU) are unsupported; there `--debug-cuda-graph` is automatically disabled with a warning.
-- **Requires `cuda-python`.** The `cuda.bindings` package must be installed (`pip install cuda-python`).
+- **Requires `cuda-python` on NVIDIA.** Stream-capture-status queries use the CUDA runtime via `cuda.bindings` (`pip install cuda-python`); the portable `torch.cuda.is_current_stream_capturing()` has proven unreliable on CUDA. On ROCm/HIP — where `cuda-python` is unavailable — the portable `torch.cuda` API (which maps to the HIP runtime) is used instead.
 - **Not compatible with memory saver mode.** Cannot be used together with `SGLANG_MEMORY_SAVER_CUDA_GRAPH`.
 
 ## Performance
@@ -174,7 +174,7 @@ For typical use cases with a small number of graph breaks, the overhead is negli
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA runtime binding utilities</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA runtime binding utilities (NVIDIA stream-capture queries)</td>
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py</code></td>

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

@@ -2323,6 +2323,11 @@ class AiterAttnBackend(AttentionBackend):
                 page_table = self.forward_metadata.swa_page_table
 
             extra_kwargs = {}
+            attn_out = getattr(forward_batch, "_attn_output", None)
+            if attn_out is not None and q.dtype != fp8_dtype:
+                extra_kwargs["out"] = attn_out.view(
+                    -1, layer.tp_q_head_num, layer.head_dim
+                )
 
             o = mha_batch_prefill_func(
                 q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),

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

@@ -1027,7 +1027,10 @@ class TritonAttnBackend(AttentionBackend):
         sinks=None,
     ):
         # TODO: reuse the buffer across layers
-        if layer.qk_head_dim != layer.v_head_dim:
+        attn_out = getattr(forward_batch, "_attn_output", None)
+        if attn_out is not None:
+            o = attn_out
+        elif layer.qk_head_dim != layer.v_head_dim:
             o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
         else:
             o = torch.empty_like(q)

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

@@ -36,7 +36,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import
     eager_on_graph,
     enable_breakable_cuda_graph,
 )
-from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils import get_bool_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
 if TYPE_CHECKING:
@@ -59,8 +59,6 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
         enable_memory_saver: bool = False,
         debug_eager: bool = False,
     ) -> None:
-        if is_hip():
-            raise RuntimeError("Breakable CUDA graph is not supported on ROCm/HIP")
         self._graphs: Dict[Any, BreakableCUDAGraph] = {}
         self._outputs: Dict[Any, Any] = {}
         self._pool = None

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -37,6 +37,7 @@ except ImportError:
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
     checkCudaErrors,
 )
+from sglang.srt.utils import is_hip
 
 logger = logging.getLogger(__name__)
 
@@ -51,7 +52,7 @@ __all__ = [
 def _check_cuda_bindings():
     if rt is None:
         raise ImportError(
-            "Breakable CUDA graph requires the 'cuda-python' package. "
+            "Breakable CUDA graph on NVIDIA requires the 'cuda-python' package. "
             "Install it with: pip install cuda-python"
         )
 
@@ -83,10 +84,16 @@ def _capture_status(stream_ptr: int) -> "rt.cudaStreamCaptureStatus":
     return status
 
 
-def _is_capturing(stream_ptr: int) -> bool:
+def _is_stream_capturing(stream: torch.cuda.Stream) -> bool:
-    _check_cuda_bindings()
+    # On ROCm/HIP, cuda-python is unavailable, so use the portable torch API
+    # (which maps to the HIP runtime). On NVIDIA, keep querying the CUDA runtime
+    # directly via cuda-python: torch.cuda.is_current_stream_capturing() has
+    # proven unreliable there, so we preserve the original behavior.
+    if is_hip():
+        with torch.cuda.stream(stream):
+            return torch.cuda.is_current_stream_capturing()
     return (
-        _capture_status(stream_ptr)
+        _capture_status(stream.cuda_stream)
         == rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
     )
 
@@ -116,10 +123,7 @@ def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
     is_other_cap = other is capturing or other.cuda_stream == cap_ptr
 
     if is_self_cap and not is_other_cap:
-        if (
+        if not _is_stream_capturing(other):
-            _capture_status(other.cuda_stream)
-            != rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
-        ):
             return
         _original_wait_stream(self, other)
         forked.discard(other)
@@ -155,9 +159,9 @@ def _weak_ref_if_tensor(x):
     mempool reclaim per-layer intermediates between segments — storage stays
     alive for each segment CUDAGraph's lifetime via its pool use_count.
 
-    weak_ref_tensors is imported lazily: the module hard-raises on
+    weak_ref_tensors is imported lazily because it hard-raises on
-    non-CUDA/NPU platforms, and we only reach this code during an active
+    platforms without a CUDA/HIP/NPU backend; we only reach this code during
-    Breakable capture (which can't happen on CPU-only runners anyway)."""
+    an active Breakable capture, which runs only on those backends."""
     if torch.is_tensor(x):
         from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
 
@@ -329,7 +333,7 @@ class BreakableCUDAGraphCapture:
         if forked:
             assert _original_wait_stream is not None
             for side in list(forked):
-                if _is_capturing(side.cuda_stream):
+                if _is_stream_capturing(side):
                     _original_wait_stream(main_stream, side)
             forked.clear()
         self.cuda_graph._segments[-1].capture_end()

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py

@@ -31,11 +31,6 @@ def _cudaGetErrorString(error):
 
 
 def checkCudaErrors(result):
-    if rt is None:
-        raise RuntimeError(
-            "cuda.bindings is not available. "
-            "Install it with: pip install cuda-python"
-        )
     if rt is None:
         raise RuntimeError(
             "cuda.bindings is not available. "

python/sglang/srt/server_args.py

@@ -1529,9 +1529,9 @@ class ServerArgs:
                 self.cuda_graph_config.prefill.backend = Backend.DISABLED
 
     def _disable_breakable_cudagraph_if_incompatible(self):
-        """Breakable (segmented capture, no torch.compile). Breakable enforces HIP
+        """Breakable (segmented capture, no torch.compile). Breakable enforces
-        / memory-saver rejection in its own __init__; config-time
+        memory-saver rejection in its own __init__; config-time rules can be
-        rules can be added here as they're discovered.
+        added here as they're discovered.
         """
         rules = [
             # MLA prefill takes a different attn-forward path under BCG (no
@@ -4351,9 +4351,9 @@ class ServerArgs:
                 )
             envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.set("0")
         if self.debug_cuda_graph:
-            if not is_cuda():
+            if not (is_cuda() or is_hip()):
                 logger.warning(
-                    "--debug-cuda-graph is not supported on non CUDA devices. "
+                    "--debug-cuda-graph is not supported on non CUDA/HIP devices. "
                     "Disabling breakable CUDA graph."
                 )
                 self.debug_cuda_graph = False

test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py

@@ -12,7 +12,7 @@ import unittest
 import torch
 
 from sglang.srt.utils import kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -24,21 +24,7 @@ from sglang.test.test_utils import (
 
 # CI Registration — large suite to fit the integration test's server startup.
 register_cuda_ci(est_time=79, stage="base-b", runner_config="1-gpu-large")
-
+register_amd_ci(est_time=120, suite="stage-c-test-large-8-gpu-amd-mi35x")
-
-def _skip_if_no_cuda(test_func):
-    return unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")(
-        test_func
-    )
-
-
-def _skip_if_no_cuda_bindings(test_func):
-    try:
-        from cuda.bindings import runtime as rt  # noqa: F401
-
-        return test_func
-    except ImportError:
-        return unittest.skip("cuda-python not installed")(test_func)
 
 
 class TestBreakableCUDAGraphBasic(CustomTestCase):
@@ -48,10 +34,6 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
     def setUpClass(cls):
         if not torch.cuda.is_available():
             raise unittest.SkipTest("CUDA not available")
-        try:
-            from cuda.bindings import runtime  # noqa: F401
-        except ImportError:
-            raise unittest.SkipTest("cuda-python not installed")
 
         from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
             BreakableCUDAGraph,
@@ -194,10 +176,6 @@ class TestCopyOutput(CustomTestCase):
     def setUpClass(cls):
         if not torch.cuda.is_available():
             raise unittest.SkipTest("CUDA not available")
-        try:
-            from cuda.bindings import runtime  # noqa: F401
-        except ImportError:
-            raise unittest.SkipTest("cuda-python not installed")
 
         from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
             _copy_output,
@@ -256,10 +234,6 @@ class TestBreakGraphHelper(CustomTestCase):
     def setUpClass(cls):
         if not torch.cuda.is_available():
             raise unittest.SkipTest("CUDA not available")
-        try:
-            from cuda.bindings import runtime  # noqa: F401
-        except ImportError:
-            raise unittest.SkipTest("cuda-python not installed")
 
         from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
             BreakableCUDAGraph,