duanyll/ktransformers repository overview

FROM node:20.16.0 AS web_compile

COPY ktransformers /home/ktransformers

RUN cd /home/ktransformers/ktransformers/website/ && \
    npm install && \
    npm run build && \
    rm -rf node_modules

FROM alpine/curl AS downloader

# Download https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu11torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl
RUN mkdir -p /tmp/download && \
    curl -L -o /tmp/download/flash_attn-2.7.3+cu11torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl \
    https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu11torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

FROM pytorch/pytorch:2.5.1-cuda11.8-cudnn9-devel

COPY ktransformers/requirements-local_chat.txt /tmp/requirements.txt

RUN pip install --no-cache-dir -r /tmp/requirements.txt

RUN --mount=type=bind,from=downloader,source=/tmp/download,target=/tmp/download \
    pip install --no-cache-dir /tmp/download/flash_attn-2.7.3+cu11torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

RUN apt update -y && \
    apt install -y --no-install-recommends numactl && \
    rm -rf /var/lib/apt/lists/*

COPY --from=web_compile /home/ktransformers /workspace/ktransformers

WORKDIR /workspace/ktransformers

ARG KTRANSFORMERS_FORCE_BUILD=TRUE
ARG TORCH_CUDA_ARCH_LIST="8.0;8.6;8.7;8.9;9.0+PTX"
ARG CPU_INSTRUCT=NATIVE
ENV USE_NUMA=1

RUN KTRANSFORMERS_FORCE_BUILD=$KTRANSFORMERS_FORCE_BUILD TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST CPU_INSTRUCT=$CPU_INSTRUCT pip install . --no-cache-dir --no-build-isolation --verbose && \
    pip cache purge

docker run --rm -it -v /path/to/weights/:/weights -v ./prompt.txt:/prompt.txt -v ./rule.yaml:/rule.yaml --gpus='"device=0"' --cap-add SYS_NICE duanyll/ktransformers:0.2.1-cu118-avx2-numa python ./ktransformers/local_chat.py --model_path /weights/DeepSeek-R1 --gguf_path /weights/DeepSeek-R1-Q4_K_M --prompt_file /prompt.txt --cpu_infer 120 --max_new_tokens 32000