[Kimi] Support kimi-k3 (#32541)

Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: Mick <mickjagger19@icloud.com> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Co-authored-by: Ke Bao <ispobaoke@gmail.com> Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com> Co-authored-by: Chunan Zeng <zcnrex@gmail.com> Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Ziyi Xu <ziyi.xu@radixark.ai> Co-authored-by: Zijie Xia <37504505+zijiexia@users.noreply.github.com> Co-authored-by: Yuwei An <ayw.sirius19@gmail.com> Co-authored-by: zhangxiaohao <1024393531@qq.com> Co-authored-by: Yangmin Li <yangminl@nvidia.com> Co-authored-by: Julien Lin <jullin@nvidia.com> Co-authored-by: Hao Phan <htphan@nvidia.com> Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: RolaoDenthu <xinyisong0111@gmail.com> Co-authored-by: pigeonsoup <32922982+pigeonsoup@users.noreply.github.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com> Co-authored-by: Lee Nau <lee.nau@gmail.com> Co-authored-by: HMING <126185151+Hearum@users.noreply.github.com> Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: Byron Hsu <byronhsu1230@gmail.com> Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Thomas Wang <thomawan@amd.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai> Co-authored-by: BBuf <xiaoyu.zhang@radixark.ai> Co-authored-by: Hanming Lu <hanminglu@meta.com> Co-authored-by: Xinyi Song <xinyis10@illinois.edu>
2026-08-04 13:22:49 -07:00
parent 0753663b8e
commit abddb1c7e9
139 changed files with 15414 additions and 911 deletions
@@ -0,0 +1,139 @@
+#!/bin/bash
+# Install the standard CUDA CI dependencies plus Kimi-K3's FlashInfer assets.
+set -euxo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+# Source (not bash) so the generic install's Python/venv selection remains
+# active while patching the FlashInfer package it just installed.
+# shellcheck disable=SC1091
+source "${SCRIPT_DIR}/ci_install_dependency.sh" "$@"
+
+TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip"
+TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192"
+TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1"
+export SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool"
+
+install_required_tools() {
+    local missing=()
+    local tool
+    for tool in patch unzip wget; do
+        command -v "${tool}" >/dev/null 2>&1 || missing+=("${tool}")
+    done
+    if [ ${#missing[@]} -eq 0 ]; then
+        return
+    fi
+
+    apt-get update || true
+    apt-get install -y --no-install-recommends "${missing[@]}"
+}
+
+cubin_pool_is_valid() {
+    [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" ] &&
+        [ "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" -type f -name '*.cubin' | wc -l)" -eq 1696 ] &&
+        [ -f "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/flashinferMetaInfo.h" ] &&
+        [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/local" ] &&
+        [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/overlay/csrc" ]
+}
+
+install_trtllm_gen_moe_cubin_pool() (
+    if cubin_pool_is_valid; then
+        echo "Reusing validated TRT-LLM Gen MoE cubin pool"
+        return
+    fi
+
+    local cubin_archive
+    local cubin_extract_dir
+    local extracted_pool
+    cubin_archive="$(mktemp /tmp/trtllm_gen_moe_cubin_pool.XXXXXX.zip)"
+    cubin_extract_dir="$(mktemp -d /tmp/trtllm_gen_moe_cubin_extract.XXXXXX)"
+    extracted_pool="${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}"
+    trap 'rm -f "${cubin_archive}"; rm -rf "${cubin_extract_dir}"' EXIT
+
+    wget --no-verbose --output-document="${cubin_archive}" \
+        "${TRTLLM_GEN_MOE_CUBIN_URL}"
+    echo "${TRTLLM_GEN_MOE_CUBIN_SHA256}  ${cubin_archive}" | \
+        sha256sum --check --strict -
+    unzip -q "${cubin_archive}" -d "${cubin_extract_dir}"
+    test "$(find "${extracted_pool}" -type f -name '*.cubin' | wc -l)" -eq 1696
+
+    rm -rf "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}"
+    mv "${extracted_pool}" "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}"
+    cubin_pool_is_valid
+)
+
+apply_flashinfer_dcp_patch() {
+    local flashinfer_dcp_patch
+    local flashinfer_site_packages
+    flashinfer_dcp_patch="${REPO_ROOT}/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt"
+    flashinfer_site_packages="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')"
+
+    sed '/^diff --git a\/tests\//,$d' "${flashinfer_dcp_patch}" | \
+        patch --dry-run --batch --forward --strip=1 \
+            --directory="${flashinfer_site_packages}"
+    sed '/^diff --git a\/tests\//,$d' "${flashinfer_dcp_patch}" | \
+        patch --batch --forward --strip=1 \
+            --directory="${flashinfer_site_packages}"
+
+    rm -rf /root/.cache/flashinfer /root/.cache/pip
+    python3 -c 'import inspect; from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla; assert "enable_dcp" in inspect.signature(trtllm_batch_decode_with_kv_cache_mla).parameters'
+}
+
+apply_transformers_symlink_patch() {
+    # transformers 5.12.1 resolves custom-code symlinks out of the HF snapshot
+    # and into blobs/, then looks for relative imports by filename in blobs/.
+    # Mirror the hot patch in docker/rocm.Dockerfile until upstream PR #46618
+    # reaches a transformers release.
+    python3 - <<'PY'
+import importlib
+import pathlib
+import tempfile
+
+import transformers.dynamic_module_utils as dynamic_module_utils
+
+marks = ["Path(resolved_module_file).resolve()", "Path(source_file).resolve()"]
+path = pathlib.Path(dynamic_module_utils.__file__)
+src = path.read_text()
+if not any(mark in src for mark in marks):
+    print("transformers dynamic_module_utils already fixed; no patch needed")
+else:
+    patched = (
+        src.replace(
+            "Path(resolved_module_file).resolve()", "Path(resolved_module_file)"
+        ).replace("Path(source_file).resolve()", "Path(source_file)")
+    )
+    assert patched != src, "FATAL: transformers symlink patch matched nothing"
+    path.write_text(patched)
+    print("patched transformers dynamic_module_utils.py (symlink hash fix)")
+
+# Exercise the exact HF-cache layout that failed in the B300 Kimi-K3 run.
+dynamic_module_utils = importlib.reload(dynamic_module_utils)
+with tempfile.TemporaryDirectory() as tmp:
+    cache = pathlib.Path(tmp) / "models--org--model"
+    blobs = cache / "blobs"
+    snapshot = cache / "snapshots" / "revision"
+    blobs.mkdir(parents=True)
+    snapshot.mkdir(parents=True)
+    (blobs / "modeling-blob").write_text("from .media_utils import VALUE\n")
+    (blobs / "media-blob").write_text("VALUE = 1\n")
+    (snapshot / "modeling.py").symlink_to("../../blobs/modeling-blob")
+    (snapshot / "media_utils.py").symlink_to("../../blobs/media-blob")
+    source_hash = dynamic_module_utils._compute_local_source_files_hash(
+        snapshot, snapshot / "modeling.py"
+    )
+    assert len(source_hash) == 16, source_hash
+PY
+}
+
+install_required_tools
+install_trtllm_gen_moe_cubin_pool
+apply_flashinfer_dcp_patch
+apply_transformers_symlink_patch
+
+# The install runs in its own shell. Persist the pool path for later workflow
+# steps so Kimi-K3 selects the FlashInfer MXFP4 MoE runner instead of failing
+# its startup validation.
+if [ -n "${GITHUB_ENV:-}" ]; then
+    echo "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" >> "${GITHUB_ENV}"
+fi