forked from Karylab-cklius/vllm
Compare commits
79
Commits
khluu-patch-1
...
v0.15.1
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1892993bc1 | ||
|
|
7d98f09b1c | ||
|
|
daa2784bb9 | ||
|
|
e4bf6ed90d | ||
|
|
611b18757e | ||
|
|
eec3546bba | ||
|
|
7c023baf58 | ||
|
|
099a787ee2 | ||
|
|
31a64c63a8 | ||
|
|
57eae2f891 | ||
|
|
f0d005864a | ||
|
|
94cbe0a328 | ||
|
|
8b45c58fe9 | ||
|
|
c7039a80b8 | ||
|
|
15ebd0cedf | ||
|
|
2915268369 | ||
|
|
d984d664cc | ||
|
|
5f45b0b7e0 | ||
|
|
a2dba556db | ||
|
|
6ff16b77f8 | ||
|
|
1ed963d43a | ||
|
|
39e8b49378 | ||
|
|
f176443446 | ||
|
|
fe18ce4d3f | ||
|
|
5f7f9ea884 | ||
|
|
7779de34da | ||
|
|
0d8ce320a2 | ||
|
|
d51e1f8b62 | ||
|
|
5042815ab6 | ||
|
|
afb390ab02 | ||
|
|
cf1167e50b | ||
|
|
11b556878b | ||
|
|
ee484b3f4b | ||
|
|
a9b53dd435 | ||
|
|
254db42ede | ||
|
|
105d104576 | ||
|
|
566cdb6cfb | ||
|
|
2f0d3ba745 | ||
|
|
edf927bc9f | ||
|
|
22aeb43007 | ||
|
|
a698e8e7ad | ||
|
|
151e5451c2 | ||
|
|
73b243463b | ||
|
|
7e67df5570 | ||
|
|
ff6c1da4e6 | ||
|
|
fcb9df99bd | ||
|
|
1ebdff412a | ||
|
|
91601ff478 | ||
|
|
d4dbb7af63 | ||
|
|
203d0bc0c2 | ||
|
|
17ab54de81 | ||
|
|
cd775bdbe0 | ||
|
|
da5e7b12be | ||
|
|
719ac592ed | ||
|
|
1209b784f2 | ||
|
|
5fa0f6efa9 | ||
|
|
bc0d291bfe | ||
|
|
9ad7f89f55 | ||
|
|
6450b536a6 | ||
|
|
0f19427db5 | ||
|
|
51931c5c9a | ||
|
|
06b557ecd9 | ||
|
|
81c2a889ce | ||
|
|
8edaf38570 | ||
|
|
5c86a89805 | ||
|
|
0ccecf8833 | ||
|
|
0b9a735e11 | ||
|
|
14d03b8ddb | ||
|
|
d0cbac5827 | ||
|
|
c0d820457a | ||
|
|
97ef11dd34 | ||
|
|
ecc3dd66cc | ||
|
|
7e1f10d562 | ||
|
|
a28b94e6ef | ||
|
|
0118cdcc02 | ||
|
|
136c499f6e | ||
|
|
ebd0a17e0e | ||
|
|
37c9859fab | ||
|
|
4561f13985 |
@@ -274,14 +274,14 @@ steps:
|
||||
- input-release-version
|
||||
- build-wheels
|
||||
|
||||
- label: "Upload release wheels to PyPI and GitHub"
|
||||
- label: "Upload release wheels to PyPI"
|
||||
depends_on:
|
||||
- block-upload-release-wheels
|
||||
id: upload-release-wheels
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload-release-wheels.sh"
|
||||
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
|
||||
|
||||
# =============================================================================
|
||||
# ROCm Release Pipeline (x86_64 only)
|
||||
@@ -638,9 +638,93 @@ steps:
|
||||
depends_on:
|
||||
- step: upload-rocm-wheels
|
||||
allow_failure: true
|
||||
- step: input-release-version
|
||||
allow_failure: true
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
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_postmerge
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm700"
|
||||
|
||||
# ROCm Job 5: Build ROCm Release Docker Image
|
||||
- label: ":rocm: :docker: Build ROCm Release Docker Image"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Download Docker image from S3 (set by build-rocm-base-wheels)
|
||||
DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')"
|
||||
if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then
|
||||
echo "ERROR: rocm-docker-image-s3-path metadata not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Downloading base image from $${DOCKER_IMAGE_S3_PATH}"
|
||||
mkdir -p artifacts/rocm-docker-image
|
||||
aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz
|
||||
|
||||
# Load base Docker image
|
||||
echo "Loading base Docker image..."
|
||||
LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load)
|
||||
BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //')
|
||||
echo "Loaded base image: $${BASE_IMAGE_TAG}"
|
||||
|
||||
# Tag and push the base image to ECR
|
||||
docker tag "$${BASE_IMAGE_TAG}" public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base
|
||||
echo "Pushed base image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base"
|
||||
|
||||
# Get GPU architectures from meta-data
|
||||
PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')"
|
||||
PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}"
|
||||
|
||||
# Build vLLM ROCm release image using cached base
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \
|
||||
--build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
|
||||
--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 "Pushed: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
@@ -11,51 +11,80 @@ fi
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
|
||||
(Optional) For CUDA 13.0:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
|
||||
(Optional) For CPU:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
\`\`\`
|
||||
|
||||
To download the wheel (by version):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
|
||||
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu129/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu130/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux1_x86_64.whl .
|
||||
\`\`\`
|
||||
|
||||
To download and upload the image:
|
||||
|
||||
\`\`\`
|
||||
Download images:
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
Tag and push images:
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
|
||||
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
|
||||
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai:rocm
|
||||
docker tag vllm/vllm-openai:rocm vllm/vllm-openai:latest-rocm
|
||||
docker tag vllm/vllm-openai:rocm vllm/vllm-openai:v${RELEASE_VERSION}-rocm
|
||||
docker push vllm/vllm-openai:latest-rocm
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-rocm
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm
|
||||
|
||||
Create multi-arch manifest:
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest
|
||||
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu130
|
||||
docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker manifest push vllm/vllm-openai:latest-cu130
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
|
||||
\`\`\`
|
||||
EOF
|
||||
|
||||
@@ -3,25 +3,32 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Generate Buildkite annotation for ROCm wheel release
|
||||
|
||||
set -ex
|
||||
|
||||
# Get build configuration from meta-data
|
||||
# Extract ROCm version dynamically from Dockerfile.rocm_base
|
||||
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.1-complete -> extracts "7.1"
|
||||
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0"
|
||||
ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown")
|
||||
PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12")
|
||||
PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151")
|
||||
|
||||
# TODO: Enable the nightly build for ROCm
|
||||
# Get release version, default to 1.0.0.dev for nightly/per-commit builds
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "")
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
RELEASE_VERSION="1.0.0.dev"
|
||||
fi
|
||||
|
||||
# S3 URLs
|
||||
S3_BUCKET="${S3_BUCKET:-vllm-wheels}"
|
||||
S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
|
||||
S3_URL="https://${S3_BUCKET}.s3.${S3_REGION}.amazonaws.com"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}"
|
||||
S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
|
||||
|
||||
# Format ROCm version for path (e.g., "7.1" -> "rocm710")
|
||||
ROCM_VERSION_PATH="rocm$(echo ${ROCM_VERSION} | tr -d '.')"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}"
|
||||
buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF
|
||||
## :rocm: ROCm Wheel Release
|
||||
|
||||
## ROCm Wheel and Docker Image Releases
|
||||
### Build Configuration
|
||||
| Setting | Value |
|
||||
|---------|-------|
|
||||
@@ -34,41 +41,72 @@ buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' <<
|
||||
### :package: Installation
|
||||
|
||||
**Install from this build (by commit):**
|
||||
\`\`\`bash
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/{rocm_variant}/
|
||||
|
||||
# Example:
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/rocm700/
|
||||
\`\`\`bash
|
||||
pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
|
||||
# Example for ROCm ${ROCM_VERSION}:
|
||||
pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
\`\`\`
|
||||
|
||||
**Install from nightly (if published):**
|
||||
|
||||
\`\`\`bash
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/
|
||||
pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
\`\`\`
|
||||
|
||||
### :floppy_disk: Download Wheels Directly
|
||||
|
||||
\`\`\`bash
|
||||
# List all ROCm wheels
|
||||
aws s3 ls s3://${S3_BUCKET}/${ROCM_PATH}/
|
||||
|
||||
aws s3 ls s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/
|
||||
# Download specific wheels
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/vllm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torch-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/triton_rocm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torchvision-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/amdsmi-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/vllm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torch-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-kernels-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl .
|
||||
\`\`\`
|
||||
|
||||
### :gear: Included Packages
|
||||
- **vllm**: vLLM with ROCm support
|
||||
- **torch**: PyTorch built for ROCm ${ROCM_VERSION}
|
||||
- **triton_rocm**: Triton built for ROCm
|
||||
- **triton**: Triton
|
||||
- **triton-kernels**: Triton kernels
|
||||
- **torchvision**: TorchVision for ROCm PyTorch
|
||||
- **torchaudio**: Torchaudio for ROCm PyTorch
|
||||
- **amdsmi**: AMD SMI Python bindings
|
||||
- **aiter**: Aiter for ROCm
|
||||
- **flash-attn**: Flash Attention for ROCm
|
||||
|
||||
### :warning: Notes
|
||||
- These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs
|
||||
- Supported GPU architectures: ${PYTORCH_ROCM_ARCH}
|
||||
- Platform: Linux x86_64 only
|
||||
|
||||
### :package: Docker Image Release
|
||||
|
||||
To download and upload the image:
|
||||
|
||||
\`\`\`
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
\`\`\`
|
||||
|
||||
EOF
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -e
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
|
||||
GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
|
||||
|
||||
echo "Release version from Buildkite: $RELEASE_VERSION"
|
||||
|
||||
if [[ -z "$GIT_VERSION" ]]; then
|
||||
echo "[FATAL] Not on a git tag, cannot create release."
|
||||
exit 1
|
||||
else
|
||||
echo "Git version for commit $BUILDKITE_COMMIT: $GIT_VERSION"
|
||||
fi
|
||||
# sanity check for version mismatch
|
||||
if [[ "$RELEASE_VERSION" != "$GIT_VERSION" ]]; then
|
||||
if [[ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]]; then
|
||||
echo "[WARNING] Force release and ignore version mismatch"
|
||||
else
|
||||
echo "[FATAL] Release version from Buildkite does not match Git version."
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
PURE_VERSION=${RELEASE_VERSION#v} # remove leading 'v'
|
||||
|
||||
# check pypi token
|
||||
if [[ -z "$PYPI_TOKEN" ]]; then
|
||||
echo "[FATAL] PYPI_TOKEN is not set."
|
||||
exit 1
|
||||
else
|
||||
export TWINE_USERNAME="__token__"
|
||||
export TWINE_PASSWORD="$PYPI_TOKEN"
|
||||
fi
|
||||
|
||||
set -x # avoid printing secrets above
|
||||
|
||||
# install twine from pypi
|
||||
python3 -m venv /tmp/vllm-release-env
|
||||
source /tmp/vllm-release-env/bin/activate
|
||||
pip install twine
|
||||
python3 -m twine --version
|
||||
|
||||
# copy release wheels to local directory
|
||||
DIST_DIR=/tmp/vllm-release-dist
|
||||
echo "Existing wheels on S3:"
|
||||
aws s3 ls "$S3_COMMIT_PREFIX"
|
||||
echo "Copying wheels to local directory"
|
||||
mkdir -p $DIST_DIR
|
||||
# include only wheels for the release version, ignore all files with "dev" or "rc" in the name (without excluding 'aarch64')
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
|
||||
echo "Wheels copied to local directory"
|
||||
# generate source tarball
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
ls -la $DIST_DIR
|
||||
|
||||
# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
|
||||
PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${PURE_VERSION}*.whl" -not -name "*+*")
|
||||
if [[ -z "$PYPI_WHEEL_FILES" ]]; then
|
||||
echo "No default variant wheels found, quitting..."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 -m twine check $PYPI_WHEEL_FILES
|
||||
python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
|
||||
echo "Wheels uploaded to PyPI"
|
||||
@@ -1,103 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -e
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
|
||||
echo "Release version from Buildkite: $RELEASE_VERSION"
|
||||
GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
|
||||
if [ -z "$GIT_VERSION" ]; then
|
||||
echo "[FATAL] Not on a git tag, cannot create release."
|
||||
exit 1
|
||||
else
|
||||
echo "Git version for commit $BUILDKITE_COMMIT: $GIT_VERSION"
|
||||
fi
|
||||
# sanity check for version mismatch
|
||||
if [ "v$RELEASE_VERSION" != "$GIT_VERSION" ]; then
|
||||
if [ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]; then
|
||||
echo "[WARNING] Force release and ignore version mismatch"
|
||||
else
|
||||
echo "[FATAL] Release version from Buildkite does not match Git version."
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# check pypi token
|
||||
if [ -z "$PYPI_TOKEN" ]; then
|
||||
echo "[FATAL] PYPI_TOKEN is not set."
|
||||
exit 1
|
||||
else
|
||||
export TWINE_USERNAME="__token__"
|
||||
export TWINE_PASSWORD="$PYPI_TOKEN"
|
||||
fi
|
||||
|
||||
# check github token
|
||||
if [ -z "$GITHUB_TOKEN" ]; then
|
||||
echo "[FATAL] GITHUB_TOKEN is not set."
|
||||
exit 1
|
||||
else
|
||||
export GH_TOKEN="$GITHUB_TOKEN"
|
||||
fi
|
||||
|
||||
set -x # avoid printing secrets above
|
||||
|
||||
# download gh CLI from github
|
||||
# Get latest gh CLI version from GitHub API
|
||||
GH_VERSION=$(curl -s https://api.github.com/repos/cli/cli/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/' | sed 's/^v//')
|
||||
if [ -z "$GH_VERSION" ]; then
|
||||
echo "[FATAL] Failed to get latest gh CLI version from GitHub"
|
||||
exit 1
|
||||
fi
|
||||
echo "Downloading gh CLI version: $GH_VERSION"
|
||||
GH_TARBALL="gh_${GH_VERSION}_linux_amd64.tar.gz"
|
||||
GH_URL="https://github.com/cli/cli/releases/download/v${GH_VERSION}/${GH_TARBALL}"
|
||||
GH_INSTALL_DIR="/tmp/gh-install"
|
||||
mkdir -p "$GH_INSTALL_DIR"
|
||||
pushd "$GH_INSTALL_DIR"
|
||||
curl -L -o "$GH_TARBALL" "$GH_URL"
|
||||
tar -xzf "$GH_TARBALL"
|
||||
GH_BIN=$(realpath $(find . -name "gh" -type f -executable | head -n 1))
|
||||
if [ -z "$GH_BIN" ]; then
|
||||
echo "[FATAL] Failed to find gh CLI executable"
|
||||
exit 1
|
||||
fi
|
||||
echo "gh CLI downloaded successfully, version: $($GH_BIN --version)"
|
||||
echo "Last 5 releases on GitHub:" # as a sanity check of gh and GH_TOKEN
|
||||
command "$GH_BIN" release list --limit 5
|
||||
popd
|
||||
|
||||
# install twine from pypi
|
||||
python3 -m venv /tmp/vllm-release-env
|
||||
source /tmp/vllm-release-env/bin/activate
|
||||
pip install twine
|
||||
python3 -m twine --version
|
||||
|
||||
# copy release wheels to local directory
|
||||
DIST_DIR=/tmp/vllm-release-dist
|
||||
echo "Existing wheels on S3:"
|
||||
aws s3 ls "$S3_COMMIT_PREFIX"
|
||||
echo "Copying wheels to local directory"
|
||||
mkdir -p $DIST_DIR
|
||||
# include only wheels for the release version, ignore all files with "dev" or "rc" in the name
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm-${RELEASE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc*" "$S3_COMMIT_PREFIX" $DIST_DIR
|
||||
echo "Wheels copied to local directory"
|
||||
# generate source tarball
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${RELEASE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
ls -la $DIST_DIR
|
||||
|
||||
|
||||
# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
|
||||
PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${RELEASE_VERSION}*.whl" -not -name "*+*")
|
||||
if [ -z "$PYPI_WHEEL_FILES" ]; then
|
||||
echo "No default variant wheels found, quitting..."
|
||||
exit 1
|
||||
fi
|
||||
python3 -m twine check $PYPI_WHEEL_FILES
|
||||
python3 -m twine --non-interactive --verbose upload $PYPI_WHEEL_FILES
|
||||
echo "Wheels uploaded to PyPI"
|
||||
|
||||
# create release on GitHub with the release version and all wheels
|
||||
command "$GH_BIN" release create $GIT_VERSION -d --latest --notes-from-tag --verify-tag $DIST_DIR/*.whl
|
||||
@@ -29,8 +29,9 @@ jobs:
|
||||
|
||||
- name: Install dependencies and build vLLM
|
||||
run: |
|
||||
uv pip install -r requirements/cpu-build.txt --index-strategy unsafe-best-match
|
||||
uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match
|
||||
uv pip install -e .
|
||||
uv pip install -e . --no-build-isolation
|
||||
env:
|
||||
CMAKE_BUILD_PARALLEL_LEVEL: 4
|
||||
|
||||
|
||||
+9
-9
@@ -377,7 +377,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# preselected input type pairs and schedules.
|
||||
# Generate sources:
|
||||
set(MARLIN_GEN_SCRIPT
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/gptq_marlin/generate_kernels.py)
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py)
|
||||
file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
|
||||
list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
|
||||
set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
|
||||
@@ -412,7 +412,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
if (MARLIN_ARCHS)
|
||||
file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_float16.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_ARCHS}")
|
||||
@@ -422,7 +422,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
|
||||
|
||||
file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_bfloat16.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
|
||||
@@ -434,7 +434,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
if (MARLIN_SM75_ARCHS)
|
||||
file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/gptq_marlin/sm75_kernel_*.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
|
||||
@@ -446,7 +446,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
if (MARLIN_FP8_ARCHS)
|
||||
file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/gptq_marlin/sm89_kernel_*.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
|
||||
@@ -459,10 +459,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
|
||||
set(MARLIN_SRCS
|
||||
"csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
|
||||
"csrc/quantization/gptq_marlin/gptq_marlin.cu"
|
||||
"csrc/quantization/gptq_marlin/marlin_int4_fp8_preprocess.cu"
|
||||
"csrc/quantization/gptq_marlin/gptq_marlin_repack.cu"
|
||||
"csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
|
||||
"csrc/quantization/marlin/marlin.cu"
|
||||
"csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu"
|
||||
"csrc/quantization/marlin/gptq_marlin_repack.cu"
|
||||
"csrc/quantization/marlin/awq_marlin_repack.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_SRCS}"
|
||||
CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
|
||||
|
||||
@@ -20,8 +20,12 @@ FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
PROVIDER_CFGS = {
|
||||
"vllm": dict(backend="vllm", enabled=True),
|
||||
"flashinfer": dict(backend="flashinfer", enabled=True),
|
||||
"vllm": dict(backend="vllm", is_sf_swizzled_layout=False, enabled=True),
|
||||
"vllm-swizzle": dict(backend="vllm", is_sf_swizzled_layout=True, enabled=True),
|
||||
"flashinfer": dict(backend="flashinfer", is_sf_swizzled_layout=False, enabled=True),
|
||||
"flashinfer-swizzle": dict(
|
||||
backend="flashinfer", is_sf_swizzled_layout=True, enabled=True
|
||||
),
|
||||
}
|
||||
|
||||
_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
|
||||
@@ -36,7 +40,7 @@ def compute_global_scale(tensor: torch.Tensor) -> torch.Tensor:
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
|
||||
x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=_enabled,
|
||||
@@ -63,19 +67,36 @@ def benchmark(batch_size, provider, N, K):
|
||||
|
||||
if cfg["backend"] == "vllm":
|
||||
# vLLM's FP4 quantization
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.scaled_fp4_quant(a, a_global_scale),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
if cfg["is_sf_swizzled_layout"]:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.scaled_fp4_quant(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
else:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.scaled_fp4_quant(
|
||||
a, a_global_scale, is_sf_swizzled_layout=False
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
elif cfg["backend"] == "flashinfer":
|
||||
# FlashInfer's FP4 quantization
|
||||
# Use is_sf_swizzled_layout=True to match vLLM's output format
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
if cfg["is_sf_swizzled_layout"]:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
else:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=False
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
|
||||
# Convert ms to us for better readability at small batch sizes
|
||||
to_us = lambda t_ms: t_ms * 1000
|
||||
@@ -92,7 +113,9 @@ def prepare_shapes(args):
|
||||
return out
|
||||
|
||||
|
||||
def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
|
||||
def _test_accuracy_once(
|
||||
M: int, K: int, dtype: torch.dtype, device: str, is_sf_swizzled_layout: bool
|
||||
):
|
||||
"""Test accuracy between vLLM and FlashInfer FP4 quantization."""
|
||||
# Create input tensor
|
||||
a = torch.randn((M, K), device=device, dtype=dtype)
|
||||
@@ -101,11 +124,13 @@ def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
|
||||
a_global_scale = compute_global_scale(a)
|
||||
|
||||
# vLLM quantization
|
||||
vllm_fp4, vllm_scale = ops.scaled_fp4_quant(a, a_global_scale)
|
||||
vllm_fp4, vllm_scale = ops.scaled_fp4_quant(
|
||||
a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
# FlashInfer quantization (with swizzled layout to match vLLM's output)
|
||||
flashinfer_fp4, flashinfer_scale = flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
|
||||
)
|
||||
flashinfer_scale = flashinfer_scale.view(torch.float8_e4m3fn)
|
||||
|
||||
@@ -114,7 +139,14 @@ def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
|
||||
vllm_fp4,
|
||||
flashinfer_fp4,
|
||||
)
|
||||
print(f"M={M}, K={K}, dtype={dtype}: PASSED")
|
||||
# Compare scales
|
||||
torch.testing.assert_close(
|
||||
vllm_scale,
|
||||
flashinfer_scale,
|
||||
)
|
||||
print(
|
||||
f"M={M}, K={K}, dtype={dtype}, is_sf_swizzled_layout={is_sf_swizzled_layout}: PASSED" # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
def test_accuracy():
|
||||
@@ -130,9 +162,10 @@ def test_accuracy():
|
||||
Ms = [1, 1024]
|
||||
Ks = [4096]
|
||||
|
||||
for M in Ms:
|
||||
for K in Ks:
|
||||
_test_accuracy_once(M, K, dtype, device)
|
||||
for is_sf_swizzled_layout in [True, False]:
|
||||
for M in Ms:
|
||||
for K in Ks:
|
||||
_test_accuracy_once(M, K, dtype, device, is_sf_swizzled_layout)
|
||||
|
||||
print("\nAll accuracy tests passed!")
|
||||
|
||||
@@ -145,7 +178,7 @@ if __name__ == "__main__":
|
||||
"--models",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=["meta-llama/Llama-3.1-8B-Instruct"],
|
||||
default=["meta-llama/Llama-3.3-70B-Instruct"],
|
||||
choices=list(WEIGHT_SHAPES.keys()),
|
||||
)
|
||||
parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
|
||||
|
||||
@@ -231,7 +231,7 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
|
||||
assert bt.w_tok_s is None
|
||||
assert bt.group_size is not None
|
||||
|
||||
fn = lambda: ops.gptq_marlin_gemm(
|
||||
fn = lambda: ops.marlin_gemm(
|
||||
a=bt.a,
|
||||
c=None,
|
||||
b_q_weight=w_q,
|
||||
|
||||
@@ -239,7 +239,7 @@ def bench_run(
|
||||
"sm_version": sm_version,
|
||||
"CUBLAS_M_THRESHOLD": CUBLAS_M_THRESHOLD,
|
||||
# Kernels
|
||||
"gptq_marlin_gemm": ops.gptq_marlin_gemm,
|
||||
"marlin_gemm": ops.marlin_gemm,
|
||||
"gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
|
||||
"gptq_marlin_repack": ops.gptq_marlin_repack,
|
||||
"allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
|
||||
@@ -263,21 +263,21 @@ def bench_run(
|
||||
|
||||
results.append(
|
||||
benchmark.Timer(
|
||||
stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501
|
||||
stmt="output = marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501
|
||||
globals=globals,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description="gptq_marlin_gemm",
|
||||
description="marlin_gemm",
|
||||
).blocked_autorange(min_run_time=min_run_time)
|
||||
)
|
||||
|
||||
results.append(
|
||||
benchmark.Timer(
|
||||
stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501
|
||||
stmt="output = marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501
|
||||
globals=globals,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description="gptq_marlin_gemm_fp32",
|
||||
description="marlin_gemm_fp32",
|
||||
).blocked_autorange(min_run_time=min_run_time)
|
||||
)
|
||||
|
||||
|
||||
@@ -13,6 +13,8 @@ endif()
|
||||
#
|
||||
# Define environment variables for special configurations
|
||||
#
|
||||
set(ENABLE_AVX2 $ENV{VLLM_CPU_AVX2})
|
||||
set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
|
||||
set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
|
||||
set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
|
||||
set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
|
||||
@@ -103,6 +105,16 @@ else()
|
||||
find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
|
||||
find_isa(${CPUINFO} "S390" S390_FOUND)
|
||||
find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
|
||||
|
||||
# Support cross-compilation by allowing override via environment variables
|
||||
if (ENABLE_AVX2)
|
||||
set(AVX2_FOUND ON)
|
||||
message(STATUS "AVX2 support enabled via VLLM_CPU_AVX2 environment variable")
|
||||
endif()
|
||||
if (ENABLE_AVX512)
|
||||
set(AVX512_FOUND ON)
|
||||
message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (AVX512_FOUND AND NOT AVX512_DISABLED)
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
|
||||
#endif
|
||||
|
||||
#include "quantization/gptq_marlin/marlin.cuh"
|
||||
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
|
||||
#include "quantization/marlin/marlin.cuh"
|
||||
#include "quantization/marlin/marlin_dtypes.cuh"
|
||||
#include "core/scalar_type.hpp"
|
||||
|
||||
#define MARLIN_KERNEL_PARAMS \
|
||||
|
||||
@@ -23,10 +23,10 @@
|
||||
#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
|
||||
#endif
|
||||
|
||||
#include "quantization/gptq_marlin/marlin.cuh"
|
||||
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
|
||||
#include "quantization/gptq_marlin/dequant.h"
|
||||
#include "quantization/gptq_marlin/marlin_mma.h"
|
||||
#include "quantization/marlin/marlin.cuh"
|
||||
#include "quantization/marlin/marlin_dtypes.cuh"
|
||||
#include "quantization/marlin/dequant.h"
|
||||
#include "quantization/marlin/marlin_mma.h"
|
||||
#include "core/scalar_type.hpp"
|
||||
|
||||
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
|
||||
|
||||
+2
-1
@@ -293,7 +293,8 @@ std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_scale,
|
||||
torch::Tensor const& input_scale);
|
||||
torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
|
||||
void scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
|
||||
@@ -27,17 +27,24 @@
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
float const* SFScale, uint32_t* out,
|
||||
uint32_t* SFout) {
|
||||
using PackedVec = PackedVec<Type>;
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols,
|
||||
int32_t num_padded_cols,
|
||||
Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out,
|
||||
uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = vllm::PackedVec<Type>;
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
@@ -49,34 +56,60 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0];
|
||||
float const SFScaleVal = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
// Input tensor row/col loops.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
|
||||
for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
colIdx += blockDim.x) {
|
||||
if (colIdx < num_padded_cols) {
|
||||
PackedVec in_vec;
|
||||
PackedVec in_vec2;
|
||||
int64_t inOffset =
|
||||
rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) +
|
||||
numCols / CVT_FP4_ELTS_PER_THREAD + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
PackedVec in_vec2 = reinterpret_cast<PackedVec const*>(in)[inOffset2];
|
||||
|
||||
// Get the output tensor offset.
|
||||
// Same as inOffset because 8 elements are packed into one uint32_t.
|
||||
int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 8],
|
||||
valid);
|
||||
} else {
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 4],
|
||||
valid);
|
||||
}
|
||||
|
||||
// Compute silu and mul
|
||||
PackedVec out_silu_mul = compute_silu_mul(in_vec, in_vec2);
|
||||
PackedVec out_silu_mul = compute_silu_mul<Type>(in_vec, in_vec2);
|
||||
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx, colIdx, numKTiles, SFout);
|
||||
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(out_silu_mul, SFScaleVal,
|
||||
sf_out);
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
out_silu_mul, SFScaleVal, sf_out);
|
||||
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -103,17 +136,23 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
|
||||
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
|
||||
|
||||
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
|
||||
|
||||
int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
|
||||
int grid_x = std::min(
|
||||
int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
|
||||
m, n, input_ptr, input_sf_ptr,
|
||||
m, n, sf_n_unpadded, input_ptr, input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
|
||||
@@ -140,8 +140,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -246,8 +246,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -21,7 +21,8 @@
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf);
|
||||
torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout);
|
||||
#endif
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
@@ -51,10 +52,12 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
#endif
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf) {
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf);
|
||||
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
|
||||
is_sf_swizzled_layout);
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
@@ -27,29 +27,23 @@
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
template <typename Int>
|
||||
__host__ __device__ inline Int round_up(Int x, Int y) {
|
||||
static_assert(std::is_integral_v<Int>,
|
||||
"round_up argument must be integral type");
|
||||
return ((x + y - 1) / y) * y;
|
||||
}
|
||||
|
||||
// Compute effective rows for grid configuration with swizzled SF layouts.
|
||||
inline int computeEffectiveRows(int m) {
|
||||
constexpr int ROW_TILE = 128;
|
||||
return round_up(m, ROW_TILE);
|
||||
}
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
float const* SFScale, uint32_t* out, uint32_t* SFout) {
|
||||
using PackedVec = PackedVec<Type>;
|
||||
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t num_padded_cols,
|
||||
Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = vllm::PackedVec<Type>;
|
||||
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
@@ -59,33 +53,31 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
int32_t const numKTiles = (numCols + 63) / 64;
|
||||
|
||||
int sf_m = round_up<int>(numRows, 128);
|
||||
int sf_n_unpadded = numCols / CVT_FP4_SF_VEC_SIZE;
|
||||
int sf_n_int = round_up<int>(sf_n_unpadded, 4) / 4;
|
||||
int num_padded_cols = sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE;
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
float const global_scale = SFScale == nullptr ? 1.0f : SFScale[0];
|
||||
float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
|
||||
// Iterate over all rows and cols including padded ones -
|
||||
// ensures we visit every single scale factor address to initialize it.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) {
|
||||
for (int colIdx = threadIdx.x;
|
||||
colIdx < num_padded_cols / CVT_FP4_ELTS_PER_THREAD;
|
||||
colIdx += blockDim.x) {
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
if (colIdx < num_padded_cols) {
|
||||
PackedVec in_vec;
|
||||
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
|
||||
// If we are outside valid rows OR outside valid columns -> Use Zeros
|
||||
if (rowIdx >= numRows || elem_idx >= numCols) {
|
||||
memset(&in_vec, 0, sizeof(PackedVec));
|
||||
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
} else {
|
||||
// Valid Region: Load actual data
|
||||
in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
}
|
||||
|
||||
auto sf_out =
|
||||
@@ -94,13 +86,85 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
rowIdx, colIdx, numKTiles, SFout);
|
||||
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, global_scale, sf_out);
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
in_vec, global_scale, sf_out);
|
||||
|
||||
// We do NOT write output for padding because the 'out' tensor is not
|
||||
// padded.
|
||||
if (rowIdx < numRows && elem_idx < numCols) {
|
||||
// Same as inOffset because 8 elements are packed into one uint32_t.
|
||||
out[inOffset] = out_val;
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols,
|
||||
int32_t sf_n_unpadded, Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out,
|
||||
uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = PackedVec<Type>;
|
||||
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
|
||||
// Iterate over all rows and cols including padded ones -
|
||||
// ensures we visit every single scale factor address to initialize it.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
|
||||
if (colIdx < sf_n_unpadded) {
|
||||
PackedVec in_vec;
|
||||
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
|
||||
// If we are outside valid rows OR outside valid columns -> Use Zeros
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
} else {
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
}
|
||||
|
||||
auto sf_out =
|
||||
sf_out_rowmajor_u8<uint32_t>(rowIdx, colIdx, sf_n_unpadded, SFout);
|
||||
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
in_vec, global_scale, sf_out);
|
||||
|
||||
// We do NOT write output for padding because the 'out' tensor is not
|
||||
// padded.
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -111,7 +175,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf) {
|
||||
torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
int32_t m = input.size(0);
|
||||
int32_t n = input.size(1);
|
||||
|
||||
@@ -129,19 +194,48 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
|
||||
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
|
||||
|
||||
// Grid, Block size. Each thread converts 8 values.
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
int effectiveRows = vllm::computeEffectiveRows(m);
|
||||
dim3 grid(std::min(effectiveRows, multiProcessorCount * numBlocksPerSM));
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
|
||||
m, n, input_ptr, input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
}
|
||||
if (is_sf_swizzled_layout) {
|
||||
int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);
|
||||
int32_t num_padded_cols =
|
||||
sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x));
|
||||
int grid_x =
|
||||
std::min(vllm::computeEffectiveRows(m),
|
||||
std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
|
||||
m, n, num_padded_cols, input_ptr, input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
} else {
|
||||
int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
|
||||
int grid_x = std::min(
|
||||
m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
|
||||
<<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, input_ptr,
|
||||
input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,9 +19,17 @@
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
#define ELTS_PER_THREAD 8
|
||||
|
||||
#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
|
||||
defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
|
||||
#define ELTS_PER_THREAD 16
|
||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 16;
|
||||
constexpr bool CVT_FP4_PACK16 = true;
|
||||
#else
|
||||
#define ELTS_PER_THREAD 8
|
||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
||||
constexpr bool CVT_FP4_PACK16 = false;
|
||||
#endif
|
||||
|
||||
constexpr int CVT_FP4_SF_VEC_SIZE = 16;
|
||||
|
||||
namespace vllm {
|
||||
@@ -68,19 +76,46 @@ struct TypeConverter<__nv_bfloat16> {
|
||||
using Type = __nv_bfloat162;
|
||||
};
|
||||
|
||||
#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
|
||||
defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
|
||||
// Define a 32 bytes packed data type.
|
||||
template <class Type>
|
||||
struct alignas(32) PackedVec {
|
||||
typename TypeConverter<Type>::Type elts[8];
|
||||
};
|
||||
#else
|
||||
// Define a 16 bytes packed data type.
|
||||
template <class Type>
|
||||
struct PackedVec {
|
||||
struct alignas(16) PackedVec {
|
||||
typename TypeConverter<Type>::Type elts[4];
|
||||
};
|
||||
#endif
|
||||
|
||||
template <>
|
||||
struct PackedVec<__nv_fp8_e4m3> {
|
||||
__nv_fp8x2_e4m3 elts[8];
|
||||
};
|
||||
|
||||
template <typename Int>
|
||||
__host__ __device__ inline Int round_up(Int x, Int y) {
|
||||
static_assert(std::is_integral_v<Int>,
|
||||
"round_up argument must be integral type");
|
||||
return ((x + y - 1) / y) * y;
|
||||
}
|
||||
|
||||
template <typename Int>
|
||||
__host__ __device__ __forceinline__ Int div_round_up(Int x, Int y) {
|
||||
return (x + y - 1) / y;
|
||||
}
|
||||
|
||||
// Compute effective rows for grid configuration with swizzled SF layouts.
|
||||
inline int computeEffectiveRows(int m) {
|
||||
constexpr int ROW_TILE = 128;
|
||||
return round_up(m, ROW_TILE);
|
||||
}
|
||||
|
||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
||||
inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
||||
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
|
||||
uint32_t val;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
@@ -101,7 +136,7 @@ inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
||||
}
|
||||
|
||||
// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
|
||||
inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
||||
__device__ __forceinline__ uint32_t fp32_vec8_to_e2m1(float2 (&array)[4]) {
|
||||
uint32_t val;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
@@ -114,20 +149,115 @@ inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
||||
"}"
|
||||
"}\n"
|
||||
: "=r"(val)
|
||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
|
||||
return val;
|
||||
}
|
||||
|
||||
struct u32x2 {
|
||||
uint32_t lo, hi;
|
||||
};
|
||||
|
||||
using fp4_packed_t = std::conditional_t<CVT_FP4_PACK16, u32x2, uint32_t>;
|
||||
|
||||
__device__ __forceinline__ u32x2 fp32_vec16_to_e2m1(float2 (&array)[8]) {
|
||||
u32x2 out;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b8 b0;\n"
|
||||
".reg .b8 b1;\n"
|
||||
".reg .b8 b2;\n"
|
||||
".reg .b8 b3;\n"
|
||||
".reg .b8 b4;\n"
|
||||
".reg .b8 b5;\n"
|
||||
".reg .b8 b6;\n"
|
||||
".reg .b8 b7;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b0, %3, %2;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b1, %5, %4;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b2, %7, %6;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b3, %9, %8;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b4, %11, %10;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b5, %13, %12;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b6, %15, %14;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b7, %17, %16;\n"
|
||||
"mov.b32 %0, {b0, b1, b2, b3};\n"
|
||||
"mov.b32 %1, {b4, b5, b6, b7};\n"
|
||||
"}\n"
|
||||
: "=r"(out.lo), "=r"(out.hi)
|
||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y),
|
||||
"f"(array[4].x), "f"(array[4].y), "f"(array[5].x), "f"(array[5].y),
|
||||
"f"(array[6].x), "f"(array[6].y), "f"(array[7].x), "f"(array[7].y));
|
||||
return out;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t pack_fp4(float2 (&v)[4]) {
|
||||
return fp32_vec8_to_e2m1(v);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ u32x2 pack_fp4(float2 (&v)[8]) {
|
||||
return fp32_vec16_to_e2m1(v);
|
||||
}
|
||||
|
||||
// Fast reciprocal.
|
||||
inline __device__ float reciprocal_approximate_ftz(float a) {
|
||||
__device__ __forceinline__ float reciprocal_approximate_ftz(float a) {
|
||||
float b;
|
||||
asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
|
||||
asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a));
|
||||
return b;
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__device__ __forceinline__ void ld128_or_zero_cg_u32(PackedVec<Type>& out,
|
||||
const void* ptr,
|
||||
bool pred) {
|
||||
uint32_t r0, r1, r2, r3;
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred pr;\n"
|
||||
" setp.ne.u32 pr, %4, 0;\n"
|
||||
" mov.u32 %0, 0;\n"
|
||||
" mov.u32 %1, 0;\n"
|
||||
" mov.u32 %2, 0;\n"
|
||||
" mov.u32 %3, 0;\n"
|
||||
" @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
|
||||
"}\n"
|
||||
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
|
||||
: "r"((int)pred), "l"(ptr));
|
||||
|
||||
*reinterpret_cast<uint4*>(&out) = uint4{r0, r1, r2, r3};
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__device__ __forceinline__ void ld256_or_zero_cg_u32(PackedVec<Type>& out,
|
||||
const void* ptr,
|
||||
bool pred) {
|
||||
uint32_t r0, r1, r2, r3, r4, r5, r6, r7;
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred pr;\n"
|
||||
" setp.ne.u32 pr, %8, 0;\n"
|
||||
" mov.u32 %0, 0;\n"
|
||||
" mov.u32 %1, 0;\n"
|
||||
" mov.u32 %2, 0;\n"
|
||||
" mov.u32 %3, 0;\n"
|
||||
" mov.u32 %4, 0;\n"
|
||||
" mov.u32 %5, 0;\n"
|
||||
" mov.u32 %6, 0;\n"
|
||||
" mov.u32 %7, 0;\n"
|
||||
" @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
|
||||
"}\n"
|
||||
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6),
|
||||
"=r"(r7)
|
||||
: "r"((int)pred), "l"(ptr));
|
||||
|
||||
reinterpret_cast<uint4*>(&out)[0] = uint4{r0, r1, r2, r3};
|
||||
reinterpret_cast<uint4*>(&out)[1] = uint4{r4, r5, r6, r7};
|
||||
}
|
||||
|
||||
// Compute SF output offset for swizzled tensor core layout.
|
||||
// SF layout: [numMTiles, numKTiles, 32, 4, 4]
|
||||
// Caller must precompute: numKTiles = (numCols + 63) / 64
|
||||
@@ -166,21 +296,41 @@ __device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
|
||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
||||
}
|
||||
|
||||
template <class SFType>
|
||||
__device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
|
||||
int packs_per_row_sf,
|
||||
SFType* SFout) {
|
||||
constexpr int PACK = CVT_FP4_ELTS_PER_THREAD;
|
||||
constexpr int THREADS_PER_SF =
|
||||
CVT_FP4_SF_VEC_SIZE / PACK; // 1 if PACK=16, 2 else PACK=8
|
||||
|
||||
if (threadIdx.x % THREADS_PER_SF != 0) return nullptr;
|
||||
|
||||
int sf_col =
|
||||
pack / THREADS_PER_SF; // PACK=16 => sf_col=pack; PACK=8 => sf_col=pack/2
|
||||
int64_t off = (int64_t)row * packs_per_row_sf + sf_col;
|
||||
|
||||
return (uint8_t*)SFout + off;
|
||||
}
|
||||
|
||||
// Quantizes the provided PackedVec into the uint32_t output
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
||||
uint8_t* SFout) {
|
||||
template <class Type, int CVT_FP4_NUM_THREADS_PER_SF, bool UE8M0_SF = false>
|
||||
__device__ __forceinline__ fp4_packed_t
|
||||
cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
|
||||
// Get absolute maximum values among the local 8 values.
|
||||
auto localMax = __habs2(vec.elts[0]);
|
||||
|
||||
// Local maximum value.
|
||||
// Local maximum value.
|
||||
#pragma unroll
|
||||
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
||||
localMax = __hmax2(localMax, __habs2(vec.elts[i]));
|
||||
}
|
||||
|
||||
// Get the absolute maximum among all 16 values (two threads).
|
||||
localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
|
||||
|
||||
if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) {
|
||||
localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax);
|
||||
}
|
||||
// Get the final absolute maximum values.
|
||||
float vecMax = float(__hmax(localMax.x, localMax.y));
|
||||
|
||||
@@ -205,18 +355,17 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
||||
// Convert back to fp32.
|
||||
SFValue = float(tmp);
|
||||
}
|
||||
|
||||
// Write the SF to global memory (STG.8).
|
||||
if (SFout) *SFout = fp8SFVal;
|
||||
|
||||
// Get the output scale.
|
||||
// Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
|
||||
// reciprocal(SFScaleVal))
|
||||
float outputScale =
|
||||
SFValue != 0 ? reciprocal_approximate_ftz(
|
||||
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
||||
: 0.0f;
|
||||
|
||||
if (SFout) {
|
||||
// Write the SF to global memory (STG.8).
|
||||
*SFout = fp8SFVal;
|
||||
}
|
||||
SFValue != 0.0f ? reciprocal_approximate_ftz(
|
||||
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
||||
: 0.0f;
|
||||
|
||||
// Convert the input to float.
|
||||
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
||||
@@ -233,10 +382,7 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
||||
}
|
||||
|
||||
// Convert to e2m1 values.
|
||||
uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
|
||||
|
||||
// Write the e2m1 values to global memory.
|
||||
return e2m1Vec;
|
||||
return pack_fp4(fp2Vals);
|
||||
}
|
||||
|
||||
// silu in float32
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <iostream>
|
||||
#include "../gptq_marlin/marlin_dtypes.cuh"
|
||||
#include "../marlin/marlin_dtypes.cuh"
|
||||
using marlin::MarlinScalarType2;
|
||||
|
||||
namespace allspark {
|
||||
|
||||
@@ -46,7 +46,7 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
|
||||
|
||||
} // namespace marlin
|
||||
|
||||
torch::Tensor gptq_marlin_gemm(
|
||||
torch::Tensor marlin_gemm(
|
||||
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
|
||||
torch::Tensor& b_q_weight,
|
||||
std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
|
||||
@@ -528,7 +528,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
|
||||
|
||||
} // namespace marlin
|
||||
|
||||
torch::Tensor gptq_marlin_gemm(
|
||||
torch::Tensor marlin_gemm(
|
||||
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
|
||||
torch::Tensor& b_q_weight,
|
||||
std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
|
||||
@@ -856,5 +856,5 @@ torch::Tensor gptq_marlin_gemm(
|
||||
#endif
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("gptq_marlin_gemm", &gptq_marlin_gemm);
|
||||
m.impl("marlin_gemm", &marlin_gemm);
|
||||
}
|
||||
@@ -303,9 +303,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
ops.impl("permute_cols", torch::kCUDA, &permute_cols);
|
||||
|
||||
// gptq_marlin Optimized Quantized GEMM for GPTQ.
|
||||
// Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4).
|
||||
ops.def(
|
||||
"gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
|
||||
"marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
|
||||
"Tensor? b_bias_or_none,Tensor b_scales, "
|
||||
"Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
|
||||
"Tensor? "
|
||||
@@ -546,7 +546,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// Compute NVFP4 block quantized tensor.
|
||||
ops.def(
|
||||
"scaled_fp4_quant(Tensor! output, Tensor input,"
|
||||
" Tensor! output_scale, Tensor input_scale) -> ()");
|
||||
" Tensor! output_scale, Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout) -> ()");
|
||||
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
|
||||
|
||||
// Compute NVFP4 experts quantization.
|
||||
|
||||
+50
-11
@@ -15,9 +15,11 @@
|
||||
# Build arguments:
|
||||
# PYTHON_VERSION=3.13|3.12 (default)|3.11|3.10
|
||||
# VLLM_CPU_DISABLE_AVX512=false (default)|true
|
||||
# VLLM_CPU_AVX512BF16=false (default)|true
|
||||
# VLLM_CPU_AVX512VNNI=false (default)|true
|
||||
# VLLM_CPU_AMXBF16=false |true (default)
|
||||
# VLLM_CPU_AVX2=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512BF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512VNNI=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AMXBF16=false (default)|true (for cross-compilation)
|
||||
#
|
||||
|
||||
######################### COMMON BASE IMAGE #########################
|
||||
@@ -54,9 +56,12 @@ ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
# Copy requirements files for installation
|
||||
COPY requirements/common.txt requirements/common.txt
|
||||
COPY requirements/cpu.txt requirements/cpu.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/common.txt,target=requirements/common.txt \
|
||||
--mount=type=bind,src=requirements/cpu.txt,target=requirements/cpu.txt \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/cpu.txt
|
||||
|
||||
@@ -88,6 +93,12 @@ ARG GIT_REPO_CHECK=0
|
||||
# Support for building with non-AVX512 vLLM: docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" ...
|
||||
ARG VLLM_CPU_DISABLE_AVX512=0
|
||||
ENV VLLM_CPU_DISABLE_AVX512=${VLLM_CPU_DISABLE_AVX512}
|
||||
# Support for cross-compilation with AVX2 ISA: docker build --build-arg VLLM_CPU_AVX2="1" ...
|
||||
ARG VLLM_CPU_AVX2=0
|
||||
ENV VLLM_CPU_AVX2=${VLLM_CPU_AVX2}
|
||||
# Support for cross-compilation with AVX512 ISA: docker build --build-arg VLLM_CPU_AVX512="1" ...
|
||||
ARG VLLM_CPU_AVX512=0
|
||||
ENV VLLM_CPU_AVX512=${VLLM_CPU_AVX512}
|
||||
# Support for building with AVX512BF16 ISA: docker build --build-arg VLLM_CPU_AVX512BF16="true" ...
|
||||
ARG VLLM_CPU_AVX512BF16=0
|
||||
ENV VLLM_CPU_AVX512BF16=${VLLM_CPU_AVX512BF16}
|
||||
@@ -100,18 +111,19 @@ ENV VLLM_CPU_AMXBF16=${VLLM_CPU_AMXBF16}
|
||||
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
# Copy build requirements
|
||||
COPY requirements/cpu-build.txt requirements/build.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/cpu-build.txt,target=requirements/build.txt \
|
||||
uv pip install -r requirements/build.txt
|
||||
|
||||
COPY . .
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/workspace/vllm/.deps,sharing=locked \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
|
||||
|
||||
######################### TEST DEPS #########################
|
||||
@@ -119,9 +131,11 @@ FROM base AS vllm-test-deps
|
||||
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
# Copy test requirements
|
||||
COPY requirements/test.in requirements/cpu-test.in
|
||||
|
||||
# TODO: Update to 2.9.0 when there is a new build for intel_extension_for_pytorch for that version
|
||||
RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
|
||||
cp requirements/test.in requirements/cpu-test.in && \
|
||||
RUN \
|
||||
sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
|
||||
remove_packages_not_supported_on_aarch64() { \
|
||||
case "$(uname -m)" in \
|
||||
@@ -200,4 +214,29 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=vllm-build,src=/workspace/vllm/dist,target=dist \
|
||||
uv pip install dist/*.whl
|
||||
|
||||
# Add labels to document build configuration
|
||||
LABEL org.opencontainers.image.title="vLLM CPU"
|
||||
LABEL org.opencontainers.image.description="vLLM inference engine for CPU platforms"
|
||||
LABEL org.opencontainers.image.vendor="vLLM Project"
|
||||
LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm"
|
||||
|
||||
# Build configuration labels
|
||||
ARG TARGETARCH
|
||||
ARG VLLM_CPU_DISABLE_AVX512
|
||||
ARG VLLM_CPU_AVX2
|
||||
ARG VLLM_CPU_AVX512
|
||||
ARG VLLM_CPU_AVX512BF16
|
||||
ARG VLLM_CPU_AVX512VNNI
|
||||
ARG VLLM_CPU_AMXBF16
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
LABEL ai.vllm.build.target-arch="${TARGETARCH}"
|
||||
LABEL ai.vllm.build.cpu-disable-avx512="${VLLM_CPU_DISABLE_AVX512:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx2="${VLLM_CPU_AVX2:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512="${VLLM_CPU_AVX512:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512bf16="${VLLM_CPU_AVX512BF16:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512vnni="${VLLM_CPU_AVX512VNNI:-false}"
|
||||
LABEL ai.vllm.build.cpu-amxbf16="${VLLM_CPU_AMXBF16:-false}"
|
||||
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
+14
-1
@@ -227,7 +227,7 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
|
||||
# This ensures setuptools_scm sees clean repo state for version detection
|
||||
RUN --mount=type=bind,source=.git,target=vllm/.git \
|
||||
cd vllm \
|
||||
&& pip install setuptools_scm \
|
||||
&& pip install setuptools_scm regex \
|
||||
&& VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
|
||||
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
|
||||
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
|
||||
@@ -342,6 +342,19 @@ RUN mkdir src && mv vllm src/vllm
|
||||
FROM base AS final
|
||||
|
||||
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Clean up sccache from release image (not needed at runtime)
|
||||
# This removes the binary and wrappers that may have been installed during build
|
||||
RUN rm -f /usr/bin/sccache || true \
|
||||
&& rm -rf /opt/sccache-wrappers || true
|
||||
|
||||
# Unset sccache environment variables for the release image
|
||||
# This prevents S3 bucket config from leaking into production images
|
||||
ENV SCCACHE_BUCKET=
|
||||
ENV SCCACHE_REGION=
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=
|
||||
ENV SCCACHE_IDLE_TIMEOUT=
|
||||
|
||||
# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
|
||||
# Manually remove it so that later steps of numpy upgrade can continue
|
||||
RUN case "$(which python3)" in \
|
||||
|
||||
@@ -13,14 +13,14 @@ For x86 CPU environment, please use the image with "-cpu" postfix. For AArch64 C
|
||||
Here is an example for docker run command for CPU. For GPUs skip setting the `ON_CPU` env var.
|
||||
|
||||
```bash
|
||||
export VLLM_COMMIT=1da94e673c257373280026f75ceb4effac80e892 # use full commit hash from the main branch
|
||||
export VLLM_COMMIT=7f42dc20bb2800d09faa72b26f25d54e26f1b694 # use full commit hash from the main branch
|
||||
export HF_TOKEN=<valid Hugging Face token>
|
||||
if [[ "$(uname -m)" == aarch64 || "$(uname -m)" == arm64 ]]; then
|
||||
IMG_SUFFIX="arm64-cpu"
|
||||
else
|
||||
IMG_SUFFIX="cpu"
|
||||
fi
|
||||
docker run -it --entrypoint /bin/bash -v /data/huggingface:/root/.cache/huggingface -e HF_TOKEN=$HF_TOKEN -e ON_ARM64_CPU=1 --shm-size=16g --name vllm-cpu-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${VLLM_COMMIT}-${IMG_SUFFIX}
|
||||
docker run -it --entrypoint /bin/bash -v /data/huggingface:/root/.cache/huggingface -e HF_TOKEN=$HF_TOKEN -e ON_CPU=1 --shm-size=16g --name vllm-cpu-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${VLLM_COMMIT}-${IMG_SUFFIX}
|
||||
```
|
||||
|
||||
Then, run below command inside the docker instance.
|
||||
|
||||
@@ -106,6 +106,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
|
||||
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`
|
||||
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
|
||||
Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
|
||||
|
||||
+139
-117
@@ -20,67 +20,6 @@ To input multi-modal data, follow this schema in [vllm.inputs.PromptType][]:
|
||||
- `prompt`: The prompt should follow the format that is documented on HuggingFace.
|
||||
- `multi_modal_data`: This is a dictionary that follows the schema defined in [vllm.multimodal.inputs.MultiModalDataDict][].
|
||||
|
||||
### Stable UUIDs for Caching (multi_modal_uuids)
|
||||
|
||||
When using multi-modal inputs, vLLM normally hashes each media item by content to enable caching across requests. You can optionally pass `multi_modal_uuids` to provide your own stable IDs for each item so caching can reuse work across requests without rehashing the raw content.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_a = Image.open("/path/to/a.jpg")
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [img_a, img_b]},
|
||||
# Provide stable IDs for caching.
|
||||
# Requirements (matched by this example):
|
||||
# - Include every modality present in multi_modal_data.
|
||||
# - For lists, provide the same number of entries.
|
||||
# - Use None to fall back to content hashing for that item.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
Using UUIDs, you can also skip sending media data entirely if you expect cache hits for respective items. Note that the request will fail if the skipped media doesn't have a corresponding UUID, or if the UUID fails to hit the cache.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [None, img_b]},
|
||||
# Since img_a is expected to be cached, we can skip sending the actual
|
||||
# image entirely.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
If both multimodal processor caching and prefix caching are disabled, user-provided `multi_modal_uuids` are ignored.
|
||||
|
||||
### Image Inputs
|
||||
|
||||
You can pass a single image to the `'image'` field of the multi-modal dictionary, as shown in the following examples:
|
||||
@@ -397,7 +336,8 @@ No manual conversion is needed - vLLM handles the channel normalization automati
|
||||
### Embedding Inputs
|
||||
|
||||
To input pre-computed embeddings belonging to a data type (i.e. image, video, or audio) directly to the language model,
|
||||
pass a tensor of shape `(num_items, feature_size, hidden_size of LM)` to the corresponding field of the multi-modal dictionary.
|
||||
pass a tensor of shape `(..., hidden_size of LM)` to the corresponding field of the multi-modal dictionary.
|
||||
The exact shape depends on the model being used.
|
||||
|
||||
You must enable this feature via `enable_mm_embeds=True`.
|
||||
|
||||
@@ -418,8 +358,7 @@ You must enable this feature via `enable_mm_embeds=True`.
|
||||
# Refer to the HuggingFace repo for the correct format to use
|
||||
prompt = "USER: <image>\nWhat is the content of this image?\nASSISTANT:"
|
||||
|
||||
# Embeddings for single image
|
||||
# torch.Tensor of shape (1, image_feature_size, hidden_size of LM)
|
||||
# For most models, `image_embeds` has shape: (num_images, image_feature_size, hidden_size)
|
||||
image_embeds = torch.load(...)
|
||||
|
||||
outputs = llm.generate({
|
||||
@@ -430,21 +369,8 @@ You must enable this feature via `enable_mm_embeds=True`.
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
For Qwen2-VL and MiniCPM-V, we accept additional parameters alongside the embeddings:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
# Construct the prompt based on your model
|
||||
prompt = ...
|
||||
|
||||
# Embeddings for multiple images
|
||||
# torch.Tensor of shape (num_images, image_feature_size, hidden_size of LM)
|
||||
image_embeds = torch.load(...)
|
||||
|
||||
# Qwen2-VL
|
||||
# Additional examples for models that require extra fields
|
||||
llm = LLM(
|
||||
"Qwen/Qwen2-VL-2B-Instruct",
|
||||
limit_mm_per_prompt={"image": 4},
|
||||
@@ -452,13 +378,15 @@ For Qwen2-VL and MiniCPM-V, we accept additional parameters alongside the embedd
|
||||
)
|
||||
mm_data = {
|
||||
"image": {
|
||||
"image_embeds": image_embeds,
|
||||
# Shape: (total_feature_size, hidden_size)
|
||||
# total_feature_size = sum(image_feature_size for image in images)
|
||||
"image_embeds": torch.load(...),
|
||||
# Shape: (num_images, 3)
|
||||
# image_grid_thw is needed to calculate positional encoding.
|
||||
"image_grid_thw": torch.load(...), # torch.Tensor of shape (1, 3),
|
||||
"image_grid_thw": torch.load(...),
|
||||
}
|
||||
}
|
||||
|
||||
# MiniCPM-V
|
||||
llm = LLM(
|
||||
"openbmb/MiniCPM-V-2_6",
|
||||
trust_remote_code=True,
|
||||
@@ -467,20 +395,14 @@ For Qwen2-VL and MiniCPM-V, we accept additional parameters alongside the embedd
|
||||
)
|
||||
mm_data = {
|
||||
"image": {
|
||||
"image_embeds": image_embeds,
|
||||
# Shape: (num_images, num_slices, hidden_size)
|
||||
# num_slices can differ for each image
|
||||
"image_embeds": [torch.load(...) for image in images],
|
||||
# Shape: (num_images, 2)
|
||||
# image_sizes is needed to calculate details of the sliced image.
|
||||
"image_sizes": [image.size for image in images], # list of image sizes
|
||||
"image_sizes": [image.size for image in images],
|
||||
}
|
||||
}
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": mm_data,
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
For Qwen3-VL, the `image_embeds` should contain both the base image embedding and deepstack features.
|
||||
@@ -501,8 +423,8 @@ You can pass pre-computed audio embeddings similar to image embeddings:
|
||||
# Refer to the HuggingFace repo for the correct format to use
|
||||
prompt = "USER: <audio>\nWhat is in this audio?\nASSISTANT:"
|
||||
|
||||
# Load pre-computed audio embeddings
|
||||
# torch.Tensor of shape (1, audio_feature_size, hidden_size of LM)
|
||||
# Load pre-computed audio embeddings, usually with shape:
|
||||
# (num_audios, audio_feature_size, hidden_size of LM)
|
||||
audio_embeds = torch.load(...)
|
||||
|
||||
outputs = llm.generate({
|
||||
@@ -515,6 +437,67 @@ You can pass pre-computed audio embeddings similar to image embeddings:
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
### Cached Inputs
|
||||
|
||||
When using multi-modal inputs, vLLM normally hashes each media item by content to enable caching across requests. You can optionally pass `multi_modal_uuids` to provide your own stable IDs for each item so caching can reuse work across requests without rehashing the raw content.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_a = Image.open("/path/to/a.jpg")
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [img_a, img_b]},
|
||||
# Provide stable IDs for caching.
|
||||
# Requirements (matched by this example):
|
||||
# - Include every modality present in multi_modal_data.
|
||||
# - For lists, provide the same number of entries.
|
||||
# - Use None to fall back to content hashing for that item.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
Using UUIDs, you can also skip sending media data entirely if you expect cache hits for respective items. Note that the request will fail if the skipped media doesn't have a corresponding UUID, or if the UUID fails to hit the cache.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [None, img_b]},
|
||||
# Since img_a is expected to be cached, we can skip sending the actual
|
||||
# image entirely.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
If both multimodal processor caching and prefix caching are disabled, user-provided `multi_modal_uuids` are ignored.
|
||||
|
||||
## Online Serving
|
||||
|
||||
Our OpenAI-compatible server accepts multi-modal data via the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat). Media inputs also support optional UUIDs users can provide to uniquely identify each media, which is used to cache the media results across requests.
|
||||
@@ -879,7 +862,11 @@ Full example: [examples/online_serving/openai_chat_completion_client_for_multimo
|
||||
### Embedding Inputs
|
||||
|
||||
To input pre-computed embeddings belonging to a data type (i.e. image, video, or audio) directly to the language model,
|
||||
pass a tensor of shape `(num_items, feature_size, hidden_size of LM)` to the corresponding field of the multi-modal dictionary.
|
||||
pass a tensor of shape `(..., hidden_size of LM)` for each item to the corresponding field of the multi-modal dictionary.
|
||||
|
||||
!!! important
|
||||
Unlike offline inference, the embeddings for each item must be passed separately
|
||||
in order for placeholder tokens to be applied correctly by the chat template.
|
||||
|
||||
You must enable this feature via the `--enable-mm-embeds` flag in `vllm serve`.
|
||||
|
||||
@@ -897,11 +884,6 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
```python
|
||||
from vllm.utils.serial_utils import tensor2base64
|
||||
|
||||
image_embedding = torch.load(...)
|
||||
grid_thw = torch.load(...) # Required by Qwen/Qwen2-VL-2B-Instruct
|
||||
|
||||
base64_image_embedding = tensor2base64(image_embedding)
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
@@ -912,29 +894,33 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
model = "llava-hf/llava-1.5-7b-hf"
|
||||
embeds = {
|
||||
"type": "image_embeds",
|
||||
"image_embeds": f"{base64_image_embedding}",
|
||||
"image_embeds": tensor2base64(torch.load(...)), # Shape: (image_feature_size, hidden_size)
|
||||
"uuid": image_url, # Optional
|
||||
}
|
||||
|
||||
# Pass additional parameters (available to Qwen2-VL and MiniCPM-V)
|
||||
|
||||
# Additional examples for models that require extra fields
|
||||
model = "Qwen/Qwen2-VL-2B-Instruct"
|
||||
embeds = {
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": f"{base64_image_embedding}", # Required
|
||||
"image_grid_thw": f"{base64_image_grid_thw}", # Required by Qwen/Qwen2-VL-2B-Instruct
|
||||
"image_embeds": tensor2base64(torch.load(...)), # Shape: (image_feature_size, hidden_size)
|
||||
"image_grid_thw": tensor2base64(torch.load(...)), # Shape: (3,)
|
||||
},
|
||||
"uuid": image_url, # Optional
|
||||
}
|
||||
|
||||
model = "openbmb/MiniCPM-V-2_6"
|
||||
embeds = {
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": f"{base64_image_embedding}", # Required
|
||||
"image_sizes": f"{base64_image_sizes}", # Required by openbmb/MiniCPM-V-2_6
|
||||
"image_embeds": tensor2base64(torch.load(...)), # Shape: (num_slices, hidden_size)
|
||||
"image_sizes": tensor2base64(torch.load(...)), # Shape: (2,)
|
||||
},
|
||||
"uuid": image_url, # Optional
|
||||
}
|
||||
|
||||
# Single image input
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
@@ -954,9 +940,55 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Multi image input
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this image?",
|
||||
},
|
||||
embeds,
|
||||
embeds,
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Multi image input (interleaved)
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
embeds,
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this image?",
|
||||
},
|
||||
embeds,
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
```
|
||||
|
||||
For Online Serving, you can also skip sending media if you expect cache hits with provided UUIDs. You can do so by sending media like this:
|
||||
### Cached Inputs
|
||||
|
||||
Just like with offline inference, you can skip sending media if you expect cache hits with provided UUIDs. You can do so by sending media like this:
|
||||
|
||||
??? code
|
||||
|
||||
@@ -990,13 +1022,3 @@ For Online Serving, you can also skip sending media if you expect cache hits wit
|
||||
},
|
||||
|
||||
```
|
||||
|
||||
!!! note
|
||||
Multiple messages can now contain `{"type": "image_embeds"}`, enabling you to pass multiple image embeddings in a single request (similar to regular images). The number of embeddings is limited by `--limit-mm-per-prompt`.
|
||||
|
||||
**Important**: The embedding shape format differs based on the number of embeddings:
|
||||
|
||||
- **Single embedding**: 3D tensor of shape `(1, feature_size, hidden_size)`
|
||||
- **Multiple embeddings**: List of 2D tensors, each of shape `(feature_size, hidden_size)`
|
||||
|
||||
If used with a model that requires additional parameters, you must also provide a tensor for each of them, e.g. `image_grid_thw`, `image_sizes`, etc.
|
||||
|
||||
@@ -184,15 +184,6 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
|
||||
--kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
|
||||
```
|
||||
|
||||
### Cross layers blocks
|
||||
|
||||
By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
|
||||
To enable this feature:
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
|
||||
```
|
||||
|
||||
## Example Scripts/Code
|
||||
|
||||
Refer to these example scripts in the vLLM repository:
|
||||
|
||||
@@ -131,7 +131,7 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
|
||||
|
||||
=== "Apple silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/cpu.arm.inc.md:build-image-from-source"
|
||||
--8<-- "docs/getting_started/installation/cpu.apple.inc.md:build-image-from-source"
|
||||
|
||||
=== "IBM Z (S390X)"
|
||||
--8<-- "docs/getting_started/installation/cpu.s390x.inc.md:build-image-from-source"
|
||||
|
||||
@@ -164,21 +164,76 @@ uv pip install dist/*.whl
|
||||
[https://gallery.ecr.aws/q9t5s3a7/vllm-cpu-release-repo](https://gallery.ecr.aws/q9t5s3a7/vllm-cpu-release-repo)
|
||||
|
||||
!!! warning
|
||||
If deploying the pre-built images on machines without `avx512f`, `avx512_bf16`, or `avx512_vnni` support, an `Illegal instruction` error may be raised. It is recommended to build images for these machines with the appropriate build arguments (e.g., `--build-arg VLLM_CPU_DISABLE_AVX512=true`, `--build-arg VLLM_CPU_AVX512BF16=false`, or `--build-arg VLLM_CPU_AVX512VNNI=false`) to disable unsupported features. Please note that without `avx512f`, AVX2 will be used and this version is not recommended because it only has basic feature support.
|
||||
If deploying the pre-built images on machines without `avx512f`, `avx512_bf16`, or `avx512_vnni` support, an `Illegal instruction` error may be raised. See the build-image-from-source section below for build arguments to match your target CPU capabilities.
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
# --8<-- [start:build-image-from-source]
|
||||
|
||||
## Building for your target CPU
|
||||
|
||||
vLLM supports building Docker images for x86 CPU platforms with automatic instruction set detection.
|
||||
|
||||
### Basic build command
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_AVX512BF16=false (default)|true \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=false (default)|true \
|
||||
--build-arg VLLM_CPU_AMXBF16=false|true (default) \
|
||||
--build-arg VLLM_CPU_DISABLE_AVX512=false (default)|true \
|
||||
--build-arg VLLM_CPU_DISABLE_AVX512=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX2=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX512=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX512BF16=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AMXBF16=<false|true (default)> \
|
||||
--tag vllm-cpu-env \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
# Launching OpenAI server
|
||||
!!! note "Instruction set auto-detection"
|
||||
By default, vLLM will auto-detect CPU instruction sets (AVX512, AVX2, etc.) from the build system's CPU flags. Build arguments like `VLLM_CPU_AVX2`, `VLLM_CPU_AVX512`, `VLLM_CPU_AVX512BF16`, `VLLM_CPU_AVX512VNNI`, and `VLLM_CPU_AMXBF16` are primarily used for **cross-compilation** or for building container images on systems that don't have the target platforms ISA:
|
||||
|
||||
- Set `VLLM_CPU_{ISA}=true` to force-enable an instruction set (for cross-compilation to target platforms with that ISA)
|
||||
- Set `VLLM_CPU_{ISA}=false` to rely on auto-detection
|
||||
- When an ISA build arg is set to `true`, vLLM will build with that instruction set regardless of the build system's CPU capabilities
|
||||
|
||||
### Build examples
|
||||
|
||||
**Example 1: Auto-detection (native build)**
|
||||
|
||||
Build on a machine with the same CPU as your target deployment:
|
||||
|
||||
```bash
|
||||
# Auto-detects all CPU features from the build system
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--tag vllm-cpu-env \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Example 2: Cross-compilation for AVX512 deployment**
|
||||
|
||||
Build an AVX512 image on any x86_64 system (even without AVX512):
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_AVX512=true \
|
||||
--build-arg VLLM_CPU_AVX512BF16=true \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=true \
|
||||
--tag vllm-cpu-avx512 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Example 3: Cross-compilation for AVX2 deployment**
|
||||
|
||||
Build an AVX2 image for older CPUs:
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_AVX2=true \
|
||||
--tag vllm-cpu-avx2 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
## Launching the OpenAI server
|
||||
|
||||
```bash
|
||||
docker run --rm \
|
||||
--security-opt seccomp=unconfined \
|
||||
--cap-add SYS_NICE \
|
||||
|
||||
@@ -118,7 +118,7 @@ There are more environment variables to control the behavior of Python-only buil
|
||||
|
||||
* `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file path of a pre-compiled wheel to use. All other logic to find the wheel will be skipped.
|
||||
* `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit hash to download the pre-compiled wheel. It can be `nightly` to use the last **already built** commit on the main branch.
|
||||
* `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cpu`. If not specified, the CUDA variant with `VLLM_MAIN_CUDA_VERSION` will be tried, then fallback to the default variant on the remote index.
|
||||
* `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cu130`, `cpu`. If not specified, the variant is auto-detected based on your system's CUDA version (from PyTorch or nvidia-smi). You can also set `VLLM_MAIN_CUDA_VERSION` to override auto-detection.
|
||||
|
||||
You can find more information about vLLM's wheels in [Install the latest code](#install-the-latest-code).
|
||||
|
||||
|
||||
@@ -31,7 +31,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/
|
||||
To install a specific version and ROCm variant of vLLM wheel.
|
||||
|
||||
```bash
|
||||
uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.1/rocm700
|
||||
```
|
||||
|
||||
!!! warning "Caveats for using `pip`"
|
||||
@@ -41,7 +41,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
If you insist on using `pip`, you have to specify the exact vLLM version and full URL of the wheel path `https://wheels.vllm.ai/rocm/<version>/<rocm-variant>` (which can be obtained from the web page).
|
||||
|
||||
```bash
|
||||
pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
pip install vllm==0.14.1+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.1/rocm700
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-wheels]
|
||||
|
||||
+47
-26
@@ -25,34 +25,55 @@ Maintainers form a hierarchy based on sustained, high-quality contributions and
|
||||
|
||||
### Core Maintainers
|
||||
|
||||
Core Maintainers function like a project planning and decision making committee. In other convention, they might be called a Technical Steering Committee (TSC). In vLLM vocabulary, they are often known as "Project Leads". They meet weekly to coordinate roadmap priorities and allocate engineering resources. Current active leads: @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, @robertgshaw2-redhat, @tlrmchlsmth, @mgoin, @njhill, @ywang96, @houseroad, @yeqcharlotte, @ApostaC
|
||||
Core Maintainers function like a project planning and decision making committee. In other convention, they might be called a Technical Steering Committee (TSC). In vLLM vocabulary, they are often known as "Project Leads". They meet weekly to coordinate roadmap priorities and allocate engineering resources.
|
||||
|
||||
The responsibilities of the core maintainers are:
|
||||
**Project Leads:**
|
||||
|
||||
* Author quarterly roadmap and responsible for each development effort.
|
||||
* Making major changes to the technical direction or scope of vLLM and vLLM projects.
|
||||
* Defining the project's release strategy.
|
||||
* Work with model providers, hardware vendors, and key users of vLLM to ensure the project is on the right track.
|
||||
- Woosuk Kwon ([@WoosukKwon](https://github.com/WoosukKwon))
|
||||
- Zhuohan Li ([@zhuohan123](https://github.com/zhuohan123))
|
||||
- Simon Mo ([@simon-mo](https://github.com/simon-mo))
|
||||
- Kaichao You ([@youkaichao](https://github.com/youkaichao))
|
||||
- Robert Shaw ([@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat))
|
||||
- Tyler Michael Smith ([@tlrmchlsmth](https://github.com/tlrmchlsmth))
|
||||
- Michael Goin ([@mgoin](https://github.com/mgoin))
|
||||
- Nick Hill ([@njhill](https://github.com/njhill))
|
||||
- Roger Wang ([@ywang96](https://github.com/ywang96))
|
||||
- Lu Fang ([@houseroad](https://github.com/houseroad))
|
||||
- Ye (Charlotte) Qi ([@yeqcharlotte](https://github.com/yeqcharlotte))
|
||||
- Yihua Cheng ([@ApostaC](https://github.com/ApostaC))
|
||||
|
||||
**Responsibilities:**
|
||||
|
||||
- Author quarterly roadmap and responsible for each development effort.
|
||||
- Making major changes to the technical direction or scope of vLLM and vLLM projects.
|
||||
- Defining the project's release strategy.
|
||||
- Work with model providers, hardware vendors, and key users of vLLM to ensure the project is on the right track.
|
||||
|
||||
### Lead Maintainers
|
||||
|
||||
While Core maintainers assume the day-to-day responsibilities of the project, Lead maintainers are responsible for the overall direction and strategy of the project. A committee of @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, and @robertgshaw2-redhat currently shares this role with divided responsibilities.
|
||||
While Core maintainers assume the day-to-day responsibilities of the project, Lead maintainers are responsible for the overall direction and strategy of the project. The following committee currently shares this role with divided responsibilities:
|
||||
|
||||
The responsibilities of the lead maintainers are:
|
||||
- Woosuk Kwon ([@WoosukKwon](https://github.com/WoosukKwon))
|
||||
- Zhuohan Li ([@zhuohan123](https://github.com/zhuohan123))
|
||||
- Simon Mo ([@simon-mo](https://github.com/simon-mo))
|
||||
- Kaichao You ([@youkaichao](https://github.com/youkaichao))
|
||||
- Robert Shaw ([@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat))
|
||||
|
||||
* Making decisions where consensus among core maintainers cannot be reached.
|
||||
* Adopting changes to the project's technical governance.
|
||||
* Organizing the voting process for new committers.
|
||||
**Responsibilities:**
|
||||
|
||||
- Making decisions where consensus among core maintainers cannot be reached.
|
||||
- Adopting changes to the project's technical governance.
|
||||
- Organizing the voting process for new committers.
|
||||
|
||||
### Committers and Area Owners
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas and help the community.
|
||||
|
||||
The responsibilities of the committers are:
|
||||
**Responsibilities:**
|
||||
|
||||
* Reviewing PRs and providing feedback.
|
||||
* Addressing issues and questions from the community.
|
||||
* Own specific areas of the codebase and development efforts: reviewing PRs, addressing issues, answering questions, improving documentation.
|
||||
- Reviewing PRs and providing feedback.
|
||||
- Addressing issues and questions from the community.
|
||||
- Own specific areas of the codebase and development efforts: reviewing PRs, addressing issues, answering questions, improving documentation.
|
||||
|
||||
Specially, committers are almost all area owners. They author subsystems, review PRs, refactor code, monitor tests, and ensure compatibility with other areas. All area owners are committers with deep expertise in that area, but not all committers own areas.
|
||||
|
||||
@@ -68,23 +89,23 @@ Any committer can nominate candidates via our private mailing list:
|
||||
|
||||
Committership is highly selective and merit based. The selection criteria requires:
|
||||
|
||||
* **Area expertise**: leading design/implementation of core subsystems, material performance or reliability improvements adopted project‑wide, or accepted RFCs that shape technical direction.
|
||||
* **Sustained contributions**: high‑quality merged contributions and reviews across releases, responsiveness to feedback, and stewardship of code health.
|
||||
* **Community leadership**: mentoring contributors, triaging issues, improving docs, and elevating project standards.
|
||||
- **Area expertise**: leading design/implementation of core subsystems, material performance or reliability improvements adopted project‑wide, or accepted RFCs that shape technical direction.
|
||||
- **Sustained contributions**: high‑quality merged contributions and reviews across releases, responsiveness to feedback, and stewardship of code health.
|
||||
- **Community leadership**: mentoring contributors, triaging issues, improving docs, and elevating project standards.
|
||||
|
||||
To further illustrate, a committer typically satisfies at least two of the following accomplishment patterns:
|
||||
|
||||
* Author of an accepted RFC or design that materially shaped project direction
|
||||
* Measurable, widely adopted performance or reliability improvement in core paths
|
||||
* Long‑term ownership of a subsystem with demonstrable quality and stability gains
|
||||
* Significant cross‑project compatibility or ecosystem enablement work (models, hardware, tooling)
|
||||
- Author of an accepted RFC or design that materially shaped project direction
|
||||
- Measurable, widely adopted performance or reliability improvement in core paths
|
||||
- Long‑term ownership of a subsystem with demonstrable quality and stability gains
|
||||
- Significant cross‑project compatibility or ecosystem enablement work (models, hardware, tooling)
|
||||
|
||||
While there isn't a quantitative bar, past committers have:
|
||||
|
||||
* Submitted approximately 30+ PRs of substantial quality and scope
|
||||
* Provided high-quality reviews of approximately 10+ substantial external contributor PRs
|
||||
* Addressed multiple issues and questions from the community in issues/forums/Slack
|
||||
* Led concentrated efforts on RFCs and their implementation, or significant performance or reliability improvements adopted project‑wide
|
||||
- Submitted approximately 30+ PRs of substantial quality and scope
|
||||
- Provided high-quality reviews of approximately 10+ substantial external contributor PRs
|
||||
- Addressed multiple issues and questions from the community in issues/forums/Slack
|
||||
- Led concentrated efforts on RFCs and their implementation, or significant performance or reliability improvements adopted project‑wide
|
||||
|
||||
### Working Groups
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
| [Intel® Xeon® 6 Processors](https://www.intel.com/content/www/us/en/products/details/processors/xeon.html) |
|
||||
| [Intel® Xeon® 5 Processors](https://www.intel.com/content/www/us/en/products/docs/processors/xeon/5th-gen-xeon-scalable-processors.html) |
|
||||
|
||||
## Supported Models
|
||||
## Recommended Models
|
||||
|
||||
### Text-only Language Models
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
| ----------------------------------------- |
|
||||
| [Intel® Arc™ Pro B-Series Graphics](https://www.intel.com/content/www/us/en/products/docs/discrete-gpus/arc/workstations/b-series/overview.html) |
|
||||
|
||||
## Supported Models
|
||||
## Recommended Models
|
||||
|
||||
### Text-only Language Models
|
||||
|
||||
|
||||
@@ -422,7 +422,7 @@ th {
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | ︎| ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 |`MiniMaxAI/MiniMax-M2`, etc. | | ✅︎ |
|
||||
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 |`MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
|
||||
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
|
||||
| `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -456,6 +456,7 @@ th {
|
||||
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
|
||||
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
|
||||
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
|
||||
| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/step-3.5-flash`, etc. | | ✅︎ |
|
||||
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
|
||||
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
|
||||
| `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -686,6 +687,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
|
||||
| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
|
||||
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
|
||||
| `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
|
||||
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
|
||||
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -39,6 +39,7 @@ Launch Claude Code with environment variables pointing to your vLLM server:
|
||||
```bash
|
||||
ANTHROPIC_BASE_URL=http://localhost:8000 \
|
||||
ANTHROPIC_API_KEY=dummy \
|
||||
ANTHROPIC_AUTH_TOKEN=dummy \
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL=my-model \
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL=my-model \
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL=my-model \
|
||||
@@ -51,6 +52,7 @@ The environment variables:
|
||||
| -------------------------------- | --------------------------------------------------------------------- |
|
||||
| `ANTHROPIC_BASE_URL` | Points to your vLLM server (default port is 8000) |
|
||||
| `ANTHROPIC_API_KEY` | Can be any value since vLLM doesn't require authentication by default |
|
||||
| `ANTHROPIC_AUTH_TOKEN` | Is required. Can be any value. |
|
||||
| `ANTHROPIC_DEFAULT_OPUS_MODEL` | Model name for Opus-tier requests |
|
||||
| `ANTHROPIC_DEFAULT_SONNET_MODEL` | Model name for Sonnet-tier requests |
|
||||
| `ANTHROPIC_DEFAULT_HAIKU_MODEL` | Model name for Haiku-tier requests |
|
||||
|
||||
@@ -112,10 +112,36 @@ def get_multi_audios_query() -> QueryResult:
|
||||
)
|
||||
|
||||
|
||||
def get_multi_images_query() -> QueryResult:
|
||||
question = "What are the differences between these two images?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_bos|><|IMAGE|><|vision_eos|>"
|
||||
"<|vision_bos|><|IMAGE|><|vision_eos|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n"
|
||||
)
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"image": [
|
||||
convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
|
||||
convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB"),
|
||||
],
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={
|
||||
"image": 2,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
query_map = {
|
||||
"mixed_modalities": get_mixed_modalities_query,
|
||||
"use_audio_in_video": get_use_audio_in_video_query,
|
||||
"multi_audios": get_multi_audios_query,
|
||||
"multi_images": get_multi_images_query,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -59,8 +59,10 @@ class PrithviMAE:
|
||||
input_data = input_data[0]
|
||||
|
||||
mm_data = {
|
||||
"pixel_values": input_data,
|
||||
"location_coords": location_coords,
|
||||
"image": {
|
||||
"pixel_values": input_data,
|
||||
"location_coords": location_coords,
|
||||
}
|
||||
}
|
||||
|
||||
prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}
|
||||
|
||||
@@ -18,48 +18,32 @@ e.g.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
with requests.get(content_url, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return {"url": f"data:image/jpeg;base64,{result}"}
|
||||
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
documents = {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
),
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": encode_base64_content_from_url(
|
||||
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
),
|
||||
},
|
||||
]
|
||||
}
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -74,23 +58,36 @@ def main(args):
|
||||
models_url = base_url + "/v1/models"
|
||||
rerank_url = base_url + "/rerank"
|
||||
|
||||
response = requests.get(models_url, headers=headers)
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
data = {
|
||||
print("Query: string & Document: list of string")
|
||||
prompt = {"model": model, "query": query, "documents": [document]}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text")
|
||||
prompt = {"model": model, "query": query, "documents": {"content": [documents[0]]}}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"documents": {"content": [documents[1]]},
|
||||
}
|
||||
response = requests.post(rerank_url, headers=headers, json=data)
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
# Check the response
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
print("Query: string & Document: image base64")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -17,48 +17,32 @@ e.g.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
with requests.get(content_url, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return {"url": f"data:image/jpeg;base64,{result}"}
|
||||
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
queries = "slm markdown"
|
||||
documents = {
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": encode_base64_content_from_url(
|
||||
"https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
|
||||
),
|
||||
},
|
||||
]
|
||||
}
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -73,15 +57,40 @@ def main(args):
|
||||
models_url = base_url + "/v1/models"
|
||||
score_url = base_url + "/score"
|
||||
|
||||
response = requests.get(models_url, headers=headers)
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
prompt = {"model": model, "queries": queries, "documents": documents}
|
||||
response = requests.post(score_url, headers=headers, json=prompt)
|
||||
print("\nPrompt when queries is string and documents is a image list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
print("Query: string & Document: string")
|
||||
prompt = {"model": model, "queries": query, "documents": document}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image base64")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+3
-2
@@ -9,7 +9,7 @@ requires = [
|
||||
"torch == 2.9.1",
|
||||
"wheel",
|
||||
"jinja2",
|
||||
"grpcio-tools>=1.76.0",
|
||||
"grpcio-tools",
|
||||
]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
@@ -132,7 +132,8 @@ python = "./.venv"
|
||||
# these files may be written in non english words
|
||||
extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*",
|
||||
"benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*"]
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
|
||||
"docs/governance/process.md"]
|
||||
ignore-hidden = true
|
||||
ignore-files = true
|
||||
ignore-dot = true
|
||||
|
||||
@@ -9,5 +9,5 @@ wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
build
|
||||
protobuf>=6.33.2
|
||||
grpcio-tools>=1.76.0
|
||||
protobuf >= 6.33.5
|
||||
grpcio-tools
|
||||
|
||||
@@ -9,9 +9,9 @@ blake3
|
||||
py-cpuinfo
|
||||
transformers >= 4.56.0, < 5
|
||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||
protobuf >= 6.30.0 # Required by LlamaTokenizer, gRPC.
|
||||
protobuf >= 6.33.5 # 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.
|
||||
aiohttp
|
||||
aiohttp >= 3.13.3
|
||||
openai >= 1.99.1 # For Responses API with reasoning content
|
||||
pydantic >= 2.12.0
|
||||
prometheus_client >= 0.18.0
|
||||
@@ -51,5 +51,5 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
|
||||
anthropic >= 0.71.0
|
||||
model-hosting-container-standards >= 0.1.13, < 1.0.0
|
||||
mcp
|
||||
grpcio>=1.76.0
|
||||
grpcio-reflection>=1.76.0
|
||||
grpcio
|
||||
grpcio-reflection
|
||||
@@ -1,2 +1,2 @@
|
||||
lmcache
|
||||
lmcache >= 0.3.9
|
||||
nixl >= 0.7.1 # Required for disaggregated prefill
|
||||
|
||||
@@ -14,7 +14,7 @@ pytest-shard==0.1.2
|
||||
# Async/HTTP dependencies
|
||||
anyio==4.6.2.post1
|
||||
# via httpx, starlette
|
||||
aiohttp==3.13.0
|
||||
aiohttp==3.13.3
|
||||
# via gpt-oss
|
||||
httpx==0.27.2
|
||||
# HTTP testing
|
||||
|
||||
@@ -12,7 +12,7 @@ affine==2.4.0
|
||||
# via rasterio
|
||||
aiohappyeyeballs==2.6.1
|
||||
# via aiohttp
|
||||
aiohttp==3.13.0
|
||||
aiohttp==3.13.3
|
||||
# via
|
||||
# aiohttp-cors
|
||||
# datasets
|
||||
|
||||
@@ -438,6 +438,49 @@ class precompiled_wheel_utils:
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def detect_system_cuda_variant() -> str:
|
||||
"""Auto-detect CUDA variant from torch, nvidia-smi, or env default."""
|
||||
|
||||
# Map CUDA major version to hosted wheel variants on wheels.vllm.ai
|
||||
supported = {12: "cu129", 13: "cu130"}
|
||||
|
||||
# Respect explicitly set VLLM_MAIN_CUDA_VERSION
|
||||
if envs.is_set("VLLM_MAIN_CUDA_VERSION"):
|
||||
v = envs.VLLM_MAIN_CUDA_VERSION
|
||||
print(f"Using VLLM_MAIN_CUDA_VERSION={v}")
|
||||
return "cu" + v.replace(".", "")[:3]
|
||||
|
||||
# Try torch.version.cuda
|
||||
cuda_version = None
|
||||
try:
|
||||
import torch
|
||||
|
||||
cuda_version = torch.version.cuda
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Try nvidia-smi
|
||||
if not cuda_version:
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["nvidia-smi"], capture_output=True, text=True, timeout=10
|
||||
)
|
||||
if m := re.search(r"CUDA Version:\s*(\d+\.\d+)", out.stdout):
|
||||
cuda_version = m.group(1)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fall back to default
|
||||
if not cuda_version:
|
||||
cuda_version = envs.VLLM_MAIN_CUDA_VERSION
|
||||
|
||||
# Map to supported variant
|
||||
major = int(cuda_version.split(".")[0])
|
||||
variant = supported.get(major, supported[max(supported)])
|
||||
print(f"Detected CUDA {cuda_version}, using variant {variant}")
|
||||
return variant
|
||||
|
||||
@staticmethod
|
||||
def find_local_rocm_wheel() -> str | None:
|
||||
"""Search for a local vllm wheel in common locations."""
|
||||
@@ -513,8 +556,8 @@ class precompiled_wheel_utils:
|
||||
1. user-specified wheel location (can be either local or remote, via
|
||||
VLLM_PRECOMPILED_WHEEL_LOCATION)
|
||||
2. user-specified variant (VLLM_PRECOMPILED_WHEEL_VARIANT) from nightly repo
|
||||
3. the variant corresponding to VLLM_MAIN_CUDA_VERSION from nightly repo
|
||||
4. the default variant from nightly repo
|
||||
or auto-detected CUDA variant based on system (torch, nvidia-smi)
|
||||
3. the default variant from nightly repo
|
||||
|
||||
If downloading from the nightly repo, the commit can be specified via
|
||||
VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch
|
||||
@@ -533,9 +576,11 @@ class precompiled_wheel_utils:
|
||||
import platform
|
||||
|
||||
arch = platform.machine()
|
||||
# try to fetch the wheel metadata from the nightly wheel repo
|
||||
main_variant = "cu" + envs.VLLM_MAIN_CUDA_VERSION.replace(".", "")
|
||||
variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", main_variant)
|
||||
# try to fetch the wheel metadata from the nightly wheel repo,
|
||||
# detecting CUDA variant from system if not specified
|
||||
variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
|
||||
if variant is None:
|
||||
variant = precompiled_wheel_utils.detect_system_cuda_variant()
|
||||
commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
|
||||
if not commit or len(commit) != 40:
|
||||
print(
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from torch._dynamo.utils import counters
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
|
||||
|
||||
|
||||
def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache):
|
||||
# Run in same process so we can access PyTorch's internal counters
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
# I'm not sure if this is going to affect the numbers
|
||||
monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
|
||||
|
||||
# Force cold compilation
|
||||
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
|
||||
|
||||
compilation_config = CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
cudagraph_mode=CUDAGraphMode.NONE, # make the model loading faster
|
||||
)
|
||||
|
||||
counters.clear()
|
||||
|
||||
_ = LLM(
|
||||
model="microsoft/Phi-tiny-MoE-instruct",
|
||||
max_model_len=256,
|
||||
load_format="dummy", # make the model loading faster
|
||||
compilation_config=compilation_config,
|
||||
num_gpu_blocks_override=8, # make the model loading faster
|
||||
)
|
||||
|
||||
# vLLM-compile cold start is special. By default, we do
|
||||
# one full dynamo capture of the entire forward pass.
|
||||
# The forward pass consists of 32 transformer layers.
|
||||
# Then, we split on the attention operation. This results in
|
||||
# 33 subgraphs (not including the attention operation).
|
||||
# The 33 subgraphs then get standalone_compile'd.
|
||||
#
|
||||
# There are actually only 3 unique subgraphs for this model
|
||||
# (all of its transformer layers are the same modulo weights);
|
||||
# this is true for most vLLM models.
|
||||
# So we test that during cold start, the aot_autograd cache
|
||||
# misses for 3 subgraphs and hits for the rest.
|
||||
assert counters["aot_autograd"]["autograd_cache_miss"] == 3
|
||||
assert counters["aot_autograd"]["autograd_cache_hit"] == 30
|
||||
@@ -8,6 +8,10 @@ import torch
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
|
||||
from vllm.compilation.backends import split_graph
|
||||
from vllm.compilation.fx_utils import find_op_nodes
|
||||
|
||||
# This import automatically registers `torch.ops.silly.attention`
|
||||
from . import silly_attention # noqa: F401
|
||||
|
||||
|
||||
def test_getitem_moved_to_producer_subgraph():
|
||||
@@ -122,3 +126,61 @@ def test_no_tuple_inputs_with_multiple_consumers():
|
||||
output_split = split_gm(new_x)
|
||||
|
||||
assert torch.allclose(output_original, output_split), "Output mismatch after split"
|
||||
|
||||
|
||||
def test_consecutive_ops_in_split():
|
||||
"""
|
||||
Test that consecutive splitting operations are grouped into the same subgraph
|
||||
"""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Define a simple model where consecutive operations create opportunities
|
||||
for splitting subgraphs.
|
||||
"""
|
||||
# Apply silly attention followed by consecutive operations
|
||||
intermediate = torch.relu(x)
|
||||
attn_inout = torch.sqrt(intermediate)
|
||||
torch.ops.silly.attention(intermediate, intermediate, attn_inout, attn_inout)
|
||||
final_result = torch.sigmoid(attn_inout)
|
||||
return final_result
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
# Create the traced FX graph for the model
|
||||
x = torch.randn(8, 4)
|
||||
|
||||
gm = make_fx(model_fn)(x)
|
||||
|
||||
# Assert presence of the expected operations in the setup
|
||||
assert (
|
||||
len(list(find_op_nodes(torch.ops.aten.relu, gm.graph))) == 1
|
||||
and len(list(find_op_nodes(torch.ops.aten.sqrt, gm.graph))) == 1
|
||||
), "Test setup failed: Expected sqrt and relu operations in the graph."
|
||||
|
||||
# Configure split operations to test
|
||||
splitting_ops = ["silly::attention", "aten::sqrt"]
|
||||
split_gm, split_items = split_graph(gm, splitting_ops)
|
||||
|
||||
# Validate the number of partitions
|
||||
assert len(split_items) == 3, (
|
||||
"Consecutive splitting operations were not grouped correctly."
|
||||
)
|
||||
|
||||
# Validate that correctness is preserved
|
||||
new_x = torch.randn(8, 4)
|
||||
output_original = gm(new_x)
|
||||
output_split = split_gm(new_x)
|
||||
assert torch.allclose(output_original, output_split), (
|
||||
"Output mismatch after splitting."
|
||||
)
|
||||
|
||||
# Check the splitting item has 2 nodes exactly (relu and attn)
|
||||
splitting_items = list(s for s in split_items if s.is_splitting_graph)
|
||||
assert len(splitting_items) == 1, "Expecting a single splitting graph"
|
||||
print(splitting_items[0].graph.graph)
|
||||
splitting_gm = splitting_items[0].graph
|
||||
assert len(splitting_gm.graph.nodes) == 4, "Expecting 4 nodes in splitting graph"
|
||||
assert [node.op for node in splitting_gm.graph.nodes] == ["placeholder"] + 2 * [
|
||||
"call_function"
|
||||
] + ["output"]
|
||||
|
||||
@@ -164,3 +164,40 @@ def test_classes_are_types():
|
||||
pass
|
||||
|
||||
assert endswith_fqname(LocalDummy, ".LocalDummy")
|
||||
|
||||
|
||||
def test_envs_compile_factors_stable():
|
||||
"""Test that envs.compile_factors() hash is stable across fresh initializations.
|
||||
|
||||
Uses subprocesses to ensure env vars with dynamic defaults (like UUIDs)
|
||||
are freshly generated each time, verifying they're properly ignored.
|
||||
"""
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
code = """
|
||||
import sys
|
||||
import logging
|
||||
logging.disable(logging.CRITICAL)
|
||||
from vllm import envs
|
||||
from vllm.config.utils import hash_factors
|
||||
print(hash_factors(envs.compile_factors()))
|
||||
"""
|
||||
|
||||
def get_hash_in_subprocess():
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
env={**dict(__import__("os").environ), "VLLM_LOGGING_LEVEL": "ERROR"},
|
||||
)
|
||||
return result.stdout.strip()
|
||||
|
||||
hash1 = get_hash_in_subprocess()
|
||||
hash2 = get_hash_in_subprocess()
|
||||
|
||||
assert hash1 == hash2, (
|
||||
"compile_factors hash differs between fresh initializations - "
|
||||
"dynamic env vars may not be properly ignored"
|
||||
)
|
||||
|
||||
@@ -992,7 +992,7 @@ async def test_mcp_tool_multi_turn(client: OpenAI, model_name: str, server):
|
||||
# First turn - make a calculation
|
||||
response1 = await client.responses.create(
|
||||
model=model_name,
|
||||
input="Calculate 123 * 456 using python and print the result.",
|
||||
input="Calculate 1234 * 4567 using python tool and print the result.",
|
||||
tools=tools,
|
||||
temperature=0.0,
|
||||
instructions=(
|
||||
|
||||
@@ -13,31 +13,10 @@ from vllm.utils.serial_utils import tensor2base64
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
|
||||
def _terratorch_dummy_messages():
|
||||
pixel_values = torch.full((6, 512, 512), 1.0, dtype=torch.float16)
|
||||
location_coords = torch.full((1, 2), 1.0, dtype=torch.float16)
|
||||
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"pixel_values": tensor2base64(pixel_values),
|
||||
"location_coords": tensor2base64(location_coords),
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
|
||||
)
|
||||
def test_single_request(model_name: str):
|
||||
def test_single_content(model_name: str):
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
@@ -59,7 +38,24 @@ def test_single_request(model_name: str):
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": _terratorch_dummy_messages(),
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"pixel_values": tensor2base64(
|
||||
torch.ones((6, 512, 512), dtype=torch.float16)
|
||||
),
|
||||
"location_coords": tensor2base64(
|
||||
torch.ones((1, 2), dtype=torch.float16)
|
||||
),
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"encoding_format": "base64",
|
||||
},
|
||||
)
|
||||
@@ -69,3 +65,87 @@ def test_single_request(model_name: str):
|
||||
|
||||
np_response = np.frombuffer(base64.b64decode(output), dtype=np.float32)
|
||||
assert len(np_response) == 524288
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["Qwen/Qwen3-VL-2B-Instruct"])
|
||||
def test_multi_content(model_name: str):
|
||||
args = [
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
"--max-model-len",
|
||||
"8192",
|
||||
"--enable-mm-embeds",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(model_name, args) as server:
|
||||
client = server.get_client()
|
||||
|
||||
# Image only
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
max_tokens=5,
|
||||
)
|
||||
|
||||
assert chat_completion.id is not None
|
||||
assert len(chat_completion.choices) == 1
|
||||
|
||||
# Interleaved text and image
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "OCR:"},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
max_tokens=5,
|
||||
)
|
||||
|
||||
assert chat_completion.id is not None
|
||||
assert len(chat_completion.choices) == 1
|
||||
|
||||
@@ -5,9 +5,9 @@ import json
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.entrypoints.test_utils import encode_base64_content_from_url
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
|
||||
MAXIMUM_VIDEOS = 1
|
||||
@@ -19,7 +19,7 @@ HF_OVERRIDES = {
|
||||
}
|
||||
input_text = "This product was excellent and exceeded my expectations"
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
image_base64 = encode_base64_content_from_url(image_url)
|
||||
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
|
||||
video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.score.protocol import ScoreResponse
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B"
|
||||
HF_OVERRIDES = {
|
||||
"architectures": ["Qwen3VLForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
}
|
||||
|
||||
query = "A cat standing in the snow."
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": query,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"8192",
|
||||
"--chat-template",
|
||||
str(VLLM_PATH / "examples/pooling/score/template/qwen3_vl_reranker.jinja"),
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer):
|
||||
queries = "What is the capital of France?"
|
||||
documents = "The capital of France is Paris."
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer):
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIServer):
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_base64_content(
|
||||
server: RemoteOpenAIServer,
|
||||
):
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
@@ -68,6 +68,16 @@ def phi3v_model_config_image_embeds():
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def qwen25omni_model_config_image_embeds():
|
||||
return ModelConfig(
|
||||
QWEN25OMNI_MODEL_ID,
|
||||
runner="generate",
|
||||
limit_mm_per_prompt={"image": 2},
|
||||
enable_mm_embeds=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def qwen2_audio_model_config():
|
||||
return ModelConfig(
|
||||
@@ -823,7 +833,8 @@ def test_parse_chat_messages_audio_embeds_with_string(
|
||||
import torch
|
||||
|
||||
# Create a sample audio embedding tensor
|
||||
audio_embedding = torch.randn(1, 128, 768)
|
||||
hidden_size = audio_embeds_model_config.get_inputs_embeds_size()
|
||||
audio_embedding = torch.randn(1, 128, hidden_size)
|
||||
|
||||
# Encode it as base64
|
||||
base64_audio_embedding = tensor2base64(audio_embedding)
|
||||
@@ -865,7 +876,8 @@ async def test_parse_chat_messages_audio_embeds_async(
|
||||
import torch
|
||||
|
||||
# Create a sample audio embedding tensor
|
||||
audio_embedding = torch.randn(1, 128, 768)
|
||||
hidden_size = audio_embeds_model_config.get_inputs_embeds_size()
|
||||
audio_embedding = torch.randn(1, 128, hidden_size)
|
||||
|
||||
# Encode it as base64
|
||||
base64_audio_embedding = tensor2base64(audio_embedding)
|
||||
@@ -908,8 +920,9 @@ def test_parse_chat_messages_multiple_image_embeds(
|
||||
can be provided in a single request, similar to regular images.
|
||||
"""
|
||||
# Create two sample image embedding tensors
|
||||
image_embedding_1 = torch.randn(256, 1024)
|
||||
image_embedding_2 = torch.randn(128, 1024)
|
||||
hidden_size = phi3v_model_config_image_embeds.get_inputs_embeds_size()
|
||||
image_embedding_1 = torch.randn(256, hidden_size)
|
||||
image_embedding_2 = torch.randn(128, hidden_size)
|
||||
|
||||
# Encode them as base64 using the convenience function
|
||||
base64_image_embedding_1 = tensor2base64(image_embedding_1)
|
||||
@@ -1022,8 +1035,9 @@ async def test_parse_chat_messages_multiple_image_embeds_async(
|
||||
This validates the AsyncMultiModalItemTracker also supports multiple embeddings.
|
||||
"""
|
||||
# Create two sample image embedding tensors
|
||||
image_embedding_1 = torch.randn(200, 768)
|
||||
image_embedding_2 = torch.randn(150, 768)
|
||||
hidden_size = phi3v_model_config_image_embeds.get_inputs_embeds_size()
|
||||
image_embedding_1 = torch.randn(200, hidden_size)
|
||||
image_embedding_2 = torch.randn(150, hidden_size)
|
||||
|
||||
# Encode them as base64 using the convenience function
|
||||
base64_image_embedding_1 = tensor2base64(image_embedding_1)
|
||||
@@ -1145,13 +1159,14 @@ def test_parse_chat_messages_empty_dict_image_embeds(
|
||||
|
||||
|
||||
def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
phi3v_model_config_image_embeds,
|
||||
qwen25omni_model_config_image_embeds,
|
||||
):
|
||||
"""Test that multiple dictionaries for image_embeds is handled without errors."""
|
||||
# Create two sample image embedding tensors
|
||||
batch_size = 2
|
||||
image_embedding_1 = torch.randn(batch_size, 256, 1024)
|
||||
image_embedding_2 = torch.randn(batch_size, 3)
|
||||
hidden_size = qwen25omni_model_config_image_embeds.get_inputs_embeds_size()
|
||||
image_embeds = torch.randn(batch_size * 220, hidden_size)
|
||||
image_grid_thw = torch.tensor([[1, 22, 40] for _ in range(batch_size)])
|
||||
|
||||
conversation, mm_data, mm_uuids = parse_chat_messages(
|
||||
[
|
||||
@@ -1161,18 +1176,20 @@ def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embedding_1": tensor2base64(p),
|
||||
"image_embedding_2": tensor2base64(i),
|
||||
"image_embeds": tensor2base64(embeds),
|
||||
"image_grid_thw": tensor2base64(grid_thw),
|
||||
},
|
||||
}
|
||||
for p, i in zip(image_embedding_1, image_embedding_2)
|
||||
for embeds, grid_thw in zip(
|
||||
image_embeds.chunk(batch_size), image_grid_thw
|
||||
)
|
||||
]
|
||||
+ [
|
||||
{"type": "text", "text": "Describe these two images."},
|
||||
],
|
||||
}
|
||||
],
|
||||
phi3v_model_config_image_embeds,
|
||||
qwen25omni_model_config_image_embeds,
|
||||
content_format="string",
|
||||
)
|
||||
|
||||
@@ -1180,7 +1197,8 @@ def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
assert conversation == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "<|image_1|>\n<|image_2|>\nDescribe these two images.",
|
||||
"content": "<|vision_start|><|IMAGE|><|vision_end|>\n"
|
||||
"<|vision_start|><|IMAGE|><|vision_end|>\nDescribe these two images.",
|
||||
}
|
||||
]
|
||||
|
||||
@@ -1191,10 +1209,10 @@ def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
assert len(mm_data["image"]) == batch_size
|
||||
|
||||
# Verify each embedding has the correct shape
|
||||
assert isinstance(mm_data["image"]["image_embedding_1"], torch.Tensor)
|
||||
assert mm_data["image"]["image_embedding_1"].shape == image_embedding_1.shape
|
||||
assert isinstance(mm_data["image"]["image_embedding_2"], torch.Tensor)
|
||||
assert mm_data["image"]["image_embedding_2"].shape == image_embedding_2.shape
|
||||
assert isinstance(mm_data["image"]["image_embeds"], torch.Tensor)
|
||||
assert mm_data["image"]["image_embeds"].shape == image_embeds.shape
|
||||
assert isinstance(mm_data["image"]["image_grid_thw"], torch.Tensor)
|
||||
assert mm_data["image"]["image_grid_thw"].shape == image_grid_thw.shape
|
||||
|
||||
# Verify UUIDs (None since we didn't provide any)
|
||||
_assert_mm_uuids(mm_uuids, batch_size, expected_uuids=[None, None])
|
||||
|
||||
@@ -1,9 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import base64
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.entrypoints.utils import sanitize_message
|
||||
|
||||
|
||||
@@ -12,11 +8,3 @@ def test_sanitize_message():
|
||||
sanitize_message("<_io.BytesIO object at 0x7a95e299e750>")
|
||||
== "<_io.BytesIO object>"
|
||||
)
|
||||
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
|
||||
with requests.get(content_url) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return {"url": f"data:image/jpeg;base64,{result}"}
|
||||
|
||||
@@ -1,8 +0,0 @@
|
||||
model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
-8
@@ -1,8 +0,0 @@
|
||||
model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
@@ -5,12 +5,10 @@ Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-marlin.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-fi-cutlass-dp-ep.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass-dp-ep.yaml
|
||||
Llama-4-Scout-BF16-fi-cutlass.yaml
|
||||
Llama-4-Scout-BF16-triton.yaml
|
||||
Mixtral-8x7B-BF16-fi-cutlass.yaml
|
||||
|
||||
@@ -17,6 +17,8 @@ from vllm.model_executor.layers.activation import (
|
||||
QuickGELU,
|
||||
SiluAndMul,
|
||||
SwigluOAIAndMul,
|
||||
SwigluStepAndMul,
|
||||
swiglustep_and_mul_triton,
|
||||
)
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
@@ -36,6 +38,7 @@ CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 e
|
||||
"gelu_tanh",
|
||||
"fatrelu",
|
||||
"swigluoai_and_mul",
|
||||
"swiglustep_and_mul",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@@ -75,9 +78,12 @@ def test_act_and_mul(
|
||||
elif activation == "swigluoai_and_mul":
|
||||
layer = SwigluOAIAndMul()
|
||||
fn = torch.ops._C.swigluoai_and_mul
|
||||
elif activation == "swiglustep_and_mul":
|
||||
layer = SwigluStepAndMul()
|
||||
fn = swiglustep_and_mul_triton
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
if activation == "swigluoai_and_mul":
|
||||
if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]:
|
||||
rtol = {
|
||||
# For fp16, change the relative tolerance from 1e-3 to 2e-3
|
||||
torch.float16: 2e-3,
|
||||
@@ -104,7 +110,7 @@ def test_act_and_mul(
|
||||
opcheck(fn, (out, x, threshold))
|
||||
elif activation == "swigluoai_and_mul":
|
||||
opcheck(fn, (out, x, layer.alpha, layer.limit))
|
||||
else:
|
||||
elif activation != "swiglustep_and_mul":
|
||||
opcheck(fn, (out, x))
|
||||
|
||||
|
||||
|
||||
@@ -715,7 +715,7 @@ def test_mixtral_moe(
|
||||
|
||||
# need to override the forward context for unittests, otherwise it assumes
|
||||
# we're running the model forward pass (the model specified in vllm_config)
|
||||
get_forward_context().remaining_moe_layers = None
|
||||
get_forward_context().all_moe_layers = None
|
||||
|
||||
# Run forward passes for both MoE blocks
|
||||
hf_states, _ = hf_moe.forward(hf_inputs)
|
||||
|
||||
@@ -107,10 +107,14 @@ def test_flashinfer_nvfp4_gemm(
|
||||
# from checkpoints are in linear scales.
|
||||
# So instead of needing to swizzle for cutlass as in modelopt.py,
|
||||
# we need to unswizzle for trtllm here.
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale, backend)
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
|
||||
a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend
|
||||
)
|
||||
is_sf_128x4_layout = not (backend == "trtllm" and m <= 32)
|
||||
|
||||
b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
|
||||
b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(
|
||||
b_dtype, b_global_scale, is_sf_swizzled_layout=True
|
||||
)
|
||||
|
||||
# get_ref_results unswizzles the scales internally.
|
||||
expected_out = get_ref_results(
|
||||
|
||||
@@ -59,7 +59,7 @@ if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"These tests require gptq_marlin_repack,"
|
||||
"marlin_int4_fp8_preprocess, gptq_marlin_24_gemm,"
|
||||
"or gptq_marlin_gemm which are not supported on ROCm.",
|
||||
"or marlin_gemm which are not supported on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
@@ -417,7 +417,7 @@ def marlin_generate_valid_test_cases():
|
||||
),
|
||||
marlin_generate_valid_test_cases(),
|
||||
)
|
||||
def test_gptq_marlin_gemm(
|
||||
def test_marlin_gemm(
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
@@ -511,7 +511,7 @@ def test_gptq_marlin_gemm(
|
||||
|
||||
output = torch.empty((size_m, size_n), dtype=dtype, device=a_input.device)
|
||||
|
||||
output = ops.gptq_marlin_gemm(
|
||||
output = ops.marlin_gemm(
|
||||
a_input,
|
||||
output,
|
||||
marlin_q_w,
|
||||
@@ -646,7 +646,7 @@ def test_marlin_gemm_subset_input():
|
||||
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
|
||||
workspace = marlin_make_workspace_new(a_input.device)
|
||||
|
||||
output = ops.gptq_marlin_gemm(
|
||||
output = ops.marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
@@ -695,7 +695,7 @@ def test_marlin_gemm_with_bias(size_m):
|
||||
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
|
||||
workspace = marlin_make_workspace_new(a_input.device)
|
||||
|
||||
output = ops.gptq_marlin_gemm(
|
||||
output = ops.marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
|
||||
@@ -27,6 +27,12 @@ PAD_SHAPES = [
|
||||
(150, 128),
|
||||
(150, 48),
|
||||
(90, 80),
|
||||
(128, 512),
|
||||
(128, 1024),
|
||||
(128, 2048),
|
||||
(64, 7168),
|
||||
(64, 7152),
|
||||
(32, 14336),
|
||||
]
|
||||
SEEDS = [42]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
@@ -173,3 +179,25 @@ def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_padded_no_sf_swizzled(pad_shape: tuple[int, int]) -> None:
|
||||
dtype = torch.float16
|
||||
set_random_seed(42)
|
||||
torch.set_default_device("cuda:0")
|
||||
|
||||
m, n = pad_shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = ops.scaled_fp4_quant(x, global_scale, is_sf_swizzled_layout=False)
|
||||
scale_ans = out_scale.to(torch.float32)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
@@ -87,6 +87,13 @@ NKM_FACTORS_WVSPLITK_FP8 = [
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
def pad_weights_fp8(weight):
|
||||
num_pad = 256 // weight.element_size()
|
||||
import torch.nn.functional as F
|
||||
|
||||
return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@@ -191,11 +198,12 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("padded", [False, True])
|
||||
@pytest.mark.skipif(
|
||||
not (current_platform.is_rocm() and current_platform.supports_fp8()),
|
||||
reason="only test for rocm fp8",
|
||||
)
|
||||
def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
|
||||
def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
A = torch.rand(n, k, device="cuda") - 0.5
|
||||
@@ -203,6 +211,8 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
|
||||
|
||||
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
|
||||
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
|
||||
if padded:
|
||||
B = pad_weights_fp8(B)
|
||||
|
||||
ref_out = torch._scaled_mm(
|
||||
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b
|
||||
@@ -222,11 +232,12 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("padded", [False, True])
|
||||
@pytest.mark.skipif(
|
||||
not (current_platform.is_rocm() and current_platform.supports_fp8()),
|
||||
reason="only test for rocm fp8",
|
||||
)
|
||||
def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
|
||||
def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
|
||||
@@ -236,6 +247,8 @@ def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
|
||||
|
||||
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
|
||||
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
|
||||
if padded:
|
||||
B = pad_weights_fp8(B)
|
||||
|
||||
ref_out = torch._scaled_mm(
|
||||
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
|
||||
|
||||
@@ -2,13 +2,11 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Integration tests for FlexAttention backend vs default backend"""
|
||||
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from packaging import version
|
||||
|
||||
from tests.utils import set_random_seed
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
@@ -27,15 +25,6 @@ MINIMUM_TORCH_VERSION = version.parse("2.7.0")
|
||||
DIRECT_BUILD_VERSION = version.parse("2.9.dev0")
|
||||
|
||||
|
||||
def set_seed(seed):
|
||||
"""Set seeds for reproducibility"""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
|
||||
reason="CUDA not available or PyTorch version < 2.7",
|
||||
@@ -57,7 +46,7 @@ def test_flex_attention_vs_default_backend(vllm_runner):
|
||||
]
|
||||
|
||||
# Run with flex attention
|
||||
set_seed(seed)
|
||||
set_random_seed(seed)
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
runner="generate",
|
||||
@@ -71,7 +60,7 @@ def test_flex_attention_vs_default_backend(vllm_runner):
|
||||
)
|
||||
|
||||
# Run with default backend
|
||||
set_seed(seed)
|
||||
set_random_seed(seed)
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
runner="generate",
|
||||
|
||||
@@ -7,14 +7,6 @@ import torch
|
||||
from ....conftest import VllmRunner
|
||||
|
||||
|
||||
def generate_test_mm_data():
|
||||
mm_data = {
|
||||
"pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
|
||||
"location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
|
||||
}
|
||||
return mm_data
|
||||
|
||||
|
||||
def _run_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
@@ -23,7 +15,12 @@ def _run_test(
|
||||
{
|
||||
# This model deals with no text input
|
||||
"prompt_token_ids": [1],
|
||||
"multi_modal_data": generate_test_mm_data(),
|
||||
"multi_modal_data": {
|
||||
"image": {
|
||||
"pixel_values": torch.ones((6, 512, 512), dtype=torch.float16),
|
||||
"location_coords": torch.ones((1, 2), dtype=torch.float16),
|
||||
}
|
||||
},
|
||||
}
|
||||
for _ in range(10)
|
||||
]
|
||||
|
||||
@@ -48,18 +48,6 @@ VideoInput: TypeAlias = (
|
||||
AudioInput = list[tuple[np.ndarray, int]]
|
||||
|
||||
|
||||
MM_OPTIONS_OVERRIDES = {
|
||||
# Qwen3-VL's default profiling video size (64x64) can cause trouble
|
||||
# after resizing, so we override it here for testing.
|
||||
"qwen3_vl": dict(
|
||||
video=VideoDummyOptions(num_frames=128, width=256, height=256),
|
||||
),
|
||||
"qwen3_vl_moe": dict(
|
||||
video=VideoDummyOptions(num_frames=128, width=256, height=256),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _resize_data(
|
||||
_data: Image.Image | np.ndarray, size_factor: float
|
||||
) -> Image.Image | np.ndarray:
|
||||
@@ -73,12 +61,12 @@ def _resize_data(
|
||||
elif is_list_of(_data, Image.Image):
|
||||
W, H = next(iter(_data)).width, next(iter(_data)).height
|
||||
T = len(_data)
|
||||
T, W, H = map(lambda x: max(int(x * size_factor), 1), (T, W, H))
|
||||
T, W, H = map(lambda x: max(int(x * size_factor), 2), (T, W, H))
|
||||
return [d.resize((W, H)) for d in _data[:T]]
|
||||
# Video input with numpy arrays
|
||||
elif isinstance(_data, np.ndarray) and _data.ndim >= 4:
|
||||
T, H, W, C = _data.shape[-4:]
|
||||
T, H, W = map(lambda x: max(int(x * size_factor), 1), (T, H, W))
|
||||
T, H, W = map(lambda x: max(int(x * size_factor), 2), (T, H, W))
|
||||
return _data[..., :T, :H, :W, :C]
|
||||
# Audio input
|
||||
elif isinstance(_data, np.ndarray) and _data.ndim == 1:
|
||||
@@ -103,8 +91,6 @@ def create_batched_mm_kwargs(
|
||||
processor: BaseMultiModalProcessor,
|
||||
size_factors: tuple[float, ...] = (1.0, 0.5, 0.25),
|
||||
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
|
||||
model_type = model_config.hf_config.model_type
|
||||
|
||||
processing_info = processor.info
|
||||
dummy_inputs = processor.dummy_inputs
|
||||
supported_mm_limits = processing_info.get_supported_mm_limits()
|
||||
@@ -115,7 +101,6 @@ def create_batched_mm_kwargs(
|
||||
processor_inputs = dummy_inputs.get_dummy_processor_inputs(
|
||||
seq_len=model_config.max_model_len,
|
||||
mm_counts=mm_counts,
|
||||
mm_options=MM_OPTIONS_OVERRIDES.get(model_type),
|
||||
)
|
||||
mm_data = processor_inputs.mm_data
|
||||
resized_mm_data = {
|
||||
|
||||
@@ -480,6 +480,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"Step1ForCausalLM": _HfExamplesInfo(
|
||||
"stepfun-ai/Step-Audio-EditX", trust_remote_code=True
|
||||
),
|
||||
"Step3p5ForCausalLM": _HfExamplesInfo(
|
||||
"stepfun-ai/step-3.5-flash", is_available_online=False
|
||||
),
|
||||
"SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
|
||||
"StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),
|
||||
"StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
|
||||
@@ -771,6 +774,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
)
|
||||
},
|
||||
),
|
||||
"KimiK25ForConditionalGeneration": _HfExamplesInfo(
|
||||
"moonshotai/Kimi-K2.5",
|
||||
trust_remote_code=True,
|
||||
is_available_online=False,
|
||||
),
|
||||
"LightOnOCRForConditionalGeneration": _HfExamplesInfo(
|
||||
"lightonai/LightOnOCR-1B-1025"
|
||||
),
|
||||
@@ -1076,6 +1084,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
"Qwen3NextMTP": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-Next-80B-A3B-Instruct", min_transformers_version="4.56.3"
|
||||
),
|
||||
"Step3p5MTP": _HfExamplesInfo(
|
||||
"stepfun-ai/Step-3.5-Flash",
|
||||
trust_remote_code=True,
|
||||
speculative_model="stepfun-ai/Step-3.5-Flash",
|
||||
is_available_online=False,
|
||||
),
|
||||
}
|
||||
|
||||
_TRANSFORMERS_BACKEND_MODELS = {
|
||||
|
||||
@@ -349,8 +349,10 @@ class PrithviMultimodalDataProcessor(IOProcessor):
|
||||
{
|
||||
"prompt_token_ids": [1],
|
||||
"multi_modal_data": {
|
||||
"pixel_values": window.to(torch.float16)[0],
|
||||
"location_coords": location_coords.to(torch.float16),
|
||||
"image": {
|
||||
"pixel_values": window.to(torch.float16)[0],
|
||||
"location_coords": location_coords.to(torch.float16),
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
@@ -20,7 +20,7 @@ MM_BEAM_WIDTHS = [2]
|
||||
MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
|
||||
|
||||
|
||||
@pytest.mark.skip_v1 # FIXME: This fails on V1 right now.
|
||||
@pytest.mark.skip_v1 # V1 engine does not yet support beam search
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
@@ -62,7 +62,7 @@ def test_beam_search_single_input(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip_v1 # FIXME: This fails on V1 right now.
|
||||
@pytest.mark.skip_v1 # V1 engine does not yet support beam search
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
|
||||
+4
-1
@@ -59,7 +59,10 @@ from vllm.tokenizers import get_tokenizer
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.mem_constants import GB_bytes
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import cuda_device_count_stateless
|
||||
from vllm.utils.torch_utils import (
|
||||
cuda_device_count_stateless,
|
||||
set_random_seed, # noqa: F401 - re-exported for use in test files
|
||||
)
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (
|
||||
create_common_attn_metadata,
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
try_backend_includes_kv_cache_update,
|
||||
try_get_attention_backend,
|
||||
)
|
||||
from vllm.config import ModelConfig
|
||||
@@ -295,6 +296,10 @@ def run_attention_backend(
|
||||
# Run forward pass
|
||||
# NOTE: The query, key, and value are already shaped correctly
|
||||
# in the calling test function.
|
||||
if not try_backend_includes_kv_cache_update(actual_backend):
|
||||
impl.do_kv_cache_update(
|
||||
mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
|
||||
)
|
||||
output = impl.forward(
|
||||
mock_layer, query, key, value, kv_cache, attn_metadata, output=output
|
||||
)
|
||||
|
||||
@@ -130,6 +130,18 @@ def try_get_attention_backend(
|
||||
raise AssertionError("unreachable") from None
|
||||
|
||||
|
||||
def try_backend_includes_kv_cache_update(
|
||||
backend: AttentionBackendEnum,
|
||||
) -> bool:
|
||||
"""Try to get the attention backend class, skipping test if not found."""
|
||||
try:
|
||||
backend_class = backend.get_class()
|
||||
return backend_class.forward_includes_kv_cache_update
|
||||
except ImportError as e:
|
||||
pytest.skip(f"{backend.name} not available: {e}")
|
||||
raise AssertionError("unreachable") from None
|
||||
|
||||
|
||||
def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
|
||||
"""Create a FullAttentionSpec from ModelParams only."""
|
||||
return FullAttentionSpec(
|
||||
|
||||
@@ -107,7 +107,10 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
|
||||
|
||||
|
||||
def make_kv_cache_config_hybrid_model(
|
||||
block_size: int, num_blocks: int, second_spec_type: str = "sliding_window"
|
||||
block_size: int,
|
||||
num_blocks: int,
|
||||
sliding_window_blocks: int,
|
||||
second_spec_type: str = "sliding_window",
|
||||
) -> KVCacheConfig:
|
||||
if second_spec_type == "sliding_window":
|
||||
second_spec = SlidingWindowSpec(
|
||||
@@ -115,7 +118,7 @@ def make_kv_cache_config_hybrid_model(
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=2 * block_size,
|
||||
sliding_window=sliding_window_blocks * block_size,
|
||||
)
|
||||
elif second_spec_type == "mamba":
|
||||
second_spec = MambaSpec(
|
||||
@@ -325,7 +328,7 @@ def test_prefill(hash_fn):
|
||||
def test_prefill_hybrid_model():
|
||||
block_size = 16
|
||||
manager = KVCacheManager(
|
||||
make_kv_cache_config_hybrid_model(block_size, 21),
|
||||
make_kv_cache_config_hybrid_model(block_size, 21, 2),
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
@@ -334,7 +337,8 @@ def test_prefill_hybrid_model():
|
||||
hash_fn = sha256
|
||||
|
||||
# Complete 3 blocks (48 tokens)
|
||||
common_token_ids = [i for i in range(3) for _ in range(block_size)]
|
||||
num_full_blocks = 3
|
||||
common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
|
||||
|
||||
# Fully cache miss
|
||||
# Incomplete 1 block (7 tokens)
|
||||
@@ -375,6 +379,7 @@ def test_prefill_hybrid_model():
|
||||
# Cache hit in the common prefix
|
||||
# Incomplete 1 block (5 tokens)
|
||||
unique_token_ids = [3] * 5
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
|
||||
assert len(req1.block_hashes) == 3
|
||||
@@ -394,34 +399,13 @@ def test_prefill_hybrid_model():
|
||||
manager.free(req0)
|
||||
manager.free(req1)
|
||||
|
||||
cached_block_hash_to_block_bak = copy.copy(
|
||||
manager.block_pool.cached_block_hash_to_block._cache
|
||||
)
|
||||
|
||||
def test_partial_request_hit(
|
||||
request_id: str,
|
||||
hash_to_evict: list[BlockHashWithGroupId],
|
||||
expect_hit_length: int,
|
||||
):
|
||||
req = make_request(
|
||||
request_id, common_token_ids + unique_token_ids, block_size, sha256
|
||||
)
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
|
||||
assert len(req.block_hashes) == 3
|
||||
assert num_computed_tokens == expect_hit_length * block_size
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
assert len(block_per_group) == num_computed_tokens // block_size
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (
|
||||
cached_block_hash_to_block_bak[hash_with_group_id]
|
||||
)
|
||||
manager.free(req)
|
||||
|
||||
# Evict the blocks outside sliding window, does not affect the hit length.
|
||||
test_partial_request_hit(
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"2",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
make_block_hash_with_group_id(block_hashes[0], 2),
|
||||
@@ -430,13 +414,23 @@ def test_prefill_hybrid_model():
|
||||
)
|
||||
|
||||
# Evict the first block of full attention, makes total cache miss.
|
||||
test_partial_request_hit(
|
||||
"3", [make_block_hash_with_group_id(block_hashes[0], 0)], 0
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"3",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[0], 0)],
|
||||
0,
|
||||
)
|
||||
|
||||
# Evict the last block of all layers, reduces the hit length to 2.
|
||||
test_partial_request_hit(
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"4",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[2], 0),
|
||||
make_block_hash_with_group_id(block_hashes[2], 1),
|
||||
@@ -446,18 +440,36 @@ def test_prefill_hybrid_model():
|
||||
)
|
||||
|
||||
# Evict the last block of full attention, reduces the hit length to 2.
|
||||
test_partial_request_hit(
|
||||
"5", [make_block_hash_with_group_id(block_hashes[2], 0)], 2
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"5",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[2], 0)],
|
||||
2,
|
||||
)
|
||||
|
||||
# Evict the last block of sliding window, reduces the hit length to 2.
|
||||
test_partial_request_hit(
|
||||
"6", [make_block_hash_with_group_id(block_hashes[2], 1)], 2
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"6",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[2], 1)],
|
||||
2,
|
||||
)
|
||||
|
||||
# Evict the last block of sliding window, reduces the hit length to 2.
|
||||
test_partial_request_hit(
|
||||
"7", [make_block_hash_with_group_id(block_hashes[2], 2)], 2
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"7",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[2], 2)],
|
||||
2,
|
||||
)
|
||||
|
||||
# Evict different set of blocks for full attention and sliding window makes
|
||||
@@ -466,8 +478,12 @@ def test_prefill_hybrid_model():
|
||||
# The cache hit length of sliding window is 2 * block_size.
|
||||
# Then it is cache miss as the two type of layers
|
||||
# have different hit length.
|
||||
test_partial_request_hit(
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"8",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[2], 0),
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
@@ -477,6 +493,214 @@ def test_prefill_hybrid_model():
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_hybrid_model_eagle():
|
||||
block_size = 16
|
||||
kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 3)
|
||||
manager = KVCacheManager(
|
||||
kv_cache_config,
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
use_eagle=True,
|
||||
)
|
||||
|
||||
hash_fn = sha256
|
||||
|
||||
# Complete 6 blocks (96 tokens)
|
||||
num_full_blocks = 6
|
||||
common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
|
||||
|
||||
# Fully cache miss
|
||||
# Incomplete 1 block (7 tokens)
|
||||
unique_token_ids = [6] * 7
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
req0 = make_request("0", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
|
||||
assert len(req0.block_hashes) == len(all_token_ids) // block_size
|
||||
assert not computed_blocks.blocks[0]
|
||||
assert num_computed_tokens == 0
|
||||
blocks = manager.allocate_slots(
|
||||
req0, len(all_token_ids), num_computed_tokens, computed_blocks
|
||||
)
|
||||
block_ids = (
|
||||
[1, 2, 3, 4, 5, 6, 7],
|
||||
[8, 9, 10, 11, 12, 13, 14],
|
||||
[15, 16, 17, 18, 19, 20, 21],
|
||||
)
|
||||
assert blocks is not None and blocks.get_block_ids() == block_ids
|
||||
|
||||
# Check full block metadata
|
||||
parent_block_hash = None
|
||||
for i, full_block_ids in enumerate(zip(*(row[:-1] for row in block_ids))):
|
||||
block_tokens = tuple(all_token_ids[i * block_size : (i + 1) * block_size])
|
||||
block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)
|
||||
for group_id, block_id in enumerate(full_block_ids):
|
||||
blk_hash = manager.block_pool.blocks[block_id].block_hash
|
||||
assert blk_hash is not None
|
||||
assert get_block_hash(blk_hash) == block_hash
|
||||
assert get_group_id(blk_hash) == group_id
|
||||
assert manager.block_pool.blocks[block_id].ref_cnt == 1
|
||||
parent_block_hash = block_hash
|
||||
|
||||
# Check partial block metadata
|
||||
for partial_block_id in (row[-1] for row in block_ids):
|
||||
assert manager.block_pool.blocks[partial_block_id].block_hash is None
|
||||
assert manager.block_pool.blocks[partial_block_id].ref_cnt == 1
|
||||
|
||||
# Cache hit in the common prefix
|
||||
# Incomplete 1 block (5 tokens)
|
||||
unique_token_ids = [6] * 5
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
req1 = make_request("1", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
|
||||
assert len(req1.block_hashes) == num_full_blocks
|
||||
assert computed_blocks.get_block_ids() == (
|
||||
[1, 2, 3, 4],
|
||||
[0, 9, 10, 11],
|
||||
[0, 16, 17, 18],
|
||||
)
|
||||
assert num_computed_tokens == 4 * block_size
|
||||
num_new_tokens = len(all_token_ids) - num_computed_tokens
|
||||
blocks = manager.allocate_slots(
|
||||
req1, num_new_tokens, num_computed_tokens, computed_blocks
|
||||
)
|
||||
assert blocks is not None and blocks.get_block_ids() == (
|
||||
[22, 23, 24],
|
||||
[25, 26, 27],
|
||||
[28, 29, 30],
|
||||
)
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
for block in block_per_group:
|
||||
if block != manager.block_pool.null_block:
|
||||
assert block.ref_cnt == 2
|
||||
|
||||
block_hashes = req1.block_hashes
|
||||
manager.free(req0)
|
||||
manager.free(req1)
|
||||
|
||||
# Evict the blocks outside sliding window, does not affect the hit length.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"2",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
make_block_hash_with_group_id(block_hashes[0], 2),
|
||||
],
|
||||
4,
|
||||
)
|
||||
|
||||
# Evict the first block of full attention, makes total cache miss.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"3",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[0], 0)],
|
||||
0,
|
||||
)
|
||||
|
||||
# Evict the last block of all layers, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"4",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[-1], 0),
|
||||
make_block_hash_with_group_id(block_hashes[-1], 1),
|
||||
make_block_hash_with_group_id(block_hashes[-1], 2),
|
||||
],
|
||||
3,
|
||||
)
|
||||
|
||||
# Evict the last block of full attention, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"5",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[-1], 0)],
|
||||
3,
|
||||
)
|
||||
|
||||
# Since the last block of full attention is dropped for eagle, evict
|
||||
# the second last block of sliding window, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"6",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[-2], 1)],
|
||||
3,
|
||||
)
|
||||
|
||||
# Since the last block of full attention is dropped for eagle, evict
|
||||
# the second last block of sliding window, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"7",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[-2], 2)],
|
||||
3,
|
||||
)
|
||||
|
||||
# Evict different set of blocks for full attention and sliding window makes
|
||||
# total cache miss.
|
||||
# The cache hit length of full attention is 4 * block_size.
|
||||
# The cache hit length of sliding window is 3 * block_size.
|
||||
# Then it is cache miss as the two type of layers
|
||||
# have different hit length.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"8",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[-1], 0),
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
make_block_hash_with_group_id(block_hashes[0], 2),
|
||||
],
|
||||
0,
|
||||
)
|
||||
|
||||
|
||||
def _test_partial_request_hit(
|
||||
manager: KVCacheManager,
|
||||
block_size: int,
|
||||
num_full_blocks,
|
||||
request_id: str,
|
||||
prompt_token_ids: list[int],
|
||||
hash_to_evict: list[BlockHashWithGroupId],
|
||||
expect_hit_length: int,
|
||||
):
|
||||
cached_block_hash_to_block_bak = copy.copy(
|
||||
manager.block_pool.cached_block_hash_to_block._cache
|
||||
)
|
||||
req = make_request(request_id, prompt_token_ids, block_size, sha256)
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
|
||||
assert len(req.block_hashes) == num_full_blocks
|
||||
assert num_computed_tokens == expect_hit_length * block_size
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
assert len(block_per_group) == num_computed_tokens // block_size
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (
|
||||
cached_block_hash_to_block_bak[hash_with_group_id]
|
||||
)
|
||||
manager.free(req)
|
||||
|
||||
|
||||
def _make_hybrid_kv_cache_config(
|
||||
block_size: int, num_blocks: int, spec_types: list[str]
|
||||
) -> KVCacheConfig:
|
||||
@@ -655,6 +879,85 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
|
||||
manager.free(req1)
|
||||
|
||||
|
||||
# Test cases with eagle enabled: Only test a single simple case for now.
|
||||
# - 2 groups: 1 full + 1 other
|
||||
_EAGLE_HYBRID_MODEL_TEST_CASES = [
|
||||
# 2 groups: 1 full + 1 other
|
||||
pytest.param(["full", "sliding_window"], 2, id="2g-full+sw"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spec_types,expect_hit_length", _EAGLE_HYBRID_MODEL_TEST_CASES)
|
||||
def test_prefill_hybrid_model_combinations_eagle(
|
||||
spec_types: list[str], expect_hit_length: int
|
||||
):
|
||||
"""
|
||||
Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE.
|
||||
More complex hybrid models with EAGLE are not yet supported (see issue #32802).
|
||||
"""
|
||||
block_size = 16
|
||||
num_groups = len(spec_types)
|
||||
# Allocate enough blocks for all groups
|
||||
num_blocks = 10 * num_groups
|
||||
|
||||
kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)
|
||||
manager = KVCacheManager(
|
||||
kv_cache_config,
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
use_eagle=True,
|
||||
)
|
||||
|
||||
hash_fn = sha256
|
||||
|
||||
# Complete 3 blocks (48 tokens)
|
||||
num_full_blocks = 4
|
||||
common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
|
||||
unique_token_ids = [4] * 7
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
|
||||
# First request: no cache hit initially
|
||||
req0 = make_request("0", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
|
||||
|
||||
assert len(req0.block_hashes) == num_full_blocks
|
||||
assert not computed_blocks.blocks[0] # No cache hit initially
|
||||
assert num_computed_tokens == 0
|
||||
|
||||
blocks = manager.allocate_slots(
|
||||
req0, len(all_token_ids), num_computed_tokens, computed_blocks
|
||||
)
|
||||
assert blocks is not None
|
||||
# Should have blocks for all groups
|
||||
assert len(blocks.get_block_ids()) == num_groups
|
||||
|
||||
# Second request: should hit cached blocks for common prefix
|
||||
all_token_ids = common_token_ids + [6] * 5
|
||||
req1 = make_request("1", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
|
||||
|
||||
# Should hit cached blocks for all groups
|
||||
assert num_computed_tokens == expect_hit_length * block_size
|
||||
assert len(computed_blocks.blocks) == num_groups
|
||||
# Verify each group has the correct number of computed blocks
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
assert len(block_per_group) == expect_hit_length
|
||||
|
||||
# Allocate and verify blocks for second request
|
||||
blocks = manager.allocate_slots(
|
||||
req1,
|
||||
len(all_token_ids) - num_computed_tokens,
|
||||
num_computed_tokens,
|
||||
computed_blocks,
|
||||
)
|
||||
assert blocks is not None
|
||||
assert len(blocks.get_block_ids()) == num_groups
|
||||
|
||||
manager.free(req0)
|
||||
manager.free(req1)
|
||||
|
||||
|
||||
def test_prefill_plp():
|
||||
"""Test prefill with APC and some prompt logprobs (plp) requests.
|
||||
|
||||
|
||||
@@ -650,9 +650,9 @@ def test_schedule_order(enable_chunked_prefill: bool):
|
||||
)
|
||||
|
||||
# long requests
|
||||
requests = create_requests(num_requests=2, num_tokens=800)
|
||||
requests = create_requests(num_requests=2, num_tokens=800, req_ids=["1", "2"])
|
||||
# short requests
|
||||
requests += create_requests(num_requests=2, num_tokens=10)
|
||||
requests += create_requests(num_requests=2, num_tokens=10, req_ids=["3", "4"])
|
||||
|
||||
for request in requests:
|
||||
scheduler.add_request(request)
|
||||
@@ -870,6 +870,66 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
|
||||
assert stats.num_accepted_tokens_per_pos == expected[3]
|
||||
|
||||
|
||||
def test_spec_decoding_stats_empty_output():
|
||||
"""Test that spec decoding stats handle empty output tokens gracefully.
|
||||
|
||||
This is a regression test for a bug where empty sampled_token_ids
|
||||
would cause num_accepted = len([]) - 1 = -1, leading to a
|
||||
ValueError when incrementing a Prometheus counter with a negative value.
|
||||
"""
|
||||
num_spec_tokens = 3
|
||||
scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens)
|
||||
requests = create_requests(num_requests=1, num_tokens=1)
|
||||
request = requests[0]
|
||||
req_id = request.request_id
|
||||
|
||||
scheduler.add_request(request)
|
||||
|
||||
# Initial schedule (prefill)
|
||||
output = scheduler.schedule()
|
||||
assert len(output.scheduled_new_reqs) == 1
|
||||
|
||||
# Complete the prefill with a sampled token
|
||||
model_runner_output = ModelRunnerOutput(
|
||||
req_ids=[req_id],
|
||||
req_id_to_index={req_id: 0},
|
||||
sampled_token_ids=[[0]],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
)
|
||||
scheduler.update_from_output(output, model_runner_output)
|
||||
|
||||
# Add draft tokens for speculation
|
||||
draft_token_ids = DraftTokenIds([req_id], [[1, 2, 3]])
|
||||
scheduler.update_draft_token_ids(draft_token_ids)
|
||||
|
||||
# Schedule the speculated tokens for validation
|
||||
output = scheduler.schedule()
|
||||
assert req_id in output.scheduled_spec_decode_tokens
|
||||
assert len(output.scheduled_spec_decode_tokens[req_id]) == 3
|
||||
|
||||
# Simulate empty output tokens (e.g., due to request abortion or error)
|
||||
# This would previously cause num_accepted = -1 and crash
|
||||
model_runner_output = ModelRunnerOutput(
|
||||
req_ids=[req_id],
|
||||
req_id_to_index={req_id: 0},
|
||||
sampled_token_ids=[[]], # Empty output tokens
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
)
|
||||
|
||||
# This should not raise an error
|
||||
engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
|
||||
|
||||
# Spec decoding stats should be None since no tokens were generated
|
||||
scheduler_stats = (
|
||||
engine_core_outputs[0].scheduler_stats if engine_core_outputs else None
|
||||
)
|
||||
assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None
|
||||
|
||||
|
||||
def _assert_right_scheduler_output(
|
||||
output: SchedulerOutput,
|
||||
num_requests: int,
|
||||
@@ -1806,6 +1866,12 @@ def test_priority_scheduling_mixed_priority_and_arrival():
|
||||
assert scheduled_req_ids == ["3", "2", "1", "0"]
|
||||
|
||||
|
||||
# This test had previously been passing due to its use of duplicate
|
||||
# request ids which resulted in incorrect behavior.
|
||||
# Now that the duplicate req ids had been fixed it fails and
|
||||
# investigation is needed into whether the priority scheduling
|
||||
# preemption logic is working as designed or not.
|
||||
@pytest.mark.skip("needs investigation")
|
||||
def test_priority_scheduling_preemption():
|
||||
"""Test that priority scheduling preempts
|
||||
lower priority requests when memory is constrained."""
|
||||
@@ -1822,7 +1888,8 @@ def test_priority_scheduling_preemption():
|
||||
num_requests=2,
|
||||
priorities=[5, 5], # Low priority
|
||||
arrival_times=[1.0, 2.0],
|
||||
num_tokens=30, # Large enough to consume significant memory
|
||||
num_tokens=30, # Large enough to consume significant memory,
|
||||
req_ids=["lo1", "lo2"],
|
||||
)
|
||||
|
||||
# Add and schedule low priority requests
|
||||
@@ -1855,6 +1922,7 @@ def test_priority_scheduling_preemption():
|
||||
priorities=[0], # High priority
|
||||
arrival_times=[3.0],
|
||||
num_tokens=30, # Large enough to require significant memory
|
||||
req_ids=["hi1"],
|
||||
)[0]
|
||||
|
||||
scheduler.add_request(high_priority_request)
|
||||
@@ -1876,13 +1944,13 @@ def test_priority_scheduling_preemption():
|
||||
output2 = scheduler.schedule()
|
||||
assert len(output2.scheduled_new_reqs) == 1
|
||||
# High priority request
|
||||
assert output2.scheduled_new_reqs[0].req_id == "0"
|
||||
assert output2.scheduled_new_reqs[0].req_id == "hi1"
|
||||
else:
|
||||
# No preemption needed - all requests fit
|
||||
# This is also valid behavior if memory allows
|
||||
assert len(output.scheduled_new_reqs) == 1
|
||||
# High priority request
|
||||
assert output.scheduled_new_reqs[0].req_id == "0"
|
||||
assert output.scheduled_new_reqs[0].req_id == "hi1"
|
||||
|
||||
|
||||
def test_priority_scheduling_no_preemption_when_space_available():
|
||||
@@ -1895,7 +1963,11 @@ def test_priority_scheduling_no_preemption_when_space_available():
|
||||
|
||||
# Add two low-priority running requests
|
||||
low_priority_requests = create_requests_with_priority(
|
||||
num_requests=2, priorities=[5, 5], arrival_times=[1.0, 2.0], num_tokens=30
|
||||
num_requests=2,
|
||||
priorities=[5, 5],
|
||||
arrival_times=[1.0, 2.0],
|
||||
num_tokens=30,
|
||||
req_ids=["lo1", "lo2"],
|
||||
)
|
||||
|
||||
for request in low_priority_requests:
|
||||
@@ -1916,7 +1988,11 @@ def test_priority_scheduling_no_preemption_when_space_available():
|
||||
|
||||
# Add high-priority request
|
||||
high_priority_request = create_requests_with_priority(
|
||||
num_requests=1, priorities=[0], arrival_times=[3.0], num_tokens=30
|
||||
num_requests=1,
|
||||
priorities=[0],
|
||||
arrival_times=[3.0],
|
||||
num_tokens=30,
|
||||
req_ids=["hi1"],
|
||||
)[0]
|
||||
|
||||
scheduler.add_request(high_priority_request)
|
||||
|
||||
@@ -0,0 +1,656 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
End-to-end tests for the streaming input feature in AsyncLLM.
|
||||
|
||||
These tests verify that:
|
||||
1. Streaming inputs work correctly with bunched inputs (queued)
|
||||
2. Streaming inputs work correctly with spaced out inputs
|
||||
3. Outputs are equivalent whether inputs are bunched or spaced
|
||||
4. Cancelling the output stream correctly aborts the session
|
||||
5. Closing the input stream correctly signals completion
|
||||
6. Queued inputs are cancelled when the session is aborted
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.engine.async_llm import AsyncLLM, StreamingInput
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
# Use a small model that doesn't require authentication for fast tests
|
||||
MODEL = "facebook/opt-125m"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="module", loop_scope="module")
|
||||
async def engine():
|
||||
"""Create an AsyncLLM engine for the test.
|
||||
|
||||
Note: Using function scope because pytest_asyncio creates a new event loop
|
||||
for each test, and the output_handler task gets cancelled between tests
|
||||
with module scope.
|
||||
"""
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=MODEL, enforce_eager=True, gpu_memory_utilization=0.7
|
||||
)
|
||||
with set_default_torch_num_threads(1):
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
try:
|
||||
yield engine
|
||||
finally:
|
||||
engine.shutdown()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
|
||||
def get_sampling_params(max_tokens: int = 20) -> SamplingParams:
|
||||
"""Create sampling params for streaming input tests."""
|
||||
return SamplingParams(
|
||||
max_tokens=max_tokens,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.DELTA,
|
||||
temperature=0.0, # Deterministic for reproducibility
|
||||
)
|
||||
|
||||
|
||||
async def collect_outputs(
|
||||
output_gen: AsyncGenerator[RequestOutput, None],
|
||||
) -> tuple[list[RequestOutput], str]:
|
||||
"""Collect all outputs from a generate call, return outputs and full text."""
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
async for output in output_gen:
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
return outputs, full_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_bunched(engine: AsyncLLM):
|
||||
"""Test streaming input where all inputs are sent at once (bunched/queued).
|
||||
|
||||
This tests the case where multiple inputs arrive before any completes.
|
||||
The inputs should be queued and processed in sequence.
|
||||
"""
|
||||
request_id = "test_bunched"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
# Create an input generator that yields all inputs quickly
|
||||
async def bunched_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# Send multiple inputs rapidly - they should be queued
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
yield StreamingInput(prompt=" Alice and I like")
|
||||
yield StreamingInput(prompt=" to code in Python")
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(
|
||||
bunched_input_generator(),
|
||||
sampling_params,
|
||||
request_id,
|
||||
)
|
||||
)
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify the final output is marked as finished
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify intermediate outputs are not marked as finished
|
||||
for output in outputs[:-1]:
|
||||
assert not output.finished, "Intermediate outputs should not be finished"
|
||||
|
||||
# Verify we generated some text
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
print(f"Bunched test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_spaced(engine: AsyncLLM):
|
||||
"""Test streaming input where inputs are spaced out.
|
||||
|
||||
This tests the case where each input completes processing before the
|
||||
next one is sent. Each chunk should be prefilled, generate tokens,
|
||||
then the next chunk should be processed.
|
||||
"""
|
||||
request_id = "test_spaced"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
# Track when each input is sent
|
||||
input_times: list[float] = []
|
||||
outputs_per_chunk: list[int] = [0, 0, 0]
|
||||
current_chunk = 0
|
||||
|
||||
async def spaced_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal current_chunk
|
||||
import time
|
||||
|
||||
# First input
|
||||
input_times.append(time.time())
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
current_chunk = 0
|
||||
|
||||
# Wait for some outputs to be generated
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second input
|
||||
input_times.append(time.time())
|
||||
current_chunk = 1
|
||||
yield StreamingInput(prompt=" Alice and I like")
|
||||
|
||||
# Wait for some outputs
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Third input
|
||||
input_times.append(time.time())
|
||||
current_chunk = 2
|
||||
yield StreamingInput(prompt=" to code in Python")
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
|
||||
async for output in engine.generate(
|
||||
spaced_input_generator(),
|
||||
sampling_params,
|
||||
request_id,
|
||||
):
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
outputs_per_chunk[current_chunk] += 1
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify the final output is marked as finished
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify we received outputs from multiple chunks
|
||||
# (with spaced inputs, we should see outputs distributed across chunks)
|
||||
chunks_with_outputs = sum(1 for c in outputs_per_chunk if c > 0)
|
||||
assert chunks_with_outputs >= 1, "Should have outputs from at least one chunk"
|
||||
|
||||
print(f"Spaced test generated: {full_text}")
|
||||
print(f"Outputs per chunk: {outputs_per_chunk}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_output_equivalence(engine: AsyncLLM):
|
||||
"""Test that bunched and spaced inputs produce equivalent outputs.
|
||||
|
||||
When the same prompts are provided either bunched or spaced,
|
||||
the final concatenated output should be the same (with deterministic
|
||||
sampling).
|
||||
"""
|
||||
prompts = ["Hello, my name is", " Bob and I work", " at Anthropic"]
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
# Test bunched inputs
|
||||
async def bunched_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
|
||||
_, bunched_text = await collect_outputs(
|
||||
engine.generate(bunched_gen(), sampling_params, "equiv_bunched")
|
||||
)
|
||||
|
||||
# Test spaced inputs (same prompts, but with delays)
|
||||
async def spaced_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
_, spaced_text = await collect_outputs(
|
||||
engine.generate(spaced_gen(), sampling_params, "equiv_spaced")
|
||||
)
|
||||
|
||||
# Both should produce the same output since we use temperature=0
|
||||
assert bunched_text == spaced_text, (
|
||||
f"Bunched and spaced should produce same output.\n"
|
||||
f"Bunched: {bunched_text!r}\n"
|
||||
f"Spaced: {spaced_text!r}"
|
||||
)
|
||||
|
||||
print(f"Equivalence test passed. Generated: {bunched_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_cancel_output_stream(engine: AsyncLLM):
|
||||
"""Test that cancelling the output stream aborts the entire session.
|
||||
|
||||
When the consumer cancels iteration over the output generator,
|
||||
the session should be aborted including any queued inputs.
|
||||
"""
|
||||
request_id = "test_cancel_output"
|
||||
sampling_params = get_sampling_params(max_tokens=1000)
|
||||
|
||||
input_completed = asyncio.Event()
|
||||
input_task_cancelled = False
|
||||
|
||||
async def slow_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_task_cancelled
|
||||
try:
|
||||
yield StreamingInput(prompt="Tell me a very long story about")
|
||||
yield StreamingInput(prompt=" a dragon and a knight")
|
||||
|
||||
# This should be cancelled before we get here
|
||||
await asyncio.sleep(10)
|
||||
yield StreamingInput(prompt=" who become friends")
|
||||
input_completed.set()
|
||||
except asyncio.CancelledError:
|
||||
input_task_cancelled = True
|
||||
raise
|
||||
|
||||
outputs_received = 0
|
||||
output_gen = engine.generate(slow_input_generator(), sampling_params, request_id)
|
||||
|
||||
# Collect a few outputs then cancel
|
||||
try:
|
||||
async for output in output_gen:
|
||||
outputs_received += 1
|
||||
if outputs_received >= 5:
|
||||
# Cancel by breaking out of the loop (generator will be GC'd)
|
||||
break
|
||||
finally:
|
||||
# Explicitly close the generator to ensure cleanup
|
||||
await output_gen.aclose()
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify we got some outputs before cancelling
|
||||
assert outputs_received >= 5, "Should have received outputs before cancel"
|
||||
|
||||
# Verify the input task was cancelled
|
||||
assert input_task_cancelled, "Input task should have been cancelled"
|
||||
|
||||
# Verify the session is properly cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after cancel"
|
||||
)
|
||||
|
||||
print(f"Cancel test passed. Received {outputs_received} outputs before cancel")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_close_signals_completion(engine: AsyncLLM):
|
||||
"""Test that closing the input stream signals completion.
|
||||
|
||||
When the input generator finishes (naturally or via return),
|
||||
the session should complete with finished=True on the last output.
|
||||
"""
|
||||
request_id = "test_close_completion"
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
input_generator_finished = False
|
||||
|
||||
async def limited_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_generator_finished
|
||||
yield StreamingInput(prompt="What is 2 + 2? The answer is")
|
||||
# Generator finishes naturally here
|
||||
input_generator_finished = True
|
||||
|
||||
outputs, _ = await collect_outputs(
|
||||
engine.generate(limited_input_generator(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
# Verify the input generator completed
|
||||
assert input_generator_finished, "Input generator should have finished"
|
||||
|
||||
# Verify we got a finished output
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests"
|
||||
)
|
||||
|
||||
print("Close completion test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_abort_queued_inputs(engine: AsyncLLM):
|
||||
"""Test that aborting the session cancels queued inputs.
|
||||
|
||||
When multiple inputs are queued and the session is aborted,
|
||||
all pending inputs should be cancelled.
|
||||
"""
|
||||
request_id = "test_abort_queued"
|
||||
# Use large max_tokens to ensure we have time to queue inputs
|
||||
sampling_params = get_sampling_params(max_tokens=2000)
|
||||
|
||||
inputs_sent = 0
|
||||
input_cancelled = False
|
||||
|
||||
async def many_inputs_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal inputs_sent, input_cancelled
|
||||
try:
|
||||
# Send several inputs to fill the queue
|
||||
for i in range(10):
|
||||
yield StreamingInput(prompt=f" Part {i}: Tell me about the number {i}.")
|
||||
inputs_sent += 1
|
||||
# Small delay to interleave with output processing
|
||||
await asyncio.sleep(0.05)
|
||||
except asyncio.CancelledError:
|
||||
input_cancelled = True
|
||||
raise
|
||||
|
||||
outputs_received = 0
|
||||
output_gen = engine.generate(many_inputs_generator(), sampling_params, request_id)
|
||||
|
||||
try:
|
||||
async for output in output_gen:
|
||||
outputs_received += 1
|
||||
# Cancel after receiving some outputs
|
||||
if outputs_received >= 10:
|
||||
break
|
||||
finally:
|
||||
await output_gen.aclose()
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify we received some outputs
|
||||
assert outputs_received >= 10, "Should have received outputs before abort"
|
||||
|
||||
# Verify the input generator was cancelled OR finished naturally
|
||||
# (it might finish naturally if all inputs were sent before cancel)
|
||||
assert input_cancelled or inputs_sent == 10, (
|
||||
f"Input generator should have been cancelled or completed. "
|
||||
f"cancelled={input_cancelled}, inputs_sent={inputs_sent}"
|
||||
)
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after abort"
|
||||
)
|
||||
|
||||
print(
|
||||
f"Abort queued test passed. Sent {inputs_sent} inputs, "
|
||||
f"received {outputs_received} outputs"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_error_propagation(engine: AsyncLLM):
|
||||
"""Test that errors in the input generator are propagated to the caller."""
|
||||
request_id = "test_error_propagation"
|
||||
sampling_params = get_sampling_params(max_tokens=20)
|
||||
|
||||
class InputError(Exception):
|
||||
pass
|
||||
|
||||
async def error_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Start with this")
|
||||
await asyncio.sleep(0.1)
|
||||
raise InputError("Simulated input error")
|
||||
|
||||
# Note: The current implementation catches exceptions and puts them
|
||||
# in the queue, so we should get the error when iterating outputs
|
||||
with pytest.raises(InputError, match="Simulated input error"):
|
||||
async for _ in engine.generate(
|
||||
error_input_generator(), sampling_params, request_id
|
||||
):
|
||||
pass
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after error"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_multiple_concurrent_sessions(engine: AsyncLLM):
|
||||
"""Test multiple concurrent streaming input sessions.
|
||||
|
||||
Multiple streaming sessions should be able to run concurrently
|
||||
without interfering with each other.
|
||||
"""
|
||||
num_sessions = 3
|
||||
results: list[tuple[str, str]] = []
|
||||
|
||||
async def run_session(session_id: int) -> tuple[str, str]:
|
||||
request_id = f"test_concurrent_{session_id}"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
prompts = [f"Session {session_id}: Hello", f" world from session {session_id}"]
|
||||
|
||||
async def input_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
_, text = await collect_outputs(
|
||||
engine.generate(input_gen(), sampling_params, request_id)
|
||||
)
|
||||
return request_id, text
|
||||
|
||||
# Run sessions concurrently
|
||||
tasks = [asyncio.create_task(run_session(i)) for i in range(num_sessions)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all sessions completed
|
||||
assert len(results) == num_sessions
|
||||
|
||||
for request_id, text in results:
|
||||
assert len(text) > 0, f"Session {request_id} should have generated text"
|
||||
print(f"{request_id}: {text}")
|
||||
|
||||
# Verify cleanup
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_per_chunk_sampling_params(engine: AsyncLLM):
|
||||
"""Test that per-chunk sampling params are respected.
|
||||
|
||||
Each StreamingInput can have its own sampling_params.
|
||||
"""
|
||||
request_id = "test_per_chunk_params"
|
||||
base_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
async def variable_params_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# First chunk with base params
|
||||
yield StreamingInput(prompt="Count to five:", sampling_params=base_params)
|
||||
|
||||
# Second chunk with different max_tokens
|
||||
chunk_params = get_sampling_params(max_tokens=5)
|
||||
yield StreamingInput(
|
||||
prompt=" Now count backwards:", sampling_params=chunk_params
|
||||
)
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(variable_params_generator(), base_params, request_id)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
assert outputs[-1].finished, "Last output should be finished"
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
|
||||
print(f"Per-chunk params test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_empty_generator(engine: AsyncLLM):
|
||||
"""Test behavior when the input generator yields nothing.
|
||||
|
||||
An empty generator should still produce a finished output.
|
||||
"""
|
||||
request_id = "test_empty_generator"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
async def empty_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# Don't yield anything
|
||||
return
|
||||
yield # Make it a generator
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
async for output in engine.generate(empty_generator(), sampling_params, request_id):
|
||||
outputs.append(output)
|
||||
|
||||
# Should still get a finished marker
|
||||
assert len(outputs) >= 1, "Should receive at least one output"
|
||||
assert outputs[-1].finished, "Should have a finished output"
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_single_chunk(engine: AsyncLLM):
|
||||
"""Test streaming input with a single chunk.
|
||||
|
||||
This is effectively the same as a regular non-streaming request,
|
||||
but using the streaming input API.
|
||||
"""
|
||||
request_id = "test_single_chunk"
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
async def single_chunk_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="What color is the sky? The sky is")
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(single_chunk_generator(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0
|
||||
assert outputs[-1].finished
|
||||
assert "blue" in full_text.lower() or len(full_text) > 0
|
||||
|
||||
print(f"Single chunk test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_reuse_request_id(engine: AsyncLLM):
|
||||
"""Test that request IDs can be reused after a session completes."""
|
||||
request_id = "test_reuse_id"
|
||||
sampling_params = get_sampling_params(max_tokens=5)
|
||||
|
||||
# First session
|
||||
async def gen1() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="First session")
|
||||
|
||||
_, text1 = await collect_outputs(
|
||||
engine.generate(gen1(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
# Second session with same ID
|
||||
async def gen2() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Second session")
|
||||
|
||||
_, text2 = await collect_outputs(
|
||||
engine.generate(gen2(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
assert len(text1) > 0
|
||||
assert len(text2) > 0
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
print(f"Reuse ID test: session 1: {text1}, session 2: {text2}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_validation_errors(engine: AsyncLLM):
|
||||
"""Test that invalid configurations raise appropriate errors."""
|
||||
|
||||
async def dummy_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="test")
|
||||
|
||||
# Test n > 1 is rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_n2 = SamplingParams(max_tokens=10, n=2)
|
||||
async for _ in engine.generate(dummy_generator(), params_n2, "test_n2"):
|
||||
pass
|
||||
|
||||
# Test FINAL_ONLY is rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_final = SamplingParams(
|
||||
max_tokens=10, output_kind=RequestOutputKind.FINAL_ONLY
|
||||
)
|
||||
async for _ in engine.generate(dummy_generator(), params_final, "test_final"):
|
||||
pass
|
||||
|
||||
# Test stop strings are rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_stop = SamplingParams(max_tokens=10, stop=["stop"])
|
||||
async for _ in engine.generate(dummy_generator(), params_stop, "test_stop"):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_delayed_generator_exit(engine: AsyncLLM):
|
||||
"""Test that output generator exits when input generator closes after outputs.
|
||||
|
||||
This tests the case where:
|
||||
1. Multiple inputs are sent and fully processed
|
||||
2. The engine has finished
|
||||
3. The input generator doesn't exit until after the engine finishes
|
||||
4. The output generator should exit properly once the input generator exits
|
||||
"""
|
||||
request_id = "test_delayed_exit"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
engine_finished_event = asyncio.Event()
|
||||
input_generator_exited = False
|
||||
finish_count = 0
|
||||
|
||||
async def delayed_exit_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_generator_exited
|
||||
# Send all inputs immediately
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
yield StreamingInput(prompt=" Alice")
|
||||
|
||||
# Wait until the engine has finished generating before exiting
|
||||
await engine_finished_event.wait()
|
||||
|
||||
# Add a small delay to ensure we're testing the "delayed exit" case
|
||||
await asyncio.sleep(0.1)
|
||||
input_generator_exited = True
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
|
||||
async for output in engine.generate(
|
||||
delayed_exit_input_generator(), sampling_params, request_id
|
||||
):
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
|
||||
# Signal when the engine finishes both input chunks (each gets a finish_reason)
|
||||
# Note: output.finished will be False while input stream is open
|
||||
if output.outputs and output.outputs[0].finish_reason is not None:
|
||||
finish_count += 1
|
||||
if finish_count == 2:
|
||||
engine_finished_event.set()
|
||||
|
||||
# Verify the input generator exited properly
|
||||
assert input_generator_exited, (
|
||||
"Input generator should have exited after engine finished"
|
||||
)
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify we generated some text
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests"
|
||||
)
|
||||
|
||||
print(f"Delayed exit test passed. Generated: {full_text}")
|
||||
@@ -5,14 +5,8 @@ import pytest
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
DeviceConfig,
|
||||
ModelConfig,
|
||||
MultiModalConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.multimodal import MultiModalRegistry, MultiModalUUIDDict
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
||||
from vllm.multimodal import MultiModalUUIDDict
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine.input_processor import InputProcessor
|
||||
|
||||
@@ -21,55 +15,26 @@ stop_pil_image = ImageAsset("stop_sign").pil_image
|
||||
baby_reading_np_ndarrays = VideoAsset("baby_reading").np_ndarrays
|
||||
|
||||
|
||||
def _mock_input_processor(
|
||||
monkeypatch, *, mm_cache_gb: float = 4.0, enable_prefix_caching: bool = True
|
||||
def _build_input_processor(
|
||||
*, mm_cache_gb: float = 4.0, enable_prefix_caching: bool = True
|
||||
) -> InputProcessor:
|
||||
"""
|
||||
Create a Processor instance with minimal configuration suitable for unit
|
||||
tests without accessing external resources.
|
||||
"""
|
||||
monkeypatch.setattr(
|
||||
ModelConfig, "try_get_generation_config", lambda self: {}, raising=True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ModelConfig, "__post_init__", lambda self, *args: None, raising=True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ModelConfig,
|
||||
"verify_with_parallel_config",
|
||||
lambda self, parallel_config: None,
|
||||
raising=True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
MultiModalRegistry,
|
||||
"processor_cache_from_config",
|
||||
lambda self, vllm_config: None,
|
||||
raising=True,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(VllmConfig, "__post_init__", lambda self: None, raising=True)
|
||||
|
||||
model_config = ModelConfig(
|
||||
tokenizer="dummy",
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
skip_tokenizer_init=True,
|
||||
max_model_len=128,
|
||||
mm_processor_cache_gb=mm_cache_gb,
|
||||
generation_config="vllm",
|
||||
)
|
||||
model_config.runner_type = "generate"
|
||||
model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=mm_cache_gb)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=CacheConfig(enable_prefix_caching=enable_prefix_caching),
|
||||
device_config=DeviceConfig(device="cpu"),
|
||||
)
|
||||
|
||||
return InputProcessor(vllm_config)
|
||||
|
||||
|
||||
def test_multi_modal_uuids_length_mismatch_raises(monkeypatch):
|
||||
input_processor = _mock_input_processor(monkeypatch)
|
||||
def test_multi_modal_uuids_length_mismatch_raises():
|
||||
input_processor = _build_input_processor()
|
||||
|
||||
prompt = {
|
||||
"prompt": "USER: <image>\nDescribe\nASSISTANT:",
|
||||
@@ -78,7 +43,7 @@ def test_multi_modal_uuids_length_mismatch_raises(monkeypatch):
|
||||
"multi_modal_uuids": {"image": ["hash_cherry"]},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="must have same length as data"):
|
||||
with pytest.raises(ValueError, match="must have same length as"):
|
||||
input_processor.process_inputs(
|
||||
request_id="req-1",
|
||||
prompt=prompt, # type: ignore[arg-type]
|
||||
@@ -86,21 +51,21 @@ def test_multi_modal_uuids_length_mismatch_raises(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_multi_modal_uuids_missing_modality_raises(monkeypatch):
|
||||
input_processor = _mock_input_processor(monkeypatch)
|
||||
def test_multi_modal_uuids_missing_modality_raises():
|
||||
input_processor = _build_input_processor()
|
||||
|
||||
prompt = {
|
||||
"prompt": "USER: <image><video>\nDescribe\nASSISTANT:",
|
||||
# Two modalities provided in data
|
||||
"multi_modal_data": {
|
||||
"image": [cherry_pil_image],
|
||||
"video": [baby_reading_np_ndarrays],
|
||||
"video": None,
|
||||
},
|
||||
# Only image uuids provided; video missing should raise
|
||||
"multi_modal_uuids": {"image": ["hash_cherry"]},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="must be provided if multi_modal_data"):
|
||||
with pytest.raises(ValueError, match="is empty but .* is missing"):
|
||||
input_processor.process_inputs(
|
||||
request_id="req-2",
|
||||
prompt=prompt, # type: ignore[arg-type]
|
||||
@@ -119,8 +84,7 @@ def test_multi_modal_uuids_missing_modality_raises(monkeypatch):
|
||||
def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
|
||||
):
|
||||
input_processor = _mock_input_processor(
|
||||
monkeypatch,
|
||||
input_processor = _build_input_processor(
|
||||
mm_cache_gb=mm_cache_gb,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
@@ -163,8 +127,8 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
# When both processor cache is 0 and prefix caching disabled, the
|
||||
# processor builds overrides from request id instead of using user UUIDs.
|
||||
input_processor = _mock_input_processor(
|
||||
monkeypatch, mm_cache_gb=0.0, enable_prefix_caching=False
|
||||
input_processor = _build_input_processor(
|
||||
mm_cache_gb=0.0, enable_prefix_caching=False
|
||||
)
|
||||
|
||||
captured: dict[str, MultiModalUUIDDict] = {}
|
||||
@@ -180,12 +144,12 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
)
|
||||
|
||||
request_id = "req-42"
|
||||
mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": "hash_video"}
|
||||
mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}
|
||||
prompt = {
|
||||
"prompt": "USER: <image><image><video>\nDescribe\nASSISTANT:",
|
||||
"multi_modal_data": {
|
||||
"image": [cherry_pil_image, stop_pil_image],
|
||||
"video": baby_reading_np_ndarrays,
|
||||
"video": [baby_reading_np_ndarrays],
|
||||
},
|
||||
"multi_modal_uuids": mm_uuids,
|
||||
}
|
||||
@@ -197,16 +161,15 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
)
|
||||
|
||||
# Expect request-id-based overrides are passed through
|
||||
mm_uuids = captured["mm_uuids"]
|
||||
assert set(mm_uuids.keys()) == {"image", "video"}
|
||||
assert len(mm_uuids["image"]) == 2
|
||||
assert len(mm_uuids["video"]) == 1
|
||||
assert mm_uuids["image"][0].startswith(f"{request_id}-image-") and mm_uuids[
|
||||
"image"
|
||||
][0].endswith("-0")
|
||||
assert mm_uuids["image"][1].startswith(f"{request_id}-image-") and mm_uuids[
|
||||
"image"
|
||||
][1].endswith("-1")
|
||||
assert mm_uuids["video"][0].startswith(f"{request_id}-video-") and mm_uuids[
|
||||
"video"
|
||||
][0].endswith("-0")
|
||||
assert captured["mm_uuids"]["image"][0].startswith(
|
||||
f"{request_id}-image-"
|
||||
) and captured["mm_uuids"]["image"][0].endswith("-0")
|
||||
assert captured["mm_uuids"]["image"][1].startswith(
|
||||
f"{request_id}-image-"
|
||||
) and captured["mm_uuids"]["image"][1].endswith("-1")
|
||||
assert captured["mm_uuids"]["video"][0].startswith(
|
||||
f"{request_id}-video-"
|
||||
) and captured["mm_uuids"]["video"][0].endswith("-0")
|
||||
|
||||
@@ -70,15 +70,28 @@ async def test_background_cancel(client: openai.AsyncOpenAI):
|
||||
assert response.status == "queued"
|
||||
|
||||
# Cancel the response before it is completed.
|
||||
# FIXME: This test can be flaky.
|
||||
await asyncio.sleep(0.5)
|
||||
# Poll until the response is no longer queued (started processing) or timeout
|
||||
loop = asyncio.get_running_loop()
|
||||
start_time = loop.time()
|
||||
max_wait_seconds = 5.0
|
||||
poll_interval = 0.1
|
||||
while loop.time() - start_time < max_wait_seconds:
|
||||
response = await client.responses.retrieve(response.id)
|
||||
if response.status != "queued":
|
||||
# Started processing or completed - try to cancel
|
||||
break
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
response = await client.responses.cancel(response.id)
|
||||
assert response.status == "cancelled"
|
||||
|
||||
# Make sure the response status remains unchanged.
|
||||
await asyncio.sleep(5)
|
||||
response = await client.responses.retrieve(response.id)
|
||||
assert response.status == "cancelled"
|
||||
# Make sure the response status remains unchanged after some time.
|
||||
max_retries = 10
|
||||
for _ in range(max_retries):
|
||||
await asyncio.sleep(0.5)
|
||||
response = await client.responses.retrieve(response.id)
|
||||
# Verify status is still cancelled
|
||||
assert response.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -34,18 +34,11 @@ else
|
||||
KV_CONFIG_HETERO_LAYOUT=''
|
||||
fi
|
||||
|
||||
CROSS_LAYERS_BLOCKS=${CROSS_LAYERS_BLOCKS:-"False"} # Default to non cross layers
|
||||
if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
|
||||
KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"cross_layers_blocks": "True"}'
|
||||
else
|
||||
KV_EXTRA_CONFIG=''
|
||||
fi
|
||||
|
||||
# Build the kv-transfer-config once
|
||||
if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
|
||||
else
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
|
||||
fi
|
||||
|
||||
# Models to run
|
||||
|
||||
@@ -86,7 +86,7 @@ class DecodeBenchTestRunner:
|
||||
self._block_hasher = get_request_block_hasher(block_size, sha256)
|
||||
|
||||
self._dummy_ctx: ForwardContext = ForwardContext(
|
||||
no_compile_layers={}, attn_metadata={}, virtual_engine=0
|
||||
no_compile_layers={}, attn_metadata={}, virtual_engine=0, slot_mapping={}
|
||||
)
|
||||
|
||||
def new_request(self, token_ids: list[int]) -> Request:
|
||||
|
||||
@@ -18,12 +18,8 @@ import ray
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import KVTransferConfig, set_current_vllm_config
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import (
|
||||
KVOutputAggregator,
|
||||
TpKVTopology,
|
||||
get_current_attn_backend,
|
||||
)
|
||||
from vllm.config import KVTransferConfig
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
|
||||
@@ -52,11 +48,8 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.output_processor import OutputProcessor
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
|
||||
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
|
||||
from vllm.v1.worker.utils import AttentionGroup
|
||||
|
||||
from .utils import create_request, create_scheduler, create_vllm_config
|
||||
|
||||
@@ -373,7 +366,6 @@ def test_kv_transfer_handshake(dist_init):
|
||||
|
||||
# Decode connector will be able to create handshake with the prefill connector.
|
||||
decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
# Here we are testing the retrieval of NIXLAgentMetadata.
|
||||
# Knowing the implementation detail, we override the add_remote_agent
|
||||
@@ -410,23 +402,6 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
|
||||
self.kv_cache_layout = kv_cache_layout
|
||||
# Mock register_kv_caches attribute needed for tests that do not call it.
|
||||
self.src_xfer_handles_by_block_size = {self.block_size: 1}
|
||||
test_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
self.kv_topo = TpKVTopology(
|
||||
tp_rank=self.tp_rank,
|
||||
engine_id=self.engine_id,
|
||||
remote_tp_size=self._tp_size, # shared state
|
||||
remote_block_size=self._block_size, # shared state
|
||||
is_mla=self.use_mla,
|
||||
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=self.attn_backend,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
self.compat_hash = compute_nixl_compatibility_hash(
|
||||
self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
|
||||
)
|
||||
|
||||
def _nixl_handshake(
|
||||
self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
|
||||
@@ -548,6 +523,7 @@ class TestNixlHandshake:
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
_before_load = time.perf_counter()
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
@@ -618,6 +594,7 @@ class TestNixlHandshake:
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
_before_load = time.perf_counter()
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
@@ -844,6 +821,7 @@ class TestNixlHandshake:
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
_before_load = time.perf_counter()
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
@@ -1006,6 +984,7 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -1391,7 +1370,6 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
|
||||
),
|
||||
),
|
||||
"TRITON_ATTN",
|
||||
"FLASHINFER",
|
||||
],
|
||||
)
|
||||
def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
@@ -1408,11 +1386,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
|
||||
vllm_config = create_vllm_config(attention_backend=attn_backend)
|
||||
|
||||
# Enable cross layers blocks
|
||||
vllm_config.kv_transfer_config.kv_connector_extra_config[
|
||||
"enable_cross_layers_blocks"
|
||||
] = True
|
||||
|
||||
# Import the appropriate backend based on the parameter
|
||||
if attn_backend == "FLASH_ATTN":
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
@@ -1422,11 +1395,49 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
|
||||
|
||||
backend_cls = RocmAttentionBackend
|
||||
else: # TRITON
|
||||
else: # TRITON_ATTN
|
||||
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
|
||||
|
||||
backend_cls = TritonAttentionBackend
|
||||
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
|
||||
# Store tensor info for validation
|
||||
|
||||
test_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
|
||||
|
||||
if is_blocks_first:
|
||||
expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
|
||||
expected_base_addrs = [
|
||||
shared_tensor.data_ptr(),
|
||||
unique_tensor.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 2
|
||||
else:
|
||||
expected_tensor_size = (
|
||||
shared_tensor[0].element_size() * shared_tensor[0].numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor[0].data_ptr(),
|
||||
shared_tensor[1].data_ptr(),
|
||||
unique_tensor[0].data_ptr(),
|
||||
unique_tensor[1].data_ptr(),
|
||||
]
|
||||
expected_num_entries = 4
|
||||
|
||||
nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
|
||||
with (
|
||||
patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
|
||||
@@ -1455,107 +1466,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
# Reassure the shutdown() check that the thread is terminated
|
||||
mock_thread.return_value.is_alive.return_value = False
|
||||
|
||||
expected_tensor_size: int
|
||||
expected_base_addrs: list[int]
|
||||
expected_num_entries: int
|
||||
kv_caches: dict[str, torch.Tensor]
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_layers = 32
|
||||
block_size = 16
|
||||
num_blocks = 8
|
||||
kv_cache_spec = AttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(
|
||||
size=kv_cache_spec.page_size_bytes * num_blocks,
|
||||
shared_by=["dummy-layer"],
|
||||
)
|
||||
for i in range(num_layers)
|
||||
],
|
||||
# allocate_uniform_kv_caches does not use this
|
||||
kv_cache_groups=[],
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
_, cross_layers_kv_cache, _ = (
|
||||
KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
|
||||
kv_cache_config=kv_cache_config,
|
||||
attn_groups=[
|
||||
[
|
||||
AttentionGroup(
|
||||
backend=backend_cls,
|
||||
layer_names=[],
|
||||
kv_cache_spec=kv_cache_spec,
|
||||
kv_cache_group_id=0,
|
||||
)
|
||||
]
|
||||
],
|
||||
cache_dtype=torch.bfloat16,
|
||||
device=torch.cuda.current_device(),
|
||||
kernel_block_sizes=[block_size],
|
||||
)
|
||||
)
|
||||
# Store tensor info for validation
|
||||
expected_tensor_size = (
|
||||
cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
cross_layers_kv_cache.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 1
|
||||
|
||||
expected_blocks_count = 8
|
||||
|
||||
kv_caches = {"all-layers": cross_layers_kv_cache}
|
||||
|
||||
else:
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
|
||||
# Store tensor info for validation
|
||||
|
||||
test_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
|
||||
|
||||
if is_blocks_first:
|
||||
expected_tensor_size = (
|
||||
shared_tensor.element_size() * shared_tensor.numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor.data_ptr(),
|
||||
unique_tensor.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 2
|
||||
else:
|
||||
expected_tensor_size = (
|
||||
shared_tensor[0].element_size() * shared_tensor[0].numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor[0].data_ptr(),
|
||||
shared_tensor[1].data_ptr(),
|
||||
unique_tensor[0].data_ptr(),
|
||||
unique_tensor[1].data_ptr(),
|
||||
]
|
||||
expected_num_entries = 4
|
||||
expected_blocks_count = 8
|
||||
|
||||
# Execute register_kv_caches
|
||||
connector.register_kv_caches(kv_caches)
|
||||
|
||||
@@ -1579,19 +1489,16 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
|
||||
|
||||
# Validate blocks_data structure and size
|
||||
expected_blocks_count = 8
|
||||
assert len(blocks_data) == expected_blocks_count, (
|
||||
f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
|
||||
)
|
||||
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_blocks = 8
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
num_blocks = 2
|
||||
if is_blocks_first:
|
||||
expected_block_len = expected_tensor_size // num_blocks // 2
|
||||
else:
|
||||
num_blocks = 2
|
||||
if is_blocks_first:
|
||||
expected_block_len = expected_tensor_size // num_blocks // 2
|
||||
else:
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
|
||||
for i, block_entry in enumerate(blocks_data):
|
||||
block_start_addr, block_len, tp_rank = block_entry
|
||||
@@ -1767,6 +1674,7 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -1917,6 +1825,7 @@ def test_transfer_failure_logging(
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
|
||||
# Capture logs from the nixl_connector logger specifically
|
||||
@@ -2017,6 +1926,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -2067,6 +1977,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -2138,17 +2049,6 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
|
||||
decode_worker = decode_connector.connector_worker
|
||||
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
remote_config_params: dict[str, Any] = {
|
||||
"model": "facebook/opt-125m",
|
||||
@@ -2171,9 +2071,7 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
)
|
||||
remote_hash = compute_nixl_compatibility_hash(
|
||||
remote_vllm_config,
|
||||
decode_worker.backend_name,
|
||||
decode_worker.kv_topo.cross_layers_blocks,
|
||||
remote_vllm_config, decode_worker.backend_name
|
||||
)
|
||||
|
||||
prefill_block_size = config_overrides.get("block_size", 16)
|
||||
@@ -2252,27 +2150,6 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
|
||||
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
|
||||
decode_worker = decode_connector.connector_worker
|
||||
|
||||
backend = get_current_attn_backend(local_vllm_config)
|
||||
test_shape = backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
decode_worker.kv_topo = TpKVTopology(
|
||||
tp_rank=decode_worker.tp_rank,
|
||||
engine_id=decode_worker.engine_id,
|
||||
remote_tp_size=decode_worker._tp_size, # shared state
|
||||
remote_block_size=decode_worker._block_size, # shared state
|
||||
is_mla=decode_worker.use_mla,
|
||||
total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=backend,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
decode_worker.compat_hash = compute_nixl_compatibility_hash(
|
||||
decode_worker.vllm_config,
|
||||
decode_worker.backend_name,
|
||||
decode_worker.kv_topo.cross_layers_blocks,
|
||||
)
|
||||
|
||||
if error_scenario == "handshake_decode_error":
|
||||
msg_bytes = b"this is not valid msgpack data"
|
||||
elif error_scenario == "handshake_validation_error":
|
||||
|
||||
@@ -209,7 +209,10 @@ class RequestRunner:
|
||||
self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
|
||||
|
||||
self._dummy_ctx: ForwardContext = ForwardContext(
|
||||
no_compile_layers={}, attn_metadata={}, virtual_engine=0
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
|
||||
def new_request(self, token_ids: list[int]):
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.utils import create_new_process_for_each_test
|
||||
from tests.utils import create_new_process_for_each_test, set_random_seed
|
||||
from tests.v1.logits_processors.utils import (
|
||||
DUMMY_LOGITPROC_ARG,
|
||||
DUMMY_LOGITPROC_FQCN,
|
||||
@@ -135,7 +134,7 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
|
||||
|
||||
# Test that logitproc info is passed to workers
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
|
||||
random.seed(40)
|
||||
set_random_seed(40)
|
||||
|
||||
# Choose LLM args based on logitproc source
|
||||
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_NONE:
|
||||
@@ -194,7 +193,7 @@ def test_custom_logitsprocs_req(monkeypatch):
|
||||
|
||||
# Test that logitproc info is passed to workers
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
|
||||
random.seed(40)
|
||||
set_random_seed(40)
|
||||
_run_test(
|
||||
{"logits_processors": [WrappedPerReqLogitsProcessor]}, logitproc_loaded=True
|
||||
)
|
||||
@@ -237,7 +236,7 @@ def test_rejects_custom_logitsprocs(
|
||||
logitproc from
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
random.seed(40)
|
||||
set_random_seed(40)
|
||||
|
||||
test_params: dict[str, dict[str, Any]] = {
|
||||
"pooling": {
|
||||
|
||||
@@ -48,7 +48,11 @@ def test_topk_impl_equivalence():
|
||||
assert torch.allclose(result1, result2)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="FIXME: This test is failing right now.")
|
||||
@pytest.mark.skip(
|
||||
reason="FlashInfer top-k/top-p renorm comparison fails; "
|
||||
"needs investigation of tolerance threshold or "
|
||||
"interface differences between Python and FlashInfer implementations"
|
||||
)
|
||||
def test_flashinfer_sampler():
|
||||
"""
|
||||
This test verifies that the FlashInfer top-k and top-p sampling
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch
|
||||
from tests.v1.attention.utils import (
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
try_backend_includes_kv_cache_update,
|
||||
try_get_attention_backend,
|
||||
)
|
||||
from vllm.config import ParallelConfig, SpeculativeConfig
|
||||
@@ -120,6 +121,14 @@ def forward_attention(
|
||||
key = k.view(-1, num_kv_heads, dim_per_head)
|
||||
value = v.view(-1, num_kv_heads, dim_per_head)
|
||||
output = torch.empty_like(query)
|
||||
if not try_backend_includes_kv_cache_update(backend):
|
||||
instance.do_kv_cache_update(
|
||||
layer=layer,
|
||||
key=key,
|
||||
value=value,
|
||||
kv_cache=kv_cache,
|
||||
slot_mapping=attn_metadata.slot_mapping,
|
||||
)
|
||||
return instance.forward(
|
||||
layer=layer,
|
||||
query=query,
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.v1.engine.async_llm import AsyncLLM, StreamingInput
|
||||
from vllm.v1.engine.output_processor import RequestOutputCollector
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_llm():
|
||||
"""Create a mock AsyncLLM with mocked dependencies."""
|
||||
# Create a minimal mock without initializing the full engine
|
||||
llm = MagicMock(spec=AsyncLLM)
|
||||
|
||||
# Mock the essential attributes
|
||||
llm.vllm_config = MagicMock()
|
||||
llm.vllm_config.cache_config.kv_sharing_fast_prefill = False
|
||||
llm.model_config = MagicMock()
|
||||
llm.model_config.max_model_len = 2048
|
||||
llm.log_requests = False
|
||||
llm.errored = False
|
||||
llm._pause_cond = asyncio.Condition()
|
||||
llm._paused = False
|
||||
|
||||
# Mock methods
|
||||
llm._run_output_handler = MagicMock()
|
||||
llm.abort = AsyncMock()
|
||||
|
||||
# Use the real generate method from AsyncLLM
|
||||
llm.generate = AsyncLLM.generate.__get__(llm, AsyncLLM)
|
||||
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_normal_flow(mock_async_llm):
|
||||
"""Test normal generation flow with streaming requests."""
|
||||
request_id = "test_request"
|
||||
prompt = "Tell me about Paris"
|
||||
sampling_params = SamplingParams(max_tokens=10)
|
||||
|
||||
# Create a mock queue with outputs
|
||||
queue = RequestOutputCollector(RequestOutputKind.FINAL_ONLY, request_id)
|
||||
output1 = RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="Tell me about Paris",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[],
|
||||
finished=False,
|
||||
)
|
||||
output2 = RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="Tell me about Paris",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[],
|
||||
finished=True,
|
||||
)
|
||||
|
||||
# Feed outputs to queue as they're consumed to avoid aggregation
|
||||
async def feed_outputs():
|
||||
queue.put(output1)
|
||||
await asyncio.sleep(1) # Let first output be consumed
|
||||
queue.put(output2)
|
||||
|
||||
asyncio.create_task(feed_outputs()) # noqa
|
||||
|
||||
# Mock add_request to return the queue
|
||||
async def mock_add_request(*args, **kwargs):
|
||||
return queue
|
||||
|
||||
mock_async_llm.add_request = mock_add_request
|
||||
|
||||
# Collect outputs from generate
|
||||
outputs = []
|
||||
async for output in mock_async_llm.generate(
|
||||
prompt=prompt,
|
||||
sampling_params=sampling_params,
|
||||
request_id=request_id,
|
||||
):
|
||||
outputs.append(output)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0].finished is False
|
||||
assert outputs[1].finished is True
|
||||
|
||||
|
||||
def make_output(request_id: str, finished: bool) -> RequestOutput:
|
||||
"""Helper to create a RequestOutput."""
|
||||
return RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[],
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_with_async_generator():
|
||||
"""Test generate with an async input generator.
|
||||
|
||||
With the new streaming input API, completion is signaled by finishing
|
||||
the input generator (not via a resumable flag). Each input chunk
|
||||
produces intermediate outputs, and the final output has finished=True.
|
||||
"""
|
||||
request_id = "test"
|
||||
sampling_params = SamplingParams(max_tokens=10)
|
||||
|
||||
llm = MagicMock(spec=AsyncLLM)
|
||||
llm.vllm_config = MagicMock()
|
||||
llm.vllm_config.cache_config.kv_sharing_fast_prefill = False
|
||||
llm.model_config = MagicMock()
|
||||
llm.model_config.max_model_len = 2048
|
||||
llm.log_requests = False
|
||||
llm.errored = False
|
||||
llm._pause_cond = asyncio.Condition()
|
||||
llm._paused = False
|
||||
llm._run_output_handler = MagicMock()
|
||||
llm.abort = AsyncMock()
|
||||
|
||||
# Bind the real generate method
|
||||
llm.generate = AsyncLLM.generate.__get__(llm, AsyncLLM)
|
||||
|
||||
# Track inputs processed
|
||||
inputs_received = []
|
||||
queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id)
|
||||
|
||||
async def mock_add_request(req_id, prompt, params, *args, **kwargs):
|
||||
# When prompt is an AsyncGenerator, process streaming inputs
|
||||
if isinstance(prompt, AsyncGenerator):
|
||||
# Process inputs in background, produce outputs
|
||||
async def handle_stream():
|
||||
async for input_chunk in prompt:
|
||||
inputs_received.append(input_chunk.prompt)
|
||||
# Each input produces an intermediate output
|
||||
queue.put(make_output(req_id, finished=False))
|
||||
await asyncio.sleep(0.01)
|
||||
# Final output when stream ends
|
||||
queue.put(make_output(req_id, finished=True))
|
||||
|
||||
asyncio.create_task(handle_stream())
|
||||
return queue
|
||||
return queue
|
||||
|
||||
llm.add_request = mock_add_request
|
||||
|
||||
async def input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Hello", sampling_params=sampling_params)
|
||||
yield StreamingInput(prompt=" world", sampling_params=sampling_params)
|
||||
|
||||
outputs = []
|
||||
async for output in llm.generate(input_generator(), sampling_params, request_id):
|
||||
outputs.append(output)
|
||||
|
||||
# Two intermediate outputs + one final output
|
||||
assert len(outputs) == 3
|
||||
assert outputs[0].finished is False
|
||||
assert outputs[1].finished is False
|
||||
assert outputs[2].finished is True
|
||||
# Both inputs were processed
|
||||
assert inputs_received == ["Hello", " world"]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user