From 10d33bd77ea8168cee4678fe43c976f88100d4ab Mon Sep 17 00:00:00 2001 From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Date: Sun, 7 Jun 2026 16:28:20 -0700 Subject: [PATCH] [AMD] Enable Piecewise CUDA Graph for AMD GPUs (#22299) --- .../srt/compilation/cuda_piecewise_backend.py | 25 +++- .../compilation/piecewise_context_manager.py | 6 + python/sglang/srt/layers/moe/hash_topk.py | 13 +- python/sglang/srt/layers/moe/topk.py | 22 +++ .../quark/schemes/quark_w4a4_mxfp4.py | 139 +++++++++++++++++- python/sglang/srt/layers/radix_attention.py | 29 ++++ .../sglang/srt/model_executor/model_runner.py | 43 +++++- .../piecewise_cuda_graph_runner.py | 78 ++++++++-- .../amd/test_deepseek_r1_mxfp4_8gpu.py | 9 +- ...test_piecewise_cuda_graph_support_1_gpu.py | 3 +- 10 files changed, 335 insertions(+), 32 deletions(-) diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index 984dac620..100aa4762 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/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) diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index bbe0d040d..620fd29e6 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/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: diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index 1e13881dc..1ca160cd2 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 120f42e95..b104e40d7 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index 6bcf21477..65c2f4a9a 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/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_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"] diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index f3697a090..d5f770fc1 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c46405511..42b428830 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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( - 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( diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index ec2f5b2c2..c567a6d6f 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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(): - 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( diff --git a/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py b/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py index f8d2490da..e426ec694 100644 --- a/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py +++ b/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): diff --git a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py b/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py index 11f4fbf1d..e5707e14a 100644 --- a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py +++ b/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):