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

[AMD] rocm 7.2 image release, PR test, Nightly Test (#17799)

Browse Source 
Co-authored-by: Alan Kao <akao@amd.com>
Co-authored-by: bingxche <Bingxu.Chen@amd.com>
Co-authored-by: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
This commit is contained in:
YC Tseng
2026-02-12 13:29:25 +08:00
committed by GitHub
parent 93ede0db19
commit 20554a0a4f
25 changed files with 2716 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/layers/layernorm.py
View File

@@ -65,11 +65,20 @@ if _is_cuda or _is_xpu:
gemma_rmsnorm,
rmsnorm,
)
_has_vllm_rms_norm = False
if _use_aiter:
from aiter import rmsnorm2d_fwd as rms_norm
from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm
_has_vllm_rms_norm = True # aiter provides the rms_norm functions
elif _is_hip:
from vllm._custom_ops import fused_add_rms_norm, rms_norm
try:
from vllm._custom_ops import fused_add_rms_norm, rms_norm
_has_vllm_rms_norm = True
except ImportError:
# Fallback: vllm not available, will use forward_native
_has_vllm_rms_norm = False
logger = logging.getLogger(__name__)
@@ -182,6 +191,10 @@ class RMSNorm(MultiPlatformOp):
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
# Fallback to native implementation if vllm is not available
if not _has_vllm_rms_norm:
return self.forward_native(x, residual, post_residual_addition)
if not x.is_contiguous():
# NOTE: Remove this if aiter kernel supports discontinuous input
x = x.contiguous()

55
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

@@ -57,11 +57,22 @@ elif _is_hip:
from aiter import moe_sum
except ImportError:
raise ImportError("aiter is required when SGLANG_USE_AITER is set to True")
else:
from vllm import _custom_ops as vllm_ops
# Note: vllm_ops is not needed for HIP when _use_aiter=False
# because the code uses moe_sum_reduce_triton as fallback (line 619)
elif _is_xpu:
from sgl_kernel import moe_sum_reduce, silu_and_mul
# Try to import vllm_ops for non-CUDA/HIP/XPU platforms
_has_vllm_ops = False
if not _is_cuda and not _is_hip and not _is_xpu:
try:
from vllm import _custom_ops as vllm_ops
_has_vllm_ops = True
except ImportError:
# Fallback: vllm not available, will use native PyTorch implementations
_has_vllm_ops = False
padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
@@ -513,9 +524,15 @@ def fused_experts_impl(
activation,
)
else:
vllm_ops.silu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
)
if _has_vllm_ops:
vllm_ops.silu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
)
else:
# Fallback: native PyTorch silu_and_mul
x = intermediate_cache1.view(-1, N)
d = x.shape[-1] // 2
intermediate_cache2.copy_(F.silu(x[..., :d]) * x[..., d:])
elif activation == "gelu" and is_gated:
assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
assert gemm1_limit is None, "gemm1_limit is not supported for gelu"
@@ -533,9 +550,15 @@ def fused_experts_impl(
activation,
)
else:
vllm_ops.gelu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
)
if _has_vllm_ops:
vllm_ops.gelu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
)
else:
# Fallback: native PyTorch gelu_and_mul
x = intermediate_cache1.view(-1, N)
d = x.shape[-1] // 2
intermediate_cache2.copy_(F.gelu(x[..., :d]) * x[..., d:])
# Activation function without multiplication
elif activation == "silu" and not is_gated:
intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))
@@ -634,10 +657,18 @@ def fused_experts_impl(
routed_scaling_factor,
)
else:
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states[begin_chunk_idx:end_chunk_idx],
)
if _has_vllm_ops:
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states[begin_chunk_idx:end_chunk_idx],
)
else:
# Fallback: use triton moe_sum_reduce when vllm is not available
moe_sum_reduce_triton(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states[begin_chunk_idx:end_chunk_idx],
routed_scaling_factor,
)
return out_hidden_states

18
python/sglang/srt/layers/moe/moe_runner/triton.py
View File

@@ -41,6 +41,7 @@ if _is_cuda or _is_hip:
from sgl_kernel import gelu_and_mul, silu_and_mul
if _is_hip:
_has_vllm = False
if _use_aiter:
try:
from aiter import moe_sum
@@ -49,7 +50,13 @@ if _is_cuda or _is_hip:
"aiter is required when SGLANG_USE_AITER is set to True"
)
else:
from vllm import _custom_ops as vllm_ops # moe_sum
try:
from vllm import _custom_ops as vllm_ops # moe_sum
_has_vllm = True
except ImportError:
# Fallback: vllm not available, will use triton moe_sum
_has_vllm = False
elif _is_cpu and _is_cpu_amx_available:
pass
elif _is_xpu:
@@ -314,11 +321,18 @@ class TritonRunnerCore(MoeRunnerCore):
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
)
else:
elif _has_vllm:
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
)
else:
# Fallback: use triton moe_sum when vllm is not available
moe_sum_reduce_triton(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
elif _is_xpu:
moe_sum_reduce(
intermediate_cache3.view(*intermediate_cache3.shape),

49
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

@@ -64,6 +64,7 @@ if _is_cuda:
enable_sgl_per_token_group_quant_8bit = False
if _is_hip:
_has_vllm = False
if _use_aiter:
try:
from aiter import ( # v0.1.3
@@ -76,8 +77,11 @@ if _is_hip:
else:
try:
import vllm._C # noqa: F401
_has_vllm = True
except ImportError:
raise ImportError("vllm is required when SGLANG_USE_AITER is set to False")
# Fallback: vllm not available, will use native PyTorch implementation
_has_vllm = False
logger = logging.getLogger(__name__)
@@ -1537,6 +1541,37 @@ Raises:
"""
if _is_hip:
def _native_dynamic_per_token_quant_fp8(output, input, scale):
"""Native PyTorch fallback for dynamic per-token FP8 quantization when vLLM is unavailable."""
M, N = input.shape
eps = 1e-12
# Compute per-token scale
absmax = input.abs().max(dim=1, keepdim=True).values
absmax = torch.clamp(absmax, min=eps)
scale_val = absmax / fp8_max
scale.copy_(scale_val)
# Quantize
output_data = torch.clamp(input / scale_val, fp8_min, fp8_max).to(fp8_dtype)
output.copy_(output_data)
def _native_dynamic_per_tensor_quant_fp8(output, input, scale):
"""Native PyTorch fallback for dynamic per-tensor FP8 quantization when vLLM is unavailable."""
eps = 1e-12
absmax = input.abs().max()
absmax = torch.clamp(absmax, min=eps)
scale_val = absmax / fp8_max
# Use copy_ instead of fill_ with .item() to avoid CPU-GPU sync
scale.view(-1).copy_(scale_val.view(-1))
# Quantize
output_data = torch.clamp(input / scale_val, fp8_min, fp8_max).to(fp8_dtype)
output.copy_(output_data)
def _native_static_quant_fp8(output, input, scale):
"""Native PyTorch fallback for static FP8 quantization when vLLM is unavailable."""
# Use tensor directly instead of .item() to avoid CPU-GPU sync
output_data = torch.clamp(input / scale, fp8_min, fp8_max).to(fp8_dtype)
output.copy_(output_data)
def scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
@@ -1557,16 +1592,20 @@ if _is_hip:
)
if _use_aiter:
dynamic_per_token_scaled_quant(output, input, scale)
else:
elif _has_vllm:
torch.ops._C.dynamic_per_token_scaled_fp8_quant(
output, input.contiguous(), scale, None
)
else:
_native_dynamic_per_token_quant_fp8(output, input, scale)
else:
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
if _use_aiter:
dynamic_per_tensor_quant(output, input, scale)
else:
elif _has_vllm:
torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
else:
_native_dynamic_per_tensor_quant_fp8(output, input, scale)
else:
# Static scaling
assert (
@@ -1574,8 +1613,10 @@ if _is_hip:
), f"Expected scalar scale, got numel={scale.numel()}"
if _use_aiter:
static_per_tensor_quant(output, input, scale)
else:
elif _has_vllm:
torch.ops._C.static_scaled_fp8_quant(output, input, scale)
else:
_native_static_quant_fp8(output, input, scale)
return output, scale

10
python/sglang/srt/layers/quantization/unquant.py
View File

@@ -224,7 +224,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
set_weight_attrs(w2_weight_bias, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if _use_aiter:
# Skip aiter weight shuffle when using non-auto MoE backend (e.g., triton, triton_kernels)
# because aiter CK kernels don't support all GEMM dimensions
_should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto()
if _should_use_aiter_moe:
layer.w13_weight = torch.nn.Parameter(
shuffle_weight(layer.w13_weight.data, (16, 16)),
requires_grad=False,
@@ -383,7 +386,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
)[0]
return StandardCombineInput(hidden_states=output)
else:
if _use_aiter:
# Skip aiter fused_moe when using non-auto MoE backend (e.g., triton, triton_kernels)
# because aiter CK kernels don't support all GEMM dimensions
_should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto()
if _should_use_aiter_moe:
assert not moe_runner_config.no_combine, "unsupported"
topk_weights, topk_ids, _ = topk_output
if moe_runner_config.apply_router_weight_on_input:

7
python/sglang/srt/server_args.py
View File

@@ -1375,6 +1375,13 @@ class ServerArgs:
logger.warning(
"Detected ROCm and MXFP4 quantization format for GPT-OSS model, enabling aiter MXFP4 MOE kernel."
)
elif is_hip() and get_bool_env_var("SGLANG_USE_AITER"):
# For GPT-OSS bf16 on ROCm with aiter, use triton backend
# because aiter CK kernel doesn't support all GEMM dimensions
self.moe_runner_backend = "triton"
logger.warning(
"Detected ROCm with SGLANG_USE_AITER for GPT-OSS bf16 model, using triton MOE kernel."
)
elif self.ep_size == 1 and is_triton_kernels_available():
self.moe_runner_backend = "triton_kernel"
logger.warning(

3
python/sglang/test/gpt_oss_common.py
View File

@@ -41,7 +41,8 @@ class BaseTestGptOss(CustomTestCase):
if model_variant == "20b":
other_args += ["--cuda-graph-max-bs", "600"]
if _is_hip:
# Respect SGLANG_USE_AITER if already set, otherwise default to "0" for HIP
if _is_hip and "SGLANG_USE_AITER" not in os.environ:
os.environ["SGLANG_USE_AITER"] = "0"
self._run_test_raw(
model=model,

6
python/sglang/test/nightly_utils.py
View File

@@ -228,6 +228,7 @@ class NightlyBenchmarkRunner:
variant: str = "",
extra_bench_args: Optional[List[str]] = None,
enable_profile: bool = True,
timeout: Optional[int] = None,
) -> Tuple[List[BenchmarkResult], bool, Optional[float]]:
"""Run a complete benchmark for a single model with server management.
@@ -247,6 +248,7 @@ class NightlyBenchmarkRunner:
variant: Optional variant suffix (e.g., "basic", "mtp")
extra_bench_args: Extra arguments for the benchmark command
enable_profile: Whether to enable profiling (default True for NVIDIA)
timeout: Optional timeout for server launch (defaults to DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH)
Returns:
Tuple of (list of BenchmarkResult objects, success_bool, avg_spec_accept_length or None)
@@ -260,7 +262,9 @@ class NightlyBenchmarkRunner:
model=model_path,
base_url=self.base_url,
other_args=other_args or [],
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
timeout=(
timeout if timeout is not None else DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
),
)
try: