Add MoRI conn.py overlay for sglang v0.5.12.post1 PD-disaggregation

sglang v0.5.12.post1 ships an unmigrated MoRI backend (legacy singular
`state_type` + flat-int wire format) that crashes hybrid-attention
models at PD-disagg startup. Concretely, three sites in
disaggregation/mori/conn.py break for any model whose state_pool is
populated as `state_item_lens: List[List[int]]` (Qwen3.5-MoE,
GLM-5, anything using state_types: List[StateType]):

  1. _register_kv_args raises `struct.error: required argument is not
     an integer` because it packs state_item_lens with naked
     struct.pack("I", item_len), expecting the legacy `List[int]`.
  2. send_state raises `state_type is 'none' but state_indices were
     provided` because it reads the legacy `kv_args.state_type`
     (singular), which the framework no longer populates for these
     models. The plural `kv_args.state_types: List[StateType]` -- the
     attribute Mooncake (mooncake/conn.py:912) and NIXL
     (nixl/conn.py:1381) already use -- is set correctly.
  3. The downstream length check at L935 and per-tensor indexing at
     L1005/1014/1096 assume flat-int `state_item_lens[i]`, so the
     run fails with `local state_item_lens count (1) does not match
     state descriptor count (90)` once registration is unblocked.

This commit carries an in-tree overlay of mori/conn.py with three
conservative fixes (flatten on the sender, state_types[0] fallback in
send_state, normalize state_item_lens once at send_state entry). The
overlay is bind-mounted into the container via job.slurm's existing
docker invocation, gated on the affected image tag and an opt-out
env var:

  - benchmarks/multi_node/amd_utils/patches/mori_conn.py     (new)
  - benchmarks/multi_node/amd_utils/patches/README.md        (new)
  - benchmarks/multi_node/amd_utils/job.slurm                (+1 hook,
    +1 auto-apply block; both no-op when DOCKER_IMAGE_NAME doesn't
    contain "v0.5.12.post1" or MORI_CONN_PATCH=skip is set)

Validated end-to-end on lmsysorg/sglang-rocm:v0.5.12.post1-rocm720-
mi35x-20260523 with Qwen3.5-397B-A17B-FP8 1P+1D TP=8 dp-attn=false:

  gsm8k strict-match     exact_match: 0.9780 +/- 0.004   PASS
  gsm8k flexible-extract exact_match: 0.9780 +/- 0.004   PASS

(Matches and slightly exceeds the 0.970 GSM8K reported by upstream
PR sgl-project/sglang#22665, which introduced state-transfer for
hybrid models but kept MoRI on the legacy singular API.)

The proper upstream fix is to migrate MoRI to the plural state_types
API that Mooncake and NIXL already use; full diff is in
scripts/sglang_disagg/docs/03-upstream-pr-proposal.md (out-of-tree
to keep this PR lean). Tracks sgl-project/sglang#21886.

Author: chunfangamd

benchmarks/multi_node/amd_utils/job.slurm

@@ -52,6 +52,32 @@ fi
 # $(pwd) is amd_utils/ (the sbatch submit dir); go up 3 levels to reach the repo root.
 export DI_REPO_DIR=$(cd "$(pwd)/../../.." && pwd)
 
+# ── In-tree sglang patches: auto-apply for known-affected images ──────
+# Some sglang releases ship known bugs in the MoRI PD-disaggregation
+# backend that crash hybrid-attention models (Qwen3.5-MoE, GLM-5,
+# anything using state_types: List[StateType]) at startup. We carry an
+# in-tree overlay of `mori/conn.py` that fixes the wire format + the
+# legacy state_type fallback (see benchmarks/multi_node/amd_utils/patches/README.md).
+#
+# Activate the overlay automatically when the image tag contains a
+# known-affected version string, unless the caller explicitly opts out
+# by setting MORI_CONN_PATCH=skip in the environment. The overlay is
+# appended to ${EXTRA_DOCKER_MOUNTS:-} so callers (CI workflow steps,
+# local drivers, etc.) can still inject additional mounts.
+#
+# Dedups: if EXTRA_DOCKER_MOUNTS already contains the target path,
+# don't double-mount (Docker rejects duplicate destinations).
+_MORI_PATCH_FILE="$DI_REPO_DIR/benchmarks/multi_node/amd_utils/patches/mori_conn.py"
+_MORI_PATCH_TARGET="/sgl-workspace/sglang/python/sglang/srt/disaggregation/mori/conn.py"
+if [[ "${MORI_CONN_PATCH:-auto}" != "skip" ]] \
+   && [[ -f "$_MORI_PATCH_FILE" ]] \
+   && [[ "${DOCKER_IMAGE_NAME:-}" == *"v0.5.12.post1"* ]] \
+   && [[ "${EXTRA_DOCKER_MOUNTS:-}" != *"$_MORI_PATCH_TARGET"* ]]; then
+    EXTRA_DOCKER_MOUNTS="${EXTRA_DOCKER_MOUNTS:-} -v ${_MORI_PATCH_FILE}:${_MORI_PATCH_TARGET}:ro"
+    export EXTRA_DOCKER_MOUNTS
+    echo "[job.slurm] auto-applied MoRI conn.py overlay: ${_MORI_PATCH_FILE}"
+fi
+
 xP="${xP:-1}" #-> Number of Prefill Workers
 yD="${yD:-1}" #-> Number of Decode Workers
 
@@ -387,6 +413,7 @@ exec \$_SUDO docker run --rm \
     -v /tmp:/run_logs \
     -v ${BENCHMARK_LOGS_DIR}:/benchmark_logs \
     -v ${DI_REPO_DIR}:${DOCKER_MOUNT_PATH} \
+    ${EXTRA_DOCKER_MOUNTS:-} \
     -e SLURM_JOB_ID=\$SLURM_JOB_ID \
     -e SLURM_JOB_NODELIST=\$SLURM_JOB_NODELIST \
     -e NNODES=\$NNODES \

benchmarks/multi_node/amd_utils/patches/README.md

@@ -0,0 +1,74 @@
+# In-tree sglang patches for the MoRI PD-disagg path
+
+This directory carries small Python overlays that get bind-mounted over
+the upstream sglang source inside the docker container at runtime.
+They are needed because some sglang releases ship known bugs in the
+MoRI disaggregation backend that block our benchmark + accuracy
+configs.
+
+The mount is wired through the `EXTRA_DOCKER_MOUNTS` env var that
+`job.slurm` consumes (an opt-in `${EXTRA_DOCKER_MOUNTS:-}` after the
+existing `-v` block). The local-test driver scripts under
+`scripts/sglang_disagg/` pre-set this env var to the path of the
+relevant overlay; CI runners that need the patch can do the same.
+
+## `mori_conn.py`
+
+Overlays
+`/sgl-workspace/sglang/python/sglang/srt/disaggregation/mori/conn.py`.
+
+Source: forked from the file shipped in
+`lmsysorg/sglang-rocm:v0.5.12.post1-rocm720-mi35x-20260523`
+(sglang [v0.5.12.post1](https://github.com/sgl-project/sglang/tree/v0.5.12.post1)).
+Three logical edits, all confined to `MoriKVReceiver.send_state` and
+`MoriKVReceiver._register_kv_args`:
+
+1. **Sender flatten** — handle the framework's nested
+   `state_item_lens: List[List[int]]` instead of crashing in the
+   naked `struct.pack("I", item_len)` (the legacy `List[int]`
+   assumption). Idempotent for legacy flat callers.
+2. **`state_type` legacy fallback** — when the legacy singular
+   `kv_args.state_type` is `'none'` but `state_mem_descs` is non-empty,
+   read `kv_args.state_types[0]` (the modern plural API that Mooncake
+   and NIXL already use). Routes `MAMBA → _send_mamba_state` and
+   `DSA/SWA → _send_swa_dsa_state` correctly.
+3. **Consumer normalization** — flatten `state_item_lens` and
+   `state_dim_per_tensor` to flat `List[int]` once at the entry of
+   `send_state`, so the existing per-tensor index arithmetic
+   (`state_item_lens[i]`) and length checks
+   (`len(state_item_lens) == len(state_mem_descs)`) keep working.
+
+Verified passing GSM8K = 0.978 ± 0.004 on Qwen3.5-397B-A17B-FP8 1P+1D
+TP=8 dp-attn=false (matches and slightly exceeds upstream
+[PR #22665](https://github.com/sgl-project/sglang/pull/22665)'s
+reported 0.970 GSM8K on the bf16 baseline).
+
+This is a stop-gap. The proper upstream fix is to migrate MoRI to the
+plural `state_types: List[StateType]` API (full design + diff in
+`scripts/sglang_disagg/docs/03-upstream-pr-proposal.md`).
+
+## How to enable
+
+```bash
+export EXTRA_DOCKER_MOUNTS="-v $DI_REPO_DIR/benchmarks/multi_node/amd_utils/patches/mori_conn.py:/sgl-workspace/sglang/python/sglang/srt/disaggregation/mori/conn.py:ro"
+```
+
+`$DI_REPO_DIR` is the InferenceX checkout root that `job.slurm`
+already mounts into the container at `/workspace`.
+
+When this env var is unset (CI default for runs that don't need the
+patch), `${EXTRA_DOCKER_MOUNTS:-}` expands to the empty string and
+container behavior is byte-identical to the unpatched path.
+
+## When to use which patch
+
+| Image / version | Need `mori_conn.py` overlay? |
+|---|---|
+| `lmsysorg/sglang-rocm:v0.5.12.post1-rocm720-mi35x-20260523` | yes (Qwen3.5-MoE-FP8, GLM-5, any hybrid model on this image) |
+| `lmsysorg/sglang-rocm:v0.5.10.post1-rocm720-mi35x-*` (used by `dsr1-fp4-*-disagg`) | not validated; same code path likely affected — try with the overlay if you hit the same `struct.error` |
+| `rocm/sgl-dev:sglang-0.5.9-rocm720-mi35x-mori-*` (used by `dsr1-fp8-*-disagg`, `glm5-*-disagg`) | predates [PR #22665](https://github.com/sgl-project/sglang/pull/22665); different code paths; **do not** apply this overlay |
+
+When upstream merges the proper fix (see
+`scripts/sglang_disagg/docs/03-upstream-pr-proposal.md`) and that
+fix lands in a published image, retire this overlay and the
+`EXTRA_DOCKER_MOUNTS` knob can stay (still useful for future patches).