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

AITER backend extension and workload optimizations (#6838)

Browse Source 
Co-authored-by: wunhuang <wunhuang@amd.com>
Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
This commit is contained in:
HAI
2025-06-05 23:00:18 -07:00
committed by GitHub
parent 562f279a2d
commit b819381fec
12 changed files with 583 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -38,11 +38,10 @@ _is_hip = is_hip()
_is_cuda = is_cuda()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
use_aiter_moe = get_bool_env_var("SGLANG_AITER_MOE")
if _is_hip and use_aiter_moe:
from aiter import gemm_a8w8_blockscale
if _use_aiter:
from aiter import gemm_a8w8_blockscale_CK
if _is_cuda:
from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm
@@ -141,7 +140,7 @@ def dispatch_w8a8_block_fp8_linear() -> Callable:
return flashinfer_gemm_w8a8_block_fp8_linear
elif CUTLASS_BLOCK_FP8_SUPPORTED:
return cutlass_w8a8_block_fp8_linear_with_fallback
elif _is_hip and use_aiter_moe:
elif _use_aiter:
return aiter_w8a8_block_fp8_linear
elif _ENABLE_JIT_DEEPGEMM:
return deepgemm_w8a8_block_fp8_linear_with_fallback
@@ -268,12 +267,9 @@ def aiter_w8a8_block_fp8_linear(
q_input, x_scale = per_token_group_quant_fp8(
input_2d, block_size[1], column_major_scales=False
)
output = torch.zeros(
[q_input.shape[0], weight.shape[0]],
dtype=input_2d.dtype,
device=q_input.device,
output = gemm_a8w8_blockscale_CK(
q_input, weight, x_scale, weight_scale, dtype=input.dtype
)
gemm_a8w8_blockscale(q_input, weight, x_scale, weight_scale, output)
if bias is not None:
output += bias