[Fix] Keep ABI-matched SGLang kernel in PR image

This commit is contained in:
Zhiyi Hong 2026-08-19 11:00:50 +08:00
parent e8ff3ce1e8
commit 7f67dfe6b3
2 changed files with 4 additions and 8 deletions

View File

@ -8,15 +8,14 @@ ENV FLASHINFER_DISABLE_VERSION_CHECK=1 \
PYTHONPATH=/sgl-workspace/sglang/python \ PYTHONPATH=/sgl-workspace/sglang/python \
PYTHONUNBUFFERED=1 PYTHONUNBUFFERED=1
COPY sglang_kernel-0.4.6.post1-cp310-abi3-manylinux2014_x86_64.whl /tmp/
COPY flashinfer_python-0.6.18-py3-none-any.whl /tmp/ COPY flashinfer_python-0.6.18-py3-none-any.whl /tmp/
RUN python3 -m pip install --no-deps --force-reinstall \ RUN python3 -m pip install --no-deps --force-reinstall \
/tmp/sglang_kernel-0.4.6.post1-cp310-abi3-manylinux2014_x86_64.whl \
/tmp/flashinfer_python-0.6.18-py3-none-any.whl && \ /tmp/flashinfer_python-0.6.18-py3-none-any.whl && \
rm -f /tmp/*.whl rm -f /tmp/*.whl
# Replace the image's older Python package with the exact Draft tree. The # Replace the image's older Python package with the exact Draft tree. Keep the
# matching 0.4.6.post1 sglang-kernel wheel is installed above. # base image's ABI-matched torch and sglang-kernel binaries: the Draft changes
# only the Python Kimi/FlashInfer integration and does not require a new kernel.
RUN rm -rf /sgl-workspace/sglang/python/sglang RUN rm -rf /sgl-workspace/sglang/python/sglang
COPY sglang/ /sgl-workspace/sglang/python/sglang/ COPY sglang/ /sgl-workspace/sglang/python/sglang/
COPY test_mxfp4_sm120_cutlass.py /opt/pr-tests/test_mxfp4_sm120_cutlass.py COPY test_mxfp4_sm120_cutlass.py /opt/pr-tests/test_mxfp4_sm120_cutlass.py

View File

@ -19,7 +19,6 @@ SGLANG_SOURCE="${SGLANG_SOURCE:-/data/hzy/src/sglang-kimi-sm120-draft-wt}"
FLASHINFER_SOURCE="${FLASHINFER_SOURCE:-/data/hzy/src/flashinfer-pr4460-kimi-test-wt}" FLASHINFER_SOURCE="${FLASHINFER_SOURCE:-/data/hzy/src/flashinfer-pr4460-kimi-test-wt}"
ARTIFACT_DIR="${ARTIFACT_DIR:-/data/hzy/artifacts/sglang-pr-kimi-sm120-300c87a}" ARTIFACT_DIR="${ARTIFACT_DIR:-/data/hzy/artifacts/sglang-pr-kimi-sm120-300c87a}"
FLASHINFER_WHEEL="${FLASHINFER_WHEEL:-${ARTIFACT_DIR}/flashinfer_python-0.6.18-py3-none-any.whl}" FLASHINFER_WHEEL="${FLASHINFER_WHEEL:-${ARTIFACT_DIR}/flashinfer_python-0.6.18-py3-none-any.whl}"
SGLANG_KERNEL_WHEEL="${SGLANG_KERNEL_WHEEL:-${ARTIFACT_DIR}/sglang_kernel-0.4.6.post1-cp310-abi3-manylinux2014_x86_64.whl}"
BUILD_CONTEXT="${BUILD_CONTEXT:-/tmp/kimi3_sm120_pr_validation_context}" BUILD_CONTEXT="${BUILD_CONTEXT:-/tmp/kimi3_sm120_pr_validation_context}"
FLASHINFER_CACHE="${FLASHINFER_CACHE:-/data/hzy/cache/flashinfer-pr-b525c51}" FLASHINFER_CACHE="${FLASHINFER_CACHE:-/data/hzy/cache/flashinfer-pr-b525c51}"
@ -79,7 +78,6 @@ check_inputs() {
local path local path
for path in "${MODEL_PATH}" "${SGLANG_SOURCE}/python/sglang" \ for path in "${MODEL_PATH}" "${SGLANG_SOURCE}/python/sglang" \
"${FLASHINFER_SOURCE}" "${FLASHINFER_WHEEL}" \ "${FLASHINFER_SOURCE}" "${FLASHINFER_WHEEL}" \
"${SGLANG_KERNEL_WHEEL}" \
"${SCRIPT_DIR}/Dockerfile.pr_validation"; do "${SCRIPT_DIR}/Dockerfile.pr_validation"; do
[[ -e "$path" ]] || { echo "ERROR: missing ${path}" >&2; exit 2; } [[ -e "$path" ]] || { echo "ERROR: missing ${path}" >&2; exit 2; }
done done
@ -103,7 +101,6 @@ prepare_build_context() {
mkdir -p "${BUILD_CONTEXT}/sglang" mkdir -p "${BUILD_CONTEXT}/sglang"
cp "${SCRIPT_DIR}/Dockerfile.pr_validation" "${BUILD_CONTEXT}/Dockerfile" cp "${SCRIPT_DIR}/Dockerfile.pr_validation" "${BUILD_CONTEXT}/Dockerfile"
cp "${FLASHINFER_WHEEL}" "${BUILD_CONTEXT}/" cp "${FLASHINFER_WHEEL}" "${BUILD_CONTEXT}/"
cp "${SGLANG_KERNEL_WHEEL}" "${BUILD_CONTEXT}/"
cp "${SGLANG_SOURCE}/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py" \ cp "${SGLANG_SOURCE}/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py" \
"${BUILD_CONTEXT}/test_mxfp4_sm120_cutlass.py" "${BUILD_CONTEXT}/test_mxfp4_sm120_cutlass.py"
tar -C "${SGLANG_SOURCE}/python/sglang" -cf - . | \ tar -C "${SGLANG_SOURCE}/python/sglang" -cf - . | \
@ -111,7 +108,7 @@ prepare_build_context() {
{ {
printf 'sglang=%s\n' "${SGLANG_COMMIT}" printf 'sglang=%s\n' "${SGLANG_COMMIT}"
printf 'flashinfer_implementation=%s\n' "${FLASHINFER_COMMIT}" printf 'flashinfer_implementation=%s\n' "${FLASHINFER_COMMIT}"
sha256sum "${FLASHINFER_WHEEL}" "${SGLANG_KERNEL_WHEEL}" sha256sum "${FLASHINFER_WHEEL}"
} >"${BUILD_CONTEXT}/source_identity.txt" } >"${BUILD_CONTEXT}/source_identity.txt"
cp "${BUILD_CONTEXT}/source_identity.txt" "${RESULT_ROOT}/build/" cp "${BUILD_CONTEXT}/source_identity.txt" "${RESULT_ROOT}/build/"
du -sh "${BUILD_CONTEXT}" | tee "${RESULT_ROOT}/build/context_size.txt" du -sh "${BUILD_CONTEXT}" | tee "${RESULT_ROOT}/build/context_size.txt"