vllm_dissag: add GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP enablement#2
Open
raviguptaamd wants to merge 6 commits into
Open
vllm_dissag: add GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP enablement#2raviguptaamd wants to merge 6 commits into
raviguptaamd wants to merge 6 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.
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
1 participant
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