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

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

Browse Source
This commit is contained in:
Michael
2026-07-01 12:56:33 -07:00
committed by GitHub
parent c312cdd3a7
commit 8361561c24
2 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/int8_kernel.py
+10 -1
View File
@@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Tuple
import torch import torch
import triton import triton
import triton.language as tl 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_cuda = is_cuda()
_is_hip = is_hip()
if _is_cuda: if _is_cuda:
# Temporary # Temporary
try: try:
@@ -36,6 +38,7 @@ def _per_token_quant_int8(
N, N,
CAL_SUM: tl.constexpr, CAL_SUM: tl.constexpr,
BLOCK: 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 # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
row_id = tl.program_id(0) row_id = tl.program_id(0)
@@ -47,6 +50,11 @@ def _per_token_quant_int8(
absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
scale_x = absmax / 127 scale_x = absmax / 127
x_q = x * (127 / absmax) x_q = x * (127 / absmax)
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) x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8)
if CAL_SUM: if CAL_SUM:
x_sum = tl.sum(x, axis=0) x_sum = tl.sum(x, axis=0)
@@ -80,6 +88,7 @@ def per_token_quant_int8(x, scale_dtype=torch.float32, cal_sum=False):
N=N, N=N,
CAL_SUM=cal_sum, CAL_SUM=cal_sum,
BLOCK=BLOCK, BLOCK=BLOCK,
IS_HIP=_is_hip,
num_warps=num_warps, num_warps=num_warps,
num_stages=1, num_stages=1,
) )
test/registered/quant/test_int8_kernel.py
+2 -1
View File
@@ -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.moe.topk import TopKConfig, select_experts
from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 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.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 from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small") 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): def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):