diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 38059cfe6..db0327b6a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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( - forward_batch.input_ids, - forward_batch.positions, - forward_batch, - **kwargs, - ) + if _is_hip and self.prefill_cuda_graph_runner is not None: + # AMD/HIP: when PCG is enabled but the batch exceeds max captured + # size, run eagerly under enable_tc_piecewise_cuda_graph() and + # set_tc_piecewise_forward_context() so that (a) Dynamo guards on + # _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( diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index e6d4d7276..e58fc2f27 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py index d543bbb82..d7e427757 100644 --- a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py +++ b/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 = ( - tqdm.tqdm(list(reversed(cuda_graph_runner.capture_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=})" + if is_hip(): + # AMD: single Dynamo trace is sufficient; the capture + # phase does per-shape JIT kernel warmup before each + # CUDA graph recording. The N-iteration loop is + # redundant and extremely slow on ROCm (~30 min). + cuda_graph_runner._run_dummy_forward( + num_tokens=cuda_graph_runner.capture_num_tokens[-1] + ) + 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 diff --git a/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py b/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py index 59e699d8c..34a3ad5b5 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py +++ b/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