[Deps] Bump FlashInfer to 0.6.18 (#36954)
This commit is contained in:
Mohammad Miadh Angkad
+1
-1
@@ -14,7 +14,7 @@ ARG PIP_DEFAULT_INDEX
|
||||
ARG UBUNTU_MIRROR
|
||||
ARG GITHUB_ARTIFACTORY=github.com
|
||||
ARG INSTALL_FLASHINFER_JIT_CACHE=0
|
||||
ARG FLASHINFER_VERSION=0.6.17
|
||||
ARG FLASHINFER_VERSION=0.6.18
|
||||
ARG MOONCAKE_VERSION=0.3.13
|
||||
ARG MSCCLPP_VERSION=sglang-v0.9.1
|
||||
|
||||
|
||||
@@ -1,169 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Kimi-K3 DeepEP patch + rebuild -- IMAGE-BUILD variant of the runtime patcher.
|
||||
#
|
||||
# Stock DeepEP (deepseek-ai@d28bd67, shipped in the base image at
|
||||
# /sgl-workspace/DeepEP) does not serve Kimi-K3:
|
||||
# - LL topk capped at 11; K3 routes top-16 -> assert at internode_ll.cu
|
||||
# - SWITCH_HIDDEN lacks 3584 (K3 latent-MoE dispatch dim)
|
||||
# - normal-internode issues an unaligned 64-bit SourceMeta access for packed
|
||||
# K3 FP8 scales when dispatch crosses the 8-rank NVL domain (EP > 8)
|
||||
# - CUDA 13 relocated the cccl headers -> the stock build cannot find them
|
||||
#
|
||||
# Differences vs the devbox runtime script (why a separate copy):
|
||||
# 1. arch comes from $TORCH_CUDA_ARCH_LIST (no GPU at image-build time; the
|
||||
# runtime script probes torch.cuda.get_device_capability which needs a GPU)
|
||||
# 2. adds the setup.py cccl include-dir fix (required to COMPILE on CUDA 13;
|
||||
# the devbox had it applied manually)
|
||||
# 3. adds the configs.cuh CPU/cycle timeout bump (cross-node init headroom)
|
||||
#
|
||||
# Idempotent (grep/count-guarded). The default wheel contains native cubins for
|
||||
# Hopper sm_90, B200 sm_100a, and GB300 sm_103a.
|
||||
set -euo pipefail
|
||||
: "${TORCH_CUDA_ARCH_LIST:=9.0;10.0a;10.3a}"
|
||||
DEEPEP_DIR="${DEEPEP_DIR:-/sgl-workspace/DeepEP}"
|
||||
DEEPEP_COMMIT="${DEEPEP_COMMIT:-d28bd676c2120573c9f1425f0c16c39faa4117e6}"
|
||||
|
||||
# PyTorch accepts CUDA architecture lists separated by spaces or semicolons.
|
||||
read -r -a CUDA_ARCHES <<< "${TORCH_CUDA_ARCH_LIST//;/ }"
|
||||
[ "${#CUDA_ARCHES[@]}" -gt 0 ] || { echo "ERROR: TORCH_CUDA_ARCH_LIST is empty"; exit 1; }
|
||||
for CUDA_ARCH in "${CUDA_ARCHES[@]}"; do
|
||||
[[ "$CUDA_ARCH" =~ ^[0-9]+\.[0-9]+a?$ ]] || {
|
||||
echo "ERROR: unsupported CUDA architecture '$CUDA_ARCH' in TORCH_CUDA_ARCH_LIST='$TORCH_CUDA_ARCH_LIST'"
|
||||
exit 1
|
||||
}
|
||||
done
|
||||
|
||||
# The base image is expected to ship the DeepEP source here; clone-pin if not.
|
||||
if [ ! -d "$DEEPEP_DIR/csrc" ]; then
|
||||
echo "== $DEEPEP_DIR missing — cloning deepseek-ai/DeepEP @ $DEEPEP_COMMIT"
|
||||
git clone --recursive https://github.com/deepseek-ai/DeepEP.git "$DEEPEP_DIR"
|
||||
git -C "$DEEPEP_DIR" checkout "$DEEPEP_COMMIT"
|
||||
git -C "$DEEPEP_DIR" submodule update --init --recursive
|
||||
fi
|
||||
cd "$DEEPEP_DIR"
|
||||
|
||||
echo "== [1/6] internode_ll.cu: LL topk caps 9/11 -> 16"
|
||||
TOPK_CAPS_EXPECTED=$(grep -Eci "kNumMaxTop[Kk] = (9|11|16)" csrc/kernels/internode_ll.cu || true)
|
||||
[ "$TOPK_CAPS_EXPECTED" -ge 2 ] || {
|
||||
echo "ERROR: expected >=2 recognized topk caps, found $TOPK_CAPS_EXPECTED"
|
||||
exit 1
|
||||
}
|
||||
sed -i 's/constexpr int kNumMaxTopK = 11;/constexpr int kNumMaxTopK = 16;/' csrc/kernels/internode_ll.cu
|
||||
sed -i 's/constexpr int kNumMaxTopK = 9;/constexpr int kNumMaxTopK = 16;/' csrc/kernels/internode_ll.cu
|
||||
sed -i 's/constexpr int kNumMaxTopk = 9;/constexpr int kNumMaxTopk = 16;/' csrc/kernels/internode_ll.cu
|
||||
sed -i 's/constexpr int kNumMaxTopk = 11;/constexpr int kNumMaxTopk = 16;/' csrc/kernels/internode_ll.cu
|
||||
TOPK_CAPS_PATCHED=$(grep -ci "kNumMaxTop[Kk] = 16" csrc/kernels/internode_ll.cu || true)
|
||||
TOPK_CAPS_UNPATCHED=$(grep -Eci "kNumMaxTop[Kk] = (9|11)" csrc/kernels/internode_ll.cu || true)
|
||||
if [ "$TOPK_CAPS_PATCHED" -ne "$TOPK_CAPS_EXPECTED" ] || [ "$TOPK_CAPS_UNPATCHED" -ne 0 ]; then
|
||||
echo "ERROR: expected $TOPK_CAPS_EXPECTED topk caps patched, found $TOPK_CAPS_PATCHED patched and $TOPK_CAPS_UNPATCHED unpatched"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "== [2/6] launch.cuh: SWITCH_HIDDEN += case 3584 (K3 latent MoE)"
|
||||
grep -q "case_macro(3584)" csrc/kernels/launch.cuh || \
|
||||
sed -i 's/case 4096: case_macro(4096);/case 3584: case_macro(3584); case 4096: case_macro(4096);/' csrc/kernels/launch.cuh
|
||||
grep -q "case_macro(3584)" csrc/kernels/launch.cuh || { echo "ERROR: hidden 3584 case not applied"; exit 1; }
|
||||
|
||||
echo "== [3/6] configs.cuh: raise CPU/cycle timeouts 100s -> 1000s (cross-node init headroom)"
|
||||
# '@' delimiter (pattern contains '#'); '$' anchor avoids the ENABLE_FAST_DEBUG '...10' line.
|
||||
sed -i 's@^#define NUM_CPU_TIMEOUT_SECS 100$@#define NUM_CPU_TIMEOUT_SECS 1000@' csrc/kernels/configs.cuh
|
||||
sed -i 's@^#define NUM_TIMEOUT_CYCLES 200000000000ull@#define NUM_TIMEOUT_CYCLES 2000000000000ull@' csrc/kernels/configs.cuh
|
||||
grep -q "NUM_CPU_TIMEOUT_SECS 1000$" csrc/kernels/configs.cuh || echo " WARN: CPU timeout not raised (layout changed?) — non-fatal"
|
||||
|
||||
echo "== [4/6] tests/test_low_latency.py: nvfp4 tolerance 0.007 -> 0.008 (topk16 noise)"
|
||||
python3 - <<'EOF'
|
||||
p = "tests/test_low_latency.py"
|
||||
s = open(p).read()
|
||||
old = "elif dispatch_use_nvfp4:\n diff_threshold = 0.007"
|
||||
new = "elif dispatch_use_nvfp4:\n diff_threshold = 0.008"
|
||||
if old in s:
|
||||
open(p, "w").write(s.replace(old, new)); print(" patched")
|
||||
else:
|
||||
print(" already patched or layout changed (skipped)")
|
||||
EOF
|
||||
|
||||
echo "== [5/6] internode.cu: 4-byte-aligned SourceMeta scalar access (EP>8 normal internode)"
|
||||
python3 - <<'EOF'
|
||||
from pathlib import Path
|
||||
path = Path("csrc/kernels/internode.cu")
|
||||
source = path.read_text()
|
||||
changed = False
|
||||
abi_old = 'EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`");'
|
||||
abi_new = 'EP_STATIC_ASSERT(sizeof(SourceMeta) == 2 * sizeof(int), "SourceMeta scalar access requires exactly two int fields");'
|
||||
if source.count(abi_old) == 2 and source.count(abi_new) == 0:
|
||||
source = source.replace(abi_old, abi_new); changed = True; print(" tightened both SourceMeta ABI assertions")
|
||||
elif source.count(abi_old) == 0 and source.count(abi_new) == 2:
|
||||
print(" ABI assertions already tightened")
|
||||
else:
|
||||
raise SystemExit(f"ERROR: unexpected SourceMeta ABI layout: old={source.count(abi_old)}, new={source.count(abi_new)}")
|
||||
replacements = (
|
||||
("sender store",
|
||||
""" // Copy source metadata into symmetric send buffer
|
||||
if (lane_id < num_topk_ranks)
|
||||
st_na_global(reinterpret_cast<SourceMeta*>(dst_send_buffers[lane_id]), src_meta);
|
||||
""",
|
||||
""" // SourceMeta may be only 4-byte aligned after packed scales.
|
||||
// Store its two int fields separately to avoid an unaligned 64-bit store.
|
||||
if (lane_id < num_topk_ranks) {
|
||||
auto meta_values = reinterpret_cast<int*>(dst_send_buffers[lane_id]);
|
||||
st_na_global(meta_values, src_meta.src_rdma_rank);
|
||||
st_na_global(meta_values + 1, src_meta.is_token_in_nvl_rank_bits);
|
||||
}
|
||||
"""),
|
||||
("forwarder load",
|
||||
""" auto src_meta = ld_nc_global(reinterpret_cast<SourceMeta*>(shifted + hidden_bytes + scale_bytes));
|
||||
""",
|
||||
""" auto src_meta_values = reinterpret_cast<const int*>(shifted + hidden_bytes + scale_bytes);
|
||||
SourceMeta src_meta;
|
||||
src_meta.src_rdma_rank = ld_nc_global(src_meta_values);
|
||||
src_meta.is_token_in_nvl_rank_bits = ld_nc_global(src_meta_values + 1);
|
||||
"""),
|
||||
("receiver load",
|
||||
""" auto meta = ld_nc_global(reinterpret_cast<SourceMeta*>(shifted + hidden_bytes + scale_bytes));
|
||||
""",
|
||||
""" auto meta_values = reinterpret_cast<const int*>(shifted + hidden_bytes + scale_bytes);
|
||||
SourceMeta meta;
|
||||
meta.src_rdma_rank = ld_nc_global(meta_values);
|
||||
meta.is_token_in_nvl_rank_bits = ld_nc_global(meta_values + 1);
|
||||
"""),
|
||||
)
|
||||
for label, old, new in replacements:
|
||||
oc, nc = source.count(old), source.count(new)
|
||||
if oc == 1 and nc == 0:
|
||||
source = source.replace(old, new); changed = True; print(f" patched {label}")
|
||||
elif oc == 0 and nc == 1:
|
||||
print(f" {label} already patched")
|
||||
else:
|
||||
raise SystemExit(f"ERROR: unexpected {label} layout: old={oc}, new={nc}")
|
||||
for unsafe in ("st_na_global(reinterpret_cast<SourceMeta*>", "ld_nc_global(reinterpret_cast<SourceMeta*>"):
|
||||
if unsafe in source:
|
||||
raise SystemExit(f"ERROR: unsafe SourceMeta access remains: {unsafe}")
|
||||
if changed:
|
||||
path.write_text(source)
|
||||
EOF
|
||||
|
||||
echo "== [5b] setup.py: add CUDA 13 cccl include dir (compile fix)"
|
||||
grep -q "/usr/local/cuda/include/cccl" setup.py || \
|
||||
sed -i "s#\( include_dirs = \['csrc/'\]\)#\1\n include_dirs.append('/usr/local/cuda/include/cccl')#" setup.py
|
||||
grep -q "/usr/local/cuda/include/cccl" setup.py || { echo "ERROR: cccl include not added to setup.py"; exit 1; }
|
||||
|
||||
echo "== [6/6] rebuild + reinstall for TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST (no GPU needed)"
|
||||
rm -rf build dist
|
||||
TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" python3 setup.py bdist_wheel
|
||||
pip install dist/*.whl --force-reinstall --no-deps
|
||||
|
||||
echo "== verify installed cubin arches"
|
||||
SO=$(find /usr/local/lib/python3*/dist-packages -maxdepth 1 -name "deep_ep_cpp*.so" | head -1)
|
||||
[ -n "$SO" ] || { echo "ERROR: installed deep_ep_cpp shared object not found"; exit 1; }
|
||||
CUBIN_LIST=$(cuobjdump --list-elf "$SO")
|
||||
VERIFIED_SMS=()
|
||||
for CUDA_ARCH in "${CUDA_ARCHES[@]}"; do
|
||||
SM="sm_${CUDA_ARCH/./}"
|
||||
grep -Eq "(^|[^[:alnum:]_])${SM}([^[:alnum:]_]|$)" <<< "$CUBIN_LIST" || {
|
||||
echo "ERROR: cubin arch $SM not found in $SO"
|
||||
printf '%s\n' "$CUBIN_LIST" | head
|
||||
exit 1
|
||||
}
|
||||
VERIFIED_SMS+=("$SM")
|
||||
done
|
||||
echo "== OK: DeepEP rebuilt (topk16 + hidden3584 + SourceMeta align + cccl) for ${VERIFIED_SMS[*]}"
|
||||
@@ -1,113 +0,0 @@
|
||||
# Kimi-K3 serving image (x86_64 / CUDA 12.9 / sm_90 + sm_100a).
|
||||
#
|
||||
# Base ships stock SGLang (editable at /sgl-workspace/sglang), DeepEP source
|
||||
# (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package,
|
||||
# and the CUDA 12.9 toolchain.
|
||||
#
|
||||
# This image adds the three Kimi-K3-specific pieces that stock lacks:
|
||||
# 1. the Kimi-K3 SGLang code (this repo), editable-installed
|
||||
# 2. DeepEP patch + rebuild:
|
||||
# topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment,
|
||||
# and cross-node timeout headroom; rebuilt for sm_90 and sm_100a only
|
||||
# 3. DeepGEMM upgrade to 0.1.5.post2:
|
||||
# official MegaMoE runtime-JIT header with Kimi-K3 SiTU support
|
||||
#
|
||||
# Build (on/for x86_64; nvcc cross-compiles the DeepEP cubin, no GPU needed):
|
||||
# docker build -f docker/kimi_k3/kimi_k3_cu12.Dockerfile \
|
||||
# --build-arg 'TORCH_CUDA_ARCH_LIST=9.0;10.0a' -t kimi-k3-cu129 .
|
||||
#
|
||||
# The FlashInfer MXFP4 MoE runner cubins are installed in the image below.
|
||||
# The runner is auto-selected on SM100; the remaining kernel sources
|
||||
# JIT-compile from the installed FlashInfer wheel on first launch and are
|
||||
# cached.
|
||||
|
||||
FROM lmsysorg/sglang:v0.5.16-cu129 AS base
|
||||
|
||||
ARG SGL_DEEP_GEMM_VERSION="0.1.5.post2"
|
||||
ARG NVIMGCODEC_VERSION="0.9.0.20"
|
||||
|
||||
# Current Kimi-K3 source auto-discovers and builds its PyO3 extensions.
|
||||
ARG RUST_VERSION="1.90.0"
|
||||
ENV RUSTUP_HOME="/usr/local/rustup" \
|
||||
CARGO_HOME="/usr/local/cargo" \
|
||||
PATH="/usr/local/cargo/bin:${PATH}"
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl && \
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --no-modify-path --profile minimal \
|
||||
--default-toolchain "${RUST_VERSION}" && \
|
||||
cargo --version && \
|
||||
rustc --version && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Build one DeepEP wheel with native cubins for Hopper and B200. This CUDA 12.9
|
||||
# recipe intentionally excludes GB300 (sm_103/sm_103a).
|
||||
ARG TORCH_CUDA_ARCH_LIST="9.0;10.0a"
|
||||
RUN set -eu; \
|
||||
for arch in $(printf '%s' "${TORCH_CUDA_ARCH_LIST}" | tr ';' ' '); do \
|
||||
case "${arch}" in \
|
||||
10.3|10.3a) \
|
||||
echo "ERROR: CUDA 12.9 image does not support SM103 (${arch})" >&2; \
|
||||
exit 1 \
|
||||
;; \
|
||||
esac; \
|
||||
done
|
||||
|
||||
# --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) ---
|
||||
# Keep the installed extension modules, but discard Rust and pip build
|
||||
# artifacts that are not used at runtime.
|
||||
RUN rm -rf /sgl-workspace/sglang && \
|
||||
git clone --branch main \
|
||||
https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \
|
||||
cd /sgl-workspace/sglang && \
|
||||
rm -rf .git && \
|
||||
test ! -e .git && \
|
||||
pip install -e python --no-deps && \
|
||||
rm -rf \
|
||||
rust/target \
|
||||
rust/sglang-grpc/target \
|
||||
rust/sglang-mm/target \
|
||||
rust/sglang-server/target \
|
||||
/usr/local/cargo/registry \
|
||||
/root/.cache/pip
|
||||
|
||||
# --- 2. DeepEP: patch (topk16 / hidden3584 / SourceMeta) + multi-arch rebuild ---
|
||||
RUN TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
|
||||
bash /sgl-workspace/sglang/docker/kimi_k3/apply_deepep_k3_patch.sh && \
|
||||
rm -rf /sgl-workspace/DeepEP/build /sgl-workspace/DeepEP/dist
|
||||
|
||||
# --- 3. DeepGEMM: upgrade to the first release with Kimi-K3 SiTU ---
|
||||
# The v0.5.16 base contains DeepGEMM 0.1.4.post1. PyPI publishes the CUDA 13
|
||||
# build, so CUDA 12.9 uses the matching official release asset.
|
||||
RUN python3 -m pip install --no-deps --force-reinstall \
|
||||
"https://github.com/sgl-project/whl/releases/download/v${SGL_DEEP_GEMM_VERSION}/sgl_deep_gemm-${SGL_DEEP_GEMM_VERSION}+cu129-py3-none-manylinux2014_x86_64.whl"
|
||||
|
||||
# High-fidelity GPU JPEG decode. The K3 processor enables nvJPEG interpolated
|
||||
# chroma upsampling through nvImageCodec and zero-copy DLPack handoff to Torch.
|
||||
RUN python3 -m pip install \
|
||||
"nvidia-nvimgcodec-cu12[all]==${NVIMGCODEC_VERSION}" && \
|
||||
rm -rf /root/.cache/pip
|
||||
|
||||
# Install the matching official FlashInfer package trio. A mixed
|
||||
# Python/cubin/JIT-cache installation fails at import time.
|
||||
# flashinfer-python and flashinfer-cubin are CUDA-independent packages; the
|
||||
# JIT-cache wheel is selected from the official CUDA 12.9 index.
|
||||
RUN python3 -m pip uninstall -y \
|
||||
flashinfer-python flashinfer-cubin flashinfer-jit-cache && \
|
||||
rm -rf /root/.cache/flashinfer /root/.cache/pip && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-python==0.6.17" && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-cubin==0.6.17" \
|
||||
--index-url https://flashinfer.ai/whl && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-jit-cache==0.6.17" \
|
||||
--index-url https://flashinfer.ai/whl/cu129 && \
|
||||
python3 -c 'from importlib.metadata import version; expected = "0.6.17"; assert version("flashinfer-python").split("+", 1)[0] == expected; assert version("flashinfer-cubin").split("+", 1)[0] == expected; assert version("flashinfer-jit-cache").startswith(expected + "+cu129"), version("flashinfer-jit-cache")' && \
|
||||
rm -rf /root/.cache/pip
|
||||
|
||||
ENV FLASHINFER_VERSION="0.6.17"
|
||||
|
||||
WORKDIR /sgl-workspace/sglang
|
||||
@@ -1,101 +0,0 @@
|
||||
# Kimi-K3 serving image (aarch64 / sm_90 + sm_100a + sm_103a).
|
||||
#
|
||||
# Base ships stock SGLang (editable at /sgl-workspace/sglang), DeepEP source
|
||||
# (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package,
|
||||
# and the CUDA 13 toolchain (nvcc + /usr/local/cuda/include/cccl).
|
||||
#
|
||||
# This image adds the three Kimi-K3-specific pieces that stock lacks:
|
||||
# 1. the Kimi-K3 SGLang code (this repo), editable-installed
|
||||
# 2. DeepEP patch + rebuild:
|
||||
# topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment,
|
||||
# cross-node timeout headroom, CUDA-13 cccl include; rebuilt for
|
||||
# sm_90, sm_100a, and sm_103a
|
||||
# 3. DeepGEMM upgrade to 0.1.5.post2:
|
||||
# official MegaMoE runtime-JIT header with Kimi-K3 SiTU support
|
||||
#
|
||||
# Build (on/for aarch64; nvcc cross-compiles the DeepEP cubin, no GPU needed):
|
||||
# docker build -f docker/kimi_k3/kimi_k3_cu13.Dockerfile \
|
||||
# --build-arg 'TORCH_CUDA_ARCH_LIST=9.0;10.0a;10.3a' -t kimi-k3 .
|
||||
#
|
||||
# The FlashInfer MXFP4 MoE runner cubins are installed in the image below.
|
||||
# The runner is auto-selected on SM100/103; the remaining kernel sources
|
||||
# JIT-compile from the installed FlashInfer wheel on first launch and are
|
||||
# cached.
|
||||
|
||||
FROM lmsysorg/sglang:v0.5.16 AS base
|
||||
|
||||
ARG SGL_DEEP_GEMM_VERSION="0.1.5.post2"
|
||||
ARG NVIMGCODEC_VERSION="0.9.0.20"
|
||||
|
||||
# Current Kimi-K3 source auto-discovers and builds its PyO3 extensions.
|
||||
ARG RUST_VERSION="1.90.0"
|
||||
ENV RUSTUP_HOME="/usr/local/rustup" \
|
||||
CARGO_HOME="/usr/local/cargo" \
|
||||
PATH="/usr/local/cargo/bin:${PATH}"
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl && \
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --no-modify-path --profile minimal \
|
||||
--default-toolchain "${RUST_VERSION}" && \
|
||||
cargo --version && \
|
||||
rustc --version && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Build one DeepEP wheel with native cubins for Hopper, B200, and GB300.
|
||||
ARG TORCH_CUDA_ARCH_LIST="9.0;10.0a;10.3a"
|
||||
|
||||
# --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) ---
|
||||
# Keep the installed extension modules, but discard Rust and pip build
|
||||
# artifacts that are not used at runtime.
|
||||
RUN rm -rf /sgl-workspace/sglang && \
|
||||
git clone --branch main \
|
||||
https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \
|
||||
cd /sgl-workspace/sglang && \
|
||||
rm -rf .git && \
|
||||
test ! -e .git && \
|
||||
pip install -e python --no-deps && \
|
||||
rm -rf \
|
||||
rust/target \
|
||||
rust/sglang-grpc/target \
|
||||
rust/sglang-mm/target \
|
||||
rust/sglang-server/target \
|
||||
/usr/local/cargo/registry \
|
||||
/root/.cache/pip
|
||||
|
||||
# --- 2. DeepEP: patch (topk16 / hidden3584 / SourceMeta / cccl) + multi-arch rebuild ---
|
||||
RUN TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
|
||||
bash /sgl-workspace/sglang/docker/kimi_k3/apply_deepep_k3_patch.sh && \
|
||||
rm -rf /sgl-workspace/DeepEP/build /sgl-workspace/DeepEP/dist
|
||||
|
||||
# --- 3. DeepGEMM: upgrade to the first release with Kimi-K3 SiTU ---
|
||||
# The v0.5.16 base contains DeepGEMM 0.1.4.post1.
|
||||
RUN python3 -m pip install --no-deps --force-reinstall \
|
||||
"sgl-deep-gemm==${SGL_DEEP_GEMM_VERSION}"
|
||||
|
||||
# High-fidelity GPU JPEG decode. The K3 processor enables nvJPEG interpolated
|
||||
# chroma upsampling through nvImageCodec and zero-copy DLPack handoff to Torch.
|
||||
RUN python3 -m pip install \
|
||||
"nvidia-nvimgcodec-cu13[all]==${NVIMGCODEC_VERSION}" && \
|
||||
rm -rf /root/.cache/pip
|
||||
|
||||
# Install the matching official FlashInfer package trio. A mixed
|
||||
# Python/cubin/JIT-cache installation fails at import time.
|
||||
RUN python3 -m pip uninstall -y \
|
||||
flashinfer-python flashinfer-cubin flashinfer-jit-cache && \
|
||||
rm -rf /root/.cache/flashinfer /root/.cache/pip && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-python==0.6.17" && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-cubin==0.6.17" \
|
||||
--index-url https://flashinfer.ai/whl && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-jit-cache==0.6.17" \
|
||||
--index-url https://flashinfer.ai/whl/cu130 && \
|
||||
python3 -c 'from importlib.metadata import version; expected = "0.6.17"; packages = ("flashinfer-python", "flashinfer-cubin", "flashinfer-jit-cache"); actual = {package: version(package).split("+", 1)[0] for package in packages}; assert all(value == expected for value in actual.values()), actual' && \
|
||||
rm -rf /root/.cache/pip
|
||||
|
||||
ENV FLASHINFER_VERSION="0.6.17"
|
||||
|
||||
WORKDIR /sgl-workspace/sglang
|
||||
@@ -149,7 +149,7 @@ Capacity levers, all in the Playground. Each trades precision or cache behavior
|
||||
|
||||
Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request — the calculator folds this in — and an unset `--max-running-requests` resets to 48 under spec (the command panel reminds you; set it explicitly to raise).
|
||||
|
||||
**MoE runner.** Leave `--moe-runner-backend` unset on Blackwell: FlashInfer MXFP4 (W4A8, official trtllm-gen SiTU kernels) is selected with the pinned FlashInfer 0.6.17 dependency; H100/H200 pin Marlin. The B200 Balanced and High-Throughput cells pin `flashinfer_mxfp4` explicitly because that is the shape they were brought up on. The published Docker images install the matching official `flashinfer-python`, `flashinfer-cubin`, and `flashinfer-jit-cache` packages.
|
||||
**MoE runner.** Leave `--moe-runner-backend` unset on Blackwell: FlashInfer MXFP4 (W4A8, official trtllm-gen SiTU kernels) is selected; H100/H200 pin Marlin. The B200 Balanced and High-Throughput cells pin `flashinfer_mxfp4` explicitly because that is the shape they were brought up on.
|
||||
|
||||
**Attention backend.** Leave all three attention knobs unset on Blackwell: K3 resolves prefill, decode, and — under DSPARK — verification as a set (`trtllm_mla` across the board; `cutedsl_mla` takes decode and verification under DCP). On the non-DCP recipes, setting any one of the three cancels the auto-resolution for the others. The B200 Balanced and High-Throughput cells pin `--decode-attention-backend cutedsl_mla`, which is what auto-resolution picks for those DCP recipes anyway — it is written out because it is the shape they were brought up on, not because it changes the resolution. H100/H200 pin `flashmla` for decode.
|
||||
|
||||
|
||||
@@ -33,7 +33,7 @@ dependencies = [
|
||||
"einops",
|
||||
"fastapi",
|
||||
"flash-attn-4>=4.0.0b18",
|
||||
"flashinfer_python[cu13]==0.6.17", # keep it aligned with jit-cache version in Dockerfile
|
||||
"flashinfer_python[cu13]==0.6.18", # keep it aligned with jit-cache version in Dockerfile
|
||||
"gguf",
|
||||
"humming-kernels[cu13]==0.1.12",
|
||||
"interegular",
|
||||
|
||||
@@ -1686,7 +1686,7 @@ def _set_envs_and_config(server_args: ServerArgs):
|
||||
if "flashinfer" in attention_backends_of(resolved_view(cfg)):
|
||||
assert_pkg_version(
|
||||
"flashinfer_python",
|
||||
"0.6.17",
|
||||
"0.6.18",
|
||||
"Please uninstall the old version and "
|
||||
"reinstall the latest version by following the instructions "
|
||||
"at https://docs.flashinfer.ai/installation.html.",
|
||||
|
||||
@@ -2157,7 +2157,7 @@ def check_pkg_version_at_least(pkg: str, min_version: str) -> bool:
|
||||
|
||||
Args:
|
||||
pkg: Package name (distribution name, e.g., "flashinfer-python")
|
||||
min_version: Minimum version required (e.g., "0.6.17")
|
||||
min_version: Minimum version required (e.g., "0.6.18")
|
||||
|
||||
Returns:
|
||||
True if package is installed and version >= min_version, False otherwise
|
||||
|
||||
@@ -313,10 +313,10 @@ def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict:
|
||||
class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase):
|
||||
"""Derived property: the envelope-strided KDA temporal view (unified memory
|
||||
/ page-major layout) must satisfy the state contract of FlashInfer
|
||||
``recurrent_kda`` (pinned ``flashinfer_python==0.6.17``), because the KDA
|
||||
flashinfer decode wrapper (``linear/kernels/kda_flashinfer.py``) passes the
|
||||
committed per-layer pool view straight into the kernel (in-place state
|
||||
update on the cu_seqlens path — no gather/scatter copy around the call).
|
||||
``recurrent_kda``, because the KDA flashinfer decode wrapper
|
||||
(``linear/kernels/kda_flashinfer.py``) passes the committed per-layer pool
|
||||
view straight into the kernel (in-place state update on the cu_seqlens path
|
||||
— no gather/scatter copy around the call).
|
||||
|
||||
The kernel compiles its state argument as a CuTe fake tensor of shape
|
||||
``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)``
|
||||
|
||||
Reference in New Issue
Block a user