ARG BASE_IMAGE=lmsysorg/sglang:kimi-k3-fiv617situ-warm
FROM ${BASE_IMAGE}

ENV FLASHINFER_DISABLE_JIT=1 \
    FLASHINFER_DISABLE_VERSION_CHECK=1 \
    PYTHONUNBUFFERED=1

RUN rm -rf \
      /usr/local/lib/python3.12/dist-packages/flashinfer \
      /usr/local/lib/python3.12/dist-packages/flashinfer_python-*.dist-info

COPY flashinfer_python-0.6.18-py3-none-any.whl /tmp/flashinfer_python-0.6.18-py3-none-any.whl
RUN python3 -m pip install --no-deps --force-reinstall \
      /tmp/flashinfer_python-0.6.18-py3-none-any.whl && \
    rm -f /tmp/flashinfer_python-0.6.18-py3-none-any.whl

# The wheel contains Python/C++ sources. These three libraries are the exact
# SM120 AOT artifacts that passed the Phase 3/4 correctness gate.
COPY aot/ /usr/local/lib/python3.12/dist-packages/flashinfer/data/aot/

# Use the complete tested SGLang Python package, not a pair of loose runtime
# mounts. This preserves API compatibility between the two changed modules and
# their surrounding Kimi/MoE integration.
COPY sglang/ /sgl-workspace/sglang/python/sglang/

# Existing Kimi image compatibility fix: tcgen05 attention-residual is valid on
# SM100/SM103, not consumer Blackwell SM120. Bake the fallback into the image.
RUN python3 -c "from pathlib import Path; p=Path('/sgl-workspace/sglang/python/sglang/srt/layers/attn_residual.py'); s=p.read_text(); old='        _FAST_SUPPORTED = major >= 10'; new='        _FAST_SUPPORTED = major in (10, 11)'; assert old in s; p.write_text(s.replace(old,new))" && \
    python3 -m compileall -q /sgl-workspace/sglang/python/sglang && \
    python3 -c "from flashinfer.fused_moe import core; assert core.CUTLASS_FUSED_MOE_SUPPORTS_SITU is True"

LABEL ai.meta-stone.purpose="Kimi-K3 SM120 FlashInfer MXFP4 Phase 5 acceptance" \
      ai.meta-stone.sglang.patch="fb929bb" \
      ai.meta-stone.flashinfer.patch="7d1a4ee"
