Compare commits

..
Author SHA1 Message Date
khluuandClaude Opus 4.6 f478d42cdb [CI] Reorganize release pipeline: separate nightly vs release sections
Reorder the pipeline into clear sections for readability:

1. Build Python Wheels (always runs)
2. ROCm Wheel Pipeline (always runs)
3. Nightly Docker Images (NIGHTLY=1 only) - CUDA/Ubuntu builds,
   multi-arch manifests, DockerHub publish, ROCm image + publish
4. Release (manual) - version input, PyPI upload, CPU image builds,
   ROCm root index

Key changes:
- Extract CPU image builds (manual/blocked) from nightly-gated group
  into their own "Build release CPU Docker images" group so they remain
  available for actual releases without NIGHTLY=1
- Move ROCm wheel jobs (1-4) up next to CUDA wheel builds
- Remove redundant per-step NIGHTLY gates inside already-gated groups
- Rename groups: "Build release Docker images" -> "Build nightly Docker
  images", "Publish release images" -> "Publish nightly images"

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-15 01:08:39 -07:00
khluuandClaude Opus 4.6 1a811d5747 [CI] Only build release Docker images when NIGHTLY=1
Gate the "Build release Docker images" group, "Publish release images"
group, and ROCm release image build behind NIGHTLY=1 to avoid expensive
image builds on every commit in the release pipeline.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-15 01:01:03 -07:00
260 changed files with 2818 additions and 11852 deletions
+2 -2
View File
@@ -46,7 +46,7 @@ steps:
- tests/models/language/pooling/
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
pytest -x -v -s tests/models/language/generation -m cpu_model
pytest -x -v -s tests/models/language/pooling -m cpu_model"
@@ -99,7 +99,7 @@ steps:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
parallelism: 3
parallelism: 2
- label: "Arm CPU Test"
depends_on: []
+347 -353
View File
@@ -1,9 +1,7 @@
steps:
- input: "Provide Release version here"
id: input-release-version
fields:
- text: "What is the release version?"
key: release-version
# =============================================================================
# Build Python Wheels (runs on every pipeline trigger)
# =============================================================================
- group: "Build Python wheels"
key: "build-wheels"
@@ -98,15 +96,257 @@ steps:
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"
# =============================================================================
# ROCm Wheel Pipeline (runs on every pipeline trigger)
# =============================================================================
- group: "Build release Docker images"
# ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
- label: ":rocm: Build ROCm Base Image & Wheels"
id: build-rocm-base-wheels
depends_on: ~
agents:
queue: cpu_queue_release
commands:
- |
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"
echo "========================================"
echo "ROCm Base Build Configuration"
echo "========================================"
echo " CACHE_KEY: $${CACHE_KEY}"
echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
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"
fi
# Scenario 1: Both cached (best case)
if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
echo ""
echo "FULL CACHE HIT - Reusing both image and 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
else
echo ""
echo " CACHE MISS - Building from scratch..."
echo ""
# Build full base image and push to ECR
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag "$${ECR_CACHE_TAG}" \
--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 \
.
# Build wheel extraction stage
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
--target debs_wheel_release \
--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
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}
.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}"
echo ""
echo " Build complete - Image and wheels cached"
fi
artifact_paths:
- "artifacts/rocm-base-wheels/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 2: Build vLLM ROCm Wheel
- label: ":python: Build vLLM ROCm Wheel - x86_64"
id: build-rocm-vllm-wheel
depends_on:
- step: build-rocm-base-wheels
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 180
commands:
# Download artifacts and prepare Docker image
- |
set -euo pipefail
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
# This fixes version detection when tags are moved/force-pushed
echo "Fetching latest tags from origin..."
git fetch --tags --force origin
# Log tag information for debugging version detection
echo "========================================"
echo "Git Tag Verification"
echo "========================================"
echo "Current HEAD: $(git rev-parse HEAD)"
echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
echo ""
echo "Recent tags (pointing to commits near HEAD):"
git tag -l --sort=-creatordate | head -5
echo "setuptools_scm version detection:"
pip install -q setuptools_scm 2>/dev/null || true
python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
echo "========================================"
# Download wheel artifacts from current build
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"
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}"
# Prepare base wheels for Docker build context
mkdir -p docker/context/base-wheels
touch docker/context/base-wheels/.keep
cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
echo "Base wheels for vLLM build:"
ls -lh docker/context/base-wheels/
echo "========================================"
echo "Building vLLM wheel with:"
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
echo "========================================"
# Build vLLM wheel using local checkout (REMOTE_VLLM=0)
DOCKER_BUILDKIT=1 docker build \
--file docker/Dockerfile.rocm \
--target export_vllm_wheel_release \
--output type=local,dest=rocm-dist \
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
--build-arg REMOTE_VLLM=0 \
--build-arg GIT_REPO_CHECK=1 \
--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 \
.
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/
echo "Final vLLM wheel:"
ls -lh artifacts/rocm-vllm-wheel/
artifact_paths:
- "artifacts/rocm-vllm-wheel/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 3: Upload Wheels to S3
- label: ":s3: Upload ROCm Wheels to S3"
id: upload-rocm-wheels
depends_on:
- step: build-rocm-vllm-wheel
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 60
commands:
# Download all wheel artifacts and run upload
- |
set -euo pipefail
# Download artifacts from current build
echo "Downloading artifacts from current build"
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
# Run upload script
bash .buildkite/scripts/upload-rocm-wheels.sh
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 4: Annotate ROCm Wheel Release
- label: ":memo: Annotate ROCm wheel release"
id: annotate-rocm-release
depends_on:
- upload-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash .buildkite/scripts/annotate-rocm-release.sh"
env:
S3_BUCKET: "vllm-wheels"
# =============================================================================
# Nightly: Build & Publish Docker Images (NIGHTLY=1 only)
# =============================================================================
- group: "Build nightly Docker images"
key: "build-release-images"
depends_on: block-build-release-images
allow_dependency_failure: true
if: build.env("NIGHTLY") == "1"
steps:
- label: "Build release image - x86_64 - CUDA 12.9"
depends_on: ~
@@ -199,44 +439,9 @@ steps:
- "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 UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --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-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
- block: "Build release image for x86_64 CPU"
key: block-cpu-release-image-build
depends_on: ~
- label: "Build release image - x86_64 - CPU"
depends_on:
- block-cpu-release-image-build
- input-release-version
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 --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 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)"
env:
DOCKER_BUILDKIT: "1"
- block: "Build release image for arm64 CPU"
key: block-arm64-cpu-release-image-build
depends_on: ~
- label: "Build release image - arm64 - CPU"
depends_on:
- block-arm64-cpu-release-image-build
- input-release-version
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 --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)"
env:
DOCKER_BUILDKIT: "1"
- group: "Publish release images"
- group: "Publish nightly images"
key: "publish-release-images"
if: build.env("NIGHTLY") == "1"
steps:
- label: "Create multi-arch manifest - CUDA 12.9"
depends_on:
@@ -298,7 +503,6 @@ steps:
- 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
commands:
@@ -316,7 +520,6 @@ steps:
- label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
depends_on:
- create-multi-arch-manifest-cuda-13-0
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
@@ -331,301 +534,11 @@ 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_release
commands:
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
# =============================================================================
# ROCm Release Pipeline (x86_64 only)
# =============================================================================
#
# vLLM version is determined by the Buildkite checkout (like CUDA pipeline).
# To build a specific version, trigger the build from that branch/tag.
#
# Environment variables for ROCm builds (set via Buildkite UI or schedule):
#
# Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base
#
# =============================================================================
# ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
- label: ":rocm: Build ROCm Base Image & Wheels"
id: build-rocm-base-wheels
depends_on: ~
agents:
queue: cpu_queue_release
commands:
- |
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"
echo "========================================"
echo "ROCm Base Build Configuration"
echo "========================================"
echo " CACHE_KEY: $${CACHE_KEY}"
echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
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"
fi
# Scenario 1: Both cached (best case)
if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
echo ""
echo "FULL CACHE HIT - Reusing both image and 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
else
echo ""
echo " CACHE MISS - Building from scratch..."
echo ""
# Build full base image and push to ECR
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag "$${ECR_CACHE_TAG}" \
--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 \
.
# Build wheel extraction stage
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
--target debs_wheel_release \
--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
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}
.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}"
echo ""
echo " Build complete - Image and wheels cached"
fi
artifact_paths:
- "artifacts/rocm-base-wheels/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 2: Build vLLM ROCm Wheel
- label: ":python: Build vLLM ROCm Wheel - x86_64"
id: build-rocm-vllm-wheel
depends_on:
- step: build-rocm-base-wheels
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 180
commands:
# Download artifacts and prepare Docker image
- |
set -euo pipefail
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
# This fixes version detection when tags are moved/force-pushed
echo "Fetching latest tags from origin..."
git fetch --tags --force origin
# Log tag information for debugging version detection
echo "========================================"
echo "Git Tag Verification"
echo "========================================"
echo "Current HEAD: $(git rev-parse HEAD)"
echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
echo ""
echo "Recent tags (pointing to commits near HEAD):"
git tag -l --sort=-creatordate | head -5
echo "setuptools_scm version detection:"
pip install -q setuptools_scm 2>/dev/null || true
python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
echo "========================================"
# Download wheel artifacts from current build
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"
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}"
# Prepare base wheels for Docker build context
mkdir -p docker/context/base-wheels
touch docker/context/base-wheels/.keep
cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
echo "Base wheels for vLLM build:"
ls -lh docker/context/base-wheels/
echo "========================================"
echo "Building vLLM wheel with:"
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
echo "========================================"
# Build vLLM wheel using local checkout (REMOTE_VLLM=0)
DOCKER_BUILDKIT=1 docker build \
--file docker/Dockerfile.rocm \
--target export_vllm_wheel_release \
--output type=local,dest=rocm-dist \
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
--build-arg REMOTE_VLLM=0 \
--build-arg GIT_REPO_CHECK=1 \
--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 \
.
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/
echo "Final vLLM wheel:"
ls -lh artifacts/rocm-vllm-wheel/
artifact_paths:
- "artifacts/rocm-vllm-wheel/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 3: Upload Wheels to S3
- label: ":s3: Upload ROCm Wheels to S3"
id: upload-rocm-wheels
depends_on:
- step: build-rocm-vllm-wheel
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 60
commands:
# Download all wheel artifacts and run upload
- |
set -euo pipefail
# Download artifacts from current build
echo "Downloading artifacts from current build"
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
# Run upload script
bash .buildkite/scripts/upload-rocm-wheels.sh
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 4: Annotate ROCm Wheel Release
- label: ":memo: Annotate ROCm wheel release"
id: annotate-rocm-release
depends_on:
- upload-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash .buildkite/scripts/annotate-rocm-release.sh"
env:
S3_BUCKET: "vllm-wheels"
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
- block: "Generate Root Index for ROCm Wheels for Release"
key: block-generate-root-index-rocm-wheels
depends_on: upload-rocm-wheels
- label: ":package: Generate Root Index for ROCm Wheels for Release"
depends_on: block-generate-root-index-rocm-wheels
id: generate-root-index-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
env:
S3_BUCKET: "vllm-wheels"
VARIANT: "rocm721"
# ROCm Job 6: Build ROCm Release Docker Image
# ROCm nightly Docker image
- label: ":docker: Build release image - x86_64 - ROCm"
id: build-rocm-release-image
if: build.env("NIGHTLY") == "1"
depends_on:
- step: block-build-release-images
allow_failure: true
- step: build-rocm-base-wheels
allow_failure: false
agents:
@@ -634,11 +547,11 @@ steps:
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
@@ -646,23 +559,23 @@ steps:
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}"
# 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 "========================================"
# Build vLLM ROCm release image using cached base
DOCKER_BUILDKIT=1 docker build \
--build-arg max_jobs=16 \
@@ -675,10 +588,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"
@@ -705,3 +618,84 @@ steps:
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
# =============================================================================
# Release: Publish Wheels & Build CPU Images (manual, requires release version)
# =============================================================================
- input: "Provide Release version here"
id: input-release-version
fields:
- text: "What is the release version?"
key: release-version
- group: "Publish release 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"
- group: "Build release CPU Docker images"
steps:
- block: "Build release image for x86_64 CPU"
key: block-cpu-release-image-build
depends_on: ~
- label: "Build release image - x86_64 - CPU"
depends_on:
- block-cpu-release-image-build
- input-release-version
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 --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 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)"
env:
DOCKER_BUILDKIT: "1"
- block: "Build release image for arm64 CPU"
key: block-arm64-cpu-release-image-build
depends_on: ~
- label: "Build release image - arm64 - CPU"
depends_on:
- block-arm64-cpu-release-image-build
- input-release-version
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 --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)"
env:
DOCKER_BUILDKIT: "1"
- block: "Generate Root Index for ROCm Wheels for Release"
key: block-generate-root-index-rocm-wheels
depends_on: upload-rocm-wheels
- label: ":package: Generate Root Index for ROCm Wheels for Release"
depends_on: block-generate-root-index-rocm-wheels
id: generate-root-index-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
env:
S3_BUCKET: "vllm-wheels"
VARIANT: "rocm721"
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
exit 1
fi
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename dp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
#echo "--- DP+TP"
#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
#server_pid=$!
#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
#vllm bench serve \
# --backend vllm \
# --dataset-name random \
# --model meta-llama/Llama-3.2-3B-Instruct \
# --num-prompts 20 \
# --result-dir ./test_results \
# --result-filename dp_pp.json \
# --save-result \
# --endpoint /v1/completions
#kill -s SIGTERM $server_pid; wait $server_pid || true
#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
#if [ "$failed_req" -ne 0 ]; then
# echo "Some requests were failed!"
# exit 1
#fi
@@ -51,7 +51,6 @@ function cpu_tests() {
set -e
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
# basic online serving
@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
# Run the image, setting --shm-size=4g for tensor parallel.
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
+1 -1
View File
@@ -769,7 +769,7 @@ steps:
- tests/kernels/helion/
- vllm/platforms/rocm.py
commands:
- pip install helion==1.0.0
- pip install helion==0.3.3
- pytest -v -s kernels/helion/
+1 -2
View File
@@ -196,8 +196,7 @@ steps:
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
- TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
- label: Distributed Tests (2 GPUs)(B200)
device: b200
+1 -1
View File
@@ -155,7 +155,7 @@ steps:
- vllm/utils/import_utils.py
- tests/kernels/helion/
commands:
- pip install helion==1.0.0
- pip install helion==0.3.3
- pytest -v -s kernels/helion/
+1 -15
View File
@@ -4,6 +4,7 @@ depends_on:
steps:
- label: Basic Models Tests (Initialization)
timeout_in_minutes: 45
device: h200_18gb
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -72,18 +73,3 @@ steps:
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
- label: Transformers Backward Compatibility Models Test
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
commands:
- pip install transformers==4.57.5
- pytest -v -s tests/models/test_initialization.py
- pytest -v -s tests/models/test_transformers.py
- pytest -v -s tests/models/multimodal/processing/
- pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/offline_inference/basic/chat.py
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
-13
View File
@@ -42,16 +42,3 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: DFlash Speculators Correctness
timeout_in_minutes: 30
device: h100
optional: true
num_devices: 1
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/model_executor/models/qwen3_dflash.py
- tests/v1/spec_decode/test_speculators_dflash.py
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
+3 -7
View File
@@ -44,9 +44,8 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/vllm/pooling_params.py @noooop @DarkLight1337
/vllm/tokenizers @DarkLight1337 @njhill
/vllm/renderers @DarkLight1337 @njhill
/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/reasoning @aarnphm @chaunceyjiang
/vllm/tool_parsers @aarnphm @chaunceyjiang
# vLLM V1
/vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -92,10 +91,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/tests/v1/kv_connector/nixl_integration @NickLucche
/tests/v1/kv_connector @ApostaC @orozery
/tests/v1/kv_offload @ApostaC @orozery
/tests/v1/determinism @yewentao256
/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/v1/determinism @yewentao256
# Transformers modeling backend
/vllm/model_executor/models/transformers @hmellor
+1
View File
@@ -15,6 +15,7 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before and after, or e2e results
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0).
</details>
**BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)
-1
View File
@@ -45,7 +45,6 @@ jobs:
- name: Smoke test vllm serve
run: |
# Start server in background
VLLM_CPU_KVCACHE_SPACE=1 \
vllm serve Qwen/Qwen3-0.6B \
--max-model-len=2K \
--load-format=dummy \
+5 -5
View File
@@ -62,14 +62,14 @@ jobs:
const prAuthor = context.payload.pull_request.user.login;
const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
per_page: 1,
});
const mergedPRCount = searchResults.total_count;
console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
const authorPRCount = searchResults.total_count;
console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
if (mergedPRCount === 0) {
if (authorPRCount === 1) {
console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
await github.rest.issues.createComment({
owner,
@@ -98,5 +98,5 @@ jobs:
].join('\n'),
});
} else {
console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
}
-15
View File
@@ -923,14 +923,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
# nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
# from cache_kernels.cu, so it belongs in _C rather than _C_stable.
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -957,12 +949,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -1240,7 +1226,6 @@ endif()
if (VLLM_GPU_LANG STREQUAL "CUDA")
include(cmake/external_projects/deepgemm.cmake)
include(cmake/external_projects/flashmla.cmake)
include(cmake/external_projects/cutlass_fa3.cmake)
include(cmake/external_projects/qutlass.cmake)
# vllm-flash-attn should be last as it overwrites some CMake functions
+14 -16
View File
@@ -30,21 +30,6 @@ else()
list(APPEND CXX_COMPILE_FLAGS
"-fopenmp"
"-DVLLM_CPU_EXTENSION")
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
endif()
if (NOT MACOSX_FOUND)
@@ -190,6 +175,20 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
if(NOT NPROC)
set(NPROC 4)
endif()
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
# Fetch and populate ACL
if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
@@ -361,7 +360,6 @@ set(VLLM_EXT_SRC
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
"csrc/cpu/activation_lut_bf16.cpp"
${VLLM_EXT_SRC})
endif()
-163
View File
@@ -1,163 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# CUTLASS FA3 MLA Sparse Attention — requires CUDA >= 12.4, SM90a
#
# Vendors the sgl-attn CUTLASS FlashAttention3 kernel from SGLang into vLLM
# as a self-contained extension (_cutlass_fa3_C). This provides a high-
# performance sparse MLA attention kernel for SM90 (Hopper) GPUs.
#
# Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
# CUTLASS: https://github.com/NVIDIA/cutlass (commit 57e3cfb4)
# Guard: CUDA >= 12.4 required for SM90a features used by FA3
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.4")
message(STATUS "Skipping CUTLASS FA3: requires CUDA >= 12.4")
# Create empty target so setup.py doesn't fail on unsupported systems
add_custom_target(_cutlass_fa3_C)
return()
endif()
# Guard: SM90 architecture required
set(CUTLASS_FA3_SUPPORT_ARCHS)
list(APPEND CUTLASS_FA3_SUPPORT_ARCHS "9.0a")
cuda_archs_loose_intersection(
CUTLASS_FA3_ARCHS "${CUTLASS_FA3_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
if(NOT CUTLASS_FA3_ARCHS)
message(STATUS "Skipping CUTLASS FA3: requires SM90 (CUDA_ARCHS=${CUDA_ARCHS})")
add_custom_target(_cutlass_fa3_C)
return()
endif()
include(FetchContent)
# Fetch sgl-attn (Flash Attention 3 kernels from SGLang)
# We only need the source files, not the build system, so we use
# FetchContent_Populate to download without building.
if (DEFINED ENV{SGL_ATTN_SRC_DIR})
set(SGL_ATTN_SRC_DIR $ENV{SGL_ATTN_SRC_DIR})
endif()
if(SGL_ATTN_SRC_DIR)
FetchContent_Declare(cutlass_fa3
SOURCE_DIR ${SGL_ATTN_SRC_DIR})
else()
FetchContent_Declare(cutlass_fa3
GIT_REPOSITORY https://github.com/sgl-project/sgl-attn.git
GIT_TAG bcf72ccc6816b36a5fae2c5a3c027604629785e0
GIT_PROGRESS TRUE
GIT_SHALLOW FALSE)
endif()
FetchContent_GetProperties(cutlass_fa3)
if(NOT cutlass_fa3_POPULATED)
FetchContent_Populate(cutlass_fa3)
endif()
message(STATUS "CUTLASS FA3 sgl-attn source: ${cutlass_fa3_SOURCE_DIR}")
# Fetch CUTLASS for FA3 (headers only, separate from vLLM's main CUTLASS
# to avoid version conflicts). Use FetchContent_Populate to avoid running
# CUTLASS's own CMakeLists.txt which would create conflicting targets.
if (DEFINED ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
set(CUTLASS_FA3_CUTLASS_SRC_DIR $ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
endif()
if(CUTLASS_FA3_CUTLASS_SRC_DIR)
FetchContent_Declare(cutlass_for_fa3
SOURCE_DIR ${CUTLASS_FA3_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(cutlass_for_fa3
GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
GIT_TAG 57e3cfb47a2d9e0d46eb6335c3dc411498efa198
GIT_PROGRESS TRUE
GIT_SHALLOW FALSE)
endif()
FetchContent_GetProperties(cutlass_for_fa3)
if(NOT cutlass_for_fa3_POPULATED)
FetchContent_Populate(cutlass_for_fa3)
endif()
message(STATUS "CUTLASS FA3 cutlass source: ${cutlass_for_fa3_SOURCE_DIR}")
set(FA3_SRC "${cutlass_fa3_SOURCE_DIR}/hopper")
# flash_api.cpp dispatches to all head dimensions + dtypes (BF16, FP16, FP8)
# at compile time. With FLASHATTENTION_DISABLE_SM8x, only SM90 instantiations
# are needed. We exclude hdimall_* (fails on CUDA 13+) and backward files.
file(GLOB FA3_INSTANTIATION_SOURCES
# BF16 instantiations
"${FA3_SRC}/instantiations/flash_fwd_hdim64_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim96_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim128_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim192_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim256_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu"
# FP16 instantiations
"${FA3_SRC}/instantiations/flash_fwd_hdim64_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim96_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim128_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim192_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim256_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu"
# FP8 (e4m3) instantiations
"${FA3_SRC}/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
set(FA3_CORE_SOURCES
"${FA3_SRC}/flash_api.cpp"
"${FA3_SRC}/flash_prepare_scheduler.cu"
"${FA3_SRC}/flash_fwd_combine.cu")
set(FA3_ALL_SOURCES
"${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_fa3_extension.cc"
${FA3_CORE_SOURCES}
${FA3_INSTANTIATION_SOURCES})
set(FA3_INCLUDE_DIRS
${FA3_SRC}
${cutlass_fa3_SOURCE_DIR}/include
${cutlass_for_fa3_SOURCE_DIR}/include
${cutlass_for_fa3_SOURCE_DIR}/tools/util/include
${CMAKE_CURRENT_SOURCE_DIR}/csrc)
# Set SM90a gencode flags for all FA3 CUDA sources
set_gencode_flags_for_srcs(
SRCS "${FA3_ALL_SOURCES}"
CUDA_ARCHS "${CUTLASS_FA3_ARCHS}")
define_extension_target(_cutlass_fa3_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${FA3_ALL_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${FA3_INCLUDE_DIRS}
USE_SABI 3
WITH_SOABI)
# FA3-specific compile options for CUDA and C++ source files:
# - C++17 required by CUTLASS
# - Fast math for performance
# - Relaxed constexpr for CUTLASS template metaprogramming
# - Disable backward pass, dropout, uneven K (not needed for inference)
# - Enable varlen-only mode (all our use cases are variable-length)
target_compile_options(_cutlass_fa3_C PRIVATE
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
$<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
target_compile_definitions(_cutlass_fa3_C PRIVATE
CUTE_USE_PACKED_TUPLE=1
CUTLASS_ENABLE_GDC_FOR_SM90
CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED
CUTLASS_ENABLE_TENSOR_CORE_MMA=1
FLASHATTENTION_DISABLE_BACKWARD
FLASHATTENTION_DISABLE_DROPOUT
FLASHATTENTION_DISABLE_UNEVEN_K
FLASHATTENTION_DISABLE_SM8x
FLASHATTENTION_VARLEN_ONLY)
message(STATUS "CUTLASS FA3 MLA Sparse: enabled for SM90 (${CUTLASS_FA3_ARCHS})")
+2 -22
View File
@@ -724,28 +724,6 @@ void reshape_and_cache_flash(
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (kv_cache_dtype == "nvfp4") {
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
// NVFP4 dispatch is compiled separately for SM100+.
extern void reshape_and_cache_nvfp4_dispatch(
torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
torch::Tensor & value_cache, torch::Tensor & slot_mapping,
torch::Tensor & k_scale, torch::Tensor & v_scale);
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
slot_mapping, k_scale, v_scale);
return;
#else
TORCH_CHECK(false,
"NVFP4 KV cache requires SM100+ (Blackwell). "
"Please rebuild vllm with a Blackwell-compatible CUDA target.");
#endif
}
// Original FP8/auto path.
int block_size = key_cache.size(1);
int64_t key_stride = key.stride(0);
@@ -763,6 +741,8 @@ void reshape_and_cache_flash(
dim3 grid(num_tokens);
dim3 block(std::min(num_heads * head_size, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
CALL_RESHAPE_AND_CACHE_FLASH);
-71
View File
@@ -1,71 +0,0 @@
#include "cpu_types.hpp"
#include <array>
#include <cstdint>
#include <mutex>
#include <string>
#include <ATen/ops/empty.h>
#include <ATen/ops/gelu.h>
#include <c10/util/BFloat16.h>
constexpr uint32_t ActivationLutSize = 1u << 16;
at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
void maybe_init_activation_lut_bf16(
uint16_t* lut, std::once_flag& once,
at::Tensor (*activation)(const at::Tensor&)) {
std::call_once(once, [&]() {
auto lut_input =
at::empty({static_cast<int64_t>(ActivationLutSize)},
at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
auto* lut_input_ptr = lut_input.data_ptr<float>();
#pragma omp parallel for
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
}
auto lut_output = activation(lut_input);
const auto* lut_output_ptr = lut_output.data_ptr<float>();
#pragma omp parallel for
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
}
});
}
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const uint16_t* lut, const char* op_name) {
TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
": input must be bfloat16");
TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
": out must be bfloat16");
TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
const auto* src =
reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
const int64_t n = input.numel();
CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
#pragma omp parallel for
for (int64_t i = 0; i < n; ++i) {
dst[i] = lut[src[i]];
}
CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
}
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const std::string& activation) {
if (activation == "gelu") {
static std::array<uint16_t, ActivationLutSize> lut{};
static std::once_flag once;
maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
activation_lut_bf16(out, input, lut.data(), "gelu_lut");
return;
}
TORCH_CHECK(false, "Unsupported activation: ", activation);
}
-3
View File
@@ -147,9 +147,6 @@ struct AttentionMetadata {
case ISA::NEON:
ss << "NEON, ";
break;
case ISA::VXE:
ss << "VXE, ";
break;
}
ss << "workitem_group_num: " << workitem_group_num
<< ", reduction_item_num: " << reduction_item_num
-16
View File
@@ -85,9 +85,6 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
at::Tensor& w_scales,
std::optional<at::Tensor> bias);
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
const std::string& activation);
torch::Tensor get_scheduler_metadata(
const int64_t num_req, const int64_t num_heads_q,
const int64_t num_heads_kv, const int64_t head_dim,
@@ -141,8 +138,6 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
void eagle_prepare_inputs_padded_kernel_impl(
const torch::Tensor& cu_num_draft_tokens,
@@ -236,15 +231,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
#if (defined(__aarch64__) && !defined(__APPLE__))
ops.def(
"activation_lut_bf16(Tensor! out, Tensor input, str activation)"
" -> ()");
ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
#endif // (defined(__aarch64__) && !defined(__APPLE__))
// Layernorm
// Apply Root Mean Square (RMS) Normalization to the input tensor.
ops.def(
@@ -433,8 +419,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
ops.def(
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
+6 -73
View File
@@ -13,80 +13,13 @@
#include "cpu/utils.hpp"
#ifdef VLLM_NUMA_DISABLED
void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
#else
void init_cpu_memory_env(std::vector<int64_t> node_ids) {
// Memory node binding
if (numa_available() != -1) {
// Concatenate all node_ids into a single comma-separated string
if (!node_ids.empty()) {
std::string node_ids_str;
for (const int node_id : node_ids) {
if (!node_ids_str.empty()) {
node_ids_str += ",";
}
node_ids_str += std::to_string(node_id);
}
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
bitmask* src_mask = numa_get_mems_allowed();
int pid = getpid();
if (mask && src_mask) {
// move all existing pages to the specified numa node.
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
int page_num = numa_migrate_pages(pid, src_mask, mask);
if (page_num == -1) {
TORCH_WARN("numa_migrate_pages failed. errno: " +
std::to_string(errno));
}
// Restrict memory allocation to the selected NUMA node(s).
// Enhances memory locality for the threads bound to those NUMA CPUs.
if (node_ids.size() > 1) {
errno = 0;
numa_set_interleave_mask(mask);
if (errno != 0) {
TORCH_WARN("numa_set_interleave_mask failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using INTERLEAVE policy for memory "
"allocation across multiple NUMA nodes (nodes: " +
node_ids_str +
"). Memory allocations will be "
"interleaved across the specified NUMA nodes.");
}
} else {
errno = 0;
numa_set_membind(mask);
if (errno != 0) {
TORCH_WARN("numa_set_membind failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using MEMBIND policy for memory "
"allocation on the NUMA nodes (" +
node_ids_str +
"). Memory allocations will be "
"strictly bound to these NUMA nodes.");
}
}
numa_set_strict(1);
numa_free_nodemask(mask);
numa_free_nodemask(src_mask);
} else {
TORCH_WARN(
"numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
std::to_string(errno));
}
}
}
std::string init_cpu_threads_env(const std::string& cpu_ids) {
return std::string(
"Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
"no effect to setup thread affinity.");
}
#endif // VLLM_NUMA_DISABLED
#endif
namespace cpu_utils {
ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
-22
View File
@@ -54,34 +54,12 @@ struct Counter {
};
inline int64_t get_available_l2_size() {
#if defined(__s390x__)
static int64_t size = []() {
uint32_t l2_cache_size = 0;
auto caps = at::cpu::get_cpu_capabilities();
auto it = caps.find("l2_cache_size");
if (it != caps.end()) {
l2_cache_size = static_cast<uint32_t>(it->second.toInt());
}
if (l2_cache_size == 0) {
long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
if (sys_l2 > 0) {
l2_cache_size = static_cast<uint32_t>(sys_l2);
}
}
if (l2_cache_size == 0) {
l2_cache_size = 256 * 1024;
}
return static_cast<int64_t>(l2_cache_size) >> 1; // use 50% of L2 cache
}();
return size;
#else
static int64_t size = []() {
auto caps = at::cpu::get_cpu_capabilities();
const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
return l2_cache_size >> 1; // use 50% of L2 cache
}();
return size;
#endif
}
template <int32_t alignment_v, typename T>
-72
View File
@@ -1,72 +0,0 @@
/* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* Vendored CUTLASS FA3 MLA attention kernel binding for vLLM.
* Based on sgl-kernel/csrc/flash_extension.cc from SGLang.
*
* This registers the FA3 forward pass as a PyTorch C++ extension under
* the _cutlass_fa3_C namespace, enabling torch.ops._cutlass_fa3_C.fwd().
*
* Original source:
* https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
* sgl-kernel/csrc/flash_extension.cc
*/
#include <Python.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <torch/all.h>
#include <torch/library.h>
#include "sgl_flash_kernel_ops.h"
TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {
/*
* CUTLASS FA3 MLA forward pass.
* Signature matches sgl-attn's mha_fwd() exactly.
*/
m.def(
"fwd(Tensor q,"
" Tensor k,"
" Tensor v,"
" Tensor? k_new,"
" Tensor? v_new,"
" Tensor? q_v,"
" Tensor? out,"
" Tensor? cu_seqlens_q,"
" Tensor? cu_seqlens_k,"
" Tensor? cu_seqlens_k_new,"
" Tensor? seqused_q,"
" Tensor? seqused_k,"
" int? max_seqlen_q,"
" int? max_seqlen_k,"
" Tensor? page_table,"
" Tensor? kv_batch_idx,"
" Tensor? leftpad_k,"
" Tensor? rotary_cos,"
" Tensor? rotary_sin,"
" Tensor? seqlens_rotary,"
" Tensor? q_descale,"
" Tensor? k_descale,"
" Tensor? v_descale,"
" float? softmax_scale,"
" bool is_causal,"
" int window_size_left,"
" int window_size_right,"
" int attention_chunk,"
" float softcap,"
" bool is_rotary_interleaved,"
" Tensor? scheduler_metadata,"
" int num_splits,"
" bool? pack_gqa,"
" int sm_margin,"
" Tensor? sinks"
") -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
}
// Python module initialization for _cutlass_fa3_C
PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
nullptr, 0, nullptr};
return PyModule_Create(&module);
}
-275
View File
@@ -1,275 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
// NVFP4 KV cache store kernel.
// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
// into the paged KV cache.
//
// Per page layout: [K_data | K_scale | V_data | V_scale]
// Both data and scale regions are contiguous per head, enabling direct
// TMA descriptor use.
//
// Reuses device functions from nvfp4_utils.cuh:
// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale
// - pack_fp4() for packing float pairs to fp4
// - reciprocal_approximate_ftz() for fast reciprocal
#define NVFP4_ENABLE_ELTS16 1
#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "dispatch_utils.h"
namespace vllm {
// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
// The swizzle pattern for HND layout is:
// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
// where T = block_size (page_size), S = scale_dim = head_size // 16.
//
// For a linear (t, s) position, the swizzled position is:
// swizzled_t = (t / 4) * 4 + (s / (S / 4))
// swizzled_s = (s % (S / 4)) * 4 + (t % 4)
__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
int scale_dim) {
int s_group = scale_dim / 4;
int swizzled_t = (t / 4) * 4 + (s / s_group);
int swizzled_s = (s % s_group) * 4 + (t % 4);
return swizzled_t * scale_dim + swizzled_s;
}
// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
//
// Takes separate data and scale cache pointers for K and V.
// Within each KV side, data and scale are separate contiguous regions.
//
// Threading: one CUDA block per token, threads process heads and
// groups of 16 elements within each head.
template <typename scalar_t>
__global__ void reshape_and_cache_nvfp4_kernel(
const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
uint8_t* __restrict__ key_data_cache, // data region for K
uint8_t* __restrict__ value_data_cache, // data region for V
uint8_t* __restrict__ key_scale_cache, // scale region for K
uint8_t* __restrict__ value_scale_cache, // scale region for V
const int64_t* __restrict__ slot_mapping, // [num_actual_tokens]
const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale
const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale
const int64_t key_stride, // key.stride(0) in elements
const int64_t value_stride, // value.stride(0) in elements
const int num_heads, const int head_size, const int block_size,
const int64_t data_block_stride, // data cache stride for dim 0
const int64_t data_head_stride, // data cache stride for heads
const int64_t data_block_offset_stride, // data cache stride for tokens
const int64_t scale_block_stride, // scale cache stride for dim 0
const int64_t scale_head_stride, // scale cache stride for heads
const int64_t scale_block_offset_stride // scale cache stride for tokens
) {
using CudaType = typename CUDATypeConverter<scalar_t>::Type;
using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8
static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
if (slot_idx < 0) return;
const int64_t block_idx = slot_idx / block_size;
const int block_offset = static_cast<int>(slot_idx % block_size);
const int scale_dim = head_size / 16;
const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
const int total_groups = num_heads * groups_per_head;
const int tid = threadIdx.x;
const int num_thread_groups = blockDim.x / THREADS_PER_SF;
const int tg_id = tid / THREADS_PER_SF;
const int tg_lane = tid % THREADS_PER_SF;
// Process both K (kv=0) and V (kv=1)
#pragma unroll
for (int kv = 0; kv < 2; kv++) {
const scalar_t* __restrict__ src = (kv == 0) ? key : value;
const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
uint8_t* __restrict__ data_cache =
(kv == 0) ? key_data_cache : value_data_cache;
uint8_t* __restrict__ sc_cache =
(kv == 0) ? key_scale_cache : value_scale_cache;
// Source pointer for this token (use actual stride, not assumed contiguous)
const CudaType* __restrict__ token_src =
reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
// Destination bases in data and scale caches for this token's block
uint8_t* __restrict__ data_block =
data_cache + block_idx * data_block_stride;
uint8_t* __restrict__ scale_block =
sc_cache + block_idx * scale_block_stride;
for (int g = tg_id; g < total_groups; g += num_thread_groups) {
const int head = g / groups_per_head;
const int group_in_head = g % groups_per_head;
// Load 16 (or 8) bf16 elements from source
PVec in_vec;
const CudaType* __restrict__ src_ptr =
token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
tg_lane * ELTS;
#pragma unroll
for (int i = 0; i < ELTS / 2; i++) {
in_vec.elts[i] = reinterpret_cast<
const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
}
// Quantize: produces packed fp4 and writes scale factor.
uint8_t sf_val;
uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
in_vec, global_scale, sf_out_ptr);
// Write packed FP4 data to data cache
uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
block_offset * data_block_offset_stride;
#if CVT_FP4_PACK16
{
// 16 elements → 8 bytes (u32x2)
int data_byte_offset = group_in_head * 8;
reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
(uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
}
#else
{
// 8 elements → 4 bytes (uint32_t)
int data_byte_offset =
group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
}
#endif
// Write block scale to scale cache.
// K (kv==0): linear layout (no swizzle).
// V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
if (sf_out_ptr != nullptr) {
int scale_idx = group_in_head;
uint8_t* __restrict__ scale_dst;
if (kv == 0) {
scale_dst = scale_block + head * scale_head_stride +
block_offset * scale_block_offset_stride + scale_idx;
} else {
int swizzled_offset =
swizzle_scale_offset(block_offset, scale_idx, scale_dim);
int swizzled_t = swizzled_offset / scale_dim;
int swizzled_s = swizzled_offset % scale_dim;
scale_dst = scale_block + head * scale_head_stride +
swizzled_t * scale_block_offset_stride + swizzled_s;
}
*scale_dst = sf_val;
}
}
}
}
} // namespace vllm
// Non-template entry point callable from cache_kernels.cu.
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
// Each KV side contains both data and scale:
// page = [K_data | K_scale | V_data | V_scale]
void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
torch::Tensor& k_scale,
torch::Tensor& v_scale) {
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
int data_dim = head_size / 2;
int scale_dim = head_size / 16;
int full_dim = data_dim + scale_dim;
// key_cache is kv_cache[:, 0] with shape
// [num_blocks, block_size, num_heads, full_dim] in logical order.
// Strides encode the physical layout (HND or NHD).
TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
TORCH_CHECK(key_cache.size(3) == full_dim,
"key_cache last dim must be data_dim + scale_dim, got ",
key_cache.size(3), " expected ", full_dim);
int block_size = key_cache.size(1);
TORCH_CHECK(head_size % 16 == 0,
"head_size must be divisible by 16 for NVFP4 KV cache");
TORCH_CHECK(block_size % 4 == 0,
"block_size must be divisible by 4 for NVFP4 KV cache swizzle");
// Detect physical layout from strides (based on full_dim).
// HND: head stride > block_offset stride.
bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
int64_t data_block_stride = key_cache.stride(0); // page_bytes
int64_t data_head_stride, data_block_offset_stride;
if (is_hnd) {
data_head_stride = (int64_t)block_size * data_dim;
data_block_offset_stride = data_dim;
} else {
data_head_stride = data_dim;
data_block_offset_stride = (int64_t)num_heads * data_dim;
}
// Page layout: [K_data | K_scale | V_data | V_scale]
// Scale follows data within each KV side.
int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
// Scale strides: same page stride, inner strides from layout.
int64_t scale_block_stride = data_block_stride;
int64_t scale_head_stride, scale_block_offset_stride;
if (is_hnd) {
scale_head_stride = (int64_t)block_size * scale_dim;
scale_block_offset_stride = scale_dim;
} else {
scale_head_stride = scale_dim;
scale_block_offset_stride = (int64_t)num_heads * scale_dim;
}
const float* k_scale_ptr = k_scale.data_ptr<float>();
const float* v_scale_ptr = v_scale.data_ptr<float>();
int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
int total_groups = num_heads * groups_per_head;
constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
num_threads = ((num_threads + 31) / 32) * 32;
dim3 grid(num_tokens);
dim3 block(num_threads);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_REDUCED_FLOATING_TYPES(
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
<<<grid, block, 0, stream>>>(
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
key_scale_ptr, value_scale_ptr,
slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
key.stride(0), value.stride(0), num_heads, head_size,
block_size, data_block_stride, data_head_stride,
data_block_offset_stride, scale_block_stride, scale_head_stride,
scale_block_offset_stride);
});
}
-1
View File
@@ -1,7 +1,6 @@
#pragma once
#include <optional>
#include <string>
#include <torch/library.h>
#include <tuple>
-45
View File
@@ -1,45 +0,0 @@
/* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
*
* Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
* Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
* NO MODIFICATIONS from the original (except removing unused macros).
*/
#pragma once
#include <ATen/ATen.h>
#include <ATen/Tensor.h>
#include <torch/library.h>
#include <torch/torch.h>
#include <vector>
#include "sgl_kernel_torch_shim.h"
/*
* From flash-attention (sgl-attn fork)
*/
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
at::Tensor q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
at::Tensor k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
at::Tensor v, // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
std::optional<at::Tensor> q_v_, // MLA value projection query
std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
std::optional<at::Tensor> cu_seqlens_k_,
std::optional<at::Tensor> cu_seqlens_k_new_,
std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
std::optional<at::Tensor> page_table_,
std::optional<at::Tensor> kv_batch_idx_,
std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
std::optional<at::Tensor> rotary_sin_,
std::optional<at::Tensor> seqlens_rotary_,
std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
bool is_causal, int64_t window_size_left, int64_t window_size_right,
int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
std::optional<bool> pack_gqa_, int64_t sm_margin,
std::optional<const at::Tensor>& sinks_);
-121
View File
@@ -1,121 +0,0 @@
/* Adapted from:
* https://github.com/neuralmagic/vllm-flash-attention/blob/90eacc1af2a7c3de62ea249e929ed5faccf38954/csrc/common/pytorch_shim.h
*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
*
* Vendored from sgl-kernel/include/sgl_kernel_torch_shim.h (commit bcf72ccc).
* Provides make_pytorch_shim() template for PyTorch op registration type
* conversion. NO MODIFICATIONS from the original.
*/
#pragma once
#include <torch/library.h>
/**
* Unfortunately, the type signatures of the flash_attn ops are not compatible
* with the PyTorch library bindings. To get around that we use
* `make_pytorch_shim` which creates a lambda that exposes the API using
* PyTorch compatible types to the types, then converts them to the types
* expected by the flash_attn ops. This shims allows us to make minimal changes
* to `flash_api.cpp` making it easier to synchronize with upstream changes.
*
* The `pytorch_library_compatible_type` struct is used to map from the
* flash_attn ops types to a PyTorch library compatible one. The main issues is
* that the following types are not support by PyTorch library bindings:
* - `int`
* - `float`
* - `std::optional<T> &`
* - `std::optional<const at::Tensor> &`
* So we convert them to (respectively):
* - `int64_t`
* - `double`
* - `const std::optional<T>&`
* - `const std::optional<at::Tensor>&`
*/
template <typename T>
struct pytorch_library_compatible_type {
using type = T;
static T convert_from_type(T arg) { return arg; }
};
template <typename T>
using pytorch_library_compatible_type_t =
typename pytorch_library_compatible_type<T>::type;
template <typename T>
T convert_from_pytorch_compatible_type(
pytorch_library_compatible_type_t<T> arg) {
return pytorch_library_compatible_type<T>::convert_from_type(arg);
}
// Map `c10::optional<T> &` -> `const c10::optional<T>&`
// (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
// the optional container)
template <typename T>
struct pytorch_library_compatible_type<c10::optional<T>&> {
using type = const c10::optional<T>&;
static c10::optional<T>& convert_from_type(const c10::optional<T>& arg) {
return const_cast<c10::optional<T>&>(arg);
}
};
// Map `c10::optional<T>` ->
// `c10::optional<pytorch_library_compatible_type_t<T>>`
// (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
template <typename T>
struct pytorch_library_compatible_type<c10::optional<T>> {
using type = c10::optional<pytorch_library_compatible_type_t<T>>;
static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
c10::optional<T> arg) {
return arg;
}
};
// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
template <>
struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
using type = const c10::optional<at::Tensor>&;
static c10::optional<const at::Tensor>& convert_from_type(
const c10::optional<at::Tensor>& arg) {
return const_cast<c10::optional<const at::Tensor>&>(
reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
}
};
// Map `int` -> `int64_t`
template <>
struct pytorch_library_compatible_type<int> {
using type = int64_t;
static int convert_from_type(int64_t arg) {
TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
"int64_t value is too large to be converted to int");
TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
"int64_t value is too small to be converted to int");
return arg;
}
};
// Map `float` -> `double`
template <>
struct pytorch_library_compatible_type<float> {
using type = double;
static float convert_from_type(double arg) {
TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
"double value is too large to be converted to float");
return arg;
}
};
//
// Shim Utils
//
template <typename Ret, typename... Args>
auto make_pytorch_shim(Ret (*fun)(Args... args)) {
return [fun](pytorch_library_compatible_type_t<Args>... args) {
return fun(convert_from_pytorch_compatible_type<Args>(args)...);
};
}
+4 -5
View File
@@ -642,7 +642,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
else \
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
fi; \
uv pip install --system accelerate modelscope \
uv pip install --system accelerate hf_transfer modelscope \
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
# ============================================================
@@ -756,10 +756,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -e tests/vllm_test_utils
# 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
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system hf_transfer
ENV HF_HUB_ENABLE_HF_TRANSFER 1
# Copy in the v1 package for testing (it isn't distributed yet)
COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
+1 -8
View File
@@ -173,8 +173,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/lint.txt && \
uv pip install -r requirements/test/cpu.txt && \
uv pip install -r requirements/dev.txt && \
pre-commit install --hook-type pre-commit --hook-type commit-msg
ENTRYPOINT ["bash"]
@@ -198,12 +197,6 @@ ADD ./.buildkite/ ./.buildkite/
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -e tests/vllm_test_utils
# 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
######################### RELEASE IMAGE #########################
FROM base AS vllm-openai
+3 -4
View File
@@ -272,10 +272,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -e tests/vllm_test_utils
# 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
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system hf_transfer
ENV HF_HUB_ENABLE_HF_TRANSFER 1
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system -r requirements/test/nightly-torch.txt
+3 -4
View File
@@ -365,10 +365,9 @@ RUN cd /vllm-workspace \
&& 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
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system hf_transfer
ENV HF_HUB_ENABLE_HF_TRANSFER=1
# install audio decode package `torchcodec` from source (required due to
# ROCm and torch version mismatch) for tests with datasets package
+34 -34
View File
@@ -42,7 +42,7 @@ FROM python-install AS pyarrow
# Build Apache Arrow
WORKDIR /tmp
RUN --mount=type=cache,target=/root/.cache/uv \
git clone https://github.com/apache/arrow.git -b maint-19.0.1 && \
git clone https://github.com/apache/arrow.git && \
cd arrow/cpp && \
mkdir release && cd release && \
cmake -DCMAKE_BUILD_TYPE=Release \
@@ -68,6 +68,19 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements-build.txt && \
python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
FROM python-install AS numa-build
# Install numactl (needed for numa.h dependency)
WORKDIR /tmp
RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \
tar -xvzf v2.0.16.tar.gz && \
cd numactl-2.0.16 && \
./autogen.sh && \
./configure && \
make
# Set include path
ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
FROM python-install AS rust
ENV CARGO_HOME=/root/.cargo
ENV RUSTUP_HOME=/root/.rustup
@@ -78,18 +91,6 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \
rustup default stable && \
rustup show
FROM python-install AS numa-build
WORKDIR /tmp
RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \
tar -xvzf v2.0.19.tar.gz && \
cd numactl-2.0.19 && \
./autogen.sh && \
./configure && \
make
# Set include path
ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
FROM python-install AS torch-vision
# Install torchvision
ARG TORCH_VISION_VERSION=v0.26.0
@@ -132,7 +133,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
cd llvm-project && mkdir build && cd build && \
uv pip install 'cmake<4' 'setuptools<70' numpy && \
uv pip install 'cmake<4' setuptools numpy && \
export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \
CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \
CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \
@@ -192,22 +193,27 @@ RUN --mount=type=cache,target=/root/.cache/uv \
cd opencv-python && \
python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist
## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330
FROM python-install AS guidance-builder
# Build Outlines Core
FROM python-install AS outlines-core-builder
WORKDIR /tmp
ENV CARGO_HOME=/root/.cargo
ENV RUSTUP_HOME=/root/.rustup
ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
COPY requirements/common.txt /tmp/requirements/common.txt
ARG OUTLINES_CORE_VERSION
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
--mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
git clone https://github.com/guidance-ai/llguidance.git && \
cd llguidance && \
git checkout s390x-fix-v2 && \
OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \
if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \
git clone https://github.com/dottxt-ai/outlines-core.git && \
cd outlines-core && \
git checkout tags/${OUTLINES_CORE_VERSION} && \
sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \
uv pip install maturin && \
python -m maturin build --release --out dist --compatibility linux
python -m maturin build --release --out dist
# # Final build stage
# Final build stage
FROM python-install AS vllm-cpu
ARG PYTHON_VERSION
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
@@ -223,12 +229,10 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/
ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
# Force pure Python protobuf to avoid s390x C++ extension crashes
ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
COPY . /workspace/vllm
WORKDIR /workspace/vllm
RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \
make -C /numactl install
# Install dependencies, including PyTorch and Apache Arrow
@@ -241,22 +245,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \
--mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \
--mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \
--mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \
ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \
--mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \
ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \
VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \
HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \
LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
uv pip install -v \
$ARROW_WHL_FILE \
OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \
uv pip install -v \
$ARROW_WHL_FILE \
$VISION_WHL_FILE \
$HF_XET_WHL_FILE \
$LLVM_WHL_FILE \
$NUMBA_WHL_FILE \
$OPENCV_WHL_FILE \
$GUIDANCE_WHL_FILE \
$OUTLINES_CORE_WHL_FILE \
--index-strategy unsafe-best-match \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt
@@ -267,10 +271,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
uv pip install "$(echo dist/*.whl)[tensorizer]"
# Remove protobuf C++ extension that crashes on s390x
RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \
/opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true
# setup non-root user for vllm
RUN umask 002 && \
/usr/sbin/useradd --uid 2000 --gid 0 vllm && \
-64
View File
@@ -37,7 +37,6 @@ th {
| HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` |
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
@@ -240,69 +239,6 @@ vllm bench serve \
--spec-bench-category "summarization"
```
#### SPEED-Bench Benchmark with Speculative Decoding
[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k).
!!! note
This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets.
First, download the dataset to a folder, using this one liner:
```bash
curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
```
The command supports also the following arguments:
- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets.
- `--output_dir`: download to a specified folder. By default, it will download to the current directory.
Start a server with speculative decoding:
```bash
vllm serve meta-llama/Llama-3.3-70B-Instruct \
--speculative-config $'{"method": "eagle3",
"num_speculative_tokens": 3,
"model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}'
```
Run all categories in the Qualitative split:
```bash
vllm bench serve \
--model meta-llama/Llama-3.3-70B-Instruct \
--dataset-name speed_bench \
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
--num-prompts -1
```
Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`.
Run only a specific category like "multilingual":
```bash
vllm bench serve \
--model meta-llama/Llama-3.3-70B-Instruct \
--dataset-name speed_bench \
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
--num-prompts -1
--speed-bench-category "multilingual"
```
Run all categories in the Throughput split (2k ISL):
```bash
vllm bench serve \
--model meta-llama/Llama-3.3-70B-Instruct \
--dataset-name speed_bench \
--speed-bench-dataset-subset throughput_2k
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
--num-prompts -1
```
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
#### Other HuggingFaceDataset Examples
```bash
+10 -3
View File
@@ -106,7 +106,6 @@ Priority is **1 = highest** (tried first).
| 2 | `FLASH_ATTN` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
**Ampere/Hopper (SM 8.x-9.x):**
@@ -116,7 +115,6 @@ Priority is **1 = highest** (tried first).
| 2 | `FLASHINFER` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
### MLA Attention (DeepSeek-style)
@@ -132,6 +130,16 @@ Priority is **1 = highest** (tried first).
| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
| 7 | `FLASHMLA_SPARSE` |
**Ampere/Hopper (SM 8.x-9.x):**
| Priority | Backend |
| -------- | ------- |
| 1 | `FLASH_ATTN_MLA` |
| 2 | `FLASHMLA` |
| 3 | `FLASHINFER_MLA` |
| 4 | `TRITON_MLA` |
| 5 | `FLASHMLA_SPARSE` |
> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
>
> **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -199,7 +207,6 @@ configuration.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+2 -2
View File
@@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us
### Grafana Dashboard
vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important:
@@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing:
- Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
- Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
- [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
- [User-facing docs](../../examples/observability/opentelemetry/README.md)
- [User-facing docs](../../examples/online_serving/opentelemetry/README.md)
- [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
- [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview)
-1
View File
@@ -16,7 +16,6 @@ The following are the supported quantization formats for vLLM:
- [INT8 W8A8](int8.md)
- [FP8 W8A8](fp8.md)
- [NVIDIA Model Optimizer](modelopt.md)
- [Online Quantization](online.md)
- [AMD Quark](quark.md)
- [Quantized KV Cache](quantized_kvcache.md)
- [TorchAO](torchao.md)
-94
View File
@@ -1,94 +0,0 @@
# Online Quantization
Online quantization lets you take a BF16/FP16 model and quantize its Linear
and MoE weights to lower precision (such as FP8) at load time, without needing
a pre-quantized checkpoint or calibration data. Weights are converted during
model loading and activations are dynamically scaled during each forward pass.
## Quick Start
Pass a scheme name to the `quantization` parameter:
```python
from vllm import LLM
# Per-tensor FP8 quantization (one scale per weight tensor)
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
```
Or with the CLI:
```bash
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
```
## Supported Schemes
| Scheme | Weight recipe | Activation recipe | Notes |
| ------ | ------------- | ------------------ | ----- |
| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
Support for additional schemes will be added in future versions of vllm.
## Advanced Configuration
For fine-grained control, use a `quantization_config` dictionary.
### Separate Schemes for Dense and MoE Layers
You can apply different quantization schemes to dense linear layers and MoE expert layers:
```python
from vllm import LLM
llm = LLM(
"ibm-granite/granite-3.0-1b-a400m-base",
quantization="fp8_per_tensor",
quantization_config={
"linear_scheme_override": "fp8_per_block",
},
)
```
Or,
```python
from vllm import LLM
llm = LLM(
"ibm-granite/granite-3.0-1b-a400m-base",
quantization="fp8_per_tensor",
quantization_config={
"moe_scheme_override": "fp8_per_block",
},
)
```
### Excluding Layers from Quantization
Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`):
```python
from vllm import LLM
llm = LLM(
"ibm-granite/granite-3.0-1b-a400m-base",
quantization="fp8_per_tensor",
quantization_config={
"ignore": [
# exact layer name
"model.layers.1.self_attn.o_proj",
# regex: skip all QKV projections
"re:.*[qkv]_proj",
],
},
)
```
!!! note
For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name.
@@ -3,15 +3,15 @@
vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform.
Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16.
Currently, the CPU implementation for s390x architecture supports FP32 datatype only.
--8<-- [end:installation]
--8<-- [start:requirements]
- OS: `Linux`
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
- SDK: `gcc/g++ >= 12.3.0` or later with Command Line Tools
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
- Build install python packages: `pyarrow`, `torch` and `torchvision`
--8<-- [end:requirements]
--8<-- [start:set-up-using-python]
@@ -24,14 +24,13 @@ Currently, there are no pre-built IBM Z CPU wheels.
--8<-- [end:pre-built-wheels]
--8<-- [start:build-wheel-from-source]
Install the following packages from the package manager before building the vLLM. For example on RHEL 9.6:
Install the following packages from the package manager before building the vLLM. For example on RHEL 9.4:
```bash
dnf install -y \
which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
which procps findutils tar vim git gcc g++ make patch make cython zlib-devel \
libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
clang llvm-devel llvm-static clang-devel
openssl-devel openblas openblas-devel wget autoconf automake libtool cmake numactl-devel
```
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
@@ -44,13 +43,13 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
Execute the following commands to build and install vLLM from source.
!!! tip
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM.
```bash
sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds
uv pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
--torch-backend auto \
-r requirements/build/cpu.txt \
-r requirements/build/cuda.txt \
-r requirements/cpu.txt \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
uv pip install dist/*.whl
@@ -58,9 +57,10 @@ Execute the following commands to build and install vLLM from source.
??? console "pip"
```bash
sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds
pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
-r requirements/build/cpu.txt \
--extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-r requirements/build/cuda.txt \
-r requirements/cpu.txt \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
pip install dist/*.whl
@@ -240,7 +240,7 @@ uv pip install vllm==${VLLM_VERSION} \
# Install dependencies
pip install --upgrade numba \
scipy \
huggingface-hub[cli] \
huggingface-hub[cli,hf_transfer] \
setuptools_scm
pip install -r requirements/rocm.txt
+6 -9
View File
@@ -14,7 +14,6 @@ Sorted alphabetically by GitHub handle:
- [@aarnphm](https://github.com/aarnphm): Structured output
- [@alexm-redhat](https://github.com/alexm-redhat): Performance
- [@ApostaC](https://github.com/ApostaC): Connectors, offloading
- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
- [@benchislett](https://github.com/benchislett): Engine core and spec decode
- [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
- [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
@@ -32,7 +31,6 @@ Sorted alphabetically by GitHub handle:
- [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
- [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed
- [@markmc](https://github.com/markmc): Observability
- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance
- [@mgoin](https://github.com/mgoin): Quantization and performance
- [@NickLucche](https://github.com/NickLucche): KV connector
- [@njhill](https://github.com/njhill): Distributed, API server, engine core
@@ -43,7 +41,6 @@ Sorted alphabetically by GitHub handle:
- [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
- [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
- [@russellb](https://github.com/russellb): Structured output, engine core, security
- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
- [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
- [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
- [@tdoublep](https://github.com/tdoublep): State space models
@@ -89,7 +86,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
- @robertgshaw2-redhat, @njhill, @russellb
- ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
- KV Connector: Connector interface and implementation for KV cache offload and transfer
- @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC
- Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks
@@ -108,7 +105,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms
- @ProExpertProg
- Attention: Attention interface for paged attention
- @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni
- @WoosukKwon, @LucasWilkinson, @heheda12345
- FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
- @tlrmchlsmth
- Quantization: Various quantization config, weight loading, and kernel.
@@ -122,7 +119,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- State space models: The state space models implementation in vLLM
- @tdoublep, @tlrmchlsmth
- Reasoning and tool calling parsers
- @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning
- @chaunceyjiang, @aarnphm
### Entrypoints
@@ -136,7 +133,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
### Features
- Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
- @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni
- @WoosukKwon, @benchislett, @luccafong
- Structured Output: The structured output implementation
- @russellb, @aarnphm
- RL: The RL related features such as collective rpc, sleep mode, etc.
@@ -156,8 +153,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
### External Kernels Integration
- FlashAttention: @LucasWilkinson, @MatthewBonanni
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
- FlashAttention: @LucasWilkinson
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
- Blackwell Kernels: @mgoin, @yewentao256
- DeepEP/DeepGEMM: @mgoin, @yewentao256
-10
View File
@@ -59,16 +59,6 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
### Pooling Types
| Pooling Tasks | Granularity | Description |
|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `CLS` pooling | Sequence-wise | For BERTlike (bidirectional selfattention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. |
| `LAST` pooling | Sequence-wise | For GPTlike (causal selfattention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output. |
| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling. |
| `ALL` pooling | Token-wise | Outputs the last_hidden_states for all input tokens. |
| `STEP` pooling | Token-wise | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids. |
### Score Types
The scoring models is designed to compute similarity scores between two input prompts. It supports three model types
-7
View File
@@ -45,7 +45,6 @@ You can compute pairwise similarity scores to build a similarity matrix using th
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
| `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | |
| `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | |
| `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ |
| `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
| `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
@@ -74,12 +73,6 @@ You can compute pairwise similarity scores to build a similarity matrix using th
!!! note
`jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights.
!!! note
`jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters
(`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the
selected adapter into the base weights at load time. Choose the task with
`--hf-overrides '{"jina_task": "<task>"}'`; the default is `retrieval`.
### Multimodal Models
!!! note
-4
View File
@@ -160,8 +160,6 @@ The following Score API parameters are supported:
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params"
```
#### Examples
@@ -372,8 +370,6 @@ The following rerank api parameters are supported:
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
--8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
--8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params"
```
#### Examples
+1 -1
View File
@@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the
Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
### Reward Models
### As Reward Models
Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
-8
View File
@@ -400,7 +400,6 @@ th {
| `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
| `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
| `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
| `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
@@ -555,7 +554,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
| `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
| `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
| `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
@@ -635,7 +633,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
<sup>+</sup> Multiple items can be inputted per text prompt for this modality.
<sup>*</sup> Only specific variants of the model support this modality (see notes below).</br>
!!! note
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
@@ -646,11 +643,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
!!! note
For `Gemma4ForConditionalGeneration`:
- audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
- The model does not ingest videos directly. However, vLLMs Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
!!! note
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
+18 -1
View File
@@ -467,11 +467,28 @@ It consists of two endpoints:
- `/tokenize` corresponds to calling `tokenizer.encode()`.
- `/detokenize` corresponds to calling `tokenizer.decode()`.
### Score API
#### Score Template
Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)).
Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template.
Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}`
This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future.
Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
### Generative Scoring API
The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes".
This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
**Requirements:**
@@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml
For detailed deployment instructions and platform-specific options, see:
- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import
- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping
- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import
- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping
## Contributing
+3 -5
View File
@@ -7,7 +7,7 @@ requests >= 2.26.0
tqdm
blake3
py-cpuinfo
transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
transformers >= 4.56.0, < 5
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
@@ -19,7 +19,7 @@ pillow # Required for image processing
prometheus-fastapi-instrumentator >= 7.0.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x" or platform_machine == "ppc64le"
outlines_core == 0.2.11
# required for outlines backend disk cache
diskcache == 5.6.3
@@ -32,14 +32,12 @@ pyzmq >= 25.0.0
msgspec
gguf >= 0.17.0
mistral_common[image] >= 1.11.0
av # required for audio in video IO
opencv-python-headless >= 4.13.0 # required for video IO
soundfile # required for audio IO
pyyaml
six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
einops # Required for Qwen2-VL.
compressed-tensors == 0.15.0.1 # required for compressed-tensors
compressed-tensors == 0.14.0.1 # required for compressed-tensors
depyf==0.20.0 # required for profiling and debugging with compilation config
cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
watchfiles # required for http server to monitor the updates of TLS files
+3 -3
View File
@@ -1,5 +1,5 @@
lmcache >= 0.3.9
nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
nixl-cu12 >= 0.7.1, <= 0.10.1
nixl-cu13 >= 0.7.1, <= 0.10.1
nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
nixl-cu12 >= 0.7.1, < 0.10.0
nixl-cu13 >= 0.7.1, < 0.10.0
mooncake-transfer-engine >= 0.3.8
+4 -3
View File
@@ -18,9 +18,10 @@ httpx
librosa # required for audio tests
vector_quantize_pytorch # required for minicpmo_26 test
vocos # required for minicpmo_26 test
peft>=0.18.1 # required for phi-4-mm test
peft>=0.15.0 # required for phi-4-mm test
pqdm
ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
resampy # required for audio tests
sentence-transformers>=5.2.0 # required for embedding tests
soundfile # required for audio tests
jiwer # required for audio tests
@@ -38,8 +39,8 @@ opencv-python-headless >= 4.13.0 # required for video test
datamodel_code_generator # required for minicpm3 test
lm-eval[api]>=0.4.11 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
transformers==4.57.5
tokenizers==0.22.0
schemathesis>=3.39.15 # Required for openai schema test.
# quantization
bitsandbytes==0.49.2
+14 -10
View File
@@ -4,7 +4,7 @@ absl-py==2.1.0
# via
# rouge-score
# tensorboard
accelerate==1.13.0
accelerate==1.0.1
# via peft
aenum==3.1.16
# via lightly
@@ -248,6 +248,7 @@ filelock==3.16.1
# huggingface-hub
# ray
# torch
# transformers
# virtualenv
fiona==1.10.1
# via torchgeo
@@ -330,7 +331,7 @@ h5py==3.13.0
# via terratorch
harfile==0.3.0
# via schemathesis
hf-xet==1.4.3
hf-xet==1.1.7
# via huggingface-hub
hiredis==3.0.0
# via tensorizer
@@ -344,10 +345,9 @@ httpx==0.27.2
# via
# -r requirements/test/cuda.in
# diffusers
# huggingface-hub
# perceptron
# schemathesis
huggingface-hub==1.10.2
huggingface-hub==0.36.2
# via
# accelerate
# datasets
@@ -555,6 +555,7 @@ numba==0.61.2
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
# librosa
# resampy
numpy==2.2.6
# via
# -r requirements/test/cuda.in
@@ -595,6 +596,7 @@ numpy==2.2.6
# pyogrio
# pywavelets
# rasterio
# resampy
# rioxarray
# rouge-score
# runai-model-streamer
@@ -754,7 +756,7 @@ pathvalidate==3.2.1
# via pytablewriter
patsy==1.0.1
# via statsmodels
peft==0.18.1
peft==0.16.0
# via -r requirements/test/cuda.in
perceptron==0.1.4
# via -r requirements/test/cuda.in
@@ -980,7 +982,7 @@ referencing==0.35.1
# via
# jsonschema
# jsonschema-specifications
regex==2026.2.28
regex==2024.9.11
# via
# diffusers
# nltk
@@ -1000,6 +1002,7 @@ requests==2.32.3
# google-api-core
# google-cloud-storage
# gpt-oss
# huggingface-hub
# lightly
# lm-eval
# mistral-common
@@ -1012,7 +1015,10 @@ requests==2.32.3
# starlette-testclient
# tacoreader
# tiktoken
# transformers
# wandb
resampy==0.4.3
# via -r requirements/test/cuda.in
responses==0.25.3
# via genai-perf
rfc3339-validator==0.1.4
@@ -1210,7 +1216,7 @@ timm==1.0.17
# segmentation-models-pytorch
# terratorch
# torchgeo
tokenizers==0.22.2
tokenizers==0.22.0
# via
# -c requirements/common.txt
# -r requirements/test/cuda.in
@@ -1289,7 +1295,7 @@ tqdm==4.67.3
# tacoreader
# terratorch
# transformers
transformers==5.5.3
transformers==4.57.5
# via
# -c requirements/common.txt
# -r requirements/test/cuda.in
@@ -1311,9 +1317,7 @@ typepy==1.3.2
typer==0.15.2
# via
# fastsafetensors
# huggingface-hub
# perceptron
# transformers
types-python-dateutil==2.9.0.20241206
# via arrow
typeshed-client==2.8.2
+2 -2
View File
@@ -29,8 +29,8 @@ opencv-python-headless >= 4.13.0 # required for video test
datamodel_code_generator # required for minicpm3 test
lm-eval[api]>=0.4.11 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
transformers==4.57.5
tokenizers==0.22.0
schemathesis>=3.39.15 # Required for openai schema test.
# quantization
bitsandbytes>=0.49.2
+4 -2
View File
@@ -23,6 +23,7 @@ vocos # required for minicpmo_26 test
peft>=0.15.0 # required for phi-4-mm test
pqdm
ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
resampy # required for audio tests
sentence-transformers>=5.2.0 # required for embedding tests
soundfile # required for audio tests
jiwer # required for audio tests
@@ -37,8 +38,8 @@ opencv-python-headless>=4.13.0 # required for video test
datamodel_code_generator # required for minicpm3 test
lm-eval[api]>=0.4.11 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
transformers==4.57.5
tokenizers==0.22.0
schemathesis>=3.39.15 # Required for openai schema test
# quantization
bitsandbytes==0.49.2
@@ -81,3 +82,4 @@ plotly # required for perf comparison html report
rapidfuzz
torchgeo==0.7.0
multiprocess==0.70.16
huggingface-hub==0.36.2
+21 -19
View File
@@ -39,7 +39,7 @@ annotated-doc==0.0.4
# typer
annotated-types==0.7.0
# via pydantic
anthropic==0.93.0
anthropic==0.89.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -76,9 +76,7 @@ attrs==26.1.0
audioread==3.0.1
# via librosa
av==16.1.0
# via
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# via -r requirements/test/rocm.in
azure-core==1.39.0
# via
# azure-identity
@@ -174,7 +172,7 @@ colorful==0.5.8
# via ray
colorlog==6.10.1
# via optuna
compressed-tensors==0.15.0.1
compressed-tensors==0.14.0.1
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -271,9 +269,9 @@ fastapi==0.135.2
# model-hosting-container-standards
fastapi-cli==0.0.24
# via fastapi
fastapi-cloud-cli==0.16.1
fastapi-cloud-cli==0.15.1
# via fastapi-cli
fastar==0.10.0
fastar==0.9.0
# via fastapi-cloud-cli
fastparquet==2026.3.0
# via genai-perf
@@ -292,6 +290,7 @@ filelock==3.25.2
# python-discovery
# ray
# torch
# transformers
# virtualenv
fiona==1.10.1
# via torchgeo
@@ -385,7 +384,7 @@ h5py==3.16.0
# via terratorch
harfile==0.4.0
# via schemathesis
hf-xet==1.4.3
hf-xet==1.4.2
# via huggingface-hub
hiredis==3.3.1
# via tensorizer
@@ -404,7 +403,6 @@ httpx==0.27.2
# diffusers
# fastapi
# fastapi-cloud-cli
# huggingface-hub
# mcp
# model-hosting-container-standards
# openai
@@ -412,8 +410,9 @@ httpx==0.27.2
# schemathesis
httpx-sse==0.4.3
# via mcp
huggingface-hub==1.10.2
huggingface-hub==0.36.2
# via
# -r requirements/test/rocm.in
# accelerate
# datasets
# diffusers
@@ -485,7 +484,7 @@ jinja2==3.1.6
# genai-perf
# lm-eval
# torch
jiter==0.14.0
jiter==0.13.0
# via
# anthropic
# openai
@@ -632,7 +631,7 @@ msgpack==1.1.2
# via
# librosa
# ray
msgspec==0.21.0
msgspec==0.20.0
# via -r requirements/test/../common.txt
mteb==2.11.5
# via -r requirements/test/rocm.in
@@ -664,6 +663,7 @@ numba==0.61.2
# -c requirements/rocm.txt
# -r requirements/test/rocm.in
# librosa
# resampy
numkong==7.1.1
# via albucore
numpy==2.2.6
@@ -709,6 +709,7 @@ numpy==2.2.6
# pytrec-eval-terrier
# pywavelets
# rasterio
# resampy
# rioxarray
# rouge-score
# runai-model-streamer
@@ -741,7 +742,7 @@ omegaconf==2.3.0
# lightning
open-clip-torch==2.32.0
# via -r requirements/test/rocm.in
openai==2.31.0
openai==2.30.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -1092,7 +1093,7 @@ python-dotenv==1.2.2
# uvicorn
python-json-logger==4.1.0
# via -r requirements/test/../common.txt
python-multipart==0.0.26
python-multipart==0.0.22
# via
# fastapi
# mcp
@@ -1179,6 +1180,7 @@ requests==2.32.5
# google-api-core
# google-cloud-storage
# gpt-oss
# huggingface-hub
# lightly
# lm-eval
# mistral-common
@@ -1192,7 +1194,10 @@ requests==2.32.5
# starlette-testclient
# tacoreader
# tiktoken
# transformers
# wandb
resampy==0.4.3
# via -r requirements/test/rocm.in
responses==0.26.0
# via genai-perf
rfc3339-validator==0.1.4
@@ -1333,7 +1338,6 @@ sortedcontainers==2.4.0
# via hypothesis
soundfile==0.13.1
# via
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# genai-perf
# librosa
@@ -1424,7 +1428,7 @@ timm==1.0.17
# segmentation-models-pytorch
# terratorch
# torchgeo
tokenizers==0.22.2
tokenizers==0.22.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -1467,7 +1471,7 @@ tqdm==4.67.3
# tacoreader
# terratorch
# transformers
transformers==5.5.3
transformers==4.57.5
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -1494,9 +1498,7 @@ typer==0.24.1
# fastapi-cli
# fastapi-cloud-cli
# fastsafetensors
# huggingface-hub
# perceptron
# transformers
typeshed-client==2.9.0
# via jsonargparse
typing-extensions==4.15.0
+1
View File
@@ -13,6 +13,7 @@ pytest-shard
absl-py
accelerate
arctic-inference
hf_transfer
lm_eval[api]
modelscope
+9 -15
View File
@@ -19,9 +19,7 @@ aiosignal==1.4.0
albumentations==1.4.6
# via -r requirements/test/xpu.in
annotated-doc==0.0.4
# via
# fastapi
# typer
# via fastapi
annotated-types==0.7.0
# via pydantic
anyio==4.13.0
@@ -66,7 +64,6 @@ click==8.3.1
# jiwer
# nltk
# schemathesis
# typer
# uvicorn
colorama==0.4.6
# via sacrebleu
@@ -115,6 +112,7 @@ filelock==3.25.2
# huggingface-hub
# modelscope
# torch
# transformers
frozenlist==1.8.0
# via
# aiohttp
@@ -135,7 +133,9 @@ h11==0.16.0
# uvicorn
harfile==0.4.0
# via schemathesis
hf-xet==1.4.3
hf-transfer==0.1.9
# via -r requirements/test/xpu.in
hf-xet==1.4.2
# via huggingface-hub
html2text==2025.4.15
# via gpt-oss
@@ -144,9 +144,8 @@ httpcore==1.0.9
httpx==0.28.1
# via
# datasets
# huggingface-hub
# schemathesis
huggingface-hub==1.10.2
huggingface-hub==0.36.2
# via
# accelerate
# datasets
@@ -516,6 +515,7 @@ requests==2.33.1
# docker
# evaluate
# gpt-oss
# huggingface-hub
# lm-eval
# mistral-common
# modelscope
@@ -524,11 +524,11 @@ requests==2.33.1
# schemathesis
# starlette-testclient
# tiktoken
# transformers
rich==14.3.3
# via
# mteb
# schemathesis
# typer
rouge-score==0.1.2
# via lm-eval
rpds-py==0.30.0
@@ -572,8 +572,6 @@ setuptools==80.10.2
# modelscope
# pytablewriter
# torch
shellingham==1.5.4
# via typer
six==1.17.0
# via
# -c requirements/common.txt
@@ -667,7 +665,7 @@ tqdm==4.67.3
# pqdm
# sentence-transformers
# transformers
transformers==5.5.3
transformers==4.57.6
# via
# -c requirements/common.txt
# sentence-transformers
@@ -678,10 +676,6 @@ typepy==1.3.4
# dataproperty
# pytablewriter
# tabledata
typer==0.24.1
# via
# huggingface-hub
# transformers
typing-extensions==4.15.0
# via
# -c requirements/common.txt
+5 -12
View File
@@ -693,12 +693,6 @@ class precompiled_wheel_utils:
flash_attn_regex = re.compile(
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
# __init__.py and flash_attn_interface.py are source-controlled
# in vllm and should not be overwritten (matches cmake exclusions)
flash_attn_files_to_skip = {
"vllm/vllm_flash_attn/__init__.py",
"vllm/vllm_flash_attn/flash_attn_interface.py",
}
triton_kernels_regex = re.compile(
r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
@@ -711,11 +705,7 @@ class precompiled_wheel_utils:
filter(lambda x: x.filename in files_to_copy, wheel.filelist)
)
file_members += list(
filter(
lambda x: flash_attn_regex.match(x.filename)
and x.filename not in flash_attn_files_to_skip,
wheel.filelist,
)
filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
)
file_members += list(
filter(
@@ -1092,7 +1082,10 @@ setup(
"instanttensor": ["instanttensor >= 0.1.5"],
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
"audio": [
"av",
"resampy",
"scipy",
"soundfile",
"mistral_common[audio]",
], # Required for audio processing
"video": [], # Kept for backwards compatibility
@@ -1101,7 +1094,7 @@ setup(
# NOTE: When updating helion version, also update CI files:
# - .buildkite/test_areas/kernels.yaml
# - .buildkite/test-amd.yaml
"helion": ["helion==1.0.0"],
"helion": ["helion==0.3.3"],
# Optional deps for gRPC server (vllm serve --grpc)
"grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"],
# Optional deps for OpenTelemetry tracing
@@ -1,126 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.config
from tests.compile.backend import TestBackend
from vllm.compilation.passes.vllm_inductor_pass import (
VllmFusionPatternMatcherPass,
VllmPatternMatcherPass,
VllmPatternReplacement,
)
from vllm.config import CompilationConfig, CompilationMode, VllmConfig
from vllm.platforms import current_platform
class ReluToAbsPattern(VllmPatternReplacement):
"""Replaces relu(x) with abs(x) — a minimal test fixture."""
@property
def pattern(self):
def _pattern(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.relu.default(x)
return _pattern
@property
def replacement(self):
def _replacement(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.abs.default(x)
return _replacement
def get_inputs(self) -> list[torch.Tensor]:
return [self.empty_fp32(4)]
class ExpToSqrtPattern(VllmPatternReplacement):
"""A second distinct pattern type — used to test uuid differentiation."""
@property
def pattern(self):
def _pattern(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.exp.default(x)
return _pattern
@property
def replacement(self):
def _replacement(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.sqrt.default(x)
return _replacement
def get_inputs(self) -> list[torch.Tensor]:
return [self.empty_fp32(4)]
class ReluFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "test_relu_fusion")
self.register(ReluToAbsPattern())
class TwoPatternFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "test_two_pattern_fusion")
self.register(ReluToAbsPattern())
self.register(ExpToSqrtPattern())
@pytest.fixture
def vllm_config():
return VllmConfig(
compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE),
)
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
def test_register_tracks_patterns(vllm_config):
"""register() appends each VllmPatternReplacement to _pattern_replacements."""
with vllm.config.set_current_vllm_config(vllm_config):
single = ReluFusionPass(vllm_config)
two = TwoPatternFusionPass(vllm_config)
assert len(single._pattern_replacements) == 1
assert len(two._pattern_replacements) == 2
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
def test_uuid_stable(vllm_config):
"""Two instances of the same pass class produce identical uuids."""
with vllm.config.set_current_vllm_config(vllm_config):
p1 = ReluFusionPass(vllm_config)
p2 = ReluFusionPass(vllm_config)
p3 = TwoPatternFusionPass(vllm_config)
assert p1.uuid() == p2.uuid()
assert p1.uuid() != p3.uuid()
assert p2.uuid() != p3.uuid()
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
@pytest.mark.parametrize("N", [1, 2, 4])
def test_matched_count_and_match_table(vllm_config, N):
"""matched_count and match_table reflect the number of matched patterns."""
class Model(torch.nn.Module):
def forward(self, *inputs):
# N independent relus
return sum(torch.relu(x) for x in inputs)
with vllm.config.set_current_vllm_config(vllm_config):
torch.set_default_device("cuda")
torch.set_default_dtype(torch.float32)
fusion_pass = ReluFusionPass(vllm_config)
backend = TestBackend(fusion_pass)
model = torch.compile(Model(), backend=backend)
inputs = [torch.rand(8) for _ in range(N)]
model(*inputs)
assert fusion_pass.matched_count == N
assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N
@@ -222,47 +222,3 @@ def test_model_specialization_with_evaluate_guards(
torch.randn(1, 10).cuda(),
is_01_specialization=True,
)
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
def test_piecewise_backend_empty_sym_shape_indices():
"""Test that PiecewiseBackend handles empty sym_shape_indices correctly.
When all inputs have static shapes (no torch.SymInt), sym_shape_indices
will be empty. The fix in PiecewiseBackend.__call__ handles this case
by using the first compiled range_entry.
"""
gc.collect()
torch.accelerator.empty_cache()
torch.accelerator.synchronize()
# Use small max_model_len and max_num_batched_tokens to encourage
# static shape compilation with empty sym_shape_indices
llm = LLM(
model="Qwen/Qwen3-0.6B",
max_model_len=512,
max_num_batched_tokens=1,
compilation_config={
"mode": CompilationMode.VLLM_COMPILE,
"dynamic_shapes_config": {
"type": DynamicShapesType.BACKED.value,
},
},
)
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
# Generate with static shape inputs
output = llm.generate("Hello, my name is", sampling_params=sampling_params)
result = output[0].outputs[0].text
assert len(result) > 0, "Should generate non-empty output"
# Generate again to verify compilation works with empty sym_shape_indices
output = llm.generate("The capital of France is", sampling_params=sampling_params)
result = output[0].outputs[0].text
assert len(result) > 0, "Should generate non-empty output on second run"
del llm
gc.collect()
torch.accelerator.empty_cache()
torch.accelerator.synchronize()
@@ -356,23 +356,6 @@
"is_multimodal_model": false,
"dtype": "torch.float32"
},
"stepfun-ai/Step-3.5-Flash": {
"architectures": [
"Step3p5ForCausalLM"
],
"model_type": "step3p5",
"text_model_type": "step3p5",
"hidden_size": 4096,
"total_num_hidden_layers": 45,
"total_num_attention_heads": 64,
"head_size": 128,
"vocab_size": 128896,
"total_num_kv_heads": 8,
"num_experts": 288,
"is_deepseek_mla": false,
"is_multimodal_model": false,
"dtype": "torch.bfloat16"
},
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": {
"architectures": [
"NemotronHForCausalLM"
-1
View File
@@ -16,7 +16,6 @@ BASE_TRUST_REMOTE_CODE_MODELS = {
"nvidia/Llama-3_3-Nemotron-Super-49B-v1",
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
"XiaomiMiMo/MiMo-7B-RL",
"stepfun-ai/Step-3.5-Flash",
# Excluded: Not available online right now
# "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1",
"meituan-longcat/LongCat-Flash-Chat",
-15
View File
@@ -364,7 +364,6 @@ class HfRunner:
model_name: str,
dtype: str = "auto",
*,
revision: str | None = None,
model_kwargs: dict[str, Any] | None = None,
trust_remote_code: bool = True,
is_sentence_transformer: bool = False,
@@ -384,7 +383,6 @@ class HfRunner:
self._init(
model_name=model_name,
dtype=dtype,
revision=revision,
model_kwargs=model_kwargs,
trust_remote_code=trust_remote_code,
is_sentence_transformer=is_sentence_transformer,
@@ -398,7 +396,6 @@ class HfRunner:
model_name: str,
dtype: str = "auto",
*,
revision: str | None = None,
model_kwargs: dict[str, Any] | None = None,
trust_remote_code: bool = True,
is_sentence_transformer: bool = False,
@@ -413,15 +410,6 @@ class HfRunner:
model_name,
trust_remote_code=trust_remote_code,
)
# HF runner should use the HF config so that it's consistent with the HF model
if self.config.__module__.startswith("vllm.transformers_utils.configs"):
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
del CONFIG_MAPPING._extra_content[self.config.model_type]
self.config = AutoConfig.from_pretrained(
model_name,
trust_remote_code=trust_remote_code,
)
self.device = self.get_default_device()
self.dtype = dtype = _get_and_verify_dtype(
self.model_name,
@@ -440,7 +428,6 @@ class HfRunner:
self.model = SentenceTransformer(
model_name,
revision=revision,
device=self.device,
model_kwargs=model_kwargs,
trust_remote_code=trust_remote_code,
@@ -451,7 +438,6 @@ class HfRunner:
self.model = CrossEncoder(
model_name,
revision=revision,
device=self.device,
automodel_args=model_kwargs,
trust_remote_code=trust_remote_code,
@@ -461,7 +447,6 @@ class HfRunner:
nn.Module,
auto_cls.from_pretrained(
model_name,
revision=revision,
trust_remote_code=trust_remote_code,
**model_kwargs,
),
@@ -41,7 +41,6 @@ def create_mock_parallel_config(
config.rank = rank
config.world_size = world_size
config.data_parallel_rank = dp_rank
config.data_parallel_index = dp_rank
return config
@@ -284,7 +283,6 @@ def inference_receive_tensor(
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = NCCLWeightTransferEngine(config, parallel_config)
@@ -668,7 +666,6 @@ def inference_receive_ipc_tensor(
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = IPCWeightTransferEngine(config, parallel_config)
-6
View File
@@ -91,12 +91,6 @@ def test_multiple_priority(llm: LLM):
outputs = llm.generate(PROMPTS, sampling_params=None, priority=[])
def test_single_prompt_priority(llm: LLM):
# Single string prompts should be normalized to one request.
outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=[0])
assert len(outputs) == 1
def test_max_model_len():
max_model_len = 20
llm = LLM(
@@ -1,150 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for tool_calls Iterable → list materialisation.
Regression tests for https://github.com/vllm-project/vllm/issues/34792.
Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral
models because:
1. The OpenAI Python SDK types tool_calls as Iterable[...] in
ChatCompletionAssistantMessageParam.
2. Pydantic v2, when validating from Python objects (not from raw JSON),
wraps Iterable fields in a one-shot lazy iterator.
3. Debug logging called model_dump_json() which consumed that iterator.
4. The Mistral tokenizer then saw empty tool_calls and raised
"ValueError: Unexpected tool call id ...".
"""
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
def _make_tool_call(tc_id: str, name: str, args: str) -> dict:
return {
"id": tc_id,
"type": "function",
"function": {"name": name, "arguments": args},
}
def _make_request(messages: list) -> ChatCompletionRequest:
return ChatCompletionRequest(
model="test-model",
messages=messages,
)
def test_tool_calls_list_preserved_after_model_dump():
"""tool_calls in assistant messages must be readable after model_dump_json.
When the request is built from Python dicts (as in the Anthropic → OpenAI
conversion path), Pydantic v2 previously wrapped the Iterable tool_calls
in a one-shot iterator. model_dump_json() consumed it, leaving subsequent
readers (e.g. the Mistral tokenizer) with an empty sequence.
"""
tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}')
messages = [
{"role": "user", "content": "What is the weather in Paris?"},
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
{
"role": "tool",
"tool_call_id": "call_abc123",
"content": '{"temperature": 20}',
},
]
req = _make_request(messages)
# Simulate debug logging: serialize the model (this was the trigger)
_ = req.model_dump_json()
# The assistant message must still have accessible tool_calls afterwards
assistant_msg = req.messages[1]
assert isinstance(assistant_msg, dict)
tool_calls = assistant_msg.get("tool_calls")
assert tool_calls is not None, "tool_calls must not be None after model_dump_json"
assert isinstance(tool_calls, list), "tool_calls must be a list"
assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json"
def test_tool_calls_from_generator_are_materialised():
"""tool_calls passed as a generator must be converted to list on validation."""
tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}')
def tool_calls_gen():
yield tool_call
messages = [
{"role": "user", "content": "Search for vllm"},
{
"role": "assistant",
"content": None,
"tool_calls": tool_calls_gen(), # one-shot generator
},
]
req = _make_request(messages)
assistant_msg = req.messages[1]
assert isinstance(assistant_msg, dict)
# Iterate twice — must not raise or return empty on second pass
tool_calls_first = list(assistant_msg.get("tool_calls", []))
tool_calls_second = list(assistant_msg.get("tool_calls", []))
assert len(tool_calls_first) == 1, "First read must return the tool call"
assert len(tool_calls_second) == 1, "Second read must also return the tool call"
def test_tool_calls_list_passthrough():
"""tool_calls already provided as a list must remain a list."""
tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}')
messages = [
{"role": "user", "content": "Calculate 2+2"},
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
]
req = _make_request(messages)
assistant_msg = req.messages[1]
assert isinstance(assistant_msg, dict)
assert isinstance(assistant_msg.get("tool_calls"), list)
def test_messages_without_tool_calls_unaffected():
"""Messages without tool_calls must be handled correctly."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"},
{"role": "assistant", "content": "Hi there!"},
]
req = _make_request(messages)
# None of the messages should have tool_calls injected
for msg in req.messages:
assert isinstance(msg, dict)
assert msg.get("tool_calls") is None or msg.get("tool_calls") == []
@pytest.mark.parametrize("num_tool_calls", [1, 3])
def test_multiple_tool_calls_materialised(num_tool_calls: int):
"""Multiple tool calls in a single message are all preserved."""
tool_calls = [
_make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}')
for i in range(num_tool_calls)
]
messages = [
{"role": "user", "content": "Do things"},
{"role": "assistant", "content": None, "tool_calls": iter(tool_calls)},
]
req = _make_request(messages)
assistant_msg = req.messages[1]
assert isinstance(assistant_msg, dict)
result_tool_calls = assistant_msg.get("tool_calls")
assert isinstance(result_tool_calls, list)
assert len(result_tool_calls) == num_tool_calls
# Verify after model_dump_json too
_ = req.model_dump_json()
assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls
+21 -99
View File
@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
from vllm.platforms import current_platform
from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
from vllm.utils.torch_utils import set_random_seed
COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
DTYPES = [torch.bfloat16, torch.float]
@@ -172,7 +172,7 @@ def test_reshape_and_cache(
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"])
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@@ -202,25 +202,6 @@ def test_reshape_and_cache_flash(
if kv_scale_type == "attn_head" and implementation != "cuda":
pytest.skip("Only CUDA implementation supports attn_head scaling.")
if kv_cache_dtype == "nvfp4":
if not current_platform.has_device_capability(100):
pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
if implementation != "cuda":
pytest.skip("NVFP4 only supports CUDA implementation.")
if kv_scale_type != "tensor":
pytest.skip("NVFP4 only supports per-tensor scaling.")
if head_size % 16 != 0:
pytest.skip("NVFP4 requires head_size divisible by 16.")
if (head_size // 16) % 4 != 0:
pytest.skip(
"NVFP4 requires (head_size // 16) divisible by 4 "
"for 4x4 block scale swizzle."
)
if block_size % 4 != 0:
pytest.skip("NVFP4 requires block_size divisible by 4.")
if dtype not in (torch.float16, torch.bfloat16):
pytest.skip("NVFP4 quantization only supports fp16/bf16 input.")
# fp8 conversion requires continugous memory buffer. Reduce the number of
# blocks and tokens to consume less memory.
num_tokens = num_tokens // 2
@@ -248,23 +229,7 @@ def test_reshape_and_cache_flash(
del key_caches
del value_caches
# For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes.
# Split views are still needed for dequant verification.
key_scale_cache = None
value_scale_cache = None
nvfp4_key_data = None
nvfp4_value_data = None
if kv_cache_dtype == "nvfp4":
(nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
(nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
value_cache
)
if kv_cache_dtype == "nvfp4":
# Global scale = amax / 448 (per-tensor)
k_scale = (key.abs().amax() / 448.0).to(torch.float32)
v_scale = (value.abs().amax() / 448.0).to(torch.float32)
elif kv_scale_type == "tensor":
if kv_scale_type == "tensor":
k_scale = (key.amax() / 64.0).to(torch.float32)
v_scale = (value.amax() / 64.0).to(torch.float32)
else: # "attn_head"
@@ -275,9 +240,8 @@ def test_reshape_and_cache_flash(
y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
return y.contiguous()
if kv_cache_dtype != "nvfp4":
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
def convert_fp8_local(output, input, scale, kv_dtype):
fp8_input = input.view(current_platform.fp8_dtype())
@@ -293,7 +257,7 @@ def test_reshape_and_cache_flash(
result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
output.copy_(result)
# Clone the KV caches (for non-nvfp4, used as reference baseline).
# Clone the KV caches.
if kv_cache_dtype == "fp8":
cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
@@ -301,27 +265,25 @@ def test_reshape_and_cache_flash(
convert_fp8_local(
cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
)
elif kv_cache_dtype != "nvfp4":
else:
cloned_key_cache = key_cache_compact.clone()
cloned_value_cache = value_cache_compact.clone()
# Call the reshape_and_cache kernel.
if implementation == "cuda":
if kv_cache_dtype != "nvfp4":
opcheck(
torch.ops._C_cache_ops.reshape_and_cache_flash,
(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
),
cond=(head_size == HEAD_SIZES[0]),
)
opcheck(
torch.ops._C_cache_ops.reshape_and_cache_flash,
(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
),
cond=(head_size == HEAD_SIZES[0]),
)
ops.reshape_and_cache_flash(
key,
value,
@@ -347,46 +309,6 @@ def test_reshape_and_cache_flash(
k_scale,
v_scale,
)
if kv_cache_dtype == "nvfp4":
# Verify NVFP4 by dequantizing the entire cache and comparing
# the written positions against original bf16 values.
# Same pattern as FP8: dequant whole cache, then extract and compare.
from tests.kernels.quantization.nvfp4_utils import (
dequant_nvfp4_kv_cache,
)
def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
# data_cache: [N, T, H, data_dim] NHD (contiguous inner dims)
# scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
# Permute to HND layout for the dequant utility.
data_hnd = data_cache.permute(0, 2, 1, 3)
scale_hnd = scale_cache.permute(0, 2, 1, 3)
result_hnd = dequant_nvfp4_kv_cache(
data_hnd, scale_hnd, global_scale, head_size, block_size
)
return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D]
result_key_cache = dequant_nvfp4_cache_nhd(
nvfp4_key_data, key_scale_cache, k_scale.item()
)
result_value_cache = dequant_nvfp4_cache_nhd(
nvfp4_value_data, value_scale_cache, v_scale.item()
)
# Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
num_slots = num_blocks * block_size
result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
torch.testing.assert_close(
result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
)
torch.testing.assert_close(
result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5
)
return
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
File diff suppressed because it is too large Load Diff
-111
View File
@@ -1,111 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
from tests.kernels.utils import opcheck
from vllm.platforms import CpuArchEnum, current_platform
from vllm.utils.torch_utils import set_random_seed
if not current_platform.is_cpu():
pytest.skip("skipping CPU-only tests", allow_module_level=True)
from vllm.model_executor.layers.activation import (
GELU,
FastGELU,
GeluAndMul,
NewGELU,
QuickGELU,
SiluAndMul,
)
DTYPES = [torch.bfloat16, torch.float32]
NUM_TOKENS = [7, 83]
D = [512, 2048]
SEEDS = [0]
@pytest.mark.parametrize(
("activation_cls", "fn"),
[
(SiluAndMul, torch.ops._C.silu_and_mul),
(GeluAndMul, torch.ops._C.gelu_and_mul),
(GeluAndMul, torch.ops._C.gelu_tanh_and_mul),
],
)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("d", D)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_cpu_act_and_mul(
default_vllm_config,
activation_cls: type[torch.nn.Module],
fn: object,
num_tokens: int,
d: int,
dtype: torch.dtype,
seed: int,
) -> None:
set_random_seed(seed)
x = torch.randn(num_tokens, 2 * d, dtype=dtype)
layer = activation_cls()
out = layer(x)
ref_out = layer.forward_native(x)
torch.testing.assert_close(
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
)
output_shape = x.shape[:-1] + (x.shape[-1] // 2,)
raw_out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
opcheck(fn, (raw_out, x))
@pytest.mark.parametrize(
("activation_cls", "fn", "op_args"),
[
(NewGELU, torch.ops._C.gelu_new, ()),
(FastGELU, torch.ops._C.gelu_fast, ()),
(QuickGELU, torch.ops._C.gelu_quick, ()),
pytest.param(
GELU,
getattr(torch.ops._C, "activation_lut_bf16", None),
("gelu",),
marks=pytest.mark.skipif(
current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
reason="activation_lut_bf16 is only built on Arm CPU",
),
),
],
)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("d", D)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_cpu_unary_activation(
default_vllm_config,
activation_cls: type[torch.nn.Module],
fn: object,
op_args: tuple[str, ...],
num_tokens: int,
d: int,
dtype: torch.dtype,
seed: int,
) -> None:
set_random_seed(seed)
x = torch.randn(num_tokens, d, dtype=dtype)
layer = activation_cls()
out = layer(x)
ref_out = layer.forward_native(x)
torch.testing.assert_close(
out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
)
# gelu with activation_lut_bf16 only makes sense for BF16
if not (activation_cls is GELU and dtype != torch.bfloat16):
raw_out = torch.empty_like(x)
opcheck(fn, (raw_out, x, *op_args))
-48
View File
@@ -36,11 +36,9 @@ from vllm.kernels.helion.register import (
)
if _HOP_AVAILABLE:
from helion._compat import supports_torch_compile_fusion
from helion._compiler._dynamo.higher_order_ops import (
helion_kernel_wrapper_mutation,
)
from torch._inductor.utils import run_and_get_code
def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -1005,49 +1003,3 @@ class TestTorchCompileHOP:
"Compiled execution result doesn't match eager execution. "
f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}"
)
@pytest.mark.skipif(
not (_HOP_AVAILABLE and supports_torch_compile_fusion()),
reason="Requires PyTorch with Helion inductor fusion support",
)
def test_inductor_backend_compiles_helion_hop(self):
"""Test torch.compile with inductor backend and Helion fusion enabled."""
configs = {"default": helion.Config(block_sizes=[4, 4])}
with dummy_kernel_registry(configs=configs) as register:
add_helion_kernel = register(
op_name="test_inductor_add_kernel",
config_picker=lambda args, keys: "default",
helion_settings=helion.Settings(
torch_compile_fusion=True, static_shapes=False
),
)(_add_kernel)
def f(x, y):
x = x * 2.0
y = y + 1.0
out = add_helion_kernel(x, y)
return out.relu()
torch._dynamo.reset()
compiled_f = torch.compile(f, backend="inductor", fullgraph=True)
x = torch.randn(4, 4, device="cuda")
y = torch.randn(4, 4, device="cuda")
compiled_result, source_codes = run_and_get_code(compiled_f, x, y)
eager_result = f(x, y)
assert torch.allclose(compiled_result, eager_result, atol=1e-5, rtol=1e-5), (
"Inductor-compiled result doesn't match eager execution. "
f"Max difference: {torch.max(torch.abs(compiled_result - eager_result))}"
)
# With fusion enabled, prologue/epilogue ops should be fused into
# a single triton kernel rather than generating separate kernels.
kernel_count = sum(code.count("@triton.jit") for code in source_codes)
assert kernel_count == 1, (
f"Expected 1 fused triton kernel, got {kernel_count}. "
"Prologue/epilogue ops were not fused into the Helion kernel."
)
+4 -48
View File
@@ -13,11 +13,6 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
selective_state_update,
)
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
MambaSpec,
)
try:
import flashinfer.mamba # noqa: F401
@@ -27,40 +22,22 @@ except ImportError:
HAS_FLASHINFER = False
def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig:
spec = MambaSpec(
block_size=16,
shapes=((16, 64),),
dtypes=(torch.float16,),
mamba_type=mamba_type,
)
return KVCacheConfig(
num_blocks=1,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)],
)
def test_default_backend_is_triton():
initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
initialize_mamba_ssu_backend(MambaConfig())
backend = get_mamba_ssu_backend()
assert isinstance(backend, TritonSSUBackend)
assert backend.name == "triton"
def test_explicit_triton_backend():
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
)
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
backend = get_mamba_ssu_backend()
assert isinstance(backend, TritonSSUBackend)
@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
def test_flashinfer_backend_init():
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu()
)
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
backend = get_mamba_ssu_backend()
assert isinstance(backend, FlashInferSSUBackend)
assert backend.name == "flashinfer"
@@ -76,25 +53,6 @@ def test_uninitialized_backend_raises():
mod._mamba_ssu_backend = old
@pytest.mark.parametrize(
"mamba_type", ["linear_attention", "gdn_attention", "short_conv"]
)
def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
old = mod._mamba_ssu_backend
mod._mamba_ssu_backend = None
try:
initialize_mamba_ssu_backend(
MambaConfig(), _kv_cache_config_with_ssu(mamba_type)
)
assert mod._mamba_ssu_backend is None
with pytest.raises(RuntimeError, match="not been initialized"):
get_mamba_ssu_backend()
finally:
mod._mamba_ssu_backend = old
@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
def test_flashinfer_import_error():
with pytest.raises(ImportError, match="FlashInfer is required"):
@@ -103,9 +61,7 @@ def test_flashinfer_import_error():
def test_triton_basic_call():
set_random_seed(0)
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
)
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
device = "cuda"
batch_size = 2
dim = 64
-8
View File
@@ -465,14 +465,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
if config.enable_eplb and config.ep_size == 1:
return False, "EPLB only works with EP+DP"
# Disable fp4 tests until flashinfer is updated or the Dockerfile is
# modified to install cublasLt.h. See #39525.
if (
config.quantization == "modelopt_fp4"
and current_platform.is_device_capability_family(100)
):
return False, "Temporarily skip until #39525 is resolved"
return True, None
-54
View File
@@ -88,60 +88,6 @@ def break_fp4_bytes(a, dtype):
return values.reshape(m, n * 2).to(dtype=dtype)
def dequant_nvfp4_kv_cache(
fp4_data: torch.Tensor,
block_scale: torch.Tensor,
global_scale: float,
head_size: int,
block_size: int,
) -> torch.Tensor:
"""Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
The input must be in HND layout so that the last two dims are
(block_size, last_dim). For NHD caches, permute to HND first.
Args:
fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4.
block_scale: [..., num_heads, block_size, head_size//16] fp8 block
scales (as uint8 or float8_e4m3fn).
global_scale: checkpoint dequant scale (k_scale or v_scale).
head_size: head dimension.
block_size: page size.
Returns:
[..., num_heads, block_size, head_size] float32.
"""
data_dim = head_size // 2
scale_dim = head_size // 16
fp4_packed = fp4_data
sf_swizzled = block_scale.view(torch.uint8)
# Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
# [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
batch_shape = sf_swizzled.shape[:-2]
T, S = block_size, scale_dim
sg = S // 4
sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
ndim = sf_reshape.ndim
# Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
# Unpack fp4
shape = fp4_packed.shape # [..., T, data_dim]
fp4_flat = fp4_packed.reshape(-1, data_dim)
fp4_vals = break_fp4_bytes(fp4_flat, torch.float32)
fp4_vals = fp4_vals.reshape(*shape[:-1], head_size)
# Dequant: fp4_val * block_scale * global_scale per 16-element group
return (
fp4_vals.reshape(*shape[:-1], scale_dim, 16)
* (sf_f32 * global_scale).unsqueeze(-1)
).reshape(*shape[:-1], head_size)
def get_nvfp4_global_scale(a: torch.Tensor):
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
-6
View File
@@ -3,7 +3,6 @@
import tempfile
from collections import OrderedDict
from importlib import reload
from unittest.mock import MagicMock
import pytest
@@ -48,11 +47,6 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):
def maybe_enable_lora_dual_stream(monkeypatch: pytest.MonkeyPatch):
if current_platform.is_cuda():
monkeypatch.setenv("VLLM_LORA_ENABLE_DUAL_STREAM", "1")
import vllm.lora.layers.base_linear
if not hasattr(vllm.lora.layers.base_linear, "lora_linear_async"):
# Reload the module to ensure the environment variable takes effect.
reload(vllm.lora.layers.base_linear)
yield
-11
View File
@@ -1,10 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from importlib.metadata import version
import pytest
from packaging.version import Version
import vllm
from vllm.assets.image import ImageAsset
@@ -13,14 +10,6 @@ from vllm.platforms import current_platform
from ..utils import multi_gpu_test
pytestmark = pytest.mark.skipif(
Version("5.0") <= Version(version("transformers")),
reason=(
"MiniCPMV custom processor uses tokenizer.im_start_id which is not "
"available on TokenizersBackend in transformers v5.0+"
),
)
MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5"
PROMPT_TEMPLATE = (
@@ -164,34 +164,6 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
assert add_perp < mul_perp
def test_kv_scale_reload(vllm_runner):
"""Test reloading a checkpoint that contains k_scale/v_scale weights."""
if not current_platform.supports_fp8():
pytest.skip(reason="Requires FP8 support")
model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV"
# Load dummy weights, then reload real checkpoint
with vllm_runner(
model_name=model,
load_format="dummy",
enable_prefix_caching=False,
max_model_len=16,
max_num_seqs=1,
) as llm:
llm.collective_rpc(
"update_config",
kwargs={"overrides": {"load_config": {"load_format": "auto"}}},
)
llm.collective_rpc("reload_weights", kwargs={"weights_path": model})
reloaded_perp = llm.generate_prompt_perplexity(
["The capital of France is the city of Paris"],
mask=["The capital of France is"],
)[0]
assert reloaded_perp < 10
@pytest.mark.parametrize(
"tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
)
-145
View File
@@ -1,145 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass, field
import pytest
import torch
from vllm.model_executor.models.keye import KeyeForConditionalGeneration
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalFieldElem,
MultiModalKwargsItem,
PlaceholderRange,
)
pytestmark = pytest.mark.skip_global_cleanup
@pytest.fixture(autouse=True, scope="module")
def _force_cpu_default_device():
original = torch.get_default_device()
torch.set_default_device("cpu")
yield
torch.set_default_device(original)
@dataclass
class DummyVisionConfig:
spatial_merge_size: int = 2
@dataclass
class DummyConfig:
vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
def make_model(config: DummyConfig) -> KeyeForConditionalGeneration:
model = object.__new__(KeyeForConditionalGeneration)
model.config = config
return model
def make_mm_feature(
*,
modality: str,
offset: int,
length: int,
grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
) -> MultiModalFeatureSpec:
field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
return MultiModalFeatureSpec(
data=MultiModalKwargsItem(
{
field_name: MultiModalFieldElem(
data=torch.tensor(grid_thw),
field=None, # HACK.
),
}
),
modality=modality,
identifier="DUMMY",
mm_position=PlaceholderRange(offset=offset, length=length),
)
def test_get_mrope_input_positions_text_only():
model = make_model(DummyConfig())
positions, delta = model.get_mrope_input_positions(
input_tokens=[11, 12, 13, 14, 15],
mm_features=[],
)
expected = torch.tensor(
[
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
]
)
assert torch.equal(positions, expected)
assert delta == 0
def test_get_mrope_input_positions_single_image():
model = make_model(DummyConfig())
mm_features = [
make_mm_feature(
modality="image",
offset=1,
length=4,
grid_thw=(1, 4, 4),
)
]
positions, delta = model.get_mrope_input_positions(
input_tokens=[10, 20, 21, 22, 23, 30, 31],
mm_features=mm_features,
)
expected = torch.tensor(
[
[0, 1, 1, 1, 1, 3, 4],
[0, 1, 1, 2, 2, 3, 4],
[0, 1, 2, 1, 2, 3, 4],
]
)
assert torch.equal(positions, expected)
assert delta == -2
def test_get_mrope_input_positions_interleaved_image_and_video():
model = make_model(DummyConfig())
mm_features = [
make_mm_feature(
modality="image",
offset=1,
length=4,
grid_thw=(1, 4, 4),
),
make_mm_feature(
modality="video",
offset=7,
length=4,
grid_thw=[(2, 4, 2)],
),
]
positions, delta = model.get_mrope_input_positions(
input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51],
mm_features=mm_features,
)
expected = torch.tensor(
[
[0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10],
[0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10],
]
)
assert torch.equal(positions, expected)
assert delta == -2
@@ -1,145 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass, field
import pytest
import torch
from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalFieldElem,
MultiModalKwargsItem,
PlaceholderRange,
)
pytestmark = pytest.mark.skip_global_cleanup
@pytest.fixture(autouse=True, scope="module")
def _force_cpu_default_device():
original = torch.get_default_device()
torch.set_default_device("cpu")
yield
torch.set_default_device(original)
@dataclass
class DummyVisionConfig:
spatial_merge_size: int = 2
@dataclass
class DummyConfig:
vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration:
model = object.__new__(KeyeVL1_5ForConditionalGeneration)
model.config = config
return model
def make_mm_feature(
*,
modality: str,
offset: int,
length: int,
grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
is_embed: list[bool] | None = None,
) -> MultiModalFeatureSpec:
field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
return MultiModalFeatureSpec(
data=MultiModalKwargsItem(
{
field_name: MultiModalFieldElem(
data=torch.tensor(grid_thw),
field=None, # HACK.
),
}
),
modality=modality,
identifier="DUMMY",
mm_position=PlaceholderRange(
offset=offset,
length=length,
is_embed=None if is_embed is None else torch.tensor(is_embed),
),
)
def test_get_mrope_input_positions_text_only():
model = make_model(DummyConfig())
positions, delta = model.get_mrope_input_positions(
input_tokens=[11, 12, 13, 14, 15],
mm_features=[],
)
expected = torch.tensor(
[
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
]
)
assert torch.equal(positions, expected)
assert delta == 0
def test_get_mrope_input_positions_single_image():
model = make_model(DummyConfig())
mm_features = [
make_mm_feature(
modality="image",
offset=1,
length=4,
grid_thw=(1, 4, 4),
)
]
positions, delta = model.get_mrope_input_positions(
input_tokens=[10, 20, 21, 22, 23, 30, 31],
mm_features=mm_features,
)
expected = torch.tensor(
[
[0, 1, 1, 1, 1, 3, 4],
[0, 1, 1, 2, 2, 3, 4],
[0, 1, 2, 1, 2, 3, 4],
]
)
assert torch.equal(positions, expected)
assert delta == -2
def test_get_mrope_input_positions_video_uses_embed_ranges():
model = make_model(DummyConfig())
mm_features = [
make_mm_feature(
modality="video",
offset=1,
length=8,
grid_thw=[(2, 4, 2)],
is_embed=[False, False, True, True, False, False, True, True],
)
]
positions, delta = model.get_mrope_input_positions(
input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41],
mm_features=mm_features,
)
expected = torch.tensor(
[
[0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
]
)
assert torch.equal(positions, expected)
assert delta == 0
@@ -1,210 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass, field
import pytest
import torch
from vllm.model_executor.models.paddleocr_vl import (
PaddleOCRVLForConditionalGeneration,
)
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalFieldElem,
MultiModalKwargsItem,
PlaceholderRange,
)
pytestmark = pytest.mark.skip_global_cleanup
@pytest.fixture(autouse=True, scope="module")
def _force_cpu_default_device():
original = torch.get_default_device()
torch.set_default_device("cpu")
yield
torch.set_default_device(original)
@dataclass
class DummyVisionConfig:
spatial_merge_size: int = 2
patch_size: int = 14
@dataclass
class DummyConfig:
image_token_id: int = 151655
video_token_id: int = 151654
vision_start_token_id: int = 151652
vision_end_token_id: int = 151653
vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
def make_model(config: DummyConfig) -> PaddleOCRVLForConditionalGeneration:
model = object.__new__(PaddleOCRVLForConditionalGeneration)
model.config = config
return model
def make_mm_feature(
*,
offset: int,
length: int,
image_grid_thw: tuple[int, int, int],
) -> MultiModalFeatureSpec:
return MultiModalFeatureSpec(
data=MultiModalKwargsItem(
{
"image_grid_thw": MultiModalFieldElem(
data=torch.tensor(image_grid_thw),
field=None,
),
}
),
modality="image",
identifier="DUMMY",
mm_position=PlaceholderRange(offset=offset, length=length),
)
def test_get_mrope_input_positions_text_only():
model = make_model(DummyConfig())
input_tokens = [11, 12, 13, 14, 15]
positions, delta = model.get_mrope_input_positions(
input_tokens=input_tokens,
mm_features=[],
)
expected = torch.tensor(
[
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
]
)
assert torch.equal(positions, expected)
assert delta == 0
def test_get_mrope_input_positions_single_image():
model = make_model(DummyConfig())
spatial_merge_size = model.config.vision_config.spatial_merge_size
t, h, w = 1, 2, 2
num_image_tokens = t * h * w
input_tokens = (
[10]
+ [model.config.vision_start_token_id]
+ [model.config.image_token_id] * num_image_tokens
+ [model.config.vision_end_token_id]
+ [30, 31]
)
mm_features = [
make_mm_feature(
offset=2, # 1 (text) + 1 (vision_start)
length=num_image_tokens,
image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
)
]
positions, delta = model.get_mrope_input_positions(
input_tokens=input_tokens,
mm_features=mm_features,
)
expected = torch.tensor(
[
[0, 1, 2, 2, 2, 2, 4, 5, 6],
[0, 1, 2, 2, 3, 3, 4, 5, 6],
[0, 1, 2, 3, 2, 3, 4, 5, 6],
]
)
assert torch.equal(positions, expected)
expected_delta = (positions.max().item() + 1) - len(input_tokens)
assert delta == expected_delta
def test_get_mrope_input_positions_multiple_images():
model = make_model(DummyConfig())
spatial_merge_size = model.config.vision_config.spatial_merge_size
t1, h1, w1 = 1, 2, 2
num1 = t1 * h1 * w1
t2, h2, w2 = 1, 1, 3
num2 = t2 * h2 * w2
input_tokens = (
[10]
+ [model.config.vision_start_token_id]
+ [model.config.image_token_id] * num1
+ [model.config.vision_end_token_id]
+ [20, 21]
+ [model.config.vision_start_token_id]
+ [model.config.image_token_id] * num2
+ [model.config.vision_end_token_id]
+ [30]
)
mm_features = [
make_mm_feature(
offset=2,
length=num1,
image_grid_thw=(t1, h1 * spatial_merge_size, w1 * spatial_merge_size),
),
make_mm_feature(
offset=2 + num1 + 1 + 2 + 1,
length=num2,
image_grid_thw=(t2, h2 * spatial_merge_size, w2 * spatial_merge_size),
),
]
positions, delta = model.get_mrope_input_positions(
input_tokens=input_tokens,
mm_features=mm_features,
)
assert positions.shape == (3, 15)
assert not torch.equal(positions[:, 2:6], torch.arange(4).expand(3, 4) + 2)
assert not torch.equal(positions[:, 10:13], torch.arange(3).expand(3, 3) + 10)
def test_get_mrope_input_positions_image_at_start():
model = make_model(DummyConfig())
spatial_merge_size = model.config.vision_config.spatial_merge_size
t, h, w = 1, 2, 2
num_tokens = t * h * w
input_tokens = (
[model.config.vision_start_token_id]
+ [model.config.image_token_id] * num_tokens
+ [model.config.vision_end_token_id]
+ [10, 11]
)
mm_features = [
make_mm_feature(
offset=1, # start token at index 0
length=num_tokens,
image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
)
]
positions, delta = model.get_mrope_input_positions(
input_tokens=input_tokens,
mm_features=mm_features,
)
expected = torch.tensor(
[
[0, 1, 1, 1, 1, 3, 4, 5],
[0, 1, 1, 2, 2, 3, 4, 5],
[0, 1, 2, 1, 2, 3, 4, 5],
]
)
assert torch.equal(positions, expected)
+18
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import tempfile
import huggingface_hub.constants
@@ -9,10 +10,26 @@ from huggingface_hub.utils import LocalEntryNotFoundError
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf,
enable_hf_transfer,
maybe_remap_kv_scale_name,
)
def test_hf_transfer_auto_activation():
if "HF_HUB_ENABLE_HF_TRANSFER" in os.environ:
# in case it is already set, we can't test the auto activation
pytest.skip("HF_HUB_ENABLE_HF_TRANSFER is set, can't test auto activation")
enable_hf_transfer()
try:
# enable hf hub transfer if available
import hf_transfer # type: ignore # noqa
HF_TRANSFER_ACTIVE = True
except ImportError:
HF_TRANSFER_ACTIVE = False
assert huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER == HF_TRANSFER_ACTIVE
def test_download_weights_from_hf():
with tempfile.TemporaryDirectory() as tmpdir:
# assert LocalEntryNotFoundError error is thrown
@@ -161,4 +178,5 @@ class TestMaybeRemapKvScaleName:
if __name__ == "__main__":
test_hf_transfer_auto_activation()
test_download_weights_from_hf()
@@ -46,7 +46,7 @@ AITER_MODEL_LIST = [
),
pytest.param(
"openai-community/gpt2", # gpt2
marks=[pytest.mark.core_model],
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
),
pytest.param("Milos/slovak-gpt-j-405M"), # gptj
pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode
@@ -100,7 +100,7 @@ AITER_MODEL_LIST = [
pytest.param("bigcode/starcoder2-3b"), # starcoder2
pytest.param(
"TitanML/tiny-mixtral", # mixtral
marks=[pytest.mark.core_model],
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
),
pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus
pytest.param(
@@ -15,7 +15,6 @@ MODELS = [
@pytest.mark.parametrize("dtype", ["bfloat16"])
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.cpu_model
def test_models(
hf_runner,
vllm_runner,
@@ -74,25 +74,10 @@ class MtebEmbedMixin(mteb.EncoderProtocol):
return sim
class HfMtebEncoder(MtebEmbedMixin):
def __init__(self, model):
self.model = model
def encode(
self,
inputs: DataLoader[mteb.types.BatchedInput],
*args,
**kwargs,
) -> np.ndarray:
sentences = [text for batch in inputs for text in batch["text"]]
return self.model.encode(sentences)
class VllmMtebEncoder(MtebEmbedMixin):
def __init__(self, vllm_model, prompt_prefix: str | None = None):
def __init__(self, vllm_model):
self.llm = vllm_model
self.rng = np.random.default_rng(seed=42)
self.prompt_prefix = prompt_prefix
def encode(
self,
@@ -102,11 +87,7 @@ class VllmMtebEncoder(MtebEmbedMixin):
) -> np.ndarray:
# Hoping to discover potential scheduling
# issues by randomizing the order.
sentences = [
self.prompt_prefix + text if self.prompt_prefix else text
for batch in inputs
for text in batch["text"]
]
sentences = [text for batch in inputs for text in batch["text"]]
r = self.rng.permutation(len(sentences))
sentences = [sentences[i] for i in r]
outputs = self.llm.embed(sentences, use_tqdm=False)
@@ -162,7 +143,6 @@ def mteb_test_embed_models(
vllm_extra_kwargs=None,
hf_model_callback=None,
atol=MTEB_EMBED_TOL,
prompt_prefix: str | None = None,
):
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
@@ -202,7 +182,7 @@ def mteb_test_embed_models(
)
vllm_main_score = run_mteb_embed_task(
VllmMtebEncoder(vllm_model, prompt_prefix=prompt_prefix), MTEB_EMBED_TASKS
VllmMtebEncoder(vllm_model), MTEB_EMBED_TASKS
)
vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
head_dtype = model_config.head_dtype
@@ -230,9 +210,7 @@ def mteb_test_embed_models(
if hf_model_callback is not None:
hf_model_callback(hf_model)
st_main_score = run_mteb_embed_task(
HfMtebEncoder(hf_model), MTEB_EMBED_TASKS
)
st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
st_dtype = next(hf_model.model.parameters()).dtype
# Check embeddings close to hf outputs
@@ -69,10 +69,7 @@ MODELS = [
attn_type="decoder",
is_prefix_caching_supported=True,
is_chunked_prefill_supported=True,
# Skip: model's custom tokenizer on HF hub is incompatible with
# transformers v5 (sets attrs before super().__init__, triggering
# AttributeError on 'verbose' in __getattr__).
enable_test=False,
enable_test=True,
),
]
@@ -72,8 +72,7 @@ MODELS = [
attn_type="encoder_only",
is_prefix_caching_supported=False,
is_chunked_prefill_supported=False,
# Skip: numerical regression with transformers v5.
enable_test=False,
enable_test=True,
),
########## ModernBertModel
EmbedModelInfo(
@@ -28,16 +28,7 @@ EMBEDDING_MODELS = [
attn_type="encoder_only",
is_prefix_caching_supported=False,
is_chunked_prefill_supported=False,
),
EmbedModelInfo(
"jinaai/jina-embeddings-v5-text-small",
mteb_score=0.794535707854956,
architecture="JinaEmbeddingsV5Model",
seq_pooling_type="LAST",
attn_type="decoder",
is_prefix_caching_supported=True,
is_chunked_prefill_supported=True,
),
)
]
RERANK_MODELS = [
@@ -55,18 +46,11 @@ RERANK_MODELS = [
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
task = "retrieval" if "v5" in model_info.name else "text-matching"
prompt_prefix: str | None = "Document: " if "v5" in model_info.name else None
def hf_model_callback(model):
model.encode = partial(model.encode, task=task)
model.encode = partial(model.encode, task="text-matching")
mteb_test_embed_models(
hf_runner,
vllm_runner,
model_info,
hf_model_callback=hf_model_callback,
prompt_prefix=prompt_prefix,
hf_runner, vllm_runner, model_info, hf_model_callback=hf_model_callback
)
@@ -74,10 +58,8 @@ def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -
def test_embed_models_correctness(
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
) -> None:
task = "retrieval" if "v5" in model_info.name else "text-matching"
def hf_model_callback(model):
model.encode = partial(model.encode, task=task)
model.encode = partial(model.encode, task="text-matching")
correctness_test_embed_models(
hf_runner,
@@ -93,10 +75,6 @@ def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
mteb_test_rerank_models(vllm_runner, model_info)
@pytest.mark.skip(
reason="jinaai/jina-embeddings-v3 custom XLMRobertaLoRA model on HF hub "
"is incompatible with transformers v5 (missing all_tied_weights_keys)"
)
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
@pytest.mark.parametrize("dtype", ["half"])
@pytest.mark.parametrize("dimensions", [16, 32])
@@ -115,14 +93,12 @@ def test_matryoshka(
# ST will strip the input texts, see test_embedding.py
example_prompts = [str(s).strip() for s in example_prompts]
task = "retrieval" if "v5" in model_info.name else "text-matching"
with hf_runner(
model_info.name,
dtype=dtype,
is_sentence_transformer=True,
) as hf_model:
hf_outputs = hf_model.encode(example_prompts, task=task)
hf_outputs = hf_model.encode(example_prompts, task="text-matching")
hf_outputs = matryoshka_fy(hf_outputs, dimensions)
with vllm_runner(
@@ -186,14 +186,7 @@ VLM_TEST_SETTINGS = {
max_num_seqs=2,
auto_cls=AutoModel,
hf_output_post_proc=model_utils.ultravox_trunc_hf_output,
marks=[
pytest.mark.core_model,
pytest.mark.cpu_model,
# TODO: Remove skip once model has been upstreamed to Transformers
pytest.mark.skip(
reason="Custom model code is not compatible with Transformers v5"
),
],
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
),
#### Transformers fallback to test
## To reduce test burden, we only test batching arbitrary image size
@@ -404,14 +397,14 @@ VLM_TEST_SETTINGS = {
"gemma4": VLMTestInfo(
models=["google/gemma-4-E2B-it"],
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
prompt_formatter=lambda img_prompt: f"<bos><|turn>user\n{img_prompt}<turn|>\n<|turn>model\n", # noqa: E501
prompt_formatter=lambda img_prompt: f"<bos><start_of_turn>user\n{img_prompt}<end_of_turn>\n<start_of_turn>model\n", # noqa: E501
single_image_prompts=IMAGE_ASSETS.prompts(
{
"stop_sign": "<|image|>What's the content in the center of the image?", # noqa: E501
"cherry_blossom": "<|image|>What is the season?",
"stop_sign": "What's the content in the center of the image?",
"cherry_blossom": "What is the season?",
}
),
multi_image_prompt="<|image|><|image|>Describe the two images in detail.", # noqa: E501
multi_image_prompt="Describe the two images in detail.",
max_model_len=4096,
max_num_seqs=2,
auto_cls=AutoModelForImageTextToText,
@@ -540,12 +533,6 @@ VLM_TEST_SETTINGS = {
max_model_len=4096,
use_tokenizer_eos=True,
patch_hf_runner=model_utils.internvl_patch_hf_runner,
# TODO: Remove skip once model has been upstreamed to Transformers
marks=[
pytest.mark.skip(
reason="Custom model code tries to access data from meta-tensor"
)
],
),
"intern_vl-video": VLMTestInfo(
models=[
@@ -558,12 +545,6 @@ VLM_TEST_SETTINGS = {
use_tokenizer_eos=True,
patch_hf_runner=model_utils.internvl_patch_hf_runner,
num_logprobs=10 if current_platform.is_rocm() else 5,
# TODO: Remove skip once model has been upstreamed to Transformers
marks=[
pytest.mark.skip(
reason="Custom model code tries to access data from meta-tensor"
)
],
),
"intern_vl-hf": VLMTestInfo(
models=["OpenGVLab/InternVL3-1B-hf"],
@@ -610,8 +591,6 @@ VLM_TEST_SETTINGS = {
hf_model_kwargs={"device_map": "auto"},
patch_hf_runner=model_utils.isaac_patch_hf_runner,
image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
# TODO: Remove skip once model has been upstreamed to Transformers
marks=[pytest.mark.skip(reason="Custom model imports deleted object")], # noqa: E501
),
"kimi_vl": VLMTestInfo(
models=["moonshotai/Kimi-VL-A3B-Instruct"],
@@ -827,12 +806,7 @@ VLM_TEST_SETTINGS = {
pytest.mark.skipif(
Version(TRANSFORMERS_VERSION) == Version("4.57.3"),
reason="This model is broken in Transformers v4.57.3",
),
pytest.mark.skipif(
Version(TRANSFORMERS_VERSION) >= Version("5.0.0"),
reason="Model's custom code uses ROPE_INIT_FUNCTIONS"
"['default'] which was removed in transformers v5",
),
)
],
),
"phi3v": VLMTestInfo(
@@ -986,12 +960,6 @@ VLM_TEST_SETTINGS = {
)
for inp in custom_inputs.different_patch_input_cases_internvl()
],
# TODO: Remove skip once model has been upstreamed to Transformers
marks=[
pytest.mark.skip(
reason="Custom model code tries to access data from meta-tensor"
)
],
),
"llava_onevision-multiple-images": VLMTestInfo(
models=["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"],

Some files were not shown because too many files have changed in this diff Show More