[AMD] WIP - Clamp disagg max_total_num_tokens to MORI's single-MR size limit

.github/workflows/nightly-amd-mi355x-disagg.yml

@@ -118,6 +118,8 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
 
       - name: Clean up prior Slurm jobs from this runner
         continue-on-error: true
@@ -217,6 +219,16 @@ jobs:
             scancel $ACTIVE_JOBS
           fi
 
+      - name: Clean up MI355X scratch
+        if: always()
+        continue-on-error: true
+        run: |
+          LOG_DIR="$HOME/.mi355x_ci/${MATRIX_CONFIG_NAME}"
+          if [ -d "$LOG_DIR" ]; then
+            echo "Removing MI355X scratch directory: $LOG_DIR"
+            rm -rf "$LOG_DIR"
+          fi
+
   collect-results:
     needs: nightly-mi355x-benchmark
     if: github.repository == 'sgl-project/sglang' && always()

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -28,8 +28,8 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None
     # currently returns incorrect output for DeepSeek-V4-Flash on ROCm/HIP
     # (MI355X), which breaks the disaggregation nightly. Keep the previous
     # (dense prefill) behavior on ROCm until the sparse kernel is validated
-    # there; an explicit env var still overrides this.
-    if is_hip() and not envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.is_set():
+    # there;
+    if is_hip():
         logger.warning(
             "Disabling SGLANG_OPT_FLASHMLA_SPARSE_PREFILL by default on ROCm/HIP "
             f"for {model_arch}; set it explicitly to override."

python/sglang/srt/environ.py

@@ -472,6 +472,16 @@ class Envs:
     # Per-transfer SLA (ms) before a KV transfer is failed; 0 disables the SLA
     # and relies on the RDMA retry-exceeded timeout only.
     SGLANG_MORI_TRANSFER_TIMEOUT_MS = EnvInt(0)
+    # Upper bound (bytes) on a single KV memory region registered with MORI's
+    # RDMA backend. DeepSeek-V4 disagg auto-sizing can grow one unified C4 KV
+    # region past the RDMA single-registration ceiling, which makes
+    # RegisterRdmaMemoryRegion fail with errno=22 (EINVAL) so the PD path never
+    # serves. When the profiled pool would exceed this, max_total_num_tokens is
+    # clamped so the largest region fits. Empirically on MI355X a ~4 GiB region
+    # still fails errno=22 while ~2.19 GB registers fine, so the default is a
+    # conservative 2 GiB; raise it if your NIC/driver supports larger MRs, or
+    # set <= 0 to disable.
+    SGLANG_MORI_MAX_MR_BYTES = EnvInt(2 * 1024 * 1024 * 1024)
 
     # AMD & ROCm
     SGLANG_USE_AITER = EnvBool(False)

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -1222,6 +1222,49 @@ def _apply_token_constraints(self: ModelRunner, token_capacity: int) -> int:
 
         return token_capacity
 
+    def _apply_mori_mr_limit(
+        self: ModelRunner, token_capacity: int, configurator
+    ) -> int:
+        """Clamp token capacity so no single KV region exceeds MORI's RDMA
+        single-MR byte limit.
+
+        DeepSeek-V4 disagg auto-sizing can grow one unified C4 KV region past
+        ibv_reg_mr's ~4 GiB single-registration ceiling; MORI then fails
+        RegisterRdmaMemoryRegion with errno=22 (EINVAL) and the PD path never
+        serves. Keep the largest region strictly below SGLANG_MORI_MAX_MR_BYTES.
+        No-op unless this is a MORI-backed disaggregation worker.
+        """
+        if self.server_args.disaggregation_mode == "null":
+            return token_capacity
+        if self.server_args.disaggregation_transfer_backend != "mori":
+            return token_capacity
+
+        per_token_bytes = configurator.largest_registered_kv_region_bytes_per_token()
+        if per_token_bytes <= 0:
+            return token_capacity
+
+        limit = envs.SGLANG_MORI_MAX_MR_BYTES.get()
+        if limit <= 0:
+            return token_capacity
+
+        # Strictly below the limit, minus one page of tokens to absorb the
+        # per-buffer page rounding applied in get_contiguous_buf_infos.
+        max_tokens = (limit - 1) // per_token_bytes - self.server_args.page_size
+        if max_tokens <= 0 or token_capacity <= max_tokens:
+            return token_capacity
+
+        logger.warning(
+            "MORI single-MR cap: clamping max_total_num_tokens %d -> %d "
+            "(largest KV region ~%d bytes/token, SGLANG_MORI_MAX_MR_BYTES=%d "
+            "= %.2f GiB). Raise the env if the NIC/driver supports larger MRs.",
+            token_capacity,
+            max_tokens,
+            per_token_bytes,
+            limit,
+            limit / (1 << 30),
+        )
+        return max_tokens
+
     def _resolve_max_num_reqs(self: ModelRunner, token_capacity: int) -> int:
         """Compute max concurrent requests (per dp worker) from the finalized
         token capacity."""
@@ -1308,6 +1351,8 @@ def _resolve_memory_pool_config(
 
         # Apply external constraints (user cap, page alignment, PP sync)
         constrained = self._apply_token_constraints(config.max_total_num_tokens)
+        # Clamp so a single KV region stays within MORI's RDMA MR size limit.
+        constrained = self._apply_mori_mr_limit(constrained, configurator)
         if constrained != config.max_total_num_tokens:
             config = configurator.calculate_pool_sizes_from_max_tokens(
                 constrained, page_size

python/sglang/srt/model_executor/pool_configurator.py

@@ -114,6 +114,17 @@ def finalize_with_max_running_requests(
     ) -> MemoryPoolConfig:
         return config
 
+    def largest_registered_kv_region_bytes_per_token(self) -> int:
+        """Upper bound, per unit of max_total_num_tokens, on the byte size of
+        the largest single KV region that a disagg transfer backend registers
+        as one RDMA memory region.
+
+        Returns 0 ("no known bound") so callers skip clamping. Only
+        architectures whose per-region registration can exceed a backend's
+        single-MR ceiling override this.
+        """
+        return 0
+
 
 class DefaultPoolConfigurator(MemoryPoolConfigurator):
     """Configurator for standard models: MHA, MLA, DSA, FP4.
@@ -648,6 +659,24 @@ def _get_bytes_per_full_token(self) -> float:
             * self.num_layers_ca4
         )
 
+    def largest_registered_kv_region_bytes_per_token(self) -> int:
+        # Registered kv_data regions are per-layer compressed KV buffers (see
+        # DeepSeekV4TokenToKVPool.get_contiguous_buf_infos). The largest is a
+        # C4 layer: with the unified bf16 layout each compressed row is
+        # head_dim*2 bytes over ~max_total_num_tokens/ratio rows, i.e.
+        # (head_dim*2 / ratio) * max_total_num_tokens. The smallest positive
+        # (densest) compress ratio dominates. This also upper-bounds the
+        # packed-FP8 non-unified layout (< head_dim*2 bytes per compressed
+        # token). Dense layers carry compress ratio 0 (no compressed region is
+        # registered for them); exclude them so they don't zero the divisor.
+        compressed_ratios = [r for r in self.compression_ratios if r > 0]
+        if not compressed_ratios:
+            return 0
+        head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        min_ratio = min(compressed_ratios)
+        bf16_bytes = 2
+        return (head_dim * bf16_bytes + min_ratio - 1) // min_ratio
+
     def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes:
         full_token = full_token // page_size * page_size
         swa_tokens = int(full_token * self.swa_ratio) // page_size * page_size

scripts/ci/slurm/launch_mi355x.sh

@@ -24,6 +24,11 @@
 #   SLURM_NODELIST     - optional explicit node pin (else scheduler chooses)
 #   SLURM_EXCLUDE      - optional comma-separated nodes to keep the scheduler
 #                        off (e.g. hosts with a broken RDMA driver)
+#   SGLANG_USE_CHECKOUT_RUNTIME
+#                      - default 1. Reinstall this workflow checkout's Python
+#                        sglang package inside each runtime container before
+#                        launching servers/bench. Set 0 to use the image's
+#                        baked-in sglang package.
 #   RUNNER_NAME        - GitHub runner name (a built-in default env var)
 #   GITHUB_RUN_ID      - GitHub Actions run id (a built-in default env var)
 #                        The allocation is named
@@ -52,6 +57,11 @@ set -x
 SLURM_PARTITION="${SLURM_PARTITION:-amd-sglang}"
 TIME_LIMIT="${TIME_LIMIT:-02:30:00}"
 MODEL_PATH="${MODEL_PATH:-${MODEL:-}}"
+SGLANG_USE_CHECKOUT_RUNTIME="${SGLANG_USE_CHECKOUT_RUNTIME:-1}"
+case "${SGLANG_USE_CHECKOUT_RUNTIME,,}" in
+    0|false|no|off) SGLANG_USE_CHECKOUT_RUNTIME=0 ;;
+    *) SGLANG_USE_CHECKOUT_RUNTIME=1 ;;
+esac
 
 if [[ -z "$MODEL_PATH" ]]; then
     echo "ERROR: set MODEL_PATH (local snapshot) or MODEL" >&2
@@ -163,6 +173,23 @@ WORKDIR="$HOME/.mi355x_ci/${MATRIX_CONFIG_NAME}"
 rm -rf "$WORKDIR"; mkdir -p "$WORKDIR"
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
+# Stage the workflow checkout on shared NFS so Slurm compute-node containers can
+# reinstall the same code SHA the workflow checked out. The container gets a
+# read-only mount and copies it to /tmp before mutating pyproject.toml.
+CHECKOUT_DOCKER_ARGS="-e SGLANG_USE_CHECKOUT_RUNTIME=$SGLANG_USE_CHECKOUT_RUNTIME"
+if [[ "$SGLANG_USE_CHECKOUT_RUNTIME" == "1" ]]; then
+    CHECKOUT_STAGE="$WORKDIR/checkout"
+    CHECKOUT_SHA="$(git -C "$GITHUB_WORKSPACE" rev-parse HEAD)"
+    echo "Staging checkout runtime: sha=$CHECKOUT_SHA -> $CHECKOUT_STAGE"
+    rm -rf "$CHECKOUT_STAGE"
+    mkdir -p "$CHECKOUT_STAGE"
+    tar --exclude='__pycache__' --exclude='*.pyc' \
+        -C "$GITHUB_WORKSPACE" -cf - . | tar -C "$CHECKOUT_STAGE" -xf -
+    CHECKOUT_DOCKER_ARGS="$CHECKOUT_DOCKER_ARGS -e SGLANG_CHECKOUT_SHA=$CHECKOUT_SHA -v $CHECKOUT_STAGE:/sglang-checkout:ro"
+else
+    echo "SGLANG_USE_CHECKOUT_RUNTIME=0; using sglang package baked into image."
+fi
+
 # Accuracy-gate helpers (written when enabled). Pre-stage the GSM8K test set on
 # shared NFS from the login node (which has internet) so the in-container eval
 # doesn't depend on compute-node connectivity; fall back to in-container
@@ -281,7 +308,7 @@ fi
 DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \
 --security-opt seccomp=unconfined \
 --device /dev/kfd --device /dev/dri --device /dev/infiniband \
--v /it-share:/it-share:ro -v $HOME:/host_home"
+-v /it-share:/it-share:ro -v $HOME:/host_home $CHECKOUT_DOCKER_ARGS"
 
 # ---------------------------------------------------------------------------
 # Write per-role scripts that srun dispatches to each compute node.
@@ -292,16 +319,109 @@ DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \
 # `${MODEL_SERVER_ARGS[@]}` refs are backslash-escaped to survive into the script
 # and expand after `source`. For DSV4 those arrays are empty and $DSV4_ENV_STR is
 # set, so the resulting docker argv is byte-identical to the pre-Kimi launcher.
+cat > "$WORKDIR/install_checkout_sglang.sh" <<'EOF'
+#!/bin/bash
+set -euo pipefail
+
+case "${SGLANG_USE_CHECKOUT_RUNTIME:-1}" in
+  0|false|False|FALSE|no|No|NO|off|Off|OFF)
+    echo "[checkout-sglang] disabled; using image-baked sglang"
+    exit 0
+    ;;
+esac
+
+CHECKOUT_SRC="${CHECKOUT_SRC:-/sglang-checkout}"
+RUNTIME_CHECKOUT="${RUNTIME_CHECKOUT:-/tmp/sglang-checkout-runtime}"
+
+if [[ ! -f "$CHECKOUT_SRC/python/sglang/version.py" ]]; then
+  echo "[checkout-sglang] ERROR: invalid checkout mount: $CHECKOUT_SRC" >&2
+  exit 1
+fi
+
+echo "[checkout-sglang] reinstalling sglang from $CHECKOUT_SRC"
+rm -rf "$RUNTIME_CHECKOUT"
+mkdir -p "$RUNTIME_CHECKOUT"
+tar --exclude='__pycache__' --exclude='*.pyc' \
+  -C "$CHECKOUT_SRC" -cf - . | tar -C "$RUNTIME_CHECKOUT" -xf -
+
+git config --global --add safe.directory "$RUNTIME_CHECKOUT" || true
+
+# The ROCm pyproject variant is the one used by AMD CI. Mutate only the private
+# /tmp copy so prefill/decode/bench never race on the read-only checkout mount.
+rm -f "$RUNTIME_CHECKOUT/python/pyproject.toml"
+cp "$RUNTIME_CHECKOUT/python/pyproject_other.toml" "$RUNTIME_CHECKOUT/python/pyproject.toml"
+for f in README.md LICENSE; do
+  if [[ -f "$RUNTIME_CHECKOUT/$f" && ! -e "$RUNTIME_CHECKOUT/python/$f" ]]; then
+    cp "$RUNTIME_CHECKOUT/$f" "$RUNTIME_CHECKOUT/python/$f"
+  fi
+done
+
+python3 -m pip uninstall -y sglang || true
+python3 -m pip install --no-deps --no-build-isolation -e "$RUNTIME_CHECKOUT/python"
+
+export RUNTIME_CHECKOUT
+export PYTHONPATH="$RUNTIME_CHECKOUT/python:${PYTHONPATH:-}"
+python3 - <<'PY'
+import importlib.metadata
+import os
+import subprocess
+import sglang
+
+checkout = os.environ["RUNTIME_CHECKOUT"]
+expected = os.path.realpath(os.path.join(checkout, "python", "sglang")) + os.sep
+actual = os.path.realpath(os.path.dirname(sglang.__file__)) + os.sep
+try:
+    sha = subprocess.check_output(
+        ["git", "-C", checkout, "rev-parse", "HEAD"], text=True
+    ).strip()
+except Exception:
+    sha = os.environ.get("SGLANG_CHECKOUT_SHA", "unknown")
+
+print(f"[checkout-sglang] sha={sha}")
+print(f"[checkout-sglang] sglang_file={sglang.__file__}")
+print(f"[checkout-sglang] sglang_version={importlib.metadata.version('sglang')}")
+if not actual.startswith(expected):
+    raise SystemExit(f"sglang did not import from checkout: {sglang.__file__}")
+PY
+EOF
+
+cat > "$WORKDIR/prefill_entry.sh" <<EOF
+#!/bin/bash
+set -euo pipefail
+CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
+source "\$CIDIR/model_flags.sh"
+bash "\$CIDIR/install_checkout_sglang.sh"
+if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
+  export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
+fi
+exec python3 -m sglang.launch_server \
+  --model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
+  $COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
+  --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
+EOF
+
+cat > "$WORKDIR/decode_entry.sh" <<EOF
+#!/bin/bash
+set -euo pipefail
+CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
+source "\$CIDIR/model_flags.sh"
+bash "\$CIDIR/install_checkout_sglang.sh"
+if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
+  export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
+fi
+exec python3 -m sglang.launch_server \
+  --model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
+  $COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
+  --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
+EOF
+
 cat > "$WORKDIR/prefill.sh" <<EOF
 #!/bin/bash
 source "$WORKDIR/model_flags.sh"
 docker rm -f mi355x_prefill 2>/dev/null || true
 docker run $DOCKER_COMMON --name mi355x_prefill \
   -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
-  $IMAGE python3 -m sglang.launch_server \
-  --model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
-  $COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
-  --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
+  $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/prefill_entry.sh
 EOF
 
 cat > "$WORKDIR/decode.sh" <<EOF
@@ -310,10 +430,7 @@ source "$WORKDIR/model_flags.sh"
 docker rm -f mi355x_decode 2>/dev/null || true
 docker run $DOCKER_COMMON --name mi355x_decode \
   -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
-  $IMAGE python3 -m sglang.launch_server \
-  --model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
-  $COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
-  --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
+  $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/decode_entry.sh
 EOF
 
 # Probe payload + validator (separate files to avoid quoting inside the
@@ -341,7 +458,13 @@ docker rm -f mi355x_bench 2>/dev/null || true
 docker run $DOCKER_COMMON --name mi355x_bench \
   -e PIP=\$PIP -e DIP=\$DIP \
   $IMAGE bash -lc '
-    export PYTHONPATH=/sgl-workspace/sglang/python:\$PYTHONPATH
+    CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
+    bash \$CIDIR/install_checkout_sglang.sh
+    if [ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]; then
+      export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
+    else
+      export PYTHONPATH=/sgl-workspace/sglang/python:\${PYTHONPATH:-}
+    fi
     echo "[wait] prefill"; for i in \$(seq 1 600); do curl -sf http://\$PIP:$PPORT/health >/dev/null && break; sleep 5; done
     echo "[wait] decode";  for i in \$(seq 1 600); do curl -sf http://\$DIP:$DPORT/health >/dev/null && break; sleep 5; done
     python3 -m sglang_router.launch_router \
@@ -351,7 +474,6 @@ docker run $DOCKER_COMMON --name mi355x_bench \
       --host 0.0.0.0 --port $LBPORT \
       --disable-circuit-breaker &
     for i in \$(seq 1 30); do curl -sf http://127.0.0.1:$LBPORT/health >/dev/null && break; sleep 2; done
-    CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
     echo "[probe] PD end-to-end check via LB"
     curl -sf -X POST http://127.0.0.1:$LBPORT/generate \
       -H "content-type: application/json" -d @\$CIDIR/probe.json > \$CIDIR/probe_out.json \

scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/1p1d.yaml

@@ -44,10 +44,13 @@ model:
 
 runtime:
   image: lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260623
-  # Kimi uses split attention backends (aiter prefill / triton decode), not a
-  # single --attention-backend.
+  # Kimi normally uses split backends (aiter prefill / triton decode). Experiment:
+  # non-MTP disagg (triton decode reading transferred KV) only reaches ~0.88, while
+  # MTP verify — which dispatch_attn_forward_method routes to the PREFILL backend
+  # (aiter) under speculative_attention_mode=prefill — passes at ~0.95. So switch
+  # decode to aiter to test whether the triton MLA decode path is what degrades it.
   prefill_attention_backend: aiter
-  decode_attention_backend: triton
+  decode_attention_backend: aiter
   # RoCE HCAs MORI uses for cross-node KV transfer.
   ib_devices: rdma0,rdma1,rdma2,rdma3
   prefill_port: 30025