diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index b85960b94..300501c6f 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -17,8 +17,8 @@ ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG TRITON_COMMIT="improve_fa_decode_3.0.0" -ARG ATER_REPO="https://github.com/HaiShaw/ater" -ARG CK_COMMITS="fa05ae" +ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG AITER_COMMIT="dev/testx" RUN git clone ${SGL_REPO} \ && cd sglang \ @@ -53,8 +53,9 @@ RUN git clone ${TRITON_REPO} \ && cd python \ && python3 setup.py install -RUN git clone ${ATER_REPO} \ - && cd ater \ +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ && git submodule update --init --recursive \ && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index e83a32767..4a944fb85 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -167,8 +167,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): ) if is_hip_ and get_bool_env_var("CK_MOE"): - import ater - from ater.fused_moe import fused_experts_ck + import aiter + from aiter.fused_moe import fused_experts_ck assert activation == "silu", f"{activation=} is not supported." diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index f5a0005a2..9a038f384 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -789,8 +789,8 @@ class Fp8MoEMethod: ) if is_hip_ and get_bool_env_var("CK_MOE"): - import ater - from ater.fused_moe import fused_experts_ck + import aiter + from aiter.fused_moe import fused_experts_ck assert activation == "silu", f"{activation=} is not supported."