vllm_dissag: add GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP enablement - #2
Open
raviguptaamd wants to merge 13 commits into
Open
vllm_dissag: add GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP enablement#2raviguptaamd wants to merge 13 commits into
raviguptaamd wants to merge 13 commits into
Conversation
Ports the GLM-5.1-FP8 WideEP enablement (from MAD-private #338) onto the unified two-axis launcher (ROCm#171). GLM-5.1 is GlmMoeDsaForCausalLM: MLA + DeepSeek Sparse Attention, served under MoRI-EP wideEP disagg at 1P1D (EP=8) and 2P2D (EP=16). Changes: - run_xPyD_models.slurm: GLM-5.1-FP8 added to VALID_MODELS, MORI_EP_VALID_MODELS, and WIDE_EP_ONLY_MODELS (DSA validated only under wideEP; reject the untested moriio+TP path). - models.yaml: GLM-5.1-FP8 entry. env: encodes the DSA recipe (KV_BLOCK_SIZE=1, VLLM_ROCM_USE_AITER_MLA=1, eager both roles, per-role mori all2all, 16GiB shmem heap). dp_flags carries the GLM tool/reasoning parsers. - connectors/moriio.sh: * wideEP serve now appends model_args (per-model models.yaml flags) in both the dry-run and real paths. Previously only the TP path did, so wideEP dp:/dp_flags were silently dropped (rixl already appended them in both paths). No-op for models with empty dp flags (DeepSeek). * connector_runtime_patch() applies the 4 GLM DSA patchers, gated on MODEL_NAME=GLM-5.1-FP8; a pure no-op for every other model. - 4 required DSA patchers (+1 optional instrument), idempotent + anchor-based + self-skipping: kernel invalid-token fix (vllm #45324), MoRIIO connector dual-KV geometry, engine per-layer transfer-offset cache, and the completion gate fix (gate on num_transfer_layers, excluding the DSA indexer caches). DeepSeek/other models are byte-identical in behavior. Verified: bash -n, py_compile, yaml lint, and DRY_RUN driver argv for GLM 1P1D/2P2D + DeepSeek no-regression. Co-Authored-By: Claude <noreply@anthropic.com>
… image Verified the 4 GLM DSA patchers against the validated target image (mori v1.2.1 / aiter 0.1.16.post3). The MoRIIO connector is restructured there and handles DSA dual-KV natively: only the invalid-token kernel fix (#45324) is still needed; the dual-KV geometry (moriio_layout.py per-layer geometry) and completion-gate (engine sealed writes_expected, which already covers "hybrid models register more caches than hooks fire") are native. apply_glm_dsa_moriio_gate_fix.py previously patched the connector half unconditionally, so on a restructured image it half-applied (injected a dead num_transfer_layers into the connector) while the engine half found no anchor. Made it atomic: skip BOTH halves unless the engine gate anchor is present. Verified the connector is left unmutated on v1.2.1, while all 4 patchers still fully apply on the older b10a9f7a image. No per-image branching — the patchers self-adapt via anchor presence. models.yaml: document the validated image and the per-image patcher behavior.
…ecedence) A DeepSeek-tuned disagg image bakes recipe knobs as container ENV (KV_BLOCK_SIZE=16, VLLM_ROCM_USE_AITER_MLA=0, VLLM_CUDAGRAPH_MODE=PIECEWISE, KV_CACHE_MEMORY_BYTES=...). The driver exported models.yaml env: with "if k in os.environ: continue" (so submit-time -e wins), but baked container ENV also satisfies that check, so models.yaml silently LOST to the image. Any model whose recipe differs from the image defaults was mis-configured — GLM-5.1 DSA needs block=1 + AITER sparse MLA on + eager, but ran with the DeepSeek block=16 / MLA-off / PIECEWISE baked values. Fix (precedence: image-baked < models.yaml < submit-time -e): - run_xPyD_models.slurm captures MODELS_YAML_PROTECT = the recipe keys the USER set at submit time (checked on the host before the slurm sets defaults), forwarded via -e. - vllm_disagg.sh: yaml env overrides existing env EXCEPT protected keys (genuine submit-time -e still wins). Falls back to legacy "skip if in env" when MODELS_YAML_PROTECT is unset (direct/no-slurm runs), so nothing regresses. Verified in-container against the mori v1.2.1 router image: with baked block=16/MLA=0/PIECEWISE, GLM-5.1-FP8 now emits --block-size 1, kv-cache-dtype fp8, cudagraph NONE. DeepSeek (empty dp flags, matching baked recipe) unaffected.
Per-role cudagraph for GLM-5.1: prefill NONE (eager — prefill cudagraph capture deadlocks on this stack), decode PIECEWISE (captures cleanly on DSA, ~3.7x ITL win vs eager: ~69ms vs ~264ms, validated). Was NONE/NONE (the conservative "cudagraph unproven" starting point); decode PIECEWISE is now proven safe. Verified in-container: prefill emits cudagraph_mode NONE, decode PIECEWISE.
On the mori v1.2.1 connector (restructured vs the b10a9f7a connector #338 ran
on), a request aborted before its KV-transfer peer handshake has peer_zmq=None.
_release_write_prefill_blocks then calls parse_moriio_zmq_address(None) ->
None.split(",") -> AttributeError, which only the ValueError branch guards, so
the EngineCore dies and cascades to all decode workers (EngineDeadError) -> decode
dead. Observed live: a canary aborted during first-token cold JIT killed decode.
The same file already guards the other call site (request_finished notify path)
with `if peer_zmq is not None`; the release path just missed it. New patcher adds
the matching guard (bail gracefully, like the existing ValueError branch). This
function does not exist in the older b10a9f7a connector, so #338 never hit it.
Idempotent + anchor-based + self-skipping; wired into the GLM connector_runtime_patch
list. No effect when peer_zmq is valid.
Investigation finding (no functional change — comment only): GLM-5.1 on this mori v1.2.1 image serves coherent + correct output up to ~14k prompt tokens; beyond ~16-18k the ROCM_AITER_MLA_SPARSE prefill indexer corrupts output (repetition collapse), reproducing vLLM #40018. This is an upstream AITER/vLLM kernel bug (partial fix #43781 is in the image; the complete prefill fix is not), NOT the MAD port. Tested and ruled out --max-model-len 32768 as a mitigation: the corruption onset is a fixed ~16-18k token count in the gather/logits kernel, not a workspace-size artifact (workspace = max_model_len*40), so capping context does not move the threshold and only limits usable range. Left at native max_model_len; documented the ~14k safe envelope and that long-context needs a newer image with the complete upstream fix.
Port of vLLM PR #47567 (+ Rohan138 PR#1 int64 hardening) as a runtime patcher. The AITER persistent sparse-MLA kernel is numerically wrong for chunked-prefill continuations; the per-token error compounds until long-context retrieval collapses (0/10 past ~20-30k). The gate adds a per-batch check in build() and routes chunked-prefill continuation steps to the correct non-persistent kernel (work_meta_data=None); decode stays on the fast persistent path (no regression). Validated on gfx942 (07032026 image): gate-ON NIAH 35k went 0/10 (repetition collapse) -> 10/10 clean, using the exact #47042 needle harness. Decode-safe by construction (UNIFORM_SINGLE_TOKEN_DECODE, query_len==1 never trips the gate). Co-Authored-By: Claude <noreply@anthropic.com>
Adds fixes/knobs from the >=8k WideEP-disagg crash investigation. The 8k crash (prefill worker dies -> DP-lockstep collapse -> 503) is NOT yet resolved; these are the validated-safe pieces plus opt-in debug levers: - apply_glm_aiter_sampling_oob_fix.py: overlays ROCm/aiter #3658 (TopK/TopP sampling last_valid_id OOB) + rejection-sampling hang cap. Real bug, fix confirmed in-kernel; does NOT fix our crash (kept: correct + needed generally). - apply_glm_dsa_indexer_warmup_fix.py: force-compile DSA indexer kernels at boot (opt-in GLM_INDEXER_WARMUP=1) so they don't JIT mid-inference; also the prewarm requirement. Makes the crash deterministic at boot when it does occur. - models.yaml: VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=64 (bounds the fp8_mqa_logits buffer) + NCCL heartbeat-watchdog knobs (fixed a separate 8k teardown). - moriio.sh: opt-in MoRI overlay (MORI_OVERLAY_DIR) + GLM_PERSIST_GATE toggle (debug) + indexer-warmup hook. - run_xPyD_models.slurm: keepalive bench mode + GLM_PERSIST_GATE/KEEPALIVE_MINS env forwarding. - keepalive_bench.sh: hold server up for external accuracy/crash probes. Co-Authored-By: Claude <noreply@anthropic.com>
Builds the GLM-5.1 DSA WideEP-disagg image from the validated #47766 source stack (persistent sparse-MLA kernel stays ON): - vLLM -> raviguptaamd/vllm @ glm5.1-dsa-wideEP_on_shikpate_06_29_customer (Shiksha 06_29 WideEP base + 4 GLM commits: DSA invalid-token -1, #47766 cache-key accuracy, aiter fused-qk probe crash-safety, MoRIIO DSA indexer KV transfer). - AITER -> STOCK ROCm/aiter @ e03fa6040 compiled from source (no fork; #47766 keeps persistent ON -> hits aiter's native gqa64 fold). Was the 0.1.16.post3 wheel. - MoRI -> ROCm/mori @ 42e895472b08 (main tip past v1.2.1; large-transfer notify). - Router unchanged (raviguptaamd/router @ ravgupta/discovery-dp-rank-roundrobin). Mirrors MAD-private #338's validated stack (1P/1D EP8 + 2P/2D EP16 NIAH PASS), ported onto the 06_29 base rather than b10a9f7a. Co-Authored-By: Claude <noreply@anthropic.com>
The post-vLLM cross-check still asserted the old 0.1.16.post3 wheel version, but AITER is now a stock source build of ROCm/aiter@e03fa6040 (reports 0.1.17.dev195+ge03fa6040). Verify the e03fa6040 commit tag is present instead of a release string. (Build reached this step with MoRI+AITER+vLLM already compiled; only the assertion failed.) Co-Authored-By: Claude <noreply@anthropic.com>
When the serving image already carries the GLM-5.1 DSA fixes in-source (the #47766 stack image from raviguptaamd/vllm@glm5.1-dsa-wideEP_on_shikpate_06_29_customer), set GLM_SKIP_PATCHERS=1 to skip ALL runtime patchers. They are redundant there, and two would actively regress a baked image: the persistent-gate patcher would turn the persistent sparse-MLA kernel OFF (opposite of #47766, which keeps it ON), and the sampling overlay would overwrite the image's stock aiter kernels. Forwarded to the container via run_xPyD -e. Co-Authored-By: Claude <noreply@anthropic.com>
… (DP-notify fixes)
…m Dockerfile Split the shared vllm_disagg image so DeepSeek-V3/R1 and GLM-5.1-FP8 no longer compete for the same Dockerfile pins: - docker/vllm_disagg_inference.ubuntu.amd.Dockerfile: reverted to origin/develop (the DSV3/R1 stack, byte-identical to upstream ROCm#171). No longer repinned to GLM. - docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile: NEW per-model image for GLM-5.1 (raviguptaamd/vllm glm5.1-dsa-wideEP_on_shik_latest + aiter e03fa6040 + mori 42e895472b08 + router). STATUS notes EP32 is a known open defect. - models.json: add card pyt_vllm_disagg_mori_glm-5.1-fp8 -> vllm_disagg_inference.glmv5.1 (MODEL_NAME=GLM-5.1-FP8, RUN_MORI=1, GLM_SKIP_PATCHERS=1). - moriio.sh: remove the runtime MoRI overlay (dev scaffolding). MoRI version is pinned by the Dockerfile MORI_REF; update+rebuild instead of hot-swapping at runtime. connector_runtime_patch() is now a pure no-op before the GLM gate, so DSV3/other models are byte-identical to upstream. - models.yaml: DSV3 dp: caps (--max-num-seqs 64 --max-model-len 32768), isolated to the DeepSeek-V3 entry, to bound the newer base's decode logits workspace (OOM fix). Future models (e.g. Kimi-2.6) add their own vllm_disagg_inference.<model> Dockerfile + card, with zero edits to DSV3 or GLM. Co-Authored-By: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds GLM-5.1-FP8 (
GlmMoeDsaForCausalLM— MLA + DeepSeek Sparse Attention) to the unified two-axis launcher for MoRI-EP WideEP disaggregated serving at 1P1D (EP=8) and 2P2D (EP=16). Stacked on thevllm-disagg-unified-launcherbranch (ROCm#171).GLM-5.1 is the first DSA-family model on this launcher; DSA adds a 2nd KV cache per layer (indexer) and a sparse-MLA decode path needing model-specific handling on top of the working DeepSeek path.
Changes
models.yaml:GLM-5.1-FP8entry — DSA recipe env (KV_BLOCK_SIZE=1,VLLM_ROCM_USE_AITER_MLA=1, prefill eager / decode PIECEWISE, per-role mori HT/LL, fp8 KV) +dp_flagsfor the glm47/glm45 tool/reasoning parsers.run_xPyD_models.slurm: GLM inVALID_MODELS/MORI_EP_VALID_MODELS/WIDE_EP_ONLY_MODELS; models.yaml-over-image-baked-ENV precedence fix (3-tier: image < yaml < submit-e).connectors/moriio.sh: appendmodel_argson the wideEP serve path (was only on the TP path — wideEPdp:flags were silently dropped); implementconnector_runtime_patch()to apply the GLM DSA patchers, gated onMODEL_NAME==GLM-5.1-FP8(no-op for every other model).Isolation
Every GLM addition is gated on
MODEL_NAME==GLM-5.1-FP8or is an additive allowlist/catalog entry. The one shared-path edit (appendmodel_argson wideEP) is a no-op for models with emptydp:flags. DeepSeek/other behavior unchanged.Validation (MI300X gfx942, mori v1.2.1 / aiter 0.1.16.post3 / router181 image)
Test plan
bash -nslurm + moriio.sh + vllm_disagg.sh;py_compilepatchers;models.yamllintDRY_RUN=1 RUN_MORI=1 MODEL_NAME=GLM-5.1-FP8emits--block-size 1 --kv-cache-dtype fp8 --enable-expert-parallel+ glm parsers; DeepSeek argv unchanged🤖 Generated with Claude Code