minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

Replace Kimi K3 DeepGEMM patch with 0.1.5.post1 (#33143)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-07-31 22:17:46 -07:00
committed by GitHub
parent 0d186f49be
commit a1344fad4e
3 changed files with 18 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/kimi_k3/apply_deepgemm_situ_patch.py
-140
View File
@@ -1,140 +0,0 @@
"""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"
)
docker/kimi_k3/kimi_k3_cu12.Dockerfile
+9 -4
View File
@@ -9,8 +9,8 @@
# 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
# 3. DeepGEMM upgrade to 0.1.5.post1:
# official MegaMoE runtime-JIT header with Kimi-K3 SiTU support
# 4. FlashInfer CuTeDSL MLA DCP patch:
# apply the seven runtime-file diffs; exclude tests absent from the wheel
#
@@ -25,6 +25,8 @@
FROM lmsysorg/sglang:v0.5.16-cu129 AS base
ARG SGL_DEEP_GEMM_VERSION="0.1.5.post1"
# Current Kimi-K3 source auto-discovers and builds its PyO3 extensions.
ARG RUST_VERSION="1.90.0"
ENV RUSTUP_HOME="/usr/local/rustup" \
@@ -79,8 +81,11 @@ 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
# --- 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"
# 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"
docker/kimi_k3/kimi_k3_cu13.Dockerfile
+9 -5
View File
@@ -4,14 +4,14 @@
# (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:
# 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,
# 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
# 3. DeepGEMM upgrade to 0.1.5.post1:
# official MegaMoE runtime-JIT header with Kimi-K3 SiTU support
# 4. FlashInfer CuTeDSL MLA DCP patch:
# apply the seven runtime-file diffs; exclude tests absent from the wheel
#
@@ -26,6 +26,8 @@
FROM lmsysorg/sglang:v0.5.16 AS base
ARG SGL_DEEP_GEMM_VERSION="0.1.5.post1"
# Current Kimi-K3 source auto-discovers and builds its PyO3 extensions.
ARG RUST_VERSION="1.90.0"
ENV RUSTUP_HOME="/usr/local/rustup" \
@@ -70,8 +72,10 @@ 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
# --- 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}"
# 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"