From c147bb4e0eefc23c0d0622df584ac08d5cbf29ca Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 3 Jul 2026 07:54:15 +0000 Subject: [PATCH 01/13] vllm_dissag: add GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP enablement Ports the GLM-5.1-FP8 WideEP enablement (from MAD-private #338) onto the unified two-axis launcher (#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 --- .../vllm_dissag/apply_glm_dsa_kernel_fix.py | 85 +++++++++ .../apply_glm_dsa_moriio_dualkv_fix.py | 176 ++++++++++++++++++ .../apply_glm_dsa_moriio_engine_fix.py | 116 ++++++++++++ .../apply_glm_dsa_moriio_gate_fix.py | 118 ++++++++++++ .../apply_glm_dsa_moriio_instrument.py | 92 +++++++++ scripts/vllm_dissag/connectors/moriio.sh | 69 ++++++- scripts/vllm_dissag/models.yaml | 48 +++++ scripts/vllm_dissag/run_xPyD_models.slurm | 7 +- 8 files changed, 702 insertions(+), 9 deletions(-) create mode 100755 scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py diff --git a/scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py b/scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py new file mode 100755 index 00000000..33f50dc4 --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_kernel_fix.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Apply the GLM-5.1 DSA sparse-attention invalid-token kernel fix (vllm #45324). + +The DSA indexer kernel `_convert_req_index_to_global_index_kernel` in + vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +maps invalid token slots to 0 instead of -1. With block-size 1 + DSA sparse MLA +that corrupts KV reads and the model emits `!!!` for every prompt. + +Upstream fix: vllm-project/vllm #45324 -- flip the 0 to -1 in the tl.where call: + is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off + is_invalid_tok | (~valid_block), -1, base * BLOCK_SIZE + inblock_off + +Design (matches launcher contract -- runs unconditionally for GLM, aborts on real +failure): + * IDEMPOTENT : if already -1, report and exit 0 (no-op). + * SELF-SKIPPING: if the file/anchor is absent (refactored or the rebase already + fixed it differently), report and exit 0 -- do NOT abort, because b10a9f7a may + carry the fix natively. We only fail on the one unambiguous bad state we can + fix and didn't, or on write failure. + * VERIFIES the post-write state. + +Usage: apply_glm_dsa_kernel_fix.py +""" +import os +import re +import sys + +REL = "v1/attention/backends/mla/rocm_aiter_mla_sparse.py" + +# Anchor is the stable right-hand side of the tl.where; the middle operand is the +# 0 (buggy) / -1 (fixed) we toggle. Whitespace-tolerant. +RE_ANY = re.compile( + r"(is_invalid_tok\s*\|\s*\(~valid_block\)\s*,\s*)(-?\d+)(\s*,\s*base\s*\*\s*BLOCK_SIZE\s*\+\s*inblock_off)" +) + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + vllm_dir = sys.argv[1] + path = os.path.join(vllm_dir, REL) + + if not os.path.isfile(path): + # File not present on this build -> nothing we can or should do. The + # rebase may use a different sparse backend layout. Do not block launch. + print(f"[glm-dsa] {REL} not found under {vllm_dir} -- skipping (assuming native/refactored).") + return 0 + + src = open(path).read() + m = RE_ANY.search(src) + if not m: + # Anchor gone (refactored / already fixed differently). Don't block. + print(f"[glm-dsa] invalid-token kernel anchor not found in {path} -- skipping (assuming native fix).") + return 0 + + cur = m.group(2) + if cur == "-1": + print(f"[glm-dsa] already fixed (kernel returns -1) in {path} -- no-op.") + return 0 + if cur != "0": + # Unexpected value -- surface it but don't guess. Treat as needs-attention. + print(f"[glm-dsa] ERROR: unexpected invalid-token return value '{cur}' (expected 0 or -1) in {path}.", + file=sys.stderr) + return 1 + + # cur == "0" : the known bug. Flip to -1. + new_src = src[:m.start(2)] + "-1" + src[m.end(2):] + try: + open(path, "w").write(new_src) + except OSError as e: + print(f"[glm-dsa] ERROR: failed to write patched {path}: {e}", file=sys.stderr) + return 1 + + # Verify. + chk = RE_ANY.search(open(path).read()) + if not chk or chk.group(2) != "-1": + print(f"[glm-dsa] ERROR: post-write verification failed in {path}.", file=sys.stderr) + return 1 + print(f"[glm-dsa] patched: invalid-token kernel now returns -1 (vllm #45324) in {path}.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py new file mode 100755 index 00000000..ac6800bf --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_dualkv_fix.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Patch the MoRIIO KV connector to handle GLM-5.1 DSA's DUAL KV cache. + +PROBLEM (root cause of the 2P2D "Reaped deferred sends / no finished_sending" stall): + GLM-5.1 (GlmMoeDsaForCausalLM -> deepseek_v2.py) has TWO KV caches per layer: + - main MLA latent KV : MLAAttentionSpec, head_size = kv_lora_rank+rope (~576) + - DSA indexer KV : DeepseekV32IndexerCache, MLAAttentionSpec head_size = index_head_dim (128) + Both are 3D ("use_mla"), but DIFFERENT latent dim -> DIFFERENT per-block byte size. + The MoRIIO connector computes ONE global geometry from `first_kv_cache` and reuses it + for every cache, so the indexer cache is transferred with the main-MLA block size -> + wrong bytes/size -> the RDMA read for that region never reconciles -> completion notify + is never produced -> decode reaps deferred sends after 60s -> request hangs. + +FIX (surgical, per-layer geometry; no behavior change for single-cache MLA/DeepSeek): + 1. register_kv_caches: size each registered region by its OWN tensor (per-cache + region_len), not the global self.block_len. Also fix the local_kv_cache_size + append to use the current cache, not a stale loop var. + 2. _compute_block_transfer_offsets: derive shape from the PER-LAYER tensor + (self.kv_caches[layer_name].shape) instead of the global self.kv_cache_shape, + so transfer_size_byte / strides match that cache. + 3. _read_blocks: compute offsets PER LAYER inside the loop (was computed once from + first_layer and reused for all layers). + +Idempotent + anchor-based: each hunk checks if already applied / anchor present; +missing anchor -> warn-and-skip (so it is safe across connector revisions). A hunk +that finds its OLD anchor but fails to apply is a hard error (would silently keep the bug). + +Usage: apply_glm_dsa_moriio_dualkv_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-dualkv] {REL} not found -- skipping (connector layout differs).") + return 0 + + src = open(path).read() + orig = src + applied = [] + + # --- Hunk 1: per-cache region_len in register_kv_caches --------------------- + h1_old = """ for cache_or_caches in kv_caches.values(): + cache_list = [cache_or_caches] if use_mla else cache_or_caches + for cache in cache_list: + base_addr = cache.data_ptr() + region_len = self.num_blocks * self.block_len + caches_data.append((base_addr, region_len, cache.device.index, "")) + kv_caches_base_addr.append(base_addr)""" + h1_new = """ for cache_or_caches in kv_caches.values(): + cache_list = [cache_or_caches] if use_mla else cache_or_caches + for cache in cache_list: + base_addr = cache.data_ptr() + # DSA dual-KV fix: size each region by its OWN tensor, not the + # global self.block_len (the DSA indexer cache has a different + # latent dim than the main MLA cache). + region_len = cache.nelement() * cache.element_size() + caches_data.append((base_addr, region_len, cache.device.index, "")) + kv_caches_base_addr.append(base_addr)""" + if "region_len = cache.nelement() * cache.element_size()" in src: + applied.append("h1 (already)") + elif h1_old in src: + src = src.replace(h1_old, h1_new, 1) + applied.append("h1") + else: + print("[glm-dualkv] WARN: h1 anchor (region_len loop) not found -- skipping h1.") + + # --- Hunk 1b: local_kv_cache_size uses current kv_cache, not stale `cache` -- + h1b_old = " self.local_kv_cache_size.append(cache.nelement() * cache.element_size())" + h1b_new = " self.local_kv_cache_size.append(kv_cache.nelement() * kv_cache.element_size())" + if h1b_new in src: + applied.append("h1b (already)") + elif h1b_old in src: + src = src.replace(h1b_old, h1b_new, 1) + applied.append("h1b") + else: + print("[glm-dualkv] WARN: h1b anchor (local_kv_cache_size) not found -- skipping h1b.") + + # --- Hunk 2: per-layer shape in _compute_block_transfer_offsets ------------- + h2_old = """ assert self.kv_cache_shape is not None, "KV caches shape not initialized" + is_mla = len(self.kv_cache_shape) == 3 + stride = self.kv_caches[layer_name].stride() + sz = self.kv_caches[layer_name].element_size() + if is_mla: + blknum, blksize, hs = self.kv_cache_shape + hn = 1 + block_stride = stride[0] + else: + _, blknum, blksize, hn, hs = self.kv_cache_shape""" + h2_new = """ # DSA dual-KV fix: use the PER-LAYER tensor shape, not the global + # self.kv_cache_shape (the DSA indexer cache differs from the main MLA). + _layer_shape = tuple(self.kv_caches[layer_name].shape) + assert len(_layer_shape) > 0, "KV caches shape not initialized" + is_mla = len(_layer_shape) == 3 + stride = self.kv_caches[layer_name].stride() + sz = self.kv_caches[layer_name].element_size() + if is_mla: + blknum, blksize, hs = _layer_shape + hn = 1 + block_stride = stride[0] + else: + _, blknum, blksize, hn, hs = _layer_shape""" + if "_layer_shape = tuple(self.kv_caches[layer_name].shape)" in src: + applied.append("h2 (already)") + elif h2_old in src: + src = src.replace(h2_old, h2_new, 1) + applied.append("h2") + else: + print("[glm-dualkv] WARN: h2 anchor (_compute_block_transfer_offsets head) not found -- skipping h2.") + + # --- Hunk 3: per-layer offsets in _read_blocks ----------------------------- + h3_old = """ first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0] + offs = self._compute_block_transfer_offsets( + first_layer, local_block_ids, remote_block_ids, remote_moriio_meta + ) + + for layer_name in self.layer_name_to_local_kv_cache_metadata: + sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( + layer_name + ) + # TODO : apply multi-session batch-read when moriio support it + transfer_status = self.moriio_wrapper.read_remote_data( + offs[2], offs[0], offs[1], sessions[sess_idx] + )""" + h3_new = """ # DSA dual-KV fix: compute offsets PER LAYER (the DSA indexer cache has a + # different per-block size than the main MLA cache, so a single offs reused + # across all layers mis-sizes the indexer transfer -> lost completion notify). + for layer_name in self.layer_name_to_local_kv_cache_metadata: + sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( + layer_name + ) + offs = self._compute_block_transfer_offsets( + layer_name, local_block_ids, remote_block_ids, remote_moriio_meta + ) + # TODO : apply multi-session batch-read when moriio support it + transfer_status = self.moriio_wrapper.read_remote_data( + offs[2], offs[0], offs[1], sessions[sess_idx] + )""" + if "compute offsets PER LAYER" in src: + applied.append("h3 (already)") + elif h3_old in src: + src = src.replace(h3_old, h3_new, 1) + applied.append("h3") + else: + print("[glm-dualkv] WARN: h3 anchor (_read_blocks first_layer offsets) not found -- skipping h3.") + + if src != orig: + try: + open(path, "w").write(src) + except OSError as e: + print(f"[glm-dualkv] ERROR: write failed for {path}: {e}", file=sys.stderr) + return 1 + print(f"[glm-dualkv] patched {path} -- hunks: {', '.join(applied)}") + else: + print(f"[glm-dualkv] no changes ({', '.join(applied) or 'nothing applied'}) for {path}") + + # py-compile sanity + try: + import py_compile + py_compile.compile(path, doraise=True) + print("[glm-dualkv] py_compile OK") + except Exception as e: # noqa: BLE001 + print(f"[glm-dualkv] ERROR: patched file fails to compile: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py new file mode 100755 index 00000000..9d6d468b --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_engine_fix.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Fix MoRIIO WRITE-path per-layer offset caching for GLM-5.1 DSA dual KV cache. + +ROOT CAUSE (proven by instrumentation, job 37594): + GLM-5.1 (GlmMoeDsaForCausalLM) registers num_layers=156 KV caches = 78 main MLA + (per-block latent dim 576) + 78 DSA indexer caches (latent dim 132). TWO geometries. + + moriio_engine.py::MoRIIOEngine._prepare_transfer_plan computes the RDMA transfer + offsets ONCE (for whatever layer arrives first) and caches them on + request_info.transfer_offset, then REUSES that single offset/size tuple for ALL 156 + layers. The 78 indexer layers (dim 132) get written with the main-MLA geometry + (dim 576) -> wrong byte size/offset -> those RDMA writes are malformed; the per-layer + write accounting (writes_done) and/or the remote completion never reconciles -> + the producer's send_notify (gated on writes_done >= num_layers) misbehaves and the + decode side never receives a clean completion -> "Reaped deferred sends / no + finished_sending after 60s" -> request hangs. + +FIX (surgical, no dataclass change): + Cache transfer offsets PER LAYER on the request_info via a dynamically-attached dict + ``_transfer_offset_by_layer`` keyed by layer_name, instead of the single + ``transfer_offset`` slot. Each of the 156 layers then transfers with its OWN geometry + (the underlying _compute_block_transfer_offsets already takes layer_name and, with the + companion dualkv patch h2, reads the per-layer tensor shape). + + Single-geometry models (DeepSeek-V3 / Hunyuan, 1 cache/layer) are unaffected: + every layer has identical geometry, so per-layer caching yields the same offsets. + +Idempotent + anchor-based. A missing anchor warns-and-skips; a found OLD anchor that +fails to apply is a hard error (would silently keep the stall). + +Usage: apply_glm_dsa_moriio_engine_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-engine] {REL} not found -- skipping (engine layout differs).") + return 0 + + src = open(path).read() + + old = """ # Compute offsets if not cached + if request_info.transfer_offset is None: + offsets = self.worker._compute_block_transfer_offsets( + task.layer_name, + task.local_block_ids, + request_info.block_ids, + remote_moriio_meta, + ) + request_info.transfer_offset = offsets + + # Get session index + layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys()) + sess_idx = layer_names.index(task.layer_name) + + local_off, remote_off, sizes = request_info.transfer_offset""" + + new = """ # DSA dual-KV fix: cache offsets PER LAYER, not once per request. GLM-5.1 has + # two cache geometries (main MLA dim 576 + DSA indexer dim 132); a single + # cached offset reused across all 156 layers mis-sizes the indexer writes and + # the completion never reconciles. Per-layer caching is identical for + # single-geometry models (DeepSeek/Hunyuan). + _off_by_layer = getattr(request_info, "_transfer_offset_by_layer", None) + if _off_by_layer is None: + _off_by_layer = {} + request_info._transfer_offset_by_layer = _off_by_layer + offsets = _off_by_layer.get(task.layer_name) + if offsets is None: + offsets = self.worker._compute_block_transfer_offsets( + task.layer_name, + task.local_block_ids, + request_info.block_ids, + remote_moriio_meta, + ) + _off_by_layer[task.layer_name] = offsets + # keep the legacy single-slot populated (first layer) for any external reader + if request_info.transfer_offset is None: + request_info.transfer_offset = offsets + + # Get session index + layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys()) + sess_idx = layer_names.index(task.layer_name) + + local_off, remote_off, sizes = offsets""" + + if "_transfer_offset_by_layer" in src: + print(f"[glm-engine] already patched (_transfer_offset_by_layer present) -- no-op.") + elif old in src: + src = src.replace(old, new, 1) + open(path, "w").write(src) + print(f"[glm-engine] patched per-layer offset caching in {path}") + else: + print(f"[glm-engine] WARN: anchor (_prepare_transfer_plan offset block) not found -- skipping (engine revision differs).") + # Not fatal: without the anchor we can't safely patch; surface clearly. + return 0 + + try: + import py_compile + py_compile.compile(path, doraise=True) + print("[glm-engine] py_compile OK") + except Exception as e: # noqa: BLE001 + print(f"[glm-engine] ERROR: compile failed: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py new file mode 100755 index 00000000..9cec5e06 --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""Fix the MoRIIO transfer-completion gate for GLM-5.1 DSA dual KV cache. + +ROOT CAUSE (PROVEN, job 37615 instrumentation): + GLM-5.1 registers num_layers=156 KV caches = 78 main MLA (model.layers.N.self_attn.attn) + + 78 DSA indexer caches (model.layers.N.self_attn.indexer.k_cache). BUT only the 78 + main-MLA layers ever go through the KV-connector save_kv_layer hook -> only they write + -> writes_done caps at 78. The indexer caches are registered (counted in num_layers) + but vLLM NEVER calls save_kv_layer for them (the DSA indexer is a separate attention + component / DeepseekV32IndexerBackend that doesn't use the connector save path; decode + recomputes indexer state from the transferred main latent KV). + + The producer completion gate (moriio_engine.py): + request_info.writes_done += 1 + if request_info.writes_done >= self.worker.num_layers: # 156, never reached + send_notify(...) + caps at writes_done=78 < num_layers=156 -> send_notify NEVER fires -> decode never + gets completion -> "Reaped deferred sends / no finished_sending after 60s" -> stall. + +FIX: + Add self.num_transfer_layers = count of caches that actually transfer (exclude + '.indexer.' caches), with a fallback to num_layers (so single-geometry models - + DeepSeek-V3 / Hunyuan, no indexer - are bit-identical). Gate completion on + num_transfer_layers instead of num_layers. num_layers itself is left unchanged + (it is also used by the Llama-4 per-layer block-window loop, which needs all caches). + +Companion to apply_glm_dsa_moriio_engine_fix.py (per-layer offset caching). This gate +fix is the primary unblocker; the offset fix is correctness insurance for the layers +that DO write (all same geometry here, but harmless). + +Idempotent + anchor-based. Patches BOTH files (connector: define the field; engine: +use it). A found-old-anchor that fails is a hard error. + +Usage: apply_glm_dsa_moriio_gate_fix.py +""" +import os +import sys + +CONN_REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +ENG_REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" + + +def patch_connector(path: str) -> int: + src = open(path).read() + old = " self.num_layers = len(self.kv_caches.keys())" + new = """ self.num_layers = len(self.kv_caches.keys()) + # DSA dual-KV fix: the producer completion gate must count only layers that + # actually transfer via save_kv_layer. GLM-5.1 registers 2 caches/layer (main + # MLA + DSA indexer), but only the main-MLA caches go through save_kv_layer; the + # '.indexer.' caches are registered yet never written. Gating on len(kv_caches) + # would never be reached. Exclude indexer caches; fall back to num_layers for + # single-geometry models (DeepSeek/Hunyuan have no indexer -> identical). + self.num_transfer_layers = ( + len([k for k in self.kv_caches.keys() if ".indexer." not in k]) + or self.num_layers + ) + logger.info( + "[moriio] completion gate: num_transfer_layers=%d (num_layers=%d)", + self.num_transfer_layers, self.num_layers, + )""" + if "self.num_transfer_layers" in src: + print(f"[glm-gate] connector already patched -- no-op.") + return 0 + if old not in src: + print(f"[glm-gate] WARN: connector anchor (num_layers=) not found -- skipping.") + return 0 + src = src.replace(old, new, 1) + open(path, "w").write(src) + print(f"[glm-gate] patched connector: defined num_transfer_layers in {path}") + return 0 + + +def patch_engine(path: str) -> int: + src = open(path).read() + old = " if request_info.writes_done >= self.worker.num_layers:" + new = """ if request_info.writes_done >= getattr( + self.worker, "num_transfer_layers", self.worker.num_layers + ):""" + if 'getattr(\n self.worker, "num_transfer_layers"' in src or "num_transfer_layers" in src: + print(f"[glm-gate] engine already patched -- no-op.") + return 0 + if old not in src: + print(f"[glm-gate] WARN: engine anchor (writes_done gate) not found -- skipping.") + return 0 + src = src.replace(old, new, 1) + open(path, "w").write(src) + print(f"[glm-gate] patched engine: gate on num_transfer_layers in {path}") + return 0 + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + base = sys.argv[1] + conn = os.path.join(base, CONN_REL) + eng = os.path.join(base, ENG_REL) + if not os.path.isfile(conn) or not os.path.isfile(eng): + print("[glm-gate] connector/engine not found -- skipping (layout differs).") + return 0 + + rc = patch_connector(conn) or patch_engine(eng) + if rc: + return rc + + try: + import py_compile + py_compile.compile(conn, doraise=True) + py_compile.compile(eng, doraise=True) + print("[glm-gate] py_compile OK (both files)") + except Exception as e: # noqa: BLE001 + print(f"[glm-gate] ERROR: compile failed: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py new file mode 100755 index 00000000..68e03beb --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_instrument.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""TEMPORARY instrumentation: does the DSA indexer cache reach the MoRIIO save path? + +Adds logging at two points in moriio_connector.py to settle the dual-KV RCA: + 1. register_kv_caches: log all kv_caches layer names + their shapes + num_layers. + -> shows whether the DeepseekV32IndexerCache is even registered, and its geometry. + 2. _write_blocks_for_req: log each distinct layer_name that actually triggers a write. + -> compare the COUNT/SET of written layers vs num_layers. If indexer layers are in + kv_caches (counted in num_layers) but never written, writes_done can never reach + num_layers -> send_notify never fires -> the stall. + +This is diagnostic only (no behavior change). Remove before any production use. +Idempotent + anchor-safe. + +Usage: apply_glm_dsa_moriio_instrument.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-instr] {REL} not found -- skipping.") + return 0 + + src = open(path).read() + orig = src + + # --- Point 1: log kv_caches inventory at num_layers assignment ------------- + a1 = " self.num_layers = len(self.kv_caches.keys())" + b1 = """ self.num_layers = len(self.kv_caches.keys()) + # [glm-instr] kv-cache inventory (dual-KV diagnosis) + try: + for _ln, _kv in self.kv_caches.items(): + logger.info("[glm-instr][register] layer=%s shape=%s dtype=%s", + _ln, tuple(_kv.shape), _kv.dtype) + logger.info("[glm-instr][register] num_layers=%d total_kv_caches=%d", + self.num_layers, len(self.kv_caches)) + except Exception as _e: # noqa: BLE001 + logger.info("[glm-instr][register] inventory log failed: %s", _e)""" + if "[glm-instr][register]" in src: + pass + elif a1 in src: + src = src.replace(a1, b1, 1) + else: + print("[glm-instr] WARN: register anchor (num_layers=) not found.") + + # --- Point 2: log each written layer in _write_blocks_for_req ------------- + a2 = " def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):" + b2 = (a2 + "\n" + ' # [glm-instr] record which layers actually trigger a KV write\n' + ' try:\n' + ' _seen = getattr(self, "_glm_instr_written_layers", None)\n' + ' if _seen is None:\n' + ' _seen = set(); self._glm_instr_written_layers = _seen\n' + ' if layer_name not in _seen:\n' + ' _seen.add(layer_name)\n' + ' logger.info("[glm-instr][write] NEW layer=%s total_written=%d/%d",\n' + ' layer_name, len(_seen), getattr(self, "num_layers", -1))\n' + ' except Exception as _e: # noqa: BLE001\n' + ' logger.info("[glm-instr][write] log failed: %s", _e)') + if "[glm-instr][write]" in src: + pass + elif a2 in src: + src = src.replace(a2, b2, 1) + else: + print("[glm-instr] WARN: _write_blocks_for_req anchor not found.") + + if src != orig: + open(path, "w").write(src) + print(f"[glm-instr] instrumented {path}") + else: + print(f"[glm-instr] already instrumented / nothing to do for {path}") + + try: + import py_compile + py_compile.compile(path, doraise=True) + print("[glm-instr] py_compile OK") + except Exception as e: # noqa: BLE001 + print(f"[glm-instr] ERROR: compile failed: {e}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 44c64d68..b13dab9a 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -124,13 +124,65 @@ _moriio_build_kv_transfer_config() { } connector_runtime_patch() { - # No-op: the MoRIIO multi-node disagg fixes (vLLM PR#39276 notify-path, #41751 LL - # split, DP-rank hash-failsafe) are committed in-source in the vLLM the image is - # built from (see the Dockerfile VLLM_REF). There is no runtime .py patcher — that - # would be a drifting duplicate of fixes that already live upstream in the fork. - # If you ever run an image WITHOUT these fixes baked, use an image that has them - # (rebuild from the pinned VLLM_REF) rather than patching a stock image at runtime. - return 0 + # MoRIIO multi-node disagg fixes (vLLM PR#39276 notify-path, #41751 LL split, + # DP-rank hash-failsafe) are committed in-source in the vLLM the image is built + # from (Dockerfile VLLM_REF). There is no generic runtime .py patcher for those — + # that would be a drifting duplicate of fixes already upstream in the fork. + # + # EXCEPTION — GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA + DSA sparse attention): + # DSA is a NEW attention family the MoRIIO connector was never built for. It adds + # a 2nd KV cache per layer (indexer) with a different geometry, which the + # single-geometry connector mis-handles -> disagg KV transfer stalls; plus a DSA + # invalid-token kernel bug (#45324) that produces `!!!`. These are model-specific + # code gaps, applied here as idempotent, anchor-based, self-skipping .py patchers + # (they no-op cleanly if the fix is native/refactored on the chosen image). Gated + # on MODEL_NAME so DeepSeek/other models are a pure no-op (byte-identical to before). + [ "${MODEL_NAME:-}" = "GLM-5.1-FP8" ] || return 0 + _glm_dsa_runtime_patch +} + +# GLM-5.1 DSA patchers (see connector_runtime_patch). Ported from MAD-private #338. +# Resolves the vLLM install dir, then applies the 4 required patchers in order, +# aborting on a hard failure (a real failure means GLM emits garbage or stalls, so +# failing at launch is correct). Patchers self-skip (rc 0) when their anchor is +# absent, so an image that already carries or refactored a fix no-ops cleanly. +_glm_dsa_runtime_patch() { + local _patch_dir="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)}" + local _vllm_dir + _vllm_dir="$(python3 -c 'import vllm, os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || true)" + if [ -z "${_vllm_dir}" ] || [ ! -d "${_vllm_dir}" ]; then + echo "Error: [glm] cannot locate vLLM install dir for DSA patchers. Aborting." >&2 + exit 1 + fi + echo "[glm] MODEL_NAME=GLM-5.1-FP8: applying DSA runtime patchers against ${_vllm_dir}" + + # Ordered list of REQUIRED patchers (all abort on hard failure). + local _p + for _p in \ + apply_glm_dsa_kernel_fix.py \ + apply_glm_dsa_moriio_dualkv_fix.py \ + apply_glm_dsa_moriio_engine_fix.py \ + apply_glm_dsa_moriio_gate_fix.py; do + local _py="${_patch_dir}/${_p}" + if [ ! -f "${_py}" ]; then + echo "Error: [glm] required patcher ${_py} not found. Aborting." >&2 + exit 1 + fi + echo "[glm] applying ${_p}" + python3 "${_py}" "${_vllm_dir}" 2>&1 || { + echo "Error: [glm] ${_p} failed — GLM-5.1 would emit garbage or stall. Aborting." >&2 + exit 1 + } + done + + # Optional diagnostic instrumentation (GLM_INSTRUMENT=1). Non-fatal. + if [ "${GLM_INSTRUMENT:-0}" = "1" ]; then + local _instr="${_patch_dir}/apply_glm_dsa_moriio_instrument.py" + if [ -f "${_instr}" ]; then + echo "[glm] applying instrumentation (GLM_INSTRUMENT=1): apply_glm_dsa_moriio_instrument.py" + python3 "${_instr}" "${_vllm_dir}" 2>&1 || echo "Warning: [glm] instrumentation failed (non-fatal)." + fi + fi } # connector_launch_worker [dp_start_rank] @@ -219,7 +271,7 @@ connector_launch_worker() { --all2all-backend "${_all2all}" \ --trust-remote-code \ --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS:-7200}" \ - "${exec_args[@]}" "${extra_args[@]}" "${kv_args[@]}" + "${exec_args[@]}" "${extra_args[@]}" "${kv_args[@]}" "${model_args[@]}" WORKER_PID=0; return 0 fi @@ -242,6 +294,7 @@ connector_launch_worker() { "${exec_args[@]}" \ "${extra_args[@]}" \ "${kv_args[@]}" \ + "${model_args[@]}" \ 2>&1 | tee /run_logs/${SLURM_JOB_ID}/${log_prefix}_NODE${NODE_RANK}.log >/dev/null & WORKER_PID=$! return 0 diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..3b721990 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -178,3 +178,51 @@ DeepSeek-R1: dp: "" decode: dp: "" + +# ============================ MoE + DSA (wideEP only) ============================ + +# GLM-5.1-FP8 (zai-org/GLM-5.1-FP8, arch GlmMoeDsaForCausalLM): MLA + DeepSeek +# Sparse Attention (DSA). 78 layers (3 dense + 75 MoE), 256 routed experts top-8 +# + 1 shared, FP8 block 128. wideEP-only (see WIDE_EP_ONLY_MODELS in the slurm). +# +# GLM differs from DeepSeek in 3 recipe-defining ways (both are MLA MoE, but GLM +# is DSA-sparse): +# - KV_BLOCK_SIZE=1 (DSA sparse indexer REQUIRES block-size 1; DS uses 16) +# - VLLM_ROCM_USE_AITER_MLA=1 (GLM MLA path ON via AITER sparse; DS sets 0) +# - eager BOTH roles (DSA + cudagraph unproven; DS runs PIECEWISE decode) +# Note: block=1 + AITER_MLA=1 are ALREADY the moriio.sh connector defaults — GLM +# keeps them; DeepSeek is the one that overrides them off. Set here explicitly so +# the recipe is self-documenting and robust to connector default changes. +# +# DSA adds a 2nd KV cache per layer (indexer); the MoRIIO connector needs the GLM +# DSA patchers (kernel #45324 + dual-KV geometry + per-layer offset + completion +# gate) applied by connector_runtime_patch in connectors/moriio.sh (gated on this +# MODEL_NAME). Those are code patches, not flags — nothing to add here for them. +# +# dp_flags carry the GLM tool/reasoning parsers (AMD GLM recipe). They are applied +# to BOTH roles by compose() in vllm_disagg.sh and reach `vllm serve` via the +# connector's model_args. block-size / kv-cache-dtype / all2all / cudagraph come +# from the env: recipe above (the connector emits them), so the dp: blocks are +# empty like the DeepSeek family. +GLM-5.1-FP8: + env: + VLLM_USE_V1: "1" + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + KV_BLOCK_SIZE: "1" + KV_CACHE_DTYPE: "fp8" + GPU_MEMORY_UTILIZATION: "0.80" + VLLM_CUDAGRAPH_MODE: "NONE" + PREFILL_CUDAGRAPH_MODE: "NONE" + DECODE_CUDAGRAPH_MODE: "NONE" + CUDAGRAPH_CAPTURE_SIZES: "1 2 4 8 16 32 64 128 256" + VLLM_ALL2ALL_BACKEND: "mori_high_throughput" + PREFILL_MORI_BACKEND: "mori_high_throughput" + DECODE_MORI_BACKEND: "mori_low_latency" + MORI_SHMEM_HEAP_SIZE: "17179869184" + dp_flags: "--tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice --chat-template-content-format string" + prefill: + dp: "" + decode: + dp: "" diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index c71fc7e8..ba3b9586 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -87,6 +87,7 @@ VALID_MODELS=( \ "DeepSeek-R1" \ "Qwen3-32B" \ "Qwen3-30B-A3B" \ + "GLM-5.1-FP8" \ ) # Models allowed for CONNECTOR=moriio WIDE_EP=1 (MoRI-EP; legacy RUN_MORI=1) @@ -94,6 +95,7 @@ MORI_EP_VALID_MODELS=( \ "DeepSeek-V3" \ "DeepSeek-V3-5layer" \ "DeepSeek-R1" \ + "GLM-5.1-FP8" \ ) # Models allowed for CONNECTOR=rixl WIDE_EP=1 EP_BACKEND=deepep (legacy RUN_DEEPEP=1) @@ -179,7 +181,10 @@ WIDE_EP="${WIDE_EP:-0}" # the MoRI-EP / DeepEP recipe (block=16, MLA off, per-role cudagraph). Running them # in TP mode is unsupported — the TP argv would double the model's own # --compilation-config and drop the mandatory +quant_fp8 op. Reject early. -WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" ) +# GLM-5.1-FP8 (GlmMoeDsaForCausalLM, MLA+DSA) is validated only under MoRI-EP +# wideEP disagg (block=1, AITER sparse MLA on, per-role all2all). The moriio+TP +# ("Stage B") path is untested for DSA, so reject WIDE_EP=0 for it too. +WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" "GLM-5.1-FP8" ) model_is_wide_ep_only() { local m="$1" for x in "${WIDE_EP_ONLY_MODELS[@]}"; do [[ "$m" == "$x" ]] && return 0; done From 3da0a33310001d8e15628ddff1ec1c8bd3653a99 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 3 Jul 2026 09:02:45 +0000 Subject: [PATCH 02/13] vllm_dissag: make GLM gate patcher atomic; verify against mori v1.2.1 image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../vllm_dissag/apply_glm_dsa_moriio_gate_fix.py | 15 +++++++++++++++ scripts/vllm_dissag/models.yaml | 7 +++++++ 2 files changed, 22 insertions(+) diff --git a/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py b/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py index 9cec5e06..92501ae9 100755 --- a/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py +++ b/scripts/vllm_dissag/apply_glm_dsa_moriio_gate_fix.py @@ -99,6 +99,21 @@ def main() -> int: print("[glm-gate] connector/engine not found -- skipping (layout differs).") return 0 + # ATOMIC: both halves (connector defines num_transfer_layers, engine gates on + # it) are needed together or not at all. On a restructured image (e.g. mori + # v1.2.1, whose engine replaced the writes_done>=num_layers gate with a sealed + # writes_expected mechanism that already handles hybrid/DSA dual-KV natively), + # the engine anchor is gone. Applying only the connector half would inject a + # dead num_transfer_layers into restructured internals. So if the engine anchor + # is absent, skip BOTH — the native gate already does the right thing. + eng_src = open(eng).read() + eng_gate_present = " if request_info.writes_done >= self.worker.num_layers:" in eng_src + eng_already = "num_transfer_layers" in eng_src + if not eng_gate_present and not eng_already: + print("[glm-gate] engine gate anchor absent (image restructured, e.g. mori " + "v1.2.1 sealed writes_expected) -- skipping BOTH halves (native gate handles DSA).") + return 0 + rc = patch_connector(conn) or patch_engine(eng) if rc: return rc diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 3b721990..d46fce6b 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -185,6 +185,13 @@ DeepSeek-R1: # Sparse Attention (DSA). 78 layers (3 dense + 75 MoE), 256 routed experts top-8 # + 1 shared, FP8 block 128. wideEP-only (see WIDE_EP_ONLY_MODELS in the slurm). # +# Validated DOCKER_IMAGE_NAME (submit-time, not set here — #171 requires it explicit): +# rocmshared/pytorch-private:vllm-wideep_06_29_2026_Shiksha_dp16_2p2d_mori_v1.2.1_aiter_v0.1.16.post3_nightlybase_mori121 +# On this image the DSA patchers self-adapt: only the invalid-token kernel fix +# (#45324) applies; the MoRIIO dual-KV geometry + completion-gate fixes are NATIVE +# (moriio_layout.py per-layer geometry + engine sealed writes_expected), so those +# patchers cleanly no-op. On the older b10a9f7a image all 4 patchers apply. +# # GLM differs from DeepSeek in 3 recipe-defining ways (both are MLA MoE, but GLM # is DSA-sparse): # - KV_BLOCK_SIZE=1 (DSA sparse indexer REQUIRES block-size 1; DS uses 16) From dcee1a682d3d82bb3fc242b48be3fd1fb26f9a19 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 3 Jul 2026 10:10:01 +0000 Subject: [PATCH 03/13] vllm_dissag: models.yaml env must override image-baked ENV (3-tier precedence) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- scripts/vllm_dissag/run_xPyD_models.slurm | 17 ++++++++++++++ scripts/vllm_dissag/vllm_disagg.sh | 28 ++++++++++++++++++----- 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index ba3b9586..7449bd0e 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -68,6 +68,22 @@ for f in "${REQUIRED_FILES[@]}"; do done echo "Running from: $(pwd)" +# ------------------------------------------------------------------------------ +# models.yaml env precedence: capture which recipe knobs the USER explicitly set +# at submit time. The driver (vllm_disagg.sh) uses this to let models.yaml `env:` +# OVERRIDE image-baked ENV defaults (e.g. a DeepSeek-tuned image bakes +# KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0, which would otherwise shadow a +# model's own recipe — GLM-5.1 DSA needs block=1 + AITER MLA on), while a genuine +# submit-time `-e VAR=...` still wins. Precedence: image-baked < models.yaml < submit -e. +# Captured HERE (before the slurm sets any defaults) so it reflects user intent only. +_RECIPE_ENV_KEYS="VLLM_USE_V1 VLLM_ROCM_USE_AITER VLLM_ROCM_USE_AITER_RMSNORM VLLM_ROCM_USE_AITER_MLA KV_BLOCK_SIZE KV_CACHE_DTYPE KV_CACHE_MEMORY_BYTES GPU_MEMORY_UTILIZATION VLLM_CUDAGRAPH_MODE PREFILL_CUDAGRAPH_MODE DECODE_CUDAGRAPH_MODE CUDAGRAPH_CAPTURE_SIZES VLLM_ALL2ALL_BACKEND PREFILL_MORI_BACKEND DECODE_MORI_BACKEND MORI_SHMEM_HEAP_SIZE" +MODELS_YAML_PROTECT="" +for _k in $_RECIPE_ENV_KEYS; do + [ -n "${!_k+x}" ] && MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT} ${_k}" +done +export MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT# }" +echo "models.yaml protect-list (submit-time overrides): '${MODELS_YAML_PROTECT}'" + # ------------------------ # Print current time in UTC and PST formats # ------------------------ @@ -588,6 +604,7 @@ docker run --rm \ ${VLLM_ALL2ALL_BACKEND:+-e VLLM_ALL2ALL_BACKEND=$VLLM_ALL2ALL_BACKEND} \ ${PREFILL_MORI_BACKEND:+-e PREFILL_MORI_BACKEND=$PREFILL_MORI_BACKEND} \ ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ + -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ diff --git a/scripts/vllm_dissag/vllm_disagg.sh b/scripts/vllm_dissag/vllm_disagg.sh index 06fbf84f..acbabbd4 100755 --- a/scripts/vllm_dissag/vllm_disagg.sh +++ b/scripts/vllm_dissag/vllm_disagg.sh @@ -160,17 +160,33 @@ MODEL_CONFIG_DECODE="" if [[ -n "$MODEL_NAME" && -f "$MODELS_YAML" ]]; then export MODELS_YAML MODEL_NAME PARALLEL_MODE # 1) Export per-model env: block FIRST (so connector ${VAR:-default} yields to it). - # Only set a var that is NOT already in the environment, so a submit-time - # `docker -e VAR=...` (already exported) WINS over the yaml value. Precedence: - # connector default < models.yaml env: < submit-time -e. + # Precedence: image-baked ENV < models.yaml env: < submit-time -e. + # models.yaml MUST override image-baked ENV: a DeepSeek-tuned disagg image + # bakes KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0 / VLLM_CUDAGRAPH_MODE= + # PIECEWISE etc. as container ENV, which would otherwise shadow a model's own + # recipe (GLM-5.1 DSA needs block=1 + AITER sparse MLA on). But a genuine + # submit-time `-e VAR=...` must still win. The slurm can tell the two apart + # (it runs on the host) and passes MODELS_YAML_PROTECT = the space-separated + # list of keys the USER set at submit; the driver protects only those. When + # MODELS_YAML_PROTECT is unset (script run directly, no slurm), fall back to + # the old "skip if in env" behavior so nothing regresses. _yaml_env="$(python3 - <<'PY' import os, yaml, shlex m = yaml.safe_load(open(os.environ["MODELS_YAML"])) or {} cfg = m.get(os.environ["MODEL_NAME"]) or {} +protect_raw = os.environ.get("MODELS_YAML_PROTECT") +have_protect = protect_raw is not None +protect = set((protect_raw or "").split()) for k, v in (cfg.get("env") or {}).items(): - # skip if already present in the environment (submit-time -e override wins) - if k in os.environ: - continue + if have_protect: + # 3-tier: yaml overrides baked ENV; only a user submit-time -e (in the + # protect-list) wins over yaml. + if k in protect: + continue + else: + # No protect-list (direct run): legacy behavior — any existing env wins. + if k in os.environ: + continue print(f'export {k}={shlex.quote(str(v))}') PY )" From 02873a702c65468f2eab4683ae2a79051da3dabe Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 3 Jul 2026 10:13:12 +0000 Subject: [PATCH 04/13] vllm_dissag: GLM-5.1 decode PIECEWISE cudagraph, prefill eager MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- scripts/vllm_dissag/models.yaml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index d46fce6b..ba74a94e 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -196,7 +196,10 @@ DeepSeek-R1: # is DSA-sparse): # - KV_BLOCK_SIZE=1 (DSA sparse indexer REQUIRES block-size 1; DS uses 16) # - VLLM_ROCM_USE_AITER_MLA=1 (GLM MLA path ON via AITER sparse; DS sets 0) -# - eager BOTH roles (DSA + cudagraph unproven; DS runs PIECEWISE decode) +# - prefill EAGER (NONE), decode PIECEWISE (prefill cudagraph capture deadlocks +# on this stack; decode PIECEWISE captures cleanly on DSA and is a ~3.7x ITL win +# (validated: ~69ms vs ~264ms eager). Global VLLM_CUDAGRAPH_MODE=NONE as the +# safe floor; per-role PREFILL=NONE / DECODE=PIECEWISE override it.) # Note: block=1 + AITER_MLA=1 are ALREADY the moriio.sh connector defaults — GLM # keeps them; DeepSeek is the one that overrides them off. Set here explicitly so # the recipe is self-documenting and robust to connector default changes. @@ -222,7 +225,7 @@ GLM-5.1-FP8: GPU_MEMORY_UTILIZATION: "0.80" VLLM_CUDAGRAPH_MODE: "NONE" PREFILL_CUDAGRAPH_MODE: "NONE" - DECODE_CUDAGRAPH_MODE: "NONE" + DECODE_CUDAGRAPH_MODE: "PIECEWISE" CUDAGRAPH_CAPTURE_SIZES: "1 2 4 8 16 32 64 128 256" VLLM_ALL2ALL_BACKEND: "mori_high_throughput" PREFILL_MORI_BACKEND: "mori_high_throughput" From a6c16f065bf5b7bb11821697c6e6bb62e8d196d4 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 3 Jul 2026 10:47:40 +0000 Subject: [PATCH 05/13] vllm_dissag: guard MoRIIO abort path against None peer_zmq (mori v1.2.1) 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. --- .../apply_glm_moriio_abort_guard_fix.py | 98 +++++++++++++++++++ scripts/vllm_dissag/connectors/moriio.sh | 3 +- 2 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py diff --git a/scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py b/scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py new file mode 100644 index 00000000..c09b4acd --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_moriio_abort_guard_fix.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Guard the MoRIIO connector abort path against a None peer_zmq (mori v1.2.1). + +ROOT CAUSE (observed on the router image, job 199144 decode crash): + When a request is ABORTED before its KV-transfer peer handshake completes, + the connector's release path runs: + + moriio_connector.py::_release_write_prefill_blocks + peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False) # -> None + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq) # None.split(",") + -> AttributeError: 'NoneType' object has no attribute 'split' + + This only catches ValueError, not the AttributeError from a None peer_zmq, so + the EngineCore dies -> cascades to all decode workers (EngineDeadError) -> decode + is dead. Triggered by any request aborted before the peer handshake (e.g. a + canary/curl that times out during first-token cold JIT). + + The SAME FILE already guards this correctly at the other call site + (request_finished / _should_notify path): `if peer_zmq is not None:` then parse, + else fall back to params. The release path just missed the guard — an + inconsistent-guard bug in the connector. + +FIX (surgical, matches the file's own existing pattern): + In _release_write_prefill_blocks, when the params don't already carry + remote_host/remote_notify_port, guard the peer_zmq lookup: if it is None, log + and return (same graceful bail the existing `except ValueError` already does for + the "missing remote notify address" case). No behavior change when peer_zmq is + valid; single-geometry / non-aborted requests are unaffected. + +Idempotent + anchor-based + self-skipping (no-ops if the anchor is absent/already +guarded, so it is safe across connector revisions and other images). A found-old +anchor that fails to apply is a hard error (would leave the crash). + +Usage: apply_glm_moriio_abort_guard_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + +# The buggy two lines: fetch peer_zmq (may be None) then parse it unguarded. +OLD = """ peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False) + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)""" + +NEW = """ peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False) + # Abort-path guard: a request aborted before the KV peer handshake + # has peer_zmq=None; parse_moriio_zmq_address(None) would raise + # AttributeError and kill the EngineCore. Bail gracefully like the + # ValueError case below (matches the guarded call site elsewhere). + if peer_zmq is None: + logger.warning( + "Cannot release WRITE prefill blocks for request %s: " + "no peer zmq address (aborted before peer handshake)", + request_id, + ) + return + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)""" + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-abort] {REL} not found -- skipping (connector layout differs).") + return 0 + + src = open(path).read() + if "no peer zmq address (aborted before peer handshake)" in src: + print("[glm-abort] already patched -- no-op.") + return 0 + if OLD not in src: + # Anchor absent: either the release path was refactored or this image + # already guards it. Do not block launch. + print("[glm-abort] release-path anchor not found -- skipping (assuming " + "native guard / refactored).") + return 0 + + src = src.replace(OLD, NEW, 1) + try: + open(path, "w").write(src) + except OSError as e: + print(f"[glm-abort] ERROR: write failed for {path}: {e}", file=sys.stderr) + return 1 + + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: # noqa: BLE001 + print(f"[glm-abort] ERROR: patched file fails to compile: {e}", file=sys.stderr) + return 1 + print(f"[glm-abort] patched _release_write_prefill_blocks None-guard in {path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index b13dab9a..c97e7a57 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -162,7 +162,8 @@ _glm_dsa_runtime_patch() { apply_glm_dsa_kernel_fix.py \ apply_glm_dsa_moriio_dualkv_fix.py \ apply_glm_dsa_moriio_engine_fix.py \ - apply_glm_dsa_moriio_gate_fix.py; do + apply_glm_dsa_moriio_gate_fix.py \ + apply_glm_moriio_abort_guard_fix.py; do local _py="${_patch_dir}/${_p}" if [ ! -f "${_py}" ]; then echo "Error: [glm] required patcher ${_py} not found. Aborting." >&2 From 273ac540df707018e9090da0e63bf8a2c726f189 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Fri, 3 Jul 2026 17:47:54 +0000 Subject: [PATCH 06/13] vllm_dissag: document GLM-5.1 long-context prefill limit (vLLM #40018) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- scripts/vllm_dissag/models.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index ba74a94e..6d635348 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -214,6 +214,17 @@ DeepSeek-R1: # connector's model_args. block-size / kv-cache-dtype / all2all / cudagraph come # from the env: recipe above (the connector emits them), so the dp: blocks are # empty like the DeepSeek family. +# +# LONG-CONTEXT CAVEAT (vLLM #40018): the ROCM_AITER_MLA_SPARSE prefill indexer +# corrupts output for prompts beyond ~16-18k tokens on this image (mori v1.2.1) — +# coherent + correct needle retrieval up to 14k, garbage (repetition collapse, +# unique-ratio ~0.1) at ~18.7k. This is an UPSTREAM kernel bug, not the MAD port. +# TESTED: pinning --max-model-len 32768 does NOT move the threshold (workspace is +# sized max_model_len*40 but the corruption onset is a fixed ~18k token count in +# the gather/logits kernel, not a buffer-scaling artifact). So no config knob +# helps; it needs the complete upstream prefill fix in a newer image. Left at the +# native max_model_len (do not cap — capping gives no accuracy benefit and only +# limits usable context). Serve prompts <~14k for correct output on this image. GLM-5.1-FP8: env: VLLM_USE_V1: "1" From 2c5d6040382c57b9d583b4d5ca83f8641f88c8f7 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Tue, 7 Jul 2026 06:41:02 +0000 Subject: [PATCH 07/13] vllm_dissag: GLM-5.1 DSA long-context accuracy fix (persistent-MLA gate) 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 --- ...pply_glm_dsa_persistent_kernel_gate_fix.py | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py diff --git a/scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py b/scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py new file mode 100644 index 00000000..adda36cf --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_persistent_kernel_gate_fix.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Gate OFF the AITER persistent sparse-MLA kernel for chunked-prefill batches. + +ROOT CAUSE (ROCm/aiter #4076, vLLM #47042 / #47567): + The AITER persistent MLA work-stealing kernel (mla_a8w8_qh16_qseqlen1_gqaratio16_ps, + taken when work_meta_data from get_mla_metadata_v1 is non-None) is NUMERICALLY WRONG + for multi-token (prefill-shaped) batches of qseqlen==1 entries. Pure decode (1 query + token) and fresh single-chunk prefills are correct; the error only appears once a + request becomes a CHUNKED-PREFILL CONTINUATION. The small per-token error COMPOUNDS + through the KV cache across chunked-prefill passes until long-context decode collapses + into repetition/garbage. Failure is gated on CHUNK COUNT, not raw context length + (verified: 22k in 2 chunks = correct, 22k in 3 chunks = garbage). + + On this image (aiter 0.1.16.post3, before the aiter-side kernel fix #3921) GLM-5.1-FP8 + DSA collapses at ~16-18k prompt tokens. The aiter kernel fix is the long-term answer + (AITERKER-132 / aiter #3921); this is the vLLM-side short-term gate (#47567), which + costs ~no perf (decode + single-chunk prefill keep the persistent path). + +FIX (port of vLLM PR #47567, adapted to this image's rocm_aiter_mla_sparse.py::build): + In ROCMAiterMLASparseMetadataBuilder.build(), detect chunked-prefill continuations + (a request with >1 query token this step whose total seq_len exceeds its query_len, + i.e. part of its context was computed in an earlier chunk) and, when ANY request in + the batch is such a continuation: + * skip the get_mla_metadata_v1 persistent-metadata launch, and + * pass work_meta_data=None to the metadata so mla_decode_fwd takes the CORRECT + non-persistent split-KV path. + Decode-only and single-chunk-prefill batches are unchanged (persistent path kept). + + Uses `seg_lengths` (per-request step query lengths, already computed at build() top) + and `common_attn_metadata.seq_lens_cpu[:num_reqs].numpy()` (total seq lens). Both are + present in this image's build(). + +Idempotent + anchor-based + self-skipping. Missing anchor -> warn+skip (safe across +image revisions / if a newer image already carries the aiter kernel fix). A found-old +anchor that fails to apply is a hard error (would silently keep the corruption). + +Usage: apply_glm_dsa_persistent_kernel_gate_fix.py +""" +import os +import sys + +REL = "v1/attention/backends/mla/rocm_aiter_mla_sparse.py" + +# Anchor 1: the persistent-metadata guard. We insert the continuation detection +# just before it and AND it into the condition. +OLD1 = """ if metadata_key != self._prev_metadata_key: + from aiter import get_mla_metadata_v1""" +NEW1 = """ # PERSISTENT-KERNEL GATE (aiter #4076 / vLLM #47567): the persistent + # sparse-MLA work-stealing kernel is numerically wrong for chunked-prefill + # continuation batches; the error compounds and breaks long-context decode. + # Fall back to the correct non-persistent path whenever any request in the + # batch is a chunked-prefill continuation (>1 query token this step AND + # total seq_len > this step's query_len). Decode + single-chunk prefills + # keep the fast persistent path -> no decode-throughput regression. + # Slice to num_reqs and cast to int64 (vLLM #47567 hardening / Rohan138 PR#1) + # so the masks cannot broadcast-mismatch under cudagraph padding. + _step_query_lens = seg_lengths[:num_reqs].astype(np.int64) + _total_seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs].numpy().astype( + np.int64 + ) + _is_chunked_continuation = (_step_query_lens > 1) & ( + _total_seq_lens > _step_query_lens + ) + _use_persistent = not bool(_is_chunked_continuation.any()) + if _use_persistent and metadata_key != self._prev_metadata_key: + from aiter import get_mla_metadata_v1""" + +# Anchor 2: the metadata construction passes the persistent buffer unconditionally. +# Gate it on _use_persistent. +OLD2 = " work_meta_data=self._mla_work_meta_data," +NEW2 = " work_meta_data=(self._mla_work_meta_data if _use_persistent else None)," + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[glm-persist] {REL} not found -- skipping (backend layout differs).") + return 0 + + src = open(path).read() + + if "_is_chunked_continuation" in src or "_use_persistent" in src: + print("[glm-persist] already patched (persistent-kernel gate present) -- no-op.") + return 0 + + # Both anchors must be present to apply safely. + if OLD1 not in src: + print("[glm-persist] WARN: persistent-metadata anchor (metadata_key guard) not " + "found -- skipping (image may already carry the aiter kernel fix, or the " + "backend was refactored).") + return 0 + if OLD2 not in src: + print("[glm-persist] ERROR: found the metadata_key guard but NOT the " + "work_meta_data=self._mla_work_meta_data assignment -- refusing partial " + "patch (would leave persistent kernel active). Aborting.", file=sys.stderr) + return 1 + + src = src.replace(OLD1, NEW1, 1) + src = src.replace(OLD2, NEW2, 1) + + try: + open(path, "w").write(src) + except OSError as e: + print(f"[glm-persist] ERROR: write failed for {path}: {e}", file=sys.stderr) + return 1 + + # Verify both edits landed. + chk = open(path).read() + if "_use_persistent = not bool(_is_chunked_continuation.any())" not in chk or \ + "if _use_persistent else None" not in chk: + print("[glm-persist] ERROR: post-write verification failed.", file=sys.stderr) + return 1 + + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: # noqa: BLE001 + print(f"[glm-persist] ERROR: patched file fails to compile: {e}", file=sys.stderr) + return 1 + + print(f"[glm-persist] patched persistent-kernel gate (aiter #4076 / vLLM #47567) in {path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 63d85833bcc0a379f95acc34fcb84c46f5f3cd62 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Tue, 7 Jul 2026 06:41:20 +0000 Subject: [PATCH 08/13] vllm_dissag: GLM-5.1 disagg crash-investigation fixes + knobs (WIP) 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 --- .../apply_glm_aiter_sampling_oob_fix.py | 148 +++++++++++ .../apply_glm_dsa_indexer_warmup_fix.py | 251 ++++++++++++++++++ scripts/vllm_dissag/connectors/moriio.sh | 40 ++- scripts/vllm_dissag/keepalive_bench.sh | 18 ++ scripts/vllm_dissag/models.yaml | 22 ++ scripts/vllm_dissag/run_xPyD_models.slurm | 7 +- 6 files changed, 484 insertions(+), 2 deletions(-) create mode 100644 scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py create mode 100755 scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py create mode 100755 scripts/vllm_dissag/keepalive_bench.sh diff --git a/scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py b/scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py new file mode 100644 index 00000000..d491b838 --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_aiter_sampling_oob_fix.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Overlay the fixed AITER sampling kernel (ROCm/aiter #3658 + hang cap) into the image. + +DEFECT 2 (the 8k prefill/decode crash): the AITER TopP/TopK sampling kernel +(csrc/cpp_itfs/sampling/sampling.cuh) has two bugs on the released aiter post3 +that this image ships: + + 1. HSA OUT-OF-BOUNDS (ROCm/aiter #3658): SamplingTempStorage::last_valid_id is + never initialized. When a probs row is all-zero / NaN (more likely at long + input, e.g. 8k), the guarded write-back (max_valid != -1) is skipped, and the + fallback `sampled_id = temp_storage.last_valid_id` reads UNINITIALIZED shared + memory -> garbage index -> `probs[row*d + sampled_id]` dereferences OOB and + HSA page-faults ("Memory access fault by GPU node-N"). Deterministic under + CUDA graph (shared-mem residue is stable across replays). This is the silent + worker death at 8k that collapses the disagg DP group -> 503. + Fix: init `last_valid_id = 0` at top of each loop iter + defensive clamp on + the loaded sampled_id before it indexes probs. + + 2. REJECTION-SAMPLING HANG: the bisection `do { ... } while(low < high)` can + spin forever when the [low,high] interval stagnates in float precision on a + degenerate (near-uniform) row -> never-completing HSA signal / hang (the + "sampler hang" that forced the skip-warmup workaround). Fix: cap the loop at + kMaxSamplingRounds=32 (float32 mantissa is exhausted well within 32 rounds, + so healthy distributions always converge via break long before the cap). + +Both fixes land in sampling.cuh. #3658 is MERGED upstream but NOT in the released +aiter post3 (this image). Source: A/B-tested by Shiksha (shikpate); staged fixed +tree at SAMPLING_FIX_DIR. + +METHOD (from Shiksha's validated in-container overlay): copy the whole patched +sampling source dir (.cuh + .py + .jinja) over the container's aiter, then purge +any compiled sampling JIT objects so the kernel recompiles from the fixed source +on next use. + +Idempotent: skips if the fix markers are already present. Model-agnostic at the +kernel level, but invoked from the GLM patch hook. Safe no-op if the staged fix +dir or the target aiter dir is absent. + +Usage: apply_glm_aiter_sampling_oob_fix.py + (vllm_install_dir arg is accepted for hook uniformity but not required; + the aiter dir is resolved via `import aiter`.) +""" +import os +import shutil +import subprocess +import sys + +FIX_DIR = os.environ.get( + "SAMPLING_FIX_DIR", + "/shared_inference/ravgupta/aiter_sampling_fix_3658/sampling_patched", +) +MARKERS = ("last_valid_id = 0", "kMaxSamplingRounds") + + +def _aiter_sampling_dir(): + """Locate the installed aiter sampling source dir (aiter_meta/csrc/...).""" + try: + import aiter # noqa: F401 + except Exception as e: # noqa: BLE001 + print(f"[sampling-fix] aiter not importable ({e}); skipping.") + return None + # The kernel source lives under aiter_meta (sibling of aiter), path is stable. + candidates = [] + try: + import aiter_meta # type: ignore + + candidates.append( + os.path.join(os.path.dirname(aiter_meta.__file__), + "csrc", "cpp_itfs", "sampling") + ) + except Exception: # noqa: BLE001 + pass + # Fallback: search site-packages. + import aiter + sp = os.path.dirname(os.path.dirname(aiter.__file__)) + candidates.append(os.path.join(sp, "aiter_meta", "csrc", "cpp_itfs", "sampling")) + for c in candidates: + if os.path.isdir(c): + return c + print(f"[sampling-fix] could not locate aiter sampling dir (tried {candidates}); skipping.") + return None + + +def main() -> int: + tgt = _aiter_sampling_dir() + if tgt is None: + return 0 # safe no-op + + tgt_cuh = os.path.join(tgt, "sampling.cuh") + if os.path.isfile(tgt_cuh): + cur = open(tgt_cuh, errors="ignore").read() + if all(m in cur for m in MARKERS): + print(f"[sampling-fix] already applied (markers present) in {tgt_cuh}.") + return 0 + + if not os.path.isdir(FIX_DIR): + print(f"[sampling-fix] WARN: staged fix dir {FIX_DIR} not found; leaving image kernel unpatched.", file=sys.stderr) + return 0 + + src_cuh = os.path.join(FIX_DIR, "sampling.cuh") + if not os.path.isfile(src_cuh) or not all(m in open(src_cuh, errors="ignore").read() for m in MARKERS): + print(f"[sampling-fix] WARN: staged {src_cuh} missing/lacks fix markers; skipping.", file=sys.stderr) + return 0 + + # Overlay the whole sampling source dir (.cuh + .py + .jinja), per Shiksha's method. + copied = [] + for fn in os.listdir(FIX_DIR): + s = os.path.join(FIX_DIR, fn) + if os.path.isfile(s): + shutil.copy2(s, os.path.join(tgt, fn)) + copied.append(fn) + print(f"[sampling-fix] overlaid #3658 + hang-cap into {tgt}: {', '.join(sorted(copied))}") + + # Verify. + cur = open(tgt_cuh, errors="ignore").read() + if not all(m in cur for m in MARKERS): + print(f"[sampling-fix] ERROR: markers still absent after overlay in {tgt_cuh}.", file=sys.stderr) + return 1 + + # Purge any compiled sampling JIT objects so the kernel recompiles from source. + purged = 0 + for base in ( + os.path.expanduser("~/.aiter"), "/root/.aiter", "/tmp/aiter", + "/opt/vllm_cache/aiter_jit", os.path.join(os.path.dirname(tgt), "..", "..", "jit"), + ): + if base and os.path.isdir(base): + try: + out = subprocess.run( + ["find", base, "-maxdepth", "6", "-iname", "*sampling_from_probs*"], + capture_output=True, text=True, timeout=60, + ) + for p in out.stdout.split(): + try: + if os.path.isdir(p): + shutil.rmtree(p, ignore_errors=True) + else: + os.remove(p) + purged += 1 + except OSError: + pass + except Exception: # noqa: BLE001 + pass + print(f"[sampling-fix] purged {purged} stale sampling JIT object(s); kernel will recompile from fixed source.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py b/scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py new file mode 100755 index 00000000..3eaef8cf --- /dev/null +++ b/scripts/vllm_dissag/apply_glm_dsa_indexer_warmup_fix.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +"""Force-compile the GLM-5.1 DSA sparse-attention indexer Triton kernels at BOOT. + +PROBLEM (root cause of the "first big prompt stalls the whole DP group" hang): + The DSA indexer's Triton kernels are seq-length specialized: + - v1/attention/ops/triton_fp8_mqa_logits.py flips `matrix_instr_nonkdim` + at seq_len<=1024 and launches with grid=[(seq_len,)] (seq_len is a + specialized kernel arg) -> a >1024-row prefill needs a *different* JIT + specialization than a <=1024-row one. + But the boot-time warmup never drives the indexer: + - profile_run() calls _dummy_run(is_profile=True) with force_attention=False + and cudagraph mode NONE -> attn_metadata stays None -> sparse_attn_indexer + takes the `sparse_attn_indexer_fake` path (see the "careful! this will be + None in dummy run" comment in layers/sparse_attn_indexer.py). The real + kernels are never compiled. + - _warmup_and_capture() only sets force_attention=True when the cudagraph + runtime mode is FULL; the DSA indexer builder reports UNIFORM_BATCH, so on + this ROCm/DP build the mixed prefill-decode graphs are PIECEWISE and + force_attention stays False. Even when attention IS forced, capture uses + uniform-decode / small mixed batches -- never a large prefill at + max_num_batched_tokens -- so the >1024 specialization is still absent. + Effect: the first >=8k prompt JIT-compiles the indexer kernel mid-inference on + whichever DP rank happens to receive it. That rank falls out of the DP lockstep + gloo all_reduce (coordinate_batch_across_dp) while it compiles -> the whole DP + group collapses. It is also a general cold-cache robustness hole. + +FIX (surgical, reuses vLLM's OWN metadata construction -- no hand-synthesized +tensors, so zero risk of a bad-input crash at boot): + 1. gpu_model_runner.py: add a `_maybe_warmup_dsa_indexer()` method. It is a + strict NO-OP unless one of the runner's attention backends is (a subclass + of) DeepseekV32IndexerBackend. When present, it runs + `_dummy_run(..., force_attention=True, cudagraph_runtime_mode=NONE)` at TWO + prefill-size regimes -- a small one (<=1024 rows) and a large one + (max_num_batched_tokens, >1024) -- so BOTH Triton specializations compile. + `force_attention=True` makes _dummy_run build a real + DeepseekV32IndexerMetadata via the normal _build_attention_metadata path + (num_prefills>0 because the default dummy batch is multi-token requests and + the indexer decode_threshold is 1), which drives the real + `sparse_attn_indexer` prefill kernels. The whole thing is wrapped in + try/except that only WARNs -- a warmup failure must never crash boot. + 2. gpu_worker.py: call it from compile_or_warm_up_model, right after the + existing warmup loop and before kernel_warmup(). At that point the KV cache + is already allocated (initialize_from_config runs before + compile_or_warm_up_model), which the forced-attention indexer path needs. + +Idempotent + anchor-based (matches the other apply_glm_* patchers): + * Each hunk self-detects if already applied (marker string present) -> no-op. + * Missing anchor -> WARN and skip that hunk (safe across vllm revisions; the + rebase may already warm the indexer natively or have refactored the site). + * Anchor found but the file does not contain the applied marker and the + replace produces no change -> hard error (would silently keep the bug). + * py_compile at the end; hard error if the patched file won't compile. + +Usage: apply_glm_dsa_indexer_warmup_fix.py +""" +import os +import sys + +RUNNER_REL = "v1/worker/gpu_model_runner.py" +WORKER_REL = "v1/worker/gpu_worker.py" + +MARKER = "glm-dsa-indexer-warmup" + +# --- Hunk A: new method inserted immediately before `def capture_model` ------- +# Anchor: the (unique) capture_model definition head in gpu_model_runner.py. +RUNNER_ANCHOR = " def capture_model(self) -> int:\n" + +RUNNER_METHOD = ''' def _maybe_warmup_dsa_indexer(self) -> None: + """Force-compile the DSA sparse-attention indexer Triton kernels at boot. + + NO-OP unless this model actually has a DeepseekV32IndexerBackend (GLM-5.1 + DSA / DeepSeek V3.2). The indexer kernels are seq-length specialized + (triton_fp8_mqa_logits flips matrix_instr_nonkdim at seq_len<=1024 and + launches grid=[(seq_len,)]), and the normal profile/warmup passes never + drive the indexer (attn_metadata is None -> the *_fake path). Without this + the first large prompt JIT-compiles mid-inference and, under DP lockstep, + stalls the whole group. We warm BOTH regimes: a small (<=1024) and a large + (max_num_batched_tokens, >1024) prefill batch, using force_attention=True + so _dummy_run builds a real indexer metadata via the standard path. + """ + # {marker} + try: + from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerBackend, + ) + except Exception: # noqa: BLE001 -- backend module absent -> not a DSA build + return + + has_indexer = False + try: + for attn_group in self._attn_group_iterator(): + backend = getattr(attn_group, "backend", None) + if backend is not None and isinstance(backend, type) and issubclass( + backend, DeepseekV32IndexerBackend + ): + has_indexer = True + break + except Exception: # noqa: BLE001 -- iterator shape changed -> stay a no-op + return + if not has_indexer: + return + + # Two prefill-size regimes so both Triton specializations compile. + # Small must be <=1024 rows; large must exceed 1024 (use the real max). + max_tokens = int(self.max_num_tokens) + small = min(512, max_tokens) + sizes = [] + for s in (small, max_tokens): + if s > 0 and s not in sizes: + sizes.append(s) + + logger.info( + "Warming up DSA indexer kernels at prefill sizes %s " + "to avoid mid-inference JIT.", + sizes, + ) + for size in sizes: + try: + self._dummy_run( + size, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + force_attention=True, + skip_eplb=True, + remove_lora=False, + ) + except Exception as e: # noqa: BLE001 -- warmup must NEVER crash boot + logger.warning( + "DSA indexer warmup at size %d failed (%s); the kernel may " + "JIT-compile on first use instead.", + size, + e, + ) + self._sync_device() + +'''.replace("{marker}", MARKER) + +# --- Hunk B: call site in gpu_worker.compile_or_warm_up_model ----------------- +WORKER_ANCHOR = ( + " self.model_runner.maybe_remove_all_loras(" + "self.model_runner.lora_config)\n" + "\n" + " # Warmup and tune the kernels used during model execution before\n" + " # cuda graph capture.\n" + " kernel_warmup(self)\n" +) + +WORKER_REPLACEMENT = ( + " self.model_runner.maybe_remove_all_loras(" + "self.model_runner.lora_config)\n" + "\n" + " # " + MARKER + ": force-compile the DSA sparse-attention indexer\n" + " # Triton kernels now (KV cache is allocated), so a large prompt\n" + " # never JIT-compiles them mid-inference and stalls DP lockstep.\n" + " # No-op unless this model has a DeepseekV32IndexerBackend.\n" + " if hasattr(self.model_runner, \"_maybe_warmup_dsa_indexer\"):\n" + " self.model_runner._maybe_warmup_dsa_indexer()\n" + "\n" + " # Warmup and tune the kernels used during model execution before\n" + " # cuda graph capture.\n" + " kernel_warmup(self)\n" +) + + +def _patch_file(path, tag, anchor, apply_fn, already_marker): + """Return 0 on success/no-op, 1 on hard error.""" + if not os.path.isfile(path): + print(f"[{tag}] {path} not found -- skipping (layout differs).") + return 0 + src = open(path).read() + if already_marker in src: + print(f"[{tag}] already applied ({already_marker} present) in {path} -- no-op.") + return 0 + if anchor not in src: + print( + f"[{tag}] WARN: anchor not found in {path} -- skipping " + "(assuming native warmup / refactor)." + ) + return 0 + new_src = apply_fn(src) + if new_src == src: + print( + f"[{tag}] ERROR: anchor found but patch produced no change in {path}.", + file=sys.stderr, + ) + return 1 + try: + open(path, "w").write(new_src) + except OSError as e: + print(f"[{tag}] ERROR: failed to write patched {path}: {e}", file=sys.stderr) + return 1 + if already_marker not in open(path).read(): + print( + f"[{tag}] ERROR: post-write verification failed in {path}.", + file=sys.stderr, + ) + return 1 + print(f"[{tag}] patched {path} -- 1 hunk.") + return 0 + + +def main() -> int: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + vllm_dir = sys.argv[1] + + runner_path = os.path.join(vllm_dir, RUNNER_REL) + worker_path = os.path.join(vllm_dir, WORKER_REL) + + rc = 0 + + # Hunk A: insert the method before capture_model. + rc |= _patch_file( + runner_path, + "glm-dsa-warmup", + RUNNER_ANCHOR, + lambda s: s.replace(RUNNER_ANCHOR, RUNNER_METHOD + RUNNER_ANCHOR, 1), + MARKER, + ) + + # Hunk B: call it from compile_or_warm_up_model. + rc |= _patch_file( + worker_path, + "glm-dsa-warmup", + WORKER_ANCHOR, + lambda s: s.replace(WORKER_ANCHOR, WORKER_REPLACEMENT, 1), + MARKER, + ) + + if rc: + return 1 + + # py-compile sanity for whichever files exist. + try: + import py_compile + + for p in (runner_path, worker_path): + if os.path.isfile(p): + py_compile.compile(p, doraise=True) + print("[glm-dsa-warmup] py_compile OK") + except Exception as e: # noqa: BLE001 + print( + f"[glm-dsa-warmup] ERROR: patched file fails to compile: {e}", + file=sys.stderr, + ) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index c97e7a57..159bd662 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -137,6 +137,26 @@ connector_runtime_patch() { # code gaps, applied here as idempotent, anchor-based, self-skipping .py patchers # (they no-op cleanly if the fix is native/refactored on the chosen image). Gated # on MODEL_NAME so DeepSeek/other models are a pure no-op (byte-identical to before). + # MoRI library bump (opt-in): if a rebuilt MoRI package is staged at + # MORI_OVERLAY_DIR (default /shared_inference/$USER/mori_built/mori), overlay it + # over the image's stock MoRI BEFORE any mori import. The stock image ships MoRI + # v1.2.1 whose KV-transfer notify/mapping fails on large (>~8k) transfers + # (decode "unmap MISS" -> deferred-write expires 60s -> prefill EngineCore crash). + # Post-1.2.1 MoRI main carries the fixes (#424 notify SENDs, #436 large-transfer + # chunking, #432 RDMA resilience). Built from source against THIS image's libpci + # (nightly wheel is ABI-incompatible: needs LIBPCI_3.8, image has 3.7). Applies to + # ALL models (MoRI is model-agnostic), so this runs before the GLM gate. + _mori_overlay="${MORI_OVERLAY_DIR:-/shared_inference/${USER_NAME:-$USER}/mori_built/mori}" + if [ -d "${_mori_overlay}" ]; then + _mori_dst="$(python3 -c 'import mori,os;print(os.path.dirname(mori.__file__))' 2>/dev/null || true)" + if [ -n "${_mori_dst}" ] && [ -d "${_mori_dst}" ]; then + echo "[mori-bump] overlaying rebuilt MoRI from ${_mori_overlay} onto ${_mori_dst}" + rm -rf "${_mori_dst}" && cp -a "${_mori_overlay}" "${_mori_dst}" && \ + echo "[mori-bump] now: $(python3 -c 'import mori;print(mori.__version__)' 2>&1 | tail -1)" || \ + echo "[mori-bump] WARN: overlay failed; continuing with stock MoRI" >&2 + fi + fi + [ "${MODEL_NAME:-}" = "GLM-5.1-FP8" ] || return 0 _glm_dsa_runtime_patch } @@ -157,13 +177,19 @@ _glm_dsa_runtime_patch() { echo "[glm] MODEL_NAME=GLM-5.1-FP8: applying DSA runtime patchers against ${_vllm_dir}" # Ordered list of REQUIRED patchers (all abort on hard failure). + # GLM_PERSIST_GATE=0 skips the persistent-MLA accuracy gate (debug only: to test + # whether the non-persistent kernel it routes to is what crashes disagg at >=8k). + local _gate_patcher="apply_glm_dsa_persistent_kernel_gate_fix.py" + [ "${GLM_PERSIST_GATE:-1}" = "0" ] && _gate_patcher="" local _p for _p in \ apply_glm_dsa_kernel_fix.py \ apply_glm_dsa_moriio_dualkv_fix.py \ apply_glm_dsa_moriio_engine_fix.py \ apply_glm_dsa_moriio_gate_fix.py \ - apply_glm_moriio_abort_guard_fix.py; do + apply_glm_moriio_abort_guard_fix.py \ + ${_gate_patcher} \ + apply_glm_aiter_sampling_oob_fix.py; do local _py="${_patch_dir}/${_p}" if [ ! -f "${_py}" ]; then echo "Error: [glm] required patcher ${_py} not found. Aborting." >&2 @@ -176,6 +202,18 @@ _glm_dsa_runtime_patch() { } done + # Optional DSA indexer boot-warmup (GLM_INDEXER_WARMUP=1). Force-compiles the DSA + # indexer kernels at boot so they never JIT mid-inference. Opt-in because it drives a + # large (>=8k) prefill forward at boot: on stacks where that forward faults it makes + # the fault DETERMINISTIC at boot (useful for debugging) rather than on first request. + if [ "${GLM_INDEXER_WARMUP:-0}" = "1" ]; then + local _warm="${_patch_dir}/apply_glm_dsa_indexer_warmup_fix.py" + if [ -f "${_warm}" ]; then + echo "[glm] applying DSA indexer boot-warmup (GLM_INDEXER_WARMUP=1)" + python3 "${_warm}" "${_vllm_dir}" 2>&1 || echo "Warning: [glm] indexer-warmup patch failed (non-fatal)." + fi + fi + # Optional diagnostic instrumentation (GLM_INSTRUMENT=1). Non-fatal. if [ "${GLM_INSTRUMENT:-0}" = "1" ]; then local _instr="${_patch_dir}/apply_glm_dsa_moriio_instrument.py" diff --git a/scripts/vllm_dissag/keepalive_bench.sh b/scripts/vllm_dissag/keepalive_bench.sh new file mode 100755 index 00000000..68d80dcb --- /dev/null +++ b/scripts/vllm_dissag/keepalive_bench.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Keepalive with LIGHT heartbeat traffic: holds the disagg server up AND sends a +# tiny request every ~20s so prefill discovery/ping stays registered (idle sleep +# lets the prefill ZMQ ping die ~2min in). Runs KEEPALIVE_MINS (default 90). +: "${KEEPALIVE_MINS:=90}" +PORT="${BENCHMARK_PORT:-30000}" +MODEL="/mnt/m2m_nobackup/models_blog/GLM-5.1-FP8" +echo "[keepalive] light-traffic hold ${KEEPALIVE_MINS}min on :${PORT}" +_end=$(( $(date +%s) + KEEPALIVE_MINS*60 )) +i=0 +while [ "$(date +%s)" -lt "$_end" ]; do + curl -s -m 30 "http://127.0.0.1:${PORT}/v1/completions" -H "Content-Type: application/json" \ + -d "{\"model\":\"${MODEL}\",\"prompt\":\"hi\",\"max_tokens\":1,\"temperature\":0}" >/dev/null 2>&1 + i=$((i+1)) + [ $((i % 3)) -eq 0 ] && echo "[keepalive] heartbeat $i, $(( (_end-$(date +%s))/60 ))min left" + sleep 20 +done +echo "[keepalive] done" diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 6d635348..60fbd2fb 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -242,6 +242,28 @@ GLM-5.1-FP8: PREFILL_MORI_BACKEND: "mori_high_throughput" DECODE_MORI_BACKEND: "mori_low_latency" MORI_SHMEM_HEAP_SIZE: "17179869184" + # DSA sparse-indexer logits-buffer cap (crash fix). The indexer prefill computes an + # M*N fp32 logits buffer; split_indexer_prefill_chunks only sub-chunks the query dim + # when M*N*4 > this budget. Default 512MB lets an 8192-token prefill build a single + # 268MB buffer + launch the fp8_mqa_logits kernel at grid=(8192,), which HARD-FAULTS + # the worker on gfx942 (silent GPU fault -> DP group collapse -> 503 at >=8k prompts). + # Capping at 64MB forces M-dim sub-chunking (~2k tokens/chunk) so the buffer and the + # kernel launch stay bounded. Root cause: vllm/v1/attention/ops/triton_fp8_mqa_logits.py + # fp8_mqa_logits_gfx942; chunking logic: mla/indexer.py split_indexer_prefill_chunks. + VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: "64" + # NCCL heartbeat watchdog: at long context (>~8k) a DP rank's sparse-MLA/MoE + # all2all collective can exceed the default HeartbeatMonitor timeout -> + # ProcessGroupNCCL::HeartbeatMonitor::runLoop() declares the rank dead and + # tears down the whole process group -> prefill EngineCore crashes -> 503. + # (Confirmed root cause of the 8k+ prefill crash; #338 EP-landmine.) Disable + # the monitor-triggered teardown and extend timeouts so long-ctx collectives + # complete instead of being watchdog-killed. + TORCH_NCCL_ENABLE_MONITORING: "0" + TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC: "1800" + TORCH_NCCL_DUMP_ON_TIMEOUT: "0" + TORCH_NCCL_BLOCKING_WAIT: "0" + TORCH_NCCL_ASYNC_ERROR_HANDLING: "1" + NCCL_IB_TIMEOUT: "22" dp_flags: "--tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice --chat-template-content-format string" prefill: dp: "" diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index 7449bd0e..531ef4d6 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -448,11 +448,14 @@ BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}" # Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. # sweep (default) -> benchmark_xPyD.sh (general concurrency sweep) # long_context -> benchmark_long_context.sh (per-shape warmup, c=1-first) +# keepalive -> keepalive_bench.sh (hold server up KEEPALIVE_MINS +# for external accuracy probes) BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" case "$BENCHMARK_SCRIPT" in sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; long_context) BENCHMARK_SCRIPT_FILE="benchmark_long_context.sh" ;; - *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context)" >&2; exit 1 ;; + keepalive) BENCHMARK_SCRIPT_FILE="keepalive_bench.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, keepalive)" >&2; exit 1 ;; esac if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 @@ -605,6 +608,7 @@ docker run --rm \ ${PREFILL_MORI_BACKEND:+-e PREFILL_MORI_BACKEND=$PREFILL_MORI_BACKEND} \ ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ + ${GLM_PERSIST_GATE:+-e GLM_PERSIST_GATE=$GLM_PERSIST_GATE} \ ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ @@ -612,6 +616,7 @@ docker run --rm \ ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \ ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \ ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \ + ${KEEPALIVE_MINS:+-e KEEPALIVE_MINS=$KEEPALIVE_MINS} \ ${PREFILL_CUDAGRAPH_MODE:+-e PREFILL_CUDAGRAPH_MODE=$PREFILL_CUDAGRAPH_MODE} \ ${DECODE_CUDAGRAPH_MODE:+-e DECODE_CUDAGRAPH_MODE=$DECODE_CUDAGRAPH_MODE} \ ${CUDAGRAPH_CAPTURE_SIZES:+-e CUDAGRAPH_CAPTURE_SIZES="$CUDAGRAPH_CAPTURE_SIZES"} \ From c8c945487a4a24a0ab108091b7e5f9911061157d Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Tue, 7 Jul 2026 07:59:57 +0000 Subject: [PATCH 09/13] docker(vllm_disagg): repin vLLM/AITER/MoRI to the #47766 GLM-5.1 stack 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 --- ...llm_disagg_inference.ubuntu.amd.Dockerfile | 55 ++++++++++++------- 1 file changed, 34 insertions(+), 21 deletions(-) diff --git a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile index 06dd18aa..ec2a991f 100644 --- a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile +++ b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile @@ -99,7 +99,9 @@ ARG NIC_COMPILATION_ARCH="cx7" # backends produced a MoRI that deadlocked at the cross-node EP all-to-all init. # ----------------------------------------------------------------------------- ARG MORI_REPO=https://github.com/ROCm/mori.git -ARG MORI_REF=v1.2.1 +# 42e895472b08: MoRI main tip past v1.2.1, validated by MAD-private #338 for GLM-5.1 +# DSA WideEP disagg (v1.2.1 large-transfer notify path was insufficient at high EP). +ARG MORI_REF=42e895472b08 ENV MORI_GPU_ARCHS=gfx942 # Newer MoRI added the UMBP subsystem which requires gRPC (grpcpp/grpcpp.h) not # present in this base; UMBP is unrelated to the EP dispatch/combine kernels, so @@ -123,29 +125,40 @@ RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ rm -rf /tmp/mori-src # ----------------------------------------------------------------------------- -# 2. AITER: install 0.1.16.post3 (prebuilt rocm7.2 wheel + flydsl 0.2.2), then -# invalidate the stale prewarmed AITER JIT cache compiled against the old .so. +# 2. AITER: build STOCK upstream ROCm/aiter @ e03fa6040 from source (NO fork, +# NO gqa64-fold patch). Under vLLM #47766 the sparse-MLA persistent path stays +# ON, so GLM's gqa=64 decode hits aiter's PRE-EXISTING persistent gqa64->16 fold +# (aiter/mla.py: `nhead in range(32,128+1,16) and persistent_mode`); the fork's +# extra non-persistent fold is never exercised, so stock is sufficient. +# Validated by MAD-private #338: 1P/1D EP8 + 2P/2D EP16 NIAH PASS on this exact +# aiter tip under #47766. Pin the exact commit (the one tested), not the release +# wheel. Then invalidate the stale prewarmed JIT cache compiled against the old .so. # ----------------------------------------------------------------------------- -ARG AITER_VERSION=0.1.16.post3 -ARG AITER_WHEEL_URL="https://github.com/ROCm/aiter/releases/download/v0.1.16.post3/amd_aiter-0.1.16.post3%2Brocm7.2.manylinux.2.28-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" -RUN echo "Bumping AITER to ${AITER_VERSION} from ${AITER_WHEEL_URL}" && \ - _W="/tmp/$(basename "${AITER_WHEEL_URL}" | sed 's/%2B/+/g')" && \ - curl -fL --retry 3 --retry-delay 2 -o "${_W}" "${AITER_WHEEL_URL}" && \ +ARG AITER_REPO=https://github.com/ROCm/aiter.git +ARG AITER_REF=e03fa6040 +RUN echo "Compiling STOCK AITER (no fork) from ${AITER_REPO}@${AITER_REF}" && \ + rm -rf /tmp/aiter-src && \ + git clone --recursive "${AITER_REPO}" /tmp/aiter-src && \ + cd /tmp/aiter-src && git checkout "${AITER_REF}" && \ + git submodule update --init --recursive && \ (pip uninstall -y amd_aiter amd-aiter aiter 2>/dev/null || true) && \ - pip install --no-deps "${_W}" && \ - pip install "flydsl==0.2.2" && \ - rm -f "${_W}" && \ + pip install --no-build-isolation --no-deps -v . && \ + pip install --no-deps -U "flydsl>=0.1.7,<0.1.9" && \ + echo "AITER_REF=${AITER_REF}@$(git rev-parse HEAD) (stock ROCm/aiter, no fork)" >> /app/versions.txt && \ + rm -rf /tmp/aiter-src && \ python3 - <<'PYEOF' -from importlib.metadata import version as v, PackageNotFoundError -vm = None -for n in ("amd-aiter", "amd_aiter", "aiter"): - try: vm = v(n); break - except PackageNotFoundError: pass -assert vm and vm.split("+", 1)[0] == "0.1.16.post3", f"AITER not 0.1.16.post3: {vm!r}" -print("AITER OK:", vm) +# Verify aiter/mla.py installed + has the persistent gqa64 fold, WITHOUT importing +# aiter/torch (torch->amdsmi->libamd_smi.so is not loadable at build: no GPU in sandbox). +import glob, pathlib +cands = glob.glob("/usr/local/lib/python*/dist-packages/aiter/mla.py") + \ + glob.glob("/usr/lib/python*/dist-packages/aiter/mla.py") +assert cands, "aiter/mla.py not found in site-packages after install" +src = pathlib.Path(cands[0]).read_text() +assert "persistent_mode" in src, f"AITER persistent fold path MISSING in {cands[0]}" +print("STOCK AITER OK (persistent gqa64 fold path present):", cands[0]) PYEOF RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" && \ - echo "AITER_VERSION=${AITER_VERSION}" >> /app/versions.txt + echo "AITER_REF=${AITER_REF} (stock)" >> /app/versions.txt # ----------------------------------------------------------------------------- # 3. vLLM: compile from source at the 06_29 validated Wide-EP WRITE-mode branch @@ -157,8 +170,8 @@ RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER J # ----------------------------------------------------------------------------- # VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the # dist-inf-cookbook mori121 image builds from). Override to your own vLLM fork/branch. -ARG VLLM_REPO=https://github.com/shikamd123/vllm.git -ARG VLLM_REF=vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer +ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git +ARG VLLM_REF=glm5.1-dsa-wideEP_on_shikpate_06_29_customer ENV VLLM_TARGET_DEVICE=rocm \ PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ MAX_JOBS=${MAX_JOBS} From 8aabc3674404e6ab8b8ec976441732656890ac80 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Tue, 7 Jul 2026 08:24:52 +0000 Subject: [PATCH 10/13] docker(vllm_disagg): fix stale AITER version assert for source build 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 --- docker/vllm_disagg_inference.ubuntu.amd.Dockerfile | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile index ec2a991f..4febc1fc 100644 --- a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile +++ b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile @@ -51,7 +51,8 @@ # # - BASE: rocm/vllm-dev:ci_base-0fcd9b99... (open ROCm 7.2 / cp312 CI base). # - MoRI -> built from ROCm/MoRI @ v1.2.1 (BUILD_UMBP=OFF). -# - AITER -> 0.1.16.post3 prebuilt rocm7.2 wheel + flydsl 0.2.2; stale JIT wiped. +# - AITER -> STOCK ROCm/aiter @ e03fa6040 compiled from source + flydsl 0.1.7-0.1.9; +# stale JIT wiped. (#47766 keeps persistent MLA ON -> aiter native gqa64 fold.) # - vLLM -> COMPILED from shikamd123/vllm @ # vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer (Wide-EP multi-pod PD, the # connector/router reference for the 2P2D DP=EP=16 topology). Full compile: it is @@ -193,7 +194,10 @@ def get(names): except PackageNotFoundError: pass return None av = get(("amd-aiter", "amd_aiter", "aiter")) -assert av and av.split("+", 1)[0] == "0.1.16.post3", f"AITER downgraded: {av!r}" +# Stock source build of ROCm/aiter@e03fa6040 reports 0.1.17.dev195+ge03fa6040. +# Verify the aiter install survived the vLLM install (present + carries the e03fa6040 +# commit tag) rather than pinning a release version string. +assert av and "e03fa6040" in av, f"AITER missing/downgraded (want e03fa6040 build): {av!r}" import mori, mori.io, mori.ops print("Post-vLLM check OK: AITER", av, "+ MoRI importable") PYEOF From 63cb42d3bfc76ffe7eed52d859d5820035fe2e50 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Tue, 7 Jul 2026 16:20:52 +0000 Subject: [PATCH 11/13] vllm_dissag: GLM_SKIP_PATCHERS switch for baked-fix images 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 --- scripts/vllm_dissag/connectors/moriio.sh | 9 +++++++++ scripts/vllm_dissag/run_xPyD_models.slurm | 1 + 2 files changed, 10 insertions(+) diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 159bd662..7d3d022b 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -167,6 +167,15 @@ connector_runtime_patch() { # failing at launch is correct). Patchers self-skip (rc 0) when their anchor is # absent, so an image that already carries or refactored a fix no-ops cleanly. _glm_dsa_runtime_patch() { + # GLM_SKIP_PATCHERS=1: the serving image already carries the GLM-5.1 DSA fixes + # in-source (e.g. the #47766 stack image built from raviguptaamd/vllm@ + # glm5.1-dsa-wideEP_on_shikpate_06_29_customer). Skip ALL runtime patchers — they + # are redundant, and the persistent-gate/sampling-overlay patchers would actively + # REGRESS a baked image (turn persistent MLA off / overwrite stock aiter kernels). + if [ "${GLM_SKIP_PATCHERS:-0}" = "1" ]; then + echo "[glm] GLM_SKIP_PATCHERS=1: image carries DSA fixes in-source; skipping runtime patchers." + return 0 + fi local _patch_dir="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)}" local _vllm_dir _vllm_dir="$(python3 -c 'import vllm, os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || true)" diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index 531ef4d6..98281f57 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -609,6 +609,7 @@ docker run --rm \ ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ ${GLM_PERSIST_GATE:+-e GLM_PERSIST_GATE=$GLM_PERSIST_GATE} \ + ${GLM_SKIP_PATCHERS:+-e GLM_SKIP_PATCHERS=$GLM_SKIP_PATCHERS} \ ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ From 9f7dbd85c7f592dc30023c425a63a276a886316c Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Tue, 7 Jul 2026 19:36:40 +0000 Subject: [PATCH 12/13] docker(vllm_disagg): repoint vLLM to glm5.1-dsa-wideEP_on_shik_latest (DP-notify fixes) --- docker/vllm_disagg_inference.ubuntu.amd.Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile index 4febc1fc..fef3995c 100644 --- a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile +++ b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile @@ -172,7 +172,7 @@ RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER J # VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the # dist-inf-cookbook mori121 image builds from). Override to your own vLLM fork/branch. ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git -ARG VLLM_REF=glm5.1-dsa-wideEP_on_shikpate_06_29_customer +ARG VLLM_REF=glm5.1-dsa-wideEP_on_shik_latest ENV VLLM_TARGET_DEVICE=rocm \ PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ MAX_JOBS=${MAX_JOBS} From 7b810b79dac4d05b91a07b885585fd64eef40b01 Mon Sep 17 00:00:00 2001 From: raviguptaamd Date: Wed, 8 Jul 2026 06:19:50 +0000 Subject: [PATCH 13/13] vllm_dissag: isolate GLM-5.1 into its own image; keep DSV3 on upstream 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 #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. Dockerfile + card, with zero edits to DSV3 or GLM. Co-Authored-By: Claude --- ...gg_inference.glmv5.1.ubuntu.amd.Dockerfile | 335 ++++++++++++++++++ ...llm_disagg_inference.ubuntu.amd.Dockerfile | 63 ++-- models.json | 32 ++ scripts/vllm_dissag/connectors/moriio.sh | 23 +- scripts/vllm_dissag/models.yaml | 11 +- 5 files changed, 402 insertions(+), 62 deletions(-) create mode 100644 docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile diff --git a/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..f61a8819 --- /dev/null +++ b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile @@ -0,0 +1,335 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# ============================================================================= +# vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile +# GLM-5.1-FP8 (MLA + DeepSeek Sparse Attention) MoRI-EP WideEP disagg image. +# PER-MODEL image, isolated from the base vllm_disagg_inference Dockerfile +# (which stays pinned to the DeepSeek-V3 / R1 stack). This split lets each model +# pin its own vLLM/AITER/MoRI without disturbing the others -- add a new +# vllm_disagg_inference..ubuntu.amd.Dockerfile per future model +# (e.g. Kimi-2.6) rather than repinning the shared DSV3 image. +# +# ALL connectors in one image: moriio (TP + MoRI-EP wideEP) + rixl (NIXL TP + +# DeepEP wideEP). = the fullsource MoRI stack, plus a UCX/RIXL/rocSHMEM/DeepEP +# transport layer gated by --build-arg WITH_NIXL (default 1 = everything). +# +# docker build -f docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile \ +# -t /vllm-disagg:glmv5.1 . +# export DOCKER_IMAGE_NAME=/vllm-disagg:glmv5.1 +# +# WITH_NIXL=1 (default) => builds UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from +# source, so all four connector combos (moriio TP/wideEP, rixl NIXL TP, DeepEP +# wideEP) are present (~+30-45 min build vs WITH_NIXL=0). +# WITH_NIXL=0 => MoRI-EP only (moriio TP/wideEP + deepep-from-base); lean, faster. +# +# STATUS (GLM-5.1-FP8 on this stack): 1P/1D EP8 + 2P/2D EP16 NIAH 2k-35k = 10/10, +# no crash; long-context accuracy fixed via vLLM #47766 (persistent sparse-MLA kept +# ON). 4P/4D EP32 is a KNOWN OPEN DEFECT: token corruption at ALL context lengths +# (garbage output even at 2k), distinct from the long-context bug; prime suspect is +# the moriep all-to-all combine at EP32 scale -> deferred to future work. Use 1P/1D +# and 2P/2D only. (BASE_IMAGE is a gated nightly; override --build-arg BASE_IMAGE=...) +# ============================================================================= +# Reconstructs the validated v1.2.1 (mori121) runtime stack by applying the recipe's +# component pins ON TOP of the open ROCm vLLM ci_base, cloning each source from +# public Git (no local build-contexts). Mirrors dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable: +# +# - BASE: rocm/vllm-dev:ci_base-0fcd9b99... (open ROCm 7.2 / cp312 CI base). +# - MoRI -> built from ROCm/MoRI @ v1.2.1 (BUILD_UMBP=OFF). +# - AITER -> STOCK ROCm/aiter @ e03fa6040 compiled from source + flydsl 0.1.7-0.1.9; +# stale JIT wiped. (#47766 keeps persistent MLA ON -> aiter native gqa64 fold.) +# - vLLM -> COMPILED from shikamd123/vllm @ +# vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer (Wide-EP multi-pod PD, the +# connector/router reference for the 2P2D DP=EP=16 topology). Full compile: it is +# a different commit than the base's, so a .py-only overlay would be ABI-mismatched. +# - RDMA fix (expandable_segments:False x2 + HSA_ENABLE_IPC_MODE_LEGACY=0) is NOT baked +# here — it lives in scripts/vllm_dissag/connectors/.env and the launcher +# forwards it via docker -e. ROCm 7.2.3 cannot dmabuf-export VMM memory, else MoRI +# RegisterRdmaMemoryRegion EFAULTs (errno 14) on the first disagg WRITE. +# - vllm-router (vllm-project/router PR#181 = DP-rank round-robin + 2P2D KV-notify +# dpfix) built in -> no external router binary needed. +# - validated recipe knobs baked as ENV. The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) are native in this vLLM (no runtime patcher). +# +# Build context = repo root: +# docker build -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile -t / . +# +# BASE_IMAGE is the open rocm/vllm-dev ci_base pinned by the validated recipe +# (dist-inf-cookbook Dockerfile.vllm.mori121_shareable). Override --build-arg +# BASE_IMAGE=... to build on a different ROCm base. vLLM compile is long (~30-60 min). +# ============================================================================= + +ARG BASE_IMAGE=rocm/vllm-dev:ci_base-0fcd9b99cc9d63202da4c858d8ebc6582c9e2491 +FROM ${BASE_IMAGE} + +ENTRYPOINT [] +WORKDIR /app + +ARG GFX_COMPILATION_ARCH="gfx942" +ARG PYTORCH_ROCM_ARCH="gfx942" +ARG MAX_JOBS=32 +# NIXL/RIXL transport for the rixl connector. Default 1 => all connectors built +# (UCX/RIXL/rocSHMEM/DeepEP). Set --build-arg WITH_NIXL=0 for a lean MoRI-EP-only image. +ARG WITH_NIXL=1 +ARG NIC_COMPILATION_ARCH="cx7" + +# ----------------------------------------------------------------------------- +# 1. MoRI: replace the base's bundled MoRI with the validated ROCm/MoRI @ v1.2.1 +# (the version for the 06_29 mori121 image, dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable). v1.2.1 carries the EP/RDMA correctness fixes +# plus the ROCm-7.2.3 dmabuf registration path used by the connector .env +# (expandable_segments:False). MoRI is JIT-built, so this swaps the JIT sources the +# kernels compile from at runtime. +# BUILD CONFIG: match the cookbook build — MORI_GPU_ARCHS=gfx942, BUILD_UMBP=OFF, +# DEFAULT NIC backends. Do NOT pass USE_IONIC=OFF / USE_BNXT=OFF: disabling NIC +# backends produced a MoRI that deadlocked at the cross-node EP all-to-all init. +# ----------------------------------------------------------------------------- +ARG MORI_REPO=https://github.com/ROCm/mori.git +# 42e895472b08: MoRI main tip past v1.2.1, validated by MAD-private #338 for GLM-5.1 +# DSA WideEP disagg (v1.2.1 large-transfer notify path was insufficient at high EP). +ARG MORI_REF=42e895472b08 +ENV MORI_GPU_ARCHS=gfx942 +# Newer MoRI added the UMBP subsystem which requires gRPC (grpcpp/grpcpp.h) not +# present in this base; UMBP is unrelated to the EP dispatch/combine kernels, so +# disable it to avoid pulling in a gRPC build dependency. +ENV BUILD_UMBP=OFF BUILD_UMBP_SPDK=OFF +# Build/install matches dist-inf-cookbook Dockerfile.vllm.mori121_shareable for v1.2.1: +# `BUILD_UMBP=OFF pip install .` (default build isolation). apt/pip build tooling kept +# for bases that lack it; harmless where already present. +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + sed -i 's|http://|https://|g' /etc/apt/sources.list.d/*.list 2>/dev/null || true && \ + apt-get update && apt-get install -y --no-install-recommends \ + git build-essential cmake ninja-build ccache libssl-dev pkg-config curl ca-certificates && \ + pip install meson==0.64.0 "pybind11[global]" tqdm prettytable && \ + pip uninstall -y amd_mori amd-mori amd-mori-nightly mori 2>/dev/null || true && \ + rm -rf /tmp/mori-src && \ + git clone --recursive "${MORI_REPO}" /tmp/mori-src && \ + cd /tmp/mori-src && git checkout "${MORI_REF}" && git submodule update --init --recursive && \ + BUILD_UMBP=OFF pip install . && \ + python3 -c "import mori, mori.io, mori.ops; print('MoRI OK at', mori.__path__[0])" && \ + mkdir -p /app && echo "MORI_REF=${MORI_REF}@$(git -C /tmp/mori-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/mori-src + +# ----------------------------------------------------------------------------- +# 2. AITER: build STOCK upstream ROCm/aiter @ e03fa6040 from source (NO fork, +# NO gqa64-fold patch). Under vLLM #47766 the sparse-MLA persistent path stays +# ON, so GLM's gqa=64 decode hits aiter's PRE-EXISTING persistent gqa64->16 fold +# (aiter/mla.py: `nhead in range(32,128+1,16) and persistent_mode`); the fork's +# extra non-persistent fold is never exercised, so stock is sufficient. +# Validated by MAD-private #338: 1P/1D EP8 + 2P/2D EP16 NIAH PASS on this exact +# aiter tip under #47766. Pin the exact commit (the one tested), not the release +# wheel. Then invalidate the stale prewarmed JIT cache compiled against the old .so. +# ----------------------------------------------------------------------------- +ARG AITER_REPO=https://github.com/ROCm/aiter.git +ARG AITER_REF=e03fa6040 +RUN echo "Compiling STOCK AITER (no fork) from ${AITER_REPO}@${AITER_REF}" && \ + rm -rf /tmp/aiter-src && \ + git clone --recursive "${AITER_REPO}" /tmp/aiter-src && \ + cd /tmp/aiter-src && git checkout "${AITER_REF}" && \ + git submodule update --init --recursive && \ + (pip uninstall -y amd_aiter amd-aiter aiter 2>/dev/null || true) && \ + pip install --no-build-isolation --no-deps -v . && \ + pip install --no-deps -U "flydsl>=0.1.7,<0.1.9" && \ + echo "AITER_REF=${AITER_REF}@$(git rev-parse HEAD) (stock ROCm/aiter, no fork)" >> /app/versions.txt && \ + rm -rf /tmp/aiter-src && \ + python3 - <<'PYEOF' +# Verify aiter/mla.py installed + has the persistent gqa64 fold, WITHOUT importing +# aiter/torch (torch->amdsmi->libamd_smi.so is not loadable at build: no GPU in sandbox). +import glob, pathlib +cands = glob.glob("/usr/local/lib/python*/dist-packages/aiter/mla.py") + \ + glob.glob("/usr/lib/python*/dist-packages/aiter/mla.py") +assert cands, "aiter/mla.py not found in site-packages after install" +src = pathlib.Path(cands[0]).read_text() +assert "persistent_mode" in src, f"AITER persistent fold path MISSING in {cands[0]}" +print("STOCK AITER OK (persistent gqa64 fold path present):", cands[0]) +PYEOF +RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" && \ + echo "AITER_REF=${AITER_REF} (stock)" >> /app/versions.txt + +# ----------------------------------------------------------------------------- +# 3. vLLM: compile from source at the 06_29 validated Wide-EP WRITE-mode branch +# (matches the published dist-inf-cookbook mori121 image). Full source compile +# (the base ships a different commit). The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) are native in this branch, so no runtime +# patcher is needed. Override VLLM_REF to rebuild a different commit; build only +# committed commits (no working-tree edits). +# ----------------------------------------------------------------------------- +# VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the +# dist-inf-cookbook mori121 image builds from). Override to your own vLLM fork/branch. +ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git +ARG VLLM_REF=glm5.1-dsa-wideEP_on_shik_latest +ENV VLLM_TARGET_DEVICE=rocm \ + PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ + MAX_JOBS=${MAX_JOBS} +RUN rm -rf /tmp/vllm-src && \ + git clone "${VLLM_REPO}" /tmp/vllm-src && \ + cd /tmp/vllm-src && git checkout "${VLLM_REF}" && \ + echo "VLLM_REF=${VLLM_REF}@$(git rev-parse HEAD)" >> /app/versions.txt && \ + pip uninstall -y vllm 2>/dev/null || true && \ + pip install --no-deps --no-build-isolation -v . && \ + python3 -c "import vllm; print('vLLM', vllm.__version__, 'from', vllm.__file__)" && \ + rm -rf /tmp/vllm-src + +# Cross-check MoRI + AITER survived the vLLM install (no silent downgrade). +RUN python3 - <<'PYEOF' +from importlib.metadata import version as v, PackageNotFoundError +def get(names): + for n in names: + try: return v(n) + except PackageNotFoundError: pass + return None +av = get(("amd-aiter", "amd_aiter", "aiter")) +# Stock source build of ROCm/aiter@e03fa6040 reports 0.1.17.dev195+ge03fa6040. +# Verify the aiter install survived the vLLM install (present + carries the e03fa6040 +# commit tag) rather than pinning a release version string. +assert av and "e03fa6040" in av, f"AITER missing/downgraded (want e03fa6040 build): {av!r}" +import mori, mori.io, mori.ops +print("Post-vLLM check OK: AITER", av, "+ MoRI importable") +PYEOF + +# ----------------------------------------------------------------------------- +# 4. vllm-router (DP-rank round-robin + MoRIIO connector) — built in, so NO +# external vllm-router binary is needed (leave ROUTER_BINARY unset). +# Source = vllm-project/router PR #181 branch, which now carries BOTH the +# round-robin DP-rank fix (11841c0d) AND the 2P2D KV-notify fix (6409ac1: +# remote_dp_rank_override + remote_dp_size). The KV-notify fix is REQUIRED: +# without it the 2P2D EP=16 run reproducibly wedges with "remote blocks never +# arrived" deferred-write expiries (decode notify targets the wrong DP rank). +# This is the exact source of the validated vllm-router-2p2d-dpfix binary. +# Pinned Rust toolchain (>=1.88: router deps time/home require rustc 1.88). +# ----------------------------------------------------------------------------- +ARG ROUTER_REPO=https://github.com/raviguptaamd/router.git +ARG ROUTER_REF=ravgupta/discovery-dp-rank-roundrobin +ARG RUST_TOOLCHAIN=1.88.0 +RUN if ! command -v cargo >/dev/null 2>&1; then \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain "${RUST_TOOLCHAIN}"; \ + fi && \ + export PATH="/root/.cargo/bin:${PATH}" && \ + rm -rf /tmp/vllm-router-src && \ + git clone --filter=blob:none "${ROUTER_REPO}" /tmp/vllm-router-src && \ + cd /tmp/vllm-router-src && git checkout "${ROUTER_REF}" && \ + cargo build --release && \ + install -m 755 target/release/vllm-router /usr/local/bin/vllm-router && \ + vllm-router --help 2>&1 | grep -q moriio && \ + echo "VLLM_ROUTER_REF=${ROUTER_REPO}@${ROUTER_REF}@$(git -C /tmp/vllm-router-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/vllm-router-src + +# ----------------------------------------------------------------------------- +# 4b. WITH_NIXL=1 (default): UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from source, +# so the rixl connector (NIXL TP + DeepEP wideEP) is present. Single guarded RUN so +# WITH_NIXL=0 skips it entirely (no layers, no cost). Build-verified on ci_base. +# ----------------------------------------------------------------------------- +ENV _ROCM_DIR=/opt/rocm \ + _UCX_SOURCE=https://github.com/ROCm/ucx.git \ + _UCX_BRANCH=da3fac2a \ + _UCX_INSTALL_DIR=/usr/local/ucx/ \ + _RIXL_SOURCE=https://github.com/ROCm/RIXL.git \ + _RIXL_BRANCH=f33a5599 \ + _RIXL_INSTALL_DIR=/usr/local/RIXL/install \ + _NIXLBENCH_INSTALL_DIR=/usr/local/RIXL +RUN if [ "${WITH_NIXL}" != "1" ]; then \ + echo "WITH_NIXL=${WITH_NIXL}: skipping UCX/RIXL/rocSHMEM/DeepEP (MoRI-EP + base DeepEP only)"; \ + else set -e && \ + echo "WITH_NIXL=1: building UCX + RIXL + rocSHMEM + DeepEP" && \ + apt-get update && apt-get install -y \ + autoconf automake libtool autogen pkg-config m4 gcc make \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core strace libgflags-dev \ + libaio-dev liburing-dev libcpprest-dev libgrpc-dev libgrpc++-dev \ + libprotobuf-dev protobuf-compiler-grpc wget && \ + pip install meson==0.64.0 "pybind11[global]" pyyaml && \ + # UCX + cd /tmp && git clone "${_UCX_SOURCE}" && cd ucx && git checkout "${_UCX_BRANCH}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${_UCX_INSTALL_DIR}" --with-rocm="${_ROCM_DIR}" \ + --disable-go --disable-java --disable-assertions --enable-mt && \ + make -j && make install && \ + # googletest (RIXL dep) + cd /tmp && wget -q https://github.com/google/googletest/archive/refs/tags/v1.14.0.tar.gz && \ + tar -xzf v1.14.0.tar.gz && cd googletest-1.14.0 && mkdir -p build && cd build && \ + cmake -DBUILD_SHARED_LIBS=on .. && make -j && make install && \ + # RIXL + python bindings + cd /tmp && git clone "${_RIXL_SOURCE}" && cd RIXL && git checkout "${_RIXL_BRANCH}" && \ + meson setup build/ --prefix="${_RIXL_INSTALL_DIR}" -Ducx_path="${_UCX_INSTALL_DIR}" \ + -Ddisable_gds_backend=true -Dcudapath_inc="${_ROCM_DIR}/include" -Dcudapath_lib="${_ROCM_DIR}/lib" && \ + cd build && ninja && ninja install && cd /tmp/RIXL && \ + pip install --config-settings=setup-args="-Dcudapath_inc=${_ROCM_DIR}/include" \ + --config-settings=setup-args="-Dcudapath_lib=${_ROCM_DIR}/lib" \ + --config-settings=setup-args="-Ducx_path=${_UCX_INSTALL_DIR}" \ + --config-settings=setup-args="-Ddisable_gds_backend=true" . && \ + # rocSHMEM (DeepEP dep) + cd /tmp && git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git && \ + cd rocm-systems && git sparse-checkout set --cone projects/rocshmem && git checkout develop && \ + mkdir -p /tmp/rocshmem-build && cd /tmp/rocshmem-build && \ + /tmp/rocm-systems/projects/rocshmem/scripts/build_configs/all_backends \ + -DUSE_EXTERNAL_MPI=OFF -DGPU_TARGETS="${GFX_COMPILATION_ARCH}" && \ + # DeepEP (build develop against the installed vLLM/torch) + cd /tmp && git clone https://github.com/ROCm/DeepEP.git && cd DeepEP && \ + PYTORCH_ROCM_ARCH="${GFX_COMPILATION_ARCH}" CFLAGS="-O3 -fPIC" \ + CXXFLAGS="-O3 -fPIC --offload-arch=${GFX_COMPILATION_ARCH}" HIP_CXX_FLAGS="-O3 -fPIC" \ + python3 setup.py --variant rocm --nic "${NIC_COMPILATION_ARCH}" build develop && \ + echo "WITH_NIXL build complete" >> /app/versions.txt && \ + rm -rf /tmp/ucx /tmp/googletest-1.14.0 /tmp/v1.14.0.tar.gz /tmp/rocm-systems /tmp/rocshmem-build; \ + fi +ENV LD_LIBRARY_PATH="/usr/local/ucx/lib:/usr/local/lib:/usr/local/RIXL/install/lib:${LD_LIBRARY_PATH}" \ + PATH="/usr/local/ucx/bin:${PATH}" + +# ----------------------------------------------------------------------------- +# 5. Cache locations (structural: WHERE the JIT/compile caches live in the image). +# These are the mount target for the launcher's persistent host JIT cache. +# ----------------------------------------------------------------------------- +# The image ships NO runtime recipe / tuning / platform ENV. By design, everything +# run-tunable is applied at launch, so this image stays a clean binary/library artifact +# and the same image serves any model/cluster without a rebuild: +# - model-serving recipe (KV_BLOCK_SIZE, KV_CACHE_DTYPE, *_CUDAGRAPH_MODE, *_MORI_BACKEND, +# GPU_MEMORY_UTILIZATION, KV_CACHE_MEMORY_BYTES, VLLM_ROCM_USE_AITER_MLA, ...) +# -> scripts/vllm_dissag/models.yaml (per-model env:, so dense vs MoE differ) +# - ROCm-7.2.3 GPU-RDMA platform env (expandable_segments:False x2, MORI_GPU_ARCHS, +# HSA_ENABLE_IPC_MODE_LEGACY=0, HSA_NO_SCRATCH_RECLAIM) and the MoRI/RDMA fabric +# tuning (MORI_RDMA_TC/SL, MORI_IB_GID_INDEX, MORI_NUM_QP_PER_PE, VLLM_MORIIO_*, ...) +# -> scripts/vllm_dissag/connectors/.env (cluster-editable, no rebuild) +# The slurm launcher forwards both via `docker -e` (platform env must reach PID 1 - +# PyTorch reads alloc-conf at import). Running this image WITHOUT the launcher: set the +# vars you need yourself (see connectors/moriio.env + models.yaml for the values). +ENV AITER_JIT_DIR=/opt/vllm_cache/aiter_jit \ + VLLM_CACHE_ROOT=/opt/vllm_cache/vllm \ + TRITON_CACHE_DIR=/opt/vllm_cache/triton \ + COMGR_CACHE_DIR=/opt/vllm_cache/comgr + +# ----------------------------------------------------------------------------- +# 6. CRITICAL: scrub build-time MoRI JIT state. The `import mori` verification +# steps above compile/lock MoRI EP kernels under /root/.mori/jit on THIS build +# host, leaving stale .hsaco.lock files (ep_internode_v1, ep_internode_v1ll, ...). +# At runtime on the cluster, MoriAll2AllManager finds those locks, waits on a +# build-in-progress whose owner PID is long gone, and DEADLOCKS at ep:0 init. +# A clean image ships /root/.mori empty -> runtime compiles fresh. +# Clearing these makes the from-source image boot clean on 2P2D/4P4D. +# ----------------------------------------------------------------------------- +RUN rm -rf /root/.mori /tmp/mori_jit_* && mkdir -p /root/.mori && \ + echo "JIT_SCRUBBED: /root/.mori + /tmp/mori_jit_* cleared at build end" >> /app/versions.txt + +RUN cat /app/versions.txt 2>/dev/null | tail -20 || true diff --git a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile index fef3995c..06dd18aa 100644 --- a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile +++ b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile @@ -51,8 +51,7 @@ # # - BASE: rocm/vllm-dev:ci_base-0fcd9b99... (open ROCm 7.2 / cp312 CI base). # - MoRI -> built from ROCm/MoRI @ v1.2.1 (BUILD_UMBP=OFF). -# - AITER -> STOCK ROCm/aiter @ e03fa6040 compiled from source + flydsl 0.1.7-0.1.9; -# stale JIT wiped. (#47766 keeps persistent MLA ON -> aiter native gqa64 fold.) +# - AITER -> 0.1.16.post3 prebuilt rocm7.2 wheel + flydsl 0.2.2; stale JIT wiped. # - vLLM -> COMPILED from shikamd123/vllm @ # vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer (Wide-EP multi-pod PD, the # connector/router reference for the 2P2D DP=EP=16 topology). Full compile: it is @@ -100,9 +99,7 @@ ARG NIC_COMPILATION_ARCH="cx7" # backends produced a MoRI that deadlocked at the cross-node EP all-to-all init. # ----------------------------------------------------------------------------- ARG MORI_REPO=https://github.com/ROCm/mori.git -# 42e895472b08: MoRI main tip past v1.2.1, validated by MAD-private #338 for GLM-5.1 -# DSA WideEP disagg (v1.2.1 large-transfer notify path was insufficient at high EP). -ARG MORI_REF=42e895472b08 +ARG MORI_REF=v1.2.1 ENV MORI_GPU_ARCHS=gfx942 # Newer MoRI added the UMBP subsystem which requires gRPC (grpcpp/grpcpp.h) not # present in this base; UMBP is unrelated to the EP dispatch/combine kernels, so @@ -126,40 +123,29 @@ RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ rm -rf /tmp/mori-src # ----------------------------------------------------------------------------- -# 2. AITER: build STOCK upstream ROCm/aiter @ e03fa6040 from source (NO fork, -# NO gqa64-fold patch). Under vLLM #47766 the sparse-MLA persistent path stays -# ON, so GLM's gqa=64 decode hits aiter's PRE-EXISTING persistent gqa64->16 fold -# (aiter/mla.py: `nhead in range(32,128+1,16) and persistent_mode`); the fork's -# extra non-persistent fold is never exercised, so stock is sufficient. -# Validated by MAD-private #338: 1P/1D EP8 + 2P/2D EP16 NIAH PASS on this exact -# aiter tip under #47766. Pin the exact commit (the one tested), not the release -# wheel. Then invalidate the stale prewarmed JIT cache compiled against the old .so. +# 2. AITER: install 0.1.16.post3 (prebuilt rocm7.2 wheel + flydsl 0.2.2), then +# invalidate the stale prewarmed AITER JIT cache compiled against the old .so. # ----------------------------------------------------------------------------- -ARG AITER_REPO=https://github.com/ROCm/aiter.git -ARG AITER_REF=e03fa6040 -RUN echo "Compiling STOCK AITER (no fork) from ${AITER_REPO}@${AITER_REF}" && \ - rm -rf /tmp/aiter-src && \ - git clone --recursive "${AITER_REPO}" /tmp/aiter-src && \ - cd /tmp/aiter-src && git checkout "${AITER_REF}" && \ - git submodule update --init --recursive && \ +ARG AITER_VERSION=0.1.16.post3 +ARG AITER_WHEEL_URL="https://github.com/ROCm/aiter/releases/download/v0.1.16.post3/amd_aiter-0.1.16.post3%2Brocm7.2.manylinux.2.28-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" +RUN echo "Bumping AITER to ${AITER_VERSION} from ${AITER_WHEEL_URL}" && \ + _W="/tmp/$(basename "${AITER_WHEEL_URL}" | sed 's/%2B/+/g')" && \ + curl -fL --retry 3 --retry-delay 2 -o "${_W}" "${AITER_WHEEL_URL}" && \ (pip uninstall -y amd_aiter amd-aiter aiter 2>/dev/null || true) && \ - pip install --no-build-isolation --no-deps -v . && \ - pip install --no-deps -U "flydsl>=0.1.7,<0.1.9" && \ - echo "AITER_REF=${AITER_REF}@$(git rev-parse HEAD) (stock ROCm/aiter, no fork)" >> /app/versions.txt && \ - rm -rf /tmp/aiter-src && \ + pip install --no-deps "${_W}" && \ + pip install "flydsl==0.2.2" && \ + rm -f "${_W}" && \ python3 - <<'PYEOF' -# Verify aiter/mla.py installed + has the persistent gqa64 fold, WITHOUT importing -# aiter/torch (torch->amdsmi->libamd_smi.so is not loadable at build: no GPU in sandbox). -import glob, pathlib -cands = glob.glob("/usr/local/lib/python*/dist-packages/aiter/mla.py") + \ - glob.glob("/usr/lib/python*/dist-packages/aiter/mla.py") -assert cands, "aiter/mla.py not found in site-packages after install" -src = pathlib.Path(cands[0]).read_text() -assert "persistent_mode" in src, f"AITER persistent fold path MISSING in {cands[0]}" -print("STOCK AITER OK (persistent gqa64 fold path present):", cands[0]) +from importlib.metadata import version as v, PackageNotFoundError +vm = None +for n in ("amd-aiter", "amd_aiter", "aiter"): + try: vm = v(n); break + except PackageNotFoundError: pass +assert vm and vm.split("+", 1)[0] == "0.1.16.post3", f"AITER not 0.1.16.post3: {vm!r}" +print("AITER OK:", vm) PYEOF RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" && \ - echo "AITER_REF=${AITER_REF} (stock)" >> /app/versions.txt + echo "AITER_VERSION=${AITER_VERSION}" >> /app/versions.txt # ----------------------------------------------------------------------------- # 3. vLLM: compile from source at the 06_29 validated Wide-EP WRITE-mode branch @@ -171,8 +157,8 @@ RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER J # ----------------------------------------------------------------------------- # VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the # dist-inf-cookbook mori121 image builds from). Override to your own vLLM fork/branch. -ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git -ARG VLLM_REF=glm5.1-dsa-wideEP_on_shik_latest +ARG VLLM_REPO=https://github.com/shikamd123/vllm.git +ARG VLLM_REF=vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer ENV VLLM_TARGET_DEVICE=rocm \ PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ MAX_JOBS=${MAX_JOBS} @@ -194,10 +180,7 @@ def get(names): except PackageNotFoundError: pass return None av = get(("amd-aiter", "amd_aiter", "aiter")) -# Stock source build of ROCm/aiter@e03fa6040 reports 0.1.17.dev195+ge03fa6040. -# Verify the aiter install survived the vLLM install (present + carries the e03fa6040 -# commit tag) rather than pinning a release version string. -assert av and "e03fa6040" in av, f"AITER missing/downgraded (want e03fa6040 build): {av!r}" +assert av and av.split("+", 1)[0] == "0.1.16.post3", f"AITER downgraded: {av!r}" import mori, mori.io, mori.ops print("Post-vLLM check OK: AITER", av, "+ MoRI importable") PYEOF diff --git a/models.json b/models.json index af914f46..689754db 100644 --- a/models.json +++ b/models.json @@ -3420,6 +3420,38 @@ }, "args": "-N 2 -n 2" }, + { + "name": "pyt_vllm_disagg_mori_glm-5.1-fp8", + "url": "", + "dockerfile": "docker/vllm_disagg_inference.glmv5.1", + "scripts": "scripts/vllm_dissag/run_xPyD_models.slurm", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "vllm", + "vllm_disagg", + "mori_ep", + "inference" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "GLM-5.1-FP8", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "GLM_SKIP_PATCHERS": "1", + "BENCHMARK_COMBINATIONS": "1024/1024" + }, + "args": "-N 2 -n 2" + }, { "name": "pyt_vllm_disagg_mori_deepseek-v3-5layer", "url": "", diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 7d3d022b..c26a0944 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -137,26 +137,9 @@ connector_runtime_patch() { # code gaps, applied here as idempotent, anchor-based, self-skipping .py patchers # (they no-op cleanly if the fix is native/refactored on the chosen image). Gated # on MODEL_NAME so DeepSeek/other models are a pure no-op (byte-identical to before). - # MoRI library bump (opt-in): if a rebuilt MoRI package is staged at - # MORI_OVERLAY_DIR (default /shared_inference/$USER/mori_built/mori), overlay it - # over the image's stock MoRI BEFORE any mori import. The stock image ships MoRI - # v1.2.1 whose KV-transfer notify/mapping fails on large (>~8k) transfers - # (decode "unmap MISS" -> deferred-write expires 60s -> prefill EngineCore crash). - # Post-1.2.1 MoRI main carries the fixes (#424 notify SENDs, #436 large-transfer - # chunking, #432 RDMA resilience). Built from source against THIS image's libpci - # (nightly wheel is ABI-incompatible: needs LIBPCI_3.8, image has 3.7). Applies to - # ALL models (MoRI is model-agnostic), so this runs before the GLM gate. - _mori_overlay="${MORI_OVERLAY_DIR:-/shared_inference/${USER_NAME:-$USER}/mori_built/mori}" - if [ -d "${_mori_overlay}" ]; then - _mori_dst="$(python3 -c 'import mori,os;print(os.path.dirname(mori.__file__))' 2>/dev/null || true)" - if [ -n "${_mori_dst}" ] && [ -d "${_mori_dst}" ]; then - echo "[mori-bump] overlaying rebuilt MoRI from ${_mori_overlay} onto ${_mori_dst}" - rm -rf "${_mori_dst}" && cp -a "${_mori_overlay}" "${_mori_dst}" && \ - echo "[mori-bump] now: $(python3 -c 'import mori;print(mori.__version__)' 2>&1 | tail -1)" || \ - echo "[mori-bump] WARN: overlay failed; continuing with stock MoRI" >&2 - fi - fi - + # The MoRI version is pinned by the Dockerfile MORI_REF (post-1.2.1 main with the + # large-transfer notify/mapping fixes #424/#436/#432 baked in); if a newer MoRI is + # needed, update MORI_REF and rebuild the image — no runtime library swap here. [ "${MODEL_NAME:-}" = "GLM-5.1-FP8" ] || return 0 _glm_dsa_runtime_patch } diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 60fbd2fb..9eec3e7d 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -160,10 +160,17 @@ _deepseek_recipe_env: &deepseek_recipe_env # above. No tp: blocks — TP is unsupported for these models. DeepSeek-V3: env: *deepseek_recipe_env + # NEWER-BASE ADAPTATION (isolated to this model): Shiksha's newer vLLM base added + # TritonMLAMetadataBuilder._reserve_attn_logits_workspace(), which pre-reserves the + # decode split-KV logits workspace at WORST CASE + # (max_num_seqs x q_heads x max_kv_splits(max_model_len) x lse_dim x fp32). With + # DSV3's defaults (max_num_seqs=256, max_model_len~163k) this reserves ~128 GiB and + # OOMs at KV init. Cap max-num-seqs + max-model-len to bound the workspace. This is a + # per-model dp: flag (model_args), fully isolated -- does NOT touch GLM/other recipes. prefill: - dp: "" + dp: "--max-num-seqs 64 --max-model-len 32768" decode: - dp: "" + dp: "--max-num-seqs 64 --max-model-len 32768" DeepSeek-V3-5layer: env: *deepseek_recipe_env