[AMD] Restore AMD piecewise CUDA graph support dropped by #23906 (#27811)

python/sglang/srt/model_executor/model_runner.py

@@ -172,6 +172,10 @@ from sglang.srt.model_executor.runner import (
 from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
     _allocate_decode_buffers,
 )
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    enable_tc_piecewise_cuda_graph,
+    set_tc_piecewise_forward_context,
+)
 from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
 from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
     RemoteInstanceWeightLoaderBackend,
@@ -2979,6 +2983,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
@@ -3313,12 +3319,37 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else contextlib.nullcontext()
         )
         with ctx:
-            ret = self.model.forward(
+            if _is_hip and self.prefill_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_tc_piecewise_cuda_graph() and
-                forward_batch,
+                # set_tc_piecewise_forward_context() so that (a) Dynamo guards on
-                **kwargs,
+                # _in_tc_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_tc_piecewise_cuda_graph(),
+                    set_tc_piecewise_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/runner/prefill_cuda_graph_runner.py

@@ -72,6 +72,8 @@ from sglang.srt.model_executor.runner_utils.buffers import (
 )
 from sglang.srt.utils import (
     get_available_gpu_memory,
+    get_bool_env_var,
+    is_hip,
     is_npu,
     log_info_on_rank0,
     require_attn_tp_gather,
@@ -82,6 +84,9 @@ from sglang.srt.utils import (
 warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
 logger = logging.getLogger(__name__)
 
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
 # Names of the static prefill input tensors a Breakable-backed prefill
 # runner owns. Each is a 1-D int64 tensor of length max_bs; captured
 # Breakable segments read from these stable addresses.
@@ -246,6 +251,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     f"this model architecture."
                 )
 
+        # --- aiter chip info pre-warming (AMD) -------------------------
+        if _use_aiter:
+            self._pre_warm_aiter_chip_info()
+
         # --- capture --------------------------------------------------
         self.device_module.synchronize()
         self.model_runner.tp_group.barrier()
@@ -266,6 +275,41 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     def _cache_loc_dtype(self):
         return torch.int64 if not is_npu() else torch.int32
 
+    _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}")
+
     @torch.no_grad()
     def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int):
         """Run forward inside the prefill set_tc_piecewise_forward_context.
@@ -765,6 +809,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
         with self.backend.replay_session():
             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
 
             if self.layer_model is not None:
                 # BCG path. The captured graph is a bs=1 replay of
@@ -793,6 +839,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                         self.moe_layers,
                         self.moe_fusions,
                         dsa_indexers=self.dsa_indexers,
+                        num_tokens=static_num_tokens,
+                        raw_num_tokens=raw_num_tokens,
                     ):
                         output = self.model_runner.model.forward(
                             static_forward_batch.input_ids,
@@ -815,6 +863,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     self.moe_layers,
                     self.moe_fusions,
                     dsa_indexers=self.dsa_indexers,
+                    num_tokens=static_num_tokens,
+                    raw_num_tokens=raw_num_tokens,
                 ):
                     output = self.backend.replay(
                         self._static_num_tokens, static_forward_batch, **kwargs

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

@@ -33,6 +33,7 @@ from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     enable_tc_piecewise_cuda_graph,
 )
+from sglang.srt.utils import is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -153,17 +154,28 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
                 )
 
                 with enable_torch_compile_warmup():
-                    compile_range = (
+                    if is_hip():
-                        tqdm.tqdm(list(reversed(cuda_graph_runner.capture_num_tokens)))
+                        # AMD: single Dynamo trace is sufficient; the capture
-                        if get_tensor_model_parallel_rank() == 0
+                        # phase does per-shape JIT kernel warmup before each
-                        else reversed(cuda_graph_runner.capture_num_tokens)
+                        # CUDA graph recording.  The N-iteration loop is
-                    )
+                        # redundant and extremely slow on ROCm (~30 min).
-                    for num_tokens in compile_range:
+                        cuda_graph_runner._run_dummy_forward(
-                        if get_tensor_model_parallel_rank() == 0:
+                            num_tokens=cuda_graph_runner.capture_num_tokens[-1]
-                            compile_range.set_description(
+                        )
-                                f"Compiling num tokens ({num_tokens=})"
+                    else:
+                        compile_range = (
+                            tqdm.tqdm(
+                                list(reversed(cuda_graph_runner.capture_num_tokens))
                             )
-                        cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
+                            if get_tensor_model_parallel_rank() == 0
+                            else reversed(cuda_graph_runner.capture_num_tokens)
+                        )
+                        for num_tokens in compile_range:
+                            if get_tensor_model_parallel_rank() == 0:
+                                compile_range.set_description(
+                                    f"Compiling num tokens ({num_tokens=})"
+                                )
+                            cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
             finally:
                 _toggle_multi_platform_ops(
                     language_model.model, reverse=True, num_tokens=16

python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py

@@ -72,6 +72,8 @@ class TcPiecewiseForwardContext:
     moe_layers: Optional[List[Any]] = field(default=None)
     moe_fusions: Optional[List[Any]] = field(default=None)
     dsa_indexers: Optional[List[Any]] = field(default=None)
+    num_tokens: Optional[int] = None
+    raw_num_tokens: Optional[int] = None
 
 
 _tc_piecewise_forward_context: Optional[TcPiecewiseForwardContext] = None
@@ -89,6 +91,8 @@ def set_tc_piecewise_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 _tc_piecewise_forward_context
     _tc_piecewise_forward_context = TcPiecewiseForwardContext(
@@ -98,6 +102,8 @@ def set_tc_piecewise_forward_context(
         moe_layers=moe_layers,
         moe_fusions=moe_fusions,
         dsa_indexers=dsa_indexers,
+        num_tokens=num_tokens,
+        raw_num_tokens=raw_num_tokens,
     )
     try:
         yield