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

unify is_cuda and is_hip (#4321)

Browse Source
This commit is contained in:
Yineng Zhang
2025-03-11 18:12:56 -07:00
committed by GitHub
parent 1cf63485c1
commit d1da58e275
18 changed files with 104 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -22,12 +22,12 @@ import torch
import triton
import triton.language as tl
from sglang.srt.utils import get_device_core_count, get_device_name, is_hip
from sglang.srt.utils import get_device_core_count, get_device_name, is_cuda, is_hip
is_hip_ = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if is_hip_ else torch.float8_e4m3fn
_is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
_is_cuda = torch.cuda.is_available() and torch.version.cuda
_is_cuda = is_cuda()
if _is_cuda:
import deep_gemm
from sgl_kernel import sgl_per_token_group_quant_fp8, sgl_per_token_quant_fp8
@@ -157,7 +157,7 @@ def per_token_group_quant_fp8(
finfo = torch.finfo(dtype)
fp8_max = finfo.max
if is_hip_:
if _is_hip:
fp8_max = 224.0
fp8_min = -fp8_max
@@ -332,7 +332,7 @@ def static_quant_fp8(
finfo = torch.finfo(dtype)
fp8_max = finfo.max
if is_hip_:
if _is_hip:
fp8_max = 224.0
fp8_min = -fp8_max
@@ -732,7 +732,7 @@ def w8a8_block_fp8_matmul(
else:
kernel = (
_w8a8_block_fp8_matmul_unrolledx4
if (is_hip_ == True and num_workgroups <= get_device_core_count())
if (_is_hip == True and num_workgroups <= get_device_core_count())
else _w8a8_block_fp8_matmul
)