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

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:
-    _check_cuda_bindings()
+def _is_stream_capturing(stream: torch.cuda.Stream) -> bool:
+    # 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 (
-            _capture_status(other.cuda_stream)
-            != rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
-        ):
+        if not _is_stream_capturing(other):
             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
-    non-CUDA/NPU platforms, and we only reach this code during an active
-    Breakable capture (which can't happen on CPU-only runners anyway)."""
+    weak_ref_tensors is imported lazily because it hard-raises on
+    platforms without a CUDA/HIP/NPU backend; we only reach this code during
+    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
-        / memory-saver rejection in its own __init__; config-time
-        rules can be added here as they're discovered.
+        """Breakable (segmented capture, no torch.compile). Breakable enforces
+        memory-saver rejection in its own __init__; config-time rules can be
+        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