Skip to content

fix(sglang): get GLM-5.2-FP8 PD + DPA + MTP + kv-aware running on MI325X - #56

Merged
jiejingzhangamd merged 8 commits into
mainfrom
llying/glm5p2_fp8_fixes
Aug 3, 2026
Merged

fix(sglang): get GLM-5.2-FP8 PD + DPA + MTP + kv-aware running on MI325X#56
jiejingzhangamd merged 8 commits into
mainfrom
llying/glm5p2_fp8_fixes

Conversation

@llying-001

@llying-001 llying-001 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Description

GLM-5.2-FP8 1P1D on MI325X (gfx942) would not start on our engine image, and once it did,
long prompts came back quietly wrong. This fixes the image foundation and the four
engine/router bugs the bring-up exposed. Validated on 2× 8×MI325X, ROCm 7.2.0, sglang
v0.5.16, GLM-5.2-FP8 1P1D over mooncake RDMA.

The three image commits come first on purpose: everything below is only validated on that base.

The KV fix is not gfx942-specific. Any PD deployment doing chunked prefill over the
mooncake transport with overlap scheduling has been transferring pages while the forward that
writes them is still running. It never raises — prompts longer than one prefill chunk are just
partially wrong, the corruption boundary lands exactly on the chunk boundary, and the same
model/backends/chunk size in an aggregated server passes. Same symptom as
sgl-project/sglang#25583, which was auto-closed with no follow-up.

Important

Running GLM-5.2 PD + DP-attention + MTP on the gfx942 image requires every leg to launch
with --json-model-override-args '{"index_share_for_mtp_iteration":false}'.
That is not a
tuning knob: it stands in for sglang_dsa patches 02b and 04, which are --fuzz=0 diffs
against v0.5.15.post1 and cannot apply to this image's v0.5.16 base. Leave IndexShare on and
the decode leg deadlocks on the first request. Every 1P1D run behind this branch set it. The
requirement is stated in the Dockerfile header, at the patch block, in
apply_sglang_dsa_patches.sh and in both patch READMEs, so a deployment inherits it from the
image rather than from this description.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Image — land first:

  • Dockerfile.sglang.gfx942 base → lmsysorg/sglang:v0.5.16-rocm720-mi30x. GLM-5.2 MTP
    needs sglang #30839 plus GLM-5.2's own GlmMoeDsaForCausalLMNextN quant resolution. This
    image deliberately does not copy patches/sglang/: v0.5.16 fixes GLM-5.2's nextn quark
    exclude upstream, while that backport's anchor survives the refactor inside DeepSeek's own
    nextn path, so applying it here would rewrite a probe that is already correct.
  • gfx942 builds Mooncake with the HIP-transport gate, i.e. runs the rebuild
    Dockerfile.sglang already did. The bundled Mooncake (upstream #2682, 01d1eb2a) prefers
    HIP IPC over RDMA, so cross-node PD dies on the first request in hipIpcOpenMemHandle and
    every bring-up had to rebuild engine.so inside the running container first. Same script and
    diffs; only the hardcoded HIP arch changed, into MC_GPU_ARCH (default gfx950, so
    Dockerfile.sglang is unaffected).
  • patches/sglang_disagg/patch_mooncake_early_send_wait_event.py, applied at build time by
    both engine images. sglang's prefill.py already records a completion event as the barrier,
    but only mori ever read it: mooncake/conn.py has no wait_event/synchronize() at all,
    and the overlap path that moves non-final chunks recorded none. The patch mirrors mori.
  • sglang_dsa patch 01 now reaches gfx942. main added that set while this branch was
    out, baked only by Dockerfile.sglang; without patch 01 the aiter (HIP) paged-MQA branch
    sizes its logits from DP-padded rows while lengths carries the real count, so
    Expected lengths.size(0) == B fires at concurrency > 1 — ROCm-specific and not substituted
    by anything at runtime. Patch 01 is re-cut from a pinned diff into a self-locating idempotent
    script (its two edit sites are byte-identical on v0.5.15.post1 and v0.5.16; the diff is
    deleted rather than kept as a second source of truth), and
    apply_sglang_dsa_patches.sh grows a DSA_PATCH_SET arm — full is byte-for-byte the old
    gfx950 behaviour, indexer is the gfx942 arm. 02a is neither carried nor substituted there:
    not re-cut for v0.5.16, and not observed on that base, which is not the same as fixed.

Engine / router:

  • dsv4 detection (infera/engine/dsv4_gfx942.py): index_topk alone meant DeepSeek-V4, so
    GLM-5.2 (glm_moe_dsa, index_topk 2048) silently got the dsv4 knobs. Trusts
    model_type/architectures first, falls back to index_topk only for a checkpoint that
    names no architecture.
  • kv-event bigram tokens (infera/router/kv_event/, rust/router/src/kv_event.rs): under
    MTP the radix key is a bigram view, so tokens arrive as overlapping (t[i], t[i+1]) pairs —
    Python hashed the pairs, Rust's int-only decode dropped them, and the cache view stayed empty
    so prefix hits never fired. Both sides now take the first element of each pair, rebuilding
    the flat slice the query side chunks.
  • NodePort-safe port allocation (infera/common/net.py): the scan could land in
    Kubernetes' NodePort range, where the port binds and works over loopback but is
    connection-refused on the node IP we advertise. Both the per-DP-rank block and the single
    port skip that window; INFERA_NODEPORT_RANGE overrides it, documented in
    manual/reference/environment.md.

Verification

  • needle 5/9 → 9/9, 29k depth sweep 4/9 → 9/9, full correctness suite passes — with the
    prompt still really split into 4 chunks, so this is a barrier and not chunked prefill quietly
    turning itself off. Retires the CHUNK=524288 + MEM_FRAC=0.80 workaround.
  • Mooncake wait-event patch script run against real sglang v0.5.16 and v0.5.15.post1
    source: byte-identical to the diff it replaces on both, patched files compile, re-run is a
    no-op. On drift it exits 1 having written nothing (a half-applied handover corrupts just the
    same) and its build loop uses set -eu, so a drifted anchor fails the build instead of
    shipping a silently broken image.
  • DSA indexer patch 01, same treatment against both real sources: reproduces the deleted diff
    byte-for-byte on v0.5.15.post1, applies and compiles on v0.5.16, no-op on re-run,
    _p1v2_trim present in the bytecode on both, and exit 1 writing nothing on a drifted anchor.
    The full arm was re-run end-to-end on a v0.5.15.post1 tree — all 8 bytecode markers, the
    nextn prerequisite and patch 2a's source marker still pass, so gfx950 is unchanged.
  • Regression tests, each verified to fail without its fix: bigram BlockStored e2e (Python)
    and decodes_sglang_bigram_batch_under_mtp (Rust), NodePort avoidance for both the block and
    the single port, dsv4 detection. Full unit suite: 1183 passed.
  • Not verified here: the two docker builds (no docker/GPU on the dev box). Only the gfx942
    Mooncake compile needs one, and that rebuild already succeeded against this exact image at
    runtime during bring-up.
  • Not established: anything about sglang_dsa 02a on gfx942, and the gfx942 arm generally is
    FP8 / dp8 / conc ≤ 128 / tilelang DSA backends — a different quantization, base and backend
    selection from the gfx950 validation. Tracked for a follow-up PR.

Checklist:

  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

Copilot AI review requested due to automatic review settings July 31, 2026 09:00

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR extends Infera’s gfx942 workflow support for GLM-5.2 FP8 with PD+DPA+MTP+kv-aware routing by hardening KV-event ingestion (including SGLang MTP bigram token views), refining model-family detection for gfx942 tuning, and improving engine networking/startup ergonomics.

Changes:

  • Add bigram-aware token flattening for SGLang KV events (Python + Rust) and an end-to-end regression test to ensure kv-aware routing hits under MTP/EAGLE.
  • Avoid allocating KV-event port blocks in the Kubernetes NodePort range (with env override + unit tests) and document related environment knobs.
  • Refine DeepSeek-V4 detection so other sparse-attention families (e.g., GLM-5.2 DSA) don’t receive dsv4 gfx942-specific knobs.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
tests/unit/router/test_kv_event_e2e.py Adds an e2e regression test covering bigram (MTP/EAGLE) KV-event token shape vs flat request tokens.
tests/unit/common/test_net_ports.py Adds unit tests ensuring allocated port blocks avoid NodePort range and remain bindable/contiguous.
tests/engine/test_dsv4_gfx942.py Adds test ensuring non-dsv4 sparse-attention families aren’t misdetected as dsv4.
rust/router/src/kv_event.rs Decodes SGLang bigram token_ids by extracting first element; adds Rust regression test.
manual/reference/environment.md Documents INFERA_SGLANG_READY_TIMEOUT and INFERA_NODEPORT_RANGE.
infera/router/kv_event/events.py Updates SGLangBlockStored.token_ids typing/docs to reflect int vs bigram-pair shapes.
infera/router/kv_event/client.py Introduces token flattening for bigram token_ids before hashing into cache view.
infera/engine/sglang/worker.py Makes SGLang readiness timeout configurable via INFERA_SGLANG_READY_TIMEOUT.
infera/engine/dsv4_gfx942.py Tightens dsv4 detection logic to avoid false positives from index_topk alone.
infera/common/net.py Skips Kubernetes NodePort range when allocating contiguous TCP port blocks; adds env override parsing + warning.
deploy/docker/patches/sglang_disagg/README.md Documents a correctness patch for chunked prefill KV transfer over mooncake PD transport.
deploy/docker/patches/sglang_disagg/mooncake_early_send_wait_event.diff Adds CUDA-event handoff and synchronization to prevent KV corruption in mooncake chunk transfer.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread infera/common/net.py
Comment thread infera/router/kv_event/client.py Outdated
@llying-001 llying-001 changed the title support glm5.2-fp8 pd+dpa+mtp+kv-aware workflow on gfx942 fix(sglang): GLM-5.2 FP8 1P1D on MI325X — v0.5.16 image base, mooncake KV barrier, and four engine/router fixes Jul 31, 2026
@llying-001 llying-001 changed the title fix(sglang): GLM-5.2 FP8 1P1D on MI325X — v0.5.16 image base, mooncake KV barrier, and four engine/router fixes fix(sglang): get GLM-5.2-FP8 PD + DPA + MTP + kv-aware running on MI325X Jul 31, 2026
@jiejingzhangamd

Copy link
Copy Markdown
Collaborator

check review and conflict, and document already re-org with model card receipt.

llying-001 and others added 7 commits August 3, 2026 03:30
Chunked prefill over the mooncake PD transport hands pages to the decode leg
while the forward that writes them is still running, so every chunk but the last
can be RDMA-read half-written. It does not crash: prompts longer than one prefill
chunk come back *partially* wrong, and on GLM-5.2 DSA a needle read returns the
first digits and then repeats `</think>`. The boundary lands exactly on the chunk
boundary, and an aggregated server with the same model/backends/chunk size passes.

sglang's `prefill.py` already records a completion event as the barrier for this,
but only the `mori` backend ever read it — `mooncake/conn.py` had no `wait_event`
or `synchronize()` at all, and the overlap-scheduling path that moves non-final
chunks never recorded one. The diff mirrors what `mori` does: carry the event on
`TransferKVChunk`, `synchronize()` in mooncake's transfer worker before reading
device memory, and record the event for the overlap non-final-chunk send.

Not DSA-specific — any PD deployment running chunked prefill over mooncake with
overlap scheduling is affected; DSA's sparse retrieval only makes it conspicuous.
Not yet wired into an image build: the diff is pinned to v0.5.16 source and the
engine bases we ship are older, so it is carried here (with the evidence and the
apply recipe) for runtime use until the base bump.

Verified on 2x 8xMI325X (gfx942), ROCm 7.2.0, sglang v0.5.16, GLM-5.2-FP8 1P1D
over mooncake RDMA: needle 5/9 -> 9/9 and a 29k depth sweep 4/9 -> 9/9, with the
prompt still really split into 4 chunks after the fix.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
`detect_dsv4` treated any config carrying `index_topk` as DeepSeek-V4, but that
key only says "sparse attention" — GLM-5.2 is `glm_moe_dsa` with `index_topk`
2048, so launching it on gfx942 silently picked up the dsv4 knobs that are not
its own: the dsv4 attention backend, shared-experts fusion disabled, and the
FlashMLA workaround.

Trust what the checkpoint calls itself instead: `model_type` prefix, else the
declared architecture, and fall back to `index_topk` only for a checkpoint that
names no architecture at all (a re-exported dsv4), which is the case the existing
`index_topk` test covers.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
With EAGLE/MTP on, SGLang keys its radix tree on bigrams, so each stored page
reports its tokens as the overlapping pairs `(t[i], t[i+1])` instead of bare ints.
Both routers hashed that as-is: the Python client fed tuples to the chunk hasher,
and the Rust decoder dropped every element because a pair is not an integer. The
result is a cache view that no request can match — `cache_hits` is pinned at 0 and
kv-aware routing degrades to load balancing on exactly the deployments that need
it most, silently, since nothing in the pipeline errors.

Take the first element of each pair: that rebuilds `t[start:end]`, the same flat
slice the query side chunks, and radix nodes split on page boundaries so the two
chunkings stay aligned. Widen the SGLang `BlockStored` schema to the union it
really is, and cover the shape in both test suites (the Python e2e publish path
and the Rust wire decoder), since the failure mode is silent.

Verified on GLM-5.2-FP8 1P1D with DP-attention and MTP: the dp0 cache view goes
from permanently empty to 333 blocks for a 21k-token prompt, and a repeat of the
same prompt reports a full prefix hit.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
`free_tcp_port_block` scans down from just below the ephemeral range, which on a
default host means it hands back a block inside 30000-32767 — the Kubernetes
NodePort range. Any Service in the cluster can claim a port in that window at any
time, and kube-proxy in IPVS mode then owns it on *every* node address with no
real server behind it here, so traffic to the node IP is swallowed before it
reaches our listener and the kernel answers RST.

The bind probe cannot see this: the port is free, it binds, and loopback works.
Only a peer dialing the address we advertise breaks, which is exactly the router
subscribing to a worker's kv-events — so the router's cache view stays empty and
kv-aware routing quietly degrades to load balancing.

Skip the window during the scan, overridable with `$INFERA_NODEPORT_RANGE` for a
cluster that moved the range (`none` drops the guard); a malformed value keeps the
default rather than dropping it.

`free_tcp_port` owes the same guarantee: it feeds advertised endpoints too — the
kv-event publisher on the single-DP path, the ATOM rendezvous port — and the
kernel only avoids ports that are *bound*, not ones an IPVS Service owns on the
node address. A default host cannot draw one, since the ephemeral range starts
one past the NodePort range, but a cluster that widened either range can, and
silently. Draw again when a kernel-chosen port lands in the window; after 16
draws the ranges overlap far enough that retrying will not converge, so fall back
to the same downward scan, extracted as `_scan_free_port_block` with no
behaviour change for a block.

Unit tests cover the default, a moved range, both opt-out paths, and the
single-port guard together with its fallback.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
….5.16 base

The MI325X image was pinned to a v0.5.15 mi30x base, which cannot run the GLM-5.2
FP8 recipe this image exists for: MTP needs sglang #30839 (IndexShare carried
across the PD legs and CUDA-graph replay) plus GLM-5.2's own
GlmMoeDsaForCausalLMNextN quant resolution, and both land in v0.5.16. Pin exactly
v0.5.16-rocm720-mi30x, the base every 1P1D run behind this branch was done on.
That also retires the reason to apply patches/sglang/ here: v0.5.16 fixes GLM-5.2's
nextn quark exclude upstream, while the backport's anchor survives the refactor
inside DeepSeek's own nextn path, so applying it on this base would rewrite a probe
that is already correct.

The new base bundles Mooncake at upstream #2682 (01d1eb2a) — the same commit the
old one did, so the shared mooncake_cpp diffs still apply — and that revision
installs a HIP IPC transport unconditionally and prefers it over RDMA, so
cross-node PD dies on the first request in hipIpcOpenMemHandle, which cannot open a
peer's handle. Dockerfile.sglang already rebuilds Mooncake in place to gate it
behind MC_ENABLE_HIP_TRANSPORT (default off) and to compile the dma-buf branch in;
gfx942 had no such block, so every MI325X PD bring-up first had to rebuild
engine.so inside the running container before it could move a single KV chunk. Run
the same script with the same diffs instead of forking them: `docker build` has no
GPU to probe, so the only thing that differs is the pinned HIP arch, which the
script now takes as MC_GPU_ARCH and still defaults to gfx950 (Dockerfile.sglang
unchanged). The rebuild stays self-verifying and skippable with BUILD_MOONCAKE=0.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
786f238 carried this fix as a hand-cut diff with an "apply it by hand" recipe,
which is no use to a deployment: the fix only counts if it is in the image the PD
legs actually run, and until it is, everything validated on top of it is validated
against something we do not ship. Wire it into both engine images.

Re-cut the diff as a self-locating idempotent script in the shape of
patches/sglang/. It anchors on source text instead of line numbers, so it survives
the offsets between sglang releases: the anchors are present in both v0.5.15.post1
(the mi35x base) and v0.5.16 (the gfx942 base), and on both trees the result is
byte-identical to the diff it replaces — which is why the diff is dropped instead of
kept alongside as a second source of truth.

The loop is deliberately not tolerant like the nextn backport loop above it: a
script here exits non-zero only when its anchors drifted, i.e. when the fix did not
go in, and an engine image that silently corrupts long prompts is worse than a
failed build. For the same reason the script is all-or-nothing — an absent or
ambiguous anchor writes nothing at all, since a half-applied handover corrupts just
the same. Once a base sglang reads the event itself, the script reports "already
present", no-ops, and can be deleted.

In Dockerfile.sglang the block now sits after the sglang_dsa set main added while
this branch was out. The two are order-independent — the DSA diffs touch decode.py,
this script touches prefill.py, mooncake/conn.py and disaggregation/common/utils.py
— and the comment says so rather than leaving the next reader to diff them. Register
the patch in patch.upstream.status.md, also new on main, which asks for the row in
the same commit as the patch; that row is honest that its "none found" was read from
the web UI rather than verified with `gh`.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
The gfx942 image runs the same GLM-5.2 PD + DP-attention + MTP topology the
sglang_dsa set exists for, but carried none of it: main added that set while this
branch was out, and it is baked only by Dockerfile.sglang. Patch 01 is not
optional there — the aiter (HIP) paged-MQA branch sizes its logits from the
DP-padded row count while `lengths` carries the real one, so `Expected
lengths.size(0) == B` fires as soon as concurrency exceeds 1. It is ROCm-specific
and independent of IndexShare, so nothing at runtime stands in for it.

It could not simply be added: 02 and 04 are `--fuzz=0` context diffs cut against
v0.5.15.post1, and `dsa_indexer.py` drifts on v0.5.16 — `_is_xpu` joins the
platform flags and `_mask_init_and_local_tokens` lands right before
`topk_transform`, failing 2 of patch 01's 4 hunks. Both of its own edit sites,
though, are byte-identical across the two releases. So re-cut patch 01 as a
self-locating idempotent script anchored on source text, the same shape as
patches/sglang/, and delete the diff rather than keep a second source of truth.
On v0.5.15.post1 the script's output is byte-identical to the diff it replaces,
which is what makes this a no-op for gfx950.

`apply_sglang_dsa_patches.sh` grows a `DSA_PATCH_SET` arm instead of being forked:
`full` (default, gfx950) is the old behaviour exactly, and `indexer` (gfx942)
applies patch 01 alone, verifies its bytecode marker, and skips the nextn
prerequisite assert that only the full arm needs. An unknown value fails the
build rather than silently patching nothing.

What gfx942 still does not have is written down rather than left to be
rediscovered. 02b and 04 are substituted at RUNTIME by
`--json-model-override-args '{"index_share_for_mtp_iteration":false}'` — the
substitution main's own README already documents — so that flag is mandatory on
this image, not tuning, and every 1P1D run behind this branch set it. Leave
IndexShare on and the decode leg deadlocks on the first request. 02a is neither
carried nor substituted: its diff was not re-cut for v0.5.16 and its failure was
not observed on that base, which is not the same as fixed. The flag is stated in
the gfx942 Dockerfile header, at the patch block, in the apply script, and in
both patch docs, because a deployment inherits a Dockerfile and a recipe — not a
PR description.

Verified against real sglang sources for both releases, no GPU needed: the script
reproduces the diff byte-for-byte on v0.5.15.post1, applies and compiles on
v0.5.16, is a no-op on re-run, leaves `_p1v2_trim` in the bytecode on both, and on
a drifted anchor exits 1 having written nothing. The full arm was re-run
end-to-end on a v0.5.15.post1 tree — all 8 bytecode markers, the nextn
prerequisite and patch 2a's source marker still pass, so gfx950 is unchanged.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@llying-001
llying-001 force-pushed the llying/glm5p2_fp8_fixes branch from 6121189 to 1380228 Compare August 3, 2026 04:52
llying-001 added a commit that referenced this pull request Aug 3, 2026
Review on #56 caught that _reserved_nodeport_range() accepted "" alongside
"none", so a manifest rendering INFERA_NODEPORT_RANGE from an unset variable
dropped the guard silently -- the opposite of what the same function already
does for a malformed value, and undocumented. Empty now means unset.

Also widen _flat_tokens() to the schema's own list[int | tuple[int, int]]
instead of a union of homogeneous lists, so the annotation matches what
BlockStored actually declares and the flat path no longer returns a list
typed as pairs. Behaviour is unchanged: a radix node is bigrams or bare ints
for its whole length, so the single leading check still decides, and the
common flat case still hands the list back uncopied.

Co-authored-by: Cursor <cursoragent@cursor.com>
Review on #56 caught that _reserved_nodeport_range() accepted "" alongside
"none", so a manifest rendering INFERA_NODEPORT_RANGE from an unset variable
dropped the guard silently -- the opposite of what the same function already
does for a malformed value, and undocumented. Empty now means unset.

Also widen _flat_tokens() to the schema's own list[int | tuple[int, int]]
instead of a union of homogeneous lists, so the annotation matches what
BlockStored actually declares and the flat path no longer returns a list
typed as pairs. Behaviour is unchanged: a radix node is bigrams or bare ints
for its whole length, so the single leading check still decides, and the
common flat case still hands the list back uncopied.

Signed-off-by: liyingli <liyingli@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@llying-001
llying-001 force-pushed the llying/glm5p2_fp8_fixes branch from 11b5905 to c178b69 Compare August 3, 2026 05:06
@jiejingzhangamd
jiejingzhangamd merged commit 9bf69f2 into main Aug 3, 2026
28 checks passed
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.

3 participants