# Usage (to build SGLang ROCm docker image):
#   docker build --build-arg SGL_BRANCH=v0.4.1.post4 -t v0.4.1.post4-rocm620 -f Dockerfile.rocm .

# default base image
ARG BASE_IMAGE="rocmshared/vllm-rocm:20241031-tuned"

FROM $BASE_IMAGE AS base
USER root

WORKDIR /sgl-workspace
ARG BUILD_TYPE=all
ARG SGL_REPO="https://github.com/sgl-project/sglang"
ENV SGL_DEFAULT="main"
ARG SGL_BRANCH=${SGL_DEFAULT}

ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
ARG TRITON_COMMIT="845d75a"

RUN git clone ${SGL_REPO} \
    && cd sglang \
    && if [ "${SGL_BRANCH}" = ${SGL_DEFAULT} ]; then \
         echo "Using ${SGL_DEFAULT}, default branch."; \
       else \
         echo "Using ${SGL_BRANCH} branch."; \
         git checkout ${SGL_BRANCH}; \
       fi \
    && if [ "$BUILD_TYPE" = "srt" ]; then \
         python -m pip --no-cache-dir install -e "python[srt_hip]"; \
       else \
         python -m pip --no-cache-dir install -e "python[all_hip]"; \
       fi

RUN cp -r /sgl-workspace/sglang /sglang
RUN python -m pip cache purge

RUN pip install IPython \
    && pip install orjson \
    && pip install python-multipart \
    && pip install torchao \
    && pip install pybind11

RUN pip uninstall -y triton
RUN git clone ${TRITON_REPO} \
    && cd triton \
    && git checkout ${TRITON_COMMIT} \
    && cd python \
    && python3 setup.py install

# Performance environment variable.

ENV HIP_FORCE_DEV_KERNARG=1
ENV SGLANG_SET_CPU_AFFINITY=1
ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1
ENV NCCL_MIN_NCHANNELS=112

ENV MOE_PADDING=1
ENV VLLM_FP8_PADDING=1
ENV VLLM_FP8_ACT_PADDING=1
ENV VLLM_FP8_WEIGHT_PADDING=1
ENV VLLM_FP8_REDUCE_CONV=1
ENV TORCHINDUCTOR_MAX_AUTOTUNE=1
ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1

CMD ["/bin/bash"]