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

docker/kimi_k3/apply_deepgemm_situ_patch.py

@@ -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,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:
+#   3. DeepGEMM upgrade to 0.1.5.post1:
-#        JIT-header sentinel (activation_clamp==0.03125 -> K3 SiTU); no rebuild
+#        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) ---
+# --- 3. DeepGEMM: upgrade to the first release with Kimi-K3 SiTU ---
-RUN python3 /sgl-workspace/sglang/docker/kimi_k3/apply_deepgemm_situ_patch.py
+# 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

@@ -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:
+#   3. DeepGEMM upgrade to 0.1.5.post1:
-#        JIT-header sentinel (activation_clamp==0.03125 -> K3 SiTU); no rebuild
+#        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) ---
+# --- 3. DeepGEMM: upgrade to the first release with Kimi-K3 SiTU ---
-RUN python3 /sgl-workspace/sglang/docker/kimi_k3/apply_deepgemm_situ_patch.py
+# 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"