[AMD] Fix int8 per-token quant Triton portability + register test for AMD nightly CI (#29694)

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,
     )

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):