diff --git a/python/sglang/srt/layers/quantization/int8_kernel.py b/python/sglang/srt/layers/quantization/int8_kernel.py index 9a122cad5..948c15753 100644 --- a/python/sglang/srt/layers/quantization/int8_kernel.py +++ b/python/sglang/srt/layers/quantization/int8_kernel.py @@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Tuple import torch import triton import triton.language as tl +from triton.language.extra import libdevice -from sglang.srt.utils import get_device_name, is_cuda +from sglang.srt.utils import get_device_name, is_cuda, is_hip _is_cuda = is_cuda() +_is_hip = is_hip() if _is_cuda: # Temporary try: @@ -36,6 +38,7 @@ def _per_token_quant_int8( N, CAL_SUM: tl.constexpr, BLOCK: tl.constexpr, + IS_HIP: tl.constexpr, ): # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 row_id = tl.program_id(0) @@ -47,7 +50,12 @@ def _per_token_quant_int8( absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) scale_x = absmax / 127 x_q = x * (127 / absmax) - x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8) + if IS_HIP: + # ROCm Triton dropped the CUDA `tl.extra.cuda.libdevice.*` shim + # (`__nv_roundf`); use the backend-agnostic libdevice instead. + x_q = libdevice.round(x_q).to(tl.int8) + else: + x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8) if CAL_SUM: x_sum = tl.sum(x, axis=0) tl.store(x_sum_ptr + row_id, x_sum.to(x_sum_ptr.dtype.element_ty)) @@ -80,6 +88,7 @@ def per_token_quant_int8(x, scale_dtype=torch.float32, cal_sum=False): N=N, CAL_SUM=cal_sum, BLOCK=BLOCK, + IS_HIP=_is_hip, num_warps=num_warps, num_stages=1, ) diff --git a/test/registered/quant/test_int8_kernel.py b/test/registered/quant/test_int8_kernel.py index 02cb34af1..3b1c06e1c 100644 --- a/test/registered/quant/test_int8_kernel.py +++ b/test/registered/quant/test_int8_kernel.py @@ -8,10 +8,11 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe from sglang.srt.layers.moe.topk import TopKConfig, select_experts from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler -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.test_utils import CustomTestCase register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True) def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):