docker: add Kimi K3 images (#32760)
This commit is contained in:
Baizhou Zhang
Executable
+169
@@ -0,0 +1,169 @@
|
||||
#!/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[*]}"
|
||||
@@ -0,0 +1,140 @@
|
||||
"""Patch deep_gemm's mega-MoE JIT header to support Kimi-K3 SiTU activation.
|
||||
|
||||
Mechanism (no _C.so rebuild needed — the kernel body is a runtime-JIT header):
|
||||
the host passes `activation='swiglu'` with the magic `activation_clamp =
|
||||
0.03125` (2^-5: exactly representable, round-trips through the host's float
|
||||
stringification, and no legitimate swiglu clamp uses it; the host asserts
|
||||
clamp >= 0 so a negative sentinel is not possible). In-kernel:
|
||||
kActivationClamp == 0.03125f selects SiTU with K3 constants baked in:
|
||||
beta = 4.0, linear_beta = 25.0 (config activation_situ_{beta,linear_beta})
|
||||
SiTU(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * (linear_beta*tanh(up/linear_beta))
|
||||
|
||||
A distinct clamp value produces a distinct JIT template instantiation, so the
|
||||
new variant compiles fresh; cached swiglu kernels are unaffected. If you edit
|
||||
the SiTU math itself, clear /root/.cache/deep_gemm first (same sentinel value
|
||||
would otherwise hit a stale cache entry).
|
||||
|
||||
Idempotent; also migrates the deprecated negative-sentinel V1 patch.
|
||||
Run on every node: python3 apply_deepgemm_situ_patch.py
|
||||
"""
|
||||
|
||||
P = "/usr/local/lib/python3.12/dist-packages/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh"
|
||||
|
||||
OLD = """ // Apply SwiGLU: silu(gate) * up
|
||||
// Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7)
|
||||
auto fp32_values = reinterpret_cast<float*>(values);
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; ++ k) {
|
||||
auto bf16_gate = __float22bfloat162_rn(make_float2(fp32_values[k * 4], fp32_values[k * 4 + 1]));
|
||||
auto bf16_up = __float22bfloat162_rn(make_float2(fp32_values[k * 4 + 2], fp32_values[k * 4 + 3]));
|
||||
|
||||
// Clamp
|
||||
if constexpr (kActivationClamp != cute::numeric_limits<float>::infinity()) {
|
||||
bf16_gate = __hmin2(bf16_gate, {kActivationClamp, kActivationClamp});
|
||||
bf16_up = __hmax2(bf16_up, {-kActivationClamp, -kActivationClamp});
|
||||
bf16_up = __hmin2(bf16_up, {kActivationClamp, kActivationClamp});
|
||||
}
|
||||
|
||||
// SwiGLU
|
||||
auto gate = __bfloat1622float2(bf16_gate);
|
||||
auto neg_gate_exp = make_float2(
|
||||
kFastMath ? __expf(-gate.x) : expf(-gate.x),
|
||||
kFastMath ? __expf(-gate.y) : expf(-gate.y));
|
||||
const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp);
|
||||
if constexpr (kFastMath) {
|
||||
gate = __fmul2_rn(gate, {math::fast_rcp(denom.x), math::fast_rcp(denom.y)});
|
||||
} else {
|
||||
gate = {gate.x / denom.x, gate.y / denom.y};
|
||||
}
|
||||
const auto up = __bfloat1622float2(bf16_up);
|
||||
activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights);
|
||||
}
|
||||
"""
|
||||
|
||||
NEW = """ // Apply activation: SwiGLU, or Kimi-K3 SiTU via sentinel
|
||||
// Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7)
|
||||
// K3-SITU-PATCH: kActivationClamp == 0.03125f (2^-5 magic;
|
||||
// host asserts clamp >= 0 so negatives can't sentinel) selects SiTU:
|
||||
// act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate)
|
||||
// up' = kSituLinearBeta * tanh(up/kSituLinearBeta)
|
||||
// K3 config constants baked in (activation_situ_{beta,linear_beta}).
|
||||
constexpr bool kUseSitu = (kActivationClamp == 0.03125f);
|
||||
constexpr float kSituBeta = 4.0f;
|
||||
constexpr float kSituLinearBeta = 25.0f;
|
||||
auto fp32_values = reinterpret_cast<float*>(values);
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; ++ k) {
|
||||
auto bf16_gate = __float22bfloat162_rn(make_float2(fp32_values[k * 4], fp32_values[k * 4 + 1]));
|
||||
auto bf16_up = __float22bfloat162_rn(make_float2(fp32_values[k * 4 + 2], fp32_values[k * 4 + 3]));
|
||||
|
||||
// Clamp (SwiGLU-with-limit only; SiTU soft-clips below)
|
||||
if constexpr (!kUseSitu && kActivationClamp != cute::numeric_limits<float>::infinity()) {
|
||||
bf16_gate = __hmin2(bf16_gate, {kActivationClamp, kActivationClamp});
|
||||
bf16_up = __hmax2(bf16_up, {-kActivationClamp, -kActivationClamp});
|
||||
bf16_up = __hmin2(bf16_up, {kActivationClamp, kActivationClamp});
|
||||
}
|
||||
|
||||
// sigmoid(gate)
|
||||
auto gate = __bfloat1622float2(bf16_gate);
|
||||
auto neg_gate_exp = make_float2(
|
||||
kFastMath ? __expf(-gate.x) : expf(-gate.x),
|
||||
kFastMath ? __expf(-gate.y) : expf(-gate.y));
|
||||
const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp);
|
||||
float2 sig;
|
||||
if constexpr (kFastMath) {
|
||||
sig = {math::fast_rcp(denom.x), math::fast_rcp(denom.y)};
|
||||
} else {
|
||||
sig = {1.0f / denom.x, 1.0f / denom.y};
|
||||
}
|
||||
auto up = __bfloat1622float2(bf16_up);
|
||||
if constexpr (kUseSitu) {
|
||||
// K3-SITU-PATCH: tanh-bounded gate, soft-clipped up
|
||||
gate = {kSituBeta * tanhf(gate.x / kSituBeta) * sig.x,
|
||||
kSituBeta * tanhf(gate.y / kSituBeta) * sig.y};
|
||||
up = {kSituLinearBeta * tanhf(up.x / kSituLinearBeta),
|
||||
kSituLinearBeta * tanhf(up.y / kSituLinearBeta)};
|
||||
} else {
|
||||
// SwiGLU: silu(gate) * up
|
||||
gate = __fmul2_rn(gate, sig);
|
||||
}
|
||||
activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights);
|
||||
}
|
||||
"""
|
||||
|
||||
V1_LINES = """ constexpr bool kUseSitu = kActivationClamp < 0.0f;
|
||||
constexpr float kSituBeta = kUseSitu ? -kActivationClamp : 1.0f;
|
||||
constexpr float kSituLinearBeta = 25.0f;"""
|
||||
V2_LINES = """ constexpr bool kUseSitu = (kActivationClamp == 0.03125f);
|
||||
constexpr float kSituBeta = 4.0f;
|
||||
constexpr float kSituLinearBeta = 25.0f;"""
|
||||
|
||||
s = open(P).read()
|
||||
if V2_LINES in s:
|
||||
print("already patched (v2)")
|
||||
elif V1_LINES in s:
|
||||
# migrate deprecated negative-sentinel v1 -> magic-value v2
|
||||
assert s.count(V1_LINES) == 1
|
||||
s = s.replace(V1_LINES, V2_LINES)
|
||||
s = s.replace(
|
||||
""" // K3-SITU-PATCH: kActivationClamp < 0 selects SiTU:
|
||||
// act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate)
|
||||
// up' = kSituLinearBeta * tanh(up/kSituLinearBeta)
|
||||
// with kSituBeta = -kActivationClamp (host passes -beta).
|
||||
""",
|
||||
""" // K3-SITU-PATCH: kActivationClamp == 0.03125f (2^-5 magic;
|
||||
// host asserts clamp >= 0 so negatives can't sentinel) selects SiTU:
|
||||
// act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate)
|
||||
// up' = kSituLinearBeta * tanh(up/kSituLinearBeta)
|
||||
// K3 config constants baked in (activation_situ_{beta,linear_beta}).
|
||||
""",
|
||||
)
|
||||
open(P, "w").write(s)
|
||||
print("migrated v1 -> v2")
|
||||
elif OLD in s:
|
||||
assert s.count(OLD) == 1
|
||||
open(P, "w").write(s.replace(OLD, NEW))
|
||||
print("patched")
|
||||
else:
|
||||
raise SystemExit(
|
||||
"ERROR: expected SwiGLU epilogue block not found — header layout changed"
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,136 @@
|
||||
# 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 four 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 mega-MoE SiTU patch:
|
||||
# JIT-header sentinel (activation_clamp==0.03125 -> K3 SiTU); no rebuild
|
||||
# 4. FlashInfer CuTeDSL MLA DCP patch:
|
||||
# apply the seven runtime-file diffs; exclude tests absent from the wheel
|
||||
#
|
||||
# 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
|
||||
|
||||
# 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 \
|
||||
unzip \
|
||||
wget && \
|
||||
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 kimi-k3 \
|
||||
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 mega-MoE: SiTU JIT-header patch (runtime-JIT, no rebuild) ---
|
||||
RUN python3 /sgl-workspace/sglang/docker/kimi_k3/apply_deepgemm_situ_patch.py
|
||||
|
||||
# Install the pinned FlashInfer MXFP4 MoE runner cubin pool.
|
||||
ARG 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"
|
||||
ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192"
|
||||
ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1"
|
||||
ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool"
|
||||
|
||||
RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \
|
||||
cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \
|
||||
wget --no-verbose --output-document="${cubin_archive}" \
|
||||
"${TRTLLM_GEN_MOE_CUBIN_URL}" && \
|
||||
echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \
|
||||
sha256sum --check --strict - && \
|
||||
mkdir -p "${cubin_extract_dir}" && \
|
||||
unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \
|
||||
test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
|
||||
mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \
|
||||
"${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
|
||||
test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \
|
||||
-type f -name '*.cubin' | wc -l)" -eq 1696 && \
|
||||
rm -f "${cubin_archive}" && \
|
||||
rm -rf "${cubin_extract_dir}"
|
||||
|
||||
# Reinstall the matching FlashInfer package trio before patching its Python
|
||||
# sources. 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.15.post1" && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-cubin==0.6.15.post1" \
|
||||
--index-url https://flashinfer.ai/whl && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-jit-cache==0.6.15.post1" \
|
||||
--index-url https://flashinfer.ai/whl/cu129 && \
|
||||
python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; 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.15.post1"
|
||||
|
||||
# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch ---
|
||||
RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/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
|
||||
|
||||
WORKDIR /sgl-workspace/sglang
|
||||
@@ -0,0 +1,125 @@
|
||||
# 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 mega-MoE SiTU patch:
|
||||
# JIT-header sentinel (activation_clamp==0.03125 -> K3 SiTU); no rebuild
|
||||
# 4. FlashInfer CuTeDSL MLA DCP patch:
|
||||
# apply the seven runtime-file diffs; exclude tests absent from the wheel
|
||||
#
|
||||
# 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
|
||||
|
||||
# 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 \
|
||||
unzip \
|
||||
wget && \
|
||||
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 kimi-k3 \
|
||||
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 mega-MoE: SiTU JIT-header patch (runtime-JIT, no rebuild) ---
|
||||
RUN python3 /sgl-workspace/sglang/docker/kimi_k3/apply_deepgemm_situ_patch.py
|
||||
|
||||
# Install the pinned FlashInfer MXFP4 MoE runner cubin pool.
|
||||
ARG 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"
|
||||
ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192"
|
||||
ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1"
|
||||
ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool"
|
||||
|
||||
RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \
|
||||
cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \
|
||||
wget --no-verbose --output-document="${cubin_archive}" \
|
||||
"${TRTLLM_GEN_MOE_CUBIN_URL}" && \
|
||||
echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \
|
||||
sha256sum --check --strict - && \
|
||||
mkdir -p "${cubin_extract_dir}" && \
|
||||
unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \
|
||||
test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
|
||||
mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \
|
||||
"${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
|
||||
test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \
|
||||
-type f -name '*.cubin' | wc -l)" -eq 1696 && \
|
||||
rm -f "${cubin_archive}" && \
|
||||
rm -rf "${cubin_extract_dir}"
|
||||
|
||||
# Reinstall the matching FlashInfer package trio before patching its Python
|
||||
# sources. 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.15.post1" && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-cubin==0.6.15.post1" \
|
||||
--index-url https://flashinfer.ai/whl && \
|
||||
python3 -m pip install --no-deps \
|
||||
"flashinfer-jit-cache==0.6.15.post1" \
|
||||
--index-url https://flashinfer.ai/whl/cu130 && \
|
||||
python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; 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.15.post1"
|
||||
|
||||
# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch ---
|
||||
RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/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
|
||||
|
||||
WORKDIR /sgl-workspace/sglang
|
||||
@@ -37,7 +37,7 @@ ENV BUILD_TRITON="0"
|
||||
ENV BUILD_LLVM="0"
|
||||
ENV BUILD_AITER_ALL="1"
|
||||
ENV BUILD_MOONCAKE="1"
|
||||
ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02"
|
||||
ENV AITER_COMMIT_DEFAULT="v0.1.19"
|
||||
|
||||
# ===============================
|
||||
# Base image 942 with rocm720 and args
|
||||
@@ -47,7 +47,7 @@ ENV BUILD_TRITON="1"
|
||||
ENV BUILD_LLVM="0"
|
||||
ENV BUILD_AITER_ALL="1"
|
||||
ENV BUILD_MOONCAKE="1"
|
||||
ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02"
|
||||
ENV AITER_COMMIT_DEFAULT="v0.1.19"
|
||||
|
||||
# ===============================
|
||||
# Base image 950 and args
|
||||
@@ -57,7 +57,7 @@ ENV BUILD_TRITON="0"
|
||||
ENV BUILD_LLVM="0"
|
||||
ENV BUILD_AITER_ALL="1"
|
||||
ENV BUILD_MOONCAKE="1"
|
||||
ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02"
|
||||
ENV AITER_COMMIT_DEFAULT="v0.1.19"
|
||||
|
||||
# ===============================
|
||||
# Base image 950 with rocm720 and args
|
||||
@@ -67,7 +67,7 @@ ENV BUILD_TRITON="1"
|
||||
ENV BUILD_LLVM="0"
|
||||
ENV BUILD_AITER_ALL="1"
|
||||
ENV BUILD_MOONCAKE="1"
|
||||
ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02"
|
||||
ENV AITER_COMMIT_DEFAULT="v0.1.19"
|
||||
|
||||
# Local source stage: with BRANCH_TYPE=local the build context is copied here and
|
||||
# used instead of git clone (mirrors docker/Dockerfile's local_src stage).
|
||||
|
||||
Reference in New Issue
Block a user