33 lines
2.4 KiB
Docker

FROM lmsysorg/sglang@sha256:28e0d26073161e49ca56eba808d264a4223804a212020f1dfe1b2405b9f8a399
# Use the official core and compile SM120 CUTLASS through its supported JIT path.
# Remove old companion packages instead of suppressing their version checks.
COPY results/official_flashinfer_0618/flashinfer_python-0.6.18-py3-none-any.whl /opt/kimi-dflash-wheels/
RUN python3 -m pip uninstall -y flashinfer-jit-cache flashinfer-cubin && \
python3 -m pip install --no-index --no-deps --force-reinstall \
/opt/kimi-dflash-wheels/flashinfer_python-0.6.18-py3-none-any.whl
# #4460 is merged on main but absent from the v0.6.18 release branch.
# Apply the unchanged official commit to Python and wheel-bundled C++ sources.
COPY flashinfer_4460_b460bc0.patch /opt/kimi-dflash-wheels/
RUN cd "$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" && \
git apply --check --include='flashinfer/**' /opt/kimi-dflash-wheels/flashinfer_4460_b460bc0.patch && \
git apply --check --directory=flashinfer/data --include='flashinfer/data/csrc/**' /opt/kimi-dflash-wheels/flashinfer_4460_b460bc0.patch && \
git apply --include='flashinfer/**' /opt/kimi-dflash-wheels/flashinfer_4460_b460bc0.patch && \
git apply --directory=flashinfer/data --include='flashinfer/data/csrc/**' /opt/kimi-dflash-wheels/flashinfer_4460_b460bc0.patch
COPY results/sglang-kimi-pp-dflash-33863-integrated.tar.gz /opt/kimi-dflash-source.tar.gz
RUN mkdir -p /opt/kimi-dflash && \
tar -xzf /opt/kimi-dflash-source.tar.gz -C /opt/kimi-dflash
ENV PYTHONPATH=/opt/kimi-dflash/python
ENV PYTHONDONTWRITEBYTECODE=1
ENV FLASHINFER_DISABLE_VERSION_CHECK=""
ENV SGLANG_SOURCE_ROOT=/opt/kimi-dflash
# Build has no GPU: check the API, not SGLang's CUDA-device availability predicate.
RUN python3 -c "import inspect, flashinfer; from flashinfer.fused_moe import cutlass_fused_moe; from flashinfer.tllm_enums import ActivationType; from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2; from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager; assert hasattr(ActivationType, 'Situ'); assert {'situ_beta', 'situ_linear_beta'} <= set(inspect.signature(cutlass_fused_moe).parameters); print('RUNTIME_IMPORT_OK', flashinfer.__version__, 'upstream SiTU backport b460bc0')"
RUN python3 -m unittest discover -s /opt/kimi-dflash/test/registered/disaggregation -v
ENTRYPOINT ["python3", "-m", "sglang.launch_server"]