ROCm: update AITER (#5816)

This commit is contained in:
HAI
2025-04-28 11:01:20 -07:00
committed by GitHub
parent 849c83a0c0
commit d364b9b0f2
7 changed files with 48 additions and 52 deletions

View File

@@ -31,7 +31,7 @@ from sglang.srt.utils import (
_is_hip = is_hip()
_is_cuda = is_cuda()
if _is_hip and get_bool_env_var("CK_MOE"):
if _is_hip and get_bool_env_var("SGLANG_AITER_MOE"):
from aiter import gemm_a8w8_blockscale
if _is_cuda:
@@ -132,7 +132,7 @@ def apply_w8a8_block_fp8_linear(
output = fp8_blockwise_scaled_mm(
q_input, weight.T, x_scale, weight_scale.T, out_dtype=input.dtype
)
elif _is_hip and get_bool_env_var("CK_MOE"):
elif _is_hip and get_bool_env_var("SGLANG_AITER_MOE"):
q_input, x_scale = per_token_group_quant_fp8(
input_2d, block_size[1], column_major_scales=False
)