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,