diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index fd33933e4c38..5d291d1b9b84 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -978,14 +978,22 @@ steps: - step: build-cpu-release-image-arm64 allow_failure: true - - label: "Publish release images to DockerHub" + - label: "Publish {{matrix}} release images to DockerHub" depends_on: - block-publish-release-images key: publish-release-images-dockerhub agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/publish-release-images.sh" + - "bash .buildkite/scripts/publish-release-images.sh {{matrix}}" + matrix: + - "cuda-13-0" + - "cuda-12-9" + - "cuda-13-0-ubuntu-24-04" + - "cuda-12-9-ubuntu-24-04" + - "rocm" + - "xpu" + - "cpu" plugins: - docker-login#v3.0.0: username: vllmbot diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index 91b5c3ace1be..dd54b1717e10 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -2,12 +2,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # -# Publish release Docker images from ECR to DockerHub. +# Publish a release Docker image family from ECR to DockerHub. # Pulls per-arch images, tags with latest and versioned tags, pushes them, # then creates and pushes multi-arch manifests. set -euo pipefail +TARGET="${1:-all}" +case "${TARGET}" in + cuda-13-0 | cuda-12-9 | cuda-13-0-ubuntu-24-04 | \ + cuda-12-9-ubuntu-24-04 | rocm | xpu | cpu | all) ;; + *) + echo "Usage: $0 {cuda-13-0|cuda-12-9|cuda-13-0-ubuntu-24-04|cuda-12-9-ubuntu-24-04|rocm|xpu|cpu|all}" + exit 2 + ;; +esac + +target_enabled() { + [ "${TARGET}" = "all" ] || [ "${TARGET}" = "$1" ] +} + RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//') if [ -z "${RELEASE_VERSION}" ]; then echo "ERROR: release-version metadata not set" @@ -15,12 +29,10 @@ if [ -z "${RELEASE_VERSION}" ]; then fi COMMIT="$BUILDKITE_COMMIT" -ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) echo "========================================" -echo "Publishing release images v${RELEASE_VERSION}" +echo "Publishing ${TARGET} release images v${RELEASE_VERSION}" echo " Commit: ${COMMIT}" -echo " ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" echo "========================================" # Login to ECR to pull staging images @@ -29,122 +41,137 @@ aws ecr-public get-login-password --region us-east-1 | \ # ---- CUDA (default: 13.0) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 -docker push vllm/vllm-openai:latest-x86_64 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker push vllm/vllm-openai:latest-aarch64 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 - -docker manifest rm vllm/vllm-openai:latest || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true -docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker manifest push vllm/vllm-openai:latest -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} +if target_enabled cuda-13-0; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai:latest-x86_64 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker push vllm/vllm-openai:latest-aarch64 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + + docker manifest rm vllm/vllm-openai:latest || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker manifest push vllm/vllm-openai:latest + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} +fi # ---- CUDA 12.9 ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 -docker push vllm/vllm-openai:latest-x86_64-cu129 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 -docker push vllm/vllm-openai:latest-aarch64-cu129 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 - -docker manifest rm vllm/vllm-openai:latest-cu129 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true -docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 -docker manifest push vllm/vllm-openai:latest-cu129 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 +if target_enabled cuda-12-9; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + docker push vllm/vllm-openai:latest-x86_64-cu129 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker push vllm/vllm-openai:latest-aarch64-cu129 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + + docker manifest rm vllm/vllm-openai:latest-cu129 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true + docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker manifest push vllm/vllm-openai:latest-cu129 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 +fi # ---- Ubuntu 24.04 (CUDA 13.0) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 -docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 -docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 - -docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true -docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 -docker manifest push vllm/vllm-openai:latest-ubuntu2404 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 +if target_enabled cuda-13-0-ubuntu-24-04; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + + docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true + docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker manifest push vllm/vllm-openai:latest-ubuntu2404 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 +fi # ---- Ubuntu 24.04 (CUDA 12.9) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 -docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 -docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 - -docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true -docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 -docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 +if target_enabled cuda-12-9-ubuntu-24-04; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + + docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true + docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 +fi # ---- ROCm ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base +if target_enabled rocm; then + ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + echo "ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" + + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker push vllm/vllm-openai-rocm:latest -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker push vllm/vllm-openai-rocm:latest + docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base -docker push vllm/vllm-openai-rocm:latest-base -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker push vllm/vllm-openai-rocm:latest-base + docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base +fi # ---- XPU ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu +if target_enabled xpu; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 -docker push vllm/vllm-openai-xpu:latest-x86_64 -docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai-xpu:latest-x86_64 + docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 -docker manifest rm vllm/vllm-openai-xpu:latest || true -docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true -docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend -docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend -docker manifest push vllm/vllm-openai-xpu:latest -docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} + docker manifest rm vllm/vllm-openai-xpu:latest || true + docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend + docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend + docker manifest push vllm/vllm-openai-xpu:latest + docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} +fi # ---- CPU ---- # CPU images are behind separate block steps and may not have been built. @@ -153,44 +180,46 @@ docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} # arch would leave `:latest-x86_64` pointing at the new release while the # `:latest` multi-arch manifest still resolves to the previous release. -CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} -CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} - -CPU_X86_AVAILABLE=false -CPU_ARM_AVAILABLE=false -docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true -docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true - -if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then - docker pull "${CPU_X86_TAG}" - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 - docker push vllm/vllm-openai-cpu:latest-x86_64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 - - docker pull "${CPU_ARM_TAG}" - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - docker push vllm/vllm-openai-cpu:latest-arm64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - - docker manifest rm vllm/vllm-openai-cpu:latest || true - docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true - docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 - docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - docker manifest push vllm/vllm-openai-cpu:latest - docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} -elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then - echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" -else - # Partial state: one arch built, the other did not. Fail loudly rather than - # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch) - # disagree on which release they point at. - echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})." - echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags." - echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended." - exit 1 +if target_enabled cpu; then + CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} + CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} + + CPU_X86_AVAILABLE=false + CPU_ARM_AVAILABLE=false + docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true + docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true + + if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then + docker pull "${CPU_X86_TAG}" + docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 + docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai-cpu:latest-x86_64 + docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + + docker pull "${CPU_ARM_TAG}" + docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 + docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker push vllm/vllm-openai-cpu:latest-arm64 + docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + + docker manifest rm vllm/vllm-openai-cpu:latest || true + docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 + docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker manifest push vllm/vllm-openai-cpu:latest + docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} + elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then + echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" + else + # Partial state: one arch built, the other did not. Fail loudly rather than + # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch) + # disagree on which release they point at. + echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})." + echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags." + echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended." + exit 1 + fi fi echo "" -echo "Successfully published release images for v${RELEASE_VERSION}" +echo "Successfully published ${TARGET} release images for v${RELEASE_VERSION}" diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 15a7d0cde6bf..be0eccd45aed 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -30,7 +30,7 @@ else() # Keep in sync with tools/install_deepgemm.sh set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") # TODO: switch to nv_dev branch after it support situ - set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9") + set(_DEEPGEMM_UPSTREAM_TAG "e21c821f39a2056d68067a466c64ddc942200106") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/csrc/cpu/sgl-kernels/fla.cpp b/csrc/cpu/sgl-kernels/fla.cpp index 26af209a4179..fec6648cec66 100644 --- a/csrc/cpu/sgl-kernels/fla.cpp +++ b/csrc/cpu/sgl-kernels/fla.cpp @@ -101,7 +101,7 @@ struct l2norm_kernel { using bVec = at::vec::Vectorized; using fVec = at::vec::Vectorized; constexpr int bVecSize = bVec::size(); - constexpr float scale = 1.f / std::sqrt(static_cast(D)); + const float scale = 1.f / std::sqrt(static_cast(D)); fVec sum_fvec0(0.f), sum_fvec1(0.f); int d = 0; @@ -155,7 +155,7 @@ struct l2norm_kernel { __m512bh va[COLS]; __m512 vrscale; - constexpr float scale = 1.f / std::sqrt(D); + const float scale = 1.f / std::sqrt(D); __m512 vscale = _mm512_set1_ps(scale); // step 1: load input and do reduce with avx512-bf16 diff --git a/docker/Dockerfile b/docker/Dockerfile index 96ed0bddf1bc..b3fd7c097ded 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -811,7 +811,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.15.post1 +ARG FLASHINFER_VERSION=0.6.16.post3 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') @@ -949,7 +949,8 @@ COPY requirements/test/cuda.txt requirements/test/cuda.txt COPY requirements/dev.txt requirements/dev.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt -RUN --mount=type=cache,target=/opt/uv/cache \ +# Source-built wheels in this cache must match the Ubuntu 22.04 ABI. +RUN --mount=type=cache,id=uv-v027-ubuntu2204,target=/opt/uv/cache \ CUDA_MAJOR="${CUDA_VERSION%%.*}"; \ if [ "$CUDA_MAJOR" -ge 12 ]; then \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ @@ -971,6 +972,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ && uv pip install --system -r requirements/dev.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi \ + && if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \ + python3 -c "from arctic_inference.suffix_decoding import _C"; \ + fi; \ fi # install development dependencies (for testing) diff --git a/docker/versions.json b/docker/versions.json index 783d80aace43..009fa9c31cae 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -71,7 +71,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.15.post1" + "default": "0.6.16.post3" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 42b05e3c306e..59271c481283 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -32,7 +32,7 @@ th { | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass | | ------- | ------------------ | ------------ | ------------- | ----- | --------------------- | --------- | -| naive | standard | all1 | G,A,T | N | 6 | [layer.py][vllm.model_executor.layers.fused_moe.runner.MoERunner] | +| naive | standard | all1 | G,A,T | N | 6 | [`MoERunner`][vllm.model_executor.layers.fused_moe.runner.moe_runner.MoERunner] | | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e434c57012df..de937ec520cd 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -364,7 +364,7 @@ th { | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ | | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | ✅︎ | ✅︎ | | `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ | -| `ExaoneMoEForCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B`, etc. | | | +| `ExaoneMoeForCausalLM` | K-EXAONE, K-EXAONE-2 | `LGAI-EXAONE/K-EXAONE-236B-A23B`, `LGAI-EXAONE/K-EXAONE-2.0-750B-A37B`, etc. | | | | `Exaone4ForCausalLM` | EXAONE-4 | `LGAI-EXAONE/EXAONE-4.0-32B`, etc. | ✅︎ | ✅︎ | | `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ | | `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ | diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 2a8b2734f8f1..af1f5f61ca6a 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4 # flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from # install_requires so the published wheel does not carry an unresolvable pin --extra-index-url https://flashinfer.ai/whl/ -flashinfer-python==0.6.15.post1 -flashinfer-cubin==0.6.15.post1 +flashinfer-python==0.6.16.post3 +flashinfer-cubin==0.6.16.post3 apache-tvm-ffi==0.1.11 tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 @@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) nvidia-cutlass-dsl[cu13]==4.6.0 -quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA +quack-kernels==0.6.1 # Required for CUTLASS DSL 4.6 by MSA # Tokenspeed_MLA for faster mla with spec decode tokenspeed-mla==0.1.8; platform_system == "Linux" diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml index 14b4f2c0d1da..d4ddfcc99fc6 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml @@ -13,5 +13,6 @@ server_args: >- --enable-expert-parallel --kv-cache-dtype fp8 env: + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" VLLM_LOGGING_LEVEL: "DEBUG" VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 703a5df163e9..c37f74eb5079 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import Mock + import pytest import torch -from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( FlashInferSSUBackend, TritonSSUBackend, @@ -69,6 +71,48 @@ def test_flashinfer_backend_init(): assert backend.name == "flashinfer" +@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed") +@pytest.mark.parametrize( + ("algorithm", "expected"), + [ + (None, "auto"), + ("auto", "auto"), + ("simple", "simple"), + ("vertical", "vertical"), + ("horizontal", "horizontal"), + ], +) +def test_flashinfer_forwards_ssu_algorithm( + algorithm: MambaSSUAlgorithm | None, + expected: MambaSSUAlgorithm, + monkeypatch, +): + import flashinfer.mamba + + kernel = Mock() + monkeypatch.setattr(flashinfer.mamba, "selective_state_update", kernel) + backend = FlashInferSSUBackend( + MambaConfig( + backend=MambaBackendEnum.FLASHINFER, + ssu_algorithm=algorithm, + ) + ) + + tensor = torch.empty(1) + backend( + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + ) + + assert kernel.call_args.kwargs["algorithm"] == expected + + def test_uninitialized_backend_raises(): import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py index b61bb498cfef..9e8faea099cb 100644 --- a/tests/models/kimi_k3/test_sequence_parallel.py +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -10,6 +10,7 @@ from torch import nn from vllm.config import ParallelConfig +from vllm.model_executor.layers.activation import SiluAndMul from vllm.models.common.ops import sequence_parallel as sp_ops from vllm.models.kimi_k3.nvidia import model as kimi_model from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp @@ -265,6 +266,84 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch): torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states) +@pytest.mark.parametrize( + ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"), + [ + (True, True, True, 8, True), + (False, True, True, 8, False), # opt-in only + (True, False, True, 8, False), # replication only exists under SP + (True, True, False, 8, False), # FusedMoE path owns the reduction + (True, True, True, 1, False), # nothing to shard + (True, True, True, 5, False), # 6144 % 5 -- would fail divide() + ], +) +def test_shard_sequence_parallel_mlp_gating( + monkeypatch, + enabled: bool, + use_sequence_parallel: bool, + eligible: bool, + tp_size: int, + expected: bool, +): + monkeypatch.setattr(kimi_model.envs, "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", enabled) + monkeypatch.setattr( + kimi_model, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + + assert ( + kimi_model.shard_sequence_parallel_mlp( + hidden_size=7168, + intermediate_size=6144, + use_sequence_parallel=use_sequence_parallel, + eligible=eligible, + ) + is expected + ) + + +def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config): + """Sharded SP MLP must reproduce the replicated result for every token. + + Each rank owns a *disjoint* token shard, so a weight shard alone cannot + finish a rank's own tokens: the ranks must gather the full token set, + compute partial sums over their intermediate shard, and reduce-scatter. + Splicing per-rank feature slices together instead silently mixes different + tokens and produces plausible-looking garbage. + """ + tp_size, hidden, intermediate, tokens_per_rank = 4, 16, 12, 3 + torch.manual_seed(0) + num_tokens = tp_size * tokens_per_rank + x = torch.randn(num_tokens, hidden) + gate_weight = torch.randn(intermediate, hidden) + up_weight = torch.randn(intermediate, hidden) + down_weight = torch.randn(hidden, intermediate) + act_fn = SiluAndMul() + + replicated = act_fn(x @ torch.cat([gate_weight, up_weight]).T) @ down_weight.T + + shard = intermediate // tp_size + # Every rank all-gathers the full token set, then computes its partial. + partials = [ + act_fn( + x + @ torch.cat( + [ + gate_weight[r * shard : (r + 1) * shard], + up_weight[r * shard : (r + 1) * shard], + ] + ).T + ) + @ down_weight[:, r * shard : (r + 1) * shard].T + for r in range(tp_size) + ] + reduced = torch.stack(partials).sum(0) + # Reduce-scatter: rank r keeps only its own token shard. + for r in range(tp_size): + mine = reduced[r * tokens_per_rank : (r + 1) * tokens_per_rank] + expected = replicated[r * tokens_per_rank : (r + 1) * tokens_per_rank] + torch.testing.assert_close(mine, expected, atol=1e-5, rtol=1e-5) + + def test_sp_all_gather_uses_custom_kernel(monkeypatch): hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) expected = torch.cat([hidden_states, hidden_states]) diff --git a/tests/models/registry.py b/tests/models/registry.py index f8d0387dff2b..4ea312071593 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -267,7 +267,7 @@ def check_available_online( "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct", trust_remote_code=True ), "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"), - "ExaoneMoEForCausalLM": _HfExamplesInfo( + "ExaoneMoeForCausalLM": _HfExamplesInfo( "LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.1.0" ), "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"), diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 838ca4495638..eef3a91de5e8 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -496,7 +496,7 @@ class DummyKernel: monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes." "inc_mxfp4_linear.init_mxfp4_linear_kernel", - lambda: DummyKernel(), + lambda **kwargs: DummyKernel(), ) from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501 @@ -614,7 +614,7 @@ def process_weights_after_loading(self, layer) -> None: monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes." "inc_mxfp4_linear.init_mxfp4_linear_kernel", - lambda: DummyKernel(), + lambda **kwargs: DummyKernel(), ) monkeypatch.setattr( "vllm.model_executor.parameter.get_tensor_model_parallel_rank", diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index cac49a5a80fa..4124085c23e3 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, + ModelOptFp8LinearMethod, ModelOptMixedPrecisionConfig, ModelOptMxFp8Config, ModelOptNvFp4Config, @@ -114,6 +115,28 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head(): assert isinstance(method, ModelOptNvFp4LinearMethod) +def test_modelopt_fp8_updates_weight_dims_after_transpose(): + layer = torch.nn.Module() + layer.register_parameter( + "weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False) + ) + layer.register_parameter( + "weight_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False) + ) + layer.register_parameter( + "input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False) + ) + + method = ModelOptFp8LinearMethod.__new__(ModelOptFp8LinearMethod) + method.fp8_linear = Mock() + method.process_weights_after_loading(layer) + + assert layer.weight.shape == (2, 3) + assert layer.weight.input_dim == 0 + assert layer.weight.output_dim == 1 + method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer) + + def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized(): config = ModelOptNvFp4Config( is_checkpoint_nvfp4_serialized=True, diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py new file mode 100644 index 000000000000..31671caf766b --- /dev/null +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mamba "align" prefill chunk splitting (`_mamba_block_aligned_split`). + +Invariant: slot `p` holds the state after exactly `(p + 1) * block_size` tokens. +State is written at chunk ends, so a chunk ending mid-block leaves its slot at +the wrong offset, and a later chunk crossing that boundary publishes it anyway. +Requests resuming from it then restore a truncated state (#43559). +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.utils.math_utils import cdiv +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.request import Request + +from .utils import create_requests + +pytestmark = pytest.mark.cpu_test + +# Mirrors the deployment where the poisoning was observed (Qwen3.6-27B): mamba +# block 1600, MTP with 3 draft tokens, prompts shorter than 2 mamba blocks. +ATTN_BLOCK_SIZE = 16 +MAMBA_BLOCK_SIZE = 1600 +NUM_SPEC = 3 +PROMPT_LEN = 2002 +MAMBA_GROUP_ID = 1 + + +def _make_hybrid_kv_cache_manager() -> KVCacheManager: + config = KVCacheConfig( + num_blocks=10000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_layer"], + FullAttentionSpec( + block_size=ATTN_BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba_layer"], + MambaSpec( + block_size=MAMBA_BLOCK_SIZE, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=NUM_SPEC, + ), + ), + ], + ) + return KVCacheManager( + config, + max_model_len=262144, + scheduler_block_size=MAMBA_BLOCK_SIZE, + hash_block_size=ATTN_BLOCK_SIZE, + enable_caching=True, + use_eagle=True, + ) + + +def _split( + request: Request, + num_new_tokens: int, + use_eagle: bool = True, + partial_hit: bool = False, +) -> int: + """Call the real `Scheduler._mamba_block_aligned_split` on a stub self.""" + stub = SimpleNamespace( + cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE), + use_eagle=use_eagle, + max_num_scheduled_tokens=16384, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + # `prefix_match_unit` finer than the block size (#46384). + mamba_partial_cache_hit=partial_hit, + hash_block_size=ATTN_BLOCK_SIZE, + ) + return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens) + + +def _run_chunked_prefill( + manager: KVCacheManager, request: Request, budgets: list[int] +) -> dict[int, int]: + """Prefill `request`, one step per entry in `budgets`. + + A zero-token split means the budget cannot fund an aligned chunk; the + scheduler defers the request to a later step, so this does too. + + Returns physical mamba block id -> the token offset of the state it holds, + mirroring the GDN kernel: the running slot ends up at the chunk end. + """ + mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID] + state_at: dict[int, int] = {} + # `budgets` fragments the first steps; afterwards the request is alone and + # gets as much as it can use, so the prefill always finishes. + for step in range(len(budgets) + 64): + computed = request.num_computed_tokens + if computed >= request.num_tokens: + break + budget = budgets[step] if step < len(budgets) else request.num_tokens + num_new = _split(request, min(request.num_tokens - computed, budget)) + if num_new == 0: + continue + assert ( + manager.allocate_slots(request, num_new, num_lookahead_tokens=NUM_SPEC) + is not None + ) + request.num_computed_tokens = computed + num_new + blocks = mamba_manager.req_to_blocks[request.request_id] + running = cdiv(request.num_computed_tokens, MAMBA_BLOCK_SIZE) - 1 + state_at[blocks[running].block_id] = request.num_computed_tokens + return state_at + + +def _count_cached_boundary_states( + manager: KVCacheManager, request: Request, state_at: dict[int, int] +) -> int: + """Assert every hash-cached mamba slot holds the state its hash claims. + + Covers both full-block snapshots (`(p + 1) * block_size`) and the + partial-tail entries align mode registers at an exact token count. + + Returns the number of cached slots checked. + """ + mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID] + checked = 0 + for pos, block in enumerate(mamba_manager.req_to_blocks[request.request_id]): + if block.is_null or block.block_hash is None: + continue + claimed = block.block_hash_num_tokens + assert state_at.get(block.block_id) == claimed, ( + f"mamba slot {pos} is hashed as state@{claimed} but holds " + f"state@{state_at.get(block.block_id)}" + ) + checked += 1 + return checked + + +def _prefill(prompt_len: int, budgets: list[int]) -> int: + manager = _make_hybrid_kv_cache_manager() + (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE) + state_at = _run_chunked_prefill(manager, request, budgets) + assert request.num_computed_tokens == prompt_len, "prefill did not complete" + return _count_cached_boundary_states(manager, request, state_at) + + +def test_fragmented_first_chunk_does_not_poison_mamba_prefix_cache() -> None: + """EAGLE zeroes `last_cache_position` for prompts under two blocks. + + Past that point any chunk end used to be accepted, so a short first chunk + (concurrent prefills sharing the budget) left slot 0 at state@364 while the + next chunk crossed 1600 and published slot 0 as state@1600. + """ + _prefill(PROMPT_LEN, budgets=[364, PROMPT_LEN]) + + +def test_fragmented_tail_chunk_does_not_poison_mamba_prefix_cache() -> None: + """Same poisoning one block in, where a hit is still cacheable. + + `last_cache_position` is 1600, so the chunk ending there is cached. The next + chunk used to be free to stop mid-block (slot 1 at state@2600) and the one + after it crossed 3200, publishing slot 1 as state@3200. + """ + assert _prefill(3602, budgets=[1600, 1000, 3602]) > 0 + + +@pytest.mark.parametrize("first_chunk", [800, 900, 1599, 1601, 2000]) +def test_intermediate_chunk_ends_stay_block_aligned(first_chunk: int) -> None: + """Every non-final prefill chunk must end on a mamba block boundary.""" + _prefill(PROMPT_LEN, budgets=[first_chunk, PROMPT_LEN, PROMPT_LEN]) + + +@pytest.mark.parametrize( + ("block_size", "prompt_len", "budgets"), + [ + # Kimi-K3-scale mamba blocks: TP8 shards the recurrent state 8 ways + # (~1.5k block), DEP16 keeps it whole (~12k block). The first budget + # walks the request up to `last_cache_position`; the second is the + # sub-block fragment that lands in the unguarded tail region. + (1536, 30000, [27648, 1024, 30000]), + (12288, 30000, [12288, 4000, 30000]), + (12288, 41000, [24576, 8000, 41000]), + ], +) +def test_poisoning_is_block_size_independent( + monkeypatch: pytest.MonkeyPatch, + block_size: int, + prompt_len: int, + budgets: list[int], +) -> None: + """The invariant is per-block, so large mamba blocks are not safer.""" + import sys + + monkeypatch.setattr(sys.modules[__name__], "MAMBA_BLOCK_SIZE", block_size) + assert _prefill(prompt_len, budgets=budgets) > 0 + + +@pytest.mark.parametrize("partial_hit", [False, True]) +@pytest.mark.parametrize("resume_at", [331, 1599, 1601, 2531, 3011]) +def test_unaligned_resume_never_runs_past_its_block( + partial_hit: bool, resume_at: int +) -> None: + """A prefill resuming mid-block must re-align before crossing a boundary. + + Reachable with a finer `prefix_match_unit` (its partial-tail stop ends a + chunk off-grid by design) and with unaligned external tokens from a KV + connector. + """ + prompt_len = 3602 + (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE) + tail_boundary = prompt_len // ATTN_BLOCK_SIZE * ATTN_BLOCK_SIZE + + pos, ends = resume_at, [] + while pos < prompt_len: + request.num_computed_tokens = pos + num_new = _split(request, prompt_len - pos, partial_hit=partial_hit) + assert num_new > 0, f"no progress at {pos}" + if pos % MAMBA_BLOCK_SIZE != 0: + block_end = (pos // MAMBA_BLOCK_SIZE + 1) * MAMBA_BLOCK_SIZE + assert pos + num_new <= block_end, ( + f"chunk [{pos}, {pos + num_new}) starts mid-block and runs past " + f"{block_end}; the slot holding state@{pos} gets hashed as " + f"state@{block_end}" + ) + pos += num_new + ends.append(pos) + + for end in ends[:-1]: + aligned = end % MAMBA_BLOCK_SIZE == 0 + assert aligned or (partial_hit and end == tail_boundary), ( + f"intermediate chunk end {end} is neither block-aligned nor the " + f"partial-tail boundary" + ) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 1d6ff5a80bee..3f80ab567302 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -8,7 +8,7 @@ set -e # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This is currently targeting nv-dev branch due to sm120 support -DEEPGEMM_GIT_REF="f5a76426fa084087169693fd0cd815223576d6e9" +DEEPGEMM_GIT_REF="e21c821f39a2056d68067a466c64ddc942200106" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/config/fault_tolerance.py b/vllm/config/fault_tolerance.py index 7ed095d204f6..9f530d35f1fd 100644 --- a/vllm/config/fault_tolerance.py +++ b/vllm/config/fault_tolerance.py @@ -16,3 +16,8 @@ class FaultToleranceConfig: instructions on how to handle the error and then recover from the fault. If vLLM does not recover during this time, the original error is raised. """ + resume_requests_after_recovery: bool = False + """If True, in-flight requests are preempted back to the waiting queue on + fault (instead of being aborted) so they are automatically re-scheduled + after recovery via ``retry``. Clients will not receive abort notifications. + """ \ No newline at end of file diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index a68842f4bfee..e8988aab1940 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum, EnumMeta -from typing import Any +from typing import Any, Literal, get_args from pydantic import field_validator @@ -30,6 +30,9 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): CPU = "cpu" +MambaSSUAlgorithm = Literal["auto", "simple", "vertical", "horizontal"] + + @config class MambaConfig: """Configuration for Mamba SSM backends.""" @@ -46,6 +49,12 @@ class MambaConfig: generation. 0 uses the Triton default. Higher values improve randomness quality at the cost of compute.""" + ssu_algorithm: MambaSSUAlgorithm | None = None + """Selective state update algorithm to use with the FlashInfer backend. + None defaults to FlashInfer's "auto" algorithm. Forced algorithms must + be supported by FlashInfer for the active GPU, state dtype, and decoding + mode.""" + @field_validator("backend", mode="before") @classmethod def validate_backend_before(cls, value: Any) -> Any: @@ -54,7 +63,25 @@ def validate_backend_before(cls, value: Any) -> Any: return MambaBackendEnum[value.upper()] return value + def validate_ssu_algorithm(self) -> None: + if self.ssu_algorithm is None: + return + valid_algorithms = get_args(MambaSSUAlgorithm) + if self.ssu_algorithm not in valid_algorithms: + valid = ", ".join(valid_algorithms) + raise ValueError( + f"Unknown Mamba SSU algorithm: '{self.ssu_algorithm}'. " + f"Valid options are: {valid}" + ) + if self.backend != MambaBackendEnum.FLASHINFER: + raise ValueError( + "Mamba SSU algorithm selection is only supported with the " + "FlashInfer backend. Please set `--mamba-backend flashinfer`, " + "or omit `--mamba-ssu-algorithm`." + ) + def __post_init__(self): + self.validate_ssu_algorithm() if self.enable_stochastic_rounding: from vllm.platforms import current_platform diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index eff9bec8ee93..05b69394f987 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -60,6 +60,12 @@ def __init__(self, expected_finished_count: int): self._recv_remaining_count = dict[str, int]() self._send_remaining_count = dict[str, int]() self._expected_finished_count = expected_finished_count + + # Pending invalid block ids from failed workers + self._pending_invalid_block_ids = set[int]() + # Pending finished sending/recving from failed workers + self._pending_finished_sending = set[str]() + self._pending_finished_recving = set[str]() @classmethod def from_connector(cls, connector: "KVConnectorBase", world_size: int): @@ -156,6 +162,15 @@ def update_finished_set( invalid_block_ids |= kv_output.invalid_block_ids + # Add pending finished sets and invalid block ids from failed workers + # (fault tolerance scenario), then clear them. + finished_sending |= self._pending_finished_sending + finished_recving |= self._pending_finished_recving + invalid_block_ids |= self._pending_invalid_block_ids + self._pending_finished_sending.clear() + self._pending_finished_recving.clear() + self._pending_invalid_block_ids.clear() + # select output of the worker specified by output_rank output = outputs[output_rank] @@ -172,6 +187,57 @@ def update_finished_set( return output + def merge_kv_connector_output( + self, kv_connector_outputs: list[KVConnectorOutput | None] + ) -> None: + """Merge KV connector outputs from failed workers into the aggregator. + + Unlike aggregate(), this method does NOT return the merged result. + It only updates the internal state (remaining count dictionaries and + invalid_block_ids), so that when all workers recover and complete + successfully in a future step, the finished requests can be properly + identified and returned. + + This is used in fault tolerance scenarios where some workers fail + but their KV transfer progress needs to be preserved for later + aggregation. + """ + + def update_remaining_count( + req_ids: set[str] | None, + remaining_count_dict: dict[str, int], + pending_finished_set: set[str], + ) -> None: + for req_id in req_ids or (): + remaining_count = remaining_count_dict.get( + req_id, self._expected_finished_count + ) + remaining_count_dict[req_id] = remaining_count - 1 + if remaining_count_dict[req_id] == 0: + # All workers have finished this request, add to pending + # finished set so it can be returned in future aggregate() + pending_finished_set.add(req_id) + del remaining_count_dict[req_id] + + for kv_output in kv_connector_outputs: + if not kv_output: + continue + + # Merge finished sending/recving into remaining count + update_remaining_count( + kv_output.finished_sending, + self._send_remaining_count, + self._pending_finished_sending, + ) + update_remaining_count( + kv_output.finished_recving, + self._recv_remaining_count, + self._pending_finished_recving, + ) + + # Merge invalid block ids + self._pending_invalid_block_ids |= kv_output.invalid_block_ids + def _make_src_and_dst_indices( src_block_ids: list[int], diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 3003d7385995..219f8ffb082a 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -77,7 +77,7 @@ from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend from vllm.config.load import SafetensorsLoadStrategy from vllm.config.lora import MaxLoRARanks -from vllm.config.mamba import MambaBackendEnum +from vllm.config.mamba import MambaBackendEnum, MambaSSUAlgorithm from vllm.config.model import ( ConvertOption, HfOverrides, @@ -708,6 +708,7 @@ class EngineArgs: use_replayssm: bool = CacheConfig.use_replayssm mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON + mamba_ssu_algorithm: MambaSSUAlgorithm | None = None enable_mamba_cache_stochastic_rounding: bool = ( MambaConfig.enable_stochastic_rounding ) @@ -962,6 +963,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: description=MambaConfig.__doc__, ) mamba_group.add_argument("--mamba-backend", **mamba_kwargs["backend"]) + mamba_group.add_argument( + "--mamba-ssu-algorithm", **mamba_kwargs["ssu_algorithm"] + ) mamba_group.add_argument( "--enable-mamba-cache-stochastic-rounding", **mamba_kwargs["enable_stochastic_rounding"], @@ -2361,6 +2365,8 @@ def create_engine_config( mamba_config.backend = MambaBackendEnum[self.mamba_backend.upper()] else: mamba_config.backend = self.mamba_backend + if self.mamba_ssu_algorithm is not None: + mamba_config.ssu_algorithm = self.mamba_ssu_algorithm if self.enable_mamba_cache_stochastic_rounding: mamba_config.enable_stochastic_rounding = ( self.enable_mamba_cache_stochastic_rounding @@ -2369,6 +2375,7 @@ def create_engine_config( mamba_config.stochastic_rounding_philox_rounds = ( self.mamba_cache_philox_rounds ) + mamba_config.validate_ssu_algorithm() # Kernel config overrides kernel_config = copy.deepcopy(self.kernel_config) diff --git a/vllm/envs.py b/vllm/envs.py index 792a2084b36d..4587139f1b9b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -196,6 +196,7 @@ ] = "relax" VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True VLLM_MOE_SKIP_PADDING: bool = True + VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1533,6 +1534,18 @@ def _resolve_rust_cli_path() -> str | None: # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that # treats topk_id == -1 as a skip sentinel "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))), + # Kimi-K3 only. Under sequence-parallel MoE the dense and shared-expert MLPs + # are replicated on every rank, so each rank streams the whole weight to + # serve its own token shard. Shard them across TP instead: the MLP then + # all-gathers the full token set, computes this rank's partial, and + # reduce-scatters (which both sums across TP and restores the sequence + # sharding). Trades weight bandwidth and resident memory for two collectives + # per layer, so it only wins at low token counts: intended for decode + # instances in a P/D disaggregated deployment, not for prefill or unified + # serving. + "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool( + int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0")) + ), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py index b3ba20863403..ef00bc57647d 100644 --- a/vllm/model_executor/kernels/linear/mxfp6/base.py +++ b/vllm/model_executor/kernels/linear/mxfp6/base.py @@ -18,7 +18,7 @@ class MxFp6LinearLayerConfig: Attributes: weight_quant_key: Identifies the weight quantization format. Can be - kMxfp6E2M3Static or kMxfp6E3M2Static. + kMxfp6E2M3Static or kMxfp6E3M2Static. activation_quant_key: Identifies the activation quantization format, or `None` when activations must not be quantized. """ diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 1e1d55ac9cb6..1451a0373c63 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1269,7 +1269,11 @@ def convert_weight_to_mxfp4_moe_kernel_format( Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends. """ - from vllm.platforms.rocm import on_gfx1250 + is_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + is_gfx1250 = on_gfx1250() if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( @@ -1442,7 +1446,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not on_gfx1250(): + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not is_gfx1250: # Initially introduced for DeepSeekV4 if w13_bias is not None: @@ -1500,7 +1504,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend in TRITON_BACKENDS or ( - mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 ): from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 1795a036b124..24bdba688e70 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -13,7 +13,7 @@ import torch -from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.logger import init_logger from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -126,8 +126,13 @@ def __init__(self, mamba_config: MambaConfig): "Please install flashinfer (>= 0.6.4): " "pip install flashinfer-python" ) from e + logger.info_once("Using FlashInfer Mamba SSU algorithm: %s", self._algorithm) self._kernel = _fi_ssu + @property + def _algorithm(self) -> MambaSSUAlgorithm: + return self._mamba_config.ssu_algorithm or "auto" + @property def name(self) -> str: return "flashinfer" @@ -157,7 +162,6 @@ def __call__( if self._mamba_config.enable_stochastic_rounding else None ) - self._kernel( state, x, @@ -180,6 +184,7 @@ def __call__( out=out, rand_seed=rand_seed, philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10, + algorithm=self._algorithm, ) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py index 8fbbc65a0b33..04b3cd65cf26 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py @@ -7,6 +7,7 @@ from torch.nn.parameter import Parameter from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.model_executor.parameter import ( GroupQuantScaleParameter, ModelWeightParameter, @@ -29,7 +30,7 @@ class INCMxfp4LinearMethod(INCLinearScheme): def __init__(self, layer_config: "INCLayerConfig") -> None: self.group_size = layer_config.group_size or 32 - self.kernel = init_mxfp4_linear_kernel() + self.kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) @classmethod def get_min_capability(cls) -> int: diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index de0b0655efa9..e3edb88ae2e0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -519,6 +519,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight, layer.weight_scale, layer.logical_widths ) layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 layer.weight_scale = Parameter(max_w_scale, requires_grad=False) layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) self.fp8_linear.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 1ee1ad75fded..e9f3d0a05376 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -38,6 +38,7 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -486,8 +487,6 @@ def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool: K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic MXFP4 backend selector does not cover; route it to AITER on gfx950. """ - from vllm.platforms import current_platform - if not current_platform.is_rocm(): return False from vllm._aiter_ops import rocm_aiter_ops @@ -732,10 +731,14 @@ def _setup_kernel( # For TRITON backends, weights are wrapped tensors from triton_kernels # that don't support .detach(). Manually assign parameters. - from vllm.platforms.rocm import on_gfx1250 + is_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + is_gfx1250 = on_gfx1250() uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or ( - self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 ) if not uses_triton_weight_format: replace_parameter(layer, "w13_weight", w13) @@ -846,10 +849,14 @@ def get_fused_moe_quant_config( w2_bias = getattr(layer, "w2_bias", None) swiglu_limit = getattr(layer, "swiglu_limit", None) - from vllm.platforms.rocm import on_gfx1250 + is_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + is_gfx1250 = on_gfx1250() if self.mxfp4_backend in TRITON_BACKENDS or ( - self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 ): # TRITON backends free w13/w2_weight_scale after swizzling; the # swizzled scales live inside the precision configs instead. diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index dc88c15fc019..8e8252f45922 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -31,7 +31,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -70,6 +70,7 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, bias: bool = False, + swiglu_limit: float | None = None, prefix: str = "", use_data_parallel: bool = False, ) -> None: @@ -95,7 +96,9 @@ def __init__( raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) - self.act_fn = SiluAndMul() + self.act_fn = ( + SiluAndMul() if swiglu_limit is None else SiluAndMulWithClamp(swiglu_limit) + ) def forward(self, x): gate_up, _ = self.gate_up_proj(x) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index a4d4e9019d68..0d7da9232d2e 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -29,19 +29,25 @@ get_pp_group, get_tensor_model_parallel_world_size, ) +from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.config import set_default_rope_theta -from .exaone4 import Exaone4Attention as ExaoneMoeAttention from .exaone4 import Exaone4GatedMLP as ExaoneMoeGatedMLP from .interfaces import SupportsLoRA, SupportsPP from .utils import ( @@ -59,6 +65,7 @@ class ExaoneMoe(nn.Module): def __init__( self, config: PretrainedConfig, + swiglu_limit: float | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", enable_eplb: bool = False, @@ -119,6 +126,7 @@ def __init__( hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=False, + swiglu_limit=swiglu_limit, prefix=f"{prefix}.shared_experts", ) else: @@ -133,6 +141,8 @@ def __init__( intermediate_size=config.moe_intermediate_size, renormalize=config.norm_topk_prob, quant_config=quant_config, + activation="silu", + swiglu_limit=swiglu_limit, use_grouped_topk=True, num_expert_group=config.n_group, topk_group=config.topk_group, @@ -157,13 +167,138 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return final_hidden_states.view(orig_shape) +class ExaoneMoeAttention(nn.Module): + def __init__( + self, + config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + is_mtp: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr(config, "head_dim", None) + if self.head_dim is None: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + is_neox_style = True + if quant_config is not None and quant_config.get_name() == "gguf": + is_neox_style = False + + layer_idx = extract_layer_index(prefix) + + if config.sliding_windows is not None: + self.sliding_window_size = config.sliding_windows[layer_idx] + + if config.layer_types[layer_idx] == "full_attention": + self.sliding_window_size = None + + if is_mtp: + self.sliding_window = ( + config.mtp_layer_types[layer_idx] == "sliding_attention" + ) + if config.mtp_sliding_windows is not None: + self.sliding_window_size = config.mtp_sliding_windows[layer_idx] + + if config.mtp_layer_types[layer_idx] == "full_attention": + self.sliding_window_size = None + + # apply rotary embeddings to every layer in full attention models + self.apply_rope_all_layers = "sliding_attention" not in config.layer_types + + set_default_rope_theta(config, default_theta=1000000) + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=is_neox_style, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=self.sliding_window_size, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + q = self.q_norm(q) + q = q.flatten(-2, -1) + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) + k = self.k_norm(k) + k = k.flatten(-2, -1) + + if self.sliding_window_size or self.apply_rope_all_layers: + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + class ExaoneMoeDecoderLayer(nn.Module): def __init__( self, config: PretrainedConfig, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, - mtp_layer: bool = None, + is_mtp: bool = None, prefix: str = "", ) -> None: super().__init__() @@ -187,17 +322,29 @@ def __init__( bias=attention_bias, cache_config=cache_config, prefix=f"{prefix}.self_attn", + is_mtp=is_mtp, ) - if config.is_moe_layer[layer_idx] and not mtp_layer: + swiglu_limits = getattr(config, "swiglu_limits", None) + if swiglu_limits is not None: + swiglu_limit = swiglu_limits[layer_idx] + swiglu_limit = swiglu_limit if swiglu_limit > 0 else None + else: + swiglu_limit = None + + if config.mlp_layer_types[layer_idx] == "sparse" and not is_mtp: self.mlp = ExaoneMoe( - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" + config=config, + swiglu_limit=swiglu_limit, + quant_config=quant_config, + prefix=f"{prefix}.mlp", ) else: self.mlp = ExaoneMoeGatedMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, + swiglu_limit=swiglu_limit, quant_config=quant_config, bias=getattr(config, "mlp_bias", False), prefix=f"{prefix}.mlp", diff --git a/vllm/model_executor/models/exaone_moe_mtp.py b/vllm/model_executor/models/exaone_moe_mtp.py index a37da487dba8..ddd570d604e7 100644 --- a/vllm/model_executor/models/exaone_moe_mtp.py +++ b/vllm/model_executor/models/exaone_moe_mtp.py @@ -81,7 +81,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config.model_config.hf_config, quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", - mtp_layer=True, + is_mtp=True, ) for idx in range(self.num_mtp_layers) ) diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index 1e5dd2cda30b..dca62d596b42 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -61,7 +61,7 @@ def wrapper(org, interpolation_mode, shape): @get_rope_shape_decorate -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, disable=current_platform.simple_compile_backend == "tpu") def get_rope_shape(org, interpolation_mode, shape): return ( F.interpolate( diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index aa2a9395de19..04af45669498 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -22,6 +22,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, ) @@ -51,6 +52,7 @@ def __init__( draft_vocab_size: int, markov_rank: int, prefix: str, + quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() self.markov_w1 = nn.Embedding(vocab_size, markov_rank) @@ -58,6 +60,7 @@ def __init__( draft_vocab_size, markov_rank, bias=False, + quant_config=quant_config, prefix=maybe_prefix(prefix, "markov_w2"), disable_tp=True, ) @@ -97,6 +100,7 @@ def __init__( draft_vocab_size, config.markov_rank, prefix=maybe_prefix(prefix, "markov_head"), + quant_config=self.quant_config, ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 3798d22b6cc9..545b257794d0 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -97,7 +97,7 @@ "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"), "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"), - "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), + "ExaoneMoeForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"), "FalconForCausalLM": ("falcon", "FalconForCausalLM"), "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 71d776d8482a..c35e1ac30c9a 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -229,12 +229,19 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: Without autotuning, FlashInfer will rely on heuristics, which may be significantly slower. - Tuning is performed only on rank 0. The resulting cache is broadcast - to every rank so all ranks dispatch the same kernel tactic. + Every rank profiles the same tactics. When distributed, per-tactic + timings are averaged over the world CPU group so all ranks select the + same tactic. """ + from flashinfer.autotuner import AutoTuner, set_autotune_process_group + import vllm.utils.flashinfer as fi_utils from vllm.distributed.parallel_state import get_world_group + world = get_world_group() + is_leader = world.rank_in_group == 0 + tuner = AutoTuner.get() + autotune_kwargs: dict = {} skip_ops = _flashinfer_autotune_skip_ops(runner) if skip_ops: @@ -244,25 +251,6 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: ) autotune_kwargs["skip_ops"] = skip_ops - use_persistent_cache = True - - # When distributed, tune on every rank so the collectives stay synchronized. - if get_world_group().world_size > 1: - use_persistent_cache = False - - if not use_persistent_cache: - with torch.inference_mode(), fi_utils.autotune(**autotune_kwargs): - runner._dummy_run( - num_tokens=runner.scheduler_config.max_num_batched_tokens, - skip_eplb=True, - is_profile=True, - ) - get_world_group().barrier() - return - - world = get_world_group() - is_leader = world.rank_in_group == 0 - cache_path = resolve_flashinfer_autotune_file(runner) if is_leader: logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) @@ -271,43 +259,40 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: # When autotuning with number of tokens m, flashinfer will autotune # operations for all number of tokens up to m, so we only need to # run with the max number of tokens. + # Randomize inputs to avoid every token pick the same experts, + # which lead to some EP ranks receiving no tokens and skipping their + # MoE kernel entirely, and cause hang due to all-reduce collective + # during synchronized autotuning. dummy_run_kwargs = dict( num_tokens=runner.scheduler_config.max_num_batched_tokens, skip_eplb=True, is_profile=True, + randomize_inputs=True, ) - with torch.inference_mode(): - if is_leader: - with fi_utils.autotune( - tune_mode=True, cache=str(cache_path), **autotune_kwargs - ): - runner._dummy_run(**dummy_run_kwargs) - else: - runner._dummy_run(**dummy_run_kwargs) - - # Broadcast autotune cache from rank 0 to all other ranks so every - # rank loads the same set of chosen tactics. - tune_results: bytes | None = None + # Read cached autotune results and broadcast to all ranks. + cached_results: bytes | None = None if is_leader and cache_path.exists(): with open(cache_path, "rb") as f: - tune_results = f.read() - - tune_results = world.broadcast_object(tune_results, src=0) - - if tune_results is None: - logger.warning_once( - "No FlashInfer autotune cache entries found; " - "falling back to default tactics." - ) - else: - write_flashinfer_autotune_cache(cache_path, tune_results) + cached_results = f.read() + cached_results = world.broadcast_object(cached_results, src=0) + if cached_results is not None: + write_flashinfer_autotune_cache(cache_path, cached_results) world.barrier() - from flashinfer.autotuner import AutoTuner + tuner.load_configs(str(cache_path)) + + group = world.cpu_group if world.world_size > 1 else None + set_autotune_process_group(group) + try: + with ( + torch.inference_mode(), + fi_utils.autotune(tune_mode=True, **autotune_kwargs), + ): + runner._dummy_run(**dummy_run_kwargs) + finally: + set_autotune_process_group(None) - AutoTuner.get().load_configs(str(cache_path)) - logger.info_once( - "FlashInfer autotune cache loaded on rank %d from %s.", - world.rank_in_group, - cache_path, - ) + if world.world_size > 1: + world.barrier() + if is_leader: + tuner.save_configs(str(cache_path)) diff --git a/vllm/models/kimi_k3/__init__.py b/vllm/models/kimi_k3/__init__.py index 54f965eadaa7..52c8a6538e66 100644 --- a/vllm/models/kimi_k3/__init__.py +++ b/vllm/models/kimi_k3/__init__.py @@ -13,13 +13,20 @@ # The NVIDIA branch is the static default that type-checkers see; the ROCm # branch overrides it at runtime (kept type-compatible via type: ignore). -if TYPE_CHECKING or not current_platform.is_rocm(): +# TPU plugins import the shared ``common`` modules through this package, but +# register their own model classes. Do not eagerly import a GPU implementation. +if TYPE_CHECKING: from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM from .nvidia.mtp import KimiK3MTP -else: +elif current_platform.device_type == "tpu": + pass +elif current_platform.is_rocm(): from .amd.linear import KimiLinearForCausalLM # type: ignore[assignment] from .amd.model import KimiK3ForConditionalGeneration # type: ignore[assignment] from .amd.mtp import KimiK3MTP # type: ignore[assignment] +else: + from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM + from .nvidia.mtp import KimiK3MTP __all__ = [ "KimiK3ForConditionalGeneration", diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 94178c9c73c2..68015827b851 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -129,7 +129,42 @@ _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256 +def shard_sequence_parallel_mlp( + hidden_size: int, + intermediate_size: int, + use_sequence_parallel: bool, + eligible: bool, +) -> bool: + """Whether to TP-shard a sequence-parallel MLP instead of replicating it. + + Opt-in via ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT``; see :class:`KimiMLP` for + the trade-off and :mod:`vllm.envs` for when it is worth enabling. + """ + enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT + if not (use_sequence_parallel and eligible and enabled): + return False + tp_size = get_tensor_model_parallel_world_size() + return ( + tp_size > 1 and intermediate_size % tp_size == 0 and hidden_size % tp_size == 0 + ) + + class KimiMLP(nn.Module): + """Dense / shared-expert MLP, optionally TP-sharded under sequence parallel. + + Under sequence parallelism each rank owns a distinct slice of the tokens, so + by default both projections are replicated (``disable_tp``) and the block + needs no collective. That makes every rank stream the entire weight to serve + its own token shard. + + With ``VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT`` the weights are TP-sharded + instead. A rank then holds only a slice of the intermediate dim, so it + cannot finish its own tokens alone: ``forward`` all-gathers the full token + set, computes this rank's partial, and reduce-scatters. The reduce-scatter + sums across TP and restores the sequence sharding in one collective, so the + block still ends with one collective per direction. + """ + def __init__( self, hidden_size: int, @@ -138,18 +173,27 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, use_sequence_parallel: bool = False, + can_shard_sequence_parallel: bool = False, prefix: str = "", activation_situ_beta: float | None = None, activation_situ_linear_beta: float | None = None, ) -> None: super().__init__() + self.shard_sequence_parallel = shard_sequence_parallel_mlp( + hidden_size, + intermediate_size, + use_sequence_parallel, + can_shard_sequence_parallel, + ) + replicate = use_sequence_parallel and not self.shard_sequence_parallel + self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, - disable_tp=use_sequence_parallel, + disable_tp=replicate, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( @@ -157,8 +201,10 @@ def __init__( hidden_size, bias=False, quant_config=quant_config, - reduce_results=reduce_results, - disable_tp=use_sequence_parallel, + # Sharded sequence parallel reduces via the reduce-scatter in + # forward(), which also restores the sequence sharding. + reduce_results=False if self.shard_sequence_parallel else reduce_results, + disable_tp=replicate, prefix=f"{prefix}.down_proj", ) if hidden_act == "silu": @@ -175,9 +221,17 @@ def __init__( ) def forward(self, x): + if self.shard_sequence_parallel: + # Each rank holds a weight shard but only its own tokens, so it + # cannot finish those tokens alone: gather the full token set, + # compute this rank's partial for all of them, then reduce-scatter, + # which sums across TP and restores the sequence sharding. + x = sp_all_gather(x) gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) + if self.shard_sequence_parallel: + x = sp_reduce_scatter(x) return x @@ -443,7 +497,7 @@ def __init__( min_moe_intermediate_per_partition = getattr( config, "min_moe_intermediate_per_partition", 256 ) - if self.tp_size > 1: + if self.tp_size > 1 and not vllm_config.parallel_config.enable_expert_parallel: moe_intermediate_per_partition = moe_intermediate_size // self.tp_size if moe_intermediate_per_partition < min_moe_intermediate_per_partition: self.padded_moe_intermediate_size = ( @@ -478,6 +532,10 @@ def __init__( quant_config=quant_config, reduce_results=False, use_sequence_parallel=use_sequence_parallel, + # Only the MegaMoE path calls the shared experts directly; the + # FusedMoE path below hands them to the runner, which fuses + # their reduction and assumes the replicated layout. + can_shard_sequence_parallel=self.use_mega_moe, prefix=f"{prefix}.shared_experts", activation_situ_beta=activation_situ_beta, activation_situ_linear_beta=activation_situ_linear_beta, @@ -800,6 +858,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.mlp", use_sequence_parallel=self.use_sequence_parallel, + can_shard_sequence_parallel=True, activation_situ_beta=config.activation_situ_beta, activation_situ_linear_beta=config.activation_situ_linear_beta, ) @@ -956,6 +1015,7 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant): "gate_up_proj": ["gate_proj", "up_proj"], "in_proj_qkvgfab": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj"], "conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"], + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 99a8b44041f9..9758fd794e0a 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -381,7 +381,8 @@ def _mamba_block_aligned_split( ) # Split only during prefill: `request.num_tokens - 1` extends this to # resumed requests replaying their output tokens. - if start >= max(request.num_prompt_tokens, request.num_tokens - 1): + prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1) + if start >= prefill_end: return num_new_tokens block_size = self.cache_config.block_size @@ -393,10 +394,12 @@ def _mamba_block_aligned_split( last_cache_position = max(last_cache_position - block_size, 0) end = start + num_new_tokens - # Until `last_cache_position`, prefer chunks ending on block - # boundaries. When a block cannot fit in any configured prefill chunk, - # allow sub-block progress and re-align at the next reachable boundary. - if end < last_cache_position: + # Invariant: slot p holds the state after exactly (p + 1) * block_size + # tokens. State is written at chunk ends, so chunk ends must be block + # aligned. Exempt: the prompt's last chunk, whose slot decode advances + # to the boundary. A block too wide for one chunk advances sub-block + # and re-aligns at the next boundary. + if end < prefill_end: max_prefill_tokens = self.max_num_scheduled_tokens long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold if long_prefill_threshold > 0: @@ -412,12 +415,9 @@ def _mamba_block_aligned_split( else 0 ) stops = ( - # Resumed mid-block (fine-grained partial hash hit): re-align to - # the block grid before running on, so the crossed boundary's - # state is materialized (unless it is past the cacheable range). - next_block_boundary - if start % block_size != 0 and next_block_boundary <= last_cache_position - else 0, + # Same invariant: a chunk starting mid-block stops at the boundary + # rather than running past it. + next_block_boundary if start % block_size != 0 else 0, # Never run past the last cacheable block boundary mid-chunk. last_cache_position, # Fine-grained hits: the prompt's partial-tail entry can only be @@ -2740,6 +2740,20 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): assert req_id in self.requests self._free_blocks(self.requests[req_id]) + def ft_update_kv_xfer_finished( + self, kv_connector_output: KVConnectorOutput | None + ) -> None: + """Fault tolerance: update KV transfer state from kv_connector_output. + + Unlike update_from_output(), this only updates the KV transfer state + (e.g. frees blocks of finished-sending requests, marks requests whose + KV has been fully received), without consuming other fields of the + model runner output. + """ + if kv_connector_output is None: + return + self._update_from_kv_xfer_finished(kv_connector_output) + def _update_requests_with_invalid_blocks( self, requests: Iterable[Request], diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 11c78c7e7453..09d7822899d2 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -651,6 +651,17 @@ def step_with_batch_queue( deferred_scheduler_output = None if self.scheduler.has_requests(): scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) + print( + f"[DEBUG] scheduler_output: " + f"total_scheduled={scheduler_output.total_num_scheduled_tokens}, " + f"num_new_reqs={len(scheduler_output.scheduled_new_reqs)}, " + f"num_cached_reqs={len(scheduler_output.scheduled_cached_reqs.req_ids)}, " + f"finished_req_ids={scheduler_output.finished_req_ids}, " + f"num_scheduled_tokens={scheduler_output.num_scheduled_tokens}, " + f"num_common_prefix_blocks={scheduler_output.num_common_prefix_blocks}, " + f"num_spec_tokens={scheduler_output.num_spec_tokens_to_schedule}" + ) + with self.log_error_detail(scheduler_output): exec_future = self.model_executor.execute_model( scheduler_output, non_block=True diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index b464f2dcf635..69f1e4dc8f5d 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -61,12 +61,24 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.executor.abstract import Executor, FailureCallback from vllm.v1.executor.vllm_net_devices import set_worker_net_device -from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput +from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.worker.worker_base import WorkerWrapperBase logger = init_logger(__name__) +class ExceptionWithKVConnectorOutput(Exception): + """Exception that carries KV connector output for fault tolerance. + + When a worker fails but KV transfer state needs to be preserved, + this exception is used to pass the kv_connector_output along with + the error information. + """ + def __init__(self, error_msg: str, kv_connector_output: KVConnectorOutput | None): + super().__init__(error_msg) + self.kv_connector_output = kv_connector_output + + class FutureWrapper(Future): def __init__( self, @@ -321,6 +333,7 @@ def register_failure_callback(self, callback: FailureCallback): def execute_model( # type: ignore[override] self, scheduler_output: SchedulerOutput, non_block: bool = False ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]: + logger.info("engine core execute_model") return self.collective_rpc( "execute_model", args=(scheduler_output,), @@ -333,6 +346,7 @@ def execute_model( # type: ignore[override] def sample_tokens( # type: ignore[override] self, grammar_output: GrammarOutput | None, non_block: bool = False ) -> ModelRunnerOutput | Future[ModelRunnerOutput]: + logger.info("engine core sample_tokens") return self.collective_rpc( "sample_tokens", args=(grammar_output,), @@ -343,6 +357,7 @@ def sample_tokens( # type: ignore[override] ) def execute_dummy_batch(self) -> None: + logger.info("engine core execute_dummy_batch") self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank) def take_draft_token_ids(self) -> DraftTokenIds | None: @@ -393,6 +408,8 @@ def collective_rpc( # type: ignore[override] def get_response(): responses = [] + has_failure = False + for mq in response_mqs: dequeue_timeout = ( None if deadline is None else max(0.0, deadline - time.monotonic()) @@ -401,12 +418,51 @@ def get_response(): status, result = mq.dequeue(timeout=dequeue_timeout) except TimeoutError as e: raise TimeoutError(f"RPC call to {method} timed out.") from e - if status != WorkerProc.ResponseStatus.SUCCESS: - raise RuntimeError( - f"Worker failed with error '{result}', please check the" - " stack trace above for the root cause" - ) - responses.append(result) + + # Whether this call carries KV aggregation. Mixed deployment + # without KV transfer often has no kv_connector_output at all, + # in which case there is nothing to preserve and we can fail + # fast instead of collecting all responses. + has_kv = kv_output_aggregator is not None + + # Only FT scenario will have FAILURE_WITH_KV_OUTPUT status + if status == WorkerProc.ResponseStatus.FAILURE_WITH_KV_OUTPUT: + if self.is_failed or not has_kv: + raise RuntimeError("Worker failed with KV connector output") + has_failure = True + assert isinstance(result, KVConnectorOutput) + responses.append(result) + elif status != WorkerProc.ResponseStatus.SUCCESS: + if ( + not self.vllm_config.parallel_config.enable_fault_tolerance + or self.is_failed + or not has_kv + ): + raise RuntimeError( + f"Worker failed with error '{result}', please check the" + " stack trace above for the root cause" + ) + has_failure = True + else: + responses.append(result) + + # If any worker failed, extract KV connector outputs, merge them, + # then raise exception + if has_failure: + kv_connector_outputs = [] + for resp in responses: + if isinstance(resp, KVConnectorOutput): + kv_connector_outputs.append(resp) + elif resp is not None and hasattr(resp, 'kv_connector_output'): + kv_output = resp.kv_connector_output + if kv_output is not None: + kv_connector_outputs.append(kv_output) + + if kv_connector_outputs and kv_output_aggregator is not None: + kv_output_aggregator.merge_kv_connector_output(kv_connector_outputs) + + raise RuntimeError("One or more workers failed") + return responses[0] if output_rank is not None else responses future = FutureWrapper( @@ -615,6 +671,10 @@ def __init__( shared_worker_lock: LockType, is_driver_worker: bool, ): + self.get_method_count = 0 + self.enqueue_output_count = 0 + self.dequeue_output_count = 0 + self.enqueue_response_count = 0 self.rank = rank wrapper = WorkerWrapperBase(rpc_rank=local_rank, global_rank=rank) # TODO: move `init_worker` to executor level as a collective rpc call @@ -946,6 +1006,7 @@ def signal_handler(signum, frame): class ResponseStatus(Enum): SUCCESS = auto() FAILURE = auto() + FAILURE_WITH_KV_OUTPUT = auto() def enqueue_output(self, output: Any): """Prepares output from the worker and enqueues it to the @@ -957,14 +1018,33 @@ def enqueue_output(self, output: Any): output = output.get_output() except Exception as e: logger.exception("Error getting async model runner output") - output = e - - if isinstance(output, Exception): + self.worker.fault_occur = True + # Check if fault tolerance is enabled + ft_enabled = self.worker.vllm_config.parallel_config.enable_fault_tolerance + if ft_enabled: + # In FT scenario, even if the forward pass throws an exception, + # we need to extract the kv_connector_output and pass it to the + # executor so that KV transfer progress is not lost. + kv_connector_output = output.get_kv_connector_output() + if kv_connector_output is not None and not kv_connector_output.is_empty(): + output = ExceptionWithKVConnectorOutput(str(e), kv_connector_output) + else: + output = e + else: + output = e + + if isinstance(output, ExceptionWithKVConnectorOutput): + # FT scenario: worker failed but we extracted KV connector state + # so that KV transfer progress is not lost. + result = (WorkerProc.ResponseStatus.FAILURE_WITH_KV_OUTPUT, output.kv_connector_output) + elif isinstance(output, Exception): result = (WorkerProc.ResponseStatus.FAILURE, str(output)) else: result = (WorkerProc.ResponseStatus.SUCCESS, output) if (response_mq := self.worker_response_mq) is not None: response_mq.enqueue(result) + self.enqueue_response_count += 1 + logger.info(f"{self.rank=} worker_response_mq enqueue, {self.enqueue_response_count=}") def handle_output(self, output: Any): """Handles output from the worker. If async scheduling is enabled, @@ -973,6 +1053,8 @@ def handle_output(self, output: Any): """ if self.use_async_scheduling: self.async_output_queue.put(output) + self.enqueue_output_count += 1 + logger.info(f"{self.rank=} async_output_queue put, {self.enqueue_output_count=}") else: self.enqueue_output(output) @@ -992,6 +1074,8 @@ def async_output_busy_loop(self): while True: output = self.async_output_queue.get() + self.dequeue_output_count += 1 + logger.info(f"{self.rank=} async_output_queue get, {self.dequeue_output_count=}") self.enqueue_output(output) def worker_busy_loop(self): @@ -1001,25 +1085,49 @@ def worker_busy_loop(self): method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue( indefinite=True ) + self.get_method_count += 1 try: if isinstance(method, str): func = getattr(self.worker, method) elif isinstance(method, bytes): func = partial(cloudpickle.loads(method), self.worker) - + logger.info(f"{self.rank=} exec {method=} {self.get_method_count=}") output = func(*args, **kwargs) - + logger.info(f"{self.rank=} exec finish {method=}, {self.get_method_count=}") if output_rank is None or self.rank == output_rank: self.handle_output(output) except Exception as e: # Notes have been introduced in python 3.11 + self.worker.fault_occur = True if hasattr(e, "add_note"): - e.add_note(traceback.format_exc()) + pass + #e.add_note(traceback.format_exc()) logger.exception("WorkerProc hit an exception.") # exception might not be serializable, so we convert it to # string, only for logging purpose. if output_rank is None or self.rank == output_rank: - self.handle_output(e) + # Check if fault tolerance is enabled + ft_enabled = self.worker.vllm_config.parallel_config.enable_fault_tolerance + if ft_enabled: + # In FT scenario, even if the forward pass throws an exception, + # we need to extract the kv_connector_output and pass it to the + # executor so that KV transfer progress is not lost. + model_runner = self.worker.model_runner + assert model_runner is not None + # Only v1 model runner has the kv_connector_output + # field; v2 stores it in the async output instead. + if hasattr(model_runner, 'kv_connector_output'): + kv_connector_output = model_runner.kv_connector_output + model_runner.kv_connector_output = None + else: + kv_connector_output = None + if kv_connector_output is not None and not kv_connector_output.is_empty(): + output = ExceptionWithKVConnectorOutput(str(e), kv_connector_output) + self.handle_output(output) + else: + self.handle_output(e) + else: + self.handle_output(e) @staticmethod def setup_proc_title_and_log_prefix(enable_ep: bool) -> None: diff --git a/vllm/v1/fault_tolerance/engine_core_sentinel.py b/vllm/v1/fault_tolerance/engine_core_sentinel.py index 1d82dd9f7435..f095aca72970 100644 --- a/vllm/v1/fault_tolerance/engine_core_sentinel.py +++ b/vllm/v1/fault_tolerance/engine_core_sentinel.py @@ -5,7 +5,8 @@ import json import threading from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast +import time import msgspec @@ -21,6 +22,7 @@ UtilityOutput, ) from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult +from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.request import RequestStatus from vllm.v1.serial_utils import UtilityResult, run_method @@ -85,10 +87,22 @@ def on_fault(self, exc: Exception): ) engine = self.engine - aborted = engine.scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED) - engine._send_abort_outputs(aborted) - if engine.batch_queue is not None: - engine.batch_queue.clear() + scheduler = cast(Scheduler, engine.scheduler) + ft_config = self.parallel_config.fault_tolerance_config + #self._clear_contaminated_blocks() + + if ft_config.resume_requests_after_recovery: + timestamp = time.monotonic() + while scheduler.running: + request = scheduler.running.pop() + scheduler._preempt_request(request, timestamp) + request.num_in_flight_tokens = 0 + request.num_stale_output_tokens = 0 + scheduler.prev_step_scheduled_req_ids.clear() + else: + aborted = scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED) + engine._send_abort_outputs(aborted) + if ( hasattr(engine.model_executor, "is_failed") and engine.model_executor.is_failed @@ -105,6 +119,67 @@ def on_fault(self, exc: Exception): ) self._push_status() + def _clean_batch_queue(self) -> None: + """Drain the batch queue instead of clearing it directly. + + The pending futures may hold kv_connector_output that must be + consumed (worker-side connectors use get-and-clear semantics), + otherwise KV transfer progress would be lost. + Note: we don't call scheduler.update_from_output() here; we only + update the KV transfer state via the kv_connector_output. + """ + engine = self.engine + if engine.batch_queue is None: + return + scheduler = cast(Scheduler, engine.scheduler) + + while engine.batch_queue: + future, _, _ = engine.batch_queue.pop() + try: + model_output = future.result() + except Exception as e: + # Failed step: kv_connector_output has already been merged + # into the executor's KVOutputAggregator (in get_response), + # so it is safe to drop the exception here. + logger.warning( + "[FT] Dropping exception from batch queue during fault " + "handling: %s", + e, + ) + continue + if model_output is not None: + scheduler.ft_update_kv_xfer_finished( + model_output.kv_connector_output + ) + engine.batch_queue.clear() + + def _clear_contaminated_blocks(self) -> None: + """Evict KV blocks possibly contaminated by the failed step.""" + + scheduler = cast(Scheduler, self.engine.scheduler) + kv_cache_manager = scheduler.kv_cache_manager + dirty_block_ids: set[int] = set() + + for request in scheduler.running: + num_written = request.num_in_flight_tokens + if num_written <= 0: + continue + start = max(request.num_computed_tokens - num_written, 0) + end = request.num_computed_tokens - 1 + + for mgr in kv_cache_manager.coordinator.single_type_managers: + blocks = mgr.req_to_blocks.get(request.request_id) + if not blocks: + continue + blk_start = start // mgr.block_size + blk_end = end // mgr.block_size + for blk in blocks[blk_start : blk_end + 1]: + if not blk.is_null: + dirty_block_ids.add(blk.block_id) + + if dirty_block_ids: + kv_cache_manager.evict_blocks(dirty_block_ids) + def _push_status(self): """Push current health to the client so it can refresh its cache.""" payload = {"id": self.engine_index, "status": self.status_type.name.lower()} @@ -129,7 +204,7 @@ def retry(self, ft_request: FaultToleranceRequest) -> FaultToleranceResult: engine.step_counter = 0 executor.collective_rpc("handle_ft_command", args=(ft_request,)) - + self._clean_batch_queue() self.status_type = EngineStatusType.HEALTHY logger.info("[FT] Engine %d status -> HEALTHY", self.engine_index) self.resumed.set() diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index abb6a2be0d95..3d26601c6b97 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -333,6 +333,15 @@ def get_output(self) -> ModelRunnerOutput: """ pass + @abstractmethod + def get_kv_connector_output(self) -> KVConnectorOutput | None: + """Get the KV connector output for this async output. + + This method should return the pending KV connector output without + consuming it, or None if there is no pending output. + """ + pass + @dataclass class DraftTokenIds: diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index a5d5e87ce178..0e860c50b813 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -8,6 +8,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager from vllm.v1.outputs import ( AsyncModelRunnerOutput, + KVConnectorOutput, LogprobsTensors, ModelRunnerOutput, PoolerOutput, @@ -89,6 +90,15 @@ def get_output(self) -> ModelRunnerOutput: return self.model_runner_output + def get_kv_connector_output(self) -> KVConnectorOutput | None: + """Get the KV connector output for this async output. + + In FT scenario, even if the forward pass throws an exception, + we need to extract the kv_connector_output and pass it to the + executor so that KV transfer progress is not lost. + """ + return self.model_runner_output.kv_connector_output + class AsyncPoolingOutput(AsyncModelRunnerOutput): def __init__( @@ -133,6 +143,17 @@ def get_output(self) -> ModelRunnerOutput: self.model_runner_output.pooler_output = pooler_output return self.model_runner_output + def get_kv_connector_output(self) -> KVConnectorOutput | None: + """Get the KV connector output for this async output. + + In FT scenario, even if the forward pass throws an exception, + we need to extract the kv_connector_output and pass it to the + executor so that KV transfer progress is not lost. + For v2 model runner, the KV connector output is stored in the + model_runner_output's kv_connector_output field. + """ + return self.model_runner_output.kv_connector_output + def async_copy_to_np(x: torch.Tensor) -> np.ndarray: return x.to("cpu", non_blocking=True).numpy() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 5083087bf072..71c250041560 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -349,6 +349,17 @@ def get_output(self) -> ModelRunnerOutput: return output + def get_kv_connector_output(self) -> KVConnectorOutput | None: + """Get the KV connector output for this async output. + + In FT scenario, even if the forward pass throws an exception, + we need to extract the kv_connector_output and pass it to the + executor so that KV transfer progress is not lost. + For v1 model runner, the KV connector output is stored in the + model_runner_output's kv_connector_output field. + """ + return self._model_runner_output.kv_connector_output + def _copy_pooler_output_to_cpu( raw_pooler_output: PoolerOutput, finished_mask: list[bool] @@ -5749,7 +5760,10 @@ def _get_nans_in_logits( @contextmanager def maybe_randomize_inputs( - self, input_ids: torch.Tensor | None, inputs_embeds: torch.Tensor | None + self, + input_ids: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + randomize_inputs: bool = False, ): """ Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set. @@ -5759,7 +5773,9 @@ def maybe_randomize_inputs( """ dp_size = self.vllm_config.parallel_config.data_parallel_size - randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1 + randomize_inputs = randomize_inputs or ( + envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1 + ) if not randomize_inputs: yield elif input_ids is not None: @@ -5836,6 +5852,7 @@ def _dummy_run( is_graph_capturing: bool = False, num_active_loras: int = 0, profile_seq_lens: int | None = None, + randomize_inputs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -6133,7 +6150,9 @@ def _dummy_run( num_tokens_across_dp[:] = num_tokens_padded with ( - self.maybe_randomize_inputs(input_ids, inputs_embeds), + self.maybe_randomize_inputs( + input_ids, inputs_embeds, randomize_inputs=randomize_inputs + ), set_forward_context( attn_metadata, self.vllm_config, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index bbab98dbd7b3..1f9567bdf5df 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -150,6 +150,10 @@ def __init__( self.elastic_ep_executor = ElasticEPScalingExecutor(self) self.worker_sentinel: WorkerSentinel | None = None + # FT: set to True once this worker raises an exception. While True, + # execute_model skips forward and only processes KV transfers. + # Reset by WorkerSentinel.retry() after recovery. + self.fault_occur = False if self.parallel_config.enable_fault_tolerance: self.worker_sentinel = WorkerSentinel(worker=self) # Buffers saved before sleep @@ -1019,6 +1023,24 @@ def sample_tokens( def execute_model( self, scheduler_output: "SchedulerOutput" ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None: + # FT: after a fault the worker state is contaminated; skip forward + # and only process KV transfer tasks (preemptions, loads/saves), + # otherwise in-flight transfers would be silently dropped. + # Returns None: no model output for this step. + if self.parallel_config.enable_fault_tolerance and self.fault_occur: + logger.info("WorkerProc is in fault tolerance mode, skip forward pass.") + if self.use_v2_model_runner: + self.model_runner.kv_connector.no_forward(scheduler_output) + elif has_kv_transfer_group(): + assert scheduler_output.kv_connector_metadata is not None + get_kv_transfer_group().handle_preemptions( + scheduler_output.kv_connector_metadata + ) + self.model_runner.kv_connector_no_forward( + scheduler_output, self.vllm_config + ) + return None + # ensure any previous non-blocking PP sends are complete if self._pp_send_work: for handle in self._pp_send_work: @@ -1163,6 +1185,10 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): self.profiler.stop() def execute_dummy_batch(self) -> None: + # FT: skip the dummy forward on a faulted worker; its state is + # contaminated and it must not join collective ops until retry(). + if self.parallel_config.enable_fault_tolerance and self.fault_occur: + return num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1) self.model_runner._dummy_run(num_tokens, uniform_decode=True) diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 7a640b875fa5..4b066a2ba723 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -93,23 +93,30 @@ def _get_kv_connector_output( # involved may be disjoint from the running requests. # Do this here to save a collective_rpc. kv_connector.start_load_kv(get_forward_context()) - try: - yield output - finally: - if wait_for_save and not defer_finalize: - kv_connector.wait_for_save() + yield output + + # Note: If _model_forward raises an exception, the code below will not + # execute. The KV transfer state (finished_sending/finished_recving/ + # invalid_block_ids) will accumulate inside kv_connector. On the next + # successful execution, get_finished() and related methods will return + # all accumulated data at once. This is acceptable because: + # 1. KVOutputAggregator's counter-based logic tolerates delayed reports + # 2. The data is not lost, just delayed + # 3. bind_connector_metadata will overwrite stale metadata + if wait_for_save and not defer_finalize: + kv_connector.wait_for_save() - output.finished_sending, output.finished_recving = ( - kv_connector.get_finished(scheduler_output.finished_req_ids) - ) - output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors() + output.finished_sending, output.finished_recving = ( + kv_connector.get_finished(scheduler_output.finished_req_ids) + ) + output.invalid_block_ids = kv_connector.get_block_ids_with_load_errors() - output.kv_connector_stats = kv_connector.get_kv_connector_stats() - output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events() - output.kv_connector_worker_meta = kv_connector.build_connector_worker_meta() + output.kv_connector_stats = kv_connector.get_kv_connector_stats() + output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events() + output.kv_connector_worker_meta = kv_connector.build_connector_worker_meta() - if not defer_finalize: - kv_connector.clear_connector_metadata() + if not defer_finalize: + kv_connector.clear_connector_metadata() @staticmethod def use_uniform_kv_cache( diff --git a/vllm/v1/worker/sentinel/gpu_worker_sentinel.py b/vllm/v1/worker/sentinel/gpu_worker_sentinel.py index 80050cf9d391..911cfbf24edf 100644 --- a/vllm/v1/worker/sentinel/gpu_worker_sentinel.py +++ b/vllm/v1/worker/sentinel/gpu_worker_sentinel.py @@ -66,10 +66,14 @@ def retry(self, ft_request: FaultToleranceRequest): self.dp_size, backend="gloo", ) + # Recovery completed; allow the worker to execute methods again. + self.worker.fault_occur = False def _clean_worker_state(self): model_runner = self.worker.model_runner model_runner.execute_model_state = None + ft_config = model_runner.vllm_config.parallel_config.fault_tolerance_config + resume = ft_config.resume_requests_after_recovery if self.worker.use_v2_model_runner: runner = cast("GPUModelRunnerV2", model_runner) for req_id in list(runner.req_states.req_id_to_index): @@ -80,8 +84,9 @@ def _clean_worker_state(self): input_batch = model_runner.input_batch cached_req_ids = list(input_batch.req_id_to_index) for req_id in cached_req_ids: - model_runner.requests.pop(req_id, None) - model_runner.num_prompt_logprobs.pop(req_id, None) + if not resume: + model_runner.requests.pop(req_id, None) + model_runner.num_prompt_logprobs.pop(req_id, None) input_batch.remove_request(req_id) input_batch.condense()