[AMD] Enable Piecewise CUDA Graph for AMD GPUs (#22299)

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -18,8 +18,11 @@ from sglang.srt.compilation.piecewise_context_manager import (
     is_in_pcg_torch_compile,
 )
 from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
+from sglang.srt.utils import is_hip
+from sglang.srt.utils.common import print_warning_once
 
 logger = logging.getLogger(__name__)
+_is_hip = is_hip()
 
 
 @dataclasses.dataclass
@@ -150,6 +153,24 @@ class CUDAPiecewiseBackend:
                 entry.num_finished_warmup += 1
                 return entry.runnable(*args)
 
+            # During normal capture (PiecewiseCudaGraphRunner.capture()),
+            # set_pcg_capture_stream() guarantees a valid stream. However,
+            # Dynamo may silently recompile on HIP/MLA serving batches whose
+            # token count exceeds the captured range. The replacement backend
+            # has no capture stream; fall back there instead of crashing while
+            # preserving the original assertion on other platforms.
+            stream = get_pcg_capture_stream()
+            if _is_hip and stream is None:
+                print_warning_once(
+                    "PCG capture stream is not set; likely a Dynamo runtime "
+                    "recompilation. Falling back to eager execution for this "
+                    "subgraph."
+                )
+                return entry.runnable(*args)
+            assert (
+                stream is not None
+            ), "PCG capture stream is not set, please check if runtime recompilation happened"
+
             if self.compile_config.get_enable_debug_mode():
                 input_addresses = [
                     x.data_ptr() for x in args if isinstance(x, torch.Tensor)
@@ -168,10 +189,6 @@ class CUDAPiecewiseBackend:
                     stack.enter_context(patch("gc.collect", lambda: None))
                     stack.enter_context(patch("torch.cuda.empty_cache", lambda: None))
                 # mind-exploding: carefully manage the reference and memory.
-                stream = get_pcg_capture_stream()
-                assert (
-                    stream is not None
-                ), "PCG capture stream is not set, please check if runtime recompilation happened"
                 with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream):
                     # `output` is managed by pytorch's cudagraph pool
                     output = entry.runnable(*args)

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -72,6 +72,8 @@ class ForwardContext:
         self.moe_layers = None
         self.moe_fusions = None
         self.dsa_indexers = None
+        self.num_tokens: Optional[int] = None
+        self.raw_num_tokens: Optional[int] = None
 
     def set_forward_batch(self, forward_batch: ForwardBatch):
         self.forward_batch = forward_batch
@@ -109,6 +111,8 @@ def set_forward_context(
     moe_layers: List[Any],
     moe_fusions: List[Any],
     dsa_indexers: Optional[List[Any]] = None,
+    num_tokens: Optional[int] = None,
+    raw_num_tokens: Optional[int] = None,
 ):
     global _forward_context
     _forward_context = ForwardContext()
@@ -119,6 +123,8 @@ def set_forward_context(
     _forward_context.set_moe_fusions(moe_fusions)
     if dsa_indexers is not None:
         _forward_context.set_dsa_indexers(dsa_indexers)
+    _forward_context.num_tokens = num_tokens
+    _forward_context.raw_num_tokens = raw_num_tokens
     try:
         yield
     finally:

python/sglang/srt/layers/moe/hash_topk.py

@@ -17,6 +17,7 @@ from sglang.srt.eplb.expert_location_dispatch import (
 from sglang.srt.layers.moe.topk import (
     StandardTopKOutput,
     _mask_topk_ids_padded_region,
+    _zero_topk_weights_padded_region,
 )
 from sglang.srt.utils import is_hip
 
@@ -175,9 +176,17 @@ class HashTopK(nn.Module):
             topk_weights = topk_weights.to(torch.float32)
 
         topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
-        _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
+        if is_hip():
+            _zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
+        else:
+            _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
         get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
         topk_output = StandardTopKOutput(
             topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
         )
-        return self._apply_deepep_waterfill(topk_output, hidden_states.shape[0])
+        topk_output = self._apply_deepep_waterfill(topk_output, hidden_states.shape[0])
+        if is_hip():
+            _zero_topk_weights_padded_region(
+                topk_output.topk_weights, num_token_non_padded
+            )
+        return topk_output

python/sglang/srt/layers/moe/topk.py

@@ -1125,6 +1125,16 @@ def _mask_topk_ids_padded_region(
         topk_ids[indices >= num_token_non_padded, :] = -1
 
 
+def _zero_topk_weights_padded_region(
+    topk_weights: torch.Tensor,
+    num_token_non_padded: Optional[torch.Tensor] = None,
+):
+    if num_token_non_padded is None:
+        return
+    indices = torch.arange(0, topk_weights.shape[0], device=topk_weights.device)
+    topk_weights[indices >= num_token_non_padded, :] = 0.0
+
+
 @torch.compile(dynamic=True, backend=get_compiler_backend())
 def _biased_grouped_topk_postprocess(
     topk_ids, expert_location_dispatch_info, num_token_non_padded
@@ -1495,6 +1505,12 @@ def _post_process_topk_ids(
             topk_ids = _biased_grouped_topk_postprocess(
                 topk_ids, expert_location_dispatch_info, num_token_non_padded
             )
+    elif _is_hip:
+        # On AMD HIP, the aiter MoE kernels do not handle topk_ids=-1 safely
+        # (negative indices cause illegal memory access). Instead, zero the
+        # routing weights for padded tokens so their MoE output contributes
+        # nothing to the hidden state after the weighted sum.
+        _zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
 
     if recorder_topk_ids is None:
         recorder_topk_ids = topk_ids
@@ -1536,6 +1552,12 @@ def _post_process_topk_ids(
             topk_config,
         )
 
+    if _is_hip:
+        # Shared-expert append/remap can introduce non-zero weights after the
+        # initial HIP padding mask above. Ensure padded tokens leave this helper
+        # with all expert weights zeroed.
+        _zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
+
     return topk_ids, topk_weights, recorder_topk_ids
 
 

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py

@@ -8,16 +8,145 @@ import torch
 from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
 from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme
 from sglang.srt.utils import is_hip
-from sglang.srt.utils.common import mxfp_supported
+from sglang.srt.utils.common import direct_register_custom_op, mxfp_supported
 
 _is_hip = is_hip()
 if _is_hip:
     from aiter.ops.triton.gemm.fused.fused_gemm_afp4wfp4_split_cat import (
-        fused_gemm_afp4wfp4_split_cat,
+        fused_gemm_afp4wfp4_split_cat as _fused_gemm_afp4wfp4_split_cat_orig,
     )
-    from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
+    from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 as _gemm_afp4wfp4_orig
-    from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import gemm_afp4wfp4_pre_quant
+    from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import (
-    from aiter.ops.triton.quant import dynamic_mxfp4_quant
+        gemm_afp4wfp4_pre_quant as _gemm_afp4wfp4_pre_quant_orig,
+    )
+    from aiter.ops.triton.quant import dynamic_mxfp4_quant as _dynamic_mxfp4_quant_orig
+
+    def _aiter_gemm_afp4wfp4(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        x_scales: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        _gemm_afp4wfp4_orig(x, w, x_scales, w_scales, y.dtype, y)
+
+    def _aiter_gemm_afp4wfp4_fake(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        x_scales: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        return None
+
+    direct_register_custom_op(
+        op_name="aiter_gemm_afp4wfp4",
+        op_func=_aiter_gemm_afp4wfp4,
+        mutates_args=["y"],
+        fake_impl=_aiter_gemm_afp4wfp4_fake,
+    )
+
+    def gemm_afp4wfp4(x, w, x_scales, w_scales, dtype, y):
+        torch.ops.sglang.aiter_gemm_afp4wfp4(x, w, x_scales, w_scales, y)
+
+    def _aiter_gemm_afp4wfp4_pre_quant(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        _gemm_afp4wfp4_pre_quant_orig(x, w, w_scales, y.dtype, y)
+
+    def _aiter_gemm_afp4wfp4_pre_quant_fake(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        return None
+
+    direct_register_custom_op(
+        op_name="aiter_gemm_afp4wfp4_pre_quant",
+        op_func=_aiter_gemm_afp4wfp4_pre_quant,
+        mutates_args=["y"],
+        fake_impl=_aiter_gemm_afp4wfp4_pre_quant_fake,
+    )
+
+    def gemm_afp4wfp4_pre_quant(x, w, w_scales, dtype, y):
+        torch.ops.sglang.aiter_gemm_afp4wfp4_pre_quant(x, w, w_scales, y)
+
+    def _aiter_dynamic_mxfp4_quant(
+        x: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _dynamic_mxfp4_quant_orig(x)
+
+    def _aiter_dynamic_mxfp4_quant_fake(
+        x: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        M, N = x.shape
+        x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
+        blockscale = torch.empty(
+            (M, (N + 31) // 32), dtype=torch.uint8, device=x.device
+        )
+        return x_fp4, blockscale
+
+    direct_register_custom_op(
+        op_name="aiter_dynamic_mxfp4_quant",
+        op_func=_aiter_dynamic_mxfp4_quant,
+        mutates_args=[],
+        fake_impl=_aiter_dynamic_mxfp4_quant_fake,
+    )
+
+    def dynamic_mxfp4_quant(x):
+        return torch.ops.sglang.aiter_dynamic_mxfp4_quant(x)
+
+    def _aiter_fused_gemm_split_cat(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        y: torch.Tensor,
+        x_scale: torch.Tensor,
+        w_scale: torch.Tensor,
+        S1: int,
+        S2: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _fused_gemm_afp4wfp4_split_cat_orig(
+            x=x,
+            w=w,
+            y=y,
+            x_scale=x_scale,
+            w_scale=w_scale,
+            S1=S1,
+            S2=S2,
+            dtype=y.dtype,
+        )
+
+    def _aiter_fused_gemm_split_cat_fake(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        y: torch.Tensor,
+        x_scale: torch.Tensor,
+        w_scale: torch.Tensor,
+        S1: int,
+        S2: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        M = x.shape[0]
+        D = y.shape[1]
+        S3 = y.shape[2]
+        c1 = torch.empty((M, D, S1 + S3), dtype=y.dtype, device=x.device)
+        c2 = torch.empty((M, D, S2), dtype=y.dtype, device=x.device)
+        return c1, c2
+
+    direct_register_custom_op(
+        op_name="aiter_fused_gemm_split_cat",
+        op_func=_aiter_fused_gemm_split_cat,
+        mutates_args=[],
+        fake_impl=_aiter_fused_gemm_split_cat_fake,
+    )
+
+    def fused_gemm_afp4wfp4_split_cat(x, w, y, x_scale, w_scale, S1, S2, dtype):
+        return torch.ops.sglang.aiter_fused_gemm_split_cat(
+            x, w, y, x_scale, w_scale, S1, S2
+        )
 
 
 __all__ = ["QuarkW4A4MXFP4"]

python/sglang/srt/layers/radix_attention.py

@@ -30,8 +30,11 @@ from sglang.srt.model_executor.breakable_cuda_graph.context import (
     is_in_breakable_cuda_graph,
 )
 from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.utils import is_hip
 from sglang.srt.utils.custom_op import register_custom_op
 
+_is_hip = is_hip()
+
 if TYPE_CHECKING:
     from sglang.srt.layers.quantization.base_config import QuantizationConfig
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -179,6 +182,13 @@ def unified_attention_with_output(
     if value is not None:
         value = value[:real_num_tokens]
 
+    # DeepSeek MLA has two RadixAttention instances per layer (attn_mqa and
+    # attn_mha) that share the same layer_id. The attention_layers list only
+    # stores attn_mqa. When the MHA path is active (save_kv_cache=False), use
+    # the companion attn_mha so the backend sees correct head/dim metadata.
+    if _is_hip and not save_kv_cache and hasattr(attention_layer, "_pcg_mha_companion"):
+        attention_layer = attention_layer._pcg_mha_companion
+
     kwargs = {}
     if q_rope is not None:
         kwargs["q_rope"] = q_rope[:real_num_tokens]
@@ -218,6 +228,25 @@ def unified_attention_with_output(
 
     if ret.data_ptr() != output.data_ptr():
         output[:real_num_tokens].view(ret.shape).copy_(ret)
+
+    if _is_hip:
+        # During PCG replay on AMD, varlen attention kernels only fill positions
+        # 0..actual_tokens-1 and leave padded positions with uninitialized
+        # garbage from torch.empty.  Zero these so garbage (NaN/Inf) does not
+        # propagate through residual connections, MoE routing, and allreduce.
+        # Use context.raw_num_tokens (pre-padding count from PCG runner)
+        # instead of forward_batch.extend_num_tokens, because
+        # extend_num_tokens is None for TARGET_VERIFY (EAGLE) batches.
+        pcg_static_tokens = context.num_tokens
+        actual_tokens = context.raw_num_tokens
+        if (
+            pcg_static_tokens is not None
+            and actual_tokens is not None
+            and pcg_static_tokens > actual_tokens
+        ):
+            first_dim = output.shape[0]
+            elems_per_token = output.numel() // first_dim
+            output.view(first_dim, elems_per_token)[actual_tokens:].zero_()
     return
 
 

python/sglang/srt/model_executor/model_runner.py

@@ -35,6 +35,10 @@ import torch.distributed as dist
 from torch import nn
 
 from sglang.jit_kernel.ngram_embedding import update_token_table_decode
+from sglang.srt.compilation.piecewise_context_manager import (
+    enable_piecewise_cuda_graph,
+    set_forward_context,
+)
 from sglang.srt.configs import (
     BailingHybridConfig,
     FalconH1Config,
@@ -2971,6 +2975,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 elif hasattr(layer.self_attn, "attn_mqa"):
                     # For DeepSeek model
                     attn_layer = layer.self_attn.attn_mqa
+                    if _is_hip and hasattr(layer.self_attn, "attn_mha"):
+                        attn_layer._pcg_mha_companion = layer.self_attn.attn_mha
             # For hybrid model
             elif hasattr(layer, "attn"):
                 attn_layer = layer.attn
@@ -3309,12 +3315,37 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else contextlib.nullcontext()
         )
         with ctx:
-            ret = self.model.forward(
+            if _is_hip and self.piecewise_cuda_graph_runner is not None:
-                forward_batch.input_ids,
+                # AMD/HIP: when PCG is enabled but the batch exceeds max captured
-                forward_batch.positions,
+                # size, run eagerly under enable_piecewise_cuda_graph() and
-                forward_batch,
+                # set_forward_context() so that (a) Dynamo guards on
-                **kwargs,
+                # _in_piecewise_cuda_graph stay consistent with the PCG-traced
-            )
+                # graph (preventing runtime recompilation) and (b) PCG-specific
+                # code paths (MoE, attention) can access their layer objects.
+                with (
+                    enable_piecewise_cuda_graph(),
+                    set_forward_context(
+                        forward_batch,
+                        self.attention_layers,
+                        getattr(self.model, "quant_config", None),
+                        self.moe_layers,
+                        self.moe_fusions,
+                        dsa_indexers=self.dsa_indexers,
+                    ),
+                ):
+                    ret = self.model.forward(
+                        forward_batch.input_ids,
+                        forward_batch.positions,
+                        forward_batch,
+                        **kwargs,
+                    )
+            else:
+                ret = self.model.forward(
+                    forward_batch.input_ids,
+                    forward_batch.positions,
+                    forward_batch,
+                    **kwargs,
+                )
         return (ret, can_run_graph)
 
     def forward_idle(

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -61,12 +61,16 @@ from sglang.srt.model_executor.forward_batch_info import (
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
 from sglang.srt.utils import (
     get_available_gpu_memory,
+    get_bool_env_var,
+    is_hip,
     is_musa,
     is_npu,
     log_info_on_rank0,
     require_gathered_buffer,
 )
 
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 # Suppress Dynamo warning about tracing through lru_cache-wrapped functions (e.g., is_arch_support_pdl).
 warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
 logger = logging.getLogger(__name__)
@@ -280,22 +284,33 @@ class PiecewiseCudaGraphRunner:
                     graph_pool=get_global_graph_memory_pool(),
                 )
 
-                with enable_piecewise_cuda_graph_compile():
+                if _is_hip:
-                    compile_range = (
+                    # AMD: single Dynamo trace is sufficient; the capture
-                        tqdm.tqdm(list(reversed(self.capture_num_tokens)))
+                    # phase does per-shape JIT kernel warmup before each
-                        if get_tensor_model_parallel_rank() == 0
+                    # CUDA graph recording.  The N-iteration loop is
-                        else reversed(self.capture_num_tokens)
+                    # redundant and extremely slow on ROCm (~30 min).
-                    )
+                    with enable_piecewise_cuda_graph_compile():
-                    for _, num_tokens in enumerate(compile_range):
+                        self.warmup_compile(num_tokens=self.capture_num_tokens[-1])
-                        if get_tensor_model_parallel_rank() == 0:
+                else:
-                            compile_range.set_description(
+                    with enable_piecewise_cuda_graph_compile():
-                                f"Compiling num tokens ({num_tokens=})"
+                        compile_range = (
-                            )
+                            tqdm.tqdm(list(reversed(self.capture_num_tokens)))
-                        self.warmup_compile(num_tokens=num_tokens)
+                            if get_tensor_model_parallel_rank() == 0
+                            else reversed(self.capture_num_tokens)
+                        )
+                        for _, num_tokens in enumerate(compile_range):
+                            if get_tensor_model_parallel_rank() == 0:
+                                compile_range.set_description(
+                                    f"Compiling num tokens ({num_tokens=})"
+                                )
+                            self.warmup_compile(num_tokens=num_tokens)
 
                 set_global_graph_memory_pool(self.device_module.graph_pool_handle())
                 set_graph_pool_id(get_global_graph_memory_pool())
 
+                if _use_aiter:
+                    self._pre_warm_aiter_chip_info()
+
                 self.device_module.synchronize()
                 self.model_runner.tp_group.barrier()
                 # Capture
@@ -303,6 +318,41 @@ class PiecewiseCudaGraphRunner:
 
         self.raw_num_tokens = 0
 
+    _aiter_chip_info_cached = False
+
+    @classmethod
+    def _pre_warm_aiter_chip_info(cls):
+        """Pre-populate aiter chip info env vars before CUDA graph capture.
+
+        aiter's get_cu_num_custom_op and get_gfx_custom_op call
+        subprocess.run(rocminfo) to query GPU info. During CUDA graph capture
+        the GPU context is locked, so rocminfo hangs indefinitely. Pre-calling
+        them here caches the results as environment variables so the subprocess
+        is never invoked during capture. Only runs once per process.
+        """
+        if cls._aiter_chip_info_cached:
+            return
+        cls._aiter_chip_info_cached = True
+
+        import os
+
+        try:
+            from aiter.jit.utils.chip_info import get_cu_num, get_gfx
+
+            if not os.environ.get("CU_NUM"):
+                cu_num = get_cu_num()
+                os.environ["CU_NUM"] = str(cu_num)
+                logger.info(f"Pre-warmed aiter CU_NUM={cu_num}")
+
+            if not os.environ.get("GPU_ARCHS"):
+                gfx = get_gfx()
+                os.environ["GPU_ARCHS"] = gfx
+                logger.info(f"Pre-warmed aiter GPU_ARCHS={gfx}")
+        except ImportError:
+            pass
+        except Exception as e:
+            logger.warning(f"Failed to pre-warm aiter chip info: {e}")
+
     def warmup_compile(self, num_tokens: int):
         """Warmup the model with a simple forward pass before CUDA graph capture."""
         registry = self.buffer_registry
@@ -723,6 +773,8 @@ class PiecewiseCudaGraphRunner:
     ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
         with enable_piecewise_cuda_graph():
             static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
+            static_num_tokens = len(static_forward_batch.input_ids)
+            raw_num_tokens = self.raw_num_tokens
             # Replay
             with set_forward_context(
                 static_forward_batch,
@@ -731,6 +783,8 @@ class PiecewiseCudaGraphRunner:
                 self.moe_layers,
                 self.moe_fusions,
                 dsa_indexers=self.dsa_indexers,
+                num_tokens=static_num_tokens,
+                raw_num_tokens=raw_num_tokens,
             ):
                 self.model_runner.attn_backend.init_forward_metadata(forward_batch)
                 output = self.model_runner.model.forward(

test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py

@@ -40,6 +40,11 @@ class TestDeepseekR1MXFP4(CustomTestCase):
             "131072",
             "--model-loader-extra-config",
             '{"enable_multithread_load": true}',
+            "--enforce-piecewise-cuda-graph",
+            "--piecewise-cuda-graph-compiler",
+            "eager",
+            "--piecewise-cuda-graph-max-tokens",
+            "8192",
         ]
         cls.process = popen_launch_server(
             cls.model,
@@ -73,7 +78,7 @@ class TestDeepseekR1MXFP4(CustomTestCase):
             write_github_step_summary(
                 f"### test_gsm8k (deepseek-r1-mxfp4)\n" f'{metrics["accuracy"]=:.3f}\n'
             )
-            self.assertGreater(metrics["accuracy"], 0.94)
+        self.assertGreater(metrics["accuracy"], 0.94)
 
     def test_bs_1_speed(self):
         args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
@@ -85,7 +90,7 @@ class TestDeepseekR1MXFP4(CustomTestCase):
             write_github_step_summary(
                 f"### test_bs_1_speed (deepseek-r1-mxfp4)\n" f"{speed=:.2f} token/s\n"
             )
-            self.assertGreater(speed, 75)
+        self.assertGreater(speed, 75)
 
 
 class TestDeepseekR1MXFP4MTP(CustomTestCase):

test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py

@@ -5,7 +5,7 @@ import torch
 from sglang import Engine
 from sglang.lang.chat_template import get_chat_template_by_model_path
 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_IMAGE_URL,
@@ -18,6 +18,7 @@ from sglang.test.test_utils import (
 
 # CI Registration
 register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-large")
+register_amd_ci(est_time=180, suite="stage-b-test-1-gpu-large-amd")
 
 
 class TestPiecewiseCudaGraphQwen25VL(CustomTestCase):