chore: vendor sglang v0.5.10 snapshot

2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
--- a/third_party/sglang/scripts/ci/amd/amd_ci_exec.sh
+++ b/third_party/sglang/scripts/ci/amd/amd_ci_exec.sh
@@ -0,0 +1,89 @@
+#!/bin/bash
+set -euo pipefail
+
+# Detect GPU family from hostname (e.g., linux-mi35x-gpu-1-xxxxx-runner-zzzzz)
+HOSTNAME_VALUE=$(hostname)
+GPU_FAMILY=""
+
+# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
+if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
+  GPU_FAMILY="${BASH_REMATCH[1]}"
+  echo "Detected GPU family from hostname: ${GPU_FAMILY}"
+else
+  echo "Warning: could not parse GPU family from '${HOSTNAME_VALUE}'"
+fi
+
+WORKDIR="/sglang-checkout/test/srt"
+declare -A ENV_MAP=(
+  [SGLANG_IS_IN_CI_AMD]=1
+  [SGLANG_IS_IN_CI]=1
+  [SGLANG_USE_AITER]=1
+)
+
+# Conditionally add GPU_ARCHS only for mi35x
+if [[ "${GPU_FAMILY}" == "mi35x" ]]; then
+  ENV_MAP[GPU_ARCHS]="gfx950"
+fi
+
+# Parse -w/--workdir and -e ENV=VAL
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    -w|--workdir)
+      WORKDIR="$2"
+      shift 2
+      ;;
+    -e)
+      IFS="=" read -r key val <<< "$2"
+      ENV_MAP["$key"]="$val"
+      shift 2
+      ;;
+    --)
+      shift
+      break
+      ;;
+    *)
+      break
+      ;;
+  esac
+done
+
+# Build final ENV_ARGS
+ENV_ARGS=()
+for key in "${!ENV_MAP[@]}"; do
+  ENV_ARGS+=("-e" "$key=${ENV_MAP[$key]}")
+done
+
+# Run docker exec with retry logic for HuggingFace network/download issues
+# When HF model downloads fail due to network timeouts or rate limits,
+# retrying with HF_HUB_OFFLINE=1 uses cached models from previous downloads.
+#
+# First attempt: normal mode (allows HF downloads)
+if docker exec \
+  -w "$WORKDIR" \
+  "${ENV_ARGS[@]}" \
+  ci_sglang "$@"; then
+  exit 0
+else
+  FIRST_EXIT_CODE=$?
+fi
+
+echo "First attempt failed with exit code $FIRST_EXIT_CODE"
+
+# Skip retry for test failures that won't be fixed by offline mode:
+#   - Exit 1: Test assertion failures (accuracy below threshold)
+#   - Exit 137 (128+9): Process killed by OOM
+#   - Exit 255: Test suite completed with test errors
+# Only retry for other exit codes (e.g., network timeouts, HF download failures)
+if [[ "$FIRST_EXIT_CODE" -eq 1 || "$FIRST_EXIT_CODE" -eq 137 || "$FIRST_EXIT_CODE" -eq 255 ]]; then
+  echo "Exit code $FIRST_EXIT_CODE indicates test failure (not network issue), not retrying"
+  exit $FIRST_EXIT_CODE
+fi
+
+echo "Retrying with HF_HUB_OFFLINE=1 (offline mode to use cached models)..."
+
+# Second attempt: force HF offline mode to avoid network timeouts
+docker exec \
+  -w "$WORKDIR" \
+  "${ENV_ARGS[@]}" \
+  -e HF_HUB_OFFLINE=1 \
+  ci_sglang "$@"
--- a/third_party/sglang/scripts/ci/amd/amd_ci_install_dependency.sh
+++ b/third_party/sglang/scripts/ci/amd/amd_ci_install_dependency.sh
@@ -0,0 +1,320 @@
+#!/bin/bash
+set -euo pipefail
+HOSTNAME_VALUE=$(hostname)
+GPU_ARCH="mi30x"   # default
+SKIP_TT_DEPS=""
+SKIP_SGLANG_BUILD=""
+SKIP_AITER_BUILD=""
+
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --skip-aiter-build) SKIP_AITER_BUILD="1"; shift;;
+    --skip-sglang-build) SKIP_SGLANG_BUILD="1"; shift;;
+    --skip-test-time-deps) SKIP_TT_DEPS="1"; shift;;
+    -h|--help)
+      echo "Usage: $0 [OPTIONS] [OPTIONAL_DEPS]"
+      echo "Options:"
+      echo "  --skip-sglang-build         Don't build checkout sglang, use what was shipped with the image"
+      echo "  --skip-aiter-build          Don't build aiter, use what was shipped with the image"
+      echo "  --skip-test-time-deps       Don't build miscellaneous dependencies"
+      exit 0
+      ;;
+    *) break ;;
+  esac
+done
+
+OPTIONAL_DEPS="${1:-}"
+
+# Build python extras
+EXTRAS="dev_hip"
+if [ -n "$OPTIONAL_DEPS" ]; then
+    EXTRAS="dev_hip,${OPTIONAL_DEPS}"
+fi
+echo "Installing python extras: [${EXTRAS}]"
+
+# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
+if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
+  GPU_ARCH="${BASH_REMATCH[1]}"
+  echo "Detected GPU architecture from hostname: ${GPU_ARCH}"
+else
+  echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}"
+fi
+
+# Install the required dependencies in CI.
+# Fix permissions on pip cache, ignore errors from concurrent access or missing temp files
+docker exec ci_sglang chown -R root:root /sgl-data/pip-cache 2>/dev/null || true
+docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache --upgrade pip
+
+# Helper function to install with retries and fallback PyPI mirror
+install_with_retry() {
+  local max_attempts=3
+  local cmd="$@"
+
+  for attempt in $(seq 1 $max_attempts); do
+    echo "Attempt $attempt/$max_attempts: $cmd"
+    if eval "$cmd"; then
+      echo "Success!"
+      return 0
+    fi
+
+    if [ $attempt -lt $max_attempts ]; then
+      echo "Failed, retrying in 5 seconds..."
+      sleep 5
+      # Try with alternative PyPI index on retry
+      if [[ "$cmd" =~ "pip install" ]] && [ $attempt -eq 2 ]; then
+        cmd="$cmd --index-url https://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com"
+        echo "Using fallback PyPI mirror: $cmd"
+      fi
+    fi
+  done
+
+  echo "Failed after $max_attempts attempts"
+  return 1
+}
+
+# Helper function to git clone with retries
+git_clone_with_retry() {
+  local repo_url="$1"
+  local dest_dir="${2:-}"
+  local branch_args="${3:-}"
+  local max_attempts=3
+
+  for attempt in $(seq 1 $max_attempts); do
+    echo "Git clone attempt $attempt/$max_attempts: $repo_url"
+
+    # prevent from partial clone
+    if [ -n "$dest_dir" ] && [ -d "$dest_dir" ]; then
+      rm -rf "$dest_dir"
+    fi
+
+    if git \
+      -c http.lowSpeedLimit=1000 \
+      -c http.lowSpeedTime=30 \
+      clone --depth 1 ${branch_args:+$branch_args} "$repo_url" "$dest_dir"; then
+      echo "Git clone succeeded."
+      return 0
+    fi
+
+    if [ $attempt -lt $max_attempts ]; then
+      echo "Git clone failed, retrying in 5 seconds..."
+      sleep 5
+    fi
+  done
+
+  echo "Git clone failed after $max_attempts attempts: $repo_url"
+  return 1
+}
+
+# Install checkout sglang
+if [ -n "$SKIP_SGLANG_BUILD" ]; then
+  echo "Didn't build checkout SGLang"
+else
+  docker exec ci_sglang pip uninstall sgl-kernel -y || true
+  docker exec ci_sglang pip uninstall sglang-kernel -y || true
+  docker exec ci_sglang pip uninstall sglang -y || true
+  # Clear Python cache to ensure latest code is used
+  docker exec ci_sglang find /opt/venv -name "*.pyc" -delete || true
+  docker exec ci_sglang find /opt/venv -name "__pycache__" -type d -exec rm -rf {} + || true
+  # Also clear cache in sglang-checkout
+  docker exec ci_sglang find /sglang-checkout -name "*.pyc" -delete || true
+  docker exec ci_sglang find /sglang-checkout -name "__pycache__" -type d -exec rm -rf {} + || true
+  docker exec -w /sglang-checkout/sgl-kernel ci_sglang bash -c "rm -f pyproject.toml && mv pyproject_rocm.toml pyproject.toml && python3 setup_rocm.py install"
+
+  docker exec ci_sglang bash -c 'rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml'
+  install_with_retry docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache -e "python[${EXTRAS}]"
+fi
+
+if [[ -n "${SKIP_TT_DEPS}" ]]; then
+  echo "Didn't build lmms_eval, human-eval, and others"
+else
+  # For lmms_evals evaluating MMMU
+  docker exec -w / ci_sglang git clone --branch v0.4.1 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git
+  install_with_retry docker exec -w /lmms-eval ci_sglang pip install --cache-dir=/sgl-data/pip-cache -e .
+
+  git_clone_with_retry https://github.com/akao-amd/human-eval.git human-eval
+  docker cp human-eval ci_sglang:/
+  install_with_retry docker exec -w /human-eval ci_sglang pip install --cache-dir=/sgl-data/pip-cache -e .
+
+  docker exec -w / ci_sglang mkdir -p /dummy-grok
+  # Create dummy grok config inline (bypasses Azure blob storage which may have auth issues)
+  mkdir -p dummy-grok
+  cat > dummy-grok/config.json << 'EOF'
+  {
+    "architectures": [
+      "Grok1ModelForCausalLM"
+    ],
+    "embedding_multiplier_scale": 78.38367176906169,
+    "output_multiplier_scale": 0.5773502691896257,
+    "vocab_size": 131072,
+    "hidden_size": 6144,
+    "intermediate_size": 32768,
+    "max_position_embeddings": 8192,
+    "num_experts_per_tok": 2,
+    "num_local_experts": 8,
+    "num_attention_heads": 48,
+    "num_hidden_layers": 64,
+    "num_key_value_heads": 8,
+    "head_dim": 128,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 10000.0,
+    "model_type": "mixtral",
+    "torch_dtype": "bfloat16"
+  }
+EOF
+  # docker exec -w / ci_sglang mkdir -p /dummy-grok
+  # mkdir -p dummy-grok && wget https://sharkpublic.blob.core.windows.net/sharkpublic/sglang/dummy_grok.json -O dummy-grok/config.json
+  # docker cp ./dummy-grok ci_sglang:/
+
+  docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache huggingface_hub[hf_xet]
+  docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache pytest
+
+  # Install cache-dit for qwen_image_t2i_cache_dit_enabled test (added in PR 16204)
+  docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache cache-dit || echo "cache-dit installation failed"
+
+  # Install accelerate for distributed training and inference support
+  docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache accelerate || echo "accelerate installation failed"
+fi
+
+if [[ -n "${SKIP_AITER_BUILD}" ]]; then
+  exit 0
+fi
+
+# Detect AITER version
+#############################################
+# Detect correct AITER_COMMIT for this runner
+# + Check mismatch
+# + Rebuild AITER if needed
+#############################################
+
+echo "[CI-AITER-CHECK] === AITER VERSION CHECK START ==="
+
+DOCKERFILE="docker/rocm.Dockerfile"
+
+# GPU_ARCH
+GPU_ARCH="${GPU_ARCH:-mi30x}"
+echo "[CI-AITER-CHECK] Runner GPU_ARCH=${GPU_ARCH}"
+
+#############################################
+# 1. Extract AITER_COMMIT from correct Dockerfile block
+#############################################
+if [[ "${GPU_ARCH}" == "mi35x" ]]; then
+    echo "[CI-AITER-CHECK] Using gfx950 block from Dockerfile..."
+    REPO_AITER_COMMIT=$(grep -F -A20 'FROM $BASE_IMAGE_950 AS gfx950' docker/rocm.Dockerfile \
+                        | grep 'AITER_COMMIT_DEFAULT=' \
+                        | head -n1 \
+                        | sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
+else
+    echo "[CI-AITER-CHECK] Using gfx942 block from Dockerfile..."
+    REPO_AITER_COMMIT=$(grep -F -A20 'FROM $BASE_IMAGE_942 AS gfx942' docker/rocm.Dockerfile \
+                        | grep 'AITER_COMMIT_DEFAULT=' \
+                        | head -n1 \
+                        | sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
+fi
+
+
+if [[ -z "${REPO_AITER_COMMIT}" ]]; then
+    echo "[CI-AITER-CHECK] ERROR: Failed to extract AITER_COMMIT from Dockerfile."
+    exit 1
+fi
+
+echo "[CI-AITER-CHECK] Dockerfile expects AITER_COMMIT=${REPO_AITER_COMMIT}"
+
+#############################################
+# 2. Check container pre-installed AITER version
+#############################################
+IMAGE_AITER_VERSION=$(docker exec ci_sglang bash -c "pip show amd-aiter 2>/dev/null | grep '^Version:' | awk '{print \$2}'" || echo "none")
+IMAGE_AITER_VERSION="v${IMAGE_AITER_VERSION}"
+echo "[CI-AITER-CHECK] AITER version inside CI image: ${IMAGE_AITER_VERSION}"
+
+#############################################
+# 3. Decide rebuild
+#############################################
+NEED_REBUILD="false"
+
+if [[ -n "${AITER_COMMIT_OVERRIDE:-}" ]]; then
+    echo "[CI-AITER-CHECK] AITER_COMMIT_OVERRIDE=${AITER_COMMIT_OVERRIDE} → forcing rebuild"
+    REPO_AITER_COMMIT="${AITER_COMMIT_OVERRIDE}"
+    NEED_REBUILD="true"
+elif [[ "${IMAGE_AITER_VERSION}" == "vnone" || "${IMAGE_AITER_VERSION}" == "v" ]]; then
+    echo "[CI-AITER-CHECK] No AITER found in image → rebuild needed"
+    NEED_REBUILD="true"
+elif [[ "${IMAGE_AITER_VERSION}" == "${REPO_AITER_COMMIT}" ]]; then
+    echo "[CI-AITER-CHECK] AITER version matches"
+elif [[ "${IMAGE_AITER_VERSION}" =~ (dev|\+g[0-9a-f]+) ]]; then
+    # Dev/patched version (contains 'dev' or git hash) → preserve it
+    echo "[CI-AITER-CHECK] Dev/patched version detected: ${IMAGE_AITER_VERSION} → skipping rebuild"
+else
+    echo "[CI-AITER-CHECK] Version mismatch: image=${IMAGE_AITER_VERSION}, repo=${REPO_AITER_COMMIT}"
+    NEED_REBUILD="true"
+fi
+
+
+#############################################
+# 4. Rebuild AITER if needed
+#############################################
+if [[ "${NEED_REBUILD}" == "true" ]]; then
+    echo "[CI-AITER-CHECK] === AITER REBUILD START ==="
+
+    # uninstall existing aiter
+    docker exec ci_sglang pip uninstall -y amd-aiter || true
+
+    # delete old aiter directory
+    docker exec ci_sglang rm -rf /sgl-workspace/aiter
+
+    # clone a fresh copy to /sgl-workspace/aiter
+    docker exec ci_sglang git clone https://github.com/ROCm/aiter.git /sgl-workspace/aiter
+
+    # checkout correct version
+    docker exec ci_sglang bash -c "
+        cd /sgl-workspace/aiter && \
+        git fetch --all && \
+        git checkout ${REPO_AITER_COMMIT} && \
+        git submodule update --init --recursive
+    "
+
+    if [[ "${GPU_ARCH}" == "mi35x" ]]; then
+        GPU_ARCH_LIST="gfx950"
+    else
+        GPU_ARCH_LIST="gfx942"
+    fi
+    echo "[CI-AITER-CHECK] GPU_ARCH_LIST=${GPU_ARCH_LIST}"
+
+    # Re-apply Dockerfile hotpatches for ROCm 7.2 (the fresh clone lost them, can be removed after triton fixed this problem)
+    ROCM_VERSION=$(docker exec ci_sglang bash -c "cat /opt/rocm/.info/version 2>/dev/null || echo unknown")
+    if [[ "${ROCM_VERSION}" == 7.2* ]]; then
+        echo "[CI-AITER-CHECK] ROCm 7.2 detected (${ROCM_VERSION}), applying AITER hotpatches..."
+        docker exec ci_sglang bash -c "
+            cd /sgl-workspace/aiter && \
+            TARGET_FILE='aiter/ops/triton/attention/pa_mqa_logits.py' && \
+            if [ -f \"\${TARGET_FILE}\" ]; then \
+                sed -i '459 s/if.*:/if False:/' \"\${TARGET_FILE}\" && \
+                echo '[CI-AITER-CHECK] Hotpatch applied to pa_mqa_logits.py'; \
+            else \
+                echo '[CI-AITER-CHECK] pa_mqa_logits.py not found, skipping hotpatch'; \
+            fi
+        "
+    else
+        echo "[CI-AITER-CHECK] ROCm version=${ROCM_VERSION}, no hotpatch needed"
+    fi
+
+    # build AITER
+    docker exec ci_sglang bash -c "
+        cd /sgl-workspace/aiter && \
+        GPU_ARCHS=${GPU_ARCH_LIST} python3 setup.py develop
+    "
+
+    echo "[CI-AITER-CHECK] === AITER REBUILD COMPLETE ==="
+fi
+
+echo "[CI-AITER-CHECK] === AITER VERSION CHECK END ==="
+
+
+# # Clear pre-built AITER kernels from Docker image to avoid segfaults
+# # The Docker image may contain pre-compiled kernels incompatible with the current environment
+# echo "Clearing pre-built AITER kernels from Docker image..."
+# docker exec ci_sglang find /sgl-workspace/aiter/aiter/jit -name "*.so" -delete 2>/dev/null || true
+# docker exec ci_sglang ls -la /sgl-workspace/aiter/aiter/jit/ 2>/dev/null || echo "jit dir empty or not found"
+
+# # Pre-build AITER kernels to avoid timeout during tests
+# echo "Warming up AITER JIT kernels..."
+# docker exec -e SGLANG_USE_AITER=1 ci_sglang python3 /sglang-checkout/scripts/ci/amd/amd_ci_warmup_aiter.py || echo "AITER warmup completed (some kernels may not be available)"
--- a/third_party/sglang/scripts/ci/amd/amd_ci_start_container.sh
+++ b/third_party/sglang/scripts/ci/amd/amd_ci_start_container.sh
@@ -0,0 +1,248 @@
+#!/bin/bash
+set -euo pipefail
+
+# Get version from git tags
+SGLANG_VERSION="v0.5.5"   # Default version, will be overridden if git tags are found
+
+# Fetch tags from origin to ensure we have the latest
+if git fetch --tags origin; then
+  # Get the latest version tag sorted by version number (e.g., v0.5.7)
+  VERSION_FROM_TAG=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1)
+  if [ -n "$VERSION_FROM_TAG" ]; then
+    SGLANG_VERSION="$VERSION_FROM_TAG"
+    echo "Using SGLang version from git tags: $SGLANG_VERSION"
+  else
+    echo "Warning: No version tags found; using default $SGLANG_VERSION" >&2
+  fi
+else
+  echo "Warning: Failed to fetch tags from origin; using default $SGLANG_VERSION" >&2
+fi
+
+
+# Default base tags (can be overridden by command line arguments)
+ROCM_VERSION="rocm700"
+DEFAULT_MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
+DEFAULT_MI35X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi35x"
+
+# Parse command line arguments
+MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}"
+MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}"
+CUSTOM_IMAGE=""
+BUILD_FROM_DOCKERFILE=""
+GPU_ARCH_BUILD=""
+
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --mi30x-base-tag) MI30X_BASE_TAG="$2"; shift 2;;
+    --mi35x-base-tag) MI35X_BASE_TAG="$2"; shift 2;;
+    --custom-image) CUSTOM_IMAGE="$2"; shift 2;;
+    --build-from-dockerfile) BUILD_FROM_DOCKERFILE="1"; shift;;
+    --gpu-arch) GPU_ARCH_BUILD="$2"; shift 2;;
+    --rocm-version)
+      ROCM_VERSION="$2"
+      MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
+      MI35X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi35x"
+      echo "Using ROCm version override: ${ROCM_VERSION}"
+      shift 2;;
+    -h|--help)
+      echo "Usage: $0 [OPTIONS]"
+      echo "Options:"
+      echo "  --mi30x-base-tag TAG       Override MI30x base image tag"
+      echo "  --mi35x-base-tag TAG       Override MI35x base image tag"
+      echo "  --custom-image IMAGE       Use a specific Docker image directly"
+      echo "  --build-from-dockerfile    Build image from docker/rocm.Dockerfile"
+      echo "  --gpu-arch ARCH            GPU architecture for Dockerfile build (e.g., gfx950-rocm720)"
+      echo "  --rocm-version VERSION     Override ROCm version for image lookup (e.g., rocm720)"
+      exit 0
+      ;;
+    *) echo "Unknown option $1"; exit 1;;
+  esac
+done
+
+
+
+# Detect GPU architecture from the Kubernetes runner hostname
+HOSTNAME_VALUE=$(hostname)
+GPU_ARCH="mi30x"   # default
+
+# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
+if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
+  GPU_ARCH="${BASH_REMATCH[1]}"
+  echo "Detected GPU architecture from hostname: ${GPU_ARCH}"
+else
+  echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}"
+fi
+
+# Normalise / collapse architectures we don't yet build specifically for
+case "${GPU_ARCH}" in
+  mi35x)
+    echo "Runner uses ${GPU_ARCH}; will fetch mi35x image."
+    ;;
+  mi30x|mi300|mi325)
+    echo "Runner uses ${GPU_ARCH}; will fetch mi30x image."
+    GPU_ARCH="mi30x"
+    ;;
+  *)
+    echo "Runner architecture '${GPU_ARCH}' unrecognised; defaulting to mi30x image." >&2
+    GPU_ARCH="mi30x"
+    ;;
+esac
+
+
+# Set up DEVICE_FLAG based on Kubernetes pod info
+if [[ -f /etc/podinfo/gha-render-devices ]]; then
+  DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
+else
+  DEVICE_FLAG="--device /dev/dri"
+fi
+
+
+# Find the latest image
+find_latest_image() {
+  local gpu_arch=$1
+  local base_tag days_back image_tag
+
+  case "${gpu_arch}" in
+      mi30x) base_tag="${MI30X_BASE_TAG}" ;;
+      mi35x) base_tag="${MI35X_BASE_TAG}" ;;
+      *)     echo "Error: unsupported GPU architecture '${gpu_arch}'" >&2; return 1 ;;
+  esac
+
+  # First, check local cache
+  for days_back in {0..6}; do
+    image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
+    local local_image="rocm/sgl-dev:${image_tag}"
+    image_id=$(docker images -q "${local_image}")
+    if [[ -n "$image_id" ]]; then
+        echo "Found cached image locally: ${local_image}" >&2
+        echo "${local_image}"
+        return 0
+    fi
+  done
+
+  # If not found locally, fall back to pulling from public registry
+  for days_back in {0..6}; do
+    image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
+    echo "Checking for image: rocm/sgl-dev:${image_tag}" >&2
+    if docker manifest inspect "rocm/sgl-dev:${image_tag}" >/dev/null 2>&1; then
+      echo "Found available image: rocm/sgl-dev:${image_tag}" >&2
+      echo "rocm/sgl-dev:${image_tag}"
+      return 0
+    fi
+  done
+
+  # If still not found, try finding any image matching ROCm+arch from remote registry
+  echo "Exact version not found. Searching remote registry for any ${ROCM_VERSION}-${gpu_arch} image…" >&2
+  for days_back in {0..6}; do
+    local target_date=$(date -d "${days_back} days ago" +%Y%m%d)
+    local remote_tags=$(curl -s "https://registry.hub.docker.com/v2/repositories/rocm/sgl-dev/tags?page_size=100&name=${ROCM_VERSION}-${gpu_arch}-${target_date}" 2>/dev/null | grep -o '"name":"[^"]*"' | cut -d'"' -f4 | head -n 1)
+    if [[ -n "$remote_tags" ]]; then
+      echo "Found available image: rocm/sgl-dev:${remote_tags}" >&2
+      echo "rocm/sgl-dev:${remote_tags}"
+      return 0
+    fi
+  done
+
+  echo "No recent images found. Searching any cached local images matching ROCm+arch…" >&2
+  local any_local
+  any_local=$(docker images --format '{{.Repository}}:{{.Tag}}' --filter "reference=rocm/sgl-dev:*${ROCM_VERSION}*${gpu_arch}*" | sort -r | head -n 1)
+  if [[ -n "$any_local" ]]; then
+      echo "Using cached fallback image: ${any_local}" >&2
+      echo "${any_local}"
+      return 0
+  fi
+
+  echo "Error: no ${gpu_arch} image found in the last 7 days for base ${base_tag}" >&2
+  echo "Using hard-coded fallback for ${ROCM_VERSION}…" >&2
+  case "${ROCM_VERSION}" in
+    rocm720)
+      if [[ "${gpu_arch}" == "mi35x" ]]; then
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm720-mi35x-20260211-preview"
+      else
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm720-mi30x-20260211-preview"
+      fi
+      ;;
+    rocm700)
+      if [[ "${gpu_arch}" == "mi35x" ]]; then
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm700-mi35x-20260211"
+      else
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm700-mi30x-20260211"
+      fi
+      ;;
+    *)
+      echo "Error: no hard-coded fallback available for ${ROCM_VERSION}" >&2
+      return 1
+      ;;
+  esac
+}
+
+# Determine which image to use
+if [[ -n "${CUSTOM_IMAGE}" ]]; then
+  # Use explicitly provided custom image
+  IMAGE="${CUSTOM_IMAGE}"
+  echo "Using custom image: ${IMAGE}"
+  docker pull "${IMAGE}"
+elif [[ -n "${BUILD_FROM_DOCKERFILE}" ]]; then
+  # Build image from Dockerfile
+  if [[ -z "${GPU_ARCH_BUILD}" ]]; then
+    echo "Error: --gpu-arch is required when using --build-from-dockerfile" >&2
+    exit 1
+  fi
+
+  DOCKERFILE_DIR="${GITHUB_WORKSPACE:-$PWD}/docker"
+  DOCKERFILE="${DOCKERFILE_DIR}/rocm.Dockerfile"
+
+  if [[ ! -f "${DOCKERFILE}" ]]; then
+    echo "Error: Dockerfile not found at ${DOCKERFILE}" >&2
+    exit 1
+  fi
+
+  IMAGE="sglang-ci:${GPU_ARCH_BUILD}-$(date +%Y%m%d)"
+  echo "Building Docker image from ${DOCKERFILE} with GPU_ARCH=${GPU_ARCH_BUILD}..."
+
+  # Pass full GPU_ARCH (e.g., gfx950-rocm720) - Dockerfile handles stripping suffix
+  docker build \
+    --build-arg GPU_ARCH="${GPU_ARCH_BUILD}" \
+    --build-arg SGL_BRANCH="main" \
+    -t "${IMAGE}" \
+    -f "${DOCKERFILE}" \
+    "${DOCKERFILE_DIR}"
+  echo "Successfully built image: ${IMAGE}"
+else
+  # Find the latest pre-built image
+  IMAGE=$(find_latest_image "${GPU_ARCH}")
+  echo "Pulling Docker image: ${IMAGE}"
+  docker pull "${IMAGE}"
+fi
+
+CACHE_HOST=/home/runner/sgl-data
+if [[ -d "$CACHE_HOST" ]]; then
+    CACHE_VOLUME="-v $CACHE_HOST:/sgl-data"
+else
+    CACHE_VOLUME=""
+fi
+
+echo "Launching container: ci_sglang"
+docker run -dt --user root --device=/dev/kfd ${DEVICE_FLAG} \
+  --ulimit nofile=65536:65536 \
+  -v "${GITHUB_WORKSPACE:-$PWD}:/sglang-checkout" \
+  $CACHE_VOLUME \
+  --group-add video \
+  --shm-size 32g \
+  --cap-add=SYS_PTRACE \
+  -e HF_TOKEN="${HF_TOKEN:-}" \
+  -e HF_HOME=/sgl-data/hf-cache \
+  -e HF_HUB_ETAG_TIMEOUT=300 \
+  -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
+  -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
+  -e MIOPEN_CUSTOM_CACHE_DIR=/sgl-data/miopen-cache \
+  -e PYTHONPATH="/opt/tilelang:${PYTHONPATH:-}" \
+  --security-opt seccomp=unconfined \
+  -w /sglang-checkout \
+  --name ci_sglang \
+  "${IMAGE}"
+
+# The checkout is owned by the runner (non-root) but the container runs as
+# root.  Git >= 2.35.2 rejects cross-user repos; mark the mount as safe so
+# setuptools-scm / vcs_versioning can resolve the package version.
+docker exec ci_sglang git config --global --add safe.directory /sglang-checkout
--- a/third_party/sglang/scripts/ci/amd/amd_ci_start_container_disagg.sh
+++ b/third_party/sglang/scripts/ci/amd/amd_ci_start_container_disagg.sh
@@ -0,0 +1,270 @@
+#!/bin/bash
+set -euo pipefail
+
+# Get version from git tags
+SGLANG_VERSION="v0.5.5"   # Default version, will be overridden if git tags are found
+
+# Fetch tags from origin to ensure we have the latest
+if git fetch --tags origin; then
+  # Get the latest version tag sorted by version number (e.g., v0.5.7)
+  VERSION_FROM_TAG=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1)
+  if [ -n "$VERSION_FROM_TAG" ]; then
+    SGLANG_VERSION="$VERSION_FROM_TAG"
+    echo "Using SGLang version from git tags: $SGLANG_VERSION"
+  else
+    echo "Warning: No version tags found; using default $SGLANG_VERSION" >&2
+  fi
+else
+  echo "Warning: Failed to fetch tags from origin; using default $SGLANG_VERSION" >&2
+fi
+
+
+# Default base tags (can be overridden by command line arguments)
+ROCM_VERSION="rocm700"
+DEFAULT_MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
+DEFAULT_MI35X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi35x"
+
+# Parse command line arguments
+MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}"
+MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}"
+
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --mi30x-base-tag) MI30X_BASE_TAG="$2"; shift 2;;
+    --mi35x-base-tag) MI35X_BASE_TAG="$2"; shift 2;;
+    --rocm-version)
+      ROCM_VERSION="$2"
+      MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
+      MI35X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi35x"
+      echo "Using ROCm version override: ${ROCM_VERSION}"
+      shift 2;;
+    -h|--help)
+      echo "Usage: $0 [--mi30x-base-tag TAG] [--mi35x-base-tag TAG] [--rocm-version VERSION]"
+      exit 0
+      ;;
+    *) echo "Unknown option $1"; exit 1;;
+  esac
+done
+
+
+
+# Detect GPU architecture from the Kubernetes runner hostname
+HOSTNAME_VALUE=$(hostname)
+GPU_ARCH="mi30x"   # default
+
+# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
+if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
+  GPU_ARCH="${BASH_REMATCH[1]}"
+  echo "Detected GPU architecture from hostname: ${GPU_ARCH}"
+else
+  echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}"
+fi
+
+# Normalise / collapse architectures we don’t yet build specifically for
+case "${GPU_ARCH}" in
+  mi35x)
+    echo "Runner uses ${GPU_ARCH}; will fetch mi35x image."
+    ;;
+  mi30x|mi300|mi325)
+    echo "Runner uses ${GPU_ARCH}; will fetch mi30x image."
+    GPU_ARCH="mi30x"
+    ;;
+  *)
+    echo "Runner architecture '${GPU_ARCH}' unrecognised; defaulting to mi30x image." >&2
+    GPU_ARCH="mi30x"
+    ;;
+esac
+
+
+# Set up DEVICE_FLAG based on Kubernetes pod info
+if [[ -f /etc/podinfo/gha-render-devices ]]; then
+  DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
+else
+  DEVICE_FLAG="--device /dev/dri"
+fi
+
+
+# Find the latest image
+find_latest_image() {
+  local gpu_arch=$1
+  local base_tag days_back image_tag
+
+  case "${gpu_arch}" in
+      mi30x) base_tag="${MI30X_BASE_TAG}" ;;
+      mi35x) base_tag="${MI35X_BASE_TAG}" ;;
+      *)     echo "Error: unsupported GPU architecture '${gpu_arch}'" >&2; return 1 ;;
+  esac
+
+  # First, check local cache
+  for days_back in {0..6}; do
+    image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
+    local local_image="rocm/sgl-dev:${image_tag}"
+    image_id=$(docker images -q "${local_image}")
+    if [[ -n "$image_id" ]]; then
+        echo "Found cached image locally: ${local_image}" >&2
+        echo "${local_image}"
+        return 0
+    fi
+  done
+
+  # If not found locally, fall back to pulling from public registry
+  for days_back in {0..6}; do
+    image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)"
+    echo "Checking for image: rocm/sgl-dev:${image_tag}" >&2
+    if docker manifest inspect "rocm/sgl-dev:${image_tag}" >/dev/null 2>&1; then
+      echo "Found available image: rocm/sgl-dev:${image_tag}" >&2
+      echo "rocm/sgl-dev:${image_tag}"
+      return 0
+    fi
+  done
+
+  # If still not found, try finding any image matching ROCm+arch from remote registry
+  echo "Exact version not found. Searching remote registry for any ${ROCM_VERSION}-${gpu_arch} image…" >&2
+  for days_back in {0..6}; do
+    local target_date=$(date -d "${days_back} days ago" +%Y%m%d)
+    local remote_tags=$(curl -s "https://registry.hub.docker.com/v2/repositories/rocm/sgl-dev/tags?page_size=100&name=${ROCM_VERSION}-${gpu_arch}-${target_date}" 2>/dev/null | grep -o '"name":"[^"]*"' | cut -d'"' -f4 | head -n 1)
+    if [[ -n "$remote_tags" ]]; then
+      echo "Found available image: rocm/sgl-dev:${remote_tags}" >&2
+      echo "rocm/sgl-dev:${remote_tags}"
+      return 0
+    fi
+  done
+
+  echo "No recent images found. Searching any cached local images matching ROCm+arch…" >&2
+  local any_local
+  any_local=$(docker images --format '{{.Repository}}:{{.Tag}}' --filter "reference=rocm/sgl-dev:*${ROCM_VERSION}*${gpu_arch}*" | sort -r | head -n 1)
+  if [[ -n "$any_local" ]]; then
+      echo "Using cached fallback image: ${any_local}" >&2
+      echo "${any_local}"
+      return 0
+  fi
+
+  echo "Error: no ${gpu_arch} image found in the last 7 days for base ${base_tag}" >&2
+  echo "Using hard-coded fallback for ${ROCM_VERSION}…" >&2
+  case "${ROCM_VERSION}" in
+    rocm720)
+      if [[ "${gpu_arch}" == "mi35x" ]]; then
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm720-mi35x-20260211-preview"
+      else
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm720-mi30x-20260211-preview"
+      fi
+      ;;
+    rocm700)
+      if [[ "${gpu_arch}" == "mi35x" ]]; then
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm700-mi35x-20260211"
+      else
+        echo "rocm/sgl-dev:v0.5.8.post1-rocm700-mi30x-20260211"
+      fi
+      ;;
+    *)
+      echo "Error: no hard-coded fallback available for ${ROCM_VERSION}" >&2
+      return 1
+      ;;
+  esac
+}
+
+# Pull and run the latest image
+IMAGE=$(find_latest_image "${GPU_ARCH}")
+echo "Pulling Docker image: ${IMAGE}"
+docker pull "${IMAGE}"
+
+CACHE_HOST=/home/runner/sgl-data
+if [[ -d "$CACHE_HOST" ]]; then
+    CACHE_VOLUME="-v $CACHE_HOST:/sgl-data"
+else
+    CACHE_VOLUME=""
+fi
+
+# Detect libionic library for RDMA support
+LIBIONIC_MOUNT=""
+IONIC_SYMLINK="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
+if [[ -L "$IONIC_SYMLINK" ]]; then
+    LIBIONIC_LIB=$(readlink -f "$IONIC_SYMLINK" 2>/dev/null)
+    if [[ -f "$LIBIONIC_LIB" ]]; then
+        echo "Found libionic library: $LIBIONIC_LIB (resolved from symlink)"
+        LIBIONIC_MOUNT="-v ${LIBIONIC_LIB}:${LIBIONIC_LIB}:ro"
+    else
+        echo "Warning: libionic symlink exists but target does not: $LIBIONIC_LIB"
+    fi
+else
+    # Fallback: try to find directly
+    LIBIONIC_FOUND=$(find /usr/lib/x86_64-linux-gnu -maxdepth 1 -name "libionic.so.*" 2>/dev/null | head -1)
+    if [[ -n "$LIBIONIC_FOUND" ]]; then
+        LIBIONIC_LIB=$(readlink -f "$LIBIONIC_FOUND" 2>/dev/null)
+        if [[ -f "$LIBIONIC_LIB" ]]; then
+            echo "Found libionic library: $LIBIONIC_LIB"
+            LIBIONIC_MOUNT="-v ${LIBIONIC_LIB}:${LIBIONIC_LIB}:ro"
+        else
+            echo "Warning: libionic found but cannot resolve real path: $LIBIONIC_FOUND"
+        fi
+    else
+        echo "Warning: libionic library not found on host, RDMA may not work"
+    fi
+fi
+
+MOUNT_ARGS=""
+
+add_mount_if_exists() {
+    local name=$1
+    local search_pattern=$2
+    local path=$(find /lib/x86_64-linux-gnu /usr/lib/x86_64-linux-gnu /lib64 /usr/lib64 -name "$search_pattern" -print -quit 2>/dev/null)
+
+    if [ -n "$path" ]; then
+        echo "Found $name at: $path"
+        MOUNT_ARGS="$MOUNT_ARGS -v $path:$path:ro"
+    else
+        echo "WARNING: Could not find $name on host! (Pattern: $search_pattern)"
+    fi
+}
+
+IONIC_LINK="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
+if [ -L "$IONIC_LINK" ]; then
+    IONIC_REAL=$(readlink -f "$IONIC_LINK")
+    if [ -f "$IONIC_REAL" ]; then
+        echo "Ionic Driver: $IONIC_REAL"
+        MOUNT_ARGS="$MOUNT_ARGS -v $IONIC_REAL:$IONIC_REAL:ro"
+    fi
+fi
+
+add_mount_if_exists "libnl-3" "libnl-3.so*"
+add_mount_if_exists "libmnl" "libmnl.so*"
+
+echo "Mount args: $MOUNT_ARGS"
+
+echo "Launching container: ci_sglang"
+docker run -dt --user root \
+  --device=/dev/kfd \
+  --device=/dev/dri \
+  ${DEVICE_FLAG} \
+  -v "${GITHUB_WORKSPACE:-$PWD}:/sglang-checkout" \
+  -v /sys/class/infiniband:/sys/class/infiniband:ro \
+  -v /sys/class/infiniband_verbs:/sys/class/infiniband_verbs:ro \
+  -v /sys/class/net:/sys/class/net:ro \
+  -v /etc/libibverbs.d:/etc/libibverbs.d:ro \
+  -v /usr/lib/x86_64-linux-gnu/libibverbs:/usr/lib/x86_64-linux-gnu/libibverbs:ro \
+  $MOUNT_ARGS \
+  $CACHE_VOLUME \
+  --privileged \
+  --network=host \
+  --ipc=host \
+  --ulimit memlock=-1 \
+  --cap-add=IPC_LOCK \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --group-add video \
+  --group-add rdma \
+  --shm-size 32g \
+  -e HF_TOKEN="${HF_TOKEN:-}" \
+  -e HF_HOME=/sgl-data/hf-cache \
+  -e HF_HUB_ETAG_TIMEOUT=300 \
+  -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
+  -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \
+  -e MIOPEN_CUSTOM_CACHE_DIR=/sgl-data/miopen-cache \
+  -w /sglang-checkout \
+  --name ci_sglang \
+  "${IMAGE}"
+
+# The checkout is owned by the runner (non-root) but the container runs as
+# root.  Git >= 2.35.2 rejects cross-user repos; mark the mount as safe so
+# setuptools-scm / vcs_versioning can resolve the package version.
+docker exec ci_sglang git config --global --add safe.directory /sglang-checkout
--- a/third_party/sglang/scripts/ci/amd/amd_ci_warmup_aiter.py
+++ b/third_party/sglang/scripts/ci/amd/amd_ci_warmup_aiter.py
@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""
+Warmup script to pre-build AITER JIT kernels.
+
+This script triggers compilation of commonly used AITER kernels by importing
+the relevant modules and calling functions with sample data. This avoids
+timeouts during actual tests when kernels need to be compiled on first use.
+
+Run this after clearing pre-built AITER kernels from the Docker image.
+"""
+
+import os
+import sys
+import time
+
+# Ensure AITER is enabled
+os.environ["SGLANG_USE_AITER"] = "1"
+
+
+def warmup_aiter_kernels():
+    """Trigger AITER JIT kernel compilation."""
+    import torch
+
+    if not torch.cuda.is_available():
+        print("CUDA/ROCm not available, skipping AITER warmup")
+        return
+
+    print("=" * 60)
+    print("AITER JIT Kernel Warmup")
+    print("=" * 60)
+
+    device = torch.device("cuda:0")
+    start_time = time.time()
+
+    # Warmup module_rmsnorm_quant (small module, ~2MB)
+    # Triggered by rmsnorm2d_fwd when hidden_size <= 8192
+    try:
+        print(
+            "\n[1/5] Warming up module_rmsnorm_quant (rmsnorm2d_fwd, hidden<=8192)..."
+        )
+        from aiter import rmsnorm2d_fwd
+
+        hidden_size = 4096
+        batch_size = 512  # Use larger batch to match CUDA graph capture
+        x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device=device)
+        weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device)
+        eps = 1e-6
+
+        # hidden_size=4096 <= 8192 -> takes rmsnorm() path -> compiles module_rmsnorm_quant
+        _ = rmsnorm2d_fwd(x, weight, eps)
+        torch.cuda.synchronize()
+        print("   module_rmsnorm_quant compiled successfully")
+    except Exception as e:
+        print(f"   module_rmsnorm_quant warmup failed: {e}")
+
+    # Warmup module_rmsnorm (large CK module, ~159MB)
+    # Triggered by rmsnorm2d_fwd_with_add (always uses CK path)
+    # NOTE: rmsnorm2d_fwd_with_add signature is:
+    #   rmsnorm2d_fwd_with_add(out, input, residual_in, residual_out, weight, epsilon)
+    try:
+        print("\n[2/5] Warming up module_rmsnorm (rmsnorm2d_fwd_with_add, CK path)...")
+        from aiter import rmsnorm2d_fwd_with_add
+
+        hidden_size = 4096
+        batch_size = 512
+        x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device=device)
+        residual_in = torch.randn(
+            batch_size, hidden_size, dtype=torch.bfloat16, device=device
+        )
+        output = torch.empty_like(x)
+        residual_out = torch.empty_like(x)
+        weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device)
+        eps = 1e-6
+
+        # This triggers JIT compilation of module_rmsnorm (CK kernels)
+        rmsnorm2d_fwd_with_add(output, x, residual_in, residual_out, weight, eps)
+        torch.cuda.synchronize()
+        print("   module_rmsnorm compiled successfully")
+    except Exception as e:
+        print(f"   module_rmsnorm warmup failed: {e}")
+
+    # Warmup module_rmsnorm via rmsnorm2d_fwd with large hidden_size (CK path)
+    # When hidden_size > 8192, rmsnorm2d_fwd takes the rmsnorm2d_fwd_ck path
+    # which also uses module_rmsnorm (already compiled in step 2, but this
+    # ensures the CK rmsnorm2d_fwd path is exercised as well)
+    try:
+        print("\n[3/5] Warming up rmsnorm2d_fwd CK path (hidden>8192)...")
+        from aiter import rmsnorm2d_fwd
+
+        hidden_size = 16384  # > 8192 to trigger rmsnorm2d_fwd_ck (module_rmsnorm)
+        batch_size = 32
+        x = torch.randn(batch_size, hidden_size, dtype=torch.bfloat16, device=device)
+        weight = torch.ones(hidden_size, dtype=torch.bfloat16, device=device)
+        eps = 1e-6
+
+        _ = rmsnorm2d_fwd(x, weight, eps)
+        torch.cuda.synchronize()
+        print("   rmsnorm2d_fwd CK path compiled successfully")
+    except Exception as e:
+        print(f"   rmsnorm2d_fwd CK path warmup skipped: {e}")
+
+    # Warmup rotary embedding kernel if available
+    try:
+        print("\n[4/5] Warming up rotary embedding kernel...")
+        from aiter import rotary_embedding
+
+        head_size = 128
+        seq_len = 32
+        num_heads = 32
+        positions = torch.arange(seq_len, device=device)
+        query = torch.randn(
+            seq_len, num_heads, head_size, dtype=torch.bfloat16, device=device
+        )
+        key = torch.randn(
+            seq_len, num_heads, head_size, dtype=torch.bfloat16, device=device
+        )
+        cos = torch.ones(seq_len, head_size // 2, dtype=torch.bfloat16, device=device)
+        sin = torch.zeros(seq_len, head_size // 2, dtype=torch.bfloat16, device=device)
+
+        _ = rotary_embedding(positions, query, key, head_size, cos, sin, True)
+        torch.cuda.synchronize()
+        print("   Rotary embedding kernel compiled successfully")
+    except Exception as e:
+        print(f"   Rotary embedding warmup skipped (may not be available): {e}")
+
+    # Warmup activation kernels if available
+    try:
+        print("\n[5/5] Warming up activation kernels...")
+        from aiter import silu_and_mul
+
+        hidden_size = 4096
+        batch_size = 512
+        x = torch.randn(
+            batch_size, hidden_size * 2, dtype=torch.bfloat16, device=device
+        )
+        out = torch.empty(batch_size, hidden_size, dtype=torch.bfloat16, device=device)
+
+        silu_and_mul(out, x)
+        torch.cuda.synchronize()
+        print("   Activation kernel compiled successfully")
+    except Exception as e:
+        print(f"   Activation warmup skipped (may not be available): {e}")
+
+    elapsed = time.time() - start_time
+    print("\n" + "=" * 60)
+    print(f"AITER warmup completed in {elapsed:.1f}s")
+    print("=" * 60 + "\n")
+
+
+if __name__ == "__main__":
+    warmup_aiter_kernels()
--- a/third_party/sglang/scripts/ci/amd/check_vram_clear.sh
+++ b/third_party/sglang/scripts/ci/amd/check_vram_clear.sh
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+check_vram_clear() {
+    local vram_threshold_percent=5  # Allow up to 5% VRAM usage
+    local memory_threshold_mb=500   # Allow up to 500MB memory usage
+
+    if command -v rocm-smi >/dev/null 2>&1; then
+        echo "Checking ROCm GPU VRAM usage..."
+        # Check if any GPU has more than threshold VRAM allocated
+        local high_usage=$(rocm-smi --showmemuse | grep -E "GPU Memory Allocated \(VRAM%\): ([6-9]|[1-9][0-9]|100)")
+        if [ -n "$high_usage" ]; then
+            echo "ERROR: VRAM usage exceeds threshold (${vram_threshold_percent}%) on some GPUs:"
+            echo "$high_usage"
+            rocm-smi --showmemuse
+            return 1
+        else
+            echo "✓ VRAM usage is within acceptable limits on all GPUs"
+            return 0
+        fi
+   fi
+}
+
+# If this script is run directly (not sourced), run the check
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
+    set -e
+    check_vram_clear
+fi
--- a/third_party/sglang/scripts/ci/amd/ensure_vram_clear.sh
+++ b/third_party/sglang/scripts/ci/amd/ensure_vram_clear.sh
@@ -0,0 +1,103 @@
+#!/bin/bash
+
+# Source the VRAM checking function
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "$SCRIPT_DIR/check_vram_clear.sh"
+
+ensure_vram_clear() {
+    local max_retries=3
+    local retry_count=0
+
+    # Stop and remove any existing ci_sglang container
+    echo "Stopping any existing ci_sglang container..."
+    docker stop ci_sglang || true
+    docker rm ci_sglang || true
+
+    # Log host information for debugging
+    echo "=== Host Information ==="
+    echo "Hostname: $(hostname)"
+    echo "Host IP: $(hostname -I 2>/dev/null || echo 'N/A')"
+    echo "Date: $(date)"
+    echo "Mode: rocm"
+    echo "========================"
+    echo "Running in ROCm mode"
+
+    # Show initial GPU status
+    echo "=== Initial GPU Memory Status ==="
+    rocm-smi --showmemuse
+    echo "=================================="
+
+    while [ $retry_count -lt $max_retries ]; do
+        echo "=== Cleanup Attempt $((retry_count + 1))/$max_retries ==="
+
+        # Clean SGLang processes
+        echo "Killing SGLang processes..."
+        pgrep -f 'sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt' | xargs -r kill -9 || true
+
+        if [ $retry_count -gt 0 ]; then
+            echo "Performing aggressive cleanup..."
+            # Kill all processes using KFD
+            rocm-smi --showpids 2>/dev/null | grep 'PID:' | awk '{print $2}' | xargs -r kill -9 2>/dev/null || true
+            # Wait a bit for cleanup to take effect
+            echo "Waiting 30 seconds for VRAM to clear..."
+            sleep 30
+        fi
+
+        # Check VRAM
+        echo "Checking VRAM status..."
+        if check_vram_clear; then
+            echo "✓ VRAM cleanup successful after $((retry_count + 1)) attempts"
+            return 0
+        else
+            echo "✗ VRAM still not clear after attempt $((retry_count + 1))"
+            retry_count=$((retry_count + 1))
+        fi
+    done
+
+    # Failed after all retries
+    echo "=== FAILED: VRAM cleanup unsuccessful after $max_retries attempts ==="
+    echo "Final GPU status:"
+    timeout 30 rocm-smi --showmemuse || echo "rocm-smi timed out"
+    echo "Processes using GPU:"
+    rocm-smi --showpids 2>/dev/null | grep -q 'PID:' || echo "No processes found using /dev/kfd"
+
+    # Print detailed information about suspicious processes
+    echo "=== Detailed Process Information ==="
+    if command -v rocm-smi >/dev/null 2>&1; then
+        # For AMD GPUs, get processes from rocm-smi --showpids
+        kfd_pids=$(rocm-smi --showpids 2>/dev/null | grep 'PID:' | awk '{print $2}' | sort -u)
+        if [ -n "$kfd_pids" ]; then
+            echo "Processes accessing /dev/kfd (AMD GPU device):"
+            for pid in $kfd_pids; do
+                if ps -p $pid -o pid,ppid,cmd --no-headers 2>/dev/null; then
+                    echo "  └─ Command line: $(ps -p $pid -o cmd --no-headers 2>/dev/null | head -1)"
+                else
+                    echo "  └─ PID $pid: Process not found or already terminated"
+                fi
+            done
+        else
+            echo "No processes found accessing /dev/kfd"
+        fi
+    fi
+
+    # Check for any remaining sglang-related processes
+    echo "Checking for any remaining sglang-related processes:"
+    sglang_procs=$(pgrep -f 'sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt' 2>/dev/null)
+    if [ -n "$sglang_procs" ]; then
+        echo "Found sglang processes still running:"
+        for pid in $sglang_procs; do
+            ps -p $pid -o pid,ppid,cmd --no-headers 2>/dev/null || echo "PID $pid not found"
+        done
+    else
+        echo "No sglang-related processes found."
+    fi
+
+    echo "=================================================================="
+    return 1
+}
+
+# If this script is run directly (not sourced), run the ensure function
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
+    set -e
+    ensure_vram_clear "$@"
+fi
--- a/third_party/sglang/scripts/ci/amd/test_rccl_multi_gpu.py
+++ b/third_party/sglang/scripts/ci/amd/test_rccl_multi_gpu.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""
+Simple RCCL test for multi-GPU communication.
+This test verifies that RCCL can initialize and communicate across multiple GPUs.
+"""
+
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+
+def test_rccl_allreduce():
+    """Test basic RCCL allreduce operation across all GPUs."""
+    if not torch.cuda.is_available():
+        print("CUDA not available, skipping test")
+        sys.exit(1)
+
+    # Initialize process group with NCCL (RCCL on AMD)
+    dist.init_process_group(backend="nccl")
+
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+
+    print(f"[Rank {rank}/{world_size}] Initialized successfully")
+
+    # Set device
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    print(f"[Rank {rank}] Device: {torch.cuda.get_device_name(device)}")
+    print(
+        f"[Rank {rank}] Device memory: {torch.cuda.get_device_properties(device).total_memory / 1e9:.2f} GB"
+    )
+
+    # Create a tensor and perform allreduce
+    tensor = torch.ones(1000, device=device) * rank
+    print(f"[Rank {rank}] Before allreduce: tensor sum = {tensor.sum().item()}")
+
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+    expected_sum = sum(range(world_size)) * 1000
+    actual_sum = tensor.sum().item()
+
+    print(
+        f"[Rank {rank}] After allreduce: tensor sum = {actual_sum}, expected = {expected_sum}"
+    )
+
+    if abs(actual_sum - expected_sum) < 0.1:
+        print(f"[Rank {rank}] ✓ RCCL allreduce test PASSED")
+        dist.destroy_process_group()
+        sys.exit(0)
+    else:
+        print(f"[Rank {rank}] ✗ RCCL allreduce test FAILED")
+        dist.destroy_process_group()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    test_rccl_allreduce()