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

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

Browse Source
This commit is contained in:
Hubert Lu
2026-06-07 16:28:20 -07:00
committed by GitHub
parent 02be2e7189
commit 10d33bd77e
10 changed files with 335 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/compilation/cuda_piecewise_backend.py
+21 -4
View File
@@ -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
+6
View File
@@ -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
+11 -2
View File
@@ -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
+22
View File
@@ -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
+134 -5
View File
@@ -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_pre_quant_atomic import gemm_afp4wfp4_pre_quant
from aiter.ops.triton.quant import dynamic_mxfp4_quant
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 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
+29
View File
@@ -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
+37 -6
View File
@@ -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(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
if _is_hip and self.piecewise_cuda_graph_runner is not None:
# AMD/HIP: when PCG is enabled but the batch exceeds max captured
# size, run eagerly under enable_piecewise_cuda_graph() and
# set_forward_context() so that (a) Dynamo guards on
# _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
+66 -12
View File
@@ -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():
compile_range = (
tqdm.tqdm(list(reversed(self.capture_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)
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).
with enable_piecewise_cuda_graph_compile():
self.warmup_compile(num_tokens=self.capture_num_tokens[-1])
else:
with enable_piecewise_cuda_graph_compile():
compile_range = (
tqdm.tqdm(list(reversed(self.capture_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(