Skip to content

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

Open
jiejingzhangamd wants to merge 15 commits into
mainfrom
feat/kimi-k3-pd
Open

feat(recipes): Kimi-K3 PD on SGLang, selectable alongside the vLLM one#71
jiejingzhangamd wants to merge 15 commits into
mainfrom
feat/kimi-k3-pd

Conversation

@jiejingzhangamd

Copy link
Copy Markdown
Collaborator

Kimi-K3 PD had one shape (vLLM, MooncakeConnector). This adds the SGLang shape and makes the choice explicit, because the two disaggregate by different mechanisms rather than being two spellings of the same thing:

Engine Mechanism Manifest Image
vLLM --kv-transfer-config with a MooncakeConnector pd/deploy.yaml vllm/vllm-openai-rocm:kimi-k3
SGLang --disaggregation-mode + Mooncake bootstrap handshake pd-sglang/deploy.yaml lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727

The manual page now presents them as a choice inside the PD tab.

The image is the load-bearing detail

The general lmsysorg/sglang tags carry no Kimi-K3 support at all — no kimi_k3.py, nothing in the registry. Only that dated -k3- build has it:

/sgl-workspace/sglang/python/sglang/srt/models/
    kimi_k3.py   kimi_k3_vl.py   kimi_k25.py

class KimiK3ForConditionalGeneration
class KimiK3LinearForCausalLM

Verifying that is less obvious than it looks, and I got it wrong twice before getting it right. The image installs sglang as an editable package, so import sglang does not resolve to the source tree, and a registry probe returns an empty list — doubly so because rocminfo fails in a GPU-less container and some model imports abort. Both of my first two checks said "no Kimi support", which would have been the wrong conclusion to act on. The files are under /sgl-workspace/sglang/python/.

Weights come from NFS, and that is not a preference

Kimi-K3 is ~1.5 TB; the decode node here has 147 GB free. Both nodes seeing one path is the only option, so the manifest points at NFS rather than a node-local directory.

Validation status — stated, not implied

Neither PD combination has been run for this model. The SGLang PD path is validated cross-node (the Qwen run in the PD-fixes PR: prefill chi2800 → decode chi2866, KV over RoCE, correct answer), but that was Qwen3-0.6B at TP1 — not Kimi-K3's size or settings. The vLLM PD path is unproven on this fabric entirely.

The status table says exactly that rather than letting the model inherit a validation it did not earn.

Depends on

The SGLang PD path needs the five fixes in the cross-node PD PR — in particular the advertise-host fix, without which the router cannot pair prefill with decode, and the Mooncake HIP gate, without which the KV transfer dies in hipIpcOpenMemHandle.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Pw7JSvdQb796xh5pEcctLz

`_private_rail_ipv4()` returns the first private RFC1918 address it finds,
skipping only loopback and docker. On a Kubernetes node the CNI bridge is also
private and — sorted by name — comes first:

    cni0  docker0  enp105s0  enp121s0  enp137s0 ...

so the scan returns `cni0` and the KV host IP gets pinned to a bridge gateway the
peer cannot reach. Observed running PD across two k3s nodes: prefill advertised
10.42.0.1 and decode 10.42.1.1, each the other node's own flannel gateway, while
the routable private rail was 10.2.122.22 / 10.2.122.47.

This only bites where PD is deployed on Kubernetes with hostNetwork, which is
exactly where it must work: hostNetwork is required for the pods to reach the RDMA
rails at all, and it is what exposes the host's virtual interfaces to this scan.
Without hostNetwork the pod sees only its own veth and the bug is invisible.

Skips interfaces that are virtual by construction — CNI plugins, container
bridges, veth pairs, tunnels — rather than adding `cni` to a two-name deny list,
since the next CNI would reintroduce it. On this host the fix moves the choice
from 10.42.0.1 to 10.2.122.22, which is the rail an independent check confirmed
routable to the peer.

The regression test models the real interface layout, including the ordering that
makes the naive scan wrong, and asserts a host with only a public address yields
None rather than a public IP. Confirmed it fails against the original two-name
skip list.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
…d host

The KV plane published `args.server_args.host` as its advertise address. With the
usual `--host 0.0.0.0` that means the router is told to fetch snapshots from
`http://0.0.0.0:8801`, which it dutifully tried for every worker, forever:

    snapshot HTTP GET failed for 137.220.61.177:30000 @ http://0.0.0.0:8801
      (/models/Qwen3-0.6B/aeb13307): All connection attempts failed

Both workers were registered and healthy the whole time, so nothing pointed at
the cause.

main() already resolves args.advertise_host from POD_IP under Kubernetes
discovery for exactly this reason — its comment says "the registered url/worker_id
and kv endpoints all derive from this" — but this path never read it. The vLLM
worker has always done `args.advertise_host or args.host`; this brings SGLang in
line.

Found while getting cross-node PD working: with the endpoints unreachable the
router could not pair prefill with decode and fell back to the mixed path,
answering "no active mixed worker" to every request.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
The vendor bases ship whatever libionic they were built with, and the ionic kernel
driver only accepts a matching ABI. The stock SGLang image has 54.0-149 (ABI 1)
against a host needing 54.0-187 (ABI 4), so libibverbs rejects every device:

    libibverbs: Driver ionic does not support the kernel ABI of 1 (supports 4 to 4)

which presents as "No RDMA devices found". Mooncake then falls back to its HIP IPC
transport, and cross-node PD dies in hipIpcOpenMemHandle — three failures deep from
a cause none of them names.

The baked images defended this twice: a pinned ABI-4 deb, plus an entrypoint that
swapped in the host build. Running a stock base keeps neither, and infera-exec is
what replaced that entrypoint — so it does the injection now.

Opt-in by mounting the host build at /host-libionic/libionic.so; a missing mount is
not an error, since aggregated serving never touches RDMA. Verified on two k3s
nodes: 0 RDMA devices before, 8 after.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Harvesting Mooncake from an engine image inherits whatever that image happened to
contain — and what they contain is the base's own build. In BOTH
rocm/infera:sglang-v0.1.1 and v0.1.2:

    site-packages/mooncake/engine.so   mtime 2026-07-15 04:57:53   gate: absent
    Mooncake/build/.../engine.so       mtime 2026-07-15 04:57:53   gate: absent

Identical timestamps, straight from the base. The B.2 HIP-transport gate landed
2026-07-21 and those images were built 07-22 and 07-31, so the rebuild never took
effect in either — despite build_mooncake_sglang.sh self-verifying the gate and
exiting 1 without it, which means the step did not run at all rather than running
and failing. Why is not determined here; the build logs would say.

The consequence is not cosmetic. Without the gate Mooncake registers the HIP
(hipIpc) transport unconditionally and prefers it for GPU-to-GPU, so cross-node PD
dies in hipIpcOpenMemHandle — an IPC handle is host-local, so a peer node can never
open one. That is every published SGLang image, not just the overlay.

So the overlay compiles Mooncake itself, in its own stage, and fails the build if
the gate is not in the result. The payload becomes something we can state
properties about rather than something we inherited. Costs ~55 s, against
rebuilding a 78 GB engine image.

Measured across two k3s nodes: "HIP transport installed" 1 -> 0, and cross-node PD
went from KVTransferError to a correct answer.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Two mounts/settings without which cross-node PD cannot work, both learned by
running it:

hostNetwork, on the workers only. The RDMA rails are host interfaces; the flannel
pod network (10.42.0.0/24) cannot reach them. Without it Mooncake has no path to
the peer while the engine looks healthy. The server is a CPU-only router and does
not get it — that would take a host port for nothing.

The host libionic, so the vendor image's provider matches the ionic kernel ABI.
Mounted as a file at /host-libionic/libionic.so, which infera-exec swaps in.

Validated end to end: prefill on chi2800, decode on chi2866, KV over RDMA
(GID index 1, ionic_6), correct answer through the router.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
The recipes said PD needed "two nodes on a mutually routable RoCE fabric" and left
it there, which read as a prerequisite nobody here met. It was met all along — the
check was wrong.

This fabric is routed L3 RoCEv2: every NIC gets its own /64, so no two hosts share
a subnet by design. An unbound ping6 picks a default source address, leaves via the
wrong rail, and reports "No route" — which is what produced the standing conclusion
that PD could not run on these nodes. Source-bound to the matching local rail, all
8 rails answer at ~0.1 ms. Mooncake binds its QP to a device and GID index, so it
takes the path the naive ping does not.

Documents the bound-ping check, since getting this wrong costs a "we can't do PD
here" that survives for a long time.

The status table now says the PD path is validated cross-node (SGLang + Mooncake
over RoCE) while being explicit that GLM-5.2's own settings have not been run that
way — those are different claims and were being collapsed into one.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Kimi-K3 PD had one shape (vLLM, MooncakeConnector). Adds the SGLang shape and
makes the choice explicit in the manual, because the two disaggregate by different
mechanisms rather than being two spellings of one thing:

    vLLM     --kv-transfer-config with a MooncakeConnector
    SGLang   --disaggregation-mode + a Mooncake bootstrap handshake

The image matters more than usual here. The general `lmsysorg/sglang` tags carry
no Kimi-K3 support at all — no kimi_k3.py, nothing in the registry — while
`lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727` has KimiK3ForConditionalGeneration,
KimiK3LinearForCausalLM and kimi_k3_vl. Checking that is less obvious than it
sounds: the image installs sglang as an editable package, so `import sglang` does
not resolve to the source tree and a registry probe reports an empty list, doubly
so because rocminfo fails in a GPU-less container. The files are under
/sgl-workspace/sglang/python/sglang/srt/models/.

Weights come from NFS rather than a node-local path: Kimi-K3 is ~1.5 TB and the
decode node here has 147 GB free, so both nodes seeing one path is the only
option, not a preference.

Neither PD combination has been run for this model yet — the SGLang PD path itself
is validated cross-node (see the Qwen run in the PD fixes), but not with Kimi-K3's
size or settings, and the vLLM PD path is unproven on this fabric entirely. The
status table says so rather than implying the model inherits that validation.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
infera's fp8 KV-cache default assigned `server_args.kv_cache_dtype` after
`ServerArgs.from_cli_args()`. Newer SGLang freezes server_args once resolved:

    AttributeError: server_args.kv_cache_dtype assigned after resolution;
    server_args is read-only -- use get_context().override(source, ...)

which crash-looped every worker on lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727,
the only SGLang build carrying Kimi-K3.

Set it on the parsed args before resolution instead. That needs no override API
and works on both the old and the new SGLang, where calling get_context().override
would only work on the new one.

Also carries the two mounts the PD recipe needs and had lost by being derived from
main rather than from the PD-fixes branch: hostNetwork on the workers, and the host
libionic. Without them the recipe reproduces the exact failure that branch fixes —
0 RDMA devices, then Mooncake down its HIP-IPC path. Verified in this run that the
injection fires: "infera-exec: replaced .../libionic.so.1.0.54.0-149 with the host
libionic build".

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Both workers hardcoded a 1800 s wait for the engine's /health. That is generous
for local NVMe and impossible for NFS: Kimi-K3 read its 96 shards in 502 s from
local disk and ~95 min from NFS with both PD nodes competing for the same mount.

The failure is worse than a slow start. The worker times out mid-load, exits 1,
restarts, and begins the load again from zero — so it never finishes, and the pod
looks like a crash loop rather than "the storage is too slow for this deadline":

    TimeoutError: SGLang not ready after 1800s

Load time tracks the storage, not the model, so this cannot have one right value
baked in. INFERA_ENGINE_READY_TIMEOUT now sets it, on both the SGLang and vLLM
workers, defaulting to the previous 1800 and falling back to it on a malformed
value rather than raising during startup.

The Kimi-K3 PD recipe sets 7200, with the measured numbers next to it so the next
person can tell whether their storage needs more.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Both PD manifests were derived from ancestors that predate what running PD
actually taught us, and each inherited the same gaps. Bringing them level:

  hostNetwork on the workers      the RDMA rails are host interfaces; the flannel
                                  pod network cannot reach them
  host libionic mount             the vendor base's provider must match the host
                                  ionic kernel ABI or libibverbs sees no devices
  /dev/infiniband + privileged    RDMA needs the device nodes and pinned memory
  INFERA_ENGINE_READY_TIMEOUT     1800 s is impossible on NFS-backed weights

The SGLang manifest additionally carries what its own bring-up cost: memory 96Gi
(a Qwen-0.6B number) against 1.5 TB of shards, a mamba state cache sized from a
negative residual, and the mem-fraction trap documented inline.

None of this is speculative — every line replaces a failure whose message pointed
somewhere other than its cause: "0 RDMA devices" for an ABI mismatch, a crash loop
for a timeout, a bare "-9" for two different problems in turn.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
The overlay could ship code the vendor base lacks, but not fix code it has. That
gap became blocking: Kimi-K3 PD on vLLM dies during KV registration because the
Mooncake connector destructures every Mamba-family layer as a 2-tuple —

    if isinstance(layer_spec, MambaSpec):
        conv, _ = cache_or_caches

which is Mamba2's shape. Kimi-K3 uses KDA linear attention with a different
number of state tensors, so every rank raises "too many values to unpack".
MambaSpec.shapes is already a variable-length tuple, so the arity assumption is
the connector's, not the spec's.

The obvious fix — bake the patch into Dockerfile.vllm — reintroduces exactly the
fork this overlay exists to avoid, and for a 42 GB image. So the payload now
carries deploy/docker/patches/vllm/ and infera-exec applies them before exec.
Each script locates its target through the installed package, is idempotent, and
no-ops when its anchor is gone, so a base that already carries a fix is untouched
and the patches survive a base bump.

Guarded rather than unconditional: skipped when vllm is not importable, skipped
by INFERA_SKIP_ENGINE_PATCHES, and a failing patch warns instead of aborting —
a mis-anchored patch must not take down a worker that would otherwise serve.

Verified in a stock vllm/vllm-openai-rocm:kimi-k3: the mamba unpack is applied,
the old line is gone, and the pre-existing blocksize/prom patches apply too — so
this also closes the gap where an overlay deployment silently lacked fixes the
baked image had.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
…s known

infera resolves the Pod IP under Kubernetes discovery and logs it —

    k8s discovery: advertising Pod IP 137.220.61.177

— but disagg_meta, which carries Mooncake's bootstrap_addr, was already built
during arg parsing, while advertise_host was still unset. So the address froze as
the bind host and the prefiller published:

    bootstrap_addr: 'http://0.0.0.0:8998'

The decoder reads that from kv_transfer_params (vLLM's own comment says decoders
must) and cannot connect:

    Failed to connect to bootstrap server http://0.0.0.0:8998
    Failed to find remote engine_id ...

Both workers stay healthy and registered; the request simply hangs until the
connector's 480 s timeout. Nothing in that path names the address.

This is the same defect as the SGLang KV-endpoint one, in a second place: an
address is derived before the host that makes it routable is known. Rebuild it
after the resolution, touching only the 0.0.0.0 form so an explicitly configured
address is left alone.

The manifest also gains --data-parallel-address $(POD_IP): that fixes the OTHER
half, where the prefiller computes its own address via
get_mooncake_bootstrap_addr() -> parallel_config.data_parallel_master_ip. Both
paths have to be right; fixing either alone leaves PD hanging.

With both, Kimi-K3 PD answers across two nodes: prefill on chi2800, decode on
chi2866, 1.5 TB MXFP4 at TP8 from local NVMe, "The capital of France is Paris."

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>

# Conflicts:
#	deploy/overlay/infera-exec
v0.2.1 predates everything cross-node PD needed. Comparing the published image
against the current build:

                        v0.2.1   v0.2.2
    libionic injection      no      yes
    runtime engine patches  no      yes   (8 patch scripts)
    native trees            2       2

So anyone following a recipe today gets an overlay that cannot bring up PD: no
libionic swap means libibverbs rejects every ionic device, and no runtime patches
means the Mooncake connector still destructures Kimi-K3's Mamba layers as a
2-tuple. Both were fixed in this branch and neither was reachable.

inferaimage/infera-overlay:v0.2.2 is public and pulls with no credentials —
verified from a clean DOCKER_CONFIG, since "we pushed it" and "a reader can get
it" have already turned out to be different things twice in this repo.

Validated with this image: Kimi-K3 PD across chi2800 and chi2866, 1.5 TB MXFP4 at
TP8 from local NVMe, ~6 min to both workers registered, 8 RDMA devices visible to
libibverbs on each side, and a correct answer through the router.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Two problems, one fix.

The recipes still carried 20 references to rocm/infera-overlay:latest — the name
I built under locally while writing #64, which 404s on Docker Hub. #67 fixed those
on a branch cut from main; this branch predates that merge and had them back. A
name that only exists on one laptop keeps coming back because a tag is a promise
nobody checks.

And the tags that do resolve are mutable. vllm/vllm-openai-rocm:kimi-k3 was
rebuilt on 2026-07-27 during this very session, which is what invalidated the
"prefix caching fails engine init" note: the tag was the same and the image was
not. A recipe that says "this was validated" has to name the thing that was
validated.

So every image is now pinned by digest:

    inferaimage/infera-overlay          @6918eff3  (v0.2.2, this session's fixes)
    vllm/vllm-openai-rocm:kimi-k3       @5aa7e626
    lmsysorg/sglang-rocm:...k3-20260727 @3c01f73f
    lmsysorg/sglang:v0.5.15...mi35x     @40e940a0

Zero mutable tags remain under examples/recipes/. Each digest is one I pulled and
ran, not one read off a page. All nine manifests still pass
`kubectl apply --dry-run=server` against the live CRD.

The readable tag is kept alongside the digest so the ref still says what it is.

Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant