minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fxmarty-amd
2026-06-10 17:43:48 -07:00
committed by GitHub
parent 5f913c1135
commit 99ab90c5b7
4 changed files with 115 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+37 -6
View File
@@ -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(
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+50
View File
@@ -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
+22 -10
View File
@@ -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
python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py
+6
View File
@@ -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