feat(recipes): Kimi-K3 PD on SGLang, selectable alongside the vLLM one

deploy/docker/patches/vllm/patch_mooncake_mamba_unpack.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+"""Mooncake connector: don't assume a Mamba layer has exactly two state tensors.
+
+register_kv_caches() destructures every MambaSpec layer as a 2-tuple:
+
+    if isinstance(layer_spec, MambaSpec):
+        conv, _ = cache_or_caches
+        cache_list = [conv]
+
+That holds for Mamba2 (conv state, SSM state) and fails for anything else.
+Kimi-K3 uses KDA linear attention, whose layers carry a different number of state
+tensors, so PD dies during KV registration for every rank at once:
+
+    mooncake_connector.py:1678 in register_kv_caches
+      conv, _ = cache_or_caches
+    ValueError: too many values to unpack (expected 2)
+
+MambaSpec.shapes is already a variable-length tuple of shapes, so the two-tensor
+assumption is the connector's, not the spec's. Take the first state tensor —
+which is what the original code kept — without constraining how many follow.
+
+Self-locating and idempotent: re-running is a no-op, and it no-ops if upstream
+fixes this, so the patch can stay in place across a base bump.
+"""
+
+import sys
+from pathlib import Path
+
+OLD = """            if isinstance(layer_spec, MambaSpec):
+                conv, _ = cache_or_caches
+                cache_list = [conv]"""
+
+NEW = """            if isinstance(layer_spec, MambaSpec):
+                # A Mamba-family layer does not necessarily carry exactly two
+                # state tensors: that is Mamba2's shape (conv, ssm). KDA linear
+                # attention (Kimi-K3) carries a different number, and `conv, _ =`
+                # then raises "too many values to unpack" on every rank. Keep the
+                # first state tensor, as before, without fixing the arity.
+                cache_list = [cache_or_caches[0]]"""
+
+
+def main() -> int:
+    import vllm
+
+    target = (
+        Path(vllm.__file__).parent
+        / "distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
+    )
+    if not target.is_file():
+        print(f"[patch] {target} not found; skipping")
+        return 0
+    src = target.read_text()
+    if NEW in src:
+        print("[patch] mooncake mamba unpack: already applied")
+        return 0
+    if OLD not in src:
+        print("[patch] mooncake mamba unpack: anchor absent (upstream changed?); skipping")
+        return 0
+    target.write_text(src.replace(OLD, NEW))
+    print(f"[patch] mooncake mamba unpack: applied to {target}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

deploy/overlay/Dockerfile.payload

@@ -112,5 +112,11 @@ COPY --from=deps312 /payload/py312 /payload/py312
 COPY --from=native312 /native /payload/native
 COPY --from=native310 /native /payload/native
 COPY --from=native312 /usr/local/bin/infera-router /payload/bin/infera-router
+# Engine patches applied at container start (see infera-exec). These fix vendor
+# code we do not otherwise ship — the overlay's whole point is not forking the
+# vendor image, and a runtime patch keeps that true where a rebuild would not.
+# Each script is self-locating, idempotent, and no-ops once upstream carries the
+# fix, so they are safe to leave in place across a base bump.
+COPY deploy/docker/patches/vllm/ /payload/patches/vllm/
 COPY deploy/overlay/infera-exec /payload/bin/infera-exec
 CMD ["sh", "-c", "cp -a /payload/. /out/ && echo 'infera overlay payload installed to /out'"]

deploy/overlay/infera-exec

@@ -110,6 +110,31 @@ if [ -e /host-libionic/libionic.so ]; then
     fi
 fi
 
+# --- vendor engine patches --------------------------------------------------
+# Fixes to the vendor's own code, applied here rather than baked into a forked
+# image — forking is exactly what this overlay exists to avoid. Each script
+# locates its target through the installed package, is idempotent, and no-ops
+# when the anchor is absent, so a base that already carries the fix is untouched
+# and the patch can stay in place across bumps.
+#
+# Currently: the Mooncake connector destructures every Mamba-family layer as a
+# 2-tuple, which is Mamba2's shape. Kimi-K3's KDA linear attention carries a
+# different number of state tensors, so PD died on every rank with "too many
+# values to unpack" during KV registration.
+#
+# Skipped entirely when INFERA_SKIP_ENGINE_PATCHES is set, and a failing patch
+# warns rather than blocking startup — a mis-anchored patch must not take down a
+# worker that would otherwise serve.
+if [ -d "$PAYLOAD_ROOT/patches/vllm" ] && [ -z "${INFERA_SKIP_ENGINE_PATCHES:-}" ]; then
+    if "$PY" -c 'import vllm' >/dev/null 2>&1; then
+        for _p in "$PAYLOAD_ROOT"/patches/vllm/*.py; do
+            [ -e "$_p" ] || continue
+            (cd / && PYTHONPATH="$PYTREE${PYTHONPATH:+:$PYTHONPATH}" "$PY" "$_p") \
+                || echo "infera-exec: patch $(basename "$_p") failed; continuing" >&2
+        done
+    fi
+fi
+
 # Python puts the working directory at sys.path[0] — ahead of PYTHONPATH. Vendor
 # images that already ship infera set WORKDIR to its parent (e.g. /opt/infera),
 # so the baked-in copy would silently shadow this payload and the overlay would

examples/k8s-deployments/glm5.2-sglang.yaml

@@ -29,7 +29,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -64,7 +64,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/kimi-k3-vllm.yaml

@@ -34,7 +34,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -69,7 +69,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/mixed-kvd-vllm.yaml

@@ -58,7 +58,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -91,7 +91,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/mixed-kvd.yaml

@@ -65,7 +65,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -98,7 +98,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/pd-1p1d-mooncake.yaml

@@ -37,7 +37,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -75,7 +75,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -132,7 +132,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/pd-kvd.yaml

@@ -81,7 +81,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -113,7 +113,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -180,7 +180,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/single-node-qwen-sglang.yaml

@@ -25,7 +25,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -60,7 +60,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/k8s-deployments/single-node-qwen-vllm.yaml

@@ -25,7 +25,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
@@ -60,7 +60,7 @@ spec:
         # is an image-tag edit rather than a rebuild of ours.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay:v0.2.2
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:

examples/recipes/README.md

@@ -42,7 +42,7 @@ forking the base for one model broke every other model.
 Build the overlay before deploying:
 
 ```bash
-docker build -f deploy/overlay/Dockerfile.payload -t inferaimage/infera-overlay:v0.2.1 .
+docker build -f deploy/overlay/Dockerfile.payload -t inferaimage/infera-overlay:v0.2.2 .
 ```
 
 The build harvests **one native tree per ABI family** — `NATIVE_IMAGE` supplies

examples/recipes/glm5.2/mixed-kvd/deploy.yaml

@@ -1,7 +1,7 @@
 # GLM-5.2-MXFP4 — mixed-kvd — Kubernetes recipe
 #
 #   base image   lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x      (stock vendor image, unmodified)
-#   overlay      inferaimage/infera-overlay:v0.2.1                     (infera + router + kvd + native deps)
+#   overlay      inferaimage/infera-overlay@sha256:6918eff34f201548a738dd592d2a1ece0627354d2e88f24a87cfa8f787a72a44                     (infera + router + kvd + native deps)
 #   engine       infera.engine.sglang   TP8
 #
 # Deploy:  kubectl apply -f examples/recipes/glm5.2/mixed-kvd/deploy.yaml
@@ -45,14 +45,14 @@ spec:
         # so following an upstream bump costs nothing in this repo.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay@sha256:6918eff34f201548a738dd592d2a1ece0627354d2e88f24a87cfa8f787a72a44
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
               - {name: overlay, mountPath: /overlay}
         containers:
           - name: main
-            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x
+            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x@sha256:40e940a0c55b87105c773d8b484616616b3a91662bfa223c48ff721d9793dc8d
             imagePullPolicy: IfNotPresent
             command: ["/overlay/bin/infera-exec","python3","-m","infera.server",
                       "--host","0.0.0.0","--port","8000",
@@ -79,14 +79,14 @@ spec:
         # so following an upstream bump costs nothing in this repo.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay@sha256:6918eff34f201548a738dd592d2a1ece0627354d2e88f24a87cfa8f787a72a44
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
               - {name: overlay, mountPath: /overlay}
         containers:
           - name: main
-            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x
+            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x@sha256:40e940a0c55b87105c773d8b484616616b3a91662bfa223c48ff721d9793dc8d
             imagePullPolicy: IfNotPresent
             command: ["/overlay/bin/infera-exec",
                       "python3","-m","infera.engine.sglang","--host","0.0.0.0","--port","30000",
@@ -136,7 +136,7 @@ spec:
           # infera.kvd out of the shared overlay. The overlay image itself is a
           # busybox carrying the payload — it has no interpreter to run kvd with.
           - name: kvd
-            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x
+            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x@sha256:40e940a0c55b87105c773d8b484616616b3a91662bfa223c48ff721d9793dc8d
             imagePullPolicy: IfNotPresent
             command: ["/overlay/bin/infera-exec","python3","-m","infera.kvd",
                       "--socket","/kvd/kvd.sock",

examples/recipes/glm5.2/mixed/deploy.yaml

@@ -1,7 +1,7 @@
 # GLM-5.2-MXFP4 — mixed — Kubernetes recipe
 #
 #   base image   lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x      (stock vendor image, unmodified)
-#   overlay      inferaimage/infera-overlay:v0.2.1                     (infera + router + kvd + native deps)
+#   overlay      inferaimage/infera-overlay@sha256:6918eff34f201548a738dd592d2a1ece0627354d2e88f24a87cfa8f787a72a44                     (infera + router + kvd + native deps)
 #   engine       infera.engine.sglang   TP8
 #
 # Deploy:  kubectl apply -f examples/recipes/glm5.2/mixed/deploy.yaml
@@ -34,14 +34,14 @@ spec:
         # so following an upstream bump costs nothing in this repo.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay@sha256:6918eff34f201548a738dd592d2a1ece0627354d2e88f24a87cfa8f787a72a44
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
               - {name: overlay, mountPath: /overlay}
         containers:
           - name: main
-            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x
+            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x@sha256:40e940a0c55b87105c773d8b484616616b3a91662bfa223c48ff721d9793dc8d
             imagePullPolicy: IfNotPresent
             command: ["/overlay/bin/infera-exec","python3","-m","infera.server",
                       "--host","0.0.0.0","--port","8000",
@@ -68,14 +68,14 @@ spec:
         # so following an upstream bump costs nothing in this repo.
         initContainers:
           - name: infera-overlay
-            image: inferaimage/infera-overlay:v0.2.1
+            image: inferaimage/infera-overlay@sha256:6918eff34f201548a738dd592d2a1ece0627354d2e88f24a87cfa8f787a72a44
             imagePullPolicy: IfNotPresent
             command: ["sh","-c","cp -a /payload/. /overlay/"]
             volumeMounts:
               - {name: overlay, mountPath: /overlay}
         containers:
           - name: main
-            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x
+            image: lmsysorg/sglang:v0.5.15.post1-rocm720-mi35x@sha256:40e940a0c55b87105c773d8b484616616b3a91662bfa223c48ff721d9793dc8d
             imagePullPolicy: IfNotPresent
             command: ["/overlay/bin/infera-exec",
                       "python3","-m","infera.engine.sglang","--host","0.0.0.0","--port","30000",