From 3af991fb3eede82357568ec28882f11fe7830365 Mon Sep 17 00:00:00 2001 From: Oguz Ulgen Date: Fri, 19 Jun 2026 07:16:00 -0700 Subject: [PATCH] [AMD] Make breakable CUDA graph run on ROCm/HIP (#28173) --- .../breakable_cuda_graph.mdx | 6 ++-- .../srt/layers/attention/aiter_backend.py | 5 ++++ .../srt/layers/attention/triton_backend.py | 5 +++- .../breakable_cuda_graph_backend.py | 4 +-- .../breakable_cuda_graph.py | 28 +++++++++-------- .../breakable_cuda_graph/cuda_utils.py | 5 ---- python/sglang/srt/server_args.py | 10 +++---- .../breakable/test_breakable_cuda_graph.py | 30 ++----------------- 8 files changed, 36 insertions(+), 57 deletions(-) diff --git a/docs_new/docs/advanced_features/breakable_cuda_graph.mdx b/docs_new/docs/advanced_features/breakable_cuda_graph.mdx index d41bbb5cd..546af6484 100644 --- a/docs_new/docs/advanced_features/breakable_cuda_graph.mdx +++ b/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. -- **Requires `cuda-python`.** The `cuda.bindings` package must be installed (`pip install cuda-python`). +- **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` 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 python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py - CUDA runtime binding utilities + CUDA runtime binding utilities (NVIDIA stream-capture queries) python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 5c2149728..9d990a9ad 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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), diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index ad1151ed4..71e3972c3 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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) diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py index e16332df5..19ee9a2d8 100644 --- a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index 114a5df6e..03c0313c9 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/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() diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py index c95471012..df86e523e 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py +++ b/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. " diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 882360b4b..3b5cf53dc 100644 --- a/python/sglang/srt/server_args.py +++ b/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 diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index 12f7f8b7e..46c9e7733 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/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") - - -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) +register_amd_ci(est_time=120, suite="stage-c-test-large-8-gpu-amd-mi35x") 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,