diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index df9b80f7f9a8..1d9662f3f5b7 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -1,16 +1,3 @@ -# CUDA architecture lists — following PyTorch RELEASE.md -# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md) -# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support -env: - CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX" - # aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13) - CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX" - CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0" - CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0" - MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl" - MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl" - MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl" - steps: - input: "Provide Release version here" id: input-release-version @@ -25,13 +12,14 @@ steps: depends_on: ~ id: build-wheel-arm64-cuda-12-9 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: + # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -39,13 +27,14 @@ steps: depends_on: ~ id: build-wheel-arm64-cuda-13-0 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: + # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" @@ -53,13 +42,12 @@ steps: depends_on: ~ id: build-wheel-arm64-cpu agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" @@ -67,13 +55,12 @@ steps: depends_on: ~ id: build-wheel-x86-cuda-12-9 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31" env: DOCKER_BUILDKIT: "1" @@ -81,13 +68,12 @@ steps: depends_on: ~ id: build-wheel-x86-cuda-13-0 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" @@ -95,265 +81,80 @@ steps: depends_on: ~ id: build-wheel-x86-cpu agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" - - label: "Generate and upload wheel indices" - depends_on: "build-wheels" - allow_dependency_failure: true - agents: - queue: cpu_queue_release - commands: - - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" - - - block: "Unblock to build release Docker images" - depends_on: ~ - key: block-build-release-images - if: build.env("NIGHTLY") != "1" - - group: "Build release Docker images" key: "build-release-images" - depends_on: block-build-release-images - allow_dependency_failure: true steps: - - label: "Build release image - x86_64 - CUDA 13.0" + - label: "Build release image - x86_64 - CUDA 12.9" depends_on: ~ id: build-release-image-x86 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=13.0.2 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" # re-tag to default image tag and push, just in case arm64 build fails - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"' - - label: "Build release image - aarch64 - CUDA 13.0" + - label: "Build release image - aarch64 - CUDA 12.9" depends_on: ~ id: build-release-image-arm64 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=13.0.2 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"' - - label: "Build release image - x86_64 - CUDA 12.9" + - label: "Build release image - x86_64 - CUDA 13.0" depends_on: ~ - id: build-release-image-x86-cuda-12-9 + id: build-release-image-x86-cuda-13-0 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" # re-tag to default image tag and push, just in case arm64 build fails - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"' + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" - - label: "Build release image - aarch64 - CUDA 12.9" - depends_on: ~ - id: build-release-image-arm64-cuda-12-9 - agents: - queue: arm64_cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"' - - - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04" - depends_on: ~ - id: build-release-image-x86-ubuntu2404 - agents: - queue: cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=13.0.2 \ - --build-arg UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"' - - - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04" - depends_on: ~ - id: build-release-image-arm64-ubuntu2404 - agents: - queue: arm64_cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=13.0.2 \ - --build-arg UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"' - - - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04" - depends_on: ~ - id: build-release-image-x86-cuda-12-9-ubuntu2404 - agents: - queue: cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' - - - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04" + - label: "Build release image - aarch64 - CUDA 13.0" depends_on: ~ - id: build-release-image-arm64-cuda-12-9-ubuntu2404 + id: build-release-image-arm64-cuda-13-0 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' + # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" - block: "Build release image for x86_64 CPU" key: block-cpu-release-image-build depends_on: ~ - label: "Build release image - x86_64 - CPU" - key: build-cpu-release-image-x86 depends_on: - block-cpu-release-image-build - input-release-version agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest" - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"' env: DOCKER_BUILDKIT: "1" @@ -362,82 +163,61 @@ steps: depends_on: ~ - label: "Build release image - arm64 - CPU" - key: build-cpu-release-image-arm64 - depends_on: + depends_on: - block-arm64-cpu-release-image-build - input-release-version agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest" - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"' env: DOCKER_BUILDKIT: "1" - group: "Publish release images" key: "publish-release-images" steps: - - label: "Create multi-arch manifest - CUDA 13.0" + - label: "Create multi-arch manifest - CUDA 12.9" depends_on: - build-release-image-x86 - build-release-image-arm64 id: create-multi-arch-manifest agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend" - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"' - - - label: "Create multi-arch manifest - CUDA 12.9" - depends_on: - - build-release-image-x86-cuda-12-9 - - build-release-image-arm64-cuda-12-9 - id: create-multi-arch-manifest-cuda-12-9 - agents: - queue: small_cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"' - - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04" + - label: "Annotate release workflow - CUDA 12.9" depends_on: - - build-release-image-x86-ubuntu2404 - - build-release-image-arm64-ubuntu2404 - id: create-multi-arch-manifest-ubuntu2404 + - create-multi-arch-manifest + id: annotate-release-workflow agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"' + - "bash .buildkite/scripts/annotate-release.sh" - - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04" + - label: "Create multi-arch manifest - CUDA 13.0" depends_on: - - build-release-image-x86-cuda-12-9-ubuntu2404 - - build-release-image-arm64-cuda-12-9-ubuntu2404 - id: create-multi-arch-manifest-cuda-12-9-ubuntu2404 + - build-release-image-x86-cuda-13-0 + - build-release-image-arm64-cuda-13-0 + id: create-multi-arch-manifest-cuda-13-0 agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"' + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" - label: "Publish nightly multi-arch image to DockerHub" depends_on: - create-multi-arch-manifest if: build.env("NIGHTLY") == "1" agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - "bash .buildkite/scripts/push-nightly-builds.sh" # Clean up old nightly builds (keep only last 14) @@ -450,16 +230,16 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" - - label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9" + - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0" depends_on: - - create-multi-arch-manifest-cuda-12-9 + - create-multi-arch-manifest-cuda-13-0 if: build.env("NIGHTLY") == "1" agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - - "bash .buildkite/scripts/push-nightly-builds.sh cu129" + - "bash .buildkite/scripts/push-nightly-builds.sh cu130" # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-" plugins: - docker-login#v3.0.0: username: vllmbot @@ -468,6 +248,24 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" + - group: "Publish wheels" + key: "publish-wheels" + steps: + - block: "Confirm update release wheels to PyPI (experimental, use with caution)?" + key: block-upload-release-wheels + depends_on: + - input-release-version + - build-wheels + + - label: "Upload release wheels to PyPI" + depends_on: + - block-upload-release-wheels + id: upload-release-wheels + agents: + queue: small_cpu_queue_postmerge + commands: + - "bash .buildkite/scripts/upload-release-wheels-pypi.sh" + # ============================================================================= # ROCm Release Pipeline (x86_64 only) # ============================================================================= @@ -476,112 +274,184 @@ steps: # To build a specific version, trigger the build from that branch/tag. # # Environment variables for ROCm builds (set via Buildkite UI or schedule): + # ROCM_PYTHON_VERSION: Python version (default: 3.12) + # PYTORCH_ROCM_ARCH: GPU architectures (default: gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151) + # ROCM_UPLOAD_WHEELS: Upload to S3 (default: false for nightly, true for releases) + # ROCM_FORCE_REBUILD: Force rebuild base wheels, ignore S3 cache (default: false) # # Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base + # (currently rocm/dev-ubuntu-22.04:7.1-complete) # # ============================================================================= + # ROCm Input Step - Collect build configuration (manual trigger only) + - input: "ROCm Wheel Release Build Configuration" + key: input-rocm-config + depends_on: ~ + if: build.source == "ui" + fields: + - text: "Python Version" + key: "rocm-python-version" + default: "3.12" + hint: "Python version (e.g., 3.12)" + - text: "GPU Architectures" + key: "rocm-pytorch-rocm-arch" + default: "gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" + hint: "Semicolon-separated GPU architectures" + - select: "Upload Wheels to S3" + key: "rocm-upload-wheels" + default: "true" + options: + - label: "No - Build only (nightly/dev)" + value: "false" + - label: "Yes - Upload to S3 (release)" + value: "true" + - select: "Force Rebuild Base Wheels" + key: "rocm-force-rebuild" + default: "false" + hint: "Ignore S3 cache and rebuild base wheels from scratch" + options: + - label: "No - Use cached wheels if available" + value: "false" + - label: "Yes - Rebuild even if cache exists" + value: "true" + # ROCm Job 1: Build ROCm Base Wheels (with S3 caching) - - label: ":rocm: Build ROCm Base Image & Wheels" + - label: ":rocm: Build ROCm Base Wheels" id: build-rocm-base-wheels - depends_on: ~ + depends_on: + - step: input-rocm-config + allow_failure: true # Allow failure so non-UI builds can proceed (input step is skipped) agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: + # Set configuration and check cache - | set -euo pipefail - # Generate cache key - CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base" + # Get values from meta-data (set by input step) or use defaults + PYTHON_VERSION="$$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo '')" + export PYTHON_VERSION="$${PYTHON_VERSION:-3.12}" + + PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" + export PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" + + # Check for force rebuild flag + ROCM_FORCE_REBUILD="$${ROCM_FORCE_REBUILD:-}" + if [ -z "$${ROCM_FORCE_REBUILD}" ]; then + ROCM_FORCE_REBUILD="$$(buildkite-agent meta-data get rocm-force-rebuild 2>/dev/null || echo '')" + fi echo "========================================" - echo "ROCm Base Build Configuration" + echo "ROCm Base Wheels Build Configuration" echo "========================================" - echo " CACHE_KEY: $${CACHE_KEY}" - echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}" + echo " PYTHON_VERSION: $${PYTHON_VERSION}" + echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}" + echo " ROCM_FORCE_REBUILD: $${ROCM_FORCE_REBUILD:-false}" echo "========================================" - - # Login to ECR - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - IMAGE_EXISTS=false - WHEELS_EXIST=false - - # Check ECR for Docker image - if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then - IMAGE_EXISTS=true - echo "ECR image cache HIT" - fi - - # Check S3 for wheels - WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check) - if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then - WHEELS_EXIST=true - echo "S3 wheels cache HIT" + # Save resolved config for later jobs + buildkite-agent meta-data set "rocm-python-version" "$${PYTHON_VERSION}" + buildkite-agent meta-data set "rocm-pytorch-rocm-arch" "$${PYTORCH_ROCM_ARCH}" + + # Check S3 cache for pre-built wheels + CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + CACHE_PATH=$$(.buildkite/scripts/cache-rocm-base-wheels.sh path) + echo "" + echo "Cache key: $${CACHE_KEY}" + echo "Cache path: $${CACHE_PATH}" + + # Save cache key for downstream jobs + buildkite-agent meta-data set "rocm-cache-key" "$${CACHE_KEY}" + + CACHE_STATUS="miss" + if [ "$${ROCM_FORCE_REBUILD}" != "true" ]; then + CACHE_STATUS=$$(.buildkite/scripts/cache-rocm-base-wheels.sh check) + else + echo "Force rebuild requested, skipping cache check" fi - - # Scenario 1: Both cached (best case) - if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then + if [ "$${CACHE_STATUS}" = "hit" ]; then echo "" - echo "FULL CACHE HIT - Reusing both image and wheels" + echo "CACHE HIT! Downloading pre-built wheels..." echo "" - - # Download wheels .buildkite/scripts/cache-rocm-base-wheels.sh download - - # Save ECR tag for downstream jobs - buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" - - # Scenario 2: Full rebuild needed + + # Set the S3 path for the cached Docker image (for Job 2 to download) + S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}" + buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + + # Mark that we used cache (for Docker image handling) + buildkite-agent meta-data set "rocm-used-cache" "true" + + echo "" + echo "Cache download complete. Skipping Docker build." + echo "Docker image will be downloaded from: $${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" else echo "" - echo " CACHE MISS - Building from scratch..." + echo "CACHE MISS. Building from scratch..." echo "" - - # Build full base image and push to ECR + + # Build full base image (for later vLLM build) DOCKER_BUILDKIT=1 docker buildx build \ --file docker/Dockerfile.rocm_base \ - --tag "$${ECR_CACHE_TAG}" \ + --tag rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} \ + --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ + --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --push \ + --load \ . - - # Build wheel extraction stage + + # Build debs_wheel_release stage for wheel extraction DOCKER_BUILDKIT=1 docker buildx build \ --file docker/Dockerfile.rocm_base \ --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \ --target debs_wheel_release \ + --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ + --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ --load \ . - - # Extract and upload wheels + + # Extract wheels from Docker image mkdir -p artifacts/rocm-base-wheels - cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) - docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/ - docker rm $${cid} - + container_id=$$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) + docker cp $${container_id}:/app/debs/. artifacts/rocm-base-wheels/ + docker rm $${container_id} + echo "Extracted base wheels:" + ls -lh artifacts/rocm-base-wheels/ + + # Upload wheels to S3 cache for future builds + echo "" + echo "Uploading wheels to S3 cache..." .buildkite/scripts/cache-rocm-base-wheels.sh upload - # Cache base docker image to ECR - docker push "$${ECR_CACHE_TAG}" - - buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" - + # Export base Docker image for reuse in vLLM build + mkdir -p artifacts/rocm-docker-image + docker save rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} | gzip > artifacts/rocm-docker-image/rocm-base-image.tar.gz + echo "Docker image size:" + ls -lh artifacts/rocm-docker-image/ + + # Upload large Docker image to S3 (also cached by cache key) + S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}" + echo "Uploading Docker image to $${S3_ARTIFACT_PATH}/" + aws s3 cp artifacts/rocm-docker-image/rocm-base-image.tar.gz "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + + # Save the S3 path for downstream jobs + buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + + # Mark that we did NOT use cache + buildkite-agent meta-data set "rocm-used-cache" "false" + echo "" - echo " Build complete - Image and wheels cached" + echo "Build complete. Wheels cached for future builds." fi - artifact_paths: - "artifacts/rocm-base-wheels/*.whl" env: @@ -595,7 +465,7 @@ steps: - step: build-rocm-base-wheels allow_failure: false agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge timeout_in_minutes: 180 commands: # Download artifacts and prepare Docker image @@ -625,25 +495,31 @@ steps: echo "Downloading wheel artifacts from current build" buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . - # Get ECR image tag from metadata (set by build-rocm-base-wheels) - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rocm-base-image-tag metadata not found" + # Download Docker image from S3 (too large for Buildkite artifacts) + DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')" + if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then + echo "ERROR: rocm-docker-image-s3-path metadata not found" echo "This should have been set by the build-rocm-base-wheels job" exit 1 fi - - echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" - - # Login to ECR - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Pull base Docker image from ECR - docker pull "$${ECR_IMAGE_TAG}" - - echo "Loaded base image: $${ECR_IMAGE_TAG}" - + echo "Downloading Docker image from $${DOCKER_IMAGE_S3_PATH}" + mkdir -p artifacts/rocm-docker-image + aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz + + # Load base Docker image and capture the tag + echo "Loading base Docker image..." + LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load) + echo "$${LOAD_OUTPUT}" + # Extract the actual loaded image tag from "Loaded image: " output + # This avoids picking up stale images (like rocm/vllm-dev:nightly) already on the agent + BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //') + if [ -z "$${BASE_IMAGE_TAG}" ]; then + echo "ERROR: Failed to extract image tag from docker load output" + echo "Load output was: $${LOAD_OUTPUT}" + exit 1 + fi + echo "Loaded base image: $${BASE_IMAGE_TAG}" + # Prepare base wheels for Docker build context mkdir -p docker/context/base-wheels touch docker/context/base-wheels/.keep @@ -651,11 +527,16 @@ steps: echo "Base wheels for vLLM build:" ls -lh docker/context/base-wheels/ + # Get GPU architectures from meta-data + PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" + PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" + echo "========================================" echo "Building vLLM wheel with:" echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" + echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}" + echo " BASE_IMAGE: $${BASE_IMAGE_TAG}" echo "========================================" # Build vLLM wheel using local checkout (REMOTE_VLLM=0) @@ -663,7 +544,8 @@ steps: --file docker/Dockerfile.rocm \ --target export_vllm_wheel_release \ --output type=local,dest=rocm-dist \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ + --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \ + --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ --build-arg REMOTE_VLLM=0 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg USE_SCCACHE=1 \ @@ -671,8 +553,10 @@ steps: --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ . + echo "Built vLLM wheel:" ls -lh rocm-dist/*.whl + # Copy wheel to artifacts directory mkdir -p artifacts/rocm-vllm-wheel cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/ @@ -691,13 +575,35 @@ steps: - step: build-rocm-vllm-wheel allow_failure: false agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge timeout_in_minutes: 60 commands: # Download all wheel artifacts and run upload - | set -euo pipefail + # Check if upload is enabled (from env var, meta-data, or release branch) + ROCM_UPLOAD_WHEELS="$${ROCM_UPLOAD_WHEELS:-}" + if [ -z "$${ROCM_UPLOAD_WHEELS}" ]; then + # Try to get from meta-data (input form) + ROCM_UPLOAD_WHEELS="$$(buildkite-agent meta-data get rocm-upload-wheels 2>/dev/null || echo '')" + fi + + echo "========================================" + echo "Upload check:" + echo " ROCM_UPLOAD_WHEELS: $${ROCM_UPLOAD_WHEELS}" + echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" + echo "========================================" + + # Skip upload if not enabled + if [ "$${ROCM_UPLOAD_WHEELS}" != "true" ]; then + echo "Skipping S3 upload (ROCM_UPLOAD_WHEELS != true, NIGHTLY != 1, not a release branch)" + echo "To enable upload, set 'Upload Wheels to S3' to 'Yes' in the build configuration" + exit 0 + fi + + echo "Upload enabled, proceeding..." + # Download artifacts from current build echo "Downloading artifacts from current build" buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . @@ -713,9 +619,12 @@ steps: - label: ":memo: Annotate ROCm wheel release" id: annotate-rocm-release depends_on: - - upload-rocm-wheels + - step: upload-rocm-wheels + allow_failure: true + - step: input-release-version + allow_failure: true agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "bash .buildkite/scripts/annotate-rocm-release.sh" env: @@ -732,60 +641,61 @@ steps: depends_on: block-generate-root-index-rocm-wheels id: generate-root-index-rocm-wheels agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" env: S3_BUCKET: "vllm-wheels" - VARIANT: "rocm722" + VARIANT: "rocm700" - # ROCm Job 6: Build ROCm Release Docker Image + # ROCm Job 5: Build ROCm Release Docker Image - label: ":docker: Build release image - x86_64 - ROCm" id: build-rocm-release-image depends_on: - - step: block-build-release-images - allow_failure: true - step: build-rocm-base-wheels allow_failure: false agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge timeout_in_minutes: 60 commands: - | set -euo pipefail - + # Login to ECR aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Get ECR image tag from metadata (set by build-rocm-base-wheels) - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rocm-base-image-tag metadata not found" - echo "This should have been set by the build-rocm-base-wheels job" + + # Download Docker image from S3 (set by build-rocm-base-wheels) + DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')" + if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then + echo "ERROR: rocm-docker-image-s3-path metadata not found" exit 1 fi - - echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" - - # Pull base Docker image from ECR - docker pull "$${ECR_IMAGE_TAG}" - - echo "Loaded base image: $${ECR_IMAGE_TAG}" - - # Pass the base image ECR tag to downstream steps (nightly publish) - buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}" - - echo "========================================" - echo "Building vLLM ROCm release image with:" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" - echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" - echo "========================================" - + + echo "Downloading base image from $${DOCKER_IMAGE_S3_PATH}" + mkdir -p artifacts/rocm-docker-image + aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz + + # Load base Docker image + echo "Loading base Docker image..." + LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load) + BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //') + echo "Loaded base image: $${BASE_IMAGE_TAG}" + + # Tag and push the base image to ECR + docker tag "$${BASE_IMAGE_TAG}" public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base + docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base + echo "Pushed base image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base" + + # Get GPU architectures from meta-data + PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" + PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" + # Build vLLM ROCm release image using cached base DOCKER_BUILDKIT=1 docker build \ --build-arg max_jobs=16 \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ + --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \ + --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ @@ -794,90 +704,10 @@ steps: --target vllm-openai \ --progress plain \ -f docker/Dockerfile.rocm . - + # Push to ECR docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm - - echo "" - echo " Successfully built and pushed ROCm release image" - echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" - echo "" + echo "Pushed: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" env: DOCKER_BUILDKIT: "1" S3_BUCKET: "vllm-wheels" - - - label: "Publish nightly ROCm image to DockerHub" - depends_on: - - build-rocm-release-image - if: build.env("NIGHTLY") == "1" - agents: - queue: small_cpu_queue_release - commands: - - "bash .buildkite/scripts/push-nightly-builds-rocm.sh" - # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm" - - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm" - plugins: - - docker-login#v3.0.0: - username: vllmbot - password-env: DOCKERHUB_TOKEN - env: - DOCKER_BUILDKIT: "1" - DOCKERHUB_USERNAME: "vllmbot" - - # ============================================================================= - # Publish to DockerHub and PyPI (at the end so all builds complete first) - # ============================================================================= - - - block: "Publish release images to DockerHub" - key: block-publish-release-images - depends_on: - - create-multi-arch-manifest - - create-multi-arch-manifest-cuda-12-9 - - create-multi-arch-manifest-ubuntu2404 - - create-multi-arch-manifest-cuda-12-9-ubuntu2404 - - build-rocm-release-image - - input-release-version - # Wait for CPU builds if their block steps were unblocked, so publish - # doesn't race the in-progress CPU build. allow_failure lets publish - # proceed when the operator legitimately leaves the CPU block steps - # unblocked or the CPU build fails. - - step: build-cpu-release-image-x86 - allow_failure: true - - step: build-cpu-release-image-arm64 - allow_failure: true - if: build.env("NIGHTLY") != "1" - - - label: "Publish 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" - plugins: - - docker-login#v3.0.0: - username: vllmbot - password-env: DOCKERHUB_TOKEN - env: - DOCKER_BUILDKIT: "1" - DOCKERHUB_USERNAME: "vllmbot" - - - group: "Publish wheels" - key: "publish-wheels" - steps: - - block: "Confirm update release wheels to PyPI (experimental, use with caution)?" - key: block-upload-release-wheels - depends_on: - - input-release-version - - build-wheels - - - label: "Upload release wheels to PyPI" - depends_on: - - block-upload-release-wheels - id: upload-release-wheels - agents: - queue: small_cpu_queue_release - commands: - - "bash .buildkite/scripts/upload-release-wheels-pypi.sh" diff --git a/.buildkite/scripts/annotate-rocm-release.sh b/.buildkite/scripts/annotate-rocm-release.sh index d66129722748..f42f77fc9a85 100755 --- a/.buildkite/scripts/annotate-rocm-release.sh +++ b/.buildkite/scripts/annotate-rocm-release.sh @@ -5,21 +5,20 @@ # Generate Buildkite annotation for ROCm wheel release set -ex -# Extract build configuration from Dockerfile.rocm_base (single source of truth) +# Get build configuration from meta-data # Extract ROCm version dynamically from Dockerfile.rocm_base # BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0" ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown") -PYTHON_VERSION=$(grep '^ARG PYTHON_VERSION=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTHON_VERSION=//') -PYTORCH_ROCM_ARCH=$(grep '^ARG PYTORCH_ROCM_ARCH=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTORCH_ROCM_ARCH=//') +PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12") +PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151") +# TODO: Enable the nightly build for ROCm # Get release version, default to 1.0.0.dev for nightly/per-commit builds RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "") if [ -z "${RELEASE_VERSION}" ]; then RELEASE_VERSION="1.0.0.dev" fi -ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - # S3 URLs S3_BUCKET="${S3_BUCKET:-vllm-wheels}" S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}" @@ -69,7 +68,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl . +aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl . \`\`\` @@ -81,7 +80,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash- - **torchvision**: TorchVision for ROCm PyTorch - **torchaudio**: Torchaudio for ROCm PyTorch - **amdsmi**: AMD SMI Python bindings -- **amd_aiter**: Aiter for ROCm +- **aiter**: Aiter for ROCm - **flash-attn**: Flash Attention for ROCm ### :warning: Notes @@ -97,7 +96,7 @@ To download and upload the image: docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base docker push vllm/vllm-openai-rocm:latest-base diff --git a/CMakeLists.txt b/CMakeLists.txt index f0b1f53af831..d9cd8466470b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -41,7 +41,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") # ROCm installation prefix. Default to /opt/rocm but allow override via # -DROCM_PATH=/your/rocm/path when invoking cmake. @@ -1267,13 +1267,24 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/torch_bindings.cpp" "csrc/rocm/skinny_gemms.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS}) + + # skinny_gemms.cu is built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy + # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude it from the + # gfx1250 build and disable its op registrations (VLLM_SKIP_SKINNY_GEMMS); + # vLLM falls back to default/Triton GEMM for those ops on gfx1250. + if(VLLM_GPU_ARCHES MATCHES "gfx1250") + message(STATUS "gfx1250 detected: excluding csrc/rocm/skinny_gemms.cu from _rocm_C") + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS") + endif() define_extension_target( _rocm_C DESTINATION vllm LANGUAGE ${VLLM_GPU_LANG} SOURCES ${VLLM_ROCM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} + COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) diff --git a/benchmarks/curl_gpt.sh b/benchmarks/curl_gpt.sh new file mode 100644 index 000000000000..51a467bfbc2f --- /dev/null +++ b/benchmarks/curl_gpt.sh @@ -0,0 +1,7 @@ +MODEL="${MODEL:-/data/models/gpt-oss-120b-w-mxfp4-a-fp8}" \ +curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{ + "model": ${MODEL}, + "prompt": "What is the capital of France, and what is it known for?", + "temperature": 0.0, + "max_tokens": 100 +}' diff --git a/benchmarks/gpt_fp4_serve.sh b/benchmarks/gpt_fp4_serve.sh new file mode 100644 index 000000000000..ef23808e95ab --- /dev/null +++ b/benchmarks/gpt_fp4_serve.sh @@ -0,0 +1,7 @@ +#ROCR_VISIBLE_DEVICE=0 \ +MODEL="${MODEL:-/data/models/gpt-oss-120b-w-mxfp4-a-fp8}" \ +HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 \ +VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 \ +VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 \ +VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 \ +vllm serve --model ${MODEL} --host localhost --port 8000 --tensor-parallel-size 1 --gpu_memory_utilization 0.7 #--compilation-config '{"mode":"None","cudagraph_mode": "FULL", "cudagraph_capture_sizes": [1]}' diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh new file mode 100644 index 000000000000..959dc77f1a92 --- /dev/null +++ b/benchmarks/vllm_smoketest.sh @@ -0,0 +1,103 @@ +#!/bin/bash +set -euo pipefail + +# add model args here +declare -A MODEL_PATHS=( + [gpt-oss-120b-mxfp4]="/data/amd/gpt-oss-120b-w-mxfp4-a-fp8" + [DeepSeek-R1-0528-MXFP4]="/data/amd/DeepSeek-R1-0528-MXFP4" +) +declare -A MODEL_SERVE_ARGS=( + [gpt-oss-120b-mxfp4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --attention-backend ROCM_AITER_UNIFIED_ATTN --compilation-config {\"mode\":\"None\",\"cudagraph_mode\":\"FULL\",\"cudagraph_capture_sizes\":[1]}" + [DeepSeek-R1-0528-MXFP4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --dtype auto --no-enable-prefix-caching --disable-uvicorn-access-log --trust-remote-code" +) +declare -A MODEL_ENV=( + [gpt-oss-120b-mxfp4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [DeepSeek-R1-0528-MXFP4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" +) + +MODEL_NAME="gpt-oss-120b-mxfp4" # default +MODEL_PATH_OVERRIDE="" # overrided path +MODEL_SET=0 # whether --model was explicitly passed +RUN_LM_EVAL=0 +PORT=8000 + +usage() { + cat <<'EOF' +Usage: ./vllm_smoketest.sh [--model NAME] [--model-path PATH] [--lm-eval] [--port PORT] [--list] + --model NAME Registry key to run (default: gpt-oss-120b-mxfp4) + --model-path PATH Override model path for this run + --lm-eval Run lm_eval after curl check + --port PORT Server port (default: 8000) + --list List available models +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL_NAME="$2"; MODEL_SET=1; shift 2 ;; + --model-path) MODEL_PATH_OVERRIDE="$2"; shift 2 ;; + --lm-eval) RUN_LM_EVAL=1; shift ;; + --port) PORT="$2"; shift 2 ;; + --list) for n in "${!MODEL_PATHS[@]}"; do echo "$n -> ${MODEL_PATHS[$n]}"; done; exit 0 ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown argument: $1" >&2; usage; exit 1 ;; + esac +done + +if [[ -n "$MODEL_PATH_OVERRIDE" && $MODEL_SET -eq 0 ]]; then + echo "ERROR: --model-path requires --model (which model's serve args/env to use)." >&2 + usage; exit 1 +fi + +MODEL="${MODEL_PATHS[$MODEL_NAME]:-$MODEL_NAME}" +[[ -n "$MODEL_PATH_OVERRIDE" ]] && MODEL="$MODEL_PATH_OVERRIDE" +SERVE_ARGS="${MODEL_SERVE_ARGS[$MODEL_NAME]:-}" +MODEL_ENV_ARGS="${MODEL_ENV[$MODEL_NAME]:-}" +echo "Model: $MODEL | Port: $PORT | lm_eval: $RUN_LM_EVAL" + +# temp +pip3 install "fastapi<0.137" + +LOG_DIR="$(pwd)/vllm_smoke_test_logs" +mkdir -p "$LOG_DIR" +echo "Logs: $LOG_DIR" + +server_pid="" +cleanup() { [[ -n "$server_pid" ]] && kill "$server_pid" 2>/dev/null && wait "$server_pid" 2>/dev/null; true; } +trap cleanup EXIT +trap 'exit 130' INT TERM + +# --- start server --- +env $MODEL_ENV_ARGS \ +vllm serve --model "$MODEL" --host localhost --port "$PORT" $SERVE_ARGS \ + > "$LOG_DIR/server.log" 2>&1 & +server_pid=$! +echo "Waiting for server (pid $server_pid)..." + +until curl -s "http://localhost:$PORT/health" &>/dev/null; do + ps -p "$server_pid" >/dev/null || { echo "ERROR: server died"; tail -n 50 "$LOG_DIR/server.log"; exit 1; } + sleep 5 +done +echo "Server ready." + +# --- check if text was generated --- +echo "=== Curl completion check ===" +http_code=$(curl -s -o "$LOG_DIR/curl_completion.log" -w '%{http_code}' \ + "http://localhost:$PORT/v1/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL\",\"prompt\":\"San Francisco is a\",\"max_tokens\":32,\"temperature\":0}") +cat "$LOG_DIR/curl_completion.log"; echo +[[ "$http_code" == "200" ]] || { echo "FAIL: HTTP $http_code" >&2; exit 1; } +grep -q '"text"' "$LOG_DIR/curl_completion.log" || { echo "FAIL: no text in response" >&2; exit 1; } +echo "PASS" + +# --- lm_eval --- +if [[ $RUN_LM_EVAL -eq 1 ]]; then + echo "=== lm_eval ===" + pip install "lm-eval[api]" + lm_eval --model local-completions \ + --model_args "model=$MODEL,base_url=http://localhost:$PORT/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False" \ + --tasks gsm8k --num_fewshot 3 --limit 100 2>&1 | tee "$LOG_DIR/lm_eval.log" +fi + +echo "Done." diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index a2170e483207..d047e21a6df5 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -81,11 +81,11 @@ union BufferResource { __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( int32x4_t srsrc, int32_t voffset, int32_t soffset, - int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32"); + int32_t aux) {} __quickreduce_device_inline__ static void buffer_store_dwordx4( int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, - int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32"); + int32_t aux) {} __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) { #if defined(__gfx942__) diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 9e6c0726d19e..5d861f63d8d9 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -2405,6 +2405,15 @@ template __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, const bit16x8& inpB, const floatx8& inpC) { +#if defined(__gfx1250__) + // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by + // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is + // unsupported on gfx1250; trap if ever launched (fail loud, not silent-wrong). + (void)inpA; + (void)inpB; + __builtin_trap(); + return inpC; +#else if constexpr (std::is_same::value) { return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); } else if constexpr (std::is_same::value) { @@ -2412,6 +2421,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, } else { static_assert(false, "unsupported 16b dtype"); } +#endif } template diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index b0b44964c243..9666428f7c2e 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -14,6 +14,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { // vLLM custom ops for rocm +// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250 +// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid +// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250. +#ifndef VLLM_SKIP_SKINNY_GEMMS // Custom gemm op for matrix-vector multiplication rocm_ops.def( "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> " @@ -38,6 +42,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "Tensor scale_a, " " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#endif // VLLM_SKIP_SKINNY_GEMMS // Custom attention op // Compute the attention between an input query and the cached diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index b2342200a682..7188d53e5ad3 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,7 +1,8 @@ # default base image ARG REMOTE_VLLM="0" ARG COMMON_WORKDIR=/app -ARG BASE_IMAGE=rocm/vllm-dev:base +ARG BASE_IMAGE=rocm/vllm-private:juan_455_npi_base +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. @@ -16,7 +17,8 @@ ARG NIC_BACKEND=all ARG AINIC_VERSION=1.117.3-hydra ARG UBUNTU_CODENAME=jammy -# Sccache configuration (only used in release pipeline) +# Sccache configuration. Release builds use this today; CI can opt in when a +# shared S3-compatible cache backend is available. ARG USE_SCCACHE ARG SCCACHE_DOWNLOAD_URL ARG SCCACHE_ENDPOINT @@ -26,15 +28,19 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG ARG_PYTORCH_ROCM_ARCH +ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} -# Install some basic utilities +# Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ - libnuma-dev -RUN python3 -m pip install --upgrade pip + build-essential libnuma-dev ccache mold +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --upgrade pip +# Note: mold is installed but not set as the system default linker because +# some packages use JIT compilation at runtime with flags mold does not support. +# Build stages opt in via LDFLAGS="-fuse-ld=mold". # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE RUN if [ "$USE_SCCACHE" != "1" ]; then \ @@ -55,6 +61,16 @@ ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy +# python binary fall back for non venv builds +ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 +# Expose paths from wheel installation +ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} +# ccache directory - persisted across layer rebuilds via cache mounts. +ENV CCACHE_DIR=/root/.cache/ccache +ENV CCACHE_COMPILERCHECK=content +# Empty by default so build steps fall back to $(nproc); CI can override. +ARG max_jobs +ENV MAX_JOBS=${max_jobs} # Install sccache if USE_SCCACHE is enabled (for release builds) ARG USE_SCCACHE @@ -98,8 +114,8 @@ WORKDIR ${COMMON_WORKDIR} FROM base AS fetch_vllm_0 ONBUILD COPY ./ vllm/ FROM base AS fetch_vllm_1 -ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" -ARG VLLM_BRANCH="main" +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip_nowhls" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ @@ -114,8 +130,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm # ----------------------- # Rust build stage # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain or protoc. Runs in parallel with the main wheel -# build for faster end-to-end builds. +# don't need the rust toolchain or protoc. FROM fetch_vllm AS rust-build ARG COMMON_WORKDIR @@ -135,28 +150,83 @@ ENV PATH="/root/.cargo/bin:${PATH}" # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 - -# Build the release binary. Cache cargo registry/git, and copy the binary out -# so it persists into the image layer for later COPY --from=rust-build. -RUN --mount=type=cache,target=/root/.cargo/registry \ - --mount=type=cache,target=/root/.cargo/git \ +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 + +# Build the release binary. Cargo's registry/git caches can be written by +# concurrent BuildKit jobs on shared workers, so lock those cache mounts while +# keeping the cache benefit. Copy the binary out so it persists into the image +# layer for later COPY --from=rust-build. +RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ - && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh + && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \ + && test -x /tmp/vllm-rs # ----------------------- -# vLLM build stages +# vLLM native build stages +# +# csrc-build intentionally copies only files that affect ROCm native extension +# compilation. That keeps unrelated CI/test/docs edits from invalidating the +# expensive HIP/C++ build layer. +FROM base AS csrc-build +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR}/vllm + +COPY requirements/rocm.txt requirements/rocm.txt +COPY requirements/common.txt requirements/common.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/rocm.txt + +# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do +# not invalidate the expensive native build layer. +COPY setup.py CMakeLists.txt ./ +COPY cmake cmake/ +COPY csrc csrc/ +COPY vllm/envs.py vllm/envs.py +COPY vllm/__init__.py vllm/__init__.py + +ENV VLLM_TARGET_DEVICE=rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" + +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + export CCACHE_BASEDIR="$PWD" \ + && echo "=== ccache stats before ROCm native build ===" \ + && (ccache --show-stats || true) \ + && (ccache --zero-stats || true) \ + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null \ + && echo "=== ccache stats after ROCm native build ===" \ + && (ccache --show-stats || true) + +# Build the full vLLM ROCm wheel by reusing the native extension wheel from +# csrc-build. This stage still rebuilds for Python/package changes, but skips +# the expensive HIP/C++ compile when native inputs are unchanged. FROM fetch_vllm AS build_vllm ARG COMMON_WORKDIR +ENV VLLM_TARGET_DEVICE=rocm + +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels # Drop the pre-built rust frontend binary into the source tree. setup.py # detects it and ships it as-is, skipping the local cargo build. COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -# Build vLLM (setup.py auto-detects sccache in PATH) -RUN cd vllm \ - && python3 -m pip install -r requirements/rocm.txt \ - && python3 setup.py clean --all \ - && python3 setup.py bdist_wheel --dist-dir=dist +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system -r requirements/rocm.txt \ + && export VLLM_USE_PRECOMPILED=1 \ + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ + && export VLLM_DOCKER_BUILD_CONTEXT=1 \ + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null FROM scratch AS export_vllm ARG COMMON_WORKDIR COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / @@ -166,6 +236,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -174,7 +245,7 @@ ARG RIXL_BRANCH="39be1de8" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" ARG UCX_BRANCH="bfb51733" ARG UCX_REPO="https://github.com/openucx/ucx.git" -ENV ROCM_PATH=/opt/rocm +# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl ENV RIXL_BENCH_HOME=/usr/local/rixl_bench @@ -196,14 +267,17 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ ibverbs-providers \ && rm -rf /var/lib/apt/lists/* -RUN uv pip install --system meson auditwheel patchelf tomlkit +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system meson auditwheel patchelf tomlkit -RUN cd /usr/local/src && \ +RUN --mount=type=cache,target=/root/.cache/ccache \ + cd /usr/local/src && \ git clone ${UCX_REPO} && \ cd ucx && \ git checkout ${UCX_BRANCH} && \ ./autogen.sh && \ mkdir build && cd build && \ + CC="ccache gcc" CXX="ccache g++" \ ../configure \ --prefix=/usr/local/ucx \ --enable-shared \ @@ -215,20 +289,23 @@ RUN cd /usr/local/src && \ --with-verbs \ --with-dm \ --enable-mt && \ - make -j && \ + make -j$(nproc) && \ make install ENV PATH=/usr/local/ucx/bin:$PATH ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} -RUN git clone ${RIXL_REPO} /opt/rixl && \ +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone ${RIXL_REPO} /opt/rixl && \ cd /opt/rixl && \ git checkout ${RIXL_BRANCH} && \ + CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ + --force-fallback-for=abseil-cpp \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ - ninja && \ + ninja -j$(nproc) && \ ninja install # Generate RIXL wheel @@ -237,6 +314,10 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \ RUN cd /opt/rixl && \ sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ contrib/build-wheel.sh && \ + # The wheel build re-runs meson via meson-python; force the bundled abseil + sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ + pyproject.toml && \ + grep -q 'force-fallback-for' pyproject.toml && \ mkdir -p /app/install && \ _ucx_install_dir=${UCX_HOME} \ ./contrib/build-wheel.sh \ @@ -245,30 +326,44 @@ RUN cd /opt/rixl && \ --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins -# DeepEP build stage -FROM base AS build_deep +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not +# invalidate the slow ROCShmem build. +FROM base AS build_rocshmem ARG ROCSHMEM_BRANCH="f0acb0c6" ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -ARG DEEPEP_BRANCH="a9ea9774" -ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" -ARG DEEPEP_NIC="cx7" +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so +# it can be linked against DeepEP without arch mismatches. ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" +# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image ENV ROCSHMEM_DIR=/opt/rocshmem -RUN git clone ${ROCSHMEM_REPO} \ +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ && git checkout ${ROCSHMEM_BRANCH} \ && mkdir -p projects/rocshmem/build \ && cd projects/rocshmem/build \ - && INSTALL_PREFIX=${ROCSHMEM_DIR} \ - ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + bash ../scripts/build_configs/all_backends \ + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF + +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. +FROM build_rocshmem AS build_deepep +ARG DEEPEP_BRANCH="a9ea9774" +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" +ARG DEEPEP_NIC="cx7" -# Build DeepEP wheel. -# DeepEP looks for rocshmem at ROCSHMEM_DIR. -RUN git clone ${DEEPEP_REPO} \ +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. +RUN --mount=type=cache,target=/root/.cache/ccache \ + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ + && git clone ${DEEPEP_REPO} \ && cd DeepEP \ && git checkout ${DEEPEP_BRANCH} \ - && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do # not force users to rebuild the long-lived Dockerfile.rocm_base image. @@ -367,8 +462,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) # This ensures setuptools_scm sees clean repo state for version detection RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,target=/root/.cache/uv \ cd vllm \ - && pip install setuptools_scm regex \ + && uv pip install --system setuptools_scm regex \ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ && echo "Detected vLLM version: ${VLLM_VERSION}" \ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt @@ -404,18 +500,20 @@ RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt # Install dependencies using custom wheels from /install -RUN cd vllm \ +RUN --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ && echo "Building vLLM with custom wheels from /install" \ - && python3 -m pip install --find-links /install -r requirements/rocm.txt \ - && python3 setup.py clean --all + && uv pip install --system --find-links /install -r requirements/rocm.txt # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt -# (setup.py auto-detects sccache in PATH) +# (setup.py auto-detects ccache/sccache in PATH) RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ + && export CCACHE_BASEDIR="$PWD" \ && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ - && python3 setup.py bdist_wheel --dist-dir=dist + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist FROM scratch AS export_vllm_wheel_release ARG COMMON_WORKDIR @@ -426,112 +524,118 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- -# Test vLLM image -FROM mori_base AS test - -RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* - -# Install vLLM using uv (inherited from base stage) -# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && uv pip install --system -r requirements/rocm.txt \ - && uv pip install --system -r requirements/test/rocm.txt \ - && pip uninstall -y vllm \ - && uv pip install --system *.whl - -# Persist the built wheel in the image so python_only_compile_rocm.sh can -# reinstall it after removing compilers. The bind-mounted /install contents -# above are not available once that RUN step completes. -COPY --from=export_vllm /*.whl /opt/vllm-wheels/ +# CI base image (Tier 1) - stable, rarely changing CI dependencies. +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers +# in the vLLM artifacts for the current commit. +FROM mori_base AS ci_base +ARG COMMON_WORKDIR -# Update rdma-core to support latest rocshmem +# Update rdma-core to support latest rocshmem. ARG DEEPEP_NIC RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ cd /tmp/rdma-core && \ mkdir -p build && cd build && \ cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ - ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ fi -# Install RIXL wheel +# Install RIXL + DeepEP wheels. RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ - uv pip install --system /rixl_install/*.whl + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ + uv pip install --system /rixl_install/*.whl /deep_install/*.whl -# Install DeepEP wheel -RUN --mount=type=bind,from=build_deep,src=/app/deep_install,target=/deep_install \ - uv pip install --system /deep_install/*.whl -COPY --from=build_deep /opt/rocshmem /opt/rocshmem +# Copy ROCShmem runtime libraries. +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem -# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries) -RUN apt-get update -q -y && apt-get install -q -y \ +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ librdmacm1 \ libibverbs1 \ ibverbs-providers \ ibverbs-utils \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* -WORKDIR /vllm-workspace -ARG COMMON_WORKDIR -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace - -# install development dependencies (for testing) -RUN cd /vllm-workspace \ - && python3 -m pip install -e tests/vllm_test_utils \ - && python3 -m pip install pytest-shard - -# enable fast downloads from hf (for testing) -ENV HF_XET_HIGH_PERFORMANCE=1 - -# increase timeout for hf downloads (for testing) -ENV HF_HUB_DOWNLOAD_TIMEOUT 60 - -# install audio decode package `torchcodec` from source (required due to -# ROCm and torch version mismatch) for tests with datasets package +# Install torchcodec from source for ROCm/torch ABI compatibility. COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh -RUN bash /tmp/install_torchcodec.sh \ +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + bash /tmp/install_torchcodec.sh \ && rm /tmp/install_torchcodec.sh \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* + && apt-get clean && rm -rf /var/lib/apt/lists/* -# Copy in the v1 package (for python-only install test group) -COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 +# Pre-install shared ROCm runtime dependencies. +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ + && rm -rf /tmp/ci-base-requirements + +# Enable fast and less brittle model downloads in tests. +ENV HF_XET_HIGH_PERFORMANCE=1 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60 -# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel +# Pre-install vLLM test dependencies. +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/rocm-test-reqs.txt + +# Rebuild fastsafetensors from source so its C++ extension is compiled with +# USE_ROCM and can detect libamdhip64.so at runtime. +RUN --mount=type=cache,target=/root/.cache/pip \ + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ + && test -n "${FASTSAFETENSORS_REQ}" \ + && python3 -m pip install --force-reinstall --no-deps \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + && rm /tmp/rocm-test-reqs.txt + +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # See: https://github.com/pytorch/pytorch/issues/169857 ENV MIOPEN_DEBUG_CONV_DIRECT=0 ENV MIOPEN_DEBUG_CONV_GEMM=0 -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. # See: https://github.com/ROCm/rocm-libraries/issues/6266 ENV HSA_ENABLE_IPC_MODE_LEGACY=1 -# Source code is used in the `python_only_compile.sh` test -# We hide it inside `src/` so that this source code -# will not be imported by other tests -RUN mkdir src && mv vllm src/vllm +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -# This is a workaround to ensure pytest exits with the correct status code in CI tests. -RUN printf '%s\n' \ - 'import os' \ - '' \ - '_exit_code = 1' \ - '' \ - 'def pytest_sessionfinish(session, exitstatus):' \ - ' global _exit_code' \ - ' _exit_code = int(exitstatus)' \ - '' \ - 'def pytest_unconfigure(config):' \ - ' import sys' \ - ' sys.stdout.flush()' \ - ' sys.stderr.flush()' \ - ' os._exit(_exit_code)' \ - > /vllm-workspace/conftest.py +# Install vllm_test_utils in ci_base for ci_base + wheel parity. +COPY tests/vllm_test_utils /tmp/vllm_test_utils +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /tmp/vllm_test_utils \ + && rm -rf /tmp/vllm_test_utils + +# ----------------------- +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. +FROM ${CI_BASE_IMAGE} AS test +ARG COMMON_WORKDIR + +# Install the vLLM wheel (--no-deps: all deps already in ci_base). +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system --no-deps *.whl + +# Store the vLLM wheel in the image for python-only install tests. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ + +WORKDIR /vllm-workspace +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# Copy in the v1 package (for python-only install test group). +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 + +# Hide source under src/ so it won't shadow the installed package in tests. +RUN mkdir src && mv vllm src/vllm # ----------------------- # Final vLLM image @@ -548,6 +652,7 @@ RUN rm -f /usr/bin/sccache || true \ # This prevents S3 bucket config from leaking into production images ENV SCCACHE_BUCKET= ENV SCCACHE_REGION= +ENV SCCACHE_ENDPOINT= ENV SCCACHE_S3_NO_CREDENTIALS= ENV SCCACHE_IDLE_TIMEOUT= diff --git a/docker/Dockerfile.rocm_1250 b/docker/Dockerfile.rocm_1250 new file mode 100644 index 000000000000..697d92dda512 --- /dev/null +++ b/docker/Dockerfile.rocm_1250 @@ -0,0 +1,114 @@ +FROM ubuntu:24.04 +ENV DEBIAN_FRONTEND=noninteractive + +# --- ROCm + PyTorch versioning (override any of these at build time) ------- +ARG GFX_ARCH=gfx1250 +ARG ROCM_VERSION=7.14.0a20260604 +ARG ROCM_BUILD=a0 +ARG TORCH_VERSION=2.10.0 +ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250/ +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" +ARG AITER_BRANCH="jpvillam/gfx1250_0604" + +ENV VLLM_TARGET_DEVICE=rocm +ENV PYTORCH_ROCM_ARCH=${GFX_ARCH} +ENV CMAKE_BUILD_TYPE=Release +ENV GPU_TARGET=${GFX_ARCH} +ENV VIRTUAL_ENV=/opt/venv + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + build-essential clang-19 lld ccache ninja-build cmake \ + git curl libcurl4-openssl-dev rsync ssh wget ca-certificates \ + python3 python3-pip python3-dev python3-venv \ + less vim libzstd-dev numactl libelf1 m4 && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + +SHELL ["/bin/bash", "-e", "-u", "-o", "pipefail", "-c"] + +RUN python3 -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH + +# COPY ffm/ /root/x/ffm/ +# ENV PATH=${ROCM_PATH}/llvm/bin:${ROCM_PATH}/bin:$PATH +# ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${ROCM_PATH}/lib/rocm_sysdeps/lib:${ROCM_PATH}/llvm/lib +# ENV HSA_MODEL_LIB=/root/x/ffm/libhsakmtmodel.so +# ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" +# ENV HSA_MODEL_ARGS=ffm_enable_time_slicing +# ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 +# ENV HSA_MODEL_NUM_THREADS=64 + +ENV HSA_ENABLE_SDMA=0 +ENV HSA_ENABLE_INTERRUPT=0 + +# Install the TheRock PyTorch wheel and the matching rocm-sdk wheels into the venv +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + "torch==${TORCH_VERSION}+rocm${ROCM_VERSION}${ROCM_BUILD:+.${ROCM_BUILD}}" \ + "rocm[libraries,devel]==${ROCM_VERSION}${ROCM_BUILD:++${ROCM_BUILD}}" && \ + rocm-sdk init + +RUN pip install \ + filelock \ + "typing-extensions>=4.10.0" \ + "sympy>=1.13.3" \ + "networkx>=2.5.1" \ + jinja2 \ + "fsspec>=0.8.5" + +# Set ROCm, Python, Paths, and Environment Variables +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python3.12/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib:${SITE_PACKAGES}/_rocm_sdk_libraries_${GFX_ARCH}/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# NPI base-image fix-ups so the wheel ROCm behaves like a normal install: +# - build amd_smi from the bundled source (provides the `amdsmi` module), +# - flip rocm_sdk's library preload off RTLD_GLOBAL (avoids symbol clashes), +RUN if [ -d "${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi" ]; then \ + cd "${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi" && pip install .; \ + fi && \ + if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi + +# Install AITER (https://github.com/ROCm/aiter): plain clone WITHOUT submodules +# (so 3rdparty/composable_kernel is absent), built with Composable Kernel disabled. +RUN git clone -b ${AITER_BRANCH} --depth 1 https://github.com/ROCm/aiter.git /app/aiter && \ + cd /app/aiter && \ + pip install packaging "cmake<4" ninja wheel setuptools_scm vcs_versioning pybind11 numpy && \ + ENABLE_CK=0 AITER_USE_SYSTEM_TRITON=1 pip install --no-build-isolation -e . + + +RUN apt-get update && \ + apt-get install -y --no-install-recommends pkg-config libnuma-dev libdrm-dev && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + +RUN pip install "setuptools>=77.0.3,<81.0.0" setuptools_scm setuptools_rust wheel packaging jinja2 + +# Build + install vLLM for gfx1250 (cloned from VLLM_REPO @ VLLM_BRANCH) +RUN git clone "${VLLM_REPO}" /app/vllm \ + && cd /app/vllm \ + && git fetch -v --prune -- origin "${VLLM_BRANCH}" \ + && git checkout FETCH_HEAD + +# Compile the C++/HIP extensions (_C, _rocm_C) for gfx1250 and editable-install. +# MAX_JOBS is a build-arg so it can be tuned per machine (docker build --build-arg). +RUN cd /app/vllm && \ + pip install -e . --no-build-isolation --no-deps -v + +# vLLM Python runtime deps. common.txt pins no torch, so the NPI ROCm torch +# installed above is left untouched. +RUN pip install -r /app/vllm/requirements/common.txt + +WORKDIR /app/vllm +ENTRYPOINT ["/bin/bash"] \ No newline at end of file diff --git a/docker/Dockerfile.rocm_1250_ffm b/docker/Dockerfile.rocm_1250_ffm new file mode 100644 index 000000000000..ada7e0ee4eb6 --- /dev/null +++ b/docker/Dockerfile.rocm_1250_ffm @@ -0,0 +1,11 @@ +# FFM overlay for gfx1250 builds +ARG BASE_IMAGE=rocm/vllm-private:juan_455_npi_test +FROM ${BASE_IMAGE} + +# extract tarball into /root/x/ffm/. +ADD ubuntu_24_04_rel.tar.gz /root/x/ffm/ +ENV HSA_MODEL_LIB=/root/x/ffm/libhsakmtmodel.so +ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" +ENV HSA_MODEL_ARGS=ffm_enable_time_slicing +ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 +ENV HSA_MODEL_NUM_THREADS=64 diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index cb6d09f5bb6c..57ac22394e4c 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,15 +1,12 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete -ARG TRITON_BRANCH="ba5c1517" -ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17 -ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.24.1" -ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG PYTORCH_AUDIO_BRANCH="v2.9.0" -ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" +ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.10.0-rocm7.14.0a20260605.a0-7.14.0a20260605-a0-nightly-ubuntu24.04-gfx1250 +ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250 +ARG TORCH_VERSION=2.10.0+rocm7.14.0a20260605.a0 +ARG TORCHVISION_VERSION=0.28.0a0+rocm7.14.0a20260605 +ARG TORCHAUDIO_VERSION=2.11.0a0+rocm7.14.0a20260605 +ARG ROCM_SDK_VERSION=7.14.0a20260605 ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.13" +ARG AITER_BRANCH="jpvillam/gfx1250_0604" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -24,14 +21,15 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -ENV ROCM_PATH=/opt/rocm -ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 +ARG PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx942;gfx950 +ENV AITER_ROCM_ARCH=gfx1250 ENV MORI_GPU_ARCHS=gfx942;gfx950 +# TODO: Unset these when support is available for gfx1250 +ENV ENABLE_CK=0 +ARG PREBUILD_KERNELS=0 + # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 @@ -46,17 +44,21 @@ ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update -y \ && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ done \ && apt-get update -y \ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version && python3 -m pip --version + && python3 --version + +ENV VIRTUAL_ENV=/opt/venv +RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* @@ -69,16 +71,61 @@ ARG SCCACHE_BUCKET_NAME ARG SCCACHE_REGION_NAME ARG SCCACHE_S3_NO_CREDENTIALS RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi + +### +### Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +### +#ARG ROCM_WHEEL_INDEX +#ARG TORCH_VERSION +#ARG TORCHVISION_VERSION +#ARG TORCHAUDIO_VERSION +#ARG ROCM_SDK_VERSION +## torch/torchvision/torchaudio must be pinned to mutually-consistent builds +## (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +## sdk version is derived from torch's own dependency pin unless overridden, +## which keeps the set consistent and avoids pip backtracking. +#RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ +# "torch==${TORCH_VERSION}" \ +# "torchvision==${TORCHVISION_VERSION}" \ +# "torchaudio==${TORCHAUDIO_VERSION}" \ +# "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ +# rocm-sdk init + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream +# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm; + +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ fi # Setup sccache for HIP compilation via HIP_CLANG_PATH @@ -87,13 +134,13 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ fi # Set sccache environment variables only when USE_SCCACHE=1 @@ -104,112 +151,13 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} -# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062 -# will be removed once we move to ROCm 7.2.3 -RUN apt-get update && apt-get install -y rocm-llvm-dev -RUN pip install CppHeaderParser -RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \ - && cd /tmp/rocm-systems \ - && git sparse-checkout init --cone \ - && git sparse-checkout set projects/hip projects/clr \ - && git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \ - && export CLR_DIR=/tmp/rocm-systems/projects/clr \ - && export HIP_DIR=/tmp/rocm-systems/projects/hip \ - && mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \ - && cmake \ - -DHIP_COMMON_DIR=$HIP_DIR \ - -DCMAKE_PREFIX_PATH="/opt/rocm/" \ - -DCLR_BUILD_HIP=ON \ - -DCLR_BUILD_OCL=OFF \ - -DHIP_PLATFORM=amd \ - .. \ - && make -j$(nproc) \ - && make install \ - && rm -rf /tmp/rocm-systems - -### -### Triton Build -### -FROM base AS build_triton -ARG TRITON_BRANCH -ARG TRITON_REPO -RUN git clone ${TRITON_REPO} -# Cherry picking the following -# https://github.com/triton-lang/triton/pull/8991 -# https://github.com/triton-lang/triton/pull/9541 -RUN cd triton \ - && git checkout ${TRITON_BRANCH} \ - && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ - && git cherry-pick 555d04f \ - && git cherry-pick dd998b6 \ - && if [ ! -f setup.py ]; then cd python; fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install -RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \ - && python3 -m build --wheel && cp dist/*.whl /app/install; fi - - ### ### AMD SMI Build ### FROM base AS build_amdsmi -RUN cd /opt/rocm/share/amd_smi \ +RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install - - -### -### Pytorch build -### -FROM base AS build_pytorch -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_REPO -ARG USE_SCCACHE - -RUN apt-get update && apt-get install -y pkg-config liblzma-dev -RUN git clone ${PYTORCH_REPO} pytorch -RUN cd pytorch && git checkout ${PYTORCH_BRANCH} -RUN cd pytorch \ - && pip install -r requirements.txt && git submodule update --init --recursive -RUN cd pytorch && python3 tools/amd_build/build_amd.py \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ - && sccache --show-stats; \ - fi \ - && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN git clone ${PYTORCH_VISION_REPO} vision -RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN git clone ${PYTORCH_AUDIO_REPO} audio -RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ - && git submodule update --init --recursive \ - && pip install -r requirements.txt \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ - && cp /app/vision/dist/*.whl /app/install \ - && cp /app/audio/dist/*.whl /app/install +RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install ### @@ -218,14 +166,17 @@ RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${MORI_REPO} -RUN cd mori \ +RUN mkdir -p /app/install; \ + if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ + echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ + else \ + git clone ${MORI_REPO} \ + && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ - && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl -RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ + && cp /app/mori/dist/*.whl /app/install; \ + fi ### @@ -235,19 +186,22 @@ FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${FA_REPO} -RUN cd flash-attention \ +RUN mkdir -p /app/install; \ + if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ + echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping FlashAttention build"; \ + else \ + git clone ${FA_REPO} \ + && cd flash-attention \ && git checkout ${FA_BRANCH} \ && git submodule update --init \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && cp dist/*.whl /app/install; \ + fi ### @@ -257,19 +211,16 @@ FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone --recursive ${AITER_REPO} +RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} RUN cd aiter \ - && git checkout ${AITER_BRANCH} \ && git submodule update --init --recursive \ && pip install -r requirements.txt RUN pip install pyyaml && cd aiter \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install @@ -283,46 +234,35 @@ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install # only includes dependencies used by wheel release pipeline FROM base AS debs_wheel_release RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs # Full debs stage - includes Mori (used by Docker releases) FROM base AS debs RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi FROM base AS final RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ pip install /install/*.whl ARG BASE_IMAGE -ARG TRITON_BRANCH -ARG TRITON_REPO -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO +ARG ROCM_WHEEL_INDEX +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH @@ -330,14 +270,11 @@ ARG AITER_REPO ARG MORI_BRANCH ARG MORI_REPO RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ - && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ + && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 589d90d1eca7..9510d2cc552e 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -4,6 +4,12 @@ Test modular OAI Triton MoE """ +from __future__ import annotations + +import json +import os +from pathlib import Path + import pytest import torch @@ -48,6 +54,52 @@ ] +def deepseek_v4_flash_moe_topology(): + """MoE sizes for the MODEL path used by ``launch_dsv4.sh``. + + Default weights path: ``/data/deepseek-ai/DeepSeek-V4-Flash/config.json``. + When that file is readable, values come from ``hidden_size``, + ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``. + Otherwise fall back to the same numeric constants. + """ + defaults = { + "hidden_size": 4096, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + } + cfg_path = Path( + os.environ.get( + "DEEPSEEK_V4_FLASH_CONFIG", + "/data/deepseek-ai/DeepSeek-V4-Flash/config.json", + ) + ) + if cfg_path.is_file(): + with cfg_path.open() as f: + cfg = json.load(f) + return { + "hidden_size": int(cfg["hidden_size"]), + "moe_intermediate_size": int(cfg["moe_intermediate_size"]), + "n_routed_experts": int(cfg["n_routed_experts"]), + "num_experts_per_tok": int(cfg["num_experts_per_tok"]), + } + return defaults + + +def scaled_deepseek_v4_flash_problem( + *, + dim_scale: int = 8, + expert_scale: int = 8, +): + """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1).""" + t = deepseek_v4_flash_moe_topology() + k = max(128, t["hidden_size"] // dim_scale) + n = max(64, t["moe_intermediate_size"] // dim_scale) + num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale) + topk = t["num_experts_per_tok"] + return k, n, num_experts, topk + + def unshuffle_weight(w: torch.Tensor): first = w[..., ::2] second = w[..., 1::2] @@ -262,3 +314,95 @@ def test_oai_triton_moe( ) assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): + """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. + + Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``, + with linear dimensions and expert count scaled down for test GPU memory. + """ + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) + set_random_seed(0) + + k, n, num_experts, topk = scaled_deepseek_v4_flash_problem() + m = 7 + dtype = torch.bfloat16 + + ( + w1, + w2, + w1_bias, + w2_bias, + w1_tri, + w2_tri, + w1_bias_tri, + w2_bias_tri, + w1_precision_config, + w2_precision_config, + ) = make_weights(dtype, k, n, num_experts) + + x = torch.randn((m, k), dtype=dtype, device="cuda") + router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) + + quant_config = mxfp4_w4a16_moe_quant_config( + w1_bias=w1_bias_tri, + w2_bias=w2_bias_tri, + w1_scale=w1_precision_config, + w2_scale=w2_precision_config, + ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size_per_partition=n, + ) + experts = UnfusedOAITritonExperts(moe_config, quant_config) + + if not UnfusedOAITritonExperts._supports_current_device(): + pytest.skip("UnfusedOAITritonExperts does not support this device") + + _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids) + assert top_k == topk + ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( + m, + N, + K, + topk, + num_experts, + num_experts, + None, + MoEActivation.SWIGLUOAI, + ) + workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda") + workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda") + output = torch.empty(out_shape, dtype=dtype, device="cuda") + + with set_current_vllm_config(VllmConfig()): + out_ref = torch_moe_impl( + x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids + ) + experts.apply( + output=output, + hidden_states=x, + w1=w1_tri, + w2=w2_tri, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SWIGLUOAI, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index ce4fc3cfbadb..f739bf901606 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -6,6 +6,7 @@ from typing import Protocol import torch +import torch.nn.functional as F from torch._ops import OpOverload from torch.distributed import ProcessGroup @@ -124,6 +125,143 @@ def wrapper(*args, **kwargs): return wrapper +def _mxfp4_moe_w1_triton_gemm_view(w1: torch.Tensor, w2: torch.Tensor) -> torch.Tensor: + """Map expert w1 to moe_gemm_a4w4 layout (stride(-2)==1 on [E, K, N]). + + vLLM stores w13 as [E, 2*inter, hidden/2] (K along dim2). aiter downcast uses the + same axis order with K-packed stride-1 on dim1; do not compose transpose+as_strided + (that overruns storage). Standard layout only needs transpose(1, 2). + """ + if w1.stride(-2) == 1: + return w1 + h = w2.shape[1] + e, a, b = w1.shape + if b * 2 == h: + return w1.transpose(1, 2) + if a * 2 == h: + kp, n_out = a, b + return w1.as_strided((e, kp, n_out), (kp * n_out, 1, kp)) + return w1.transpose(1, 2) + + +def _mxfp4_moe_w2_triton_gemm_view(w2: torch.Tensor) -> torch.Tensor: + """Map vLLM w2 [E, model_dim, inter/2] to moe_gemm_a4w4 layout.""" + if w2.stride(-2) == 1: + return w2 + return w2.transpose(1, 2) + + +def _silu_and_mul_glu(x: torch.Tensor) -> torch.Tensor: + """SiLU(gate) * up for g1u1 tensors with last dim 2 * inter (matches CK ``silu_and_mul``).""" + d = x.shape[-1] // 2 + gate, up = x[..., :d], x[..., d:] + return (F.silu(gate.to(torch.float32)) * up.to(torch.float32)).to(dtype=x.dtype) + + +def _mxfp4_moe_weight_as_uint8(w: torch.Tensor) -> torch.Tensor: + """Triton moe_gemm_a4w4 JIT cannot canonicalize torch float4 dtypes (KeyError).""" + if w.dtype == torch.uint8: + return w + fp4 = getattr(torch, "float4_e2m1fn_x2", None) + if fp4 is not None and w.dtype == fp4: + return w.view(torch.uint8) + name = getattr(w.dtype, "name", str(w.dtype)) + if "float4" in name or "e2m1" in name: + return w.view(torch.uint8) + return w + + +def _rocm_aiter_fused_moe_triton_gemm_a4w4( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weight: torch.Tensor, + topk_ids: torch.Tensor, + activation_method: int, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + output_dtype: torch.dtype | None, +) -> torch.Tensor: + from aiter import ActivationType + from aiter.fused_moe import get_inter_dim + from aiter.ops.triton.moe.moe_op_gemm_a4w4 import moe_gemm_a4w4, mxfp4_quant + from aiter.ops.triton.moe.moe_routing.routing import routing + + if activation_method != int(ActivationType.Silu): + raise RuntimeError( + "Triton moe_gemm_a4w4 path supports only Silu/SwiGLU (ActivationType.Silu)." + ) + + device = hidden_states.device + out_dtype = output_dtype or hidden_states.dtype + # FP32 matmul outputs for numerical headroom before the second mxfp4_quant. + gemm_out_dt = torch.float32 + m, _ = topk_ids.shape + num_experts = w1.shape[0] + + logits = torch.full( + (m, num_experts), -1e9, device=device, dtype=torch.float32 + ) + tid = topk_ids.long().clamp(min=0, max=num_experts - 1) + logits.scatter_( + 1, + tid, + torch.log(topk_weight.to(torch.float32).clamp(min=1e-20)), + ) + + rdata, gather_idx, scatter_idx = routing(logits, topk_ids.shape[1]) + gate_scal = rdata.gate_scal + + e, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape) + is_g1u1 = inter_dim != w1.shape[1] + if not is_g1u1: + raise RuntimeError( + "Triton moe_gemm_a4w4 path expects SwiGLU (w1 dim1 == 2 * inter)." + ) + + w1_v = _mxfp4_moe_weight_as_uint8(_mxfp4_moe_w1_triton_gemm_view(w1, w2)) + w2_v = _mxfp4_moe_weight_as_uint8(_mxfp4_moe_w2_triton_gemm_view(w2)) + + x_q, x_s = mxfp4_quant(hidden_states.to(out_dtype)) + # Raw W1@x per (token, expert), then CK-style silu_and_mul (not Triton ``apply_swiglu``). + # ``scatter_indx=None`` keeps expanded rows until after activation; higher VRAM than fused swiglu. + mid_raw = moe_gemm_a4w4( + x_q, + w1_v, + x_s, + w1_scale, + x_static_scale=None, + quant_static_scale=None, + bias=None, + routing_data=rdata, + gather_indx=gather_idx, + scatter_indx=None, + gammas=None, + swizzle_mx_scale=None, + out_dtype=gemm_out_dt, + apply_swiglu=False, + ) + mid = _silu_and_mul_glu(mid_raw) + mid_q, mid_s = mxfp4_quant(mid.to(out_dtype)) + out = moe_gemm_a4w4( + mid_q, + w2_v, + mid_s, + w2_scale, + x_static_scale=None, + quant_static_scale=None, + bias=None, + routing_data=rdata, + gather_indx=None, + scatter_indx=scatter_idx, + gammas=gate_scal, + swizzle_mx_scale=None, + out_dtype=gemm_out_dt, + apply_swiglu=False, + ) + return out.to(out_dtype) + + def _rocm_aiter_fused_moe_impl( hidden_states: torch.Tensor, w1: torch.Tensor, @@ -151,6 +289,30 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) + m_tokens = hidden_states.shape[0] + use_triton_a4w4 = ( + envs.VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4 + and quant_type == QuantType.per_1x32 + and expert_mask is None + and not doweight_stage1 + and num_local_tokens is None + and a1_scale is None + and a2_scale is None + and m_tokens >= envs.VLLM_ROCM_AITER_TRITON_MOE_MIN_TOKENS + ) + if use_triton_a4w4: + return _rocm_aiter_fused_moe_triton_gemm_a4w4( + hidden_states, + w1, + w2, + topk_weight, + topk_ids, + activation_method, + w1_scale, + w2_scale, + output_dtype, + ) + return fused_moe( hidden_states, w1, diff --git a/vllm/envs.py b/vllm/envs.py index 9b4a6790b6cf..17f7d4e04363 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -126,6 +126,7 @@ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True + VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4: bool = False VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True VLLM_ROCM_MOE_PADDING: bool = True @@ -1153,6 +1154,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1") ), + # ROCm AITER fused MoE: use Triton moe_gemm_a4w4 (MXFP4) instead of aiter.fused_moe. + "VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4": lambda: ( + os.getenv("VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4", "False").lower() + in ("true", "1") + ), # use rocm skinny gemms "VLLM_ROCM_USE_SKINNY_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1") diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 3906a7e057ca..7f09703f9b20 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -46,11 +46,34 @@ def aiter_triton_kernel_w4a8_moe_forward( and quant_config.use_mxfp4_w4a8 and rocm_aiter_ops.is_enabled() ) - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + # TODO: (JPVILLAM) This causes a tl compile error on 1250. + # Need to figure out why this is a problem and sync with triton team + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize ) + + # gfx1250: aiter's in-kernel gather is numerically broken (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, + # sorted row i reads source token gather_idx[i] // n_expts_act, so this + # reproduces the in-kernel gather exactly (manual gather -> maxrel ~5e-3). + # gfx950 keeps the (working) in-kernel gather. + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + return triton_kernel_fused_mxfp4_w4a8_experts( None, hidden_states, @@ -124,6 +147,15 @@ def triton_kernel_fused_mxfp4_w4a8_experts( hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale ) + # gfx1250 stores the MXFP4 weight scale unswizzled (StridedLayout, see + # mxfp4_utils._swizzle_mxfp4) because the gfx1250 moe_gemm_a8w4 reads a + # CDNA4-swizzled scale as garbage (validated on the FFM sim: CDNA4_SCALE -> + # maxrel ~7e4, plain/None -> ~6e-3); pass swizzle_mx_scale=None there. + # gfx950 uses the CDNA4 swizzle layout. + from vllm.platforms.rocm import on_gfx1250 + + mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" + intermediate_cache1 = moe_gemm_a8w4( hidden_states, w1.storage.data, @@ -135,7 +167,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( routing_data, gather_indx=gather_indx, gammas=gammas if apply_router_weight_on_input else None, - swizzle_mx_scale="CDNA4_SCALE", + swizzle_mx_scale=mx_scale_swizzle, out_dtype=torch.float8_e4m3fn, apply_swiglu=True, alpha=swiglu_alpha, @@ -155,7 +187,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( routing_data, scatter_indx=scatter_indx, gammas=None if apply_router_weight_on_input else gammas, - swizzle_mx_scale="CDNA4_SCALE", + swizzle_mx_scale=mx_scale_swizzle, unpadded_N=unpadded_N_w2, unpadded_K=unpadded_K_w2, ) @@ -193,12 +225,16 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_current_device() -> bool: - # Requires AITER and GFX950 + # Requires AITER and a supported AMD arch. gfx950 (CDNA4) uses the + # in-kernel gather + CDNA4 scale swizzle; gfx1250 routes through the + # same moe_gemm_a8w4 kernel with a manual gather and unswizzled scales + # (see aiter_triton_kernel_w4a8_moe_forward / the swizzle handling in + # triton_kernel_fused_mxfp4_w4a8_experts). if not rocm_aiter_ops.is_enabled(): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - return on_gfx950() + return on_gfx950() or on_gfx1250() @staticmethod def _supports_no_act_and_mul() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 98265abf7c8f..96093acff84a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import replace + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -577,6 +580,117 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx +def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_states: torch.Tensor, + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int = 1, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). + + Uses ``aiter.ops.triton.quant_moe.downcast_to_static_fp8`` with a scalar scale: + for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when + single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or + ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static + FP8 path). + + Returns: + (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for + FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + if not quant_config.use_mxfp4_w4a16: + return hidden_states, None + try: + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return hidden_states, None + if not rocm_aiter_ops.is_enabled(): + return hidden_states, None + + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + + lhs_scale = None + if qp is not None and qp.flex_ctx.lhs_data.scale is not None: + s = qp.flex_ctx.lhs_data.scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = hidden_states.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + + return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale + + +def _mxfp4_w4a8_unpadded_dims( + moe_cfg: FusedMoEConfig, +) -> tuple[int | None, int | None, int | None, int | None]: + """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" + if ( + moe_cfg.intermediate_size_per_partition_unpadded is None + or moe_cfg.hidden_dim_unpadded is None + ): + return None, None, None, None + unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 + unpadded_k_w1 = moe_cfg.hidden_dim_unpadded + unpadded_n_w2 = moe_cfg.hidden_dim_unpadded + unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded + return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 + + +def _mxfp4_w4a8_resolve_lhs_scale( + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int, + activations: torch.Tensor, +) -> torch.Tensor: + """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. + + Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise + ``max(|activations|) / 448`` (e4m3), matching + ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + lhs_scale = None + if qp is not None and qp.flex_ctx is not None: + ld = qp.flex_ctx.lhs_data + if ld is not None and ld.scale is not None and ld.scale.numel() == 1: + lhs_scale = ld.scale + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = activations.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + return lhs_scale + + +def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + precision_config, + lhs_scale: torch.Tensor | None, +): + """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" + if precision_config is None or lhs_scale is None: + return precision_config + from triton_kernels.numerics import InFlexData + + new_flex = replace( + precision_config.flex_ctx, + lhs_data=InFlexData(scale=lhs_scale), + ) + return replace(precision_config, flex_ctx=new_flex) + + class BaseOAITritonExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: @@ -816,6 +930,193 @@ def activation( else: super().activation(activation, output, input) + def _try_apply_aiter_w4a8( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + apply_router_weight_on_input: bool, + ) -> torch.Tensor | None: + """ROCm/gfx1250 MXFP4 MoE via aiter's W4A8 ``moe_gemm_a8w4``. + + The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls + ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and + fails to compile on gfx1250. aiter's gluon ``moe_gemm_a8w4`` (FP8 act x MXFP4 + weight) does run on gfx1250, so we route the MoE through it. + + Routing is built aiter-native directly from ``topk_ids``/``topk_weights`` so we + never touch the diverged vendored ``triton_kernels`` routing structs (whose + ``RoutingData``/``GatherIndx``/``ExptData`` field names and layout differ from + aiter's). This mirrors ``triton_kernel_fused_mxfp4_w4a8_experts`` + (aiter_mxfp4_w4a8_moe.py) but resolves the FP8 LHS scales dynamically because + DeepSeek-V4 is ``activation_scheme: dynamic`` (no calibrated static scale). + + aiter's ``routing`` always softmaxes the gate weights, so we feed + ``log(topk_weights)``: softmax-over-selected then recovers + ``topk_weights / sum(topk_weights)``. DeepSeek-V4 weights are + norm_topk_prob-normalized then scaled by ``routed_scaling_factor`` (so they sum + to that factor, not 1); we restore the exact weighting with a per-token output + rescale by the weight sum (== 1 for plain renormalize routing, a no-op there). + + Returns ``output`` on success, or ``None`` if aiter is unavailable (caller then + falls back to ``matmul_ogs``). + """ + try: + from aiter.ops.triton.moe.moe_routing.routing import ( + routing as aiter_routing, + ) + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + M = hidden_states.shape[0] + num_experts = w1.shape[0] + topk = topk_ids.size(1) + + # Reconstruct dense gating logits from the sparse topk selection. log(weight) + # padded with a large-negative sentinel makes aiter's softmax reproduce both + # this exact expert selection and weight = topk_weights / per-token-sum. + tw = topk_weights.to(torch.float32) + logits = torch.full( + (M, num_experts), -1e30, device=hidden_states.device, dtype=torch.float32 + ) + logits.scatter_( + 1, + topk_ids.long().clamp(min=0, max=num_experts - 1), + torch.log(tw.clamp(min=1e-20)), + ) + logger.info( + "[aiter-w4a8] ENTER M=%d K=%d num_experts=%d topk=%d " + "w1=%s w2=%s w1_wscale=%s w2_wscale=%s", + M, + hidden_states.shape[1], + num_experts, + topk, + tuple(w1.storage.data.shape), + tuple(w2.storage.data.shape), + tuple(quant_config.w1_precision.weight_scale.storage.data.shape), + tuple(quant_config.w2_precision.weight_scale.storage.data.shape), + ) + routing_data, gather_idx, scatter_idx = aiter_routing( + logits, topk, sm_first=False + ) + gammas = routing_data.gate_scal + logger.info( + "[aiter-w4a8] routing OK block_m=%s gate_scal=%s gather=%s scatter=%s", + getattr(routing_data, "block_m", None), + tuple(gammas.shape), + tuple(gather_idx.shape), + tuple(scatter_idx.shape), + ) + + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + # FP8 LHS scales (static if calibrated, else dynamic amax/448). GEMM1 emits FP8 + # directly (apply_swiglu=True, out_dtype fp8) quantized with GEMM2's LHS scale, + # so resolve GEMM2's scale up front (hidden states as the amax proxy) and reuse + # it as GEMM2's input scale so quantize/dequantize are consistent. + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=1, activations=hidden_states + ) + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=2, activations=hidden_states + ) + # aiter's in-kernel gather is numerically broken on gfx1250 (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, sorted + # row i reads source token gather_idx[i] // n_expts_act, so this reproduces the + # in-kernel gather exactly (validated: manual gather -> maxrel ~5e-3). + gather_src = gather_idx.to(torch.long) // topk + hidden_sorted = hidden_states[gather_src] + hidden_fp8 = downcast_to_static_fp8(hidden_sorted, gemm1_lhs_scale) + logger.info( + "[aiter-w4a8] scales g1_lhs=%s g2_lhs=%s hidden_fp8=%s; calling gemm1 " + "(unpadded N/K w1=%s/%s w2=%s/%s, alpha=%.4f limit=%.2f)", + float(gemm1_lhs_scale), + float(gemm2_lhs_scale), + tuple(hidden_fp8.shape), + unpadded_n_w1, + unpadded_k_w1, + unpadded_n_w2, + unpadded_k_w2, + swiglu_alpha, + swiglu_limit, + ) + + intermediate_cache1 = moe_gemm_a8w4( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale, + quant_config.w1_bias, + routing_data, + gather_indx=None, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + out_dtype=torch.float8_e4m3fn, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + logger.info( + "[aiter-w4a8] gemm1 OK ic1=%s %s; calling gemm2", + tuple(intermediate_cache1.shape), + intermediate_cache1.dtype, + ) + intermediate_cache3 = moe_gemm_a8w4( + intermediate_cache1, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + + # Restore the per-token weight sum (e.g. routed_scaling_factor) that the + # log+softmax normalization divided out. Output weighting (gammas on GEMM2) is + # linear in the gate, so a single post-scale is exact. + logger.info( + "[aiter-w4a8] gemm2 OK ic3=%s %s; rescaling + writing output=%s", + tuple(intermediate_cache3.shape), + intermediate_cache3.dtype, + tuple(output.shape), + ) + out = intermediate_cache3.to(output.dtype) * tw.sum(dim=1, keepdim=True) + output.copy_(out.reshape(output.shape)) + logger.info("[aiter-w4a8] DONE") + return output + def apply( self, output: torch.Tensor, @@ -874,6 +1175,40 @@ def apply( if global_num_experts == -1: global_num_experts = E + # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from + # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs + # unswizzle. Output-side router weighting only (the post-scale below assumes + # it); LoRA still uses the matmul_ogs path further down. + if ( + (quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16) + and not apply_router_weight_on_input + and self._lora_context is None + and rocm_aiter_ops.is_enabled() + ): + if ( + self._try_apply_aiter_w4a8( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + apply_router_weight_on_input, + ) + is not None + ): + return + + # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids + # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for + # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). + # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not + # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA + # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below + # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA + # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. + # Note that the output tensor might be in workspace13 intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) @@ -882,18 +1217,108 @@ def apply( gammas = routing_data.gate_scal if routing_data else None - matmul_ogs( - hidden_states, - w1, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=quant_config.w1_precision, - gammas=gammas if apply_router_weight_on_input else None, - fused_activation=None, - y=intermediate_cache1, + hidden_bf16_for_lora = hidden_states + + _moe_gemm_a8w4_fn = None + _downcast_static_fp8_fn = None + # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's + # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also + # serve the w4a16 case by dynamically quantizing the bf16 activations to + # FP8 here and passing the scale. This routes the MoE through aiter's + # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. + _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 + if _mxfp4_weights and rocm_aiter_ops.is_enabled(): + try: + from aiter.ops.triton.moe_op_gemm_a8w4 import ( + moe_gemm_a8w4 as _moe_gemm_a8w4_fn, + ) + from aiter.ops.triton.quant_moe import ( + downcast_to_static_fp8 as _downcast_static_fp8_fn, + ) + except ImportError: + pass + use_aiter_unfused_a8w4 = ( + _mxfp4_weights + and _moe_gemm_a8w4_fn is not None + and _downcast_static_fp8_fn is not None ) + if use_aiter_unfused_a8w4: + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=1, + activations=hidden_bf16_for_lora, + ) + # GEMM1 passes W2's static scale through to the kernel; resolve it before + # GEMM2 activations exist (fallback uses token hidden states for amax). + gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=hidden_bf16_for_lora, + ) + hidden_fp8 = _downcast_static_fp8_fn( + hidden_bf16_for_lora, + gemm1_lhs_scale, + ) + gemm1_out = _moe_gemm_a8w4_fn( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale_proxy, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale="CDNA4_SCALE", + out_dtype=torch.bfloat16, + apply_swiglu=False, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) + else: + hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_bf16_for_lora, + quant_config, + ) + ) + w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w1_precision, + gemm1_lhs_scale_fp8, + ) + + matmul_ogs( + hidden_for_gemm1, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w1_precision_for_gemm, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=None, + y=intermediate_cache1, + ) + # w13 LoRA: gather the activation input from expert-sorted # intermediate_cache1, then add the LoRA delta in-place on that copy # before passing it to activation — exactly mirroring the old @@ -914,7 +1339,7 @@ def apply( ) = self.apply_w13_lora( lora_context, y=act_input, - x=hidden_states, + x=hidden_bf16_for_lora, topk_ids=global_topk_ids, topk_weights=topk_weights, expert_map=expert_map, @@ -935,16 +1360,56 @@ def apply( # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. routing_data.n_expts_act = 1 - matmul_ogs( - intermediate_cache2[gather_indx.src_indx], - w2, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - precision_config=quant_config.w2_precision, - gammas=None if apply_router_weight_on_input else gammas, - y=intermediate_cache3, - ) + gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] + if use_aiter_unfused_a8w4: + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=gemm2_input_bf16, + ) + gemm2_fp8 = _downcast_static_fp8_fn( + gemm2_input_bf16, + gemm2_lhs_scale, + ) + gemm2_out = _moe_gemm_a8w4_fn( + gemm2_fp8, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale="CDNA4_SCALE", + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) + else: + gemm2_input, gemm2_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + gemm2_input_bf16, + quant_config, + gemm_num=2, + ) + ) + w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w2_precision, + gemm2_lhs_scale_fp8, + ) + + matmul_ogs( + gemm2_input, + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=w2_precision_for_gemm, + gammas=None if apply_router_weight_on_input else gammas, + y=intermediate_cache3, + ) # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 51b7b29551d0..04944c8696c7 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -59,6 +59,12 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): scale_layout = CDNA4MXScaleLayout else: + # gfx1250 (and other ROCm archs): the vendored triton_kernels has no + # gfx1250 scale-swizzle, and the gfx1250 aiter `moe_gemm_a8w4` kernel + # reads a CDNA4-swizzled scale as garbage (validated on the FFM sim: + # CDNA4_SCALE -> maxrel ~7e4, plain/None -> ~6e-3). Keep the scale + # unswizzled (StridedLayout) and pass swizzle_mx_scale=None in the + # W4A8 path (gpt_oss_triton_kernels_moe._try_apply_aiter_w4a8). scale_layout = StridedLayout else: value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout( diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index dbc05273f305..5c0b0c86aeb7 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950 + from vllm.platforms.rocm import on_gfx1250, on_gfx1x, on_gfx9, on_gfx950 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -172,6 +172,9 @@ def rocm_unquantized_gemm_impl( use_skinny = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) + # gfx1250 excludes the skinny-GEMM kernels (wvSplitK/LLMM1) from the + # build (gfx9/gfx11 ISA); fall back to torch GEMM there. + and not on_gfx1250() and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index d12db96c5d46..2e6f066fc20a 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -384,6 +384,14 @@ def _load_weights_mxfp4( if is_pp_missing_parameter(name, self): continue + # Skip checkpoint weights for layers dropped by an hf_overrides + # `num_hidden_layers` prune: the model only builds params for the + # layers it keeps, so extra-layer weights have no destination param. + if "layers." in name: + idx_str = name.split("layers.", 1)[1].split(".", 1)[0] + if idx_str.isdigit() and int(idx_str) >= self.config.num_hidden_layers: + continue + if ".w13_weight_scale" in name: # Handle MLP gate and up projection weights scale if use_ep: @@ -606,6 +614,13 @@ def _get_moe_weight_dtype(layer_id: int = 0) -> str | None: if is_pp_missing_parameter(name, self): continue + # Skip checkpoint weights for layers dropped by an hf_overrides + # `num_hidden_layers` prune (mirrors _load_weights_mxfp4). + if "layers." in name: + idx_str = name.split("layers.", 1)[1].split(".", 1)[0] + if idx_str.isdigit() and int(idx_str) >= self.config.num_hidden_layers: + continue + layer_id, expert_id, fused_name = None, None, None moe_quant_method = None if "experts" in name: diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 84318a8107d3..c9cce402f1d7 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -794,6 +794,16 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: expert_mapping = self.get_expert_mapping() for name, loaded_weight in weights: + # Skip checkpoint weights for transformer layers pruned by a + # num_hidden_layers hf_override (keeps load + run tractable on the + # FFM simulator). Names here are already vLLM-mapped ("layers.N."). + if "layers." in name: + try: + _li = int(name.split("layers.", 1)[1].split(".", 1)[0]) + if _li >= self.config.num_hidden_layers: + continue + except (IndexError, ValueError): + pass for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -850,6 +860,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: elif "attn_sink" in name: if is_pp_missing_parameter(name, self): continue + if name not in params_dict: + # Layer pruned via a num_hidden_layers hf_override; skip + # its checkpoint weights instead of KeyError-ing. + continue narrow_weight = loaded_weight[head_rank_start:head_rank_end] n = narrow_weight.shape[0] params_dict[name][:n].copy_(narrow_weight) @@ -858,6 +872,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: else: if is_pp_missing_parameter(name, self): continue + if name not in params_dict: + # Layer pruned via a num_hidden_layers hf_override; skip. + continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index bc66178f74bf..e01be21a0f1c 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -16,6 +16,7 @@ _compress_kv_sparse_attn_cutedsl, _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + _fused_kv_compress_norm_rope_insert_sparse_attn, _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl, _norm_rope_insert_sparse_attn_cutedsl, ) @@ -242,13 +243,24 @@ def __init__( assert not use_fp4_cache, ( "MXFP4 cache is only supported for indexer (head=128)" ) - self._use_cutedsl_sparse_compressor = True - self._use_cutedsl_fused_sparse_compressor = self.compress_ratio == 4 - self._compress_kernel = _compress_kv_sparse_attn_cutedsl - self._norm_rope_store_kernel = _norm_rope_insert_sparse_attn_cutedsl - self._fused_sparse_kernel = ( - _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl - ) + # The CuTeDSL (CUTLASS) compressor kernels are NVIDIA-only (`import + # cutlass`). On ROCm/gfx1250 use the equivalent pure-Triton fused kernel + # (`_fused_kv_compress_norm_rope_insert_sparse_attn`, head=512 / nope=448 + # FP8 + rope=64 bf16), which shares the same call signature as the indexer + # Triton path below. + if current_platform.is_rocm(): + self._use_cutedsl_sparse_compressor = False + self._fused_kernel = ( + _fused_kv_compress_norm_rope_insert_sparse_attn + ) + else: + self._use_cutedsl_sparse_compressor = True + self._use_cutedsl_fused_sparse_compressor = self.compress_ratio == 4 + self._compress_kernel = _compress_kv_sparse_attn_cutedsl + self._norm_rope_store_kernel = _norm_rope_insert_sparse_attn_cutedsl + self._fused_sparse_kernel = ( + _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl + ) self._quant_block = 64 self._token_stride = self.nope_head_dim + self.rope_head_dim * 2 self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index 645da0a1fe97..84113b5c49ed 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -126,6 +126,20 @@ def rocm_platform_plugin() -> str | None: except Exception as e: logger.debug("ROCm platform is not available because: %s", str(e)) + # Fallback: amdsmi queries the kernel driver/sysfs, which does not see GPUs + # exposed only through the HSA layer (e.g. the FFM pre-silicon simulator). + # Trust torch's HIP runtime detection in that case. + if not is_rocm: + try: + import torch + + if torch.version.hip is not None and torch.cuda.device_count() > 0: + is_rocm = True + logger.debug("ROCm platform detected via torch.version.hip " + "(amdsmi unavailable).") + except Exception as e: + logger.debug("torch HIP ROCm fallback failed: %s", str(e)) + return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 114d236f1319..710c13f2b56f 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -35,8 +35,14 @@ amdsmi_topo_get_link_type, amdsmi_topo_get_numa_node_number, ) + + _AMDSMI_AVAILABLE = True except ImportError as e: logger.warning("Failed to import from amdsmi with %r", e) + # amdsmi is unavailable on the FFM pre-silicon simulator (and would report + # the host's real GPUs, not the sim). amdsmi-decorated methods fall back to + # torch when this is False. + _AMDSMI_AVAILABLE = False try: import vllm._C # noqa: F401 @@ -143,6 +149,10 @@ def _sync_hip_cuda_env_vars(): def with_amdsmi_context(fn): @wraps(fn) def wrapper(*args, **kwargs): + if not _AMDSMI_AVAILABLE: + # amdsmi unavailable (e.g. FFM simulator): skip init/shutdown and let + # the wrapped function use its own non-amdsmi (torch) fallback. + return fn(*args, **kwargs) amdsmi_init() try: return fn(*args, **kwargs) @@ -175,11 +185,13 @@ def _get_gcn_arch() -> str: return _query_gcn_arch_from_amdsmi() except Exception as e: logger.debug("Failed to get GCN arch via amdsmi: %s", e) - logger.warning_once( - "Failed to get GCN arch via amdsmi, falling back to torch.cuda. " - "This will initialize CUDA and may cause " - "issues if CUDA_VISIBLE_DEVICES is not set yet." - ) + # NOTE: use logger.debug, not warning_once, here. This runs at module + # load while resolving the platform; warning_once imports + # vllm.distributed, which causes a circular import before + # current_platform is bound. This path is taken on the FFM simulator, + # where amdsmi is unavailable (and would report the host's real gfx950 + # cards rather than the simulated gfx1250) — torch.cuda below is correct. + logger.debug("Failed to get GCN arch via amdsmi, falling back to torch.cuda.") # Ultimate fallback: use torch.cuda (will initialize CUDA) return torch.cuda.get_device_properties("cuda").gcnArchName @@ -191,11 +203,17 @@ def _get_gcn_arch() -> str: _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"]) _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) -_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) -_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) +_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950", "gfx1250"]) +_ON_GFX9 = any( + arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950", "gfx1250"] +) # TODO(JPVILLAM): Bubblegum patch to unlock gptoss _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH +# any( +# arch in _GCN_ARCH for arch in ["gfx950", "gfx1250"] +# ) # TODO(JPVILLAM): Bubblegum patch to unlock DSR1 +_ON_GFX1250 = "gfx1250" in _GCN_ARCH def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: @@ -297,6 +315,10 @@ def on_gfx950() -> bool: return _ON_GFX950 +def on_gfx1250() -> bool: + return _ON_GFX1250 + + @cache def use_rocm_custom_paged_attention( qtype: torch.dtype, @@ -670,6 +692,9 @@ def is_fully_connected(cls, physical_device_ids: list[int]) -> bool: @with_amdsmi_context @lru_cache(maxsize=8) def get_device_name(cls, device_id: int = 0) -> str: + if not _AMDSMI_AVAILABLE: + # FFM simulator: amdsmi can't see the simulated GPU; use torch. + return torch.cuda.get_device_name(device_id) physical_device_id = cls.device_id_to_physical_device_id(device_id) handle = amdsmi_get_processor_handles()[physical_device_id] asic_info = amdsmi_get_gpu_asic_info(handle) @@ -801,7 +826,7 @@ def get_device_communicator_cls(cls) -> str: @classmethod def supports_mx(cls) -> bool: - return any(gfx in _GCN_ARCH for gfx in ["gfx95"]) + return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"]) @classmethod def supports_fp8(cls) -> bool: diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 0d6a3d298b66..3a619ac7bcc5 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -30,17 +30,22 @@ flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[assignment] get_scheduler_metadata = xpu_ops.get_scheduler_metadata # type: ignore[assignment] elif current_platform.is_rocm(): + # On ROCm we use AITER's Triton flash-attention; the upstream flash-attn + # package is not installed/available. (Same source as aiter_triton_mla.py.) try: - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] + from aiter.ops.triton.mha import ( # type: ignore[no-redef] + flash_attn_varlen_func, + ) - # Mark that upstream flash-attn is available on ROCm + # A working flash_attn_varlen_func is available on ROCm (via AITER). _ROCM_FLASH_ATTN_AVAILABLE = True except ImportError: def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc] raise ImportError( - "ROCm platform requires upstream flash-attn " - "to be installed. Please install flash-attn first." + "ROCm platform requires AITER's Triton MHA " + "(aiter.ops.triton.mha.flash_attn_varlen_func). " + "Please install aiter." ) # ROCm doesn't use scheduler metadata (FA3 feature), provide stub @@ -231,7 +236,7 @@ def is_flash_attn_varlen_func_available() -> bool: Platform-specific sources: - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func - XPU: xpu_ops.flash_attn_varlen_func - - ROCm: upstream flash_attn.flash_attn_varlen_func (if available) + - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py) which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 582c6a17cb4e..0b93285e4e12 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -50,7 +50,6 @@ from vllm.tracing import instrument from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling -from vllm.utils.torch_utils import set_random_seed from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.outputs import ( @@ -289,7 +288,7 @@ def init_device(self): logger.info_once("Using V2 Model Runner") # Set random seed. - set_random_seed(self.model_config.seed) + # set_random_seed(self.model_config.seed) # Now take memory snapshot after NCCL is initialized gc.collect() @@ -709,7 +708,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. - set_random_seed(self.model_config.seed) + # set_random_seed(self.model_config.seed) # All warmup is done — start monitoring for unexpected JIT # compilations that would cause latency spikes during inference.