Bump FlashInfer to 0.6.17 and remove Kimi K3 workarounds (#33997)

This commit is contained in:
Mohammad Miadh Angkad
2026-08-12 02:17:26 -07:00
committed by GitHub
parent 2d76d537e5
commit 00e57d74f0
19 changed files with 84 additions and 6496 deletions
+2 -83
View File
@@ -1,85 +1,14 @@
#!/bin/bash
# Install the standard CUDA CI dependencies plus Kimi-K3's FlashInfer assets.
# Install the standard CUDA CI dependencies plus Kimi-K3's Transformers fix.
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.
# active while applying the Transformers compatibility fix.
# 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/.
@@ -126,14 +55,4 @@ with tempfile.TemporaryDirectory() as tmp:
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